Skip to main content

codei_sdk/
lib.rs

1//! Programmatic SDK for CodeI.
2
3use std::path::PathBuf;
4use std::sync::{Arc, RwLock};
5
6use codei_agent::{AgentEvent, AgentLoop, TurnOutcome};
7use codei_config::{load, LoadOptions, ResolvedConfig};
8use codei_llm::create_provider;
9use codei_mcp::McpManager;
10use codei_session::{Session, SessionStore};
11use codei_tools::{handler_for_policy, ApprovalPolicy, ToolContext};
12use tokio::sync::mpsc;
13
14/// Result of a single agent run.
15#[derive(Debug, Clone)]
16pub struct RunResult {
17    pub session_id: String,
18    pub outcome: TurnOutcome,
19}
20
21/// Builder for [`CodeiClient`].
22pub struct CodeiClientBuilder {
23    cwd: Option<PathBuf>,
24    model: Option<String>,
25    provider: Option<String>,
26    auto_approve: bool,
27}
28
29impl Default for CodeiClientBuilder {
30    fn default() -> Self {
31        Self {
32            cwd: None,
33            model: None,
34            provider: None,
35            auto_approve: true,
36        }
37    }
38}
39
40impl CodeiClientBuilder {
41    pub fn new() -> Self {
42        Self::default()
43    }
44
45    pub fn cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
46        self.cwd = Some(cwd.into());
47        self
48    }
49
50    pub fn model(mut self, model: impl Into<String>) -> Self {
51        self.model = Some(model.into());
52        self
53    }
54
55    pub fn provider(mut self, provider: impl Into<String>) -> Self {
56        self.provider = Some(provider.into());
57        self
58    }
59
60    pub fn auto_approve(mut self, yes: bool) -> Self {
61        self.auto_approve = yes;
62        self
63    }
64
65    pub async fn build(self) -> Result<CodeiClient, SdkError> {
66        let resolved = load(&LoadOptions {
67            cwd: self.cwd,
68            model: self.model.clone(),
69            provider: self.provider.clone(),
70            language: None,
71        })
72        .map_err(SdkError::Config)?;
73
74        let config = Arc::new(resolved);
75        let provider_name = self
76            .provider
77            .unwrap_or_else(|| config.config.defaults.provider.clone());
78        let provider = create_provider(&config).map_err(SdkError::Llm)?;
79        let model = Arc::new(RwLock::new(
80            self.model
81                .unwrap_or_else(|| config.config.defaults.model.clone()),
82        ));
83
84        Ok(CodeiClient {
85            config,
86            provider,
87            provider_name,
88            model,
89            auto_approve: self.auto_approve,
90        })
91    }
92}
93
94/// Programmatic entry point for running CodeI agents.
95pub struct CodeiClient {
96    config: Arc<ResolvedConfig>,
97    provider: Arc<dyn codei_llm::LlmProvider>,
98    provider_name: String,
99    model: Arc<RwLock<String>>,
100    auto_approve: bool,
101}
102
103impl CodeiClient {
104    pub fn builder() -> CodeiClientBuilder {
105        CodeiClientBuilder::new()
106    }
107
108    /// Run a prompt and invoke `on_event` for each agent event.
109    pub async fn run_with_handler<F>(
110        &self,
111        prompt: &str,
112        mut on_event: F,
113    ) -> Result<RunResult, SdkError>
114    where
115        F: FnMut(AgentEvent),
116    {
117        let (tx, mut rx) = mpsc::unbounded_channel();
118        let policy = if self.auto_approve {
119            ApprovalPolicy::Never
120        } else {
121            ApprovalPolicy::OnDestructive
122        };
123        let tool_ctx = ToolContext {
124            cwd: self.config.cwd.clone(),
125            config: Arc::clone(&self.config),
126            approval: Arc::from(handler_for_policy(policy)),
127        };
128        let mcp = McpManager::connect_optional().await;
129        let agent = AgentLoop::new(
130            Arc::clone(&self.config),
131            Arc::clone(&self.model),
132            Arc::clone(&self.provider),
133            self.provider_name.clone(),
134            tool_ctx,
135            mcp,
136            Some(tx),
137        );
138
139        let mut session = Session::new(self.config.cwd.clone());
140        let store = SessionStore::open_for_config(&self.config.config.session)
141            .map_err(SdkError::Session)?;
142        let prompt = prompt.to_string();
143
144        let session_id = session.id.clone();
145        let agent_task = async {
146            agent
147                .run_turn(&mut session, &prompt, &store)
148                .await
149                .map_err(SdkError::Agent)
150        };
151        tokio::pin!(agent_task);
152
153        let mut outcome = TurnOutcome::default();
154        loop {
155            tokio::select! {
156                event = rx.recv() => {
157                    match event {
158                        Some(AgentEvent::TurnComplete { usage }) => {
159                            outcome.usage = usage;
160                            on_event(AgentEvent::TurnComplete { usage: outcome.usage });
161                            break;
162                        }
163                        Some(other) => on_event(other),
164                        None => break,
165                    }
166                }
167                result = &mut agent_task => {
168                    outcome = result?;
169                    break;
170                }
171            }
172        }
173
174        Ok(RunResult {
175            session_id,
176            outcome,
177        })
178    }
179
180    /// Convenience wrapper that collects events only for completion.
181    pub async fn run(&self, prompt: &str) -> Result<RunResult, SdkError> {
182        self.run_with_handler(prompt, |_| {}).await
183    }
184}
185
186#[derive(Debug, thiserror::Error)]
187pub enum SdkError {
188    #[error("config error: {0}")]
189    Config(#[from] codei_config::ConfigError),
190    #[error("llm error: {0}")]
191    Llm(#[from] codei_llm::LlmError),
192    #[error("session error: {0}")]
193    Session(#[from] codei_session::SessionError),
194    #[error("agent error: {0}")]
195    Agent(#[from] codei_agent::AgentError),
196}