mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-10-07 16:19:09 +00:00
Adam Moussa
7b26ab9516
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.)
103 lines
3.9 KiB
Python
103 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
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
|
|
|
|
|
|
def _make_request(messages: list[object], model: object | None = None) -> MagicMock:
|
|
request = MagicMock()
|
|
request.model = model or MagicMock(spec=ChatAnthropic)
|
|
request.messages = messages
|
|
return request
|
|
|
|
|
|
class TestSanitizeThinkingBlocksMiddleware:
|
|
def test_drops_empty_thinking_block_for_anthropic(self) -> None:
|
|
message = AIMessage(
|
|
content=[
|
|
{"type": "thinking", "signature": "abc", "thinking": ""},
|
|
{"type": "text", "text": "ok"},
|
|
]
|
|
)
|
|
request = _make_request([message])
|
|
response = MagicMock()
|
|
|
|
def handler(req: object) -> object:
|
|
assert req is request
|
|
return response
|
|
|
|
result = SanitizeThinkingBlocksMiddleware().wrap_model_call(request, handler)
|
|
|
|
assert result is response
|
|
assert message.content == [{"type": "text", "text": "ok"}]
|
|
|
|
def test_preserves_non_empty_thinking_block_for_anthropic(self) -> None:
|
|
thinking_block = {"type": "thinking", "signature": "abc", "thinking": "reasoning"}
|
|
text_block = {"type": "text", "text": "ok"}
|
|
message = AIMessage(content=[thinking_block, text_block])
|
|
request = _make_request([message])
|
|
|
|
SanitizeThinkingBlocksMiddleware().wrap_model_call(request, lambda req: MagicMock())
|
|
|
|
assert message.content == [thinking_block, text_block]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_drops_missing_thinking_block_for_anthropic(self) -> None:
|
|
message = AIMessage(
|
|
content=[
|
|
{"type": "thinking", "signature": "abc"},
|
|
{"type": "text", "text": "ok"},
|
|
]
|
|
)
|
|
request = _make_request([HumanMessage(content="hi"), message])
|
|
response = MagicMock()
|
|
|
|
async def handler(req: object) -> object:
|
|
assert req is request
|
|
return response
|
|
|
|
result = await SanitizeThinkingBlocksMiddleware().awrap_model_call(request, handler)
|
|
|
|
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"}])
|
|
request = _make_request([message], model=MagicMock())
|
|
|
|
SanitizeThinkingBlocksMiddleware().wrap_model_call(request, lambda req: MagicMock())
|
|
|
|
assert message.content == [thinking_block, {"type": "text", "text": "ok"}]
|