adk_graph/

agent.rs

//! GraphAgent - ADK Agent integration for graph workflows
//!
//! Provides a builder pattern similar to LlmAgent and RealtimeAgent.

use crate::checkpoint::Checkpointer;
use crate::deferred::DeferredNodeConfig;
use crate::edge::{END, Edge, EdgeTarget, START};
use crate::error::{GraphError, Result};
use crate::graph::{CompiledGraph, StateGraph};
use crate::node::{ExecutionConfig, FunctionNode, Node, NodeContext, NodeOutput};
use crate::state::{State, StateSchema};
use crate::stream::{StreamEvent, StreamMode};
use crate::timeout::TimeoutPolicy;
use adk_core::{
    Agent, AgentCapabilities, AgentRelationshipKind, AgentTopology, AgentTopologyMember,
    AgentTopologyRelationship, Content, Event, EventStream, InvocationContext,
};
use async_trait::async_trait;
use serde_json::json;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

/// Type alias for callbacks
pub type BeforeAgentCallback = Arc<
    dyn Fn(Arc<dyn InvocationContext>) -> Pin<Box<dyn Future<Output = adk_core::Result<()>> + Send>>
        + Send
        + Sync,
>;

pub type AfterAgentCallback = Arc<
    dyn Fn(
            Arc<dyn InvocationContext>,
            Event,
        ) -> Pin<Box<dyn Future<Output = adk_core::Result<()>> + Send>>
        + Send
        + Sync,
>;

/// Type alias for input mapper function
pub type InputMapper = Arc<dyn Fn(&dyn InvocationContext) -> State + Send + Sync>;

/// Type alias for output mapper function
pub type OutputMapper = Arc<dyn Fn(&State) -> Vec<Event> + Send + Sync>;

/// GraphAgent wraps a CompiledGraph as an ADK Agent
///
/// # Turns and checkpoints
///
/// Each [`Agent::run`] call is one turn on a thread whose id is the session id.
/// With a checkpointer configured:
///
/// | Thread's latest checkpoint | What the turn does |
/// |----------------------------|--------------------|
/// | None | Runs the graph from its entry nodes. |
/// | Paused (non-empty frontier) | Resumes the paused run, with the turn's input merged into its state. |
/// | Finished (empty frontier) | Runs the graph again from its entry nodes, starting from the previous turn's final state with the turn's input merged on top. |
///
/// The recursion limit applies to each turn separately. [`GraphAgent::invoke`]
/// keeps the direct-graph behaviour, where re-invoking a finished thread returns
/// its final state without running anything.
pub struct GraphAgent {
    name: String,
    description: String,
    graph: Arc<CompiledGraph>,
    /// Map InvocationContext to graph input state
    input_mapper: InputMapper,
    /// Map graph output state to ADK Events
    output_mapper: OutputMapper,
    /// Before agent callback
    before_callback: Option<BeforeAgentCallback>,
    /// After agent callback
    after_callback: Option<AfterAgentCallback>,
}

impl GraphAgent {
    /// Create a new GraphAgent builder
    pub fn builder(name: &str) -> GraphAgentBuilder {
        GraphAgentBuilder::new(name)
    }

    /// Create directly from a compiled graph
    pub fn from_graph(name: &str, graph: CompiledGraph) -> Self {
        Self {
            name: name.to_string(),
            description: String::new(),
            graph: Arc::new(graph),
            input_mapper: Arc::new(default_input_mapper),
            output_mapper: Arc::new(default_output_mapper),
            before_callback: None,
            after_callback: None,
        }
    }

    /// Build a `GraphAgent` from a `WorkflowSchema`.
    ///
    /// Delegates to `schema.build_graph()` to construct the graph from the
    /// workflow schema's action nodes, edges, and conditions.
    #[cfg(feature = "action")]
    pub fn from_workflow_schema(
        name: &str,
        schema: &crate::workflow::WorkflowSchema,
    ) -> Result<Self> {
        schema.build_graph(name)
    }

