fix(agent): scrub model-supplied _owner on the qualified mcp__ image path (P2, #4123 review)

The wrapper path (_call_mcp_tool) injected the trusted _owner only for the builtin
"generate_image" name. A model could instead issue the fully-qualified
mcp__image_gen__generate_image directly, which the generic mcp__ dispatch branch
passed straight to MCP with the model's JSON args — letting a model-supplied _owner
through to resolve another user's private image endpoint/API key and tag the gallery
row under the spoofed owner.

Centralise the contract in _apply_trusted_owner(name, args, owner): drop any
model-supplied _owner (server-side-only field), then inject the trusted caller owner
for owner-scoped tools — keyed by both the builtin name and its qualified mcp__ name.
Both dispatch paths now call it, so they enforce identically.

Tests: qualified-name and builtin-name spoof regressions on the real helper, a strip
on a non-owner-scoped mcp__ tool, the no-owner case, and a guard that both dispatch
paths call the helper. Full suite green (3544 passed).

Co-Authored-By: Claude <noreply@anthropic.com>

Author: nsgds

src/tool_execution.py

@@ -357,6 +357,31 @@ def _parse_qualified_mcp_args(tool: str, content: str) -> tuple[Dict, Optional[s
         return {}, None
     return parsed, None
 
+# Builtin tools that must run under the TRUSTED, server-side caller owner — never a
+# model-supplied one. The image generator resolves the caller's private endpoints/keys
+# and tags the gallery row with the owner, so `_owner` is a server-only arg.
+_OWNER_SCOPED_TOOLS = {"generate_image"}
+# Their fully-qualified mcp__ names, so a model that calls the tool by its qualified
+# name through the generic mcp__ dispatch path gets the same enforcement.
+_OWNER_SCOPED_QUALIFIED = {
+    f"mcp__{_MCP_TOOL_MAP[t][0]}__{_MCP_TOOL_MAP[t][1]}" for t in _OWNER_SCOPED_TOOLS
+}
+
+
+def _apply_trusted_owner(name: str, args: Dict, owner: Optional[str]) -> Dict:
+    """Enforce the trusted-owner contract for owner-scoped tools on BOTH dispatch
+    paths — the builtin name (via _call_mcp_tool) and a model-issued fully-qualified
+    ``mcp__…`` call. `_owner` is server-side-only: drop any model-supplied value, then
+    inject the trusted caller owner for owner-scoped tools. Without this a direct
+    ``mcp__image_gen__generate_image`` call could spoof `_owner` to resolve another
+    user's image endpoint/key and mis-tag the gallery row (#4123 review)."""
+    if not isinstance(args, dict):
+        return args
+    args = {k: v for k, v in args.items() if k != "_owner"}
+    if owner and (name in _OWNER_SCOPED_TOOLS or name in _OWNER_SCOPED_QUALIFIED):
+        args["_owner"] = owner
+    return args
+
 
 def _parse_generate_image(content: str) -> Dict:
     # Native tool-callers (supports_tools=True endpoints) deliver args as a JSON
@@ -479,14 +504,11 @@ async def _call_mcp_tool(
     server_id, tool_name = _MCP_TOOL_MAP[tool]
     qualified = f"mcp__{server_id}__{tool_name}"
     args = _build_mcp_args(tool, content)
-    # Inject the TRUSTED, server-side owner for owner-scoped MCP tools. The image
-    # generator must resolve the *caller's* endpoints (not whichever endpoint
-    # happens to resolve first) and tag the gallery row with the owner. This is
-    # set here — after arg parsing, never in the tool schema — so the model can't
-    # control or spoof it (a model-supplied "_owner" is dropped by _build_mcp_args
-    # and overwritten here).
-    if owner and tool == "generate_image":
-        args = {**args, "_owner": owner}
+    # Enforce the trusted, server-side owner for owner-scoped tools (the image
+    # generator must resolve the *caller's* endpoints and tag the gallery row with
+    # the owner, never a model-supplied value). Centralised in _apply_trusted_owner
+    # so the generic mcp__ path applies the identical scrub-and-inject (#4123 review).
+    args = _apply_trusted_owner(tool, args, owner)
     result = await mcp.call_tool(qualified, args)
 
     # If MCP server not connected, try direct fallback
@@ -961,6 +983,10 @@ async def _execute_tool_block_impl(
             if parse_error:
                 result = {"error": parse_error, "exit_code": 1}
             else:
+                # A model can issue a fully-qualified mcp__… call directly; apply the
+                # trusted-owner scrub/inject the builtin path uses so it can't spoof
+                # `_owner` on an owner-scoped builtin (e.g. mcp__image_gen__generate_image).
+                args = _apply_trusted_owner(tool, args, owner)
                 if tool.startswith("mcp__email__") and owner:
                     args = dict(args)
                     args[_EMAIL_MCP_OWNER_ARG] = owner

tests/test_generate_image_owner_scope.py

@@ -94,11 +94,43 @@ def test_mcp_image_server_threads_trusted_owner():
     assert "owner=owner," in src                       # GalleryImage tagged with the owner
 
 
-def test_bridge_injects_trusted_owner_for_generate_image_only():
-    """_call_mcp_tool injects `_owner` only for generate_image, gated by tool name."""
-    src = inspect.getsource(te._call_mcp_tool)
-    assert 'tool == "generate_image"' in src
-    assert '"_owner": owner' in src
+def test_bridge_applies_trusted_owner_via_helper():
+    """Both dispatch paths delegate owner enforcement to _apply_trusted_owner so they
+    scrub/inject identically: the builtin-name wrapper (_call_mcp_tool) and the generic
+    qualified mcp__ branch (_execute_tool_block_impl). Guards the #4123 P2 — a direct
+    `mcp__image_gen__generate_image` call must not bypass owner injection."""
+    assert "_apply_trusted_owner(tool, args, owner)" in inspect.getsource(te._call_mcp_tool)
+    assert "_apply_trusted_owner(tool, args, owner)" in inspect.getsource(te._execute_tool_block_impl)
+
+
+def test_apply_trusted_owner_qualified_image_call_cannot_spoof_owner():
+    """P2 regression (qualified name): a model-issued `mcp__image_gen__generate_image`
+    call carrying a spoofed `_owner` has it dropped and replaced with the trusted owner."""
+    qualified = "mcp__image_gen__generate_image"
+    assert qualified in te._OWNER_SCOPED_QUALIFIED
+    out = te._apply_trusted_owner(qualified, {"prompt": "a cat", "_owner": "attacker"}, "alice")
+    assert out["_owner"] == "alice"          # trusted owner wins
+    assert out["prompt"] == "a cat"          # real args preserved
+
+
+def test_apply_trusted_owner_builtin_name_cannot_spoof_owner():
+    """P1 path (builtin name): same scrub/inject for the friendly `generate_image` name."""
+    out = te._apply_trusted_owner("generate_image", {"_owner": "attacker"}, "alice")
+    assert out["_owner"] == "alice"
+
+
+def test_apply_trusted_owner_strips_owner_from_non_scoped_tool():
+    """A model-supplied `_owner` is never honoured for a non-owner-scoped mcp__ tool —
+    it's stripped, not re-injected (the field is server-side-only for any tool)."""
+    out = te._apply_trusted_owner("mcp__filesystem__read_file", {"path": "/x", "_owner": "attacker"}, "alice")
+    assert "_owner" not in out
+    assert out["path"] == "/x"
+
+
+def test_apply_trusted_owner_no_owner_strips_spoofed():
+    """With no trusted owner to inject, a spoofed `_owner` is still dropped."""
+    out = te._apply_trusted_owner("mcp__image_gen__generate_image", {"_owner": "attacker"}, None)
+    assert "_owner" not in out
 
 
 def test_resolve_model_isolates_private_image_endpoint_by_owner():