Skip to main content

codex_code_mode_protocol/
session.rs

1use std::fmt;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5
6use serde::Deserialize;
7use serde::Serialize;
8use serde_json::Value as JsonValue;
9use tokio::sync::oneshot;
10use tokio_util::sync::CancellationToken;
11
12use crate::CodeModeNestedToolCall;
13use crate::ExecuteRequest;
14use crate::RuntimeResponse;
15use crate::WaitOutcome;
16use crate::WaitRequest;
17
18pub type CodeModeSessionResultFuture<'a, T> =
19    Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>;
20pub type CodeModeSessionProviderFuture<'a> =
21    CodeModeSessionResultFuture<'a, Arc<dyn CodeModeSession>>;
22pub type ToolInvocationFuture<'a> =
23    Pin<Box<dyn Future<Output = Result<JsonValue, String>> + Send + 'a>>;
24pub type NotificationFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;
25
26#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
27pub struct CellId(String);
28
29impl CellId {
30    pub fn new(value: String) -> Self {
31        Self(value)
32    }
33
34    pub fn as_str(&self) -> &str {
35        &self.0
36    }
37}
38
39impl AsRef<str> for CellId {
40    fn as_ref(&self) -> &str {
41        self.as_str()
42    }
43}
44
45impl fmt::Display for CellId {
46    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47        formatter.write_str(self.as_str())
48    }
49}
50
51pub struct StartedCell {
52    pub cell_id: CellId,
53    initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>,
54}
55
56impl StartedCell {
57    pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver<RuntimeResponse>) -> Self {
58        Self {
59            cell_id,
60            initial_response: Box::pin(async move {
61                initial_response_rx
62                    .await
63                    .map_err(|_| "exec runtime ended unexpectedly".to_string())
64            }),
65        }
66    }
67
68    pub fn from_result_receiver(
69        cell_id: CellId,
70        initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
71    ) -> Self {
72        Self {
73            cell_id,
74            initial_response: Box::pin(async move {
75                initial_response_rx
76                    .await
77                    .map_err(|_| "exec runtime ended unexpectedly".to_string())?
78            }),
79        }
80    }
81
82    pub async fn initial_response(self) -> Result<RuntimeResponse, String> {
83        self.initial_response.await
84    }
85}
86
87/// Host callbacks used by a code-mode session while cells are executing.
88pub trait CodeModeSessionDelegate: Send + Sync {
89    fn invoke_tool<'a>(
90        &'a self,
91        invocation: CodeModeNestedToolCall,
92        cancellation_token: CancellationToken,
93    ) -> ToolInvocationFuture<'a>;
94
95    fn notify<'a>(
96        &'a self,
97        call_id: String,
98        cell_id: CellId,
99        text: String,
100        cancellation_token: CancellationToken,
101    ) -> NotificationFuture<'a>;
102
103    /// Releases delegate state associated with a cell after it reaches a terminal state.
104    fn cell_closed(&self, cell_id: &CellId);
105}
106
107/// A durable code-mode session owned by one Codex thread.
108///
109/// Cells executed in the same session share stored values. Separate sessions
110/// must keep those values isolated. Implementations may execute cells
111/// in-process or remotely.
112pub trait CodeModeSession: Send + Sync {
113    fn execute<'a>(
114        &'a self,
115        request: ExecuteRequest,
116    ) -> CodeModeSessionResultFuture<'a, StartedCell>;
117
118    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
119
120    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
121
122    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
123}
124
125/// Creates code-mode sessions for Codex threads.
126///
127/// Implementations may share a remote host process across all sessions created
128/// by one provider.
129pub trait CodeModeSessionProvider: Send + Sync {
130    fn create_session<'a>(
131        &'a self,
132        delegate: Arc<dyn CodeModeSessionDelegate>,
133    ) -> CodeModeSessionProviderFuture<'a>;
134}
135
136#[cfg(test)]
137#[path = "session_tests.rs"]
138mod tests;