Skip to main content

codex_runtime/runtime/client/
session.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::Arc;
3
4use tokio::sync::Mutex;
5
6use crate::runtime::api::{PromptRunError, PromptRunParams, PromptRunResult};
7use crate::runtime::core::Runtime;
8use crate::runtime::errors::RpcError;
9use crate::runtime::hooks::merge_hook_configs;
10
11use super::profile::{profile_to_prompt_params_with_hooks, session_prompt_params};
12use super::{RunProfile, SessionConfig};
13
14#[derive(Clone)]
15pub struct Session {
16    runtime: Runtime,
17    pub thread_id: String,
18    pub config: SessionConfig,
19    tool_use_loop_started: Arc<AtomicBool>,
20    closed: Arc<AtomicBool>,
21    close_result: Arc<Mutex<Option<Result<(), RpcError>>>>,
22}
23
24impl Session {
25    pub(super) fn new(
26        runtime: Runtime,
27        thread_id: String,
28        config: SessionConfig,
29        tool_use_loop_started: Arc<AtomicBool>,
30    ) -> Self {
31        Self {
32            runtime,
33            thread_id,
34            config,
35            tool_use_loop_started,
36            closed: Arc::new(AtomicBool::new(false)),
37            close_result: Arc::new(Mutex::new(None)),
38        }
39    }
40
41    /// Returns true when this local session handle is closed.
42    /// Allocation: none. Complexity: O(1).
43    pub fn is_closed(&self) -> bool {
44        self.closed.load(Ordering::Acquire)
45    }
46
47    /// Continue this session with one prompt.
48    /// Side effects: sends turn/start RPC calls on one already-loaded thread.
49    /// Allocation: PromptRunParams clone payloads (cwd/model/sandbox/attachments). Complexity: O(n), n = attachment count + prompt length.
50    pub async fn ask(&self, prompt: impl Into<String>) -> Result<PromptRunResult, PromptRunError> {
51        ensure_session_open_for_prompt(self.is_closed())?;
52        self.ensure_tool_use_hook_loop(self.config.hooks.has_pre_tool_use_hooks());
53        self.runtime
54            .run_prompt_on_loaded_thread_with_hooks(
55                &self.thread_id,
56                session_prompt_params(&self.config, prompt),
57                Some(&self.config.hooks),
58            )
59            .await
60    }
61
62    /// Continue this session with one prompt while overriding selected turn options.
63    /// Side effects: sends turn/start RPC calls on one already-loaded thread.
64    /// Allocation: depends on caller-provided params. Complexity: O(1) wrapper.
65    pub async fn ask_with(
66        &self,
67        params: PromptRunParams,
68    ) -> Result<PromptRunResult, PromptRunError> {
69        ensure_session_open_for_prompt(self.is_closed())?;
70        self.ensure_tool_use_hook_loop(self.config.hooks.has_pre_tool_use_hooks());
71        self.runtime
72            .run_prompt_on_loaded_thread_with_hooks(
73                &self.thread_id,
74                params,
75                Some(&self.config.hooks),
76            )
77            .await
78    }
79
80    /// Continue this session with one prompt using one explicit profile override.
81    /// Side effects: sends turn/start RPC calls on one already-loaded thread.
82    /// Allocation: moves profile-owned Strings/vectors + one prompt String. Complexity: O(n), n = attachment count + field sizes.
83    pub async fn ask_with_profile(
84        &self,
85        prompt: impl Into<String>,
86        profile: RunProfile,
87    ) -> Result<PromptRunResult, PromptRunError> {
88        ensure_session_open_for_prompt(self.is_closed())?;
89        let (params, profile_hooks) =
90            profile_to_prompt_params_with_hooks(self.config.cwd.clone(), prompt, profile);
91        let merged_hooks = merge_hook_configs(&self.config.hooks, &profile_hooks);
92        self.ensure_tool_use_hook_loop(merged_hooks.has_pre_tool_use_hooks());
93        self.runtime
94            .run_prompt_on_loaded_thread_with_hooks(&self.thread_id, params, Some(&merged_hooks))
95            .await
96    }
97
98    /// Return current session default profile snapshot.
99    /// Allocation: clones Strings/attachments. Complexity: O(n), n = attachment count + string sizes.
100    pub fn profile(&self) -> RunProfile {
101        self.config.profile()
102    }
103
104    /// Interrupt one in-flight turn in this session.
105    /// Side effects: sends turn/interrupt RPC call to app-server.
106    /// Allocation: one small JSON payload in runtime layer. Complexity: O(1).
107    pub async fn interrupt_turn(&self, turn_id: &str) -> Result<(), RpcError> {
108        ensure_session_open_for_rpc(self.is_closed())?;
109        self.runtime.turn_interrupt(&self.thread_id, turn_id).await
110    }
111
112    /// Archive this session on server side.
113    /// Side effects: sends thread/archive RPC call to app-server.
114    /// Allocation: one small JSON payload in runtime layer. Complexity: O(1).
115    pub async fn close(&self) -> Result<(), RpcError> {
116        let mut guard = self.close_result.lock().await;
117        if let Some(result) = guard.as_ref() {
118            return result.clone();
119        }
120
121        self.closed.store(true, Ordering::Release);
122        let result = self.runtime.thread_archive(&self.thread_id).await;
123        *guard = Some(result.clone());
124        result
125    }
126
127    fn ensure_tool_use_hook_loop(&self, needs_loop: bool) {
128        if !needs_loop {
129            return;
130        }
131        if self
132            .tool_use_loop_started
133            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
134            .is_ok()
135        {
136            tokio::spawn(
137                crate::runtime::api::tool_use_hooks::run_tool_use_approval_loop(
138                    self.runtime.clone(),
139                ),
140            );
141        }
142    }
143}
144
145pub(super) fn ensure_session_open_for_prompt(closed: bool) -> Result<(), PromptRunError> {
146    if closed {
147        return Err(PromptRunError::Rpc(RpcError::InvalidRequest(
148            "session is closed".to_owned(),
149        )));
150    }
151    Ok(())
152}
153
154pub(super) fn ensure_session_open_for_rpc(closed: bool) -> Result<(), RpcError> {
155    if closed {
156        return Err(RpcError::InvalidRequest("session is closed".to_owned()));
157    }
158    Ok(())
159}