From ab12d288701db8b34a6ba9165e6e15feb4299589 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Wed, 25 Mar 2026 09:15:55 +0800 Subject: [PATCH 01/15] feat(renderer): edit SessionItem style --- .../components/WindowSideBarSessionItem.vue | 92 +++++++------------ 1 file changed, 34 insertions(+), 58 deletions(-) diff --git a/src/renderer/src/components/WindowSideBarSessionItem.vue b/src/renderer/src/components/WindowSideBarSessionItem.vue index fe98d12bc5..ac777ee089 100644 --- a/src/renderer/src/components/WindowSideBarSessionItem.vue +++ b/src/renderer/src/components/WindowSideBarSessionItem.vue @@ -1,64 +1,37 @@ From 6d84b0895019e34d366c13e504470739b20eb3c1 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Wed, 25 Mar 2026 10:23:09 +0800 Subject: [PATCH 02/15] feat(renderer): replace SessionItem context menu with hover interaction --- .../components/WindowSideBarSessionItem.vue | 92 ++++++++++++------- 1 file changed, 58 insertions(+), 34 deletions(-) diff --git a/src/renderer/src/components/WindowSideBarSessionItem.vue b/src/renderer/src/components/WindowSideBarSessionItem.vue index ac777ee089..fe98d12bc5 100644 --- a/src/renderer/src/components/WindowSideBarSessionItem.vue +++ b/src/renderer/src/components/WindowSideBarSessionItem.vue @@ -1,37 +1,64 @@ From 49d7851cd1b7c116ef73d1134921ecf019f1912b Mon Sep 17 00:00:00 2001 From: xiao-test Date: Wed, 25 Mar 2026 11:34:21 +0800 Subject: [PATCH 03/15] feat(renderer): replace SessionItem context menu with hover interaction --- .../components/WindowSideBarSessionItem.vue | 92 +++++++------------ 1 file changed, 34 insertions(+), 58 deletions(-) diff --git a/src/renderer/src/components/WindowSideBarSessionItem.vue b/src/renderer/src/components/WindowSideBarSessionItem.vue index fe98d12bc5..ac777ee089 100644 --- a/src/renderer/src/components/WindowSideBarSessionItem.vue +++ b/src/renderer/src/components/WindowSideBarSessionItem.vue @@ -1,64 +1,37 @@ From b27897a5dc15fa8ce7b2e7c9fd206bf01c324e34 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Tue, 8 Sep 2026 11:40:42 +0800 Subject: [PATCH 04/15] fix(images): cache generated and tool-returned image previews to disk --- .../deepchat/runtime/deferredToolExecutor.ts | 27 +++-- src/main/agent/deepchat/runtime/dispatch.ts | 27 +++-- src/main/lib/toolCallImagePreviews.ts | 34 ++++++ src/main/platform/imageCache.ts | 10 +- .../agentTools/agentImageGenerationTool.ts | 51 ++++++++- src/main/tool/agentTools/agentToolManager.ts | 3 +- test/main/lib/toolCallImagePreviews.test.ts | 65 +++++++++++ test/main/platform/imageCache.test.ts | 6 +- .../agentImageGenerationTool.test.ts | 102 ++++++++++++++---- 9 files changed, 275 insertions(+), 50 deletions(-) diff --git a/src/main/agent/deepchat/runtime/deferredToolExecutor.ts b/src/main/agent/deepchat/runtime/deferredToolExecutor.ts index 1d7869b212..0d1ff5ac7f 100644 --- a/src/main/agent/deepchat/runtime/deferredToolExecutor.ts +++ b/src/main/agent/deepchat/runtime/deferredToolExecutor.ts @@ -9,7 +9,10 @@ import type { } from '@shared/types/core/mcp' import type { ToolExecutionPort, ToolResultPort } from '@/agent/deepchat/loop/ports' import { awaitWithAbort } from '@/lib/awaitWithAbort' -import { extractToolCallImagePreviews } from '@/lib/toolCallImagePreviews' +import { + cacheToolCallImagePreviews, + extractToolCallImagePreviews +} from '@/lib/toolCallImagePreviews' import { CommittedToolOutcomeProjectionError, ExecutionJournalCorruptionError, @@ -657,15 +660,19 @@ export class DeferredToolExecutor { ) ) } - const imagePreviews = - rawData.imagePreviews ?? - (await extractToolCallImagePreviews({ - toolName, - toolArgs: toolCall.params || '{}', - content: rawData.content, - cacheImage: this.dependencies.cacheImage, - signal: deferredAbortSignal - })) + const imagePreviews = await cacheToolCallImagePreviews({ + imagePreviews: + rawData.imagePreviews ?? + (await extractToolCallImagePreviews({ + toolName, + toolArgs: toolCall.params || '{}', + content: rawData.content, + cacheImage: this.dependencies.cacheImage, + signal: deferredAbortSignal + })), + cacheImage: this.dependencies.cacheImage, + signal: deferredAbortSignal + }) throwIfAbortRequested(deferredAbortSignal) const normalizedContent = await this.dependencies.toolResultPort.normalize({ sessionId, diff --git a/src/main/agent/deepchat/runtime/dispatch.ts b/src/main/agent/deepchat/runtime/dispatch.ts index 9b89913e73..823506d822 100644 --- a/src/main/agent/deepchat/runtime/dispatch.ts +++ b/src/main/agent/deepchat/runtime/dispatch.ts @@ -73,7 +73,10 @@ import { buildAssistantResponseMarkdown, extractWaitingInteraction } from './sessionUpdates' -import { extractToolCallImagePreviews } from '@/lib/toolCallImagePreviews' +import { + cacheToolCallImagePreviews, + extractToolCallImagePreviews +} from '@/lib/toolCallImagePreviews' import { selectToolBatchExecutionMode } from './toolExecutionPolicy' import { resolveToolPermissionMode } from '@/tool/permission/permissionMode' import { segmentAssistantBlocksByProviderReplay } from './providerReplaySegments' @@ -2485,15 +2488,19 @@ async function runToolCall(params: { const subagentState = extractSubagentToolState(toolRawData) const rawResponseText = toolResponseToText(toolRawData.content) - const imagePreviews = - toolRawData.imagePreviews ?? - (await extractToolCallImagePreviews({ - toolName: completedToolCall.name, - toolArgs: completedToolCall.arguments, - content: toolRawData.content, - cacheImage: controls?.cacheImage, - signal: io.abortSignal - })) + const imagePreviews = await cacheToolCallImagePreviews({ + imagePreviews: + toolRawData.imagePreviews ?? + (await extractToolCallImagePreviews({ + toolName: completedToolCall.name, + toolArgs: completedToolCall.arguments, + content: toolRawData.content, + cacheImage: controls?.cacheImage, + signal: io.abortSignal + })), + cacheImage: controls?.cacheImage, + signal: io.abortSignal + }) toolRawData = { ...toolRawData, diff --git a/src/main/lib/toolCallImagePreviews.ts b/src/main/lib/toolCallImagePreviews.ts index be3fbd7a8c..0657147b85 100644 --- a/src/main/lib/toolCallImagePreviews.ts +++ b/src/main/lib/toolCallImagePreviews.ts @@ -350,3 +350,37 @@ export async function extractToolCallImagePreviews( ): Promise { return (await prepareToolCallImageContent(params)).imagePreviews } + +/** + * Backstop for tools that return `imagePreviews` directly (bypassing extraction): rewrites any + * inline base64 data URL to an on-disk `imgcache://` reference so multi-MB payloads never reach + * message persistence, IPC, or the renderer. Previews that cannot be cached are left unchanged. + */ +export async function cacheToolCallImagePreviews(params: { + imagePreviews: ToolCallImagePreview[] + cacheImage?: (data: string) => Promise + signal?: AbortSignal +}): Promise { + const { imagePreviews, cacheImage, signal } = params + if (!cacheImage || imagePreviews.length === 0) { + return imagePreviews + } + + let changed = false + const resolved: ToolCallImagePreview[] = [] + for (const preview of imagePreviews) { + const data = preview.data?.trim() + if (!data || !data.toLowerCase().startsWith('data:image/')) { + resolved.push(preview) + continue + } + const cached = await cachePreviewData(data, cacheImage, signal) + if (!cached) { + resolved.push(preview) + continue + } + changed = true + resolved.push({ ...preview, data: cached, mimeType: inferMimeType(cached, preview.mimeType) }) + } + return changed ? resolved : imagePreviews +} diff --git a/src/main/platform/imageCache.ts b/src/main/platform/imageCache.ts index adf20424d1..7414bc848e 100644 --- a/src/main/platform/imageCache.ts +++ b/src/main/platform/imageCache.ts @@ -7,7 +7,13 @@ import { nanoid } from 'nanoid' import axios, { type AxiosRequestConfig } from 'axios' const IMGCACHE_URL_PREFIX = 'imgcache://' -const MAX_CACHED_IMAGE_BYTES = 8 * 1024 * 1024 +// Write side: how large a generated/downloaded image may be and still land on disk. Anything +// larger stays an inline base64 payload, which then flows through the main process and renderer. +const MAX_CACHED_IMAGE_BYTES = 32 * 1024 * 1024 +// Read side: how large a cached image may be when expanded back into a base64 data URL as model +// or MCP tool input. Providers reject much smaller payloads; keep this tight independently of +// the on-disk cache budget. +const MAX_CACHED_IMAGE_INPUT_BYTES = 8 * 1024 * 1024 const IMAGE_CACHE_TIMEOUT_MS = 10_000 const MAX_IMAGE_REDIRECTS = 5 const REDIRECT_STATUS_CODES = new Set([301, 302, 303, 307, 308]) @@ -332,7 +338,7 @@ export async function resolveCachedImageDataUrl( if (!fileStat.isFile() || fileStat.isSymbolicLink()) { throw new Error('Cached image reference is not a regular file') } - if (fileStat.size > MAX_CACHED_IMAGE_BYTES) { + if (fileStat.size > MAX_CACHED_IMAGE_INPUT_BYTES) { throw new Error('Cached image exceeds the MCP image input limit') } diff --git a/src/main/tool/agentTools/agentImageGenerationTool.ts b/src/main/tool/agentTools/agentImageGenerationTool.ts index 8043b6a59c..0ceba071cb 100644 --- a/src/main/tool/agentTools/agentImageGenerationTool.ts +++ b/src/main/tool/agentTools/agentImageGenerationTool.ts @@ -66,6 +66,17 @@ type ImageGenerationModelSelection = { modelId: string } +// Inline base64 payloads below this size are allowed to pass through uncached; anything larger +// that could not be written to the image cache fails the tool call instead of flowing multi-MB +// strings through message persistence, IPC and the renderer. +const MAX_INLINE_IMAGE_BASE64_CHARS = 2 * 1024 * 1024 + +const estimateDataUrlBytes = (dataUrl: string): number => { + const commaIndex = dataUrl.indexOf(',') + const base64Length = commaIndex === -1 ? dataUrl.length : dataUrl.length - commaIndex - 1 + return Math.floor((base64Length * 3) / 4) +} + type AgentImageGenerationToolCallResult = { content: string rawData: { @@ -83,6 +94,7 @@ export class AgentImageGenerationTool { agentSettings: Pick sessions: AgentToolSessionPort provider: AgentProviderToolPort + cacheImage?: (data: string) => Promise } ) {} @@ -157,7 +169,13 @@ export class AgentImageGenerationTool { imageOptions, { signal: options?.signal } ) - const imagePreviews = result.images.map((image, index) => ({ + const images = await Promise.all( + result.images.map(async (image) => ({ + mimeType: image.mimeType, + data: await this.cacheGeneratedImageData(image.data, image.mimeType) + })) + ) + const imagePreviews = images.map((image, index) => ({ id: `generated-image-${index + 1}`, data: image.data, mimeType: image.mimeType, @@ -198,6 +216,37 @@ export class AgentImageGenerationTool { } } + private async cacheGeneratedImageData(data: string, mimeType: string): Promise { + const trimmed = data.trim() + if (trimmed.toLowerCase().startsWith('imgcache://')) { + return trimmed + } + + const source = + trimmed.startsWith('data:') || /^https?:\/\//i.test(trimmed) + ? trimmed + : `data:${mimeType || 'image/png'};base64,${trimmed}` + + let resolved = source + if (this.options.cacheImage) { + try { + resolved = await this.options.cacheImage(source) + } catch (error) { + logger.warn('[AgentImageGenerationTool] Failed to cache generated image', { error }) + } + } + + if ( + resolved.startsWith('data:') && + estimateDataUrlBytes(resolved) > MAX_INLINE_IMAGE_BASE64_CHARS + ) { + throw new Error( + 'Generated image could not be written to the image cache and is too large to return inline.' + ) + } + return resolved + } + private async resolveImageGenerationModel( conversationId?: string, options: { strict?: boolean; reportDiagnostics?: boolean } = {} diff --git a/src/main/tool/agentTools/agentToolManager.ts b/src/main/tool/agentTools/agentToolManager.ts index ddd4590d75..f70b223352 100644 --- a/src/main/tool/agentTools/agentToolManager.ts +++ b/src/main/tool/agentTools/agentToolManager.ts @@ -541,7 +541,8 @@ export class AgentToolManager { providerSettings: this.providerSettings, agentSettings: this.agentSettings, sessions: this.dependencies.sessions, - provider: this.dependencies.provider + provider: this.dependencies.provider, + cacheImage: this.dependencies.cacheImage }) this.planTool = new AgentPlanTool() this.tapeToolHandler = new AgentTapeToolHandler( diff --git a/test/main/lib/toolCallImagePreviews.test.ts b/test/main/lib/toolCallImagePreviews.test.ts index 4c6f08c64e..3ea29e786b 100644 --- a/test/main/lib/toolCallImagePreviews.test.ts +++ b/test/main/lib/toolCallImagePreviews.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from 'vitest' import { + cacheToolCallImagePreviews, extractToolCallImagePreviews, prepareToolCallImageContent } from '@/lib/toolCallImagePreviews' @@ -288,3 +289,67 @@ describe('extractToolCallImagePreviews', () => { } }) }) + +describe('cacheToolCallImagePreviews', () => { + it('rewrites inline base64 previews to imgcache references', async () => { + const cacheImage = vi.fn(async () => 'imgcache://cached.png') + + const previews = await cacheToolCallImagePreviews({ + imagePreviews: [ + { + id: 'tool_output-1', + data: 'data:image/png;base64,AAAA', + mimeType: 'image/png', + source: 'tool_output' + } + ], + cacheImage + }) + + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(previews).toEqual([ + { + id: 'tool_output-1', + data: 'imgcache://cached.png', + mimeType: 'image/png', + source: 'tool_output' + } + ]) + }) + + it('leaves references and uncacheable previews unchanged', async () => { + const cacheImage = vi.fn(async (data: string) => data) + const input = [ + { + id: 'tool_output-1', + data: 'imgcache://already-cached.png', + mimeType: 'image/png', + source: 'tool_output' as const + }, + { + id: 'tool_output-2', + data: 'data:image/png;base64,AAAA', + mimeType: 'image/png', + source: 'tool_output' as const + } + ] + + const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + + expect(cacheImage).toHaveBeenCalledOnce() + expect(previews).toBe(input) + }) + + it('returns the input unchanged without a cacheImage function', async () => { + const input = [ + { + id: 'tool_output-1', + data: 'data:image/png;base64,AAAA', + mimeType: 'image/png', + source: 'tool_output' as const + } + ] + + await expect(cacheToolCallImagePreviews({ imagePreviews: input })).resolves.toBe(input) + }) +}) diff --git a/test/main/platform/imageCache.test.ts b/test/main/platform/imageCache.test.ts index d4333dac79..26ffc75457 100644 --- a/test/main/platform/imageCache.test.ts +++ b/test/main/platform/imageCache.test.ts @@ -95,8 +95,8 @@ describe('imageCache', () => { expect(axiosMock).toHaveBeenCalledWith( expect.objectContaining({ maxRedirects: 0, - maxContentLength: 8 * 1024 * 1024, - maxBodyLength: 8 * 1024 * 1024, + maxContentLength: 32 * 1024 * 1024, + maxBodyLength: 32 * 1024 * 1024, signal: expect.any(AbortSignal) }) ) @@ -124,7 +124,7 @@ describe('imageCache', () => { axiosMock.mockResolvedValueOnce({ status: 200, headers: { 'content-type': 'image/png' }, - data: Buffer.alloc(8 * 1024 * 1024 + 1) + data: Buffer.alloc(32 * 1024 * 1024 + 1) }) await expect(cacheImage(sourceUrl, { allowPrivateNetwork: true })).resolves.toBe(sourceUrl) diff --git a/test/main/tool/agentTools/agentImageGenerationTool.test.ts b/test/main/tool/agentTools/agentImageGenerationTool.test.ts index 767a11aa19..bb5ad00f59 100644 --- a/test/main/tool/agentTools/agentImageGenerationTool.test.ts +++ b/test/main/tool/agentTools/agentImageGenerationTool.test.ts @@ -23,28 +23,8 @@ describe('Agent image generation tool', () => { let resolveConversationSessionInfo: ReturnType let manager: AgentToolManager - beforeEach(() => { - vi.clearAllMocks() - generateImageStandalone = vi.fn() - resolveConversationSessionInfo = vi.fn().mockResolvedValue({ - agentId: 'deepchat', - agentType: 'deepchat' - }) - providerSettings = { - resolveDeepChatAgentConfig: vi.fn().mockResolvedValue({ - imageGenerationModel: { providerId: 'openai', modelId: 'gpt-image-1' } - }), - getModelConfig: vi.fn().mockReturnValue({ - type: ModelType.ImageGeneration, - apiEndpoint: ApiEndpointType.Image, - vision: false, - functionCall: false, - reasoning: false, - maxTokens: 1024, - contextLength: 4096 - }) - } - manager = new AgentToolManager({ + const buildManager = (cacheImage?: (data: string) => Promise) => + new AgentToolManager({ skillSettings: { isEnabled: () => false } as any, settings: { get: vi.fn() }, commandPermissionHandler: new CommandPermissionService(), @@ -76,9 +56,33 @@ describe('Agent image generation tool', () => { createSettingsWindow: vi.fn(), sendToWindow: vi.fn().mockReturnValue(true), getApprovedFilePaths: vi.fn().mockReturnValue([]), - consumeSettingsApproval: vi.fn().mockReturnValue(false) + consumeSettingsApproval: vi.fn().mockReturnValue(false), + ...(cacheImage ? { cacheImage } : {}) }) }) + + beforeEach(() => { + vi.clearAllMocks() + generateImageStandalone = vi.fn() + resolveConversationSessionInfo = vi.fn().mockResolvedValue({ + agentId: 'deepchat', + agentType: 'deepchat' + }) + providerSettings = { + resolveDeepChatAgentConfig: vi.fn().mockResolvedValue({ + imageGenerationModel: { providerId: 'openai', modelId: 'gpt-image-1' } + }), + getModelConfig: vi.fn().mockReturnValue({ + type: ModelType.ImageGeneration, + apiEndpoint: ApiEndpointType.Image, + vision: false, + functionCall: false, + reasoning: false, + maxTokens: 1024, + contextLength: 4096 + }) + } + manager = buildManager() }) it('shows image_generate in settings context without a conversation', async () => { @@ -146,6 +150,58 @@ describe('Agent image generation tool', () => { expect(result.rawData.toolResult.ok).toBe(true) }) + it('caches raw base64 image data to an imgcache reference', async () => { + const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') + manager = buildManager(cacheImage) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: 'aGVsbG8=', mimeType: 'image/png' }] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=') + expect(result.rawData.imagePreviews).toEqual([ + { + id: 'generated-image-1', + data: 'imgcache://cached.png', + mimeType: 'image/png', + title: 'Generated image 1', + source: 'tool_output' + } + ]) + expect(result.rawData.toolResult.ok).toBe(true) + }) + + it('fails the tool call when a large image cannot be cached', async () => { + manager = buildManager(vi.fn(async (data: string) => data)) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [ + { data: `data:image/png;base64,${'A'.repeat(3 * 1024 * 1024)}`, mimeType: 'image/png' } + ] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + expect(result.rawData.imagePreviews).toBeUndefined() + }) + it('returns a recoverable tool error when no image model is configured', async () => { providerSettings.resolveDeepChatAgentConfig.mockResolvedValueOnce({}) From c747cfa47c4ce6e467fd191e8a87d3918677c24b Mon Sep 17 00:00:00 2001 From: zhangmo8 Date: Tue, 8 Sep 2026 18:22:36 +0800 Subject: [PATCH 05/15] fix(images): address review on cache signal and limits --- .../agent/deepchat/harness/runtimeServices.ts | 3 +- .../deepchat/runtime/deepChatLoopRunner.ts | 3 +- .../deepchat/runtime/deferredToolExecutor.ts | 3 +- src/main/agent/deepchat/runtime/types.ts | 3 +- src/main/app/composition.ts | 4 +- src/main/lib/toolCallImagePreviews.ts | 74 +++++++++--- src/main/platform/imageCache.ts | 9 +- .../agentTools/agentImageGenerationTool.ts | 39 ++++-- src/main/tool/runtimePorts.ts | 3 +- test/main/lib/toolCallImagePreviews.test.ts | 94 +++++++++++++-- test/main/platform/imageCache.test.ts | 9 ++ .../agentImageGenerationTool.test.ts | 112 +++++++++++++++++- 12 files changed, 316 insertions(+), 40 deletions(-) diff --git a/src/main/agent/deepchat/harness/runtimeServices.ts b/src/main/agent/deepchat/harness/runtimeServices.ts index fbbb1091cd..78e3f43f97 100644 --- a/src/main/agent/deepchat/harness/runtimeServices.ts +++ b/src/main/agent/deepchat/harness/runtimeServices.ts @@ -14,6 +14,7 @@ import type { SessionData } from '@/session/data' import type { SessionDatabase } from '@/session/data/database' import type { SessionPermissionPort, SessionUiPort } from '@/session/contracts' import type { SkillSettingsPort } from '@/skill/settings' +import type { CacheImageOptions } from '@/platform/imageCache' import type { AcpAgentInstanceDependencyFactory } from '@/agent/acp/instance' import type { DeepChatAgentRuntime } from '@/agent/deepchat/instance/deepChatAgentRuntime' import type { CommandShellService } from '@/agent/shared/process/commandShellService' @@ -81,7 +82,7 @@ export interface DeepChatHarnessDependencies { sessionUiPort: SessionUiPort memoryPort: MemoryRuntimePort getMemoryIngestionProjection(): MemoryIngestionProjection - cacheImage(data: string): Promise + cacheImage(data: string, options?: CacheImageOptions): Promise skillService: DeepChatHarnessSkillPort skillSettings: SkillSettingsPort traceSettings: AgentTraceSettingsPort diff --git a/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts b/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts index 7473fe536f..5272bf64b9 100644 --- a/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts +++ b/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts @@ -1,6 +1,7 @@ import type { PluginContextPort } from '@shared/types/userPlugin' import { projectPluginContext } from './pluginContext' import type { ProviderModelResolutionPort } from '@/provider/settings' +import type { CacheImageOptions } from '@/platform/imageCache' import logger from '@shared/logger' import type { AssistantMessageBlock, @@ -453,7 +454,7 @@ export interface DeepChatLoopRunnerPorts { memoryIngestionObserver: MemoryIngestionObserver toolExecutionPort: ToolExecutionPort toolResultPort: ToolResultPort - cacheImage(data: string): Promise + cacheImage(data: string, options?: CacheImageOptions): Promise registry: SessionScopeRegistry sessionSettings: Pick promptAssembly: Pick diff --git a/src/main/agent/deepchat/runtime/deferredToolExecutor.ts b/src/main/agent/deepchat/runtime/deferredToolExecutor.ts index 0d1ff5ac7f..ae222fffba 100644 --- a/src/main/agent/deepchat/runtime/deferredToolExecutor.ts +++ b/src/main/agent/deepchat/runtime/deferredToolExecutor.ts @@ -8,6 +8,7 @@ import type { ToolOutcomeProjection } from '@shared/types/core/mcp' import type { ToolExecutionPort, ToolResultPort } from '@/agent/deepchat/loop/ports' +import type { CacheImageOptions } from '@/platform/imageCache' import { awaitWithAbort } from '@/lib/awaitWithAbort' import { cacheToolCallImagePreviews, @@ -82,7 +83,7 @@ export interface DeferredToolExecutorDependencies { toolExecutionPort: ToolExecutionPort toolResultPort: ToolResultPort toolResolver: DeepChatToolResolver - cacheImage(data: string): Promise + cacheImage(data: string, options?: CacheImageOptions): Promise runLifecycle: Pick< RunLifecycleCoordinator, 'registerDeferredToolController' | 'clearDeferredToolController' | 'getAbortSignal' diff --git a/src/main/agent/deepchat/runtime/types.ts b/src/main/agent/deepchat/runtime/types.ts index 7e8293ce1f..851b4450d1 100644 --- a/src/main/agent/deepchat/runtime/types.ts +++ b/src/main/agent/deepchat/runtime/types.ts @@ -20,6 +20,7 @@ import type { DeepChatProviderAttemptIdentity } from '@shared/types/provider-att import type { DeepchatEventName } from '@shared/contracts/events' import type { DeepChatInternalSessionUpdate } from './sessionUpdates' import type { SessionTranscript } from '@/session/data/transcript' +import type { CacheImageOptions } from '@/platform/imageCache' import type { AgentPlanSnapshot, AgentPlanTerminalReason } from '@shared/types/agent-plan' import type { LoopRun } from '@/agent/deepchat/loop/loopRun' import type { @@ -206,7 +207,7 @@ export interface ProcessControlCollaborators { operation: ExecutionOperationIdentity outcomeEntryId: number }) => Promise | void - cacheImage?: (data: string) => Promise + cacheImage?: (data: string, options?: CacheImageOptions) => Promise } export interface ProcessInternalDiagnostics { diff --git a/src/main/app/composition.ts b/src/main/app/composition.ts index 622baf5330..257996f617 100644 --- a/src/main/app/composition.ts +++ b/src/main/app/composition.ts @@ -1701,7 +1701,7 @@ export async function createMainProcessControl(dependencies: { generateImageStandalone: (providerId, prompt, modelId, imageOptions, options) => providerRuntime.generateImageStandalone(providerId, prompt, modelId, imageOptions, options) }, - cacheImage: (data) => deviceService.cacheImage(data), + cacheImage: (data, options) => deviceService.cacheImage(data, options), desktop: { createSettingsWindow: () => windowPresenter.createSettingsWindow(), sendToWindow: (windowId, channel, ...args) => @@ -1882,7 +1882,7 @@ export async function createMainProcessControl(dependencies: { sessionUiPort, memoryPort: memoryService, getMemoryIngestionProjection: () => memoryDatabase.ingestionProjectionTable, - cacheImage: (data) => deviceService.cacheImage(data), + cacheImage: (data, options) => deviceService.cacheImage(data, options), runJournalObserver: emitRunJournalObservation, skillService: skillService, skillSettings, diff --git a/src/main/lib/toolCallImagePreviews.ts b/src/main/lib/toolCallImagePreviews.ts index 0657147b85..208d3d053a 100644 --- a/src/main/lib/toolCallImagePreviews.ts +++ b/src/main/lib/toolCallImagePreviews.ts @@ -1,6 +1,9 @@ import type { MCPContentItem, ToolCallImagePreview } from '@shared/types/core/mcp' +import type { CacheImageOptions } from '@/platform/imageCache' import { awaitWithAbort } from './awaitWithAbort' +export type CacheImageCallback = (data: string, options?: CacheImageOptions) => Promise + type ImagePreviewInput = { data: string mimeType: string @@ -13,7 +16,7 @@ type ExtractToolCallImagePreviewsParams = { toolName?: string toolArgs?: string content: string | MCPContentItem[] - cacheImage?: (data: string) => Promise + cacheImage?: CacheImageCallback signal?: AbortSignal } @@ -93,10 +96,10 @@ function extractEmbeddedHttpImageReferences(content: string): string[] { function normalizeImagePayload(data: string, mimeType: string): string { const trimmed = data.trim() if ( - trimmed.startsWith('data:image/') || - trimmed.startsWith('imgcache://') || - trimmed.startsWith('http://') || - trimmed.startsWith('https://') + trimmed.toLowerCase().startsWith('data:image/') || + trimmed.toLowerCase().startsWith('imgcache://') || + trimmed.toLowerCase().startsWith('http://') || + trimmed.toLowerCase().startsWith('https://') ) { return trimmed } @@ -104,9 +107,25 @@ function normalizeImagePayload(data: string, mimeType: string): string { return `data:${mimeType || 'image/png'};base64,${trimmed}` } +/** + * Normalizes a preview payload before it is handed to the image cache: lowercases a case-insensitive + * `data:` scheme so `imageCache`'s base64 router recognizes it, and wraps bare base64 into a full + * data URL. Non-base64 references (HTTP(S) URLs, `imgcache://`) are returned unchanged. + */ +function normalizeBackstopPreviewData(data: string, mimeType: string): string { + const trimmed = data.trim() + if (/^data:/i.test(trimmed)) { + return `data:${trimmed.slice('data:'.length)}` + } + if (/^https?:\/\//i.test(trimmed) || trimmed.toLowerCase().startsWith('imgcache://')) { + return trimmed + } + return `data:${mimeType || 'image/png'};base64,${trimmed}` +} + async function cachePreviewData( data: string, - cacheImage?: (data: string) => Promise, + cacheImage?: CacheImageCallback, signal?: AbortSignal ): Promise { if (data.trim().toLowerCase().startsWith('imgcache://')) { @@ -118,7 +137,7 @@ async function cachePreviewData( try { signal?.throwIfAborted() - const cachedData = await awaitWithAbort(cacheImage(data), signal) + const cachedData = await awaitWithAbort(cacheImage(data, { signal }), signal) const cachedDataTrimmed = cachedData.trim() return cachedDataTrimmed.toLowerCase().startsWith('imgcache://') ? cachedDataTrimmed : undefined } catch (error) { @@ -353,12 +372,18 @@ export async function extractToolCallImagePreviews( /** * Backstop for tools that return `imagePreviews` directly (bypassing extraction): rewrites any - * inline base64 data URL to an on-disk `imgcache://` reference so multi-MB payloads never reach - * message persistence, IPC, or the renderer. Previews that cannot be cached are left unchanged. + * inline base64 payload (data URL or bare base64, in any casing) to an on-disk `imgcache://` + * reference so multi-MB payloads never reach message persistence, IPC, or the renderer. Previews + * that cannot be cached are left unchanged. Mirroring the extraction path, at most + * `MAX_TOOL_CALL_IMAGE_PREVIEWS` distinct payloads are written per call and duplicates are dropped. + * + * Note: like the extraction path this only rewrites the `imagePreviews` array — a tool that also + * embeds the same base64 in its textual `content` is expected to use the extraction path (which + * rewrites content references) instead of returning `imagePreviews` directly. */ export async function cacheToolCallImagePreviews(params: { imagePreviews: ToolCallImagePreview[] - cacheImage?: (data: string) => Promise + cacheImage?: CacheImageCallback signal?: AbortSignal }): Promise { const { imagePreviews, cacheImage, signal } = params @@ -368,17 +393,38 @@ export async function cacheToolCallImagePreviews(params: { let changed = false const resolved: ToolCallImagePreview[] = [] + const seenInputs = new Set() + const seen = new Set() for (const preview of imagePreviews) { - const data = preview.data?.trim() - if (!data || !data.toLowerCase().startsWith('data:image/')) { + const rawData = preview.data?.trim() + if ( + !rawData || + rawData.toLowerCase().startsWith('imgcache://') || + /^https?:\/\//i.test(rawData) + ) { resolved.push(preview) continue } - const cached = await cachePreviewData(data, cacheImage, signal) - if (!cached) { + const normalized = normalizeBackstopPreviewData(rawData, preview.mimeType) + if (!/^data:image\//i.test(normalized)) { + resolved.push(preview) + continue + } + const inputKey = normalized.toLowerCase() + if (seenInputs.has(inputKey)) { + continue + } + if (seenInputs.size >= MAX_TOOL_CALL_IMAGE_PREVIEWS) { + resolved.push(preview) + continue + } + seenInputs.add(inputKey) + const cached = await cachePreviewData(normalized, cacheImage, signal) + if (!cached || seen.has(cached)) { resolved.push(preview) continue } + seen.add(cached) changed = true resolved.push({ ...preview, data: cached, mimeType: inferMimeType(cached, preview.mimeType) }) } diff --git a/src/main/platform/imageCache.ts b/src/main/platform/imageCache.ts index 7414bc848e..9b4069d74f 100644 --- a/src/main/platform/imageCache.ts +++ b/src/main/platform/imageCache.ts @@ -9,6 +9,11 @@ import axios, { type AxiosRequestConfig } from 'axios' const IMGCACHE_URL_PREFIX = 'imgcache://' // Write side: how large a generated/downloaded image may be and still land on disk. Anything // larger stays an inline base64 payload, which then flows through the main process and renderer. +// NOTE: this intentionally exceeds the read-side budget below — an image cached in the +// 8–32 MiB band is displayable via its `imgcache://` reference but can never be expanded back +// into model/MCP input (oversized inline payloads are instead rejected at the tool boundary, +// e.g. `cacheGeneratedImageData`). The asymmetry is deliberate: writes are cheap and bounded, +// while reads must stay within provider input limits. const MAX_CACHED_IMAGE_BYTES = 32 * 1024 * 1024 // Read side: how large a cached image may be when expanded back into a base64 data URL as model // or MCP tool input. Providers reject much smaller payloads; keep this tight independently of @@ -242,7 +247,7 @@ async function cacheImageFromBase64( ): Promise { try { signal?.throwIfAborted() - const matches = base64Data.match(/^data:([^;]+);base64,(.*)$/) + const matches = base64Data.match(/^data:([^;]+);base64,(.*)$/i) if (!matches || matches.length !== 3) { console.warn('无效的Base64图片数据') return base64Data @@ -297,7 +302,7 @@ export async function cacheImage( if (imageData.startsWith('http://') || imageData.startsWith('https://')) { return cacheImageFromUrl(imageData, cacheDir, fileName, options) } - if (imageData.startsWith('data:image/')) { + if (/^data:image\//i.test(imageData)) { return cacheImageFromBase64(imageData, cacheDir, fileName, options.signal) } console.warn('不支持的图片格式') diff --git a/src/main/tool/agentTools/agentImageGenerationTool.ts b/src/main/tool/agentTools/agentImageGenerationTool.ts index 0ceba071cb..8cd40b5e20 100644 --- a/src/main/tool/agentTools/agentImageGenerationTool.ts +++ b/src/main/tool/agentTools/agentImageGenerationTool.ts @@ -22,6 +22,7 @@ import { IMAGE_GENERATION_TOOL_SERVER_NAME } from '@shared/agentImageGenerationTool' import logger from '@shared/logger' +import type { CacheImageCallback } from '@/lib/toolCallImagePreviews' import type { AgentProviderToolPort, AgentToolSessionPort } from '../runtimePorts' import type { AgentSettingsPort } from '@/agent/settings' @@ -68,13 +69,13 @@ type ImageGenerationModelSelection = { // Inline base64 payloads below this size are allowed to pass through uncached; anything larger // that could not be written to the image cache fails the tool call instead of flowing multi-MB -// strings through message persistence, IPC and the renderer. +// strings through message persistence, IPC and the renderer. The limit applies to the encoded +// base64 character length of the payload (not its decoded byte size). const MAX_INLINE_IMAGE_BASE64_CHARS = 2 * 1024 * 1024 -const estimateDataUrlBytes = (dataUrl: string): number => { +const estimateBase64PayloadChars = (dataUrl: string): number => { const commaIndex = dataUrl.indexOf(',') - const base64Length = commaIndex === -1 ? dataUrl.length : dataUrl.length - commaIndex - 1 - return Math.floor((base64Length * 3) / 4) + return commaIndex === -1 ? dataUrl.length : dataUrl.length - commaIndex - 1 } type AgentImageGenerationToolCallResult = { @@ -94,7 +95,7 @@ export class AgentImageGenerationTool { agentSettings: Pick sessions: AgentToolSessionPort provider: AgentProviderToolPort - cacheImage?: (data: string) => Promise + cacheImage?: CacheImageCallback } ) {} @@ -172,7 +173,7 @@ export class AgentImageGenerationTool { const images = await Promise.all( result.images.map(async (image) => ({ mimeType: image.mimeType, - data: await this.cacheGeneratedImageData(image.data, image.mimeType) + data: await this.cacheGeneratedImageData(image.data, image.mimeType, options?.signal) })) ) const imagePreviews = images.map((image, index) => ({ @@ -216,7 +217,11 @@ export class AgentImageGenerationTool { } } - private async cacheGeneratedImageData(data: string, mimeType: string): Promise { + private async cacheGeneratedImageData( + data: string, + mimeType: string, + signal?: AbortSignal + ): Promise { const trimmed = data.trim() if (trimmed.toLowerCase().startsWith('imgcache://')) { return trimmed @@ -230,15 +235,31 @@ export class AgentImageGenerationTool { let resolved = source if (this.options.cacheImage) { try { - resolved = await this.options.cacheImage(source) + // Provider-returned HTTP(S) URLs are cached with private-network access disabled and the + // tool-call abort signal forwarded, so cancellation cannot leave an unmanaged download + // running. `allowPrivateNetwork: false` is passed unconditionally so an omitted signal + // cannot re-enable private-network access. + resolved = await this.options.cacheImage(source, { + signal, + allowPrivateNetwork: false + }) } catch (error) { + if (signal?.aborted) throw error logger.warn('[AgentImageGenerationTool] Failed to cache generated image', { error }) } } + signal?.throwIfAborted() + + // A provider-returned HTTP(S) URL that could not be written to the cache must not flow through + // the tool-result pipeline as an unmanaged remote reference; fail the call instead so the + // agent can surface a recoverable error. + if (/^https?:\/\//i.test(resolved)) { + throw new Error('Generated image URL could not be written to the image cache.') + } if ( resolved.startsWith('data:') && - estimateDataUrlBytes(resolved) > MAX_INLINE_IMAGE_BASE64_CHARS + estimateBase64PayloadChars(resolved) > MAX_INLINE_IMAGE_BASE64_CHARS ) { throw new Error( 'Generated image could not be written to the image cache and is too large to return inline.' diff --git a/src/main/tool/runtimePorts.ts b/src/main/tool/runtimePorts.ts index b2d05948d9..60efdab8fb 100644 --- a/src/main/tool/runtimePorts.ts +++ b/src/main/tool/runtimePorts.ts @@ -28,6 +28,7 @@ import type { } from '@shared/orchestration/liveDelegation' import type { AgentInvocationAdmissionPort } from '@/agent/invocationAdmission' import type { SkillServicePort } from '@shared/types/skill' +import type { CacheImageOptions } from '@/platform/imageCache' import type { AgentMemoryCategory } from '@shared/types/agent-memory' import type { MemoryCommandResult } from '@shared/contracts/routes/memory.routes' import type { SessionRuntimeUpdate } from '@/session/runtimeEvents' @@ -275,5 +276,5 @@ export interface AgentToolDependencies { provider: AgentProviderToolPort desktop: AgentDesktopToolPort permissions: AgentToolPermissionPort - cacheImage(data: string): Promise + cacheImage(data: string, options?: CacheImageOptions): Promise } diff --git a/test/main/lib/toolCallImagePreviews.test.ts b/test/main/lib/toolCallImagePreviews.test.ts index 3ea29e786b..67a87effe7 100644 --- a/test/main/lib/toolCallImagePreviews.test.ts +++ b/test/main/lib/toolCallImagePreviews.test.ts @@ -19,7 +19,7 @@ describe('extractToolCallImagePreviews', () => { }) expect(cacheImage).toHaveBeenCalledOnce() - expect(cacheImage).toHaveBeenCalledWith(sourceUrl) + expect(cacheImage).toHaveBeenCalledWith(sourceUrl, { signal: undefined }) expect(prepared.content).toEqual([ { type: 'text', text: 'Success. Image URL(s): imgcache://output.jpg' }, { type: 'text', text: 'Reference: imgcache://output.jpg' } @@ -116,7 +116,9 @@ describe('extractToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA', { + signal: undefined + }) expect(previews).toEqual([ { id: 'mcp_image-1', @@ -159,7 +161,9 @@ describe('extractToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('https://example.com/output.webp') + expect(cacheImage).toHaveBeenCalledWith('https://example.com/output.webp', { + signal: undefined + }) expect(previews).toEqual([ { id: 'tool_output-1', @@ -180,7 +184,9 @@ describe('extractToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA', { + signal: undefined + }) expect(previews).toEqual([ { id: 'mcp_image-1', @@ -198,7 +204,9 @@ describe('extractToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA', { + signal: undefined + }) expect(previews).toEqual([ { id: 'mcp_image-1', @@ -216,7 +224,9 @@ describe('extractToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA', { + signal: undefined + }) expect(previews).toEqual([ { id: 'mcp_image-1', @@ -306,7 +316,9 @@ describe('cacheToolCallImagePreviews', () => { cacheImage }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,AAAA', { + signal: undefined + }) expect(previews).toEqual([ { id: 'tool_output-1', @@ -352,4 +364,72 @@ describe('cacheToolCallImagePreviews', () => { await expect(cacheToolCallImagePreviews({ imagePreviews: input })).resolves.toBe(input) }) + + it('normalizes an uppercase data URL prefix before caching', async () => { + const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') + const input = [ + { + id: 'tool_output-1', + data: 'DATA:IMAGE/PNG;BASE64,QUFBQQ==', + mimeType: 'image/png', + source: 'tool_output' as const + } + ] + + const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + + expect(cacheImage).toHaveBeenCalledWith('data:IMAGE/PNG;BASE64,QUFBQQ==', { + signal: undefined + }) + expect(previews).toEqual([{ ...input[0], data: 'imgcache://cached.png' }]) + }) + + it('wraps and caches bare base64 previews', async () => { + const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') + const input = [ + { + id: 'tool_output-1', + data: 'aGVsbG8=', + mimeType: 'image/png', + source: 'tool_output' as const + } + ] + + const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { + signal: undefined + }) + expect(previews).toEqual([{ ...input[0], data: 'imgcache://cached.png' }]) + }) + + it('dedupes identical payloads and caps cache writes', async () => { + const cacheImage = vi.fn(async (data: string) => `imgcache://cached-${data.length}.png`) + const duplicate = { + id: 'tool_output-1', + data: 'data:image/png;base64,QUFBQQ==', + mimeType: 'image/png', + source: 'tool_output' as const + } + const input = [ + duplicate, + { ...duplicate, id: 'tool_output-dup' }, + ...Array.from({ length: 6 }, (_, index) => ({ + id: `tool_output-${index + 2}`, + data: `data:image/png;base64,${'A'.repeat(index + 1)}`, + mimeType: 'image/png', + source: 'tool_output' as const + })) + ] + + const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + + expect(cacheImage).toHaveBeenCalledTimes(4) + expect(previews).toHaveLength(7) + expect(previews.find((preview) => preview.id === 'tool_output-dup')).toBeUndefined() + expect(previews.filter((preview) => preview.data?.startsWith('imgcache://'))).toHaveLength(4) + expect( + previews.filter((preview) => preview.data?.startsWith('data:image/')) + ).toHaveLength(3) + }) }) diff --git a/test/main/platform/imageCache.test.ts b/test/main/platform/imageCache.test.ts index 26ffc75457..31f546b1be 100644 --- a/test/main/platform/imageCache.test.ts +++ b/test/main/platform/imageCache.test.ts @@ -153,6 +153,15 @@ describe('imageCache', () => { expect(requestSignal?.aborted).toBe(true) }) + it('caches base64 data URLs with an uppercase scheme prefix', async () => { + const cached = await cacheImage('DATA:IMAGE/PNG;BASE64,aW1hZ2U=') + + expect(cached).toMatch(/^imgcache:\/\/.+\.png$/) + await expect( + fs.readFile(path.join(electronMock.userDataPath, 'images', cached.slice('imgcache://'.length))) + ).resolves.toEqual(Buffer.from('image')) + }) + it('resolves a cached image to a MIME-correct data URL', async () => { await fs.writeFile(path.join(electronMock.userDataPath, 'images', 'generated.png'), 'image') diff --git a/test/main/tool/agentTools/agentImageGenerationTool.test.ts b/test/main/tool/agentTools/agentImageGenerationTool.test.ts index bb5ad00f59..627b06d43e 100644 --- a/test/main/tool/agentTools/agentImageGenerationTool.test.ts +++ b/test/main/tool/agentTools/agentImageGenerationTool.test.ts @@ -165,7 +165,62 @@ describe('Agent image generation tool', () => { 'conv-1' )) as any - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { + signal: undefined, + allowPrivateNetwork: false + }) + expect(result.rawData.imagePreviews).toEqual([ + { + id: 'generated-image-1', + data: 'imgcache://cached.png', + mimeType: 'image/png', + title: 'Generated image 1', + source: 'tool_output' + } + ]) + expect(result.rawData.toolResult.ok).toBe(true) + }) + + it('fails the tool call when a generated HTTP image URL cannot be cached', async () => { + manager = buildManager(vi.fn(async (data: string) => data)) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: 'https://example.com/generated.png', mimeType: 'image/png' }] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + }) + + it('caches generated HTTP image URLs with private-network access disabled', async () => { + const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') + manager = buildManager(cacheImage) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: 'https://example.com/generated.png', mimeType: 'image/png' }] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(cacheImage).toHaveBeenCalledWith('https://example.com/generated.png', { + signal: undefined, + allowPrivateNetwork: false + }) expect(result.rawData.imagePreviews).toEqual([ { id: 'generated-image-1', @@ -178,6 +233,61 @@ describe('Agent image generation tool', () => { expect(result.rawData.toolResult.ok).toBe(true) }) + it('fails the tool call above the inline base64 character limit even when decoded bytes are smaller', async () => { + manager = buildManager(vi.fn(async (data: string) => data)) + // Encoded character length is just above the 2 MiB limit while the decoded payload is only + // ~1.5 MiB — the limit is enforced on the encoded base64 payload, not decoded byte size. + const payload = 'A'.repeat(2 * 1024 * 1024 + 1) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: `data:image/png;base64,${payload}`, mimeType: 'image/png' }] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + }) + + it('does not return a success result when the cache write is aborted', async () => { + const controller = new AbortController() + const cacheImage = vi.fn(async () => { + controller.abort() + throw new DOMException('The operation was aborted.', 'AbortError') + }) + manager = buildManager(cacheImage) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: 'aGVsbG8=', mimeType: 'image/png' }] + }) + + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1', + { signal: controller.signal } + )) as any + + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { + signal: controller.signal, + allowPrivateNetwork: false + }) + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + }) + it('fails the tool call when a large image cannot be cached', async () => { manager = buildManager(vi.fn(async (data: string) => data)) generateImageStandalone.mockResolvedValue({ From 6687499976d0b1a01044ca8e197bf6e6357ca83a Mon Sep 17 00:00:00 2001 From: zhangmo8 Date: Tue, 8 Sep 2026 18:27:24 +0800 Subject: [PATCH 06/15] style(images): apply oxfmt to cache tests --- test/main/lib/toolCallImagePreviews.test.ts | 4 +--- test/main/platform/imageCache.test.ts | 4 +++- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/main/lib/toolCallImagePreviews.test.ts b/test/main/lib/toolCallImagePreviews.test.ts index 67a87effe7..4864a45930 100644 --- a/test/main/lib/toolCallImagePreviews.test.ts +++ b/test/main/lib/toolCallImagePreviews.test.ts @@ -428,8 +428,6 @@ describe('cacheToolCallImagePreviews', () => { expect(previews).toHaveLength(7) expect(previews.find((preview) => preview.id === 'tool_output-dup')).toBeUndefined() expect(previews.filter((preview) => preview.data?.startsWith('imgcache://'))).toHaveLength(4) - expect( - previews.filter((preview) => preview.data?.startsWith('data:image/')) - ).toHaveLength(3) + expect(previews.filter((preview) => preview.data?.startsWith('data:image/'))).toHaveLength(3) }) }) diff --git a/test/main/platform/imageCache.test.ts b/test/main/platform/imageCache.test.ts index 31f546b1be..6e64910f27 100644 --- a/test/main/platform/imageCache.test.ts +++ b/test/main/platform/imageCache.test.ts @@ -158,7 +158,9 @@ describe('imageCache', () => { expect(cached).toMatch(/^imgcache:\/\/.+\.png$/) await expect( - fs.readFile(path.join(electronMock.userDataPath, 'images', cached.slice('imgcache://'.length))) + fs.readFile( + path.join(electronMock.userDataPath, 'images', cached.slice('imgcache://'.length)) + ) ).resolves.toEqual(Buffer.from('image')) }) From 4a2cd3a4a06d8a0f4119e25ebecf74bc5ec9e7fa Mon Sep 17 00:00:00 2001 From: zerob13 Date: Wed, 9 Sep 2026 11:11:34 +0800 Subject: [PATCH 07/15] fix(images): preserve previews and cancellation --- src/main/lib/toolCallImagePreviews.ts | 25 ++- src/main/platform/imageCache.ts | 2 +- src/main/provider/aiSdk/runtime.ts | 7 +- src/main/provider/aiSdk/streamAdapter.ts | 8 +- .../agentTools/agentImageGenerationTool.ts | 15 +- src/main/tool/agentTools/agentToolManager.ts | 4 +- test/main/lib/toolCallImagePreviews.test.ts | 67 ++++--- test/main/platform/imageCache.test.ts | 84 ++++++--- test/main/provider/aiSdkStreamAdapter.test.ts | 29 ++- .../agentImageGenerationTool.test.ts | 174 ++++++++++-------- 10 files changed, 267 insertions(+), 148 deletions(-) diff --git a/src/main/lib/toolCallImagePreviews.ts b/src/main/lib/toolCallImagePreviews.ts index 208d3d053a..891edb933f 100644 --- a/src/main/lib/toolCallImagePreviews.ts +++ b/src/main/lib/toolCallImagePreviews.ts @@ -373,13 +373,12 @@ export async function extractToolCallImagePreviews( /** * Backstop for tools that return `imagePreviews` directly (bypassing extraction): rewrites any * inline base64 payload (data URL or bare base64, in any casing) to an on-disk `imgcache://` - * reference so multi-MB payloads never reach message persistence, IPC, or the renderer. Previews - * that cannot be cached are left unchanged. Mirroring the extraction path, at most - * `MAX_TOOL_CALL_IMAGE_PREVIEWS` distinct payloads are written per call and duplicates are dropped. + * reference to reduce payloads reaching message persistence, IPC, and the renderer. Previews + * that cannot be cached are left unchanged. At most `MAX_TOOL_CALL_IMAGE_PREVIEWS` distinct + * inline payloads are processed per call; duplicates and excess inline previews are dropped. * - * Note: like the extraction path this only rewrites the `imagePreviews` array — a tool that also - * embeds the same base64 in its textual `content` is expected to use the extraction path (which - * rewrites content references) instead of returning `imagePreviews` directly. + * Tools returning `imagePreviews` directly must keep inline image data out of textual `content`; + * this helper only rewrites the preview array. */ export async function cacheToolCallImagePreviews(params: { imagePreviews: ToolCallImagePreview[] @@ -387,6 +386,7 @@ export async function cacheToolCallImagePreviews(params: { signal?: AbortSignal }): Promise { const { imagePreviews, cacheImage, signal } = params + signal?.throwIfAborted() if (!cacheImage || imagePreviews.length === 0) { return imagePreviews } @@ -396,6 +396,7 @@ export async function cacheToolCallImagePreviews(params: { const seenInputs = new Set() const seen = new Set() for (const preview of imagePreviews) { + signal?.throwIfAborted() const rawData = preview.data?.trim() if ( !rawData || @@ -410,23 +411,29 @@ export async function cacheToolCallImagePreviews(params: { resolved.push(preview) continue } - const inputKey = normalized.toLowerCase() + const inputKey = normalized if (seenInputs.has(inputKey)) { + changed = true continue } if (seenInputs.size >= MAX_TOOL_CALL_IMAGE_PREVIEWS) { - resolved.push(preview) + changed = true continue } seenInputs.add(inputKey) const cached = await cachePreviewData(normalized, cacheImage, signal) - if (!cached || seen.has(cached)) { + if (!cached) { resolved.push(preview) continue } + if (seen.has(cached)) { + changed = true + continue + } seen.add(cached) changed = true resolved.push({ ...preview, data: cached, mimeType: inferMimeType(cached, preview.mimeType) }) } + signal?.throwIfAborted() return changed ? resolved : imagePreviews } diff --git a/src/main/platform/imageCache.ts b/src/main/platform/imageCache.ts index 9b4069d74f..bfcfdc254c 100644 --- a/src/main/platform/imageCache.ts +++ b/src/main/platform/imageCache.ts @@ -299,7 +299,7 @@ export async function cacheImage( if (!fs.existsSync(cacheDir)) fs.mkdirSync(cacheDir, { recursive: true }) const fileName = `img_${Date.now()}_${nanoid(8)}` - if (imageData.startsWith('http://') || imageData.startsWith('https://')) { + if (/^https?:\/\//i.test(imageData)) { return cacheImageFromUrl(imageData, cacheDir, fileName, options) } if (/^data:image\//i.test(imageData)) { diff --git a/src/main/provider/aiSdk/runtime.ts b/src/main/provider/aiSdk/runtime.ts index c5809a6c3c..5244460197 100644 --- a/src/main/provider/aiSdk/runtime.ts +++ b/src/main/provider/aiSdk/runtime.ts @@ -1441,7 +1441,8 @@ export async function* runAiSdkCoreStream( ) const dataUrl = `data:${mimeType};base64,${base64}` - const cachedAudio = await cacheImage(dataUrl) + const cachedAudio = await cacheImage(dataUrl, { signal }) + signal?.throwIfAborted() yield { type: 'image_data', image_data: { @@ -1546,7 +1547,8 @@ export async function* runAiSdkCoreStream( for (const image of result.images) { const dataUrl = `data:${image.mediaType};base64,${image.base64}` - const cachedImage = await cacheImage(dataUrl) + const cachedImage = await cacheImage(dataUrl, { signal: requestSignal }) + requestSignal?.throwIfAborted() yield { type: 'image_data', image_data: { @@ -1605,6 +1607,7 @@ export async function* runAiSdkCoreStream( yield* adaptAiSdkStream(result.stream, { supportsNativeTools: runtime.supportsNativeTools, cacheImage, + signal: requestSignal, projectRawChunk: runtime.providerAdapter?.projectRawChunk }) } diff --git a/src/main/provider/aiSdk/streamAdapter.ts b/src/main/provider/aiSdk/streamAdapter.ts index b9bfe00d54..ecbf85901d 100644 --- a/src/main/provider/aiSdk/streamAdapter.ts +++ b/src/main/provider/aiSdk/streamAdapter.ts @@ -1,6 +1,7 @@ import { createStreamEvent, type LLMCoreStreamEvent } from '@shared/types/core/llm-events' import type { ChatMessageProviderOptions } from '@shared/types/core/chat-message' import type { ToolSet, TextStreamPart } from 'ai' +import type { CacheImageOptions } from '@/platform/imageCache' import { parseLegacyFunctionCalls } from './toolProtocol' import { extractProviderFailureMetadata } from '../providerFailure' @@ -89,7 +90,8 @@ function toProviderOptions(value: unknown): ChatMessageProviderOptions | undefin export interface AdaptAiSdkStreamOptions { supportsNativeTools: boolean - cacheImage?: (data: string) => Promise + cacheImage?: (data: string, options?: CacheImageOptions) => Promise + signal?: AbortSignal projectRawChunk?: (rawValue: unknown) => LLMCoreStreamEvent | null } @@ -322,12 +324,14 @@ export async function* adaptAiSdkStream( if (options.cacheImage) { try { - cachedImage = await options.cacheImage(dataUrl) + cachedImage = await options.cacheImage(dataUrl, { signal: options.signal }) } catch (error) { + options.signal?.throwIfAborted() console.warn('[AI SDK Stream Adapter] Failed to cache image part:', error) } } + options.signal?.throwIfAborted() yield createStreamEvent.imageData({ data: cachedImage, mimeType: mediaType diff --git a/src/main/tool/agentTools/agentImageGenerationTool.ts b/src/main/tool/agentTools/agentImageGenerationTool.ts index 8cd40b5e20..87b391fe30 100644 --- a/src/main/tool/agentTools/agentImageGenerationTool.ts +++ b/src/main/tool/agentTools/agentImageGenerationTool.ts @@ -23,6 +23,7 @@ import { } from '@shared/agentImageGenerationTool' import logger from '@shared/logger' import type { CacheImageCallback } from '@/lib/toolCallImagePreviews' +import { awaitWithAbort } from '@/lib/awaitWithAbort' import type { AgentProviderToolPort, AgentToolSessionPort } from '../runtimePorts' import type { AgentSettingsPort } from '@/agent/settings' @@ -176,6 +177,7 @@ export class AgentImageGenerationTool { data: await this.cacheGeneratedImageData(image.data, image.mimeType, options?.signal) })) ) + options?.signal?.throwIfAborted() const imagePreviews = images.map((image, index) => ({ id: `generated-image-${index + 1}`, data: image.data, @@ -222,13 +224,14 @@ export class AgentImageGenerationTool { mimeType: string, signal?: AbortSignal ): Promise { + signal?.throwIfAborted() const trimmed = data.trim() if (trimmed.toLowerCase().startsWith('imgcache://')) { return trimmed } const source = - trimmed.startsWith('data:') || /^https?:\/\//i.test(trimmed) + /^data:/i.test(trimmed) || /^https?:\/\//i.test(trimmed) ? trimmed : `data:${mimeType || 'image/png'};base64,${trimmed}` @@ -239,10 +242,10 @@ export class AgentImageGenerationTool { // tool-call abort signal forwarded, so cancellation cannot leave an unmanaged download // running. `allowPrivateNetwork: false` is passed unconditionally so an omitted signal // cannot re-enable private-network access. - resolved = await this.options.cacheImage(source, { - signal, - allowPrivateNetwork: false - }) + resolved = await awaitWithAbort( + this.options.cacheImage(source, { signal, allowPrivateNetwork: false }), + signal + ) } catch (error) { if (signal?.aborted) throw error logger.warn('[AgentImageGenerationTool] Failed to cache generated image', { error }) @@ -258,7 +261,7 @@ export class AgentImageGenerationTool { } if ( - resolved.startsWith('data:') && + /^data:/i.test(resolved) && estimateBase64PayloadChars(resolved) > MAX_INLINE_IMAGE_BASE64_CHARS ) { throw new Error( diff --git a/src/main/tool/agentTools/agentToolManager.ts b/src/main/tool/agentTools/agentToolManager.ts index f70b223352..4279471ef9 100644 --- a/src/main/tool/agentTools/agentToolManager.ts +++ b/src/main/tool/agentTools/agentToolManager.ts @@ -2481,13 +2481,15 @@ export class AgentToolManager { let previewData: string | undefined let dispatchCommitFailed = false try { - const cachedPreviewData = await this.dependencies.cacheImage(dataUrl) + const cachedPreviewData = await this.dependencies.cacheImage(dataUrl, { signal }) if (cachedPreviewData && !cachedPreviewData.startsWith('data:image/')) { previewData = cachedPreviewData } } catch (error) { + throwIfAbortRequested(signal) logger.warn('[AgentToolManager] Failed to cache image preview', { filePath, error }) } + throwIfAbortRequested(signal) const imagePreviews: ToolCallImagePreview[] = [ { id: 'file_read-1', diff --git a/test/main/lib/toolCallImagePreviews.test.ts b/test/main/lib/toolCallImagePreviews.test.ts index 4864a45930..9286614c96 100644 --- a/test/main/lib/toolCallImagePreviews.test.ts +++ b/test/main/lib/toolCallImagePreviews.test.ts @@ -403,31 +403,54 @@ describe('cacheToolCallImagePreviews', () => { expect(previews).toEqual([{ ...input[0], data: 'imgcache://cached.png' }]) }) - it('dedupes identical payloads and caps cache writes', async () => { - const cacheImage = vi.fn(async (data: string) => `imgcache://cached-${data.length}.png`) - const duplicate = { - id: 'tool_output-1', - data: 'data:image/png;base64,QUFBQQ==', - mimeType: 'image/png', - source: 'tool_output' as const - } - const input = [ - duplicate, - { ...duplicate, id: 'tool_output-dup' }, - ...Array.from({ length: 6 }, (_, index) => ({ - id: `tool_output-${index + 2}`, - data: `data:image/png;base64,${'A'.repeat(index + 1)}`, + it.each([true, false])( + 'dedupes and caps inline previews when caching succeeds: %s', + async (succeeds) => { + const cacheImage = vi.fn(async (data: string) => + succeeds ? `imgcache://cached-${data.length}.png` : data + ) + const duplicate = { + id: 'tool_output-1', + data: 'data:image/png;base64,QUFBQQ==', mimeType: 'image/png', source: 'tool_output' as const - })) - ] + } + const input = [ + duplicate, + { ...duplicate, id: 'tool_output-dup' }, + ...Array.from({ length: 6 }, (_, index) => ({ + id: `tool_output-${index + 2}`, + data: `data:image/png;base64,${'A'.repeat(index + 1)}`, + mimeType: 'image/png', + source: 'tool_output' as const + })) + ] + + const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + + expect(cacheImage).toHaveBeenCalledTimes(4) + expect(previews).toHaveLength(4) + expect(previews.find((preview) => preview.id === 'tool_output-dup')).toBeUndefined() + expect( + previews.every((preview) => + preview.data?.startsWith(succeeds ? 'imgcache://' : 'data:image/') + ) + ).toBe(true) + } + ) - const previews = await cacheToolCallImagePreviews({ imagePreviews: input, cacheImage }) + it('drops duplicate cache results without restoring inline data', async () => { + const previews = await cacheToolCallImagePreviews({ + imagePreviews: ['AAAA', 'BBBB'].map((data, index) => ({ + id: `image-${index}`, + data: `data:image/png;base64,${data}`, + mimeType: 'image/png', + source: 'tool_output' + })), + cacheImage: async () => 'imgcache://same.png' + }) - expect(cacheImage).toHaveBeenCalledTimes(4) - expect(previews).toHaveLength(7) - expect(previews.find((preview) => preview.id === 'tool_output-dup')).toBeUndefined() - expect(previews.filter((preview) => preview.data?.startsWith('imgcache://'))).toHaveLength(4) - expect(previews.filter((preview) => preview.data?.startsWith('data:image/'))).toHaveLength(3) + expect(previews).toHaveLength(1) + expect(previews[0].data).toBe('imgcache://same.png') }) }) diff --git a/test/main/platform/imageCache.test.ts b/test/main/platform/imageCache.test.ts index 6e64910f27..3c20e356c0 100644 --- a/test/main/platform/imageCache.test.ts +++ b/test/main/platform/imageCache.test.ts @@ -21,6 +21,7 @@ vi.mock('electron', () => ({ vi.mock('axios', () => ({ default: axiosMock })) import { cacheImage, resolveCachedImageDataUrl } from '@/platform/imageCache' +import { cacheToolCallImagePreviews } from '@/lib/toolCallImagePreviews' describe('imageCache', () => { const tempDirectories: string[] = [] @@ -42,6 +43,40 @@ describe('imageCache', () => { vi.mocked(console.error).mockRestore() }) + it('preserves distinct valid images whose base64 differs only by case', async () => { + const black = Buffer.alloc(58) + black.write('BM') + black.writeUInt32LE(58, 2) + black.writeUInt32LE(54, 10) + black.writeUInt32LE(40, 14) + black.writeInt32LE(1, 18) + black.writeInt32LE(1, 22) + black.writeUInt16LE(1, 26) + black.writeUInt16LE(24, 28) + black.writeUInt32LE(4, 34) + const blue = Buffer.from(black) + blue[54] = 104 + const dataUrls = [black, blue].map((data) => `data:image/bmp;base64,${data.toString('base64')}`) + expect(dataUrls[0]).not.toBe(dataUrls[1]) + expect(dataUrls[0].toLowerCase()).toBe(dataUrls[1].toLowerCase()) + + const previews = await cacheToolCallImagePreviews({ + imagePreviews: dataUrls.map((data, index) => ({ + id: `image-${index}`, + data, + mimeType: 'image/bmp', + source: 'tool_output' + })), + cacheImage + }) + + expect(previews).toHaveLength(2) + expect(previews[0].data).not.toBe(previews[1].data) + await expect( + Promise.all(previews.map((preview) => resolveCachedImageDataUrl(preview.data!))) + ).resolves.toEqual(dataUrls) + }) + it('blocks a public image redirect to a private-network address', async () => { const sourceUrl = 'https://8.8.8.8/generated.png' axiosMock.mockResolvedValueOnce({ @@ -80,32 +115,35 @@ describe('imageCache', () => { await expect(fs.readdir(path.join(electronMock.userDataPath, 'images'))).resolves.toEqual([]) }) - it('bounds network responses and preserves supported MIME types', async () => { - axiosMock.mockResolvedValueOnce({ - status: 200, - headers: { 'content-type': 'image/avif' }, - data: Buffer.from('0000001866747970617669660000000061766966', 'hex') - }) - - const cached = await cacheImage('http://127.0.0.1/generated.avif', { - allowPrivateNetwork: true - }) + it.each(['http', 'HTTP'])( + 'caches %s URLs with bounded responses and supported MIME types', + async (scheme) => { + axiosMock.mockResolvedValueOnce({ + status: 200, + headers: { 'content-type': 'image/avif' }, + data: Buffer.from('0000001866747970617669660000000061766966', 'hex') + }) - expect(cached).toMatch(/^imgcache:\/\/.+\.avif$/) - expect(axiosMock).toHaveBeenCalledWith( - expect.objectContaining({ - maxRedirects: 0, - maxContentLength: 32 * 1024 * 1024, - maxBodyLength: 32 * 1024 * 1024, - signal: expect.any(AbortSignal) + const cached = await cacheImage(`${scheme}://127.0.0.1/generated.avif`, { + allowPrivateNetwork: true }) - ) - await expect( - fs.readFile( - path.join(electronMock.userDataPath, 'images', cached.slice('imgcache://'.length)) + + expect(cached).toMatch(/^imgcache:\/\/.+\.avif$/) + expect(axiosMock).toHaveBeenCalledWith( + expect.objectContaining({ + maxRedirects: 0, + maxContentLength: 32 * 1024 * 1024, + maxBodyLength: 32 * 1024 * 1024, + signal: expect.any(AbortSignal) + }) ) - ).resolves.toEqual(Buffer.from('0000001866747970617669660000000061766966', 'hex')) - }) + await expect( + fs.readFile( + path.join(electronMock.userDataPath, 'images', cached.slice('imgcache://'.length)) + ) + ).resolves.toEqual(Buffer.from('0000001866747970617669660000000061766966', 'hex')) + } + ) it('does not cache non-image HTTP responses', async () => { const sourceUrl = 'http://127.0.0.1/generated.png' diff --git a/test/main/provider/aiSdkStreamAdapter.test.ts b/test/main/provider/aiSdkStreamAdapter.test.ts index f7928583cf..49562ffa4f 100644 --- a/test/main/provider/aiSdkStreamAdapter.test.ts +++ b/test/main/provider/aiSdkStreamAdapter.test.ts @@ -541,7 +541,7 @@ describe('AI SDK stream adapter', () => { { supportsNativeTools: true, cacheImage } ) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,ZmFrZQ==') + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,ZmFrZQ==', { signal: undefined }) expect(events[0]).toEqual({ type: 'image_data', image_data: { @@ -552,6 +552,29 @@ describe('AI SDK stream adapter', () => { expect(events[2]).toEqual({ type: 'stop', stop_reason: 'complete' }) }) + it.each([true, false])( + 'does not emit an image after cancellation when caching rejects: %s', + async (rejects) => { + const controller = new AbortController() + const cacheImage = vi.fn(async () => { + controller.abort() + if (rejects) throw controller.signal.reason + return 'imgcache://late.png' + }) + + await expect( + collectEvents([{ type: 'file', file: { mediaType: 'image/png', base64: 'ZmFrZQ==' } }], { + supportsNativeTools: true, + cacheImage, + signal: controller.signal + }) + ).rejects.toMatchObject({ name: 'AbortError' }) + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,ZmFrZQ==', { + signal: controller.signal + }) + } + ) + it('falls back to the original image data url when image caching fails', async () => { const cacheImage = vi.fn().mockRejectedValue(new Error('cache failed')) const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}) @@ -579,7 +602,9 @@ describe('AI SDK stream adapter', () => { { supportsNativeTools: true, cacheImage } ) - expect(cacheImage).toHaveBeenCalledWith('data:image/jpeg;base64,YWJjZA==') + expect(cacheImage).toHaveBeenCalledWith('data:image/jpeg;base64,YWJjZA==', { + signal: undefined + }) expect(warnSpy).toHaveBeenCalled() expect(events[0]).toEqual({ type: 'image_data', diff --git a/test/main/tool/agentTools/agentImageGenerationTool.test.ts b/test/main/tool/agentTools/agentImageGenerationTool.test.ts index 627b06d43e..8f991c3650 100644 --- a/test/main/tool/agentTools/agentImageGenerationTool.test.ts +++ b/test/main/tool/agentTools/agentImageGenerationTool.test.ts @@ -150,36 +150,42 @@ describe('Agent image generation tool', () => { expect(result.rawData.toolResult.ok).toBe(true) }) - it('caches raw base64 image data to an imgcache reference', async () => { - const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') - manager = buildManager(cacheImage) - generateImageStandalone.mockResolvedValue({ - providerId: 'openai', - modelId: 'gpt-image-1', - images: [{ data: 'aGVsbG8=', mimeType: 'image/png' }] - }) - - const result = (await manager.callTool( - IMAGE_GENERATE_TOOL_NAME, - { prompt: 'A warm sunset over the ocean' }, - 'conv-1' - )) as any + it.each(['aGVsbG8=', 'DATA:IMAGE/PNG;BASE64,aGVsbG8='])( + 'caches generated image data without corrupting %s', + async (data) => { + const cacheImage = vi.fn().mockResolvedValue('imgcache://cached.png') + manager = buildManager(cacheImage) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data, mimeType: 'image/png' }] + }) - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { - signal: undefined, - allowPrivateNetwork: false - }) - expect(result.rawData.imagePreviews).toEqual([ - { - id: 'generated-image-1', - data: 'imgcache://cached.png', - mimeType: 'image/png', - title: 'Generated image 1', - source: 'tool_output' - } - ]) - expect(result.rawData.toolResult.ok).toBe(true) - }) + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(cacheImage).toHaveBeenCalledWith( + data.includes(':') ? data : `data:image/png;base64,${data}`, + { + signal: undefined, + allowPrivateNetwork: false + } + ) + expect(result.rawData.imagePreviews).toEqual([ + { + id: 'generated-image-1', + data: 'imgcache://cached.png', + mimeType: 'image/png', + title: 'Generated image 1', + source: 'tool_output' + } + ]) + expect(result.rawData.toolResult.ok).toBe(true) + } + ) it('fails the tool call when a generated HTTP image URL cannot be cached', async () => { manager = buildManager(vi.fn(async (data: string) => data)) @@ -233,60 +239,68 @@ describe('Agent image generation tool', () => { expect(result.rawData.toolResult.ok).toBe(true) }) - it('fails the tool call above the inline base64 character limit even when decoded bytes are smaller', async () => { - manager = buildManager(vi.fn(async (data: string) => data)) - // Encoded character length is just above the 2 MiB limit while the decoded payload is only - // ~1.5 MiB — the limit is enforced on the encoded base64 payload, not decoded byte size. - const payload = 'A'.repeat(2 * 1024 * 1024 + 1) - generateImageStandalone.mockResolvedValue({ - providerId: 'openai', - modelId: 'gpt-image-1', - images: [{ data: `data:image/png;base64,${payload}`, mimeType: 'image/png' }] - }) - - const result = (await manager.callTool( - IMAGE_GENERATE_TOOL_NAME, - { prompt: 'A warm sunset over the ocean' }, - 'conv-1' - )) as any - - expect(result.rawData.isError).toBe(true) - expect(result.rawData.toolResult.error).toMatchObject({ - code: 'IMAGE_GENERATION_FAILED', - recoverable: true - }) - }) + it.each(['data:image/png;base64,', 'DATA:IMAGE/PNG;BASE64,'])( + 'rejects oversized encoded payloads with prefix %s', + async (prefix) => { + manager = buildManager(vi.fn(async (data: string) => data)) + // Encoded character length is just above the 2 MiB limit while the decoded payload is only + // ~1.5 MiB — the limit is enforced on the encoded base64 payload, not decoded byte size. + const payload = 'A'.repeat(2 * 1024 * 1024 + 1) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: `${prefix}${payload}`, mimeType: 'image/png' }] + }) - it('does not return a success result when the cache write is aborted', async () => { - const controller = new AbortController() - const cacheImage = vi.fn(async () => { - controller.abort() - throw new DOMException('The operation was aborted.', 'AbortError') - }) - manager = buildManager(cacheImage) - generateImageStandalone.mockResolvedValue({ - providerId: 'openai', - modelId: 'gpt-image-1', - images: [{ data: 'aGVsbG8=', mimeType: 'image/png' }] - }) + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1' + )) as any + + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + } + ) + + it.each(['reject', 'pending', 'resolve'])( + 'fails promptly when an aborted cache write would %s', + async (outcome) => { + const controller = new AbortController() + const cacheImage = vi.fn(async () => { + controller.abort() + if (outcome === 'pending') return new Promise(() => {}) + if (outcome === 'resolve') return 'imgcache://late.png' + throw new DOMException('The operation was aborted.', 'AbortError') + }) + manager = buildManager(cacheImage) + generateImageStandalone.mockResolvedValue({ + providerId: 'openai', + modelId: 'gpt-image-1', + images: [{ data: 'aGVsbG8=', mimeType: 'image/png' }] + }) - const result = (await manager.callTool( - IMAGE_GENERATE_TOOL_NAME, - { prompt: 'A warm sunset over the ocean' }, - 'conv-1', - { signal: controller.signal } - )) as any + const result = (await manager.callTool( + IMAGE_GENERATE_TOOL_NAME, + { prompt: 'A warm sunset over the ocean' }, + 'conv-1', + { signal: controller.signal } + )) as any - expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { - signal: controller.signal, - allowPrivateNetwork: false - }) - expect(result.rawData.isError).toBe(true) - expect(result.rawData.toolResult.error).toMatchObject({ - code: 'IMAGE_GENERATION_FAILED', - recoverable: true - }) - }) + expect(cacheImage).toHaveBeenCalledWith('data:image/png;base64,aGVsbG8=', { + signal: controller.signal, + allowPrivateNetwork: false + }) + expect(result.rawData.isError).toBe(true) + expect(result.rawData.toolResult.error).toMatchObject({ + code: 'IMAGE_GENERATION_FAILED', + recoverable: true + }) + } + ) it('fails the tool call when a large image cannot be cached', async () => { manager = buildManager(vi.fn(async (data: string) => data)) From 3f59ad24108cb1b2976e1321af1a03131604c904 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Thu, 10 Sep 2026 11:29:11 +0800 Subject: [PATCH 08/15] feat(agent): reap orphaned child processes Add a shared child-process registry (record/clear/reap primitives) that persists launch records with pid, owner pid, command-line fingerprint and start timestamp. On startup each subsystem reaps stale records only after identity attestation (process start time plus command-line fingerprint) to guard against pid reuse, and skips records still owned by a live process. Wire it into background exec sessions (utility host boot reaping, record on session start, clear on finalization), MCP stdio servers (record after connect, clear on close/force-terminate, reap during McpService.initialize) and ACP agent processes (record on spawn, clear on kill/exit, reap on manager construction). --- .../agent/acp/runtime/acpProcessManager.ts | 27 ++ .../process/backgroundExecSessionManager.ts | 12 + .../process/backgroundExecUtilityHost.ts | 5 + .../shared/process/childProcessRegistry.ts | 401 ++++++++++++++++++ src/main/mcp/index.ts | 12 + src/main/mcp/mcpClient.ts | 16 + .../acp/runtime/acpProcessManager.test.ts | 109 +++++ .../acpProcessManagerCapabilities.test.ts | 10 + .../backgroundExecSessionManager.test.ts | 52 ++- .../process/backgroundExecUtilityHost.test.ts | 44 +- .../process/childProcessRegistry.test.ts | 308 ++++++++++++++ test/main/mcp/mcpClient.test.ts | 57 +++ test/main/mcp/mcpService.test.ts | 11 + 13 files changed, 1060 insertions(+), 4 deletions(-) create mode 100644 src/main/agent/shared/process/childProcessRegistry.ts create mode 100644 test/main/agent/shared/process/childProcessRegistry.test.ts diff --git a/src/main/agent/acp/runtime/acpProcessManager.ts b/src/main/agent/acp/runtime/acpProcessManager.ts index 1807bc4446..bd8af02c87 100644 --- a/src/main/agent/acp/runtime/acpProcessManager.ts +++ b/src/main/agent/acp/runtime/acpProcessManager.ts @@ -29,6 +29,7 @@ import { mergeCommandEnvironment, setPathEntriesOnEnv } from '@/agent/shared/process/shellEnvHelper' +import { childProcessRegistry } from '@/agent/shared/process/childProcessRegistry' import { RuntimeHelper } from '@/lib/runtimeHelper' import { ToolchainService } from '@/toolchains' import { @@ -271,6 +272,7 @@ export class AcpProcessManager implements AgentProcessManager() private readonly initializingChildren = new Set() private readonly terminatedChildren = new WeakSet() + private readonly launchRecordIds = new WeakMap() private readonly disposedHandles = new WeakSet() private readonly shutdownController = new AbortController() private shuttingDown = false @@ -284,6 +286,9 @@ export class AcpProcessManager implements AgentProcessManager { + console.warn('[ACP] Failed to reap stale agent processes:', error) + }) } getTerminalSnapshot(terminalId: string): schema.TerminalOutputResponse | null { @@ -1258,6 +1263,7 @@ export class AcpProcessManager implements AgentProcessManager void> { diff --git a/src/main/agent/shared/process/backgroundExecSessionManager.ts b/src/main/agent/shared/process/backgroundExecSessionManager.ts index b557b7bac7..533e32e8c5 100644 --- a/src/main/agent/shared/process/backgroundExecSessionManager.ts +++ b/src/main/agent/shared/process/backgroundExecSessionManager.ts @@ -14,6 +14,7 @@ import { } from './shellOutputEncoding' import { describeSpawnFailure, resolveUsableSpawnCwd } from './spawnGuard' import { terminateProcessTree } from './processTree' +import { childProcessRegistry } from './childProcessRegistry' import { resolveSessionDir } from '@/agent/shared/storage/sessionPaths' import { assertSkillExecutionPackageTreeIntact, @@ -358,6 +359,16 @@ export class BackgroundExecSessionManager { } this.sessions.get(conversationId)!.set(sessionId, session) + if (typeof child.pid === 'number') { + childProcessRegistry.record({ + subsystem: 'background-exec', + recordId: sessionId, + pid: child.pid, + commandLine: [executable, ...args], + cwd: spawnCwd + }) + } + logger.info(`[BackgroundExec] Started session ${sessionId} for conversation ${conversationId}`) return { sessionId, status: 'running' } @@ -1013,6 +1024,7 @@ export class BackgroundExecSessionManager { clearTimeout(session.closeWatchdogId) session.closeWatchdogId = undefined } + childProcessRegistry.clear('background-exec', session.sessionId) try { session.flushOutputDecoders?.() await session.outputWriteQueue.catch((error) => { diff --git a/src/main/agent/shared/process/backgroundExecUtilityHost.ts b/src/main/agent/shared/process/backgroundExecUtilityHost.ts index 6c434bb2a1..7ae68f9528 100644 --- a/src/main/agent/shared/process/backgroundExecUtilityHost.ts +++ b/src/main/agent/shared/process/backgroundExecUtilityHost.ts @@ -3,6 +3,8 @@ import { type BackgroundExecRpcRequest, type BackgroundExecRpcResponse } from './backgroundExecSessionManager' +import { childProcessRegistry } from './childProcessRegistry' +import logger from './backgroundExecLogger' const EXEC_UTILITY_HOST_ARG = '--deepchat-exec-utility-host' @@ -105,6 +107,9 @@ export function runBackgroundExecUtilityHostIfRequested(): boolean { } const manager = new BackgroundExecSessionManager() + void childProcessRegistry.reapStaleOnce('background-exec').catch((error) => { + logger.warn('[BackgroundExec] Failed to reap stale child processes:', error) + }) const keepAliveIntervalId = setInterval(() => {}, 2 ** 31 - 1) parentPort.start?.() diff --git a/src/main/agent/shared/process/childProcessRegistry.ts b/src/main/agent/shared/process/childProcessRegistry.ts new file mode 100644 index 0000000000..ea8ad05f12 --- /dev/null +++ b/src/main/agent/shared/process/childProcessRegistry.ts @@ -0,0 +1,401 @@ +import { spawn } from 'child_process' +import { createHash } from 'crypto' +import fs from 'fs' +import os from 'os' +import path from 'path' +import { terminateProcessTreeByPid } from './processTree' + +const RECORD_VERSION = 1 +const PROCESS_START_TOLERANCE_MS = 60_000 +const DEFAULT_MAX_RECORD_AGE_MS = 7 * 24 * 60 * 60 * 1000 + +export type ChildProcessSubsystem = 'background-exec' | 'mcp-stdio' | 'acp-agent' + +export interface ChildProcessLaunchRecord { + version: number + subsystem: string + recordId: string + pid: number + ownerPid: number + commandLine: string[] + cwd?: string + recordedAt: number +} + +export interface ObservedProcessIdentity { + alive: boolean + commandLine?: string + startedAtMs?: number +} + +export type ChildProcessAttester = ( + record: ChildProcessLaunchRecord, + observed: ObservedProcessIdentity +) => boolean + +export interface ReapStaleChildProcessesResult { + reaped: string[] + refused: string[] + cleared: string[] + skipped: string[] +} + +export interface ReapStaleChildProcessesOptions { + attester?: ChildProcessAttester + excludeRecordIds?: ReadonlySet + maxRecordAgeMs?: number +} + +export interface ChildProcessRegistryOptions { + rootDir?: string + now?: () => number + isAlive?: (pid: number) => boolean + observe?: (pid: number) => Promise + terminate?: (pid: number) => Promise + log?: (message: string, ...args: unknown[]) => void +} + +function defaultRegistryRoot(): string { + return path.join(os.homedir(), '.deepchat', 'child-processes') +} + +function defaultIsAlive(pid: number): boolean { + try { + process.kill(pid, 0) + return true + } catch (error) { + return (error as NodeJS.ErrnoException).code === 'EPERM' + } +} + +async function runAndCapture( + command: string, + args: string[] +): Promise<{ code: number | null; stdout: string }> { + return await new Promise((resolve) => { + let stdout = '' + try { + const child = spawn(command, args, { + stdio: ['ignore', 'pipe', 'ignore'], + ...(process.platform === 'win32' ? { windowsHide: true } : {}) + }) + child.stdout?.on('data', (chunk: Buffer | string) => { + stdout += chunk.toString() + }) + child.on('error', () => resolve({ code: null, stdout: '' })) + child.on('close', (code) => resolve({ code, stdout })) + } catch { + resolve({ code: null, stdout: '' }) + } + }) +} + +async function observePosix(pid: number): Promise { + const { stdout } = await runAndCapture('ps', ['-p', `${pid}`, '-o', 'etimes=', '-o', 'command=']) + const trimmed = stdout.trim() + if (!trimmed) { + return { alive: false } + } + const match = /^(\d+)\s+([\s\S]*)$/.exec(trimmed) + if (!match) { + return { alive: true } + } + const elapsedSeconds = Number.parseInt(match[1], 10) + if (!Number.isFinite(elapsedSeconds)) { + return { alive: true } + } + return { + alive: true, + commandLine: match[2], + startedAtMs: Date.now() - elapsedSeconds * 1000 + } +} + +async function observeWindows(pid: number): Promise { + const command = [ + 'Get-CimInstance Win32_Process -Filter "ProcessId=', + `${pid}`, + '" | Select-Object -Property CreationDate,CommandLine | ConvertTo-Json -Compress' + ].join('') + const { stdout } = await runAndCapture('powershell', [ + '-NoProfile', + '-NonInteractive', + '-Command', + command + ]) + const trimmed = stdout.trim() + if (!trimmed) { + return { alive: true } + } + try { + const parsed = JSON.parse(trimmed) as { CreationDate?: string; CommandLine?: string | null } + const startedAtMs = parsed.CreationDate ? Date.parse(parsed.CreationDate) : Number.NaN + return { + alive: true, + ...(parsed.CommandLine ? { commandLine: parsed.CommandLine } : {}), + ...(Number.isFinite(startedAtMs) ? { startedAtMs } : {}) + } + } catch { + return { alive: true } + } +} + +async function defaultObserve(pid: number): Promise { + if (!defaultIsAlive(pid)) { + return { alive: false } + } + return process.platform === 'win32' ? await observeWindows(pid) : await observePosix(pid) +} + +export function defaultChildProcessAttester( + record: ChildProcessLaunchRecord, + observed: ObservedProcessIdentity +): boolean { + if (!observed.alive) { + return false + } + if ( + typeof observed.startedAtMs !== 'number' || + Math.abs(observed.startedAtMs - record.recordedAt) > PROCESS_START_TOLERANCE_MS + ) { + return false + } + if (!observed.commandLine) { + return false + } + return record.commandLine.every( + (segment) => segment.length > 0 && observed.commandLine!.includes(segment) + ) +} + +function isLaunchRecord(value: unknown): value is ChildProcessLaunchRecord { + if (!value || typeof value !== 'object') { + return false + } + const record = value as Record + return ( + typeof record.recordId === 'string' && + typeof record.subsystem === 'string' && + Number.isSafeInteger(record.pid) && + (record.pid as number) > 0 && + Number.isSafeInteger(record.ownerPid) && + Array.isArray(record.commandLine) && + record.commandLine.every((segment) => typeof segment === 'string') && + Number.isFinite(record.recordedAt) + ) +} + +function recordFileName(recordId: string): string { + const sanitized = recordId.replace(/[^a-zA-Z0-9_.-]/g, '_').slice(0, 80) || 'record' + const fingerprint = createHash('sha1').update(recordId).digest('hex').slice(0, 8) + return `${sanitized}_${fingerprint}.json` +} + +export class ChildProcessRegistry { + private readonly rootDir: string + private readonly now: () => number + private readonly isAlive: (pid: number) => boolean + private readonly observe: (pid: number) => Promise + private readonly terminate: (pid: number) => Promise + private readonly log: (message: string, ...args: unknown[]) => void + private readonly reapedSubsystems = new Set() + private readonly inflightReaps = new Map>() + + constructor(options: ChildProcessRegistryOptions = {}) { + this.rootDir = options.rootDir ?? defaultRegistryRoot() + this.now = options.now ?? Date.now + this.isAlive = options.isAlive ?? defaultIsAlive + this.observe = options.observe ?? defaultObserve + this.terminate = + options.terminate ?? ((pid: number) => terminateProcessTreeByPid(pid, { graceMs: 2000 })) + this.log = options.log ?? ((message, ...args) => console.warn(message, ...args)) + } + + record(entry: { + subsystem: ChildProcessSubsystem + recordId: string + pid: number + commandLine: string[] + cwd?: string + }): void { + if (!Number.isSafeInteger(entry.pid) || entry.pid <= 0 || !entry.recordId) { + return + } + try { + const record: ChildProcessLaunchRecord = { + version: RECORD_VERSION, + subsystem: entry.subsystem, + recordId: entry.recordId, + pid: entry.pid, + ownerPid: process.pid, + commandLine: [...entry.commandLine], + ...(entry.cwd ? { cwd: entry.cwd } : {}), + recordedAt: this.now() + } + const filePath = this.recordPath(entry.subsystem, entry.recordId) + fs.mkdirSync(path.dirname(filePath), { recursive: true }) + const tempPath = `${filePath}.${process.pid}.tmp` + fs.writeFileSync(tempPath, JSON.stringify(record), 'utf-8') + fs.renameSync(tempPath, filePath) + } catch (error) { + this.log(`[ChildProcessRegistry] Failed to record launch ${entry.recordId}:`, error) + } + } + + clear(subsystem: ChildProcessSubsystem, recordId: string): void { + if (!recordId) { + return + } + try { + fs.unlinkSync(this.recordPath(subsystem, recordId)) + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== 'ENOENT') { + this.log(`[ChildProcessRegistry] Failed to clear launch record ${recordId}:`, error) + } + } + } + + list(subsystem: ChildProcessSubsystem): ChildProcessLaunchRecord[] { + const subsystemDir = path.join(this.rootDir, subsystem) + let entries: fs.Dirent[] + try { + entries = fs.readdirSync(subsystemDir, { withFileTypes: true }) + } catch { + return [] + } + + const records: ChildProcessLaunchRecord[] = [] + for (const entry of entries) { + if (!entry.isFile() || !entry.name.endsWith('.json')) { + continue + } + const filePath = path.join(subsystemDir, entry.name) + try { + const parsed: unknown = JSON.parse(fs.readFileSync(filePath, 'utf-8')) + if (isLaunchRecord(parsed)) { + records.push(parsed) + } else { + fs.unlinkSync(filePath) + } + } catch { + try { + fs.unlinkSync(filePath) + } catch { + // Unreadable record files are removed best-effort. + } + } + } + return records + } + + async reapStale( + subsystem: ChildProcessSubsystem, + options: ReapStaleChildProcessesOptions = {} + ): Promise { + const inflight = this.inflightReaps.get(subsystem) + if (inflight) { + return await inflight + } + const reapPromise = this.reapStaleRecords(subsystem, options).finally(() => { + this.inflightReaps.delete(subsystem) + }) + this.inflightReaps.set(subsystem, reapPromise) + return await reapPromise + } + + async reapStaleOnce( + subsystem: ChildProcessSubsystem, + options: ReapStaleChildProcessesOptions = {} + ): Promise { + if (this.reapedSubsystems.has(subsystem)) { + return null + } + this.reapedSubsystems.add(subsystem) + return await this.reapStale(subsystem, options) + } + + private async reapStaleRecords( + subsystem: ChildProcessSubsystem, + options: ReapStaleChildProcessesOptions + ): Promise { + const result: ReapStaleChildProcessesResult = { + reaped: [], + refused: [], + cleared: [], + skipped: [] + } + const attester = options.attester ?? defaultChildProcessAttester + const maxRecordAgeMs = options.maxRecordAgeMs ?? DEFAULT_MAX_RECORD_AGE_MS + + for (const record of this.list(subsystem)) { + if (options.excludeRecordIds?.has(record.recordId)) { + result.skipped.push(record.recordId) + continue + } + if (this.now() - record.recordedAt > maxRecordAgeMs) { + this.clear(subsystem, record.recordId) + result.cleared.push(record.recordId) + continue + } + if (record.pid === process.pid) { + result.refused.push(record.recordId) + continue + } + if (!this.isAlive(record.pid)) { + this.clear(subsystem, record.recordId) + result.cleared.push(record.recordId) + continue + } + if (record.ownerPid !== process.pid && this.isAlive(record.ownerPid)) { + result.skipped.push(record.recordId) + continue + } + + const observed = await this.observe(record.pid) + if (!observed.alive) { + this.clear(subsystem, record.recordId) + result.cleared.push(record.recordId) + continue + } + + let attested = false + try { + attested = attester(record, observed) + } catch (error) { + this.log(`[ChildProcessRegistry] Attester failed for ${record.recordId}:`, error) + } + + if (!attested) { + result.refused.push(record.recordId) + this.log( + `[ChildProcessRegistry] Refusing unattested cleanup of pid ${record.pid} (${record.recordId})` + ) + if (observed.commandLine !== undefined && observed.startedAtMs !== undefined) { + // Identity was fully checked and did not match: the recorded process + // is gone and the pid now belongs to an unrelated process. + this.clear(subsystem, record.recordId) + } + continue + } + + const terminated = await this.terminate(record.pid).catch((error) => { + this.log(`[ChildProcessRegistry] Failed to terminate pid ${record.pid}:`, error) + return false + }) + if (terminated) { + this.clear(subsystem, record.recordId) + result.reaped.push(record.recordId) + } else { + result.refused.push(record.recordId) + } + } + + return result + } + + private recordPath(subsystem: string, recordId: string): string { + return path.join(this.rootDir, subsystem, recordFileName(recordId)) + } +} + +export const childProcessRegistry = new ChildProcessRegistry() diff --git a/src/main/mcp/index.ts b/src/main/mcp/index.ts index bf0fa08111..ee59d1fe8c 100644 --- a/src/main/mcp/index.ts +++ b/src/main/mcp/index.ts @@ -59,6 +59,7 @@ import { McpAppHost } from './apps/appHost' import { hasMcpIdentityBearingChange } from './serverIdentity' import type { CacheImageOptions } from '@/platform/imageCache' import { awaitWithAbort } from '@/lib/awaitWithAbort' +import { childProcessRegistry } from '@/agent/shared/process/childProcessRegistry' type McpToolAccessContext = { enabledTools?: string[] @@ -333,6 +334,17 @@ export class McpService implements McpServicePort { } try { + void childProcessRegistry + .reapStaleOnce('mcp-stdio') + .then((result) => { + if (result && result.reaped.length > 0) { + logger.info(`[MCP] Reaped stale stdio server processes: ${result.reaped.join(', ')}`) + } + }) + .catch((error) => + console.error('[MCP] Failed to reap stale stdio server processes:', error) + ) + // Load configuration const [servers, enabledServers, mcpEnabled] = await Promise.all([ this.mcpSettings.getMcpServers(), diff --git a/src/main/mcp/mcpClient.ts b/src/main/mcp/mcpClient.ts index 0717047a3c..e4ab1072ac 100644 --- a/src/main/mcp/mcpClient.ts +++ b/src/main/mcp/mcpClient.ts @@ -40,6 +40,7 @@ import { RuntimeHelper } from '@/lib/runtimeHelper' import { ToolchainService } from '@/toolchains' import { getPathEntriesFromEnv, setPathEntriesOnEnv } from '@/agent/shared/process/shellEnvHelper' import { terminateProcessTreeByPid } from '@/agent/shared/process/processTree' +import { childProcessRegistry } from '@/agent/shared/process/childProcessRegistry' import { awaitWithAbort } from '@/lib/awaitWithAbort' import type { McpOAuthManager } from './mcpOAuthManager' import type { ChatMessage } from '@shared/types/core/chat-message' @@ -256,6 +257,7 @@ export class McpClient { private isConnected: boolean = false private connectionTimeout: NodeJS.Timeout | null = null private stdioPidForShutdown?: number + private stdioCommandLine?: string[] private connectPromise: Promise | null = null private startupAttempt = 0 private lifecycleStatus: McpServerLifecycleStatus = 'stopped' @@ -641,6 +643,7 @@ export class McpClient { ? resolveBinding(this.serverConfig.cwd) : undefined const pluginRoot = this.serverConfig.source === 'plugin' ? env.PLUGIN_ROOT : undefined + this.stdioCommandLine = [command, ...args] this.transport = new StdioClientTransport({ command, args, @@ -788,6 +791,16 @@ export class McpClient { } console.info(`MCP server ${this.serverName} connected successfully`) + const stdioPid = this.getStdioPid() + if (stdioPid && this.stdioCommandLine) { + childProcessRegistry.record({ + subsystem: 'mcp-stdio', + recordId: this.serverName, + pid: stdioPid, + commandLine: this.stdioCommandLine + }) + } + this.emitServerStatusChanged('connected', { phase, attempt }) } catch (error) { // 清除超时 @@ -869,6 +882,7 @@ export class McpClient { // 重置状态 this.client = null this.isConnected = false + this.stdioCommandLine = undefined if (options.emitStopped) { this.emitServerStatusChanged('stopped', { reason: 'shutdown' }) @@ -891,6 +905,7 @@ export class McpClient { try { await terminateProcessTreeByPid(pid, { graceMs: 2000 }) console.warn(`[MCP] Force terminated stdio process tree for ${this.serverName}: ${reason}`) + childProcessRegistry.clear('mcp-stdio', this.serverName) return true } catch (error) { console.warn( @@ -913,6 +928,7 @@ export class McpClient { } catch (error) { console.error(`Failed to terminate MCP stdio process tree for ${this.serverName}:`, error) } + childProcessRegistry.clear('mcp-stdio', this.serverName) } if (this.stdioPidForShutdown === pid) { this.stdioPidForShutdown = undefined diff --git a/test/main/agent/acp/runtime/acpProcessManager.test.ts b/test/main/agent/acp/runtime/acpProcessManager.test.ts index a98d9a140a..5a4cd414d5 100644 --- a/test/main/agent/acp/runtime/acpProcessManager.test.ts +++ b/test/main/agent/acp/runtime/acpProcessManager.test.ts @@ -1,6 +1,7 @@ import { EventEmitter } from 'events' import * as fs from 'fs' import path from 'path' +import { PassThrough } from 'node:stream' import { beforeEach, describe, expect, it, vi } from 'vitest' import spawn from 'cross-spawn' import * as shellEnvHelper from '@/agent/shared/process/shellEnvHelper' @@ -12,6 +13,12 @@ import { ToolchainService } from '@/toolchains' const publishDeepchatEventMock = vi.hoisted(() => vi.fn()) +const childProcessRegistryMock = vi.hoisted(() => ({ + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) +})) + vi.mock('electron', () => ({ app: { getVersion: vi.fn(() => '0.0.0-test'), @@ -23,6 +30,10 @@ vi.mock('cross-spawn', () => ({ default: vi.fn() })) +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: childProcessRegistryMock +})) + vi.mock('@/agent/shared/process/shellEnvHelper', async (importOriginal) => { const actual = await importOriginal() return { @@ -1081,3 +1092,101 @@ describe('AcpProcessManager config cache fallback', () => { expect(onExit).toHaveBeenCalledTimes(1) }) }) + +describe('AcpProcessManager child process launch records', () => { + class MockStreamingChild extends EventEmitter { + stdout = new PassThrough() + stderr = new PassThrough() + stdin = new PassThrough() + pid = 1234 + killed = false + exitCode = null + signalCode = null + kill = vi.fn(() => true) + } + + const createManager = () => + new AcpProcessManager({ + publishEvent: publishDeepchatEventMock, + providerId: 'acp', + resolveLaunchSpec: vi.fn() + }) + + const agent = { id: 'agent-1', name: 'Agent One', command: 'agent' } + const launch = { + command: 'agent', + args: ['--acp'], + env: {}, + cwd: '/tmp/workspace' + } + + beforeEach(() => { + childProcessRegistryMock.record.mockClear() + childProcessRegistryMock.clear.mockClear() + childProcessRegistryMock.reapStaleOnce.mockClear() + vi.mocked(spawn).mockClear() + }) + + it('reaps stale agent process records on construction', () => { + createManager() + + expect(childProcessRegistryMock.reapStaleOnce).toHaveBeenCalledWith('acp-agent') + }) + + it('records the launch after spawning an agent process', () => { + const manager = createManager() + const child = new MockSpawnedChild() + vi.mocked(spawn).mockReturnValue(child as never) + + ;(manager as any).spawnAgentProcess(agent, launch) + + expect(childProcessRegistryMock.record).toHaveBeenCalledWith({ + subsystem: 'acp-agent', + recordId: 'agent-1:1234', + pid: 1234, + commandLine: ['agent', '--acp'], + cwd: '/tmp/workspace' + }) + }) + + it('clears the launch record when the child is killed', () => { + const manager = createManager() + const child = new MockSpawnedChild() + vi.mocked(spawn).mockReturnValue(child as never) + + const spawned = (manager as any).spawnAgentProcess(agent, launch) + Object.defineProperty(spawned, 'pid', { value: undefined }) + ;(manager as any).killChild(spawned) + + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('acp-agent', 'agent-1:1234') + }) + + it('clears the launch record when the agent process exits', async () => { + const manager = createManager() + const child = new MockStreamingChild() + vi.mocked(spawn).mockReturnValue(child as never) + + const spawned = (manager as any).spawnAgentProcess(agent, launch) + const initPromise = (manager as any).initializeSpawnedProcess( + spawned, + agent, + '/tmp/workspace', + { + agentId: 'agent-1', + source: 'manual', + distributionType: 'manual', + command: 'agent', + args: [], + env: {} + }, + 'signature', + launch + ) + const rejection = expect(initPromise).rejects.toThrow('exited during initialization') + + spawned.emit('exit', 1, null) + await rejection + + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('acp-agent', 'agent-1:1234') + }) +}) diff --git a/test/main/agent/acp/runtime/acpProcessManagerCapabilities.test.ts b/test/main/agent/acp/runtime/acpProcessManagerCapabilities.test.ts index c1f5fa6ef3..0c3cb6b2d3 100644 --- a/test/main/agent/acp/runtime/acpProcessManagerCapabilities.test.ts +++ b/test/main/agent/acp/runtime/acpProcessManagerCapabilities.test.ts @@ -43,6 +43,16 @@ vi.mock('electron', () => ({ } })) +const childProcessRegistryMock = vi.hoisted(() => ({ + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) +})) + +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: childProcessRegistryMock +})) + class MockChild extends EventEmitter { stdout = new PassThrough() stderr = new PassThrough() diff --git a/test/main/agent/shared/process/backgroundExecSessionManager.test.ts b/test/main/agent/shared/process/backgroundExecSessionManager.test.ts index fdb03245e3..aa4aaead6e 100644 --- a/test/main/agent/shared/process/backgroundExecSessionManager.test.ts +++ b/test/main/agent/shared/process/backgroundExecSessionManager.test.ts @@ -8,12 +8,18 @@ const { mockUtilityProcessFork, mockAssertPackageTree, mockCleanupPackageTree, - mockTerminateProcessTree + mockTerminateProcessTree, + registryMock } = vi.hoisted(() => ({ mockUtilityProcessFork: vi.fn(), mockAssertPackageTree: vi.fn(), mockCleanupPackageTree: vi.fn(), - mockTerminateProcessTree: vi.fn() + mockTerminateProcessTree: vi.fn(), + registryMock: { + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) + } })) vi.mock('child_process', () => ({ @@ -53,6 +59,10 @@ vi.mock('@/agent/shared/process/processTree', () => ({ terminateProcessTree: mockTerminateProcessTree })) +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: registryMock +})) + import { BackgroundExecSessionManager, backgroundExecSessionManager @@ -126,6 +136,8 @@ describe('BackgroundExecSessionManager', () => { mockAssertPackageTree.mockResolvedValue(undefined) mockCleanupPackageTree.mockResolvedValue(undefined) mockTerminateProcessTree.mockResolvedValue(true) + registryMock.record.mockReset() + registryMock.clear.mockReset() vi.spyOn(fs, 'existsSync').mockReturnValue(true) vi.spyOn(fs, 'statSync').mockImplementation((candidate) => String(candidate).includes('workspace') ? mockStats('directory') : mockStats('file') @@ -332,6 +344,42 @@ describe('BackgroundExecSessionManager', () => { }) }) + it('records the launch in the child process registry when a session starts', async () => { + const child = new MockChildProcess() + vi.mocked(spawn).mockReturnValue(child as never) + + const started = await manager.start('conv-1', 'echo test', '/workspace', { + commandShell: PLATFORM_COMMAND_SHELL, + timeout: 0 + }) + + expect(registryMock.record).toHaveBeenCalledWith({ + subsystem: 'background-exec', + recordId: started.sessionId, + pid: child.pid, + commandLine: expect.arrayContaining([PLATFORM_COMMAND_SHELL.executable]), + cwd: expect.stringMatching(/[\\/]workspace$/) + }) + }) + + it('clears the registry record when the session process closes', async () => { + const child = new MockChildProcess() + vi.mocked(spawn).mockReturnValue(child as never) + + const started = await manager.start('conv-1', 'echo test', '/workspace', { + commandShell: PLATFORM_COMMAND_SHELL, + timeout: 0 + }) + + child.stdout.emit('end') + child.stderr.emit('end') + child.emit('close', 0, null) + + await vi.waitFor(() => + expect(registryMock.clear).toHaveBeenCalledWith('background-exec', started.sessionId) + ) + }) + it('bounds completion when an exited process never closes inherited stdio', async () => { vi.useFakeTimers() const child = new MockChildProcess() diff --git a/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts b/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts index 1d322acb92..b26d651aa3 100644 --- a/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts +++ b/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts @@ -1,6 +1,20 @@ import path from 'node:path' -import { describe, expect, it, vi } from 'vitest' -import { getParentPortMessagePayload } from '@/agent/shared/process/backgroundExecUtilityHost' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const registryMock = vi.hoisted(() => ({ + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) +})) + +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: registryMock +})) + +import { + getParentPortMessagePayload, + runBackgroundExecUtilityHostIfRequested +} from '@/agent/shared/process/backgroundExecUtilityHost' import type { BackgroundExecRpcRequest } from '@/agent/shared/process/backgroundExecSessionManager' describe('backgroundExecUtilityHost', () => { @@ -11,6 +25,18 @@ describe('backgroundExecUtilityHost', () => { args: ['conversation-1'] } + beforeEach(() => { + registryMock.record.mockReset() + registryMock.clear.mockReset() + registryMock.reapStaleOnce.mockReset().mockResolvedValue(null) + }) + + afterEach(() => { + vi.useRealTimers() + delete process.env.DEEPCHAT_EXEC_UTILITY_HOST + delete (process as NodeJS.Process & { parentPort?: unknown }).parentPort + }) + it('keeps raw RPC payloads for unit-test and mock callers', () => { expect(getParentPortMessagePayload(request)).toBe(request) }) @@ -19,6 +45,20 @@ describe('backgroundExecUtilityHost', () => { expect(getParentPortMessagePayload({ data: request })).toBe(request) }) + it('reaps stale background-exec child processes when the host runs', () => { + vi.useFakeTimers() + process.env.DEEPCHAT_EXEC_UTILITY_HOST = '1' + const parentPort = { + postMessage: vi.fn(), + on: vi.fn(), + start: vi.fn() + } + Object.defineProperty(process, 'parentPort', { configurable: true, value: parentPort }) + + expect(runBackgroundExecUtilityHostIfRequested()).toBe(true) + expect(registryMock.reapStaleOnce).toHaveBeenCalledWith('background-exec') + }) + it('keeps shell environment helper on the utility-safe logger', async () => { const { readFileSync } = await vi.importActual('node:fs') const source = readFileSync( diff --git a/test/main/agent/shared/process/childProcessRegistry.test.ts b/test/main/agent/shared/process/childProcessRegistry.test.ts new file mode 100644 index 0000000000..6222d5762e --- /dev/null +++ b/test/main/agent/shared/process/childProcessRegistry.test.ts @@ -0,0 +1,308 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('fs', async () => { + const actual = await vi.importActual('fs') + return { __esModule: true, ...actual, default: actual } +}) + +vi.mock('path', async () => { + const actual = await vi.importActual('path') + return { __esModule: true, ...actual, default: actual } +}) + +import fs from 'fs' +import os from 'os' +import path from 'path' +import { createHash } from 'crypto' +import { + ChildProcessRegistry, + defaultChildProcessAttester, + type ChildProcessLaunchRecord, + type ObservedProcessIdentity +} from '@/agent/shared/process/childProcessRegistry' + +function makeRecord(overrides: Partial = {}): ChildProcessLaunchRecord { + return { + version: 1, + subsystem: 'background-exec', + recordId: 'bg_test', + pid: 4321, + ownerPid: 9999, + commandLine: ['/bin/zsh', '-c', 'npm run dev'], + recordedAt: 1_000_000, + ...overrides + } +} + +describe('ChildProcessRegistry', () => { + let rootDir: string + let alivePids: Set + let observations: Map + let terminate: ReturnType + let registry: ChildProcessRegistry + + const writeRecord = (subsystem: string, record: ChildProcessLaunchRecord) => { + const dir = path.join(rootDir, subsystem) + fs.mkdirSync(dir, { recursive: true }) + const sanitized = record.recordId.replace(/[^a-zA-Z0-9_.-]/g, '_').slice(0, 80) || 'record' + const fingerprint = createHash('sha1').update(record.recordId).digest('hex').slice(0, 8) + fs.writeFileSync(path.join(dir, `${sanitized}_${fingerprint}.json`), JSON.stringify(record)) + } + + const recordFileCount = (subsystem: string) => { + const dir = path.join(rootDir, subsystem) + return fs.existsSync(dir) ? fs.readdirSync(dir).filter((f) => f.endsWith('.json')).length : 0 + } + + beforeEach(() => { + rootDir = fs.mkdtempSync(path.join(os.tmpdir(), 'child-process-registry-')) + alivePids = new Set() + observations = new Map() + terminate = vi.fn().mockResolvedValue(true) + registry = new ChildProcessRegistry({ + rootDir, + now: () => 1_060_000, + isAlive: (pid) => alivePids.has(pid), + observe: async (pid) => observations.get(pid) ?? { alive: alivePids.has(pid) }, + terminate, + log: () => {} + }) + }) + + afterEach(() => { + fs.rmSync(rootDir, { recursive: true, force: true }) + }) + + it('persists launch records and lists them back', () => { + registry.record({ + subsystem: 'background-exec', + recordId: 'bg_session/1', + pid: 4321, + commandLine: ['/bin/zsh', '-c', 'npm run dev'], + cwd: '/tmp/work' + }) + + const records = registry.list('background-exec') + expect(records).toHaveLength(1) + expect(records[0]).toMatchObject({ + recordId: 'bg_session/1', + pid: 4321, + ownerPid: process.pid, + commandLine: ['/bin/zsh', '-c', 'npm run dev'], + cwd: '/tmp/work', + recordedAt: 1_060_000 + }) + }) + + it('clears records and tolerates missing files', () => { + registry.record({ + subsystem: 'mcp-stdio', + recordId: 'server-a', + pid: 4321, + commandLine: ['npx', 'server-a'] + }) + expect(recordFileCount('mcp-stdio')).toBe(1) + + registry.clear('mcp-stdio', 'server-a') + expect(recordFileCount('mcp-stdio')).toBe(0) + expect(() => registry.clear('mcp-stdio', 'server-a')).not.toThrow() + }) + + it('drops unreadable record files during list', () => { + const dir = path.join(rootDir, 'mcp-stdio') + fs.mkdirSync(dir, { recursive: true }) + fs.writeFileSync(path.join(dir, 'corrupt.json'), '{nope') + expect(registry.list('mcp-stdio')).toEqual([]) + expect(recordFileCount('mcp-stdio')).toBe(0) + }) + + it('clears records whose process is already gone without terminating', async () => { + writeRecord('background-exec', makeRecord()) + + const result = await registry.reapStale('background-exec') + + expect(result.cleared).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(0) + }) + + it('reaps an attested orphan from a dead owner', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { + alive: true, + commandLine: '/bin/zsh -c npm run dev', + startedAtMs: record.recordedAt - 2000 + }) + + const result = await registry.reapStale('background-exec') + + expect(result.reaped).toEqual(['bg_test']) + expect(terminate).toHaveBeenCalledWith(record.pid) + expect(recordFileCount('background-exec')).toBe(0) + }) + + it('skips records still owned by a live foreign process', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + alivePids.add(record.ownerPid) + + const result = await registry.reapStale('background-exec') + + expect(result.skipped).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(1) + }) + + it('refuses to kill a reused pid whose command line does not match', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { + alive: true, + commandLine: '/usr/libexec/LoginWindowUnrelated', + startedAtMs: record.recordedAt - 2000 + }) + + const result = await registry.reapStale('background-exec') + + expect(result.refused).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + // The recorded process is demonstrably gone; the stale record is dropped. + expect(recordFileCount('background-exec')).toBe(0) + }) + + it('refuses to kill a reused pid whose start time is outside the tolerance window', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { + alive: true, + commandLine: '/bin/zsh -c npm run dev', + startedAtMs: record.recordedAt + 10 * 60 * 1000 + }) + + const result = await registry.reapStale('background-exec') + + expect(result.refused).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(0) + }) + + it('keeps the record when identity cannot be observed', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { alive: true }) + + const result = await registry.reapStale('background-exec') + + expect(result.refused).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(1) + }) + + it('never terminates the current process', async () => { + const record = makeRecord({ pid: process.pid }) + writeRecord('background-exec', record) + + const result = await registry.reapStale('background-exec') + + expect(result.refused).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + }) + + it('clears records older than the maximum record age', async () => { + const record = makeRecord({ recordedAt: 1_000_000 }) + writeRecord('background-exec', record) + alivePids.add(record.pid) + + const result = await registry.reapStale('background-exec', { maxRecordAgeMs: 1000 }) + + expect(result.cleared).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(0) + }) + + it('keeps the record when termination fails', async () => { + const record = makeRecord() + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { + alive: true, + commandLine: '/bin/zsh -c npm run dev', + startedAtMs: record.recordedAt + }) + terminate.mockResolvedValue(false) + + const result = await registry.reapStale('background-exec') + + expect(result.refused).toEqual(['bg_test']) + expect(recordFileCount('background-exec')).toBe(1) + }) + + it('honors excluded record ids and custom attesters', async () => { + const excluded = makeRecord({ recordId: 'keep-me', pid: 5001 }) + const vetted = makeRecord({ recordId: 'custom', pid: 5002 }) + writeRecord('acp-agent', excluded) + writeRecord('acp-agent', vetted) + alivePids.add(5001).add(5002) + observations.set(5002, { alive: true }) + const attester = vi.fn().mockReturnValue(true) + + const result = await registry.reapStale('acp-agent', { + excludeRecordIds: new Set(['keep-me']), + attester + }) + + expect(result.skipped).toEqual(['keep-me']) + expect(result.reaped).toEqual(['custom']) + expect(attester).toHaveBeenCalledWith(vetted, { alive: true }) + expect(terminate).toHaveBeenCalledWith(5002) + expect(terminate).not.toHaveBeenCalledWith(5001) + }) + + it('reaps a subsystem only once per process', async () => { + writeRecord('mcp-stdio', makeRecord({ subsystem: 'mcp-stdio' })) + + const first = await registry.reapStaleOnce('mcp-stdio') + const second = await registry.reapStaleOnce('mcp-stdio') + + expect(first).not.toBeNull() + expect(second).toBeNull() + }) +}) + +describe('defaultChildProcessAttester', () => { + const record = makeRecord() + + it('accepts matching start time and command line fingerprint', () => { + expect( + defaultChildProcessAttester(record, { + alive: true, + commandLine: '/bin/zsh -c npm run dev', + startedAtMs: record.recordedAt + 30_000 + }) + ).toBe(true) + }) + + it('rejects dead processes, missing observations and mismatched fingerprints', () => { + expect(defaultChildProcessAttester(record, { alive: false })).toBe(false) + expect(defaultChildProcessAttester(record, { alive: true })).toBe(false) + expect( + defaultChildProcessAttester(record, { + alive: true, + commandLine: '/bin/zsh -c npm run dev' + }) + ).toBe(false) + expect( + defaultChildProcessAttester(record, { + alive: true, + commandLine: '/bin/zsh -c npm test', + startedAtMs: record.recordedAt + }) + ).toBe(false) + }) +}) diff --git a/test/main/mcp/mcpClient.test.ts b/test/main/mcp/mcpClient.test.ts index f1fee1f8c7..30dbf830e0 100644 --- a/test/main/mcp/mcpClient.test.ts +++ b/test/main/mcp/mcpClient.test.ts @@ -15,6 +15,15 @@ import { const fsExistsSyncMock = vi.hoisted(() => vi.fn()) const terminateProcessTreeMock = vi.hoisted(() => vi.fn().mockResolvedValue(true)) +const childProcessRegistryMock = vi.hoisted(() => ({ + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) +})) + +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: childProcessRegistryMock +})) // Mock electron modules vi.mock('electron', () => ({ @@ -707,6 +716,54 @@ describe('McpClient Runtime Command Processing Tests', () => { }) }) + describe('Child process registry', () => { + it('records the stdio launch after a successful connect', async () => { + const pid = 321 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = pid + } as any) + const client = createMcpClient('registry-test', { + type: 'stdio', + command: 'node', + args: ['server.js'] + }) + + await client.connect() + + expect(childProcessRegistryMock.record).toHaveBeenCalledWith({ + subsystem: 'mcp-stdio', + recordId: 'registry-test', + pid, + commandLine: ['node', 'server.js'] + }) + }) + + it('clears the record when the stdio server disconnects', async () => { + const pid = 654 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = pid + } as any) + const client = createMcpClient('registry-test', { + type: 'stdio', + command: 'node', + args: ['server.js'] + }) + + await client.connect() + await client.disconnect() + + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', 'registry-test') + }) + }) + describe('Version negotiation', () => { it('retries one timed-out HTTP version negotiation probe', async () => { const client = createMcpClient('remote-server', { diff --git a/test/main/mcp/mcpService.test.ts b/test/main/mcp/mcpService.test.ts index 6fb5bc0090..1401b772a3 100644 --- a/test/main/mcp/mcpService.test.ts +++ b/test/main/mcp/mcpService.test.ts @@ -64,6 +64,16 @@ vi.mock('../../../src/main/mcp/mcprouterManager', () => ({ McpRouterManager: vi.fn().mockImplementation(() => ({})) })) +const childProcessRegistryMock = vi.hoisted(() => ({ + record: vi.fn(), + clear: vi.fn(), + reapStaleOnce: vi.fn().mockResolvedValue(null) +})) + +vi.mock('@/agent/shared/process/childProcessRegistry', () => ({ + childProcessRegistry: childProcessRegistryMock +})) + import { McpService } from '../../../src/main/mcp' import { ToolManager } from '../../../src/main/mcp/toolManager' import type { CacheImageOptions } from '../../../src/main/platform/imageCache' @@ -135,6 +145,7 @@ describe('McpService', () => { toolManagerMocks.getAllToolDefinitions.mockResolvedValue([]) toolManagerMocks.snapshotCachedToolDefinitions.mockReturnValue({ state: 'uninitialized' }) toolManagerMocks.callTool.mockReset() + childProcessRegistryMock.reapStaleOnce.mockResolvedValue(null) }) it('caches embedded MCP image URLs before exposing the tool result to the model', async () => { From 32fdf3d4054f9ff72cf7a942b292fbd377e782d0 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Thu, 10 Sep 2026 13:45:48 +0800 Subject: [PATCH 09/15] fix(agent): harden orphaned child process recovery - ACP: keep launch records until the child exit event instead of clearing them at kill time, so unconfirmed terminations remain reapable on next startup - background-exec: skip registry record cleanup when process-tree termination is unconfirmed; clear only on the child close path - MCP stdio: check the terminateProcessTreeByPid boolean result and only clear the launch record after confirmed termination, preserving the pid/record for force-kill or startup reaping - MCP stdio: bind a unique record id (serverName + UUID) to each transport instance so overlapping same-name clients never overwrite or delete each other's recovery record - registry: resolve the record root from the configured userData directory (propagated after app.setPath override) instead of hard-coding the host home directory - tests: add regression coverage for unconfirmed termination and overlapping same-name MCP clients --- .../agent/acp/runtime/acpProcessManager.ts | 4 +- .../process/backgroundExecSessionManager.ts | 23 ++- .../shared/process/childProcessRegistry.ts | 17 +- src/main/appMain.ts | 3 + src/main/mcp/mcpClient.ts | 144 +++++++++++---- .../acp/runtime/acpProcessManager.test.ts | 57 +++++- .../process/childProcessRegistry.test.ts | 23 +++ test/main/mcp/mcpClient.test.ts | 168 +++++++++++++++++- 8 files changed, 380 insertions(+), 59 deletions(-) diff --git a/src/main/agent/acp/runtime/acpProcessManager.ts b/src/main/agent/acp/runtime/acpProcessManager.ts index bd8af02c87..7853ce93d8 100644 --- a/src/main/agent/acp/runtime/acpProcessManager.ts +++ b/src/main/agent/acp/runtime/acpProcessManager.ts @@ -1263,7 +1263,6 @@ export class AcpProcessManager implements AgentProcessManager this.clearLaunchRecord(child)) } return child @@ -2425,8 +2425,6 @@ export class AcpProcessManager implements AgentProcessManager void> { diff --git a/src/main/agent/shared/process/backgroundExecSessionManager.ts b/src/main/agent/shared/process/backgroundExecSessionManager.ts index 533e32e8c5..439b58443e 100644 --- a/src/main/agent/shared/process/backgroundExecSessionManager.ts +++ b/src/main/agent/shared/process/backgroundExecSessionManager.ts @@ -120,6 +120,7 @@ interface BackgroundSession { closePromise: Promise resolveClose: () => void closeSettled: boolean + registryRecordCleared: boolean closeWatchdogId?: NodeJS.Timeout finalizationPromise?: Promise flushOutputDecoders?: () => void @@ -340,6 +341,7 @@ export class BackgroundExecSessionManager { closePromise, resolveClose, closeSettled: false, + registryRecordCleared: false, timedOut: false, ...(ownedSkillExecutionPackageTree ? { ownedSkillExecutionPackageTree } : {}) } @@ -770,6 +772,7 @@ export class BackgroundExecSessionManager { clearTimeout(session.closeWatchdogId) session.closeWatchdogId = undefined } + this.clearRegistryRecord(session) if (session.closeSettled || session.finalizationPromise) return this.recordProcessTerminalState(session, code, signal) void this.finalizeSession(session, code, signal) @@ -826,7 +829,7 @@ export class BackgroundExecSessionManager { session.child.stdin?.destroy() session.child.unref() session.exitCode = undefined - await this.finalizeSession(session, null, 'SIGKILL') + await this.finalizeSession(session, null, 'SIGKILL', { clearRegistryRecord: false }) } await session.closePromise @@ -1009,22 +1012,32 @@ export class BackgroundExecSessionManager { private async finalizeSession( session: BackgroundSession, code: number | null, - signal: NodeJS.Signals | null + signal: NodeJS.Signals | null, + options?: { clearRegistryRecord?: boolean } ): Promise { - session.finalizationPromise ??= this.completeSessionFinalization(session, code, signal) + session.finalizationPromise ??= this.completeSessionFinalization(session, code, signal, options) await session.finalizationPromise } + private clearRegistryRecord(session: BackgroundSession): void { + if (session.registryRecordCleared) return + session.registryRecordCleared = true + childProcessRegistry.clear('background-exec', session.sessionId) + } + private async completeSessionFinalization( session: BackgroundSession, code: number | null, - signal: NodeJS.Signals | null + signal: NodeJS.Signals | null, + options?: { clearRegistryRecord?: boolean } ): Promise { if (session.closeWatchdogId) { clearTimeout(session.closeWatchdogId) session.closeWatchdogId = undefined } - childProcessRegistry.clear('background-exec', session.sessionId) + if (options?.clearRegistryRecord !== false) { + this.clearRegistryRecord(session) + } try { session.flushOutputDecoders?.() await session.outputWriteQueue.catch((error) => { diff --git a/src/main/agent/shared/process/childProcessRegistry.ts b/src/main/agent/shared/process/childProcessRegistry.ts index ea8ad05f12..b370b185da 100644 --- a/src/main/agent/shared/process/childProcessRegistry.ts +++ b/src/main/agent/shared/process/childProcessRegistry.ts @@ -56,7 +56,10 @@ export interface ChildProcessRegistryOptions { } function defaultRegistryRoot(): string { - return path.join(os.homedir(), '.deepchat', 'child-processes') + const userDataDir = + process.env.DEEPCHAT_E2E_USER_DATA_DIR?.trim() || process.env.DEEPCHAT_USER_DATA_DIR?.trim() + const baseDir = userDataDir || path.join(os.homedir(), '.deepchat') + return path.join(baseDir, 'child-processes') } function defaultIsAlive(pid: number): boolean { @@ -192,7 +195,8 @@ function recordFileName(recordId: string): string { } export class ChildProcessRegistry { - private readonly rootDir: string + private readonly configuredRootDir?: string + private resolvedRootDir?: string private readonly now: () => number private readonly isAlive: (pid: number) => boolean private readonly observe: (pid: number) => Promise @@ -202,7 +206,7 @@ export class ChildProcessRegistry { private readonly inflightReaps = new Map>() constructor(options: ChildProcessRegistryOptions = {}) { - this.rootDir = options.rootDir ?? defaultRegistryRoot() + this.configuredRootDir = options.rootDir this.now = options.now ?? Date.now this.isAlive = options.isAlive ?? defaultIsAlive this.observe = options.observe ?? defaultObserve @@ -393,6 +397,13 @@ export class ChildProcessRegistry { return result } + private get rootDir(): string { + // Resolved on first use so the default root reflects the userData path configured + // during app startup (including the DEEPCHAT_E2E_USER_DATA_DIR override). + this.resolvedRootDir ??= this.configuredRootDir ?? defaultRegistryRoot() + return this.resolvedRootDir + } + private recordPath(subsystem: string, recordId: string): string { return path.join(this.rootDir, subsystem, recordFileName(recordId)) } diff --git a/src/main/appMain.ts b/src/main/appMain.ts index ccf18f917b..dc0b1d372c 100644 --- a/src/main/appMain.ts +++ b/src/main/appMain.ts @@ -29,6 +29,9 @@ export function startApp(): void { if (e2eUserDataDir) { app.setPath('userData', e2eUserDataDir) } + // Propagate the effective userData dir so child process registry records (and forked + // utility processes, which cannot access the electron app module) stay isolated with it. + process.env.DEEPCHAT_USER_DATA_DIR = app.getPath('userData') app.setName(APP_NAME) if (process.platform === 'darwin') { diff --git a/src/main/mcp/mcpClient.ts b/src/main/mcp/mcpClient.ts index e4ab1072ac..730beb8b70 100644 --- a/src/main/mcp/mcpClient.ts +++ b/src/main/mcp/mcpClient.ts @@ -1,7 +1,10 @@ import { resolveMcpEnvironmentBinding } from './environmentBindings' import type { ProviderSettingsPort } from '@/provider/settings' import logger from '@shared/logger' -import { StdioClientTransport } from '@modelcontextprotocol/client/stdio' +import { + StdioClientTransport, + type StdioServerParameters +} from '@modelcontextprotocol/client/stdio' import { Client, InMemoryTransport, @@ -249,6 +252,38 @@ const withUnsupportedCapabilityFallback = async ( } } +// Records the spawned server process in the child process registry as soon as +// start() spawns it, so failed, cancelled, or hard-timed-out connects still +// leave a reapable record with a recordedAt close to the process start time. +// Client.connect() calls start() itself, so never start this transport manually. +class RegistryRecordedStdioTransport extends StdioClientTransport { + readonly registryRecordId: string + private readonly registryCommandLine: string[] + + constructor( + params: StdioServerParameters, + registryRecordId: string, + registryCommandLine: string[] + ) { + super(params) + this.registryRecordId = registryRecordId + this.registryCommandLine = registryCommandLine + } + + override async start(): Promise { + await super.start() + const pid = this.pid + if (typeof pid === 'number' && pid > 0) { + childProcessRegistry.record({ + subsystem: 'mcp-stdio', + recordId: this.registryRecordId, + pid, + commandLine: this.registryCommandLine + }) + } + } +} + export class McpClient { private client: Client | null = null private transport: Transport | null = null @@ -257,7 +292,7 @@ export class McpClient { private isConnected: boolean = false private connectionTimeout: NodeJS.Timeout | null = null private stdioPidForShutdown?: number - private stdioCommandLine?: string[] + private stdioRegistryRecordId?: string private connectPromise: Promise | null = null private startupAttempt = 0 private lifecycleStatus: McpServerLifecycleStatus = 'stopped' @@ -643,16 +678,23 @@ export class McpClient { ? resolveBinding(this.serverConfig.cwd) : undefined const pluginRoot = this.serverConfig.source === 'plugin' ? env.PLUGIN_ROOT : undefined - this.stdioCommandLine = [command, ...args] - this.transport = new StdioClientTransport({ - command, - args, - env, - stderr: 'pipe', - cwd: - configuredCwd && pluginRoot ? path.resolve(pluginRoot, configuredCwd) : configuredCwd, - maxBufferSize: MCP_STDIO_MAX_BUFFER_BYTES - }) + // Each transport instance gets a unique registry record id so overlapping + // reconnects of the same server never overwrite or delete each other's record. + const registryRecordId = `${this.serverName}:${randomUUID()}` + this.stdioRegistryRecordId = registryRecordId + this.transport = new RegistryRecordedStdioTransport( + { + command, + args, + env, + stderr: 'pipe', + cwd: + configuredCwd && pluginRoot ? path.resolve(pluginRoot, configuredCwd) : configuredCwd, + maxBufferSize: MCP_STDIO_MAX_BUFFER_BYTES + }, + registryRecordId, + [command, ...args] + ) ;(this.transport as StdioClientTransport).stderr?.on('data', (data) => { console.info('mcp StdioClientTransport error', this.serverName, data.toString()) }) @@ -791,16 +833,6 @@ export class McpClient { } console.info(`MCP server ${this.serverName} connected successfully`) - const stdioPid = this.getStdioPid() - if (stdioPid && this.stdioCommandLine) { - childProcessRegistry.record({ - subsystem: 'mcp-stdio', - recordId: this.serverName, - pid: stdioPid, - commandLine: this.stdioCommandLine - }) - } - this.emitServerStatusChanged('connected', { phase, attempt }) } catch (error) { // 清除超时 @@ -870,6 +902,8 @@ export class McpClient { // 关闭transport const transport = this.transport this.stdioPidForShutdown = this.getStdioPid(transport) ?? this.stdioPidForShutdown + this.stdioRegistryRecordId = + this.getStdioRegistryRecordId(transport) ?? this.stdioRegistryRecordId this.transport = null if (transport) { try { @@ -882,7 +916,6 @@ export class McpClient { // 重置状态 this.client = null this.isConnected = false - this.stdioCommandLine = undefined if (options.emitStopped) { this.emitServerStatusChanged('stopped', { reason: 'shutdown' }) @@ -896,42 +929,75 @@ export class McpClient { return transport.pid ?? undefined } + private getStdioRegistryRecordId( + transport: Transport | null = this.transport + ): string | undefined { + if (!(transport instanceof RegistryRecordedStdioTransport)) { + return undefined + } + return transport.registryRecordId + } + + // Terminates the stdio process tree and only clears its launch record when + // termination is confirmed, so a still-running process keeps its recovery record. + private async terminateStdioProcessTreeAndClearRecord( + pid: number, + recordId: string | undefined + ): Promise { + let terminated = false + try { + terminated = await terminateProcessTreeByPid(pid, { graceMs: 2000 }) + } catch (error) { + console.error(`Failed to terminate MCP stdio process tree for ${this.serverName}:`, error) + } + if (!terminated) { + return false + } + if (recordId) { + childProcessRegistry.clear('mcp-stdio', recordId) + } + if (this.stdioPidForShutdown === pid) { + this.stdioPidForShutdown = undefined + } + if (recordId && this.stdioRegistryRecordId === recordId) { + this.stdioRegistryRecordId = undefined + } + return true + } + async forceTerminateStdioProcessTree(reason: string): Promise { const pid = this.getStdioPid() ?? this.stdioPidForShutdown if (!pid) { return false } - try { - await terminateProcessTreeByPid(pid, { graceMs: 2000 }) - console.warn(`[MCP] Force terminated stdio process tree for ${this.serverName}: ${reason}`) - childProcessRegistry.clear('mcp-stdio', this.serverName) - return true - } catch (error) { + const recordId = this.getStdioRegistryRecordId() ?? this.stdioRegistryRecordId + const terminated = await this.terminateStdioProcessTreeAndClearRecord(pid, recordId) + if (!terminated) { console.warn( - `Failed to force terminate MCP stdio process tree for ${this.serverName}:`, - error + `[MCP] Could not confirm force termination of stdio process tree for ${this.serverName}; keeping launch record for recovery: ${reason}` ) return false } + console.warn(`[MCP] Force terminated stdio process tree for ${this.serverName}: ${reason}`) + return true } private async closeTransport(transport: Transport): Promise { const pid = this.getStdioPid(transport) + const recordId = this.getStdioRegistryRecordId(transport) try { await transport.close() } finally { if (pid) { - try { - await terminateProcessTreeByPid(pid, { graceMs: 2000 }) - } catch (error) { - console.error(`Failed to terminate MCP stdio process tree for ${this.serverName}:`, error) + const terminated = await this.terminateStdioProcessTreeAndClearRecord(pid, recordId) + if (!terminated) { + // Termination is unconfirmed: keep the pid and launch record so the + // process can still be force-terminated or reaped on next startup. + this.stdioPidForShutdown = pid + this.stdioRegistryRecordId = recordId ?? this.stdioRegistryRecordId } - childProcessRegistry.clear('mcp-stdio', this.serverName) - } - if (this.stdioPidForShutdown === pid) { - this.stdioPidForShutdown = undefined } } } diff --git a/test/main/agent/acp/runtime/acpProcessManager.test.ts b/test/main/agent/acp/runtime/acpProcessManager.test.ts index 5a4cd414d5..92cba4c722 100644 --- a/test/main/agent/acp/runtime/acpProcessManager.test.ts +++ b/test/main/agent/acp/runtime/acpProcessManager.test.ts @@ -1149,7 +1149,7 @@ describe('AcpProcessManager child process launch records', () => { }) }) - it('clears the launch record when the child is killed', () => { + it('keeps the launch record when kill is requested and clears it on exit', () => { const manager = createManager() const child = new MockSpawnedChild() vi.mocked(spawn).mockReturnValue(child as never) @@ -1158,9 +1158,64 @@ describe('AcpProcessManager child process launch records', () => { Object.defineProperty(spawned, 'pid', { value: undefined }) ;(manager as any).killChild(spawned) + expect(childProcessRegistryMock.clear).not.toHaveBeenCalled() + + spawned.emit('exit', null, 'SIGTERM') + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('acp-agent', 'agent-1:1234') }) + it('clears the launch record for a late-spawn kill before initialization completes', async () => { + const manager = createManager() + const child = new MockSpawnedChild() + vi.mocked(spawn).mockReturnValue(child as never) + vi.spyOn(manager as any, 'materializeAgentLaunch').mockResolvedValue(launch) + const killSpy = vi.spyOn(process, 'kill').mockReturnValue(true) + const originalSpawn = (manager as any).spawnAgentProcess.bind(manager) + vi.spyOn(manager as any, 'spawnAgentProcess').mockImplementation( + (agentArg: unknown, launchArg: unknown) => { + const spawned = originalSpawn(agentArg, launchArg) + ;(manager as any).shuttingDown = true + return spawned + } + ) + + try { + await expect( + (manager as any).spawnProcessOnce( + agent, + '/tmp/workspace', + { + agentId: 'agent-1', + source: 'manual', + distributionType: 'manual', + command: 'agent', + args: [], + env: {} + }, + 'signature', + undefined + ) + ).rejects.toThrow('shutting down') + + expect(childProcessRegistryMock.record).toHaveBeenCalledWith({ + subsystem: 'acp-agent', + recordId: 'agent-1:1234', + pid: 1234, + commandLine: ['agent', '--acp'], + cwd: '/tmp/workspace' + }) + expect(child.kill).toHaveBeenCalledOnce() + expect(childProcessRegistryMock.clear).not.toHaveBeenCalled() + + child.emit('exit', null, 'SIGTERM') + + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('acp-agent', 'agent-1:1234') + } finally { + killSpy.mockRestore() + } + }) + it('clears the launch record when the agent process exits', async () => { const manager = createManager() const child = new MockStreamingChild() diff --git a/test/main/agent/shared/process/childProcessRegistry.test.ts b/test/main/agent/shared/process/childProcessRegistry.test.ts index 6222d5762e..0d0d25df86 100644 --- a/test/main/agent/shared/process/childProcessRegistry.test.ts +++ b/test/main/agent/shared/process/childProcessRegistry.test.ts @@ -273,6 +273,29 @@ describe('ChildProcessRegistry', () => { expect(first).not.toBeNull() expect(second).toBeNull() }) + + it('resolves the default root from DEEPCHAT_USER_DATA_DIR at first use', () => { + const userDataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'registry-userdata-')) + const previous = process.env.DEEPCHAT_USER_DATA_DIR + process.env.DEEPCHAT_USER_DATA_DIR = userDataDir + try { + const defaultRegistry = new ChildProcessRegistry({ log: () => {} }) + defaultRegistry.record({ + subsystem: 'mcp-stdio', + recordId: 'srv', + pid: 4321, + commandLine: ['npx', 'srv'] + }) + expect(fs.readdirSync(path.join(userDataDir, 'child-processes', 'mcp-stdio'))).toHaveLength(1) + } finally { + if (previous === undefined) { + delete process.env.DEEPCHAT_USER_DATA_DIR + } else { + process.env.DEEPCHAT_USER_DATA_DIR = previous + } + fs.rmSync(userDataDir, { recursive: true, force: true }) + } + }) }) describe('defaultChildProcessAttester', () => { diff --git a/test/main/mcp/mcpClient.test.ts b/test/main/mcp/mcpClient.test.ts index 30dbf830e0..8743c1d559 100644 --- a/test/main/mcp/mcpClient.test.ts +++ b/test/main/mcp/mcpClient.test.ts @@ -122,7 +122,9 @@ const createLargeToolCatalog = () => })) const createSdkToolClient = (tools: unknown[], era: 'modern' | 'legacy') => ({ - connect: vi.fn().mockResolvedValue(undefined), + connect: vi.fn().mockImplementation(async (transport: any) => { + await transport?.start?.() + }), callTool: vi.fn().mockResolvedValue({ content: [] }), listTools: vi.fn().mockResolvedValue({ tools }), listPrompts: vi.fn(), @@ -145,7 +147,9 @@ vi.mock('@modelcontextprotocol/client', async (importOriginal) => { return { ...actual, Client: vi.fn().mockImplementation(() => ({ - connect: vi.fn().mockResolvedValue(undefined), + connect: vi.fn().mockImplementation(async (transport: any) => { + await transport?.start?.() + }), callTool: vi.fn(), listTools: vi.fn(), listPrompts: vi.fn(), @@ -164,14 +168,18 @@ vi.mock('@modelcontextprotocol/client', async (importOriginal) => { } }) -vi.mock('@modelcontextprotocol/client/stdio', () => ({ - StdioClientTransport: vi.fn().mockImplementation(() => ({ +vi.mock('@modelcontextprotocol/client/stdio', () => { + const StdioClientTransport = vi.fn().mockImplementation(() => ({ stderr: { on: vi.fn() }, close: vi.fn() })) -})) + // The production subclass calls super.start(); the real SDK's Client.connect() + // invokes transport.start(), which the Client mock mirrors below. + StdioClientTransport.prototype.start = vi.fn().mockResolvedValue(undefined) + return { StdioClientTransport } +}) describe('McpClient Runtime Command Processing Tests', () => { let mockFsExistsSync: any @@ -210,7 +218,9 @@ describe('McpClient Runtime Command Processing Tests', () => { vi.mocked(Client).mockImplementation( () => ({ - connect: vi.fn().mockResolvedValue(undefined), + connect: vi.fn().mockImplementation(async (transport: any) => { + await transport?.start?.() + }), callTool: vi.fn(), listTools: vi.fn(), listPrompts: vi.fn(), @@ -736,12 +746,59 @@ describe('McpClient Runtime Command Processing Tests', () => { expect(childProcessRegistryMock.record).toHaveBeenCalledWith({ subsystem: 'mcp-stdio', - recordId: 'registry-test', + recordId: expect.stringMatching(/^registry-test:.+/), pid, commandLine: ['node', 'server.js'] }) }) + it('records the spawned process even when connect fails after start', async () => { + const pid = 432 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = pid + } as any) + vi.mocked(Client).mockImplementationOnce( + () => + ({ + connect: vi.fn().mockImplementation(async (transport: any) => { + await transport?.start?.() + throw new Error('handshake failed') + }), + callTool: vi.fn(), + listTools: vi.fn(), + listPrompts: vi.fn(), + getPrompt: vi.fn(), + listResources: vi.fn(), + readResource: vi.fn(), + setNotificationHandler: vi.fn(), + setRequestHandler: vi.fn(), + getProtocolEra: vi.fn(() => 'modern') + }) as any + ) + const client = createMcpClient('failing-server', { + type: 'stdio', + command: 'node', + args: ['server.js'] + }) + + // test/setup.ts restores all mocks after each test, so the hoisted default + // implementation is gone; set the termination result explicitly per test. + terminateProcessTreeMock.mockResolvedValue(true) + await expect(client.connect()).rejects.toThrow('handshake failed') + + expect(childProcessRegistryMock.record).toHaveBeenCalledWith({ + subsystem: 'mcp-stdio', + recordId: expect.stringMatching(/^failing-server:.+/), + pid, + commandLine: ['node', 'server.js'] + }) + const recordId = childProcessRegistryMock.record.mock.calls[0][0].recordId + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', recordId) + }) it('clears the record when the stdio server disconnects', async () => { const pid = 654 vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { @@ -757,10 +814,105 @@ describe('McpClient Runtime Command Processing Tests', () => { args: ['server.js'] }) + terminateProcessTreeMock.mockResolvedValue(true) + await client.connect() + await client.disconnect() + + const recordId = childProcessRegistryMock.record.mock.calls[0][0].recordId + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', recordId) + }) + + it('keeps the record when process-tree termination is unconfirmed on disconnect', async () => { + const pid = 765 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = pid + } as any) + const client = createMcpClient('unconfirmed-server', { + type: 'stdio', + command: 'node', + args: ['server.js'] + }) + await client.connect() + const recordId = childProcessRegistryMock.record.mock.calls[0][0].recordId + + terminateProcessTreeMock.mockResolvedValueOnce(false) await client.disconnect() - expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', 'registry-test') + expect(childProcessRegistryMock.clear).not.toHaveBeenCalled() + + // A later confirmed force termination clears the preserved record. + terminateProcessTreeMock.mockResolvedValueOnce(true) + await expect(client.forceTerminateStdioProcessTree('test cleanup')).resolves.toBe(true) + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', recordId) + }) + + it('returns false and keeps the record when force termination is unconfirmed', async () => { + const pid = 876 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = pid + } as any) + const client = createMcpClient('force-server', { + type: 'stdio', + command: 'node', + args: ['server.js'] + }) + + await client.connect() + const recordId = childProcessRegistryMock.record.mock.calls[0][0].recordId + + terminateProcessTreeMock.mockResolvedValueOnce(false) + await expect(client.forceTerminateStdioProcessTree('test')).resolves.toBe(false) + expect(childProcessRegistryMock.clear).not.toHaveBeenCalledWith('mcp-stdio', recordId) + }) + + it('keeps per-instance records for overlapping same-name clients', async () => { + const firstPid = 111 + const secondPid = 222 + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = firstPid + } as any) + vi.mocked(StdioClientTransport).mockImplementationOnce(function (this: any) { + this.stderr = { + on: vi.fn() + } + this.close = vi.fn().mockResolvedValue(undefined) + this.pid = secondPid + } as any) + const serverConfig = { + type: 'stdio', + command: 'node', + args: ['server.js'] + } + const firstClient = createMcpClient('overlap-server', serverConfig) + const secondClient = createMcpClient('overlap-server', serverConfig) + + await firstClient.connect() + await secondClient.connect() + + const firstRecordId = childProcessRegistryMock.record.mock.calls[0][0].recordId + const secondRecordId = childProcessRegistryMock.record.mock.calls[1][0].recordId + expect(firstRecordId).toMatch(/^overlap-server:.+/) + expect(secondRecordId).toMatch(/^overlap-server:.+/) + expect(firstRecordId).not.toBe(secondRecordId) + + // Closing the older client must not delete the newer instance's record. + terminateProcessTreeMock.mockResolvedValue(true) + await firstClient.disconnect() + expect(childProcessRegistryMock.clear).toHaveBeenCalledWith('mcp-stdio', firstRecordId) + expect(childProcessRegistryMock.clear).not.toHaveBeenCalledWith('mcp-stdio', secondRecordId) }) }) From ec03951d0c555e86535f88f0fc0f7007f429a1ba Mon Sep 17 00:00:00 2001 From: zerob13 Date: Thu, 10 Sep 2026 16:41:45 +0800 Subject: [PATCH 10/15] fix(agent): attest and await orphan recovery --- .../agent/acp/runtime/acpProcessManager.ts | 6 +- .../process/backgroundExecUtilityHost.ts | 4 +- .../shared/process/childProcessRegistry.ts | 82 ++++++--- src/main/mcp/index.ts | 2 +- .../acp/runtime/acpProcessManager.test.ts | 19 ++- .../process/backgroundExecUtilityHost.test.ts | 25 +++ .../process/childProcessRegistry.test.ts | 155 ++++++++++++++++-- test/main/mcp/mcpService.test.ts | 15 +- 8 files changed, 266 insertions(+), 42 deletions(-) diff --git a/src/main/agent/acp/runtime/acpProcessManager.ts b/src/main/agent/acp/runtime/acpProcessManager.ts index 7853ce93d8..4f87167875 100644 --- a/src/main/agent/acp/runtime/acpProcessManager.ts +++ b/src/main/agent/acp/runtime/acpProcessManager.ts @@ -286,9 +286,6 @@ export class AcpProcessManager implements AgentProcessManager { - console.warn('[ACP] Failed to reap stale agent processes:', error) - }) } getTerminalSnapshot(terminalId: string): schema.TerminalOutputResponse | null { @@ -1207,6 +1204,9 @@ export class AcpProcessManager implements AgentProcessManager { + await childProcessRegistry.reapStaleOnce('acp-agent')?.catch((error) => { + console.warn('[ACP] Failed to reap stale agent processes:', error) + }) this.assertAcceptingProcesses() const materializedLaunch = await this.materializeAgentLaunch( agent, diff --git a/src/main/agent/shared/process/backgroundExecUtilityHost.ts b/src/main/agent/shared/process/backgroundExecUtilityHost.ts index 7ae68f9528..92f713d952 100644 --- a/src/main/agent/shared/process/backgroundExecUtilityHost.ts +++ b/src/main/agent/shared/process/backgroundExecUtilityHost.ts @@ -107,7 +107,7 @@ export function runBackgroundExecUtilityHostIfRequested(): boolean { } const manager = new BackgroundExecSessionManager() - void childProcessRegistry.reapStaleOnce('background-exec').catch((error) => { + const recovery = childProcessRegistry.reapStaleOnce('background-exec').catch((error) => { logger.warn('[BackgroundExec] Failed to reap stale child processes:', error) }) const keepAliveIntervalId = setInterval(() => {}, 2 ** 31 - 1) @@ -118,7 +118,7 @@ export function runBackgroundExecUtilityHostIfRequested(): boolean { if (!isBackgroundExecRpcRequest(request)) { return } - void handleRequest(manager, parentPort, request) + void recovery.then(() => handleRequest(manager, parentPort, request)) }) process.once('beforeExit', () => { diff --git a/src/main/agent/shared/process/childProcessRegistry.ts b/src/main/agent/shared/process/childProcessRegistry.ts index b370b185da..f2508c015c 100644 --- a/src/main/agent/shared/process/childProcessRegistry.ts +++ b/src/main/agent/shared/process/childProcessRegistry.ts @@ -6,7 +6,6 @@ import path from 'path' import { terminateProcessTreeByPid } from './processTree' const RECORD_VERSION = 1 -const PROCESS_START_TOLERANCE_MS = 60_000 const DEFAULT_MAX_RECORD_AGE_MS = 7 * 24 * 60 * 60 * 1000 export type ChildProcessSubsystem = 'background-exec' | 'mcp-stdio' | 'acp-agent' @@ -20,6 +19,7 @@ export interface ChildProcessLaunchRecord { commandLine: string[] cwd?: string recordedAt: number + startedAtMs?: number } export interface ObservedProcessIdentity { @@ -80,6 +80,9 @@ async function runAndCapture( try { const child = spawn(command, args, { stdio: ['ignore', 'pipe', 'ignore'], + env: { ...process.env, LC_ALL: 'C', TZ: 'UTC' }, + timeout: 5000, + killSignal: 'SIGKILL', ...(process.platform === 'win32' ? { windowsHide: true } : {}) }) child.stdout?.on('data', (chunk: Buffer | string) => { @@ -94,23 +97,31 @@ async function runAndCapture( } async function observePosix(pid: number): Promise { - const { stdout } = await runAndCapture('ps', ['-p', `${pid}`, '-o', 'etimes=', '-o', 'command=']) + const { code, stdout } = await runAndCapture('ps', [ + '-ww', + '-p', + `${pid}`, + '-o', + 'lstart=', + '-o', + 'command=' + ]) const trimmed = stdout.trim() - if (!trimmed) { - return { alive: false } + if (code !== 0 || !trimmed) { + return { alive: defaultIsAlive(pid) } } - const match = /^(\d+)\s+([\s\S]*)$/.exec(trimmed) + const match = /^(\w{3}\s+\w{3}\s+\d+\s+\d{2}:\d{2}:\d{2}\s+\d{4})\s+([\s\S]*)$/.exec(trimmed) if (!match) { return { alive: true } } - const elapsedSeconds = Number.parseInt(match[1], 10) - if (!Number.isFinite(elapsedSeconds)) { + const startedAtMs = Date.parse(`${match[1]} GMT`) + if (!Number.isFinite(startedAtMs)) { return { alive: true } } return { alive: true, commandLine: match[2], - startedAtMs: Date.now() - elapsedSeconds * 1000 + startedAtMs } } @@ -118,7 +129,8 @@ async function observeWindows(pid: number): Promise { const command = [ 'Get-CimInstance Win32_Process -Filter "ProcessId=', `${pid}`, - '" | Select-Object -Property CreationDate,CommandLine | ConvertTo-Json -Compress' + '" | Select-Object CommandLine,@{Name="CreationDate";Expression={', + '$_.CreationDate.ToUniversalTime().ToString("o")}} | ConvertTo-Json -Compress' ].join('') const { stdout } = await runAndCapture('powershell', [ '-NoProfile', @@ -158,8 +170,9 @@ export function defaultChildProcessAttester( return false } if ( - typeof observed.startedAtMs !== 'number' || - Math.abs(observed.startedAtMs - record.recordedAt) > PROCESS_START_TOLERANCE_MS + !Number.isFinite(record.startedAtMs) || + !Number.isFinite(observed.startedAtMs) || + observed.startedAtMs !== record.startedAtMs ) { return false } @@ -204,6 +217,7 @@ export class ChildProcessRegistry { private readonly log: (message: string, ...args: unknown[]) => void private readonly reapedSubsystems = new Set() private readonly inflightReaps = new Map>() + private readonly pendingRecords = new Map() constructor(options: ChildProcessRegistryOptions = {}) { this.configuredRootDir = options.rootDir @@ -215,18 +229,20 @@ export class ChildProcessRegistry { this.log = options.log ?? ((message, ...args) => console.warn(message, ...args)) } - record(entry: { + async record(entry: { subsystem: ChildProcessSubsystem recordId: string pid: number commandLine: string[] cwd?: string - }): void { + }): Promise { if (!Number.isSafeInteger(entry.pid) || entry.pid <= 0 || !entry.recordId) { return } + const key = `${entry.subsystem}:${entry.recordId}` + let record: ChildProcessLaunchRecord | undefined try { - const record: ChildProcessLaunchRecord = { + record = { version: RECORD_VERSION, subsystem: entry.subsystem, recordId: entry.recordId, @@ -236,17 +252,29 @@ export class ChildProcessRegistry { ...(entry.cwd ? { cwd: entry.cwd } : {}), recordedAt: this.now() } - const filePath = this.recordPath(entry.subsystem, entry.recordId) - fs.mkdirSync(path.dirname(filePath), { recursive: true }) - const tempPath = `${filePath}.${process.pid}.tmp` - fs.writeFileSync(tempPath, JSON.stringify(record), 'utf-8') - fs.renameSync(tempPath, filePath) + this.pendingRecords.set(key, record) + this.writeRecord(record) + const observed = await this.observe(entry.pid) + // An exit or a newer launch may have cleared/replaced this record while ps ran. + if ( + this.pendingRecords.get(key) === record && + observed.alive && + typeof observed.startedAtMs === 'number' && + Number.isFinite(observed.startedAtMs) && + observed.startedAtMs <= record.recordedAt + ) { + record.startedAtMs = observed.startedAtMs + this.writeRecord(record) + } } catch (error) { this.log(`[ChildProcessRegistry] Failed to record launch ${entry.recordId}:`, error) + } finally { + if (this.pendingRecords.get(key) === record) this.pendingRecords.delete(key) } } clear(subsystem: ChildProcessSubsystem, recordId: string): void { + this.pendingRecords.delete(`${subsystem}:${recordId}`) if (!recordId) { return } @@ -311,6 +339,8 @@ export class ChildProcessRegistry { subsystem: ChildProcessSubsystem, options: ReapStaleChildProcessesOptions = {} ): Promise { + const inflight = this.inflightReaps.get(subsystem) + if (inflight) return await inflight if (this.reapedSubsystems.has(subsystem)) { return null } @@ -363,8 +393,12 @@ export class ChildProcessRegistry { } let attested = false + if (!Number.isFinite(record.startedAtMs) || !Number.isFinite(observed.startedAtMs)) { + result.refused.push(record.recordId) + continue + } try { - attested = attester(record, observed) + attested = record.startedAtMs === observed.startedAtMs && attester(record, observed) } catch (error) { this.log(`[ChildProcessRegistry] Attester failed for ${record.recordId}:`, error) } @@ -404,6 +438,14 @@ export class ChildProcessRegistry { return this.resolvedRootDir } + private writeRecord(record: ChildProcessLaunchRecord): void { + const filePath = this.recordPath(record.subsystem, record.recordId) + fs.mkdirSync(path.dirname(filePath), { recursive: true }) + const tempPath = `${filePath}.${process.pid}.tmp` + fs.writeFileSync(tempPath, JSON.stringify(record), 'utf-8') + fs.renameSync(tempPath, filePath) + } + private recordPath(subsystem: string, recordId: string): string { return path.join(this.rootDir, subsystem, recordFileName(recordId)) } diff --git a/src/main/mcp/index.ts b/src/main/mcp/index.ts index ee59d1fe8c..3ed248e506 100644 --- a/src/main/mcp/index.ts +++ b/src/main/mcp/index.ts @@ -334,7 +334,7 @@ export class McpService implements McpServicePort { } try { - void childProcessRegistry + await childProcessRegistry .reapStaleOnce('mcp-stdio') .then((result) => { if (result && result.reaped.length > 0) { diff --git a/test/main/agent/acp/runtime/acpProcessManager.test.ts b/test/main/agent/acp/runtime/acpProcessManager.test.ts index 92cba4c722..b2b482abd4 100644 --- a/test/main/agent/acp/runtime/acpProcessManager.test.ts +++ b/test/main/agent/acp/runtime/acpProcessManager.test.ts @@ -1127,10 +1127,23 @@ describe('AcpProcessManager child process launch records', () => { vi.mocked(spawn).mockClear() }) - it('reaps stale agent process records on construction', () => { - createManager() - + it('waits for orphan recovery before spawning an agent', async () => { + let finishRecovery!: () => void + childProcessRegistryMock.reapStaleOnce.mockImplementationOnce(() => new Promise((resolve) => { + finishRecovery = resolve + })) + const manager = createManager() + const child = new MockSpawnedChild() + vi.mocked(spawn).mockReturnValue(child as never) + vi.spyOn(manager as any, 'materializeAgentLaunch').mockResolvedValue(launch) + vi.spyOn(manager as any, 'initializeSpawnedProcess').mockResolvedValue({}) + const pending = (manager as any).spawnProcessOnce(agent, launch.cwd, {}, 'signature', undefined) + await Promise.resolve() expect(childProcessRegistryMock.reapStaleOnce).toHaveBeenCalledWith('acp-agent') + expect(spawn).not.toHaveBeenCalled() + finishRecovery() + await pending + expect(spawn).toHaveBeenCalledOnce() }) it('records the launch after spawning an agent process', () => { diff --git a/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts b/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts index b26d651aa3..9e16f58a46 100644 --- a/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts +++ b/test/main/agent/shared/process/backgroundExecUtilityHost.test.ts @@ -59,6 +59,31 @@ describe('backgroundExecUtilityHost', () => { expect(registryMock.reapStaleOnce).toHaveBeenCalledWith('background-exec') }) + it('waits for orphan recovery before handling RPC requests', async () => { + vi.useFakeTimers() + let finishRecovery!: () => void + registryMock.reapStaleOnce.mockImplementation( + () => + new Promise((resolve) => { + finishRecovery = resolve + }) + ) + process.env.DEEPCHAT_EXEC_UTILITY_HOST = '1' + const parentPort = { postMessage: vi.fn(), on: vi.fn(), start: vi.fn() } + Object.defineProperty(process, 'parentPort', { configurable: true, value: parentPort }) + + runBackgroundExecUtilityHostIfRequested() + const onMessage = parentPort.on.mock.calls[0][1] + onMessage({ data: request }) + await vi.advanceTimersByTimeAsync(0) + expect(parentPort.postMessage).not.toHaveBeenCalled() + finishRecovery() + await vi.advanceTimersByTimeAsync(0) + expect(parentPort.postMessage).toHaveBeenCalledWith( + expect.objectContaining({ id: 'rpc-1', ok: true }) + ) + }) + it('keeps shell environment helper on the utility-safe logger', async () => { const { readFileSync } = await vi.importActual('node:fs') const source = readFileSync( diff --git a/test/main/agent/shared/process/childProcessRegistry.test.ts b/test/main/agent/shared/process/childProcessRegistry.test.ts index 0d0d25df86..cce71aa8bc 100644 --- a/test/main/agent/shared/process/childProcessRegistry.test.ts +++ b/test/main/agent/shared/process/childProcessRegistry.test.ts @@ -10,10 +10,16 @@ vi.mock('path', async () => { return { __esModule: true, ...actual, default: actual } }) +vi.mock('child_process', async (importOriginal) => { + const actual = await importOriginal() + return { ...actual, spawn: vi.fn(actual.spawn) } +}) + import fs from 'fs' import os from 'os' import path from 'path' import { createHash } from 'crypto' +import { spawn } from 'child_process' import { ChildProcessRegistry, defaultChildProcessAttester, @@ -30,6 +36,7 @@ function makeRecord(overrides: Partial = {}): ChildPro ownerPid: 9999, commandLine: ['/bin/zsh', '-c', 'npm run dev'], recordedAt: 1_000_000, + startedAtMs: 998_000, ...overrides } } @@ -73,8 +80,9 @@ describe('ChildProcessRegistry', () => { fs.rmSync(rootDir, { recursive: true, force: true }) }) - it('persists launch records and lists them back', () => { - registry.record({ + it('persists the observed process start time independently of the write time', async () => { + observations.set(4321, { alive: true, startedAtMs: 998_000 }) + await registry.record({ subsystem: 'background-exec', recordId: 'bg_session/1', pid: 4321, @@ -90,7 +98,8 @@ describe('ChildProcessRegistry', () => { ownerPid: process.pid, commandLine: ['/bin/zsh', '-c', 'npm run dev'], cwd: '/tmp/work', - recordedAt: 1_060_000 + recordedAt: 1_060_000, + startedAtMs: 998_000 }) }) @@ -174,14 +183,14 @@ describe('ChildProcessRegistry', () => { expect(recordFileCount('background-exec')).toBe(0) }) - it('refuses to kill a reused pid whose start time is outside the tolerance window', async () => { + it('refuses a reused pid with the same command even within the old tolerance window', async () => { const record = makeRecord() writeRecord('background-exec', record) alivePids.add(record.pid) observations.set(record.pid, { alive: true, commandLine: '/bin/zsh -c npm run dev', - startedAtMs: record.recordedAt + 10 * 60 * 1000 + startedAtMs: record.startedAtMs! + 1000 }) const result = await registry.reapStale('background-exec') @@ -204,6 +213,104 @@ describe('ChildProcessRegistry', () => { expect(recordFileCount('background-exec')).toBe(1) }) + it('keeps legacy records without a persisted start time without terminating', async () => { + const record = makeRecord({ startedAtMs: undefined }) + writeRecord('background-exec', record) + alivePids.add(record.pid) + observations.set(record.pid, { + alive: true, + commandLine: '/bin/zsh -c npm run dev', + startedAtMs: 998_000 + }) + + expect((await registry.reapStale('background-exec')).refused).toEqual(['bg_test']) + expect(terminate).not.toHaveBeenCalled() + expect(recordFileCount('background-exec')).toBe(1) + }) + + it.each(['clear', 'replace'] as const)( + 'does not restore an old record when observation completes after %s', + async (action) => { + let resolveObservation!: (identity: ObservedProcessIdentity) => void + const observe = vi + .fn() + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveObservation = resolve + }) + ) + .mockResolvedValue({ alive: true, startedAtMs: 999_000 }) + const delayedRegistry = new ChildProcessRegistry({ rootDir, observe, log: () => {} }) + const entry = { + subsystem: 'mcp-stdio' as const, + recordId: 'server', + pid: 4321, + commandLine: ['node'] + } + const pending = delayedRegistry.record(entry) + if (action === 'clear') { + delayedRegistry.clear(entry.subsystem, entry.recordId) + } else { + await delayedRegistry.record({ ...entry, pid: 4322 }) + } + resolveObservation({ alive: true, startedAtMs: 998_000 }) + await pending + + const records = delayedRegistry.list('mcp-stdio') + if (action === 'clear') expect(records).toEqual([]) + else expect(records).toMatchObject([{ pid: 4322, startedAtMs: 999_000 }]) + } + ) + + it('observes a live process consistently across scans without terminating it', async () => { + const realRegistry = new ChildProcessRegistry({ rootDir, terminate, log: () => {} }) + await realRegistry.record({ + subsystem: 'mcp-stdio', + recordId: 'live-process', + pid: process.pid, + commandLine: [process.execPath] + }) + const [record] = realRegistry.list('mcp-stdio') + expect(record.startedAtMs).toEqual(expect.any(Number)) + // A separate observer must report the same OS timestamp, not a Date.now()-based estimate. + await realRegistry.record({ + subsystem: 'mcp-stdio', + recordId: 'second-observation', + pid: process.pid, + commandLine: [process.execPath] + }) + expect(realRegistry.list('mcp-stdio').map((entry) => entry.startedAtMs)).toEqual([ + record.startedAtMs, + record.startedAtMs + ]) + expect(terminate).not.toHaveBeenCalled() + }) + + it('retains a live process record when the OS query fails', async () => { + const record = makeRecord({ pid: process.ppid }) + writeRecord('mcp-stdio', { ...record, subsystem: 'mcp-stdio' }) + const actual = await vi.importActual('child_process') + vi.mocked(spawn).mockImplementationOnce(() => + actual.spawn(process.execPath, ['-e', 'process.exit(1)'], { + stdio: ['ignore', 'pipe', 'ignore'] + }) + ) + const failedQueryRegistry = new ChildProcessRegistry({ + rootDir, + now: () => 1_060_000, + isAlive: (pid) => pid === record.pid, + terminate, + log: () => {} + }) + + const result = await failedQueryRegistry.reapStale('mcp-stdio') + + expect(result.refused).toEqual([record.recordId]) + expect(recordFileCount('mcp-stdio')).toBe(1) + expect(terminate).not.toHaveBeenCalled() + }) + it('never terminates the current process', async () => { const record = makeRecord({ pid: process.pid }) writeRecord('background-exec', record) @@ -233,7 +340,7 @@ describe('ChildProcessRegistry', () => { observations.set(record.pid, { alive: true, commandLine: '/bin/zsh -c npm run dev', - startedAtMs: record.recordedAt + startedAtMs: record.startedAtMs }) terminate.mockResolvedValue(false) @@ -249,7 +356,7 @@ describe('ChildProcessRegistry', () => { writeRecord('acp-agent', excluded) writeRecord('acp-agent', vetted) alivePids.add(5001).add(5002) - observations.set(5002, { alive: true }) + observations.set(5002, { alive: true, startedAtMs: vetted.startedAtMs }) const attester = vi.fn().mockReturnValue(true) const result = await registry.reapStale('acp-agent', { @@ -259,7 +366,7 @@ describe('ChildProcessRegistry', () => { expect(result.skipped).toEqual(['keep-me']) expect(result.reaped).toEqual(['custom']) - expect(attester).toHaveBeenCalledWith(vetted, { alive: true }) + expect(attester).toHaveBeenCalledWith(vetted, { alive: true, startedAtMs: vetted.startedAtMs }) expect(terminate).toHaveBeenCalledWith(5002) expect(terminate).not.toHaveBeenCalledWith(5001) }) @@ -274,13 +381,39 @@ describe('ChildProcessRegistry', () => { expect(second).toBeNull() }) - it('resolves the default root from DEEPCHAT_USER_DATA_DIR at first use', () => { + it('waits for in-flight recovery before allowing another startup to proceed', async () => { + writeRecord('mcp-stdio', makeRecord({ subsystem: 'mcp-stdio' })) + alivePids.add(4321) + observations.set(4321, { + alive: true, + startedAtMs: 998_000, + commandLine: '/bin/zsh -c npm run dev' + }) + let finishTermination!: (terminated: boolean) => void + terminate.mockImplementation( + () => + new Promise((resolve) => { + finishTermination = resolve + }) + ) + const first = registry.reapStaleOnce('mcp-stdio') + const ready = vi.fn() + const second = registry.reapStaleOnce('mcp-stdio').then(ready) + await Promise.resolve() + expect(ready).not.toHaveBeenCalled() + finishTermination(true) + await Promise.all([first, second]) + expect(ready).toHaveBeenCalledWith(expect.objectContaining({ reaped: ['bg_test'] })) + expect(terminate).toHaveBeenCalledOnce() + }) + + it('resolves the default root from DEEPCHAT_USER_DATA_DIR at first use', async () => { const userDataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'registry-userdata-')) const previous = process.env.DEEPCHAT_USER_DATA_DIR process.env.DEEPCHAT_USER_DATA_DIR = userDataDir try { const defaultRegistry = new ChildProcessRegistry({ log: () => {} }) - defaultRegistry.record({ + await defaultRegistry.record({ subsystem: 'mcp-stdio', recordId: 'srv', pid: 4321, @@ -306,7 +439,7 @@ describe('defaultChildProcessAttester', () => { defaultChildProcessAttester(record, { alive: true, commandLine: '/bin/zsh -c npm run dev', - startedAtMs: record.recordedAt + 30_000 + startedAtMs: record.startedAtMs }) ).toBe(true) }) diff --git a/test/main/mcp/mcpService.test.ts b/test/main/mcp/mcpService.test.ts index 1401b772a3..d6b6ac82c6 100644 --- a/test/main/mcp/mcpService.test.ts +++ b/test/main/mcp/mcpService.test.ts @@ -563,7 +563,14 @@ describe('McpService', () => { expect(serverManagerMocks.startServer).not.toHaveBeenCalled() }) - it('does not wait for hanging enabled servers during initialization', async () => { + it('waits for orphan recovery but not hanging enabled servers during initialization', async () => { + let finishRecovery!: () => void + childProcessRegistryMock.reapStaleOnce.mockImplementationOnce( + () => + new Promise((resolve) => { + finishRecovery = resolve + }) + ) const providerSettings = createProviderSettings( true, false, @@ -584,8 +591,12 @@ describe('McpService', () => { serverManagerMocks.startServer.mockImplementation(() => new Promise(() => {})) serverManagerMocks.isServerActive.mockReturnValue(true) + const initialization = presenter.initialize() + await vi.advanceTimersByTimeAsync(0) + expect(serverManagerMocks.startServer).not.toHaveBeenCalled() + finishRecovery() const result = Promise.race([ - presenter.initialize().then(() => 'initialized'), + initialization.then(() => 'initialized'), new Promise((resolve) => setTimeout(() => resolve('blocked'), 1)) ]) await vi.advanceTimersByTimeAsync(1) From 9a46189c6e25a59dfcb7cb13b2a0ca20289e410d Mon Sep 17 00:00:00 2001 From: xiao-test Date: Fri, 11 Sep 2026 16:47:24 +0800 Subject: [PATCH 11/15] fix(renderer): bind store events at setup scope --- src/renderer/src/stores/cliApproval.ts | 42 +++++++++--------- src/renderer/src/stores/dialog.ts | 10 +++-- src/renderer/src/stores/floatingButton.ts | 8 ++-- src/renderer/src/stores/language.ts | 8 ++-- src/renderer/src/stores/mcp.ts | 24 +++++----- src/renderer/src/stores/mcpAppConsent.ts | 46 ++++++++++---------- src/renderer/src/stores/mcpElicitation.ts | 27 +++++++----- src/renderer/src/stores/mcpSampling.ts | 24 +++++----- src/renderer/src/stores/shortcutKey.ts | 8 ++-- src/renderer/src/stores/uiSettingsStore.ts | 20 +++++---- test/renderer/stores/mcpElicitation.test.ts | 13 +++--- test/renderer/stores/mcpStore.test.ts | 20 +++------ test/renderer/stores/uiSettingsStore.test.ts | 1 + 13 files changed, 131 insertions(+), 120 deletions(-) diff --git a/src/renderer/src/stores/cliApproval.ts b/src/renderer/src/stores/cliApproval.ts index cf8139829f..e82e6f9d32 100644 --- a/src/renderer/src/stores/cliApproval.ts +++ b/src/renderer/src/stores/cliApproval.ts @@ -1,4 +1,4 @@ -import { computed, onMounted, onUnmounted, ref, shallowRef } from 'vue' +import { computed, getCurrentScope, onScopeDispose, ref, shallowRef } from 'vue' import { defineStore } from 'pinia' import { createApprovalClient } from '@api/ApprovalClient' import type { DeepchatEventPayload } from '@shared/contracts/events' @@ -33,25 +33,27 @@ export const useCliApprovalStore = defineStore('cliApproval', () => { } } - onMounted(() => { - cleanups.push( - client.onRequested((next) => { - if (queue.value.some((entry) => entry.requestId === next.requestId)) return - if (queue.value.length >= MAX_PENDING_CLI_APPROVALS) { - void client.resolve(next.requestId, 'denied').catch((error) => { - console.error('[CLI Approval] Failed to reject queued approval:', error) - }) - return - } - queue.value = [...queue.value, next] - }), - client.onClosed(({ requestId }) => remove(requestId)) - ) - }) - - onUnmounted(() => { - while (cleanups.length > 0) cleanups.pop()?.() - }) + // Subscribe at store setup top level (not in a component lifecycle hook) so global + // approval events are not lost when the first consuming component unmounts. + cleanups.push( + client.onRequested((next) => { + if (queue.value.some((entry) => entry.requestId === next.requestId)) return + if (queue.value.length >= MAX_PENDING_CLI_APPROVALS) { + void client.resolve(next.requestId, 'denied').catch((error) => { + console.error('[CLI Approval] Failed to reject queued approval:', error) + }) + return + } + queue.value = [...queue.value, next] + }), + client.onClosed(({ requestId }) => remove(requestId)) + ) + + if (getCurrentScope()) { + onScopeDispose(() => { + while (cleanups.length > 0) cleanups.pop()?.() + }) + } return { request, diff --git a/src/renderer/src/stores/dialog.ts b/src/renderer/src/stores/dialog.ts index 06c1a9ae87..bcce5002c0 100644 --- a/src/renderer/src/stores/dialog.ts +++ b/src/renderer/src/stores/dialog.ts @@ -1,7 +1,7 @@ import { createDialogClient } from '@api/DialogClient' import type { DialogRequest, DialogResponse } from '@shared/types/dialog' import { defineStore } from 'pinia' -import { onMounted, onUnmounted, ref } from 'vue' +import { getCurrentScope, onScopeDispose, ref } from 'vue' export const useDialogStore = defineStore('dialog', () => { const dialogClient = createDialogClient() @@ -106,8 +106,12 @@ export const useDialogStore = defineStore('dialog', () => { } } - onMounted(setupUpdateListener) - onUnmounted(removeUpdateListener) + // Subscribe at store setup top level (not in a component lifecycle hook) so the + // global dialog listener is not lost when the first consuming component unmounts. + setupUpdateListener() + if (getCurrentScope()) { + onScopeDispose(removeUpdateListener) + } return { timeoutMilliseconds, diff --git a/src/renderer/src/stores/floatingButton.ts b/src/renderer/src/stores/floatingButton.ts index c88d040598..8c1bdb5028 100644 --- a/src/renderer/src/stores/floatingButton.ts +++ b/src/renderer/src/stores/floatingButton.ts @@ -1,5 +1,5 @@ import { defineStore } from 'pinia' -import { ref, onMounted, onScopeDispose } from 'vue' +import { ref, onScopeDispose } from 'vue' import { createConfigClient } from '../../api/ConfigClient' export const useFloatingButtonStore = defineStore('floatingButton', () => { @@ -73,10 +73,8 @@ export const useFloatingButtonStore = defineStore('floatingButton', () => { return task } - // 在组件挂载时初始化 - onMounted(() => { - void initializeState() - }) + // 在 store setup 顶层初始化（不依赖组件生命周期钩子） + void initializeState() onScopeDispose(() => { removeFloatingButtonListener?.() diff --git a/src/renderer/src/stores/language.ts b/src/renderer/src/stores/language.ts index caafa9eff3..a09f9bec2c 100644 --- a/src/renderer/src/stores/language.ts +++ b/src/renderer/src/stores/language.ts @@ -1,5 +1,5 @@ import { defineStore } from 'pinia' -import { onMounted, onScopeDispose, shallowRef } from 'vue' +import { onScopeDispose, shallowRef } from 'vue' import { useI18n } from 'vue-i18n' import { createConfigClient } from '@api/ConfigClient' @@ -88,9 +88,9 @@ export const useLanguageStore = defineStore('language', () => { await applyLanguageState(languageState, revision) } - onMounted(async () => { - await initLanguage() - }) + // Initialize at store setup top level (not in a component lifecycle hook); + // languageInitialization dedupes concurrent calls. + void initLanguage() onScopeDispose(() => { removeLanguageListener?.() diff --git a/src/renderer/src/stores/mcp.ts b/src/renderer/src/stores/mcp.ts index 5faed84194..6c54daabfb 100644 --- a/src/renderer/src/stores/mcp.ts +++ b/src/renderer/src/stores/mcp.ts @@ -1,4 +1,4 @@ -import { ref, computed, onMounted, onUnmounted, watch } from 'vue' +import { ref, computed, getCurrentScope, onScopeDispose, watch } from 'vue' import { defineStore } from 'pinia' import { createMcpClient } from '@api/McpClient' import { createConfigClient } from '../../api/ConfigClient' @@ -1183,17 +1183,17 @@ export const useMcpStore = defineStore('mcp', () => { } } - // 立即初始化 - onMounted(async () => { - await init() - }) - - onUnmounted(() => { - while (eventCleanups.length > 0) { - eventCleanups.pop()?.() - } - eventsBound = false - }) + // 立即初始化：订阅与数据加载下沉到 store setup 顶层（不挂在组件生命周期钩子上）， + // 避免首个消费组件卸载后丢失全局事件。 + void init() + if (getCurrentScope()) { + onScopeDispose(() => { + while (eventCleanups.length > 0) { + eventCleanups.pop()?.() + } + eventsBound = false + }) + } // 获取NPM Registry状态 const getNpmRegistryStatus = async () => { diff --git a/src/renderer/src/stores/mcpAppConsent.ts b/src/renderer/src/stores/mcpAppConsent.ts index 8fa760199b..6fbeef1af8 100644 --- a/src/renderer/src/stores/mcpAppConsent.ts +++ b/src/renderer/src/stores/mcpAppConsent.ts @@ -1,4 +1,4 @@ -import { computed, onMounted, onUnmounted, ref } from 'vue' +import { computed, getCurrentScope, onScopeDispose, ref } from 'vue' import { defineStore } from 'pinia' import { createMcpClient } from '@api/McpClient' import type { McpAppConsentRequestPayload } from '@shared/types/mcp' @@ -29,28 +29,30 @@ export const useMcpAppConsentStore = defineStore('mcpAppConsent', () => { } } - onMounted(() => { - eventCleanups.push( - mcpClient.onAppConsentRequest(({ request: next }) => { - if (queue.value.some((entry) => entry.requestId === next.requestId)) { - return - } - if (queue.value.length >= MAX_PENDING_APP_CONSENTS) { - void mcpClient.submitAppConsent(next.requestId, false).catch((error) => { - console.error('[MCP Apps] Failed to reject queued consent:', error) - }) - return - } - queue.value.push(next) - }) - ) - }) + // Subscribe at store setup top level (not in a component lifecycle hook) so global + // consent events are not lost when the first consuming component unmounts. + eventCleanups.push( + mcpClient.onAppConsentRequest(({ request: next }) => { + if (queue.value.some((entry) => entry.requestId === next.requestId)) { + return + } + if (queue.value.length >= MAX_PENDING_APP_CONSENTS) { + void mcpClient.submitAppConsent(next.requestId, false).catch((error) => { + console.error('[MCP Apps] Failed to reject queued consent:', error) + }) + return + } + queue.value.push(next) + }) + ) - onUnmounted(() => { - while (eventCleanups.length > 0) { - eventCleanups.pop()?.() - } - }) + if (getCurrentScope()) { + onScopeDispose(() => { + while (eventCleanups.length > 0) { + eventCleanups.pop()?.() + } + }) + } return { request, diff --git a/src/renderer/src/stores/mcpElicitation.ts b/src/renderer/src/stores/mcpElicitation.ts index 5313558233..173f1f2b96 100644 --- a/src/renderer/src/stores/mcpElicitation.ts +++ b/src/renderer/src/stores/mcpElicitation.ts @@ -1,4 +1,4 @@ -import { computed, onMounted, onUnmounted, ref } from 'vue' +import { computed, getCurrentScope, onScopeDispose, ref } from 'vue' import { defineStore } from 'pinia' import { createMcpClient } from '@api/McpClient' import { createBrowserClient } from '@api/BrowserClient' @@ -353,8 +353,10 @@ export const useMcpElicitationStore = defineStore('mcpElicitation', () => { } } - onMounted(() => { - eventCleanups.push( + // Subscribe at store setup top level (not in a component lifecycle hook) so the global + // elicitation listeners are not lost when the first consuming component unmounts. + const registerEvents = () => { + const cleanups = [ mcpClient.onElicitationRequest(({ request: next }) => queueOrOpenRequest(next)), mcpClient.onElicitationDecision(({ decision }) => { finishRequest(decision.requestId) @@ -362,14 +364,19 @@ export const useMcpElicitationStore = defineStore('mcpElicitation', () => { mcpClient.onElicitationCancelled(({ requestId }) => { finishRequest(requestId) }) - ) - }) + ].filter((cleanup): cleanup is () => void => typeof cleanup === 'function') - onUnmounted(() => { - while (eventCleanups.length > 0) { - eventCleanups.pop()?.() - } - }) + eventCleanups.push(...cleanups) + } + + registerEvents() + if (getCurrentScope()) { + onScopeDispose(() => { + while (eventCleanups.length > 0) { + eventCleanups.pop()?.() + } + }) + } return { request, diff --git a/src/renderer/src/stores/mcpSampling.ts b/src/renderer/src/stores/mcpSampling.ts index 3641b6f45f..b5ced4b36e 100644 --- a/src/renderer/src/stores/mcpSampling.ts +++ b/src/renderer/src/stores/mcpSampling.ts @@ -1,5 +1,5 @@ import { defineStore } from 'pinia' -import { ref, computed, onMounted, onUnmounted } from 'vue' +import { ref, computed, getCurrentScope, onScopeDispose } from 'vue' import { createMcpClient } from '@api/McpClient' import type { McpSamplingDecision, McpSamplingRequestPayload } from '@shared/types/mcp' import type { RENDERER_MODEL_META } from '@shared/types/provider' @@ -433,17 +433,19 @@ export const useMcpSamplingStore = defineStore('mcpSampling', () => { } } - onMounted(() => { - eventCleanups.push(mcpClient.onSamplingRequest(handleSamplingRequest)) - eventCleanups.push(mcpClient.onSamplingCancelled(handleSamplingCancelled)) - eventCleanups.push(mcpClient.onSamplingDecision(handleSamplingDecision)) - }) + // Subscribe at store setup top level (not in a component lifecycle hook) so global + // sampling events are not lost when the first consuming component unmounts. + eventCleanups.push(mcpClient.onSamplingRequest(handleSamplingRequest)) + eventCleanups.push(mcpClient.onSamplingCancelled(handleSamplingCancelled)) + eventCleanups.push(mcpClient.onSamplingDecision(handleSamplingDecision)) - onUnmounted(() => { - while (eventCleanups.length > 0) { - eventCleanups.pop()?.() - } - }) + if (getCurrentScope()) { + onScopeDispose(() => { + while (eventCleanups.length > 0) { + eventCleanups.pop()?.() + } + }) + } return { request, diff --git a/src/renderer/src/stores/shortcutKey.ts b/src/renderer/src/stores/shortcutKey.ts index 8939e2ee39..1a54697047 100644 --- a/src/renderer/src/stores/shortcutKey.ts +++ b/src/renderer/src/stores/shortcutKey.ts @@ -1,5 +1,5 @@ import { defineStore } from 'pinia' -import { onMounted, ref } from 'vue' +import { ref } from 'vue' import type { ShortcutKeySetting } from '@shared/types/desktop' import { createShortcutClient } from '@api/ShortcutClient' import { createConfigClient } from '../../api/ConfigClient' @@ -32,9 +32,9 @@ export const useShortcutKeyStore = defineStore('shortcutKey', () => { await shortcutClient.destroy() } - onMounted(async () => { - await loadShortcutKeys() - }) + // Load at store setup top level (not in a component lifecycle hook) so the data is + // fetched regardless of which component first uses the store. No cleanup semantics. + void loadShortcutKeys() return { shortcutKeys, diff --git a/src/renderer/src/stores/uiSettingsStore.ts b/src/renderer/src/stores/uiSettingsStore.ts index 8475f36b3d..065cfcb8e5 100644 --- a/src/renderer/src/stores/uiSettingsStore.ts +++ b/src/renderer/src/stores/uiSettingsStore.ts @@ -1,4 +1,4 @@ -import { computed, onBeforeUnmount, onMounted, ref } from 'vue' +import { computed, getCurrentScope, onScopeDispose, ref } from 'vue' import { defineStore } from 'pinia' import type { SettingsChange, SettingsSnapshotValues } from '@shared/contracts/routes' import { buildFontStack, DEFAULT_CODE_FONT_STACK, DEFAULT_TEXT_FONT_STACK } from '@/lib/fontStack' @@ -371,15 +371,17 @@ export const useUiSettingsStore = defineStore('uiSettings', () => { }) } - onMounted(() => { - void loadSettings() - setupListeners() - }) + // Load and subscribe at store setup top level (not in a component lifecycle hook) so + // the global settings listener is not lost when the first consuming component unmounts. + void loadSettings() + setupListeners() - onBeforeUnmount(() => { - unsubscribeFromSettings?.() - unsubscribeFromSettings = null - }) + if (getCurrentScope()) { + onScopeDispose(() => { + unsubscribeFromSettings?.() + unsubscribeFromSettings = null + }) + } return { fontSizeLevel, diff --git a/test/renderer/stores/mcpElicitation.test.ts b/test/renderer/stores/mcpElicitation.test.ts index eb8f5825c6..8979013760 100644 --- a/test/renderer/stores/mcpElicitation.test.ts +++ b/test/renderer/stores/mcpElicitation.test.ts @@ -3,12 +3,6 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' const submitElicitationDecisionMock = vi.hoisted(() => vi.fn().mockResolvedValue(undefined)) const onElicitationRequestMock = vi.hoisted(() => vi.fn()) -vi.mock('vue', async (importOriginal) => ({ - ...(await importOriginal()), - onMounted: (callback: () => void) => callback(), - onUnmounted: vi.fn() -})) - vi.mock('@api/McpClient', () => ({ createMcpClient: () => ({ submitElicitationDecision: submitElicitationDecisionMock, @@ -71,4 +65,11 @@ describe('MCP elicitation store', () => { expect(content['__proto__']).toBe('prototype-value') expect(content.toString).toBe('method-value') }) + + it('disposes cleanly when a listener registration returns no cleanup', async () => { + const store = await setupStore() + + expect(onElicitationRequestMock).toHaveBeenCalledTimes(1) + expect(() => store.$dispose()).not.toThrow() + }) }) diff --git a/test/renderer/stores/mcpStore.test.ts b/test/renderer/stores/mcpStore.test.ts index b8c73a9809..0a5d308e4b 100644 --- a/test/renderer/stores/mcpStore.test.ts +++ b/test/renderer/stores/mcpStore.test.ts @@ -1,11 +1,11 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' +import { flushPromises } from '@vue/test-utils' const setMcpServerEnabledMutate = vi.hoisted(() => vi.fn()) const addMcpServerMutate = vi.hoisted(() => vi.fn()) const updateMcpServerMutate = vi.hoisted(() => vi.fn()) const removeMcpServerMutate = vi.hoisted(() => vi.fn()) const configRefetch = vi.hoisted(() => vi.fn()) -const mountedCallbacks = vi.hoisted(() => [] as Array<() => Promise>) const mcpClientMock = vi.hoisted(() => ({ getMcpServers: vi.fn().mockResolvedValue({}), @@ -51,16 +51,6 @@ const createQueryState = () => ({ refetch: vi.fn(async () => ({ status: 'success', data: undefined })) }) -vi.mock('vue', async () => { - const actual = await vi.importActual('vue') - return { - ...actual, - onMounted: vi.fn((callback: () => Promise) => { - mountedCallbacks.push(callback) - }) - } -}) - vi.mock('@api/McpClient', () => ({ createMcpClient: vi.fn(() => mcpClientMock) })) @@ -108,13 +98,16 @@ const setupStore = async () => { const { createPinia, setActivePinia } = await vi.importActual('pinia') setActivePinia(createPinia()) const { useMcpStore } = await import('@/stores/mcp') - return useMcpStore() + const store = useMcpStore() + // The store subscribes and loads its data at setup top level, so let that initial pass settle + // before tests mutate the store state directly. + await flushPromises() + return store } describe('useMcpStore', () => { beforeEach(() => { vi.clearAllMocks() - mountedCallbacks.length = 0 setMcpServerEnabledMutate.mockReset() addMcpServerMutate.mockReset() updateMcpServerMutate.mockReset() @@ -306,7 +299,6 @@ describe('useMcpStore', () => { }) ) const store = await setupStore() - await mountedCallbacks[0]() store.config = { mcpServers: { demo: { diff --git a/test/renderer/stores/uiSettingsStore.test.ts b/test/renderer/stores/uiSettingsStore.test.ts index 789a535f0c..4878d7a15e 100644 --- a/test/renderer/stores/uiSettingsStore.test.ts +++ b/test/renderer/stores/uiSettingsStore.test.ts @@ -141,6 +141,7 @@ describe('uiSettingsStore', () => { expect(store.launchAtLoginEnabled).toBe(false) mountedWrappers = mountedWrappers.filter((candidate) => candidate !== wrapper) wrapper.unmount() + store.$dispose() expect(unsubscribe).toHaveBeenCalledTimes(1) }) From ad8295aee1072b3f9876efa996f9b063065ae58f Mon Sep 17 00:00:00 2001 From: xiao-test Date: Sat, 12 Sep 2026 18:15:57 +0800 Subject: [PATCH 12/15] perf(stream): dedupe snapshot serialization Renderer applyStreamingBlocksToMessage ran JSON.stringify on every 120ms snapshot regardless of whether main had actually changed blocks. DB flush is 600ms (4:1), so most snapshots repeat unchanged content. Add a monotonic blocksRevision counter to StreamState that bumps at every site that previously set dirty. Emit it as `revision` on chat.stream.updated (plus a per-instance rateLimitRevision counter on the rate_limit path). The renderer compares the new revision to the last applied one with an O(1) int compare and skips re-serialization when it did not advance. Sites replaced (state.dirty = true -> markStreamChanged): accumulator (10), dispatch (9 + 2 finalize-marked paths), process (5 incl. the normalizeInheritedUnresolvedBlocks special case), providerPermissionCoordinator (1), acp adapters (1). Schema uses nonnegative int because echo.flush() emits revision 0 on a no-change round. Lifecycle cleanup for the appliedStreamRevision Map: clear, purgeSessionTracking, applyPersistedMessageRecords, and commitSessionView replay (forced full re-fold across view swaps where the swapped-in record may lag the live fold). Tests cover accumulator bumps, echo payload revisions, the renderer fast path with a JSON.stringify spy (5 snapshots with 2 duplicate revisions yield 6 stringify calls instead of 10), and the persisted-record-cleanup path so a recycled stream re-folds instead of being short-circuited. --- src/main/agent/acp/compatibility/adapters.ts | 3 +- .../agent/deepchat/runtime/accumulator.ts | 22 ++-- .../deepchat/runtime/deepChatLoopRunner.ts | 3 + src/main/agent/deepchat/runtime/dispatch.ts | 42 ++++--- src/main/agent/deepchat/runtime/echo.ts | 1 + src/main/agent/deepchat/runtime/process.ts | 13 +- .../runtime/providerPermissionCoordinator.ts | 3 +- src/main/agent/deepchat/runtime/types.ts | 9 +- src/renderer/src/stores/ui/message.ts | 69 +++++++++-- src/renderer/src/stores/ui/messageIpc.ts | 19 ++- src/renderer/src/stores/ui/stream.ts | 7 +- src/shared/contracts/events/chat.events.ts | 1 + .../deepchat/runtime/accumulator.test.ts | 25 ++++ test/main/agent/deepchat/runtime/echo.test.ts | 63 +++++++++- test/renderer/stores/messageStore.test.ts | 111 ++++++++++++++++++ 15 files changed, 337 insertions(+), 54 deletions(-) diff --git a/src/main/agent/acp/compatibility/adapters.ts b/src/main/agent/acp/compatibility/adapters.ts index fd92e3fb47..28a467cff8 100644 --- a/src/main/agent/acp/compatibility/adapters.ts +++ b/src/main/agent/acp/compatibility/adapters.ts @@ -20,6 +20,7 @@ import { import { createAcpPromptTerminalEvents } from '@/agent/acp/runtime/acpContentMapper' import { createState, + markStreamChanged, type DeepChatEventPublisher, type DeepChatSessionUpdatePublisher, type IoParams, @@ -159,7 +160,7 @@ export class AcpCompatibilityProjectionAdapter implements AcpCompatibilityProjec ? block.extra.permissionType : 'all' markStreamingProviderPermissionResolved(block, granted, permissionType) - state.stream.dirty = true + markStreamChanged(state.stream) this.flushIfDirty(state) } diff --git a/src/main/agent/deepchat/runtime/accumulator.ts b/src/main/agent/deepchat/runtime/accumulator.ts index bbbfe4b3c9..55ab0acf5e 100644 --- a/src/main/agent/deepchat/runtime/accumulator.ts +++ b/src/main/agent/deepchat/runtime/accumulator.ts @@ -4,7 +4,7 @@ import type { ProviderUrlSourcePayload } from '@shared/types/core/llm-events' import type { ChatMessageProviderOptions } from '@shared/types/core/chat-message' -import type { StreamState } from './types' +import { markStreamChanged, type StreamState } from './types' const MAX_VISIBLE_SEARCH_PAGES = 6 @@ -144,7 +144,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void if (state.firstTokenTime === null) state.firstTokenTime = Date.now() const block = getCurrentBlock(state.blocks, 'content', event.provider_options) block.content += event.content - state.dirty = true + markStreamChanged(state) break } case 'reasoning': { @@ -167,12 +167,12 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void } const reasoningTime = block.reasoning_time as { start: number; end: number } updateReasoningMetadata(state, reasoningTime.start, reasoningTime.end) - state.dirty = true + markStreamChanged(state) break } case 'plan': { if (finalizeTrailingPendingNarrativeBlocks(state.blocks)) { - state.dirty = true + markStreamChanged(state) } const revision = event.revision ?? (state.latestAgentPlanSnapshot?.revision ?? 0) + 1 state.latestAgentPlanSnapshot = { @@ -209,7 +209,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void executionOwner: event.tool_call_execution_owner ?? 'deepchat', providerOptions: event.provider_options }) - state.dirty = true + markStreamChanged(state) break } case 'tool_call_chunk': { @@ -229,7 +229,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void } } } - state.dirty = true + markStreamChanged(state) } break } @@ -265,7 +265,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void }) } state.pendingToolCalls.delete(event.tool_call_id) - state.dirty = true + markStreamChanged(state) } break } @@ -296,12 +296,12 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void } } state.blocks.push(block) - state.dirty = true + markStreamChanged(state) break } case 'provider_url_source': { if (appendProviderUrlSource(state.blocks, event.provider_url_source)) { - state.dirty = true + markStreamChanged(state) } break } @@ -318,7 +318,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void } } state.blocks.push(block) - state.dirty = true + markStreamChanged(state) break } case 'usage': { @@ -354,7 +354,7 @@ export function accumulate(state: StreamState, event: LLMCoreStreamEvent): void if (block.status === 'pending') block.status = 'error' } state.stopReason = 'error' - state.dirty = true + markStreamChanged(state) break } default: diff --git a/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts b/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts index 5272bf64b9..89ddf35b4d 100644 --- a/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts +++ b/src/main/agent/deepchat/runtime/deepChatLoopRunner.ts @@ -661,6 +661,7 @@ export function buildTapeViewSelection( export class DeepChatLoopRunner { private readonly toolSurfaceAdapterHistory = new ToolSurfaceAdapterHistory() + private rateLimitRevision = 0 constructor(private readonly ports: DeepChatLoopRunnerPorts) {} @@ -2726,6 +2727,7 @@ export class DeepChatLoopRunner { sessionId, messageId, updatedAt: Date.now(), + revision: ++this.rateLimitRevision, blocks: cloneBlocksForRenderer([block]) }) } @@ -2737,6 +2739,7 @@ export class DeepChatLoopRunner { sessionId, messageId, updatedAt: Date.now(), + revision: ++this.rateLimitRevision, blocks: [] }) } diff --git a/src/main/agent/deepchat/runtime/dispatch.ts b/src/main/agent/deepchat/runtime/dispatch.ts index 823506d822..cce3e77449 100644 --- a/src/main/agent/deepchat/runtime/dispatch.ts +++ b/src/main/agent/deepchat/runtime/dispatch.ts @@ -32,6 +32,7 @@ import type { ToolCallResult, ToolDispatchCollaborators } from './types' +import { markStreamChanged } from './types' import type { ChatMessage, ChatMessageProviderOptions, @@ -1253,7 +1254,7 @@ function finalizePendingNarrativeBeforeToolSettlement(state: StreamState): void } finalizeTrailingPendingNarrativeBlocks(state.blocks) - state.dirty = true + markStreamChanged(state) } function applyFinalizedToolResults(params: { @@ -1387,7 +1388,7 @@ function applyFinalizedToolResults(params: { } } - state.dirty = true + markStreamChanged(state) return interactions } @@ -1673,7 +1674,7 @@ async function reviewAutoApproveAction(params: { } if (setToolCallAutoApproveReviewing(batchToolCallBlocks, execution.completedToolCall.id, true)) { - state.dirty = true + markStreamChanged(state) rendererFlushHandle.flush() } try { @@ -1705,7 +1706,7 @@ async function reviewAutoApproveAction(params: { if ( setToolCallAutoApproveReviewing(batchToolCallBlocks, execution.completedToolCall.id, false) ) { - state.dirty = true + markStreamChanged(state) rendererFlushHandle.flush() } } @@ -1772,7 +1773,7 @@ function appendPermissionActionBlock( ...(permission.rememberable === false ? { rememberable: false } : {}) } }) - state.dirty = true + markStreamChanged(state) return { type: 'permission', origin, @@ -1830,7 +1831,7 @@ function appendQuestionActionBlock( ...extra } }) - state.dirty = true + markStreamChanged(state) return { type: 'question', origin, @@ -1896,8 +1897,8 @@ function appendSkillDraftQuestionActionBlock( ) } -function flushBlocksToRenderer(io: IoParams, blocks: AssistantMessageBlock[]): void { - const renderedBlocks = cloneBlocksForRenderer(blocks) +function flushBlocksToRenderer(io: IoParams, state: StreamState): void { + const renderedBlocks = cloneBlocksForRenderer(state.blocks) io.publishEvent('chat.stream.updated', { kind: 'snapshot', requestId: io.requestId, @@ -1906,6 +1907,7 @@ function flushBlocksToRenderer(io: IoParams, blocks: AssistantMessageBlock[]): v providerId: io.providerId, modelId: io.modelId, updatedAt: Date.now(), + revision: state.blocksRevision, blocks: renderedBlocks }) @@ -1914,10 +1916,10 @@ function flushBlocksToRenderer(io: IoParams, blocks: AssistantMessageBlock[]): v kind: 'blocks', updatedAt: Date.now(), messageId: io.messageId, - previewMarkdown: buildAssistantPreviewMarkdown(blocks), - responseMarkdown: buildAssistantResponseMarkdown(blocks), - deliverySegments: buildAssistantDeliverySegments(io.messageId, blocks), - waitingInteraction: extractWaitingInteraction(blocks, io.messageId) + previewMarkdown: buildAssistantPreviewMarkdown(state.blocks), + responseMarkdown: buildAssistantResponseMarkdown(state.blocks), + deliverySegments: buildAssistantDeliverySegments(io.messageId, state.blocks), + waitingInteraction: extractWaitingInteraction(state.blocks, io.messageId) }) } @@ -2130,7 +2132,7 @@ async function runToolCall(params: { } state.latestAgentPlanSnapshot = snapshot publishPlanUpdated(io, snapshot) - state.dirty = true + markStreamChanged(state) scheduleRendererFlush(state, rendererFlushHandle) return } @@ -2149,7 +2151,7 @@ async function runToolCall(params: { update.responseMarkdown, update.progressJson ) - state.dirty = true + markStreamChanged(state) scheduleRendererFlush(state, rendererFlushHandle) } @@ -3204,7 +3206,7 @@ export async function settleToolBatch( content: errorText }) updateToolCallBlock(batchToolCallBlocks, tc.id, errorText, true) - state.dirty = true + markStreamChanged(state) batchState.committedResultCallIds.add(tc.id) executed += 1 persistToolExecutionState(io, state, rendererFlushHandle) @@ -3534,7 +3536,7 @@ export function finalizePaused(state: StreamState, io: IoParams): void { stampGenerationTiming(state) io.messageStore.updateAssistantContent(io.messageId, state.blocks, JSON.stringify(state.metadata)) - flushBlocksToRenderer(io, state.blocks) + flushBlocksToRenderer(io, state) io.publishEvent('chat.stream.completed', { requestId: io.requestId, sessionId: io.sessionId, @@ -3547,6 +3549,7 @@ export function finalize(state: StreamState, io: IoParams): void { for (const block of state.blocks) { if (block.status === 'pending') block.status = 'success' } + markStreamChanged(state) stampPlanTerminalIfOpen(state, io, state.planTerminalReason) stampGenerationTiming(state) @@ -3556,7 +3559,7 @@ export function finalize(state: StreamState, io: IoParams): void { state.blocks, JSON.stringify(state.metadata) ) - flushBlocksToRenderer(io, state.blocks) + flushBlocksToRenderer(io, state) io.publishEvent('chat.stream.completed', { requestId: io.requestId, sessionId: io.sessionId, @@ -3568,6 +3571,7 @@ export function finalize(state: StreamState, io: IoParams): void { export function finalizeError(state: StreamState, io: IoParams, error: unknown): void { const errorMessage = error instanceof Error ? error.message : String(error) state.blocks = buildTerminalErrorBlocks(state.blocks, errorMessage) + markStreamChanged(state) stampPlanTerminalIfOpen( state, io, @@ -3577,7 +3581,7 @@ export function finalizeError(state: StreamState, io: IoParams, error: unknown): stampGenerationTiming(state) io.messageStore.setMessageError(io.messageId, state.blocks, JSON.stringify(state.metadata)) - flushBlocksToRenderer(io, state.blocks) + flushBlocksToRenderer(io, state) io.publishEvent('chat.stream.failed', { requestId: io.requestId, sessionId: io.sessionId, @@ -3596,5 +3600,5 @@ export function persistAbortExceptionPlanState(state: StreamState, io: IoParams) } io.messageStore.updateAssistantContent(io.messageId, state.blocks) - flushBlocksToRenderer(io, state.blocks) + flushBlocksToRenderer(io, state) } diff --git a/src/main/agent/deepchat/runtime/echo.ts b/src/main/agent/deepchat/runtime/echo.ts index 98db4f5480..5d8fcc1b8a 100644 --- a/src/main/agent/deepchat/runtime/echo.ts +++ b/src/main/agent/deepchat/runtime/echo.ts @@ -23,6 +23,7 @@ export function startEcho(state: StreamState, io: IoParams): EchoHandle { providerId: io.providerId, modelId: io.modelId, updatedAt: Date.now(), + revision: state.blocksRevision, blocks: renderedBlocks }) } diff --git a/src/main/agent/deepchat/runtime/process.ts b/src/main/agent/deepchat/runtime/process.ts index 67c04c428d..4b8c2c1356 100644 --- a/src/main/agent/deepchat/runtime/process.ts +++ b/src/main/agent/deepchat/runtime/process.ts @@ -13,6 +13,7 @@ import type { StreamState, ToolCallResult } from './types' +import { markStreamChanged } from './types' import { accumulate, commitRoundUsage, finalizeTrailingPendingNarrativeBlocks } from './accumulator' import { startEcho } from './echo' import { @@ -250,7 +251,7 @@ function markUnexecutedToolCallsForLimit(state: StreamState): void { ...block.extra, toolCallSkippedReason: 'max_tool_calls' } - state.dirty = true + markStreamChanged(state) } } @@ -364,7 +365,7 @@ function markOtherTruncatedToolCallsIncomplete( ...block.extra, toolCallIncompleteReason: 'max_tokens' } - state.dirty = true + markStreamChanged(state) } } @@ -655,7 +656,7 @@ export function appendStreamingProviderPermissionBlock( } state.blocks.push(actionBlock) - state.dirty = true + markStreamChanged(state) return { actionBlock, @@ -971,7 +972,9 @@ export async function processStream(params: ProcessParams): Promise 0) { state.blocks = JSON.parse(JSON.stringify(initialBlocks)) as typeof state.blocks - state.dirty = normalizeInheritedUnresolvedBlocks(state.blocks) || state.dirty + if (normalizeInheritedUnresolvedBlocks(state.blocks)) { + markStreamChanged(state) + } } state.metadata.runId = run.runId const echo = startEcho(state, io) @@ -1249,7 +1252,7 @@ export async function processStream(params: ProcessParams): Promise void @@ -363,6 +364,12 @@ export function createState(): StreamState { stopReason: null, roundUsage: null, toolCallCount: 0, - dirty: false + dirty: false, + blocksRevision: 0 } } + +export function markStreamChanged(state: StreamState): void { + state.dirty = true + state.blocksRevision += 1 +} diff --git a/src/renderer/src/stores/ui/message.ts b/src/renderer/src/stores/ui/message.ts index 4e7331f1c1..c781235db1 100644 --- a/src/renderer/src/stores/ui/message.ts +++ b/src/renderer/src/stores/ui/message.ts @@ -56,6 +56,10 @@ export const useMessageStore = defineStore('message', () => { const currentStreamRequestId = toStoreStateRef(streamStateStore, 'currentStreamRequestId') const currentStreamMessageId = toStoreStateRef(streamStateStore, 'currentStreamMessageId') const currentStreamMetadata = toStoreStateRef(streamStateStore, 'currentStreamMetadata') + const currentStreamBlocksRevision = toStoreStateRef( + streamStateStore, + 'currentStreamBlocksRevision' + ) const streamRevision = toStoreStateRef(streamStateStore, 'streamRevision') // --- State --- @@ -81,6 +85,7 @@ export const useMessageStore = defineStore('message', () => { // Stream message ids currently being hydrated into the cache as a placeholder // record (before the backend persists them). Prevents re-entrant duplicate inserts. const hydratingStreamMessageIds = new Set() + const appliedStreamRevision = new Map() let latestLoadRequestId = 0 let latestHistoryRequestId = 0 let latestLoadSessionId: string | null = null @@ -581,11 +586,13 @@ export const useMessageStore = defineStore('message', () => { streamMessageId && !isEphemeralStreamMessageId(streamMessageId) ) { + appliedStreamRevision.delete(streamMessageId) applyStreamingBlocksToMessage( streamMessageId, view.sessionId, streamingBlocks.value as AssistantMessageBlock[], - currentStreamMetadata.value ?? undefined + currentStreamMetadata.value ?? undefined, + currentStreamBlocksRevision.value ) } } @@ -938,6 +945,9 @@ export const useMessageStore = defineStore('message', () => { markLiveMessageViewMutation(sessionId) for (const record of changedRecords) { parsedMessageCache.delete(record.id) + // Persisted data replaces the live-folded record; drop the applied + // revision so a later snapshot re-evaluates content from scratch. + appliedStreamRevision.delete(record.id) upsertMessageRecord(record) } lastPersistedRevision.value += 1 @@ -958,6 +968,7 @@ export const useMessageStore = defineStore('message', () => { historyLoadError.value = false parsedMessageCache.clear() hydratingStreamMessageIds.clear() + appliedStreamRevision.clear() recentSessionViews.clear() messageMutationRevisions.clear() recentViewInvalidationRevisions.clear() @@ -990,17 +1001,28 @@ export const useMessageStore = defineStore('message', () => { messageId: string, conversationId: string, blocks: AssistantMessageBlock[], - metadata?: { providerId?: string; modelId?: string } + metadata?: { providerId?: string; modelId?: string }, + revision?: number ): void { if (committedSessionId.value !== conversationId) return - const serializedBlocks = JSON.stringify(blocks) - const serializedMetadata = JSON.stringify({ - ...(metadata?.providerId ? { provider: metadata.providerId } : {}), - ...(metadata?.modelId ? { model: metadata.modelId } : {}) - }) const existing = messageCache.value.get(messageId) if (existing) { if (existing.sessionId !== conversationId) return + + const lastRevision = appliedStreamRevision.get(messageId) ?? 0 + if (revision !== undefined && revision <= lastRevision && existing.status === 'pending') { + cacheStreamingAssistantBlocks(existing, blocks) + return + } + + const serializedBlocks = JSON.stringify(blocks) + const serializedMetadata = JSON.stringify({ + ...(metadata?.providerId ? { provider: metadata.providerId } : {}), + ...(metadata?.modelId ? { model: metadata.modelId } : {}) + }) + if (revision !== undefined) { + appliedStreamRevision.set(messageId, revision) + } const nextMetadata = serializedMetadata === '{}' ? existing.metadata : serializedMetadata if ( existing.content === serializedBlocks && @@ -1023,6 +1045,11 @@ export const useMessageStore = defineStore('message', () => { return } + const serializedBlocks = JSON.stringify(blocks) + const serializedMetadata = JSON.stringify({ + ...(metadata?.providerId ? { provider: metadata.providerId } : {}), + ...(metadata?.modelId ? { model: metadata.modelId } : {}) + }) if (hydratingStreamMessageIds.has(messageId)) return hydratingStreamMessageIds.add(messageId) markLiveMessageViewMutation(conversationId) @@ -1041,6 +1068,9 @@ export const useMessageStore = defineStore('message', () => { createdAt: now, updatedAt: now } + if (revision !== undefined) { + appliedStreamRevision.set(messageId, revision) + } upsertMessageRecord(nextRecord) cacheStreamingAssistantBlocks(nextRecord, blocks) hydratingStreamMessageIds.delete(messageId) @@ -1052,8 +1082,24 @@ export const useMessageStore = defineStore('message', () => { sessionId: currentStreamSessionId.value, requestId: currentStreamRequestId.value }), - setStreamingState: ({ sessionId, requestId, messageId, updatedAt, blocks, metadata }) => { - streamStateStore.setStream(sessionId, blocks, messageId, metadata, requestId, updatedAt) + setStreamingState: ({ + sessionId, + requestId, + messageId, + updatedAt, + revision, + blocks, + metadata + }) => { + streamStateStore.setStream( + sessionId, + blocks, + messageId, + metadata, + requestId, + updatedAt, + revision + ) }, clearStreamingState, loadMessages, @@ -1065,6 +1111,11 @@ export const useMessageStore = defineStore('message', () => { registerStoreCleanup(messageIpcBinding.cleanup) function purgeSessionTracking(sessionId: string): void { + for (const [id, record] of messageCache.value) { + if (record.sessionId === sessionId) { + appliedStreamRevision.delete(id) + } + } recentSessionViews.delete(sessionId) messageMutationRevisions.delete(sessionId) recentViewInvalidationRevisions.delete(sessionId) diff --git a/src/renderer/src/stores/ui/messageIpc.ts b/src/renderer/src/stores/ui/messageIpc.ts index 8cfb42479b..0d7953f838 100644 --- a/src/renderer/src/stores/ui/messageIpc.ts +++ b/src/renderer/src/stores/ui/messageIpc.ts @@ -13,6 +13,7 @@ interface BindMessageStoreIpcOptions { requestId: string messageId?: string updatedAt: number + revision: number blocks: AssistantMessageBlock[] metadata?: { providerId?: string; modelId?: string } }) => void @@ -24,7 +25,8 @@ interface BindMessageStoreIpcOptions { messageId: string, sessionId: string, blocks: AssistantMessageBlock[], - metadata?: { providerId?: string; modelId?: string } + metadata?: { providerId?: string; modelId?: string }, + revision?: number ) => void isEphemeralStreamMessageId: (messageId: string) => boolean } @@ -180,6 +182,7 @@ export function bindMessageStoreIpc(options: BindMessageStoreIpcOptions): Messag requestId: payload.requestId, messageId: streamMessageId, updatedAt: payload.updatedAt, + revision: payload.revision, blocks, metadata: { providerId: payload.providerId, @@ -192,10 +195,16 @@ export function bindMessageStoreIpc(options: BindMessageStoreIpcOptions): Messag options.applyStreamingBlocksToMessage && !options.isEphemeralStreamMessageId(streamMessageId) ) { - options.applyStreamingBlocksToMessage(streamMessageId, payload.sessionId, blocks, { - providerId: payload.providerId, - modelId: payload.modelId - }) + options.applyStreamingBlocksToMessage( + streamMessageId, + payload.sessionId, + blocks, + { + providerId: payload.providerId, + modelId: payload.modelId + }, + payload.revision + ) } }), chatClient.onStreamCompleted((payload) => { diff --git a/src/renderer/src/stores/ui/stream.ts b/src/renderer/src/stores/ui/stream.ts index 9cc7704ec0..6f458f3d5e 100644 --- a/src/renderer/src/stores/ui/stream.ts +++ b/src/renderer/src/stores/ui/stream.ts @@ -13,6 +13,7 @@ export const useStreamStateStore = defineStore('streamState', () => { const currentStreamUpdatedAt = ref(0) const currentStreamMetadata = ref<{ providerId?: string; modelId?: string } | null>(null) const streamRevision = ref(0) + const currentStreamBlocksRevision = ref(0) function setStream( sessionId: string, @@ -20,7 +21,8 @@ export const useStreamStateStore = defineStore('streamState', () => { messageId?: string, metadata?: { providerId?: string; modelId?: string }, requestId?: string, - updatedAt?: number + updatedAt?: number, + blocksRevision?: number ): void { isStreaming.value = true currentStreamSessionId.value = sessionId @@ -28,6 +30,7 @@ export const useStreamStateStore = defineStore('streamState', () => { currentStreamMessageId.value = messageId ?? null currentStreamUpdatedAt.value = updatedAt ?? 0 currentStreamMetadata.value = metadata ?? null + currentStreamBlocksRevision.value = blocksRevision ?? 0 streamingBlocks.value = blocks streamRevision.value += 1 } @@ -40,6 +43,7 @@ export const useStreamStateStore = defineStore('streamState', () => { currentStreamMessageId.value = null currentStreamUpdatedAt.value = 0 currentStreamMetadata.value = null + currentStreamBlocksRevision.value = 0 streamRevision.value += 1 } @@ -51,6 +55,7 @@ export const useStreamStateStore = defineStore('streamState', () => { currentStreamMessageId, currentStreamUpdatedAt, currentStreamMetadata, + currentStreamBlocksRevision, streamRevision, setStream, clearStreamingState diff --git a/src/shared/contracts/events/chat.events.ts b/src/shared/contracts/events/chat.events.ts index 82666c854f..8d4349588e 100644 --- a/src/shared/contracts/events/chat.events.ts +++ b/src/shared/contracts/events/chat.events.ts @@ -17,6 +17,7 @@ export const chatStreamUpdatedEvent = defineEventContract({ providerId: z.string().optional(), modelId: z.string().optional(), updatedAt: TimestampMsSchema, + revision: z.number().int().nonnegative(), blocks: z.array(AssistantMessageBlockSchema) }) }) diff --git a/test/main/agent/deepchat/runtime/accumulator.test.ts b/test/main/agent/deepchat/runtime/accumulator.test.ts index 33c87e320a..c2cd8b96d8 100644 --- a/test/main/agent/deepchat/runtime/accumulator.test.ts +++ b/test/main/agent/deepchat/runtime/accumulator.test.ts @@ -473,6 +473,31 @@ describe('accumulate', () => { expect(state.dirty).toBe(false) }) + it('bumps blocksRevision exactly when block content changes', () => { + expect(state.blocksRevision).toBe(0) + + accumulate(state, { type: 'text', content: 'first' }) + expect(state.blocksRevision).toBe(1) + + accumulate(state, { type: 'text', content: ' second' }) + expect(state.blocksRevision).toBe(2) + + state.dirty = false + accumulate(state, { + type: 'usage', + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 } + }) + accumulate(state, { type: 'stop', stop_reason: 'tool_use' }) + expect(state.blocksRevision).toBe(2) + + accumulate(state, { + type: 'tool_call_start', + tool_call_id: 'tc1', + tool_call_name: 'search' + }) + expect(state.blocksRevision).toBe(3) + }) + it('sets firstTokenTime once on first text event', () => { expect(state.firstTokenTime).toBeNull() diff --git a/test/main/agent/deepchat/runtime/echo.test.ts b/test/main/agent/deepchat/runtime/echo.test.ts index 8a46e3e06d..7dd95cc24b 100644 --- a/test/main/agent/deepchat/runtime/echo.test.ts +++ b/test/main/agent/deepchat/runtime/echo.test.ts @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' import type { StreamState, IoParams } from '@/agent/deepchat/runtime/types' -import { createState } from '@/agent/deepchat/runtime/types' +import { createState, markStreamChanged } from '@/agent/deepchat/runtime/types' vi.mock('@/events', () => ({ STREAM_EVENTS: { @@ -11,6 +11,7 @@ vi.mock('@/events', () => ({ })) import { startEcho } from '@/agent/deepchat/runtime/echo' +import { accumulate } from '@/agent/deepchat/runtime/accumulator' import { cloneBlocksForRenderer } from '@/session/clientMessageProjection' const publishDeepchatEvent = vi.fn() @@ -180,6 +181,66 @@ describe('echo', () => { echo.stop() }) + it('emits the current blocksRevision with every renderer snapshot', () => { + const echo = startEcho(state, io) + + state.blocks.push({ type: 'content', content: 'hi', status: 'pending', timestamp: Date.now() }) + markStreamChanged(state) + + echo.flush() + expect(publishDeepchatEvent).toHaveBeenCalledWith( + 'chat.stream.updated', + expect.objectContaining({ revision: 1 }) + ) + + echo.flush() + expect(publishDeepchatEvent).toHaveBeenLastCalledWith( + 'chat.stream.updated', + expect.objectContaining({ revision: 1 }) + ) + + state.blocks.push({ type: 'content', content: 'more', status: 'pending', timestamp: Date.now() }) + markStreamChanged(state) + echo.flush() + expect(publishDeepchatEvent).toHaveBeenLastCalledWith( + 'chat.stream.updated', + expect.objectContaining({ revision: 2 }) + ) + + echo.stop() + }) + + it('drives a multi-token stream end-to-end with monotonic revisions per dirty bump', () => { + const echo = startEcho(state, io) + + accumulate(state, { type: 'text', content: 'Hello ' }) + accumulate(state, { type: 'text', content: 'world' }) + echo.schedule() + vi.advanceTimersByTime(130) + let flushes = getStreamUpdatedCalls() + expect(flushes).toHaveLength(1) + expect(flushes[0]?.[1]).toMatchObject({ revision: 2 }) + + vi.advanceTimersByTime(150) + accumulate(state, { type: 'text', content: '!' }) + echo.schedule() + vi.advanceTimersByTime(130) + flushes = getStreamUpdatedCalls() + expect(flushes).toHaveLength(2) + expect(flushes[1]?.[1]).toMatchObject({ revision: 3 }) + + echo.flush() + echo.flush() + flushes = getStreamUpdatedCalls() + expect(flushes).toHaveLength(4) + expect(flushes[2]?.[1]).toMatchObject({ revision: 3 }) + expect(flushes[3]?.[1]).toMatchObject({ revision: 3 }) + + expect(state.blocksRevision).toBe(3) + + echo.stop() + }) + it('rescheduleRenderer() resets the renderer flush window from the latest interaction', () => { const echo = startEcho(state, io) diff --git a/test/renderer/stores/messageStore.test.ts b/test/renderer/stores/messageStore.test.ts index f378e21fa4..38b6ea2820 100644 --- a/test/renderer/stores/messageStore.test.ts +++ b/test/renderer/stores/messageStore.test.ts @@ -1601,4 +1601,115 @@ describe('messageStore', () => { const updatedBlocks = store.getAssistantMessageBlocks(store.messages.value[0]!) expect(updatedBlocks[0]).not.toBe(firstBlocks[0]) }) + + it('skips duplicate snapshots by revision and applies bumped revisions', async () => { + const { store, streamListeners } = await setupStore() + await store.loadMessages('s1') + + const emit = (revision: number, text: string, updatedAt: number) => + streamListeners.updated[0]({ + sessionId: 's1', + requestId: 'm1', + messageId: 'm1', + providerId: 'acp', + modelId: 'dimcode', + updatedAt, + revision, + blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] + }) + + emit(1, 'hello', 1) + const firstRecord = store.messageCache.value.get('m1')! + expect(firstRecord).toBeDefined() + expect(firstRecord.content).toContain('hello') + + emit(1, 'hello-again', 1) + const afterDuplicate = store.messageCache.value.get('m1')! + expect(afterDuplicate.content).toContain('hello') + expect(afterDuplicate.content).not.toContain('hello-again') + expect(afterDuplicate.updatedAt).toBe(firstRecord.updatedAt) + + emit(2, 'hello-again', 2) + const afterAdvance = store.messageCache.value.get('m1')! + expect(afterAdvance.content).toContain('hello-again') + }) + + it('skips JSON.stringify when the stream revision did not advance (quantified savings)', async () => { + const { store, streamListeners } = await setupStore() + await store.loadMessages('s1') + + const emit = (revision: number, text: string, updatedAt: number) => + streamListeners.updated[0]({ + sessionId: 's1', + requestId: 'm1', + messageId: 'm1', + providerId: 'acp', + modelId: 'dimcode', + updatedAt, + revision, + blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] + }) + + const stringifySpy = vi.spyOn(JSON, 'stringify') + const baseline = stringifySpy.mock.calls.length + + emit(1, 'a', 1) + emit(1, 'a', 1) + emit(2, 'ab', 2) + emit(2, 'ab', 2) + emit(3, 'abc', 3) + + const stringifyCalls = stringifySpy.mock.calls.length - baseline + stringifySpy.mockRestore() + + expect(stringifyCalls).toBe(6) + expect(stringifyCalls).toBeLessThan(10) + expect(store.messageCache.value.get('m1')?.content).toContain('abc') + }) + + it('drops the applied revision on persisted record arrival so a recycled stream re-folds', async () => { + const { store, streamListeners, messageListeners } = await setupStore() + await store.loadMessages('s1') + + const emit = (revision: number, text: string, updatedAt: number) => + streamListeners.updated[0]({ + sessionId: 's1', + requestId: 'm1', + messageId: 'm1', + providerId: 'acp', + modelId: 'dimcode', + updatedAt, + revision, + blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] + }) + + emit(1, 'first', 1) + emit(2, 'second', 2) + + messageListeners[0]({ + sessionId: 's1', + messages: [ + { + id: 'm1', + sessionId: 's1', + orderSeq: 1, + role: 'assistant' as const, + content: JSON.stringify([ + { type: 'content', content: 'persisted', status: 'success', timestamp: 5 } + ]), + status: 'success' as const, + isContextEdge: 0, + metadata: '{}', + traceCount: 0, + hasNestedExecutionAudit: false, + createdAt: 1, + updatedAt: Date.now() + 10_000 + } + ] + }) + expect(store.messageCache.value.get('m1')?.content).toContain('persisted') + + emit(1, 'recycled-stream', 6) + expect(store.messageCache.value.get('m1')?.content).toContain('recycled-stream') + }) }) From 19fa87db6a05f2c68c190a6f8f9531162a98f8f6 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Sun, 13 Sep 2026 15:17:49 +0800 Subject: [PATCH 13/15] fix(stream): apply the first revision-0 snapshot The message store's dedupe guard defaulted a missing applied revision to 0, so the first revision-0 snapshot after a resume (initialBlocks restore) was swallowed: only the parsed live-block cache refreshed while the pending record kept stale content, leaving the chat view stuck until a later revision bump. Only compare revisions when one was actually recorded, keeping "untracked" distinct from 0. Adds a renderer regression test for the first-snapshot update plus continued dedupe, and syncs main fixtures with the required revision field on chat.stream.updated. --- src/renderer/src/stores/ui/message.ts | 9 +++- test/main/events/typedEventHub.test.ts | 6 +++ test/main/routes/contracts.test.ts | 1 + test/renderer/stores/messageStore.test.ts | 52 +++++++++++++++++++++++ 4 files changed, 66 insertions(+), 2 deletions(-) diff --git a/src/renderer/src/stores/ui/message.ts b/src/renderer/src/stores/ui/message.ts index c781235db1..6ca4ad4203 100644 --- a/src/renderer/src/stores/ui/message.ts +++ b/src/renderer/src/stores/ui/message.ts @@ -1009,8 +1009,13 @@ export const useMessageStore = defineStore('message', () => { if (existing) { if (existing.sessionId !== conversationId) return - const lastRevision = appliedStreamRevision.get(messageId) ?? 0 - if (revision !== undefined && revision <= lastRevision && existing.status === 'pending') { + const lastRevision = appliedStreamRevision.get(messageId) + if ( + revision !== undefined && + lastRevision !== undefined && + revision <= lastRevision && + existing.status === 'pending' + ) { cacheStreamingAssistantBlocks(existing, blocks) return } diff --git a/test/main/events/typedEventHub.test.ts b/test/main/events/typedEventHub.test.ts index aef32a461f..3a9d2cba33 100644 --- a/test/main/events/typedEventHub.test.ts +++ b/test/main/events/typedEventHub.test.ts @@ -187,6 +187,7 @@ describe('TypedEventHub', () => { sessionId: 'run-1', messageId: 'message-1', updatedAt: 1, + revision: 0, blocks: [] }, { kind: 'run', runId: 'run-1' } @@ -201,6 +202,7 @@ describe('TypedEventHub', () => { sessionId: 'run-1', messageId: 'message-1', updatedAt: 2, + revision: 1, blocks: [] }, { kind: 'run', runId: 'run-1' } @@ -270,6 +272,7 @@ describe('SessionEventRouter', () => { sessionId: 'cli-run', messageId: 'message-1', updatedAt: 123, + revision: 0, blocks: [] }) @@ -282,6 +285,7 @@ describe('SessionEventRouter', () => { sessionId: 'cli-run', messageId: 'message-1', updatedAt: 123, + revision: 0, blocks: [] } }) @@ -305,6 +309,7 @@ describe('SessionEventRouter', () => { sessionId: 'session-1', messageId: 'message-1', updatedAt: 123, + revision: 0, blocks: [] } router.publish('chat.stream.updated', payload) @@ -336,6 +341,7 @@ describe('SessionEventRouter', () => { sessionId: 'session-1', messageId: 'message-1', updatedAt: 123, + revision: 0, blocks: [] }) diff --git a/test/main/routes/contracts.test.ts b/test/main/routes/contracts.test.ts index 7a47372492..d08ce96f34 100644 --- a/test/main/routes/contracts.test.ts +++ b/test/main/routes/contracts.test.ts @@ -2277,6 +2277,7 @@ describe('main kernel contracts', () => { providerId: 'acp', modelId: 'dimcode', updatedAt: Date.now(), + revision: 0, blocks: [ { type: 'content', diff --git a/test/renderer/stores/messageStore.test.ts b/test/renderer/stores/messageStore.test.ts index 38b6ea2820..72eb9a2186 100644 --- a/test/renderer/stores/messageStore.test.ts +++ b/test/renderer/stores/messageStore.test.ts @@ -1634,6 +1634,58 @@ describe('messageStore', () => { expect(afterAdvance.content).toContain('hello-again') }) + it('applies the first revision-zero snapshot to an existing pending record', async () => { + const { store, sessionClient, streamListeners } = await setupStore() + sessionClient.restore.mockResolvedValueOnce({ + session: { id: 's1' }, + nextCursor: null, + hasMore: false, + messages: [ + { + id: 'm1', + sessionId: 's1', + orderSeq: 1, + role: 'assistant' as const, + content: JSON.stringify([ + { type: 'content', content: 'stale-initial', status: 'pending', timestamp: 1 } + ]), + status: 'pending' as const, + isContextEdge: 0, + metadata: '{}', + traceCount: 0, + createdAt: 1, + updatedAt: 1 + } + ] + }) + await store.loadMessages('s1') + expect(store.messageCache.value.get('m1')?.content).toContain('stale-initial') + + const emit = (revision: number, text: string, updatedAt: number) => + streamListeners.updated[0]({ + sessionId: 's1', + requestId: 'm1', + messageId: 'm1', + providerId: 'acp', + modelId: 'dimcode', + updatedAt, + revision, + blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] + }) + + // No recorded revision yet: a first revision-0 snapshot must update the record + // instead of being treated as already applied. + emit(0, 'fresh-snapshot', 2) + const updated = store.messageCache.value.get('m1')! + expect(updated.content).toContain('fresh-snapshot') + + // Once revision 0 is recorded, duplicates are still deduped. + emit(0, 'duplicate-snapshot', 3) + const afterDuplicate = store.messageCache.value.get('m1')! + expect(afterDuplicate.content).toContain('fresh-snapshot') + expect(afterDuplicate.updatedAt).toBe(updated.updatedAt) + }) + it('skips JSON.stringify when the stream revision did not advance (quantified savings)', async () => { const { store, streamListeners } = await setupStore() await store.loadMessages('s1') From 39090580b6e8cd75b5aec639473e436b6f23a31a Mon Sep 17 00:00:00 2001 From: xiao-test Date: Mon, 14 Sep 2026 00:05:51 +0800 Subject: [PATCH 14/15] fix(stream): advance and scope stream revisions Two revision-contract breaks let the renderer dedupe real updates. Main: closing a previous provider-attempt narrative flipped a block from pending to success without marking the stream changed, so a deferred flush could publish the mutation under a stale revision. The close helper now reports whether it changed anything and the event loop marks only then; stripTrailingErrorBlock also marks after popping so every block mutation carries a revision bump. Renderer: appliedStreamRevision was keyed by messageId alone, so a new request or resume reusing the same message id had its first revision-0 snapshot deduped against the previous request's revision. Entries now record the owning requestId and revisions only compare within the same stream identity, with the IPC binding passing the payload's requestId through. Adds regression tests on both sides: revision advance on attempt close, and cross-request message-id reuse. --- src/main/agent/deepchat/runtime/process.ts | 14 +++++--- src/renderer/src/stores/ui/message.ts | 21 +++++++---- src/renderer/src/stores/ui/messageIpc.ts | 6 ++-- .../agent/deepchat/runtime/process.test.ts | 25 +++++++++++++ test/renderer/stores/messageStore.test.ts | 36 +++++++++++++++++++ 5 files changed, 89 insertions(+), 13 deletions(-) diff --git a/src/main/agent/deepchat/runtime/process.ts b/src/main/agent/deepchat/runtime/process.ts index 4b8c2c1356..4a3a7a2e43 100644 --- a/src/main/agent/deepchat/runtime/process.ts +++ b/src/main/agent/deepchat/runtime/process.ts @@ -142,6 +142,7 @@ function stripTrailingErrorBlock(state: StreamState, message: string): void { const lastBlock = state.blocks[state.blocks.length - 1] if (lastBlock?.type === 'error' && lastBlock.content === message) { state.blocks.pop() + markStreamChanged(state) } } @@ -184,15 +185,15 @@ function stampProviderAttemptIdentity( function closePreviousProviderAttemptNarrative( blocks: AssistantMessageBlock[], identity: DeepChatProviderAttemptIdentity | null -): void { - if (!identity) return +): boolean { + if (!identity) return false const last = blocks[blocks.length - 1] if ( !last || last.status !== 'pending' || (last.type !== 'content' && last.type !== 'reasoning_content') ) { - return + return false } const previousLogicalRound = last.extra?.providerLogicalRound const previousRequestSeq = last.extra?.providerRequestSeq @@ -205,9 +206,10 @@ function closePreviousProviderAttemptNarrative( previousRequestSeq === identity.requestSeq && previousPhysicalAttempt === identity.physicalAttempt) ) { - return + return false } last.status = 'success' + return true } function stampRunOutcome( @@ -1228,7 +1230,9 @@ export async function processStream(params: ProcessParams): Promise { // Stream message ids currently being hydrated into the cache as a placeholder // record (before the backend persists them). Prevents re-entrant duplicate inserts. const hydratingStreamMessageIds = new Set() - const appliedStreamRevision = new Map() + // Applied stream revisions are scoped to the stream identity (requestId): a new + // request or resume reusing the same message id restarts from revision 0 and must + // not be deduped against the previous request's revision. + const appliedStreamRevision = new Map() let latestLoadRequestId = 0 let latestHistoryRequestId = 0 let latestLoadSessionId: string | null = null @@ -592,7 +595,8 @@ export const useMessageStore = defineStore('message', () => { view.sessionId, streamingBlocks.value as AssistantMessageBlock[], currentStreamMetadata.value ?? undefined, - currentStreamBlocksRevision.value + currentStreamBlocksRevision.value, + currentStreamRequestId.value ?? undefined ) } } @@ -1002,14 +1006,19 @@ export const useMessageStore = defineStore('message', () => { conversationId: string, blocks: AssistantMessageBlock[], metadata?: { providerId?: string; modelId?: string }, - revision?: number + revision?: number, + requestId?: string ): void { if (committedSessionId.value !== conversationId) return const existing = messageCache.value.get(messageId) if (existing) { if (existing.sessionId !== conversationId) return - const lastRevision = appliedStreamRevision.get(messageId) + const lastApplied = appliedStreamRevision.get(messageId) + const lastRevision = + lastApplied && (requestId === undefined || lastApplied.requestId === requestId) + ? lastApplied.revision + : undefined if ( revision !== undefined && lastRevision !== undefined && @@ -1026,7 +1035,7 @@ export const useMessageStore = defineStore('message', () => { ...(metadata?.modelId ? { model: metadata.modelId } : {}) }) if (revision !== undefined) { - appliedStreamRevision.set(messageId, revision) + appliedStreamRevision.set(messageId, { requestId: requestId ?? null, revision }) } const nextMetadata = serializedMetadata === '{}' ? existing.metadata : serializedMetadata if ( @@ -1074,7 +1083,7 @@ export const useMessageStore = defineStore('message', () => { updatedAt: now } if (revision !== undefined) { - appliedStreamRevision.set(messageId, revision) + appliedStreamRevision.set(messageId, { requestId: requestId ?? null, revision }) } upsertMessageRecord(nextRecord) cacheStreamingAssistantBlocks(nextRecord, blocks) diff --git a/src/renderer/src/stores/ui/messageIpc.ts b/src/renderer/src/stores/ui/messageIpc.ts index 0d7953f838..84366e1b6d 100644 --- a/src/renderer/src/stores/ui/messageIpc.ts +++ b/src/renderer/src/stores/ui/messageIpc.ts @@ -26,7 +26,8 @@ interface BindMessageStoreIpcOptions { sessionId: string, blocks: AssistantMessageBlock[], metadata?: { providerId?: string; modelId?: string }, - revision?: number + revision?: number, + requestId?: string ) => void isEphemeralStreamMessageId: (messageId: string) => boolean } @@ -203,7 +204,8 @@ export function bindMessageStoreIpc(options: BindMessageStoreIpcOptions): Messag providerId: payload.providerId, modelId: payload.modelId }, - payload.revision + payload.revision, + payload.requestId ) } }), diff --git a/test/main/agent/deepchat/runtime/process.test.ts b/test/main/agent/deepchat/runtime/process.test.ts index fcc4884529..320dacc1de 100644 --- a/test/main/agent/deepchat/runtime/process.test.ts +++ b/test/main/agent/deepchat/runtime/process.test.ts @@ -644,6 +644,31 @@ describe('processStream', () => { ]) }) + it('advances the stream revision when an attempt change closes the previous narrative', async () => { + let identity = { logicalRound: 1, requestSeq: 2, physicalAttempt: 1 } + const coreStream = vi.fn(async function* () { + yield { type: 'text', content: 'Partial first attempt' } as LLMCoreStreamEvent + identity = { logicalRound: 1, requestSeq: 2, physicalAttempt: 2 } + // No further content: the narrative close is the only mutation on this event. + yield { type: 'stop', stop_reason: 'complete' } as LLMCoreStreamEvent + }) as unknown as ProcessParams['coreStream'] + const params = createParams({ + coreStream, + providerAttemptIdentity: () => identity + }) + + await expect(processStream(params)).resolves.toMatchObject({ status: 'completed' }) + + const state = params.run.streamState + expect(state.blocks).toEqual([ + expect.objectContaining({ type: 'content', status: 'success' }) + ]) + // text accumulate bumps once, the narrative close must bump again, and the + // terminal finalize adds its own mark — a close that does not advance the + // revision leaves a deferred flush deduped by the renderer. + expect(state.blocksRevision).toBe(3) + }) + it('persists normalized provider search results with the assistant message', async () => { const providerReplayJson = createDeepSeekReplayJson() const resultRow = { diff --git a/test/renderer/stores/messageStore.test.ts b/test/renderer/stores/messageStore.test.ts index 72eb9a2186..06640b8738 100644 --- a/test/renderer/stores/messageStore.test.ts +++ b/test/renderer/stores/messageStore.test.ts @@ -1686,6 +1686,42 @@ describe('messageStore', () => { expect(afterDuplicate.updatedAt).toBe(updated.updatedAt) }) + it('applies a revision-zero snapshot when a new request reuses the message id', async () => { + const { store, streamListeners } = await setupStore() + await store.loadMessages('s1') + + const emit = (requestId: string, revision: number, text: string, updatedAt: number) => + streamListeners.updated[0]({ + sessionId: 's1', + requestId, + messageId: 'm1', + providerId: 'acp', + modelId: 'dimcode', + updatedAt, + revision, + blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] + }) + + emit('req-a', 1, 'request-a', 1) + emit('req-a', 2, 'request-a-more', 2) + expect(store.messageCache.value.get('m1')?.content).toContain('request-a-more') + + // A new request (or a resume) reuses the same message id and restarts from + // revision 0: it must not be deduped against request A's revision. + emit('req-b', 0, 'request-b', 3) + const reused = store.messageCache.value.get('m1')! + expect(reused.content).toContain('request-b') + + // The new request's revisions are tracked independently afterwards. + emit('req-b', 0, 'request-b-duplicate', 4) + const afterDuplicate = store.messageCache.value.get('m1')! + expect(afterDuplicate.content).toContain('request-b') + expect(afterDuplicate.updatedAt).toBe(reused.updatedAt) + + emit('req-b', 1, 'request-b-more', 5) + expect(store.messageCache.value.get('m1')?.content).toContain('request-b-more') + }) + it('skips JSON.stringify when the stream revision did not advance (quantified savings)', async () => { const { store, streamListeners } = await setupStore() await store.loadMessages('s1') From 95aa76d494071652bf7d1c8d466e686bfad95403 Mon Sep 17 00:00:00 2001 From: xiao-test Date: Mon, 14 Sep 2026 09:45:18 +0800 Subject: [PATCH 15/15] test(stream): drop stringify call-count assertion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The quantified-savings test pinned JSON.stringify to exactly 6 calls, coupling the suite to the current serialization layout rather than the dedupe behavior. Any harmless refactor — moving serialization, adding metadata encoding, swapping implementations — would fail it while the user-visible behavior stayed correct. The behavioral contract (same revision keeps content, bumped revision updates it) is already covered by the duplicate-snapshot test, so the whole quantified case goes away instead of keeping a now-redundant shell. --- test/renderer/stores/messageStore.test.ts | 33 ----------------------- 1 file changed, 33 deletions(-) diff --git a/test/renderer/stores/messageStore.test.ts b/test/renderer/stores/messageStore.test.ts index 06640b8738..85626508a4 100644 --- a/test/renderer/stores/messageStore.test.ts +++ b/test/renderer/stores/messageStore.test.ts @@ -1722,39 +1722,6 @@ describe('messageStore', () => { expect(store.messageCache.value.get('m1')?.content).toContain('request-b-more') }) - it('skips JSON.stringify when the stream revision did not advance (quantified savings)', async () => { - const { store, streamListeners } = await setupStore() - await store.loadMessages('s1') - - const emit = (revision: number, text: string, updatedAt: number) => - streamListeners.updated[0]({ - sessionId: 's1', - requestId: 'm1', - messageId: 'm1', - providerId: 'acp', - modelId: 'dimcode', - updatedAt, - revision, - blocks: [{ type: 'content', content: text, status: 'pending', timestamp: updatedAt }] - }) - - const stringifySpy = vi.spyOn(JSON, 'stringify') - const baseline = stringifySpy.mock.calls.length - - emit(1, 'a', 1) - emit(1, 'a', 1) - emit(2, 'ab', 2) - emit(2, 'ab', 2) - emit(3, 'abc', 3) - - const stringifyCalls = stringifySpy.mock.calls.length - baseline - stringifySpy.mockRestore() - - expect(stringifyCalls).toBe(6) - expect(stringifyCalls).toBeLessThan(10) - expect(store.messageCache.value.get('m1')?.content).toContain('abc') - }) - it('drops the applied revision on persisted record arrival so a recycled stream re-folds', async () => { const { store, streamListeners, messageListeners } = await setupStore() await store.loadMessages('s1')