diff --git a/server/__tests__/task-orchestrator.test.ts b/server/__tests__/task-orchestrator.test.ts index 6bee3465..a3e026f4 100644 --- a/server/__tests__/task-orchestrator.test.ts +++ b/server/__tests__/task-orchestrator.test.ts @@ -839,6 +839,49 @@ describe('TaskOrchestrator', () => { expect(deps.setTaskContext).not.toHaveBeenCalled(); }); + it('does not call setTaskContext (pinned) for spawned sessions', async () => { + const spawnSession = vi.fn().mockResolvedValue('spawned-client-1'); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + const orch = new TaskOrchestrator(deps); + + const goal = store.create({ title: 'Goal' }); + store.create({ + title: 'Spawn task', + parentId: goal.id, + sessionPolicy: 'spawn', + }); + + orch.start(goal.id); + + await vi.waitFor(() => { + expect(spawnSession).toHaveBeenCalled(); + }); + + // Spawned sessions get taskContext via startChat options, not setTaskContext + expect(deps.setTaskContext).not.toHaveBeenCalled(); + }); + + it('falls back to setTaskContext (pinned) when spawn returns null', async () => { + const spawnSession = vi.fn().mockResolvedValue(null); + const deps = createTestDeps(store); + deps.spawnSession = spawnSession; + const orch = new TaskOrchestrator(deps); + + const goal = store.create({ title: 'Goal' }); + const task = store.create({ + title: 'Spawn task', + parentId: goal.id, + sessionPolicy: 'spawn', + }); + + orch.start(goal.id); + + await vi.waitFor(() => { + expect(deps.setTaskContext).toHaveBeenCalledWith(task.id, goal.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/chat.ts b/server/chat.ts index 50579b77..0a8e681b 100644 --- a/server/chat.ts +++ b/server/chat.ts @@ -745,6 +745,7 @@ export async function startChat( onSessionResolved?: (sessionId: string) => void; telosTaskId?: string; agentName?: string; + taskContext?: { currentTaskId: string; goalId: string }; }, ) { return withSpanAsync( @@ -775,6 +776,7 @@ async function _startChatInner( onSessionResolved?: (sessionId: string) => void; telosTaskId?: string; agentName?: string; + taskContext?: { currentTaskId: string; goalId: string }; }, ) { const abortController = new AbortController(); @@ -862,6 +864,9 @@ async function _startChatInner( const session = registry.get(clientId)!; session.model = options.model ?? session.model; session.inputQueue = inputQueue as { push: (msg: unknown) => void; close: () => void }; + if (options.taskContext) { + session.taskContext = options.taskContext; + } _onSessionChange?.(clientId, 'start'); // Session state machine: mark CREATED (Phase 1 — write only, no behavior change) diff --git a/server/index.ts b/server/index.ts index 8cb4640d..77d0bc7e 100644 --- a/server/index.ts +++ b/server/index.ts @@ -255,6 +255,7 @@ const orchestrator = new TaskOrchestrator({ mode: 'agent', isolation: true, telosTaskId: goalId, + taskContext: { currentTaskId: taskId, goalId }, onSessionResolved: (sessionId) => { log.info('spawned headless session resolved', { taskId, sessionId, clientId }); sseRegistry.broadcast('sessions_changed', {});