mirror of
https://github.com/openai/codex.git
synced 2026-10-29 19:08:43 +00:00
Expose session id to exec environments
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user