diff --git a/packages/harness/src/session-registry.ts b/packages/harness/src/session-registry.ts index 8f8678ac..31759bce 100644 --- a/packages/harness/src/session-registry.ts +++ b/packages/harness/src/session-registry.ts @@ -155,6 +155,10 @@ export class SessionRegistry { return this.sessions.get(clientId); } + keys(): IterableIterator { + return this.sessions.keys(); + } + entries(): IterableIterator<[string, ManagedSession]> { return this.sessions.entries(); } diff --git a/server/__tests__/task-orchestrator.test.ts b/server/__tests__/task-orchestrator.test.ts index a3e026f4..c722a05b 100644 --- a/server/__tests__/task-orchestrator.test.ts +++ b/server/__tests__/task-orchestrator.test.ts @@ -312,18 +312,19 @@ describe('TaskOrchestrator', () => { describe('orphan detection', () => { it('reclaims orphaned tasks during tick', () => { - // Create deps with getActiveSessionIds + // getActiveSessionIds returns clientIds (not SDK sessionIds) + // to match what setSessionId stores on tasks const depsWithOrphan = createTestDeps(store); - depsWithOrphan.getActiveSessionIds = () => new Set(['alive-session']); + depsWithOrphan.getActiveSessionIds = () => new Set(['alive-client']); const orch = new TaskOrchestrator(depsWithOrphan); const goal = store.create({ title: 'Goal' }); const c1 = store.create({ title: 'Orphan', parentId: goal.id }); store.create({ title: 'Next', parentId: goal.id }); - // Simulate c1 assigned to dead session + // Simulate c1 assigned to dead session (clientId not in active set) store.update(c1.id, { status: 'active' }); - store.setSessionId(c1.id, 'dead-session'); + store.setSessionId(c1.id, 'dead-client'); orch.start(goal.id); @@ -335,7 +336,7 @@ describe('TaskOrchestrator', () => { it('does not reclaim tasks with alive sessions', () => { const depsWithOrphan = createTestDeps(store); - depsWithOrphan.getActiveSessionIds = () => new Set(['alive-session']); + depsWithOrphan.getActiveSessionIds = () => new Set(['alive-client']); const orch = new TaskOrchestrator(depsWithOrphan); const goal = store.create({ title: 'Goal' }); @@ -343,13 +344,42 @@ describe('TaskOrchestrator', () => { const c2 = store.create({ title: 'Next', parentId: goal.id }); store.update(c1.id, { status: 'active' }); - store.setSessionId(c1.id, 'alive-session'); + store.setSessionId(c1.id, 'alive-client'); orch.start(goal.id); // c1 is alive, so tick should skip it and pick c2 expect(orch.getStatus().activeTaskId).toBe(c2.id); }); + + it('reclaims spawned task sessions using clientId matching', () => { + // Simulates the real scenario: task.session_id stores clientId + // (e.g. 'task:abc123'), active set contains clientIds from registry + const depsWithOrphan = createTestDeps(store); + depsWithOrphan.getActiveSessionIds = () => new Set(['task:alive-wt']); + const orch = new TaskOrchestrator(depsWithOrphan); + + const goal = store.create({ title: 'Goal' }); + const orphan = store.create({ title: 'Dead spawn', parentId: goal.id }); + const alive = store.create({ title: 'Alive spawn', parentId: goal.id }); + store.create({ title: 'Pending', parentId: goal.id }); + + // Orphan: session ended, clientId no longer in registry + store.update(orphan.id, { status: 'active' }); + store.setSessionId(orphan.id, 'task:dead-wt'); + + // Alive: session still running + store.update(alive.id, { status: 'active' }); + store.setSessionId(alive.id, 'task:alive-wt'); + + orch.start(goal.id); + + // Orphan should be reclaimed and re-dispatched + expect(store.get(orphan.id)!.sessionId).toBeNull(); + // Alive should NOT be reclaimed + expect(store.get(alive.id)!.sessionId).toBe('task:alive-wt'); + expect(store.get(alive.id)!.status).toBe('active'); + }); }); describe('task approval', () => { @@ -882,6 +912,149 @@ describe('TaskOrchestrator', () => { }); }); + it('tick after spawn session completes advances workflow', async () => { + const spawnSession = vi.fn().mockResolvedValue('task:spawned-1'); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + let activeClients = new Set(['task:spawned-1']); + deps.getActiveSessionIds = () => activeClients; + const orch = new TaskOrchestrator(deps); + + const goal = store.create({ title: 'Goal' }); + const t1 = store.create({ title: 'Spawn task', parentId: goal.id }); + const t2 = store.create({ title: 'Next task', parentId: goal.id, sessionPolicy: 'reuse' }); + + orch.start(goal.id); + + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalled(); + }); + + // Agent completed the task during the session + store.update(t1.id, { status: 'done' }); + store.cascadeStatus(t1.id); + + // .finally() fires: session removed from registry, tick() called directly + activeClients = new Set(); + orch.tick(); + + // t2 should now be active (workflow advanced) + expect(store.get(t2.id)!.status).toBe('active'); + expect(orch.getStatus().activeTaskId).toBe(t2.id); + }); + + it('resume after spawn session dies reclaims unfinished task', async () => { + // Primary scenario for .finally(): session crashes without calling + // TaskComplete. The task stays active but the session is gone. + // After all tasks are spawned, the orchestrator pauses (no pending left). + // .finally() calls resume() → tick() → orphan reclaim → re-dispatch. + const spawnSession = vi.fn().mockResolvedValue('task:spawned-1'); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + let activeClients = new Set(['task:spawned-1']); + deps.getActiveSessionIds = () => activeClients; + const orch = new TaskOrchestrator(deps); + + const goal = store.create({ title: 'Goal' }); + store.create({ title: 'Spawn task', parentId: goal.id }); + + orch.start(goal.id); + + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalledTimes(1); + }); + + // Orchestrator should be paused (all tasks active, none pending) + expect(orch.getStatus().state).toBe('paused'); + + // Verify the task is active and has sessionId set (pre-condition for orphan reclaim) + const task = store.getChildren(goal.id)[0]; + expect(store.get(task.id)!.status).toBe('active'); + expect(store.get(task.id)!.sessionId).toBe('task:spawned-1'); + + // .finally() fires: session removed from registry, resume() called + activeClients = new Set(); + orch.resume(); + + // Orphan reclaim detected the dead session and re-dispatched the task + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalledTimes(2); + }); + }); + + it('tick is a no-op when orchestrator is stopped before session ends', async () => { + // Edge case: orchestrator stopped (idle) while a spawned session is + // still in-flight. .finally() fires → tick() → returns early (guarded + // by state !== 'running'). No crash, no side effects. + const spawnSession = vi.fn().mockResolvedValue('task:spawned-1'); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + // Return spawned clientId while session is alive + let activeClients = new Set(['task:spawned-1']); + deps.getActiveSessionIds = () => activeClients; + const orch = new TaskOrchestrator(deps); + + const goal = store.create({ title: 'Goal' }); + const t1 = store.create({ title: 'Spawn task', parentId: goal.id }); + + orch.start(goal.id); + + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalledTimes(1); + }); + + // User stops the orchestrator while session is still running + orch.stop(); + expect(orch.getStatus().state).toBe('idle'); + + // .finally() fires after session ends — tick is a no-op in idle state + activeClients = new Set(); + orch.tick(); + + expect(orch.getStatus().state).toBe('idle'); + expect(store.get(t1.id)?.status).toBe('active'); // unchanged + expect(spawnSession).toHaveBeenCalledTimes(1); // no re-dispatch + }); + + it('does not advance old goal when orchestrator moved to a new goal', async () => { + const spawnSession = vi.fn().mockResolvedValue('task:spawned-1'); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + let activeClients = new Set(['task:spawned-1']); + deps.getActiveSessionIds = () => activeClients; + const orch = new TaskOrchestrator(deps); + + const goal1 = store.create({ title: 'Goal 1' }); + store.create({ title: 'Spawn task', parentId: goal1.id }); + + const goal2 = store.create({ title: 'Goal 2' }); + const g2task = store.create({ + title: 'Reuse task', + parentId: goal2.id, + sessionPolicy: 'reuse', + }); + + orch.start(goal1.id); + + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalledTimes(1); + }); + + // User switches to a new goal while spawned session is still running + orch.stop(); + orch.start(goal2.id); + expect(orch.getStatus().goalId).toBe(goal2.id); + expect(orch.getStatus().activeTaskId).toBe(g2task.id); + + // Old spawned session ends — tick should NOT interfere with goal2 + activeClients = new Set(); + orch.tick(); + + // Goal2 state should be unchanged + expect(orch.getStatus().goalId).toBe(goal2.id); + expect(orch.getStatus().activeTaskId).toBe(g2task.id); + }); + it('reuse policy tasks use pinned session as before', () => { const spawnSession = vi.fn().mockResolvedValue('spawned-client-1'); const deps = createTestDeps(store); diff --git a/server/index.ts b/server/index.ts index 77d0bc7e..c1d56bb1 100644 --- a/server/index.ts +++ b/server/index.ts @@ -235,15 +235,12 @@ const orchestrator = new TaskOrchestrator({ sseRegistry.broadcast('task_state', data); }, getActiveSessionIds: () => { - const ids = new Set(); - for (const [clientId] of registry.entries()) { - const session = registry.get(clientId); - if (session?.sessionId) ids.add(session.sessionId); - } - return ids; + // Return clientIds — task.session_id stores clientId, not SDK sessionId. + // Using sessionId here caused orphan detection to never match spawned tasks. + return new Set(registry.keys()); }, spawnSession: async (taskId: string, prompt: string, goalId: string) => { - const clientId = `headless:${generateWtId()}`; + const clientId = `task:${generateWtId()}`; try { const transport = new NullTransport(); @@ -257,16 +254,49 @@ const orchestrator = new TaskOrchestrator({ telosTaskId: goalId, taskContext: { currentTaskId: taskId, goalId }, onSessionResolved: (sessionId) => { - log.info('spawned headless session resolved', { taskId, sessionId, clientId }); + log.info('task session resolved', { taskId, sessionId, clientId }); sseRegistry.broadcast('sessions_changed', {}); }, - }).catch((err) => { - log.error('spawned session failed', { - taskId, - clientId, - error: (err as Error).message, + }) + .catch((err) => { + log.error('task session failed', { + taskId, + clientId, + error: (err as Error).message, + }); + }) + .finally(() => { + // Session ended (success or failure) — clean up registry entry + // (no WS close to trigger normal removal) and advance the + // orchestrator so orphan reclaim picks up unfinished tasks. + registry.remove(clientId); + // Only advance if the orchestrator is still on the same goal. + // A user may have paused manually or started a new goal while + // this session was in-flight — don't override that. + if (!orchestratorRef) return; + try { + const status = orchestratorRef.getStatus(); + if (status.goalId === goalId) { + if (status.state === 'paused') { + orchestratorRef.resume(); + } else { + orchestratorRef.tick(); + } + } else { + log.info('task session ended after loop stopped or goal changed', { + taskId, + clientId, + goalId, + }); + } + } catch (err: unknown) { + log.error('orchestrator advance failed after task session end', { + taskId, + clientId, + error: err instanceof Error ? err.message : String(err), + }); + } }); - }); return clientId; } catch (err) {