Python: Fix toolbox consent flow in hosted agent (#6249)

* Fix toolbox consent flow in hosted agent * Resolve conflict * Make unused tool as comment * Fix tests
2026-06-16 21:04:09 +08:00 · 2026-06-04 13:28:59 -07:00
parent bb9ed63a34
commit dbc312a78a
5 changed files with 183 additions and 110 deletions
@@ -11,7 +11,7 @@ import tempfile
 import threading
 from collections.abc import AsyncIterable, AsyncIterator, Generator, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, suppress
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, dataclass, is_dataclass
 from pathlib import Path
 from typing import Protocol, cast

@@ -264,28 +264,73 @@ def _checkpoint_storage_for_context(root: str, context_id: str) -> FileCheckpoin

 # Foundry Toolbox Auth integration
 # Consent-URL error code returned by the Foundry MCP gateway when calling `/list`
-CONSENT_ERROR_CODE = -32007
+CONSENT_ERROR_CODE = -32006


-def consent_url_from_error(exc: BaseException) -> str | None:
-    """Return the consent URL when ``exc`` wraps a Foundry MCP gateway consent error.
+@dataclass
+class ConsentError:
+    name: str
+    consent_url: str

-    The Agent Framework MCP layer surfaces gateway consent failures by wrapping the underlying
-    ``McpError`` inside an :class:`AgentFrameworkException` (typically a ``ToolExecutionException``
-    raised from ``MCPStreamableHTTPTool.__aenter__``). This helper inspects ``exc.args`` for a
-    wrapped ``McpError`` whose ``error.code`` is :data:`CONSENT_ERROR_CODE`; when found, the
-    consent link the gateway returned in ``error.message`` is returned. Returns ``None`` for
-    anything else, so callers can do ``if (url := consent_url_from_error(ex)) is None: raise``.
+
+def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None:
+    """Return the consent URLs when ``exc`` wraps Foundry MCP gateway consent errors.

    Args:
        exc: The exception to inspect.

    Returns:
-        The consent URL if ``exc`` wraps a consent ``McpError``, otherwise ``None``.
+        The consent URL(s) extracted from the error, or ``None`` if no consent error was found.
    """
    inner_exception = next((arg for arg in exc.args if isinstance(arg, McpError)), None)
    if inner_exception is not None and inner_exception.error.code == CONSENT_ERROR_CODE:
-        return inner_exception.error.message
+        # Parse the error message
+        # The error message is structured with the following format:
+        # "tools/list failed for 1 tool source(s), succeeded for 0 tool source(s) {"errors":[{"name": ..."
+        # where the second part is a JSON string that can be deserialized into an object with the following shape:
+        # ruff: disable[ERA001]
+        # {
+        #   "errors" : [
+        #       {
+        #           "name": "Name of the MCP tool that requires consent",
+        #           "type" : "mcp",
+        #           "error": {
+        #               "code": "CONSENT_REQUIRED",
+        #               "message": consent_url,
+        #           }
+        #       }
+        #   ]
+        # }
+        # ruff: enable[ERA001]
+        try:
+            consent_errors: list[ConsentError] = []
+            error_message_start = inner_exception.error.message.find("{")
+            if error_message_start == -1:
+                logger.warning("Consent error message does not contain JSON: %s", inner_exception.error.message)
+                return None
+            consent_details_json = inner_exception.error.message[error_message_start:]
+            consent_details = json.loads(consent_details_json)
+            if "errors" not in consent_details or not isinstance(consent_details["errors"], list):
+                logger.warning("Consent error message JSON does not contain 'errors' list: %s", consent_details_json)
+                return None
+            for error in consent_details["errors"]:
+                if (
+                    isinstance(error, dict)
+                    and error.get("type") == "mcp"  # type: ignore
+                    and "error" in error
+                    and isinstance(error["error"], dict)
+                    and error["error"].get("code") == "CONSENT_REQUIRED"  # type: ignore
+                    and "message" in error["error"]
+                ):
+                    consent_url = error["error"]["message"]  # type: ignore
+                    if isinstance(consent_url, str):
+                        consent_errors.append(ConsentError(name=error.get("name", "Unknown"), consent_url=consent_url))  # type: ignore
+                    else:
+                        logger.warning("Consent URL in error message is not a valid URL: %s", consent_url)  # type: ignore
+            if consent_errors:
+                return consent_errors
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse consent details JSON: %s", inner_exception.error.message)
    return None


@@ -448,18 +493,19 @@ class ResponsesHostServer(ResponsesAgentServerHost):
        try:
            await self._ensure_agent_ready()
        except AgentFrameworkException as ex:
-            consent_url = consent_url_from_error(ex)
-            if consent_url is None:
+            consent_errors = consent_url_from_error(ex)
+            if consent_errors is None:
                raise
-            logger.warning("OAuth consent required for Foundry MCP gateway.")
-            oauth_item = OAuthConsentRequestOutputItem(
-                id=IdGenerator.new_id("oacr"),
-                consent_link=consent_url,
-                server_label="Foundry Toolbox",
-            )
-            builder = response_event_stream.add_output_item(oauth_item.id)
-            yield builder.emit_added(oauth_item)
-            yield builder.emit_done(oauth_item)
+            for consent_error in consent_errors:
+                logger.warning("Consent URL for tool '%s': %s", consent_error.name, consent_error.consent_url)
+                oauth_item = OAuthConsentRequestOutputItem(
+                    id=IdGenerator.new_id("oacr"),
+                    consent_link=consent_error.consent_url,
+                    server_label=consent_error.name,
+                )
+                builder = response_event_stream.add_output_item(oauth_item.id)
+                yield builder.emit_added(oauth_item)
+                yield builder.emit_done(oauth_item)
            yield response_event_stream.emit_completed()
            return

@@ -39,6 +39,7 @@ from agent_framework_foundry_hosting import ResponsesHostServer
 from agent_framework_foundry_hosting._responses import (
    _AZURE_RESPONSES_MESSAGE_ROLE_TYPE,  # pyright: ignore[reportPrivateUsage]
    CONSENT_ERROR_CODE,
+    ConsentError,
    FileBasedFunctionApprovalStorage,  # pyright: ignore[reportPrivateUsage]
    InMemoryFunctionApprovalStorage,  # pyright: ignore[reportPrivateUsage]
    _item_to_message,  # pyright: ignore[reportPrivateUsage]
@@ -3260,7 +3261,10 @@ class TestCheckpointContextPathValidation:
 # region Agent lifecycle (lazy entry & OAuth consent surfacing)


-def _make_consent_error(url: str = "https://consent.example.com/auth") -> Exception:
+def _make_consent_error(
+    url: str = "https://consent.example.com/auth",
+    name: str = "Foundry Toolbox",
+) -> Exception:
    """Build an exception wrapping a Foundry MCP gateway consent error.

    Mirrors the real-world wrapping produced by ``MCPStreamableHTTPTool.__aenter__``,
@@ -3268,17 +3272,34 @@ def _make_consent_error(url: str = "https://consent.example.com/auth") -> Except
    ``ToolExecutionException`` (an ``AgentFrameworkException`` subclass) with the
    original error attached via ``inner_exception``. ``consent_url_from_error``
    then finds the wrapped ``McpError`` in ``exc.args``.
+
+    The McpError message uses the structured Foundry MCP gateway format:
+    a human-readable prefix followed by a JSON document describing each
+    failed tool source and its consent URL.
    """
    from agent_framework.exceptions import ToolExecutionException

-    inner = McpError(ErrorData(code=CONSENT_ERROR_CODE, message=url))
+    payload = json.dumps({
+        "errors": [
+            {
+                "name": name,
+                "type": "mcp",
+                "error": {
+                    "code": "CONSENT_REQUIRED",
+                    "message": url,
+                },
+            }
+        ]
+    })
+    message = f"tools/list failed for 1 tool source(s), succeeded for 0 tool source(s) {payload}"
+    inner = McpError(ErrorData(code=CONSENT_ERROR_CODE, message=message))
    return ToolExecutionException("MCP consent required", inner_exception=inner)


 class TestConsentUrlFromError:
    def test_returns_consent_url_when_inner_arg_is_consent_mcp_error(self) -> None:
-        exc = _make_consent_error("https://example.com/consent")
-        assert consent_url_from_error(exc) == "https://example.com/consent"
+        exc = _make_consent_error("https://example.com/consent", name="my-tool")
+        assert consent_url_from_error(exc) == [ConsentError(name="my-tool", consent_url="https://example.com/consent")]

    def test_returns_none_when_no_mcp_error_in_args(self) -> None:
        assert consent_url_from_error(Exception("boom")) is None
@@ -3295,6 +3316,13 @@ class TestConsentUrlFromError:
        bare = McpError(ErrorData(code=CONSENT_ERROR_CODE, message="https://x"))
        assert consent_url_from_error(bare) is None

+    def test_returns_none_when_message_has_no_json(self) -> None:
+        from agent_framework.exceptions import ToolExecutionException
+
+        inner = McpError(ErrorData(code=CONSENT_ERROR_CODE, message="no json here"))
+        exc = ToolExecutionException("MCP consent required", inner_exception=inner)
+        assert consent_url_from_error(exc) is None
+

 class TestAgentLifecycle:
    async def test_agent_entered_lazily_on_first_request(self) -> None: