From 9a0e6093f33da34f716f5f31b019ae4f00ac9e63 Mon Sep 17 00:00:00 2001 From: digua Date: Mon, 1 Jun 2026 21:45:45 +0800 Subject: [PATCH] fix(ai): harden shared SSE streaming --- apps/cli/src/ai/agent-stream-runner.ts | 24 ++++++++++++----- apps/cli/src/ai/agent.ts | 6 ++++- apps/desktop/main/ai/agent-stream-runner.ts | 25 +++--------------- apps/desktop/main/internal-api.ts | 20 ++++++++++---- .../src/routes/web/ai-agent-stream.ts | 26 ++++++++++++------- .../src/routes/web/ai-llm-stream.ts | 12 +++++++-- src/services/ai-stream/service.ts | 16 +++++++----- 7 files changed, 76 insertions(+), 53 deletions(-) diff --git a/apps/cli/src/ai/agent-stream-runner.ts b/apps/cli/src/ai/agent-stream-runner.ts index d19f3c8d..78aec675 100644 --- a/apps/cli/src/ai/agent-stream-runner.ts +++ b/apps/cli/src/ai/agent-stream-runner.ts @@ -35,6 +35,8 @@ export function createCliRunAgentStream( chatType, locale, assistantId, + skillId, + enableAutoSkill, compressionConfig, ownerInfo, mentionedMembers, @@ -61,17 +63,26 @@ export function createCliRunAgentStream( ? adaptToolsForAgent(AGENT_TOOL_REGISTRY, () => ({ db, sessionId, locale }), { maxToolResultTokens }) : [] - const skillManager = new SkillManager(aiDataDir) - skillManager.init() + const skillMgr = new SkillManager(aiDataDir) + skillMgr.init() const toolNames = agentTools.map((t: { name: string }) => t.name) - const skillMenu = skillManager.getSkillMenu(chatType ?? 'group', toolNames) - if (skillMenu) { + let resolvedSkillDef: { name: string; prompt: string } | undefined + let resolvedSkillMenu: string | undefined + if (skillId) { + const def = skillMgr.getSkillConfig(skillId) + if (def) resolvedSkillDef = { name: def.name, prompt: def.prompt } + } else if (enableAutoSkill) { + const menu = skillMgr.getSkillMenu(chatType ?? 'group', toolNames) + if (menu) resolvedSkillMenu = menu + } + + if (resolvedSkillMenu) { const activateSkillTool = createActivateSkillTool({ chatType: chatType ?? 'group', allowedTools: toolNames, locale, - getSkillConfig: (id) => skillManager.getSkillConfig(id), + getSkillConfig: (id) => skillMgr.getSkillConfig(id), }) agentTools.push(activateSkillTool as any) } @@ -113,7 +124,8 @@ export function createCliRunAgentStream( chatType, locale, assistantSystemPrompt, - skillMenu, + skillMenu: resolvedSkillMenu, + skillDef: resolvedSkillDef, compressionConfig: resolvedCompression, tools: agentTools, aiDataDir, diff --git a/apps/cli/src/ai/agent.ts b/apps/cli/src/ai/agent.ts index 61980ed0..b1edfab5 100644 --- a/apps/cli/src/ai/agent.ts +++ b/apps/cli/src/ai/agent.ts @@ -37,6 +37,7 @@ export interface RunAgentOptions { locale?: string assistantSystemPrompt?: string skillMenu?: string | null + skillDef?: { name: string; prompt: string } compressionConfig?: CompressionConfig tools?: AgentTool[] aiDataDir: string @@ -58,6 +59,7 @@ export async function runServerAgent(options: RunAgentOptions): Promise { locale = 'zh-CN', assistantSystemPrompt, skillMenu, + skillDef, compressionConfig, tools = [], aiDataDir, @@ -83,7 +85,9 @@ export async function runServerAgent(options: RunAgentOptions): Promise { const t = createAiTranslate(locale) let skillCtx: { skillDef?: { name: string; prompt: string }; skillMenu?: string } | undefined - if (skillMenu) { + if (skillDef) { + skillCtx = { skillDef } + } else if (skillMenu) { skillCtx = { skillMenu } } diff --git a/apps/desktop/main/ai/agent-stream-runner.ts b/apps/desktop/main/ai/agent-stream-runner.ts index 918bd536..f825c0f1 100644 --- a/apps/desktop/main/ai/agent-stream-runner.ts +++ b/apps/desktop/main/ai/agent-stream-runner.ts @@ -86,7 +86,6 @@ export function createElectronRunAgentStream(): ( const activeAIConfig = getDefaultAssistantConfig() if (!activeAIConfig) { onEvent({ type: 'error', error: { name: 'ConfigError', message: t('llm.notConfigured') } }) - onEvent({ type: 'done', isFinished: true }) return } const piModel = buildPiModel(activeAIConfig) @@ -211,38 +210,20 @@ export function createElectronRunAgentStream(): ( ) try { - const result = await agent.executeStream(userMessage, (chunk: AgentStreamChunk) => { + await agent.executeStream(userMessage, (chunk: AgentStreamChunk) => { if (abortSignal.aborted) return onEvent(chunk as SharedAgentStreamChunk) }) - - if (abortSignal.aborted) { - onEvent({ - type: 'done', - isFinished: true, - usage: result.totalUsage, - }) - return - } - - onEvent({ - type: 'done', - isFinished: true, - usage: result.totalUsage, - }) } catch (error) { - if (error instanceof Error && error.name === 'AbortError') { - onEvent({ type: 'done', isFinished: true }) - return - } + if (error instanceof Error && error.name === 'AbortError') return const serializedError = serializeError(error, activeAIConfig.provider) serializedError.friendlyMessage = formatAIError(error, { providerName: resolveProviderName(activeAIConfig.provider), rawErrorLabel: t('llm.rawErrorLabel'), }) + if (!serializedError.url && activeAIConfig.baseUrl) serializedError.url = activeAIConfig.baseUrl aiLogger.error('AgentStream', `Agent execution error: ${requestId}`, serializedError) onEvent({ type: 'error', error: serializedError, isFinished: true }) - onEvent({ type: 'done', isFinished: true }) } } } diff --git a/apps/desktop/main/internal-api.ts b/apps/desktop/main/internal-api.ts index 185d9317..07bc10bf 100644 --- a/apps/desktop/main/internal-api.ts +++ b/apps/desktop/main/internal-api.ts @@ -162,6 +162,10 @@ export async function startInternalServer(pathProvider: PathProvider): Promise { + reply.header(name, value) + reply.raw.setHeader(name, value) + } newServer.addHook('onRequest', (request, reply, done) => { const origin = request.headers.origin if (!origin) { @@ -170,15 +174,21 @@ export async function startInternalServer(pathProvider: PathProvider): Promise { + let emittedDone = false + + const safeSendSSE = (event: string, data: unknown) => { + if (reply.raw.writableEnded || reply.raw.destroyed) return reply.raw.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`) } - sendSSE('meta', { requestId }) + const safeEnd = () => { + if (reply.raw.writableEnded || reply.raw.destroyed) return + reply.raw.end() + } + + safeSendSSE('meta', { requestId }) reply.raw.on('close', () => { if (!abortController.signal.aborted) { @@ -37,20 +45,18 @@ export function registerAiAgentStreamRoutes(server: FastifyInstance, ctx: HttpRo await runAgentStream( request.body, (chunk) => { - sendSSE(chunk.type, chunk) - if (chunk.type === 'done') { - activeAgentAborts.delete(requestId) - reply.raw.end() - } + if (chunk.type === 'done') emittedDone = true + safeSendSSE(chunk.type, chunk) }, abortController.signal ) } catch (error) { const msg = error instanceof Error ? error.message : String(error) - sendSSE('error', { type: 'error', error: { name: 'ServerError', message: msg } }) - sendSSE('done', { type: 'done', isFinished: true }) + safeSendSSE('error', { type: 'error', error: { name: 'ServerError', message: msg } }) + } finally { + if (!emittedDone) safeSendSSE('done', { type: 'done', isFinished: true }) activeAgentAborts.delete(requestId) - reply.raw.end() + safeEnd() } }) diff --git a/packages/http-routes/src/routes/web/ai-llm-stream.ts b/packages/http-routes/src/routes/web/ai-llm-stream.ts index e4b19ffd..924143ab 100644 --- a/packages/http-routes/src/routes/web/ai-llm-stream.ts +++ b/packages/http-routes/src/routes/web/ai-llm-stream.ts @@ -20,6 +20,7 @@ export function registerAiLlmStreamRoutes(server: FastifyInstance, ctx: HttpRout } const piModel = buildPiModel(llmConfig) + const abortController = new AbortController() reply.raw.writeHead(200, { 'Content-Type': 'text/event-stream', @@ -27,7 +28,12 @@ export function registerAiLlmStreamRoutes(server: FastifyInstance, ctx: HttpRout Connection: 'keep-alive', }) + reply.raw.on('close', () => { + if (!abortController.signal.aborted) abortController.abort() + }) + const sendChunk = (data: unknown) => { + if (reply.raw.writableEnded || reply.raw.destroyed) return reply.raw.write(`event: chunk\ndata: ${JSON.stringify(data)}\n\n`) } @@ -39,12 +45,14 @@ export function registerAiLlmStreamRoutes(server: FastifyInstance, ctx: HttpRout temperature: options?.temperature, maxTokens: options?.maxTokens, onChunk: sendChunk, + abortSignal: abortController.signal, }) } catch (error) { + if (abortController.signal.aborted) return const msg = error instanceof Error ? error.message : String(error) sendChunk({ content: '', isFinished: true, finishReason: 'error', error: msg }) + } finally { + if (!reply.raw.writableEnded && !reply.raw.destroyed) reply.raw.end() } - - reply.raw.end() }) } diff --git a/src/services/ai-stream/service.ts b/src/services/ai-stream/service.ts index 6ed76e42..96aa3057 100644 --- a/src/services/ai-stream/service.ts +++ b/src/services/ai-stream/service.ts @@ -12,24 +12,28 @@ export function useLlmStreamService() { options?: { temperature?: number; maxTokens?: number }, onChunk?: (chunk: LlmStreamChunk) => void ): Promise<{ success: boolean; error?: string }> { + let streamError: string | undefined try { await fetchSSE({ url: '/_web/ai/llm/chat-stream', body: { messages, options }, onEvent: ({ data }) => { - if (!onChunk) return try { const parsed = JSON.parse(data) as LlmStreamChunk - onChunk(parsed) + if (parsed.finishReason === 'error' && parsed.error) { + streamError = parsed.error + } + onChunk?.(parsed) } catch { // skip malformed JSON } }, }) + if (streamError) return { success: false, error: streamError } return { success: true } } catch (error) { const msg = error instanceof Error ? error.message : String(error) - return { success: false, error: msg } + return { success: false, error: streamError ?? msg } } }, } @@ -56,7 +60,7 @@ export function useAgentStreamService() { let resultContent = '' const toolsUsed: string[] = [] let toolRounds = 0 - let lastUsage: AgentStreamResult['result'] extends { totalUsage?: infer U } ? U : undefined + let lastUsage: import('./types').TokenUsage | undefined let hasError = false let lastError: AgentStreamResult['error'] @@ -82,9 +86,7 @@ export function useAgentStreamService() { if (chunk.content) resultContent += chunk.content break case 'tool_start': - if (chunk.toolName && !toolsUsed.includes(chunk.toolName)) { - toolsUsed.push(chunk.toolName) - } + if (chunk.toolName) toolsUsed.push(chunk.toolName) break case 'status': if (chunk.status) toolRounds = chunk.status.round