import { describe, expect, it, vi } from 'vitest';
import { Agent } from '../../agent/agent.js';
import { MockMessageModel } from '../../__fixtures__/mock-message-model.js';
import { collectGenerator } from '../../__fixtures__/model-test-helpers.js';
import { createCancellableAgent } from '../../__fixtures__/agent-helpers.js';
import { BeforeNodeCallEvent, MultiAgentInitializedEvent } from '../events.js';
import { TextBlock } from '../../types/messages.js';
import { Status, MultiAgentState } from '../state.js';
import { AgentNode } from '../nodes.js';
import { Swarm } from '../swarm.js';
import { SessionManager } from '../../session/session-manager.js';
import { MockSnapshotStorage } from '../../__fixtures__/mock-storage-provider.js';
/**
 * Creates an agent that produces a structured output handoff via the strands_structured_output tool.
 * The agent exits after the structured output tool succeeds (early-exit behavior).
 */
function createHandoffAgent(agentId, handoff, description = `Agent ${agentId}`) {
    const model = new MockMessageModel().addTurn({
        type: 'toolUseBlock',
        name: 'strands_structured_output',
        toolUseId: 'tool-1',
        input: handoff,
    });
    return new Agent({ model, printer: false, id: agentId, description });
}
/**
 * Creates a simple agent that produces a final response (no handoff).
 */
