From 446b771049d41617c66d7475d292b010e9165c71 Mon Sep 17 00:00:00 2001 From: Krish Chainani Date: Tue, 15 Sep 2026 00:14:24 +0000 Subject: [PATCH] Add attachment upload and resolution APIs and pass stores into sessions (#45556) ## What changed Replace `AttachmentStore::persist` with `upload` and `resolve`. Uploads return inline bytes or a file ID; resolution returns optional file metadata and a download URL only when a minimum URL lifetime is requested. Add image dimensions, digest, size, MIME type, and categorized errors to the API. Make `InlineAttachmentStore` return the original bytes and report `NotFound` when resolving file IDs. Redact attachment bytes and file URLs in debug output. Pass the configured image store from the thread manager into sessions and inherit it in delegated sessions. Allow `TestCodexBuilder` to accept a custom image store. ## Testing Update unit coverage for preserving PNG and JPEG bytes, redacting bytes and URLs in debug output, and rejecting file resolution in the inline store. GitOrigin-RevId: f313e0048925998d84394c7e5a278e0d39d86a45 --- codex-rs/Cargo.lock | 3 +- codex-rs/attachment-store/Cargo.toml | 4 +- codex-rs/attachment-store/src/lib.rs | 206 ++++++++++++++---- codex-rs/attachment-store/src/lib_tests.rs | 95 ++++---- codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/session.rs | 3 + codex-rs/core/src/session/tests.rs | 5 + .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 1 + codex-rs/core/tests/common/Cargo.toml | 1 + codex-rs/core/tests/common/test_codex.rs | 10 +- 13 files changed, 247 insertions(+), 89 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 14d6dc3e5e..a278dff942 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2481,8 +2481,6 @@ dependencies = [ name = "codex-attachment-store" version = "0.0.0" dependencies = [ - "base64 0.22.1", - "codex-utils-image", "pretty_assertions", "serde", "tokio", @@ -5497,6 +5495,7 @@ dependencies = [ "assert_cmd", "base64 0.22.1", "codex-arg0", + "codex-attachment-store", "codex-config", "codex-core", "codex-exec-server", diff --git a/codex-rs/attachment-store/Cargo.toml b/codex-rs/attachment-store/Cargo.toml index 285dce0af0..f2a1f4e5b9 100644 --- a/codex-rs/attachment-store/Cargo.toml +++ b/codex-rs/attachment-store/Cargo.toml @@ -13,10 +13,8 @@ path = "src/lib.rs" workspace = true [dependencies] -codex-utils-image = { workspace = true } serde = { workspace = true, features = ["derive"] } [dev-dependencies] -base64 = { workspace = true } pretty_assertions = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt"] } +tokio = { workspace = true, features = ["macros", "rt", "sync"] } diff --git a/codex-rs/attachment-store/src/lib.rs b/codex-rs/attachment-store/src/lib.rs index 201614e4ff..df586a9fd1 100644 --- a/codex-rs/attachment-store/src/lib.rs +++ b/codex-rs/attachment-store/src/lib.rs @@ -4,91 +4,211 @@ use std::error::Error; use std::fmt; use std::future::Future; use std::pin::Pin; +use std::time::Duration; -use codex_utils_image::data_url_from_bytes; use serde::Deserialize; use serde::Serialize; /// Result returned by [`AttachmentStore`] operations. pub type AttachmentStoreResult = Result; -/// Future returned by [`AttachmentStore::persist`]. -pub type AttachmentStoreFuture<'a> = - Pin> + Send + 'a>>; +/// Future returned by [`AttachmentStore::upload`]. +pub type UploadFuture<'a> = + Pin> + Send + 'a>>; -/// Stores attachment bytes and returns durable references to them. +/// Future returned by [`AttachmentStore::resolve`]. +pub type ResolveFuture<'a> = + Pin> + Send + 'a>>; + +/// Uploads attachments and resolves metadata for stored files. /// -/// Implementations perform any required I/O before resolving the returned future. A successful -/// reference must remain valid when callers persist it or pass it across subsystem boundaries. +/// Implementations choose whether uploaded attachments remain inline or become file references. +/// File references returned by [`AttachmentStore::upload`] must be resolvable by +/// [`AttachmentStore::resolve`]. pub trait AttachmentStore: Send + Sync { - /// Persists `data` and returns a durable reference to it. - fn persist<'a>( - &'a self, - data: &'a [u8], - metadata: &'a AttachmentMetadata, - ) -> AttachmentStoreFuture<'a>; + /// Uploads an attachment and returns the representation callers should retain. + fn upload(&self, request: UploadRequest) -> UploadFuture<'_>; + + /// Resolves metadata and, when requested, a fresh URL for a stored attachment. + fn resolve<'a>(&'a self, request: ResolveRequest<'a>) -> ResolveFuture<'a>; } -/// Metadata used to identify and interpret an attachment. -#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] -pub struct AttachmentMetadata { - /// Name associated with the attachment. - pub file_name: String, - /// IANA media type used to interpret the attachment bytes. - pub media_type: String, +/// Attachment data supplied to [`AttachmentStore::upload`]. +#[derive(Clone, Eq, PartialEq)] +pub struct UploadRequest { + /// Optional name associated with the attachment. + pub file_name: Option, + /// Attachment bytes to persist. + pub data: Vec, } -/// Durable reference to bytes stored by an [`AttachmentStore`]. -#[derive(Clone, Deserialize, Eq, PartialEq, Serialize)] -pub struct AttachmentRef { - /// Identifier assigned to an externally stored file. - pub file_id: Option, - /// URL referring to the attachment bytes, such as a data URL or - /// `sediment://`. - pub url: String, -} - -impl fmt::Debug for AttachmentRef { +impl fmt::Debug for UploadRequest { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter - .debug_struct("AttachmentRef") - .field("file_id", &self.file_id) - .field("url", &"") + .debug_struct("UploadRequest") + .field("file_name", &self.file_name) + .field("data", &"") .finish() } } +/// Parameters supplied to [`AttachmentStore::resolve`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ResolveRequest<'a> { + /// Identifier assigned by the attachment store. + pub file_id: &'a str, + /// Minimum remaining lifetime required for the returned file URL. + /// + /// When absent, implementations must omit the file URL and may return cached metadata. When + /// present, implementations must return a URL that remains valid for at least this duration + /// after the resolve operation completes. + pub download_url_ttl: Option, +} + +/// Metadata used to validate and interpret an attachment. +#[derive(Clone, Default, Deserialize, Eq, PartialEq, Serialize)] +pub struct AttachmentMetadata { + /// Name associated with the attachment. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub file_name: Option, + /// MD5 digest of the attachment bytes. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub digest: Option, + /// Length of the attachment byte payload. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub size_bytes: Option, + /// IANA media type used to interpret the attachment bytes. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mime_type: Option, + /// Metadata specific to the attachment's format. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub format_specific: Option, + /// URL for reading the attachment bytes when requested during resolution. + /// + /// Implementations must omit this field when [`ResolveRequest::download_url_ttl`] is absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub file_url: Option, +} + +impl fmt::Debug for AttachmentMetadata { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let file_url = self.file_url.as_ref().map(|_| ""); + formatter + .debug_struct("AttachmentMetadata") + .field("file_name", &self.file_name) + .field("digest", &self.digest) + .field("size_bytes", &self.size_bytes) + .field("mime_type", &self.mime_type) + .field("format_specific", &self.format_specific) + .field("file_url", &file_url) + .finish() + } +} + +/// Metadata specific to an attachment format. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(tag = "kind", rename_all = "lowercase")] +pub enum FormatSpecificMetadata { + /// Metadata for an image attachment. + Image(ImageMetadata), +} + +/// Metadata specific to an image attachment. +#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +pub struct ImageMetadata { + /// Image width in pixels. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub width: Option, + /// Image height in pixels. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub height: Option, +} + +/// Representation callers should retain after uploading an attachment. +#[derive(Clone, Deserialize, Eq, PartialEq, Serialize)] +pub enum UploadResult { + /// Attachment data remains inline. + Inline { + /// Original attachment bytes passed to [`AttachmentStore::upload`]. + bytes: Vec, + }, + /// Attachment data is available through an external file reference. + File { + /// Identifier assigned by the attachment store. + file_id: String, + }, +} + +impl fmt::Debug for UploadResult { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Inline { .. } => formatter + .debug_struct("Inline") + .field("bytes", &"") + .finish(), + Self::File { file_id } => formatter + .debug_struct("File") + .field("file_id", file_id) + .finish(), + } + } +} + /// Leaves attachments inline instead of storing them. pub struct InlineAttachmentStore; impl AttachmentStore for InlineAttachmentStore { - fn persist<'a>( - &'a self, - data: &'a [u8], - metadata: &'a AttachmentMetadata, - ) -> AttachmentStoreFuture<'a> { + fn upload(&self, request: UploadRequest) -> UploadFuture<'_> { Box::pin(async move { - Ok(AttachmentRef { - file_id: None, - url: data_url_from_bytes(&metadata.media_type, data), + Ok(UploadResult::Inline { + bytes: request.data, }) }) } + + fn resolve<'a>(&'a self, request: ResolveRequest<'a>) -> ResolveFuture<'a> { + let file_id = request.file_id; + Box::pin(async move { + Err(AttachmentStoreError::new( + AttachmentStoreErrorKind::NotFound, + format!("attachment `{file_id}` was not found"), + )) + }) + } +} + +/// Category of failure returned by an [`AttachmentStore`]. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum AttachmentStoreErrorKind { + /// The requested attachment could not be resolved. + NotFound, + /// The attachment data or metadata is invalid. + InvalidAttachment, + /// The backing store could not complete the operation. + Backend, } /// Error returned by an [`AttachmentStore`] implementation. #[derive(Clone, Debug, Eq, PartialEq)] pub struct AttachmentStoreError { + kind: AttachmentStoreErrorKind, message: String, } impl AttachmentStoreError { /// Creates an error without exposing backend-specific error types. - pub fn new(message: impl Into) -> Self { + pub fn new(kind: AttachmentStoreErrorKind, message: impl Into) -> Self { Self { + kind, message: message.into(), } } + + /// Returns the category callers can use to handle this error. + #[must_use] + pub const fn kind(&self) -> AttachmentStoreErrorKind { + self.kind + } } impl fmt::Display for AttachmentStoreError { diff --git a/codex-rs/attachment-store/src/lib_tests.rs b/codex-rs/attachment-store/src/lib_tests.rs index 800d66b6bd..f6fe878541 100644 --- a/codex-rs/attachment-store/src/lib_tests.rs +++ b/codex-rs/attachment-store/src/lib_tests.rs @@ -1,60 +1,75 @@ -use base64::Engine; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use pretty_assertions::assert_eq; use super::*; -/// Debug output retains the Sediment file ID while redacting credential-bearing URLs. +/// Debug output redacts attachment bytes in requests and results. #[test] -fn attachment_ref_debug_redacts_url() { - let attachment = AttachmentRef { - file_id: Some("sediment-file-id".to_string()), - url: "https://attachments.test/file?signed=secret".to_string(), +fn attachment_debug_output_redacts_bytes() { + let request = UploadRequest { + file_name: Some("image.png".to_string()), + data: b"secret".to_vec(), + }; + let result = UploadResult::Inline { + bytes: b"secret".to_vec(), }; assert_eq!( - format!("{attachment:?}"), - r#"AttachmentRef { file_id: Some("sediment-file-id"), url: "" }"# + (format!("{request:?}"), format!("{result:?}")), + ( + r#"UploadRequest { file_name: Some("image.png"), data: "" }"#.to_string(), + r#"Inline { bytes: "" }"#.to_string(), + ) ); } -/// The inline store preserves binary, text, PNG, and JPEG attachment bytes. +/// Debug output redacts credential-bearing file URLs in resolved metadata. +#[test] +fn attachment_metadata_debug_output_redacts_file_url() { + let metadata = AttachmentMetadata { + file_url: Some("https://attachments.test/file?signed=secret".to_string()), + ..AttachmentMetadata::default() + }; + + let debug = format!("{metadata:?}"); + + assert!(!debug.contains("signed=secret")); + assert!(debug.contains("")); +} + +/// The inline store preserves PNG and JPEG attachment bytes. #[tokio::test] -async fn inline_store_preserves_attachment_bytes() { - let cases: [(&str, &str, &[u8]); 4] = [ - ( - "attachment.bin", - "application/octet-stream", - b"\x00\x01\x7f\x80\xfe\xff", - ), - ("note.txt", "text/plain", b"hello\n"), - ("image.png", "image/png", b"\x89PNG\r\n\x1a\n"), - ( - "image.jpg", - "image/jpeg", - b"\xff\xd8\xff\xe0JFIF\x00\xff\xd9", - ), +async fn inline_store_preserves_image_bytes() { + let cases: [(&str, &[u8]); 2] = [ + ("image.png", b"\x89PNG\r\n\x1a\n"), + ("image.jpg", b"\xff\xd8\xff\xe0JFIF\x00\xff\xd9"), ]; - for (file_name, media_type, data) in cases { - let metadata = AttachmentMetadata { - file_name: file_name.to_string(), - media_type: media_type.to_string(), - }; + for (file_name, data) in cases { let attachment = InlineAttachmentStore - .persist(data, &metadata) + .upload(UploadRequest { + file_name: Some(file_name.to_string()), + data: data.to_vec(), + }) .await .expect("inline attachment"); - let (data_url_metadata, encoded) = - attachment.url.split_once(',').expect("data URL payload"); - let decoded = BASE64_STANDARD - .decode(encoded) - .expect("valid base64 payload"); - let expected_data_url_metadata = format!("data:{media_type};base64"); + let UploadResult::Inline { bytes } = attachment else { + panic!("inline store returned a file reference"); + }; - assert_eq!( - (attachment.file_id, data_url_metadata, decoded), - (None, expected_data_url_metadata.as_str(), data.to_vec()) - ); + assert_eq!(bytes, data); } } + +/// The inline store reports file references as missing because it never uploads them. +#[tokio::test] +async fn inline_store_cannot_resolve_file_references() { + let error = InlineAttachmentStore + .resolve(ResolveRequest { + file_id: "file_123", + download_url_ttl: None, + }) + .await + .expect_err("inline store cannot resolve file references"); + + assert_eq!(error.kind(), AttachmentStoreErrorKind::NotFound); +} diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9592ed32eb..51007b553f 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -127,6 +127,7 @@ pub(crate) async fn run_codex_thread_interactive( client_mcp_extensions: parent_session.services.client_mcp_extensions.clone(), reserved_thread_id: None, analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), + image_store: Arc::clone(&parent_session.services.image_store), thread_store: Arc::clone(&parent_session.services.thread_store), attestation_provider: parent_session.services.attestation_provider.clone(), external_time_provider: Some(Arc::clone(&parent_session.services.time_provider)), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index df065f8feb..affecf42d5 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -64,6 +64,7 @@ use codex_analytics::ImagePreparationMetadata; use codex_analytics::SubAgentThreadStartedInput; use codex_analytics::TurnCodexErrorFact; use codex_async_utils::OrCancelExt; +use codex_attachment_store::AttachmentStore; use codex_connectors::connector_runtime_context_key; use codex_context_fragments::RenderedFragment; use codex_exec_server::Environment; @@ -468,6 +469,7 @@ pub(crate) struct SessionSpawnArgs { pub(crate) client_mcp_extensions: ClientMcpExtensions, pub(crate) reserved_thread_id: Option, pub(crate) analytics_events_client: Option, + pub(crate) image_store: Arc, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, pub(crate) external_time_provider: Option>, @@ -570,6 +572,7 @@ impl Session { client_mcp_extensions, reserved_thread_id, analytics_events_client, + image_store, thread_store, attestation_provider, external_time_provider, @@ -890,6 +893,7 @@ impl Session { environment_manager, inherited_environments, analytics_events_client, + image_store, thread_store, parent_rollout_thread_trace, attestation_provider, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 66c82df206..80eb5f70ad 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -21,6 +21,7 @@ use crate::shell_snapshot::ShellSnapshot; use crate::shell_snapshot::SnapshotCredentialBrokerState; use crate::state::ActiveTurn; use crate::turn_metadata::ExecutionMetadata; +use codex_attachment_store::AttachmentStore; use codex_extension_api::ExtensionDataInit; use codex_http_client::ClientRouteClass; use codex_http_client::RouteAwareClientPool; @@ -753,6 +754,7 @@ impl Session { environment_manager: Arc, inherited_environments: Option, analytics_events_client: Option, + image_store: Arc, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, @@ -1636,6 +1638,7 @@ impl Session { network_approval: Arc::clone(&network_approval), state_db: state_db_ctx.clone(), live_thread: live_thread.clone(), + image_store, thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), time_provider, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 56ca71705d..f2d047d0ea 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5824,6 +5824,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, + crate::passthrough_image_store(), Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, @@ -6061,6 +6062,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { network_approval: Arc::clone(&network_approval), state_db: None, live_thread: None, + image_store: crate::passthrough_image_store(), thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, @@ -6309,6 +6311,7 @@ async fn make_session_with_config_and_rx( environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, + crate::passthrough_image_store(), Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, @@ -6439,6 +6442,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( environment_manager, /*inherited_environments*/ None, /*analytics_events_client*/ None, + crate::passthrough_image_store(), Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), Some( @@ -8239,6 +8243,7 @@ where network_approval: Arc::clone(&network_approval), state_db: state_db.clone(), live_thread: None, + image_store: crate::passthrough_image_store(), thread_store: Arc::new(codex_thread_store::LocalThreadStore::new( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), state_db, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index a96e75472c..72ae069562 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1302,6 +1302,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { client_mcp_extensions: ClientMcpExtensions::default(), reserved_thread_id: None, analytics_events_client: None, + image_store: crate::thread_manager::passthrough_image_store(), thread_store, attestation_provider: None, external_time_provider: None, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 6248b2c977..84c660116e 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -22,6 +22,7 @@ use crate::unified_exec::UnifiedExecProcessManager; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; +use codex_attachment_store::AttachmentStore; use codex_core_plugins::PluginsManager; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; @@ -88,6 +89,7 @@ pub(crate) struct SessionServices { pub(crate) network_approval: Arc, pub(crate) state_db: Option, pub(crate) live_thread: Option, + pub(crate) image_store: Arc, pub(crate) thread_store: Arc, pub(crate) attestation_provider: Option>, pub(crate) time_provider: Arc, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 7449838971..5b3cd9b299 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -2191,6 +2191,7 @@ impl ThreadManagerState { client_mcp_extensions, reserved_thread_id, analytics_events_client: self.analytics_events_client.clone(), + image_store: Arc::clone(&self.image_store), thread_store: Arc::clone(&self.thread_store), attestation_provider: self.attestation_provider.clone(), external_time_provider: self.external_time_provider.clone(), diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 864bb83458..ad4eab6088 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -16,6 +16,7 @@ anyhow = { workspace = true } assert_cmd = { workspace = true } base64 = { workspace = true } codex-arg0 = { workspace = true } +codex-attachment-store = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 65f3aa9e6d..e51b7adf52 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -13,6 +13,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; +use codex_attachment_store::AttachmentStore; use codex_config::CloudConfigBundleLoader; use codex_core::CodexThread; pub use codex_core::StartThreadOptions; @@ -342,6 +343,7 @@ pub struct TestCodexBuilder { history_mode: Option, models_manager: Option, thread_store: Option>, + image_store: Arc, } impl TestCodexBuilder { @@ -368,6 +370,11 @@ impl TestCodexBuilder { self } + pub fn with_image_store(mut self, image_store: Arc) -> Self { + self.image_store = image_store; + self + } + pub fn with_model(self, model: &str) -> Self { let new_model = model.to_string(); self.with_config(move |config| { @@ -739,7 +746,7 @@ impl TestCodexBuilder { Arc::new(extensions.build()), user_instructions_provider, /*analytics_events_client*/ None, - codex_core::passthrough_image_store(), + Arc::clone(&self.image_store), Arc::clone(&thread_store), codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()), installation_id, @@ -1403,6 +1410,7 @@ pub fn test_codex() -> TestCodexBuilder { history_mode: None, models_manager: None, thread_store: None, + image_store: codex_core::passthrough_image_store(), } }