fix(bedrock): restore locking for extra_headers isolation

Author: kkauy

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -673,7 +673,8 @@ async def _call_bedrock(
         params: ConverseRequestTypeDef,
         stream: bool,
     ) -> ConverseResponseTypeDef | ConverseStreamResponseTypeDef:
-        return await self._call_bedrock_unlocked(params=params, stream=stream)
+        async with self._extra_headers_lock:
+            return await self._call_bedrock_unlocked(params=params, stream=stream)
 
     async def _call_with_extra_headers(
         self,

tests/models/test_bedrock.py

@@ -4,6 +4,7 @@
 from datetime import date, datetime, timezone
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 from typing_extensions import TypedDict
@@ -124,6 +125,88 @@ def _bedrock_model_with_client_error(error: ClientError) -> BedrockConverseModel
         provider=_StubBedrockProvider(_StubBedrockClient(error)),
     )
 
+
+def test_register_extra_headers_injects_headers():
+    client = MagicMock()
+    client.meta.events = MagicMock()
+
+    headers = {'X-Test': '123'}
+
+    handler, events = BedrockConverseModel._register_extra_headers(client, headers, stream=False)
+
+    request = MagicMock()
+    request.headers = {}
+
+    handler(request)
+
+    assert request.headers['X-Test'] == '123'
+    client.meta.events.register.assert_called_once()
+    assert events == ['before-send.bedrock-runtime.Converse']
+
+
+async def test_call_with_extra_headers_registers_and_unregisters():
+    model = BedrockConverseModel.__new__(BedrockConverseModel)
+
+    model.client = MagicMock()
+
+    import anyio
+
+    model._extra_headers_lock = anyio.Lock()
+
+    model._register_extra_headers = MagicMock(return_value=('handler', ['event']))
+    model._unregister_extra_headers = MagicMock()
+
+    model._call_bedrock_unlocked = AsyncMock(return_value='ok')
+
+    result = await model._call_with_extra_headers(
+        params={'foo': 'bar'},
+        extra_headers={'X-Test': '123'},
+        stream=False,
+    )
+
+    assert result == 'ok'
+    model._register_extra_headers.assert_called_once()
+    model._unregister_extra_headers.assert_called_once()
+
+
+async def test_call_with_extra_headers_unregisters_on_exception():
+    model = BedrockConverseModel.__new__(BedrockConverseModel)
+
+    model.client = MagicMock()
+
+    import anyio
+
+    model._extra_headers_lock = anyio.Lock()
+
+    model._register_extra_headers = MagicMock(return_value=('handler', ['event']))
+    model._unregister_extra_headers = MagicMock()
+
+    async def raise_error(*args, **kwargs):
+        raise RuntimeError('fail')
+
+    model._call_bedrock_unlocked = raise_error
+
+    with pytest.raises(RuntimeError):
+        await model._call_with_extra_headers(
+            params={'foo': 'bar'},
+            extra_headers={'X-Test': '123'},
+            stream=False,
+        )
+
+    model._unregister_extra_headers.assert_called_once()
+
+
+def test_register_extra_headers_stream_event_name():
+    client = MagicMock()
+    client.meta.events = MagicMock()
+
+    headers = {'X-Test': '123'}
+
+    handler, events = BedrockConverseModel._register_extra_headers(client, headers, stream=True)
+
+    assert events == ['before-send.bedrock-runtime.ConverseStream']
+
+
 async def test_bedrock_model(allow_model_requests: None, bedrock_provider: BedrockProvider):
     model = BedrockConverseModel('us.amazon.nova-micro-v1:0', provider=bedrock_provider)
     assert model.base_url == 'https://bedrock-runtime.us-east-1.amazonaws.com'
@@ -168,6 +251,7 @@ async def test_bedrock_model(allow_model_requests: None, bedrock_provider: Bedro
         ]
     )
 
+
 @pytest.mark.vcr()
 async def test_bedrock_model_usage_limit_exceeded(
     allow_model_requests: None,