codex/codex-rs/code-mode/src/grpc_session/reconnect.rs

use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::PoisonError;
use std::sync::atomic::Ordering;

use codex_code_mode_protocol::CodeModeOutputTruncation;
use codex_code_mode_protocol::CodeModeSession;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CodeModeSessionOptions;
use codex_code_mode_protocol::CodeModeSessionResultFuture;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::StartedCell;
use codex_code_mode_protocol::TerminateRequest;
use codex_code_mode_protocol::WaitOutcome;
use codex_code_mode_protocol::WaitRequest;
use tokio::sync::Semaphore;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;

use super::GrpcCodeModeSession;
use super::GrpcCodeModeSessionProvider;
use super::SHUTDOWN_ERROR;
use super::generation;
use super::generation::GenerationDelegate;
use crate::remote_session::ShutdownResultReceiver;
use crate::remote_session::wait_for_watch;

pub(super) struct ReconnectableSession {
    inner: Arc<ReconnectInner>,
}

struct ReconnectInner {
    provider: GrpcCodeModeSessionProvider,
    options: CodeModeSessionOptions,
    binding: Mutex<Option<SessionBinding>>,
    // The negotiated mode belongs to the logical session, including after
    // shutdown releases the binding. Never pin a mode from an unpublished open.
    output_truncation: OnceLock<CodeModeOutputTruncation>,
    opening_permit: Semaphore,
    shutdown_requested: CancellationToken,
    shutdown_result: Mutex<Option<ShutdownResultReceiver>>,
}

#[derive(Clone)]
struct SessionBinding {
    session: Arc<GrpcCodeModeSession>,
    generation: u64,
}

impl ReconnectableSession {
    pub(super) fn new(
        provider: GrpcCodeModeSessionProvider,
        options: CodeModeSessionOptions,
    ) -> Self {
        Self {
            inner: Arc::new(ReconnectInner {
                provider,
                options,
                binding: Mutex::new(None),
                output_truncation: OnceLock::new(),
                opening_permit: Semaphore::new(/*permits*/ 1),
                shutdown_requested: CancellationToken::new(),
                shutdown_result: Mutex::new(None),
            }),
        }
    }

    pub(super) async fn initialize(&self) -> Result<(), String> {
        self.inner.get_or_open_binding().await.map(|_| ())
    }
}

impl CodeModeSession for ReconnectableSession {
    fn output_truncation(&self) -> CodeModeOutputTruncation {
        self.inner
            .output_truncation
            .get()
            .copied()
            .unwrap_or_default()
    }

    fn execute<'a>(
        &'a self,
        request: ExecuteRequest,
        delegate: Arc<dyn CodeModeSessionDelegate>,
        preempt: Option<CancellationToken>,
    ) -> CodeModeSessionResultFuture<'a, StartedCell> {
        // Capture before polling so queued calls cannot silently cross a runtime reset.
        let generation = self.inner.binding().map(|binding| binding.generation);
        Box::pin(async move {
            let binding = self.inner.binding_for_request(generation).await?;
            let delegate = Arc::new(GenerationDelegate {
                delegate,
                generation: binding.generation,
            });
            let started = binding.session.execute(request, delegate, preempt).await?;
            Ok(generation::public_started_cell(binding.generation, started))
        })
    }

    fn wait<'a>(
        &'a self,
        request: WaitRequest,
        preempt: Option<CancellationToken>,
    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
        let generation = self.inner.binding().map(|binding| binding.generation);
        Box::pin(async move {
            let binding = self.inner.binding_for_request(generation).await?;
            let request = WaitRequest {
                max_output_tokens: request.max_output_tokens,
                cell_id: generation::remote_cell_id(binding.generation, &request.cell_id)?,
                yield_time_ms: request.yield_time_ms,
            };
            let outcome = binding.session.wait(request, preempt).await?;
            Ok(generation::public_wait_outcome(binding.generation, outcome))
        })
    }

    fn terminate<'a>(
        &'a self,
        request: TerminateRequest,
    ) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
        Box::pin(async move {
            // Cleanup must not replace the runtime and hide its loss from the next request.
            let binding = self
                .inner
                .binding()
                .ok_or_else(|| SHUTDOWN_ERROR.to_string())?;
            let cell_id = generation::remote_cell_id(binding.generation, &request.cell_id)?;
            let outcome = binding
                .session
                .terminate(TerminateRequest { cell_id, ..request })
                .await?;
            Ok(generation::public_wait_outcome(binding.generation, outcome))
        })
    }

    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
        Box::pin(wait_for_watch(self.inner.request_shutdown()))
    }
}

