codex_runtime/runtime/client/
session.rs1use 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 pub fn is_closed(&self) -> bool {
44 self.closed.load(Ordering::Acquire)
45 }
46
47 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 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 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 pub fn profile(&self) -> RunProfile {
101 self.config.profile()
102 }
103
104 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 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}