diff options
Diffstat (limited to 'packages/api')
| -rw-r--r-- | packages/api/src/agent-manager.ts | 37 | ||||
| -rw-r--r-- | packages/api/tests/agent-manager.test.ts | 156 |
2 files changed, 190 insertions, 3 deletions
diff --git a/packages/api/src/agent-manager.ts b/packages/api/src/agent-manager.ts index efdd732..92caf81 100644 --- a/packages/api/src/agent-manager.ts +++ b/packages/api/src/agent-manager.ts @@ -37,6 +37,7 @@ import { resolveApiKey, type SkillDefinition, type SystemChunkKind, + type TabStatusSnapshot, TaskList, updateMessage, validateConfig, @@ -711,10 +712,40 @@ export class AgentManager { return this.tabAgents.get(tabId)?.status ?? "idle"; } - getAllStatuses(): Record<string, AgentStatus> { - const result: Record<string, AgentStatus> = {}; + /** + * Snapshot of every tab the manager is currently tracking. Sent on WS + * connect and via GET /status so a freshly-loaded frontend can + * reconstruct any in-flight assistant turn without missing the chunks + * that arrived before its WS handshake completed. + * + * For each running tab, the snapshot includes: + * - status: "running" + * - currentChunks: a defensive shallow copy of `tabAgent.currentChunks` + * (the live chunk array the streaming loop appends to). The + * consumer owns this copy and may mutate it freely. + * - currentAssistantId: the DB id of the in-flight assistant message + * row. The frontend aligns its local assistant message id with + * this so the next `done` event lands on the right message. + * + * For idle/error tabs, only `status` is present. Tabs not in + * `this.tabAgents` (e.g. tabs in the DB that have never been touched + * since server start) are absent from the returned record — the + * caller infers their status from the DB row (always "idle" at rest). + */ + getAllStatuses(): Record<string, TabStatusSnapshot> { + const result: Record<string, TabStatusSnapshot> = {}; for (const [tabId, tabAgent] of this.tabAgents.entries()) { - result[tabId] = tabAgent.status; + const snap: TabStatusSnapshot = { status: tabAgent.status }; + if (tabAgent.status === "running") { + if (tabAgent.currentChunks) { + // Defensive shallow copy: callers may serialize/mutate. + snap.currentChunks = [...tabAgent.currentChunks]; + } + if (tabAgent.currentAssistantId) { + snap.currentAssistantId = tabAgent.currentAssistantId; + } + } + result[tabId] = snap; } return result; } diff --git a/packages/api/tests/agent-manager.test.ts b/packages/api/tests/agent-manager.test.ts index 71d43d8..ba14cad 100644 --- a/packages/api/tests/agent-manager.test.ts +++ b/packages/api/tests/agent-manager.test.ts @@ -658,4 +658,160 @@ describe("AgentManager", () => { expect(init[0]).toMatchObject({ role: "user", chunks: [{ type: "text", text: "q1" }] }); expect(init[1]).toMatchObject({ role: "assistant", chunks: [{ type: "text", text: "a1" }] }); }); + + // ─── getAllStatuses snapshot shape (for browser-reopen restore) ──── + // + // The snapshot enriches the legacy `Record<string, AgentStatus>` shape + // with per-tab in-flight context so a fresh frontend can render the + // streaming assistant message correctly after a reload. + + it("getAllStatuses returns an empty record when no tabs are tracked", () => { + const manager = new AgentManager(); + expect(manager.getAllStatuses()).toEqual({}); + }); + + it("getAllStatuses returns { status } for an idle tab (no currentChunks/currentAssistantId)", async () => { + const manager = new AgentManager(); + // Drive a full turn so the tab gets registered; default mock run + // settles back to idle by the time `await` resolves. + await manager.processMessage("tab-idle", "hi"); + const snap = manager.getAllStatuses(); + expect(snap["tab-idle"]).toBeDefined(); + expect(snap["tab-idle"]?.status).toBe("idle"); + expect(snap["tab-idle"]).not.toHaveProperty("currentChunks"); + expect(snap["tab-idle"]).not.toHaveProperty("currentAssistantId"); + }); + + it("getAllStatuses includes currentChunks and currentAssistantId for a running tab", () => { + const manager = new AgentManager(); + // Reach into the private map to set up a synthetic running state. + // Justification: there is no public API to enter a sustained + // "running" state without actually streaming, and we want to + // assert the snapshot shape — not the streaming pipeline. + const inner = manager as unknown as { + tabAgents: Map< + string, + { + agent: null; + status: "running" | "idle" | "error"; + keyId: null; + modelId: null; + taskList: { onChange: (cb: unknown) => void }; + messageQueue: unknown[]; + queueListeners: unknown[]; + shellStore: unknown; + transcriptStore: unknown; + currentChunks: Array<{ type: string; text?: string }> | null; + currentAssistantId: string | null; + } + >; + }; + inner.tabAgents.set("tab-running", { + agent: null, + status: "running", + keyId: null, + modelId: null, + taskList: { onChange: () => {} }, + messageQueue: [], + queueListeners: [], + shellStore: {}, + transcriptStore: {}, + currentChunks: [ + { type: "thinking", text: "let me think" }, + { type: "text", text: "partial answer" }, + ], + currentAssistantId: "assistant-msg-id-7", + }); + + const snap = manager.getAllStatuses(); + expect(snap["tab-running"]).toBeDefined(); + expect(snap["tab-running"]?.status).toBe("running"); + expect(snap["tab-running"]?.currentAssistantId).toBe("assistant-msg-id-7"); + expect(snap["tab-running"]?.currentChunks).toEqual([ + { type: "thinking", text: "let me think" }, + { type: "text", text: "partial answer" }, + ]); + }); + + it("getAllStatuses defensively copies currentChunks (mutating the snapshot doesn't affect the live array)", () => { + const manager = new AgentManager(); + const inner = manager as unknown as { + tabAgents: Map< + string, + { + agent: null; + status: "running"; + keyId: null; + modelId: null; + taskList: { onChange: (cb: unknown) => void }; + messageQueue: unknown[]; + queueListeners: unknown[]; + shellStore: unknown; + transcriptStore: unknown; + currentChunks: Array<{ type: string; text?: string }>; + currentAssistantId: string; + } + >; + }; + const liveChunks = [{ type: "text", text: "live" }]; + inner.tabAgents.set("tab-copy", { + agent: null, + status: "running", + keyId: null, + modelId: null, + taskList: { onChange: () => {} }, + messageQueue: [], + queueListeners: [], + shellStore: {}, + transcriptStore: {}, + currentChunks: liveChunks, + currentAssistantId: "msg-x", + }); + + const snap = manager.getAllStatuses(); + // Mutate the snapshot's array + snap["tab-copy"]?.currentChunks?.push({ type: "text", text: "polluted" }); + // Live array must be untouched + expect(liveChunks).toEqual([{ type: "text", text: "live" }]); + }); + + it("getAllStatuses omits currentChunks when a running tab has none yet", () => { + const manager = new AgentManager(); + const inner = manager as unknown as { + tabAgents: Map< + string, + { + agent: null; + status: "running"; + keyId: null; + modelId: null; + taskList: { onChange: (cb: unknown) => void }; + messageQueue: unknown[]; + queueListeners: unknown[]; + shellStore: unknown; + transcriptStore: unknown; + currentChunks: null; + currentAssistantId: null; + } + >; + }; + inner.tabAgents.set("tab-early", { + agent: null, + status: "running", + keyId: null, + modelId: null, + taskList: { onChange: () => {} }, + messageQueue: [], + queueListeners: [], + shellStore: {}, + transcriptStore: {}, + currentChunks: null, + currentAssistantId: null, + }); + + const snap = manager.getAllStatuses(); + expect(snap["tab-early"]?.status).toBe("running"); + expect(snap["tab-early"]).not.toHaveProperty("currentChunks"); + expect(snap["tab-early"]).not.toHaveProperty("currentAssistantId"); + }); }); |