用 Rust 从零构建 AI Agent（四）：状态机与多步任务

#[async_trait]
trait Node<S>: Send + Sync + 'static
where
    S: Send + Sync + 'static,
{
    fn name(&self) -> &'static str;
    async fn run(&self, state: &mut S) -> Result<NextStep, GraphError>;
}

enum NextStep {
    Goto(&'static str),
    Halt,
    Interrupt { reason: String },
}

struct Graph<S> {
    nodes: HashMap<&'static str, Box<dyn Node<S>>>,
    entry: &'static str,
    max_steps: usize,
}

impl<S: Send + Sync + 'static> Graph<S> {
    async fn run_with_checkpoint<C: Checkpointer<S>>(
        &self,
        run_id: &str,
        initial: S,
        ckpt: &C,
    ) -> Result<Outcome<S>, GraphError> {
        let (mut state, mut current) = match ckpt.load(run_id).await? {
            Some(ck) => (ck.state, ck.next_node),
            None => (initial, self.entry),
        };

        for _ in 0..self.max_steps {
            let node = self.nodes.get(current)
                .ok_or(GraphError::UnknownNode(current))?;
            let next = node.run(&mut state).await?;

            match next {
                NextStep::Goto(name) => {
                    ckpt.save(run_id, name, &state).await?;
                    current = name;
                }
                NextStep::Halt => {
                    ckpt.delete(run_id).await?;
                    return Ok(Outcome::Success(state));
                }
                NextStep::Interrupt { reason } => {
                    ckpt.save(run_id, current, &state).await?;
                    return Ok(Outcome::Interrupted { state, reason });
                }
            }
        }

        Err(GraphError::MaxStepsExceeded)
    }
}

Outcome::Interrupted { state, reason, .. } => {
    println!("
=== Draft ===

{}", state.draft.as_deref().unwrap_or(""));
    println!("
=== Editor's notes ===

{}",
             state.critique.as_deref().unwrap_or(""));
    if prompt_yes_no(&reason) {
        let outcome = graph
            .run_with_checkpoint(RUN_ID, state, &ckpt)
            .await?;
        // ... print final answer ...
    } else {
        eprintln!("[abort] Checkpoint kept; rerun to resume.");
    }
}

const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS checkpoints (
    run_id     TEXT PRIMARY KEY,
    next_node  TEXT NOT NULL,
    state_json TEXT NOT NULL,
    updated_at INTEGER NOT NULL
);
"#;

let mut graph = Graph::<State>::new("draft").retry_attempts(3);

struct Draft { http: reqwest::Client, api_key: String }

#[async_trait]
impl Node<State> for Draft {
    fn name(&self) -> &'static str { "draft" }
    async fn run(&self, s: &mut State) -> Result<NextStep, GraphError> {
        let system = "You are Eugene, a careful technical writer.                       Write a single short answer to the user's question.                       Plain prose, no bullet lists, fewer than 200 words.";
        let answer = call_claude(&self.http, &self.api_key, system,
                                 &format!("Write a draft answer to: {}", s.request)).await?;
        s.draft = Some(answer);
        Ok(NextStep::Goto("review"))
    }
}

struct Review { http: reqwest::Client, api_key: String }

#[async_trait]
impl Node<State> for Review {
    fn name(&self) -> &'static str { "review" }
    async fn run(&self, s: &mut State) -> Result<NextStep, GraphError> {
        let system = "You are a sharp editor. Read the draft and list the                       three biggest weaknesses in one sentence each.";
        let critique = call_claude(&self.http, &self.api_key, system,
                                   s.draft.as_deref().unwrap_or("")).await?;
        s.critique = Some(critique);
        Ok(NextStep::Interrupt {
            reason: "Draft and critique ready. Approve revision?".into(),
        })
    }
}

struct Revise { http: reqwest::Client, api_key: String }

#[async_trait]
impl Node<State> for Revise {
    fn name(&self) -> &'static str { "revise" }
    async fn run(&self, s: &mut State) -> Result<NextStep, GraphError> {
        let system = "You are Eugene. Rewrite the draft using the editor's                       notes. Keep the same tone. Plain prose, fewer than 200 words.";
        let improved = call_claude(&self.http, &self.api_key, system,
                                   &format!("Draft:
{}

Notes:
{}",
                                           s.draft.as_deref().unwrap_or(""),
                                           s.critique.as_deref().unwrap_or(""))).await?;
        s.draft = Some(improved);
        Ok(NextStep::Halt)
    }
}

#[async_trait]
pub trait NodeHook<S>: Send + Sync + 'static {
    async fn before(
        &self,
        node_name: &str,
        required: PermissionMode,
        state: &S,
    ) -> Result<HookOutcome, GraphError> {
        Ok(HookOutcome::Proceed)
    }

    async fn after(
        &self,
        node_name: &str,
        state: &S,
        result: NextStep,
    ) -> Result<NextStep, GraphError> {
        Ok(result)
    }
}

enum HookOutcome {
    Proceed,
    Modify(Value),
    Deny(String),
    Replace(String),
}

pub enum PermissionMode {
    Plan,
    Default,
    AcceptEdits,
    BypassPermissions,
}

impl PermissionMode {
    pub fn level(self) -> u8 { /* Plan=0 < Default=1 < AcceptEdits=2 < Bypass=3 */ }
    pub fn permits(self, required: PermissionMode) -> bool {
        self.level() >= required.level()
    }
}

#[async_trait]
pub trait Node<S> {
    fn name(&self) -> &'static str;
    fn required_mode(&self) -> PermissionMode { PermissionMode::Plan }
    async fn run(&self, state: &mut S) -> Result<NextStep, GraphError>;
}

impl<S: Send + Sync + 'static> NodeHook<S> for ModeGate {
    async fn before(&self, node_name: &str, required: PermissionMode, state: &S)
        -> Result<HookOutcome, GraphError> {
        if self.current_mode.permits(required) {
            Ok(HookOutcome::Proceed)
        } else {
            Ok(HookOutcome::Deny(format!(
                "Node '{}' requires {:?} mode; current mode is {:?}.",
                node_name, required, self.current_mode
            )))
        }
    }
}