diff --git a/package-lock.json b/package-lock.json index d758d26..06b7034 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@jackchen_me/open-multi-agent", - "version": "0.2.0", + "version": "1.0.1", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@jackchen_me/open-multi-agent", - "version": "0.2.0", + "version": "1.0.1", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.52.0", diff --git a/src/orchestrator/orchestrator.ts b/src/orchestrator/orchestrator.ts index 0acedb6..1526827 100644 --- a/src/orchestrator/orchestrator.ts +++ b/src/orchestrator/orchestrator.ts @@ -264,6 +264,8 @@ interface RunContext { readonly config: OrchestratorConfig /** Trace run ID, present when `onTrace` is configured. */ readonly runId?: string + /** AbortSignal for run-level cancellation. Checked between task dispatch rounds. */ + readonly abortSignal?: AbortSignal } /** @@ -295,6 +297,15 @@ async function executeQueue( : undefined while (true) { + // Check for cancellation before each dispatch round. + if (ctx.abortSignal?.aborted) { + // Mark all remaining pending tasks as skipped. + for (const t of queue.getByStatus('pending')) { + queue.update(t.id, { status: 'skipped' as TaskStatus }) + } + break + } + // Re-run auto-assignment each iteration so tasks that were unblocked since // the last round (and thus have no assignee yet) get assigned before dispatch. scheduler.autoAssign(queue, team.getAgents()) @@ -360,8 +371,8 @@ async function executeQueue( // Build trace context for this task's agent run const traceOptions: Partial | undefined = config.onTrace - ? { onTrace: config.onTrace, runId: ctx.runId ?? '', taskId: task.id, traceAgent: assignee } - : undefined + ? { onTrace: config.onTrace, runId: ctx.runId ?? '', taskId: task.id, traceAgent: assignee, abortSignal: ctx.abortSignal } + : ctx.abortSignal ? { abortSignal: ctx.abortSignal } : undefined const taskStartMs = config.onTrace ? Date.now() : 0 let retryCount = 0 @@ -638,7 +649,7 @@ export class OpenMultiAgent { * @param team - A team created via {@link createTeam} (or `new Team(...)`). * @param goal - High-level natural-language goal for the team. */ - async runTeam(team: Team, goal: string): Promise { + async runTeam(team: Team, goal: string, options?: { abortSignal?: AbortSignal }): Promise { const agentConfigs = team.getAgents() // ------------------------------------------------------------------ @@ -665,8 +676,8 @@ export class OpenMultiAgent { }) const decompTraceOptions: Partial | undefined = this.config.onTrace - ? { onTrace: this.config.onTrace, runId: runId ?? '', traceAgent: 'coordinator' } - : undefined + ? { onTrace: this.config.onTrace, runId: runId ?? '', traceAgent: 'coordinator', abortSignal: options?.abortSignal } + : options?.abortSignal ? { abortSignal: options.abortSignal } : undefined const decompositionResult = await coordinatorAgent.run(decompositionPrompt, decompTraceOptions) const agentResults = new Map() agentResults.set('coordinator:decompose', decompositionResult) @@ -712,6 +723,7 @@ export class OpenMultiAgent { agentResults, config: this.config, runId, + abortSignal: options?.abortSignal, } await executeQueue(queue, ctx) @@ -764,6 +776,7 @@ export class OpenMultiAgent { retryDelayMs?: number retryBackoff?: number }>, + options?: { abortSignal?: AbortSignal }, ): Promise { const agentConfigs = team.getAgents() const queue = new TaskQueue() @@ -794,6 +807,7 @@ export class OpenMultiAgent { agentResults, config: this.config, runId: this.config.onTrace ? generateRunId() : undefined, + abortSignal: options?.abortSignal, } await executeQueue(queue, ctx) diff --git a/tests/abort-signal.test.ts b/tests/abort-signal.test.ts new file mode 100644 index 0000000..4a86f14 --- /dev/null +++ b/tests/abort-signal.test.ts @@ -0,0 +1,107 @@ +import { describe, it, expect, vi } from 'vitest' +import { OpenMultiAgent } from '../src/orchestrator/orchestrator.js' +import { Team } from '../src/team/team.js' + +describe('AbortSignal support for runTeam and runTasks', () => { + it('runTeam should accept an abortSignal option', async () => { + const orchestrator = new OpenMultiAgent({ + defaultModel: 'test-model', + defaultProvider: 'openai', + }) + + // Verify the API accepts the option without throwing + const controller = new AbortController() + const team = new Team('test', { + name: 'test', + agents: [ + { name: 'agent1', model: 'test-model', systemPrompt: 'test' }, + ], + }) + + // Abort immediately so the run won't actually execute LLM calls + controller.abort() + + // runTeam should return gracefully (no unhandled rejection) + const result = await orchestrator.runTeam(team, 'test goal', { + abortSignal: controller.signal, + }) + + // With immediate abort, coordinator may or may not have run, + // but the function should not throw. + expect(result).toBeDefined() + expect(result.agentResults).toBeInstanceOf(Map) + }) + + it('runTasks should accept an abortSignal option', async () => { + const orchestrator = new OpenMultiAgent({ + defaultModel: 'test-model', + defaultProvider: 'openai', + }) + + const controller = new AbortController() + const team = new Team('test', { + name: 'test', + agents: [ + { name: 'agent1', model: 'test-model', systemPrompt: 'test' }, + ], + }) + + controller.abort() + + const result = await orchestrator.runTasks(team, [ + { title: 'task1', description: 'do something', assignee: 'agent1' }, + ], { abortSignal: controller.signal }) + + expect(result).toBeDefined() + expect(result.agentResults).toBeInstanceOf(Map) + }) + + it('pre-aborted signal should skip pending tasks', async () => { + const orchestrator = new OpenMultiAgent({ + defaultModel: 'test-model', + defaultProvider: 'openai', + }) + + const controller = new AbortController() + controller.abort() + + const team = new Team('test', { + name: 'test', + agents: [ + { name: 'agent1', model: 'test-model', systemPrompt: 'test' }, + ], + }) + + const result = await orchestrator.runTasks(team, [ + { title: 'task1', description: 'first', assignee: 'agent1' }, + { title: 'task2', description: 'second', assignee: 'agent1' }, + ], { abortSignal: controller.signal }) + + // No agent runs should complete since signal was already aborted + expect(result).toBeDefined() + }) + + it('runTeam and runTasks work without abortSignal (backward compat)', async () => { + const orchestrator = new OpenMultiAgent({ + defaultModel: 'test-model', + defaultProvider: 'openai', + }) + + const team = new Team('test', { + name: 'test', + agents: [ + { name: 'agent1', model: 'test-model', systemPrompt: 'test' }, + ], + }) + + // These should not throw even without abortSignal + const promise1 = orchestrator.runTeam(team, 'goal') + const promise2 = orchestrator.runTasks(team, [ + { title: 'task1', description: 'do something', assignee: 'agent1' }, + ]) + + // Both return promises (won't resolve without real LLM, but API is correct) + expect(promise1).toBeInstanceOf(Promise) + expect(promise2).toBeInstanceOf(Promise) + }) +})