Merge origin/main into guardian prewarm branch

This commit is contained in:
Joe Gershenson
2026-06-11 17:53:14 -07:00
22 changed files with 477 additions and 35 deletions
+1
View File
@@ -1106,6 +1106,7 @@ impl Session {
}
/// Flush rollout writes and return the final durability-barrier result.
#[instrument(name = "session.flush_rollout", level = "trace", skip_all)]
pub(crate) async fn flush_rollout(&self) -> std::io::Result<()> {
if let Some(live_thread) = self.live_thread() {
live_thread.flush().await.map_err(std::io::Error::other)
+1
View File
@@ -1764,6 +1764,7 @@ async fn handle_assistant_item_done_in_plan_mode(
false
}
#[instrument(level = "trace", skip_all)]
async fn drain_in_flight(
in_flight: &mut FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
sess: Arc<Session>,
@@ -6,6 +6,7 @@ use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tokio_util::task::AbortOnDropHandle;
use tracing::info;
use tracing::instrument;
use tracing::warn;
use crate::client::ModelClientSession;
@@ -53,6 +54,7 @@ impl SessionStartupPrewarmHandle {
let _ = self.task.await;
}
#[instrument(name = "startup_prewarm.resolve", level = "trace", skip_all)]
async fn resolve(
self,
session_telemetry: &SessionTelemetry,
+2
View File
@@ -19,6 +19,7 @@ use tracing::Span;
use tracing::field;
use tracing::info_span;
use tracing::trace;
use tracing::trace_span;
use tracing::warn;
use crate::codex_thread::BackgroundTerminalInfo;
@@ -400,6 +401,7 @@ impl Session {
task_input,
task_cancellation_token.child_token(),
)
.instrument(trace_span!("session_task.run"))
.await;
let sess = session_ctx.clone_session();
if let Err(err) = sess.flush_rollout().await {
@@ -5,18 +5,21 @@ use codex_protocol::items::TurnItem;
use codex_tools::ConversationHistory;
use codex_tools::ExtensionTurnItem;
use codex_tools::ToolCall as ExtensionToolCall;
use codex_tools::ToolEnvironment;
use codex_tools::ToolName;
use codex_tools::ToolSearchInfo;
use codex_tools::ToolSpec;
use codex_tools::TurnItemEmissionFuture;
use codex_tools::TurnItemEmitter;
use crate::sandboxing::SandboxPermissions;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::stream_events_utils::TurnItemContributorPolicy;
use crate::stream_events_utils::finalize_turn_item;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolPayload;
use crate::tools::handlers::apply_granted_turn_permissions;
use crate::tools::registry::CoreToolRuntime;
use crate::tools::registry::ToolExecutor;
@@ -109,6 +112,26 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {
let conversation_history =
ConversationHistory::new(invocation.session.clone_history().await.into_raw_items());
let mut environments = Vec::with_capacity(invocation.turn.environments.turn_environments.len());
for environment in &invocation.turn.environments.turn_environments {
let additional_permissions = apply_granted_turn_permissions(
invocation.session.as_ref(),
&environment.environment_id,
environment.cwd.as_path(),
SandboxPermissions::UseDefault,
/*additional_permissions*/ None,
)
.await
.additional_permissions;
environments.push(ToolEnvironment {
environment_id: environment.environment_id.clone(),
cwd: environment.cwd.clone(),
file_system: environment.environment.get_filesystem(),
file_system_sandbox_context: invocation
.turn
.file_system_sandbox_context(additional_permissions, &environment.cwd),
});
}
ExtensionToolCall {
turn_id: invocation.turn.sub_id.clone(),
call_id: invocation.call_id.clone(),
@@ -120,6 +143,7 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {
session: Arc::downgrade(&invocation.session),
turn: Arc::downgrade(&invocation.turn),
}),
environments,
payload: invocation.payload.clone(),
}
}
@@ -0,0 +1,305 @@
use std::sync::Arc;
use anyhow::Context;
use anyhow::Result;
use codex_core::config::Config;
use codex_core::config::Constrained;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_image_generation_extension::install as install_image_generation_extension;
use codex_login::CodexAuth;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::models::FileSystemPermissions;
use codex_protocol::models::PermissionProfile;
use codex_protocol::openai_models::InputModality;
use codex_protocol::permissions::FileSystemAccessMode;
use codex_protocol::permissions::FileSystemPath;
use codex_protocol::permissions::FileSystemSandboxEntry;
use codex_protocol::permissions::FileSystemSandboxPolicy;
use codex_protocol::permissions::NetworkSandboxPolicy;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::request_permissions::PermissionGrantScope;
use codex_protocol::request_permissions::RequestPermissionProfile;
use codex_protocol::request_permissions::RequestPermissionsResponse;
use codex_protocol::user_input::UserInput;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::skip_if_sandbox;
use core_test_support::test_codex::local_selections;
use core_test_support::test_codex::test_codex;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::json;
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
const TINY_PNG_BYTES: &[u8] = &[
137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6, 0,
0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1,
122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130,
];
fn image_generation_extensions(auth: &CodexAuth) -> Arc<ExtensionRegistry<Config>> {
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
install_image_generation_extension(&mut extension_builder, auth_manager);
Arc::new(extension_builder.build())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let extensions = image_generation_extensions(&auth);
let mut builder = test_codex()
.with_auth(auth)
.with_extensions(extensions)
.with_model_info_override("gpt-5.4", |model_info| {
model_info.use_responses_lite = true;
model_info.input_modalities = vec![InputModality::Text, InputModality::Image];
})
.with_config(|config| {
assert!(config.web_search_mode.set(WebSearchMode::Live).is_ok());
assert!(config.features.enable(Feature::ImageGeneration).is_ok());
assert!(config.features.disable(Feature::ImageGenExt).is_ok());
});
let test = builder.build(&server).await?;
let denied_path = test.config.cwd.join("denied.png");
std::fs::write(&denied_path, b"not readable")?;
let call_id = "image-edit-denied";
let response_mock = responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_function_call_with_namespace(
call_id,
"image_gen",
"imagegen",
&json!({
"prompt": "edit the image",
"referenced_image_paths": [denied_path.display().to_string()],
})
.to_string(),
),
responses::ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_response_created("resp-2"),
responses::ev_assistant_message("msg-1", "done"),
responses::ev_completed("resp-2"),
]),
],
)
.await;
let mut file_system_sandbox_policy = FileSystemSandboxPolicy::default();
file_system_sandbox_policy
.entries
.push(FileSystemSandboxEntry {
path: FileSystemPath::Path {
path: denied_path.clone(),
},
access: FileSystemAccessMode::Deny,
});
let permission_profile = PermissionProfile::from_runtime_permissions(
&file_system_sandbox_policy,
NetworkSandboxPolicy::Restricted,
);
test.submit_turn_with_permission_profile("edit the denied image", permission_profile)
.await?;
let request = response_mock
.last_request()
.context("missing request containing extension output")?;
let output = request
.function_call_output_content_and_success(call_id)
.and_then(|(content, _)| content)
.context("extension error text should be present")?;
assert!(
output.starts_with(&format!(
"unable to read referenced image at `{}`:",
denied_path.display()
)),
"unexpected extension error: {output}"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn extension_tool_uses_granted_turn_permissions() -> Result<()> {
skip_if_no_network!(Ok(()));
skip_if_sandbox!(Ok(()));
let server = responses::start_mock_server().await;
Mock::given(method("POST"))
.and(path("/v1/images/edits"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"created": 1,
"data": [{"b64_json": "cG5n"}],
})))
.expect(1)
.mount(&server)
.await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let extensions = image_generation_extensions(&auth);
let base_permission_profile = PermissionProfile::workspace_write_with(
&[],
NetworkSandboxPolicy::Restricted,
/*exclude_tmpdir_env_var*/ true,
/*exclude_slash_tmp*/ true,
);
let permission_profile_for_config = base_permission_profile.clone();
let mut builder = test_codex()
.with_auth(auth)
.with_extensions(extensions)
.with_model_info_override("gpt-5.4", |model_info| {
model_info.use_responses_lite = true;
model_info.input_modalities = vec![InputModality::Text, InputModality::Image];
})
.with_config(move |config| {
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest);
config
.permissions
.set_permission_profile(permission_profile_for_config)
.expect("set permission profile");
assert!(config.web_search_mode.set(WebSearchMode::Live).is_ok());
assert!(config.features.enable(Feature::ImageGeneration).is_ok());
assert!(config.features.disable(Feature::ImageGenExt).is_ok());
assert!(
config
.features
.enable(Feature::RequestPermissionsTool)
.is_ok()
);
});
let test = builder.build(&server).await?;
let image_dir = tempfile::tempdir()?;
let image_path = image_dir.path().canonicalize()?.join("granted.png");
std::fs::write(&image_path, TINY_PNG_BYTES)?;
let requested_permissions = RequestPermissionProfile {
file_system: Some(FileSystemPermissions::from_read_write_roots(
Some(vec![image_dir.path().canonicalize()?.try_into()?]),
Some(Vec::new()),
)),
..RequestPermissionProfile::default()
};
let permissions_call_id = "permissions-call";
let image_call_id = "image-edit-granted";
let response_mock = responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_response_created("resp-1"),
responses::ev_function_call(
permissions_call_id,
"request_permissions",
&serde_json::to_string(&json!({
"reason": "Read an image outside the workspace",
"permissions": requested_permissions,
}))?,
),
responses::ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_response_created("resp-2"),
responses::ev_function_call_with_namespace(
image_call_id,
"image_gen",
"imagegen",
&json!({
"prompt": "edit the image",
"referenced_image_paths": [image_path.display().to_string()],
})
.to_string(),
),
responses::ev_completed("resp-2"),
]),
responses::sse(vec![
responses::ev_response_created("resp-3"),
responses::ev_assistant_message("msg-1", "done"),
responses::ev_completed("resp-3"),
]),
],
)
.await;
let (sandbox_policy, permission_profile) =
turn_permission_fields(base_permission_profile, test.config.cwd.as_path());
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "request access and edit the image".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
environments: Some(local_selections(test.config.cwd.clone())),
approval_policy: Some(AskForApproval::OnRequest),
approvals_reviewer: Some(ApprovalsReviewer::User),
sandbox_policy: Some(sandbox_policy),
permission_profile,
collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {
mode: codex_protocol::config_types::ModeKind::Default,
settings: codex_protocol::config_types::Settings {
model: test.session_configured.model.clone(),
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
},
})
.await?;
let event = wait_for_event(&test.codex, |event| {
matches!(
event,
EventMsg::RequestPermissions(_) | EventMsg::TurnComplete(_)
)
})
.await;
let EventMsg::RequestPermissions(request) = event else {
panic!("expected request_permissions before turn completion");
};
assert_eq!(request.call_id, permissions_call_id);
test.codex
.submit(Op::RequestPermissionsResponse {
id: permissions_call_id.to_string(),
response: RequestPermissionsResponse {
permissions: request.permissions,
scope: PermissionGrantScope::Turn,
strict_auto_review: false,
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let request = response_mock
.last_request()
.context("missing request containing extension output")?;
let output = request.function_call_output(image_call_id);
let image = &output["output"][0];
assert_eq!(image["type"], "input_image");
assert_eq!(image["image_url"], "data:image/png;base64,cG5n");
Ok(())
}
+2
View File
@@ -51,6 +51,8 @@ mod compact_resume_fork;
mod deprecation_notice;
mod exec;
mod exec_policy;
#[cfg(not(target_os = "windows"))]
mod extension_sandbox;
mod fork_thread;
#[cfg(not(target_os = "windows"))]
mod guardian_review;