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

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { useImageWorkspaceStore } from '@/stores/image-workspace';
import type {
  DesignGenerationTask,
  DesignGenerationTasksSnapshotEvent,
  DesignGenerationTaskUpdatedEvent,
  DesignWorkspace,
  DesignWorkspaceBootstrap,
} from '../../shared/image-workspace';

const fetchImageWorkspaceMock = vi.hoisted(() => vi.fn());
const fetchImageWorkspaceProjectMock = vi.hoisted(() => vi.fn());
const fetchImageWorkspaceTasksMock = vi.hoisted(() => vi.fn());
const openImageWorkspaceTaskEventsMock = vi.hoisted(() => vi.fn());

vi.mock('@/lib/image-workspace', async (importOriginal) => {
  const actual = await importOriginal<typeof import('@/lib/image-workspace')>();
  return {
    ...actual,
    fetchImageWorkspace: (...args: unknown[]) => fetchImageWorkspaceMock(...args),
    fetchImageWorkspaceProject: (...args: unknown[]) => fetchImageWorkspaceProjectMock(...args),
    fetchImageWorkspaceTasks: (...args: unknown[]) => fetchImageWorkspaceTasksMock(...args),
    openImageWorkspaceTaskEvents: (...args: unknown[]) => openImageWorkspaceTaskEventsMock(...args),
  };
});

type EventListener = (event: MessageEvent<string>) => void;

class MockEventSource {
  onopen: ((event: Event) => void) | null = null;
  onerror: ((event: Event) => void) | null = null;
  readonly close = vi.fn();
  private readonly listeners = new Map<string, Set<EventListener>>();

  addEventListener(type: string, listener: EventListener): void {
    const listeners = this.listeners.get(type) ?? new Set<EventListener>();
    listeners.add(listener);
    this.listeners.set(type, listeners);
  }

  emit(type: string, payload: unknown): void {
    const event = { data: JSON.stringify(payload) } as MessageEvent<string>;
    for (const listener of this.listeners.get(type) ?? []) listener(event);
  }
}

function deferred<T>(): { promise: Promise<T>; resolve(value: T): void } {
  let resolve!: (value: T) => void;
  const promise = new Promise<T>((accept) => {
    resolve = accept;
  });
  return { promise, resolve };
}

const task: DesignGenerationTask = {
  taskId: 'task-one',
  workspaceId: 'workspace-one',
  medium: 'image',
  status: 'queued',
  briefVersion: 1,
  briefSummary: '海洋公益海报',
  quoteId: 'quote-one',
  quotedDesignPoints: 1,
  failureCode: null,
  resultAssets: [],
  createdAt: '2026-08-02T10:00:00Z',
  updatedAt: '2026-08-02T10:00:00Z',
};

function workspace(workspaceId = 'workspace-one', viewRevision = 1): DesignWorkspace {
  return {
    workspaceId,
    title: workspaceId,
    turnRevision: 1,
    viewRevision,
    phase: 'shaping',
    brief: {
      version: 1,
      status: 'ready',
      medium: 'image',
      summary: '海洋公益海报',
      ready: true,
      missingDecision: null,
    },
    updatedAt: '2026-08-02T10:00:00Z',
    messages: [],
  };
}

function bootstrap(workspaceIds = ['workspace-one']): DesignWorkspaceBootstrap {
  return {
    capabilities: { conversation: true, generation: true, image: true, video: true },
    workspaces: workspaceIds.map((workspaceId) => {
      const { messages: _messages, ...summary } = workspace(workspaceId);
      return summary;
    }),
  };
}

function taskEvent(
  workspaceViewRevision: number,
  status: DesignGenerationTask['status'],
): DesignGenerationTaskUpdatedEvent {
  return {
    id: `session-one:${workspaceViewRevision}`,
    type: 'design.generation_task.updated',
    workspaceId: 'workspace-one',
    workspaceViewRevision,
    generationTask: {
      ...task,
      status,
      updatedAt: `2026-08-02T10:0${workspaceViewRevision}:00Z`,
    },
  };
}

function taskSnapshotEvent(
  workspaceViewRevision: number,
  status: DesignGenerationTask['status'],
): DesignGenerationTasksSnapshotEvent {
  return {
    id: `session-one:${workspaceViewRevision}`,
    type: 'design.generation_tasks.snapshot',
    workspaceId: 'workspace-one',
    workspaceViewRevision,
    generationTasks: [{
      ...task,
      status,
      updatedAt: `2026-08-02T10:0${workspaceViewRevision}:00Z`,
    }],
  };
}

