Merge branch 'main' into dmontagu/fix-evals-token-double-counting

DouweM · web-flow · commit 6066bd98a67c · 2025-11-25T10:02:40.000-06:00
diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py
@@ -1040,15 +1040,16 @@ async def client_streams(
             sse_read_timeout=self.read_timeout,
         )
 
-        if self.http_client is not None:  # pragma: no cover
-
-            def httpx_client_factory(
+        if self.http_client is not None:
+            # TODO: Clean up once https://github.com/modelcontextprotocol/python-sdk/pull/1177 lands.
+            @asynccontextmanager
+            async def httpx_client_factory(
                 headers: dict[str, str] | None = None,
                 timeout: httpx.Timeout | None = None,
                 auth: httpx.Auth | None = None,
-            ) -> httpx.AsyncClient:
+            ) -> AsyncIterator[httpx.AsyncClient]:
                 assert self.http_client is not None
-                return self.http_client
+                yield self.http_client
 
             async with transport_client_partial(httpx_client_factory=httpx_client_factory) as (
                 read_stream,
diff --git a/pydantic_ai_slim/pydantic_ai/models/openai.py b/pydantic_ai_slim/pydantic_ai/models/openai.py
@@ -518,7 +518,7 @@ async def _completions_create(
             extra_headers = model_settings.get('extra_headers', {})
             extra_headers.setdefault('User-Agent', get_user_agent())
             return await self.client.chat.completions.create(
-                model=self._model_name,
+                model=self.model_name,
                 messages=openai_messages,
                 parallel_tool_calls=model_settings.get('parallel_tool_calls', OMIT),
                 tools=tools or OMIT,
diff --git a/pydantic_ai_slim/pydantic_ai/models/openrouter.py b/pydantic_ai_slim/pydantic_ai/models/openrouter.py
@@ -244,7 +244,11 @@ class _BaseReasoningDetail(BaseModel, frozen=True):
     """Common fields shared across all reasoning detail types."""
 
     id: str | None = None
-    format: Literal['unknown', 'openai-responses-v1', 'anthropic-claude-v1', 'xai-responses-v1'] | None
+    format: (
+        Literal['unknown', 'openai-responses-v1', 'anthropic-claude-v1', 'xai-responses-v1', 'google-gemini-v1']
+        | str
+        | None
+    )
     index: int | None
     type: Literal['reasoning.text', 'reasoning.summary', 'reasoning.encrypted']
 
@@ -610,13 +614,21 @@ def _map_thinking_delta(self, choice: chat_completion_chunk.Choice) -> Iterable[
         assert isinstance(choice, _OpenRouterChunkChoice)
 
         if reasoning_details := choice.delta.reasoning_details:
-            for detail in reasoning_details:
+            for i, detail in enumerate(reasoning_details):
                 thinking_part = _from_reasoning_detail(detail)
+                # Use unique vendor_part_id for each reasoning detail type to prevent
+                # different detail types (e.g., reasoning.text, reasoning.encrypted)
+                # from being incorrectly merged into a single ThinkingPart.
+                # This is required for Gemini 3 Pro which returns multiple reasoning
+                # detail types that must be preserved separately for thought_signature handling.
+                vendor_id = f'reasoning_detail_{detail.type}_{i}'
                 yield self._parts_manager.handle_thinking_delta(
-                    vendor_part_id='reasoning_detail',
+                    vendor_part_id=vendor_id,
                     id=thinking_part.id,
                     content=thinking_part.content,
+                    signature=thinking_part.signature,
                     provider_name=self._provider_name,
+                    provider_details=thinking_part.provider_details,
                 )
         else:
             return super()._map_thinking_delta(choice)
diff --git a/tests/cassettes/test_mcp/test_custom_http_client_not_closed.yaml b/tests/cassettes/test_mcp/test_custom_http_client_not_closed.yaml
@@ -0,0 +1,270 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json, text/event-stream
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '165'
+      content-type:
+      - application/json
+      host:
+      - mcp.deepwiki.com
+    method: POST
+    parsed_body:
+      id: 0
+      jsonrpc: '2.0'
+      method: initialize
+      params:
+        capabilities:
+          sampling: {}
+        clientInfo:
+          name: mcp
+          version: 0.1.0
+        protocolVersion: '2025-06-18'
+    uri: https://mcp.deepwiki.com/mcp
+  response:
+    body:
+      string: |+
+        event: message
+        data: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"tools":{"listChanged":true}},"serverInfo":{"name":"DeepWiki","version":"0.0.1"}}}
+
+        event: ping
+        data: ping
+
+    headers:
+      access-control-allow-headers:
+      - Content-Type, mcp-session-id, mcp-protocol-version
+      access-control-allow-methods:
+      - GET, POST, OPTIONS
+      access-control-allow-origin:
+      - '*'
+      access-control-expose-headers:
+      - mcp-session-id
+      access-control-max-age:
+      - '86400'
+      alt-svc:
+      - h3=":443"; ma=86400
+      cache-control:
+      - no-cache
+      connection:
+      - keep-alive
+      content-type:
+      - text/event-stream
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=rMn85WgDL0%2B8DcP3GhSk08gOMLOCdd9ylzYlWV%2BPazzXbOlw7hBD5mlea3Yzxs%2BdtHeX7t1OSADQzo5YxQq55wD0eI0FTdoN0UBf2SM%2F7zsjJieuzN%2FodKifnGLW%2FgU5j%2F%2F%2FHA%3D%3D"}]}'
+      transfer-encoding:
+      - chunked
+      vary:
+      - accept-encoding
+    status:
+      code: 200
+      message: OK
+- request:
+    body: ''
+    headers:
+      accept:
+      - application/json, text/event-stream, text/event-stream
+      accept-encoding:
+      - gzip, deflate
+      cache-control:
+      - no-store
+      connection:
+      - keep-alive
+      content-type:
+      - application/json
+      host:
+      - mcp.deepwiki.com
+      mcp-protocol-version:
+      - '2025-03-26'
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+    method: GET
+    uri: https://mcp.deepwiki.com/mcp
+  response:
+    body:
+      string: '{"jsonrpc":"2.0","error":{"code":-32000,"message":"Method not allowed"},"id":null}'
+    headers:
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '82'
+      content-type:
+      - text/plain;charset=UTF-8
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=PJNTRlKrI8lL0aiZCc1uJTiDVehFnTLMMt8VdPxyqnbI%2F0JHNcJ7zP3QWIemrZXXzvU8Vd%2Bei6J8F5UuykGqFqdTH%2BXrjvH5bM9Is8%2F%2BWmpiJxekG6tEGDnwy9qFv8xegxo8jQ%3D%3D"}]}'
+      vary:
+      - accept-encoding
+    status:
+      code: 405
+      message: Method Not Allowed
+- request:
+    headers:
+      accept:
+      - application/json, text/event-stream
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '54'
+      content-type:
+      - application/json
+      host:
+      - mcp.deepwiki.com
+      mcp-protocol-version:
+      - '2025-03-26'
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+    method: POST
+    parsed_body:
+      jsonrpc: '2.0'
+      method: notifications/initialized
+    uri: https://mcp.deepwiki.com/mcp
+  response:
+    body:
+      string: ''
+    headers:
+      access-control-allow-headers:
+      - Content-Type, mcp-session-id, mcp-protocol-version
+      access-control-allow-methods:
+      - GET, POST, OPTIONS
+      access-control-allow-origin:
+      - '*'
+      access-control-expose-headers:
+      - mcp-session-id
+      access-control-max-age:
+      - '86400'
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '0'
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=YOAoKhRcD%2FYSYhguoqnWMHg%2FDamN%2BTbORh3ouJR1UoKYxtW%2BOW%2BAHCtlYMucysWlEeSIhTC5vm5N999APwF5Rwu1Z8liG5LqCo4ixNr69BVK1ssPS4i24x%2FyVLYXc0HTo9dkIg%3D%3D"}]}'
+      vary:
+      - accept-encoding
+    status:
+      code: 202
+      message: Accepted
+- request:
+    headers:
+      accept:
+      - application/json, text/event-stream
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '46'
+      content-type:
+      - application/json
+      host:
+      - mcp.deepwiki.com
+      mcp-protocol-version:
+      - '2025-03-26'
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+    method: POST
+    parsed_body:
+      id: 1
+      jsonrpc: '2.0'
+      method: tools/list
+    uri: https://mcp.deepwiki.com/mcp
+  response:
+    body:
+      string: |+
+        event: message
+        data: {"jsonrpc":"2.0","id":1,"result":{"tools":[{"name":"read_wiki_structure","description":"Get a list of documentation topics for a GitHub repository","inputSchema":{"type":"object","properties":{"repoName":{"type":"string","description":"GitHub repository: owner/repo (e.g. \"facebook/react\")"}},"required":["repoName"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}},{"name":"read_wiki_contents","description":"View documentation about a GitHub repository","inputSchema":{"type":"object","properties":{"repoName":{"type":"string","description":"GitHub repository: owner/repo (e.g. \"facebook/react\")"}},"required":["repoName"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}},{"name":"ask_question","description":"Ask any question about a GitHub repository","inputSchema":{"type":"object","properties":{"repoName":{"type":"string","description":"GitHub repository: owner/repo (e.g. \"facebook/react\")"},"question":{"type":"string","description":"The question to ask about the repository"}},"required":["repoName","question"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}}]}}
+
+        event: ping
+        data: ping
+
+    headers:
+      access-control-allow-headers:
+      - Content-Type, mcp-session-id, mcp-protocol-version
+      access-control-allow-methods:
+      - GET, POST, OPTIONS
+      access-control-allow-origin:
+      - '*'
+      access-control-expose-headers:
+      - mcp-session-id
+      access-control-max-age:
+      - '86400'
+      alt-svc:
+      - h3=":443"; ma=86400
+      cache-control:
+      - no-cache
+      connection:
+      - keep-alive
+      content-type:
+      - text/event-stream
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=rTAkw2F56hOqBG74w2PLRaPXDJf5XzFASzgxvvmTIrJZmebUbzATwXgvs3W6GMYCXnwMdcOUJcN05oKIJ6hq8XoPuu%2F97yD66drIfQOhQjT0HG2W4Kd1btBaRf0stfwOdus64w%3D%3D"}]}'
+      transfer-encoding:
+      - chunked
+      vary:
+      - accept-encoding
+    status:
+      code: 200
+      message: OK
+- request:
+    body: ''
+    headers:
+      accept:
+      - application/json, text/event-stream
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-type:
+      - application/json
+      host:
+      - mcp.deepwiki.com
+      mcp-protocol-version:
+      - '2025-03-26'
+      mcp-session-id:
+      - 79634f20860e529d0253e303afbce2571a0543c155a087250648960da3c67059
+    method: DELETE
+    uri: https://mcp.deepwiki.com/mcp
+  response:
+    body:
+      string: '{"jsonrpc":"2.0","error":{"code":-32000,"message":"Method not allowed"},"id":null}'
+    headers:
+      alt-svc:
+      - h3=":443"; ma=86400
+      connection:
+      - keep-alive
+      content-length:
+      - '82'
+      content-type:
+      - text/plain;charset=UTF-8
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=0AMi1Wp%2BN%2FfpGPKL%2BRDf5fdSr0BLZQRV0fR7plR%2FnrWUjtsC6slmDj71mtsoXeq%2B3YtYdqo0zDvuI59BSy9HcqeWAIau1X6Ng3C5K%2BzFGL1hIr4135Ivh%2BaiFo17vSsis7rEDQ%3D%3D"}]}'
+      vary:
+      - accept-encoding
+    status:
+      code: 405
+      message: Method Not Allowed
+version: 1
+...
diff --git a/tests/models/test_openrouter.py b/tests/models/test_openrouter.py
@@ -112,23 +112,23 @@ async def test_openrouter_stream_with_reasoning(allow_model_requests: None, open
 
         thinking_event_start = chunks[0]
         assert isinstance(thinking_event_start, PartStartEvent)
-        assert thinking_event_start.part == snapshot(
-            ThinkingPart(
-                content='',
-                id='rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f',
-                provider_name='openrouter',
-            )
-        )
+        thinking_part = thinking_event_start.part
+        assert isinstance(thinking_part, ThinkingPart)
+        assert thinking_part.id == 'rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f'
+        assert thinking_part.content == ''
+        assert thinking_part.provider_name == 'openrouter'
+        # After fix: signature and provider_details are now properly preserved
+        assert thinking_part.signature is not None
+        assert thinking_part.provider_details is not None
+        assert thinking_part.provider_details['type'] == 'reasoning.encrypted'
+        assert thinking_part.provider_details['format'] == 'openai-responses-v1'
 
         thinking_event_end = chunks[1]
         assert isinstance(thinking_event_end, PartEndEvent)
-        assert thinking_event_end.part == snapshot(
-            ThinkingPart(
-                content='',
-                id='rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f',
-                provider_name='openrouter',
-            )
-        )
+        thinking_part_end = thinking_event_end.part
+        assert isinstance(thinking_part_end, ThinkingPart)
+        assert thinking_part_end.id == 'rs_0aa4f2c435e6d1dc0169082486816c8193a029b5fc4ef1764f'
+        assert thinking_part_end.signature is not None
 
 
 async def test_openrouter_stream_error(allow_model_requests: None, openrouter_api_key: str) -> None:
diff --git a/tests/test_mcp.py b/tests/test_mcp.py