    /// Get the underlying compiled graph
    pub fn graph(&self) -> &CompiledGraph {
        &self.graph
    }

    /// Execute the graph directly (bypassing Agent trait)
    pub async fn invoke(&self, input: State, config: ExecutionConfig) -> Result<State> {
        self.graph.invoke(input, config).await
    }

    /// Stream execution
    pub fn stream(
        &self,
        input: State,
        config: ExecutionConfig,
        mode: StreamMode,
    ) -> impl futures::Stream<Item = Result<StreamEvent>> + '_ {
        self.graph.stream(input, config, mode)
    }
}

#[async_trait]
impl Agent for GraphAgent {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        &self.description
    }

    fn sub_agents(&self) -> &[Arc<dyn Agent>] {
        &[]
    }

    fn supports_agent_transfer(&self) -> bool {
        false
    }

    fn capabilities(&self) -> AgentCapabilities {
        AgentCapabilities {
            checkpoint_resume: self.graph.has_checkpointer(),
            shared_state: true,
            invocation_metadata: true,
            ..AgentCapabilities::default()
        }
    }

    fn topology(&self) -> Option<AgentTopology> {
        let mut nodes = self.graph.nodes.values().collect::<Vec<_>>();
        nodes.sort_by_key(|node| node.name().to_string());
        let mut members = vec![AgentTopologyMember {
            name: self.name.clone(),
            description: self.description.clone(),
            coordinator: true,
            capabilities: self.capabilities(),
        }];
        members.extend(nodes.into_iter().map(|node| AgentTopologyMember {
            name: node.name().to_string(),
            description: node.description().to_string(),
            coordinator: false,
            capabilities: node.capabilities(),
        }));

        let mut relationships = self
            .graph
            .get_entry_nodes()
            .into_iter()
            .map(|entry| AgentTopologyRelationship {
                from: self.name.clone(),
                to: entry,
                kind: AgentRelationshipKind::Flow,
            })
            .collect::<Vec<_>>();
        for edge in &self.graph.edges {
            match edge {
                Edge::Direct { source, target: EdgeTarget::Node(target) } => {
                    relationships.push(AgentTopologyRelationship {
                        from: source.clone(),
                        to: target.clone(),
                        kind: AgentRelationshipKind::Flow,
                    });
                }
                Edge::Conditional { source, targets, .. } => {
                    relationships.extend(targets.values().filter_map(|target| {
                        target.node_name().map(|target| AgentTopologyRelationship {
                            from: source.clone(),
                            to: target.to_string(),
                            kind: AgentRelationshipKind::Flow,
                        })
                    }));
                }
                Edge::Direct { target: EdgeTarget::End, .. } | Edge::Entry { .. } => {}
            }
        }
        relationships.sort_by(|left, right| (&left.from, &left.to).cmp(&(&right.from, &right.to)));
        relationships.dedup_by(|left, right| left.from == right.from && left.to == right.to);

        Some(AgentTopology {
            root: self.name.clone(),
            coordinator: self.name.clone(),
            members,
            relationships,
        })
    }

    async fn run(&self, ctx: Arc<dyn InvocationContext>) -> adk_core::Result<EventStream> {
        // Call before callback
        if let Some(callback) = &self.before_callback {
            callback(ctx.clone()).await?;
        }

        // Map context to input state
        let input = (self.input_mapper)(ctx.as_ref());

        // Create execution config from context, carrying the invocation so an
        // `AgentNode` inside the graph presents this run's identity, services, and
        // cancellation rather than a synthetic standalone context.
        let config = ExecutionConfig::new(ctx.session_id()).with_parent_context(ctx.clone());

        // Execute graph
        let graph = self.graph.clone();
        let output_mapper = self.output_mapper.clone();
        let after_callback = self.after_callback.clone();
        let ctx_clone = ctx.clone();

        let stream = async_stream::stream! {
            // Each call is one turn on the session's thread: a paused thread
            // resumes, and a finished one runs again from its entry nodes with the
            // previous turn's state carried over.
            let mut executor =
                crate::executor::PregelExecutor::new(&graph, config).restarting_completed_threads();
            match executor.run(input).await {
                Ok(state) => {
                    let events = output_mapper(&state);
                    for event in events {
                        // Call after callback for each event
                        if let Some(callback) = &after_callback
                            && let Err(e) = callback(ctx_clone.clone(), event.clone()).await {
                                yield Err(e);
                                return;
                            }
                        yield Ok(event);
                    }
                }
                Err(GraphError::Interrupted(interrupt)) => {
                    // The `Agent` trait yields events, so an interrupt cannot be
                    // returned as an error without ending the invocation. Emit one
                    // event carrying the structured pause so a caller can read the
                    // node, the payload, and the checkpoint to resume from.
                    let tool_confirmation = crate::interrupt::GraphToolConfirmationPause::from_interrupted_execution(&interrupt);
                    let payload = tool_confirmation.clone().map_or_else(
                        || crate::interrupt::GraphInterruptPayload::new(
                            &interrupt.interrupt,
                            &interrupt.thread_id,
                            &interrupt.checkpoint_id,
                        ),
                        crate::interrupt::GraphInterruptPayload::from_tool_confirmation_pause,
                    );
                    let mut event = Event::new("graph_interrupted");
                    if let Some(pause) = tool_confirmation {
                        event.llm_response.interrupted = true;
                        event.llm_response.turn_complete = true;
                        event.actions.tool_confirmation = Some(pause.request);
                    }
                    event.set_content(
                        Content::new("assistant").with_text(interrupt.interrupt.to_string()),
                    );
                    event.provider_metadata.insert(
                        crate::interrupt::INTERRUPT_METADATA_KEY.to_string(),
                        payload.to_metadata_value(),
                    );
                    yield Ok(event);
                }
                Err(e) => {
                    yield Err(adk_core::AdkError::agent(e.to_string()));
                }
            }
        };

        Ok(Box::pin(stream))
    }
}

