makelore/tests/unit/image-workspace-route.test.ts

import type { IncomingMessage, ServerResponse } from 'node:http';
import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { Readable, Writable } from 'node:stream';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { handleImageWorkspaceRoutes } from '@electron/api/routes/image-workspace';
import type { HostApiContext } from '@electron/api/context';

const electronMocks = vi.hoisted(() => ({
  getPath: vi.fn(),
  showSaveDialog: vi.fn(),
}));

vi.mock('electron', () => ({
  app: { getPath: electronMocks.getPath },
  dialog: { showSaveDialog: electronMocks.showSaveDialog },
}));

function createResponse() {
  const chunks: string[] = [];
  const res = {
    statusCode: 0,
    headersSent: false,
    setHeader: vi.fn(),
    end: vi.fn((chunk?: string) => {
      if (chunk) chunks.push(chunk);
    }),
    destroy: vi.fn(),
    on: vi.fn(),
    once: vi.fn(),
    emit: vi.fn(),
    write: vi.fn(),
  } as unknown as ServerResponse;
  return {
    res,
    get statusCode() {
      return res.statusCode;
    },
    json: () => JSON.parse(chunks.join('')) as Record<string, unknown>,
  };
}

function createRequest(method: string, body?: unknown): IncomingMessage {
  const request = Readable.from(
    body === undefined ? [] : [JSON.stringify(body)],
  ) as unknown as IncomingMessage;
  request.method = method;
  request.headers = body === undefined ? {} : { 'content-type': 'application/json' };
  return request;
}

class MediaResponse extends Writable {
  statusCode = 0;
  readonly chunks: Buffer[] = [];
  readonly headers = new Map<string, string>();

  setHeader(name: string, value: string | number | readonly string[]): this {
    this.headers.set(name.toLowerCase(), String(value));
    return this;
  }

  _write(
    chunk: Buffer,
    _encoding: BufferEncoding,
    callback: (error?: Error | null) => void,
  ): void {
    this.chunks.push(Buffer.from(chunk));
    callback();
  }
}

const bootstrap = {
  capabilities: {
    conversation: true,
    generation: true,
    image: true,
    video: true,
  },
  workspaces: [],
};

