File size: 6,199 Bytes
ea39c0e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use serde::Deserialize;
use serde::Serialize;
use serde_json::Value as JsonValue;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use crate::CodeModeNestedToolCall;
use crate::ExecuteRequest;
use crate::RuntimeResponse;
use crate::WaitOutcome;
use crate::WaitRequest;
pub type CodeModeSessionResultFuture<'a, T> =
Pin<Box<dyn Future<Output = Result<T, String>> + Send + 'a>>;
pub type CodeModeSessionProviderFuture<'a> =
CodeModeSessionResultFuture<'a, Arc<dyn CodeModeSession>>;
pub type ToolInvocationFuture<'a> =
Pin<Box<dyn Future<Output = Result<JsonValue, String>> + Send + 'a>>;
pub type NotificationFuture<'a> = Pin<Box<dyn Future<Output = Result<(), String>> + Send + 'a>>;
/// Optional resource limits shared by every cell in one code-mode session.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct CodeModeSessionCellExecutionLimits {
pub max_yield_time_ms: Option<u64>,
pub max_heap_size_bytes: Option<usize>,
}
#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)]
pub struct CellId(String);
impl CellId {
pub fn new(value: String) -> Self {
Self(value)
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl AsRef<str> for CellId {
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl fmt::Display for CellId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(self.as_str())
}
}
pub struct StartedCell {
pub cell_id: CellId,
initial_response: CodeModeSessionResultFuture<'static, RuntimeResponse>,
}
impl StartedCell {
pub fn new(cell_id: CellId, initial_response_rx: oneshot::Receiver<RuntimeResponse>) -> Self {
Self::from_future(cell_id, async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())
})
}
pub fn from_result_receiver(
cell_id: CellId,
initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
) -> Self {
Self::from_future(cell_id, async move {
initial_response_rx
.await
.map_err(|_| "exec runtime ended unexpectedly".to_string())?
})
}
pub fn from_future(
cell_id: CellId,
initial_response: impl Future<Output = Result<RuntimeResponse, String>> + Send + 'static,
) -> Self {
Self {
cell_id,
initial_response: Box::pin(initial_response),
}
}
pub async fn initial_response(self) -> Result<RuntimeResponse, String> {
self.initial_response.await
}
}
/// Host callbacks owned by one code-mode execution.
///
/// The session retains the supplied delegate while starting and running the cell,
/// including across yields, and releases it through its existing close/cancel paths.
pub trait CodeModeSessionDelegate: Send + Sync {
fn invoke_tool<'a>(
&'a self,
invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a>;
fn notify<'a>(
&'a self,
call_id: String,
cell_id: CellId,
text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a>;
/// Releases delegate state associated with a cell after it reaches a terminal state.
fn cell_closed(&self, cell_id: &CellId);
}
/// A session delegate for clients that do not expose nested tools or notifications.
pub struct NoopCodeModeSessionDelegate;
impl CodeModeSessionDelegate for NoopCodeModeSessionDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
cancellation_token.cancelled().await;
Err("code mode nested tools are unavailable".to_string())
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
/// A durable code-mode session owned by one Codex thread.
///
/// Cells executed in the same session share stored values. Separate sessions
/// must keep those values isolated. Implementations may execute cells
/// in-process or remotely.
pub trait CodeModeSession: Send + Sync {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionResultFuture<'a, StartedCell>;
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome>;
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()>;
}
/// Creates code-mode sessions for Codex threads.
///
/// Implementations may share a remote host process across all sessions created
/// by one provider.
pub trait CodeModeSessionProvider: Send + Sync {
/// Reports whether this provider can execute code without starting its host.
fn availability(&self) -> Result<(), String> {
Ok(())
}
fn create_session(&self) -> CodeModeSessionProviderFuture<'_>;
/// Creates a session whose cells share the supplied execution limits.
///
/// Existing providers remain compatible with unlimited sessions, but must
/// explicitly implement this method before accepting non-default limits.
fn create_session_with_limits<'a>(
&'a self,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
if limits == CodeModeSessionCellExecutionLimits::default() {
self.create_session()
} else {
Box::pin(async {
Err("code-mode session provider does not support resource limits".to_string())
})
}
}
}
#[cfg(test)]
#[path = "session_tests.rs"]
mod tests;
|