From 2b3c401848d3157d94cba399428a24a190c79771 Mon Sep 17 00:00:00 2001 From: Giles Odigwe <79032838+giles17@users.noreply.github.com> Date: Tue, 3 Mar 2026 15:02:03 -0800 Subject: [PATCH] Python: Fix: Parse oauth_consent_request events in Azure AI client (#4197) * Fix: Parse oauth_consent_request events in Azure AI client (#3950) When Azure AI Agent Service returns an oauth_consent_request output item for OAuth-protected MCP tools, the base OpenAI responses parser drops it (hits case _ default branch). This causes agent runs to complete silently with zero content. Changes: - Add oauth_consent_request ContentType and Content.from_oauth_consent_request() factory with consent_link field and user_input_request=True - Override _parse_response_from_openai and _parse_chunk_from_openai in RawAzureAIClient to intercept Azure-specific oauth_consent_request items - Add _emit_oauth_consent helper in AG-UI to emit CustomEvent for frontends - Add tests proving base parser drops the event and Azure AI override catches it Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * addressed comment * addressed comments * addressed comments --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agent_framework_ag_ui/_run_common.py | 11 +++ python/packages/ag-ui/tests/ag_ui/test_run.py | 24 +++++ .../agent_framework_azure_ai/_client.py | 62 ++++++++++++ .../azure-ai/tests/test_azure_ai_client.py | 99 +++++++++++++++++++ .../packages/core/agent_framework/_types.py | 36 +++++++ python/packages/core/tests/core/test_types.py | 27 +++++ 6 files changed, 259 insertions(+) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 997c375ed1..d8cf236add 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -372,6 +372,15 @@ def _emit_usage(content: Content) -> list[BaseEvent]: return [CustomEvent(name="usage", value=usage_details)] +def _emit_oauth_consent(content: Content) -> list[BaseEvent]: + """Emit an OAuth consent request as a custom event so frontends can render a consent link.""" + return ( + [CustomEvent(name="oauth_consent_request", value={"consent_link": content.consent_link})] + if content.consent_link + else [] + ) + + def _emit_content( content: Any, flow: FlowState, @@ -391,5 +400,7 @@ def _emit_content( return _emit_approval_request(content, flow, predictive_handler, require_confirmation) if content_type == "usage": return _emit_usage(content) + if content_type == "oauth_consent_request": + return _emit_oauth_consent(content) logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type) return [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 73c9648c02..e0771e1b7e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -4,6 +4,7 @@ import pytest from ag_ui.core import ( + CustomEvent, TextMessageEndEvent, TextMessageStartEvent, ToolCallArgsEvent, @@ -871,3 +872,26 @@ class TestTextMessageEventBalancing: assert len(start_events) == 2 assert len(end_events) == 2 + + +def test_emit_oauth_consent_request(): + """Test that oauth_consent_request content emits a CustomEvent.""" + content = Content.from_oauth_consent_request( + consent_link="https://login.microsoftonline.com/consent", + ) + flow = FlowState() + events = _emit_content(content, flow) + + assert len(events) == 1 + assert isinstance(events[0], CustomEvent) + assert events[0].name == "oauth_consent_request" + assert events[0].value == {"consent_link": "https://login.microsoftonline.com/consent"} + + +def test_emit_oauth_consent_request_no_link(): + """Test that oauth_consent_request without a consent_link emits no events.""" + content = Content("oauth_consent_request") + flow = FlowState() + events = _emit_content(content, flow) + + assert len(events) == 0 diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py index 37e5e1fbcc..265dd6c2a6 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py @@ -588,6 +588,68 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ """Get the current conversation ID from chat options or kwargs.""" return options.get("conversation_id") or kwargs.get("conversation_id") or self.conversation_id + @override + def _parse_response_from_openai( + self, + response: Any, + options: dict[str, Any], + ) -> ChatResponse: + """Parse an Azure AI Responses API response, handling Azure-specific output item types.""" + result = super()._parse_response_from_openai(response, options) + + if result.messages: + for item in response.output: + if item.type == "oauth_consent_request": + consent_link = item.consent_link + if consent_link and not consent_link.startswith("https://"): + logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", item) + consent_link = "" + if consent_link: + result.messages[0].contents.append( + Content.from_oauth_consent_request( + consent_link=consent_link, + raw_representation=item, + ) + ) + else: + logger.warning("Received oauth_consent_request output without consent_link: %s", item) + + return result + + @override + def _parse_chunk_from_openai( + self, + event: Any, + options: dict[str, Any], + function_call_ids: dict[int, tuple[str, str]], + ) -> ChatResponseUpdate: + """Parse an Azure AI streaming event, handling Azure-specific event types.""" + # Intercept output_item.added events for Azure-specific item types + if event.type == "response.output_item.added" and event.item.type == "oauth_consent_request": + event_item = event.item + consent_link = event_item.consent_link + if consent_link and not consent_link.startswith("https://"): + logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", event_item) + consent_link = "" + contents: list[Content] = [] + if consent_link: + contents.append( + Content.from_oauth_consent_request( + consent_link=consent_link, + raw_representation=event_item, + ) + ) + else: + logger.warning("Received oauth_consent_request output without consent_link: %s", event_item) + return ChatResponseUpdate( + contents=contents, + role="assistant", + model_id=self.model_id, + raw_representation=event, + ) + + return super()._parse_chunk_from_openai(event, options, function_call_ids) + def _prepare_messages_for_azure_ai(self, messages: Sequence[Message]) -> tuple[list[Message], str | None]: """Prepare input from messages and convert system/developer messages to instructions.""" result: list[Message] = [] diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py index 5ddeae6783..a62cd5628f 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_client.py @@ -2174,4 +2174,103 @@ def test_build_url_citation_content_with_dict(mock_project_client: MagicMock) -> assert "get_url" not in ann.get("additional_properties", {}) +# region OAuth Consent + + +def test_parse_chunk_with_oauth_consent_request(mock_project_client: MagicMock) -> None: + """Test that a streaming oauth_consent_request output item is parsed into oauth_consent_request content. + + This reproduces the bug from issue #3950 where the event was logged as "Unparsed event" + and silently discarded, causing the agent run to complete with zero content. + """ + client = AzureAIClient(project_client=mock_project_client, agent_name="test") + chat_options: dict[str, Any] = {} + function_call_ids: dict[int, tuple[str, str]] = {} + + mock_item = MagicMock() + mock_item.type = "oauth_consent_request" + mock_item.consent_link = "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123" + + mock_event = MagicMock() + mock_event.type = "response.output_item.added" + mock_event.item = mock_item + mock_event.output_index = 0 + + update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids) + + assert len(update.contents) == 1 + consent_content = update.contents[0] + assert consent_content.type == "oauth_consent_request" + assert consent_content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123" + assert consent_content.user_input_request is True + + +def test_parse_response_with_oauth_consent_output_item(mock_project_client: MagicMock) -> None: + """Test that a non-streaming oauth_consent_request output item is parsed correctly.""" + client = AzureAIClient(project_client=mock_project_client, agent_name="test") + + mock_item = MagicMock() + mock_item.type = "oauth_consent_request" + mock_item.consent_link = "https://login.microsoftonline.com/consent?code=abc" + + mock_response = MagicMock() + mock_response.output = [mock_item] + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.id = "resp-oauth-1" + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.usage = None + mock_response.status = "completed" + + response = client._parse_response_from_openai(mock_response, {}) + + assert len(response.messages) > 0 + consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"] + assert len(consent_contents) == 1 + assert consent_contents[0].consent_link == "https://login.microsoftonline.com/consent?code=abc" + + +def test_parse_chunk_oauth_consent_no_link(mock_project_client: MagicMock) -> None: + """Test that a streaming oauth_consent_request with no consent_link produces empty contents.""" + client = AzureAIClient(project_client=mock_project_client, agent_name="test") + + mock_item = MagicMock() + mock_item.type = "oauth_consent_request" + mock_item.consent_link = "" + + mock_event = MagicMock() + mock_event.type = "response.output_item.added" + mock_event.item = mock_item + mock_event.output_index = 0 + + update = client._parse_chunk_from_openai(mock_event, {}, {}) + + assert not any(c.type == "oauth_consent_request" for c in update.contents) + + +def test_parse_response_oauth_consent_no_link(mock_project_client: MagicMock) -> None: + """Test that a non-streaming oauth_consent_request with no consent_link appends no content.""" + client = AzureAIClient(project_client=mock_project_client, agent_name="test") + + mock_item = MagicMock() + mock_item.type = "oauth_consent_request" + mock_item.consent_link = None + + mock_response = MagicMock() + mock_response.output = [mock_item] + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.id = "resp-oauth-2" + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.usage = None + mock_response.status = "completed" + + response = client._parse_response_from_openai(mock_response, {}) + + consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"] + assert len(consent_contents) == 0 + + # endregion diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index beed97834c..ee0e813d27 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -345,6 +345,7 @@ ContentType = Literal[ "shell_command_output", "function_approval_request", "function_approval_response", + "oauth_consent_request", ] @@ -498,6 +499,8 @@ class Content: function_call: Content | None = None, user_input_request: bool | None = None, approved: bool | None = None, + # OAuth consent fields + consent_link: str | None = None, # Common fields annotations: Sequence[Annotation] | None = None, additional_properties: MutableMapping[str, Any] | None = None, @@ -546,6 +549,7 @@ class Content: self.function_call = function_call self.user_input_request = user_input_request self.approved = approved + self.consent_link = consent_link @classmethod def from_text( @@ -1122,6 +1126,37 @@ class Content: raw_representation=raw_representation, ) + @classmethod + def from_oauth_consent_request( + cls: type[ContentT], + consent_link: str, + *, + annotations: Sequence[Annotation] | None = None, + additional_properties: MutableMapping[str, Any] | None = None, + raw_representation: Any = None, + ) -> ContentT: + """Create OAuth consent request content. + + Args: + consent_link: The URL the user must visit to complete OAuth consent. + + Keyword Args: + annotations: Optional annotations. + additional_properties: Optional additional properties. + raw_representation: Optional raw representation from the provider. + + Returns: + A new Content instance with type ``oauth_consent_request``. + """ + return cls( + "oauth_consent_request", + consent_link=consent_link, + user_input_request=True, + annotations=annotations, + additional_properties=additional_properties, + raw_representation=raw_representation, + ) + def to_function_approval_response( self, approved: bool, @@ -1176,6 +1211,7 @@ class Content: "user_input_request", "approved", "id", + "consent_link", "additional_properties", ) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index c858ff1e3f..bcf3a6891b 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -3424,3 +3424,30 @@ class TestResponseStreamEdgeCases: # endregion + + +# region OAuth Consent Content + + +def test_oauth_consent_request_creation(): + """Test Content.from_oauth_consent_request creates the correct content.""" + content = Content.from_oauth_consent_request( + consent_link="https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc", + ) + assert content.type == "oauth_consent_request" + assert content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc" + assert content.user_input_request is True + + +def test_oauth_consent_request_serialization_roundtrip(): + """Test that oauth_consent_request content serializes and includes consent_link.""" + content = Content.from_oauth_consent_request( + consent_link="https://login.microsoftonline.com/consent", + ) + d = content.to_dict() + assert d["type"] == "oauth_consent_request" + assert d["consent_link"] == "https://login.microsoftonline.com/consent" + assert d["user_input_request"] is True + + +# endregion