mirror of
https://github.com/openai/codex.git
synced 2026-11-03 19:53:25 +00:00
refactor generated image artifact store seam
This commit is contained in:
@@ -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<u8>,
|
||||
) -> ArtifactWriteFuture<'_>;
|
||||
}
|
||||
|
||||
/// Future returned by artifact writes.
|
||||
pub type ArtifactWriteFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = Result<AbsolutePathBuf>> + 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<u8>,
|
||||
) -> 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
|
||||
}
|
||||
@@ -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),
|
||||
}))
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<AnalyticsEventsClient>,
|
||||
pub(crate) thread_store: Arc<dyn ThreadStore>,
|
||||
pub(crate) artifact_store: Arc<dyn ArtifactStore>,
|
||||
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
pub(crate) inherited_multi_agent_version: Option<MultiAgentVersion>,
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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<EnvironmentManager>,
|
||||
analytics_events_client: Option<AnalyticsEventsClient>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
artifact_store: Arc<dyn ArtifactStore>,
|
||||
parent_rollout_thread_trace: ThreadTraceContext,
|
||||
attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
multi_agent_version: Option<MultiAgentVersion>,
|
||||
@@ -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)),
|
||||
|
||||
@@ -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<std::sync::Mutex<Vec<Vec<u8>>>>,
|
||||
}
|
||||
|
||||
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<u8>,
|
||||
) -> 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(), "<image_id>");
|
||||
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;
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
@@ -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<StateDbHandle>,
|
||||
pub(crate) live_thread: Option<LiveThread>,
|
||||
pub(crate) thread_store: Arc<dyn ThreadStore>,
|
||||
pub(crate) artifact_store: Arc<dyn ArtifactStore>,
|
||||
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
/// Session-scoped model client shared across turns.
|
||||
pub(crate) model_client: ModelClient,
|
||||
|
||||
@@ -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, "<image_id>");
|
||||
let image_output_path = sess
|
||||
.services
|
||||
.artifact_store
|
||||
.generated_image_path(&session_id, "<image_id>");
|
||||
let image_output_dir = image_output_path
|
||||
.parent()
|
||||
.unwrap_or_else(|| turn_context.config.codex_home.clone());
|
||||
|
||||
@@ -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,<svg/>",
|
||||
|
||||
@@ -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<McpManager>,
|
||||
extensions: Arc<ExtensionRegistry<Config>>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
artifact_store: Arc<dyn ArtifactStore>,
|
||||
attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
session_source: SessionSource,
|
||||
installation_id: String,
|
||||
@@ -262,6 +265,36 @@ impl ThreadManager {
|
||||
state_db: Option<StateDbHandle>,
|
||||
installation_id: String,
|
||||
attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
) -> 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<AuthManager>,
|
||||
session_source: SessionSource,
|
||||
environment_manager: Arc<EnvironmentManager>,
|
||||
extensions: Arc<ExtensionRegistry<Config>>,
|
||||
analytics_events_client: Option<AnalyticsEventsClient>,
|
||||
thread_store: Arc<dyn ThreadStore>,
|
||||
artifact_store: Arc<dyn ArtifactStore>,
|
||||
state_db: Option<StateDbHandle>,
|
||||
installation_id: String,
|
||||
attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
||||
) -> 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<dyn ArtifactStore> =
|
||||
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,
|
||||
}))
|
||||
|
||||
@@ -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<u8>,
|
||||
) -> 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<dyn ArtifactStore> = 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 = [
|
||||
|
||||
Reference in New Issue
Block a user