import { MODEL_GROUPS, buildInitialMockLayoutViewModel, groupWorkspaceRepositories, nowMs, providerAgent, randomReply, removeFileTreePath, slugify, uid, } from "../workspace-model.js"; import { DEFAULT_WORKSPACE_MODEL_ID, workspaceAgentForModel } from "@sandbox-agent/foundry-shared"; import type { TaskWorkspaceAddSessionResponse, TaskWorkspaceChangeModelInput, TaskWorkspaceCreateTaskInput, TaskWorkspaceCreateTaskResponse, TaskWorkspaceDiffInput, TaskWorkspaceRenameInput, TaskWorkspaceRenameSessionInput, TaskWorkspaceSelectInput, TaskWorkspaceSetSessionUnreadInput, TaskWorkspaceSendMessageInput, TaskWorkspaceSnapshot, TaskWorkspaceSessionInput, TaskWorkspaceUpdateDraftInput, WorkspaceSession as AgentSession, WorkspaceTask as Task, WorkspaceTranscriptEvent as TranscriptEvent, } from "@sandbox-agent/foundry-shared"; import type { TaskWorkspaceClient } from "../workspace-client.js"; function buildTranscriptEvent(params: { sessionId: string; sender: "client" | "agent"; createdAt: number; payload: unknown; eventIndex: number; }): TranscriptEvent { return { id: uid(), sessionId: params.sessionId, sender: params.sender, createdAt: params.createdAt, payload: params.payload, connectionId: "mock-connection", eventIndex: params.eventIndex, }; } class MockWorkspaceStore implements TaskWorkspaceClient { private snapshot = buildInitialMockLayoutViewModel(); private listeners = new Set<() => void>(); private pendingTimers = new Map>(); getSnapshot(): TaskWorkspaceSnapshot { return this.snapshot; } subscribe(listener: () => void): () => void { this.listeners.add(listener); return () => { this.listeners.delete(listener); }; } async createTask(input: TaskWorkspaceCreateTaskInput): Promise { const id = uid(); const sessionId = `session-${id}`; const repo = this.snapshot.repos.find((candidate) => candidate.id === input.repoId); if (!repo) { throw new Error(`Cannot create mock task for unknown repo ${input.repoId}`); } const nextTask: Task = { id, repoId: repo.id, title: input.title?.trim() || "New Task", status: "init_enqueue_provision", repoName: repo.label, updatedAtMs: nowMs(), branch: input.branch?.trim() || null, pullRequest: null, activeSessionId: sessionId, sessions: [ { id: sessionId, sessionId: sessionId, sessionName: "Session 1", agent: workspaceAgentForModel(input.model ?? DEFAULT_WORKSPACE_MODEL_ID, MODEL_GROUPS), model: input.model ?? DEFAULT_WORKSPACE_MODEL_ID, status: "idle", thinkingSinceMs: null, unread: false, created: false, draft: { text: "", attachments: [], updatedAtMs: null }, transcript: [], }, ], fileChanges: [], diffs: {}, fileTree: [], minutesUsed: 0, }; this.updateState((current) => ({ ...current, tasks: [nextTask, ...current.tasks], })); return { taskId: id, sessionId }; } async markTaskUnread(input: TaskWorkspaceSelectInput): Promise { this.updateTask(input.taskId, (task) => { const targetSession = task.sessions[task.sessions.length - 1] ?? null; if (!targetSession) { return task; } return { ...task, sessions: task.sessions.map((session) => (session.id === targetSession.id ? { ...session, unread: true } : session)), }; }); } async renameTask(input: TaskWorkspaceRenameInput): Promise { const value = input.value.trim(); if (!value) { throw new Error(`Cannot rename task ${input.taskId} to an empty title`); } this.updateTask(input.taskId, (task) => ({ ...task, title: value, updatedAtMs: nowMs() })); } async archiveTask(input: TaskWorkspaceSelectInput): Promise { this.updateTask(input.taskId, (task) => ({ ...task, status: "archived", updatedAtMs: nowMs() })); } async publishPr(input: TaskWorkspaceSelectInput): Promise { const nextPrNumber = Math.max(0, ...this.snapshot.tasks.map((task) => task.pullRequest?.number ?? 0)) + 1; this.updateTask(input.taskId, (task) => ({ ...task, updatedAtMs: nowMs(), pullRequest: { number: nextPrNumber, title: task.title, state: "open", url: `https://example.test/pr/${nextPrNumber}`, headRefName: task.branch ?? `task/${task.id}`, baseRefName: "main", repoFullName: task.repoName, authorLogin: "mock", isDraft: false, updatedAtMs: nowMs(), }, })); } async revertFile(input: TaskWorkspaceDiffInput): Promise { this.updateTask(input.taskId, (task) => { const file = task.fileChanges.find((entry) => entry.path === input.path); const nextDiffs = { ...task.diffs }; delete nextDiffs[input.path]; return { ...task, fileChanges: task.fileChanges.filter((entry) => entry.path !== input.path), diffs: nextDiffs, fileTree: file?.type === "A" ? removeFileTreePath(task.fileTree, input.path) : task.fileTree, }; }); } async updateDraft(input: TaskWorkspaceUpdateDraftInput): Promise { this.assertSession(input.taskId, input.sessionId); this.updateTask(input.taskId, (task) => ({ ...task, updatedAtMs: nowMs(), sessions: task.sessions.map((tab) => tab.id === input.sessionId ? { ...tab, draft: { text: input.text, attachments: input.attachments, updatedAtMs: nowMs(), }, } : tab, ), })); } async sendMessage(input: TaskWorkspaceSendMessageInput): Promise { const text = input.text.trim(); if (!text) { throw new Error(`Cannot send an empty mock prompt for task ${input.taskId}`); } this.assertSession(input.taskId, input.sessionId); const startedAtMs = nowMs(); this.updateTask(input.taskId, (currentTask) => { const isFirstOnTask = String(currentTask.status).startsWith("init_"); const newTitle = isFirstOnTask ? (text.length > 50 ? `${text.slice(0, 47)}...` : text) : currentTask.title; const newBranch = isFirstOnTask ? `feat/${slugify(newTitle)}` : currentTask.branch; const userMessageLines = [text, ...input.attachments.map((attachment) => `@ ${attachment.filePath}:${attachment.lineNumber}`)]; const userEvent = buildTranscriptEvent({ sessionId: input.sessionId, sender: "client", createdAt: startedAtMs, eventIndex: candidateEventIndex(currentTask, input.sessionId), payload: { method: "session/prompt", params: { prompt: userMessageLines.map((line) => ({ type: "text", text: line })), }, }, }); return { ...currentTask, title: newTitle, branch: newBranch, status: "running", updatedAtMs: startedAtMs, sessions: currentTask.sessions.map((candidate) => candidate.id === input.sessionId ? { ...candidate, created: true, status: "running", unread: false, thinkingSinceMs: startedAtMs, draft: { text: "", attachments: [], updatedAtMs: startedAtMs }, transcript: [...candidate.transcript, userEvent], } : candidate, ), }; }); const existingTimer = this.pendingTimers.get(input.sessionId); if (existingTimer) { clearTimeout(existingTimer); } const timer = setTimeout(() => { const task = this.requireTask(input.taskId); this.requireSession(task, input.sessionId); const completedAtMs = nowMs(); const replyEvent = buildTranscriptEvent({ sessionId: input.sessionId, sender: "agent", createdAt: completedAtMs, eventIndex: candidateEventIndex(task, input.sessionId), payload: { result: { text: randomReply(), durationMs: completedAtMs - startedAtMs, }, }, }); this.updateTask(input.taskId, (currentTask) => { const updatedTabs = currentTask.sessions.map((candidate) => { if (candidate.id !== input.sessionId) { return candidate; } return { ...candidate, status: "idle" as const, thinkingSinceMs: null, unread: true, transcript: [...candidate.transcript, replyEvent], }; }); const anyRunning = updatedTabs.some((candidate) => candidate.status === "running"); return { ...currentTask, updatedAtMs: completedAtMs, sessions: updatedTabs, status: currentTask.status === "archived" ? "archived" : anyRunning ? "running" : "idle", }; }); this.pendingTimers.delete(input.sessionId); }, 2_500); this.pendingTimers.set(input.sessionId, timer); } async stopAgent(input: TaskWorkspaceSessionInput): Promise { this.assertSession(input.taskId, input.sessionId); const existing = this.pendingTimers.get(input.sessionId); if (existing) { clearTimeout(existing); this.pendingTimers.delete(input.sessionId); } this.updateTask(input.taskId, (currentTask) => { const updatedTabs = currentTask.sessions.map((candidate) => candidate.id === input.sessionId ? { ...candidate, status: "idle" as const, thinkingSinceMs: null } : candidate, ); const anyRunning = updatedTabs.some((candidate) => candidate.status === "running"); return { ...currentTask, updatedAtMs: nowMs(), sessions: updatedTabs, status: currentTask.status === "archived" ? "archived" : anyRunning ? "running" : "idle", }; }); } async selectSession(input: TaskWorkspaceSessionInput): Promise { this.assertSession(input.taskId, input.sessionId); this.updateTask(input.taskId, (currentTask) => ({ ...currentTask, activeSessionId: input.sessionId, })); } async setSessionUnread(input: TaskWorkspaceSetSessionUnreadInput): Promise { this.updateTask(input.taskId, (currentTask) => ({ ...currentTask, sessions: currentTask.sessions.map((candidate) => (candidate.id === input.sessionId ? { ...candidate, unread: input.unread } : candidate)), })); } async renameSession(input: TaskWorkspaceRenameSessionInput): Promise { const title = input.title.trim(); if (!title) { throw new Error(`Cannot rename session ${input.sessionId} to an empty title`); } this.updateTask(input.taskId, (currentTask) => ({ ...currentTask, sessions: currentTask.sessions.map((candidate) => (candidate.id === input.sessionId ? { ...candidate, sessionName: title } : candidate)), })); } async closeSession(input: TaskWorkspaceSessionInput): Promise { this.updateTask(input.taskId, (currentTask) => { if (currentTask.sessions.length <= 1) { return currentTask; } return { ...currentTask, activeSessionId: currentTask.activeSessionId === input.sessionId ? (currentTask.sessions.find((candidate) => candidate.id !== input.sessionId)?.id ?? null) : currentTask.activeSessionId, sessions: currentTask.sessions.filter((candidate) => candidate.id !== input.sessionId), }; }); } async addSession(input: TaskWorkspaceSelectInput): Promise { this.assertTask(input.taskId); const nextSessionId = uid(); const nextSession: AgentSession = { id: nextSessionId, sessionId: nextSessionId, sandboxSessionId: null, sessionName: `Session ${this.requireTask(input.taskId).sessions.length + 1}`, agent: workspaceAgentForModel(DEFAULT_WORKSPACE_MODEL_ID, MODEL_GROUPS), model: DEFAULT_WORKSPACE_MODEL_ID, status: "idle", thinkingSinceMs: null, unread: false, created: false, draft: { text: "", attachments: [], updatedAtMs: null }, transcript: [], }; this.updateTask(input.taskId, (currentTask) => ({ ...currentTask, updatedAtMs: nowMs(), activeSessionId: nextSession.id, sessions: [...currentTask.sessions, nextSession], })); return { sessionId: nextSession.id }; } async changeModel(input: TaskWorkspaceChangeModelInput): Promise { const group = MODEL_GROUPS.find((candidate) => candidate.models.some((entry) => entry.id === input.model)); if (!group) { throw new Error(`Unable to resolve model provider for ${input.model}`); } this.updateTask(input.taskId, (currentTask) => ({ ...currentTask, sessions: currentTask.sessions.map((candidate) => candidate.id === input.sessionId ? { ...candidate, model: input.model, agent: workspaceAgentForModel(input.model, MODEL_GROUPS) } : candidate, ), })); } private updateState(updater: (current: TaskWorkspaceSnapshot) => TaskWorkspaceSnapshot): void { const nextSnapshot = updater(this.snapshot); this.snapshot = { ...nextSnapshot, repositories: groupWorkspaceRepositories(nextSnapshot.repos, nextSnapshot.tasks), }; this.notify(); } private updateTask(taskId: string, updater: (task: Task) => Task): void { this.assertTask(taskId); this.updateState((current) => ({ ...current, tasks: current.tasks.map((task) => (task.id === taskId ? updater(task) : task)), })); } private notify(): void { for (const listener of this.listeners) { listener(); } } private assertTask(taskId: string): void { this.requireTask(taskId); } private assertSession(taskId: string, sessionId: string): void { const task = this.requireTask(taskId); this.requireSession(task, sessionId); } private requireTask(taskId: string): Task { const task = this.snapshot.tasks.find((candidate) => candidate.id === taskId); if (!task) { throw new Error(`Unable to find mock task ${taskId}`); } return task; } private requireSession(task: Task, sessionId: string): AgentSession { const session = task.sessions.find((candidate) => candidate.id === sessionId); if (!session) { throw new Error(`Unable to find mock session ${sessionId} in task ${task.id}`); } return session; } } function candidateEventIndex(task: Task, sessionId: string): number { const session = task.sessions.find((candidate) => candidate.id === sessionId); return (session?.transcript.length ?? 0) + 1; } let sharedMockWorkspaceClient: TaskWorkspaceClient | null = null; export function getSharedMockWorkspaceClient(): TaskWorkspaceClient { if (!sharedMockWorkspaceClient) { sharedMockWorkspaceClient = new MockWorkspaceStore(); } return sharedMockWorkspaceClient; }