openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-20 12:47:38 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
78e7825a47675b2c0dd6261ef4b3906cc1da26f1
codex/codex-rs/core/src/client_tests.rs
ningyi-oai 3a589370a4 Enable app tool result metadata with analytics controls (#46010) 
## What changed

Capture result metadata for host-owned app calls when analytics and executed tool call metadata recording are enabled. Respect both `analytics.enabled = false` and host-disabled analytics, using the prepared call's ownership to determine eligibility.

Disable incremental WebSocket request reuse when raw result metadata or its call binding changes, so late results attached to an already-sent output are included in the next request. Continue allowing reuse when only other metadata changes or result metadata is filtered out for the endpoint.

## Testing

Extend coverage for direct and code-mode calls, analytics settings and host overrides, runtime recording changes, and prepared call ownership. Add metadata comparison and WebSocket reuse tests, and verify that raw response item notifications continue to omit result metadata.

GitOrigin-RevId: 411843fbff3749ee96fd6f0ebc591362942a4555
2026-09-16 19:32:16 +00:00

1734 lines
60 KiB
Rust
Raw Blame History

use super::AuthRequestTelemetryContext;
use super::ModelClient;
use super::PendingUnauthorizedRetry;
use super::Prompt;
use super::UnauthorizedRecoveryExecution;
use super::X_CODEX_INSTALLATION_ID_HEADER;
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
use super::X_CODEX_TURN_METADATA_HEADER;
use super::X_CODEX_WINDOW_ID_HEADER;
use super::X_OPENAI_SUBAGENT_HEADER;
use crate::AttestationContext;
use crate::AttestationProvider;
use crate::GenerateAttestationFuture;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::test_support::TestCodexResponsesRequestKind;
use crate::test_support::responses_metadata as test_responses_metadata;
use base64::Engine;
use codex_api::AgentIdentityTelemetry;
use codex_api::ApiError;
use codex_api::ResponseEvent;
use codex_api::TransportError;
use codex_http_client::HttpClientFactory;
use codex_http_client::OutboundProxyPolicy;
use codex_login::AuthManager;
use codex_login::CodexAuth;
use codex_login::auth::AgentIdentityAuthPolicy;
use codex_model_provider::BearerAuthProvider;
use codex_model_provider::ModelProvider;
use codex_model_provider::ModelProviderFuture;
use codex_model_provider::ProviderAccountResult;
use codex_model_provider::ProviderAuthRecoveryMessages;
use codex_model_provider::ProviderUnauthorizedRecovery;
use codex_model_provider::SharedModelProvider;
use codex_model_provider::create_model_provider;
use codex_model_provider_info::CHATGPT_CODEX_BASE_URL;
use codex_model_provider_info::ModelProviderInfo;
use codex_model_provider_info::WireApi;
use codex_model_provider_info::create_oss_provider_with_base_url;
use codex_models_manager::manager::SharedModelsManager;
use codex_otel::SessionTelemetry;
use codex_protocol::ThreadId;
use codex_protocol::auth::AuthMode;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ExecutedToolCall;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::ToolResultMetadata;
use codex_protocol::models::ToolResultSource;
use codex_protocol::models::ToolResultSources;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::openai_models::ModelsResponse;
use codex_protocol::openai_models::ReasoningEffort;
use codex_protocol::openai_models::ReasoningEffortPreset;
use codex_protocol::protocol::InternalSessionSource;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::InferenceTraceAttempt;
use codex_rollout_trace::InferenceTraceContext;
use codex_rollout_trace::RawTraceEventPayload;
use codex_rollout_trace::RolloutTrace;
use codex_rollout_trace::TraceWriter;
use codex_rollout_trace::replay_bundle;
use futures::StreamExt;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::task::Context;
use std::task::Poll;
use std::time::Duration;
use tempfile::TempDir;
use tokio::sync::Notify;
use tracing::Event;
use tracing::Subscriber;
use tracing::field::Visit;
use tracing_subscriber::Layer;
use tracing_subscriber::layer::Context as LayerContext;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::util::SubscriberInitExt;
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
fn test_model_client(session_source: SessionSource) -> ModelClient {
test_model_client_with_thread_id(ThreadId::new(), session_source)
}
fn test_model_client_with_thread_id(
thread_id: ThreadId,
session_source: SessionSource,
) -> ModelClient {
let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
ModelClient::new(
/*auth_manager*/ None,
AgentIdentityAuthPolicy::JwtOnly,
thread_id,
provider,
session_source,
"test_originator".to_string(),
/*model_verbosity*/ None,
/*content_item_kinds_enabled*/ true,
/*enable_request_compression*/ false,
/*include_timing_metrics*/ false,
/*beta_features_header*/ None,
/*concurrent_reasoning_summaries_enabled*/ false,
/*attestation_provider*/ None,
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
codex_model_provider::WorkspaceRoutingContext::new(
"https://chatgpt.com/backend-api".into(),
),
)
}
fn test_model_provider() -> SharedModelProvider {
test_model_client(SessionSource::Cli).state.provider.clone()
}
#[tokio::test]
async fn workspace_routed_http_rejects_redirects_without_a_routing_header() {
use codex_client::HttpTransport;
use codex_login::WorkspaceRouting;
use codex_login::WorkspaceRoutingRequest;
use codex_login::WorkspaceRoutingResolver;
struct Routing(Option<&'static str>);
impl WorkspaceRoutingResolver for Routing {
fn resolve(
&self,
_request: WorkspaceRoutingRequest,
) -> Pin<
Box<
dyn std::future::Future<Output = std::io::Result<Option<WorkspaceRouting>>>
+ Send
+ '_,
>,
> {
Box::pin(async move {
Ok(self.0.map(|override_value| WorkspaceRouting {
chatgpt_account_id: "account_id".into(),
backend_origin: "https://gov.chatgpt.com".into(),
account_routing_override: override_value.into(),
}))
})
}
}
for routing_override in [Some("NO_CONSTRAINT"), Some("us_cr"), None] {
let origin = MockServer::start().await;
let destination = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(
ResponseTemplate::new(/*status*/ 307)
.insert_header("location", format!("{}/responses", destination.uri())),
)
.mount(&origin)
.await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(/*status*/ 200))
.mount(&destination)
.await;
let manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let resolver: Arc<dyn WorkspaceRoutingResolver> = Arc::new(Routing(routing_override));
manager.set_workspace_routing_resolver(Arc::downgrade(&resolver));
let mut client = test_model_client(SessionSource::Exec);
Arc::get_mut(&mut client.state).unwrap().provider = create_model_provider(
ModelProviderInfo::create_openai_provider(/*base_url*/ None),
Some(manager),
);
let mut setup = client
.current_client_setup(super::ClientRouting::Workspace)
.await
.unwrap();
if routing_override.is_some() {
assert_eq!(
setup.api_provider.base_url,
"https://gov.chatgpt.com/backend-api/codex"
);
}
// Exercise the resolved route's redirect policy against loopback HTTP servers.
setup.api_provider.base_url = origin.uri();
let transport = client
.build_api_transport(&setup.api_provider, "/responses", setup.redirect_policy)
.unwrap();
let request = setup
.api_provider
.build_request(http::Method::POST, "/responses")
.with_json(&json!({"input": "workspace content"}));
let result = transport.execute(request).await;
if routing_override.is_some() {
assert!(
matches!(
result,
Err(TransportError::Http {
status: http::StatusCode::TEMPORARY_REDIRECT,
..
})
),
"workspace redirect must be rejected: {routing_override:?}"
);
} else {
assert_eq!(result.unwrap().status, http::StatusCode::OK);
}
assert_eq!(
destination.received_requests().await.unwrap().len(),
usize::from(routing_override.is_none())
);
}
}
#[derive(Debug)]
enum SetupRefresh {
Command(PathBuf),
ChatGpt {
home: PathBuf,
token: String,
workspace: String,
},
}
#[derive(Debug)]
struct SetupRefreshProvider {
inner: SharedModelProvider,
refresh: SetupRefresh,
setup_calls: AtomicUsize,
}
impl ModelProvider for SetupRefreshProvider {
fn info(&self) -> &ModelProviderInfo {
self.inner.info()
}
fn auth_manager(&self) -> Option<Arc<AuthManager>> {
self.inner.auth_manager()
}
fn auth(&self) -> ModelProviderFuture<'_, Option<CodexAuth>> {
self.inner.auth()
}
fn account_state(&self) -> ProviderAccountResult {
self.inner.account_state()
}
fn api_provider(
&self,
) -> ModelProviderFuture<'_, codex_protocol::error::Result<codex_api::Provider>> {
Box::pin(async move {
self.setup_calls.fetch_add(1, Ordering::SeqCst);
let manager = self.inner.auth_manager().expect("auth manager");
match &self.refresh {
SetupRefresh::Command(token_path) => {
std::fs::write(token_path, "refreshed-token")?;
manager
.refresh_token_from_authority()
.await
.expect("refresh command token");
}
SetupRefresh::ChatGpt {
home,
token,
workspace,
} => {
codex_login::auth::login_with_chatgpt_auth_tokens(
home, token, workspace, /*chatgpt_plan_type*/ None,
)?;
manager.reload().await;
}
}
self.inner.api_provider().await
})
}
fn models_manager(
&self,
codex_home: PathBuf,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
self.inner.models_manager(codex_home, config_model_catalog)
}
}
#[tokio::test]
async fn client_setup_accepts_command_credential_refresh() {
for routing in [
super::ClientRouting::Workspace,
super::ClientRouting::ConfiguredProvider,
] {
let tempdir = TempDir::new().unwrap();
let token_path = tempdir.path().join("token.txt");
std::fs::write(&token_path, "initial-token").unwrap();
let mut info = test_model_provider().info().clone();
info.auth = Some(codex_protocol::config_types::ModelProviderAuthInfo {
command: if cfg!(windows) { "cmd.exe" } else { "cat" }.into(),
args: if cfg!(windows) {
vec!["/D", "/C", "type", "token.txt"]
} else {
vec!["token.txt"]
}
.into_iter()
.map(Into::into)
.collect(),
timeout_ms: std::num::NonZeroU64::new(/*n*/ 5_000).unwrap(),
refresh_interval_ms: 60_000,
cwd: tempdir.path().try_into().unwrap(),
});
let provider = Arc::new(SetupRefreshProvider {
inner: create_model_provider(info, /*auth_manager*/ None),
refresh: SetupRefresh::Command(token_path),
setup_calls: AtomicUsize::new(/*v*/ 0),
});
let manager = provider.auth_manager().unwrap();
let mut client = test_model_client(SessionSource::Exec);
Arc::get_mut(&mut client.state).unwrap().provider = provider;
let setup = client.current_client_setup(routing).await.unwrap();
let mut headers = http::HeaderMap::new();
setup.api_auth.add_auth_headers(&mut headers);
assert_eq!(
headers.get(http::header::AUTHORIZATION).unwrap(),
"Bearer refreshed-token"
);
let refreshed_revision = Some(*manager.auth_change_receiver().borrow());
assert_ne!(
codex_model_provider::ResponsesConnectionKey::new(
&setup.api_provider,
setup.auth_revision
),
codex_model_provider::ResponsesConnectionKey::new(
&setup.api_provider,
refreshed_revision
),
);
assert_ne!(setup.auth_owner_generation, client.auth_owner_generation());
}
}
#[tokio::test]
async fn client_setup_rebuilds_chatgpt_refresh_but_rejects_account_switches() {
for (user, workspace, expected_calls) in [
("user-a", "workspace-a", 2),
("user-b", "workspace-a", 1),
("user-a", "workspace-b", 1),
] {
let token = |user: &str, revision: &str| {
let claims =
json!({"jti": revision, "https://api.openai.com/auth": {"chatgpt_user_id": user}});
let payload =
base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string());
format!("header.{payload}.signature")
};
let home = TempDir::new().unwrap();
let initial = CodexAuth::from_external_chatgpt_tokens(
&token("user-a", "initial"),
"workspace-a",
/*chatgpt_plan_type*/ None,
)
.unwrap();
let manager =
AuthManager::from_auth_for_testing_with_home(initial, home.path().to_path_buf());
let refreshed_token = token(user, "refreshed");
let mut info = test_model_provider().info().clone();
info.requires_openai_auth = true;
let provider = Arc::new(SetupRefreshProvider {
inner: create_model_provider(info, Some(manager.clone())),
refresh: SetupRefresh::ChatGpt {
home: home.path().to_path_buf(),
token: refreshed_token.clone(),
workspace: workspace.into(),
},
setup_calls: AtomicUsize::new(/*v*/ 0),
});
let mut client = test_model_client(SessionSource::Exec);
Arc::get_mut(&mut client.state).unwrap().provider = provider.clone();
let result = client
.current_client_setup(super::ClientRouting::ConfiguredProvider)
.await;
if expected_calls == 2 {
let setup = result.unwrap();
let mut headers = http::HeaderMap::new();
setup.api_auth.add_auth_headers(&mut headers);
assert_eq!(
headers.get(http::header::AUTHORIZATION).unwrap(),
&format!("Bearer {refreshed_token}")
);
assert_eq!(setup.auth.unwrap().get_token().unwrap(), refreshed_token);
assert_eq!(
(setup.auth_revision, setup.auth_owner_generation),
(Some(*manager.auth_change_receiver().borrow()), Some(0))
);
} else {
assert_eq!(
result.err().expect("account switch must fail").to_string(),
"account changed while preparing model request"
);
}
assert_eq!(provider.setup_calls.load(Ordering::SeqCst), expected_calls);
}
}
fn test_responses_metadata_for_client(
client: &ModelClient,
turn_id: Option<&str>,
window_id: String,
parent_thread_id: Option<ThreadId>,
request_kind: TestCodexResponsesRequestKind,
) -> CodexResponsesMetadata {
let thread_id = client.state.thread_id.to_string();
test_responses_metadata(
TEST_INSTALLATION_ID,
&thread_id,
&thread_id,
turn_id,
window_id,
&client.state.session_source,
parent_thread_id,
request_kind,
)
}
fn test_model_info() -> ModelInfo {
serde_json::from_value(json!({
"slug": "gpt-test",
"display_name": "gpt-test",
"description": "desc",
"default_reasoning_level": "medium",
"supported_reasoning_levels": [
{"effort": "medium", "description": "medium"}
],
"shell_type": "shell_command",
"visibility": "list",
"supported_in_api": true,
"priority": 1,
"upgrade": null,
"model_messages": null,
"support_verbosity": false,
"default_verbosity": null,
"apply_patch_tool_type": null,
"truncation_policy": {"mode": "bytes", "limit": 10000},
"supports_image_detail_original": false,
"context_window": 272000,
"auto_compact_token_limit": null,
"experimental_supported_tools": []
}))
.expect("deserialize test model info")
}
fn output_with_tool_result_metadata(metadata: ToolResultMetadata) -> ResponseItem {
let mut call = ExecutedToolCall::new("test_tool".to_string(), json!({ "query": "keep" }));
call.set_tool_result_sources(ToolResultSources::new(vec![ToolResultSource {
r#type: "test_resource".to_string(),
id: "R1".to_string(),
}]));
call.set_tool_result_metadata(metadata);
let mut output = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: "tool-call".to_string(),
output: FunctionCallOutputPayload::from_text("unchanged tool result".to_string()),
});
output.append_executed_tool_calls(vec![call]);
output.mark_tool_calls_complete();
output
}
#[test]
fn responses_request_limits_raw_tool_metadata_to_resolved_first_party_https_endpoint()
-> anyhow::Result<()> {
let provider =
ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".to_string()));
let mut api_provider = provider.to_api_provider(/*auth_mode*/ None)?;
let mut client = test_model_client(SessionSource::Cli);
Arc::get_mut(&mut client.state)
.expect("test client should have unique session state")
.provider = create_model_provider(provider, /*auth_manager*/ None);
let output = output_with_tool_result_metadata(ToolResultMetadata::new(&json!({
"private": { "resource": "raw-result-metadata" },
})));
let without_raw_metadata = output_with_tool_result_metadata(ToolResultMetadata::default());
let prompt = Prompt {
input: vec![output.clone()],
..Default::default()
};
let responses_metadata = test_responses_metadata_for_client(
&client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
);
for (base_url, allowed) in [
("https://api.openai.com/v1", true),
("https://chatgpt.com/backend-api/codex", true),
("https://api.chatgpt-staging.com/v1", true),
("https://proxy.example.com/v1", false),
("http://api.openai.com/v1", false),
("https://api.openai.com.evil.example/v1", false),
("https://chatgpt.com.evil.example/v1", false),
("https://api.openai.com@proxy.example.com/v1", false),
("not a URL", false),
] {
api_provider.base_url = base_url.to_string();
for responses_lite in [false, true] {
let mut model = test_model_info();
model.use_responses_lite = responses_lite;
let mut request = client.build_responses_request(
&prompt,
&model,
/*effort*/ None,
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&responses_metadata,
)?;
ModelClient::filter_tool_result_metadata(&mut request.input, &api_provider);
assert_eq!(
request.input.last(),
Some(if allowed {
&output
} else {
&without_raw_metadata
}),
"resolved endpoint: {base_url}, responses_lite: {responses_lite}",
);
assert_eq!(prompt.input, vec![output.clone()]);
}
}
Ok(())
}
#[test]
fn websocket_incremental_reuse_tracks_raw_result_metadata() -> anyhow::Result<()> {
let provider =
ModelProviderInfo::create_openai_provider(Some("https://api.openai.com/v1".to_string()));
let mut api_provider = provider.to_api_provider(/*auth_mode*/ None)?;
let mut client = test_model_client(SessionSource::Cli);
Arc::get_mut(&mut client.state)
.expect("test client should have unique session state")
.provider = create_model_provider(provider, /*auth_manager*/ None);
let responses_metadata = test_responses_metadata_for_client(
&client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
);
for (scenario, base_url, previous_metadata, current_metadata, expect_incremental) in [
(
"late_result",
"https://api.openai.com/v1",
None,
Some("first"),
false,
),
(
"unchanged_result",
"https://api.openai.com/v1",
Some("first"),
Some("first"),
true,
),
(
"changed_result",
"https://api.openai.com/v1",
Some("first"),
Some("second"),
false,
),
(
"ordinary_metadata_only",
"https://api.openai.com/v1",
None,
None,
true,
),
(
"filtered_result",
"https://proxy.example.com/v1",
None,
Some("first"),
true,
),
] {
let [mut previous_output, mut current_output] =
[previous_metadata, current_metadata].map(|metadata| {
let mut call =
ExecutedToolCall::new("apps_tool".to_string(), json!({ "query": "same" }));
if let Some(id) = metadata {
call.set_tool_result_metadata(ToolResultMetadata::new(&json!({ "id": id })));
}
let mut output = ResponseItem::from(ResponseInputItem::CustomToolCallOutput {
call_id: "exec-call".to_string(),
name: None,
output: FunctionCallOutputPayload::from_text(
"Script running with cell ID cell".to_string(),
),
});
output.append_executed_tool_calls(vec![call]);
output.set_tool_call_cell_id("exec-call");
output
});
previous_output.set_turn_id_if_missing("previous-turn");
current_output.set_turn_id_if_missing("current-turn");
let mut previous = client.build_responses_request(
&Prompt {
input: vec![previous_output],
..Default::default()
},
&test_model_info(),
/*effort*/ None,
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&responses_metadata,
)?;
let follow_up = ResponseItem::from(ResponseInputItem::FunctionCallOutput {
call_id: "wait-call".to_string(),
output: FunctionCallOutputPayload::from_text("done".to_string()),
});
let mut current = previous.clone();
current.input = vec![current_output, follow_up.clone()];
api_provider.base_url = base_url.to_string();
ModelClient::filter_tool_result_metadata(&mut previous.input, &api_provider);
ModelClient::filter_tool_result_metadata(&mut current.input, &api_provider);
let mut session = client.new_session();
session.websocket_session.last_request = Some(previous);
let (sender, receiver) = tokio::sync::oneshot::channel();
sender
.send(super::LastResponse {
response_id: "previous-response".to_string(),
items_added: Vec::new(),
})
.unwrap();
session.websocket_session.last_response_rx = Some(receiver);
let (incremental, from_warmup) = session.prepare_websocket_request(&current);
assert_eq!(
incremental,
expect_incremental.then_some(("previous-response".to_string(), vec![follow_up])),
"{scenario}",
);
assert!(!from_warmup);
}
Ok(())
}
#[tokio::test]
async fn responses_http_omits_raw_tool_metadata_for_openai_named_custom_endpoint()
-> anyhow::Result<()> {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(
ResponseTemplate::new(/*status*/ 200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-1\"}}\n\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-1\"}}\n\n",
)),
)
.expect(/*requests*/ 1)
.mount(&server)
.await;
let mut provider =
ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri())));
provider.requires_openai_auth = false;
provider.supports_websockets = false;
let mut client = test_model_client(SessionSource::Cli);
Arc::get_mut(&mut client.state)
.expect("test client should have unique session state")
.provider = create_model_provider(provider, /*auth_manager*/ None);
let output = output_with_tool_result_metadata(ToolResultMetadata::new(&json!({
"private": "raw-result-metadata",
})));
let prompt = Prompt {
input: vec![output.clone()],
..Default::default()
};
let responses_metadata = test_responses_metadata_for_client(
&client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
);
let mut session = client.new_session();
let mut stream = session
.stream(
&prompt,
&test_model_info(),
&test_session_telemetry(),
/*effort*/ None,
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&responses_metadata,
&InferenceTraceContext::disabled(),
)
.await?;
let mut completed = false;
while let Some(event) = stream.next().await {
if let ResponseEvent::Completed { response_id, .. } = event? {
assert_eq!(response_id, "resp-1");
completed = true;
}
}
assert!(completed);
let requests = server.received_requests().await.expect("received requests");
assert_eq!(requests.len(), 1);
let body: serde_json::Value = serde_json::from_slice(&requests[0].body)?;
assert_eq!(
body["input"],
serde_json::to_value(vec![output_with_tool_result_metadata(
ToolResultMetadata::default(),
)])?,
);
assert_eq!(prompt.input, vec![output]);
Ok(())
}
#[test]
fn responses_lite_prefix_ids_track_thread_and_payload() -> anyhow::Result<()> {
let thread_id = ThreadId::new();
let client = test_model_client_with_thread_id(thread_id, SessionSource::Cli);
let mut model = test_model_info();
model.use_responses_lite = true;
let mut prompt = Prompt {
base_instructions: BaseInstructions {
text: "base instructions".to_string(),
provenance: None,
},
..Default::default()
};
let build = |client: &ModelClient, prompt: &Prompt| {
client.build_responses_request(
prompt,
&model,
/*effort*/ None,
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&test_responses_metadata_for_client(
client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
),
)
};
let original = build(&client, &prompt)?;
assert_eq!(build(&client, &prompt)?, original);
prompt.base_instructions.text.push_str(" with an update");
let changed_instructions = build(&client, &prompt)?;
assert_eq!(changed_instructions.input[0], original.input[0]);
assert_ne!(changed_instructions.input[1].id(), original.input[1].id());
prompt.tools = vec![codex_tools::ToolSpec::Freeform(codex_tools::FreeformTool {
name: "exec".to_string(),
description: "Execute JavaScript.".to_string(),
defer_loading: None,
format: codex_tools::FreeformToolFormat {
r#type: "grammar".to_string(),
syntax: "lark".to_string(),
definition: "start: /.+/".to_string(),
},
})]
.into();
let changed_tools = build(&client, &prompt)?;
assert_ne!(
changed_tools.input[0].id(),
changed_instructions.input[0].id()
);
assert_eq!(changed_tools.input[1], changed_instructions.input[1]);
let independent = build(
&test_model_client_with_thread_id(ThreadId::new(), SessionSource::Cli),
&prompt,
)?;
assert_ne!(independent.input[0].id(), changed_tools.input[0].id());
assert_ne!(independent.input[1].id(), changed_tools.input[1].id());
Ok(())
}
fn test_session_telemetry() -> SessionTelemetry {
SessionTelemetry::new(
ThreadId::new(),
"gpt-test",
"gpt-test",
/*account_id*/ None,
/*account_email*/ None,
/*auth_mode*/ None,
"test-originator".to_string(),
/*log_user_prompts*/ false,
"test-terminal".to_string(),
SessionSource::Cli,
)
}
fn spawned_session_source() -> SessionSource {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: ThreadId::new(),
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: None,
})
}
fn reasoning_effort_in_request(
model_info: &ModelInfo,
session_source: SessionSource,
effort: ReasoningEffort,
) -> ReasoningEffort {
let client = test_model_client(session_source);
client
.build_responses_request(
&Prompt::default(),
model_info,
Some(effort),
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&test_responses_metadata_for_client(
&client,
/*turn_id*/ None,
format!("{}:0", client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::Turn,
),
)
.expect("build responses request")
.reasoning
.expect("request should include reasoning")
.effort
.expect("request should include reasoning effort")
}
#[test]
fn reasoning_effort_for_requests_uses_multi_agent_override_for_ultra() {
let mut model_info = test_model_info();
model_info.multi_agent_reasoning_effort = Some(ReasoningEffort::High);
model_info
.supported_reasoning_levels
.push(ReasoningEffortPreset {
effort: ReasoningEffort::High,
description: "high".to_string(),
});
let actual = [SessionSource::Cli, spawned_session_source()].map(|session_source| {
reasoning_effort_in_request(&model_info, session_source, ReasoningEffort::Ultra)
});
assert_eq!(actual, [ReasoningEffort::High, ReasoningEffort::High]);
}
#[test]
fn reasoning_effort_for_requests_falls_back_for_missing_or_invalid_override() {
let mut model_info = test_model_info();
model_info.supported_reasoning_levels = vec![
ReasoningEffortPreset {
effort: ReasoningEffort::Low,
description: "low".to_string(),
},
ReasoningEffortPreset {
effort: ReasoningEffort::XHigh,
description: "xhigh".to_string(),
},
ReasoningEffortPreset {
effort: ReasoningEffort::Ultra,
description: "ultra".to_string(),
},
];
let actual = [
None,
Some(ReasoningEffort::Ultra),
Some(ReasoningEffort::High),
]
.map(|multi_agent_reasoning_effort| {
model_info.multi_agent_reasoning_effort = multi_agent_reasoning_effort;
reasoning_effort_in_request(&model_info, SessionSource::Cli, ReasoningEffort::Ultra)
});
assert_eq!(
actual,
[
ReasoningEffort::XHigh,
ReasoningEffort::XHigh,
ReasoningEffort::XHigh,
]
);
model_info.multi_agent_reasoning_effort = None;
model_info.supported_reasoning_levels.insert(
1,
ReasoningEffortPreset {
effort: ReasoningEffort::Max,
description: "max".to_string(),
},
);
assert_eq!(
reasoning_effort_in_request(&model_info, SessionSource::Cli, ReasoningEffort::Ultra),
ReasoningEffort::Max
);
model_info.supported_reasoning_levels.clear();
assert_eq!(
reasoning_effort_in_request(&model_info, SessionSource::Cli, ReasoningEffort::Ultra),
ReasoningEffort::Medium
);
}
#[test]
fn reasoning_effort_for_requests_preserves_non_ultra_and_persistent_behavior() {
let mut model_info = test_model_info();
model_info.multi_agent_reasoning_effort = Some(ReasoningEffort::Low);
assert_eq!(
(
reasoning_effort_in_request(&model_info, SessionSource::Cli, ReasoningEffort::High,),
reasoning_effort_in_request(
&model_info,
SessionSource::Cli,
ReasoningEffort::Persistent,
),
),
(
ReasoningEffort::High,
ReasoningEffort::Custom("disabled".to_string()),
)
);
}
#[derive(Default)]
struct TagCollectorVisitor {
tags: BTreeMap<String, String>,
}
impl Visit for TagCollectorVisitor {
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
self.tags
.insert(field.name().to_string(), value.to_string());
}
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
self.tags
.insert(field.name().to_string(), format!("{value:?}"));
}
}
#[derive(Clone)]
struct TagCollectorLayer {
tags: Arc<Mutex<BTreeMap<String, String>>>,
}
impl<S> Layer<S> for TagCollectorLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_event(&self, event: &Event<'_>, _ctx: LayerContext<'_, S>) {
if event.metadata().target() != "feedback_tags" {
return;
}
let mut visitor = TagCollectorVisitor::default();
event.record(&mut visitor);
self.tags.lock().unwrap().extend(visitor.tags);
}
}
fn started_inference_attempt(temp: &TempDir) -> anyhow::Result<InferenceTraceAttempt> {
let writer = Arc::new(TraceWriter::create(
temp.path(),
"trace-1".to_string(),
"rollout-1".to_string(),
"thread-root".to_string(),
)?);
writer.append(RawTraceEventPayload::ThreadStarted {
thread_id: "thread-root".to_string(),
agent_path: "/root".to_string(),
metadata_payload: None,
})?;
writer.append(RawTraceEventPayload::CodexTurnStarted {
codex_turn_id: "turn-1".to_string(),
thread_id: "thread-root".to_string(),
})?;
let inference_trace = InferenceTraceContext::enabled(
writer,
"thread-root".to_string(),
"turn-1".to_string(),
"gpt-test".to_string(),
"test-provider".to_string(),
);
let attempt = inference_trace.start_attempt();
attempt.record_started(&json!({
"model": "gpt-test",
"input": [{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "hello"}]
}],
}));
Ok(attempt)
}
fn output_message(id: &str, text: &str) -> ResponseItem {
ResponseItem::Message {
id: Some(codex_protocol::ResponseItemId::with_suffix("msg", id)),
role: "assistant".to_string(),
content: vec![ContentItem::OutputText {
text: text.to_string(),
}],
phase: None,
internal_chat_message_metadata_passthrough: None,
}
}
async fn replay_until_cancelled(temp: &TempDir) -> anyhow::Result<RolloutTrace> {
let mut rollout = replay_bundle(temp.path())?;
for _ in 0..50 {
let inference = rollout
.inference_calls
.values()
.next()
.expect("inference should be reduced");
if inference.execution.status == ExecutionStatus::Cancelled {
return Ok(rollout);
}
tokio::time::sleep(Duration::from_millis(10)).await;
rollout = replay_bundle(temp.path())?;
}
Ok(rollout)
}
struct NotifyAfterEventStream {
events: VecDeque<ResponseEvent>,
yielded: usize,
notify_after: usize,
notify: Arc<Notify>,
}
impl futures::Stream for NotifyAfterEventStream {
type Item = std::result::Result<ResponseEvent, ApiError>;
fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let Some(event) = self.events.pop_front() else {
return Poll::Pending;
};
self.yielded += 1;
if self.yielded == self.notify_after {
self.notify.notify_one();
}
Poll::Ready(Some(Ok(event)))
}
}
#[test]
fn build_subagent_headers_sets_other_subagent_label() {
let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other(
"memory_consolidation".to_string(),
)));
let headers = client.build_subagent_headers();
let value = headers
.get(X_OPENAI_SUBAGENT_HEADER)
.and_then(|value| value.to_str().ok());
assert_eq!(value, Some("memory_consolidation"));
}
#[test]
fn internal_session_prompt_cache_key_is_scoped_to_parent_thread() {
let parent_thread_id = ThreadId::new();
let client = test_model_client(SessionSource::Internal(InternalSessionSource::Guardian));
let metadata = test_responses_metadata_for_client(
&client,
Some("turn-123"),
"window-1".to_string(),
Some(parent_thread_id),
TestCodexResponsesRequestKind::Turn,
);
assert_eq!(
client.prompt_cache_key(&metadata),
format!("guardian:{parent_thread_id}")
);
}
#[test]
fn build_subagent_headers_sets_internal_memory_consolidation_label() {
let client = test_model_client(SessionSource::Internal(
InternalSessionSource::MemoryConsolidation,
));
let headers = client.build_subagent_headers();
let value = headers
.get(X_OPENAI_SUBAGENT_HEADER)
.and_then(|value| value.to_str().ok());
assert_eq!(value, Some("memory_consolidation"));
assert_eq!(
headers.get("originator"),
Some(&http::HeaderValue::from_static("test_originator"))
);
}
#[test]
fn build_ws_client_metadata_includes_window_lineage_and_turn_metadata() {
let parent_thread_id = ThreadId::new();
let client = test_model_client(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 2,
agent_path: None,
agent_nickname: None,
agent_role: None,
}));
let thread_id = client.state.thread_id.to_string();
let expected_window_id = format!("{thread_id}:1");
let responses_metadata = test_responses_metadata_for_client(
&client,
Some("turn-123"),
expected_window_id.clone(),
Some(parent_thread_id),
TestCodexResponsesRequestKind::Turn,
);
let client_metadata =
client.build_ws_client_metadata(&responses_metadata, /*use_responses_lite*/ false);
let parent_thread_id = parent_thread_id.to_string();
let turn_metadata: serde_json::Value = serde_json::from_str(
client_metadata
.get(X_CODEX_TURN_METADATA_HEADER)
.expect("turn metadata"),
)
.expect("valid turn metadata");
for (client_key, metadata_key, expected) in [
(
X_CODEX_INSTALLATION_ID_HEADER,
"installation_id",
"11111111-1111-4111-8111-111111111111",
),
("session_id", "session_id", thread_id.as_str()),
("thread_id", "thread_id", thread_id.as_str()),
("turn_id", "turn_id", "turn-123"),
(
X_CODEX_WINDOW_ID_HEADER,
"window_id",
expected_window_id.as_str(),
),
(
X_CODEX_PARENT_THREAD_ID_HEADER,
"parent_thread_id",
parent_thread_id.as_str(),
),
] {
assert_eq!(
client_metadata.get(client_key).map(String::as_str),
Some(expected)
);
assert_eq!(turn_metadata[metadata_key].as_str(), Some(expected));
}
assert_eq!(
client_metadata
.get(X_OPENAI_SUBAGENT_HEADER)
.map(String::as_str),
Some("collab_spawn")
);
}
#[tokio::test]
async fn summarize_memories_returns_empty_for_empty_input() {
let client = test_model_client(SessionSource::Cli);
let model_info = test_model_info();
let session_telemetry = test_session_telemetry();
let output = client
.summarize_memories(
Vec::new(),
&model_info,
/*effort*/ None,
&session_telemetry,
)
.await
.expect("empty summarize request should succeed");
assert_eq!(output.len(), 0);
}
#[tokio::test]
async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Result<()> {
let temp = TempDir::new()?;
let attempt = started_inference_attempt(&temp)?;
// The provider has produced one complete output item, but no terminal
// response.completed event. The harness has enough information to keep this
// item in history, so the trace should preserve it when the stream is
// abandoned.
let item = output_message("1", "partial answer");
let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone(item))])
.chain(futures::stream::pending());
let (mut stream, _) = super::map_response_events(
/*upstream_request_id*/ None,
api_stream,
test_session_telemetry(),
attempt,
test_model_provider(),
);
let observed = stream
.next()
.await
.expect("mapped stream should yield output item")?;
assert!(matches!(observed, ResponseEvent::OutputItemDone(_)));
// Dropping the consumer is how turn interruption/preemption stops polling
// the provider stream. The mapper task observes that drop asynchronously
// and records cancellation using the output items it has already seen.
drop(stream);
// Cancellation is recorded by the mapper task after Drop wakes it, so the
// replay may need a short wait before the terminal event appears on disk.
let rollout = replay_until_cancelled(&temp).await?;
let inference = rollout
.inference_calls
.values()
.next()
.expect("inference should be reduced");
assert_eq!(inference.execution.status, ExecutionStatus::Cancelled);
assert_eq!(inference.response_item_ids.len(), 1);
assert_eq!(rollout.raw_payloads.len(), 2);
Ok(())
}
#[tokio::test]
async fn response_stream_records_last_model_feedback_ids() {
let tags = Arc::new(Mutex::new(BTreeMap::new()));
let _guard = tracing_subscriber::registry()
.with(TagCollectorLayer { tags: tags.clone() })
.set_default();
let api_stream = futures::stream::iter([
Ok(ResponseEvent::Created { response_id: None }),
Ok(ResponseEvent::Completed {
response_id: "resp-123".to_string(),
token_usage: None,
usage_metadata: None,
end_turn: Some(true),
}),
]);
let (mut stream, _) = super::map_response_events(
Some("req-123".to_string()),
api_stream,
test_session_telemetry(),
InferenceTraceAttempt::disabled(),
test_model_provider(),
);
while stream.next().await.is_some() {}
let tags = tags.lock().unwrap().clone();
assert_eq!(
tags.get("last_model_request_id").map(String::as_str),
Some("\"req-123\"")
);
assert_eq!(
tags.get("last_model_response_id").map(String::as_str),
Some("\"resp-123\"")
);
}
#[tokio::test]
async fn bedrock_unauthorized_error_uses_provider_mapping() {
let provider = create_model_provider(
ModelProviderInfo::create_amazon_bedrock_provider(/*aws*/ None),
/*auth_manager*/ None,
);
let mut auth_recovery = None;
let mut provider_auth_recovery_attempted = false;
let url = "https://bedrock-mantle.us-east-2.api.aws/openai/v1/responses";
let error = super::handle_unauthorized(
TransportError::Http {
status: http::StatusCode::UNAUTHORIZED,
url: Some(url.to_string()),
headers: None,
body: Some(
"Signature expired: 20260609T133205Z is now earlier than 20260614T062525Z"
.to_string(),
),
},
&mut auth_recovery,
&mut provider_auth_recovery_attempted,
&test_session_telemetry(),
&provider,
/*event_sender*/ None,
/*turn_id*/ None,
)
.await
.expect_err("expired Bedrock signature should fail");
assert_eq!(
error.to_string(),
format!(
"Amazon Bedrock rejected the request because its AWS signature has expired. Refresh your AWS credentials and retry. If `AWS_BEARER_TOKEN_BEDROCK` is set, update or unset it, then restart Codex, url: {url}"
)
);
}
#[derive(Debug)]
struct TestRecoveryProvider {
inner: SharedModelProvider,
should_fail: bool,
attempts: Arc<AtomicUsize>,
}
impl ModelProvider for TestRecoveryProvider {
fn info(&self) -> &ModelProviderInfo {
self.inner.info()
}
fn auth_manager(&self) -> Option<Arc<AuthManager>> {
None
}
fn auth(&self) -> ModelProviderFuture<'_, Option<CodexAuth>> {
self.inner.auth()
}
fn account_state(&self) -> ProviderAccountResult {
self.inner.account_state()
}
fn auth_recovery_messages(&self) -> Option<ProviderAuthRecoveryMessages> {
Some(ProviderAuthRecoveryMessages {
started: "Refreshing provider authentication.",
succeeded: "Provider authentication recovered.",
})
}
fn recover_from_unauthorized(
&self,
) -> ModelProviderFuture<'_, codex_protocol::error::Result<ProviderUnauthorizedRecovery>> {
self.attempts.fetch_add(1, Ordering::Relaxed);
Box::pin(async move {
if self.should_fail {
Err(CodexErr::Io(std::io::Error::other(
"provider recovery failed",
)))
} else {
Ok(ProviderUnauthorizedRecovery::Recovered)
}
})
}
fn models_manager(
&self,
codex_home: PathBuf,
config_model_catalog: Option<ModelsResponse>,
) -> SharedModelsManager {
self.inner.models_manager(codex_home, config_model_catalog)
}
}
#[tokio::test]
async fn provider_owned_auth_recovery_is_bounded_and_preserves_unauthorized_failures() {
for should_fail in [false, true] {
let attempts = Arc::new(AtomicUsize::new(0));
let provider: SharedModelProvider = Arc::new(TestRecoveryProvider {
inner: test_model_provider(),
should_fail,
attempts: Arc::clone(&attempts),
});
assert!(provider.auth_manager().is_none());
let unauthorized = || TransportError::Http {
status: http::StatusCode::UNAUTHORIZED,
url: Some("https://example.com/v1/responses".to_string()),
headers: None,
body: Some("unauthorized".to_string()),
};
let mut auth_recovery = None;
let mut provider_auth_recovery_attempted = false;
let telemetry = test_session_telemetry();
let (event_sender, event_receiver) = async_channel::unbounded();
let result = super::handle_unauthorized(
unauthorized(),
&mut auth_recovery,
&mut provider_auth_recovery_attempted,
&telemetry,
&provider,
Some(&event_sender),
Some("turn-1"),
)
.await;
let error = if should_fail {
result.expect_err("failed provider recovery should return the original error")
} else {
let recovered = result.expect("provider recovery should succeed without AuthManager");
assert_eq!(
(recovered.mode, recovered.phase),
("provider", "provider_refresh")
);
super::handle_unauthorized(
unauthorized(),
&mut auth_recovery,
&mut provider_auth_recovery_attempted,
&telemetry,
&provider,
Some(&event_sender),
Some("turn-1"),
)
.await
.expect_err("provider recovery should not run more than once")
};
match error.details() {
CodexErrorDetails::UnexpectedStatus(response) => {
assert_eq!(response.status, http::StatusCode::UNAUTHORIZED);
assert_eq!(response.body, "unauthorized");
}
other => panic!("unexpected error after provider recovery: {other}"),
}
assert_eq!(attempts.load(Ordering::Relaxed), 1);
let events = std::iter::from_fn(|| event_receiver.try_recv().ok())
.map(|event| serde_json::to_value(event).expect("recovery event should serialize"))
.collect::<Vec<_>>();
let mut expected = vec![json!({
"id": "turn-1",
"msg": {
"type": "auth_recovery_started",
"provider": provider.info().name,
"message": "Refreshing provider authentication.",
}
})];
if !should_fail {
expected.push(json!({
"id": "turn-1",
"msg": {
"type": "auth_recovery_completed",
"provider": provider.info().name,
"message": "Provider authentication recovered.",
}
}));
}
assert_eq!(events, expected);
}
}
#[tokio::test]
async fn dropped_backpressured_response_stream_traces_cancelled_partial_output()
-> anyhow::Result<()> {
let temp = TempDir::new()?;
let attempt = started_inference_attempt(&temp)?;
let backpressured_item_yielded = Arc::new(Notify::new());
let mut events = VecDeque::new();
for _ in 0..super::RESPONSE_STREAM_CHANNEL_CAPACITY {
events.push_back(ResponseEvent::Created { response_id: None });
}
events.push_back(ResponseEvent::OutputItemDone(output_message(
"1",
"partial answer",
)));
let api_stream = NotifyAfterEventStream {
events,
yielded: 0,
notify_after: super::RESPONSE_STREAM_CHANNEL_CAPACITY + 1,
notify: Arc::clone(&backpressured_item_yielded),
};
let (stream, _) = super::map_response_events(
/*upstream_request_id*/ None,
api_stream,
test_session_telemetry(),
attempt,
test_model_provider(),
);
// Fill the mapper channel with non-terminal events, then yield one output
// item. The mapper has observed that item and is blocked trying to send it
// downstream, so dropping the consumer covers the send-failure path rather
// than the `consumer_dropped` select branch.
backpressured_item_yielded.notified().await;
drop(stream);
let rollout = replay_until_cancelled(&temp).await?;
let inference = rollout
.inference_calls
.values()
.next()
.expect("inference should be reduced");
assert_eq!(inference.execution.status, ExecutionStatus::Cancelled);
assert_eq!(inference.response_item_ids.len(), 1);
assert_eq!(rollout.raw_payloads.len(), 2);
Ok(())
}
#[test]
fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
let auth_context = AuthRequestTelemetryContext::new(
Some(AuthMode::Chatgpt),
&BearerAuthProvider::for_test(Some("access-token"), Some("workspace-123")),
/*agent_identity_telemetry*/ None,
PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution {
mode: "managed",
phase: "refresh_token",
}),
);
assert_eq!(auth_context.auth_mode, Some("Chatgpt"));
assert!(auth_context.auth_header_attached);
assert_eq!(auth_context.auth_header_name, Some("authorization"));
assert!(auth_context.retry_after_unauthorized);
assert_eq!(auth_context.recovery_mode, Some("managed"));
assert_eq!(auth_context.recovery_phase, Some("refresh_token"));
}
#[test]
fn auth_request_telemetry_context_tracks_agent_identity_ids() {
let auth_context = AuthRequestTelemetryContext::new(
Some(AuthMode::Chatgpt),
&BearerAuthProvider::for_test(/*token*/ None, /*account_id*/ None),
Some(AgentIdentityTelemetry {
agent_id: "agent-runtime-context".to_string(),
task_id: "task-run-context".to_string(),
}),
PendingUnauthorizedRetry::default(),
);
assert_eq!(
auth_context.agent_identity_telemetry(),
Some(&AgentIdentityTelemetry {
agent_id: "agent-runtime-context".to_string(),
task_id: "task-run-context".to_string(),
})
);
}
fn model_client_with_counting_attestation(
include_attestation: bool,
) -> (ModelClient, Arc<AtomicUsize>) {
#[derive(Debug)]
struct CountingAttestationProvider {
calls: Arc<AtomicUsize>,
}
impl AttestationProvider for CountingAttestationProvider {
fn header_for_request(
&self,
_context: AttestationContext,
) -> GenerateAttestationFuture<'_> {
let calls = self.calls.clone();
Box::pin(async move {
let call = calls.fetch_add(1, Ordering::Relaxed) + 1;
Some(http::HeaderValue::from_bytes(format!("v1.header-{call}").as_bytes()).unwrap())
})
}
}
let attestation_calls = Arc::new(AtomicUsize::new(0));
let (auth_manager, provider) = if include_attestation {
(
Some(AuthManager::from_auth_for_testing(
CodexAuth::create_dummy_chatgpt_auth_for_testing(),
)),
ModelProviderInfo::create_openai_provider(Some(CHATGPT_CODEX_BASE_URL.to_string())),
)
} else {
(
None,
create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses),
)
};
let model_client = ModelClient::new(
auth_manager,
AgentIdentityAuthPolicy::JwtOnly,
ThreadId::new(),
provider,
SessionSource::Exec,
"test_originator".to_string(),
/*model_verbosity*/ None,
/*content_item_kinds_enabled*/ true,
/*enable_request_compression*/ false,
/*include_timing_metrics*/ false,
/*beta_features_header*/ None,
/*concurrent_reasoning_summaries_enabled*/ false,
Some(Arc::new(CountingAttestationProvider {
calls: attestation_calls.clone(),
})),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
codex_model_provider::WorkspaceRoutingContext::new(
"https://chatgpt.com/backend-api".into(),
),
);
(model_client, attestation_calls)
}
#[test]
fn thread_responses_headers_are_scoped_to_model_and_backend_auth() {
let (mut model_client, _) =
model_client_with_counting_attestation(/*include_attestation*/ true);
let headers = http::HeaderMap::from_iter([(
http::HeaderName::from_static("x-custom-request"),
http::HeaderValue::from_static("example"),
)]);
model_client.codex_responses_headers = Some(Arc::new(crate::CodexResponsesHeaders {
model: "selected-model".to_owned(),
headers: headers.clone(),
}));
let chatgpt_auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let api_key_auth = CodexAuth::from_api_key("test-api-key");
for (auth, model, expected) in [
(Some(&chatgpt_auth), "selected-model", headers),
(Some(&chatgpt_auth), "other-model", http::HeaderMap::new()),
(
Some(&api_key_auth),
"selected-model",
http::HeaderMap::new(),
),
(None, "selected-model", http::HeaderMap::new()),
] {
assert_eq!(model_client.responses_headers(auth, model), expected);
}
Arc::get_mut(&mut model_client.state)
.expect("test client should have unique session state")
.provider = create_model_provider(
ModelProviderInfo::create_openai_provider(Some("https://proxy.example.com/v1".to_owned())),
Some(AuthManager::from_auth_for_testing(chatgpt_auth.clone())),
);
assert_eq!(
model_client.responses_headers(Some(&chatgpt_auth), "selected-model"),
http::HeaderMap::new(),
);
}
#[test_case::test_case(/*cache_key*/ None; "own_cache")]
#[test_case::test_case(Some("parent-session"); "inherited_cache")]
#[tokio::test]
async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses(
cache_key: Option<&str>,
) {
let (mut model_client, attestation_calls) =
model_client_with_counting_attestation(/*include_attestation*/ true);
let responses_metadata = test_responses_metadata_for_client(
&model_client,
/*turn_id*/ None,
format!("{}:0", model_client.state.thread_id),
/*parent_thread_id*/ None,
TestCodexResponsesRequestKind::WebsocketConnection,
);
model_client.prompt_cache_key_override = cache_key.map(str::to_string);
let headers = model_client
.build_websocket_headers(&responses_metadata)
.await;
assert_eq!(
headers
.get(crate::attestation::X_OAI_ATTESTATION_HEADER)
.and_then(|value| value.to_str().ok()),
Some("v1.header-1"),
);
assert_eq!(attestation_calls.load(Ordering::Relaxed), 1);
assert_eq!(
headers["session-id"],
cache_key.unwrap_or(&responses_metadata.session_id)
);
assert_eq!(headers["thread-id"], responses_metadata.thread_id);
}
#[tokio::test]
async fn existing_call_sideband_headers_include_attestation() {
let (model_client, attestation_calls) =
model_client_with_counting_attestation(/*include_attestation*/ true);
let headers = model_client
.realtime_sideband_headers(http::HeaderMap::new())
.await
.expect("existing call sideband headers should build");
assert_eq!(
headers
.get(crate::attestation::X_OAI_ATTESTATION_HEADER)
.and_then(|value| value.to_str().ok()),
Some("v1.header-1"),
);
assert_eq!(attestation_calls.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn non_chatgpt_codex_endpoints_omit_attestation_generation() {
let (model_client, attestation_calls) =
model_client_with_counting_attestation(/*include_attestation*/ false);
let mut response_headers = http::HeaderMap::new();
if let Some(header_value) = model_client.generate_attestation_header_for().await {
response_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
}
let mut compaction_headers = http::HeaderMap::new();
if let Some(header_value) = model_client.generate_attestation_header_for().await {
compaction_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
}
let mut realtime_headers = http::HeaderMap::new();
if let Some(header_value) = model_client.generate_attestation_header_for().await {
realtime_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
}
assert_eq!(
response_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER),
None,
);
assert_eq!(
compaction_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER),
None,
);
assert_eq!(
realtime_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER),
None,
);
assert_eq!(attestation_calls.load(Ordering::Relaxed), 0);
}
Reference in New Issue View Git Blame Copy Permalink