diff --git a/packages/server/src/server/agent/agent-manager.test.ts b/packages/server/src/server/agent/agent-manager.test.ts index 1d1ae6853..9639d40e3 100644 --- a/packages/server/src/server/agent/agent-manager.test.ts +++ b/packages/server/src/server/agent/agent-manager.test.ts @@ -712,7 +712,7 @@ describe("AgentManager", () => { expect(afterReload?.config?.title).toBeUndefined(); }); - test("resumeAgentFromPersistence seeds live helpers from durable rows and skips provider replay", async () => { + test("resumeAgentFromPersistence reads durable helpers without loading committed rows into live memory", async () => { const workdir = mkdtempSync(join(tmpdir(), "agent-manager-durable-seed-")); const storagePath = join(workdir, "agents"); const dataDir = join(workdir, "db"); @@ -803,13 +803,8 @@ describe("AgentManager", () => { const resumed = await manager.resumeAgentFromPersistence(handle, undefined, snapshot.id); expect(resumed.id).toBe(snapshot.id); - expect(manager.getTimeline(snapshot.id)).toEqual([ - { - type: "assistant_message", - text: "durable only", - }, - ]); - expect(manager.getLastAssistantMessage(snapshot.id)).toBe("durable only"); + expect(manager.getTimeline(snapshot.id)).toEqual([]); + await expect(manager.getLastAssistantMessage(snapshot.id)).resolves.toBe("durable only"); await expect(manager.getTimelineRows(snapshot.id)).resolves.toEqual([ { seq: 1, @@ -824,12 +819,7 @@ describe("AgentManager", () => { await manager.hydrateTimelineFromProvider(snapshot.id); expect(historyReplayCount).toBe(0); - expect(manager.getTimeline(snapshot.id)).toEqual([ - { - type: "assistant_message", - text: "durable only", - }, - ]); + expect(manager.getTimeline(snapshot.id)).toEqual([]); await manager.closeAgent(snapshot.id); await manager.deleteCommittedTimeline(snapshot.id); @@ -1261,7 +1251,7 @@ describe("AgentManager", () => { await durableTimelineStore.bulkInsert(snapshot.id, [durableOnlyRow]); expect(manager.getTimeline(snapshot.id)).toEqual([]); - expect(manager.getLastAssistantMessage(snapshot.id)).toBeNull(); + await expect(manager.getLastAssistantMessage(snapshot.id)).resolves.toBe("durable only"); await expect(manager.getTimelineRows(snapshot.id)).resolves.toEqual([durableOnlyRow]); await expect( manager.fetchTimeline(snapshot.id, { diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index 28279bac6..41f8f88cf 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -279,6 +279,7 @@ export class AgentManager { private readonly clients = new Map(); private readonly agents = new Map(); private readonly timelineStore = new InMemoryAgentTimelineStore(); + private readonly sessionEventTails = new Map>(); private readonly pendingForegroundRuns = new Map(); private readonly subscribers = new Set(); private readonly idFactory: () => string; @@ -526,7 +527,8 @@ export class AgentManager { } // Reconstruct an agent from provider persistence. When a durable timeline - // store is configured, committed history is seeded from the durable store. + // store is configured, the live timeline buffer only seeds seq metadata from + // the durable store instead of loading committed history back into memory. // Tests without a durable timeline store can still call // hydrateTimelineFromProvider() for backward compatibility. async resumeAgentFromPersistence( @@ -1417,20 +1419,27 @@ export class AgentManager { await this.durableTimelineStore.deleteAgent(agentId); } - getLastAssistantMessage(agentId: string): string | null { + async getLastAssistantMessage(agentId: string): Promise { const agent = this.agents.get(agentId); if (!agent) { return null; } - return this.timelineStore.getLastAssistantMessage(agentId); + return await this.getLastAssistantMessageFromStores(agentId); } private getLastAssistantMessageFromTimeline( timeline: readonly AgentTimelineItem[], ): string | null { + return this.getLastAssistantMessageSegmentFromTimeline(timeline)?.text ?? null; + } + + private getLastAssistantMessageSegmentFromTimeline( + timeline: readonly AgentTimelineItem[], + ): { text: string; startsAtBeginning: boolean } | null { // Collect the last contiguous assistant messages (Claude streams chunks) const chunks: string[] = []; + let startsAtBeginning = false; for (let i = timeline.length - 1; i >= 0; i--) { const item = timeline[i]; if (item.type !== "assistant_message") { @@ -1440,13 +1449,65 @@ export class AgentManager { continue; } chunks.push(item.text); + startsAtBeginning = i === 0; } if (!chunks.length) { return null; } - return chunks.reverse().join(""); + return { + text: chunks.reverse().join(""), + startsAtBeginning, + }; + } + + private async getLastAssistantMessageFromStores(agentId: string): Promise { + const liveTimeline = this.timelineStore.getItems(agentId); + const liveSegment = this.getLastAssistantMessageSegmentFromTimeline(liveTimeline); + if (!this.durableTimelineStore) { + return liveSegment?.text ?? null; + } + + if (!liveSegment) { + return await this.durableTimelineStore.getLastAssistantMessage(agentId); + } + + if (!liveSegment.startsAtBeginning) { + return liveSegment.text; + } + + const lastDurableItem = await this.durableTimelineStore.getLastItem(agentId); + if (lastDurableItem?.type !== "assistant_message") { + return liveSegment.text; + } + + const durableMessage = await this.durableTimelineStore.getLastAssistantMessage(agentId); + return durableMessage ? `${durableMessage}${liveSegment.text}` : liveSegment.text; + } + + private async getLastItemFromStores(agentId: string): Promise { + const lastLiveItem = this.timelineStore.getLastItem(agentId); + if (lastLiveItem) { + return lastLiveItem; + } + if (!this.durableTimelineStore) { + return null; + } + return await this.durableTimelineStore.getLastItem(agentId); + } + + private async hasCommittedUserMessageFromStores( + agentId: string, + options: { messageId: string; text: string }, + ): Promise { + if (this.timelineStore.hasCommittedUserMessage(agentId, options)) { + return true; + } + if (!this.durableTimelineStore) { + return false; + } + return await this.durableTimelineStore.hasCommittedUserMessage(agentId, options); } async waitForAgentEvent( @@ -1466,7 +1527,7 @@ export class AgentManager { return { status: snapshot.lifecycle, permission: immediatePermission, - lastMessage: this.getLastAssistantMessage(agentId), + lastMessage: await this.getLastAssistantMessage(agentId), }; } @@ -1477,14 +1538,14 @@ export class AgentManager { return { status: initialStatus, permission: null, - lastMessage: this.getLastAssistantMessage(agentId), + lastMessage: await this.getLastAssistantMessage(agentId), }; } if (waitForActive && !initialBusy && !hasForegroundTurn) { return { status: initialStatus, permission: null, - lastMessage: this.getLastAssistantMessage(agentId), + lastMessage: await this.getLastAssistantMessage(agentId), }; } @@ -1503,6 +1564,7 @@ export class AgentManager { let currentStatus: AgentLifecycleStatus = initialStatus; let hasStarted = initialBusy || hasForegroundTurn; let terminalStatusOverride: AgentLifecycleStatus | null = null; + let finished = false; // Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them let unsubscribe: (() => void) | null = null; @@ -1531,12 +1593,20 @@ export class AgentManager { }; const finish = (permission: AgentPermissionRequest | null) => { + if (finished) { + return; + } + finished = true; cleanup(); - resolve({ - status: currentStatus, - permission, - lastMessage: this.getLastAssistantMessage(agentId), - }); + void this.getLastAssistantMessage(agentId) + .then((lastMessage) => { + resolve({ + status: currentStatus, + permission, + lastMessage, + }); + }) + .catch(reject); }; // Bug #3 Fix: Set up abort handler BEFORE subscription @@ -1711,18 +1781,9 @@ export class AgentManager { return { timestamp: now.toISOString() }; } - const rows = await this.durableTimelineStore.getCommittedRows(agentId); - if (rows.length === 0) { - return { - nextSeq: 1, - timestamp: now.toISOString(), - }; - } - return { - rows, - nextSeq: rows[rows.length - 1]!.seq + 1, - timestamp: rows[rows.length - 1]!.timestamp, + nextSeq: (await this.durableTimelineStore.getLatestCommittedSeq(agentId)) + 1, + timestamp: now.toISOString(), }; } @@ -1732,16 +1793,42 @@ export class AgentManager { } const agentId = agent.id; const unsubscribe = agent.session.subscribe((event: AgentStreamEvent) => { - const current = this.agents.get(agentId); - if (!current) { - return; - } - this.dispatchSessionEvent(current, event); + this.enqueueSessionEvent(agentId, event); }); agent.unsubscribeSession = unsubscribe; } - private dispatchSessionEvent(agent: ActiveManagedAgent, event: AgentStreamEvent): void { + private enqueueSessionEvent(agentId: string, event: AgentStreamEvent): void { + const previous = this.sessionEventTails.get(agentId) ?? Promise.resolve(); + const next = previous + .catch(() => undefined) + .then(async () => { + const current = this.agents.get(agentId); + if (!current) { + return; + } + await this.dispatchSessionEvent(current, event); + }) + .catch((err) => { + this.logger.error( + { err, agentId, eventType: event.type }, + "Failed to process session event", + ); + }); + + this.sessionEventTails.set(agentId, next); + this.trackBackgroundTask(next); + void next.finally(() => { + if (this.sessionEventTails.get(agentId) === next) { + this.sessionEventTails.delete(agentId); + } + }); + } + + private async dispatchSessionEvent( + agent: ActiveManagedAgent, + event: AgentStreamEvent, + ): Promise { const turnId = (event as { turnId?: string }).turnId; const matchingWaiters = turnId == null @@ -1750,7 +1837,7 @@ export class AgentManager { (waiter) => waiter.turnId === turnId && !waiter.settled, ); - this.handleStreamEvent(agent, event); + await this.handleStreamEvent(agent, event); for (const waiter of matchingWaiters) { waiter.callback(event); @@ -1960,14 +2047,14 @@ export class AgentManager { } } - private handleStreamEvent( + private async handleStreamEvent( agent: ActiveManagedAgent, event: AgentStreamEvent, options?: { fromHistory?: boolean; canonicalUserMessagesById?: ReadonlyMap; }, - ): void { + ): Promise { const eventTurnId = (event as { turnId?: string }).turnId; const isForegroundEvent = Boolean( eventTurnId && agent.activeForegroundTurnId === eventTurnId, @@ -2017,7 +2104,7 @@ export class AgentManager { const eventText = event.item.text; if (eventMessageId) { if ( - this.timelineStore.hasCommittedUserMessage(agent.id, { + await this.hasCommittedUserMessageFromStores(agent.id, { messageId: eventMessageId, text: eventText, }) @@ -2103,7 +2190,7 @@ export class AgentManager { agent.lifecycle = "error"; } agent.lastError = event.error; - this.appendSystemErrorTimelineMessage( + await this.appendSystemErrorTimelineMessage( agent, event.provider, this.formatTurnFailedMessage(event), @@ -2209,7 +2296,7 @@ export class AgentManager { } } - private appendSystemErrorTimelineMessage( + private async appendSystemErrorTimelineMessage( agent: ActiveManagedAgent, provider: AgentProvider, message: string, @@ -2217,7 +2304,7 @@ export class AgentManager { fromHistory?: boolean; canonicalUserMessagesById?: ReadonlyMap; }, - ): void { + ): Promise { if (options?.fromHistory) { return; } @@ -2228,7 +2315,7 @@ export class AgentManager { } const text = `${SYSTEM_ERROR_PREFIX} ${normalized}`; - const lastItem = this.timelineStore.getLastItem(agent.id); + const lastItem = await this.getLastItemFromStores(agent.id); if (lastItem?.type === "assistant_message" && lastItem.text === text) { return; } diff --git a/packages/server/src/server/agent/agent-timeline-store-types.ts b/packages/server/src/server/agent/agent-timeline-store-types.ts index 8882c1c65..1543baa18 100644 --- a/packages/server/src/server/agent/agent-timeline-store-types.ts +++ b/packages/server/src/server/agent/agent-timeline-store-types.ts @@ -49,6 +49,12 @@ export interface AgentTimelineStore { ): Promise; getLatestCommittedSeq(agentId: string): Promise; getCommittedRows(agentId: string): Promise; + getLastItem(agentId: string): Promise; + getLastAssistantMessage(agentId: string): Promise; + hasCommittedUserMessage( + agentId: string, + options: { messageId: string; text: string }, + ): Promise; deleteAgent(agentId: string): Promise; bulkInsert(agentId: string, rows: readonly AgentTimelineRow[]): Promise; } diff --git a/packages/server/src/server/db/db-agent-timeline-store.test.ts b/packages/server/src/server/db/db-agent-timeline-store.test.ts index 3d586701d..2de172121 100644 --- a/packages/server/src/server/db/db-agent-timeline-store.test.ts +++ b/packages/server/src/server/db/db-agent-timeline-store.test.ts @@ -211,4 +211,56 @@ describe("DbAgentTimelineStore", () => { }), ]); }); + + test("getLastItem returns the latest committed item", async () => { + await store.bulkInsert("agent-1", [ + createRow(1, createTimelineItem("user_message", "1")), + createRow(2, createTimelineItem("assistant_message", "2")), + ]); + + await expect(store.getLastItem("agent-1")).resolves.toEqual( + createTimelineItem("assistant_message", "2"), + ); + await expect(store.getLastItem("missing-agent")).resolves.toBeNull(); + }); + + test("getLastAssistantMessage assembles the latest contiguous assistant chunks", async () => { + await store.bulkInsert("agent-1", [ + createRow(1, createTimelineItem("assistant_message", "1")), + createRow(2, createTimelineItem("assistant_message", "2")), + createRow(3, { type: "reasoning", text: "separator-1" }), + createRow(4, createTimelineItem("assistant_message", "4")), + createRow(5, createTimelineItem("assistant_message", "5")), + createRow(6, { type: "reasoning", text: "separator-2" }), + ]); + + await expect(store.getLastAssistantMessage("agent-1")).resolves.toBe("assistant-4assistant-5"); + await expect(store.getLastAssistantMessage("missing-agent")).resolves.toBeNull(); + }); + + test("hasCommittedUserMessage matches by normalized messageId and text", async () => { + await store.bulkInsert("agent-1", [ + createRow(1, createTimelineItem("user_message", "1")), + createRow(2, createTimelineItem("assistant_message", "2")), + ]); + + await expect( + store.hasCommittedUserMessage("agent-1", { + messageId: " message-1 ", + text: "user-1", + }), + ).resolves.toBe(true); + await expect( + store.hasCommittedUserMessage("agent-1", { + messageId: "message-1", + text: "different", + }), + ).resolves.toBe(false); + await expect( + store.hasCommittedUserMessage("agent-1", { + messageId: " ", + text: "user-1", + }), + ).resolves.toBe(false); + }); }); diff --git a/packages/server/src/server/db/db-agent-timeline-store.ts b/packages/server/src/server/db/db-agent-timeline-store.ts index 6a6e78da7..e3baca57d 100644 --- a/packages/server/src/server/db/db-agent-timeline-store.ts +++ b/packages/server/src/server/db/db-agent-timeline-store.ts @@ -16,6 +16,14 @@ type AgentTimelineRowInsert = typeof agentTimelineRows.$inferInsert; const DEFAULT_TIMELINE_FETCH_LIMIT = 200; +function normalizeTimelineMessageId(messageId: string | undefined): string | undefined { + if (typeof messageId !== "string") { + return undefined; + } + const normalized = messageId.trim(); + return normalized.length > 0 ? normalized : undefined; +} + function toTimelineRow(row: AgentTimelineRowRecord): AgentTimelineRow { return { seq: row.seq, @@ -185,6 +193,76 @@ export class DbAgentTimelineStore implements AgentTimelineStore { return rows.map(toTimelineRow); } + async getLastItem(agentId: string): Promise { + const [row] = await this.db + .select({ item: agentTimelineRows.item }) + .from(agentTimelineRows) + .where(eq(agentTimelineRows.agentId, agentId)) + .orderBy(desc(agentTimelineRows.seq)) + .limit(1); + return row?.item ?? null; + } + + async getLastAssistantMessage(agentId: string): Promise { + const rows = await this.db + .select({ + seq: agentTimelineRows.seq, + item: agentTimelineRows.item, + }) + .from(agentTimelineRows) + .where( + and( + eq(agentTimelineRows.agentId, agentId), + eq(agentTimelineRows.itemKind, "assistant_message"), + ), + ) + .orderBy(desc(agentTimelineRows.seq)); + + if (rows.length === 0) { + return null; + } + + const chunks: string[] = []; + let previousSeq: number | null = null; + for (const row of rows) { + if (previousSeq !== null && row.seq !== previousSeq - 1) { + break; + } + if (row.item.type !== "assistant_message") { + break; + } + chunks.push(row.item.text); + previousSeq = row.seq; + } + + return chunks.length > 0 ? chunks.reverse().join("") : null; + } + + async hasCommittedUserMessage( + agentId: string, + options: { messageId: string; text: string }, + ): Promise { + const messageId = normalizeTimelineMessageId(options.messageId); + if (!messageId) { + return false; + } + + const [row] = await this.db + .select({ seq: agentTimelineRows.seq }) + .from(agentTimelineRows) + .where( + and( + eq(agentTimelineRows.agentId, agentId), + eq(agentTimelineRows.itemKind, "user_message"), + sql`${agentTimelineRows.item} ->> 'messageId' = ${messageId}`, + sql`${agentTimelineRows.item} ->> 'text' = ${options.text}`, + ), + ) + .limit(1); + + return row !== undefined; + } + async deleteAgent(agentId: string): Promise { await this.db.delete(agentTimelineRows).where(eq(agentTimelineRows.agentId, agentId)); } diff --git a/packages/server/src/server/websocket-server.notifications.test.ts b/packages/server/src/server/websocket-server.notifications.test.ts index 2744447db..57101a5b5 100644 --- a/packages/server/src/server/websocket-server.notifications.test.ts +++ b/packages/server/src/server/websocket-server.notifications.test.ts @@ -63,7 +63,7 @@ function createServer(agentManagerOverrides?: Record) { const agentManager = { setAgentAttentionCallback: vi.fn(), getAgent: vi.fn(() => null), - getLastAssistantMessage: vi.fn(() => null), + getLastAssistantMessage: vi.fn(async () => null), ...agentManagerOverrides, }; @@ -93,9 +93,9 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => { vi.clearAllMocks(); }); - it("uses assistant preview text for push notifications with markdown removed", () => { + it("uses assistant preview text for push notifications with markdown removed", async () => { const getLastAssistantMessage = vi.fn( - () => "**Done**. Updated `README.md` and [link](https://example.com).", + async () => "**Done**. Updated `README.md` and [link](https://example.com).", ); const { server } = createServer({ getAgent: vi.fn(() => ({ @@ -106,7 +106,7 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => { getLastAssistantMessage, }); - (server as any).broadcastAgentAttention({ + await (server as any).broadcastAgentAttention({ agentId: "agent-1", provider: "claude", reason: "finished", @@ -124,8 +124,8 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => { expect(getLastAssistantMessage).toHaveBeenCalledWith("agent-1"); }); - it("sends push notifications regardless of UI label presence", () => { - const getLastAssistantMessage = vi.fn(() => "Done."); + it("sends push notifications regardless of UI label presence", async () => { + const getLastAssistantMessage = vi.fn(async () => "Done."); const { server } = createServer({ getAgent: vi.fn(() => ({ config: { title: null }, @@ -136,7 +136,7 @@ describe("VoiceAssistantWebSocketServer notification payloads", () => { getLastAssistantMessage, }); - (server as any).broadcastAgentAttention({ + await (server as any).broadcastAgentAttention({ agentId: "agent-2", provider: "claude", reason: "finished", diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 73816592c..2528f3195 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -342,7 +342,9 @@ export class VoiceAssistantWebSocketServer { this.pushService = new PushService(pushLogger, this.pushTokenStore); this.agentManager.setAgentAttentionCallback((params) => { - this.broadcastAgentAttention(params); + void this.broadcastAgentAttention(params).catch((err) => { + this.logger.warn({ err, agentId: params.agentId }, "Failed to broadcast agent attention"); + }); }); const { allowedOrigins, allowedHosts } = wsConfig; @@ -1239,11 +1241,11 @@ export class VoiceAssistantWebSocketServer { }; } - private broadcastAgentAttention(params: { + private async broadcastAgentAttention(params: { agentId: string; provider: AgentProvider; reason: "finished" | "error" | "permission"; - }): void { + }): Promise { const clientEntries: Array<{ ws: WebSocketLike; state: ClientAttentionState; @@ -1258,11 +1260,12 @@ export class VoiceAssistantWebSocketServer { const allStates = clientEntries.map((e) => e.state); const agent = this.agentManager.getAgent(params.agentId); + const assistantMessage = await this.agentManager.getLastAssistantMessage(params.agentId); const notification = buildAgentAttentionNotificationPayload({ reason: params.reason, serverId: this.serverId, agentId: params.agentId, - assistantMessage: this.agentManager.getLastAssistantMessage(params.agentId), + assistantMessage, permissionRequest: agent ? findLatestPermissionRequest(agent.pendingPermissions) : null, });