diff --git a/src/node/services/agentSession.postCompactionRefresh.test.ts b/src/node/services/agentSession.postCompactionRefresh.test.ts new file mode 100644 index 0000000000..d9e2fcf875 --- /dev/null +++ b/src/node/services/agentSession.postCompactionRefresh.test.ts @@ -0,0 +1,168 @@ +import { describe, expect, test, mock } from "bun:test"; +import { AgentSession } from "./agentSession"; +import type { Config } from "@/node/config"; +import type { HistoryService } from "./historyService"; +import type { PartialService } from "./partialService"; +import type { AIService } from "./aiService"; +import type { InitStateManager } from "./initStateManager"; +import type { BackgroundProcessManager } from "./backgroundProcessManager"; + +// NOTE: These tests focus on the event wiring (tool-call-end -> callback). +// The actual post-compaction state computation is covered elsewhere. + +describe("AgentSession post-compaction refresh trigger", () => { + test("triggers callback on file_edit_* tool-call-end when experiment enabled", () => { + const handlers = new Map void>(); + + const aiService: AIService = { + on(eventName: string | symbol, listener: (...args: unknown[]) => void) { + handlers.set(String(eventName), listener); + return this; + }, + off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + + const historyService: HistoryService = { + getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })), + } as unknown as HistoryService; + + const initStateManager: InitStateManager = { + on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + } as unknown as InitStateManager; + + const backgroundProcessManager: BackgroundProcessManager = { + setMessageQueued: mock(() => undefined), + cleanup: mock(() => Promise.resolve()), + } as unknown as BackgroundProcessManager; + + const config: Config = { srcDir: "/tmp" } as unknown as Config; + const partialService: PartialService = {} as unknown as PartialService; + + const onPostCompactionStateChange = mock(() => undefined); + + const session = new AgentSession({ + workspaceId: "ws", + config, + historyService, + partialService, + aiService, + initStateManager, + backgroundProcessManager, + onPostCompactionStateChange, + }); + + // Enable the experiment gate (normally set during sendMessage()). + (session as unknown as { postCompactionContextEnabled: boolean }).postCompactionContextEnabled = + true; + + const toolEnd = handlers.get("tool-call-end"); + expect(toolEnd).toBeDefined(); + + toolEnd!({ + type: "tool-call-end", + workspaceId: "ws", + messageId: "m1", + toolCallId: "t1b", + toolName: "file_edit_replace_lines", + result: {}, + timestamp: Date.now(), + }); + + toolEnd!({ + type: "tool-call-end", + workspaceId: "ws", + messageId: "m1", + toolCallId: "t1", + toolName: "file_edit_insert", + result: {}, + timestamp: Date.now(), + }); + + toolEnd!({ + type: "tool-call-end", + workspaceId: "ws", + messageId: "m1", + toolCallId: "t2", + toolName: "bash", + result: {}, + timestamp: Date.now(), + }); + + expect(onPostCompactionStateChange).toHaveBeenCalledTimes(2); + + session.dispose(); + }); + + test("does not trigger callback when experiment disabled", () => { + const handlers = new Map void>(); + + const aiService: AIService = { + on(eventName: string | symbol, listener: (...args: unknown[]) => void) { + handlers.set(String(eventName), listener); + return this; + }, + off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + stopStream: mock(() => Promise.resolve({ success: true as const, data: undefined })), + } as unknown as AIService; + + const initStateManager: InitStateManager = { + on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + } as unknown as InitStateManager; + + const backgroundProcessManager: BackgroundProcessManager = { + setMessageQueued: mock(() => undefined), + cleanup: mock(() => Promise.resolve()), + } as unknown as BackgroundProcessManager; + + const config: Config = { srcDir: "/tmp" } as unknown as Config; + const historyService: HistoryService = { + getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })), + } as unknown as HistoryService; + const partialService: PartialService = {} as unknown as PartialService; + + const onPostCompactionStateChange = mock(() => undefined); + + const session = new AgentSession({ + workspaceId: "ws", + config, + historyService, + partialService, + aiService, + initStateManager, + backgroundProcessManager, + onPostCompactionStateChange, + }); + + const toolEnd = handlers.get("tool-call-end"); + expect(toolEnd).toBeDefined(); + + toolEnd!({ + type: "tool-call-end", + workspaceId: "ws", + messageId: "m1", + toolCallId: "t1", + toolName: "file_edit_replace_string", + result: {}, + timestamp: Date.now(), + }); + + expect(onPostCompactionStateChange).toHaveBeenCalledTimes(0); + + session.dispose(); + }); +}); diff --git a/src/node/services/agentSession.ts b/src/node/services/agentSession.ts index 57dffc0fbc..ae2f4fdb12 100644 --- a/src/node/services/agentSession.ts +++ b/src/node/services/agentSession.ts @@ -90,6 +90,8 @@ interface AgentSessionOptions { backgroundProcessManager: BackgroundProcessManager; /** Called after compaction completes to trigger metadata refresh */ onCompactionComplete?: () => void; + /** Called when post-compaction context state may have changed (plan/file edits) */ + onPostCompactionStateChange?: () => void; } export class AgentSession { @@ -101,6 +103,7 @@ export class AgentSession { private readonly initStateManager: InitStateManager; private readonly backgroundProcessManager: BackgroundProcessManager; private readonly onCompactionComplete?: () => void; + private readonly onPostCompactionStateChange?: () => void; private readonly emitter = new EventEmitter(); private readonly aiListeners: Array<{ event: string; handler: (...args: unknown[]) => void }> = []; @@ -127,6 +130,11 @@ export class AgentSession { * Used to enable the cooldown-based attachment injection. */ private compactionOccurred = false; + /** + * Cache the last-known experiment state so we don't spam metadata refresh + * when post-compaction context is disabled. + */ + private postCompactionContextEnabled = false; constructor(options: AgentSessionOptions) { assert(options, "AgentSession requires options"); @@ -139,6 +147,7 @@ export class AgentSession { initStateManager, backgroundProcessManager, onCompactionComplete, + onPostCompactionStateChange, } = options; assert(typeof workspaceId === "string", "workspaceId must be a string"); @@ -153,6 +162,7 @@ export class AgentSession { this.initStateManager = initStateManager; this.backgroundProcessManager = backgroundProcessManager; this.onCompactionComplete = onCompactionComplete; + this.onPostCompactionStateChange = onPostCompactionStateChange; this.compactionHandler = new CompactionHandler({ workspaceId: this.workspaceId, @@ -536,6 +546,10 @@ export class AgentSession { } const historyResult = await this.historyService.getHistory(this.workspaceId); + // Cache whether post-compaction context is enabled for this session. + // Used to decide whether tool-call-end should trigger metadata refresh. + this.postCompactionContextEnabled = Boolean(options?.experiments?.postCompactionContext); + if (!historyResult.success) { return Err(createUnknownSendMessageError(historyResult.error)); } @@ -601,6 +615,18 @@ export class AgentSession { forward("tool-call-delta", (payload) => this.emitChatEvent(payload)); forward("tool-call-end", (payload) => { this.emitChatEvent(payload); + + // If post-compaction context is enabled, certain tools can change what should + // be displayed/injected (plan writes, tracked file diffs). Trigger a metadata + // refresh so the right sidebar updates without requiring an experiment toggle. + if ( + this.postCompactionContextEnabled && + payload.type === "tool-call-end" && + (payload.toolName === "propose_plan" || payload.toolName.startsWith("file_edit_")) + ) { + this.onPostCompactionStateChange?.(); + } + // Tool call completed: auto-send queued messages this.sendQueuedMessages(); }); diff --git a/src/node/services/workspaceService.test.ts b/src/node/services/workspaceService.test.ts index 49cda4476f..3df4320f22 100644 --- a/src/node/services/workspaceService.test.ts +++ b/src/node/services/workspaceService.test.ts @@ -6,6 +6,7 @@ import type { PartialService } from "./partialService"; import type { AIService } from "./aiService"; import type { InitStateManager } from "./initStateManager"; import type { ExtensionMetadataService } from "./ExtensionMetadataService"; +import type { FrontendWorkspaceMetadata } from "@/common/types/workspace"; import type { BackgroundProcessManager } from "./backgroundProcessManager"; // Helper to access private renamingWorkspaces set @@ -14,6 +15,8 @@ function addToRenamingWorkspaces(service: WorkspaceService, workspaceId: string) (service as any).renamingWorkspaces.add(workspaceId); } +// NOTE: This test file uses bun:test mocks (not Jest). + describe("WorkspaceService rename lock", () => { let workspaceService: WorkspaceService; let mockAIService: AIService; @@ -116,3 +119,116 @@ describe("WorkspaceService rename lock", () => { } }); }); + +describe("WorkspaceService post-compaction metadata refresh", () => { + let workspaceService: WorkspaceService; + + beforeEach(() => { + const aiService: AIService = { + isStreaming: mock(() => false), + getWorkspaceMetadata: mock(() => + Promise.resolve({ success: false as const, error: "not found" }) + ), + on(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + off(_eventName: string | symbol, _listener: (...args: unknown[]) => void) { + return this; + }, + } as unknown as AIService; + + const mockHistoryService: Partial = { + getHistory: mock(() => Promise.resolve({ success: true as const, data: [] })), + appendToHistory: mock(() => Promise.resolve({ success: true as const, data: undefined })), + }; + + const mockConfig: Partial = { + srcDir: "/tmp/test", + getSessionDir: mock(() => "/tmp/test/sessions"), + generateStableId: mock(() => "test-id"), + findWorkspace: mock(() => null), + }; + + const mockPartialService: Partial = { + commitToHistory: mock(() => Promise.resolve({ success: true as const, data: undefined })), + }; + + const mockInitStateManager: Partial = {}; + const mockExtensionMetadataService: Partial = {}; + const mockBackgroundProcessManager: Partial = { + cleanup: mock(() => Promise.resolve()), + }; + + workspaceService = new WorkspaceService( + mockConfig as Config, + mockHistoryService as HistoryService, + mockPartialService as PartialService, + aiService, + mockInitStateManager as InitStateManager, + mockExtensionMetadataService as ExtensionMetadataService, + mockBackgroundProcessManager as BackgroundProcessManager + ); + }); + + test("debounces multiple refresh requests into a single metadata emit", async () => { + const workspaceId = "ws-post-compaction"; + + const emitMetadata = mock(() => undefined); + + interface WorkspaceServiceTestAccess { + sessions: Map void }>; + getInfo: (workspaceId: string) => Promise; + getPostCompactionState: (workspaceId: string) => Promise<{ + planPath: string | null; + trackedFilePaths: string[]; + excludedItems: string[]; + }>; + schedulePostCompactionMetadataRefresh: (workspaceId: string) => void; + } + + const svc = workspaceService as unknown as WorkspaceServiceTestAccess; + svc.sessions.set(workspaceId, { emitMetadata }); + + const fakeMetadata: FrontendWorkspaceMetadata = { + id: workspaceId, + name: "ws", + projectName: "proj", + projectPath: "/tmp/proj", + namedWorkspacePath: "/tmp/proj/ws", + runtimeConfig: { type: "local", srcBaseDir: "/tmp" }, + }; + + const getInfoMock: WorkspaceServiceTestAccess["getInfo"] = mock(() => + Promise.resolve(fakeMetadata) + ); + + const postCompactionState = { + planPath: "~/.mux/plans/cmux/plan.md", + trackedFilePaths: ["/tmp/proj/file.ts"], + excludedItems: [], + }; + + const getPostCompactionStateMock: WorkspaceServiceTestAccess["getPostCompactionState"] = mock( + () => Promise.resolve(postCompactionState) + ); + + svc.getInfo = getInfoMock; + svc.getPostCompactionState = getPostCompactionStateMock; + + svc.schedulePostCompactionMetadataRefresh(workspaceId); + svc.schedulePostCompactionMetadataRefresh(workspaceId); + svc.schedulePostCompactionMetadataRefresh(workspaceId); + + // Debounce is short, but use a safe buffer. + await new Promise((resolve) => setTimeout(resolve, 150)); + + expect(getInfoMock).toHaveBeenCalledTimes(1); + expect(getPostCompactionStateMock).toHaveBeenCalledTimes(1); + expect(emitMetadata).toHaveBeenCalledTimes(1); + + const enriched = (emitMetadata as ReturnType).mock.calls[0][0] as { + postCompaction?: { planPath: string | null }; + }; + expect(enriched.postCompaction?.planPath).toBe(postCompactionState.planPath); + }); +}); diff --git a/src/node/services/workspaceService.ts b/src/node/services/workspaceService.ts index b87d0175c9..64cf56fa9b 100644 --- a/src/node/services/workspaceService.ts +++ b/src/node/services/workspaceService.ts @@ -53,6 +53,9 @@ import { movePlanFile } from "@/node/utils/runtime/helpers"; /** Maximum number of retry attempts when workspace name collides */ const MAX_WORKSPACE_NAME_COLLISION_RETRIES = 3; +// Keep short to feel instant, but debounce bursts of file_edit_* tool calls. +const POST_COMPACTION_METADATA_REFRESH_DEBOUNCE_MS = 100; + /** * Checks if an error indicates a workspace name collision */ @@ -90,6 +93,9 @@ export class WorkspaceService extends EventEmitter { string, { chat: () => void; metadata: () => void } >(); + + // Debounce post-compaction metadata refreshes (file_edit_* can fire rapidly) + private readonly postCompactionRefreshTimers = new Map(); // Tracks workspaces currently being renamed to prevent streaming during rename private readonly renamingWorkspaces = new Set(); @@ -224,6 +230,45 @@ export class WorkspaceService extends EventEmitter { }; } + private schedulePostCompactionMetadataRefresh(workspaceId: string): void { + assert(typeof workspaceId === "string", "workspaceId must be a string"); + const trimmed = workspaceId.trim(); + assert(trimmed.length > 0, "workspaceId must not be empty"); + + const existing = this.postCompactionRefreshTimers.get(trimmed); + if (existing) { + clearTimeout(existing); + } + + const timer = setTimeout(() => { + this.postCompactionRefreshTimers.delete(trimmed); + void this.emitPostCompactionMetadata(trimmed); + }, POST_COMPACTION_METADATA_REFRESH_DEBOUNCE_MS); + + this.postCompactionRefreshTimers.set(trimmed, timer); + } + + private async emitPostCompactionMetadata(workspaceId: string): Promise { + try { + const session = this.sessions.get(workspaceId); + if (!session) { + return; + } + + const metadata = await this.getInfo(workspaceId); + if (!metadata) { + return; + } + + const postCompaction = await this.getPostCompactionState(workspaceId); + const enrichedMetadata = { ...metadata, postCompaction }; + session.emitMetadata(enrichedMetadata); + } catch (error) { + // Workspace runtime unavailable (e.g., SSH unreachable) - skip emitting post-compaction state. + log.debug("Failed to emit post-compaction metadata", { workspaceId, error }); + } + } + public getOrCreateSession(workspaceId: string): AgentSession { assert(typeof workspaceId === "string", "workspaceId must be a string"); const trimmed = workspaceId.trim(); @@ -243,16 +288,10 @@ export class WorkspaceService extends EventEmitter { initStateManager: this.initStateManager, backgroundProcessManager: this.backgroundProcessManager, onCompactionComplete: () => { - // Emit updated metadata with postCompaction state after compaction - void (async () => { - const metadata = await this.getInfo(trimmed); - if (metadata) { - const postCompaction = await this.getPostCompactionState(trimmed); - const enrichedMetadata = { ...metadata, postCompaction }; - // Look up session from map (guaranteed to exist by the time callback runs) - this.sessions.get(trimmed)?.emitMetadata(enrichedMetadata); - } - })(); + this.schedulePostCompactionMetadataRefresh(trimmed); + }, + onPostCompactionStateChange: () => { + this.schedulePostCompactionMetadataRefresh(trimmed); }, }); @@ -277,20 +316,27 @@ export class WorkspaceService extends EventEmitter { } public disposeSession(workspaceId: string): void { - const session = this.sessions.get(workspaceId); + const trimmed = workspaceId.trim(); + const session = this.sessions.get(trimmed); + const refreshTimer = this.postCompactionRefreshTimers.get(trimmed); + if (refreshTimer) { + clearTimeout(refreshTimer); + this.postCompactionRefreshTimers.delete(trimmed); + } + if (!session) { return; } - const subscriptions = this.sessionSubscriptions.get(workspaceId); + const subscriptions = this.sessionSubscriptions.get(trimmed); if (subscriptions) { subscriptions.chat(); subscriptions.metadata(); - this.sessionSubscriptions.delete(workspaceId); + this.sessionSubscriptions.delete(trimmed); } session.dispose(); - this.sessions.delete(workspaceId); + this.sessions.delete(trimmed); } /** @@ -1084,12 +1130,14 @@ export class WorkspaceService extends EventEmitter { }); try { - // Use exec to delete files since runtime doesn't have a deleteFile method - // Delete both paths in one command for efficiency - await runtime.exec(`rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { + // Use exec to delete files since runtime doesn't have a deleteFile method. + // Delete both paths in one command for efficiency. + const execStream = await runtime.exec(`rm -f ${quotedPlanPath} ${quotedLegacyPlanPath}`, { cwd: metadata.projectPath, timeout: 10, }); + // Wait for completion so callers can rely on the plan file actually being removed. + await execStream.exitCode; } catch { // Plan files don't exist or can't be deleted - ignore }