Genarrative/server-rs/crates/api-server/src/editor_agent/context.rs

use crate::editor_agent::utils::IntoDataKey;
use crate::editor_project::{current_utc_micros, map_editor_project_error};
use crate::http_error::AppError;
use crate::state::AppState;
use platform_editor_agent::agent::asset::{ImageId, ImageMetadata};
use platform_editor_agent::agent::tools::context::EditorToolContext;
use shared_contracts::editor_agent::{
    EditorAgentAttachmentSource, EditorAgentConversationMessagesDocument,
};
use spacetime_client::{
    EditorAgentConversationRecord, EditorAssetLibraryRecord, EditorProjectGetRecordInput,
    EditorProjectRecord,
};
use std::collections::HashMap;

pub async fn build_tool_context(
    state: &AppState,
    conversation: &EditorAgentConversationRecord,
    document: &EditorAgentConversationMessagesDocument,
) -> Result<EditorToolContext, AppError> {
    let needs_project = document.messages.iter().any(|message| {
        message
            .attachments
            .iter()
            .any(|attachment| attachment.source == EditorAgentAttachmentSource::CanvasResource)
            || message.tool_call.as_ref().is_some_and(|tool_call| {
                tool_call.images.iter().any(|image| {
                    image
                        .resource_id
                        .as_deref()
                        .is_some_and(|id| !id.trim().is_empty())
                })
            })
    });
    let needs_library = document.messages.iter().any(|message| {
        message
            .attachments
            .iter()
            .any(|attachment| attachment.source == EditorAgentAttachmentSource::LibraryAsset)
    });

    let project = if needs_project {
        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 {
        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
    };

    Ok(build_tool_context_from_authoritative_records(
        document,
        project.as_ref(),
        library.as_ref(),
    ))
}

fn build_tool_context_from_authoritative_records(
    document: &EditorAgentConversationMessagesDocument,
    project: Option<&EditorProjectRecord>,
    library: Option<&EditorAssetLibraryRecord>,
) -> EditorToolContext {
    let mut images: HashMap<ImageId, ImageMetadata> = HashMap::new();

    for msg in document.messages.iter().rev() {
        for a in &msg.attachments {
            let data_key = a.clone().into_data_key();
            let image_id = ImageId::from_data_key(&data_key);
            let metadata = ImageMetadata {
                data_key,
                reference_id: Some(a.reference_id.clone()),
                asset_kind: match a.source {
                    EditorAgentAttachmentSource::CanvasResource => project
                        .and_then(|project| {
                            project
                                .resources
                                .iter()
                                .find(|resource| resource.resource_id == a.reference_id)
                        })
                        .and_then(|resource| resource.asset_kind.clone()),
                    EditorAgentAttachmentSource::LibraryAsset => library
                        .and_then(|library| {
                            library
                                .assets
                                .iter()
                                .find(|asset| asset.asset_id == a.reference_id)
                        })
                        .and_then(|asset| asset.asset_kind.clone()),
                },
                image_src: a.image_src.clone(),
                object_key: a.object_key.clone(),
                thumbnail_src: a.thumbnail_src.clone(),
                label: a.label.clone(),
                width: a.width,
                height: a.height,
            };
            images.entry(image_id).or_insert(metadata);
        }

        if let Some(tc) = &msg.tool_call {
            for img in &tc.images {
                let data_key = img.clone().into_data_key();
                let image_id = ImageId::from_data_key(&data_key);
                let metadata = ImageMetadata {
                    data_key,
                    reference_id: img.resource_id.clone(),
                    asset_kind: img.resource_id.as_deref().and_then(|resource_id| {
                        project
                            .and_then(|project| {
                                project
                                    .resources
                                    .iter()
                                    .find(|resource| resource.resource_id == resource_id)
                            })
                            .and_then(|resource| resource.asset_kind.clone())
                    }),
                    image_src: img.image_src.clone(),
                    object_key: img.object_key.clone(),
                    thumbnail_src: img.thumbnail_src.clone(),
                    label: None,
                    width: img.width,
                    height: img.height,
                };
                images.entry(image_id).or_insert(metadata);
            }
        }
    }

    EditorToolContext { images }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::editor_agent::tool::editor_agent_tool;
    use platform_editor_agent::agent::tools::generate_icon_spritesheet::GenerateIconSpritesheetTool;
    use platform_editor_agent::framework::tool::Tool;
    use serde_json::json;
    use shared_contracts::editor_agent::{
        EditorAgentAttachmentRef, EditorAgentAttachmentSource, EditorAgentMessage,
        EditorAgentMessageRole,
    };
    use spacetime_client::{
        EditorCanvasRecord, EditorCanvasViewportRecord, EditorProjectResourceRecord,
    };

    fn message(id: usize, label: &str, image_src: &str) -> EditorAgentMessage {
        EditorAgentMessage {
            id,
            client_message_id: Some(format!("message-{id}")),
            role: EditorAgentMessageRole::User,
            text: String::new(),
            attachments: vec![EditorAgentAttachmentRef {
                source: EditorAgentAttachmentSource::CanvasResource,
                reference_id: "resource-1".to_string(),
                object_key: Some("generated/reference.png".to_string()),
                image_src: image_src.to_string(),
                thumbnail_src: Some(format!("{image_src}?thumbnail=1")),
                label: Some(label.to_string()),
                width: Some(640),
                height: Some(480),
            }],
            tool_call: None,
            created_at: "2026-07-23T00:00:00Z".to_string(),
        }
    }

    fn project_with_resource_asset_kind(asset_kind: Option<&str>) -> EditorProjectRecord {
        let viewport = EditorCanvasViewportRecord {
            x: 0.0,
            y: 0.0,
            scale: 1.0,
        };
        EditorProjectRecord {
            project_id: "project-1".to_string(),
            owner_user_id: "user-1".to_string(),
            title: "测试工程".to_string(),
            canvas: EditorCanvasRecord {
                canvas_id: "canvas-1".to_string(),
                project_id: "project-1".to_string(),
                title: "测试画布".to_string(),
                viewport: viewport.clone(),
                layers: json!([]),
                revision: 0,
                layout_storage_version: 2,
                background_color: None,
                created_at: "2026-08-08T00:00:00Z".to_string(),
                updated_at: "2026-08-08T00:00:00Z".to_string(),
            },
            viewport,
            layers: json!([]),
            resources: vec![EditorProjectResourceRecord {
                resource_id: "resource-1".to_string(),
                project_id: "project-1".to_string(),
                owner_user_id: "user-1".to_string(),
                asset_object_id: Some("object-1".to_string()),
                image_src: "/api/assets/read/current.png".to_string(),
                object_key: Some("generated/reference.png".to_string()),
                width: 640,
                height: 480,
                source_type: "generated".to_string(),
                prompt: None,
                actual_prompt: None,
                model: None,
                provider: None,
                task_id: None,
                source_resource_id: None,
                asset_kind: asset_kind.map(str::to_string),
                generation_inputs: None,
                public_showcase_enabled: false,
                created_at: "2026-08-08T00:00:00Z".to_string(),
                updated_at: "2026-08-08T00:00:00Z".to_string(),
                image_sequence_frames: None,
                image_sequence_duration_ms: None,
            }],
            created_at: "2026-08-08T00:00:00Z".to_string(),
            updated_at: "2026-08-08T00:00:00Z".to_string(),
        }
    }

    #[test]
    fn tool_context_keeps_complete_metadata_from_the_latest_image_reference() {
        let document = EditorAgentConversationMessagesDocument {
            version: 2,
            conversation_id: "conversation-1".to_string(),
            messages: vec![
                message(0, "旧名称", "/api/assets/read/old.png"),
                message(1, "最新名称", "/api/assets/read/current.png"),
            ],
        };

        let context = build_tool_context_from_authoritative_records(&document, None, None);
        let image_id = ImageId::from_data_key("generated/reference.png");
        let metadata = context
            .image_metadata(&image_id)
            .expect("latest image metadata should be present");

        assert_eq!(metadata.data_key, "generated/reference.png");
        assert_eq!(metadata.image_src, "/api/assets/read/current.png");
        assert_eq!(metadata.label.as_deref(), Some("最新名称"));
        assert_eq!(metadata.width, Some(640));
        assert_eq!(metadata.height, Some(480));
    }

    #[test]
    fn icon_spritesheet_validation_uses_authoritative_project_asset_kind() {
        let document = EditorAgentConversationMessagesDocument {
            version: 2,
            conversation_id: "conversation-1".to_string(),
            messages: vec![message(0, "主规范", "/api/assets/read/current.png")],
        };
        let args = json!({
            "reference_image_id": ImageId::from_data_key("generated/reference.png").id,
            "icon_descriptions": ["背包"]
        });

        let ordinary_project = project_with_resource_asset_kind(None);
        let ordinary_context =
            build_tool_context_from_authoritative_records(&document, Some(&ordinary_project), None);
        let ordinary_tool = editor_agent_tool(GenerateIconSpritesheetTool::NAME, &ordinary_context)
            .expect("icon spritesheet tool should resolve");
        let error = ordinary_tool
            .validate_args(&args)
            .expect_err("ordinary registered image must be rejected before confirmation");
        assert!(error.to_string().contains("icon-spec"));

        let icon_spec_project = project_with_resource_asset_kind(Some("icon-spec"));
        let icon_spec_context = build_tool_context_from_authoritative_records(
            &document,
            Some(&icon_spec_project),
            None,
        );
        let icon_spec_tool =
            editor_agent_tool(GenerateIconSpritesheetTool::NAME, &icon_spec_context)
                .expect("icon spritesheet tool should resolve");
        assert!(icon_spec_tool.validate_args(&args).is_ok());
    }
}