/// Default input mapper - extracts content from InvocationContext
fn default_input_mapper(ctx: &dyn InvocationContext) -> State {
    let mut state = State::new();

    // Get user content
    let content = ctx.user_content();
    let text: String = content.parts.iter().filter_map(|p| p.text()).collect::<Vec<_>>().join("\n");

    if !text.is_empty() {
        state.insert("input".to_string(), json!(text));
        state.insert("messages".to_string(), json!([{"role": "user", "content": text}]));
    }

    // Add session ID
    state.insert("session_id".to_string(), json!(ctx.session_id()));

    state
}

/// Default output mapper - creates events from state
fn default_output_mapper(state: &State) -> Vec<Event> {
    let mut events = Vec::new();

    // Try to get output from common fields
    let output_text = state
        .get("output")
        .and_then(|v| v.as_str())
        .or_else(|| state.get("result").and_then(|v| v.as_str()))
        .or_else(|| {
            state
                .get("messages")
                .and_then(|v| v.as_array())
                .and_then(|arr| arr.last())
                .and_then(|msg| msg.get("content"))
                .and_then(|c| c.as_str())
        });

    let text = if let Some(text) = output_text {
        text.to_string()
    } else {
        // Return the full state as JSON
        serde_json::to_string_pretty(state).unwrap_or_default()
    };

    let mut event = Event::new("graph_output");
    event.set_content(Content::new("assistant").with_text(&text));
    events.push(event);

    events
}

