feat: add AWS Bedrock support for auxiliary LLM tasks

agent/auxiliary_client.py

@@ -521,6 +521,95 @@ def __init__(self, adapter: _AsyncCodexCompletionsAdapter):
 class AsyncCodexAuxiliaryClient:
     """Async-compatible wrapper matching AsyncOpenAI.chat.completions.create()."""
 
+
+# ---------------------------------------------------------------------------
+# AWS Bedrock Converse API → OpenAI-compatible wrapper for auxiliary tasks
+# ---------------------------------------------------------------------------
+
+
+class _BedrockCompletionsAdapter:
+    """Translate ``chat.completions.create()`` kwargs to Bedrock Converse API.
+
+    Uses :func:`agent.bedrock_adapter.call_converse` which already handles
+    message/tool format conversion and returns OpenAI-compatible objects.
+    """
+
+    def __init__(self, region: str, default_model: str):
+        self._region = region
+        self._default_model = default_model
+
+    def create(self, **kwargs) -> Any:
+        from agent.bedrock_adapter import call_converse
+
+        model = kwargs.get("model") or self._default_model
+        messages = kwargs.get("messages", [])
+        max_tokens = (kwargs.get("max_tokens")
+                      or kwargs.get("max_completion_tokens")
+                      or 4096)
+        temperature = kwargs.get("temperature")
+        return call_converse(
+            region=self._region,
+            model=model,
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+        )
+
+
+class _BedrockChatShim:
+    """Namespace shim so ``client.chat.completions`` resolves."""
+
+    def __init__(self, adapter: _BedrockCompletionsAdapter):
+        self.completions = adapter
+
+
+class BedrockAuxiliaryClient:
+    """OpenAI-client-compatible wrapper that routes through Bedrock Converse API.
+
+    Consumers can call ``client.chat.completions.create(**kwargs)`` as normal.
+    Also exposes ``.api_key`` and ``.base_url`` for introspection by async
+    wrappers and logging code.
+    """
+
+    def __init__(self, region: str, default_model: str):
+        adapter = _BedrockCompletionsAdapter(region, default_model)
+        self.chat = _BedrockChatShim(adapter)
+        self.api_key = "aws-sdk"
+        self.base_url = f"https://bedrock-runtime.{region}.amazonaws.com"
+
+    def close(self):
+        pass
+
+
+class _AsyncBedrockCompletionsAdapter:
+    """Async version — wraps the sync adapter via :func:`asyncio.to_thread`."""
+
+    def __init__(self, sync_adapter: _BedrockCompletionsAdapter):
+        self._sync = sync_adapter
+
+    async def create(self, **kwargs) -> Any:
+        import asyncio
+        return await asyncio.to_thread(self._sync.create, **kwargs)
+
+
+class _AsyncBedrockChatShim:
+    """Namespace shim for the async variant."""
+
+    def __init__(self, adapter: _AsyncBedrockCompletionsAdapter):
+        self.completions = adapter
+
+
+class AsyncBedrockAuxiliaryClient:
+    """Async-compatible wrapper matching ``AsyncOpenAI.chat.completions.create()``."""
+
+    def __init__(self, sync_wrapper: "BedrockAuxiliaryClient"):
+        sync_adapter = sync_wrapper.chat.completions
+        async_adapter = _AsyncBedrockCompletionsAdapter(sync_adapter)
+        self.chat = _AsyncBedrockChatShim(async_adapter)
+        self.api_key = sync_wrapper.api_key
+        self.base_url = sync_wrapper.base_url
+
+
     def __init__(self, sync_wrapper: "CodexAuxiliaryClient"):
         sync_adapter = sync_wrapper.chat.completions
         async_adapter = _AsyncCodexCompletionsAdapter(sync_adapter)
@@ -1348,6 +1437,8 @@ def _to_async_client(sync_client, model: str):
         return AsyncCodexAuxiliaryClient(sync_client), model
     if isinstance(sync_client, AnthropicAuxiliaryClient):
         return AsyncAnthropicAuxiliaryClient(sync_client), model
+    if isinstance(sync_client, BedrockAuxiliaryClient):
+        return AsyncBedrockAuxiliaryClient(sync_client), model
     try:
         from agent.copilot_acp_client import CopilotACPClient
         if isinstance(sync_client, CopilotACPClient):
@@ -1722,6 +1813,45 @@ def _wrap_if_needed(client_obj, final_model_str: str, base_url_str: str = ""):
                        "directly supported, try 'auto'", provider)
         return None, None
 
