hanzoai/dev
2
0
Fork 0
mirror of https://github.com/hanzoai/dev.git synced 2026-10-25 03:32:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): persist image generation response metadata

Browse Source
This commit is contained in:
James Peter
2026-04-24 11:53:34 +10:00
parent 2d5bb06b73
commit 168f965cc6
6 changed files with 197 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
code-rs/core/src/chat_completions.rs
+3
View File
@@ -1116,6 +1116,9 @@ where
Poll::Ready(Some(Ok(ResponseEvent::ModelsEtag(etag)))) => {
return Poll::Ready(Some(Ok(ResponseEvent::ModelsEtag(etag))));
}
Poll::Ready(Some(Ok(ResponseEvent::ResponseHeaders(headers)))) => {
return Poll::Ready(Some(Ok(ResponseEvent::ResponseHeaders(headers))));
}
Poll::Ready(Some(Ok(ResponseEvent::Completed {
response_id,
token_usage,
code-rs/core/src/client.rs
+21 -1
View File
@@ -915,6 +915,15 @@ impl ModelClient {
Ok(Ok((mut ws_stream, response))) => {
let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(1600);
let response_headers = header_map_to_json(response.headers());
if tx_event
.send(Ok(ResponseEvent::ResponseHeaders(response_headers)))
.await
.is_err()
{
debug!("receiver dropped response headers event");
}
if let Some(value) = response
.headers()
.get(X_CODEX_TURN_STATE_HEADER)
@@ -1386,12 +1395,22 @@ impl ModelClient {
"x_request_id": resp.headers()
.get("x-request-id")
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.unwrap_or_default(),
"headers": header_map_to_json(resp.headers()),
}),
);
}
let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(1600);
let response_headers = header_map_to_json(resp.headers());
if tx_event
.send(Ok(ResponseEvent::ResponseHeaders(response_headers)))
.await
.is_err()
{
debug!("receiver dropped response headers event");
}
if let Some(snapshot) = parse_rate_limit_snapshot(resp.headers()) {
debug!(
"rate limit headers:\n{}",
@@ -1686,6 +1705,7 @@ impl ModelClient {
"error",
&serde_json::json!({
"status": status.as_u16(),
"headers": header_map_to_json(&headers),
"body": body_text
}),
);
code-rs/core/src/client_common.rs
+1
View File
@@ -256,6 +256,7 @@ pub enum ResponseEvent {
response_id: Option<String>,
response_model: Option<String>,
},
ResponseHeaders(serde_json::Value),
OutputItemDone { item: ResponseItem, sequence_number: Option<u64>, output_index: Option<u32> },
/// Indicates that the server will include reasoning content on this stream.
///
code-rs/core/src/codex/streaming.rs
+123 -1
View File
@@ -74,6 +74,26 @@ const AUTO_CONTEXT_JUDGE_DEVELOPER_MESSAGE: &str = concat!(
"only when nearby context appears genuinely essential to finishing the active thread correctly."
);
#[derive(Clone, Debug, Default)]
struct ImageGenerationTurnMetadata {
requested_model: String,
latest_response_model: Option<String>,
response_headers: Option<serde_json::Value>,
}
#[derive(serde::Serialize)]
struct ImageGenerationSidecar<'a> {
call_id: &'a str,
status: &'a str,
revised_prompt: Option<&'a str>,
artifact_path: String,
requested_model: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
latest_response_model: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
response_headers: Option<&'a serde_json::Value>,
}
/// A series of Turns in response to user input.
pub(super) struct AgentTask {
sess: Arc<Session>,
@@ -4015,6 +4035,7 @@ async fn try_run_turn(
.clone()
.unwrap_or_else(|| sess.client.get_model());
let mut latest_response_model: Option<String> = None;
let mut latest_response_headers: Option<serde_json::Value> = None;
let mut stream = match sess.client.clone().stream(&prompt).await {
Ok(stream) => stream,
Err(e) => {
@@ -4105,6 +4126,9 @@ async fn try_run_turn(
);
}
ResponseEvent::ServerReasoningIncluded(_included) => {}
ResponseEvent::ResponseHeaders(headers) => {
latest_response_headers = Some(headers);
}
ResponseEvent::OutputItemDone { item, sequence_number, output_index } => {
let (item, rollout_ids) = crate::memories::sanitize_response_item(item);
if !rollout_ids.is_empty() {
@@ -4114,7 +4138,21 @@ async fn try_run_turn(
});
}
let response =
handle_response_item(sess, turn_diff_tracker, sub_id, item.clone(), sequence_number, output_index, attempt_req).await?;
handle_response_item(
sess,
turn_diff_tracker,
sub_id,
item.clone(),
sequence_number,
output_index,
attempt_req,
&ImageGenerationTurnMetadata {
requested_model: requested_model.clone(),
latest_response_model: latest_response_model.clone(),
response_headers: latest_response_headers.clone(),
},
)
.await?;
// Save into scratchpad so we can seed a retry if the stream drops later.
sess.scratchpad_push(&item, &response, &sub_id);
@@ -4300,6 +4338,7 @@ async fn handle_response_item(
seq_hint: Option<u64>,
output_index: Option<u32>,
attempt_req: u64,
image_generation_metadata: &ImageGenerationTurnMetadata,
) -> CodexResult<Option<ResponseInputItem>> {
debug!(?item, "Output item");
let output = match item {
@@ -4490,6 +4529,7 @@ async fn handle_response_item(
seq_hint,
output_index,
attempt_req,
image_generation_metadata,
)
.await;
None
@@ -4510,6 +4550,7 @@ async fn handle_image_generation_call(
seq_hint: Option<u64>,
output_index: Option<u32>,
attempt_req: u64,
metadata: &ImageGenerationTurnMetadata,
) {
let order = crate::protocol::OrderMeta {
request_ordinal: attempt_req,
@@ -4569,6 +4610,22 @@ async fn handle_image_generation_call(
}
};
if let Some(path) = saved_path.as_ref()
&& let Err(err) = save_image_generation_sidecar(
path,
&call_id,
&status,
revised_prompt.as_deref(),
metadata,
)
.await
{
warn!(
"failed to save image generation metadata sidecar for {}: {err}",
path.display()
);
}
let end = sess.make_event_with_order(
sub_id,
EventMsg::ImageGenerationEnd(crate::protocol::ImageGenerationEndEvent {
@@ -4631,6 +4688,31 @@ async fn save_image_generation_result(
.map_err(|err| err.to_string())
}
async fn save_image_generation_sidecar(
artifact_path: &code_utils_absolute_path::AbsolutePathBuf,
call_id: &str,
status: &str,
revised_prompt: Option<&str>,
metadata: &ImageGenerationTurnMetadata,
) -> std::result::Result<code_utils_absolute_path::AbsolutePathBuf, String> {
let sidecar_path = artifact_path.as_path().with_extension("metadata.json");
let sidecar = ImageGenerationSidecar {
call_id,
status,
revised_prompt,
artifact_path: artifact_path.display().to_string(),
requested_model: &metadata.requested_model,
latest_response_model: metadata.latest_response_model.as_deref(),
response_headers: metadata.response_headers.as_ref(),
};
let json = serde_json::to_vec_pretty(&sidecar).map_err(|err| err.to_string())?;
tokio::fs::write(&sidecar_path, json)
.await
.map_err(|err| err.to_string())?;
code_utils_absolute_path::AbsolutePathBuf::from_absolute_path(sidecar_path)
.map_err(|err| err.to_string())
}
fn web_search_query(query: &Option<String>, queries: &Option<Vec<String>>) -> Option<String> {
if let Some(value) = query.clone().filter(|q| !q.is_empty()) {
return Some(value);
@@ -14107,6 +14189,8 @@ mod tests {
is_context_overflow_stream_error,
is_usage_limit_stream_error,
save_image_generation_result,
save_image_generation_sidecar,
ImageGenerationTurnMetadata,
spark_fallback_model,
TRUNCATION_MARKER,
};
@@ -14153,6 +14237,44 @@ mod tests {
assert_eq!(std::fs::read(saved_path.as_path()).expect("saved file"), b"foo");
}
#[tokio::test]
async fn save_image_generation_sidecar_writes_metadata() {
let dir = TempDir::new().expect("tempdir");
let saved_path = save_image_generation_result(dir.path(), "session-1", "ig_123", "Zm9v")
.await
.expect("image should save");
let metadata = ImageGenerationTurnMetadata {
requested_model: "gpt-5.4".to_string(),
latest_response_model: Some("gpt-5.4-2026-04-01".to_string()),
response_headers: Some(serde_json::json!({
"x-request-id": ["req_123"],
})),
};
let sidecar_path = save_image_generation_sidecar(
&saved_path,
"ig_123",
"completed",
Some("A tiny square"),
&metadata,
)
.await
.expect("metadata should save");
assert_eq!(
sidecar_path.as_path(),
saved_path.as_path().with_extension("metadata.json")
);
let sidecar: Value = serde_json::from_slice(
&std::fs::read(sidecar_path.as_path()).expect("sidecar file"),
)
.expect("sidecar json");
assert_eq!(sidecar["call_id"], "ig_123");
assert_eq!(sidecar["requested_model"], "gpt-5.4");
assert_eq!(sidecar["latest_response_model"], "gpt-5.4-2026-04-01");
assert_eq!(sidecar["response_headers"]["x-request-id"][0], "req_123");
}
#[tokio::test]
async fn save_image_generation_result_rejects_non_standard_base64() {
let dir = TempDir::new().expect("tempdir");
code-rs/core/src/model_family.rs
+18 -5
View File
@@ -4,6 +4,7 @@ use crate::config_types::ReasoningEffort;
use crate::config_types::ReasoningSummary;
use crate::tool_apply_patch::ApplyPatchToolType;
use code_protocol::openai_models::ConfigShellToolType;
use code_protocol::openai_models::InputModality;
use code_protocol::openai_models::ModelInfo;
use code_protocol::openai_models::ModelsResponse;
use code_protocol::openai_models::TruncationMode;
@@ -37,7 +38,6 @@ const CONTEXT_WINDOW_96K: u64 = 96_000;
const CONTEXT_WINDOW_16K: u64 = 16_385;
const CONTEXT_WINDOW_1M: u64 = 1_047_576;
const MAX_OUTPUT_DEFAULT: u64 = 128_000;
const IMAGE_GENERATION_TOOL: &str = "image_generation";
static UPSTREAM_MODELS: Lazy<Vec<ModelInfo>> = Lazy::new(|| {
serde_json::from_str::<ModelsResponse>(include_str!("../../../codex-rs/models-manager/models.json"))
@@ -225,10 +225,7 @@ fn apply_upstream_model_overrides(mut family: ModelFamily) -> ModelFamily {
family.supports_parallel_tool_calls = model_info.supports_parallel_tool_calls;
family.web_search_tool_type = model_info.web_search_tool_type;
family.supports_image_detail_original = model_info.supports_image_detail_original;
family.supports_image_generation = model_info
.experimental_supported_tools
.iter()
.any(|tool| tool == IMAGE_GENERATION_TOOL);
family.supports_image_generation = supports_image_generation(model_info);
family.uses_local_shell_tool = matches!(model_info.shell_type, ConfigShellToolType::Local);
family.uses_shell_command_tool =
matches!(model_info.shell_type, ConfigShellToolType::ShellCommand);
@@ -481,6 +478,22 @@ pub fn derive_default_model_family(model: &str) -> ModelFamily {
})
}
fn supports_image_generation(model_info: &ModelInfo) -> bool {
model_info.input_modalities.contains(&InputModality::Image)
}
#[cfg(test)]
mod tests {
use super::find_family_for_model;
#[test]
fn image_generation_support_tracks_image_input_modality() {
let family = find_family_for_model("gpt-5.4").expect("known upstream model");
assert!(family.supports_image_generation);
}
}
impl ModelFamily {
/// Token limit at which we should automatically compact, if known.
pub fn auto_compact_token_limit(&self) -> Option<i64> {
code-rs/core/src/remote_models/mod.rs
+31 -5
View File
@@ -7,6 +7,7 @@ use code_app_server_protocol::AuthMode;
use code_protocol::config_types::ReasoningSummary as ProtocolReasoningSummary;
use code_protocol::config_types::Personality as ProtocolPersonality;
use code_protocol::openai_models::ApplyPatchToolType as ProtocolApplyPatchToolType;
use code_protocol::openai_models::InputModality;
use code_protocol::openai_models::ModelInfo;
use code_protocol::openai_models::ModelsResponse;
use code_protocol::openai_models::ReasoningEffort as ProtocolReasoningEffort;
@@ -31,7 +32,6 @@ const MODEL_CACHE_FILE: &str = "models_cache.json";
const DEFAULT_MODEL_CACHE_TTL: Duration = Duration::from_secs(300);
const REMOTE_MODELS_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
const CODEX_AUTO_BALANCED_MODEL: &str = "codex-auto-balanced";
const IMAGE_GENERATION_TOOL: &str = "image_generation";
#[derive(Debug, Default, Clone)]
struct RemoteModelsState {
@@ -420,10 +420,7 @@ fn apply_model_info_overrides_with_personality(
family.web_search_tool_type = map_web_search_tool_type(info.web_search_tool_type);
family.supports_image_detail_original = info.supports_image_detail_original;
family.supports_image_generation = info
.experimental_supported_tools
.iter()
.any(|tool| tool == IMAGE_GENERATION_TOOL);
family.supports_image_generation = supports_image_generation(info);
family.additional_speed_tiers = info.additional_speed_tiers.clone();
family.supports_search_tool = info.supports_search_tool;
@@ -445,9 +442,12 @@ fn apply_model_info_overrides_with_personality(
#[cfg(test)]
mod tests {
use super::apply_model_info_overrides_with_personality;
use super::derive_default_model_family;
use super::find_remote_model_info;
use code_protocol::config_types::ReasoningSummary;
use code_protocol::openai_models::ConfigShellToolType;
use code_protocol::openai_models::InputModality;
use code_protocol::openai_models::ModelInfo;
use code_protocol::openai_models::ModelVisibility;
use code_protocol::openai_models::TruncationPolicyConfig;
@@ -506,6 +506,28 @@ mod tests {
assert!(find_remote_model_info(&models, "foo/bar/gpt-5.3-codex").is_none());
}
#[test]
fn image_generation_support_tracks_image_input_modality() {
let mut family = derive_default_model_family("gpt-5.4");
family.supports_image_generation = false;
let mut info = model("gpt-5.4");
info.input_modalities = vec![InputModality::Text, InputModality::Image];
let family = apply_model_info_overrides_with_personality(&info, family, None);
assert!(family.supports_image_generation);
let mut text_only_info = info;
text_only_info.input_modalities = vec![InputModality::Text];
let text_only_family = apply_model_info_overrides_with_personality(
&text_only_info,
derive_default_model_family("gpt-5.4"),
None,
);
assert!(!text_only_family.supports_image_generation);
}
}
fn map_web_search_tool_type(tool_type: WebSearchToolType) -> WebSearchToolType {
@@ -515,6 +537,10 @@ fn map_web_search_tool_type(tool_type: WebSearchToolType) -> WebSearchToolType {
}
}
fn supports_image_generation(info: &ModelInfo) -> bool {
info.input_modalities.contains(&InputModality::Image)
}
fn map_personality(personality: ConfigPersonality) -> ProtocolPersonality {
match personality {
ConfigPersonality::None => ProtocolPersonality::None,