mirror of
https://github.com/openai/codex.git
synced 2026-10-20 17:38:45 +00:00
Add process-owned code mode session
This commit is contained in:
Channing Conger
@@ -1,9 +1,12 @@
|
||||
mod cell_actor;
|
||||
mod remote_session;
|
||||
mod runtime;
|
||||
mod service;
|
||||
mod session_runtime;
|
||||
|
||||
pub use codex_code_mode_protocol::*;
|
||||
pub use remote_session::ProcessOwnedCodeModeSession;
|
||||
pub use remote_session::ProcessOwnedCodeModeSessionProvider;
|
||||
pub use service::InProcessCodeModeSession;
|
||||
pub use service::InProcessCodeModeSessionProvider;
|
||||
pub use service::NoopCodeModeSessionDelegate;
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex as StdMutex;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use codex_code_mode_protocol::CellId;
|
||||
use codex_code_mode_protocol::CodeModeSession;
|
||||
use codex_code_mode_protocol::CodeModeSessionDelegate;
|
||||
use codex_code_mode_protocol::CodeModeSessionProvider;
|
||||
use codex_code_mode_protocol::CodeModeSessionProviderFuture;
|
||||
use codex_code_mode_protocol::CodeModeSessionResultFuture;
|
||||
use codex_code_mode_protocol::ExecuteRequest;
|
||||
use codex_code_mode_protocol::StartedCell;
|
||||
use codex_code_mode_protocol::WaitOutcome;
|
||||
use codex_code_mode_protocol::WaitRequest;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use self::connection::Connection;
|
||||
use crate::NoopCodeModeSessionDelegate;
|
||||
|
||||
mod connection;
|
||||
|
||||
const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH";
|
||||
|
||||
/// Creates code-mode sessions backed by one lazily spawned process host.
|
||||
pub struct ProcessOwnedCodeModeSessionProvider {
|
||||
host_program: PathBuf,
|
||||
process_host: StdMutex<Option<Arc<OwnedProcessHost>>>,
|
||||
}
|
||||
|
||||
impl ProcessOwnedCodeModeSessionProvider {
|
||||
pub fn with_host_program(host_program: PathBuf) -> Self {
|
||||
Self {
|
||||
host_program,
|
||||
process_host: StdMutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn process_host(&self) -> Arc<OwnedProcessHost> {
|
||||
let mut process_host = self
|
||||
.process_host
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if let Some(process_host) = process_host.as_ref() {
|
||||
return Arc::clone(process_host);
|
||||
}
|
||||
|
||||
let new_process_host = Arc::new(OwnedProcessHost::new(self.host_program.clone()));
|
||||
*process_host = Some(Arc::clone(&new_process_host));
|
||||
new_process_host
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProcessOwnedCodeModeSessionProvider {
|
||||
fn default() -> Self {
|
||||
Self::with_host_program(default_host_program())
|
||||
}
|
||||
}
|
||||
|
||||
impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
|
||||
fn create_session<'a>(
|
||||
&'a self,
|
||||
delegate: Arc<dyn CodeModeSessionDelegate>,
|
||||
) -> CodeModeSessionProviderFuture<'a> {
|
||||
let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host());
|
||||
Box::pin(async move {
|
||||
session.connection().await?;
|
||||
let session: Arc<dyn CodeModeSession> = Arc::new(session);
|
||||
Ok(session)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct OwnedProcessHost {
|
||||
host_program: PathBuf,
|
||||
connection: StdMutex<Option<Arc<Connection>>>,
|
||||
spawn_permit: Semaphore,
|
||||
next_session_id: AtomicU64,
|
||||
}
|
||||
|
||||
impl OwnedProcessHost {
|
||||
fn new(host_program: PathBuf) -> Self {
|
||||
Self {
|
||||
host_program,
|
||||
connection: StdMutex::new(None),
|
||||
spawn_permit: Semaphore::new(/*permits*/ 1),
|
||||
next_session_id: AtomicU64::new(1),
|
||||
}
|
||||
}
|
||||
|
||||
async fn connection(&self) -> Result<Arc<Connection>, String> {
|
||||
if let Some(connection) = self.live_connection() {
|
||||
return Ok(connection);
|
||||
}
|
||||
|
||||
let _spawn_permit = self
|
||||
.spawn_permit
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| "code-mode host spawn coordinator closed".to_string())?;
|
||||
if let Some(connection) = self.live_connection() {
|
||||
return Ok(connection);
|
||||
}
|
||||
let new_connection = Arc::new(Connection::spawn(&self.host_program).await?);
|
||||
*self
|
||||
.connection
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::clone(&new_connection));
|
||||
Ok(new_connection)
|
||||
}
|
||||
|
||||
fn live_connection(&self) -> Option<Arc<Connection>> {
|
||||
self.connection
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.as_ref()
|
||||
.filter(|connection| connection.is_alive())
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn allocate_session_id(&self) -> SessionId {
|
||||
let value = self.next_session_id.fetch_add(1, Ordering::Relaxed);
|
||||
match SessionId::new(format!("session-{value}")) {
|
||||
Ok(session_id) => session_id,
|
||||
Err(_) => unreachable!("a generated code-mode session ID is nonempty"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum SessionState {
|
||||
New,
|
||||
Open(Arc<Connection>),
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// A logical code-mode session assigned to a process-owned host.
|
||||
pub struct ProcessOwnedCodeModeSession {
|
||||
process_host: Arc<OwnedProcessHost>,
|
||||
session_id: SessionId,
|
||||
delegate: Arc<dyn CodeModeSessionDelegate>,
|
||||
state: StdMutex<SessionState>,
|
||||
transition_permit: Semaphore,
|
||||
}
|
||||
|
||||
impl ProcessOwnedCodeModeSession {
|
||||
pub fn new() -> Self {
|
||||
Self::with_process_host(
|
||||
Arc::new(NoopCodeModeSessionDelegate),
|
||||
Arc::new(OwnedProcessHost::new(default_host_program())),
|
||||
)
|
||||
}
|
||||
|
||||
fn with_process_host(
|
||||
delegate: Arc<dyn CodeModeSessionDelegate>,
|
||||
process_host: Arc<OwnedProcessHost>,
|
||||
) -> Self {
|
||||
let session_id = process_host.allocate_session_id();
|
||||
Self {
|
||||
process_host,
|
||||
session_id,
|
||||
delegate,
|
||||
state: StdMutex::new(SessionState::New),
|
||||
transition_permit: Semaphore::new(/*permits*/ 1),
|
||||
}
|
||||
}
|
||||
|
||||
async fn connection(&self) -> Result<Arc<Connection>, String> {
|
||||
if let Some(connection) = self.current_connection()? {
|
||||
return Ok(connection);
|
||||
}
|
||||
|
||||
let _transition_permit = self
|
||||
.transition_permit
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| "code-mode session transition coordinator closed".to_string())?;
|
||||
if let Some(connection) = self.current_connection()? {
|
||||
return Ok(connection);
|
||||
}
|
||||
let connection = self.process_host.connection().await?;
|
||||
connection
|
||||
.open_session(self.session_id.clone(), Arc::clone(&self.delegate))
|
||||
.await?;
|
||||
*self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner) =
|
||||
SessionState::Open(Arc::clone(&connection));
|
||||
Ok(connection)
|
||||
}
|
||||
|
||||
fn current_connection(&self) -> Result<Option<Arc<Connection>>, String> {
|
||||
match &*self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
{
|
||||
SessionState::New => Ok(None),
|
||||
SessionState::Open(connection) => Ok(Some(Arc::clone(connection))),
|
||||
SessionState::Shutdown => Err("code mode session is shutting down".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute(&self, request: ExecuteRequest) -> Result<StartedCell, String> {
|
||||
self.connection()
|
||||
.await?
|
||||
.execute(self.session_id.clone(), request)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
|
||||
self.connection()
|
||||
.await?
|
||||
.wait(self.session_id.clone(), request)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
|
||||
self.connection()
|
||||
.await?
|
||||
.terminate(self.session_id.clone(), cell_id)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn shutdown(&self) -> Result<(), String> {
|
||||
let transition_permit = self
|
||||
.transition_permit
|
||||
.acquire()
|
||||
.await
|
||||
.map_err(|_| "code-mode session transition coordinator closed".to_string())?;
|
||||
let connection = {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
match std::mem::replace(&mut *state, SessionState::Shutdown) {
|
||||
SessionState::Open(connection) => connection,
|
||||
SessionState::New | SessionState::Shutdown => return Ok(()),
|
||||
}
|
||||
};
|
||||
drop(transition_permit);
|
||||
connection.shutdown_session(self.session_id.clone()).await
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ProcessOwnedCodeModeSession {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl CodeModeSession for ProcessOwnedCodeModeSession {
|
||||
fn execute<'a>(
|
||||
&'a self,
|
||||
request: ExecuteRequest,
|
||||
) -> CodeModeSessionResultFuture<'a, StartedCell> {
|
||||
Box::pin(ProcessOwnedCodeModeSession::execute(self, request))
|
||||
}
|
||||
|
||||
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
|
||||
Box::pin(ProcessOwnedCodeModeSession::wait(self, request))
|
||||
}
|
||||
|
||||
fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
|
||||
Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id))
|
||||
}
|
||||
|
||||
fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
|
||||
Box::pin(ProcessOwnedCodeModeSession::shutdown(self))
|
||||
}
|
||||
}
|
||||
|
||||
fn default_host_program() -> PathBuf {
|
||||
if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) {
|
||||
return PathBuf::from(path);
|
||||
}
|
||||
let executable_name = if cfg!(windows) {
|
||||
"codex-code-mode-host.exe"
|
||||
} else {
|
||||
"codex-code-mode-host"
|
||||
};
|
||||
if let Ok(current_exe) = std::env::current_exe()
|
||||
&& let Some(parent) = current_exe.parent()
|
||||
{
|
||||
let sibling = parent.join(executable_name);
|
||||
if sibling.is_file() {
|
||||
return sibling;
|
||||
}
|
||||
}
|
||||
PathBuf::from(executable_name)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "remote_session_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,405 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::AtomicI64;
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
use codex_code_mode_protocol::CellId;
|
||||
use codex_code_mode_protocol::CodeModeSessionDelegate;
|
||||
use codex_code_mode_protocol::ExecuteRequest;
|
||||
use codex_code_mode_protocol::RuntimeResponse;
|
||||
use codex_code_mode_protocol::StartedCell;
|
||||
use codex_code_mode_protocol::WaitOutcome;
|
||||
use codex_code_mode_protocol::WaitRequest;
|
||||
use codex_code_mode_protocol::host::CapabilitySet;
|
||||
use codex_code_mode_protocol::host::ClientHello;
|
||||
use codex_code_mode_protocol::host::ClientToHost;
|
||||
use codex_code_mode_protocol::host::DelegateRequestId;
|
||||
use codex_code_mode_protocol::host::FramedReader;
|
||||
use codex_code_mode_protocol::host::FramedWriter;
|
||||
use codex_code_mode_protocol::host::HostRequest;
|
||||
use codex_code_mode_protocol::host::HostResponse;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::ProtocolVersion;
|
||||
use codex_code_mode_protocol::host::RequestId;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use codex_code_mode_protocol::host::SupportedProtocolVersions;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::debug;
|
||||
use tracing::warn;
|
||||
|
||||
use self::reader::drive_reader;
|
||||
|
||||
mod reader;
|
||||
|
||||
const IPC_CHANNEL_CAPACITY: usize = 128;
|
||||
|
||||
pub(super) struct Connection {
|
||||
state: Arc<ConnectionState>,
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
impl Connection {
|
||||
pub(super) async fn spawn(host_program: &Path) -> Result<Self, String> {
|
||||
let mut command = Command::new(host_program);
|
||||
#[cfg(unix)]
|
||||
command.process_group(0);
|
||||
let mut child = command
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true)
|
||||
.spawn()
|
||||
.map_err(|err| {
|
||||
format!(
|
||||
"failed to spawn code-mode host {}: {err}",
|
||||
host_program.display()
|
||||
)
|
||||
})?;
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| "spawned code-mode host has no stdin".to_string())?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| "spawned code-mode host has no stdout".to_string())?;
|
||||
let stderr = child.stderr.take();
|
||||
let mut reader = FramedReader::new(stdout);
|
||||
let mut writer = FramedWriter::new(stdin);
|
||||
|
||||
let hello = ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])
|
||||
.map_err(|err| err.to_string())?,
|
||||
CapabilitySet::empty(),
|
||||
CapabilitySet::empty(),
|
||||
)
|
||||
.map_err(|err| err.to_string())?;
|
||||
writer
|
||||
.write(&ClientToHost::ClientHello(hello))
|
||||
.await
|
||||
.map_err(|err| format!("failed to write code-mode host hello: {err}"))?;
|
||||
match reader
|
||||
.read::<HostToClient>()
|
||||
.await
|
||||
.map_err(|err| format!("failed to read code-mode host hello: {err}"))?
|
||||
{
|
||||
Some(HostToClient::HostHello(hello))
|
||||
if hello.selected_version() == ProtocolVersion::V1 => {}
|
||||
Some(HostToClient::HandshakeRejected { reason }) => {
|
||||
return Err(format!("code-mode host rejected the handshake: {reason:?}"));
|
||||
}
|
||||
Some(message) => {
|
||||
return Err(format!(
|
||||
"code-mode host returned an invalid handshake response: {message:?}"
|
||||
));
|
||||
}
|
||||
None => return Err("code-mode host exited during handshake".to_string()),
|
||||
}
|
||||
|
||||
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
|
||||
let cancellation = CancellationToken::new();
|
||||
let state = Arc::new(ConnectionState::new(outgoing_tx, cancellation.clone()));
|
||||
|
||||
let writer_state = Arc::downgrade(&state);
|
||||
let writer_cancellation = cancellation.clone();
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = writer_cancellation.cancelled() => break,
|
||||
message = outgoing_rx.recv() => {
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
if let Err(err) = writer.write(&message).await {
|
||||
if let Some(state) = writer_state.upgrade() {
|
||||
state.fail(format!(
|
||||
"failed to write code-mode host message: {err}"
|
||||
)).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let reader_state = Arc::downgrade(&state);
|
||||
let reader_cancellation = cancellation.clone();
|
||||
tokio::spawn(async move {
|
||||
drive_reader(reader, reader_state, reader_cancellation).await;
|
||||
});
|
||||
|
||||
if let Some(stderr) = stderr {
|
||||
tokio::spawn(async move {
|
||||
let mut lines = BufReader::new(stderr).lines();
|
||||
loop {
|
||||
match lines.next_line().await {
|
||||
Ok(Some(line)) => debug!("code-mode host stderr: {line}"),
|
||||
Ok(None) => break,
|
||||
Err(err) => {
|
||||
warn!("failed to read code-mode host stderr: {err}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let supervisor_state = Arc::downgrade(&state);
|
||||
let supervisor_cancellation = cancellation.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::select! {
|
||||
result = child.wait() => {
|
||||
let reason = match result {
|
||||
Ok(status) => format!("code-mode host exited with status {status}"),
|
||||
Err(err) => format!("failed waiting for code-mode host: {err}"),
|
||||
};
|
||||
if let Some(state) = supervisor_state.upgrade() {
|
||||
state.fail(reason).await;
|
||||
}
|
||||
}
|
||||
_ = supervisor_cancellation.cancelled() => {
|
||||
let _ = child.start_kill();
|
||||
let _ = child.wait().await;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
state,
|
||||
cancellation,
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn is_alive(&self) -> bool {
|
||||
self.state.alive.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub(super) async fn open_session(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
delegate: Arc<dyn CodeModeSessionDelegate>,
|
||||
) -> Result<(), String> {
|
||||
self.state
|
||||
.delegates
|
||||
.lock()
|
||||
.await
|
||||
.insert(session_id.clone(), delegate);
|
||||
let response = self
|
||||
.request(HostRequest::OpenSession {
|
||||
session_id: session_id.clone(),
|
||||
})
|
||||
.await;
|
||||
match response {
|
||||
Ok(HostResponse::SessionReady {
|
||||
session_id: ready_session_id,
|
||||
}) if ready_session_id == session_id => Ok(()),
|
||||
Ok(_) => {
|
||||
self.state.delegates.lock().await.remove(&session_id);
|
||||
Err("code-mode host returned an invalid open-session response".to_string())
|
||||
}
|
||||
Err(err) => {
|
||||
self.state.delegates.lock().await.remove(&session_id);
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn execute(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
request: ExecuteRequest,
|
||||
) -> Result<StartedCell, String> {
|
||||
let request = request
|
||||
.try_into()
|
||||
.map_err(|err| format!("failed to encode code-mode execute request: {err}"))?;
|
||||
let id = RequestId::new(self.state.next_request_id.fetch_add(1, Ordering::Relaxed));
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
let (initial_tx, initial_rx) = oneshot::channel();
|
||||
self.state.pending.lock().await.insert(id, response_tx);
|
||||
self.state
|
||||
.initial_responses
|
||||
.lock()
|
||||
.await
|
||||
.insert(id, initial_tx);
|
||||
if let Err(err) = self
|
||||
.state
|
||||
.send(ClientToHost::Request {
|
||||
id,
|
||||
request: HostRequest::Execute {
|
||||
session_id,
|
||||
request,
|
||||
},
|
||||
})
|
||||
.await
|
||||
{
|
||||
self.state.pending.lock().await.remove(&id);
|
||||
self.state.initial_responses.lock().await.remove(&id);
|
||||
return Err(err);
|
||||
}
|
||||
let response = match response_rx.await {
|
||||
Ok(Ok(response)) => response,
|
||||
Ok(Err(err)) => {
|
||||
self.state.initial_responses.lock().await.remove(&id);
|
||||
return Err(err);
|
||||
}
|
||||
Err(_) => {
|
||||
self.state.initial_responses.lock().await.remove(&id);
|
||||
return Err(self.state.failure_message());
|
||||
}
|
||||
};
|
||||
match response {
|
||||
HostResponse::ExecutionStarted { cell_id } => Ok(StartedCell::from_result_receiver(
|
||||
cell_id.into(),
|
||||
initial_rx,
|
||||
)),
|
||||
_ => {
|
||||
self.state.initial_responses.lock().await.remove(&id);
|
||||
Err("code-mode host returned an invalid execute response".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn wait(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
request: WaitRequest,
|
||||
) -> Result<WaitOutcome, String> {
|
||||
match self
|
||||
.request(HostRequest::Wait {
|
||||
session_id,
|
||||
request: request.into(),
|
||||
})
|
||||
.await?
|
||||
{
|
||||
HostResponse::WaitCompleted { outcome } => Ok(outcome.into()),
|
||||
_ => Err("code-mode host returned an invalid wait response".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn terminate(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
cell_id: CellId,
|
||||
) -> Result<WaitOutcome, String> {
|
||||
match self
|
||||
.request(HostRequest::Terminate {
|
||||
session_id,
|
||||
cell_id: cell_id.into(),
|
||||
})
|
||||
.await?
|
||||
{
|
||||
HostResponse::WaitCompleted { outcome } => Ok(outcome.into()),
|
||||
_ => Err("code-mode host returned an invalid terminate response".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn shutdown_session(&self, session_id: SessionId) -> Result<(), String> {
|
||||
let response = self
|
||||
.request(HostRequest::ShutdownSession {
|
||||
session_id: session_id.clone(),
|
||||
})
|
||||
.await;
|
||||
self.state.delegates.lock().await.remove(&session_id);
|
||||
match response? {
|
||||
HostResponse::SessionClosed {
|
||||
session_id: closed_session_id,
|
||||
} if closed_session_id == session_id => Ok(()),
|
||||
_ => Err("code-mode host returned an invalid shutdown response".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn request(&self, request: HostRequest) -> Result<HostResponse, String> {
|
||||
let id = RequestId::new(self.state.next_request_id.fetch_add(1, Ordering::Relaxed));
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
self.state.pending.lock().await.insert(id, response_tx);
|
||||
if let Err(err) = self.state.send(ClientToHost::Request { id, request }).await {
|
||||
self.state.pending.lock().await.remove(&id);
|
||||
return Err(err);
|
||||
}
|
||||
response_rx
|
||||
.await
|
||||
.map_err(|_| self.state.failure_message())?
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Connection {
|
||||
fn drop(&mut self) {
|
||||
self.cancellation.cancel();
|
||||
}
|
||||
}
|
||||
|
||||
struct ConnectionState {
|
||||
outgoing_tx: mpsc::Sender<ClientToHost>,
|
||||
pending: Mutex<HashMap<RequestId, oneshot::Sender<Result<HostResponse, String>>>>,
|
||||
initial_responses: Mutex<HashMap<RequestId, oneshot::Sender<Result<RuntimeResponse, String>>>>,
|
||||
delegates: Mutex<HashMap<SessionId, Arc<dyn CodeModeSessionDelegate>>>,
|
||||
delegate_cancellations: Mutex<HashMap<DelegateRequestId, CancellationToken>>,
|
||||
next_request_id: AtomicI64,
|
||||
alive: AtomicBool,
|
||||
failure: std::sync::Mutex<Option<String>>,
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
impl ConnectionState {
|
||||
fn new(outgoing_tx: mpsc::Sender<ClientToHost>, cancellation: CancellationToken) -> Self {
|
||||
Self {
|
||||
outgoing_tx,
|
||||
pending: Mutex::new(HashMap::new()),
|
||||
initial_responses: Mutex::new(HashMap::new()),
|
||||
delegates: Mutex::new(HashMap::new()),
|
||||
delegate_cancellations: Mutex::new(HashMap::new()),
|
||||
next_request_id: AtomicI64::new(1),
|
||||
alive: AtomicBool::new(true),
|
||||
failure: std::sync::Mutex::new(None),
|
||||
cancellation,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(&self, message: ClientToHost) -> Result<(), String> {
|
||||
if !self.alive.load(Ordering::Acquire) {
|
||||
return Err(self.failure_message());
|
||||
}
|
||||
self.outgoing_tx
|
||||
.send(message)
|
||||
.await
|
||||
.map_err(|_| self.failure_message())
|
||||
}
|
||||
|
||||
fn failure_message(&self) -> String {
|
||||
self.failure
|
||||
.lock()
|
||||
.ok()
|
||||
.and_then(|failure| failure.clone())
|
||||
.unwrap_or_else(|| "code-mode host connection closed".to_string())
|
||||
}
|
||||
|
||||
async fn fail(&self, reason: String) {
|
||||
if !self.alive.swap(false, Ordering::AcqRel) {
|
||||
return;
|
||||
}
|
||||
if let Ok(mut failure) = self.failure.lock() {
|
||||
*failure = Some(reason.clone());
|
||||
}
|
||||
self.cancellation.cancel();
|
||||
for (_, sender) in self.pending.lock().await.drain() {
|
||||
let _ = sender.send(Err(reason.clone()));
|
||||
}
|
||||
for (_, sender) in self.initial_responses.lock().await.drain() {
|
||||
let _ = sender.send(Err(reason.clone()));
|
||||
}
|
||||
self.delegates.lock().await.clear();
|
||||
for (_, cancellation) in self.delegate_cancellations.lock().await.drain() {
|
||||
cancellation.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::Weak;
|
||||
|
||||
use codex_code_mode_protocol::host::ClientToHost;
|
||||
use codex_code_mode_protocol::host::DelegateRequest;
|
||||
use codex_code_mode_protocol::host::DelegateResponse;
|
||||
use codex_code_mode_protocol::host::FramedReader;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::WireResult;
|
||||
use tokio::process::ChildStdout;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use super::ConnectionState;
|
||||
|
||||
pub(super) async fn drive_reader(
|
||||
mut reader: FramedReader<ChildStdout>,
|
||||
state: Weak<ConnectionState>,
|
||||
cancellation: CancellationToken,
|
||||
) {
|
||||
loop {
|
||||
let message = tokio::select! {
|
||||
_ = cancellation.cancelled() => return,
|
||||
result = reader.read::<HostToClient>() => result,
|
||||
};
|
||||
match message {
|
||||
Ok(Some(message)) => {
|
||||
let Some(state) = state.upgrade() else {
|
||||
return;
|
||||
};
|
||||
handle_host_message(state, message).await;
|
||||
}
|
||||
Ok(None) => {
|
||||
if let Some(state) = state.upgrade() {
|
||||
state
|
||||
.fail("code-mode host closed its stdout".to_string())
|
||||
.await;
|
||||
}
|
||||
return;
|
||||
}
|
||||
Err(err) => {
|
||||
if let Some(state) = state.upgrade() {
|
||||
state
|
||||
.fail(format!("failed to read code-mode host message: {err}"))
|
||||
.await;
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_host_message(state: Arc<ConnectionState>, message: HostToClient) {
|
||||
match message {
|
||||
HostToClient::Response { id, result } => {
|
||||
if let Some(sender) = state.pending.lock().await.remove(&id) {
|
||||
let _ = sender.send(result.into_result());
|
||||
}
|
||||
}
|
||||
HostToClient::InitialResponse { id, result } => {
|
||||
if let Some(sender) = state.initial_responses.lock().await.remove(&id) {
|
||||
let _ = sender.send(result.into_result().map(Into::into));
|
||||
}
|
||||
}
|
||||
HostToClient::DelegateRequest {
|
||||
id,
|
||||
session_id,
|
||||
request,
|
||||
} => {
|
||||
let delegate = state.delegates.lock().await.get(&session_id).cloned();
|
||||
let Some(delegate) = delegate else {
|
||||
let _ = state
|
||||
.send(ClientToHost::DelegateResponse {
|
||||
id,
|
||||
result: WireResult::Err {
|
||||
message: format!("unknown code-mode session {session_id}"),
|
||||
},
|
||||
})
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
let cancellation = CancellationToken::new();
|
||||
state
|
||||
.delegate_cancellations
|
||||
.lock()
|
||||
.await
|
||||
.insert(id, cancellation.clone());
|
||||
tokio::spawn(async move {
|
||||
let result = match request {
|
||||
DelegateRequest::InvokeTool { invocation } => delegate
|
||||
.invoke_tool(invocation.into(), cancellation)
|
||||
.await
|
||||
.map(|result| DelegateResponse::ToolResult { result }),
|
||||
DelegateRequest::Notify {
|
||||
call_id,
|
||||
cell_id,
|
||||
text,
|
||||
} => delegate
|
||||
.notify(call_id, cell_id.into(), text, cancellation)
|
||||
.await
|
||||
.map(|()| DelegateResponse::NotificationDelivered),
|
||||
};
|
||||
state.delegate_cancellations.lock().await.remove(&id);
|
||||
let _ = state
|
||||
.send(ClientToHost::DelegateResponse {
|
||||
id,
|
||||
result: WireResult::from_result(result),
|
||||
})
|
||||
.await;
|
||||
});
|
||||
}
|
||||
HostToClient::CancelDelegateRequest { id } => {
|
||||
if let Some(cancellation) = state.delegate_cancellations.lock().await.remove(&id) {
|
||||
cancellation.cancel();
|
||||
}
|
||||
}
|
||||
HostToClient::CellClosed {
|
||||
session_id,
|
||||
cell_id,
|
||||
} => {
|
||||
if let Some(delegate) = state.delegates.lock().await.get(&session_id).cloned() {
|
||||
let cell_id = cell_id.into();
|
||||
delegate.cell_closed(&cell_id);
|
||||
}
|
||||
}
|
||||
HostToClient::HostHello(_) | HostToClient::HandshakeRejected { .. } => {
|
||||
state
|
||||
.fail("code-mode host sent a second handshake response".to_string())
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_code_mode_protocol::CodeModeSessionProvider;
|
||||
|
||||
use super::ProcessOwnedCodeModeSession;
|
||||
use super::ProcessOwnedCodeModeSessionProvider;
|
||||
use crate::NoopCodeModeSessionDelegate;
|
||||
|
||||
#[test]
|
||||
fn provider_reuses_its_live_process_host() {
|
||||
let provider = ProcessOwnedCodeModeSessionProvider::default();
|
||||
|
||||
let first = provider.process_host();
|
||||
let second = provider.process_host();
|
||||
|
||||
assert!(Arc::ptr_eq(&first, &second));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_reports_host_spawn_failure() {
|
||||
let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
|
||||
"codex-code-mode-host-does-not-exist".into(),
|
||||
);
|
||||
|
||||
let error = provider
|
||||
.create_session(Arc::new(NoopCodeModeSessionDelegate))
|
||||
.await
|
||||
.err()
|
||||
.expect("session creation should fail");
|
||||
|
||||
assert!(error.contains("failed to spawn code-mode host"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_before_open_does_not_spawn_the_host() {
|
||||
let session = ProcessOwnedCodeModeSession::new();
|
||||
|
||||
session.shutdown().await.expect("shutdown session");
|
||||
let error = session
|
||||
.execute(codex_code_mode_protocol::ExecuteRequest {
|
||||
tool_call_id: "call-1".to_string(),
|
||||
enabled_tools: Vec::new(),
|
||||
source: "text('unreachable')".to_string(),
|
||||
yield_time_ms: None,
|
||||
max_output_tokens: None,
|
||||
})
|
||||
.await
|
||||
.err()
|
||||
.expect("shutdown session should reject execution");
|
||||
|
||||
assert_eq!(error, "code mode session is shutting down");
|
||||
}
|
||||
Reference in New Issue
Block a user