function createFinalAgent(agentId, message, description = `Agent ${agentId}`) {
    return createHandoffAgent(agentId, { message }, description);
}
describe('Swarm', () => {
    describe('constructor', () => {
        it('defaults id to "swarm"', () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            expect(swarm.id).toBe('swarm');
        });
        it('accepts a custom id', () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
                id: 'my-swarm',
            });
            expect(swarm.id).toBe('my-swarm');
        });
        it('accepts AgentNodeOptions with per-node config', () => {
            const swarm = new Swarm({
                nodes: [{ agent: createFinalAgent('a', 'hi') }],
                start: 'a',
            });
            expect(swarm.nodes.get('a')).toBeInstanceOf(AgentNode);
        });
        it('defaults start to the first node when not specified', () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('first', 'hi'), createFinalAgent('second', 'bye')],
            });
            expect(swarm.start.id).toBe('first');
        });
        it('throws when start references unknown agent', () => {
            expect(() => new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'missing',
            })).toThrow('start=<missing> | start references unknown agent');
        });
        it('throws when nodes list is empty', () => {
            expect(() => new Swarm({ nodes: [] })).toThrow('nodes list is empty');
        });
        it('throws on duplicate agent ids', () => {
            const agent = createFinalAgent('a', 'hi');
            expect(() => new Swarm({
                nodes: [agent, agent],
                start: 'a',
            })).toThrow('agent_id=<a> | duplicate agent id');
        });
        it('throws when maxSteps < 1', () => {
            expect(() => new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
                maxSteps: 0,
            })).toThrow('max_steps=<0> | must be at least 1');
        });
        it('defaults maxSteps, timeout, and nodeTimeout to Infinity', () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            expect(swarm.config.maxSteps).toBe(Infinity);
            expect(swarm.config.timeout).toBe(Infinity);
            expect(swarm.config.nodeTimeout).toBe(Infinity);
        });
        it('throws when timeout < 1', () => {
            expect(() => new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
                timeout: 0,
            })).toThrow('timeout=<0> | must be at least 1');
        });
        it('throws when nodeTimeout < 1', () => {
            expect(() => new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
                nodeTimeout: 0,
            })).toThrow('node_timeout=<0> | must be at least 1');
        });
    });
    describe('invoke', () => {
        it('returns completed result with content and duration', async () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'final answer')],
                start: 'a',
            });
            const result = await swarm.invoke('hello');
            expect(result).toEqual(expect.objectContaining({
                status: Status.COMPLETED,
                duration: expect.any(Number),
                content: [expect.objectContaining({ type: 'textBlock', text: 'final answer' })],
            }));
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a']);
            expect(result.results[0]?.structuredOutput).toEqual({ message: 'final answer' });
        });
        it('hands off from A to B and returns final output', async () => {
            const swarm = new Swarm({
                nodes: [
                    createHandoffAgent('a', { agentId: 'b', message: 'please handle this' }),
                    createFinalAgent('b', 'done by b'),
                ],
                start: 'a',
            });
            const result = await swarm.invoke('start');
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'b']);
        });
        it('chains handoffs across multiple agents (A → B → C)', async () => {
            const swarm = new Swarm({
                nodes: [
                    createHandoffAgent('a', { agentId: 'b', message: 'go to b' }),
                    createHandoffAgent('b', { agentId: 'c', message: 'go to c' }),
                    createFinalAgent('c', 'final from c'),
                ],
                start: 'a',
            });
            const result = await swarm.invoke('start');
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'b', 'c']);
        });
        it('passes serialized context in handoff input', async () => {
            const contextData = { key: 'value', num: 42 };
            const agentB = createFinalAgent('b', 'done');
            const streamSpy = vi.spyOn(agentB, 'stream');
            const swarm = new Swarm({
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'handle this', context: contextData }), agentB],
                start: 'a',
            });
            await swarm.invoke('start');
            expect(streamSpy).toHaveBeenCalled();
            const args = streamSpy.mock.calls[0][0];
            const texts = args.map((b) => b.text);
            expect(texts).toContainEqual('handle this');
            expect(texts).toContainEqual(expect.stringContaining(JSON.stringify(contextData, null, 2)));
        });
        it('excludes current agent from handoff schema', async () => {
            const agentA = createHandoffAgent('a', { agentId: 'b', message: 'go to b' });
            const agentB = createFinalAgent('b', 'done');
            const streamSpyA = vi.spyOn(agentA, 'stream');
            const streamSpyB = vi.spyOn(agentB, 'stream');
            const swarm = new Swarm({
                nodes: [agentA, agentB],
                start: 'a',
            });
            await swarm.invoke('start');
            // Agent A's handoff schema allows B but rejects A
            const schemaA = streamSpyA.mock.calls[0][1].structuredOutputSchema;
            expect(schemaA.parse({ agentId: 'b', message: 'ok' })).toStrictEqual({ agentId: 'b', message: 'ok' });
            expect(() => schemaA.parse({ agentId: 'a', message: 'ok' })).toThrow();
            // Agent B's handoff schema allows A but rejects B
            const schemaB = streamSpyB.mock.calls[0][1].structuredOutputSchema;
            expect(schemaB.parse({ agentId: 'a', message: 'ok' })).toStrictEqual({ agentId: 'a', message: 'ok' });
            expect(() => schemaB.parse({ agentId: 'b', message: 'ok' })).toThrow();
        });
        it('throws when maxSteps is exceeded', async () => {
            const swarm = new Swarm({
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'to b' }), createFinalAgent('b', 'done')],
                start: 'a',
                maxSteps: 1,
            });
            await expect(swarm.invoke('start')).rejects.toThrow('swarm reached step limit');
        });
        it('does not throw when swarm completes normally using exactly maxSteps', async () => {
            const swarm = new Swarm({
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'to b' }), createFinalAgent('b', 'done by b')],
                start: 'a',
                maxSteps: 2,
            });
            const result = await swarm.invoke('start');
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'b']);
        });
        it('throws when a node exceeds nodeTimeout', async () => {
            const swarm = new Swarm({
                nodes: [{ agent: createCancellableAgent('slow', 100) }],
                start: 'slow',
                nodeTimeout: 20,
            });
            await expect(swarm.invoke('go')).rejects.toThrow(/node_timeout=<20>, node_id=<slow>/);
        });
        it('applies per-node timeout over nodeTimeout', async () => {
            const swarm = new Swarm({
                nodes: [{ agent: createCancellableAgent('slow', 100), timeout: 15 }],
                start: 'slow',
                nodeTimeout: 10_000,
            });
            await expect(swarm.invoke('go')).rejects.toThrow(/node_timeout=<15>, node_id=<slow>/);
        });
        it('does not throw when nodeTimeout is Infinity', async () => {
            const swarm = new Swarm({
                nodes: [{ agent: createCancellableAgent('a', 20) }],
                start: 'a',
                nodeTimeout: Infinity,
            });
            const result = await swarm.invoke('go');
            expect(result.status).toBe(Status.COMPLETED);
        });
        it('per-node timeout of Infinity disables a finite nodeTimeout', async () => {
            const swarm = new Swarm({
                nodes: [{ agent: createCancellableAgent('slow', 30), timeout: Infinity }],
                start: 'slow',
                nodeTimeout: 10,
            });
            const result = await swarm.invoke('go');
            expect(result.status).toBe(Status.COMPLETED);
        });
        it('throws when timeout is exceeded between steps', async () => {
            const swarm = new Swarm({
                nodes: [
                    { agent: createCancellableAgent('a', 30, { agentId: 'b', message: 'to b' }) },
                    { agent: createCancellableAgent('b', 30) },
                ],
                start: 'a',
                timeout: 20,
            });
            await expect(swarm.invoke('go')).rejects.toThrow(/timeout=<20>/);
        });
        it('aborts an in-flight node when the swarm timeout expires mid-step', async () => {
            const swarm = new Swarm({
                nodes: [{ agent: createCancellableAgent('slow', 200) }],
                start: 'slow',
                timeout: 20,
            });
            await expect(swarm.invoke('go')).rejects.toThrow(/timeout=<20>/);
        });
        it('returns cancelled result with custom message when cancel is a string', async () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            swarm.addHook(BeforeNodeCallEvent, (event) => {
                event.cancel = 'agent not ready';
            });
            const { items, result } = await collectGenerator(swarm.stream('go'));
            expect(result.status).toBe(Status.CANCELLED);
            const cancelEvent = items.find((e) => e.type === 'nodeCancelEvent');
            expect(cancelEvent).toEqual(expect.objectContaining({ nodeId: 'a', state: expect.any(MultiAgentState), message: 'agent not ready' }));
        });
        it('returns failed result when agent throws', async () => {
            const model = new MockMessageModel().addTurn(new Error('agent exploded'));
            const agent = new Agent({ model, printer: false, id: 'a', description: 'Agent a' });
            const swarm = new Swarm({
                nodes: [{ agent }],
                start: 'a',
            });
            const result = await swarm.invoke('go');
            expect(result.status).toBe(Status.FAILED);
            expect(result.results).toHaveLength(1);
            expect(result.results[0]).toEqual(expect.objectContaining({ nodeId: 'a', status: Status.FAILED }));
        });
        it('calls initialize only once across invocations', async () => {
            let callCount = 0;
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            swarm.addHook(MultiAgentInitializedEvent, () => {
                callCount++;
            });
            await swarm.invoke('first');
            await swarm.invoke('second');
            expect(callCount).toBe(1);
        });
        it('preserves agent messages and state after execution', async () => {
            const agent = createFinalAgent('a', 'reply');
            const messagesBefore = [...agent.messages];
            const stateBefore = agent.appState.getAll();
            const swarm = new Swarm({
                nodes: [agent],
                start: 'a',
            });
            await swarm.invoke('hello');
            expect(agent.messages).toStrictEqual(messagesBefore);
            expect(agent.appState.getAll()).toStrictEqual(stateBefore);
        });
    });
    describe('stream', () => {
        it('yields lifecycle events in correct order for single agent', async () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'reply')],
                start: 'a',
            });
            const { items, result } = await collectGenerator(swarm.stream('go'));
            const eventTypes = items.map((e) => e.type);
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a']);
            expect(eventTypes).toStrictEqual([
                'beforeMultiAgentInvocationEvent',
                'beforeNodeCallEvent',
                // nodeStreamUpdateEvents from agent execution
                ...eventTypes.filter((t) => t === 'nodeStreamUpdateEvent'),
                'nodeResultEvent',
                'afterNodeCallEvent',
                'afterMultiAgentInvocationEvent',
                'multiAgentResultEvent',
            ]);
        });
        it('yields handoff event between agents', async () => {
            const swarm = new Swarm({
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'go' }), createFinalAgent('b', 'done')],
                start: 'a',
            });
            const { items } = await collectGenerator(swarm.stream('start'));
            const handoffEvents = items.filter((e) => e.type === 'multiAgentHandoffEvent');
            expect(handoffEvents).toHaveLength(1);
            expect(handoffEvents[0]).toEqual(expect.objectContaining({
                type: 'multiAgentHandoffEvent',
                source: 'a',
                targets: ['b'],
                state: expect.any(MultiAgentState),
            }));
        });
        it('returns cancelled result with default message when cancel is true', async () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            swarm.addHook(BeforeNodeCallEvent, (event) => {
                event.cancel = true;
            });
            const { items, result } = await collectGenerator(swarm.stream('go'));
            expect(result.status).toBe(Status.CANCELLED);
            expect(result.results).toHaveLength(1);
            expect(result.results[0]).toEqual(expect.objectContaining({ nodeId: 'a', status: Status.CANCELLED, duration: 0 }));
            const cancelEvent = items.find((e) => e.type === 'nodeCancelEvent');
            expect(cancelEvent).toEqual(expect.objectContaining({ nodeId: 'a', state: expect.any(MultiAgentState), message: 'node cancelled by hook' }));
        });
        it('returns cancelled result with custom message when cancel is a string', async () => {
            const swarm = new Swarm({
                nodes: [createFinalAgent('a', 'hi')],
                start: 'a',
            });
            swarm.addHook(BeforeNodeCallEvent, (event) => {
                event.cancel = 'agent not ready';
            });
            const { items, result } = await collectGenerator(swarm.stream('go'));
            expect(result.status).toBe(Status.CANCELLED);
            const cancelEvent = items.find((e) => e.type === 'nodeCancelEvent');
            expect(cancelEvent).toEqual(expect.objectContaining({ nodeId: 'a', state: expect.any(MultiAgentState), message: 'agent not ready' }));
        });
    });
    describe('resume with session manager', () => {
        function makeResumeSwarm(storage, options = {}) {
            const sessionManager = new SessionManager({
                sessionId: 'test-session',
                storage: { snapshot: storage },
            });
            const swarm = new Swarm({
                id: 'my-swarm',
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'go to b' }), createFinalAgent('b', 'done by b')],
                start: 'a',
                plugins: [sessionManager],
                ...options,
            });
            return swarm;
        }
        it('resumes from the pending handoff target after a crash (A→B stopped, resumes at B)', async () => {
            const storage = new MockSnapshotStorage();
            const swarm1 = makeResumeSwarm(storage, { maxSteps: 1 });
            await expect(swarm1.invoke('start')).rejects.toThrow('swarm reached step limit');
            const swarm2 = makeResumeSwarm(storage);
            const result = await swarm2.invoke('start');
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'b']);
        });
        it('starts fresh when the previous run completed normally (no pending handoff)', async () => {
            const storage = new MockSnapshotStorage();
            const sessionManager1 = new SessionManager({ sessionId: 'test-session', storage: { snapshot: storage } });
            const swarm1 = new Swarm({
                id: 'my-swarm',
                nodes: [createFinalAgent('a', 'all done'), createFinalAgent('b', 'done by b')],
                start: 'a',
                plugins: [sessionManager1],
            });
            const result1 = await swarm1.invoke('start');
            expect(result1.status).toBe(Status.COMPLETED);
            expect(result1.results.map((r) => r.nodeId)).toStrictEqual(['a']);
            const result2 = await swarm1.invoke('start');
            expect(result2.status).toBe(Status.COMPLETED);
            expect(result2.results.map((r) => r.nodeId)).toStrictEqual(['a']);
        });
        it('carries forward steps count from the previous invocation', async () => {
            const storage = new MockSnapshotStorage();
            const swarm1 = makeResumeSwarm(storage, { maxSteps: 1 });
            await expect(swarm1.invoke('start')).rejects.toThrow('swarm reached step limit');
            const swarm2 = makeResumeSwarm(storage, { maxSteps: 2 });
            const result = await swarm2.invoke('start');
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'b']);
        });
        it('passes the last handoff context to the resumed node', async () => {
            const storage = new MockSnapshotStorage();
            const handoffContext = { research: 'quantum computing basics' };
            const sessionManager1 = new SessionManager({ sessionId: 'test-session', storage: { snapshot: storage } });
            const swarm1 = new Swarm({
                id: 'my-swarm',
                nodes: [
                    createHandoffAgent('a', { agentId: 'b', message: 'write this up', context: handoffContext }),
                    createFinalAgent('b', 'done'),
                ],
                start: 'a',
                maxSteps: 1,
                plugins: [sessionManager1],
            });
            await expect(swarm1.invoke('start')).rejects.toThrow('swarm reached step limit');
            const sessionManager2 = new SessionManager({ sessionId: 'test-session', storage: { snapshot: storage } });
            const agentB = createFinalAgent('b', 'done');
            const streamSpy = vi.spyOn(agentB, 'stream');
            const swarm2 = new Swarm({
                id: 'my-swarm',
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'write this up', context: handoffContext }), agentB],
                start: 'a',
                plugins: [sessionManager2],
            });
            await swarm2.invoke('start');
            expect(streamSpy).toHaveBeenCalled();
            const args = streamSpy.mock.calls[0][0];
            const texts = args.map((b) => b.text);
            expect(texts).toContainEqual('write this up');
            expect(texts).toContainEqual(expect.stringContaining(JSON.stringify(handoffContext, null, 2)));
        });
        it('starts fresh when the resume target agent was removed from the swarm', async () => {
            const storage = new MockSnapshotStorage();
            // First invocation: A hands off to B, maxSteps=1 stops
            const sessionManager1 = new SessionManager({ sessionId: 'test-session', storage: { snapshot: storage } });
            const swarm1 = new Swarm({
                id: 'my-swarm',
                nodes: [createHandoffAgent('a', { agentId: 'b', message: 'go to b' }), createFinalAgent('b', 'done by b')],
                start: 'a',
                maxSteps: 1,
                plugins: [sessionManager1],
            });
            await expect(swarm1.invoke('start')).rejects.toThrow('swarm reached step limit');
            // Second invocation: swarm reconfigured — B removed, C added
            const sessionManager2 = new SessionManager({ sessionId: 'test-session', storage: { snapshot: storage } });
            const swarm2 = new Swarm({
                id: 'my-swarm',
                nodes: [createFinalAgent('a', 'fresh start'), createFinalAgent('c', 'done by c')],
                start: 'a',
                plugins: [sessionManager2],
            });
            const result = await swarm2.invoke('start');
            // B no longer exists, so _findResumeNode falls back to start node A
            expect(result.status).toBe(Status.COMPLETED);
            expect(result.results.map((r) => r.nodeId)).toStrictEqual(['a', 'a']);
        });
    });
});
//# sourceMappingURL=swarm.test.js.map