+    if pconfig.auth_type == "aws_sdk":
+        # AWS Bedrock — wrap the Converse API in an OpenAI-compatible client.
+        try:
+            from agent.bedrock_adapter import has_aws_credentials
+
+            if not has_aws_credentials():
+                logger.debug(
+                    "resolve_provider_client: bedrock requested but no AWS "
+                    "credentials available"
+                )
+                return None, None
+
+            # Region priority: env var > explicit base_url > hardcoded default.
+            _region = (
+                os.getenv("AWS_DEFAULT_REGION")
+                or os.getenv("AWS_REGION")
+                or "us-east-1"
+            )
+            import re as _re
+            _base = str(explicit_base_url or "")
+            _match = _re.search(r"bedrock-runtime\.([a-z0-9-]+)\.", _base)
+            if _match:
+                _region = _match.group(1)
+
+            _default_model = model or "apac.anthropic.claude-sonnet-4-20250514-v1:0"
+            _client = BedrockAuxiliaryClient(_region, _default_model)
+            logger.debug(
+                "resolve_provider_client: bedrock (%s, %s)", _region, _default_model
+            )
+            if async_mode:
+                return AsyncBedrockAuxiliaryClient(_client), _default_model
+            return _client, _default_model
+        except Exception:
+            logger.warning(
+                "resolve_provider_client: bedrock aws_sdk init failed",
+                exc_info=True,
+            )
+            return None, None
+
     logger.warning("resolve_provider_client: unhandled auth_type %s for %s",
                    pconfig.auth_type, provider)
     return None, None

tests/agent/test_bedrock_auxiliary.py

