From 7b26ab95162217c454b789f83686184354e70eb2 Mon Sep 17 00:00:00 2001 From: Adam Moussa Date: Mon, 29 Jun 2026 15:34:06 -0400 Subject: [PATCH] =?UTF-8?q?fix(bedrock):=20security-review=20NITs=20?= =?UTF-8?q?=E2=80=94=20region=20resolution,=20error=20sanitization,=20reas?= =?UTF-8?q?oning-block=20strip?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From /sh-security-review (all confirmed-low): - model.py: resolve region from AWS_REGION OR AWS_DEFAULT_REGION (matches validate_local_dev_llm_config) so the validated region is the one actually used. - model_fallback.py: sanitize Bedrock AccessDenied/ResourceNotFound errors to the error code only, so the role ARN + account id in the raw botocore message never reach logs or the user channel (CWE-209). - sanitize_thinking_blocks.py: also strip empty Bedrock reasoning_content blocks (Converse emits reasoning_content, not thinking) so the middleware is not a no-op on Bedrock; + unit tests. (Empty blocks replay fine today; defensive.) --- agent/middleware/model_fallback.py | 12 ++++++++ agent/middleware/sanitize_thinking_blocks.py | 31 ++++++++++++++------ agent/utils/model.py | 7 ++++- tests/test_sanitize_thinking_blocks.py | 27 +++++++++++++++++ 4 files changed, 67 insertions(+), 10 deletions(-) diff --git a/agent/middleware/model_fallback.py b/agent/middleware/model_fallback.py index 5ad08458..41a2cb5c 100644 --- a/agent/middleware/model_fallback.py +++ b/agent/middleware/model_fallback.py @@ -101,6 +101,18 @@ def _provider_access_error_message(exc: BaseException) -> str | None: "Choose a different model or update the workspace's OpenAI access and retry." ) + # Bedrock access/lookup failures embed the caller's role ARN and account id in the + # raw botocore message; surface only the error code so identifiers never reach logs + # or the user-facing channel. + if isinstance(exc, ClientError): + code = exc.response.get("Error", {}).get("Code", "") + if code in {"AccessDeniedException", "ResourceNotFoundException"}: + return ( + "The selected Bedrock model is not available to this deployment " + f"(Bedrock error: {code}). Verify the model's inference-profile access and " + "IAM permissions, choose a different model, and retry." + ) + return None diff --git a/agent/middleware/sanitize_thinking_blocks.py b/agent/middleware/sanitize_thinking_blocks.py index c757f5bc..6e40beb9 100644 --- a/agent/middleware/sanitize_thinking_blocks.py +++ b/agent/middleware/sanitize_thinking_blocks.py @@ -29,19 +29,32 @@ def _is_chat_anthropic(model: object) -> bool: return False +def _is_empty_thinking_block(block: object) -> bool: + """True for an empty Anthropic ``thinking`` block or an empty Bedrock + ``reasoning_content`` block (Bedrock Converse emits the latter shape).""" + if not isinstance(block, dict): + return False + block_type = block.get("type") + if block_type == "thinking": + return not block.get("thinking") + if block_type == "reasoning_content": + payload = block.get("reasoning_content") + if isinstance(payload, dict): + return not ( + payload.get("text") + or payload.get("signature") + or payload.get("redactedContent") + or payload.get("redacted_content") + ) + return not payload + return False + + def _sanitize_messages(messages: list[Any]) -> None: for message in messages: if not isinstance(message, AIMessage) or not isinstance(message.content, list): continue - content = [ - block - for block in message.content - if not ( - isinstance(block, dict) - and block.get("type") == "thinking" - and not block.get("thinking") - ) - ] + content = [block for block in message.content if not _is_empty_thinking_block(block)] if len(content) != len(message.content): message.content = content diff --git a/agent/utils/model.py b/agent/utils/model.py index 6fc469d4..cadf1dbd 100644 --- a/agent/utils/model.py +++ b/agent/utils/model.py @@ -60,7 +60,12 @@ def make_model(model_id: str, **kwargs: Unpack[ModelKwargs]): model_kwargs["base_url"] = OPENAI_RESPONSES_WS_BASE_URL model_kwargs["use_responses_api"] = True elif model_id.startswith("bedrock_converse:"): - model_kwargs.setdefault("region_name", os.environ.get("AWS_REGION", "us-east-1")) + # Resolve region with the same precedence validate_local_dev_llm_config accepts + # (AWS_REGION or AWS_DEFAULT_REGION), so the validated value is the one actually used. + model_kwargs.setdefault( + "region_name", + os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1", + ) return init_chat_model(model=model_id, **model_kwargs) diff --git a/tests/test_sanitize_thinking_blocks.py b/tests/test_sanitize_thinking_blocks.py index 7145fe73..5e8a4e96 100644 --- a/tests/test_sanitize_thinking_blocks.py +++ b/tests/test_sanitize_thinking_blocks.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock import pytest from langchain_anthropic import ChatAnthropic +from langchain_aws import ChatBedrockConverse from langchain_core.messages import AIMessage, HumanMessage from agent.middleware.sanitize_thinking_blocks import SanitizeThinkingBlocksMiddleware @@ -66,6 +67,32 @@ class TestSanitizeThinkingBlocksMiddleware: assert result is response assert message.content == [{"type": "text", "text": "ok"}] + def test_drops_empty_reasoning_content_block_for_bedrock(self) -> None: + message = AIMessage( + content=[ + {"type": "reasoning_content", "reasoning_content": {"text": "", "signature": ""}}, + {"type": "text", "text": "ok"}, + ] + ) + request = _make_request([message], model=MagicMock(spec=ChatBedrockConverse)) + + SanitizeThinkingBlocksMiddleware().wrap_model_call(request, lambda req: MagicMock()) + + assert message.content == [{"type": "text", "text": "ok"}] + + def test_preserves_non_empty_reasoning_content_block_for_bedrock(self) -> None: + reasoning_block = { + "type": "reasoning_content", + "reasoning_content": {"text": "because", "signature": "sig"}, + } + text_block = {"type": "text", "text": "ok"} + message = AIMessage(content=[reasoning_block, text_block]) + request = _make_request([message], model=MagicMock(spec=ChatBedrockConverse)) + + SanitizeThinkingBlocksMiddleware().wrap_model_call(request, lambda req: MagicMock()) + + assert message.content == [reasoning_block, text_block] + def test_ignores_non_anthropic_models(self) -> None: thinking_block = {"type": "thinking", "signature": "abc", "thinking": ""} message = AIMessage(content=[thinking_block, {"type": "text", "text": "ok"}])