1use 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#[derive(Debug, Clone)]
16pub struct RunResult {
17 pub session_id: String,
18 pub outcome: TurnOutcome,
19}
20
21pub 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
94pub 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 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 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}