refactor: simplify handoff image_urls data flow

Author: zouyonghe

astrbot/core/astr_agent_tool_exec.py

@@ -68,38 +68,26 @@ def _is_supported_image_ref(cls, image_ref: str) -> bool:
         return ext in cls._ALLOWED_IMAGE_EXTENSIONS
 
     @classmethod
-    def _coerce_image_urls(cls, image_urls: T.Any) -> list[T.Any]:
-        if image_urls is None:
-            return []
-        if isinstance(image_urls, str):
-            return [image_urls]
-        if isinstance(image_urls, (Sequence, AbstractSet)) and not isinstance(
-            image_urls, (str, bytes, bytearray)
+    async def _collect_handoff_image_urls(
+        cls,
+        run_context: ContextWrapper[AstrAgentContext],
+        image_urls_raw: T.Any,
+    ) -> list[str]:
+        candidates: list[T.Any] = []
+        if image_urls_raw is None:
+            pass
+        elif isinstance(image_urls_raw, str):
+            candidates.append(image_urls_raw)
+        elif isinstance(image_urls_raw, (Sequence, AbstractSet)) and not isinstance(
+            image_urls_raw, (str, bytes, bytearray)
         ):
-            return list(image_urls)
-        logger.warning(
-            "Unsupported image_urls type in handoff tool args: %s",
-            type(image_urls).__name__,
-        )
-        return []
-
-    @classmethod
-    def _filter_supported_image_urls(cls, candidates: list[T.Any]) -> list[str]:
-        normalized = normalize_and_dedupe_strings(candidates)
-        sanitized = [item for item in normalized if cls._is_supported_image_ref(item)]
-        dropped_count = len(normalized) - len(sanitized)
-        if dropped_count > 0:
+            candidates.extend(image_urls_raw)
+        else:
             logger.warning(
-                "Dropped %d invalid image_urls entries in handoff tool args.",
-                dropped_count,
+                "Unsupported image_urls type in handoff tool args: %s",
+                type(image_urls_raw).__name__,
             )
-        return sanitized
 
-    @classmethod
-    async def _iter_event_image_paths(
-        cls, run_context: ContextWrapper[AstrAgentContext]
-    ) -> list[str]:
-        paths: list[str] = []
         event = getattr(run_context.context, "event", None)
         message_obj = getattr(event, "message_obj", None)
         message = getattr(message_obj, "message", None)
@@ -110,26 +98,24 @@ async def _iter_event_image_paths(
                 try:
                     path = await component.convert_to_file_path()
                     if path and cls._is_supported_image_ref(path):
-                        paths.append(path)
+                        candidates.append(path)
                 except Exception as e:
                     logger.error(
                         "Failed to convert handoff image component at index %d: %s",
                         idx,
                         e,
                         exc_info=True,
                     )
-        return paths
 
-    @classmethod
-    async def _prepare_handoff_image_urls(
-        cls,
-        run_context: ContextWrapper[AstrAgentContext],
-        image_urls: T.Any,
-    ) -> list[str]:
-        candidates = cls._coerce_image_urls(image_urls)
-        event_paths = await cls._iter_event_image_paths(run_context)
-        candidates.extend(event_paths)
-        return cls._filter_supported_image_urls(candidates)
+        normalized = normalize_and_dedupe_strings(candidates)
+        sanitized = [item for item in normalized if cls._is_supported_image_ref(item)]
+        dropped_count = len(normalized) - len(sanitized)
+        if dropped_count > 0:
+            logger.warning(
+                "Dropped %d invalid image_urls entries in handoff tool args.",
+                dropped_count,
+            )
+        return sanitized
 
     @classmethod
     async def execute(cls, tool, run_context, **tool_args):
@@ -151,7 +137,7 @@ async def execute(cls, tool, run_context, **tool_args):
                 ):
                     yield r
                 return
-            async for r in cls._execute_handoff(tool, run_context, tool_args):
+            async for r in cls._execute_handoff(tool, run_context, **tool_args):
                 yield r
             return
 
@@ -254,14 +240,17 @@ async def _execute_handoff(
         cls,
         tool: HandoffTool,
         run_context: ContextWrapper[AstrAgentContext],
-        tool_args: dict[str, T.Any],
+        **tool_args: T.Any,
     ):
         input_ = tool_args.get("input")
-        image_urls = await cls._prepare_handoff_image_urls(
+        image_urls = await cls._collect_handoff_image_urls(
             run_context,
             tool_args.get("image_urls"),
         )
-        tool_args["image_urls"] = image_urls
+        effective_tool_args: dict[str, T.Any] = {
+            **tool_args,
+            "image_urls": image_urls,
+        }
 
         # Build handoff toolset from registered tools plus runtime computer tools.
         toolset = cls._build_handoff_toolset(run_context, tool.agent.tools)
@@ -295,7 +284,7 @@ async def _execute_handoff(
             event=event,
             chat_provider_id=prov_id,
             prompt=input_,
-            image_urls=image_urls,
+            image_urls=effective_tool_args["image_urls"],
             system_prompt=tool.agent.instructions,
             tools=toolset,
             contexts=contexts,
@@ -360,9 +349,8 @@ async def _do_handoff_background(
     ) -> None:
         """Run the subagent handoff and, on completion, wake the main agent."""
         result_text = ""
-        prepared_tool_args = dict(tool_args)
         try:
-            async for r in cls._execute_handoff(tool, run_context, prepared_tool_args):
+            async for r in cls._execute_handoff(tool, run_context, **tool_args):
                 if isinstance(r, mcp.types.CallToolResult):
                     for content in r.content:
                         if isinstance(content, mcp.types.TextContent):
@@ -379,7 +367,7 @@ async def _do_handoff_background(
             task_id=task_id,
             tool_name=tool.name,
             result_text=result_text,
-            tool_args=prepared_tool_args,
+            tool_args=tool_args,
             note=(
                 event.get_extra("background_note")
                 or f"Background task for subagent '{tool.agent.name}' finished."

tests/unit/test_astr_agent_tool_exec.py

@@ -30,7 +30,7 @@ def _build_run_context(message_components: list[object] | None = None):
 
 
 @pytest.mark.asyncio
-async def test_prepare_handoff_image_urls_normalizes_filters_and_appends_event_image(
+async def test_collect_handoff_image_urls_normalizes_filters_and_appends_event_image(
     monkeypatch: pytest.MonkeyPatch,
 ):
     async def _fake_convert_to_file_path(self):
@@ -46,7 +46,7 @@ async def _fake_convert_to_file_path(self):
         123,
     )
 
-    image_urls = await FunctionToolExecutor._prepare_handoff_image_urls(
+    image_urls = await FunctionToolExecutor._collect_handoff_image_urls(
         run_context,
         image_urls_input,
     )
@@ -59,7 +59,7 @@ async def _fake_convert_to_file_path(self):
 
 
 @pytest.mark.asyncio
-async def test_prepare_handoff_image_urls_skips_failed_event_image_conversion(
+async def test_collect_handoff_image_urls_skips_failed_event_image_conversion(
     monkeypatch: pytest.MonkeyPatch,
 ):
     async def _fake_convert_to_file_path(self):
@@ -68,7 +68,7 @@ async def _fake_convert_to_file_path(self):
     monkeypatch.setattr(Image, "convert_to_file_path", _fake_convert_to_file_path)
 
     run_context = _build_run_context([Image(file="file:///tmp/original.png")])
-    image_urls = await FunctionToolExecutor._prepare_handoff_image_urls(
+    image_urls = await FunctionToolExecutor._collect_handoff_image_urls(
         run_context,
         ["https://example.com/a.png"],
     )
@@ -82,23 +82,14 @@ async def test_do_handoff_background_reports_prepared_image_urls(
 ):
     captured: dict = {}
 
-    async def _unexpected_prepare(cls, run_context, image_urls):
-        raise AssertionError("background path should not pre-prepare image urls")
-
-    async def _fake_execute_handoff(cls, tool, run_context, tool_args):
-        tool_args["image_urls"] = ["prepared://image.png"]
+    async def _fake_execute_handoff(cls, tool, run_context, **tool_args):
         yield mcp.types.CallToolResult(
             content=[mcp.types.TextContent(type="text", text="ok")]
         )
 
     async def _fake_wake(cls, run_context, **kwargs):
         captured.update(kwargs)
 
-    monkeypatch.setattr(
-        FunctionToolExecutor,
-        "_prepare_handoff_image_urls",
-        classmethod(_unexpected_prepare),
-    )
     monkeypatch.setattr(
         FunctionToolExecutor,
         "_execute_handoff",
@@ -119,4 +110,4 @@ async def _fake_wake(cls, run_context, **kwargs):
         image_urls="https://example.com/raw.png",
     )
 
-    assert captured["tool_args"]["image_urls"] == ["prepared://image.png"]
+    assert captured["tool_args"]["image_urls"] == "https://example.com/raw.png"