codex_code_mode_protocol/
session.rs1use 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
87pub 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 fn cell_closed(&self, cell_id: &CellId);
105}
106
107pub 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
125pub 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;