From e51fcb4026bebd0e7a8634fe0f421b0adb766304 Mon Sep 17 00:00:00 2001 From: brooksc Date: Sun, 24 May 2026 00:47:27 -0700 Subject: [PATCH] fix(coordinator): make control handoff authoritative --- electron/ipc/register.ts | 3 +- electron/mcp/coordinator.test.ts | 30 +++++++++++++++++++- electron/mcp/coordinator.ts | 14 +++++++--- src/store/taskStatus.ts | 10 +++++-- src/store/tasks.test.ts | 48 ++++++++++++++++++++++++++++---- src/store/tasks.ts | 16 ++++++++--- 6 files changed, 103 insertions(+), 18 deletions(-) diff --git a/electron/ipc/register.ts b/electron/ipc/register.ts index f6b4151c..f561f4a2 100644 --- a/electron/ipc/register.ts +++ b/electron/ipc/register.ts @@ -1174,7 +1174,8 @@ export function registerAllHandlers(win: BrowserWindow): void { if (args.controlledBy !== 'coordinator' && args.controlledBy !== 'human') { throw new Error(`Invalid controlledBy: ${String(args.controlledBy)}`); } - coordinator?.setTaskControl(args.taskId, args.controlledBy); + if (!coordinator) throw new Error('coordinator mode not initialized'); + coordinator.setTaskControl(args.taskId, args.controlledBy); }, ); diff --git a/electron/mcp/coordinator.test.ts b/electron/mcp/coordinator.test.ts index 180f92f1..7c353f74 100644 --- a/electron/mcp/coordinator.test.ts +++ b/electron/mcp/coordinator.test.ts @@ -729,6 +729,12 @@ describe('Coordinator waitForIdle', () => { await expect(coordinator.waitForIdle('task-1')).resolves.toEqual({ reason: 'human_control' }); }); + it('rejects control changes for unknown tasks', () => { + expect(() => coordinator.setTaskControl('missing-task', 'human')).toThrow( + 'Task not found: missing-task', + ); + }); + it('rejects after timeout when task never idles', async () => { vi.useFakeTimers(); await coordinator.createTask({ name: 'test', prompt: 'do', coordinatorTaskId: 'coord-1' }); @@ -753,6 +759,25 @@ describe('Coordinator waitForIdle', () => { coordinator.setTaskControl('task-1', 'coordinator'); await expect(waitPromise).resolves.toEqual({ reason: 'idle' }); }); + + it('notifies coordinator when releasing control after waitForIdle was interrupted', async () => { + await coordinator.createTask({ name: 'my-task', prompt: 'do', coordinatorTaskId: 'coord-1' }); + const waitPromise = coordinator.waitForIdle('task-1'); + + coordinator.setTaskControl('task-1', 'human'); + await expect(waitPromise).resolves.toEqual({ reason: 'human_control' }); + mockNotifyRenderer.mockClear(); + + coordinator.setTaskControl('task-1', 'coordinator'); + + expect(mockNotifyRenderer).toHaveBeenCalledWith( + 'mcp_coordinator_notification_staged', + expect.objectContaining({ + coordinatorTaskId: 'coord-1', + text: expect.stringContaining('"my-task" has been returned to coordinator control'), + }), + ); + }); }); // ─── waitForSignalDone tests ────────────────────────────────────────────────── @@ -2462,10 +2487,11 @@ describe('Coordinator removeCoordinatedTask', () => { expect(vi.mocked(unsubscribeFromAgent)).toHaveBeenCalledWith(agentId, expect.any(Function)); }); - it('cleans up internal resource maps (subscribers, tailBuffers, decoders, controlMap, blockedByHumanControl)', async () => { + it('cleans up internal resource maps (subscribers, tailBuffers, decoders, controlMap, human control markers)', async () => { await coordinator.createTask({ name: 'test', prompt: 'do', coordinatorTaskId: 'coord-1' }); const agentId = getAgentId(); coordinator.setTaskControl('task-1', 'human'); + await coordinator.waitForIdle('task-1'); await coordinator.sendPrompt('task-1', 'hello').catch(() => {}); coordinator.removeCoordinatedTask('task-1'); @@ -2476,12 +2502,14 @@ describe('Coordinator removeCoordinatedTask', () => { decoders: Map; controlMap: Map; blockedByHumanControl: Set; + interruptedByHumanControl: Set; }; expect(c.subscribers.has(agentId)).toBe(false); expect(c.tailBuffers.has(agentId)).toBe(false); expect(c.decoders.has(agentId)).toBe(false); expect(c.controlMap.has('task-1')).toBe(false); expect(c.blockedByHumanControl.has('task-1')).toBe(false); + expect(c.interruptedByHumanControl.has('task-1')).toBe(false); }); it('deregister detaches child task state and preserves review only after prompt delivery', async () => { diff --git a/electron/mcp/coordinator.ts b/electron/mcp/coordinator.ts index 647f1d0b..133a3fa6 100644 --- a/electron/mcp/coordinator.ts +++ b/electron/mcp/coordinator.ts @@ -76,6 +76,7 @@ export class Coordinator { private decoders = new Map(); private controlMap = new Map(); private blockedByHumanControl = new Set(); + private interruptedByHumanControl = new Set(); private closingTaskIds = new Set(); private activeSignalWaitCounts = new Map(); private recentlyDelivered = new ReplayCache(); @@ -158,14 +159,14 @@ export class Coordinator { setTaskControl(taskId: string, who: 'coordinator' | 'human'): void { if (!this.tasks.has(taskId)) { - console.warn(`setTaskControl: unknown taskId ${taskId}`); - return; + throw new Error(`Task not found: ${taskId}`); } this.controlMap.set(taskId, who); if (who === 'human') { // Resolve any pending idle waiters immediately — human has taken over const resolvers = this.idleResolvers.get(taskId); if (resolvers?.length) { + this.interruptedByHumanControl.add(taskId); for (const resolve of resolvers) resolve({ reason: 'human_control' }); this.idleResolvers.delete(taskId); } @@ -177,9 +178,10 @@ export class Coordinator { for (const resolve of resolvers) resolve({ reason: 'idle' }); this.idleResolvers.delete(taskId); } - // Notify coordinator if it tried to send a prompt while blocked - if (this.blockedByHumanControl.has(taskId)) { + // Notify coordinator if it tried to send a prompt or had an idle wait interrupted. + if (this.blockedByHumanControl.has(taskId) || this.interruptedByHumanControl.has(taskId)) { this.blockedByHumanControl.delete(taskId); + this.interruptedByHumanControl.delete(taskId); const task = this.tasks.get(taskId); const coordinator = task ? this.coordinators.get(task.coordinatorTaskId) : null; if (task && coordinator) { @@ -811,6 +813,7 @@ export class Coordinator { const task = this.tasks.get(taskId); if (!task) return Promise.reject(new Error(`Task not found: ${taskId}`)); if (this.controlMap.get(taskId) === 'human') { + this.interruptedByHumanControl.add(taskId); return Promise.resolve({ reason: 'human_control' }); // resolve immediately — caller gets control-change event instead } if (task.status === 'exited') return Promise.resolve({ reason: 'exited' }); @@ -1037,6 +1040,7 @@ export class Coordinator { this.tasks.delete(taskId); this.blockedByHumanControl.delete(taskId); + this.interruptedByHumanControl.delete(taskId); this.controlMap.delete(taskId); } @@ -1121,6 +1125,7 @@ export class Coordinator { this.tasks.delete(taskId); this.controlMap.delete(taskId); this.blockedByHumanControl.delete(taskId); + this.interruptedByHumanControl.delete(taskId); this.closingTaskIds.delete(taskId); // Notify renderer @@ -1419,6 +1424,7 @@ export class Coordinator { // Transfer control to human so the user can decide what to do with orphaned tasks this.controlMap.set(taskId, 'human'); this.blockedByHumanControl.delete(taskId); + this.interruptedByHumanControl.delete(taskId); if (task.mcpConfigPath) { try { diff --git a/src/store/taskStatus.ts b/src/store/taskStatus.ts index b2f71645..37430157 100644 --- a/src/store/taskStatus.ts +++ b/src/store/taskStatus.ts @@ -629,8 +629,14 @@ function tryAutoTrust(agentId: string, rawTail: string): boolean { // auto-trust suppressed the question state, release it back to coordinator. const taskId = state.taskId; if (taskId && isAutoTrustForced(agentId) && store.tasks[taskId]?.controlledBy === 'human') { - setStore('tasks', taskId, 'controlledBy', 'coordinator'); - invoke(IPC.MCP_ControlChanged, { taskId, controlledBy: 'coordinator' }).catch(() => {}); + invoke(IPC.MCP_ControlChanged, { taskId, controlledBy: 'coordinator' }) + .then(() => setStore('tasks', taskId, 'controlledBy', 'coordinator')) + .catch((err) => { + logWarn('tasks.autoTrust', 'MCP_ControlChanged failed during auto-trust release', { + taskId, + err, + }); + }); } // Cooldown: ignore trust patterns for 1s so the same dialog // isn't re-matched while the PTY output transitions. diff --git a/src/store/tasks.test.ts b/src/store/tasks.test.ts index 4c3b1b20..0365f641 100644 --- a/src/store/tasks.test.ts +++ b/src/store/tasks.test.ts @@ -2,11 +2,14 @@ import { describe, expect, it, vi, beforeEach } from 'vitest'; import { IPC } from '../../electron/ipc/channels'; // Hoisted so these refs are available both in vi.mock() factories and in test bodies. -const { mockInvoke, mockIsAgentBracketedPasteEnabled, mockSetStore } = vi.hoisted(() => ({ - mockInvoke: vi.fn(), - mockIsAgentBracketedPasteEnabled: vi.fn(), - mockSetStore: vi.fn(), -})); +const { mockInvoke, mockIsAgentBracketedPasteEnabled, mockSaveState, mockSetStore } = vi.hoisted( + () => ({ + mockInvoke: vi.fn(), + mockIsAgentBracketedPasteEnabled: vi.fn(), + mockSaveState: vi.fn(), + mockSetStore: vi.fn(), + }), +); // ─── Coordinator test infrastructure ───────────────────────────────────────── @@ -82,7 +85,7 @@ vi.mock('./core', () => ({ })); vi.mock('../lib/ipc', () => ({ Channel: vi.fn(), invoke: mockInvoke })); -vi.mock('./persistence', () => ({ saveState: vi.fn() })); +vi.mock('./persistence', () => ({ saveState: mockSaveState })); vi.mock('./focus', () => ({ setTaskFocusedPanel: vi.fn() })); vi.mock('./projects', () => ({ getProject: vi.fn(), @@ -156,6 +159,7 @@ beforeEach(() => { mockCollapsedTaskOrder = []; mockProjects = []; mockInvoke.mockResolvedValue(undefined); + mockSaveState.mockResolvedValue(undefined); }); // ─── Coordinator tests ──────────────────────────────────────────────────────── @@ -200,6 +204,38 @@ describe('coordinator controlledBy state machine (item 9: UI disabled-state regr it('9d: setTaskControl is a no-op for unknown taskId (no crash)', () => { expect(() => setTaskControl('nonexistent-task', 'human')).not.toThrow(); + expect(mockInvoke).not.toHaveBeenCalledWith(IPC.MCP_ControlChanged, expect.anything()); + }); + + it('rolls back sub-task control when backend acknowledgement fails', async () => { + taskCreatedHandler(baseEvent); + mockInvoke.mockRejectedValueOnce(new Error('coordinator missing')); + + setTaskControl('sub-task-1', 'human'); + expect(mockTasks['sub-task-1'].controlledBy).toBe('human'); + + await Promise.resolve(); + await Promise.resolve(); + + expect(mockTasks['sub-task-1'].controlledBy).toBe('coordinator'); + expect(mockSaveState).toHaveBeenCalledTimes(1); + }); + + it('persists sub-task control only after backend acknowledgement succeeds', async () => { + taskCreatedHandler(baseEvent); + + setTaskControl('sub-task-1', 'human'); + expect(mockSaveState).not.toHaveBeenCalled(); + + await Promise.resolve(); + await Promise.resolve(); + + expect(mockTasks['sub-task-1'].controlledBy).toBe('human'); + expect(mockInvoke).toHaveBeenCalledWith(IPC.MCP_ControlChanged, { + taskId: 'sub-task-1', + controlledBy: 'human', + }); + expect(mockSaveState).toHaveBeenCalledTimes(1); }); it('9e: removing a coordinator task leaves no entry in mockTasks', () => { diff --git a/src/store/tasks.ts b/src/store/tasks.ts index 66ddf733..0ec3ec7f 100644 --- a/src/store/tasks.ts +++ b/src/store/tasks.ts @@ -1333,17 +1333,25 @@ export function retryTaskMcpStartup(taskId: string): Promise { export function setTaskControl(taskId: string, who: 'coordinator' | 'human'): void { const task = store.tasks[taskId]; + if (!task) return; const prev = task?.controlledBy; setStore('tasks', taskId, 'controlledBy', who); // Coordinator tasks manage their own control state in the frontend only. // Sub-tasks need to notify the backend Coordinator so it can gate send_prompt. - if (!task?.coordinatorMode) { - invoke(IPC.MCP_ControlChanged, { taskId, controlledBy: who }).catch((err: unknown) => { + if (task.coordinatorMode) { + void saveState(); + return; + } + + invoke(IPC.MCP_ControlChanged, { taskId, controlledBy: who }) + .then(() => { + void saveState(); + }) + .catch((err: unknown) => { console.warn('[tasks] setTaskControl IPC failed, rolling back controlledBy', err); setStore('tasks', taskId, 'controlledBy', prev); + void saveState(); }); - } - void saveState(); } export function setPlanContent(