incurs_codemode/
codemode.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use incurs::command::RequestContext;
6use incurs::tool::{ToolCallControl, ToolEvent, ToolEventSink};
7use serde_json::Value;
8use tokio::sync::Mutex;
9use tokio_util::sync::CancellationToken;
10
11use crate::{
12    ArtifactStore, CapabilitySnapshot, Clock, CodeExecutor, CodeModeRuntime, Connector,
13    DispatchRequest, DispatchSession, ExecutionEvent, ExecutionHost, ExecutionState,
14    ExecutionStatus, RuntimeStore, SearchOutput, SystemClock, ToolContext,
15};
16
17/// Transport context inherited by one Code Mode execution pass.
18#[derive(Clone, Default)]
19pub struct CodeModeRunOptions {
20    /// Cooperative cancellation signal for the active pass.
21    pub cancellation: CancellationToken,
22    /// Request metadata inherited by connector calls.
23    pub request: Option<RequestContext>,
24}
25
26struct RuntimeEventSink {
27    runtime: Arc<CodeModeRuntime>,
28    execution_id: String,
29    clock: Arc<dyn Clock>,
30}
31
32#[async_trait]
33impl ToolEventSink for RuntimeEventSink {
34    async fn emit(&self, event: ToolEvent) {
35        let at = self.clock.now_ms();
36        let event = match event {
37            ToolEvent::Progress { message, fraction } => ExecutionEvent::Progress {
38                message,
39                fraction,
40                at,
41            },
42            ToolEvent::Log { level, message } => ExecutionEvent::Log { level, message, at },
43            ToolEvent::Chunk { data } => ExecutionEvent::Chunk { data, at },
44        };
45        let _ = self.runtime.event(&self.execution_id, event).await;
46    }
47}
48
49/// Generic Code Mode lifecycle over interchangeable execution and storage adapters.
50pub struct CodeMode {
51    runtime: Arc<CodeModeRuntime>,
52    executor: Box<dyn CodeExecutor>,
53    connectors: Vec<Arc<dyn Connector>>,
54    clock: Arc<dyn Clock>,
55    contexts: Mutex<HashMap<String, ToolContext>>,
56}
57
58impl CodeMode {
59    /// Creates a native Code Mode runtime using the system clock.
60    pub fn new(
61        store: Arc<dyn RuntimeStore>,
62        executor: impl CodeExecutor + 'static,
63        connectors: Vec<Arc<dyn Connector>>,
64    ) -> Self {
65        Self::with_clock(store, executor, connectors, SystemClock)
66    }
67
68    /// Creates Code Mode with an explicit oversized-value artifact store.
69    pub fn with_artifact_store(
70        store: Arc<dyn RuntimeStore>,
71        artifacts: Arc<dyn ArtifactStore>,
72        executor: impl CodeExecutor + 'static,
73        connectors: Vec<Arc<dyn Connector>>,
74    ) -> Self {
75        Self {
76            runtime: Arc::new(CodeModeRuntime::with_artifacts(store, artifacts)),
77            executor: Box::new(executor),
78            connectors,
79            clock: Arc::new(SystemClock),
80            contexts: Mutex::new(HashMap::new()),
81        }
82    }
83
84    /// Creates Code Mode with an explicit platform clock.
85    pub fn with_clock(
86        store: Arc<dyn RuntimeStore>,
87        executor: impl CodeExecutor + 'static,
88        connectors: Vec<Arc<dyn Connector>>,
89        clock: impl Clock + 'static,
90    ) -> Self {
91        Self {
92            runtime: Arc::new(CodeModeRuntime::new(store)),
93            executor: Box::new(executor),
94            connectors,
95            clock: Arc::new(clock),
96            contexts: Mutex::new(HashMap::new()),
97        }
98    }
99
100    /// Creates Code Mode with explicit platform clock and artifact storage.
101    pub fn with_clock_and_artifact_store(
102        store: Arc<dyn RuntimeStore>,
103        artifacts: Arc<dyn ArtifactStore>,
104        executor: impl CodeExecutor + 'static,
105        connectors: Vec<Arc<dyn Connector>>,
106        clock: impl Clock + 'static,
107    ) -> Self {
108        Self {
109            runtime: Arc::new(CodeModeRuntime::with_artifacts(store, artifacts)),
110            executor: Box::new(executor),
111            connectors,
112            clock: Arc::new(clock),
113            contexts: Mutex::new(HashMap::new()),
114        }
115    }
116
117    /// Returns the portable runtime for approvals, history, and snippets.
118    pub fn runtime(&self) -> Arc<CodeModeRuntime> {
119        Arc::clone(&self.runtime)
120    }
121
122    /// Returns model-facing JavaScript declarations and connector guidance.
123    pub async fn instructions(&self) -> Result<String, String> {
124        let mut descriptions = Vec::new();
125        for connector in &self.connectors {
126            descriptions.push(connector.describe().await?);
127        }
128        let mut sections = descriptions
129            .iter()
130            .filter_map(|connector| {
131                connector
132                    .instructions
133                    .as_ref()
134                    .map(|instructions| format!("## {}\n\n{instructions}", connector.name))
135            })
136            .collect::<Vec<_>>();
137        sections.extend(descriptions.iter().map(crate::generate_types));
138        Ok(sections.join("\n\n"))
139    }
140
141    /// Searches current connector methods and saved snippets.
142    pub async fn search(&self, query: &str) -> Result<SearchOutput, String> {
143        let mut descriptions = Vec::new();
144        for connector in &self.connectors {
145            descriptions.push(connector.describe().await?);
146        }
147        let snippets = self
148            .runtime
149            .snippets()
150            .await
151            .map_err(|error| error.to_string())?;
152        Ok(crate::search(query, &descriptions, &snippets))
153    }
154
155    /// Returns one execution by identifier.
156    pub async fn execution(&self, execution_id: &str) -> Result<ExecutionState, String> {
157        self.require(execution_id).await
158    }
159
160    /// Returns one bounded execution snapshot with artifact references intact.
161    pub async fn execution_snapshot(&self, execution_id: &str) -> Result<ExecutionState, String> {
162        self.runtime
163            .execution_snapshot(execution_id)
164            .await
165            .map_err(|error| error.to_string())?
166            .ok_or_else(|| format!("Execution \"{execution_id}\" not found"))
167    }
168
169    /// Returns one oversized value owned by an execution.
170    pub async fn artifact(&self, execution_id: &str, artifact_id: &str) -> Result<Value, String> {
171        self.runtime
172            .artifact(execution_id, artifact_id)
173            .await
174            .map_err(|error| error.to_string())?
175            .ok_or_else(|| {
176                format!("Artifact \"{artifact_id}\" not found for execution \"{execution_id}\"")
177            })
178    }
179
180    /// Returns ordered retained events for one execution.
181    pub async fn events(&self, execution_id: &str) -> Result<Vec<ExecutionEvent>, String> {
182        Ok(self.require(execution_id).await?.events)
183    }
184
185    /// Cancels a running or paused execution.
186    pub async fn cancel(&self, execution_id: &str) -> Result<ExecutionState, String> {
187        if let Some(context) = self.contexts.lock().await.get(execution_id) {
188            context.control.cancellation.cancel();
189        }
190        self.runtime
191            .cancel(execution_id, self.clock.now_ms())
192            .await
193            .map_err(|error| error.to_string())?;
194        self.require(execution_id).await
195    }
196
197    /// Creates a durable running execution without driving its first pass.
198    pub async fn start(&self, code: &str) -> Result<ExecutionState, String> {
199        let mut descriptions = Vec::new();
200        for connector in &self.connectors {
201            descriptions.push(connector.describe().await?);
202        }
203        let capabilities = CapabilitySnapshot::new(descriptions)?;
204        let id = self
205            .runtime
206            .begin_with_capabilities(code, capabilities, self.clock.now_ms())
207            .await
208            .map_err(|error| error.to_string())?;
209        self.require(&id).await
210    }
211
212    /// Starts and drives a new execution until it completes, fails, or pauses.
213    pub async fn execute(&self, code: &str) -> Result<ExecutionState, String> {
214        self.execute_with(code, CodeModeRunOptions::default()).await
215    }
216
217    /// Starts and drives a new execution with transport request context.
218    pub async fn execute_with(
219        &self,
220        code: &str,
221        options: CodeModeRunOptions,
222    ) -> Result<ExecutionState, String> {
223        let state = self.start(code).await?;
224        self.drive_with(&state.id, options).await
225    }
226
227    /// Resumes a paused execution and drives the next replay pass.
228    pub async fn resume(&self, execution_id: &str) -> Result<ExecutionState, String> {
229        self.resume_with(execution_id, CodeModeRunOptions::default())
230            .await
231    }
232
233    /// Resumes a paused execution with transport request context.
234    pub async fn resume_with(
235        &self,
236        execution_id: &str,
237        options: CodeModeRunOptions,
238    ) -> Result<ExecutionState, String> {
239        self.runtime
240            .resume(execution_id, self.clock.now_ms())
241            .await
242            .map_err(|error| error.to_string())?;
243        self.drive_with(execution_id, options).await
244    }
245
246    /// Approves a pending action and resumes only if this call won the race.
247    pub async fn approve(&self, execution_id: &str, seq: u64) -> Result<ExecutionState, String> {
248        self.approve_with(execution_id, seq, CodeModeRunOptions::default())
249            .await
250    }
251
252    /// Approves a pending action and resumes with transport request context.
253    pub async fn approve_with(
254        &self,
255        execution_id: &str,
256        seq: u64,
257        options: CodeModeRunOptions,
258    ) -> Result<ExecutionState, String> {
259        if !self
260            .runtime
261            .approve(execution_id, seq, self.clock.now_ms())
262            .await
263            .map_err(|error| error.to_string())?
264        {
265            return self.require(execution_id).await;
266        }
267        self.drive_with(execution_id, options).await
268    }
269
270    /// Rejects a pending action and notifies connector lifecycle hooks.
271    pub async fn reject(&self, execution_id: &str, seq: u64) -> Result<ExecutionState, String> {
272        if self
273            .runtime
274            .reject(execution_id, seq, self.clock.now_ms())
275            .await
276            .map_err(|error| error.to_string())?
277        {
278            self.notify_execution_end(execution_id, "rejected").await;
279        }
280        self.require(execution_id).await
281    }
282
283    /// Compensates applied connector actions in reverse order.
284    pub async fn rollback(&self, execution_id: &str) -> Result<ExecutionState, String> {
285        for action in self
286            .runtime
287            .actions_to_revert(execution_id)
288            .await
289            .map_err(|error| error.to_string())?
290        {
291            let connector = self
292                .connector(&action.connector)
293                .await?
294                .ok_or_else(|| format!("Connector \"{}\" not found", action.connector))?;
295            if !connector
296                .revert(
297                    &action.method,
298                    action.arguments,
299                    action.result.unwrap_or(Value::Null),
300                    &ToolContext {
301                        execution_id: execution_id.to_string(),
302                        control: Default::default(),
303                        request: None,
304                    },
305                )
306                .await?
307            {
308                return Err(format!(
309                    "{}.{} did not compensate step {}",
310                    action.connector, action.method, action.seq
311                ));
312            }
313            self.runtime
314                .mark_reverted(execution_id, action.seq, self.clock.now_ms())
315                .await
316                .map_err(|error| error.to_string())?;
317        }
318        self.runtime
319            .finish_rollback(execution_id, self.clock.now_ms())
320            .await
321            .map_err(|error| error.to_string())?;
322        self.notify_execution_end(execution_id, "rolled_back").await;
323        self.require(execution_id).await
324    }
325
326    /// Expires stale live executions and releases connector resources.
327    pub async fn expire(&self, max_age_ms: u64) -> Result<Vec<String>, String> {
328        let ids = self
329            .runtime
330            .expire(self.clock.now_ms(), max_age_ms)
331            .await
332            .map_err(|error| error.to_string())?;
333        for id in &ids {
334            let status = self.require(id).await?.status;
335            self.notify_execution_end(
336                id,
337                if status == ExecutionStatus::Rejected {
338                    "rejected"
339                } else {
340                    "error"
341                },
342            )
343            .await;
344        }
345        Ok(ids)
346    }
347
348    /// Handles one program-to-host dispatch request.
349    pub async fn dispatch(&self, request: DispatchRequest) -> Result<Value, String> {
350        match request {
351            DispatchRequest::Call {
352                execution_id,
353                seq,
354                connector,
355                method,
356                arguments,
357            } => {
358                let session = self.session(&execution_id).await?;
359                serde_json::to_value(
360                    session
361                        .call_at(seq, &connector, &method, arguments, self.clock.now_ms())
362                        .await,
363                )
364                .map_err(|error| error.to_string())
365            }
366            DispatchRequest::BeginStep {
367                execution_id,
368                seq,
369                name,
370            } => {
371                let session = self.session(&execution_id).await?;
372                serde_json::to_value(
373                    session
374                        .begin_step_at(seq, &name, self.clock.now_ms())
375                        .await
376                        .map_err(|error| error.to_string())?,
377                )
378                .map_err(|error| error.to_string())
379            }
380            DispatchRequest::RecordStep {
381                execution_id,
382                seq,
383                result,
384            } => {
385                self.runtime
386                    .record_result(&execution_id, seq, result, self.clock.now_ms())
387                    .await
388                    .map_err(|error| error.to_string())?;
389                Ok(serde_json::json!({ "ok": true }))
390            }
391        }
392    }
393
394    /// Drives one running execution pass with transport request context.
395    pub async fn drive_with(
396        &self,
397        execution_id: &str,
398        options: CodeModeRunOptions,
399    ) -> Result<ExecutionState, String> {
400        let state = self.require(execution_id).await?;
401        if matches!(
402            state.status,
403            ExecutionStatus::Completed
404                | ExecutionStatus::Error
405                | ExecutionStatus::Rejected
406                | ExecutionStatus::RolledBack
407                | ExecutionStatus::Cancelled
408        ) {
409            return Ok(state);
410        }
411        let context = self.context(execution_id, options);
412        let cancellation = context.control.cancellation.clone();
413        self.contexts
414            .lock()
415            .await
416            .insert(execution_id.to_string(), context.clone());
417        let session = match self.session_with_context(execution_id, context).await {
418            Ok(session) => session,
419            Err(error) => {
420                self.contexts.lock().await.remove(execution_id);
421                return Err(error);
422            }
423        };
424        let descriptions = session.descriptions();
425        let execution = self.executor.execute(
426            &state.code,
427            &descriptions,
428            execution_id,
429            Arc::new(ExecutionHost::new(
430                Arc::clone(&session),
431                Arc::clone(&self.clock),
432            )),
433        );
434        tokio::pin!(execution);
435        let response = tokio::select! {
436            _ = cancellation.cancelled() => None,
437            response = &mut execution => Some(response),
438        };
439        self.contexts.lock().await.remove(execution_id);
440        let Some(response) = response else {
441            self.runtime
442                .cancel(execution_id, self.clock.now_ms())
443                .await
444                .map_err(|error| error.to_string())?;
445            session.pass_ended("cancelled").await;
446            session.execution_ended("cancelled").await;
447            return self.require(execution_id).await;
448        };
449        let response = match response {
450            Ok(response) => response,
451            Err(error) => {
452                self.runtime
453                    .fail(execution_id, error, Vec::new(), self.clock.now_ms())
454                    .await
455                    .map_err(|error| error.to_string())?;
456                session.pass_ended("error").await;
457                session.execution_ended("error").await;
458                return self.require(execution_id).await;
459            }
460        };
461        let current = self.require(execution_id).await?;
462        if current.status == ExecutionStatus::Paused {
463            session.pass_ended("paused").await;
464            return Ok(current);
465        }
466        if current.status == ExecutionStatus::Error {
467            session.pass_ended("error").await;
468            session.execution_ended("error").await;
469            return Ok(current);
470        }
471        if let Some(error) = response.error {
472            self.runtime
473                .fail(execution_id, error, response.logs, self.clock.now_ms())
474                .await
475                .map_err(|error| error.to_string())?;
476            session.pass_ended("error").await;
477            session.execution_ended("error").await;
478        } else {
479            self.runtime
480                .complete(
481                    execution_id,
482                    response.result.unwrap_or(Value::Null),
483                    response.logs,
484                    self.clock.now_ms(),
485                )
486                .await
487                .map_err(|error| error.to_string())?;
488            session.pass_ended("completed").await;
489            session.execution_ended("completed").await;
490        }
491        self.require(execution_id).await
492    }
493
494    async fn session(&self, execution_id: &str) -> Result<Arc<DispatchSession>, String> {
495        let context = self
496            .contexts
497            .lock()
498            .await
499            .get(execution_id)
500            .cloned()
501            .unwrap_or_else(|| self.context(execution_id, CodeModeRunOptions::default()));
502        self.session_with_context(execution_id, context).await
503    }
504
505    async fn session_with_context(
506        &self,
507        execution_id: &str,
508        context: ToolContext,
509    ) -> Result<Arc<DispatchSession>, String> {
510        let state = self.require(execution_id).await?;
511        let descriptions = if let Some(capabilities) = state.capabilities {
512            capabilities.connectors
513        } else {
514            let mut descriptions = Vec::new();
515            for connector in &self.connectors {
516                descriptions.push(connector.describe().await?);
517            }
518            descriptions
519        };
520        Ok(Arc::new(
521            DispatchSession::new_with_descriptions_and_context(
522                Arc::clone(&self.runtime),
523                context,
524                self.connectors.clone(),
525                descriptions,
526            )
527            .await?,
528        ))
529    }
530
531    fn context(&self, execution_id: &str, options: CodeModeRunOptions) -> ToolContext {
532        ToolContext {
533            execution_id: execution_id.to_string(),
534            control: ToolCallControl {
535                cancellation: options.cancellation,
536                events: Some(Arc::new(RuntimeEventSink {
537                    runtime: Arc::clone(&self.runtime),
538                    execution_id: execution_id.to_string(),
539                    clock: Arc::clone(&self.clock),
540                })),
541            },
542            request: options.request,
543        }
544    }
545
546    async fn require(&self, execution_id: &str) -> Result<ExecutionState, String> {
547        self.runtime
548            .execution(execution_id)
549            .await
550            .map_err(|error| error.to_string())?
551            .ok_or_else(|| format!("Execution \"{execution_id}\" not found"))
552    }
553
554    async fn connector(&self, name: &str) -> Result<Option<&Arc<dyn Connector>>, String> {
555        for connector in &self.connectors {
556            if connector.describe().await?.name == name {
557                return Ok(Some(connector));
558            }
559        }
560        Ok(None)
561    }
562
563    async fn notify_execution_end(&self, execution_id: &str, status: &str) {
564        for connector in &self.connectors {
565            connector.execution_ended(execution_id, status).await;
566        }
567    }
568}
incurs_codemode/codemode.rs

incurs_codemode/
codemode.rs
