From e4ee3e64f4e5e7b704d61467c7b93c30e99017f4 Mon Sep 17 00:00:00 2001 From: Harivansh Rathi Date: Sun, 8 Mar 2026 22:12:03 -0700 Subject: [PATCH] single chat sot --- .../coding-agent/src/core/gateway/index.ts | 1 + .../src/core/gateway/internal-types.ts | 4 + .../coding-agent/src/core/gateway/runtime.ts | 159 +++++++------- .../src/core/gateway/session-state.ts | 195 ++++++++++++++++++ .../coding-agent/src/core/gateway/types.ts | 5 + .../test/gateway-session-state.test.mjs | 163 +++++++++++++++ 6 files changed, 443 insertions(+), 84 deletions(-) create mode 100644 packages/coding-agent/src/core/gateway/session-state.ts create mode 100644 packages/coding-agent/test/gateway-session-state.test.mjs diff --git a/packages/coding-agent/src/core/gateway/index.ts b/packages/coding-agent/src/core/gateway/index.ts index f66ac61..179e06a 100644 --- a/packages/coding-agent/src/core/gateway/index.ts +++ b/packages/coding-agent/src/core/gateway/index.ts @@ -12,6 +12,7 @@ export type { GatewayMessageResult, GatewayRuntimeOptions, GatewaySessionFactory, + GatewaySessionState, GatewaySessionSnapshot, HistoryMessage, HistoryPart, diff --git a/packages/coding-agent/src/core/gateway/internal-types.ts b/packages/coding-agent/src/core/gateway/internal-types.ts index 8c34e28..467d13b 100644 --- a/packages/coding-agent/src/core/gateway/internal-types.ts +++ b/packages/coding-agent/src/core/gateway/internal-types.ts @@ -1,9 +1,11 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentSession } from "../agent-session.js"; import type { GatewayMessageRequest, GatewayMessageResult, GatewaySessionSnapshot, } from "./types.js"; +import type { GatewayTransientToolResult } from "./session-state.js"; export interface GatewayQueuedMessage { request: GatewayMessageRequest; @@ -60,6 +62,8 @@ export interface ManagedGatewaySession { session: AgentSession; queue: GatewayQueuedMessage[]; processing: boolean; + activeAssistantMessage: AgentMessage | null; + pendingToolResults: GatewayTransientToolResult[]; createdAt: number; lastActiveAt: number; listeners: Set<(event: GatewayEvent) => void>; diff --git a/packages/coding-agent/src/core/gateway/runtime.ts b/packages/coding-agent/src/core/gateway/runtime.ts index 8623047..afcd89e 100644 --- a/packages/coding-agent/src/core/gateway/runtime.ts +++ b/packages/coding-agent/src/core/gateway/runtime.ts @@ -7,7 +7,6 @@ import { import { rm } from "node:fs/promises"; import { join } from "node:path"; import { URL } from "node:url"; -import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentSession, AgentSessionEvent } from "../agent-session.js"; import { extractMessageText, getLastAssistantText } from "./helpers.js"; import { @@ -24,6 +23,7 @@ import type { GatewayMessageResult, GatewayRuntimeOptions, GatewaySessionFactory, + GatewaySessionState, GatewaySessionSnapshot, HistoryMessage, HistoryPart, @@ -36,6 +36,10 @@ import { extractUserText, finishVercelStream, } from "./vercel-ai-stream.js"; +import { + buildGatewaySessionStateMessages, + messageContentToHistoryParts, +} from "./session-state.js"; export { createGatewaySessionManager, sanitizeSessionKey, @@ -47,6 +51,7 @@ export type { GatewayMessageResult, GatewayRuntimeOptions, GatewaySessionFactory, + GatewaySessionState, GatewaySessionSnapshot, HistoryMessage, HistoryPart, @@ -246,6 +251,8 @@ export class GatewayRuntime { this.rejectQueuedMessages(managedSession, "Session reset"); await managedSession.session.newSession(); managedSession.processing = false; + managedSession.activeAssistantMessage = null; + managedSession.pendingToolResults = []; managedSession.lastActiveAt = Date.now(); this.emitState(managedSession); return; @@ -284,6 +291,8 @@ export class GatewayRuntime { session, queue: [], processing: false, + activeAssistantMessage: null, + pendingToolResults: [], createdAt: Date.now(), lastActiveAt: Date.now(), listeners: new Set(), @@ -359,6 +368,8 @@ export class GatewayRuntime { } finally { queued.onFinish?.(); managedSession.processing = false; + managedSession.activeAssistantMessage = null; + managedSession.pendingToolResults = []; managedSession.lastActiveAt = Date.now(); this.emitState(managedSession); if (managedSession.queue.length > 0) { @@ -396,18 +407,24 @@ export class GatewayRuntime { ): void { switch (event.type) { case "turn_start": + managedSession.lastActiveAt = Date.now(); this.emit(managedSession, { type: "turn_start", sessionKey: managedSession.sessionKey, }); return; case "turn_end": + managedSession.lastActiveAt = Date.now(); this.emit(managedSession, { type: "turn_end", sessionKey: managedSession.sessionKey, }); return; case "message_start": + managedSession.lastActiveAt = Date.now(); + if (event.message.role === "assistant") { + managedSession.activeAssistantMessage = event.message; + } this.emit(managedSession, { type: "message_start", sessionKey: managedSession.sessionKey, @@ -415,6 +432,10 @@ export class GatewayRuntime { }); return; case "message_update": + managedSession.lastActiveAt = Date.now(); + if (event.message.role === "assistant") { + managedSession.activeAssistantMessage = event.message; + } switch (event.assistantMessageEvent.type) { case "text_delta": this.emit(managedSession, { @@ -435,15 +456,32 @@ export class GatewayRuntime { } return; case "message_end": + managedSession.lastActiveAt = Date.now(); if (event.message.role === "assistant") { + managedSession.activeAssistantMessage = null; this.emit(managedSession, { type: "message_complete", sessionKey: managedSession.sessionKey, text: extractMessageText(event.message), }); + return; + } + if (event.message.role === "toolResult") { + const toolCallId = + typeof (event.message as { toolCallId?: unknown }).toolCallId === + "string" + ? ((event.message as { toolCallId: string }).toolCallId ?? "") + : ""; + if (toolCallId) { + managedSession.pendingToolResults = + managedSession.pendingToolResults.filter( + (entry) => entry.toolCallId !== toolCallId, + ); + } } return; case "tool_execution_start": + managedSession.lastActiveAt = Date.now(); this.emit(managedSession, { type: "tool_start", sessionKey: managedSession.sessionKey, @@ -453,6 +491,7 @@ export class GatewayRuntime { }); return; case "tool_execution_update": + managedSession.lastActiveAt = Date.now(); this.emit(managedSession, { type: "tool_update", sessionKey: managedSession.sessionKey, @@ -462,6 +501,19 @@ export class GatewayRuntime { }); return; case "tool_execution_end": + managedSession.lastActiveAt = Date.now(); + managedSession.pendingToolResults = [ + ...managedSession.pendingToolResults.filter( + (entry) => entry.toolCallId !== event.toolCallId, + ), + { + toolCallId: event.toolCallId, + toolName: event.toolName, + result: event.result, + isError: event.isError, + timestamp: Date.now(), + }, + ]; this.emit(managedSession, { type: "tool_complete", sessionKey: managedSession.sessionKey, @@ -491,6 +543,20 @@ export class GatewayRuntime { }); } + private createSessionState( + managedSession: ManagedGatewaySession, + ): GatewaySessionState { + return { + session: this.createSnapshot(managedSession), + messages: buildGatewaySessionStateMessages({ + sessionKey: managedSession.sessionKey, + rawMessages: managedSession.session.messages, + activeAssistantMessage: managedSession.activeAssistantMessage, + pendingToolResults: managedSession.pendingToolResults, + }), + }; + } + private createSnapshot( managedSession: ManagedGatewaySession, ): GatewaySessionSnapshot { @@ -731,7 +797,7 @@ export class GatewayRuntime { } const sessionMatch = path.match( - /^\/sessions\/([^/]+)(?:\/(events|messages|abort|reset|chat|history|model|reload))?$/, + /^\/sessions\/([^/]+)(?:\/(events|messages|abort|reset|chat|history|model|reload|state))?$/, ); if (!sessionMatch) { this.writeJson(response, 404, { error: "Not found" }); @@ -809,6 +875,12 @@ export class GatewayRuntime { return; } + if (action === "state" && method === "GET") { + const session = await this.ensureSession(sessionKey); + this.writeJson(response, 200, this.createSessionState(session)); + return; + } + if (action === "model" && method === "POST") { const body = await this.readJsonBody(request); const provider = typeof body.provider === "string" ? body.provider : ""; @@ -1080,7 +1152,7 @@ export class GatewayRuntime { messages.push({ id: `${msg.timestamp}-${msg.role}-${index}`, role: msg.role, - parts: this.messageContentToParts(msg), + parts: messageContentToHistoryParts(msg), timestamp: msg.timestamp, }); } @@ -1173,87 +1245,6 @@ export class GatewayRuntime { managed.session.settingsManager.reload(); } - private messageContentToParts(msg: AgentMessage): HistoryPart[] { - if (msg.role === "user") { - const content = msg.content; - if (typeof content === "string") { - return [{ type: "text", text: content }]; - } - if (Array.isArray(content)) { - return content - .filter( - (c): c is { type: "text"; text: string } => - typeof c === "object" && c !== null && c.type === "text", - ) - .map((c) => ({ type: "text" as const, text: c.text })); - } - return []; - } - - if (msg.role === "assistant") { - const content = msg.content; - if (!Array.isArray(content)) return []; - const parts: HistoryPart[] = []; - for (const c of content) { - if (typeof c !== "object" || c === null) continue; - if (c.type === "text") { - parts.push({ - type: "text", - text: (c as { type: "text"; text: string }).text, - }); - } else if (c.type === "thinking") { - parts.push({ - type: "reasoning", - text: (c as { type: "thinking"; thinking: string }).thinking, - }); - } else if (c.type === "toolCall") { - const tc = c as { - type: "toolCall"; - id: string; - name: string; - arguments: unknown; - }; - parts.push({ - type: "tool-invocation", - toolCallId: tc.id, - toolName: tc.name, - args: tc.arguments, - state: "call", - }); - } - } - return parts; - } - - if (msg.role === "toolResult") { - const tr = msg as { - role: "toolResult"; - toolCallId: string; - toolName: string; - content: unknown; - isError: boolean; - }; - const textParts = Array.isArray(tr.content) - ? (tr.content as { type: string; text?: string }[]) - .filter((c) => c.type === "text" && typeof c.text === "string") - .map((c) => c.text as string) - .join("") - : ""; - return [ - { - type: "tool-invocation", - toolCallId: tr.toolCallId, - toolName: tr.toolName, - args: undefined, - state: tr.isError ? "error" : "result", - result: textParts, - }, - ]; - } - - return []; - } - getGatewaySessionDir(sessionKey: string): string { return join(this.sessionDirRoot, sanitizeSessionKey(sessionKey)); } diff --git a/packages/coding-agent/src/core/gateway/session-state.ts b/packages/coding-agent/src/core/gateway/session-state.ts new file mode 100644 index 0000000..fee5ef7 --- /dev/null +++ b/packages/coding-agent/src/core/gateway/session-state.ts @@ -0,0 +1,195 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { HistoryMessage, HistoryPart } from "./types.js"; + +export interface GatewayTransientToolResult { + toolCallId: string; + toolName: string; + result: unknown; + isError: boolean; + timestamp: number; +} + +function isSupportedHistoryRole( + role: AgentMessage["role"], +): role is "user" | "assistant" | "toolResult" { + return role === "user" || role === "assistant" || role === "toolResult"; +} + +function historyMessageId(message: AgentMessage, index: number): string { + return `${message.timestamp}-${message.role}-${index}`; +} + +function transientAssistantId( + sessionKey: string, + message: AgentMessage | null, +): string { + return `draft:${sessionKey}:${message?.timestamp ?? 0}`; +} + +function transientToolResultId(sessionKey: string, toolCallId: string): string { + return `draft-tool:${sessionKey}:${toolCallId}`; +} + +export function messageContentToHistoryParts(msg: AgentMessage): HistoryPart[] { + if (msg.role === "user") { + const content = msg.content; + if (typeof content === "string") { + return [{ type: "text", text: content }]; + } + if (Array.isArray(content)) { + return content + .filter( + (contentPart): contentPart is { type: "text"; text: string } => + typeof contentPart === "object" && + contentPart !== null && + contentPart.type === "text", + ) + .map((contentPart) => ({ + type: "text" as const, + text: contentPart.text, + })); + } + return []; + } + + if (msg.role === "assistant") { + const content = msg.content; + if (!Array.isArray(content)) return []; + const parts: HistoryPart[] = []; + for (const contentPart of content) { + if (typeof contentPart !== "object" || contentPart === null) { + continue; + } + if (contentPart.type === "text") { + parts.push({ + type: "text", + text: (contentPart as { type: "text"; text: string }).text, + }); + } else if (contentPart.type === "thinking") { + parts.push({ + type: "reasoning", + text: (contentPart as { type: "thinking"; thinking: string }) + .thinking, + }); + } else if (contentPart.type === "toolCall") { + const toolCall = contentPart as { + type: "toolCall"; + id: string; + name: string; + arguments: unknown; + }; + parts.push({ + type: "tool-invocation", + toolCallId: toolCall.id, + toolName: toolCall.name, + args: toolCall.arguments, + state: "call", + }); + } + } + return parts; + } + + if (msg.role === "toolResult") { + const toolResult = msg as { + role: "toolResult"; + toolCallId: string; + toolName: string; + content: unknown; + isError: boolean; + }; + const textParts = Array.isArray(toolResult.content) + ? (toolResult.content as { type: string; text?: string }[]) + .filter((contentPart) => { + return ( + contentPart.type === "text" && + typeof contentPart.text === "string" + ); + }) + .map((contentPart) => contentPart.text as string) + .join("") + : ""; + + return [ + { + type: "tool-invocation", + toolCallId: toolResult.toolCallId, + toolName: toolResult.toolName, + args: undefined, + state: toolResult.isError ? "error" : "result", + result: textParts, + }, + ]; + } + + return []; +} + +export function buildGatewaySessionStateMessages(params: { + sessionKey: string; + rawMessages: AgentMessage[]; + activeAssistantMessage: AgentMessage | null; + pendingToolResults: GatewayTransientToolResult[]; +}): HistoryMessage[] { + const { + sessionKey, + rawMessages, + activeAssistantMessage, + pendingToolResults, + } = params; + const messages: HistoryMessage[] = []; + const persistedToolCallIds = new Set(); + + for (const [index, message] of rawMessages.entries()) { + if (!isSupportedHistoryRole(message.role)) { + continue; + } + + if ( + message.role === "toolResult" && + typeof (message as { toolCallId?: unknown }).toolCallId === "string" + ) { + persistedToolCallIds.add((message as { toolCallId: string }).toolCallId); + } + + messages.push({ + id: historyMessageId(message, index), + role: message.role, + parts: messageContentToHistoryParts(message), + timestamp: message.timestamp, + }); + } + + if (activeAssistantMessage?.role === "assistant") { + messages.push({ + id: transientAssistantId(sessionKey, activeAssistantMessage), + role: "assistant", + parts: messageContentToHistoryParts(activeAssistantMessage), + timestamp: activeAssistantMessage.timestamp ?? Date.now(), + }); + } + + for (const pendingToolResult of pendingToolResults) { + if (persistedToolCallIds.has(pendingToolResult.toolCallId)) { + continue; + } + + messages.push({ + id: transientToolResultId(sessionKey, pendingToolResult.toolCallId), + role: "toolResult", + parts: [ + { + type: "tool-invocation", + toolCallId: pendingToolResult.toolCallId, + toolName: pendingToolResult.toolName, + args: undefined, + state: pendingToolResult.isError ? "error" : "result", + result: pendingToolResult.result, + }, + ], + timestamp: pendingToolResult.timestamp, + }); + } + + return messages; +} diff --git a/packages/coding-agent/src/core/gateway/types.ts b/packages/coding-agent/src/core/gateway/types.ts index 410717b..49803f6 100644 --- a/packages/coding-agent/src/core/gateway/types.ts +++ b/packages/coding-agent/src/core/gateway/types.ts @@ -48,6 +48,11 @@ export interface GatewaySessionSnapshot { updatedAt: number; } +export interface GatewaySessionState { + session: GatewaySessionSnapshot; + messages: HistoryMessage[]; +} + export interface ModelInfo { provider: string; modelId: string; diff --git a/packages/coding-agent/test/gateway-session-state.test.mjs b/packages/coding-agent/test/gateway-session-state.test.mjs new file mode 100644 index 0000000..670b692 --- /dev/null +++ b/packages/coding-agent/test/gateway-session-state.test.mjs @@ -0,0 +1,163 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { + buildGatewaySessionStateMessages, + messageContentToHistoryParts, +} from "../src/core/gateway/session-state.ts"; + +test("messageContentToHistoryParts converts assistant text, reasoning, and tool calls", () => { + const parts = messageContentToHistoryParts({ + role: "assistant", + timestamp: 123, + content: [ + { type: "text", text: "hello" }, + { type: "thinking", thinking: "working" }, + { + type: "toolCall", + id: "tool-1", + name: "bash", + arguments: { cmd: "pwd" }, + }, + ], + }); + + assert.deepEqual(parts, [ + { type: "text", text: "hello" }, + { type: "reasoning", text: "working" }, + { + type: "tool-invocation", + toolCallId: "tool-1", + toolName: "bash", + args: { cmd: "pwd" }, + state: "call", + }, + ]); +}); + +test("buildGatewaySessionStateMessages includes the active assistant draft while a run is still processing", () => { + const messages = buildGatewaySessionStateMessages({ + sessionKey: "agent:test:chat", + rawMessages: [ + { + role: "user", + timestamp: 100, + content: "build a todo app", + }, + ], + activeAssistantMessage: { + role: "assistant", + timestamp: 200, + content: [ + { type: "text", text: "Working on it" }, + { + type: "toolCall", + id: "tool-1", + name: "write", + arguments: { filePath: "app.tsx" }, + }, + ], + }, + pendingToolResults: [], + }); + + assert.deepEqual(messages, [ + { + id: "100-user-0", + role: "user", + parts: [{ type: "text", text: "build a todo app" }], + timestamp: 100, + }, + { + id: "draft:agent:test:chat:200", + role: "assistant", + parts: [ + { type: "text", text: "Working on it" }, + { + type: "tool-invocation", + toolCallId: "tool-1", + toolName: "write", + args: { filePath: "app.tsx" }, + state: "call", + }, + ], + timestamp: 200, + }, + ]); +}); + +test("buildGatewaySessionStateMessages keeps transient tool results until persisted history catches up", () => { + const messages = buildGatewaySessionStateMessages({ + sessionKey: "agent:test:chat", + rawMessages: [ + { + role: "user", + timestamp: 100, + content: "ship it", + }, + { + role: "toolResult", + timestamp: 210, + toolCallId: "tool-1", + toolName: "bash", + content: [{ type: "text", text: "done" }], + isError: false, + }, + ], + activeAssistantMessage: null, + pendingToolResults: [ + { + toolCallId: "tool-1", + toolName: "bash", + result: { stdout: "done" }, + isError: false, + timestamp: 220, + }, + { + toolCallId: "tool-2", + toolName: "write", + result: { path: "README.md" }, + isError: false, + timestamp: 240, + }, + ], + }); + + assert.deepEqual(messages, [ + { + id: "100-user-0", + role: "user", + parts: [{ type: "text", text: "ship it" }], + timestamp: 100, + }, + { + id: "210-toolResult-1", + role: "toolResult", + parts: [ + { + type: "tool-invocation", + toolCallId: "tool-1", + toolName: "bash", + args: undefined, + state: "result", + result: "done", + }, + ], + timestamp: 210, + }, + { + id: "draft-tool:agent:test:chat:tool-2", + role: "toolResult", + parts: [ + { + type: "tool-invocation", + toolCallId: "tool-2", + toolName: "write", + args: undefined, + state: "result", + result: { path: "README.md" }, + }, + ], + timestamp: 240, + }, + ]); +});