import { Agent } from '../agent/agent.js';
import { dropStaleInterruptedResult } from './multiagent.js';
import { NodeStreamUpdateEvent, NodeResultEvent } from './events.js';
import { NodeResult, Status } from './state.js';
import { logger } from '../logging/logger.js';
import { normalizeError } from '../errors.js';
import { omitUndefined } from '../types/json.js';
/**
 * Abstract base class for all multi-agent orchestration nodes.
 *
 * Uses the template method pattern: {@link stream} handles orchestration
 * boilerplate (duration measurement, status tracking, error capture) and
 * delegates to {@link handle} for node-specific execution logic.
 */
export class Node {
    type = 'node';
    /** Unique identifier for this node within the orchestration. */
    id;
    /** Per-node configuration. */
    config;
    /**
     * @param id - Unique identifier for this node within the orchestration
     * @param config - Per-node configuration
     */
    constructor(id, config) {
        this.id = id;
        this.config = config;
    }
    /**
     * Execute the node. Handles duration measurement, error capture,
     * and delegates to handle() for node-specific logic.
     *
     * @param input - Input to pass to the node (string or content blocks)
     * @param state - The current multi-agent state
     * @param options - Per-invocation options from the orchestrator
     * @returns Async generator yielding streaming events and returning a NodeResult
     */
    async *stream(input, state, options) {
        const nodeState = state.node(this.id);
        // Resuming from INTERRUPTED: drop the stale result so the fresh one replaces it.
        dropStaleInterruptedResult(this.id, nodeState, state);
        nodeState.status = Status.EXECUTING;
        nodeState.startTime = Date.now();
        // Resolve invocationState once — the same reference is threaded into handle()
        // and into NodeResultEvent so callbacks see one object for the whole node run.
        const invocationState = options?.invocationState ?? {};
        const resolvedOptions = { ...options, invocationState };
        let result;
        try {
            const update = yield* this.handle(input, state, resolvedOptions);
            const defaultStatus = update.interrupts && update.interrupts.length > 0 ? Status.INTERRUPTED : Status.COMPLETED;
            result = new NodeResult({
                nodeId: this.id,
                status: defaultStatus,
                duration: Date.now() - nodeState.startTime,
                content: [],
                ...update,
            });
        }
        catch (error) {
            // Orchestrator cancellation (short-circuit or external) maps thrown errors to
            // CANCELLED — node was stopped, not broken.
            const status = options?.cancelSignal?.aborted ? Status.CANCELLED : Status.FAILED;
            result = new NodeResult({
                nodeId: this.id,
                status,
                duration: Date.now() - nodeState.startTime,
                error: normalizeError(error),
            });
            if (status === Status.FAILED) {
                logger.warn(`node_id=<${this.id}>, error=<${result.error?.message}> | node execution failed`);
            }
        }
        finally {
            nodeState.status = result.status;
            nodeState.results.push(result);
            nodeState.interrupts = result.interrupts ?? [];
            // Clear the stored snapshot on non-INTERRUPTED terminal states; `handle()`
            // repopulates it above if this run itself interrupted.
            if (result.status !== Status.INTERRUPTED) {
                delete nodeState.interruptedSnapshot;
            }
        }
        yield new NodeResultEvent({
            nodeId: this.id,
            nodeType: this.type,
            state,
            result,
            invocationState,
        });
        return result;
    }
}
/**
 * Node that wraps an {@link InvokableAgent} instance for multi-agent orchestration.
 *
 * By default, when the wrapped agent is an {@link Agent} instance, its internal
 * state is snapshot/restored around each execution so it remains unchanged
 * after the node completes. Pass `preserveContext: true` to opt out and let the
 * wrapped agent accumulate state across node executions.
 */
