From 0ebb0582b0cb8714913c397acb129a9b3bfa2a86 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A8=8B=E5=BA=8F=E5=91=98=E9=98=BF=E6=B1=9F=28Relakkes?= =?UTF-8?q?=29?= Date: Wed, 3 Jun 2026 11:25:17 +0800 Subject: [PATCH] fix: stream agent slash progress to desktop (#653) The desktop /agent slash path runs before the normal query loop, so the old implementation waited for the forked command to finish and only surfaced the final local-command stdout. Emit foreground slash-agent task events from /agent, drain SDK events while QueryEngine waits for pre-query slash processing, and allow those current-turn events through the WebSocket pre-send mute gate. Constraint: /agent is a forked slash command that does not produce ordinary assistant deltas until its final stdout is synthesized. Rejected: Frontend-only loading state | would not expose real forked-agent progress or tool/tool-result updates. Confidence: high Scope-risk: moderate Directive: Keep slash_agent task events scoped to /agent; ordinary slash command stdout and /goal event handling should remain on the existing local-command paths. Tested: bun test src/QueryEngine.test.ts src/utils/processUserInput/processSlashCommand.test.ts src/server/__tests__/ws-memory-events.test.ts Tested: cd desktop && bun run test -- --run src/stores/chatStore.test.ts src/components/chat/MessageList.test.tsx Not-tested: bun run check:server is blocked by expired quarantine review entries: server:cron-scheduler, server:providers-real, server:tasks, server:e2e:business-flow, server:e2e:full-flow --- src/QueryEngine.test.ts | 125 ++++++++++++++++ src/QueryEngine.ts | 43 +++++- src/server/__tests__/ws-memory-events.test.ts | 39 +++++ src/server/ws/handler.ts | 14 ++ .../processSlashCommand.test.ts | 38 +++++ .../processUserInput/processSlashCommand.tsx | 140 ++++++++++++++++++ src/utils/sdkEventQueue.ts | 2 + 7 files changed, 394 insertions(+), 7 deletions(-) create mode 100644 src/QueryEngine.test.ts diff --git a/src/QueryEngine.test.ts b/src/QueryEngine.test.ts new file mode 100644 index 00000000..a06e3087 --- /dev/null +++ b/src/QueryEngine.test.ts @@ -0,0 +1,125 @@ +import { beforeEach, describe, expect, mock, test } from 'bun:test' +import { setIsInteractive } from './bootstrap/state.js' +import { drainSdkEvents, enqueueSdkEvent } from './utils/sdkEventQueue.js' + +process.env.ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY ?? 'test-key' + +type Deferred = { + promise: Promise + resolve: (value: T) => void +} + +function createDeferred(): Deferred { + let resolve!: (value: T) => void + const promise = new Promise(value => { + resolve = value + }) + return { promise, resolve } +} + +const processUserInputMock = mock(async () => ({ + messages: [], + shouldQuery: false, + resultText: 'done', +})) + +mock.module('./utils/processUserInput/processUserInput.js', () => ({ + processUserInput: processUserInputMock, +})) + +mock.module('./utils/queryContext.js', () => ({ + fetchSystemPromptParts: mock(async () => ({ + defaultSystemPrompt: [], + userContext: {}, + systemContext: {}, + })), +})) + +mock.module('./utils/messages/systemInit.js', () => ({ + buildSystemInitMessage: mock(() => ({ + type: 'system', + subtype: 'init', + session_id: 'test-session', + uuid: 'init-message', + model: 'test-model', + })), + sdkCompatToolName: (name: string) => name, +})) + +const { QueryEngine } = await import('./QueryEngine.js') + +function makeEngine() { + return new QueryEngine({ + cwd: process.cwd(), + tools: [], + commands: [], + mcpClients: [], + agents: [], + canUseTool: async () => ({ behavior: 'allow' }) as never, + getAppState: () => ({ + fastMode: false, + mcp: { clients: [], tools: [] }, + toolPermissionContext: { + mode: 'default', + alwaysAllowRules: { command: [] }, + additionalWorkingDirectories: new Map(), + }, + }) as never, + setAppState: () => {}, + readFileCache: new Map() as never, + }) +} + +function timeout(ms: number, value: T): Promise { + return new Promise(resolve => setTimeout(() => resolve(value), ms)) +} + +describe('QueryEngine slash command event streaming', () => { + beforeEach(() => { + processUserInputMock.mockClear() + drainSdkEvents() + setIsInteractive(false) + }) + + test('drains SDK task events while pre-query slash command processing is still pending', async () => { + const deferred = createDeferred<{ + messages: [] + shouldQuery: false + resultText: string + }>() + processUserInputMock.mockImplementationOnce(async () => { + enqueueSdkEvent({ + type: 'system', + subtype: 'task_started', + task_id: 'slash-agent-task', + description: 'Agent Explore', + task_type: 'slash_agent', + prompt: 'inspect recent changes', + }) + return deferred.promise + }) + + const iterator = makeEngine().submitMessage('/agent Explore inspect recent changes') + const first = await Promise.race([ + iterator.next(), + timeout(350, null), + ]) + + expect(first).not.toBeNull() + expect(first?.value).toMatchObject({ + type: 'system', + subtype: 'task_started', + task_id: 'slash-agent-task', + description: 'Agent Explore', + }) + + deferred.resolve({ + messages: [], + shouldQuery: false, + resultText: 'done', + }) + for await (const _message of iterator) { + // Drain the generator after releasing the pending slash command. + } + }) +}) diff --git a/src/QueryEngine.ts b/src/QueryEngine.ts index 2e6e7660..f65f6d2c 100644 --- a/src/QueryEngine.ts +++ b/src/QueryEngine.ts @@ -69,14 +69,17 @@ import { import { loadAllPluginsCacheOnly } from './utils/plugins/pluginLoader.js' import { type ProcessUserInputContext, + type ProcessUserInputBaseResult, processUserInput, } from './utils/processUserInput/processUserInput.js' import { fetchSystemPromptParts } from './utils/queryContext.js' import { setCwd } from './utils/Shell.js' +import { drainSdkEvents } from './utils/sdkEventQueue.js' import { flushSessionStorage, recordTranscript, } from './utils/sessionStorage.js' +import { sleep } from './utils/sleep.js' import { asSystemPrompt } from './utils/systemPromptType.js' import { resolveThemeSetting } from './utils/systemTheme.js' import { @@ -119,6 +122,8 @@ const getCoordinatorUserContext: ( : () => ({}) /* eslint-enable @typescript-eslint/no-require-imports */ +const PRE_QUERY_SDK_EVENT_DRAIN_INTERVAL_MS = 100 + // Dead code elimination: conditional import for snip compaction /* eslint-disable @typescript-eslint/no-require-imports */ const snipModule = feature('HISTORY_SNIP') @@ -409,13 +414,7 @@ export class QueryEngine { } } - const { - messages: messagesFromUserInput, - shouldQuery, - allowedTools, - model: modelFromUserInput, - resultText, - } = await processUserInput({ + const processUserInputPromise = processUserInput({ input: prompt, mode: 'prompt', setToolJSX: () => {}, @@ -428,6 +427,36 @@ export class QueryEngine { isMeta: options?.isMeta, querySource: 'sdk', }) + .then( + value => ({ ok: true as const, value }), + error => ({ ok: false as const, error }), + ) + + let processUserInputOutcome: Awaited | null = null + while (!processUserInputOutcome) { + const outcome = await Promise.race([ + processUserInputPromise, + sleep(PRE_QUERY_SDK_EVENT_DRAIN_INTERVAL_MS).then(() => null), + ]) + for (const event of drainSdkEvents()) { + yield event + } + if (outcome) { + processUserInputOutcome = outcome + } + } + + if (!processUserInputOutcome.ok) { + throw processUserInputOutcome.error + } + + const { + messages: messagesFromUserInput, + shouldQuery, + allowedTools, + model: modelFromUserInput, + resultText, + }: ProcessUserInputBaseResult = processUserInputOutcome.value // Push new messages, including user input and any attachments this.mutableMessages.push(...messagesFromUserInput) diff --git a/src/server/__tests__/ws-memory-events.test.ts b/src/server/__tests__/ws-memory-events.test.ts index 334a9449..5ca7d3b7 100644 --- a/src/server/__tests__/ws-memory-events.test.ts +++ b/src/server/__tests__/ws-memory-events.test.ts @@ -394,6 +394,45 @@ describe('WebSocket goal command events', () => { content: 'late unrelated output', })).toBe(false) }) + + it('allows current /agent task events through the pre-turn mute gate', () => { + const shouldForward = createCurrentTurnLocalCommandForwarder( + parseSlashCommand('/agent Explore inspect recent changes'), + ) + + expect(shouldForward({ + type: 'system', + subtype: 'task_started', + task_type: 'slash_agent', + task_id: 'slash-agent-1', + description: 'Agent Explore', + })).toBe(true) + expect(shouldForward({ + type: 'system', + subtype: 'task_progress', + task_type: 'slash_agent', + task_id: 'slash-agent-1', + summary: 'Using Read', + })).toBe(true) + expect(shouldForward({ + type: 'system', + subtype: 'task_notification', + task_type: 'slash_agent', + task_id: 'slash-agent-1', + status: 'completed', + })).toBe(true) + + const unrelatedCommandForwarder = createCurrentTurnLocalCommandForwarder( + parseSlashCommand('/goal ship the smoke test'), + ) + expect(unrelatedCommandForwarder({ + type: 'system', + subtype: 'task_started', + task_type: 'slash_agent', + task_id: 'slash-agent-1', + description: 'Agent Explore', + })).toBe(false) + }) }) describe('WebSocket stream event translation', () => { diff --git a/src/server/ws/handler.ts b/src/server/ws/handler.ts index 02bc6760..005719c7 100644 --- a/src/server/ws/handler.ts +++ b/src/server/ws/handler.ts @@ -1809,6 +1809,9 @@ export function createCurrentTurnLocalCommandForwarder( let awaitingCurrentTurnLocalCommandOutput = false return (cliMsg: any) => { + if (command?.commandName === 'agent' && isSlashAgentTaskEvent(cliMsg)) { + return true + } if (command && isMatchingCurrentTurnLocalCommand(cliMsg, command)) { awaitingCurrentTurnLocalCommandOutput = true return true @@ -1834,6 +1837,17 @@ export function createCurrentTurnLocalCommandForwarder( } } +function isSlashAgentTaskEvent(cliMsg: any): boolean { + if (cliMsg?.type !== 'system' || cliMsg?.task_type !== 'slash_agent') { + return false + } + return ( + cliMsg.subtype === 'task_started' || + cliMsg.subtype === 'task_progress' || + cliMsg.subtype === 'task_notification' + ) +} + function isMatchingCurrentTurnLocalCommand( cliMsg: any, command: NonNullable>, diff --git a/src/utils/processUserInput/processSlashCommand.test.ts b/src/utils/processUserInput/processSlashCommand.test.ts index 607a605a..9a91828c 100644 --- a/src/utils/processUserInput/processSlashCommand.test.ts +++ b/src/utils/processUserInput/processSlashCommand.test.ts @@ -1,8 +1,10 @@ import { beforeEach, describe, expect, mock, test } from 'bun:test' +import { setIsInteractive } from '../../bootstrap/state.js' import agentCommand from '../../commands/agent.js' import type { ToolUseContext } from '../../Tool.js' import type { AgentDefinition } from '../../tools/AgentTool/loadAgentsDir.js' import { createAssistantMessage } from '../messages.js' +import { drainSdkEvents } from '../sdkEventQueue.js' process.env.ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY ?? 'test-key' @@ -49,6 +51,8 @@ function makeContext(activeAgents: AgentDefinition[]): ToolUseContext { describe('/agent slash command processing', () => { beforeEach(() => { runAgentMock.mockClear() + drainSdkEvents() + setIsInteractive(true) }) test('runs the selected agent with only the prompt body', async () => { @@ -98,4 +102,38 @@ describe('/agent slash command processing', () => { ), ).toBe(true) }) + + test('emits foreground agent task events for desktop streaming', async () => { + setIsInteractive(false) + + await processSlashCommand( + '/agent debugger fix failing tests', + [], + [], + [], + makeContext([makeAgent('general-purpose'), makeAgent('debugger')]), + () => {}, + ) + + const events = drainSdkEvents() + expect(events.map(event => event.subtype)).toEqual([ + 'task_started', + 'task_progress', + 'task_notification', + ]) + expect(events[0]).toMatchObject({ + type: 'system', + subtype: 'task_started', + description: 'Agent debugger', + task_type: 'slash_agent', + prompt: 'fix failing tests', + }) + expect(events[2]).toMatchObject({ + type: 'system', + subtype: 'task_notification', + status: 'completed', + summary: 'Agent debugger completed', + result: 'debugger result', + }) + }) }) diff --git a/src/utils/processUserInput/processSlashCommand.tsx b/src/utils/processUserInput/processSlashCommand.tsx index 985940fd..eadaebb5 100644 --- a/src/utils/processUserInput/processSlashCommand.tsx +++ b/src/utils/processUserInput/processSlashCommand.tsx @@ -39,6 +39,7 @@ import { hasPermissionsToUseTool } from '../permissions/permissions.js'; import { isOfficialMarketplaceName, parsePluginIdentifier } from '../plugins/pluginIdentifier.js'; import { isRestrictedToPluginOnly, isSourceAdminTrusted } from '../settings/pluginOnlyPolicy.js'; import { parseSlashCommand } from '../slashCommandParsing.js'; +import { enqueueSdkEvent } from '../sdkEventQueue.js'; import { sleep } from '../sleep.js'; import { recordSkillUsage } from '../suggestions/skillUsageTracking.js'; import { logOTelEvent, redactIfDisabled } from '../telemetry/events.js'; @@ -197,6 +198,26 @@ async function executeForkedSlashCommand(command: CommandBase & PromptCommand, a const progressMessages: ProgressMessage[] = []; const parentToolUseID = `forked-command-${command.name}`; let toolUseCounter = 0; + const shouldEmitAgentTaskEvents = command.name === 'agent' && dynamicAgentCommand !== null; + const sdkTaskId = shouldEmitAgentTaskEvents ? `slash-agent-${agentId}` : null; + const sdkToolUseId = shouldEmitAgentTaskEvents ? `${parentToolUseID}-${agentId}` : undefined; + const sdkDescription = dynamicAgentCommand ? `Agent ${dynamicAgentCommand.agentType}` : `/${getCommandName(command)}`; + const sdkPrompt = dynamicAgentCommand?.prompt ?? args; + const sdkStartTime = Date.now(); + let sdkProgressTokens = 0; + let sdkToolUses = 0; + + if (sdkTaskId) { + enqueueSdkEvent({ + type: 'system', + subtype: 'task_started', + task_id: sdkTaskId, + tool_use_id: sdkToolUseId, + description: sdkDescription, + task_type: 'slash_agent', + prompt: sdkPrompt + }); + } // Helper to create a progress message from an agent message const createProgressMessage = (message: AssistantMessage | NormalizedUserMessage): ProgressMessage => { @@ -229,6 +250,28 @@ async function executeForkedSlashCommand(command: CommandBase & PromptCommand, a }); }; + const emitAgentTaskProgress = (message: AssistantMessage | NormalizedUserMessage): void => { + if (!sdkTaskId) return; + const summary = summarizeAgentProgressMessage(message); + sdkProgressTokens += getApproximateMessageTokenCount(message); + sdkToolUses += countAgentToolUses(message); + enqueueSdkEvent({ + type: 'system', + subtype: 'task_progress', + task_id: sdkTaskId, + tool_use_id: sdkToolUseId, + task_type: 'slash_agent', + description: sdkDescription, + summary: summary.summary, + last_tool_name: summary.lastToolName, + usage: { + total_tokens: sdkProgressTokens, + tool_uses: sdkToolUses, + duration_ms: Date.now() - sdkStartTime + } + }); + }; + // Show initial "Initializing…" state updateProgress(); @@ -260,6 +303,7 @@ async function executeForkedSlashCommand(command: CommandBase & PromptCommand, a const normalizedMsg = normalizedNew[0]; if (normalizedMsg && normalizedMsg.type === 'assistant') { progressMessages.push(createProgressMessage(message)); + emitAgentTaskProgress(message); updateProgress(); } } @@ -269,10 +313,31 @@ async function executeForkedSlashCommand(command: CommandBase & PromptCommand, a const normalizedMsg = normalizedNew[0]; if (normalizedMsg && normalizedMsg.type === 'user') { progressMessages.push(createProgressMessage(normalizedMsg)); + emitAgentTaskProgress(normalizedMsg); updateProgress(); } } } + } catch (err) { + if (sdkTaskId) { + enqueueSdkEvent({ + type: 'system', + subtype: 'task_notification', + task_id: sdkTaskId, + tool_use_id: sdkToolUseId, + task_type: 'slash_agent', + status: 'failed', + output_file: '', + summary: `${sdkDescription} failed`, + result: err instanceof Error ? err.message : String(err), + usage: { + total_tokens: sdkProgressTokens, + tool_uses: sdkToolUses, + duration_ms: Date.now() - sdkStartTime + } + }); + } + throw err; } finally { // Clear the progress display setToolJSX(null); @@ -285,6 +350,25 @@ async function executeForkedSlashCommand(command: CommandBase & PromptCommand, a resultText = `[ANT-ONLY] API calls: ${getDisplayPath(getDumpPromptsPath(agentId))}\n${resultText}`; } + if (sdkTaskId) { + enqueueSdkEvent({ + type: 'system', + subtype: 'task_notification', + task_id: sdkTaskId, + tool_use_id: sdkToolUseId, + task_type: 'slash_agent', + status: 'completed', + output_file: '', + summary: `${sdkDescription} completed`, + result: resultText, + usage: { + total_tokens: sdkProgressTokens, + tool_uses: sdkToolUses, + duration_ms: Date.now() - sdkStartTime + } + }); + } + // Return the result as a user message (simulates the agent's output) const messages: UserMessage[] = [createUserMessage({ content: prepareUserContent({ @@ -314,6 +398,62 @@ export function looksLikeCommand(commandName: string): boolean { // If it contains other characters, it's probably a file path or other input return !/[^a-zA-Z0-9:\-_]/.test(commandName); } + +function summarizeAgentProgressMessage(message: AssistantMessage | NormalizedUserMessage): { summary: string; lastToolName?: string } { + const content = message.message.content; + const blocks = Array.isArray(content) ? content : [{ + type: 'text', + text: typeof content === 'string' ? content : '' + }]; + let firstText: string | null = null; + let lastToolName: string | undefined; + + for (const block of blocks) { + if (!block || typeof block !== 'object') continue; + const typedBlock = block as { type?: unknown; text?: unknown; name?: unknown }; + if (typedBlock.type === 'tool_use' && typeof typedBlock.name === 'string') { + lastToolName = typedBlock.name; + } + if (!firstText && typedBlock.type === 'text' && typeof typedBlock.text === 'string') { + firstText = typedBlock.text.trim(); + } + } + + if (lastToolName) return { summary: `Using ${lastToolName}`, lastToolName }; + if (firstText) return { summary: truncateProgressSummary(firstText) }; + return { summary: 'Agent is working' }; +} + +function getApproximateMessageTokenCount(message: AssistantMessage | NormalizedUserMessage): number { + const content = message.message.content; + if (typeof content === 'string') return Math.ceil(content.length / 4); + if (!Array.isArray(content)) return 0; + const textLength = content.reduce((total, block) => { + if (!block || typeof block !== 'object') return total; + const text = (block as { text?: unknown; content?: unknown }).text; + if (typeof text === 'string') return total + text.length; + const nestedContent = (block as { content?: unknown }).content; + if (typeof nestedContent === 'string') return total + nestedContent.length; + return total; + }, 0); + return Math.ceil(textLength / 4); +} + +function countAgentToolUses(message: AssistantMessage | NormalizedUserMessage): number { + const content = message.message.content; + if (!Array.isArray(content)) return 0; + return content.filter(block => { + if (!block || typeof block !== 'object') return false; + return (block as { type?: unknown }).type === 'tool_use'; + }).length; +} + +function truncateProgressSummary(text: string): string { + const normalized = text.replace(/\s+/g, ' ').trim(); + if (normalized.length <= 120) return normalized; + return `${normalized.slice(0, 117)}...`; +} + export async function processSlashCommand(inputString: string, precedingInputBlocks: ContentBlockParam[], imageContentBlocks: ContentBlockParam[], attachmentMessages: AttachmentMessage[], context: ProcessUserInputContext, setToolJSX: SetToolJSXFn, uuid?: string, isAlreadyProcessing?: boolean, canUseTool?: CanUseToolFn): Promise { const parsed = parseSlashCommand(inputString); if (!parsed) { diff --git a/src/utils/sdkEventQueue.ts b/src/utils/sdkEventQueue.ts index abdb112f..05975c74 100644 --- a/src/utils/sdkEventQueue.ts +++ b/src/utils/sdkEventQueue.ts @@ -20,6 +20,7 @@ type TaskProgressEvent = { task_id: string tool_use_id?: string description: string + task_type?: string usage: { total_tokens: number tool_uses: number @@ -43,6 +44,7 @@ type TaskNotificationSdkEvent = { subtype: 'task_notification' task_id: string tool_use_id?: string + task_type?: string status: 'completed' | 'failed' | 'stopped' output_file: string summary: string