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.)
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
"""Middleware that removes malformed Anthropic thinking blocks before model calls."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
from langchain.agents.middleware import AgentMiddleware
|
|
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
|
from langchain_anthropic import ChatAnthropic
|
|
from langchain_aws import ChatBedrockConverse
|
|
from langchain_core.messages import AIMessage
|
|
|
|
|
|
def _is_chat_anthropic(model: object) -> bool:
|
|
seen: set[int] = set()
|
|
current = model
|
|
for _ in range(10):
|
|
if isinstance(current, (ChatAnthropic, ChatBedrockConverse)):
|
|
return True
|
|
current_id = id(current)
|
|
if current_id in seen:
|
|
return False
|
|
seen.add(current_id)
|
|
bound = getattr(current, "bound", None)
|
|
if bound is None or bound is current:
|
|
return False
|
|
current = bound
|
|
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 _is_empty_thinking_block(block)]
|
|
if len(content) != len(message.content):
|
|
message.content = content
|
|
|
|
|
|
class SanitizeThinkingBlocksMiddleware(AgentMiddleware):
|
|
"""Drop empty Anthropic thinking blocks before provider validation."""
|
|
|
|
def wrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], ModelResponse],
|
|
) -> ModelCallResult:
|
|
if _is_chat_anthropic(request.model):
|
|
_sanitize_messages(request.messages)
|
|
return handler(request)
|
|
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
|
) -> Any:
|
|
if _is_chat_anthropic(request.model):
|
|
_sanitize_messages(request.messages)
|
|
return await handler(request)
|