diff --git a/ARCHITECTURE.zh-CN.md b/ARCHITECTURE.zh-CN.md index fd37f426..a64797ee 100644 --- a/ARCHITECTURE.zh-CN.md +++ b/ARCHITECTURE.zh-CN.md @@ -256,6 +256,19 @@ async fn handler( 新增事件必须遵循上述两级 camelCase 规范， 现有不一致的事件在相关模块迭代时逐步统一。 +### ACP 工具输出清洗 + +ACP Agent 的工具调用事件在 `aionui-ai-agent` 翻译层进入统一 +`AgentStreamEvent`。该边界必须保证 WebSocket 和 SQLite 消息内容 +是可控大小，不能把工具返回的大型二进制或 inline base64 原样透传。 + +Codex 图片生成工具可能同时返回 `saved_path` 和 `raw_output.result` +中的 PNG/JPEG/WebP base64。翻译层会在转发和持久化之前移除 +`result`，保留 `saved_path`、`image.path`、`result_omitted`、 +`result_bytes` 等小型结构化字段；如果图片已经落盘，则把工具状态 +归一化为 `completed`。前端应通过文件路径按需加载图片，不应依赖 +inline base64。 + ## 数据层 ### Repository Trait 模式 diff --git a/crates/aionui-ai-agent/src/protocol/events/mod.rs b/crates/aionui-ai-agent/src/protocol/events/mod.rs index 45a33a67..9b506e33 100644 --- a/crates/aionui-ai-agent/src/protocol/events/mod.rs +++ b/crates/aionui-ai-agent/src/protocol/events/mod.rs @@ -328,6 +328,182 @@ mod tests { assert!(json["data"]["update"].get("rawInput").is_none()); } + #[test] + fn codex_image_tool_update_omits_base64_result() { + let large_png_base64 = format!("iVBORw0KGgo{}", "A".repeat(128 * 1024)); + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "ig_test_image", + ToolCallUpdateFields::new() + .status(SdkToolCallStatus::InProgress) + .raw_output(json!({ + "call_id": "ig_test_image", + "status": "generating", + "saved_path": "/Users/test/.codex/generated_images/session/ig_test_image.png", + "revised_prompt": "一只小猫", + "result": large_png_base64 + })), + )), + ); + + let events = session_notification_to_events(¬if); + assert_eq!(events.len(), 1); + let json = serde_json::to_value(&events[0]).unwrap(); + let raw_output = &json["data"]["update"]["rawOutput"]; + + assert_eq!( + raw_output["saved_path"], + "/Users/test/.codex/generated_images/session/ig_test_image.png" + ); + assert_eq!( + raw_output["image"]["path"], + "/Users/test/.codex/generated_images/session/ig_test_image.png" + ); + assert_eq!(raw_output["result_omitted"], true); + assert!(raw_output.get("result").is_none()); + } + + #[test] + fn codex_image_tool_update_with_saved_path_is_completed() { + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "ig_done_image", + ToolCallUpdateFields::new() + .status(SdkToolCallStatus::InProgress) + .raw_output(json!({ + "call_id": "ig_done_image", + "status": "generating", + "saved_path": "/Users/test/.codex/generated_images/session/ig_done_image.png", + "result": format!("iVBORw0KGgo{}", "A".repeat(128 * 1024)) + })), + )), + ); + + let events = session_notification_to_events(¬if); + assert_eq!(events.len(), 1); + let json = serde_json::to_value(&events[0]).unwrap(); + + assert_eq!(json["data"]["update"]["status"], "completed"); + assert_eq!(json["data"]["update"]["rawOutput"]["status"], "completed"); + } + + #[test] + fn codex_image_tool_update_keeps_small_text_result_in_progress() { + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "small-result", + ToolCallUpdateFields::new() + .status(SdkToolCallStatus::InProgress) + .raw_output(json!({ + "status": "generating", + "saved_path": "/tmp/result.txt", + "result": "not an inline image" + })), + )), + ); + + let events = session_notification_to_events(¬if); + let json = serde_json::to_value(&events[0]).unwrap(); + let raw_output = &json["data"]["update"]["rawOutput"]; + + assert_eq!(json["data"]["update"]["status"], "in_progress"); + assert_eq!(raw_output["result"], "not an inline image"); + assert!(raw_output.get("image").is_none()); + assert!(raw_output.get("result_omitted").is_none()); + } + + #[test] + fn codex_image_tool_update_detects_image_mime_types_from_saved_path() { + let cases = [ + ("photo.jpg", "image/jpeg", "/9j/"), + ("photo.webp", "image/webp", "UklGR"), + ("photo.gif", "image/gif", "data:image/gif;base64,"), + ("photo.bin", "image/png", "iVBORw0KGgo"), + ]; + + for (file_name, expected_mime, prefix) in cases { + let saved_path = format!("/Users/test/.codex/generated_images/session/{file_name}"); + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "ig_mime", + ToolCallUpdateFields::new().raw_output(json!({ + "saved_path": saved_path, + "result": format!("{prefix}{}", "A".repeat(128 * 1024)) + })), + )), + ); + + let events = session_notification_to_events(¬if); + let json = serde_json::to_value(&events[0]).unwrap(); + + assert_eq!(json["data"]["update"]["rawOutput"]["image"]["mime_type"], expected_mime); + assert!(json["data"]["update"]["rawOutput"].get("result").is_none()); + } + } + + #[test] + fn codex_image_tool_update_omits_base64_without_saved_path() { + let large_png_base64 = format!("iVBORw0KGgo{}", "A".repeat(128 * 1024)); + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "ig_no_path", + ToolCallUpdateFields::new() + .status(SdkToolCallStatus::InProgress) + .raw_output(json!({ + "call_id": "ig_no_path", + "status": "generating", + "result": large_png_base64 + })), + )), + ); + + let events = session_notification_to_events(¬if); + assert_eq!(events.len(), 1); + let json = serde_json::to_value(&events[0]).unwrap(); + let raw_output = &json["data"]["update"]["rawOutput"]; + + // Oversized base64 must be stripped even though Codex did not save the file. + assert!(raw_output.get("result").is_none()); + assert_eq!(raw_output["result_omitted"], true); + // No saved_path means we cannot offer a path-based preview, so no image object. + assert!(raw_output.get("image").is_none()); + // Without a saved image the status must pass through unchanged. + assert_eq!(json["data"]["update"]["status"], "in_progress"); + } + + #[test] + fn codex_image_tool_update_preserves_failed_status() { + let notif = SessionNotification::new( + "sess-1", + SessionUpdate::ToolCallUpdate(SdkToolCallUpdate::new( + "ig_failed_image", + ToolCallUpdateFields::new() + .status(SdkToolCallStatus::Failed) + .raw_output(json!({ + "call_id": "ig_failed_image", + "status": "failed", + "saved_path": "/Users/test/.codex/generated_images/session/ig_failed_image.png", + "result": format!("iVBORw0KGgo{}", "A".repeat(128 * 1024)) + })), + )), + ); + + let events = session_notification_to_events(¬if); + assert_eq!(events.len(), 1); + let json = serde_json::to_value(&events[0]).unwrap(); + + // A terminal `failed` status must never be rewritten to `completed`. + assert_eq!(json["data"]["update"]["status"], "failed"); + assert_eq!(json["data"]["update"]["rawOutput"]["status"], "failed"); + // The base64 payload is still stripped regardless of the failure. + assert!(json["data"]["update"]["rawOutput"].get("result").is_none()); + } + #[test] fn permission_request_maps_to_snake_case_event_data() { let request = RequestPermissionRequest::new( diff --git a/crates/aionui-ai-agent/src/protocol/events/translate.rs b/crates/aionui-ai-agent/src/protocol/events/translate.rs index 3aff1263..b8ee16dd 100644 --- a/crates/aionui-ai-agent/src/protocol/events/translate.rs +++ b/crates/aionui-ai-agent/src/protocol/events/translate.rs @@ -63,16 +63,20 @@ pub(crate) fn session_notification_to_events(notif: &SessionNotification) -> Vec } SessionUpdate::ToolCallUpdate(tcu) => { + let mut raw_output = sanitize_raw_output(tcu.fields.raw_output.clone()); + let status = normalize_tool_status(tcu.fields.status.as_ref(), raw_output.as_ref()); + normalize_raw_output_status(&mut raw_output, status.as_ref()); + events.push(AgentStreamEvent::AcpToolCall(AcpToolCallEventData { session_id, update: AcpToolCallUpdateData { session_update: AcpToolCallSessionUpdateKind::ToolCallUpdate, tool_call_id: tcu.tool_call_id.to_string(), - status: tcu.fields.status.as_ref().map(map_sdk_tool_status), + status, title: tcu.fields.title.clone(), kind: tcu.fields.kind.as_ref().map(map_sdk_tool_kind), raw_input: tcu.fields.raw_input.clone(), - raw_output: tcu.fields.raw_output.clone(), + raw_output, content: tcu .fields .content @@ -157,6 +161,112 @@ fn map_sdk_tool_status(sdk: &SdkToolCallStatus) -> AcpToolCallStatus { } } +const ACP_RAW_OUTPUT_INLINE_IMAGE_LIMIT: usize = 64 * 1024; + +fn sanitize_raw_output(raw_output: Option) -> Option { + let mut value = raw_output?; + sanitize_inline_image_result(&mut value); + Some(value) +} + +fn sanitize_inline_image_result(value: &mut serde_json::Value) { + let Some(obj) = value.as_object_mut() else { + return; + }; + + let saved_path = obj.get("saved_path").and_then(|v| v.as_str()).map(str::to_owned); + // Strip any oversized inline-image `result` regardless of whether the image was + // saved to disk. Older codex versions and interrupted/failed generations may emit + // the multi-MB base64 without a `saved_path`; that payload must never reach the + // WebSocket broadcast or SQLite either. `saved_path` only decides whether we attach + // the structured `image { path, mime_type, source }` object below. + let should_omit = obj + .get("result") + .and_then(|v| v.as_str()) + .map(is_probably_inline_image_result) + .unwrap_or(false); + + if !should_omit { + return; + } + + let result_len = obj.get("result").and_then(|v| v.as_str()).map(str::len).unwrap_or(0); + obj.remove("result"); + obj.insert("result_omitted".to_owned(), serde_json::Value::Bool(true)); + obj.insert( + "result_omitted_reason".to_owned(), + serde_json::Value::String("image_base64".to_owned()), + ); + obj.insert( + "result_bytes".to_owned(), + serde_json::Value::Number(serde_json::Number::from(result_len)), + ); + + if let Some(path) = saved_path { + let mime_type = mime_type_from_image_path(&path); + obj.insert( + "image".to_owned(), + serde_json::json!({ + "path": path, + "mime_type": mime_type, + "source": "codex_image_generation" + }), + ); + } +} + +fn is_probably_inline_image_result(value: &str) -> bool { + value.len() > ACP_RAW_OUTPUT_INLINE_IMAGE_LIMIT + && (value.starts_with("iVBORw0KGgo") + || value.starts_with("/9j/") + || value.starts_with("UklGR") + || value.starts_with("data:image/")) +} + +fn mime_type_from_image_path(path: &str) -> &'static str { + let lower = path.to_ascii_lowercase(); + if lower.ends_with(".jpg") || lower.ends_with(".jpeg") { + "image/jpeg" + } else if lower.ends_with(".webp") { + "image/webp" + } else if lower.ends_with(".gif") { + "image/gif" + } else { + "image/png" + } +} + +fn normalize_tool_status( + sdk_status: Option<&SdkToolCallStatus>, + raw_output: Option<&serde_json::Value>, +) -> Option { + let image_saved = raw_output + .and_then(|v| v.get("image")) + .and_then(|v| v.get("path")) + .and_then(|v| v.as_str()) + .is_some(); + + // Only force `completed` when the image is on disk AND the agent did not already + // report a terminal status. Codex stalls by leaving the final event as + // `generating`/`in_progress`, but a genuine `failed` must be preserved as-is. + match (image_saved, sdk_status.map(map_sdk_tool_status)) { + (true, None | Some(AcpToolCallStatus::Pending | AcpToolCallStatus::InProgress)) => { + Some(AcpToolCallStatus::Completed) + } + (_, status) => status, + } +} + +fn normalize_raw_output_status(raw_output: &mut Option, status: Option<&AcpToolCallStatus>) { + let Some(AcpToolCallStatus::Completed) = status else { + return; + }; + let Some(obj) = raw_output.as_mut().and_then(|v| v.as_object_mut()) else { + return; + }; + obj.insert("status".to_owned(), serde_json::Value::String("completed".to_owned())); +} + fn map_sdk_tool_kind(kind: &SdkToolKind) -> AcpToolCallKind { match kind { SdkToolKind::Read | SdkToolKind::Search => AcpToolCallKind::Read, diff --git a/crates/aionui-conversation/src/stream_relay.rs b/crates/aionui-conversation/src/stream_relay.rs index ff353bab..542933e2 100644 --- a/crates/aionui-conversation/src/stream_relay.rs +++ b/crates/aionui-conversation/src/stream_relay.rs @@ -1444,6 +1444,94 @@ mod tests { ); } + #[tokio::test] + async fn run_acp_image_tool_call_update_persists_finish_without_base64() { + use aionui_ai_agent::protocol::events::tool_call::{ + AcpToolCallEventData, AcpToolCallKind, AcpToolCallSessionUpdateKind, AcpToolCallStatus, + AcpToolCallUpdateData, + }; + + let repo = Arc::new(RecordingRepo::new()); + let bus = Arc::new(aionui_realtime::BroadcastEventBus::new(64)); + let (tx, _) = broadcast::channel(64); + + let relay = StreamRelay::new( + "conv-1".into(), + "asst-1".into(), + "turn-1".into(), + "user-1".into(), + repo.clone(), + bus.clone(), + None, + ); + + let rx = tx.subscribe(); + + tx.send(AgentStreamEvent::AcpToolCall(AcpToolCallEventData { + session_id: "sess-1".into(), + update: AcpToolCallUpdateData { + session_update: AcpToolCallSessionUpdateKind::ToolCall, + tool_call_id: "ig_test_image".into(), + status: Some(AcpToolCallStatus::InProgress), + title: Some("Image generation".into()), + kind: Some(AcpToolCallKind::Execute), + raw_input: Some(json!({"prompt": "一只小猫"})), + raw_output: None, + content: None, + locations: None, + }, + meta: None, + })) + .unwrap(); + + tx.send(AgentStreamEvent::AcpToolCall(AcpToolCallEventData { + session_id: "sess-1".into(), + update: AcpToolCallUpdateData { + session_update: AcpToolCallSessionUpdateKind::ToolCallUpdate, + tool_call_id: "ig_test_image".into(), + status: Some(AcpToolCallStatus::Completed), + title: None, + kind: Some(AcpToolCallKind::Execute), + raw_input: None, + raw_output: Some(json!({ + "saved_path": "/Users/test/.codex/generated_images/session/ig_test_image.png", + "image": { + "path": "/Users/test/.codex/generated_images/session/ig_test_image.png", + "mime_type": "image/png", + "source": "codex_image_generation" + }, + "result_omitted": true, + "result_omitted_reason": "image_base64", + "result_bytes": 131_083 + })), + content: None, + locations: None, + }, + meta: None, + })) + .unwrap(); + + tx.send(AgentStreamEvent::Finish(FinishEventData::default())).unwrap(); + + relay.consume(rx).await; + + let updates = repo.take_updates(); + let acp_update = updates.iter().find(|(id, _)| id == "ig_test_image"); + assert!(acp_update.is_some()); + let (_, upd) = acp_update.unwrap(); + assert_eq!(upd.status, Some(Some("finish".to_owned()))); + + let content = upd.content.as_deref().unwrap(); + assert!(!content.contains("iVBORw0KGgo")); + assert!(content.contains("result_omitted")); + + let merged: serde_json::Value = serde_json::from_str(content).unwrap(); + assert_eq!( + merged["update"]["raw_output"]["image"]["path"], + "/Users/test/.codex/generated_images/session/ig_test_image.png" + ); + } + #[tokio::test] async fn run_tool_group_persists_message() { use aionui_ai_agent::protocol::events::tool_call::{ToolCallStatus, ToolGroupEntry};