Skip to main content

incurs_codemode/
codemode.rs

1use std::collections::{HashMap, HashSet};
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, Semaphore};
9use tokio_util::sync::CancellationToken;
10
11use crate::{
12    ArtifactStore, CapabilitySnapshot, Clock, CodeExecutor, CodeModeRuntime, Connector,
13    ConnectorDescription, DispatchRequest, DispatchSession, ExecutionEvent, ExecutionHost,
14    ExecutionState, 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    pass_gates: Mutex<HashMap<String, Arc<Semaphore>>>,
57    active_rollbacks: Mutex<HashSet<String>>,
58}
59
60impl CodeMode {
61    /// Creates a native Code Mode runtime using the system clock.
62    pub fn new(
63        store: Arc<dyn RuntimeStore>,
64        executor: impl CodeExecutor + 'static,
65        connectors: Vec<Arc<dyn Connector>>,
66    ) -> Self {
67        Self::with_clock(store, executor, connectors, SystemClock)
68    }
69
70    /// Creates Code Mode with an explicit oversized-value artifact store.
71    pub fn with_artifact_store(
72        store: Arc<dyn RuntimeStore>,
73        artifacts: Arc<dyn ArtifactStore>,
74        executor: impl CodeExecutor + 'static,
75        connectors: Vec<Arc<dyn Connector>>,
76    ) -> Self {
77        Self {
78            runtime: Arc::new(CodeModeRuntime::with_artifacts(store, artifacts)),
79            executor: Box::new(executor),
80            connectors,
81            clock: Arc::new(SystemClock),
82            contexts: Mutex::new(HashMap::new()),
83            pass_gates: Mutex::new(HashMap::new()),
84            active_rollbacks: Mutex::new(HashSet::new()),
85        }
86    }
87
88    /// Creates Code Mode with an explicit platform clock.
89    pub fn with_clock(
90        store: Arc<dyn RuntimeStore>,
91        executor: impl CodeExecutor + 'static,
92        connectors: Vec<Arc<dyn Connector>>,
93        clock: impl Clock + 'static,
94    ) -> Self {
95        Self {
96            runtime: Arc::new(CodeModeRuntime::new(store)),
97            executor: Box::new(executor),
98            connectors,
99            clock: Arc::new(clock),
100            contexts: Mutex::new(HashMap::new()),
101            pass_gates: Mutex::new(HashMap::new()),
102            active_rollbacks: Mutex::new(HashSet::new()),
103        }
104    }
105
106    /// Creates Code Mode with explicit platform clock and artifact storage.
107    pub fn with_clock_and_artifact_store(
108        store: Arc<dyn RuntimeStore>,
109        artifacts: Arc<dyn ArtifactStore>,
110        executor: impl CodeExecutor + 'static,
111        connectors: Vec<Arc<dyn Connector>>,
112        clock: impl Clock + 'static,
113    ) -> Self {
114        Self {
115            runtime: Arc::new(CodeModeRuntime::with_artifacts(store, artifacts)),
116            executor: Box::new(executor),
117            connectors,
118            clock: Arc::new(clock),
119            contexts: Mutex::new(HashMap::new()),
120            pass_gates: Mutex::new(HashMap::new()),
121            active_rollbacks: Mutex::new(HashSet::new()),
122        }
123    }
124
125    /// Returns the portable runtime for approvals, history, and snippets.
126    pub fn runtime(&self) -> Arc<CodeModeRuntime> {
127        Arc::clone(&self.runtime)
128    }
129
130    /// Returns model-facing JavaScript declarations and connector guidance.
131    pub async fn instructions(&self) -> Result<String, String> {
132        let mut descriptions = Vec::new();
133        for connector in &self.connectors {
134            descriptions.push(connector.describe().await?);
135        }
136        let mut sections = descriptions
137            .iter()
138            .filter_map(|connector| {
139                connector
140                    .instructions
141                    .as_ref()
142                    .map(|instructions| format!("## {}\n\n{instructions}", connector.name))
143            })
144            .collect::<Vec<_>>();
145        sections.extend(descriptions.iter().map(crate::generate_types));
146        Ok(sections.join("\n\n"))
147    }
148
149    /// Searches current connector methods and saved snippets.
150    pub async fn search(&self, query: &str) -> Result<SearchOutput, String> {
151        let mut descriptions = Vec::new();
152        for connector in &self.connectors {
153            descriptions.push(connector.describe().await?);
154        }
155        let snippets = self
156            .runtime
157            .snippets()
158            .await
159            .map_err(|error| error.to_string())?;
160        Ok(crate::search(query, &descriptions, &snippets))
161    }
162
163    /// Returns one execution by identifier.
164    pub async fn execution(&self, execution_id: &str) -> Result<ExecutionState, String> {
165        self.require(execution_id).await
166    }
167
168    /// Returns one bounded execution snapshot with artifact references intact.
169    pub async fn execution_snapshot(&self, execution_id: &str) -> Result<ExecutionState, String> {
170        self.runtime
171            .execution_snapshot(execution_id)
172            .await
173            .map_err(|error| error.to_string())?
174            .ok_or_else(|| format!("Execution \"{execution_id}\" not found"))
175    }
176
177    /// Returns one oversized value owned by an execution.
178    pub async fn artifact(&self, execution_id: &str, artifact_id: &str) -> Result<Value, String> {
179        self.runtime
180            .artifact(execution_id, artifact_id)
181            .await
182            .map_err(|error| error.to_string())?
183            .ok_or_else(|| {
184                format!("Artifact \"{artifact_id}\" not found for execution \"{execution_id}\"")
185            })
186    }
187
188    /// Returns ordered retained events for one execution.
189    pub async fn events(&self, execution_id: &str) -> Result<Vec<ExecutionEvent>, String> {
190        Ok(self.require(execution_id).await?.events)
191    }
192
193    /// Cancels a running or paused execution.
194    pub async fn cancel(&self, execution_id: &str) -> Result<ExecutionState, String> {
195        let pass_active = self.contexts.lock().await.contains_key(execution_id);
196        if let Some(context) = self.contexts.lock().await.get(execution_id) {
197            context.control.cancellation.cancel();
198        }
199        let changed = self
200            .runtime
201            .cancel(execution_id, self.clock.now_ms())
202            .await
203            .map_err(|error| error.to_string())?;
204        if changed && !pass_active {
205            self.notify_execution_end(execution_id, "cancelled").await;
206        }
207        self.require(execution_id).await
208    }
209
210    /// Names every connector without contacting any of them.
211    ///
212    /// The capability snapshot and the program's bindings both need to know which
213    /// namespaces exist, and nothing more. Building them from `describe` meant
214    /// connecting to every configured server before the program ran; a name-only
215    /// description carries exactly what they use, and the tools are resolved per
216    /// namespace on first call.
217    fn namespace_stubs(&self) -> Vec<ConnectorDescription> {
218        self.connectors
219            .iter()
220            .map(|connector| ConnectorDescription {
221                name: connector.name().to_string(),
222                instructions: None,
223                tools: Vec::new(),
224            })
225            .collect()
226    }
227
228    /// Creates a durable running execution without driving its first pass.
229    pub async fn start(&self, code: &str) -> Result<ExecutionState, String> {
230        let capabilities = CapabilitySnapshot::new(self.namespace_stubs())?;
231        let id = self
232            .runtime
233            .begin_with_capabilities(code, capabilities, self.clock.now_ms())
234            .await
235            .map_err(|error| error.to_string())?;
236        self.require(&id).await
237    }
238
239    /// Starts and drives a new execution until it completes, fails, or pauses.
240    pub async fn execute(&self, code: &str) -> Result<ExecutionState, String> {
241        self.execute_with(code, CodeModeRunOptions::default()).await
242    }
243
244    /// Starts and drives a new execution with transport request context.
245    pub async fn execute_with(
246        &self,
247        code: &str,
248        options: CodeModeRunOptions,
249    ) -> Result<ExecutionState, String> {
250        let state = self.start(code).await?;
251        self.drive_with(&state.id, options).await
252    }
253
254    /// Resumes a paused execution and drives the next replay pass.
255    pub async fn resume(&self, execution_id: &str) -> Result<ExecutionState, String> {
256        self.resume_with(execution_id, CodeModeRunOptions::default())
257            .await
258    }
259
260    /// Resumes a paused execution with transport request context.
261    pub async fn resume_with(
262        &self,
263        execution_id: &str,
264        options: CodeModeRunOptions,
265    ) -> Result<ExecutionState, String> {
266        self.runtime
267            .resume(execution_id, self.clock.now_ms())
268            .await
269            .map_err(|error| error.to_string())?;
270        self.drive_with(execution_id, options).await
271    }
272
273    /// Approves a pending action and resumes only if this call won the race.
274    pub async fn approve(&self, execution_id: &str, seq: u64) -> Result<ExecutionState, String> {
275        self.approve_with(execution_id, seq, CodeModeRunOptions::default())
276            .await
277    }
278
279    /// Approves a pending action and resumes with transport request context.
280    pub async fn approve_with(
281        &self,
282        execution_id: &str,
283        seq: u64,
284        options: CodeModeRunOptions,
285    ) -> Result<ExecutionState, String> {
286        if !self
287            .runtime
288            .approve(execution_id, seq, self.clock.now_ms())
289            .await
290            .map_err(|error| error.to_string())?
291        {
292            return self.require(execution_id).await;
293        }
294        self.drive_with(execution_id, options).await
295    }
296
297    /// Rejects a pending action and notifies connector lifecycle hooks.
298    pub async fn reject(&self, execution_id: &str, seq: u64) -> Result<ExecutionState, String> {
299        if self
300            .runtime
301            .reject(execution_id, seq, self.clock.now_ms())
302            .await
303            .map_err(|error| error.to_string())?
304        {
305            self.notify_execution_end(execution_id, "rejected").await;
306        }
307        self.require(execution_id).await
308    }
309
310    /// Compensates applied connector actions in reverse order.
311    pub async fn rollback(&self, execution_id: &str) -> Result<ExecutionState, String> {
312        {
313            let mut active = self.active_rollbacks.lock().await;
314            if !active.insert(execution_id.to_string()) {
315                return Err(format!(
316                    "Execution \"{execution_id}\" is already rolling back"
317                ));
318            }
319        }
320        let result = self.rollback_inner(execution_id).await;
321        self.active_rollbacks.lock().await.remove(execution_id);
322        result
323    }
324
325    async fn rollback_inner(&self, execution_id: &str) -> Result<ExecutionState, String> {
326        for action in self
327            .runtime
328            .actions_to_revert(execution_id)
329            .await
330            .map_err(|error| error.to_string())?
331        {
332            let connector = self
333                .connector(&action.connector)
334                .await?
335                .ok_or_else(|| format!("Connector \"{}\" not found", action.connector))?;
336            if !connector
337                .revert(
338                    &action.method,
339                    action.arguments,
340                    action.result.unwrap_or(Value::Null),
341                    &ToolContext {
342                        execution_id: execution_id.to_string(),
343                        control: Default::default(),
344                        request: None,
345                    },
346                )
347                .await?
348            {
349                return Err(format!(
350                    "{}.{} did not compensate step {}",
351                    action.connector, action.method, action.seq
352                ));
353            }
354            self.runtime
355                .mark_reverted(execution_id, action.seq, self.clock.now_ms())
356                .await
357                .map_err(|error| error.to_string())?;
358        }
359        self.runtime
360            .finish_rollback(execution_id, self.clock.now_ms())
361            .await
362            .map_err(|error| error.to_string())?;
363        self.notify_execution_end(execution_id, "rolled_back").await;
364        self.require(execution_id).await
365    }
366
367    /// Expires stale live executions and releases connector resources.
368    pub async fn expire(&self, max_age_ms: u64) -> Result<Vec<String>, String> {
369        let ids = self
370            .runtime
371            .expire(self.clock.now_ms(), max_age_ms)
372            .await
373            .map_err(|error| error.to_string())?;
374        for id in &ids {
375            let status = self.require(id).await?.status;
376            self.notify_execution_end(
377                id,
378                if status == ExecutionStatus::Rejected {
379                    "rejected"
380                } else {
381                    "error"
382                },
383            )
384            .await;
385        }
386        Ok(ids)
387    }
388
389    /// Handles one program-to-host dispatch request.
390    pub async fn dispatch(&self, request: DispatchRequest) -> Result<Value, String> {
391        match request {
392            DispatchRequest::Call {
393                execution_id,
394                seq,
395                connector,
396                method,
397                arguments,
398            } => {
399                let session = self.session(&execution_id).await?;
400                serde_json::to_value(
401                    session
402                        .call_at(seq, &connector, &method, arguments, self.clock.now_ms())
403                        .await,
404                )
405                .map_err(|error| error.to_string())
406            }
407            DispatchRequest::BeginStep {
408                execution_id,
409                seq,
410                name,
411            } => {
412                let session = self.session(&execution_id).await?;
413                serde_json::to_value(
414                    session
415                        .begin_step_at(seq, &name, self.clock.now_ms())
416                        .await
417                        .map_err(|error| error.to_string())?,
418                )
419                .map_err(|error| error.to_string())
420            }
421            DispatchRequest::RecordStep {
422                execution_id,
423                seq,
424                result,
425            } => {
426                let _ = self.active_context(&execution_id).await?;
427                self.runtime
428                    .record_result(&execution_id, seq, result, self.clock.now_ms())
429                    .await
430                    .map_err(|error| error.to_string())?;
431                Ok(serde_json::json!({ "ok": true }))
432            }
433        }
434    }
435
436    /// Drives one running execution pass with transport request context.
437    pub async fn drive_with(
438        &self,
439        execution_id: &str,
440        options: CodeModeRunOptions,
441    ) -> Result<ExecutionState, String> {
442        let gate = {
443            let mut gates = self.pass_gates.lock().await;
444            Arc::clone(
445                gates
446                    .entry(execution_id.to_string())
447                    .or_insert_with(|| Arc::new(Semaphore::new(1))),
448            )
449        };
450        let _permit = gate
451            .acquire_owned()
452            .await
453            .map_err(|_| "Code Mode execution gate closed".to_string())?;
454        self.drive_pass(execution_id, options).await
455    }
456
457    async fn drive_pass(
458        &self,
459        execution_id: &str,
460        options: CodeModeRunOptions,
461    ) -> Result<ExecutionState, String> {
462        let state = self.require(execution_id).await?;
463        if matches!(
464            state.status,
465            ExecutionStatus::Completed
466                | ExecutionStatus::Error
467                | ExecutionStatus::Rejected
468                | ExecutionStatus::RolledBack
469                | ExecutionStatus::Cancelled
470        ) {
471            return Ok(state);
472        }
473        let context = self.context(execution_id, options);
474        let cancellation = context.control.cancellation.clone();
475        self.contexts
476            .lock()
477            .await
478            .insert(execution_id.to_string(), context.clone());
479        let session = match self.session_with_context(execution_id, context).await {
480            Ok(session) => session,
481            Err(error) => {
482                self.contexts.lock().await.remove(execution_id);
483                return Err(error);
484            }
485        };
486        let descriptions = session.descriptions();
487        let execution = self.executor.execute(
488            &state.code,
489            &descriptions,
490            execution_id,
491            Arc::new(ExecutionHost::new(
492                Arc::clone(&session),
493                Arc::clone(&self.clock),
494            )),
495        );
496        tokio::pin!(execution);
497        let response = tokio::select! {
498            _ = cancellation.cancelled() => None,
499            response = &mut execution => Some(response),
500        };
501        self.contexts.lock().await.remove(execution_id);
502        let Some(response) = response else {
503            self.runtime
504                .cancel(execution_id, self.clock.now_ms())
505                .await
506                .map_err(|error| error.to_string())?;
507            session.pass_ended("cancelled").await;
508            session.execution_ended("cancelled").await;
509            return self.require(execution_id).await;
510        };
511        let response = match response {
512            Ok(response) => response,
513            Err(error) => {
514                self.runtime
515                    .fail(execution_id, error, Vec::new(), self.clock.now_ms())
516                    .await
517                    .map_err(|error| error.to_string())?;
518                session.pass_ended("error").await;
519                session.execution_ended("error").await;
520                return self.require(execution_id).await;
521            }
522        };
523        let current = self.require(execution_id).await?;
524        if current.status == ExecutionStatus::Paused {
525            session.pass_ended("paused").await;
526            return Ok(current);
527        }
528        if current.status == ExecutionStatus::Error {
529            session.pass_ended("error").await;
530            session.execution_ended("error").await;
531            return Ok(current);
532        }
533        if current.status == ExecutionStatus::Cancelled {
534            session.pass_ended("cancelled").await;
535            session.execution_ended("cancelled").await;
536            return Ok(current);
537        }
538        if let Some(error) = response.error {
539            self.runtime
540                .fail(execution_id, error, response.logs, self.clock.now_ms())
541                .await
542                .map_err(|error| error.to_string())?;
543            session.pass_ended("error").await;
544            session.execution_ended("error").await;
545        } else {
546            self.runtime
547                .complete(
548                    execution_id,
549                    response.result.unwrap_or(Value::Null),
550                    response.logs,
551                    self.clock.now_ms(),
552                )
553                .await
554                .map_err(|error| error.to_string())?;
555            session.pass_ended("completed").await;
556            session.execution_ended("completed").await;
557        }
558        self.require(execution_id).await
559    }
560
561    async fn session(&self, execution_id: &str) -> Result<Arc<DispatchSession>, String> {
562        let context = self.active_context(execution_id).await?;
563        self.session_with_context(execution_id, context).await
564    }
565
566    async fn active_context(&self, execution_id: &str) -> Result<ToolContext, String> {
567        self.contexts
568            .lock()
569            .await
570            .get(execution_id)
571            .cloned()
572            .ok_or_else(|| format!("Execution \"{execution_id}\" does not have an active pass"))
573    }
574
575    async fn session_with_context(
576        &self,
577        execution_id: &str,
578        context: ToolContext,
579    ) -> Result<Arc<DispatchSession>, String> {
580        let state = self.require(execution_id).await?;
581        let descriptions = match state.capabilities {
582            Some(capabilities) => capabilities.connectors,
583            None => self.namespace_stubs(),
584        };
585        Ok(Arc::new(
586            DispatchSession::new_with_descriptions_and_context(
587                Arc::clone(&self.runtime),
588                context,
589                self.connectors.clone(),
590                descriptions,
591            )
592            .await?,
593        ))
594    }
595
596    fn context(&self, execution_id: &str, options: CodeModeRunOptions) -> ToolContext {
597        ToolContext {
598            execution_id: execution_id.to_string(),
599            control: ToolCallControl {
600                cancellation: options.cancellation,
601                events: Some(Arc::new(RuntimeEventSink {
602                    runtime: Arc::clone(&self.runtime),
603                    execution_id: execution_id.to_string(),
604                    clock: Arc::clone(&self.clock),
605                })),
606            },
607            request: options.request,
608        }
609    }
610
611    async fn require(&self, execution_id: &str) -> Result<ExecutionState, String> {
612        self.runtime
613            .execution(execution_id)
614            .await
615            .map_err(|error| error.to_string())?
616            .ok_or_else(|| format!("Execution \"{execution_id}\" not found"))
617    }
618
619    async fn connector(&self, name: &str) -> Result<Option<&Arc<dyn Connector>>, String> {
620        for connector in &self.connectors {
621            if connector.describe().await?.name == name {
622                return Ok(Some(connector));
623            }
624        }
625        Ok(None)
626    }
627
628    async fn notify_execution_end(&self, execution_id: &str, status: &str) {
629        for connector in &self.connectors {
630            connector.execution_ended(execution_id, status).await;
631        }
632    }
633}
634
635#[cfg(test)]
636mod tests {
637    use std::sync::atomic::{AtomicUsize, Ordering};
638
639    use serde_json::json;
640    use tokio::sync::Notify;
641
642    use super::*;
643    use crate::{
644        ConnectorDescription, ConnectorTool, ExecuteResult, MemoryStore, ReplayPolicy,
645        ToolAnnotations, ToolPolicy,
646    };
647
648    struct TestConnector {
649        calls: AtomicUsize,
650        ended: Mutex<Vec<String>>,
651        started: Notify,
652        release: Notify,
653        blocked: bool,
654    }
655
656    impl TestConnector {
657        fn new(blocked: bool) -> Self {
658            Self {
659                calls: AtomicUsize::new(0),
660                ended: Mutex::new(Vec::new()),
661                started: Notify::new(),
662                release: Notify::new(),
663                blocked,
664            }
665        }
666    }
667
668    #[async_trait]
669    impl Connector for TestConnector {
670        fn name(&self) -> &str {
671            "test"
672        }
673
674        async fn describe(&self) -> Result<ConnectorDescription, String> {
675            Ok(ConnectorDescription {
676                name: "test".to_string(),
677                instructions: None,
678                tools: vec![ConnectorTool {
679                    name: "run".to_string(),
680                    description: None,
681                    input_schema: json!({"type": "object"}),
682                    output_schema: Some(json!({"type": "integer"})),
683                    instructions: None,
684                    examples: Vec::new(),
685                    annotations: ToolAnnotations {
686                        read_only: Some(true),
687                        ..ToolAnnotations::default()
688                    },
689                    policy: ToolPolicy {
690                        requires_approval: false,
691                        replay: ReplayPolicy::Log,
692                    },
693                }],
694            })
695        }
696
697        async fn execute(
698            &self,
699            _method: &str,
700            _arguments: Value,
701            _context: &ToolContext,
702        ) -> Result<Value, String> {
703            self.calls.fetch_add(1, Ordering::SeqCst);
704            self.started.notify_one();
705            if self.blocked {
706                self.release.notified().await;
707            }
708            Ok(json!(42))
709        }
710
711        async fn execution_ended(&self, _execution_id: &str, status: &str) {
712            self.ended.lock().await.push(status.to_string());
713        }
714    }
715
716    struct HostExecutor;
717
718    #[async_trait(?Send)]
719    impl CodeExecutor for HostExecutor {
720        async fn execute(
721            &self,
722            _code: &str,
723            _connectors: &[ConnectorDescription],
724            _execution_id: &str,
725            host: Arc<ExecutionHost>,
726        ) -> Result<ExecuteResult, String> {
727            let response = host.call(0, "test", "run", json!({})).await;
728            Ok(ExecuteResult {
729                result: response.result,
730                error: response.message,
731                logs: Vec::new(),
732            })
733        }
734    }
735
736    #[tokio::test(flavor = "current_thread")]
737    async fn dispatch_requires_an_active_pass() {
738        let connector = Arc::new(TestConnector::new(false));
739        let code_mode = CodeMode::new(
740            Arc::new(MemoryStore::default()),
741            HostExecutor,
742            vec![connector.clone()],
743        );
744        let state = code_mode.start("ignored").await.unwrap();
745
746        let error = code_mode
747            .dispatch(DispatchRequest::Call {
748                execution_id: state.id.clone(),
749                seq: 0,
750                connector: "test".to_string(),
751                method: "run".to_string(),
752                arguments: json!({}),
753            })
754            .await
755            .unwrap_err();
756
757        assert!(error.contains("active pass"));
758        assert_eq!(connector.calls.load(Ordering::SeqCst), 0);
759        assert_eq!(
760            code_mode.execution(&state.id).await.unwrap().status,
761            ExecutionStatus::Running
762        );
763    }
764
765    #[tokio::test(flavor = "current_thread")]
766    async fn cancelling_an_idle_execution_ends_connector_lifecycle_once() {
767        let connector = Arc::new(TestConnector::new(false));
768        let code_mode = CodeMode::new(
769            Arc::new(MemoryStore::default()),
770            HostExecutor,
771            vec![connector.clone()],
772        );
773        let state = code_mode.start("ignored").await.unwrap();
774
775        code_mode.cancel(&state.id).await.unwrap();
776        code_mode.cancel(&state.id).await.unwrap();
777
778        assert_eq!(&*connector.ended.lock().await, &["cancelled"]);
779    }
780
781    #[tokio::test(flavor = "current_thread")]
782    async fn concurrent_drive_calls_execute_one_pass() {
783        let connector = Arc::new(TestConnector::new(true));
784        let code_mode = CodeMode::new(
785            Arc::new(MemoryStore::default()),
786            HostExecutor,
787            vec![connector.clone()],
788        );
789        let state = code_mode.start("ignored").await.unwrap();
790        let first = code_mode.drive_with(&state.id, CodeModeRunOptions::default());
791        let second = code_mode.drive_with(&state.id, CodeModeRunOptions::default());
792        let release = async {
793            connector.started.notified().await;
794            tokio::task::yield_now().await;
795            connector.release.notify_waiters();
796        };
797
798        let (first, second, ()) = tokio::join!(first, second, release);
799
800        assert_eq!(first.unwrap().status, ExecutionStatus::Completed);
801        assert_eq!(second.unwrap().status, ExecutionStatus::Completed);
802        assert_eq!(connector.calls.load(Ordering::SeqCst), 1);
803    }
804
805    #[tokio::test(flavor = "current_thread")]
806    async fn rollback_rejects_live_executions() {
807        let code_mode = CodeMode::new(
808            Arc::new(MemoryStore::default()),
809            HostExecutor,
810            vec![Arc::new(TestConnector::new(false))],
811        );
812        let state = code_mode.start("ignored").await.unwrap();
813
814        let error = code_mode.rollback(&state.id).await.unwrap_err();
815
816        assert!(error.contains("not rollback eligible"));
817        assert_eq!(
818            code_mode.execution(&state.id).await.unwrap().status,
819            ExecutionStatus::Running
820        );
821    }
822}