From bd968b8266b8580214ce2f705fdc8592222c1ab3 Mon Sep 17 00:00:00 2001 From: DoSun Date: Tue, 19 May 2026 21:24:29 +0800 Subject: [PATCH] feat: messages API returns parts with accurate running status for active subagents --- .../src/tools/AgentTool/runAgent.ts | 17 ++++ src/QueryEngine.ts | 5 +- src/server/routes/message.ts | 81 ++++++++++++++++++- src/server/routes/session.ts | 11 ++- src/server/sessionMessageRouter.ts | 6 ++ src/utils/sdkEventQueue.ts | 12 ++- src/utils/sessionStorage.ts | 8 ++ src/utils/task/framework.ts | 1 + src/utils/task/sdkProgress.ts | 1 + 9 files changed, 134 insertions(+), 8 deletions(-) diff --git a/packages/builtin-tools/src/tools/AgentTool/runAgent.ts b/packages/builtin-tools/src/tools/AgentTool/runAgent.ts index 7857b7445..ea04a84ac 100644 --- a/packages/builtin-tools/src/tools/AgentTool/runAgent.ts +++ b/packages/builtin-tools/src/tools/AgentTool/runAgent.ts @@ -747,6 +747,21 @@ export async function* runAgent({ agentType: agentDefinition.agentType, ...(worktreePath && { worktreePath }), ...(description && { description }), + parent_session_id: getSessionId(), + prompt: promptMessages + .map((m: Message) => { + const content = (m as Record).content + if (typeof content === 'string') return content + if (Array.isArray(content)) { + return content + .filter((b: Record) => b.type === 'text') + .map((b: Record) => b.text) + .join('') + } + return '' + }) + .filter(Boolean) + .slice(-1)[0] ?? undefined, }).catch(_err => logForDebugging(`Failed to write agent metadata: ${_err}`)) // Track the last recorded message UUID for parent chain continuity @@ -812,6 +827,7 @@ export async function* runAgent({ ) break } + message.agent_id = agentId yield message as Message continue } @@ -828,6 +844,7 @@ export async function* runAgent({ if (message.type !== 'progress') { lastRecordedUuid = message.uuid } + message.agent_id = agentId yield message } } diff --git a/src/QueryEngine.ts b/src/QueryEngine.ts index af807a122..1a8878461 100644 --- a/src/QueryEngine.ts +++ b/src/QueryEngine.ts @@ -789,9 +789,10 @@ export class QueryEngine { } switch (message.type) { - case 'tombstone': - // Tombstone messages are control signals for removing messages, skip them + case 'tombstone': { + yield message as SDKMessage break + } case 'assistant': { // Capture stop_reason if already set (synthetic messages). For // streamed responses, this is null at content_block_stop time; diff --git a/src/server/routes/message.ts b/src/server/routes/message.ts index 7edaecd3e..eb40506d1 100644 --- a/src/server/routes/message.ts +++ b/src/server/routes/message.ts @@ -4,9 +4,60 @@ import { notFound } from '../errors.js' import { readSessionMessages, readSessionTodos, + readSessionTasks, readSessionDiff, type SessionMessage, + type TaskInfo, + type MessagePart, + decomposeMessageToParts, } from '../transcriptReader.js' +import { getSubagentProgress } from '../sessionMessageRouter.js' + +function findParentCwd(sessionManager: SessionManager, id: string): string | undefined { + for (const handle of sessionManager.getAllSessions()) { + if (handle.activeSubagents.has(id)) { + return handle.spawnCwd ?? handle.cwd + } + } + return undefined +} + +function patchActiveSubagentsInMessages( + messages: SessionMessage[], + activeSubagents: ReadonlyMap, +): void { + if (activeSubagents.size === 0) return + + const runningCallIDs = new Set() + const progressByCallID = new Map() + + for (const [, info] of activeSubagents) { + if (info.toolUseId) { + runningCallIDs.add(info.toolUseId) + const progress = getSubagentProgress(info.agentId) + if (progress?.progressLines?.length) { + progressByCallID.set(info.toolUseId, progress.progressLines) + } + } + } + + if (runningCallIDs.size === 0) return + + for (const msg of messages) { + if (!Array.isArray(msg.parts)) continue + for (const part of msg.parts) { + if (part.type !== 'tool') continue + const toolPart = part as MessagePart & { type: 'tool'; callID: string; state: Record } + if (!runningCallIDs.has(toolPart.callID)) continue + toolPart.state = { + ...toolPart.state, + status: 'running', + progress: progressByCallID.get(toolPart.callID) ?? [], + time: { start: (toolPart.state as any).time?.start ?? Date.now() }, + } as any + } + } +} function dedupeMessages(disk: SessionMessage[], memory: SessionMessage[]): SessionMessage[] { if (disk.length === 0) return memory @@ -20,6 +71,8 @@ export function createMessageRoutes(sessionManager: SessionManager): Hono { .get('/session/:sessionID/message', async c => { const id = c.req.param('sessionID') const handle = sessionManager.getSession(id) + const parentCwd = handle ? undefined : findParentCwd(sessionManager, id) + const effectiveCwd = handle?.spawnCwd ?? handle?.cwd ?? parentCwd const url = new URL(c.req.url) const limit = parseInt(url.searchParams.get('limit') ?? '50', 10) @@ -28,7 +81,7 @@ export function createMessageRoutes(sessionManager: SessionManager): Hono { const { messages: diskMessages, nextCursor } = await readSessionMessages({ sessionId: id, - cwd: handle?.spawnCwd ?? handle?.cwd, + cwd: effectiveCwd, limit, before, includeSystem, @@ -38,12 +91,25 @@ export function createMessageRoutes(sessionManager: SessionManager): Hono { ? dedupeMessages(diskMessages, [...handle.messageBuffer]) : diskMessages + const tombstoned = handle?.tombstonedUuids + const filtered = tombstoned && tombstoned.size > 0 + ? messages.filter(m => !tombstoned.has(m.uuid)) + : messages + + const decomposed = filtered.map(decomposeMessageToParts) + + if (handle) { + patchActiveSubagentsInMessages(decomposed, handle.activeSubagents) + } + + const result = decomposed + if (nextCursor) { c.header('Link', `<${url.pathname}?limit=${limit}&before=${nextCursor}>; rel="prev"`) c.header('X-Next-Cursor', nextCursor) } - return c.json({ messages }) + return c.json({ messages: result }) }) .get('/session/:sessionID/todo', async c => { const id = c.req.param('sessionID') @@ -56,6 +122,17 @@ export function createMessageRoutes(sessionManager: SessionManager): Hono { return c.json({ todos }) }) + .get('/session/:sessionID/tasks', async c => { + const id = c.req.param('sessionID') + const handle = sessionManager.getSession(id) + + const tasks = await readSessionTasks({ + sessionId: id, + cwd: handle?.spawnCwd ?? handle?.cwd, + }) + + return c.json({ tasks }) + }) .get('/session/:sessionID/diff', async c => { const id = c.req.param('sessionID') const handle = sessionManager.getSession(id) diff --git a/src/server/routes/session.ts b/src/server/routes/session.ts index 5f13d1339..bf55377a6 100644 --- a/src/server/routes/session.ts +++ b/src/server/routes/session.ts @@ -528,7 +528,16 @@ export function createSessionRoutes( .post('/session/:sessionID/abort', async c => { const id = c.req.param('sessionID') const handle = sessionManager.getSession(id) - if (!handle) return c.json({ aborted: true }) + if (!handle) { + // Even when no in-memory handle exists, emit idle status so that + // a prior prompt_async busy event (which fires before the handle + // is created) is not left dangling forever. + eventBus.publish('session.status', { + sessionID: id, + status: { type: 'idle' }, + }) + return c.json({ aborted: true }) + } await handle.abort() return c.json({ aborted: true }) }) diff --git a/src/server/sessionMessageRouter.ts b/src/server/sessionMessageRouter.ts index b518d6a33..7cc502d05 100644 --- a/src/server/sessionMessageRouter.ts +++ b/src/server/sessionMessageRouter.ts @@ -147,6 +147,12 @@ export type MessageRouterCtx = { addTombstonedUuid(uuid: string): void } +export function getSubagentProgress(agentId: string): { progressLines: string[]; mainToolUseID: string } | undefined { + const state = subagentToolState.get(agentId) + if (!state) return undefined + return { progressLines: state.progressLines, mainToolUseID: state.mainToolUseID } +} + export function routeMessage(msg: StdoutMessage, ctx: MessageRouterCtx): void { if ( ctx.getStatus() === 'starting' && diff --git a/src/utils/sdkEventQueue.ts b/src/utils/sdkEventQueue.ts index 9880c8961..f66ef3d33 100644 --- a/src/utils/sdkEventQueue.ts +++ b/src/utils/sdkEventQueue.ts @@ -12,6 +12,10 @@ type TaskStartedEvent = { task_type?: string workflow_name?: string prompt?: string + /** Explicit agent_id for virtual session mapping in serve mode. + * Same value as task_id but named for clarity. Enables the serve layer + * to construct virtual child sessions without guessing the relationship. */ + agent_id?: string } type TaskProgressEvent = { @@ -27,10 +31,9 @@ type TaskProgressEvent = { } last_tool_name?: string summary?: string - // Delta batch of workflow state changes. Clients upsert by - // `${type}:${index}` then group by phaseIndex to rebuild the phase tree, - // same fold as collectFromEvents + groupByPhase in PhaseProgress.tsx. workflow_progress?: SdkWorkflowProgress[] + /** Explicit agent_id for virtual session mapping in serve mode. */ + agent_id?: string } // Emitted when a foreground agent completes without being backgrounded. @@ -51,6 +54,8 @@ type TaskNotificationSdkEvent = { tool_uses: number duration_ms: number } + /** Explicit agent_id for virtual session mapping in serve mode. */ + agent_id?: string } // SessionStateChangedEvent is now emitted directly by sessionState.ts's @@ -122,5 +127,6 @@ export function emitTaskTerminatedSdk( output_file: opts?.outputFile ?? '', summary: opts?.summary ?? '', usage: opts?.usage, + agent_id: taskId, }) } diff --git a/src/utils/sessionStorage.ts b/src/utils/sessionStorage.ts index bfa892662..d8e601d1c 100644 --- a/src/utils/sessionStorage.ts +++ b/src/utils/sessionStorage.ts @@ -269,6 +269,14 @@ export type AgentMetadata = { * resumed agent's notification can show the original description instead * of a placeholder. Optional — older metadata files lack this field. */ description?: string + /** Parent session ID for virtual session hierarchy in serve mode. + * Enables the serve layer to reconstruct subagent → parent relationships + * on restart. Optional — older metadata files lack this field. */ + parent_session_id?: string + /** Agent-specific prompt for the subagent. Persisted so the serve layer + * can construct a virtual session with the original prompt context. + * Optional — older metadata files lack this field. */ + prompt?: string } /** diff --git a/src/utils/task/framework.ts b/src/utils/task/framework.ts index 7bb509968..7a569b853 100644 --- a/src/utils/task/framework.ts +++ b/src/utils/task/framework.ts @@ -113,6 +113,7 @@ export function registerTask(task: TaskState, setAppState: SetAppState): void { ? (task.workflowName as string | undefined) : undefined, prompt: 'prompt' in task ? (task.prompt as string) : undefined, + agent_id: task.id, }) } diff --git a/src/utils/task/sdkProgress.ts b/src/utils/task/sdkProgress.ts index 1430df295..c7e2da5a2 100644 --- a/src/utils/task/sdkProgress.ts +++ b/src/utils/task/sdkProgress.ts @@ -32,5 +32,6 @@ export function emitTaskProgress(params: { last_tool_name: params.lastToolName, summary: params.summary, workflow_progress: params.workflowProgress, + agent_id: params.taskId, }) }