/// Builder for GraphAgent
pub struct GraphAgentBuilder {
    name: String,
    description: String,
    schema: StateSchema,
    nodes: Vec<Arc<dyn Node>>,
    edges: Vec<Edge>,
    checkpointer: Option<Arc<dyn Checkpointer>>,
    interrupt_before: Vec<String>,
    interrupt_after: Vec<String>,
    recursion_limit: usize,
    max_concurrency: Option<usize>,
    input_mapper: Option<InputMapper>,
    output_mapper: Option<OutputMapper>,
    before_callback: Option<BeforeAgentCallback>,
    after_callback: Option<AfterAgentCallback>,
    timeout_policies: HashMap<String, TimeoutPolicy>,
    default_timeout: Option<TimeoutPolicy>,
    deferred_configs: HashMap<String, DeferredNodeConfig>,
    #[cfg(feature = "node-cache")]
    cache_policies: HashMap<String, crate::cache::NodeCachePolicy>,
}

impl GraphAgentBuilder {
    /// Create a new builder
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            description: String::new(),
            schema: StateSchema::simple(&["input", "output", "messages"]),
            nodes: vec![],
            edges: vec![],
            checkpointer: None,
            interrupt_before: vec![],
            interrupt_after: vec![],
            recursion_limit: 100,
            max_concurrency: None,
            input_mapper: None,
            output_mapper: None,
            before_callback: None,
            after_callback: None,
            timeout_policies: HashMap::new(),
            default_timeout: None,
            deferred_configs: HashMap::new(),
            #[cfg(feature = "node-cache")]
            cache_policies: HashMap::new(),
        }
    }

    /// Set description
    pub fn description(mut self, desc: &str) -> Self {
        self.description = desc.to_string();
        self
    }

    /// Set state schema
    pub fn state_schema(mut self, schema: StateSchema) -> Self {
        self.schema = schema;
        self
    }

    /// Add channels to state schema
    pub fn channels(mut self, channels: &[&str]) -> Self {
        self.schema = StateSchema::simple(channels);
        self
    }

    /// Add a node
    pub fn node<N: Node + 'static>(mut self, node: N) -> Self {
        self.nodes.push(Arc::new(node));
        self
    }

    /// Add a function as a node
    pub fn node_fn<F, Fut>(mut self, name: &str, func: F) -> Self
    where
        F: Fn(NodeContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<NodeOutput>> + Send + 'static,
    {
        self.nodes.push(Arc::new(FunctionNode::new(name, func)));
        self
    }

    /// Add a direct edge
    pub fn edge(mut self, source: &str, target: &str) -> Self {
        let target =
            if target == END { EdgeTarget::End } else { EdgeTarget::Node(target.to_string()) };

        if source == START {
            let entry_idx = self.edges.iter().position(|e| matches!(e, Edge::Entry { .. }));
            match entry_idx {
                Some(idx) => {
                    if let Edge::Entry { targets } = &mut self.edges[idx]
                        && let EdgeTarget::Node(node) = &target
                        && !targets.contains(node)
                    {
                        targets.push(node.clone());
                    }
                }
                None => {
                    if let EdgeTarget::Node(node) = target {
                        self.edges.push(Edge::Entry { targets: vec![node] });
                    }
                }
            }
        } else {
            self.edges.push(Edge::Direct { source: source.to_string(), target });
        }

        self
    }

    /// Add a conditional edge
    pub fn conditional_edge<F, I>(mut self, source: &str, router: F, targets: I) -> Self
    where
        F: Fn(&State) -> String + Send + Sync + 'static,
        I: IntoIterator<Item = (&'static str, &'static str)>,
    {
        let targets_map: HashMap<String, EdgeTarget> = targets
            .into_iter()
            .map(|(k, v)| {
                let target =
                    if v == END { EdgeTarget::End } else { EdgeTarget::Node(v.to_string()) };
                (k.to_string(), target)
            })
            .collect();

        self.edges.push(Edge::Conditional {
            source: source.to_string(),
            router: Arc::new(router),
            targets: targets_map,
        });

        self
    }

    /// Set checkpointer
    pub fn checkpointer<C: Checkpointer + 'static>(mut self, checkpointer: C) -> Self {
        self.checkpointer = Some(Arc::new(checkpointer));
        self
    }

    /// Set checkpointer with Arc
    pub fn checkpointer_arc(mut self, checkpointer: Arc<dyn Checkpointer>) -> Self {
        self.checkpointer = Some(checkpointer);
        self
    }

    /// Set nodes to interrupt before
    pub fn interrupt_before(mut self, nodes: &[&str]) -> Self {
        self.interrupt_before = nodes.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Set nodes to interrupt after
    pub fn interrupt_after(mut self, nodes: &[&str]) -> Self {
        self.interrupt_after = nodes.iter().map(|s| s.to_string()).collect();
        self
    }

    /// Set recursion limit
    /// Cap how many nodes execute concurrently within one super-step.
    ///
    /// See [`CompiledGraph::with_max_concurrency`](crate::graph::CompiledGraph::with_max_concurrency).
    pub fn max_concurrency(mut self, limit: usize) -> Self {
        self.max_concurrency = Some(limit.max(1));
        self
    }

    pub fn recursion_limit(mut self, limit: usize) -> Self {
        self.recursion_limit = limit;
        self
    }

    /// Set a timeout policy for a specific node.
    ///
    /// The policy is applied when the named node executes, enforcing
    /// wall-clock and/or idle timeouts with the configured recovery action.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use std::time::Duration;
    /// use adk_graph::timeout::{TimeoutPolicy, OnTimeout};
    ///
    /// let agent = GraphAgent::builder("my_graph")
    ///     .node_timeout("slow_node", TimeoutPolicy {
    ///         run_timeout: Some(Duration::from_secs(10)),
    ///         idle_timeout: None,
    ///         on_timeout: OnTimeout::Fail,
    ///     })
    ///     .build()?;
    /// ```
    pub fn node_timeout(mut self, node_name: &str, policy: TimeoutPolicy) -> Self {
        self.timeout_policies.insert(node_name.to_string(), policy);
        self
    }

    /// Set a default timeout policy applied to all nodes without an explicit override.
    ///
    /// Nodes that have a per-node policy set via [`node_timeout`](Self::node_timeout)
    /// will use their specific policy instead of this default.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use std::time::Duration;
    /// use adk_graph::timeout::{TimeoutPolicy, OnTimeout};
    ///
    /// let agent = GraphAgent::builder("my_graph")
    ///     .default_timeout(TimeoutPolicy {
    ///         run_timeout: Some(Duration::from_secs(30)),
    ///         idle_timeout: Some(Duration::from_secs(5)),
    ///         on_timeout: OnTimeout::Skip,
    ///     })
    ///     .build()?;
    /// ```
    pub fn default_timeout(mut self, policy: TimeoutPolicy) -> Self {
        self.default_timeout = Some(policy);
        self
    }

    /// Add a deferred (fan-in barrier) node to the graph.
    ///
    /// A deferred node waits for all upstream parallel paths to complete before
    /// executing. The provided function is wrapped as a [`FunctionNode`] and the
    /// [`DeferredNodeConfig`] controls how upstream outputs are merged and how
    /// long the node waits for all paths.
    ///
    /// # Arguments
    ///
    /// * `name` - The name of the deferred node.
    /// * `func` - The async function to execute once all upstream paths complete.
    /// * `config` - Configuration controlling merge strategy and fan-in timeout.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use std::time::Duration;
    /// use adk_graph::deferred::{DeferredNodeConfig, MergeStrategy};
    /// use adk_graph::node::NodeOutput;
    ///
    /// let agent = GraphAgent::builder("scatter_gather")
    ///     .deferred_node("aggregator", |_ctx| async {
    ///         Ok(NodeOutput::new().with_update("status", serde_json::json!("merged")))
    ///     }, DeferredNodeConfig {
    ///         merge_strategy: MergeStrategy::Collect,
    ///         fan_in_timeout: Some(Duration::from_secs(30)),
    ///     })
    ///     .build()?;
    /// ```
    /// Configure fan-in for a node already added with [`node`](Self::node).
    ///
    /// [`deferred_node`](Self::deferred_node) both adds and configures a node, so
    /// a custom `Node` added through `node` had no way to set a merge strategy or
    /// a fan-in timeout.
    ///
    /// A node reached by more than one unconditional edge is deferred
    /// automatically; this overrides that default.
    pub fn mark_deferred(mut self, name: &str, config: DeferredNodeConfig) -> Self {
        self.deferred_configs.insert(name.to_string(), config);
        self
    }

    pub fn deferred_node<F, Fut>(mut self, name: &str, func: F, config: DeferredNodeConfig) -> Self
    where
        F: Fn(NodeContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<NodeOutput>> + Send + 'static,
    {
        self.nodes.push(Arc::new(FunctionNode::new(name, func)));
        self.deferred_configs.insert(name.to_string(), config);
        self
    }

    /// Set a cache policy for a specific node.
    ///
    /// When a node has a cache policy, its execution results are cached keyed
    /// by a blake3 hash of the node name and input state. Subsequent executions
    /// with identical inputs return the cached result without re-executing the
    /// node.
    ///
    /// # Arguments
    ///
    /// * `name` — the name of the node to cache
    /// * `policy` — the cache policy specifying backend and TTL
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use std::time::Duration;
    /// use adk_graph::cache::{CacheBackend, NodeCachePolicy};
    ///
    /// let agent = GraphAgent::builder("cached_graph")
    ///     .node_cache("expensive_node", NodeCachePolicy {
    ///         backend: CacheBackend::InMemory { max_entries: 128 },
    ///         ttl: Some(Duration::from_secs(300)),
    ///     })
    ///     .build()?;
    /// ```
    #[cfg(feature = "node-cache")]
    pub fn node_cache(mut self, name: &str, policy: crate::cache::NodeCachePolicy) -> Self {
        self.cache_policies.insert(name.to_string(), policy);
        self
    }

    /// Set custom input mapper
    pub fn input_mapper<F>(mut self, mapper: F) -> Self
    where
        F: Fn(&dyn InvocationContext) -> State + Send + Sync + 'static,
    {
        self.input_mapper = Some(Arc::new(mapper));
        self
    }

    /// Set custom output mapper
    pub fn output_mapper<F>(mut self, mapper: F) -> Self
    where
        F: Fn(&State) -> Vec<Event> + Send + Sync + 'static,
    {
        self.output_mapper = Some(Arc::new(mapper));
        self
    }

    /// Set before agent callback
    pub fn before_agent_callback<F, Fut>(mut self, callback: F) -> Self
    where
        F: Fn(Arc<dyn InvocationContext>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = adk_core::Result<()>> + Send + 'static,
    {
        self.before_callback = Some(Arc::new(move |ctx| Box::pin(callback(ctx))));
        self
    }

    /// Set after agent callback
    ///
    /// Note: The callback receives a cloned Event to avoid lifetime issues.
    pub fn after_agent_callback<F, Fut>(mut self, callback: F) -> Self
    where
        F: Fn(Arc<dyn InvocationContext>, Event) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = adk_core::Result<()>> + Send + 'static,
    {
        self.after_callback = Some(Arc::new(move |ctx, event| {
            let event_clone = event.clone();
            Box::pin(callback(ctx, event_clone))
        }));
        self
    }

    /// Add an action node to the graph.
    ///
    /// Wraps the `ActionNodeConfig` in an `ActionNodeExecutor` and registers it
    /// as a node. If the config is a `SwitchNodeConfig`, conditional edges are
    /// also auto-registered from the switch conditions.
    #[cfg(feature = "action")]
    pub fn action_node(mut self, config: adk_action::ActionNodeConfig) -> Self {
        use crate::action::ActionNodeExecutor;

        // If this is a Switch node, register conditional edges
        if let adk_action::ActionNodeConfig::Switch(ref switch_config) = config {
            let conditions = switch_config.conditions.clone();
            let eval_mode = switch_config.evaluation_mode.clone();
            let default_branch = switch_config.default_branch.clone();
            let source = config.standard().id.clone();

            let mut targets_map: HashMap<String, EdgeTarget> = HashMap::new();
            for condition in &conditions {
                targets_map.insert(
                    condition.output_port.clone(),
                    EdgeTarget::Node(condition.output_port.clone()),
                );
            }
            if let Some(ref default) = default_branch {
                let target = if default == END {
                    EdgeTarget::End
                } else {
                    EdgeTarget::Node(default.clone())
                };
                targets_map.insert(default.clone(), target);
            }
            targets_map.insert(END.to_string(), EdgeTarget::End);

            let router = Arc::new(move |state: &State| -> String {
                match crate::action::switch::evaluate_switch_conditions(
                    &conditions,
                    state,
                    &eval_mode,
                    default_branch.as_deref(),
                ) {
                    Ok(ports) => ports.into_iter().next().unwrap_or_else(|| END.to_string()),
                    Err(_) => END.to_string(),
                }
            });

            self.edges.push(Edge::Conditional { source, router, targets: targets_map });
        }

        let executor = ActionNodeExecutor::new(config);
        self.nodes.push(Arc::new(executor));
        self
    }

    /// Build the GraphAgent
    pub fn build(self) -> Result<GraphAgent> {
        // Build the graph
        let mut graph = StateGraph::new(self.schema);

        // Add nodes
        for node in self.nodes {
            graph.nodes.insert(node.name().to_string(), node);
        }

        // Add edges
        graph.edges = self.edges;

        // Compile
        let mut compiled = graph.compile()?;

        // Configure
        if let Some(cp) = self.checkpointer {
            compiled.checkpointer = Some(cp);
        }
        compiled.interrupt_before = self.interrupt_before.into_iter().collect();
        compiled.interrupt_after = self.interrupt_after.into_iter().collect();
        compiled.recursion_limit = self.recursion_limit;
        compiled.max_concurrency = self.max_concurrency;
        compiled.timeout_policies = self.timeout_policies;
        compiled.default_timeout = self.default_timeout;
        compiled.deferred_configs = self.deferred_configs;

        #[cfg(feature = "node-cache")]
        {
            compiled.node_caches = self
                .cache_policies
                .iter()
                .map(|(name, policy)| {
                    (name.clone(), Arc::new(crate::cache::NodeCache::from_policy(policy)))
                })
                .collect();
            compiled.cache_policies = self.cache_policies;
        }

        Ok(GraphAgent {
            name: self.name,
            description: self.description,
            graph: Arc::new(compiled),
            input_mapper: self.input_mapper.unwrap_or(Arc::new(default_input_mapper)),
            output_mapper: self.output_mapper.unwrap_or(Arc::new(default_output_mapper)),
            before_callback: self.before_callback,
            after_callback: self.after_callback,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[tokio::test]
    async fn test_graph_agent_builder() {
        let agent = GraphAgent::builder("test")
            .description("Test agent")
            .channels(&["value"])
            .node_fn("set", |_ctx| async { Ok(NodeOutput::new().with_update("value", json!(42))) })
            .edge(START, "set")
            .edge("set", END)
            .build()
            .unwrap();

        assert_eq!(agent.name(), "test");
        assert_eq!(agent.description(), "Test agent");

        // Test direct invocation
        let result = agent.invoke(State::new(), ExecutionConfig::new("test")).await.unwrap();

        assert_eq!(result.get("value"), Some(&json!(42)));

        let topology = agent.topology().expect("graph topology");
        assert_eq!(topology.root, "test");
        assert_eq!(topology.coordinator, "test");
        assert_eq!(topology.members.len(), 2);
        assert_eq!(topology.relationships.len(), 1);
        assert_eq!(topology.relationships[0].from, "test");
        assert_eq!(topology.relationships[0].to, "set");
        assert_eq!(topology.relationships[0].kind, AgentRelationshipKind::Flow);
    }
}