Attach hosted app context to file uploads (#38101)

## What changed

- Include the connector ID, action name, and model in file creation requests for hosted app tool calls.
- Use the file size returned by upload finalization when available, while falling back to the local size for older server responses.
- Keep upload finalization requests empty for compatibility with older servers.

## Testing

- Cover hosted app upload metadata and legacy finalization responses.

GitOrigin-RevId: 1bc834736db4fe3515f8224965cde0e5d6955ca3
This commit is contained in:
wencongli-oai
2026-08-12 01:45:20 +00:00
committed by copyberry
parent ca4d532b2a
commit c909d1bc04
5 changed files with 149 additions and 15 deletions
+85 -4
View File
@@ -22,6 +22,13 @@ const OPENAI_FILE_FINALIZE_TIMEOUT: Duration = Duration::from_secs(30);
const OPENAI_FILE_FINALIZE_RETRY_DELAY: Duration = Duration::from_millis(250);
const OPENAI_FILE_USE_CASE: &str = "codex";
#[derive(Debug)]
pub struct HostedFileUploadContext {
pub connector_id: String,
pub action_name: String,
pub model: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UploadedOpenAiFile {
pub file_id: String,
@@ -101,6 +108,8 @@ struct DownloadLinkResponse {
file_name: Option<String>,
mime_type: Option<String>,
error_message: Option<String>,
#[serde(default)]
file_size_bytes: Option<u64>,
}
pub fn openai_file_uri(file_id: &str) -> String {
@@ -114,6 +123,7 @@ pub async fn upload_openai_file(
file_name: String,
file_size_bytes: u64,
contents: impl Stream<Item = std::io::Result<Bytes>> + Send + 'static,
hosted_upload: Option<&HostedFileUploadContext>,
) -> Result<UploadedOpenAiFile, OpenAiFileError> {
if file_size_bytes > OPENAI_FILE_UPLOAD_LIMIT_BYTES {
return Err(OpenAiFileError::FileTooLarge {
@@ -124,12 +134,25 @@ pub async fn upload_openai_file(
}
let create_url = format!("{}/files", base_url.trim_end_matches('/'));
let create_response = authorized_request(client_pool, auth, Method::POST, &create_url)
.json(&serde_json::json!({
let create_request = serde_json::json!({
"file_name": file_name.as_str(),
"file_size": file_size_bytes,
"use_case": OPENAI_FILE_USE_CASE,
});
let request = authorized_request(client_pool, auth, Method::POST, &create_url);
let create_request = match hosted_upload {
Some(context) => serde_json::json!({
"file_name": file_name.as_str(),
"file_size": file_size_bytes,
"use_case": OPENAI_FILE_USE_CASE,
}))
"codex_connector_id": context.connector_id,
"codex_action_name": context.action_name,
"codex_model": context.model,
}),
None => create_request,
};
let create_response = request
.json(&create_request)
.send()
.await
.map_err(|source| OpenAiFileError::Request {
@@ -233,10 +256,11 @@ pub async fn upload_openai_file(
base_url.trim_end_matches('/'),
create_payload.file_id,
);
let finalize_request = serde_json::json!({});
let finalize_started_at = Instant::now();
loop {
let finalize_response = authorized_request(client_pool, auth, Method::POST, &finalize_url)
.json(&serde_json::json!({}))
.json(&finalize_request)
.send()
.await
.map_err(|source| OpenAiFileError::Request {
@@ -260,6 +284,7 @@ pub async fn upload_openai_file(
match finalize_payload.status.as_str() {
"success" => {
let file_size_bytes = finalize_payload.file_size_bytes.unwrap_or(file_size_bytes);
return Ok(UploadedOpenAiFile {
file_id: create_payload.file_id.clone(),
uri: openai_file_uri(&create_payload.file_id),
@@ -425,6 +450,7 @@ mod tests {
"hello.txt".to_string(),
/*file_size_bytes*/ 5,
contents,
/*hosted_upload*/ None,
)
.await
.expect("upload succeeds");
@@ -440,6 +466,58 @@ mod tests {
assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn upload_hosted_app_preserves_empty_finalization_for_older_servers() {
let server = MockServer::start().await;
let hosted_upload = HostedFileUploadContext {
connector_id: "library".to_string(),
action_name: "create_library_file".to_string(),
model: "gpt-work".to_string(),
};
Mock::given(method("POST"))
.and(path("/backend-api/files"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"file_id": "file_pdf",
"upload_url": format!("{}/upload/file_pdf", server.uri()),
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("PUT"))
.and(path("/upload/file_pdf"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/files/file_pdf/uploaded"))
.and(body_json(serde_json::json!({})))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"status": "success",
"download_url": format!("{}/download/file_pdf", server.uri()),
"file_name": "report.pdf",
"mime_type": "application/pdf",
})))
.expect(1)
.mount(&server)
.await;
let uploaded = upload_openai_file(
&base_url_for(&server),
&chatgpt_auth(),
&default_http_client_pool(),
"report.pdf".to_string(),
/*file_size_bytes*/ 8,
futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"%PDF-1.4"))]),
Some(&hosted_upload),
)
.await
.expect("older servers retain the normal upload behavior");
assert_eq!(uploaded.file_size_bytes, 8);
server.verify().await;
}
#[tokio::test]
async fn upload_openai_file_reuses_client_pool_across_uploads() {
let server = MockServer::start().await;
@@ -497,6 +575,7 @@ mod tests {
file_name.to_string(),
u64::try_from(contents.len()).expect("file size should fit in a u64"),
contents_stream,
/*hosted_upload*/ None,
)
.await
.expect("upload succeeds with the shared client pool");
@@ -556,6 +635,7 @@ mod tests {
"hello.txt".to_string(),
/*file_size_bytes*/ 5,
futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"hello"))]),
/*hosted_upload*/ None,
)
.await
.expect_err("blob response failure should be returned");
@@ -593,6 +673,7 @@ mod tests {
"hello.txt".to_string(),
/*file_size_bytes*/ 5,
futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"hello"))]),
/*hosted_upload*/ None,
)
.await
.expect_err("blob transport failure should be returned");
+1
View File
@@ -71,6 +71,7 @@ pub use crate::endpoint::ResponsesWebsocketProbe;
pub use crate::endpoint::SearchClient;
pub use crate::endpoint::session_update_session_json;
pub use crate::error::ApiError;
pub use crate::files::HostedFileUploadContext;
pub use crate::files::OPENAI_FILE_UPLOAD_LIMIT_BYTES;
pub use crate::files::upload_openai_file;
pub use crate::images::ImageBackground;
+37 -10
View File
@@ -13,6 +13,7 @@
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use codex_api::HostedFileUploadContext;
use codex_api::OPENAI_FILE_UPLOAD_LIMIT_BYTES;
use codex_api::upload_openai_file;
use codex_http_client::RouteAwareClientPool;
@@ -20,11 +21,17 @@ use codex_login::CodexAuth;
use serde_json::Value as JsonValue;
use std::collections::HashMap;
struct FileArgumentLocation<'a> {
field_name: &'a str,
index: Option<usize>,
}
pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
sess: &Session,
step_context: &StepContext,
arguments_value: Option<JsonValue>,
openai_file_input_optional_fields: Option<&HashMap<String, Vec<String>>>,
hosted_upload: Option<&HostedFileUploadContext>,
) -> Result<Option<JsonValue>, String> {
let Some(openai_file_input_optional_fields) = openai_file_input_optional_fields else {
return Ok(arguments_value);
@@ -50,6 +57,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
field_name,
optional_fields,
value,
hosted_upload,
)
.await?
else {
@@ -72,6 +80,7 @@ async fn rewrite_argument_value_for_openai_files(
field_name: &str,
optional_fields: &[String],
value: &JsonValue,
hosted_upload: Option<&HostedFileUploadContext>,
) -> Result<Option<JsonValue>, String> {
match value {
JsonValue::String(file_path) => {
@@ -79,10 +88,13 @@ async fn rewrite_argument_value_for_openai_files(
step_context,
client_pool,
auth,
field_name,
/*index*/ None,
FileArgumentLocation {
field_name,
index: None,
},
optional_fields,
file_path,
hosted_upload,
)
.await?;
Ok(Some(rewritten))
@@ -97,10 +109,13 @@ async fn rewrite_argument_value_for_openai_files(
step_context,
client_pool,
auth,
field_name,
Some(index),
FileArgumentLocation {
field_name,
index: Some(index),
},
optional_fields,
file_path,
hosted_upload,
)
.await?;
rewritten_values.push(rewritten);
@@ -115,11 +130,12 @@ async fn build_uploaded_argument_value(
step_context: &StepContext,
client_pool: &RouteAwareClientPool,
auth: Option<&CodexAuth>,
field_name: &str,
index: Option<usize>,
argument: FileArgumentLocation<'_>,
optional_fields: &[String],
file_path: &str,
hosted_upload: Option<&HostedFileUploadContext>,
) -> Result<JsonValue, String> {
let FileArgumentLocation { field_name, index } = argument;
let contextualize_error = |error: String| match index {
Some(index) => {
format!("failed to upload `{file_path}` for `{field_name}[{index}]`: {error}")
@@ -184,6 +200,7 @@ async fn build_uploaded_argument_value(
file_name,
metadata.size,
contents,
hosted_upload,
)
.await
.map_err(|error| contextualize_error(error.to_string()))?;
@@ -255,6 +272,7 @@ mod tests {
&step_context,
arguments.clone(),
/*openai_file_input_optional_fields*/ None,
/*hosted_upload*/ None,
)
.await
.expect("rewrite should succeed");
@@ -353,10 +371,13 @@ mod tests {
&step_context,
&session.services.openai_file_upload_client_pool,
Some(&auth),
"file",
/*index*/ None,
FileArgumentLocation {
field_name: "file",
index: None,
},
&["mime_type".to_string(), "file_name".to_string()],
"file_report.csv",
/*hosted_upload*/ None,
)
.await
.expect("rewrite should upload the local file");
@@ -388,10 +409,13 @@ mod tests {
&step_context,
&session.services.openai_file_upload_client_pool,
Some(&auth),
"file",
/*index*/ None,
FileArgumentLocation {
field_name: "file",
index: None,
},
&[],
"oversized.bin",
/*hosted_upload*/ None,
)
.await
.expect_err("oversized file should be rejected");
@@ -465,6 +489,7 @@ mod tests {
"file",
&[],
&serde_json::json!("file_report.csv"),
/*hosted_upload*/ None,
)
.await
.expect("rewrite should succeed");
@@ -578,6 +603,7 @@ mod tests {
"files",
&[],
&serde_json::json!(["one.csv", "two.csv"]),
/*hosted_upload*/ None,
)
.await
.expect("rewrite should succeed");
@@ -611,6 +637,7 @@ mod tests {
"file": "/definitely/missing/file.csv",
})),
Some(&HashMap::from([("file".to_string(), Vec::new())])),
/*hosted_upload*/ None,
)
.await
.expect_err("missing file should fail");
+11
View File
@@ -23,6 +23,7 @@ use crate::turn_metadata::McpTurnMetadataContext;
use codex_analytics::AppInvocation;
use codex_analytics::InvocationType;
use codex_analytics::build_track_events_context;
use codex_api::HostedFileUploadContext;
use codex_config::ConfigLayerSource;
use codex_config::types::AppToolApproval;
use codex_config::types::ApprovalsReviewer;
@@ -434,11 +435,21 @@ async fn handle_approved_mcp_tool_call(
}
maybe_mark_thread_memory_mode_polluted(sess, turn_context, &prepared_call)
.await;
let hosted_upload = item_metadata
.connector_id
.as_ref()
.zip(item_metadata.action_name.as_ref())
.map(|(connector_id, action_name)| HostedFileUploadContext {
connector_id: connector_id.clone(),
action_name: action_name.clone(),
model: turn_context.model_info.slug.clone(),
});
let rewritten_arguments = rewrite_mcp_tool_arguments_for_openai_files(
sess,
step_context,
arguments_value,
metadata.openai_file_input_optional_fields.as_ref(),
hosted_upload.as_ref(),
)
.await
.map_err(anyhow::Error::msg)?;
+15 -1
View File
@@ -35,6 +35,7 @@ use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_json;
use wiremock::matchers::body_partial_json;
use wiremock::matchers::header;
use wiremock::matchers::method;
use wiremock::matchers::path;
@@ -101,7 +102,7 @@ async fn mount_file_upload_mocks(server: &MockServer, file_size_bytes: u64) {
Mock::given(method("POST"))
.and(path("/files"))
.and(header("chatgpt-account-id", "account_id"))
.and(body_json(json!({
.and(body_partial_json(json!({
"file_name": "report.txt",
"file_size": file_size_bytes,
"use_case": "codex",
@@ -233,6 +234,19 @@ async fn codex_apps_file_params_omit_fields_absent_from_tool_schema() -> Result<
"connector_id": "calendar",
}))
);
let requests = server.received_requests().await.expect("capture requests");
let upload_request = requests
.iter()
.find(|request| request.url.path() == "/files")
.expect("app tool should create a Files upload");
let upload_body: Value =
serde_json::from_slice(&upload_request.body).expect("Files request should be JSON");
assert_eq!(
upload_body.get("codex_connector_id"),
Some(&json!("calendar"))
);
assert_eq!(upload_body.get("upload_source"), None);
assert_eq!(upload_body.get("store_in_library"), None);
server.verify().await;
Ok(())