From 6f54dcc62b912415ae22b134a7414301e4ba2aa5 Mon Sep 17 00:00:00 2001 From: digua Date: Mon, 8 Jun 2026 21:38:48 +0800 Subject: [PATCH] feat(cli): expose full agent stream events --- apps/cli/src/ai/agent.ts | 1 + apps/cli/src/ai/chat-command.test.ts | 172 +++++++++++++++++- apps/cli/src/ai/chat-command.ts | 82 ++++++++- apps/cli/src/cli.ts | 2 + apps/desktop/main/ai/agent/index.ts | 1 + .../src/ai/agent/event-handler.ts | 14 +- src/services/ai-stream/types.ts | 28 ++- 7 files changed, 290 insertions(+), 10 deletions(-) diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index f7c7cdcd..0e47398d 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -191,6 +191,7 @@ export async function runServerAgent(options: RunAgentOptions): Promise { availableToolCount: tools.length, shadowOnly: true, }) + onEvent({ type: 'route', routeDecision }) let effectiveSystemPrompt = systemPrompt if (routeDecision.route === 'planned_execution') { diff --git a/apps/cli/src/ai/chat-command.test.ts b/apps/cli/src/ai/chat-command.test.ts index c2b91a1c..a41a8fd7 100644 --- a/apps/cli/src/ai/chat-command.test.ts +++ b/apps/cli/src/ai/chat-command.test.ts @@ -2,7 +2,7 @@ import { describe, it } from 'node:test' import assert from 'node:assert/strict' import { Readable } from 'node:stream' import { Writable } from 'node:stream' -import type { AIChatManager, DatabaseManager } from '@openchatlab/node-runtime' +import type { AIChatManager, ContentBlock, DatabaseManager, TokenUsageData } from '@openchatlab/node-runtime' import { resolveAIChatTarget, runChatCommand, runChatTurn } from './chat-command' function createDbManager(sessionIds: string[]): DatabaseManager { @@ -18,7 +18,13 @@ function createAIChatManager( const chats = new Map( existing.map((chat) => [chat.id, { ...chat, title: null, assistantId: chat.assistantId ?? 'general_cn' }]) ) - const messages: Array<{ aiChatId: string; role: string; content: string }> = [] + const messages: Array<{ + aiChatId: string + role: string + content: string + contentBlocks?: ContentBlock[] + tokenUsage?: TokenUsageData + }> = [] return { getAIChat: (aiChatId: string) => chats.get(aiChatId) ?? null, @@ -28,8 +34,22 @@ function createAIChatManager( chats.set(id, chat) return chat }, - addMessage: (aiChatId: string, role: string, content: string) => { - messages.push({ aiChatId, role, content }) + addMessage: ( + aiChatId: string, + role: string, + content: string, + _dataKeywords?: string[], + _dataMessageCount?: number, + contentBlocks?: ContentBlock[], + tokenUsage?: TokenUsageData + ) => { + messages.push({ + aiChatId, + role, + content, + ...(contentBlocks ? { contentBlocks } : {}), + ...(tokenUsage ? { tokenUsage } : {}), + }) return { id: `msg_${messages.length}`, aiChatId, role, content, timestamp: 1 } }, __messages: messages, @@ -150,7 +170,149 @@ describe('runChatTurn', () => { assert.equal(stdout.text(), '') assert.deepEqual((aiChatManager as unknown as { __messages: unknown[] }).__messages, [ { aiChatId: 'ai_chat_1', role: 'user', content: 'hello' }, - { aiChatId: 'ai_chat_1', role: 'assistant', content: 'hi' }, + { + aiChatId: 'ai_chat_1', + role: 'assistant', + content: 'hi', + tokenUsage: { promptTokens: 1, completionTokens: 1, totalTokens: 2 }, + }, + ]) + }) + + it('can include all agent events and persist plan content blocks', async () => { + const stdout = new MemoryWritable() + const aiChatManager = createAIChatManager() + + const result = await runChatTurn( + { sessionId: 'session-1', question: '分析过去一年话题趋势', json: true, includeEvents: true }, + { + dbManager: createDbManager(['session-1']), + pathProvider: {} as never, + aiChatManager, + stdout, + createRunAgentStream: () => async (_params, onEvent) => { + onEvent({ + type: 'route', + routeDecision: { + route: 'planned_execution', + confidence: 0.91, + reason: 'Complex long-range trend analysis.', + source: 'rule', + }, + }) + onEvent({ + type: 'plan', + plan: { + type: 'plan', + version: 1, + status: 'created', + plan: { + version: 1, + title: '年度话题趋势', + route: 'planned_execution', + intent: 'trend', + steps: [{ goal: '按季度检索', suggestedTools: ['search_messages'], evidenceNeeded: '季度证据' }], + successCriteria: ['覆盖至少三个季度'], + }, + }, + }) + onEvent({ type: 'content', content: '年度趋势如下。' }) + onEvent({ + type: 'done', + isFinished: true, + usage: { promptTokens: 3, completionTokens: 5, totalTokens: 8 }, + }) + }, + } + ) + + assert.equal(result.answer, '年度趋势如下。') + assert.equal(result.events?.length, 4) + assert.equal(result.events?.[0]?.type, 'route') + assert.equal(result.events?.[0]?.routeDecision?.route, 'planned_execution') + assert.equal(result.events?.[1]?.type, 'plan') + assert.equal(result.contentBlocks?.[0]?.type, 'plan') + assert.equal(result.contentBlocks?.[0]?.status, 'done') + assert.deepEqual((aiChatManager as unknown as { __messages: unknown[] }).__messages, [ + { aiChatId: 'ai_chat_1', role: 'user', content: '分析过去一年话题趋势' }, + { + aiChatId: 'ai_chat_1', + role: 'assistant', + content: '年度趋势如下。', + contentBlocks: result.contentBlocks, + tokenUsage: { promptTokens: 3, completionTokens: 5, totalTokens: 8 }, + }, + ]) + }) + + it('persists streamed thinking and answer text blocks for full CLI replay', async () => { + const stdout = new MemoryWritable() + const aiChatManager = createAIChatManager() + + const result = await runChatTurn( + { sessionId: 'session-1', question: '分析过去一年话题趋势', json: true, includeEvents: true }, + { + dbManager: createDbManager(['session-1']), + pathProvider: {} as never, + aiChatManager, + stdout, + createRunAgentStream: () => async (_params, onEvent) => { + onEvent({ + type: 'plan', + plan: { + type: 'plan', + version: 1, + status: 'created', + plan: { + version: 1, + title: '年度话题趋势', + route: 'planned_execution', + intent: 'trend', + steps: [{ goal: '按季度检索', suggestedTools: ['search_messages'], evidenceNeeded: '季度证据' }], + successCriteria: ['覆盖至少三个季度'], + }, + }, + }) + onEvent({ type: 'think', thinkTag: 'thinking', content: '先理解问题，' }) + onEvent({ type: 'think', thinkTag: 'thinking', content: '再整理证据。' }) + onEvent({ type: 'think', thinkTag: 'thinking', content: '', thinkDurationMs: 1200 }) + onEvent({ type: 'content', content: '年度趋势如下。' }) + onEvent({ + type: 'done', + isFinished: true, + usage: { promptTokens: 3, completionTokens: 5, totalTokens: 8 }, + }) + }, + } + ) + + assert.equal(result.answer, '年度趋势如下。') + assert.equal(result.events?.filter((event) => event.type === 'think').length, 3) + assert.deepEqual( + result.contentBlocks?.map((block) => block.type), + ['plan', 'think', 'text'] + ) + assert.equal(result.contentBlocks?.[0]?.type, 'plan') + assert.equal(result.contentBlocks?.[0]?.status, 'done') + assert.deepEqual(result.contentBlocks?.[1], { + type: 'think', + tag: 'thinking', + text: '先理解问题，再整理证据。', + durationMs: 1200, + }) + assert.deepEqual(result.contentBlocks?.[2], { + type: 'text', + text: '年度趋势如下。', + }) + assert.deepEqual((aiChatManager as unknown as { __messages: unknown[] }).__messages, [ + { aiChatId: 'ai_chat_1', role: 'user', content: '分析过去一年话题趋势' }, + { + aiChatId: 'ai_chat_1', + role: 'assistant', + content: '年度趋势如下。', + contentBlocks: result.contentBlocks, + tokenUsage: { promptTokens: 3, completionTokens: 5, totalTokens: 8 }, + }, ]) }) diff --git a/apps/cli/src/ai/chat-command.ts b/apps/cli/src/ai/chat-command.ts index 81e02525..10ed1c0f 100644 --- a/apps/cli/src/ai/chat-command.ts +++ b/apps/cli/src/ai/chat-command.ts @@ -1,6 +1,13 @@ import { createInterface } from 'node:readline/promises' import type { Readable, Writable } from 'node:stream' -import type { DatabaseManager, AIChatManager, AgentStreamChunk, TokenUsageData } from '@openchatlab/node-runtime' +import type { + DatabaseManager, + AIChatManager, + AgentStreamChunk, + ContentBlock, + PlanContentBlock, + TokenUsageData, +} from '@openchatlab/node-runtime' import type { PathProvider } from '@openchatlab/core' import { createCliRunAgentStream } from './agent-stream-runner' @@ -11,6 +18,7 @@ export interface ChatCommandOptions { json?: boolean stream?: boolean locale?: string + includeEvents?: boolean } export interface ChatCommandDeps { @@ -35,6 +43,8 @@ export interface ChatTurnResult { aiChatId: string question: string answer: string + events?: AgentStreamChunk[] + contentBlocks?: ContentBlock[] usage: { durationMs: number tokenUsage: TokenUsageData | null @@ -106,7 +116,7 @@ export function resolveAIChatTarget( export async function runChatTurn( options: Required> & - Pick, + Pick, deps: ChatCommandDeps ): Promise { const stdout = deps.stdout ?? process.stdout @@ -115,9 +125,53 @@ export async function runChatTurn( let answer = '' let tokenUsage: TokenUsageData | null = null let streamError: Error | null = null + const events: AgentStreamChunk[] = [] + const contentBlocks: ContentBlock[] = [] + let hasReplayContentBlocks = false const runAgentStream = (deps.createRunAgentStream ?? createCliRunAgentStream)(deps.dbManager, deps.aiChatManager) + // 中文注释：CLI 历史回放依赖 contentBlocks 的时序；只要出现计划/思考块， + // 就同步保留后续正文 text block，避免 UI 使用 blocks 渲染时丢失最终回答。 + const appendTextBlock = (text: string) => { + if (!text) return + const lastBlock = contentBlocks[contentBlocks.length - 1] + if (lastBlock?.type === 'text') { + lastBlock.text += text + } else { + contentBlocks.push({ type: 'text', text }) + } + } + + const appendThinkBlock = (text: string, tag = 'thinking', durationMs?: number) => { + if (!text && durationMs === undefined) return + const lastBlock = contentBlocks[contentBlocks.length - 1] + let targetBlock: ContentBlock | undefined + + if (lastBlock?.type === 'think' && lastBlock.tag === tag) { + lastBlock.text += text + targetBlock = lastBlock + } else if (text.trim().length > 0) { + targetBlock = { type: 'think', tag, text } + contentBlocks.push(targetBlock) + } else if (durationMs !== undefined) { + for (let index = contentBlocks.length - 1; index >= 0; index--) { + const block = contentBlocks[index] + if (block.type === 'think' && block.tag === tag) { + targetBlock = block + break + } + } + } + + if (durationMs !== undefined && targetBlock?.type === 'think') { + targetBlock.durationMs = durationMs + } + if (targetBlock?.type === 'think') { + hasReplayContentBlocks = true + } + } + await runAgentStream( { userMessage: options.question, @@ -128,16 +182,36 @@ export async function runChatTurn( locale: options.locale ?? 'zh-CN', }, (chunk: AgentStreamChunk) => { + if (options.includeEvents) { + events.push(JSON.parse(JSON.stringify(chunk)) as AgentStreamChunk) + } if (chunk.type === 'error') { streamError = createAgentStreamError(chunk.error) return } + if (chunk.type === 'plan' && chunk.plan) { + hasReplayContentBlocks = true + contentBlocks.push(JSON.parse(JSON.stringify(chunk.plan)) as PlanContentBlock) + return + } + if (chunk.type === 'think') { + appendThinkBlock(chunk.content ?? '', chunk.thinkTag, chunk.thinkDurationMs) + return + } if (chunk.type === 'content' && chunk.content) { answer += chunk.content + appendTextBlock(chunk.content) if (!options.json && options.stream !== false) write(stdout, chunk.content) } if (chunk.type === 'done' && chunk.usage) { tokenUsage = chunk.usage + for (let index = contentBlocks.length - 1; index >= 0; index--) { + const block = contentBlocks[index] + if (block.type === 'plan') { + block.status = 'done' + break + } + } } }, new AbortController().signal @@ -156,7 +230,7 @@ export async function runChatTurn( answer, undefined, undefined, - undefined, + hasReplayContentBlocks ? contentBlocks : undefined, tokenUsage ?? undefined ) @@ -165,6 +239,8 @@ export async function runChatTurn( aiChatId: target.aiChatId, question: options.question, answer, + ...(options.includeEvents ? { events } : {}), + ...(hasReplayContentBlocks ? { contentBlocks } : {}), usage: { durationMs: Date.now() - startedAt, tokenUsage, diff --git a/apps/cli/src/cli.ts b/apps/cli/src/cli.ts index 79ced883..d7b81b78 100644 --- a/apps/cli/src/cli.ts +++ b/apps/cli/src/cli.ts @@ -308,6 +308,7 @@ program .option('--ai-chat-id ', 'Existing AI chat ID to continue') .option('-q, --question ', 'Question to ask') .option('--json', 'Output structured JSON') + .option('--include-events', 'Include all agent stream chunks in JSON output') .option('--no-stream', 'Disable streaming output') .option('--locale ', 'AI response locale', 'zh-CN') .action(async (options) => { @@ -324,6 +325,7 @@ program json: !!options.json, stream: options.stream, locale: options.locale, + includeEvents: !!options.includeEvents, }, { dbManager, pathProvider, aiChatManager } ) diff --git a/apps/desktop/main/ai/agent/index.ts b/apps/desktop/main/ai/agent/index.ts index 7754ef29..b4f67948 100644 --- a/apps/desktop/main/ai/agent/index.ts +++ b/apps/desktop/main/ai/agent/index.ts @@ -156,6 +156,7 @@ export class Agent { availableToolCount: piTools.length, shadowOnly: true, }) + onChunk({ type: 'route', routeDecision }) let effectiveSystemPrompt = systemPrompt if (routeDecision.route === 'planned_execution') { diff --git a/packages/node-runtime/src/ai/agent/event-handler.ts b/packages/node-runtime/src/ai/agent/event-handler.ts index f712b43d..d887acd8 100644 --- a/packages/node-runtime/src/ai/agent/event-handler.ts +++ b/packages/node-runtime/src/ai/agent/event-handler.ts @@ -8,6 +8,7 @@ import type { Message as PiMessage } from '@earendil-works/pi-ai' import type { AgentCoreEvent } from './types' import type { PlanContentBlock } from './planning-types' +import type { RouteDecision } from './routing-types' // ==================== Shared types ==================== @@ -28,7 +29,17 @@ export interface AgentRuntimeStatus { } export interface AgentStreamChunk { - type: 'content' | 'think' | 'tool_start' | 'tool_result' | 'status' | 'compression_done' | 'plan' | 'done' | 'error' + type: + | 'content' + | 'think' + | 'tool_start' + | 'tool_result' + | 'status' + | 'compression_done' + | 'route' + | 'plan' + | 'done' + | 'error' content?: string thinkTag?: string thinkDurationMs?: number @@ -39,6 +50,7 @@ export interface AgentStreamChunk { isFinished?: boolean usage?: TokenUsage status?: AgentRuntimeStatus + routeDecision?: RouteDecision plan?: PlanContentBlock compressionResult?: { summaryContent: string diff --git a/src/services/ai-stream/types.ts b/src/services/ai-stream/types.ts index d0078256..41dc56f8 100644 --- a/src/services/ai-stream/types.ts +++ b/src/services/ai-stream/types.ts @@ -10,7 +10,17 @@ export interface LlmStreamChunk { } export interface AgentStreamChunk { - type: 'content' | 'think' | 'tool_start' | 'tool_result' | 'status' | 'compression_done' | 'plan' | 'done' | 'error' + type: + | 'content' + | 'think' + | 'tool_start' + | 'tool_result' + | 'status' + | 'compression_done' + | 'route' + | 'plan' + | 'done' + | 'error' content?: string thinkTag?: string thinkDurationMs?: number @@ -21,6 +31,7 @@ export interface AgentStreamChunk { isFinished?: boolean usage?: TokenUsage status?: AgentRuntimeStatus + routeDecision?: RouteDecision plan?: PlanContentBlock compressionResult?: { summaryContent: string @@ -30,6 +41,21 @@ export interface AgentStreamChunk { } } +export type RequestRoute = 'direct_response' | 'tool_assisted' | 'planned_execution' +export type RouteDecisionSource = 'rule' | 'llm' + +export interface RouteDecision { + route: RequestRoute + confidence: number + reason: string + source: RouteDecisionSource + usage?: { + promptTokens?: number + completionTokens?: number + totalTokens?: number + } +} + export interface TokenUsage { promptTokens: number completionTokens: number