export class AgentNode extends Node {
    type = 'agentNode';
    _agent;
    /**
     * Per-node wall-clock ceiling in milliseconds. When set, overrides the orchestrator's
     * `nodeTimeout` for this node. Undefined means "fall back to the orchestrator's setting."
     * See {@link AgentNodeOptions.timeout}.
     */
    timeout;
    /**
     * Whether the wrapped agent retains state across node executions.
     * See {@link AgentNodeOptions.preserveContext}.
     */
    preserveContext;
    constructor(options) {
        const { agent, timeout, preserveContext, ...config } = options;
        super(agent.id, {
            ...config,
            ...(agent.description !== undefined && { description: agent.description }),
        });
        this._agent = agent;
        if (timeout !== undefined) {
            if (timeout < 1) {
                throw new Error(`timeout=<${timeout}>, node_id=<${agent.id}> | must be at least 1`);
            }
            this.timeout = timeout;
        }
        if (preserveContext && !(agent instanceof Agent)) {
            throw new Error(`node_id=<${agent.id}> | preserveContext=true requires an Agent instance; non-Agent InvokableAgents cannot be snapshotted`);
        }
        this.preserveContext = preserveContext ?? false;
    }
    get agent() {
        return this._agent;
    }
    /**
     * Executes the wrapped agent, yielding each agent streaming event
     * wrapped in a {@link NodeStreamUpdateEvent}.
     *
     * @param input - Input to pass to the agent
     * @param state - The current multi-agent state
     * @param options - Per-invocation options from the orchestrator
     * @returns Async generator yielding streaming events and returning the agent's content blocks
     */
    async *handle(input, state, options) {
        // Resolve once per handle() call — Node.stream() normally supplies this;
        // handle() is public API, so direct callers get per-call state.
        const invocationState = options?.invocationState ?? {};
        // Only Agent instances support snapshot/restore for state isolation.
        // When `preserveContext` is set, skip the snapshot/restore cycle so the agent
        // accumulates state across node executions.
        const isAgent = this._agent instanceof Agent;
        const preRunSnapshot = !this.preserveContext && isAgent ? this._agent.takeSnapshot({ preset: 'session' }) : undefined;
        // Rehydrate agent state from a prior INTERRUPTED run (messages + interrupt state).
        // Independent of `preserveContext`: a paused run always resumes from where it left off.
        const nodeState = state.node(this.id);
        if (isAgent && nodeState?.interruptedSnapshot) {
            this._agent.loadSnapshot(nodeState.interruptedSnapshot);
        }
        try {
            const invokeOptions = {
                ...(options?.structuredOutputSchema && { structuredOutputSchema: options.structuredOutputSchema }),
                ...(options?.cancelSignal && { cancelSignal: options.cancelSignal }),
                invocationState,
            };
            const gen = this._agent.stream(input, invokeOptions);
            let next = await gen.next();
            while (!next.done) {
                yield new NodeStreamUpdateEvent({
                    nodeId: this.id,
                    nodeType: this.type,
                    state,
                    inner: isAgent
                        ? { source: 'agent', event: next.value }
                        : { source: 'custom', event: next.value },
                    invocationState,
                });
                next = await gen.next();
            }
            const agentResult = next.value;
            const interrupted = agentResult.stopReason === 'interrupt' && agentResult.interrupts && agentResult.interrupts.length > 0;
            // Capture post-interrupt state for the next resume cycle. Only Agent instances
            // are snapshottable.
            if (interrupted && isAgent && nodeState) {
                nodeState.interruptedSnapshot = this._agent.takeSnapshot({ preset: 'session' });
            }
            return omitUndefined({
                content: agentResult.lastMessage.content,
                structuredOutput: 'structuredOutput' in agentResult ? agentResult.structuredOutput : undefined,
                usage: agentResult.metrics?.accumulatedUsage,
                interrupts: interrupted ? agentResult.interrupts : undefined,
            });
        }
        finally {
            // Restore pre-run state — keeps the agent observably unchanged across runs.
            if (preRunSnapshot) {
                ;
                this._agent.loadSnapshot(preRunSnapshot);
            }
        }
    }
}
/**
 * Node that wraps a multi-agent orchestrator (e.g. Graph) for nested composition.
 *
 * Inner {@link NodeStreamUpdateEvent}s pass through to preserve the original
 * node's identity. All other events are wrapped in a new {@link NodeStreamUpdateEvent}
 * tagged with this node's identity.
 */
export class MultiAgentNode extends Node {
    type = 'multiAgentNode';
    _orchestrator;
    constructor(options) {
        const { orchestrator, ...config } = options;
        super(orchestrator.id, config);
        this._orchestrator = orchestrator;
    }
    get orchestrator() {
        return this._orchestrator;
    }
    /**
     * Executes the wrapped orchestrator. Inner {@link NodeStreamUpdateEvent}s
     * pass through as-is; all other events are wrapped in a new
     * {@link NodeStreamUpdateEvent} tagged with this node's identity.
     *
     * @param input - Input to pass to the orchestrator
     * @param state - The current multi-agent state
     * @param options - Per-invocation options. `invocationState` is forwarded to the
     *   nested orchestrator; `structuredOutputSchema` is not applicable here.
     * @returns Async generator yielding streaming events and returning the orchestrator's content
     */
    async *handle(input, state, options) {
        // Resolve once per handle() call — Node.stream() normally supplies this;
        // handle() is public API, so direct callers get per-call state.
        const invocationState = options?.invocationState ?? {};
        const gen = this._orchestrator.stream(input, {
            invocationState,
            ...(options?.cancelSignal && { cancelSignal: options.cancelSignal }),
        });
        let next = await gen.next();
        while (!next.done) {
            const event = next.value;
            if (event.type === 'nodeStreamUpdateEvent') {
                yield event;
            }
            else {
                yield new NodeStreamUpdateEvent({
                    nodeId: this.id,
                    nodeType: this.type,
                    state,
                    inner: { source: 'multiAgent', event },
                    invocationState,
                });
            }
            next = await gen.next();
        }
        const innerResult = next.value;
        const interrupted = innerResult.interrupts && innerResult.interrupts.length > 0;
        return omitUndefined({
            content: innerResult.content,
            usage: innerResult.usage,
            status: innerResult.status !== Status.COMPLETED ? innerResult.status : undefined,
            error: innerResult.error,
            interrupts: interrupted ? innerResult.interrupts : undefined,
        });
    }
}
//# sourceMappingURL=nodes.js.map