From c55a90d23eb57e851af360f5b7bf0105f9fb8483 Mon Sep 17 00:00:00 2001 From: kvtodev Date: Fri, 10 Jul 2026 19:06:39 +0800 Subject: [PATCH] centralize data key query instead of using obj key. --- .../crates/api-server/src/editor_agent/api.rs | 67 +++++++++++++++++-- .../src/editor_agent/editor_tools/common.rs | 6 ++ .../editor_agent/editor_tools/edit_image.rs | 42 ++++++++---- .../editor_tools/generate_character.rs | 8 ++- .../editor_tools/generate_icon_spritesheet.rs | 34 ++++++++-- .../editor_tools/generate_image.rs | 41 ++++++++---- .../editor_tools/generate_video.rs | 20 +++++- .../api-server/src/editor_agent/utils.rs | 52 +++++++++++--- 8 files changed, 216 insertions(+), 54 deletions(-) diff --git a/server-rs/crates/api-server/src/editor_agent/api.rs b/server-rs/crates/api-server/src/editor_agent/api.rs index dcc166820..fd06c99c6 100644 --- a/server-rs/crates/api-server/src/editor_agent/api.rs +++ b/server-rs/crates/api-server/src/editor_agent/api.rs @@ -48,7 +48,7 @@ use crate::editor_agent::editor_tools::generate_sound_effect::{ use crate::editor_agent::editor_tools::generate_ui_design::GenerateUiDesignTool; use crate::editor_agent::editor_tools::generate_video::{GenerateVideoTool, GenerateVideoToolArgs}; use crate::editor_agent::utils::{ - ImageId, ImageMetadata, IntoImageId, build_editor_agent_canvas_completion, + ImageId, ImageMetadata, IntoDataKey, IntoImageId, build_editor_agent_canvas_completion, conversation_detail_from_record, conversation_summary_from_record, editor_agent_bad_request, empty_messages_document, ensure_editor_project_access, normalize_editor_agent_attachments, now_rfc3339, read_messages_document, require_editor_agent_sidebar_enabled, @@ -270,18 +270,22 @@ fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Edi for msg in document.messages.iter().rev() { if let Some(tc) = &msg.tool_call { for img in &tc.images { - let image_id = ImageId { - id: img.object_key.clone().unwrap_or_default(), + let data_key = img.clone().into_data_key(); + let image_id = ImageId::from_data_key(&data_key); + let metadata = ImageMetadata { + tag: String::new(), + data_key, }; - let metadata = ImageMetadata { tag: String::new() }; images.insert(image_id, metadata); } } for a in &msg.attachments { - let image_id = ImageId { - id: a.object_key.clone().unwrap_or_default(), + let data_key = a.clone().into_data_key(); + let image_id = ImageId::from_data_key(&data_key); + let metadata = ImageMetadata { + tag: String::new(), + data_key, }; - let metadata = ImageMetadata { tag: String::new() }; images.insert(image_id, metadata); } } @@ -289,6 +293,55 @@ fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Edi EditorToolContext { images } } +#[cfg(test)] +mod tool_context_tests { + use super::*; + use shared_contracts::editor_agent::{ + EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, EditorAgentGeneratedImage, + EditorAgentMessageRole, EditorAgentToolCallStatus, + }; + + #[test] + fn tool_context_maps_hashed_ids_back_to_real_data_keys() { + let data_key = "generated-images/proj-1/output.png"; + let document = EditorAgentConversationMessagesDocument { + version: EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, + conversation_id: "editor-agent-conversation-1".to_string(), + messages: vec![EditorAgentMessage { + id: 0, + role: EditorAgentMessageRole::System, + text: String::new(), + attachments: Vec::new(), + tool_call: Some(EditorAgentToolCall { + tool_name: "generate-image".to_string(), + summary: String::new(), + status: EditorAgentToolCallStatus::Completed, + args: json!({}), + images: vec![EditorAgentGeneratedImage { + resource_id: None, + object_key: Some(data_key.to_string()), + asset_object_id: None, + image_src: "/generated-images/proj-1/output.png".to_string(), + thumbnail_src: None, + width: None, + height: None, + }], + videos: Vec::new(), + audios: Vec::new(), + error: None, + }), + created_at: String::new(), + }], + }; + + let context = build_tool_context(&document); + let image_id = ImageId::from_data_key(data_key); + + assert_eq!(context.image_data_key(&image_id), Some(data_key)); + assert!(!image_id.id.contains(data_key)); + } +} + #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] struct EditorAgentConversationDeleteResponse { diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs index 0c99dd43f..3cbdfc000 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs @@ -13,4 +13,10 @@ impl EditorToolContext { pub fn contains_image(&self, image_id: &ImageId) -> bool { self.images.contains_key(image_id) } + + pub fn image_data_key(&self, image_id: &ImageId) -> Option<&str> { + self.images + .get(image_id) + .map(|metadata| metadata.data_key.as_str()) + } } diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs index 2758665a6..d9fd2f5bb 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs @@ -178,16 +178,14 @@ impl EditImageTool { let name = Self::NAME.to_string(); let arg_json = serde_json::to_string(&args).unwrap_or_default(); - // TODO use the trait - let image_id: ImageId = ImageId { - id: result - .object_key - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(|value| value.trim_start_matches('/').to_string()) - .unwrap_or_else(|| result.image_src.clone()), - }; + let data_key = result + .object_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| value.trim_start_matches('/').to_string()) + .unwrap_or_else(|| result.image_src.clone()); + let image_id = ImageId::from_data_key(data_key); format!( "[tool_call:{name}] args: {arg_json} output: edited result saved as image: {image_id}" ) @@ -206,13 +204,22 @@ impl EditImageTool { source_resource_id: Option, canvas_completion: Option, ) -> Result { - // TODO id should not be assumed as src key - let source_image_src = args.object_image_id.id; - let reference_image_srcs: Vec = args + let source_image_src = self + .context + .image_data_key(&args.object_image_id) + .ok_or_else(|| image_not_found(&args.object_image_id))? + .to_string(); + let reference_image_srcs: Result, AppError> = args .reference_image_ids .into_iter() - .map(|id| id.id) + .map(|image_id| { + self.context + .image_data_key(&image_id) + .map(ToOwned::to_owned) + .ok_or_else(|| image_not_found(&image_id)) + }) .collect(); + let reference_image_srcs = reference_image_srcs?; let result = edit_editor_image_for_owner( state, @@ -256,3 +263,10 @@ impl EditImageTool { }) } } + +fn image_not_found(image_id: &ImageId) -> AppError { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": format!("asset {image_id} not found in context"), + })) +} diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs index 903a7c595..d088f4980 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_character.rs @@ -107,10 +107,14 @@ pub(crate) fn format_tool_message( result: &EditorImageGenerationResult, ) -> String { let args = serde_json::to_string(args).unwrap_or_default(); - let image_id = result + let data_key = result .object_key .as_deref() - .unwrap_or(result.image_src.as_str()); + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| value.trim_start_matches('/').to_string()) + .unwrap_or_else(|| result.image_src.clone()); + let image_id = crate::editor_agent::utils::ImageId::from_data_key(data_key); format!( "[tool_call:{tool_name}] args: {args} output: generated result saved as image: {image_id}" ) diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs index aad22b529..cde4c38f5 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_icon_spritesheet.rs @@ -170,12 +170,21 @@ impl GenerateIconSpritesheetTool { request_context, caller, EditorIconSpritesheetGenerationRequest { - reference_image_src: args.reference_image_id.id, + reference_image_src: self + .context + .image_data_key(&args.reference_image_id) + .ok_or_else(|| image_not_found(&args.reference_image_id))? + .to_string(), reference_image_srcs: Some( args.reference_image_ids .into_iter() - .map(|id| id.id) - .collect(), + .map(|image_id| { + self.context + .image_data_key(&image_id) + .map(ToOwned::to_owned) + .ok_or_else(|| image_not_found(&image_id)) + }) + .collect::, _>>()?, ), icon_descriptions: args.icon_descriptions, model: None, @@ -212,10 +221,27 @@ impl GenerateIconSpritesheetTool { result: &EditorIconSpritesheetResult, ) -> String { let args = serde_json::to_string(args).unwrap_or_default(); + let data_key = result + .spritesheet_resource + .as_ref() + .and_then(|resource| resource.get("objectKey")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| value.trim_start_matches('/').to_string()) + .unwrap_or_else(|| result.spritesheet_image_src.clone()); + let image_id = ImageId::from_data_key(data_key); format!( "[tool_call:{}] args: {args} output: generated spritesheet saved as image: {}", Self::NAME, - result.spritesheet_image_src + image_id ) } } + +fn image_not_found(image_id: &ImageId) -> AppError { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": format!("asset {image_id} not found in context"), + })) +} diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs index 03a383910..ea0272f6a 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_image.rs @@ -98,7 +98,7 @@ impl Tool for GenerateImageTool { async move { self.validate_args(&args)?; Ok(GenerateImageToolOutput { - message: "this tool call is pending user confirmation. if all is pending, just end this turn".to_string(), + message: "this tool call is pending user confirmation.".to_string(), }) } } @@ -154,13 +154,14 @@ impl GenerateImageTool { result: &EditorImageGenerationResult, ) -> String { let args = serde_json::to_string(args).unwrap_or_default(); - let image_id = result + let data_key = result .object_key .as_deref() .map(str::trim) .filter(|value| !value.is_empty()) .map(|value| value.trim_start_matches('/').to_string()) .unwrap_or_else(|| result.image_src.clone()); + let image_id = ImageId::from_data_key(data_key); format!( "[tool_call:{}] args: {args} output: generated result saved as image: {image_id}", Self::NAME, @@ -231,12 +232,7 @@ impl GenerateImageTool { .then(|| "birefnet".to_string()), aspect_ratio: args.aspect_ratio, image_size: args.image_size, - reference_image_srcs: Some( - args.reference_image_ids - .into_iter() - .map(|image_id| image_id.id) - .collect(), - ), + reference_image_srcs: Some(self.image_data_keys(args.reference_image_ids)?), project_id: Some(project_id), asset_kind: Some(asset_kind.to_string()), generation_inputs, @@ -267,6 +263,23 @@ impl GenerateImageTool { })) }) } + + fn image_data_keys(&self, image_ids: Vec) -> Result, AppError> { + image_ids + .into_iter() + .map(|image_id| { + self.context + .image_data_key(&image_id) + .map(ToOwned::to_owned) + .ok_or_else(|| { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": format!("asset {image_id} not found in context"), + })) + }) + }) + .collect() + } } #[cfg(test)] @@ -275,13 +288,15 @@ mod tests { use std::collections::HashMap; fn tool() -> GenerateImageTool { + let image_id = ImageId::from_data_key("reference.png"); GenerateImageTool { context: EditorToolContext { images: HashMap::from([( - ImageId { - id: "reference.png".to_string(), + image_id, + crate::editor_agent::utils::ImageMetadata { + tag: String::new(), + data_key: "reference.png".to_string(), }, - crate::editor_agent::utils::ImageMetadata { tag: String::new() }, )]), }, } @@ -292,9 +307,7 @@ mod tests { let output = tool() .call(GenerateImageToolArgs { prompt: "雨夜的霓虹街道".to_string(), - reference_image_ids: vec![ImageId { - id: "reference.png".to_string(), - }], + reference_image_ids: vec![ImageId::from_data_key("reference.png")], aspect_ratio: Some("16:9".to_string()), image_size: Some("2K".to_string()), }) diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs index fe93647c6..27d26609e 100644 --- a/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/generate_video.rs @@ -125,11 +125,17 @@ impl GenerateVideoTool { json!({ "provider": "editor-agent", "message": "视频提示词不能为空" }), )); } - let reference_image_srcs = args + let reference_image_srcs: Result, AppError> = args .reference_image_ids .into_iter() - .map(|id| id.id) + .map(|image_id| { + self.context + .image_data_key(&image_id) + .map(ToOwned::to_owned) + .ok_or_else(|| image_not_found(&image_id)) + }) .collect(); + let reference_image_srcs = reference_image_srcs?; let response = generate_editor_video_for_owner( state.clone(), request_context.clone(), @@ -164,11 +170,19 @@ impl GenerateVideoTool { args: &GenerateVideoToolArgs, result: &EditorVideoGenerateResponse, ) -> String { + let video_id = ImageId::from_data_key(result.video_src.as_str()); format!( "[tool_call:{}] args: {} output: generated video saved as: {}", Self::NAME, serde_json::to_string(args).unwrap_or_default(), - result.video_src + video_id ) } } + +fn image_not_found(image_id: &ImageId) -> AppError { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": format!("asset {image_id} not found in context"), + })) +} diff --git a/server-rs/crates/api-server/src/editor_agent/utils.rs b/server-rs/crates/api-server/src/editor_agent/utils.rs index e2a633f54..cc5ee9d24 100644 --- a/server-rs/crates/api-server/src/editor_agent/utils.rs +++ b/server-rs/crates/api-server/src/editor_agent/utils.rs @@ -8,12 +8,13 @@ use platform_oss::{ }; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; use shared_contracts::assets::{ EditorCanvasGenerationCompletionPayload, EditorCanvasGenerationPlaceholderPayload, }; use shared_contracts::editor_agent::{ EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, EditorAgentAttachmentRef, EditorAgentAttachmentSource, - EditorAgentConversationDetail, EditorAgentConversationMessagesDocument, + EditorAgentConversationDetail, EditorAgentConversationMessagesDocument, EditorAgentGeneratedImage, EditorAgentConversationSummary, EditorAgentMessage, }; use shared_kernel::{normalize_optional_string, normalize_required_string}; @@ -39,6 +40,16 @@ impl IntoDataKey for EditorAgentAttachmentRef { .unwrap_or_else(|| self.image_src.clone()) } } +impl IntoDataKey for EditorAgentGeneratedImage { + fn into_data_key(self) -> String { + self.object_key + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(|value| value.trim_start_matches('/').to_string()) + .unwrap_or(self.image_src) + } +} impl IntoDataKey for EditorProjectResourceRecord { fn into_data_key(self) -> String { self.object_key @@ -66,23 +77,17 @@ pub trait IntoImageId { } impl IntoImageId for EditorAgentAttachmentRef { fn into_image_id(self) -> ImageId { - ImageId { - id: self.into_data_key(), - } + ImageId::from_data_key(self.into_data_key()) } } impl IntoImageId for EditorProjectResourceRecord { fn into_image_id(self) -> ImageId { - ImageId { - id: self.into_data_key(), - } + ImageId::from_data_key(self.into_data_key()) } } impl IntoImageId for EditorAssetRecord { fn into_image_id(self) -> ImageId { - ImageId { - id: self.into_data_key(), - } + ImageId::from_data_key(self.into_data_key()) } } @@ -107,6 +112,16 @@ pub struct ImageId { pub id: String, } +impl ImageId { + /// Keep data keys server-side while giving the model a stable image reference. + pub fn from_data_key(data_key: impl AsRef) -> Self { + let digest = Sha256::digest(data_key.as_ref().as_bytes()); + Self { + id: format!("sha256:{digest:x}"), + } + } +} + impl Serialize for ImageId { fn serialize(&self, serializer: S) -> Result { self.id.serialize(serializer) @@ -122,6 +137,7 @@ impl<'de> Deserialize<'de> for ImageId { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ImageMetadata { pub tag: String, + pub data_key: String, } impl Display for ImageId { @@ -130,6 +146,22 @@ impl Display for ImageId { } } +#[cfg(test)] +mod image_id_tests { + use super::ImageId; + + #[test] + fn image_id_is_a_stable_hash_of_the_data_key() { + let data_key = "editor-projects/proj-1/image.png"; + let image_id = ImageId::from_data_key(data_key); + + assert_eq!(image_id, ImageId::from_data_key(data_key)); + assert_ne!(image_id, ImageId::from_data_key("another-image.png")); + assert!(image_id.id.starts_with("sha256:")); + assert!(!image_id.id.contains(data_key)); + } +} + type EditorAgentConversationLockMap = Mutex>>>; static EDITOR_AGENT_CONVERSATION_LOCKS: OnceLock = OnceLock::new();