Expose session id to exec environments

This commit is contained in:
Ilan Bigio
2025-10-21 20:13:04 -07:00
parent 53cadb4df6
commit 3930bf6b66
5 changed files with 91 additions and 8 deletions
+4
View File
@@ -602,6 +602,10 @@ impl Session {
self.tx_event.clone()
}
pub(crate) fn conversation_id(&self) -> ConversationId {
self.conversation_id
}
fn next_internal_sub_id(&self) -> String {
let id = self
.next_internal_sub_id
+2
View File
@@ -4,6 +4,8 @@ use crate::config_types::ShellEnvironmentPolicyInherit;
use std::collections::HashMap;
use std::collections::HashSet;
pub const CODEX_SESSION_ID_ENV_VAR: &str = "CODEX_SESSION_ID";
/// Construct an environment map based on the rules in the specified policy. The
/// resulting map can be passed directly to `Command::envs()` after calling
/// `env_clear()` to ensure no unintended variables are leaked to the spawned
+44 -4
View File
@@ -2,8 +2,10 @@ use async_trait::async_trait;
use codex_protocol::models::ShellToolCallParams;
use std::sync::Arc;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::exec::ExecParams;
use crate::exec_env::CODEX_SESSION_ID_ENV_VAR;
use crate::exec_env::create_env;
use crate::function_tool::FunctionCallError;
use crate::tools::context::ToolInvocation;
@@ -16,12 +18,22 @@ use crate::tools::registry::ToolKind;
pub struct ShellHandler;
impl ShellHandler {
fn to_exec_params(params: ShellToolCallParams, turn_context: &TurnContext) -> ExecParams {
fn to_exec_params(
params: ShellToolCallParams,
session: &Session,
turn_context: &TurnContext,
) -> ExecParams {
let mut env = create_env(&turn_context.shell_environment_policy);
env.insert(
CODEX_SESSION_ID_ENV_VAR.to_string(),
session.conversation_id().to_string(),
);
ExecParams {
command: params.command,
cwd: turn_context.resolve_path(params.workdir.clone()),
timeout_ms: params.timeout_ms,
env: create_env(&turn_context.shell_environment_policy),
env,
with_escalated_permissions: params.with_escalated_permissions,
justification: params.justification,
arg0: None,
@@ -60,7 +72,7 @@ impl ToolHandler for ShellHandler {
"failed to parse function arguments: {e:?}"
))
})?;
let exec_params = Self::to_exec_params(params, turn.as_ref());
let exec_params = Self::to_exec_params(params, session.as_ref(), turn.as_ref());
let content = handle_container_exec_with_params(
tool_name.as_str(),
exec_params,
@@ -76,7 +88,7 @@ impl ToolHandler for ShellHandler {
})
}
ToolPayload::LocalShell { params } => {
let exec_params = Self::to_exec_params(params, turn.as_ref());
let exec_params = Self::to_exec_params(params, session.as_ref(), turn.as_ref());
let content = handle_container_exec_with_params(
tool_name.as_str(),
exec_params,
@@ -97,3 +109,31 @@ impl ToolHandler for ShellHandler {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::codex::make_session_and_context;
use pretty_assertions::assert_eq;
#[test]
fn to_exec_params_includes_session_id() {
let (session, turn) = make_session_and_context();
let expected_session_id = session.conversation_id().to_string();
let params = ShellToolCallParams {
command: vec!["echo".to_string()],
workdir: None,
timeout_ms: None,
with_escalated_permissions: None,
justification: None,
};
let exec_params = ShellHandler::to_exec_params(params, &session, &turn);
assert_eq!(
exec_params.env.get(CODEX_SESSION_ID_ENV_VAR),
Some(&expected_session_id)
);
}
}
+35
View File
@@ -226,6 +226,41 @@ mod tests {
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unified_exec_exposes_session_id_env_var() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
let (session, turn) = test_session_and_turn();
let expected_session_id = session.conversation_id().to_string();
let open_shell = run_unified_exec_request(
&session,
&turn,
None,
vec!["bash".to_string(), "-i".to_string()],
Some(2_500),
)
.await?;
let session_id = open_shell.session_id.expect("expected session id");
let env_output = run_unified_exec_request(
&session,
&turn,
Some(session_id),
vec!["echo $CODEX_SESSION_ID\n".to_string()],
Some(2_500),
)
.await?;
assert!(
env_output.output.contains(&expected_session_id),
"shell output did not include session id: {}",
env_output.output
);
Ok(())
}
#[tokio::test]
async fn unified_exec_timeouts() -> anyhow::Result<()> {
skip_if_sandbox!(Ok(()));
@@ -5,6 +5,7 @@ use tokio::sync::mpsc;
use tokio::time::Duration;
use tokio::time::Instant;
use crate::exec_env::CODEX_SESSION_ID_ENV_VAR;
use crate::exec_env::create_env;
use crate::sandboxing::ExecEnv;
use crate::tools::events::ToolEmitter;
@@ -109,11 +110,12 @@ impl UnifiedExecSessionManager {
) -> Result<UnifiedExecSession, UnifiedExecError> {
let mut orchestrator = ToolOrchestrator::new();
let mut runtime = UnifiedExecRuntime::new(self);
let req = UnifiedExecToolRequest::new(
command,
context.turn.cwd.clone(),
create_env(&context.turn.shell_environment_policy),
let mut env = create_env(&context.turn.shell_environment_policy);
env.insert(
CODEX_SESSION_ID_ENV_VAR.to_string(),
context.session.conversation_id().to_string(),
);
let req = UnifiedExecToolRequest::new(command, context.turn.cwd.clone(), env);
let tool_ctx = ToolCtx {
session: context.session,
turn: context.turn,