Merge 022c5dc997 into sapling-pr-archive-bolinfest

This commit is contained in:
Michael Bolin
2025-07-21 21:55:36 -07:00
committed by GitHub
17 changed files with 392 additions and 48 deletions

View File

@@ -370,11 +370,25 @@ export function isSafeCommand(
reason: "View file with line numbers",
group: "Reading files",
};
case "rg":
case "rg": {
// Certain ripgrep options execute external commands or invoke other
// processes, so we must reject them.
const isUnsafe = command.some((arg: string) =>
UNSAFE_OPTIONS_FOR_RIPGREP_WITHOUT_ARGS.has(arg) ||
[...UNSAFE_OPTIONS_FOR_RIPGREP_WITH_ARGS].some(
(opt) => arg === opt || arg.startsWith(`${opt}=`),
),
);
if (isUnsafe) {
break;
}
return {
reason: "Ripgrep search",
group: "Searching",
};
}
case "find": {
// Certain options to `find` allow executing arbitrary processes, so we
// cannot auto-approve them.
@@ -495,6 +509,22 @@ const UNSAFE_OPTIONS_FOR_FIND_COMMAND: ReadonlySet<string> = new Set([
"-fprintf",
]);
// Ripgrep options that are considered unsafe because they may execute
// arbitrary commands or spawn auxiliary processes.
const UNSAFE_OPTIONS_FOR_RIPGREP_WITH_ARGS: ReadonlySet<string> = new Set([
// Executes an arbitrary command for each matching file.
"--pre",
// Allows custom hostname command which could leak environment details.
"--hostname-bin",
]);
const UNSAFE_OPTIONS_FOR_RIPGREP_WITHOUT_ARGS: ReadonlySet<string> = new Set([
// Enables searching inside archives which triggers external decompression
// utilities reject out of an abundance of caution.
"--search-zip",
"-z",
]);
// ---------------- Helper utilities for complex shell expressions -----------------
// A conservative allow-list of bash operators that do not, on their own, cause

View File

@@ -44,6 +44,14 @@ describe("canAutoApprove()", () => {
group: "Navigating",
runInSandbox: false,
});
// Ripgrep safe invocation.
expect(check(["rg", "TODO"])).toEqual({
type: "auto-approve",
reason: "Ripgrep search",
group: "Searching",
runInSandbox: false,
});
});
test("simple safe commands within a `bash -lc` call", () => {
@@ -67,6 +75,24 @@ describe("canAutoApprove()", () => {
});
});
test("ripgrep unsafe flags", () => {
// Flags that do not take arguments
expect(check(["rg", "--search-zip", "TODO"])).toEqual({ type: "ask-user" });
expect(check(["rg", "-z", "TODO"])).toEqual({ type: "ask-user" });
// Flags that take arguments (provided separately)
expect(check(["rg", "--pre", "cat", "TODO"])).toEqual({ type: "ask-user" });
expect(check(["rg", "--hostname-bin", "hostname", "TODO"])).toEqual({
type: "ask-user",
});
// Flags that take arguments in = form
expect(check(["rg", "--pre=cat", "TODO"])).toEqual({ type: "ask-user" });
expect(check(["rg", "--hostname-bin=hostname", "TODO"])).toEqual({
type: "ask-user",
});
});
test("bash -lc commands with unsafe redirects", () => {
expect(check(["bash", "-lc", "echo hello > file.txt"])).toEqual({
type: "ask-user",

1
codex-rs/Cargo.lock generated
View File

@@ -807,6 +807,7 @@ dependencies = [
"toml 0.9.1",
"tracing",
"tracing-subscriber",
"uuid",
"wiremock",
]

View File

@@ -35,7 +35,7 @@ pub async fn run_main(opts: ProtoCli) -> anyhow::Result<()> {
let config = Config::load_with_cli_overrides(overrides_vec, ConfigOverrides::default())?;
let ctrl_c = notify_on_sigint();
let (codex, _init_id) = Codex::spawn(config, ctrl_c.clone()).await?;
let (codex, _init_id, _session_id) = Codex::spawn(config, ctrl_c.clone()).await?;
let codex = Arc::new(codex);
// Task that reads JSON lines from stdin and forwards to Submission Queue

View File

@@ -101,7 +101,7 @@ impl Codex {
/// Spawn a new [`Codex`] and initialize the session. Returns the instance
/// of `Codex` and the ID of the `SessionInitialized` event that was
/// submitted to start the session.
pub async fn spawn(config: Config, ctrl_c: Arc<Notify>) -> CodexResult<(Codex, String)> {
pub async fn spawn(config: Config, ctrl_c: Arc<Notify>) -> CodexResult<(Codex, String, Uuid)> {
// experimental resume path (undocumented)
let resume_path = config.experimental_resume.clone();
info!("resume_path: {resume_path:?}");
@@ -124,7 +124,12 @@ impl Codex {
};
let config = Arc::new(config);
tokio::spawn(submission_loop(config, rx_sub, tx_event, ctrl_c));
// Generate a unique ID for the lifetime of this Codex session.
let session_id = Uuid::new_v4();
tokio::spawn(submission_loop(
session_id, config, rx_sub, tx_event, ctrl_c,
));
let codex = Codex {
next_id: AtomicU64::new(0),
tx_sub,
@@ -132,7 +137,7 @@ impl Codex {
};
let init_id = codex.submit(configure_session).await?;
Ok((codex, init_id))
Ok((codex, init_id, session_id))
}
/// Submit the `op` wrapped in a `Submission` with a unique ID.
@@ -521,14 +526,12 @@ impl AgentTask {
}
async fn submission_loop(
mut session_id: Uuid,
config: Arc<Config>,
rx_sub: Receiver<Submission>,
tx_event: Sender<Event>,
ctrl_c: Arc<Notify>,
) {
// Generate a unique ID for the lifetime of this Codex session.
let mut session_id = Uuid::new_v4();
let mut sess: Option<Arc<Session>> = None;
// shorthand - send an event when there is no active session
let send_no_session_event = |sub_id: String| async {

View File

@@ -6,15 +6,16 @@ use crate::protocol::Event;
use crate::protocol::EventMsg;
use crate::util::notify_on_sigint;
use tokio::sync::Notify;
use uuid::Uuid;
/// Spawn a new [`Codex`] and initialize the session.
///
/// Returns the wrapped [`Codex`] **and** the `SessionInitialized` event that
/// is received as a response to the initial `ConfigureSession` submission so
/// that callers can surface the information to the UI.
pub async fn init_codex(config: Config) -> anyhow::Result<(Codex, Event, Arc<Notify>)> {
pub async fn init_codex(config: Config) -> anyhow::Result<(Codex, Event, Arc<Notify>, Uuid)> {
let ctrl_c = notify_on_sigint();
let (codex, init_id) = Codex::spawn(config, ctrl_c.clone()).await?;
let (codex, init_id, session_id) = Codex::spawn(config, ctrl_c.clone()).await?;
// The first event must be `SessionInitialized`. Validate and forward it to
// the caller so that they can display it in the conversation history.
@@ -33,5 +34,5 @@ pub async fn init_codex(config: Config) -> anyhow::Result<(Codex, Event, Arc<Not
));
}
Ok((codex, event, ctrl_c))
Ok((codex, event, ctrl_c, session_id))
}

View File

@@ -268,6 +268,40 @@ mod tests {
}
}
#[test]
fn ripgrep_rules() {
// Safe ripgrep invocations none of the unsafe flags are present.
assert!(is_safe_to_call_with_exec(&vec_str(&[
"rg",
"Cargo.toml",
"-n"
])));
// Unsafe flags that do not take an argument (present verbatim).
for args in [
vec_str(&["rg", "--search-zip", "files"]),
vec_str(&["rg", "-z", "files"]),
] {
assert!(
!is_safe_to_call_with_exec(&args),
"expected {args:?} to be considered unsafe due to zip-search flag",
);
}
// Unsafe flags that expect a value, provided in both split and = forms.
for args in [
vec_str(&["rg", "--pre", "pwned", "files"]),
vec_str(&["rg", "--pre=pwned", "files"]),
vec_str(&["rg", "--hostname-bin", "pwned", "files"]),
vec_str(&["rg", "--hostname-bin=pwned", "files"]),
] {
assert!(
!is_safe_to_call_with_exec(&args),
"expected {args:?} to be considered unsafe due to external-command flag",
);
}
}
#[test]
fn bash_lc_safe_examples() {
assert!(is_known_safe_command(&vec_str(&["bash", "-lc", "ls"])));

View File

@@ -75,7 +75,7 @@ async fn includes_session_id_and_model_headers_in_request() {
let mut config = load_default_config_for_test(&codex_home);
config.model_provider = model_provider;
let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new());
let (codex, _init_id) = Codex::spawn(config, ctrl_c.clone()).await.unwrap();
let (codex, _init_id, _session_id) = Codex::spawn(config, ctrl_c.clone()).await.unwrap();
codex
.submit(Op::UserInput {

View File

@@ -49,7 +49,8 @@ async fn spawn_codex() -> Result<Codex, CodexErr> {
let mut config = load_default_config_for_test(&codex_home);
config.model_provider.request_max_retries = Some(2);
config.model_provider.stream_max_retries = Some(2);
let (agent, _init_id) = Codex::spawn(config, std::sync::Arc::new(Notify::new())).await?;
let (agent, _init_id, _session_id) =
Codex::spawn(config, std::sync::Arc::new(Notify::new())).await?;
Ok(agent)
}

View File

@@ -113,7 +113,7 @@ async fn keeps_previous_response_id_between_tasks() {
let mut config = load_default_config_for_test(&codex_home);
config.model_provider = model_provider;
let ctrl_c = std::sync::Arc::new(tokio::sync::Notify::new());
let (codex, _init_id) = Codex::spawn(config, ctrl_c.clone()).await.unwrap();
let (codex, _init_id, _session_id) = Codex::spawn(config, ctrl_c.clone()).await.unwrap();
// Task 1 triggers first request (no previous_response_id)
codex

View File

@@ -95,7 +95,7 @@ async fn retries_on_early_close() {
let codex_home = TempDir::new().unwrap();
let mut config = load_default_config_for_test(&codex_home);
config.model_provider = model_provider;
let (codex, _init_id) = Codex::spawn(config, ctrl_c).await.unwrap();
let (codex, _init_id, _session_id) = Codex::spawn(config, ctrl_c).await.unwrap();
codex
.submit(Op::UserInput {

View File

@@ -153,7 +153,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option<PathBuf>) -> any
.with_writer(std::io::stderr)
.try_init();
let (codex_wrapper, event, ctrl_c) = codex_wrapper::init_codex(config).await?;
let (codex_wrapper, event, ctrl_c, _session_id) = codex_wrapper::init_codex(config).await?;
let codex = Arc::new(codex_wrapper);
info!("Codex initialized with event: {event:?}");

View File

@@ -33,6 +33,7 @@ tokio = { version = "1", features = [
"rt-multi-thread",
"signal",
] }
uuid = { version = "1", features = ["serde", "v4"] }
[dev-dependencies]
assert_cmd = "2"

View File

@@ -160,6 +160,47 @@ impl CodexToolCallParam {
}
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CodexToolCallReplyParam {
/// The *session id* for this conversation.
pub session_id: String,
/// The *next user prompt* to continue the Codex conversation.
pub prompt: String,
}
/// Builds a `Tool` definition for the `codex-reply` tool-call.
pub(crate) fn create_tool_for_codex_tool_call_reply_param() -> Tool {
let schema = SchemaSettings::draft2019_09()
.with(|s| {
s.inline_subschemas = true;
s.option_add_null_type = false;
})
.into_generator()
.into_root_schema_for::<CodexToolCallReplyParam>();
#[expect(clippy::expect_used)]
let schema_value =
serde_json::to_value(&schema).expect("Codex reply tool schema should serialise to JSON");
let tool_input_schema =
serde_json::from_value::<ToolInputSchema>(schema_value).unwrap_or_else(|e| {
panic!("failed to create Tool from schema: {e}");
});
Tool {
name: "codex-reply".to_string(),
title: Some("Codex Reply".to_string()),
input_schema: tool_input_schema,
output_schema: None,
description: Some(
"Continue a Codex session by providing the session id and prompt.".to_string(),
),
annotations: None,
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -235,4 +276,34 @@ mod tests {
});
assert_eq!(expected_tool_json, tool_json);
}
#[test]
fn verify_codex_tool_reply_json_schema() {
let tool = create_tool_for_codex_tool_call_reply_param();
#[expect(clippy::expect_used)]
let tool_json = serde_json::to_value(&tool).expect("tool serializes");
let expected_tool_json = serde_json::json!({
"description": "Continue a Codex session by providing the session id and prompt.",
"inputSchema": {
"properties": {
"prompt": {
"description": "The *next user prompt* to continue the Codex conversation.",
"type": "string"
},
"sessionId": {
"description": "The *session id* for this conversation.",
"type": "string"
},
},
"required": [
"prompt",
"sessionId",
],
"type": "object",
},
"name": "codex-reply",
"title": "Codex Reply",
});
assert_eq!(expected_tool_json, tool_json);
}
}

View File

@@ -2,6 +2,7 @@
//! Tokio task. Separated from `message_processor.rs` to keep that file small
//! and to make future feature-growth easier to manage.
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
@@ -27,7 +28,9 @@ use mcp_types::TextContent;
use serde::Deserialize;
use serde::Serialize;
use serde_json::json;
use tokio::sync::Mutex;
use tracing::error;
use uuid::Uuid;
use crate::outgoing_message::OutgoingMessageSender;
@@ -42,8 +45,9 @@ pub async fn run_codex_tool_session(
initial_prompt: String,
config: CodexConfig,
outgoing: Arc<OutgoingMessageSender>,
session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
) {
let (codex, first_event, _ctrl_c) = match init_codex(config).await {
let (codex, first_event, _ctrl_c, session_id) = match init_codex(config).await {
Ok(res) => res,
Err(e) => {
let result = CallToolResult {
@@ -61,6 +65,11 @@ pub async fn run_codex_tool_session(
};
let codex = Arc::new(codex);
// update the session map so we can retrieve the session in a reply, and then drop it, since
// we no longer need it for this function
session_map.lock().await.insert(session_id, codex.clone());
drop(session_map);
// Send initial SessionConfigured event.
outgoing.send_event_as_notification(&first_event).await;
@@ -85,6 +94,37 @@ pub async fn run_codex_tool_session(
tracing::error!("Failed to submit initial prompt: {e}");
}
run_codex_tool_session_inner(codex, outgoing, id).await;
}
pub async fn run_codex_tool_session_reply(
codex: Arc<Codex>,
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
prompt: String,
) {
if let Err(e) = codex
.submit(Op::UserInput {
items: vec![InputItem::Text { text: prompt }],
})
.await
{
tracing::error!("Failed to submit user input: {e}");
}
run_codex_tool_session_inner(codex, outgoing, request_id).await;
}
async fn run_codex_tool_session_inner(
codex: Arc<Codex>,
outgoing: Arc<OutgoingMessageSender>,
request_id: RequestId,
) {
let sub_id = match &request_id {
RequestId::String(s) => s.clone(),
RequestId::Integer(n) => n.to_string(),
};
// Stream events until the task needs to pause for user interaction or
// completes.
loop {
@@ -128,7 +168,7 @@ pub async fn run_codex_tool_session(
outgoing
.send_error(
id.clone(),
request_id.clone(),
JSONRPCErrorError {
code: INVALID_PARAMS_ERROR_CODE,
message,
@@ -168,7 +208,9 @@ pub async fn run_codex_tool_session(
is_error: None,
structured_content: None,
};
outgoing.send_response(id.clone(), result.into()).await;
outgoing
.send_response(request_id.clone(), result.into())
.await;
// Continue, don't break so the session continues.
continue;
}
@@ -186,7 +228,9 @@ pub async fn run_codex_tool_session(
is_error: None,
structured_content: None,
};
outgoing.send_response(id.clone(), result.into()).await;
outgoing
.send_response(request_id.clone(), result.into())
.await;
break;
}
EventMsg::SessionConfigured(_) => {
@@ -234,7 +278,9 @@ pub async fn run_codex_tool_session(
// structured way.
structured_content: None,
};
outgoing.send_response(id.clone(), result.into()).await;
outgoing
.send_response(request_id.clone(), result.into())
.await;
break;
}
}

View File

@@ -1,10 +1,14 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use crate::codex_tool_config::CodexToolCallParam;
use crate::codex_tool_config::CodexToolCallReplyParam;
use crate::codex_tool_config::create_tool_for_codex_tool_call_param;
use crate::codex_tool_config::create_tool_for_codex_tool_call_reply_param;
use crate::outgoing_message::OutgoingMessageSender;
use codex_core::Codex;
use codex_core::config::Config as CodexConfig;
use mcp_types::CallToolRequestParams;
use mcp_types::CallToolResult;
@@ -22,12 +26,15 @@ use mcp_types::ServerCapabilitiesTools;
use mcp_types::ServerNotification;
use mcp_types::TextContent;
use serde_json::json;
use tokio::sync::Mutex;
use tokio::task;
use uuid::Uuid;
pub(crate) struct MessageProcessor {
outgoing: Arc<OutgoingMessageSender>,
initialized: bool,
codex_linux_sandbox_exe: Option<PathBuf>,
session_map: Arc<Mutex<HashMap<Uuid, Arc<Codex>>>>,
}
impl MessageProcessor {
@@ -41,6 +48,7 @@ impl MessageProcessor {
outgoing: Arc::new(outgoing),
initialized: false,
codex_linux_sandbox_exe,
session_map: Arc::new(Mutex::new(HashMap::new())),
}
}
@@ -272,7 +280,10 @@ impl MessageProcessor {
) {
tracing::trace!("tools/list -> {params:?}");
let result = ListToolsResult {
tools: vec![create_tool_for_codex_tool_call_param()],
tools: vec![
create_tool_for_codex_tool_call_param(),
create_tool_for_codex_tool_call_reply_param(),
],
next_cursor: None,
};
@@ -288,23 +299,29 @@ impl MessageProcessor {
tracing::info!("tools/call -> params: {:?}", params);
let CallToolRequestParams { name, arguments } = params;
// We only support the "codex" tool for now.
if name != "codex" {
// Tool not found return error result so the LLM can react.
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_string(),
text: format!("Unknown tool '{name}'"),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
self.send_response::<mcp_types::CallToolRequest>(id, result)
.await;
return;
match name.as_str() {
"codex" => self.handle_tool_call_codex(id, arguments).await,
"codex-reply" => {
self.handle_tool_call_codex_session_reply(id, arguments)
.await
}
_ => {
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_string(),
text: format!("Unknown tool '{name}'"),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
self.send_response::<mcp_types::CallToolRequest>(id, result)
.await;
}
}
}
async fn handle_tool_call_codex(&self, id: RequestId, arguments: Option<serde_json::Value>) {
let (initial_prompt, config): (String, CodexConfig) = match arguments {
Some(json_val) => match serde_json::from_value::<CodexToolCallParam>(json_val) {
Ok(tool_cfg) => match tool_cfg.into_config(self.codex_linux_sandbox_exe.clone()) {
@@ -359,15 +376,127 @@ impl MessageProcessor {
}
};
// Clone outgoing sender to move into async task.
// Clone outgoing and session map to move into async task.
let outgoing = self.outgoing.clone();
let session_map = self.session_map.clone();
// Spawn an async task to handle the Codex session so that we do not
// block the synchronous message-processing loop.
task::spawn(async move {
// Run the Codex session and stream events back to the client.
crate::codex_tool_runner::run_codex_tool_session(id, initial_prompt, config, outgoing)
.await;
crate::codex_tool_runner::run_codex_tool_session(
id,
initial_prompt,
config,
outgoing,
session_map,
)
.await;
});
}
async fn handle_tool_call_codex_session_reply(
&self,
request_id: RequestId,
arguments: Option<serde_json::Value>,
) {
tracing::info!("tools/call -> params: {:?}", arguments);
// parse arguments
let CodexToolCallReplyParam { session_id, prompt } = match arguments {
Some(json_val) => match serde_json::from_value::<CodexToolCallReplyParam>(json_val) {
Ok(params) => params,
Err(e) => {
tracing::error!("Failed to parse Codex tool call reply parameters: {e}");
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: format!("Failed to parse configuration for Codex tool: {e}"),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
.await;
return;
}
},
None => {
tracing::error!(
"Missing arguments for codex-reply tool-call; the `session_id` and `prompt` fields are required."
);
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: "Missing arguments for codex-reply tool-call; the `session_id` and `prompt` fields are required.".to_owned(),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
.await;
return;
}
};
let session_id = match Uuid::parse_str(&session_id) {
Ok(id) => id,
Err(e) => {
tracing::error!("Failed to parse session_id: {e}");
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: format!("Failed to parse session_id: {e}"),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
self.send_response::<mcp_types::CallToolRequest>(request_id, result)
.await;
return;
}
};
// load codex from session map
let session_map_mutex = Arc::clone(&self.session_map);
// Clone outgoing and session map to move into async task.
let outgoing = self.outgoing.clone();
// Spawn an async task to handle the Codex session so that we do not
// block the synchronous message-processing loop.
task::spawn(async move {
let session_map = session_map_mutex.lock().await;
let codex = match session_map.get(&session_id) {
Some(codex) => codex,
None => {
tracing::warn!("Session not found for session_id: {session_id}");
let result = CallToolResult {
content: vec![ContentBlock::TextContent(TextContent {
r#type: "text".to_owned(),
text: format!("Session not found for session_id: {session_id}"),
annotations: None,
})],
is_error: Some(true),
structured_content: None,
};
// unwrap_or_default is fine here because we know the result is valid JSON
outgoing
.send_response(request_id, serde_json::to_value(result).unwrap_or_default())
.await;
return;
}
};
crate::codex_tool_runner::run_codex_tool_session_reply(
codex.clone(),
outgoing,
request_id,
prompt.clone(),
)
.await;
});
}

View File

@@ -96,14 +96,15 @@ impl ChatWidget<'_> {
// Create the Codex asynchronously so the UI loads as quickly as possible.
let config_for_agent_loop = config.clone();
tokio::spawn(async move {
let (codex, session_event, _ctrl_c) = match init_codex(config_for_agent_loop).await {
Ok(vals) => vals,
Err(e) => {
// TODO: surface this error to the user.
tracing::error!("failed to initialize codex: {e}");
return;
}
};
let (codex, session_event, _ctrl_c, _session_id) =
match init_codex(config_for_agent_loop).await {
Ok(vals) => vals,
Err(e) => {
// TODO: surface this error to the user.
tracing::error!("failed to initialize codex: {e}");
return;
}
};
// Forward the captured `SessionInitialized` event that was consumed
// inside `init_codex()` so it can be rendered in the UI.