diff --git a/codex-rs/core/src/artifact_store.rs b/codex-rs/core/src/artifact_store.rs new file mode 100644 index 0000000000..51d4aebf56 --- /dev/null +++ b/codex-rs/core/src/artifact_store.rs @@ -0,0 +1,80 @@ +use std::future::Future; +use std::pin::Pin; + +use codex_protocol::error::Result; +use codex_utils_absolute_path::AbsolutePathBuf; + +const GENERATED_IMAGE_ARTIFACTS_DIR: &str = "generated_images"; + +/// Stores generated artifacts and materializes executor-readable local paths. +pub trait ArtifactStore: Send + Sync + 'static { + /// Returns the executor-readable path exposed to generated image instructions. + fn generated_image_path(&self, session_id: &str, call_id: &str) -> AbsolutePathBuf; + + /// Persists a generated image payload and returns its executor-readable path. + fn write_generated_image( + &self, + session_id: &str, + call_id: &str, + bytes: Vec, + ) -> ArtifactWriteFuture<'_>; +} + +/// Future returned by artifact writes. +pub type ArtifactWriteFuture<'a> = + Pin> + Send + 'a>>; + +#[derive(Clone)] +pub struct LocalArtifactStore { + codex_home: AbsolutePathBuf, +} + +impl LocalArtifactStore { + pub fn from_codex_home(codex_home: &AbsolutePathBuf) -> Self { + Self { + codex_home: codex_home.clone(), + } + } +} + +impl ArtifactStore for LocalArtifactStore { + fn generated_image_path(&self, session_id: &str, call_id: &str) -> AbsolutePathBuf { + self.codex_home + .join(GENERATED_IMAGE_ARTIFACTS_DIR) + .join(sanitize_artifact_path_component(session_id)) + .join(format!("{}.png", sanitize_artifact_path_component(call_id))) + } + + fn write_generated_image( + &self, + session_id: &str, + call_id: &str, + bytes: Vec, + ) -> ArtifactWriteFuture<'_> { + let path = self.generated_image_path(session_id, call_id); + Box::pin(async move { + if let Some(parent) = path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + tokio::fs::write(&path, bytes).await?; + Ok(path) + }) + } +} + +fn sanitize_artifact_path_component(value: &str) -> String { + let mut sanitized: String = value + .chars() + .map(|ch| { + if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { + ch + } else { + '_' + } + }) + .collect(); + if sanitized.is_empty() { + sanitized = "generated_image".to_string(); + } + sanitized +} diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 678a6e5d10..69d37ad558 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -105,6 +105,7 @@ pub(crate) async fn run_codex_thread_interactive( environment_selections: parent_ctx.environments.clone(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), + artifact_store: Arc::clone(&parent_session.services.artifact_store), attestation_provider: parent_session.services.attestation_provider.clone(), inherited_multi_agent_version: Some(MultiAgentVersion::Disabled), })) diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 31ac529556..64a6175f29 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -24,6 +24,7 @@ pub use codex_thread::CodexThreadSettingsOverrides; pub use codex_thread::ThreadConfigSnapshot; pub use session::turn_context::TurnContext; mod agent; +pub mod artifact_store; mod attestation; mod codex_delegate; mod command_canonicalization; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5a92c9c825..f15f8e8b32 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -13,6 +13,7 @@ use crate::agent::AgentControl; use crate::agent::AgentStatus; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; +use crate::artifact_store::ArtifactStore; use crate::attestation::AttestationProvider; use crate::build_available_skills; use crate::compact; @@ -419,6 +420,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, + pub(crate) artifact_store: Arc, pub(crate) attestation_provider: Option>, pub(crate) inherited_multi_agent_version: Option, } @@ -499,6 +501,7 @@ impl Codex { environment_selections, analytics_events_client, thread_store, + artifact_store, attestation_provider, inherited_multi_agent_version, } = args; @@ -645,6 +648,7 @@ impl Codex { environment_manager, analytics_events_client, thread_store, + artifact_store, parent_rollout_thread_trace, attestation_provider, multi_agent_version, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 05773abb8c..7ef3899065 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1,5 +1,6 @@ use super::input_queue::InputQueue; use super::*; +use crate::artifact_store::ArtifactStore; use crate::config::ConstraintError; use crate::goals::GoalRuntimeState; use crate::skills::SkillError; @@ -504,6 +505,7 @@ impl Session { environment_manager: Arc, analytics_events_client: Option, thread_store: Arc, + artifact_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, attestation_provider: Option>, multi_agent_version: Option, @@ -1030,6 +1032,7 @@ impl Session { state_db: state_db_ctx.clone(), live_thread: live_thread_init.as_ref().cloned(), thread_store: Arc::clone(&thread_store), + artifact_store, attestation_provider: attestation_provider.clone(), model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index a3e5558971..93496e851c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1,5 +1,7 @@ use super::turn_context::TurnEnvironment; use super::*; +use crate::artifact_store::ArtifactStore; +use crate::artifact_store::ArtifactWriteFuture; use crate::config::ConfigBuilder; use crate::config::ConfigOverrides; use crate::config::test_config; @@ -141,6 +143,7 @@ use codex_protocol::protocol::W3cTraceContext; use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; use codex_rmcp_client::ElicitationAction; +use codex_utils_absolute_path::AbsolutePathBuf; use core_test_support::PathBufExt; use core_test_support::PathExt; use core_test_support::context_snapshot; @@ -190,6 +193,32 @@ struct InstructionsTestCase { expects_apply_patch_description: bool, } +#[derive(Clone)] +struct FakeArtifactStore { + writes: Arc>>>, +} + +impl ArtifactStore for FakeArtifactStore { + fn generated_image_path(&self, session_id: &str, call_id: &str) -> AbsolutePathBuf { + let path = format!("/fake/{session_id}/{call_id}.png"); + test_path_buf(&path).abs() + } + + fn write_generated_image( + &self, + session_id: &str, + call_id: &str, + bytes: Vec, + ) -> ArtifactWriteFuture<'_> { + let writes = Arc::clone(&self.writes); + let path = self.generated_image_path(session_id, call_id); + Box::pin(async move { + writes.lock().expect("lock writes").push(bytes); + Ok(path) + }) + } +} + fn user_message(text: &str) -> ResponseItem { ResponseItem::Message { id: None, @@ -4603,6 +4632,9 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, )), + Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -4763,6 +4795,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, )), + artifact_store: Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )), attestation_provider: None, model_client: ModelClient::new( Some(auth_manager.clone()), @@ -4953,6 +4988,9 @@ async fn make_session_with_config_and_rx( codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, )), + Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -5065,6 +5103,9 @@ async fn make_session_with_history_source_and_agent_control_and_rx( .expect("state db should initialize"), ), )), + Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )), codex_rollout_trace::ThreadTraceContext::disabled(), /*attestation_provider*/ None, Some(config.multi_agent_version_from_features()), @@ -6853,6 +6894,9 @@ where codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), state_db, )), + artifact_store: Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )), attestation_provider: None, model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), @@ -7844,6 +7888,69 @@ async fn handle_output_item_done_records_image_save_history_message() { let _ = std::fs::remove_file(&expected_saved_path); } +#[tokio::test] +async fn handle_output_item_done_uses_injected_artifact_store() { + let (mut session, turn_context) = make_session_and_context().await; + let writes = Arc::new(std::sync::Mutex::new(Vec::new())); + session.services.artifact_store = Arc::new(FakeArtifactStore { + writes: Arc::clone(&writes), + }); + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + let call_id = "ig_injected_store"; + let expected_saved_path = session + .services + .artifact_store + .generated_image_path(&session.conversation_id.to_string(), call_id); + let item = ResponseItem::ImageGenerationCall { + id: call_id.to_string(), + status: "completed".to_string(), + revised_prompt: Some("a tiny blue square".to_string()), + result: "Zm9v".to_string(), + }; + + let mut ctx = HandleOutputCtx { + sess: Arc::clone(&session), + turn_context: Arc::clone(&turn_context), + turn_store: Arc::new(codex_extension_api::ExtensionData::new( + turn_context.sub_id.clone(), + )), + tool_runtime: test_tool_runtime(Arc::clone(&session), Arc::clone(&turn_context)), + cancellation_token: CancellationToken::new(), + }; + handle_output_item_done(&mut ctx, item.clone(), /*previously_active_item*/ None) + .await + .expect("image generation item should succeed"); + + assert_eq!( + writes.lock().expect("lock writes").as_slice(), + &[b"foo".to_vec()] + ); + let image_output_path = session + .services + .artifact_store + .generated_image_path(&session.conversation_id.to_string(), ""); + let image_output_dir = image_output_path + .parent() + .expect("generated image path should have a parent"); + let image_message: ResponseItem = crate::context::ContextualUserFragment::into( + crate::context::ImageGenerationInstructions::new( + image_output_dir.display(), + image_output_path.display(), + ), + ); + let history = session.clone_history().await; + assert_eq!(history.raw_items(), &[image_message, item]); + assert_eq!( + expected_saved_path, + test_path_buf(&format!( + "/fake/{}/ig_injected_store.png", + session.conversation_id + )) + .abs() + ); +} + #[tokio::test] async fn handle_output_item_done_skips_image_save_message_when_save_fails() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index d8b46c9999..21ac52a195 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -702,6 +702,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { codex_thread_store::LocalThreadStoreConfig::from_config(&config), /*state_db*/ None, )); + let artifact_store = Arc::new(crate::artifact_store::LocalArtifactStore::from_codex_home( + &config.codex_home, + )); let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs { config, @@ -733,6 +736,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { }, analytics_events_client: None, thread_store, + artifact_store, attestation_provider: None, inherited_multi_agent_version: None, }) diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 59def9c7ab..728d5dca8f 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::SkillsManager; use crate::agent::AgentControl; +use crate::artifact_store::ArtifactStore; use crate::attestation::AttestationProvider; use crate::client::ModelClient; use crate::config::NetworkProxyAuditMetadata; @@ -74,6 +75,7 @@ pub(crate) struct SessionServices { pub(crate) state_db: Option, pub(crate) live_thread: Option, pub(crate) thread_store: Arc, + pub(crate) artifact_store: Arc, pub(crate) attestation_provider: Option>, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 7914c8d482..a27b55a5b6 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -36,34 +36,17 @@ use tracing::debug; use tracing::instrument; use tracing::warn; -const GENERATED_IMAGE_ARTIFACTS_DIR: &str = "generated_images"; +use crate::artifact_store::ArtifactStore; +#[cfg(test)] +use crate::artifact_store::LocalArtifactStore; +#[cfg(test)] pub(crate) fn image_generation_artifact_path( codex_home: &AbsolutePathBuf, session_id: &str, call_id: &str, ) -> AbsolutePathBuf { - let sanitize = |value: &str| { - let mut sanitized: String = value - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch - } else { - '_' - } - }) - .collect(); - if sanitized.is_empty() { - sanitized = "generated_image".to_string(); - } - sanitized - }; - - codex_home - .join(GENERATED_IMAGE_ARTIFACTS_DIR) - .join(sanitize(session_id)) - .join(format!("{}.png", sanitize(call_id))) + LocalArtifactStore::from_codex_home(codex_home).generated_image_path(session_id, call_id) } fn strip_hidden_assistant_markup(text: &str, plan_mode: bool) -> String { @@ -108,7 +91,7 @@ pub(crate) fn raw_assistant_output_text_from_item(item: &ResponseItem) -> Option } async fn save_image_generation_result( - codex_home: &AbsolutePathBuf, + artifact_store: &dyn ArtifactStore, session_id: &str, call_id: &str, result: &str, @@ -118,12 +101,9 @@ async fn save_image_generation_result( .map_err(|err| { CodexErr::InvalidRequest(format!("invalid image generation payload: {err}")) })?; - let path = image_generation_artifact_path(codex_home, session_id, call_id); - if let Some(parent) = path.parent() { - tokio::fs::create_dir_all(parent).await?; - } - tokio::fs::write(&path, bytes).await?; - Ok(path) + artifact_store + .write_generated_image(session_id, call_id, bytes) + .await } pub(crate) async fn persist_image_generation_item( @@ -134,7 +114,7 @@ pub(crate) async fn persist_image_generation_item( image_item.saved_path = None; let session_id = sess.conversation_id.to_string(); match save_image_generation_result( - &turn_context.config.codex_home, + sess.services.artifact_store.as_ref(), &session_id, &image_item.id, &image_item.result, @@ -146,11 +126,10 @@ pub(crate) async fn persist_image_generation_item( Some(path) } Err(err) => { - let output_path = image_generation_artifact_path( - &turn_context.config.codex_home, - &session_id, - &image_item.id, - ); + let output_path = sess + .services + .artifact_store + .generated_image_path(&session_id, &image_item.id); let output_dir = output_path .parent() .unwrap_or_else(|| turn_context.config.codex_home.clone()); @@ -173,8 +152,10 @@ async fn record_image_generation_instructions( return; } let session_id = sess.conversation_id.to_string(); - let image_output_path = - image_generation_artifact_path(&turn_context.config.codex_home, &session_id, ""); + let image_output_path = sess + .services + .artifact_store + .generated_image_path(&session_id, ""); let image_output_dir = image_output_path .parent() .unwrap_or_else(|| turn_context.config.codex_home.clone()); diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 9c4be94576..641183da6f 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -8,6 +8,7 @@ use super::image_generation_artifact_path; use super::last_assistant_message_from_item; use super::response_item_may_include_external_context; use super::save_image_generation_result; +use crate::artifact_store::LocalArtifactStore; use crate::session::tests::make_session_and_context; use crate::tools::ToolRouter; use crate::tools::parallel::ToolCallRuntime; @@ -424,10 +425,14 @@ async fn save_image_generation_result_saves_base64_to_png_in_codex_home() { let expected_path = image_generation_artifact_path(&codex_home, "session-1", "ig_save_base64"); let _ = std::fs::remove_file(&expected_path); - let saved_path = - save_image_generation_result(&codex_home, "session-1", "ig_save_base64", "Zm9v") - .await - .expect("image should be saved"); + let saved_path = save_image_generation_result( + &LocalArtifactStore::from_codex_home(&codex_home), + "session-1", + "ig_save_base64", + "Zm9v", + ) + .await + .expect("image should be saved"); assert_eq!(saved_path, expected_path); assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo"); @@ -440,9 +445,14 @@ async fn save_image_generation_result_rejects_data_url_payload() { let codex_home = tempfile::tempdir().expect("create codex home"); let codex_home = codex_home.path().abs(); - let err = save_image_generation_result(&codex_home, "session-1", "ig_456", result) - .await - .expect_err("data url payload should error"); + let err = save_image_generation_result( + &LocalArtifactStore::from_codex_home(&codex_home), + "session-1", + "ig_456", + result, + ) + .await + .expect_err("data url payload should error"); assert!(matches!(err, CodexErr::InvalidRequest(_))); } @@ -459,9 +469,14 @@ async fn save_image_generation_result_overwrites_existing_file() { .expect("create image output dir"); std::fs::write(&existing_path, b"existing").expect("seed existing image"); - let saved_path = save_image_generation_result(&codex_home, "session-1", "ig_overwrite", "Zm9v") - .await - .expect("image should be saved"); + let saved_path = save_image_generation_result( + &LocalArtifactStore::from_codex_home(&codex_home), + "session-1", + "ig_overwrite", + "Zm9v", + ) + .await + .expect("image should be saved"); assert_eq!(saved_path, existing_path); assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo"); @@ -475,9 +490,14 @@ async fn save_image_generation_result_sanitizes_call_id_for_codex_home_output_pa let expected_path = image_generation_artifact_path(&codex_home, "session-1", "../ig/.."); let _ = std::fs::remove_file(&expected_path); - let saved_path = save_image_generation_result(&codex_home, "session-1", "../ig/..", "Zm9v") - .await - .expect("image should be saved"); + let saved_path = save_image_generation_result( + &LocalArtifactStore::from_codex_home(&codex_home), + "session-1", + "../ig/..", + "Zm9v", + ) + .await + .expect("image should be saved"); assert_eq!(saved_path, expected_path); assert_eq!(std::fs::read(&saved_path).expect("saved file"), b"foo"); @@ -488,9 +508,14 @@ async fn save_image_generation_result_sanitizes_call_id_for_codex_home_output_pa async fn save_image_generation_result_rejects_non_standard_base64() { let codex_home = tempfile::tempdir().expect("create codex home"); let codex_home = codex_home.path().abs(); - let err = save_image_generation_result(&codex_home, "session-1", "ig_urlsafe", "_-8") - .await - .expect_err("non-standard base64 should error"); + let err = save_image_generation_result( + &LocalArtifactStore::from_codex_home(&codex_home), + "session-1", + "ig_urlsafe", + "_-8", + ) + .await + .expect_err("non-standard base64 should error"); assert!(matches!(err, CodexErr::InvalidRequest(_))); } @@ -499,7 +524,7 @@ async fn save_image_generation_result_rejects_non_base64_data_urls() { let codex_home = tempfile::tempdir().expect("create codex home"); let codex_home = codex_home.path().abs(); let err = save_image_generation_result( - &codex_home, + &LocalArtifactStore::from_codex_home(&codex_home), "session-1", "ig_svg", "data:image/svg+xml,", diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d130a35815..49c2299b3f 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,5 +1,7 @@ use crate::SkillsManager; use crate::agent::AgentControl; +use crate::artifact_store::ArtifactStore; +use crate::artifact_store::LocalArtifactStore; use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; @@ -208,6 +210,7 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, extensions: Arc>, thread_store: Arc, + artifact_store: Arc, attestation_provider: Option>, session_source: SessionSource, installation_id: String, @@ -262,6 +265,36 @@ impl ThreadManager { state_db: Option, installation_id: String, attestation_provider: Option>, + ) -> Self { + Self::new_with_artifact_store( + config, + auth_manager, + session_source, + environment_manager, + extensions, + analytics_events_client, + thread_store, + Arc::new(LocalArtifactStore::from_codex_home(&config.codex_home)), + state_db, + installation_id, + attestation_provider, + ) + } + + #[allow(clippy::too_many_arguments)] + /// Constructs a thread manager with an explicit generated artifact backend. + pub fn new_with_artifact_store( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + environment_manager: Arc, + extensions: Arc>, + analytics_events_client: Option, + thread_store: Arc, + artifact_store: Arc, + state_db: Option, + installation_id: String, + attestation_provider: Option>, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -287,6 +320,7 @@ impl ThreadManager { mcp_manager, extensions, thread_store, + artifact_store, attestation_provider, auth_manager, session_source, @@ -362,7 +396,7 @@ impl ThreadManager { )); let mcp_manager = Arc::new(McpManager::new(Arc::clone(&plugins_manager))); let skills_manager = Arc::new(SkillsManager::new_with_restriction_product( - skills_codex_home, + skills_codex_home.clone(), /*bundled_skills_enabled*/ true, restriction_product, )); @@ -376,6 +410,8 @@ impl ThreadManager { }, state_db.clone(), )); + let artifact_store: Arc = + Arc::new(LocalArtifactStore::from_codex_home(&skills_codex_home)); Self { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), @@ -388,6 +424,7 @@ impl ThreadManager { mcp_manager, extensions: empty_extension_registry(), thread_store, + artifact_store, attestation_provider: None, auth_manager, session_source: SessionSource::Exec, @@ -1334,6 +1371,7 @@ impl ThreadManagerState { environment_selections, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), + artifact_store: Arc::clone(&self.artifact_store), attestation_provider: self.attestation_provider.clone(), inherited_multi_agent_version: multi_agent_version, })) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 52a0f8eb3b..a7d27ad56f 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,4 +1,6 @@ use super::*; +use crate::artifact_store::ArtifactStore; +use crate::artifact_store::ArtifactWriteFuture; use crate::config::test_config; use crate::init_state_db; use crate::installation_id::INSTALLATION_ID_FILENAME; @@ -33,6 +35,23 @@ use wiremock::MockServer; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; +struct FakeArtifactStore; + +impl ArtifactStore for FakeArtifactStore { + fn generated_image_path(&self, _session_id: &str, _call_id: &str) -> AbsolutePathBuf { + unreachable!("generated image path should not be requested") + } + + fn write_generated_image( + &self, + _session_id: &str, + _call_id: &str, + _bytes: Vec, + ) -> ArtifactWriteFuture<'_> { + unreachable!("generated image write should not be requested") + } +} + fn user_msg(text: &str) -> ResponseItem { ResponseItem::Message { id: None, @@ -64,6 +83,31 @@ fn developer_interrupted_marker() -> ResponseItem { .expect("developer interrupted marker should be enabled") } +#[tokio::test] +async fn start_thread_uses_injected_artifact_store() { + let config = test_config().await; + let artifact_store: Arc = Arc::new(FakeArtifactStore); + let manager = ThreadManager::new_with_artifact_store( + &config, + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + empty_extension_registry(), + /*analytics_events_client*/ None, + thread_store_from_config(&config, /*state_db*/ None), + Arc::clone(&artifact_store), + /*state_db*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + ); + let thread = manager.start_thread(config).await.expect("start thread"); + + assert!(Arc::ptr_eq( + &thread.thread.codex.session.services.artifact_store, + &artifact_store + )); +} + #[test] fn truncates_before_requested_user_message() { let items = [