codei_sdk/
turn.rs

1use std::sync::Arc;
2
3use codei_agent::{AgentError, AgentEvent, AgentLoop, TurnOutcome};
4use codei_config::ResolvedConfig;
5use codei_session::{Session, SessionStore};
6use codei_tools::{handler_for_policy, ApprovalPolicy, ToolContext};
7use tokio::sync::mpsc::{self, UnboundedSender};
8use tokio::task::JoinHandle;
9
10use crate::error::SdkError;
11use crate::runtime::AgentRuntime;
12use crate::session::SessionHandle;
13
14pub fn approval_policy(auto_approve: bool) -> ApprovalPolicy {
15    if auto_approve {
16        ApprovalPolicy::Never
17    } else {
18        ApprovalPolicy::OnDestructive
19    }
20}
21
22pub fn tool_context(config: &Arc<ResolvedConfig>, policy: ApprovalPolicy) -> ToolContext {
23    ToolContext {
24        cwd: config.cwd.clone(),
25        config: Arc::clone(config),
26        approval: Arc::from(handler_for_policy(policy)),
27    }
28}
29
30pub fn agent_loop(
31    runtime: &AgentRuntime,
32    tool_ctx: ToolContext,
33    events: Option<UnboundedSender<AgentEvent>>,
34) -> AgentLoop {
35    AgentLoop::new(
36        Arc::clone(&runtime.config),
37        Arc::clone(&runtime.model),
38        Arc::clone(&runtime.provider),
39        runtime.provider_name.clone(),
40        tool_ctx,
41        runtime.mcp.clone(),
42        events,
43    )
44}
45
46pub async fn run_turn(
47    agent: &AgentLoop,
48    session: &mut Session,
49    prompt: &str,
50    store: &SessionStore,
51) -> Result<TurnOutcome, AgentError> {
52    agent.run_turn(session, prompt, store).await
53}
54
55pub async fn run_turn_with_events<F>(
56    runtime: &AgentRuntime,
57    session: &mut Session,
58    prompt: &str,
59    policy: ApprovalPolicy,
60    mut on_event: F,
61) -> Result<TurnOutcome, SdkError>
62where
63    F: FnMut(AgentEvent),
64{
65    let (tx, mut rx) = mpsc::unbounded_channel();
66    let agent = agent_loop(runtime, tool_context(&runtime.config, policy), Some(tx));
67    let prompt = prompt.to_string();
68    let store = runtime.store.clone();
69
70    let mut agent_task = Box::pin(async {
71        agent
72            .run_turn(session, &prompt, store.as_ref())
73            .await
74            .map_err(SdkError::Agent)
75    });
76
77    let mut outcome = TurnOutcome::default();
78    loop {
79        tokio::select! {
80            event = rx.recv() => {
81                match event {
82                    Some(AgentEvent::TurnComplete { usage }) => {
83                        outcome.usage = usage;
84                        on_event(AgentEvent::TurnComplete { usage: outcome.usage });
85                        break;
86                    }
87                    Some(other) => on_event(other),
88                    None => break,
89                }
90            }
91            result = &mut agent_task => {
92                outcome = result?;
93                break;
94            }
95        }
96    }
97
98    Ok(outcome)
99}
100
101pub fn spawn_turn(
102    handle: Arc<SessionHandle>,
103    prompt: String,
104    policy: ApprovalPolicy,
105    tx: UnboundedSender<AgentEvent>,
106) -> JoinHandle<()> {
107    tokio::spawn(async move {
108        let _guard = handle.turn_lock.lock().await;
109        let tool_ctx = tool_context(&handle.runtime.config, policy);
110        let agent = agent_loop(&handle.runtime, tool_ctx, Some(tx.clone()));
111        let session = Arc::clone(&handle.session);
112        let store = Arc::clone(&handle.runtime.store);
113        let mut session = session.write().await;
114        if let Err(err) = agent.run_turn(&mut session, &prompt, store.as_ref()).await {
115            let _ = tx.send(AgentEvent::Error {
116                message: err.to_string(),
117            });
118        }
119    })
120}
codei_sdk/turn.rs

codei_sdk/
turn.rs
