From db2e8cca8dfe1fa63344a3cce4d876d425965f07 Mon Sep 17 00:00:00 2001 From: kvtodev Date: Fri, 10 Jul 2026 08:42:21 +0800 Subject: [PATCH] reimpl the agent --- .../api-server/src/editor_agent/agent.rs | 211 +++++++ .../crates/api-server/src/editor_agent/api.rs | 207 +++++++ .../src/editor_agent/editor_tools/common.rs | 16 + .../editor_agent/editor_tools/edit_image.rs | 245 ++++++++ .../src/editor_agent/editor_tools/mod.rs | 2 + .../crates/api-server/src/editor_agent/mod.rs | 4 + .../api-server/src/editor_agent/utils.rs | 560 ++++++++++++++++++ 7 files changed, 1245 insertions(+) create mode 100644 server-rs/crates/api-server/src/editor_agent/agent.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/api.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/mod.rs create mode 100644 server-rs/crates/api-server/src/editor_agent/utils.rs diff --git a/server-rs/crates/api-server/src/editor_agent/agent.rs b/server-rs/crates/api-server/src/editor_agent/agent.rs new file mode 100644 index 000000000..bf1aefc8f --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/agent.rs @@ -0,0 +1,211 @@ +use serde_json::Value; +use module_editor_agent::agent::agent::{Agent, LlmApiAdaptor}; +use module_editor_agent::agent::agent_builder::AgentBuilder; +use module_editor_agent::agent::error::PromptError; +use module_editor_agent::agent::hook::Hook; +use module_editor_agent::agent::memory::AgentMemory; +use module_editor_agent::agent::run::ToolCallFlow; +use module_editor_agent::agent::tool::{Tool, ToolCall, ToolDyn}; +use platform_llm::{LlmClient, LlmMessage}; + +struct LlmCompletionModel { + client: LlmClient, +} + +impl LlmApiAdaptor for LlmCompletionModel { + async fn complete(&self, messages: &[LlmMessage]) -> Result { + use platform_llm::LlmTextRequest; + let request = LlmTextRequest::new(messages.to_vec()).with_request_timeout_ms(30_000); + let response = self + .client + .request_text(request) + .await + .map_err(|e| PromptError::CompletionError(e.to_string()))?; + Ok(response.content) + } + + fn tool_result_message(&self, tool_name: &str, output: &str) -> LlmMessage { + LlmMessage::system(format!("Tool '{tool_name}' returned: {output}")) + } + + fn build_assistant_message(&self, text: &str) -> LlmMessage { + LlmMessage::assistant(text) + } +} + +pub struct LlmChatAgentBuilder { + client: Option, + system_prompt_parts: Vec, + tools: Vec>, + hooks: Vec>, + max_turns: usize, + memory_data: Option>>, + context: Option, +} + +impl AgentBuilder for LlmChatAgentBuilder { + type Client = LlmClient; + + fn new() -> Self { + Self { + client: None, + system_prompt_parts: Vec::new(), + tools: Vec::new(), + hooks: Vec::new(), + max_turns: 10, + memory_data: None, + context: None, + } + } + + fn with_client(mut self, client: LlmClient) -> Self { + self.client = Some(client); + self + } + + fn system_prompt(mut self, system_prompt: impl Into) -> Self { + self.system_prompt_parts.push(system_prompt.into()); + self + } + + fn tool(mut self, tool: impl Tool + Send + Sync + 'static) -> Self { + self.tools.push(Box::new(tool)); + self + } + + fn add_hook(mut self, hook: impl Hook + 'static) -> Self { + self.hooks.push(Box::new(hook)); + self + } + + fn max_turns(mut self, n: usize) -> Self { + self.max_turns = n; + self + } + + fn memory(mut self, memory: impl AgentMemory + 'static) -> Self { + self.memory_data = Some(Box::new(memory)); + self + } + + fn context(mut self, context: Value) -> Self { + self.context = Some(context); + self + } + + fn build(self) -> Agent { + let model = LlmCompletionModel { + client: self.client.expect("call .with_client() first"), + }; + let mut agent = Agent::new(model); + let tool_specs = self + .tools + .iter() + .map(|tool| ToolPromptSpec { + name: tool.tool_name().to_string(), + description: tool.description(), + parameters: tool.parameters(), + }) + .collect::>(); + let base_prompt = self.system_prompt_parts.join("\n\n"); + let system_prompt = build_tools_system_prompt(&base_prompt, &tool_specs); + agent.tools = self.tools; + agent.hooks = self.hooks; + agent.default_max_turns = self.max_turns; + agent.memory = self.memory_data; + agent.context = self.context; + agent.system_prompt = Some(LlmMessage::system(&system_prompt)); + agent + } +} + +struct ToolPromptSpec { + name: String, + description: String, + parameters: Value, +} + +fn build_tools_system_prompt(base_prompt: &str, tool_specs: &[ToolPromptSpec]) -> String { + let mut prompt = String::new(); + prompt.push_str(base_prompt); + prompt.push_str("\n\nYou have access to the following tools.\n\n"); + + if tool_specs.is_empty() { + prompt.push_str("(No tools available.)\n"); + } else { + prompt.push_str("## JSON Response Format\n"); + prompt.push_str( + "When you need to use a tool, respond with valid JSON only (no markdown fences):\n", + ); + prompt.push_str("{\n"); + prompt.push_str(" \"reply_text\": \"your message to the user\",\n"); + prompt.push_str(" \"tool_calls\": [\n {\n"); + prompt.push_str(" \"tool_name\": \"tool_name_here\",\n"); + prompt.push_str(" \"args\": { \"argument_name\": \"argument_value\" }\n"); + prompt.push_str(" }\n ]\n"); + prompt.push_str("}\n\n"); + prompt.push_str("If you don't need to use a tool, respond with:\n"); + prompt.push_str("{\n"); + prompt.push_str(" \"reply_text\": \"your message\",\n"); + prompt.push_str(" \"tool_calls\": []\n"); + prompt.push_str("}\n\n"); + prompt.push_str("## Available Tools\n\n"); + + for tool in tool_specs { + prompt.push_str(&format!("- {}\n", tool.name)); + prompt.push_str(&format!(" Description: {}\n", tool.description)); + prompt.push_str(" Arguments JSON Schema:\n"); + let parameters = serde_json::to_string_pretty(&tool.parameters) + .unwrap_or_else(|_| tool.parameters.to_string()); + prompt.push_str(¶meters); + prompt.push_str("\n"); + } + + prompt.push_str( + "Use `` in `reply_text` when you want to stop the conversation turn and ", + ); + prompt.push_str( + "return your response immediately without expecting further tool execution. ", + ); + prompt.push_str("For example: `\"reply_text\": \"Task complete. \"`."); + prompt.push_str("\n"); + prompt.push_str("IMPORTANT: Always respond with valid JSON only. "); + prompt.push_str( + "Do not wrap the JSON in markdown code fences or add extra text outside the JSON.", + ); + } + + prompt +} + + +// +// async fn run_chat_agent(config: LlmConfig) -> Result<(), LlmError> { +// let client = LlmClient::new(config)?; +// let agent = LlmChatAgentBuilder::new() +// .with_client(client) +// .system_prompt("You are a helpful assistant with an echo tool.") +// .build(); +// +// let names = agent +// .tools() +// .iter() +// .map(|t| t.tool_name().to_string()) +// .collect(); +// let tool_hook = ToolValidationHook::new(names.clone()); +// +// let output = agent +// .prompt(LlmMessage::user( +// "Use the echo tool to echo 'Hello from the JSON harness!', then tell me what it said.", +// )) +// .max_turns(5) +// .add_hook(tool_hook) +// .await +// .expect("agent should succeed"); +// +// println!("--- Final Agent Response ---"); +// println!("{}", output.text); +// println!("-----------------------------"); +// +// Ok(()) +// } diff --git a/server-rs/crates/api-server/src/editor_agent/api.rs b/server-rs/crates/api-server/src/editor_agent/api.rs new file mode 100644 index 000000000..b25ca7054 --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/api.rs @@ -0,0 +1,207 @@ +use std::collections::HashMap; + +use axum::extract::{Path, State}; +use axum::{Extension, Json}; +use module_editor_agent::agent::agent_builder::AgentBuilder; +use module_editor_agent::agent::error::PromptError; +use module_editor_agent::agent::memory::VecMemory; +use module_editor_agent::agent::run::PromptOutput; +use module_editor_agent::derive_conversation_title; +use platform_llm::LlmMessage; +use serde_json::{Value, json}; +use shared_contracts::editor_agent::{ + EditorAgentConversationMessagesDocument, EditorAgentMessage, EditorAgentMessageRole, + EditorAgentToolCall, EditorAgentToolCallStatus, StreamEditorAgentMessageRequest, +}; +use spacetime_client::EditorAgentConversationTouchRecordInput; + +use crate::auth::AuthenticatedAccessToken; +use crate::editor_agent::agent::LlmChatAgentBuilder; +use crate::editor_agent::editor_tools::edit_image::EditImageTool; +use crate::editor_agent::utils::{EditorAgentMessageResponse, ImageId, ImageMetadata, now_rfc3339, require_editor_agent_sidebar_enabled, write_messages_document, normalize_editor_agent_attachments, read_messages_document}; +use crate::editor_project::current_utc_micros; +use crate::http_error::AppError; +use crate::request_context::RequestContext; +use crate::state::AppState; + +pub async fn editor_agent_message( + State(state): State, + Path(conversation_id): Path, + Extension(_request_context): Extension, + Extension(authenticated): Extension, + Json(payload): Json, +) -> Result, AppError> { + let owner_user_id = authenticated.claims().user_id().to_string(); + require_editor_agent_sidebar_enabled(&state, owner_user_id.as_str()).await?; + // Load conversation & attachments + let conversation = state + .spacetime_client() + .get_editor_agent_conversation(conversation_id.clone(), owner_user_id.clone()) + .await + .map_err(|e| { + AppError::from_status(axum::http::StatusCode::NOT_FOUND) + .with_details(json!({ "message": format!("conversation not found: {e}") })) + })?; + + let attachments = normalize_editor_agent_attachments( + &state, + &conversation, + payload.attachments.as_slice(), + ) + .await?; + + let conversation_lock = crate::editor_agent::utils::editor_agent_conversation_lock( + conversation.conversation_id.as_str(), + ); + let _conversation_lock_guard = conversation_lock.lock_owned().await; + + let mut document: EditorAgentConversationMessagesDocument = + read_messages_document(&state, &conversation).await?; + + // Build conversation history as LlmMessage vec + let previous_messages: Vec = document + .messages + .iter() + .map(|msg| match msg.role { + EditorAgentMessageRole::User => LlmMessage::user(&msg.text), + EditorAgentMessageRole::Assistant => LlmMessage::assistant(&msg.text), + EditorAgentMessageRole::System => LlmMessage::system(&msg.text), + }) + .collect(); + + // Save user message to document + let was_empty = document.messages.is_empty(); + let user_now = now_rfc3339(); + let user_message = EditorAgentMessage { + id: document.messages.len(), + role: EditorAgentMessageRole::User, + text: payload.text.trim().to_string(), + attachments, + tool_call: None, + created_at: user_now, + }; + document.messages.push(user_message.clone()); + write_messages_document(&state, &conversation, &document).await?; + + // Touch conversation (set title if first message) + if was_empty { + let title = derive_conversation_title(user_message.text.as_str()); + let _ = state + .spacetime_client() + .touch_editor_agent_conversation(EditorAgentConversationTouchRecordInput { + conversation_id: conversation.conversation_id.clone(), + owner_user_id: conversation.owner_user_id.clone(), + title: Some(title), + updated_at_micros: current_utc_micros(), + }) + .await; + } + + // Build tool context from document + let tool_context_value = build_tool_context(&document); + + // Build and run agent + let llm_client = state.llm_client().ok_or_else(|| { + AppError::from_status(axum::http::StatusCode::SERVICE_UNAVAILABLE) + .with_details(json!({ "message": "LLM client not configured" })) + })?; + let llm_client = llm_client.clone(); + + let memory = VecMemory::new(previous_messages); + + let agent = LlmChatAgentBuilder::new() + .with_client(llm_client) + .tool(EditImageTool {}) + .max_turns(3) + .memory(memory) + .context(tool_context_value) + .build(); + + let agent_result = agent.prompt(LlmMessage::user("")).await; + + let assistant_now = now_rfc3339(); + + match build_delta_messages(agent_result, &assistant_now, document.messages.len()) { + Err(err) => Ok(Json(EditorAgentMessageResponse { + delta_messages: vec![], + error_message: Some(err.to_string()), + })), + Ok(delta_messages) => { + for msg in &delta_messages { + document.messages.push(msg.clone()); + } + write_messages_document(&state, &conversation, &document).await?; + + Ok(Json(EditorAgentMessageResponse { + delta_messages, + error_message: None, + })) + } + } +} + +fn build_delta_messages( + result: Result, PromptError>, + created_at: &str, + messages_offset: usize, +) -> Result, PromptError> { + let outputs = result?; + let mut messages = Vec::with_capacity(outputs.len()); + + for (i, out) in outputs.into_iter().enumerate() { + let absolute_idx = messages_offset + i; + match out { + PromptOutput::Text(text) => { + messages.push(EditorAgentMessage { + id: absolute_idx, + role: EditorAgentMessageRole::Assistant, + text, + attachments: Vec::new(), + tool_call: None, + created_at: created_at.to_string(), + }); + } + PromptOutput::Tool(tco) => { + let summary = tco.tool_call.args.to_string(); + messages.push(EditorAgentMessage { + id: absolute_idx, + role: EditorAgentMessageRole::System, + text: tco.message, + attachments: Vec::new(), + tool_call: Some(EditorAgentToolCall { + tool_name: tco.tool_call.name, + summary, + status: EditorAgentToolCallStatus::PendingConfirmation, + args: tco.tool_call.args, + images: Vec::new(), + error: None, + }), + created_at: created_at.to_string(), + }); + } + } + } + + Ok(messages) +} + +fn build_tool_context(document: &EditorAgentConversationMessagesDocument) -> Value { + let mut images: HashMap = HashMap::new(); + + for msg in document.messages.iter().rev() { + if let Some(tc) = &msg.tool_call { + if !tc.images.is_empty() { + for img in &tc.images { + let image_id = ImageId { + id: img.object_key.clone().unwrap_or_default(), + }; + let metadata = ImageMetadata { tag: String::new() }; + images.insert(image_id, metadata); + } + break; + } + } + } + + json!({ "images": images }) +} 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 new file mode 100644 index 000000000..6690b3f29 --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/common.rs @@ -0,0 +1,16 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use crate::editor_agent::utils::{ImageId, ImageMetadata}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EditorToolContext { + pub images: HashMap, +} + +impl EditorToolContext { + /// Check if an image with the given ID exists in the context. + pub fn contains_image(&self, image_id: &ImageId) -> bool { + self.images.contains_key(image_id) + } +} + 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 new file mode 100644 index 000000000..2d715a9f1 --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/edit_image.rs @@ -0,0 +1,245 @@ +use crate::editor_agent::editor_tools::common::EditorToolContext; +use crate::editor_agent::utils::ImageId; +use crate::editor_project::{ + EditorGenerationCaller, EditorImageEditRequest, edit_editor_image_for_owner, +}; +use crate::http_error::AppError; +use crate::request_context::RequestContext; +use crate::state::AppState; +use axum::http::StatusCode; +use module_editor_agent::agent::tool::{Tool, ToolFailure, ToolFailureKind}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use shared_contracts::api::ApiSuccessEnvelope; +use shared_contracts::assets::EditorCanvasGenerationCompletionPayload; +use std::error::Error; +use std::fmt::Display; + +pub struct EditImageTool {} + +#[derive(Debug, Clone)] +pub enum EditImageError { + ObjectImageNotProvided, + PromptNotProvided, + AssetNotFound(ImageId), + ContextParseError(String), +} + +impl Display for EditImageError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + EditImageError::ObjectImageNotProvided => { + write!(f, "object image not provided") + } + EditImageError::PromptNotProvided => write!(f, "prompt not provided"), + EditImageError::AssetNotFound(image_id) => { + write!(f, "asset {image_id} not found in context") + } + EditImageError::ContextParseError(msg) => { + write!(f, "failed to parse tool context: {msg}") + } + } + } +} + +impl Error for EditImageError {} + +#[derive(Debug, Clone, Deserialize)] +pub struct EditImageToolArgs { + pub object_image_id: ImageId, + #[serde(default)] + pub reference_image_ids: Vec, + pub prompt: String, + // #[serde(default)] + // pub tag: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EditImageToolOutput { + pub message: String, +} + +impl Tool for EditImageTool { + const NAME: &'static str = "edit-image"; + type Error = EditImageError; + type Args = EditImageToolArgs; + type Output = EditImageToolOutput; + + fn description(&self) -> String { + "根据用户提供的文本描述，对选定的图片进行编辑修改，例如调整颜色、替换元素、修改风格等。" + .to_string() + } + + fn parameters(&self) -> Value { + json!({ + "type": "object", + "properties": { + "object_image_id": { + "type": "string", + "description": "要修改的目标图片。" + }, + "reference_image_ids": { + "type": "array", + "items": { "type": "string" }, + "description": "修改参考图 ID 列表（可选的，用于提供风格或元素参考）。" + }, + "prompt": { + "type": "string", + "description": "编辑提示词，描述希望如何修改图片。例如「把背景换成红色」、「把人物改成坐着」。" + }, + // "tag": { + // "type": "string", + // "description": "为新生成的图片添加标签，用于后续在上下文中引用。" + // } + }, + "required": ["object_image_id", "prompt"], + "additionalProperties": false + }) + } + + fn call( + &self, + args: Self::Args, + context: Value, + ) -> impl Future> + Send { + async move { + Self::validate_context_images(&args, &context)?; + if let Some(error) = Self::validate_args(&args) { + return Err(error); + } + Ok(EditImageToolOutput { + message: "this tool call is pending user confirmation. if all is pending, just end this turn".to_string(), + }) + } + } + + fn classify_error(&self, error: &Self::Error) -> ToolFailure { + match error { + EditImageError::ContextParseError(_) => { + ToolFailure::new(ToolFailureKind::Internal, error.to_string()) + } + EditImageError::AssetNotFound(_) => { + ToolFailure::new(ToolFailureKind::NotFound, error.to_string()) + } + EditImageError::ObjectImageNotProvided | EditImageError::PromptNotProvided => { + ToolFailure::invalid_args(error.to_string()) + } + } + } +} + +/// Typed result deserialized from `edit_editor_image_for_owner` response. +/// +/// Mirrors `EditorImageGenerationResponse` but uses `String` instead of `&'static str` +/// and `Value` for nested payload types so that `#[derive(Deserialize)]` works. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct EditorImageEditResult { + pub image_src: String, + pub object_key: Option, + pub asset_object_id: Option, + pub width: u32, + pub height: u32, + pub source_type: String, + pub prompt: String, + pub actual_prompt: Option, + pub model: String, + pub provider: String, + pub task_id: String, + pub resource: Option, + pub asset: Option, + pub project: Option, +} + +impl EditImageTool { + /// Validate the semantic correctness of the arguments. + fn validate_args(args: &EditImageToolArgs) -> Option { + if args.prompt.trim().is_empty() { + return Some(EditImageError::PromptNotProvided); + } + None + } + + /// Validate that all referenced images exist in the context. + fn validate_context_images( + args: &EditImageToolArgs, + context: &Value, + ) -> Result<(), EditImageError> { + let tool_context: EditorToolContext = serde_json::from_value(context.clone()) + .map_err(|e| EditImageError::ContextParseError(e.to_string()))?; + + if !tool_context.contains_image(&args.object_image_id) { + return Err(EditImageError::AssetNotFound(args.object_image_id.clone())); + } + + for ref_id in &args.reference_image_ids { + if !tool_context.contains_image(ref_id) { + return Err(EditImageError::AssetNotFound(ref_id.clone())); + } + } + + Ok(()) + } + + pub async fn execute( + &self, + state: &AppState, + request_context: &RequestContext, + caller: EditorGenerationCaller, + args: EditImageToolArgs, + model: Option, + project_id: Option, + generation_inputs: Option, + asset_label: Option, + source_resource_id: Option, + canvas_completion: Option, + ) -> Result { + let source_image_src = args.object_image_id.id; + let reference_image_srcs: Vec = args + .reference_image_ids + .into_iter() + .map(|id| id.id) + .collect(); + + let result = edit_editor_image_for_owner( + state, + request_context, + caller, + EditorImageEditRequest { + prompt: args.prompt, + source_image_src, + size: None, + model, + reference_image_srcs: Some(reference_image_srcs), + project_id, + asset_kind: Some("editor_agent_edit_image".to_string()), + generation_inputs, + asset_folder_id: Some("project".to_string()), + asset_label, + source_resource_id, + target_layer_id: None, + canvas_completion, + }, + ) + .await? + .0; + + let data = if request_context.wants_envelope() { + serde_json::from_value::>(result) + .map_err(|e| { + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "message": format!("failed to parse success envelope: {e}"), + })) + })? + .data + } else { + result + }; + + serde_json::from_value::(data).map_err(|e| { + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "message": format!("failed to deserialize edit image result: {e}"), + })) + }) + } +} diff --git a/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs b/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs new file mode 100644 index 000000000..7feaa9a4a --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/editor_tools/mod.rs @@ -0,0 +1,2 @@ +pub mod common; +pub mod edit_image; diff --git a/server-rs/crates/api-server/src/editor_agent/mod.rs b/server-rs/crates/api-server/src/editor_agent/mod.rs new file mode 100644 index 000000000..7f3bcb29e --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/mod.rs @@ -0,0 +1,4 @@ +mod editor_tools; +mod utils; +pub mod api; +mod agent; diff --git a/server-rs/crates/api-server/src/editor_agent/utils.rs b/server-rs/crates/api-server/src/editor_agent/utils.rs new file mode 100644 index 000000000..1af7f7e60 --- /dev/null +++ b/server-rs/crates/api-server/src/editor_agent/utils.rs @@ -0,0 +1,560 @@ +use crate::editor_project::{current_utc_micros, map_editor_project_error}; +use crate::http_error::AppError; +use crate::platform_errors::map_oss_error; +use crate::state::AppState; +use axum::http::StatusCode; +use platform_oss::{LegacyAssetPrefix, OssObjectAccess, OssPutObjectRequest, OssSignedGetObjectUrlRequest}; +use serde::{Deserialize, Serialize}; +use serde_json::json; +use shared_contracts::editor_agent::{ + EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, EditorAgentAttachmentRef, EditorAgentAttachmentSource, + EditorAgentConversationDetail, EditorAgentConversationMessagesDocument, + EditorAgentConversationSummary, EditorAgentMessage, +}; +use shared_kernel::{normalize_optional_string, normalize_required_string}; +use spacetime_client::{ + EditorAgentConversationRecord, EditorAssetLibraryRecord, EditorAssetRecord, + EditorProjectGetRecordInput, EditorProjectRecord, EditorProjectResourceRecord, +}; +use std::collections::BTreeMap; +use std::fmt::Display; +use std::sync::{Arc, Mutex, OnceLock}; + +trait IntoDataKey { + fn into_data_key(self) -> String; +} +impl IntoDataKey for EditorAgentAttachmentRef { + 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()) + .or_else(|| normalize_required_string(self.reference_id.as_str())) + .unwrap_or_else(|| self.image_src.clone()) + } +} +impl IntoDataKey for EditorProjectResourceRecord { + 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_else(|| self.image_src.clone()) + } +} + +impl IntoDataKey for EditorAssetRecord { + 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_else(|| self.image_src.clone()) + } +} + +trait IntoImageId { + fn into_image_id(self) -> ImageId; +} +impl IntoImageId for EditorAgentAttachmentRef { + fn into_image_id(self) -> ImageId { + ImageId { + id: self.into_data_key(), + } + } +} +impl IntoImageId for EditorProjectResourceRecord { + fn into_image_id(self) -> ImageId { + ImageId { + id: self.into_data_key(), + } + } +} +impl IntoImageId for EditorAssetRecord { + fn into_image_id(self) -> ImageId { + ImageId { + id: self.into_data_key(), + } + } +} + +trait IntoResourceId { + fn into_resource_id(self) -> String; +} + +impl IntoResourceId for EditorProjectResourceRecord { + fn into_resource_id(self) -> String { + self.resource_id + } +} + +impl IntoResourceId for EditorAssetRecord { + fn into_resource_id(self) -> String { + self.asset_id + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ImageId { + pub id: String, +} + +impl Serialize for ImageId { + fn serialize(&self, serializer: S) -> Result { + self.id.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for ImageId { + fn deserialize>(deserializer: D) -> Result { + String::deserialize(deserializer).map(|id| ImageId { id }) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageMetadata { + pub tag: String, +} + +impl Display for ImageId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.id) + } +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct EditorAgentMessageResponse { + pub delta_messages: Vec, + pub error_message: Option, +} + +type EditorAgentConversationLockMap = Mutex>>>; +static EDITOR_AGENT_CONVERSATION_LOCKS: OnceLock = OnceLock::new(); + +pub(crate) fn editor_agent_conversation_lock(conversation_id: &str) -> Arc> { + let locks = EDITOR_AGENT_CONVERSATION_LOCKS.get_or_init(|| Mutex::new(BTreeMap::new())); + let mut locks = locks + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + locks + .entry(conversation_id.to_string()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() +} +pub fn now_rfc3339() -> String { + shared_kernel::format_rfc3339(time::OffsetDateTime::now_utc()) + .unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_string()) +} + +const EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES: usize = 2 * 1024 * 1024; +pub async fn require_editor_agent_sidebar_enabled( + state: &AppState, + owner_user_id: &str, +) -> Result<(), AppError> { + match state + .is_image_editor_agent_sidebar_enabled_for_user(Some(owner_user_id)) + .await + { + Ok(true) => Ok(()), + Ok(false) => Err(editor_agent_sidebar_unavailable()), + Err(error) => Err(AppError::from_status(StatusCode::BAD_GATEWAY) + .with_message("读取画布 Agent 灰度配置失败") + .with_details(json!({ + "provider": "spacetimedb", + "message": error.to_string(), + }))), + } +} + +fn editor_agent_sidebar_unavailable() -> AppError { + AppError::from_status(StatusCode::SERVICE_UNAVAILABLE) + .with_message("画布 Agent 暂不可用") + .with_details(json!({ + "provider": "editor-agent", + "reason": "image_editor_agent_sidebar_disabled", + "gateKey": module_runtime::IMAGE_EDITOR_AGENT_SIDEBAR_GATE_KEY, + })) +} +// const EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES: usize = 2 * 1024 * 1024; +pub fn editor_agent_bad_request(message: impl Into) -> AppError { + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": message.into(), + })) +} + +fn editor_agent_oss_unavailable() -> AppError { + AppError::from_status(StatusCode::SERVICE_UNAVAILABLE).with_details(json!({ + "provider": "aliyun-oss", + "reason": "OSS is not configured for editor agent conversations", + })) +} + +fn editor_agent_oss_read_error(message: impl Into) -> AppError { + AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({ + "provider": "aliyun-oss", + "message": message.into(), + })) +} + +fn editor_agent_messages_document_too_large() -> AppError { + AppError::from_status(StatusCode::PAYLOAD_TOO_LARGE).with_details(json!({ + "provider": "editor-agent", + "message": "message document is too large", + "maxBytes": EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES, + })) +} +pub async fn write_messages_document( + state: &AppState, + conversation: &EditorAgentConversationRecord, + document: &EditorAgentConversationMessagesDocument, +) -> Result<(), AppError> { + if document.conversation_id != conversation.conversation_id { + return Err( + AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({ + "provider": "editor-agent", + "message": "message document conversationId does not match metadata", + })), + ); + } + let body = serde_json::to_vec(document).map_err(|error| { + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "provider": "editor-agent", + "message": format!("failed to serialize message document: {error}"), + })) + })?; + if body.len() > EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES { + return Err(editor_agent_messages_document_too_large()); + } + + let oss_client = state + .oss_client() + .ok_or_else(editor_agent_oss_unavailable)?; + let put_result = oss_client + .put_object( + &reqwest::Client::new(), + OssPutObjectRequest { + prefix: LegacyAssetPrefix::EditorAgent, + path_segments: Vec::new(), + file_name: format!("{}.json", conversation.conversation_id), + content_type: Some("application/json; charset=utf-8".to_string()), + access: OssObjectAccess::Private, + metadata: BTreeMap::from([ + ( + "conversation-id".to_string(), + conversation.conversation_id.clone(), + ), + ("project-id".to_string(), conversation.project_id.clone()), + ( + "owner-user-id".to_string(), + conversation.owner_user_id.clone(), + ), + ]), + body, + }, + ) + .await + .map_err(|error| map_oss_error(error, "aliyun-oss"))?; + if put_result.object_key != conversation.messages_object_key { + return Err( + AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({ + "provider": "editor-agent", + "message": "OSS object key mismatch while writing message document", + "expectedObjectKey": conversation.messages_object_key, + "actualObjectKey": put_result.object_key, + })), + ); + } + Ok(()) +} + +const EDITOR_AGENT_MESSAGES_READ_EXPIRE_SECONDS: u64 = 60; + +pub(crate) async fn read_messages_document( + state: &AppState, + conversation: &EditorAgentConversationRecord, +) -> Result { + let oss_client = state + .oss_client() + .ok_or_else(editor_agent_oss_unavailable)?; + let signed = oss_client + .sign_internal_get_object_url(OssSignedGetObjectUrlRequest { + object_key: conversation.messages_object_key.clone(), + expire_seconds: Some(EDITOR_AGENT_MESSAGES_READ_EXPIRE_SECONDS), + }) + .map_err(|error| map_oss_error(error, "aliyun-oss"))?; + let response = reqwest::Client::new() + .get(signed.signed_url.as_str()) + .send() + .await + .map_err(|error| editor_agent_oss_read_error(error.to_string()))?; + if response.status() == StatusCode::NOT_FOUND { + return Ok(empty_messages_document( + conversation.conversation_id.as_str(), + )); + } + if !response.status().is_success() { + return Err(editor_agent_oss_read_error(format!( + "OSS returned non-success status {}", + response.status().as_u16() + ))); + } + if response + .content_length() + .is_some_and(|size| size > EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES as u64) + { + return Err(editor_agent_messages_document_too_large()); + } + let bytes = response + .bytes() + .await + .map_err(|error| editor_agent_oss_read_error(error.to_string()))?; + if bytes.is_empty() { + return Ok(empty_messages_document( + conversation.conversation_id.as_str(), + )); + } + if bytes.len() > EDITOR_AGENT_MESSAGES_DOCUMENT_MAX_BYTES { + return Err(editor_agent_messages_document_too_large()); + } + let document: EditorAgentConversationMessagesDocument = serde_json::from_slice(&bytes) + .map_err(|error| { + editor_agent_oss_read_error(format!("message document JSON invalid: {error}")) + })?; + if document.conversation_id != conversation.conversation_id { + return Err( + AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({ + "provider": "editor-agent", + "message": "message document conversationId does not match metadata", + "conversationId": conversation.conversation_id, + "documentConversationId": document.conversation_id, + })), + ); + } + Ok(document) +} + + +pub async fn ensure_editor_project_access( + state: &AppState, + project_id: &str, + owner_user_id: &str, +) -> Result<(), AppError> { + state + .spacetime_client() + .get_editor_project(EditorProjectGetRecordInput { + project_id: project_id.to_string(), + owner_user_id: owner_user_id.to_string(), + }) + .await + .map(|_| ()) + .map_err(map_editor_project_error) +} + +pub async fn normalize_editor_agent_attachments( + state: &AppState, + conversation: &EditorAgentConversationRecord, + attachments: &[EditorAgentAttachmentRef], +) -> Result, AppError> { + if attachments.is_empty() { + return Ok(Vec::new()); + } + + let needs_canvas_resources = attachments + .iter() + .any(|attachment| attachment.source == EditorAgentAttachmentSource::CanvasResource); + let needs_library_assets = attachments + .iter() + .any(|attachment| attachment.source == EditorAgentAttachmentSource::LibraryAsset); + + let project = if needs_canvas_resources { + Some( + state + .spacetime_client() + .get_editor_project(EditorProjectGetRecordInput { + project_id: conversation.project_id.clone(), + owner_user_id: conversation.owner_user_id.clone(), + }) + .await + .map_err(map_editor_project_error)?, + ) + } else { + None + }; + let library = if needs_library_assets { + Some( + state + .spacetime_client() + .get_editor_asset_library(conversation.owner_user_id.clone(), current_utc_micros()) + .await + .map_err(map_editor_project_error)?, + ) + } else { + None + }; + + attachments + .iter() + .map(|attachment| { + normalize_editor_agent_attachment( + conversation, + project.as_ref(), + library.as_ref(), + attachment, + ) + }) + .collect() +} + +fn normalize_editor_agent_attachment( + conversation: &EditorAgentConversationRecord, + project: Option<&EditorProjectRecord>, + library: Option<&EditorAssetLibraryRecord>, + attachment: &EditorAgentAttachmentRef, +) -> Result { + let reference_id = normalize_required_string(attachment.reference_id.as_str()) + .ok_or_else(|| editor_agent_bad_request("attachment.referenceId is required"))?; + match attachment.source { + EditorAgentAttachmentSource::CanvasResource => { + let project = project.ok_or_else(|| { + editor_agent_bad_request("canvas resource attachment project context missing") + })?; + let resource = project + .resources + .iter() + .find(|resource| resource.resource_id == reference_id) + .ok_or_else(|| { + editor_agent_bad_request(format!( + "canvas resource attachment not found in current project: {reference_id}" + )) + })?; + normalize_canvas_resource_attachment(conversation, attachment, resource) + } + EditorAgentAttachmentSource::LibraryAsset => { + let library = library.ok_or_else(|| { + editor_agent_bad_request("library asset attachment context missing") + })?; + let asset = library + .assets + .iter() + .find(|asset| asset.asset_id == reference_id) + .ok_or_else(|| { + editor_agent_bad_request(format!( + "library asset attachment not found for current user: {reference_id}" + )) + })?; + normalize_library_asset_attachment(attachment, asset) + } + } +} + +pub fn normalize_canvas_resource_attachment( + conversation: &EditorAgentConversationRecord, + attachment: &EditorAgentAttachmentRef, + resource: &EditorProjectResourceRecord, +) -> Result { + if resource.project_id != conversation.project_id + || resource.owner_user_id != conversation.owner_user_id + { + return Err(editor_agent_bad_request( + "canvas resource attachment does not belong to this conversation project", + )); + } + validate_attachment_object_key( + attachment.object_key.as_deref(), + resource.object_key.as_deref(), + resource.resource_id.as_str(), + )?; + + Ok(EditorAgentAttachmentRef { + source: EditorAgentAttachmentSource::CanvasResource, + reference_id: resource.resource_id.clone(), + object_key: resource.object_key.clone(), + image_src: resource.image_src.clone(), + thumbnail_src: None, + label: normalize_optional_string(attachment.label.clone()), + width: Some(resource.width), + height: Some(resource.height), + }) +} + +pub fn normalize_library_asset_attachment( + attachment: &EditorAgentAttachmentRef, + asset: &EditorAssetRecord, +) -> Result { + validate_attachment_object_key( + attachment.object_key.as_deref(), + asset.object_key.as_deref(), + asset.asset_id.as_str(), + )?; + + Ok(EditorAgentAttachmentRef { + source: EditorAgentAttachmentSource::LibraryAsset, + reference_id: asset.asset_id.clone(), + object_key: asset.object_key.clone(), + image_src: asset.image_src.clone(), + thumbnail_src: asset.thumbnail_src.clone(), + label: normalize_optional_string(attachment.label.clone()) + .or_else(|| Some(asset.label.clone())), + width: Some(asset.width), + height: Some(asset.height), + }) +} + +fn validate_attachment_object_key( + submitted_object_key: Option<&str>, + stored_object_key: Option<&str>, + reference_id: &str, +) -> Result<(), AppError> { + let Some(submitted_object_key) = submitted_object_key.and_then(normalize_required_string) + else { + return Ok(()); + }; + let Some(stored_object_key) = stored_object_key.and_then(normalize_required_string) else { + return Err(editor_agent_bad_request(format!( + "attachment objectKey is not available for reference: {reference_id}" + ))); + }; + if submitted_object_key != stored_object_key { + return Err(editor_agent_bad_request(format!( + "attachment objectKey does not match reference: {reference_id}" + ))); + } + Ok(()) +} + +pub fn conversation_summary_from_record( + conversation: EditorAgentConversationRecord, +) -> EditorAgentConversationSummary { + EditorAgentConversationSummary { + conversation_id: conversation.conversation_id, + project_id: conversation.project_id, + updated_at: conversation.updated_at, + } +} + +pub fn conversation_detail_from_record( + conversation: EditorAgentConversationRecord, + messages: Vec, +) -> EditorAgentConversationDetail { + EditorAgentConversationDetail { + conversation_id: conversation.conversation_id, + project_id: conversation.project_id, + title: conversation.title, + created_at: conversation.created_at, + updated_at: conversation.updated_at, + messages, + } +} + +pub fn empty_messages_document(conversation_id: &str) -> EditorAgentConversationMessagesDocument { + EditorAgentConversationMessagesDocument { + version: EDITOR_AGENT_MESSAGES_DOCUMENT_VERSION, + conversation_id: conversation_id.to_string(), + messages: Vec::new(), + } +}