diff --git a/packages/cli/src/acp-integration/acpAgent.test.ts b/packages/cli/src/acp-integration/acpAgent.test.ts index cd82042534a..f1c859d5a1b 100644 --- a/packages/cli/src/acp-integration/acpAgent.test.ts +++ b/packages/cli/src/acp-integration/acpAgent.test.ts @@ -356,6 +356,25 @@ vi.mock('@qwen-code/qwen-code-core', () => ({ }, runManagedRememberByAgent: mockRunManagedRememberByAgent, runManagedAutoMemoryDream: mockRunManagedAutoMemoryDream, + refreshMemoryInstruction: vi.fn( + async (config: { + refreshHierarchicalMemory?: () => Promise; + getGeminiClient?: () => + | { refreshSystemInstruction?: () => Promise } + | undefined; + }) => { + try { + await config.refreshHierarchicalMemory?.(); + } catch { + // Best-effort, matching the real helper. + } + try { + await config.getGeminiClient?.()?.refreshSystemInstruction?.(); + } catch { + // Best-effort, matching the real helper. + } + }, + ), clearCachedCredentialFile: vi.fn(), getAllGeminiMdFilenames: vi.fn(() => ['QWEN.md', 'AGENTS.md']), getAutoMemoryRoot: vi.fn( @@ -3835,9 +3854,15 @@ describe('QwenAgent MCP SSE/HTTP support', () => { }); it('runs workspace memory remember without requiring a session', async () => { + const refreshHierarchicalMemory = vi.fn().mockResolvedValue(undefined); + const refreshSystemInstruction = vi.fn().mockResolvedValue(undefined); Object.assign(mockConfig, { isManagedMemoryAvailable: vi.fn().mockReturnValue(true), getProjectRoot: vi.fn().mockReturnValue('/workspace'), + refreshHierarchicalMemory, + getGeminiClient: vi.fn().mockReturnValue({ + refreshSystemInstruction, + }), }); mockRunManagedRememberByAgent.mockResolvedValue({ summary: 'saved', @@ -3876,6 +3901,242 @@ describe('QwenAgent MCP SSE/HTTP support', () => { abortSignal: expect.any(AbortSignal), }), ); + expect(refreshHierarchicalMemory).not.toHaveBeenCalled(); + expect(refreshSystemInstruction).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('refreshes live sessions after workspace memory remember', async () => { + const sessionRefreshHierarchicalMemory = vi + .fn() + .mockResolvedValue(undefined); + const sessionRefreshSystemInstruction = vi + .fn() + .mockResolvedValue(undefined); + const innerConfig = { + ...makeInnerConfig(), + getSessionId: vi.fn().mockReturnValue('remember-session'), + refreshHierarchicalMemory: sessionRefreshHierarchicalMemory, + getGeminiClient: vi.fn().mockReturnValue({ + isInitialized: vi.fn().mockReturnValue(true), + initialize: vi.fn().mockResolvedValue(undefined), + waitForMcpReady: vi.fn().mockResolvedValue(undefined), + refreshSystemInstruction: sessionRefreshSystemInstruction, + }), + }; + vi.mocked(loadSettings).mockReturnValue(makeSessionSettings()); + vi.mocked(loadCliConfig).mockResolvedValue( + innerConfig as unknown as Config, + ); + vi.mocked(Session).mockImplementation( + () => + ({ + getId: vi.fn().mockReturnValue('remember-session'), + getConfig: vi.fn().mockReturnValue(innerConfig), + sendAvailableCommandsUpdate: vi.fn().mockResolvedValue(undefined), + replayHistory: vi.fn().mockResolvedValue(undefined), + installRewriter: vi.fn(), + startCronScheduler: vi.fn(), + dispose: vi.fn(), + }) as unknown as InstanceType, + ); + vi.mocked(buildAvailableCommandsSnapshot).mockResolvedValue({ + availableCommands: [], + availableSkills: [], + }); + + Object.assign(mockConfig, { + isManagedMemoryAvailable: vi.fn().mockReturnValue(true), + getProjectRoot: vi.fn().mockReturnValue('/workspace'), + }); + mockRunManagedRememberByAgent.mockResolvedValue({ + summary: 'saved', + filesTouched: ['/mem/MEMORY.md'], + touchedScopes: ['project'], + }); + + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.newSession({ cwd: '/workspace', mcpServers: [] }); + + await expect( + agent.extMethod(SERVE_CONTROL_EXT_METHODS.workspaceMemoryRemember, { + content: 'Remember the workspace uses vitest.', + }), + ).resolves.toEqual({ + summary: 'saved', + filesTouched: ['/mem/MEMORY.md'], + touchedScopes: ['project'], + }); + expect(sessionRefreshHierarchicalMemory).toHaveBeenCalledTimes(1); + expect(sessionRefreshSystemInstruction).toHaveBeenCalledTimes(1); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('does not refresh live sessions when workspace memory remember writes nothing', async () => { + const sessionRefreshHierarchicalMemory = vi + .fn() + .mockResolvedValue(undefined); + const sessionRefreshSystemInstruction = vi + .fn() + .mockResolvedValue(undefined); + const innerConfig = { + ...makeInnerConfig(), + getSessionId: vi.fn().mockReturnValue('remember-noop-session'), + refreshHierarchicalMemory: sessionRefreshHierarchicalMemory, + getGeminiClient: vi.fn().mockReturnValue({ + isInitialized: vi.fn().mockReturnValue(true), + initialize: vi.fn().mockResolvedValue(undefined), + waitForMcpReady: vi.fn().mockResolvedValue(undefined), + refreshSystemInstruction: sessionRefreshSystemInstruction, + }), + }; + vi.mocked(loadSettings).mockReturnValue(makeSessionSettings()); + vi.mocked(loadCliConfig).mockResolvedValue( + innerConfig as unknown as Config, + ); + vi.mocked(Session).mockImplementation( + () => + ({ + getId: vi.fn().mockReturnValue('remember-noop-session'), + getConfig: vi.fn().mockReturnValue(innerConfig), + sendAvailableCommandsUpdate: vi.fn().mockResolvedValue(undefined), + replayHistory: vi.fn().mockResolvedValue(undefined), + installRewriter: vi.fn(), + startCronScheduler: vi.fn(), + dispose: vi.fn(), + }) as unknown as InstanceType, + ); + vi.mocked(buildAvailableCommandsSnapshot).mockResolvedValue({ + availableCommands: [], + availableSkills: [], + }); + + Object.assign(mockConfig, { + isManagedMemoryAvailable: vi.fn().mockReturnValue(true), + getProjectRoot: vi.fn().mockReturnValue('/workspace'), + }); + mockRunManagedRememberByAgent.mockResolvedValue({ + summary: 'No memory files updated.', + filesTouched: [], + touchedScopes: [], + }); + + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.newSession({ cwd: '/workspace', mcpServers: [] }); + + await expect( + agent.extMethod(SERVE_CONTROL_EXT_METHODS.workspaceMemoryRemember, { + content: 'Remember the workspace uses vitest.', + }), + ).resolves.toEqual({ + summary: 'No memory files updated.', + filesTouched: [], + touchedScopes: [], + }); + expect(sessionRefreshHierarchicalMemory).not.toHaveBeenCalled(); + expect(sessionRefreshSystemInstruction).not.toHaveBeenCalled(); + + mockConnectionState.resolve(); + await agentPromise; + }); + + it('keeps workspace memory remember successful when live session refresh fails', async () => { + const sessionRefreshHierarchicalMemory = vi + .fn() + .mockRejectedValue(new Error('memory refresh failed')); + const sessionRefreshSystemInstruction = vi + .fn() + .mockRejectedValue(new Error('system instruction refresh failed')); + const innerConfig = { + ...makeInnerConfig(), + getSessionId: vi.fn().mockReturnValue('remember-fail-session'), + refreshHierarchicalMemory: sessionRefreshHierarchicalMemory, + getGeminiClient: vi.fn().mockReturnValue({ + isInitialized: vi.fn().mockReturnValue(true), + initialize: vi.fn().mockResolvedValue(undefined), + waitForMcpReady: vi.fn().mockResolvedValue(undefined), + refreshSystemInstruction: sessionRefreshSystemInstruction, + }), + }; + vi.mocked(loadSettings).mockReturnValue(makeSessionSettings()); + vi.mocked(loadCliConfig).mockResolvedValue( + innerConfig as unknown as Config, + ); + vi.mocked(Session).mockImplementation( + () => + ({ + getId: vi.fn().mockReturnValue('remember-fail-session'), + getConfig: vi.fn().mockReturnValue(innerConfig), + sendAvailableCommandsUpdate: vi.fn().mockResolvedValue(undefined), + replayHistory: vi.fn().mockResolvedValue(undefined), + installRewriter: vi.fn(), + startCronScheduler: vi.fn(), + dispose: vi.fn(), + }) as unknown as InstanceType, + ); + vi.mocked(buildAvailableCommandsSnapshot).mockResolvedValue({ + availableCommands: [], + availableSkills: [], + }); + + Object.assign(mockConfig, { + isManagedMemoryAvailable: vi.fn().mockReturnValue(true), + getProjectRoot: vi.fn().mockReturnValue('/workspace'), + }); + mockRunManagedRememberByAgent.mockResolvedValue({ + summary: 'saved', + filesTouched: ['/mem/MEMORY.md'], + touchedScopes: ['project'], + }); + + const agentPromise = runAcpAgent( + mockConfig, + makeSessionSettings(), + mockArgv, + ); + await vi.waitFor(() => expect(capturedAgentFactory).toBeDefined()); + const agent = capturedAgentFactory!({ + get closed() { + return mockConnectionState.promise; + }, + }) as AgentLike; + await agent.newSession({ cwd: '/workspace', mcpServers: [] }); + + await expect( + agent.extMethod(SERVE_CONTROL_EXT_METHODS.workspaceMemoryRemember, { + content: 'Remember the workspace uses vitest.', + }), + ).resolves.toEqual({ + summary: 'saved', + filesTouched: ['/mem/MEMORY.md'], + touchedScopes: ['project'], + }); + expect(sessionRefreshHierarchicalMemory).toHaveBeenCalledTimes(1); + expect(sessionRefreshSystemInstruction).toHaveBeenCalledTimes(1); mockConnectionState.resolve(); await agentPromise; diff --git a/packages/cli/src/acp-integration/acpAgent.ts b/packages/cli/src/acp-integration/acpAgent.ts index 304cbc9b553..a4d9a1de458 100644 --- a/packages/cli/src/acp-integration/acpAgent.ts +++ b/packages/cli/src/acp-integration/acpAgent.ts @@ -77,6 +77,7 @@ import { normalizeEventPayload, normalizeSnapshotPayload, startEventLoopLagMonitor, + refreshMemoryInstruction, type AgentParams, type ApprovalMode, type ChatRecord, @@ -4972,6 +4973,22 @@ class QwenAgent implements Agent { return session; } + private async refreshLiveSessionMemoryInstructions( + logContext: string, + ): Promise { + const sessions = [...this.sessions.values()]; + if (sessions.length === 0) { + return; + } + await Promise.all( + sessions.map((session) => + refreshMemoryInstruction(session.getConfig(), { + logContext: `${logContext} session ${session.getId()}`, + }), + ), + ); + } + private buildSessionContextStatus( sessionId: string, ): ServeSessionContextStatus { @@ -6226,6 +6243,11 @@ class QwenAgent implements Agent { contextMode, abortSignal: childSignal, }); + if (result.filesTouched.length > 0) { + await this.refreshLiveSessionMemoryInstructions( + 'workspace memory remember', + ); + } return result as unknown as Record; } catch (err) { if (err instanceof RequestError) { diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 018b0ac0926..49ac44397be 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -43,6 +43,7 @@ import { MessageType } from '../../ui/types.js'; const debugLoggerWarnSpy = vi.hoisted(() => vi.fn()); const debugLoggerDebugSpy = vi.hoisted(() => vi.fn()); const runVisionBridgeSpy = vi.hoisted(() => vi.fn()); +const refreshMemoryAfterManagedWriteSpy = vi.hoisted(() => vi.fn()); const transcribeVoiceAudioSpy = vi.hoisted(() => vi.fn()); // Records every LoopTickResolver construction's deps so a test can assert what // Session computed (e.g. the home confinement root) without a private-field peek. @@ -62,6 +63,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { generatePromptSuggestion: vi.fn(), logPromptSuggestion: vi.fn(), runVisionBridge: runVisionBridgeSpy, + refreshMemoryAfterManagedWrite: refreshMemoryAfterManagedWriteSpy, // Transparent recording wrapper: records the constructor deps, then behaves // exactly like the real resolver (subclass → instanceof + methods preserved). LoopTickResolver: class extends actual.LoopTickResolver { @@ -378,6 +380,8 @@ describe('Session', () => { beforeEach(() => { runVisionBridgeSpy.mockReset(); + refreshMemoryAfterManagedWriteSpy.mockReset(); + refreshMemoryAfterManagedWriteSpy.mockResolvedValue(false); transcribeVoiceAudioSpy.mockReset(); currentModel = 'qwen3-code-plus'; currentAuthType = AuthType.USE_OPENAI; @@ -10306,6 +10310,43 @@ describe('Session', () => { }; } + it('refreshes managed memory instructions after successful ACP tool writes', async () => { + const execute = vi.fn().mockResolvedValue({ + llmContent: 'wrote memory', + returnDisplay: 'wrote memory', + }); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + mockToolRegistry.getTool.mockReturnValue( + mockAllowedTool(core.ToolNames.WRITE_FILE, execute), + ); + + const result = await ( + session as unknown as ToolCallInternals + ).runToolCalls(new AbortController().signal, 'prompt-memory-write', [ + { + id: 'write_memory', + name: core.ToolNames.WRITE_FILE, + args: { file_path: '/workspace/.qwen/memory/project.md' }, + }, + ]); + + expect(result.stopAfterPermissionCancel).toBe(false); + expect(refreshMemoryAfterManagedWriteSpy).toHaveBeenCalledTimes(1); + expect(refreshMemoryAfterManagedWriteSpy).toHaveBeenCalledWith( + mockConfig, + [ + { + toolName: core.ToolNames.WRITE_FILE, + args: { file_path: '/workspace/.qwen/memory/project.md' }, + status: 'success', + }, + ], + { + logContext: 'ACP session test-session-id memory tool batch', + }, + ); + }); + it('does not fire PostToolBatch hooks from the ACP session path', async () => { const messageBus = { request: vi.fn().mockImplementation(async (request) => ({ diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index c06c622fed1..9b90f45350d 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -35,6 +35,7 @@ import type { LoopTickResult, ToolArtifact, VisionBridgeResult, + MemoryWriteCandidate, } from '@qwen-code/qwen-code-core'; import { AuthType, @@ -115,6 +116,7 @@ import { LoopDetectedEvent, LoopType, acquireSleepInhibitor, + refreshMemoryAfterManagedWrite, clearGoalTerminalObserver, setGoalTerminalObserver, sessionIdContext, @@ -242,6 +244,7 @@ type RunToolResult = { stopAfterPermissionCancel: boolean; repeatedDuplicateProviderToolCall?: boolean; loopDetected?: boolean; + memoryWriteCandidates?: MemoryWriteCandidate[]; }; type DaemonToolLoopState = { @@ -4148,6 +4151,17 @@ export class Session implements SessionContext { parts.push(await recordSkippedToolCall(remainingCall, message)); } }; + const memoryWriteCandidates: MemoryWriteCandidate[] = []; + const collectMemoryWriteCandidates = (result: RunToolResult): void => { + if (result.memoryWriteCandidates) { + memoryWriteCandidates.push(...result.memoryWriteCandidates); + } + }; + const refreshMemoryIfNeeded = async (): Promise => { + await refreshMemoryAfterManagedWrite(this.config, memoryWriteCandidates, { + logContext: `ACP session ${this.sessionId} memory tool batch`, + }); + }; // Bounded-concurrency runner: matches core's `runConcurrently` // behaviour (`coreToolScheduler.ts:1506`), capped by // `QWEN_CODE_MAX_TOOL_CONCURRENCY` (default 10). Results are returned @@ -4281,103 +4295,117 @@ export class Session implements SessionContext { }; const parts: Part[] = []; - for (const batch of batches) { - if (batch.kind === 'duplicate') { - await emitDuplicateBatch(batch); - parts.push(...batch.response.responseParts); - continue; - } - if (batch.concurrent && batch.calls.length > 1) { - const batchAbortController = new AbortController(); - let batchStopAfterPermissionCancel = false; - const propagateAbort = () => { - batchAbortController.abort(abortSignal.reason); - }; - if (abortSignal.aborted) { - propagateAbort(); - } else { - abortSignal.addEventListener('abort', propagateAbort, { - once: true, - }); - } - const stopBatchAfterPermissionCancel = () => { - batchStopAfterPermissionCancel = true; - batchAbortController.abort(USER_CANCEL_ABORT_REASON); - }; - let results: RunToolResult[]; - try { - results = await runBounded( - batch.calls, - batchAbortController.signal, - stopBatchAfterPermissionCancel, - () => batchAbortController.abort('loop_detected'), - () => batchStopAfterPermissionCancel, - ); - } finally { - abortSignal.removeEventListener('abort', propagateAbort); - } - let shouldStop = false; - let shouldStopForLoop = false; - for (const r of results) { - parts.push(...r.parts); - shouldStop ||= r.stopAfterPermissionCancel; - shouldStopForLoop ||= r.loopDetected === true; + try { + for (const batch of batches) { + if (batch.kind === 'duplicate') { + await emitDuplicateBatch(batch); + parts.push(...batch.response.responseParts); + continue; } - if (shouldStopForLoop) { - await appendSkippedAfter( - parts, - batch.calls[batch.calls.length - 1], - LOOP_DETECTED_SKIP_MESSAGE, - ); - return { - parts, - stopAfterPermissionCancel: false, - loopDetected: true, + if (batch.concurrent && batch.calls.length > 1) { + const batchAbortController = new AbortController(); + let batchStopAfterPermissionCancel = false; + const propagateAbort = () => { + batchAbortController.abort(abortSignal.reason); }; - } - if (shouldStop) { - await appendSkippedAfter(parts, batch.calls[batch.calls.length - 1]); - return { - parts, - stopAfterPermissionCancel: true, - repeatedDuplicateProviderToolCall: false, + if (abortSignal.aborted) { + propagateAbort(); + } else { + abortSignal.addEventListener('abort', propagateAbort, { + once: true, + }); + } + const stopBatchAfterPermissionCancel = () => { + batchStopAfterPermissionCancel = true; + batchAbortController.abort(USER_CANCEL_ABORT_REASON); }; - } - } else { - for (const fc of batch.calls) { - const r = await this.runTool( - abortSignal, - promptId, - fc, - undefined, - toolLoopState, - recordSkippedToolCall, - ); - parts.push(...r.parts); - if (r.loopDetected) { - await appendSkippedAfter(parts, fc, LOOP_DETECTED_SKIP_MESSAGE); + let results: RunToolResult[]; + try { + results = await runBounded( + batch.calls, + batchAbortController.signal, + stopBatchAfterPermissionCancel, + () => batchAbortController.abort('loop_detected'), + () => batchStopAfterPermissionCancel, + ); + } finally { + abortSignal.removeEventListener('abort', propagateAbort); + } + let shouldStop = false; + let shouldStopForLoop = false; + for (const r of results) { + parts.push(...r.parts); + collectMemoryWriteCandidates(r); + shouldStop ||= r.stopAfterPermissionCancel; + shouldStopForLoop ||= r.loopDetected === true; + } + if (shouldStopForLoop) { + await appendSkippedAfter( + parts, + batch.calls[batch.calls.length - 1], + LOOP_DETECTED_SKIP_MESSAGE, + ); return { parts, stopAfterPermissionCancel: false, loopDetected: true, + memoryWriteCandidates, }; } - if (r.stopAfterPermissionCancel) { - await appendSkippedAfter(parts, fc); + if (shouldStop) { + await appendSkippedAfter( + parts, + batch.calls[batch.calls.length - 1], + ); return { parts, stopAfterPermissionCancel: true, repeatedDuplicateProviderToolCall: false, + memoryWriteCandidates, }; } + } else { + for (const fc of batch.calls) { + const r = await this.runTool( + abortSignal, + promptId, + fc, + undefined, + toolLoopState, + recordSkippedToolCall, + ); + parts.push(...r.parts); + collectMemoryWriteCandidates(r); + if (r.loopDetected) { + await appendSkippedAfter(parts, fc, LOOP_DETECTED_SKIP_MESSAGE); + return { + parts, + stopAfterPermissionCancel: false, + loopDetected: true, + memoryWriteCandidates, + }; + } + if (r.stopAfterPermissionCancel) { + await appendSkippedAfter(parts, fc); + return { + parts, + stopAfterPermissionCancel: true, + repeatedDuplicateProviderToolCall: false, + memoryWriteCandidates, + }; + } + } } } + return { + parts, + stopAfterPermissionCancel: false, + repeatedDuplicateProviderToolCall: false, + memoryWriteCandidates, + }; + } finally { + await refreshMemoryIfNeeded(); } - return { - parts, - stopAfterPermissionCancel: false, - repeatedDuplicateProviderToolCall: false, - }; } /** @@ -5337,6 +5365,16 @@ export class Session implements SessionContext { return { parts: responseParts, stopAfterPermissionCancel: nestedPermissionCancelled, + memoryWriteCandidates: + status === 'success' + ? [ + { + toolName, + args, + status, + }, + ] + : undefined, }; } catch (e) { // Ensure cleanup on error diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 74576c3e0bf..22f8086c461 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -94,6 +94,7 @@ const mockGetActiveGoal = vi.hoisted(() => vi.fn()); const mockActiveGoalEquals = vi.hoisted(() => vi.fn()); const mockSetActiveGoal = vi.hoisted(() => vi.fn()); const mockClearActiveGoal = vi.hoisted(() => vi.fn()); +const mockRefreshMemoryAfterManagedWrite = vi.hoisted(() => vi.fn()); vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { const actualCoreModule = (await importOriginal()) as any; @@ -109,6 +110,7 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => { setActiveGoal: mockSetActiveGoal, clearActiveGoal: mockClearActiveGoal, runVisionBridge: mockRunVisionBridge, + refreshMemoryAfterManagedWrite: mockRefreshMemoryAfterManagedWrite, }; }); @@ -172,6 +174,7 @@ describe('useGeminiStream', () => { beforeEach(() => { vi.clearAllMocks(); // Clear mocks before each test + mockRefreshMemoryAfterManagedWrite.mockResolvedValue(false); mockGetActiveGoal.mockReturnValue(undefined); mockActiveGoalEquals.mockReturnValue(false); vi.mocked(findLastSafeSplitPoint).mockImplementation( @@ -5854,7 +5857,7 @@ describe('useGeminiStream', () => { describe('Memory Refresh on save_memory', () => { it('should call performMemoryRefresh when a save_memory tool call completes successfully', async () => { - const mockPerformMemoryRefresh = vi.fn(); + const mockPerformMemoryRefresh = vi.fn().mockResolvedValue(undefined); const completedToolCall: TrackedCompletedToolCall = { request: { callId: 'save-mem-call-1', @@ -5928,6 +5931,203 @@ describe('useGeminiStream', () => { expect(mockPerformMemoryRefresh).toHaveBeenCalledTimes(1); }); }); + + it('refreshes managed-memory instructions after interactive memory file writes', async () => { + const completedToolCall: TrackedCompletedToolCall = { + request: { + callId: 'write-memory-call-1', + name: 'write_file', + args: { file_path: '/workspace/.qwen/memory/project.md' }, + isClientInitiated: false, + prompt_id: 'prompt-id-memory-write', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'write-memory-call-1', + responseParts: [{ text: 'Wrote memory' }], + resultDisplay: 'Wrote memory', + error: undefined, + errorType: undefined, + }, + tool: { + name: 'write_file', + displayName: 'write_file', + description: 'Writes files', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + true, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([completedToolCall]); + } + }); + + expect(mockRefreshMemoryAfterManagedWrite).toHaveBeenCalledWith( + mockConfig, + [ + { + toolName: 'write_file', + args: { file_path: '/workspace/.qwen/memory/project.md' }, + status: 'success', + }, + ], + { logContext: 'interactive memory tool batch' }, + ); + }); + + it('does not run the legacy save_memory refresh when managed-memory writes refresh the batch', async () => { + mockRefreshMemoryAfterManagedWrite.mockResolvedValueOnce(true); + const mockPerformMemoryRefresh = vi.fn(); + const saveMemoryToolCall: TrackedCompletedToolCall = { + request: { + callId: 'save-mem-call-1', + name: 'save_memory', + args: { fact: 'test' }, + isClientInitiated: true, + prompt_id: 'prompt-id-save-memory', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'save-mem-call-1', + responseParts: [{ text: 'Memory saved' }], + resultDisplay: 'Success: Memory saved', + error: undefined, + errorType: undefined, + }, + tool: { + name: 'save_memory', + displayName: 'save_memory', + description: 'Saves memory', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + }; + const writeMemoryToolCall: TrackedCompletedToolCall = { + request: { + callId: 'write-memory-call-1', + name: 'write_file', + args: { file_path: '/workspace/.qwen/memory/project.md' }, + isClientInitiated: false, + prompt_id: 'prompt-id-memory-write', + }, + status: 'success', + responseSubmittedToGemini: false, + response: { + callId: 'write-memory-call-1', + responseParts: [{ text: 'Wrote memory' }], + resultDisplay: 'Wrote memory', + error: undefined, + errorType: undefined, + }, + tool: { + name: 'write_file', + displayName: 'write_file', + description: 'Writes files', + build: vi.fn(), + } as any, + invocation: { + getDescription: () => `Mock description`, + } as unknown as AnyToolInvocation, + }; + + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + true, + mockLoadedSettings, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + mockPerformMemoryRefresh, + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([saveMemoryToolCall, writeMemoryToolCall]); + } + }); + + expect(mockRefreshMemoryAfterManagedWrite).toHaveBeenCalledWith( + mockConfig, + [ + { + toolName: 'save_memory', + args: { fact: 'test' }, + status: 'success', + }, + { + toolName: 'write_file', + args: { file_path: '/workspace/.qwen/memory/project.md' }, + status: 'success', + }, + ], + { logContext: 'interactive memory tool batch' }, + ); + expect(mockPerformMemoryRefresh).not.toHaveBeenCalled(); + }); }); describe('Error Handling', () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index fdc158fe0e2..6b22edd681b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -64,6 +64,7 @@ import { markDuplicateProviderToolCallResponseSent, findRepeatedDuplicateProviderToolCall, AutonomousLoopTickResolver, + refreshMemoryAfterManagedWrite, } from '@qwen-code/qwen-code-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -2925,20 +2926,33 @@ export const useGeminiStream = ( !processedMemoryToolsRef.current.has(t.request.callId), ); + const geminiTools = completedAndReadyToSubmitTools.filter( + (t) => + !t.request.isClientInitiated && + !historyCallIdsWithResponse.has(t.request.callId), + ); + const didRefreshManagedMemory = await refreshMemoryAfterManagedWrite( + config, + completedAndReadyToSubmitTools.map((toolCall) => ({ + toolName: toolCall.request.name, + args: toolCall.request.args as Record, + status: toolCall.status, + })), + { logContext: 'interactive memory tool batch' }, + ); if (newSuccessfulMemorySaves.length > 0) { - // Perform the refresh only if there are new ones. - void performMemoryRefresh(); + if (!didRefreshManagedMemory) { + // Perform the legacy save_memory refresh only when the managed-memory + // write refresh did not already rebuild and publish a fresher state. + void performMemoryRefresh().catch((err) => { + debugLogger.warn(`save_memory refresh failed: ${err}`); + }); + } // Mark them as processed so we don't do this again on the next render. newSuccessfulMemorySaves.forEach((t) => processedMemoryToolsRef.current.add(t.request.callId), ); } - - const geminiTools = completedAndReadyToSubmitTools.filter( - (t) => - !t.request.isClientInitiated && - !historyCallIdsWithResponse.has(t.request.callId), - ); const completedCallIds = new Set( completedAndReadyToSubmitTools.map( (toolCall) => toolCall.request.callId, diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 3c6e776ec43..50f1511bdc3 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -320,6 +320,7 @@ export * from './memory/store.js'; export * from './memory/const.js'; export * from './memory/channel-memory.js'; export * from './memory/remember.js'; +export * from './memory/refresh.js'; export * from './memory/dream.js'; // Issue : write helper for hierarchical context files, // re-exported so the `qwen serve` daemon can mutate workspace memory diff --git a/packages/core/src/memory/refresh.test.ts b/packages/core/src/memory/refresh.test.ts new file mode 100644 index 00000000000..bcea11bfc2f --- /dev/null +++ b/packages/core/src/memory/refresh.test.ts @@ -0,0 +1,254 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs/promises'; +import * as os from 'node:os'; +import * as path from 'node:path'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { Config } from '../config/config.js'; +import { + clearAutoMemoryRootCache, + getAutoMemoryRoot, + getUserAutoMemoryRoot, +} from './paths.js'; +import { + rebuildManagedAutoMemoryIndex, + rebuildUserAutoMemoryIndex, +} from './indexer.js'; +import { + didWriteManagedMemory, + refreshMemoryAfterManagedWrite, + refreshMemoryInstruction, +} from './refresh.js'; + +vi.mock('./indexer.js', () => ({ + rebuildManagedAutoMemoryIndex: vi.fn(), + rebuildUserAutoMemoryIndex: vi.fn(), +})); + +function createConfig(projectRoot: string, managed = true): Config { + return { + isManagedMemoryAvailable: vi.fn().mockReturnValue(managed), + getProjectRoot: vi.fn().mockReturnValue(projectRoot), + refreshHierarchicalMemory: vi.fn().mockResolvedValue(undefined), + getGeminiClient: vi.fn().mockReturnValue({ + refreshSystemInstruction: vi.fn().mockResolvedValue(undefined), + }), + } as unknown as Config; +} + +describe('managed memory refresh helper', () => { + const originalMemoryBase = process.env['QWEN_CODE_MEMORY_BASE_DIR']; + let tempDir: string; + let projectRoot: string; + + beforeEach(async () => { + tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'memory-refresh-')); + projectRoot = path.join(tempDir, 'project'); + await fs.mkdir(projectRoot, { recursive: true }); + process.env['QWEN_CODE_MEMORY_BASE_DIR'] = path.join(tempDir, 'memory'); + clearAutoMemoryRootCache(); + vi.mocked(rebuildManagedAutoMemoryIndex).mockReset(); + vi.mocked(rebuildUserAutoMemoryIndex).mockReset(); + vi.mocked(rebuildManagedAutoMemoryIndex).mockResolvedValue(''); + vi.mocked(rebuildUserAutoMemoryIndex).mockResolvedValue(''); + }); + + afterEach(async () => { + if (originalMemoryBase === undefined) { + delete process.env['QWEN_CODE_MEMORY_BASE_DIR']; + } else { + process.env['QWEN_CODE_MEMORY_BASE_DIR'] = originalMemoryBase; + } + clearAutoMemoryRootCache(); + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('detects successful private managed-memory writes only', () => { + const memoryFile = path.join(getAutoMemoryRoot(projectRoot), 'project.md'); + + expect( + didWriteManagedMemory( + [ + { + toolName: 'write_file', + args: { file_path: memoryFile }, + status: 'success', + }, + ], + projectRoot, + ), + ).toBe(true); + expect( + didWriteManagedMemory( + [ + { + toolName: 'edit', + args: { file_path: memoryFile }, + status: 'error', + }, + ], + projectRoot, + ), + ).toBe(false); + expect( + didWriteManagedMemory( + [ + { + toolName: 'write_file', + args: { file_path: path.join(projectRoot, 'src/file.ts') }, + status: 'success', + }, + ], + projectRoot, + ), + ).toBe(false); + expect( + didWriteManagedMemory( + [ + { + toolName: 'write_file', + args: { + file_path: path.join( + projectRoot, + '.qwen', + 'team-memory', + 'shared.md', + ), + }, + status: 'success', + }, + ], + projectRoot, + ), + ).toBe(false); + }); + + it('supports legacy edit names and alternate file path arguments', () => { + const memoryFile = path.join(getAutoMemoryRoot(projectRoot), 'project.md'); + + expect( + didWriteManagedMemory( + [ + { + toolName: 'replace', + args: { target_file: memoryFile }, + status: 'success', + }, + ], + projectRoot, + ), + ).toBe(true); + }); + + it('rebuilds touched indexes before refreshing the live instruction', async () => { + const config = createConfig(projectRoot); + const projectFile = path.join(getAutoMemoryRoot(projectRoot), 'project.md'); + const userFile = path.join(getUserAutoMemoryRoot(), 'user.md'); + + await expect( + refreshMemoryAfterManagedWrite(config, [ + { + toolName: 'write_file', + args: { file_path: projectFile }, + status: 'success', + }, + { toolName: 'edit', args: { file_path: userFile }, status: 'success' }, + ]), + ).resolves.toBe(true); + + expect(rebuildManagedAutoMemoryIndex).toHaveBeenCalledWith(projectRoot); + expect(rebuildUserAutoMemoryIndex).toHaveBeenCalledTimes(1); + expect(config.refreshHierarchicalMemory).toHaveBeenCalledTimes(1); + expect( + config.getGeminiClient().refreshSystemInstruction, + ).toHaveBeenCalledTimes(1); + expect( + vi.mocked(rebuildManagedAutoMemoryIndex).mock.invocationCallOrder[0], + ).toBeLessThan( + vi.mocked(config.refreshHierarchicalMemory).mock.invocationCallOrder[0], + ); + }); + + it('keeps refreshing when index rebuild fails', async () => { + vi.mocked(rebuildManagedAutoMemoryIndex).mockRejectedValueOnce( + new Error('index failed'), + ); + const config = createConfig(projectRoot); + + await expect( + refreshMemoryAfterManagedWrite(config, [ + { + toolName: 'write_file', + args: { + file_path: path.join(getAutoMemoryRoot(projectRoot), 'x.md'), + }, + status: 'success', + }, + ]), + ).resolves.toBe(true); + + expect(config.refreshHierarchicalMemory).toHaveBeenCalledTimes(1); + expect( + config.getGeminiClient().refreshSystemInstruction, + ).toHaveBeenCalledTimes(1); + }); + + it('keeps refreshing the system instruction when hierarchical refresh fails', async () => { + const config = createConfig(projectRoot); + vi.mocked(config.refreshHierarchicalMemory).mockRejectedValueOnce( + new Error('hierarchical refresh failed'), + ); + + await expect(refreshMemoryInstruction(config)).resolves.toBeUndefined(); + + expect(config.refreshHierarchicalMemory).toHaveBeenCalledTimes(1); + expect( + config.getGeminiClient().refreshSystemInstruction, + ).toHaveBeenCalledTimes(1); + }); + + it('returns false without refreshing when managed memory is unavailable', async () => { + const config = createConfig(projectRoot, false); + + await expect( + refreshMemoryAfterManagedWrite(config, [ + { + toolName: 'write_file', + args: { + file_path: path.join(getAutoMemoryRoot(projectRoot), 'x.md'), + }, + status: 'success', + }, + ]), + ).resolves.toBe(false); + + expect(config.refreshHierarchicalMemory).not.toHaveBeenCalled(); + expect(rebuildManagedAutoMemoryIndex).not.toHaveBeenCalled(); + }); + + it('returns false when refresh guard evaluation throws', async () => { + const config = createConfig(projectRoot); + vi.mocked(config.getProjectRoot).mockImplementationOnce(() => { + throw new Error('project root unavailable'); + }); + + await expect( + refreshMemoryAfterManagedWrite(config, [ + { + toolName: 'write_file', + args: { + file_path: path.join(getAutoMemoryRoot(projectRoot), 'x.md'), + }, + status: 'success', + }, + ]), + ).resolves.toBe(false); + + expect(config.refreshHierarchicalMemory).not.toHaveBeenCalled(); + expect(rebuildManagedAutoMemoryIndex).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/memory/refresh.ts b/packages/core/src/memory/refresh.ts new file mode 100644 index 00000000000..4de4ae733e3 --- /dev/null +++ b/packages/core/src/memory/refresh.ts @@ -0,0 +1,175 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as path from 'node:path'; +import type { Config } from '../config/config.js'; +import { ToolNames, ToolNamesMigration } from '../tools/tool-names.js'; +import { createDebugLogger } from '../utils/debugLogger.js'; +import { isAllowedMemoryPath } from './memory-scoped-agent-config.js'; +import { + rebuildManagedAutoMemoryIndex, + rebuildUserAutoMemoryIndex, +} from './indexer.js'; + +const debugLogger = createDebugLogger('AUTO_MEMORY_REFRESH'); + +export interface MemoryWriteCandidate { + toolName: string; + args?: Record; + status?: string; +} + +export interface RefreshMemoryAfterWriteOptions { + logContext?: string; +} + +const WRITE_TOOL_NAMES = new Set([ + ToolNames.WRITE_FILE, + ToolNames.EDIT, +]); + +function canonicalToolName(toolName: string): string { + return (ToolNamesMigration as Record)[toolName] ?? toolName; +} + +function candidateFilePath( + args: Record | undefined, +): string | undefined { + const value = args?.['file_path'] ?? args?.['path'] ?? args?.['target_file']; + return typeof value === 'string' && value.length > 0 ? value : undefined; +} + +function resolveCandidatePath(filePath: string, projectRoot: string): string { + return path.resolve(projectRoot, filePath); +} + +function isSuccessfulWrite(candidate: MemoryWriteCandidate): boolean { + if (candidate.status !== undefined && candidate.status !== 'success') { + return false; + } + return WRITE_TOOL_NAMES.has(canonicalToolName(candidate.toolName)); +} + +function logPrefix( + options?: Pick, +) { + return options?.logContext ? `${options.logContext}: ` : ''; +} + +function classifyWrittenMemoryScope( + candidate: MemoryWriteCandidate, + projectRoot: string, +): 'project' | 'user' | undefined { + if (!isSuccessfulWrite(candidate)) { + return undefined; + } + const filePath = candidateFilePath(candidate.args); + if (!filePath) { + return undefined; + } + const resolved = resolveCandidatePath(filePath, projectRoot); + if (!isAllowedMemoryPath(resolved, projectRoot)) { + return undefined; + } + return isAllowedMemoryPath(resolved, projectRoot, { + includeUserMemory: false, + }) + ? 'project' + : 'user'; +} + +export function didWriteManagedMemory( + candidates: readonly MemoryWriteCandidate[], + projectRoot: string, +): boolean { + return candidates.some( + (candidate) => + classifyWrittenMemoryScope(candidate, projectRoot) !== undefined, + ); +} + +async function rebuildWrittenMemoryIndexes( + candidates: readonly MemoryWriteCandidate[], + projectRoot: string, + options?: Pick, +): Promise { + let wroteProjectMemory = false; + let wroteUserMemory = false; + + for (const candidate of candidates) { + const scope = classifyWrittenMemoryScope(candidate, projectRoot); + wroteProjectMemory ||= scope === 'project'; + wroteUserMemory ||= scope === 'user'; + } + + await Promise.all([ + wroteProjectMemory + ? rebuildManagedAutoMemoryIndex(projectRoot).catch((err: unknown) => { + debugLogger.warn( + `${logPrefix(options)}rebuildManagedAutoMemoryIndex failed: ${err}`, + ); + }) + : Promise.resolve(), + wroteUserMemory + ? rebuildUserAutoMemoryIndex().catch((err: unknown) => { + debugLogger.warn( + `${logPrefix(options)}rebuildUserAutoMemoryIndex failed: ${err}`, + ); + }) + : Promise.resolve(), + ]); +} + +export async function refreshMemoryInstruction( + config: Config, + options?: Pick, +): Promise { + try { + await config.refreshHierarchicalMemory(); + } catch (err) { + debugLogger.warn( + `${logPrefix(options)}refreshHierarchicalMemory failed: ${err}`, + ); + } + + try { + await config.getGeminiClient()?.refreshSystemInstruction(); + } catch (err) { + debugLogger.warn( + `${logPrefix(options)}refreshSystemInstruction failed: ${err}`, + ); + } +} + +export async function refreshMemoryAfterManagedWrite( + config: Config, + candidates: readonly MemoryWriteCandidate[], + options: RefreshMemoryAfterWriteOptions = {}, +): Promise { + try { + if (typeof config.isManagedMemoryAvailable !== 'function') { + return false; + } + if (!config.isManagedMemoryAvailable()) { + return false; + } + + const projectRoot = config.getProjectRoot(); + if (!didWriteManagedMemory(candidates, projectRoot)) { + return false; + } + + await rebuildWrittenMemoryIndexes(candidates, projectRoot, options); + + await refreshMemoryInstruction(config, options); + return true; + } catch (err) { + debugLogger.warn( + `${logPrefix(options)}refreshMemoryAfterManagedWrite failed: ${err}`, + ); + return false; + } +}