From a2926674ddd12a715f20ca373dfc7116df1df96d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E5=BE=B7=E5=AE=87?= <kvtodev@outlook.com>
Date: Wed, 15 Jul 2026 18:08:35 +0800
Subject: [PATCH] restore context and image_args functions as test utils

---
 .../api-server/src/editor_agent/pricing.rs    | 45 ++++++++++++++-----
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/server-rs/crates/api-server/src/editor_agent/pricing.rs b/server-rs/crates/api-server/src/editor_agent/pricing.rs
index 1cf0606c0..9824f9ed6 100644
--- a/server-rs/crates/api-server/src/editor_agent/pricing.rs
+++ b/server-rs/crates/api-server/src/editor_agent/pricing.rs
@@ -1,24 +1,47 @@
-use std::collections::HashMap;
-use platform_editor_agent::framework::tool::Tool;
-use platform_image::GPT_IMAGE_2_MODEL;
-use platform_editor_agent::agent::tools::context::EditorToolContext;
+use crate::editor_generation_config::{
+    EditorGenerationPricingConfig, load_editor_generation_pricing_from_paths,
+};
 #[cfg(test)]
 use platform_editor_agent::agent::asset::ImageId;
+use platform_editor_agent::agent::tools::context::EditorToolContext;
 use platform_editor_agent::agent::tools::edit_image::{EditImageTool, EditImageToolArgs};
 use platform_editor_agent::agent::tools::generate_background_music::{
-  GenerateBackgroundMusicTool, GenerateBackgroundMusicToolArgs,
+    GenerateBackgroundMusicTool, GenerateBackgroundMusicToolArgs,
 };
 use platform_editor_agent::agent::tools::generate_character::GenerateCharacterTool;
 use platform_editor_agent::agent::tools::generate_icon_spritesheet::{
-  GenerateIconSpritesheetTool, GenerateIconSpritesheetToolArgs,
+    GenerateIconSpritesheetTool, GenerateIconSpritesheetToolArgs,
+};
+use platform_editor_agent::agent::tools::generate_image::{
+    GenerateImageTool, GenerateImageToolArgs,
 };
-use platform_editor_agent::agent::tools::generate_image::{GenerateImageTool, GenerateImageToolArgs};
 use platform_editor_agent::agent::tools::generate_sound_effect::{
-  GenerateSoundEffectTool, GenerateSoundEffectToolArgs,
+    GenerateSoundEffectTool, GenerateSoundEffectToolArgs,
 };
 use platform_editor_agent::agent::tools::generate_ui_design::GenerateUiDesignTool;
-use platform_editor_agent::agent::tools::generate_video::{GenerateVideoTool, GenerateVideoToolArgs};
-use crate::editor_generation_config::{load_editor_generation_pricing_from_paths, EditorGenerationPricingConfig};
+use platform_editor_agent::agent::tools::generate_video::{
+    GenerateVideoTool, GenerateVideoToolArgs,
+};
+use platform_editor_agent::framework::tool::Tool;
+use platform_image::GPT_IMAGE_2_MODEL;
+use std::collections::HashMap;
+
+#[cfg(test)]
+fn context() -> EditorToolContext {
+    EditorToolContext {
+        images: HashMap::new(),
+    }
+}
+
+#[cfg(test)]
+fn image_args(image_size: Option<&str>) -> GenerateImageToolArgs {
+    GenerateImageToolArgs {
+        prompt: "生成图片".to_string(),
+        reference_image_ids: Vec::new(),
+        aspect_ratio: Some("1:1".to_string()),
+        image_size: image_size.map(ToOwned::to_owned),
+    }
+}
 
 #[test]
 fn every_editor_agent_tool_exposes_argument_based_pricing() {
@@ -105,7 +128,6 @@ fn every_editor_agent_tool_exposes_argument_based_pricing() {
             &GenerateSoundEffectToolArgs {
                 prompt: "按钮点击声".to_string(),
                 duration: None,
-                model: None,
             },
         ),
         5
@@ -169,7 +191,6 @@ fn pricing_uses_the_supplied_runtime_snapshot() {
             &GenerateSoundEffectToolArgs {
                 prompt: "sound".to_string(),
                 duration: None,
-                model: None,
             },
         ),
         19