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_default().map_err(SdkError::Session)?;
141        let prompt = prompt.to_string();
142
143        let session_id = session.id.clone();
144        let agent_task = async {
145            agent
146                .run_turn(&mut session, &prompt, &store)
147                .await
148                .map_err(SdkError::Agent)
149        };
150        tokio::pin!(agent_task);
151
152        let mut outcome = TurnOutcome::default();
153        loop {
154            tokio::select! {
155                event = rx.recv() => {
156                    match event {
157                        Some(AgentEvent::TurnComplete { usage }) => {
158                            outcome.usage = usage;
159                            on_event(AgentEvent::TurnComplete { usage: outcome.usage });
160                            break;
161                        }
162                        Some(other) => on_event(other),
163                        None => break,
164                    }
165                }
166                result = &mut agent_task => {
167                    outcome = result?;
168                    break;
169                }
170            }
171        }
172
173        Ok(RunResult {
174            session_id,
175            outcome,
176        })
177    }
178
179    /// Convenience wrapper that collects events only for completion.
180    pub async fn run(&self, prompt: &str) -> Result<RunResult, SdkError> {
181        self.run_with_handler(prompt, |_| {}).await
182    }
183}
184
185#[derive(Debug, thiserror::Error)]
186pub enum SdkError {
187    #[error("config error: {0}")]
188    Config(#[from] codei_config::ConfigError),
189    #[error("llm error: {0}")]
190    Llm(#[from] codei_llm::LlmError),
191    #[error("session error: {0}")]
192    Session(#[from] codei_session::SessionError),
193    #[error("agent error: {0}")]
194    Agent(#[from] codei_agent::AgentError),
195}