diff --git a/.gitignore b/.gitignore
index 679a6d8a..c7039374 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
.env.local
credentials.json
state.json
+venv/
# IDE
.vscode/
diff --git a/kiro/auth.py b/kiro/auth.py
index 61ffe02b..a56f0300 100644
--- a/kiro/auth.py
+++ b/kiro/auth.py
@@ -48,7 +48,7 @@
get_kiro_q_host,
get_aws_sso_oidc_url,
)
-from kiro.utils import get_machine_fingerprint
+from kiro.utils import get_machine_fingerprint, detect_kiro_agent_profile_arn
# Supported SQLite token keys (searched in priority order)
@@ -949,7 +949,16 @@ async def force_refresh(self) -> str:
@property
def profile_arn(self) -> Optional[str]:
"""AWS CodeWhisperer profile ARN."""
- return self._profile_arn
+ if self._profile_arn:
+ return self._profile_arn
+
+ # Fallback to auto-detect from Kiro Agent IDE's profile.json
+ detected = detect_kiro_agent_profile_arn()
+ if detected:
+ self._profile_arn = detected
+ return detected
+
+ return None
@property
def region(self) -> str:
diff --git a/kiro/converters_anthropic.py b/kiro/converters_anthropic.py
index ea297dc1..71d85381 100644
--- a/kiro/converters_anthropic.py
+++ b/kiro/converters_anthropic.py
@@ -24,7 +24,7 @@
to the unified format used by converters_core.py.
"""
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
from loguru import logger
@@ -45,6 +45,47 @@
)
+def _render_web_search_result_block(block: Any) -> str:
+ """
+ Render a web_search_tool_result content block as readable text.
+
+ When a client echoes a prior web_search turn back to the gateway, the raw
+ result blocks carry the grounding the assistant used. We fold them into text
+ so that context survives conversion to Kiro (which has no native
+ web_search_tool_result representation).
+
+ Args:
+ block: A web_search_tool_result block (dict or Pydantic model).
+
+ Returns:
+ A readable summary string, or "" if there is nothing to render.
+ """
+ content = block.get("content") if isinstance(block, dict) else getattr(block, "content", None)
+
+ # Failure case: the native error object.
+ if isinstance(content, dict):
+ if content.get("type") == "web_search_tool_result_error":
+ return "\n\n(Previous web search was unavailable.)\n\n"
+ return ""
+
+ if not isinstance(content, list):
+ return ""
+
+ lines = []
+ for item in content:
+ item_dict = item if isinstance(item, dict) else getattr(item, "__dict__", {})
+ title = item_dict.get("title", "")
+ url = item_dict.get("url", "")
+ snippet = item_dict.get("encrypted_content", "") or item_dict.get("snippet", "")
+ if title or url or snippet:
+ lines.append(f"- {title} ({url})\n {snippet}".rstrip())
+
+ if not lines:
+ return ""
+
+ return "\n\nPrevious search results:\n" + "\n".join(lines) + "\n\n"
+
+
def convert_anthropic_content_to_text(content: Any) -> str:
"""
Extracts text content from Anthropic message content.
@@ -53,6 +94,10 @@ def convert_anthropic_content_to_text(content: Any) -> str:
- String: "Hello, world!"
- List of content blocks: [{"type": "text", "text": "Hello"}]
+ Server-side web_search blocks (server_tool_use / web_search_tool_result) that
+ a client echoes back from a prior turn are folded into text so the search
+ grounding is preserved through conversion to Kiro.
+
Args:
content: Anthropic message content
@@ -66,10 +111,16 @@ def convert_anthropic_content_to_text(content: Any) -> str:
text_parts = []
for block in content:
if isinstance(block, dict):
- if block.get("type") == "text":
+ block_type = block.get("type")
+ if block_type == "text":
text_parts.append(block.get("text", ""))
- elif hasattr(block, "type") and block.type == "text":
- text_parts.append(block.text)
+ elif block_type == "web_search_tool_result":
+ text_parts.append(_render_web_search_result_block(block))
+ elif hasattr(block, "type"):
+ if block.type == "text":
+ text_parts.append(getattr(block, "text", ""))
+ elif block.type == "web_search_tool_result":
+ text_parts.append(_render_web_search_result_block(block))
return "".join(text_parts)
return str(content) if content else ""
@@ -426,6 +477,55 @@ def extract_thinking_config_from_anthropic(request: AnthropicMessagesRequest) ->
return ThinkingConfig(enabled=True, budget_tokens=None)
+def separate_inline_system_messages(
+ messages: List[AnthropicMessage],
+) -> Tuple[List[str], List[AnthropicMessage]]:
+ """
+ Separates inline ``system`` role messages from conversation turns.
+
+ Some clients (notably Claude Code) inject messages with ``role == "system"``
+ directly into the ``messages`` array (e.g. SessionStart hook context), even
+ though the Anthropic spec keeps ``system`` as a separate top-level field.
+
+ To stay consistent with the OpenAI adapter (which extracts all ``system``
+ messages into the system prompt, see ``convert_openai_messages_to_unified``),
+ we peel these inline system messages out here and return their text so the
+ caller can merge it into the system prompt. This keeps system instructions
+ out of the conversation history and avoids turning them into ``user`` turns
+ (which would trigger synthetic alternation placeholders downstream).
+
+ Args:
+ messages: List of Anthropic messages (possibly containing inline system)
+
+ Returns:
+ Tuple of:
+ - List of extracted system text fragments, in original order
+ - List of remaining conversation messages (user/assistant and any other
+ non-system roles, which are normalized later by the core layer)
+
+ Example:
+ >>> msgs = [
+ ... AnthropicMessage(role="user", content="Hi"),
+ ... AnthropicMessage(role="system", content="Hook context"),
+ ... ]
+ >>> parts, convo = separate_inline_system_messages(msgs)
+ >>> parts
+ ['Hook context']
+ >>> [m.role for m in convo]
+ ['user']
+ """
+ inline_system_parts: List[str] = []
+ conversation_messages: List[AnthropicMessage] = []
+
+ for msg in messages:
+ if msg.role == "system":
+ inline_system_parts.append(convert_anthropic_content_to_text(msg.content))
+ else:
+ conversation_messages.append(msg)
+
+ return inline_system_parts, conversation_messages
+
+
def anthropic_to_kiro(
request: AnthropicMessagesRequest, conversation_id: str, profile_arn: str
) -> dict:
@@ -435,7 +535,9 @@ def anthropic_to_kiro(
This is the main entry point for Anthropic → Kiro conversion.
Key differences from OpenAI:
- - System prompt is a separate field (not in messages)
+ - System prompt is a separate top-level field. Inline ``system`` messages in
+ the ``messages`` array (sent by some clients like Claude Code) are merged
+ into that system prompt, matching the OpenAI adapter's behavior.
- Content can be string or list of content blocks
- Tool format uses input_schema instead of parameters
@@ -450,15 +552,37 @@ def anthropic_to_kiro(
Raises:
ValueError: If there are no messages to send
"""
- # Convert messages to unified format
- unified_messages = convert_anthropic_messages(request.messages)
+ # Separate inline system messages (e.g. Claude Code hook context) from the
+ # actual conversation turns. This mirrors the OpenAI adapter, which extracts
+ # all system messages into the system prompt rather than the history.
+ inline_system_parts, conversation_messages = separate_inline_system_messages(
+ request.messages
+ )
+
+ if inline_system_parts:
+ logger.debug(
+ f"Merged {len(inline_system_parts)} inline system message(s) into system prompt"
+ )
+
+ # Convert conversation messages to unified format
+ unified_messages = convert_anthropic_messages(conversation_messages)
# Convert tools to unified format
unified_tools = convert_anthropic_tools(request.tools)
- # System prompt is already separate in Anthropic format!
- # It can be a string or list of content blocks (for prompt caching)
+ # System prompt comes from the top-level field first, then any inline system
+ # messages are appended (preserving Claude Code's intended ordering).
system_prompt = extract_system_prompt(request.system)
+ if inline_system_parts:
+ inline_system_text = "\n\n".join(
+ part for part in inline_system_parts if part
+ )
+ if inline_system_text:
+ system_prompt = (
+ f"{system_prompt}\n\n{inline_system_text}"
+ if system_prompt
+ else inline_system_text
+ )
# Get model ID for Kiro API (normalizes + resolves hidden models)
# Pass-through principle: we normalize and send to Kiro, Kiro decides if valid
@@ -470,7 +594,7 @@ def anthropic_to_kiro(
logger.debug(
f"Converting Anthropic request: model={request.model} -> {model_id}, "
f"messages={len(unified_messages)}, tools={len(unified_tools) if unified_tools else 0}, "
- f"system_prompt_length={len(system_prompt)}, "
+ f"system_prompt_length={len(system_prompt) if system_prompt else 0}, "
f"thinking_enabled={thinking_config.enabled}, thinking_budget={thinking_config.budget_tokens}"
)
diff --git a/kiro/mcp_tools.py b/kiro/mcp_tools.py
index c908ec55..6ee62ed3 100644
--- a/kiro/mcp_tools.py
+++ b/kiro/mcp_tools.py
@@ -147,21 +147,43 @@ async def call_kiro_mcp_api(
try:
token = await auth_manager.get_access_token()
- # EXACT headers from architecture
- headers = {
- "Authorization": f"Bearer {token}",
- "x-amzn-codewhisperer-optout": "false",
- "Content-Type": "application/json"
- }
+ # Endpoint selection for web_search MCP:
+ # - kiro-cli / AWS SSO OIDC accounts must call the Amazon Q endpoint
+ # (https://q.{region}.amazonaws.com/mcp) and pass the profile ARN via the
+ # x-amzn-kiro-profile-arn header. The runtime.kiro.dev host rejects MCP
+ # web_search with 403 "User is not authorized" / 400 "profileArn is required".
+ # - Other auth types keep the original q_host behavior.
+ from kiro.auth import AuthType
+
+ profile_arn = getattr(auth_manager, "profile_arn", None)
+ is_sso = getattr(auth_manager, "auth_type", None) == AuthType.AWS_SSO_OIDC
+
+ if is_sso and profile_arn:
+ # Derive region from the profile ARN (arn:aws:codewhisperer:{region}:...).
+ arn_parts = profile_arn.split(":")
+ region = arn_parts[3] if len(arn_parts) > 3 and arn_parts[3] else auth_manager.region
+ mcp_url = f"https://q.{region}.amazonaws.com/mcp"
+ headers = {
+ "Authorization": f"Bearer {token}",
+ "x-amzn-kiro-profile-arn": profile_arn,
+ "x-amzn-codewhisperer-optout": "false",
+ "Content-Type": "application/x-amz-json-1.0",
+ }
+ else:
+ headers = {
+ "Authorization": f"Bearer {token}",
+ "x-amzn-codewhisperer-optout": "false",
+ "Content-Type": "application/json",
+ }
+ mcp_url = f"{auth_manager.q_host}/mcp"
- mcp_url = f"{auth_manager.q_host}/mcp"
logger.debug(f"Calling MCP API: {mcp_url}")
async with httpx.AsyncClient(timeout=60.0) as client:
response = await client.post(mcp_url, json=mcp_request, headers=headers)
if response.status_code != 200:
- logger.error(f"MCP API error: {response.status_code}")
+ logger.error(f"MCP API error: {response.status_code} - {response.text[:300]}")
return None, None
mcp_response = response.json()
diff --git a/kiro/models_anthropic.py b/kiro/models_anthropic.py
index c63d60ba..14dd7113 100644
--- a/kiro/models_anthropic.py
+++ b/kiro/models_anthropic.py
@@ -162,7 +162,53 @@ class ImageContentBlock(BaseModel):
source: Union[Base64ImageSource, URLImageSource]
-# Union type for all content blocks (including images and thinking)
+class ServerToolUseContentBlock(BaseModel):
+ """
+ Server-side tool use block (e.g. web_search), in Anthropic format.
+
+ The gateway emits these for web_search (Path A and Path B). When a client
+ like Claude Code echoes the conversation back on the next turn, the assistant
+ message contains this block, so we must accept it on input as well.
+
+ Attributes:
+ type: Always "server_tool_use"
+ id: Tool use ID (links to the matching web_search_tool_result)
+ name: Server-side tool name (e.g. "web_search")
+ input: Tool input (e.g. {"query": "..."})
+ """
+
+ type: Literal["server_tool_use"] = "server_tool_use"
+ id: str
+ name: str
+ input: Dict[str, Any] = {}
+
+ model_config = {"extra": "allow"}
+
+
+class WebSearchToolResultContentBlock(BaseModel):
+ """
+ web_search tool result block, in Anthropic format.
+
+ Emitted by the gateway alongside server_tool_use and echoed back by clients
+ on subsequent turns. The content is either a list of web_search_result items
+ (success) or a web_search_tool_result_error object (failure), so it is kept
+ permissive here.
+
+ Attributes:
+ type: Always "web_search_tool_result"
+ tool_use_id: ID linking back to the server_tool_use block
+ content: List of result items, an error object, or a string
+ """
+
+ type: Literal["web_search_tool_result"] = "web_search_tool_result"
+ tool_use_id: str
+ content: Union[str, List[Dict[str, Any]], Dict[str, Any], None] = None
+
+ model_config = {"extra": "allow"}
+
+
+# Union type for all content blocks (including images, thinking, and the
+# server-side web_search blocks the gateway both emits and accepts back)
ContentBlock = Union[
TextContentBlock,
ThinkingContentBlock,
@@ -170,6 +216,8 @@ class ImageContentBlock(BaseModel):
ToolUseContentBlock,
ToolResultContentBlock,
ToolReferenceContentBlock,
+ ServerToolUseContentBlock,
+ WebSearchToolResultContentBlock,
]
@@ -182,12 +230,28 @@ class AnthropicMessage(BaseModel):
"""
Message in Anthropic format.
+ Although the Anthropic spec only documents ``user`` and ``assistant`` roles
+ in the ``messages`` array (``system`` is a separate top-level field), some
+ clients (notably Claude Code) inject inline messages with non-standard roles
+ such as ``system`` mid-conversation. Rejecting these with a 422 would break
+ those clients, so the role is accepted as a free-form string here, mirroring
+ the OpenAI ``ChatMessage`` model.
+
+ Downstream handling (in ``anthropic_to_kiro``):
+ - Inline ``system`` messages are hoisted into the top-level system prompt,
+ matching the OpenAI adapter's behavior (keeps system instructions out of
+ the conversation history).
+ - Any other non-standard role (e.g. ``developer``) is normalized to ``user``
+ by ``normalize_message_roles`` in ``converters_core.py``.
+
Attributes:
- role: Message role (user or assistant)
+ role: Message role. Standard values are ``user`` and ``assistant``;
+ inline ``system`` is hoisted to the system prompt, and other
+ non-standard roles are normalized to ``user`` downstream.
content: Message content (string or list of content blocks)
"""
- role: Literal["user", "assistant"]
+ role: str
content: Union[str, List[ContentBlock]]
model_config = {"extra": "allow"}
@@ -568,3 +632,10 @@ class AnthropicErrorResponse(BaseModel):
type: Literal["error"] = "error"
error: AnthropicErrorDetail
+
+
+# Rebuild models that use forward references in Union types.
+# ToolResultContentBlock.content references "ToolReferenceContentBlock" as a string —
+# Pydantic v2 needs an explicit rebuild to resolve it after all classes are defined.
+ToolResultContentBlock.model_rebuild()
+AnthropicMessage.model_rebuild()
diff --git a/kiro/parsers.py b/kiro/parsers.py
index 9bd7f433..e3f6703f 100644
--- a/kiro/parsers.py
+++ b/kiro/parsers.py
@@ -439,11 +439,19 @@ def _finalize_tool_call(self) -> None:
f"This is a Kiro API limitation. "
f"{'Model will be notified automatically about truncation.' if TRUNCATION_RECOVERY else 'Set TRUNCATION_RECOVERY=true in .env to auto-notify model about truncation.'}"
)
+
+ # Supply mock parameter values matching required tool schemas to bypass client-side validation
+ fallback_args = {}
+ if tool_name == "Write":
+ fallback_args = {"file_path": "TRUNCATED_BY_API", "content": ""}
+ elif tool_name == "Edit":
+ fallback_args = {"file_path": "TRUNCATED_BY_API", "edits": []}
+
+ self.current_tool_call['function']['arguments'] = json.dumps(fallback_args)
else:
# Regular JSON parse error
logger.warning(f"Failed to parse tool '{tool_name}' arguments: {e}. Raw: {args[:200]}")
-
- self.current_tool_call['function']['arguments'] = "{}"
+ self.current_tool_call['function']['arguments'] = "{}"
else:
# Empty string - use empty object
# This is normal behavior for duplicate tool calls from Kiro
diff --git a/kiro/utils.py b/kiro/utils.py
index 1ebe46b0..3e85a57b 100644
--- a/kiro/utils.py
+++ b/kiro/utils.py
@@ -26,8 +26,11 @@
import hashlib
import json
+import os
+import sys
import uuid
-from typing import TYPE_CHECKING, List, Dict, Any
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Dict, Any, Optional
from loguru import logger
@@ -170,4 +173,56 @@ def generate_tool_call_id() -> str:
Returns:
ID in format "call_{uuid_hex[:8]}"
"""
- return f"call_{uuid.uuid4().hex[:8]}"
\ No newline at end of file
+ return f"call_{uuid.uuid4().hex[:8]}"
+
+
+def detect_kiro_agent_profile_arn() -> Optional[str]:
+ """
+ Attempts to auto-detect the profile ARN from Kiro Agent IDE's global storage.
+
+ Checks %APPDATA%/Kiro/User/globalStorage/kiro.kiroagent/profile.json (on Windows)
+ or the platform-specific equivalents.
+
+ Returns:
+ Auto-detected profile ARN string, or None if not found/invalid
+ """
+ try:
+ home = Path.home()
+ if os.name == 'nt': # Windows
+ appdata = os.getenv('APPDATA')
+ if not appdata:
+ return None
+ kiro_dir = Path(appdata) / "Kiro"
+ elif sys.platform == 'darwin': # macOS
+ kiro_dir = home / "Library" / "Application Support" / "Kiro"
+ else: # Linux/Unix
+ kiro_dir = home / ".config" / "Kiro"
+
+ # 1. Check the direct path first
+ profile_path = kiro_dir / "User" / "globalStorage" / "kiro.kiroagent" / "profile.json"
+ if profile_path.exists():
+ with open(profile_path, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+ arn = data.get("profileArn") or data.get("profile_arn") or data.get("arn")
+ if arn:
+ logger.info(f"Auto-detected Kiro Agent profile ARN: {arn}")
+ return arn
+
+ # 2. Fall back to recursive search in kiro.kiroagent folder if it exists
+ kiro_agent_dir = kiro_dir / "User" / "globalStorage" / "kiro.kiroagent"
+ if kiro_agent_dir.exists():
+ for root, _, files in os.walk(kiro_agent_dir):
+ if "profile.json" in files:
+ file_path = Path(root) / "profile.json"
+ try:
+ with open(file_path, 'r', encoding='utf-8') as f:
+ data = json.load(f)
+ arn = data.get("profileArn") or data.get("profile_arn") or data.get("arn")
+ if arn:
+ logger.info(f"Auto-detected Kiro Agent profile ARN from subdirectory: {arn}")
+ return arn
+ except Exception as inner_e:
+ logger.debug(f"Failed to read Kiro Agent profile.json at {file_path}: {inner_e}")
+ except Exception as e:
+ logger.warning(f"Failed to auto-detect Kiro Agent profile ARN: {e}")
+ return None
\ No newline at end of file
diff --git a/tests/unit/test_auth_manager.py b/tests/unit/test_auth_manager.py
index 30226c6c..79b2b8fb 100644
--- a/tests/unit/test_auth_manager.py
+++ b/tests/unit/test_auth_manager.py
@@ -500,6 +500,45 @@ def test_fingerprint_property(self):
print(f"fingerprint: {manager.fingerprint}")
assert len(manager.fingerprint) == 64
+ def test_profile_arn_fallback_auto_detection(self):
+ """
+ What it does: Verifies profile_arn fallback to detect_kiro_agent_profile_arn.
+ Purpose: Ensure Kiro Agent IDE profile.json is read if no profile_arn is set.
+ """
+ print("Setup: Creating KiroAuthManager without profile_arn...")
+ manager = KiroAuthManager(refresh_token="test")
+
+ print("Action: Accessing profile_arn with mocked detect_kiro_agent_profile_arn...")
+ mock_arn = "arn:aws:codewhisperer:us-east-1:1111:profile/auto-detected"
+ with patch("kiro.auth.detect_kiro_agent_profile_arn", return_value=mock_arn) as mock_detect:
+ arn = manager.profile_arn
+ mock_detect.assert_called_once()
+
+ print(f"Comparing profile_arn: Expected '{mock_arn}', Got '{arn}'")
+ assert arn == mock_arn
+
+ # Verify it caches the value in memory
+ print("Verification: Subsequent accesses do not call detect again...")
+ with patch("kiro.auth.detect_kiro_agent_profile_arn") as mock_detect2:
+ arn2 = manager.profile_arn
+ mock_detect2.assert_not_called()
+ assert arn2 == mock_arn
+
+ def test_profile_arn_returns_none_if_no_detection(self):
+ """
+ What it does: Verifies profile_arn returns None when no profile_arn is set
+ and auto-detection fails (returns None).
+ Purpose: Prevent crash and return None as standard default.
+ """
+ print("Setup: Creating KiroAuthManager without profile_arn...")
+ manager = KiroAuthManager(refresh_token="test")
+
+ print("Action: Accessing profile_arn when auto-detection returns None...")
+ with patch("kiro.auth.detect_kiro_agent_profile_arn", return_value=None):
+ arn = manager.profile_arn
+
+ assert arn is None
+
# =============================================================================
# Tests for AuthType enum
diff --git a/tests/unit/test_mcp_tools.py b/tests/unit/test_mcp_tools.py
index 34d22ffd..84b8cb02 100644
--- a/tests/unit/test_mcp_tools.py
+++ b/tests/unit/test_mcp_tools.py
@@ -252,6 +252,111 @@ async def test_mcp_api_json_decode_error(self, mock_auth_manager):
assert results is None
+# ==================================================================================================
+# Tests for MCP Endpoint Selection (host + profile ARN header)
+# ==================================================================================================
+
+class TestMCPEndpointSelection:
+ """
+ Tests for endpoint/header selection in call_kiro_mcp_api.
+
+ The Kiro MCP web_search endpoint behaves differently depending on auth type:
+ - AWS SSO OIDC (kiro-cli) accounts must call the Amazon Q endpoint
+ (https://q.{region}.amazonaws.com/mcp) and pass the profile ARN in the
+ x-amzn-kiro-profile-arn header. The runtime.kiro.dev host rejects these
+ requests with 400 "profileArn is required" / 403 "not authorized".
+ - Other auth types (Kiro Desktop) keep the q_host endpoint without that header.
+ """
+
+ @staticmethod
+ def _make_mock_client():
+ """Build a mocked httpx.AsyncClient returning a minimal successful response."""
+ mock_response_data = {
+ "id": "web_search_tooluse_abc_1_xyz",
+ "jsonrpc": "2.0",
+ "result": {
+ "content": [{
+ "type": "text",
+ "text": json.dumps({"results": [], "totalResults": 0, "query": "q"})
+ }],
+ "isError": False
+ }
+ }
+ mock_response = Mock()
+ mock_response.status_code = 200
+ mock_response.json = Mock(return_value=mock_response_data)
+
+ mock_post = AsyncMock(return_value=mock_response)
+ mock_client = AsyncMock()
+ mock_client.__aenter__.return_value.post = mock_post
+ return mock_client, mock_post
+
+ @pytest.mark.asyncio
+ async def test_sso_account_uses_amazon_q_host_and_profile_arn_header(self):
+ """
+ What it does: Verifies AWS SSO OIDC accounts call the Amazon Q endpoint
+ with the profile ARN header and the x-amz-json content type.
+ Purpose: Prevent regression of the 403/400 web_search failure for kiro-cli.
+ """
+ from kiro.auth import KiroAuthManager
+ from datetime import datetime, timezone
+
+ # AWS SSO OIDC is detected when client_id + client_secret are present.
+ manager = KiroAuthManager(
+ refresh_token="test_refresh_token",
+ profile_arn="arn:aws:codewhisperer:eu-central-1:123456789012:profile/EXAMPLE",
+ region="us-east-1",
+ client_id="test_client_id",
+ client_secret="test_client_secret",
+ )
+ manager._access_token = "test_access_token"
+ manager._expires_at = datetime.now(timezone.utc).replace(year=2099)
+
+ mock_client, mock_post = self._make_mock_client()
+ with patch("kiro.mcp_tools.httpx.AsyncClient", return_value=mock_client):
+ await call_kiro_mcp_api("test query", manager)
+
+ called_url = mock_post.call_args.args[0]
+ called_headers = mock_post.call_args.kwargs["headers"]
+
+ # Region is derived from the profile ARN (eu-central-1), not the SSO region.
+ assert called_url == "https://q.eu-central-1.amazonaws.com/mcp"
+ assert called_headers["x-amzn-kiro-profile-arn"] == \
+ "arn:aws:codewhisperer:eu-central-1:123456789012:profile/EXAMPLE"
+ assert called_headers["Content-Type"] == "application/x-amz-json-1.0"
+ assert called_headers["Authorization"] == "Bearer test_access_token"
+
+ @pytest.mark.asyncio
+ async def test_desktop_account_uses_q_host_without_profile_arn_header(self):
+ """
+ What it does: Verifies non-SSO (Kiro Desktop) accounts keep the q_host
+ endpoint and do NOT send the profile ARN header.
+ Purpose: Ensure the SSO-specific routing does not affect Desktop auth.
+ """
+ from kiro.auth import KiroAuthManager
+ from datetime import datetime, timezone
+
+ # No client_id/secret -> Kiro Desktop auth type.
+ manager = KiroAuthManager(
+ refresh_token="test_refresh_token",
+ profile_arn="arn:aws:codewhisperer:us-east-1:123456789012:profile/EXAMPLE",
+ region="us-east-1",
+ )
+ manager._access_token = "test_access_token"
+ manager._expires_at = datetime.now(timezone.utc).replace(year=2099)
+
+ mock_client, mock_post = self._make_mock_client()
+ with patch("kiro.mcp_tools.httpx.AsyncClient", return_value=mock_client):
+ await call_kiro_mcp_api("test query", manager)
+
+ called_url = mock_post.call_args.args[0]
+ called_headers = mock_post.call_args.kwargs["headers"]
+
+ assert called_url == f"{manager.q_host}/mcp"
+ assert "x-amzn-kiro-profile-arn" not in called_headers
+ assert called_headers["Content-Type"] == "application/json"
+
+
# ==================================================================================================
# Tests for Search Summary Generation
# ==================================================================================================
diff --git a/tests/unit/test_parsers.py b/tests/unit/test_parsers.py
index 4e981aa2..e038fba6 100644
--- a/tests/unit/test_parsers.py
+++ b/tests/unit/test_parsers.py
@@ -5,6 +5,7 @@
Tests the parsing logic for AWS SSE stream from Kiro API.
"""
+import json
import pytest
from kiro.parsers import (
@@ -820,6 +821,54 @@ def test_finalize_clears_current_tool_call(self, aws_event_parser):
print(f"current_tool_call after finalization: {aws_event_parser.current_tool_call}")
assert aws_event_parser.current_tool_call is None
+ def test_finalize_with_truncated_write_tool(self, aws_event_parser):
+ """
+ What it does: Tests finalization of truncated Write tool call.
+ Goal: Ensure it returns schema-valid arguments with file_path and content.
+ """
+ print("Setup: Truncated Write tool call...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_write_truncated",
+ "type": "function",
+ "function": {
+ "name": "Write",
+ "arguments": '{"file_path": "test.txt", "content": "hello world' # Missing closing quote & brace
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ args = json.loads(aws_event_parser.tool_calls[0]["function"]["arguments"])
+ assert args["file_path"] == "TRUNCATED_BY_API"
+ assert args["content"] == ""
+
+ def test_finalize_with_truncated_edit_tool(self, aws_event_parser):
+ """
+ What it does: Tests finalization of truncated Edit tool call.
+ Goal: Ensure it returns schema-valid arguments with file_path and edits list.
+ """
+ print("Setup: Truncated Edit tool call...")
+ aws_event_parser.current_tool_call = {
+ "id": "call_edit_truncated",
+ "type": "function",
+ "function": {
+ "name": "Edit",
+ "arguments": '{"file_path": "test.txt", "edits": [' # Missing braces
+ }
+ }
+
+ print("Action: Finalizing tool call...")
+ aws_event_parser._finalize_tool_call()
+
+ print(f"Result: {aws_event_parser.tool_calls}")
+ assert len(aws_event_parser.tool_calls) == 1
+ args = json.loads(aws_event_parser.tool_calls[0]["function"]["arguments"])
+ assert args["file_path"] == "TRUNCATED_BY_API"
+ assert args["edits"] == []
+
class TestAwsEventStreamParserEdgeCases:
"""Tests for edge cases."""
diff --git a/tests/unit/test_routes_anthropic.py b/tests/unit/test_routes_anthropic.py
index d156f66a..bdec516c 100644
--- a/tests/unit/test_routes_anthropic.py
+++ b/tests/unit/test_routes_anthropic.py
@@ -337,25 +337,34 @@ def test_validates_invalid_json(self, test_client, valid_proxy_api_key):
print(f"Status: {response.status_code}")
assert response.status_code == 422
- def test_validates_invalid_role(self, test_client, valid_proxy_api_key):
+ def test_accepts_non_standard_role(self, test_client, valid_proxy_api_key):
"""
- What it does: Verifies invalid message role is rejected.
- Purpose: Anthropic model strictly validates role (only 'user' or 'assistant').
+ What it does: Verifies a non-standard message role (e.g. inline 'system')
+ is NOT rejected at the validation layer.
+ Purpose: Some clients (notably Claude Code) inject inline 'system' role
+ messages into the messages array. These previously triggered a
+ 422 literal_error. The role field is now permissive and such
+ roles are normalized to 'user' downstream, so validation must
+ pass (any non-422 status).
"""
- print("Action: POST /v1/messages with invalid role...")
+ print("Action: POST /v1/messages with inline system role...")
response = test_client.post(
"/v1/messages",
headers={"x-api-key": valid_proxy_api_key},
json={
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
- "messages": [{"role": "invalid_role", "content": "Hello"}]
+ "messages": [
+ {"role": "user", "content": "Hello"},
+ {"role": "system", "content": "Inline system reminder"},
+ {"role": "user", "content": "Continue"},
+ ]
}
)
print(f"Status: {response.status_code}")
- # Anthropic model strictly validates role - only 'user' or 'assistant' allowed
- assert response.status_code == 422
+ # Role is now permissive - must pass validation (normalized to 'user' downstream)
+ assert response.status_code != 422
def test_accepts_valid_request_format(self, test_client, valid_proxy_api_key):
"""