@@ -0,0 +1,213 @@
+"""Tests for BedrockAuxiliaryClient and aws_sdk auth_type in resolve_provider_client.
+
+Covers:
+  - BedrockAuxiliaryClient wrapper creates correct client structure
+  - _BedrockCompletionsAdapter delegates to bedrock_adapter.call_converse
+  - AsyncBedrockAuxiliaryClient wraps sync adapter via asyncio.to_thread
+  - resolve_provider_client handles aws_sdk auth_type
+  - Region resolution priority: env var > explicit base_url > default
+  - _to_async_client recognizes BedrockAuxiliaryClient
+  - Graceful fallback when AWS credentials are unavailable
+"""
+
+import asyncio
+import os
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from agent.auxiliary_client import (
+    BedrockAuxiliaryClient,
+    AsyncBedrockAuxiliaryClient,
+    _BedrockCompletionsAdapter,
+    resolve_provider_client,
+    _to_async_client,
+)
+
+
+_FAKE_RESPONSE = SimpleNamespace(
+    choices=[SimpleNamespace(message=SimpleNamespace(content="OK"))],
+    model="test-model",
+    usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
+)
+
+
+# ---------------------------------------------------------------------------
+# BedrockAuxiliaryClient structure
+# ---------------------------------------------------------------------------
+
+
+class TestBedrockAuxiliaryClient:
+    """BedrockAuxiliaryClient exposes the standard OpenAI-compatible interface."""
+
+    def test_has_chat_completions_interface(self):
+        client = BedrockAuxiliaryClient("us-east-1", "some-model")
+        assert hasattr(client, "chat")
+        assert hasattr(client.chat, "completions")
+        assert hasattr(client.chat.completions, "create")
+
+    def test_exposes_api_key_and_base_url(self):
+        client = BedrockAuxiliaryClient("ap-southeast-1", "model-id")
+        assert client.api_key == "aws-sdk"
+        assert "ap-southeast-1" in client.base_url
+
+    def test_close_is_noop(self):
+        client = BedrockAuxiliaryClient("us-east-1", "m")
+        client.close()  # should not raise
+
+
+# ---------------------------------------------------------------------------
+# _BedrockCompletionsAdapter
+# ---------------------------------------------------------------------------
+
+
+class TestBedrockCompletionsAdapter:
+    """Adapter translates chat.completions.create() kwargs to Bedrock Converse."""
+
+    @patch("agent.bedrock_adapter.call_converse")
+    def test_delegates_to_call_converse(self, mock_converse):
+        mock_converse.return_value = _FAKE_RESPONSE
+        adapter = _BedrockCompletionsAdapter("ap-southeast-1", "default-model")
+        result = adapter.create(
+            model="explicit-model",
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=100,
+            temperature=0.5,
+        )
+        mock_converse.assert_called_once_with(
+            region="ap-southeast-1",
+            model="explicit-model",
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=100,
+            temperature=0.5,
+        )
+        assert result.choices[0].message.content == "OK"
+
+    @patch("agent.bedrock_adapter.call_converse")
+    def test_uses_default_model_when_none(self, mock_converse):
+        mock_converse.return_value = _FAKE_RESPONSE
+        adapter = _BedrockCompletionsAdapter("us-east-1", "my-default")
+        adapter.create(messages=[{"role": "user", "content": "hi"}])
+        assert mock_converse.call_args.kwargs["model"] == "my-default"
+
+    @patch("agent.bedrock_adapter.call_converse")
+    def test_max_completion_tokens_fallback(self, mock_converse):
+        mock_converse.return_value = _FAKE_RESPONSE
+        adapter = _BedrockCompletionsAdapter("us-east-1", "m")
+        adapter.create(messages=[], max_completion_tokens=2048)
+        assert mock_converse.call_args.kwargs["max_tokens"] == 2048
+
+
+# ---------------------------------------------------------------------------
+# AsyncBedrockAuxiliaryClient
+# ---------------------------------------------------------------------------
+
+
+class TestAsyncBedrockAuxiliaryClient:
+    """Async wrapper delegates to sync adapter via asyncio.to_thread."""
+
+    def test_has_async_interface(self):
+        sync = BedrockAuxiliaryClient("us-east-1", "m")
+        async_client = AsyncBedrockAuxiliaryClient(sync)
+        assert hasattr(async_client, "chat")
+        assert hasattr(async_client.chat, "completions")
+        assert hasattr(async_client.chat.completions, "create")
+
+    def test_preserves_api_key_and_base_url(self):
+        sync = BedrockAuxiliaryClient("eu-west-1", "m")
+        async_client = AsyncBedrockAuxiliaryClient(sync)
+        assert async_client.api_key == "aws-sdk"
+        assert "eu-west-1" in async_client.base_url
+
+    @patch("agent.bedrock_adapter.call_converse", return_value=_FAKE_RESPONSE)
+    def test_async_create_calls_sync(self, _mock):
+        sync = BedrockAuxiliaryClient("us-east-1", "m")
+        async_client = AsyncBedrockAuxiliaryClient(sync)
+        result = asyncio.get_event_loop().run_until_complete(
+            async_client.chat.completions.create(
+                model="m", messages=[{"role": "user", "content": "hi"}]
+            )
+        )
+        assert result.choices[0].message.content == "OK"
+
+
+# ---------------------------------------------------------------------------
+# _to_async_client integration
+# ---------------------------------------------------------------------------
+
+
+class TestToAsyncClientBedrock:
+    """_to_async_client recognizes BedrockAuxiliaryClient."""
+
+    def test_returns_async_bedrock_wrapper(self):
+        sync = BedrockAuxiliaryClient("us-east-1", "m")
+        async_client, model = _to_async_client(sync, "my-model")
+        assert isinstance(async_client, AsyncBedrockAuxiliaryClient)
+        assert model == "my-model"
+
+
+# ---------------------------------------------------------------------------
+# resolve_provider_client with aws_sdk auth_type
+# ---------------------------------------------------------------------------
+
+
+class TestResolveProviderClientBedrock:
+    """resolve_provider_client handles the bedrock provider (aws_sdk auth)."""
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_returns_bedrock_client_when_creds_available(self, _mock, monkeypatch):
+        monkeypatch.setenv("AWS_DEFAULT_REGION", "ap-southeast-1")
+        client, model = resolve_provider_client(
+            "bedrock", "apac.anthropic.claude-sonnet-4-20250514-v1:0"
+        )
+        assert isinstance(client, BedrockAuxiliaryClient)
+        assert "ap-southeast-1" in client.base_url
+        assert model == "apac.anthropic.claude-sonnet-4-20250514-v1:0"
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=False)
+    def test_returns_none_when_no_creds(self, _mock):
+        client, model = resolve_provider_client("bedrock", "some-model")
+        assert client is None
+        assert model is None
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_region_from_env_var(self, _mock, monkeypatch):
+        monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-central-1")
+        client, _ = resolve_provider_client("bedrock", "m")
+        assert "eu-central-1" in client.base_url
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_region_from_explicit_base_url(self, _mock, monkeypatch):
+        monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False)
+        monkeypatch.delenv("AWS_REGION", raising=False)
+        client, _ = resolve_provider_client(
+            "bedrock", "m",
+            explicit_base_url="https://bedrock-runtime.us-west-2.amazonaws.com",
+        )
+        assert "us-west-2" in client.base_url
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_region_defaults_to_us_east_1(self, _mock, monkeypatch):
+        monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False)
+        monkeypatch.delenv("AWS_REGION", raising=False)
+        client, _ = resolve_provider_client("bedrock", "m")
+        assert "us-east-1" in client.base_url
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_default_model_when_none_provided(self, _mock, monkeypatch):
+        monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
+        _, model = resolve_provider_client("bedrock", None)
+        assert model is not None
+        assert "claude" in model.lower() or "anthropic" in model.lower()
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", return_value=True)
+    def test_async_mode_returns_async_client(self, _mock, monkeypatch):
+        monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
+        client, model = resolve_provider_client("bedrock", "m", async_mode=True)
+        assert isinstance(client, AsyncBedrockAuxiliaryClient)
+
+    @patch("agent.bedrock_adapter.has_aws_credentials", side_effect=ImportError("no boto3"))
+    def test_graceful_on_import_error(self, _mock):
+        client, model = resolve_provider_client("bedrock", "m")
+        assert client is None