describe('AI design Main route boundary', () => {
  let temporaryDirectory: string | null = null;

  beforeEach(() => {
    electronMocks.getPath.mockReset();
    electronMocks.getPath.mockReturnValue(tmpdir());
    electronMocks.showSaveDialog.mockReset();
  });

  afterEach(async () => {
    if (temporaryDirectory) {
      await rm(temporaryDirectory, { recursive: true, force: true });
      temporaryDirectory = null;
    }
  });

  it('does not claim unrelated routes', async () => {
    const response = createResponse();
    const handled = await handleImageWorkspaceRoutes(
      { method: 'GET' } as IncomingMessage,
      response.res,
      new URL('http://127.0.0.1/api/works/projects'),
      {} as never,
    );

    expect(handled).toBe(false);
    expect(response.res.end).not.toHaveBeenCalled();
  });

  it('returns a stable unavailable response only when Main has no module', async () => {
    const response = createResponse();
    await handleImageWorkspaceRoutes(
      { method: 'GET' } as IncomingMessage,
      response.res,
      new URL('http://127.0.0.1/api/works/image-workspace'),
      {} as never,
    );

    expect(response.statusCode).toBe(501);
    expect(response.json()).toMatchObject({
      success: false,
      code: 'IMAGE_WORKSPACE_UNAVAILABLE',
      error: 'AI 设计暂不可用',
    });
  });

  it('maps Host routes to the deep Workspace module instead of provider-shaped calls', async () => {
    const workspace = {
      bootstrap: vi.fn().mockResolvedValue(bootstrap),
      createWorkspace: vi.fn().mockResolvedValue({ workspaceId: 'workspace-one' }),
      createConversation: vi.fn(),
      getConversation: vi.fn(),
      renameWorkspace: vi.fn(),
      getWorkspace: vi.fn(),
      submitMessage: vi.fn().mockResolvedValue({ workspaceId: 'workspace/one' }),
      confirmGeneration: vi.fn().mockResolvedValue({ workspaceId: 'workspace/one' }),
      listTasks: vi.fn().mockResolvedValue([]),
      uploadAsset: vi.fn().mockResolvedValue({ assetId: 'asset-uploaded' }),
      openAssetContent: vi.fn(),
      getCapabilities: vi.fn(),
      reset: vi.fn().mockResolvedValue(bootstrap),
    };
    const ctx = { imageWorkspace: workspace } as unknown as HostApiContext;

    const getResponse = createResponse();
    await handleImageWorkspaceRoutes(
      createRequest('GET'),
      getResponse.res,
      new URL('http://127.0.0.1/api/works/image-workspace'),
      ctx,
    );
    expect(getResponse.json()).toMatchObject({ success: true, data: bootstrap });

    const createResponseBody = createResponse();
    await handleImageWorkspaceRoutes(
      createRequest('POST', {
        clientWorkspaceId: 'client-one',
        title: '角色设计',
      }),
      createResponseBody.res,
      new URL('http://127.0.0.1/api/works/image-workspace/workspaces'),
      ctx,
    );
    expect(workspace.createWorkspace).toHaveBeenCalledWith({
      clientWorkspaceId: 'client-one',
      title: '角色设计',
    });

    const messageResponse = createResponse();
    await handleImageWorkspaceRoutes(
      createRequest('POST', {
        clientTurnId: 'turn-one',
        expectedTurnRevision: 3,
        message: '继续编辑',
        attachmentAssetIds: [],
        modelId: 'must-be-ignored',
      }),
      messageResponse.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/conversations/conversation%2Fone/messages',
      ),
      ctx,
    );
    expect(workspace.submitMessage).toHaveBeenCalledWith({
      workspaceId: 'workspace/one',
      conversationId: 'conversation/one',
      clientTurnId: 'turn-one',
      expectedTurnRevision: 3,
      message: '继续编辑',
      attachmentAssetIds: [],
    });

    const uploadResponse = createResponse();
    await handleImageWorkspaceRoutes(
      createRequest('POST', {
        fileName: 'poster.webp',
        mimeType: 'image/webp',
        dataBase64: 'AQID',
      }),
      uploadResponse.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/assets',
      ),
      ctx,
    );
    expect(workspace.uploadAsset).toHaveBeenCalledWith({
      workspaceId: 'workspace/one',
      fileName: 'poster.webp',
      mimeType: 'image/webp',
      bytes: Buffer.from([1, 2, 3]),
    });
    expect(uploadResponse.json()).toMatchObject({
      success: true,
      data: { assetId: 'asset-uploaded' },
    });

    const confirmResponse = createResponse();
    await handleImageWorkspaceRoutes(
      createRequest('POST', {
        clientTurnId: 'turn-two',
        expectedTurnRevision: 4,
      }),
      confirmResponse.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/conversations/conversation%2Fone/quotes/quote%2Fone/confirm',
      ),
      ctx,
    );
    expect(workspace.confirmGeneration).toHaveBeenCalledWith({
      workspaceId: 'workspace/one',
      conversationId: 'conversation/one',
      quoteId: 'quote/one',
      clientTurnId: 'turn-two',
      expectedTurnRevision: 4,
    });
  });

  it('streams private asset bytes and preserves Range response headers', async () => {
    const openAssetContent = vi.fn().mockResolvedValue(new Response('partial', {
      status: 206,
      headers: {
        'Accept-Ranges': 'bytes',
        'Content-Range': 'bytes 0-6/100',
        'Content-Type': 'video/mp4',
      },
    }));
    const ctx = {
      imageWorkspace: { openAssetContent },
    } as unknown as HostApiContext;
    const request = createRequest('GET');
    request.headers = { range: 'bytes=0-6' };
    const response = new MediaResponse();

    const handled = await handleImageWorkspaceRoutes(
      request,
      response as unknown as ServerResponse,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace-one/assets/asset-one/content',
      ),
      ctx,
    );

    expect(handled).toBe(true);
    expect(openAssetContent).toHaveBeenCalledWith(
      'workspace-one',
      'asset-one',
      'bytes=0-6',
    );
    expect(response.statusCode).toBe(206);
    expect(response.headers.get('content-range')).toBe('bytes 0-6/100');
    expect(Buffer.concat(response.chunks).toString()).toBe('partial');
  });

  it('streams a private image asset into the native save location', async () => {
    temporaryDirectory = await mkdtemp(join(tmpdir(), 'makelore-design-download-'));
    const savedPath = join(temporaryDirectory, 'design.png');
    electronMocks.showSaveDialog.mockResolvedValue({ canceled: false, filePath: savedPath });
    const openAssetContent = vi.fn().mockResolvedValue(new Response('image-bytes', {
      status: 200,
      headers: { 'Content-Type': 'image/png' },
    }));
    const response = createResponse();

    await handleImageWorkspaceRoutes(
      createRequest('POST', { defaultFileName: 'Ocean poster.png' }),
      response.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/assets/asset%2Fone/download',
      ),
      { imageWorkspace: { openAssetContent } } as unknown as HostApiContext,
    );

    expect(electronMocks.showSaveDialog).toHaveBeenCalledWith(expect.objectContaining({
      defaultPath: expect.stringContaining('Ocean poster.png'),
    }));
    expect(openAssetContent).toHaveBeenCalledWith('workspace/one', 'asset/one');
    expect(response.json()).toMatchObject({
      success: true,
      data: { status: 'saved' },
    });
    await expect(readFile(savedPath, 'utf8')).resolves.toBe('image-bytes');
  });

  it('preserves an existing destination when the private image stream fails', async () => {
    temporaryDirectory = await mkdtemp(join(tmpdir(), 'makelore-design-download-'));
    const savedPath = join(temporaryDirectory, 'existing.png');
    await writeFile(savedPath, 'original');
    electronMocks.showSaveDialog.mockResolvedValue({ canceled: false, filePath: savedPath });
    const body = new ReadableStream<Uint8Array>({
      start(controller) {
        controller.enqueue(new TextEncoder().encode('partial'));
        controller.error(new Error('connection lost'));
      },
    });
    const response = createResponse();

    await handleImageWorkspaceRoutes(
      createRequest('POST', { defaultFileName: 'Ocean poster.png' }),
      response.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace-one/assets/asset-one/download',
      ),
      {
        imageWorkspace: {
          openAssetContent: vi.fn().mockResolvedValue(new Response(body, {
            status: 200,
            headers: { 'Content-Type': 'image/png' },
          })),
        },
      } as unknown as HostApiContext,
    );

    expect(response.json()).toMatchObject({
      success: false,
      code: 'DESIGN_ASSET_SAVE_FAILED',
    });
    await expect(readFile(savedPath, 'utf8')).resolves.toBe('original');
  });

  it('does not fetch private bytes when the native save dialog is cancelled', async () => {
    electronMocks.showSaveDialog.mockResolvedValue({ canceled: true });
    const openAssetContent = vi.fn();
    const response = createResponse();

    await handleImageWorkspaceRoutes(
      createRequest('POST', { defaultFileName: 'Ocean poster.png' }),
      response.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace-one/assets/asset-one/download',
      ),
      { imageWorkspace: { openAssetContent } } as unknown as HostApiContext,
    );

    expect(openAssetContent).not.toHaveBeenCalled();
    expect(response.json()).toMatchObject({
      success: true,
      data: { status: 'cancelled' },
    });
  });

  it('rejects an oversized image upload before reading it into the Main process', async () => {
    const uploadAsset = vi.fn();
    const request = createRequest('POST', {});
    request.headers['content-length'] = String(14 * 1024 * 1024);
    const response = createResponse();

    await handleImageWorkspaceRoutes(
      request,
      response.res,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace-one/assets',
      ),
      { imageWorkspace: { uploadAsset } } as unknown as HostApiContext,
    );

    expect(response.statusCode).toBe(413);
    expect(response.json()).toMatchObject({
      success: false,
      code: 'design_asset_upload_too_large',
    });
    expect(uploadAsset).not.toHaveBeenCalled();
  });

  it('relays normalized task events over local SSE and forwards the opaque resume cursor', async () => {
    const close = vi.fn();
    const openWorkspaceEvents = vi.fn().mockResolvedValue({
      events: (async function* () {
        yield {
          id: 'session-one:7',
          type: 'design.generation_task.updated' as const,
          workspaceId: 'workspace/one',
          workspaceViewRevision: 8,
          generationTask: {
            taskId: 'task-live',
            workspaceId: 'workspace/one',
            conversationId: 'conversation/one',
            medium: 'image' as const,
            status: 'succeeded' as const,
            briefVersion: 1,
            briefSummary: '海洋公益海报',
            quoteId: 'quote-live',
            quotedDesignPoints: 1,
            failureCode: null,
            resultAssets: [],
            createdAt: '2026-08-02T10:00:00Z',
            updatedAt: '2026-08-02T10:02:00Z',
          },
        };
      })(),
      close,
    });
    const request = createRequest('GET');
    request.headers = { 'last-event-id': 'session-one:6' };
    const response = new MediaResponse();

    const handled = await handleImageWorkspaceRoutes(
      request,
      response as unknown as ServerResponse,
      new URL(
        'http://127.0.0.1/api/works/image-workspace/workspaces/workspace%2Fone/conversations/conversation%2Fone/events',
      ),
      { imageWorkspace: { openWorkspaceEvents } } as unknown as HostApiContext,
    );

    expect(handled).toBe(true);
    expect(openWorkspaceEvents).toHaveBeenCalledWith({
      workspaceId: 'workspace/one',
      conversationId: 'conversation/one',
      afterEventId: 'session-one:6',
    });
    expect(response.statusCode).toBe(200);
    expect(response.headers.get('content-type')).toBe('text/event-stream; charset=utf-8');
    expect(response.headers.get('x-accel-buffering')).toBe('no');
    expect(Buffer.concat(response.chunks).toString()).toContain([
      'id: session-one:7',
      'event: design.generation_task.updated',
      'data: {"id":"session-one:7"',
    ].join('\n'));
    expect(close).toHaveBeenCalledOnce();
  });
});