impl Drop for ReconnectableSession {
    fn drop(&mut self) {
        if tokio::runtime::Handle::try_current().is_ok() {
            self.inner.request_shutdown();
        }
    }
}

impl ReconnectInner {
    async fn binding_for_request(&self, generation: Option<u64>) -> Result<SessionBinding, String> {
        let binding = self.get_or_open_binding().await?;
        if generation != Some(binding.generation) {
            return Err(concat!(
                "The exec runtime was replaced. Stored values and running cells were lost. ",
                "Prior external tool operations may have completed; inspect their results before retrying. ",
                "You may continue with a fresh exec call.",
            ).into());
        }
        Ok(binding)
    }

    async fn get_or_open_binding(&self) -> Result<SessionBinding, String> {
        if self.shutdown_requested.is_cancelled() {
            return Err(SHUTDOWN_ERROR.to_string());
        }
        if let Some(binding) = self.live_binding() {
            return Ok(binding);
        }

        let _opening_permit = tokio::select! {
            biased;
            _ = self.shutdown_requested.cancelled() => {
                return Err(SHUTDOWN_ERROR.to_string());
            }
            permit = self.opening_permit.acquire() => permit
                .map_err(|_| "gRPC code-mode session opening coordinator closed".to_string())?,
        };
        if self.shutdown_requested.is_cancelled() {
            return Err(SHUTDOWN_ERROR.to_string());
        }
        if let Some(binding) = self.live_binding() {
            return Ok(binding);
        }

        if let Some(binding) = self.binding() {
            wait_for_watch(binding.session.inner.request_shutdown()).await?;
        }

        let generation = self
            .provider
            .next_generation
            .fetch_add(1, Ordering::Relaxed);
        let session = tokio::select! {
            biased;
            _ = self.shutdown_requested.cancelled() => {
                return Err(SHUTDOWN_ERROR.to_string());
            }
            session = self.provider.open_binding(&self.options, self.output_truncation.get().copied()) => session?,
        };
        let binding = SessionBinding {
            session,
            generation,
        };
        let published = {
            let mut current = self.binding.lock().unwrap_or_else(PoisonError::into_inner);
            if self.shutdown_requested.is_cancelled() {
                false
            } else {
                *current = Some(binding.clone());
                self.output_truncation
                    .get_or_init(|| binding.session.output_truncation());
                true
            }
        };
        if !published {
            let _ = wait_for_watch(binding.session.inner.request_shutdown()).await;
            return Err(SHUTDOWN_ERROR.to_string());
        }
        Ok(binding)
    }

    fn live_binding(&self) -> Option<SessionBinding> {
        self.binding()
            .filter(|binding| !binding.session.inner.stopped.is_cancelled())
    }

    fn binding(&self) -> Option<SessionBinding> {
        self.binding
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .clone()
    }

    fn request_shutdown(self: &Arc<Self>) -> ShutdownResultReceiver {
        {
            let binding = self.binding.lock().unwrap_or_else(PoisonError::into_inner);
            self.shutdown_requested.cancel();
            if let Some(binding) = binding.as_ref() {
                binding.session.inner.request_shutdown();
            }
        }
        let mut result = self
            .shutdown_result
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        if let Some(receiver) = result.as_ref() {
            return receiver.clone();
        }

        let (sender, receiver) = watch::channel(None);
        *result = Some(receiver.clone());
        let inner = Arc::clone(self);
        tokio::spawn(async move {
            let opening_permit = match inner.opening_permit.acquire().await {
                Ok(permit) => permit,
                Err(_) => {
                    sender.send_replace(Some(Err(
                        "gRPC code-mode session opening coordinator closed".to_string(),
                    )));
                    return;
                }
            };
            let binding = inner
                .binding
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .take();
            drop(opening_permit);
            let result = match binding {
                Some(binding) => wait_for_watch(binding.session.inner.request_shutdown()).await,
                None => Ok(()),
            };
            sender.send_replace(Some(result));
        });
        receiver
    }
}