describe('AI design task event store', () => {
  beforeEach(() => {
    vi.useRealTimers();
    vi.clearAllMocks();
    useImageWorkspaceStore.getState().reset();
    fetchImageWorkspaceMock.mockResolvedValue(bootstrap());
    fetchImageWorkspaceProjectMock.mockResolvedValue(workspace());
    fetchImageWorkspaceTasksMock.mockResolvedValue([task]);
  });

  afterEach(() => {
    useImageWorkspaceStore.getState().reset();
    vi.useRealTimers();
  });

  it('keeps a revision-5 snapshot over revision 3, then upserts a newer task by task id', async () => {
    const source = new MockEventSource();
    openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource);
    fetchImageWorkspaceProjectMock.mockResolvedValue(workspace('workspace-one', 5));

    await useImageWorkspaceStore.getState().load();
    await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock)
      .toHaveBeenCalledWith('workspace-one'));
    source.onopen?.(new Event('open'));
    source.emit('design.generation_task.updated', taskEvent(3, 'failed'));
    expect(useImageWorkspaceStore.getState().tasks[0].status).toBe('queued');
    source.emit('design.generation_tasks.snapshot', taskSnapshotEvent(6, 'running'));
    source.emit('design.generation_task.updated', taskEvent(5, 'failed'));

    expect(useImageWorkspaceStore.getState()).toMatchObject({
      taskStreamState: 'connected',
      workspace: { viewRevision: 6 },
      tasks: [{ taskId: 'task-one', status: 'running' }],
    });
  });

  it('closes the previous stream when switching workspaces and on reset', async () => {
    const first = new MockEventSource();
    const second = new MockEventSource();
    openImageWorkspaceTaskEventsMock
      .mockResolvedValueOnce(first as unknown as EventSource)
      .mockResolvedValueOnce(second as unknown as EventSource);
    fetchImageWorkspaceMock.mockResolvedValue(bootstrap(['workspace-one', 'workspace-two']));
    fetchImageWorkspaceProjectMock.mockImplementation((workspaceId: string) => (
      Promise.resolve(workspace(workspaceId))
    ));

    await useImageWorkspaceStore.getState().load();
    await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledTimes(1));
    await useImageWorkspaceStore.getState().selectProject('workspace-two');
    await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledTimes(2));

    expect(first.close).toHaveBeenCalledOnce();
    useImageWorkspaceStore.getState().reset();
    expect(second.close).toHaveBeenCalledOnce();
    expect(useImageWorkspaceStore.getState().taskStreamState).toBe('idle');
  });

  it('does not let a slow previous Workspace selection overwrite the latest one', async () => {
    const rootSource = new MockEventSource();
    const latestSource = new MockEventSource();
    const slowWorkspace = deferred<DesignWorkspace>();
    openImageWorkspaceTaskEventsMock
      .mockResolvedValueOnce(rootSource as unknown as EventSource)
      .mockResolvedValueOnce(latestSource as unknown as EventSource);
    fetchImageWorkspaceMock.mockResolvedValue(bootstrap([
      'workspace-root',
      'workspace-one',
      'workspace-two',
    ]));
    fetchImageWorkspaceProjectMock.mockImplementation((workspaceId: string) => (
      workspaceId === 'workspace-one'
        ? slowWorkspace.promise
        : Promise.resolve(workspace(workspaceId))
    ));
    fetchImageWorkspaceTasksMock.mockResolvedValue([]);

    await useImageWorkspaceStore.getState().load();
    await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock)
      .toHaveBeenCalledWith('workspace-root'));
    const staleSelection = useImageWorkspaceStore.getState().selectProject('workspace-one');
    await vi.waitFor(() => expect(fetchImageWorkspaceProjectMock)
      .toHaveBeenCalledWith('workspace-one'));
    await useImageWorkspaceStore.getState().selectProject('workspace-two');
    slowWorkspace.resolve(workspace('workspace-one'));
    await staleSelection;

    expect(useImageWorkspaceStore.getState().activeWorkspaceId).toBe('workspace-two');
    expect(useImageWorkspaceStore.getState().workspace?.workspaceId).toBe('workspace-two');
    expect(openImageWorkspaceTaskEventsMock).not.toHaveBeenCalledWith('workspace-one');
  });

  it('advances each task watermark after a degraded REST snapshot', async () => {
    const source = new MockEventSource();
    openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource);

    await useImageWorkspaceStore.getState().load();
    await vi.waitFor(() => expect(openImageWorkspaceTaskEventsMock).toHaveBeenCalledOnce());
    fetchImageWorkspaceProjectMock.mockResolvedValue(workspace('workspace-one', 8));
    fetchImageWorkspaceTasksMock.mockResolvedValue([{
      ...task,
      status: 'succeeded',
    }]);
    await useImageWorkspaceStore.getState().refreshWorkspace();
    await useImageWorkspaceStore.getState().refreshTasks();
    source.emit('design.generation_task.updated', taskEvent(7, 'failed'));

    expect(useImageWorkspaceStore.getState().tasks[0].status).toBe('succeeded');
  });

  it('polls slowly only while the event stream is degraded', async () => {
    vi.useFakeTimers();
    const source = new MockEventSource();
    openImageWorkspaceTaskEventsMock.mockResolvedValue(source as unknown as EventSource);

    await useImageWorkspaceStore.getState().load();
    await vi.waitFor(() => expect(source.onerror).not.toBeNull());
    source.onerror?.(new Event('error'));
    await vi.advanceTimersByTimeAsync(14_999);
    expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(1);
    await vi.advanceTimersByTimeAsync(1);
    expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(2);

    source.onopen?.(new Event('open'));
    await vi.advanceTimersByTimeAsync(15_000);
    expect(fetchImageWorkspaceTasksMock).toHaveBeenCalledTimes(2);
  });
});