- +
- +
- +
@@ -1248,13 +1248,13 @@

CodeBuddy2API 管理面板

-
-
总凭证数量
+
Total credentials
- 0 个有效 + 0 valid
- +
@@ -1264,15 +1264,15 @@

CodeBuddy2API 管理面板

-
运行中
-
服务运行状态
+
Running
+
Service status
- 运行时长加载中... + Loading uptime...
- -
+ +
@@ -1282,13 +1282,13 @@

CodeBuddy2API 管理面板

-
-
API 服务端点
+
API service endpoint
- 点击复制链接 + Click to copy the link
- +
@@ -1303,26 +1303,26 @@

CodeBuddy2API 管理面板

0
-
总 API 调用次数
+
Total API calls
- 持续增长 + Steadily growing
-

模型使用统计

+

Model usage statistics

- - + + @@ -1334,14 +1334,14 @@

模型使用统计

-

凭证使用统计

+

Credential usage statistics

模型名称使用次数Model nameUsage count
- - + + @@ -1352,130 +1352,130 @@

凭证使用统计

- +
- +
-

自动获取认证

+

Automatic authentication

- 点击下方按钮自动启动CodeBuddy OAuth2认证流程，系统将自动获取并保存您的认证凭证。 + Click the button below to automatically start the CodeBuddy OAuth2 authentication flow. The system will automatically obtain and save your credentials.

- - + +
- +
-

手动添加凭证

+

Manually add a credential

- + + placeholder="Enter user ID (optional)">
-

已保存的凭证

+

Saved credentials

- - + +
-
加载当前状态...
+
Loading current status...
- +
-
加载中...
+
Loading...
- +
-

聊天完成测试

+

Chat completion test

- +
- +
- +
- +
- 点击"发送测试"查看API响应... + Click "Send test" to see the API response...
-

API 使用示例

+

API usage examples

-

curl 示例:

+

curl example:

curl -X POST "http://127.0.0.1:8001/codebuddy/v1/chat/completions" \ -H "Authorization: Bearer YOUR_PASSWORD" \ -H "Content-Type: application/json" \ @@ -1523,7 +1523,7 @@

curl 示例:

] }'
-

Python 示例:

+

Python example:

import openai client = openai.OpenAI( @@ -1542,25 +1542,25 @@

Python 示例:

- +
-

服务配置

+

Service configuration

- +
-
加载配置中...
+
Loading configuration...
- 注意：部分设置（如端口号）需要重启服务后才能生效。 + Note: some settings (such as the port number) require a service restart to take effect.
@@ -1568,29 +1568,29 @@

服务配置

- +
diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..78c5011 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +asyncio_mode = auto +testpaths = tests diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..dfeaa1e --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +pytest==7.4.4 +pytest-asyncio==0.23.8 diff --git a/src/auth.py b/src/auth.py index a68cc99..e778748 100644 --- a/src/auth.py +++ b/src/auth.py @@ -1,24 +1,99 @@ """ -Authentication module for CodeBuddy2API +Client and admin authentication for CodeBuddy2API. """ -from fastapi import HTTPException, Depends -from fastapi.security import HTTPBearer -from config import get_server_password +import secrets +from dataclasses import dataclass, field +from typing import Optional -security = HTTPBearer() +from fastapi import Depends, HTTPException +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from config import get_admin_password, get_client_auth_mode, get_server_password -def authenticate(credentials = Depends(security)) -> str: - """验证用户身份""" - password = get_server_password() +security = HTTPBearer(auto_error=False) + + +@dataclass(repr=False) +class ClientAuthContext: + """Request-scoped inference authentication result.""" + + mode: str + passthrough_key: Optional[str] = field(default=None, repr=False) + + +def _extract_bearer( + credentials: Optional[HTTPAuthorizationCredentials], +) -> str: + if credentials is None or credentials.scheme.lower() != "bearer": + raise HTTPException( + status_code=401, + detail="Authorization header must use a non-empty Bearer token", + headers={"WWW-Authenticate": "Bearer"}, + ) + token = credentials.credentials.strip() + if not token: + raise HTTPException( + status_code=401, + detail="Authorization header must use a non-empty Bearer token", + headers={"WWW-Authenticate": "Bearer"}, + ) + return token + + +def _matches(token: str, expected: Optional[str]) -> bool: + return bool(expected) and secrets.compare_digest(token, expected) + + +def authenticate_inference( + credentials: Optional[HTTPAuthorizationCredentials] = Depends(security), +) -> ClientAuthContext: + """Resolve relay/passthrough/hybrid inference authentication.""" + token = _extract_bearer(credentials) + try: + mode = get_client_auth_mode() + except ValueError as exc: + raise HTTPException( + status_code=500, detail="Client authentication mode is misconfigured" + ) from exc + + relay_password = get_server_password() + if mode == "relay": + if not relay_password: + raise HTTPException( + status_code=500, + detail="CODEBUDDY_PASSWORD is not configured on the server.", + ) + if not _matches(token, relay_password): + raise HTTPException(status_code=403, detail="Invalid relay password") + return ClientAuthContext(mode="relay") + + if mode == "hybrid" and _matches(token, relay_password): + return ClientAuthContext(mode="relay") + + if mode == "hybrid" and not relay_password: + raise HTTPException( + status_code=500, + detail="CODEBUDDY_PASSWORD is required when hybrid mode is enabled", + ) + + return ClientAuthContext(mode="passthrough", passthrough_key=token) + + +def authenticate_admin( + credentials: Optional[HTTPAuthorizationCredentials] = Depends(security), +) -> str: + """Authenticate dashboard and management endpoints only.""" + token = _extract_bearer(credentials) + password = get_admin_password() if not password: raise HTTPException( - status_code=500, - detail="CODEBUDDY_PASSWORD is not configured on the server." + status_code=500, + detail="CODEBUDDY_ADMIN_PASSWORD or CODEBUDDY_PASSWORD is required", ) - - token = credentials.credentials - if token != password: - raise HTTPException(status_code=403, detail="Invalid password") - - return token \ No newline at end of file + if not _matches(token, password): + raise HTTPException(status_code=403, detail="Invalid admin password") + return token + + +# Backward-compatible import name for existing admin routes. +authenticate = authenticate_admin diff --git a/src/codebuddy_api_client.py b/src/codebuddy_api_client.py index 2eb543a..a85b05e 100644 --- a/src/codebuddy_api_client.py +++ b/src/codebuddy_api_client.py @@ -1,5 +1,5 @@ """ -CodeBuddy API Client - 直接调用CodeBuddy API +CodeBuddy API Client - Calls the CodeBuddy API directly. """ import json import time @@ -13,88 +13,87 @@ class CodeBuddyAPIClient: - """CodeBuddy API客户端""" - + """CodeBuddy API client.""" + def __init__(self): from config import get_codebuddy_api_endpoint self.base_url = get_codebuddy_api_endpoint() - self.api_endpoint = self.base_url # 直接使用base_url，不需要plugin前缀 - + self.api_endpoint = self.base_url # Use base_url directly; no plugin prefix is required. + def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> List[Dict]: - """将OpenAI格式消息转换为CodeBuddy格式""" + """Convert OpenAI-format messages to CodeBuddy format.""" codebuddy_messages = [] - - # 过滤掉包含错误信息的消息，防止触发11128渠道检测 + + # Filter messages containing errors to avoid triggering channel detection 11128. filtered_messages = [] for msg in openai_messages: content = msg.get("content", "") - # 跳过包含API错误信息的助手消息 - if (msg.get("role") == "assistant" and - isinstance(content, str) and + # Skip assistant messages containing API error text. + if (msg.get("role") == "assistant" and + isinstance(content, str) and ("Error: API error" in content or "API error:" in content)): continue filtered_messages.append(msg) - - # CodeBuddy要求至少2条消息，如果只有1条用户消息，添加系统消息 + + # CodeBuddy requires at least two messages. Add a system message when only one user message exists. if len(filtered_messages) == 1 and filtered_messages[0].get("role") == "user": system_msg = { "role": "system", "content": "You are a helpful assistant." } codebuddy_messages.append(system_msg) - + for msg in filtered_messages: role = msg.get("role", "user") content = msg.get("content", "") - + logger.debug(f"[DEBUG] Processing message - role: {role}, content type: {type(content)}") - - # 处理特殊的tool角色，转换为user角色 + + # Convert the special tool role to a user role. if role == "tool": role = "user" - logger.info(f"[ROLE_CONVERSION] Converting 'tool' role to 'user'") - - # 检查是否包含工具调用相关内容 + logger.info("[ROLE_CONVERSION] Converting 'tool' role to 'user'") + + # Detect tool-call content. has_tool_content = False - - # 检查字符串化的JSON内容 + + # Parse stringified JSON content. if isinstance(content, str) and content.startswith('[{') and content.endswith('}]'): try: parsed_content = json.loads(content) if isinstance(parsed_content, list): content = parsed_content - logger.info(f"[JSON_PARSE] Parsed stringified JSON content") + logger.info("[JSON_PARSE] Parsed stringified JSON content") except json.JSONDecodeError: pass - + if isinstance(content, list): for item in content: if isinstance(item, dict) and item.get("type") in ["tool_result", "tool_use"]: has_tool_content = True break - + if has_tool_content: - # 包含工具调用内容，保持结构化格式 + # Preserve structured tool-call content. logger.info(f"[TOOL_CONTENT] Preserving structured content for role: {role}") - - # 确保工具结果有正确的toolUseId + + # Ensure every tool result has a valid toolUseId. processed_content = [] for item in content: if isinstance(item, dict): if item.get("type") == "tool_result": - # 确保toolUseId存在且有效 tool_use_id = item.get("toolUseId") or item.get("tool_use_id") or item.get("id") if not tool_use_id: - # 生成一个有效的toolUseId + # Generate a valid toolUseId. tool_use_id = f"tool_{uuid.uuid4().hex[:8]}" logger.warning(f"[TOOL_RESULT] Missing toolUseId, generated: {tool_use_id}") - - # 确保toolUseId符合正则表达式要求 [a-zA-Z0-9_-]+ + + # Ensure toolUseId matches [a-zA-Z0-9_-]+. if not tool_use_id or not all(c.isalnum() or c in '_-' for c in tool_use_id): tool_use_id = f"tool_{uuid.uuid4().hex[:8]}" logger.warning(f"[TOOL_RESULT] Invalid toolUseId format, regenerated: {tool_use_id}") - - # 标准化工具结果格式 + + # Normalize the tool-result format. tool_result = { "type": "tool_result", "toolUseId": tool_use_id, @@ -103,7 +102,7 @@ def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> L processed_content.append(tool_result) logger.info(f"[TOOL_RESULT] Processed tool result with toolUseId: {tool_use_id}") elif item.get("type") == "tool_use": - # 确保工具使用有正确的id + # Ensure every tool use has an ID. tool_id = item.get("id") or f"tool_{uuid.uuid4().hex[:8]}" tool_use = { "type": "tool_use", @@ -114,12 +113,11 @@ def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> L processed_content.append(tool_use) logger.info(f"[TOOL_USE] Processed tool use with id: {tool_id}") elif item.get("type") == "text": - # 处理纯文本内容 + # Preserve plain text content. processed_content.append(item) else: - # 其他类型，可能是工具结果的简化格式 + # Convert simplified tool-result formats when possible. if "text" in item and not item.get("type"): - # 可能是工具结果，转换为标准格式 tool_use_id = f"tool_{uuid.uuid4().hex[:8]}" tool_result = { "type": "tool_result", @@ -132,13 +130,13 @@ def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> L processed_content.append(item) else: processed_content.append(item) - + codebuddy_msg = { "role": role, "content": processed_content } else: - # 普通文本内容，转换为字符串 + # Convert ordinary text content to a string. if isinstance(content, str): text_content = content elif isinstance(content, list): @@ -161,9 +159,9 @@ def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> L "role": role, "content": text_content } - + codebuddy_messages.append(codebuddy_msg) - + return codebuddy_messages def generate_codebuddy_headers( @@ -173,12 +171,15 @@ def generate_codebuddy_headers( conversation_id: Optional[str] = None, conversation_request_id: Optional[str] = None, conversation_message_id: Optional[str] = None, - request_id: Optional[str] = None + request_id: Optional[str] = None, + api_key_header: str = "bearer" ) -> Dict[str, str]: """ - 生成CodeBuddy API所需的完整请求头。 - 优先使用传入的会话ID，如果未提供则随机生成。 + Generate the complete header set required by the CodeBuddy API. + Prefer supplied conversation IDs and generate missing IDs automatically. """ + if api_key_header not in {"x-api-key", "bearer", "both"}: + raise ValueError("api_key_header must be x-api-key, bearer, or both") headers = { 'Host': 'www.codebuddy.ai', 'Accept': 'application/json', @@ -199,14 +200,17 @@ def generate_codebuddy_headers( 'X-IDE-Type': 'CLI', 'X-IDE-Name': 'CLI', 'X-IDE-Version': '1.0.7', - 'Authorization': f'Bearer {bearer_token}', 'X-Domain': 'www.codebuddy.ai', 'User-Agent': 'CLI/1.0.7 CodeBuddy/1.0.7', 'X-Product': 'SaaS', 'X-User-Id': user_id or 'b5be3a67-237e-4ee6-9b9a-0b9ecd7b454b' } + if api_key_header in {"x-api-key", "both"}: + headers["X-API-Key"] = bearer_token + if api_key_header in {"bearer", "both"}: + headers["Authorization"] = f"Bearer {bearer_token}" return headers -# 全局客户端实例 -codebuddy_api_client = CodeBuddyAPIClient() \ No newline at end of file +# Global client instance. +codebuddy_api_client = CodeBuddyAPIClient() diff --git a/src/codebuddy_api_key_manager.py b/src/codebuddy_api_key_manager.py new file mode 100644 index 0000000..26109ee --- /dev/null +++ b/src/codebuddy_api_key_manager.py @@ -0,0 +1,300 @@ +""" +CodeBuddy API Key Manager - TXT key pool, rotation, and runtime state +""" +import asyncio +import hashlib +import logging +import os +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Callable, Dict, List, Optional, Set + +from config import ( + get_codebuddy_api_key_cooldown_seconds, + get_codebuddy_api_key_reload_interval, + get_codebuddy_api_key_rotation, + get_codebuddy_api_keys_file, +) + +logger = logging.getLogger(__name__) + + +class ApiKeyConfigurationError(Exception): + """The API key file configuration is invalid.""" + + +@dataclass(repr=False) +class ApiKeySelection: + """A request-scoped key that must never be serialized or logged.""" + + key: str = field(repr=False) + key_id: str + masked_key: str + + +@dataclass(repr=False) +class _ApiKeyState: + key: str = field(repr=False) + key_id: str + masked_key: str + status: str = "active" + request_count: int = 0 + error_count: int = 0 + last_used_at: Optional[float] = None + cooldown_until: Optional[float] = None + last_error: Optional[str] = None + + +class CodeBuddyApiKeyManager: + """Concurrency-safe CodeBuddy upstream API key pool.""" + + def __init__( + self, + file_path: Optional[str] = None, + cooldown_seconds: Optional[int] = None, + reload_interval: Optional[int] = None, + clock: Callable[[], float] = time.time, + ): + self._file_path = file_path + self._cooldown_seconds = cooldown_seconds + self._reload_interval = reload_interval + self._clock = clock + self._states: List[_ApiKeyState] = [] + self._cursor = 0 + self._lock = asyncio.Lock() + self._reload_task: Optional[asyncio.Task] = None + + @staticmethod + def parse_text(content: str) -> List[str]: + """Parse TXT content while preserving order and removing duplicates.""" + keys = [] + seen = set() + for raw_line in content.splitlines(): + key = raw_line.strip() + if not key or key.startswith("#") or key in seen: + continue + seen.add(key) + keys.append(key) + return keys + + @staticmethod + def _key_id(key: str) -> str: + return hashlib.sha256(key.encode("utf-8")).hexdigest() + + @staticmethod + def mask_key(key: str) -> str: + if len(key) < 9: + return "****" + return f"{key[:4]}...{key[-4:]}" + + def _get_file_path(self) -> str: + path = self._file_path if self._file_path is not None else get_codebuddy_api_keys_file() + if not path: + raise ApiKeyConfigurationError("CODEBUDDY_API_KEYS_FILE is not configured") + return os.path.abspath(os.path.expanduser(path)) + + def _get_cooldown_seconds(self) -> int: + if self._cooldown_seconds is not None: + return self._cooldown_seconds + return get_codebuddy_api_key_cooldown_seconds() + + def _get_reload_interval(self) -> int: + if self._reload_interval is not None: + return self._reload_interval + return get_codebuddy_api_key_reload_interval() + + def _refresh_expired_cooldowns(self, now: float) -> None: + for state in self._states: + if ( + state.status == "cooldown" + and state.cooldown_until is not None + and state.cooldown_until <= now + ): + state.status = "active" + state.cooldown_until = None + + def _read_keys(self) -> tuple[List[str], bool]: + try: + get_codebuddy_api_key_rotation() + path = self._get_file_path() + with open(path, "r", encoding="utf-8") as file: + return self.parse_text(file.read()), True + except (OSError, ValueError, ApiKeyConfigurationError) as exc: + if isinstance(exc, ValueError): + logger.error("Invalid API key rotation configuration") + else: + logger.warning("CodeBuddy API key file is unavailable") + return [], False + + async def reload(self) -> Dict[str, object]: + """Force a TXT file reload and merge the runtime state.""" + keys, source_available = await asyncio.to_thread(self._read_keys) + async with self._lock: + if not source_available: + return { + "loaded": len(self._states), + "added": 0, + "removed": 0, + "available": bool(self._states), + "reloaded": False, + } + existing = {state.key_id: state for state in self._states} + new_states = [] + added = 0 + for key in keys: + key_id = self._key_id(key) + state = existing.get(key_id) + if state is None: + state = _ApiKeyState( + key=key, + key_id=key_id, + masked_key=self.mask_key(key), + ) + added += 1 + new_states.append(state) + + new_ids = {state.key_id for state in new_states} + removed = sum(1 for state in self._states if state.key_id not in new_ids) + self._states = new_states + self._cursor = self._cursor % len(self._states) if self._states else 0 + logger.info( + "CodeBuddy API key pool reloaded: loaded=%d, added=%d, removed=%d", + len(self._states), + added, + removed, + ) + return { + "loaded": len(self._states), + "added": added, + "removed": removed, + "available": bool(self._states), + } + + async def start_periodic_reload(self) -> None: + """Perform the initial load and start configured background reloads.""" + await self.reload() + interval = self._get_reload_interval() + if interval <= 0 or (self._reload_task and not self._reload_task.done()): + return + self._reload_task = asyncio.create_task( + self._periodic_reload_loop(), name="codebuddy-api-key-reload" + ) + + async def _periodic_reload_loop(self) -> None: + try: + while True: + interval = self._get_reload_interval() + if interval <= 0: + return + await asyncio.sleep(interval) + await self.reload() + except asyncio.CancelledError: + raise + + async def stop_periodic_reload(self) -> None: + task = self._reload_task + self._reload_task = None + if task is None: + return + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def eligible_count(self) -> int: + async with self._lock: + self._refresh_expired_cooldowns(self._clock()) + return sum(1 for state in self._states if state.status == "active") + + async def total_count(self) -> int: + async with self._lock: + return len(self._states) + + async def acquire(self, excluded_ids: Optional[Set[str]] = None) -> Optional[ApiKeySelection]: + """Select a key not yet tried by this request using round-robin.""" + excluded_ids = excluded_ids or set() + async with self._lock: + now = self._clock() + self._refresh_expired_cooldowns(now) + if not self._states: + return None + + for offset in range(len(self._states)): + index = (self._cursor + offset) % len(self._states) + state = self._states[index] + if state.status != "active" or state.key_id in excluded_ids: + continue + + self._cursor = (index + 1) % len(self._states) + state.request_count += 1 + state.last_used_at = now + return ApiKeySelection( + key=state.key, + key_id=state.key_id, + masked_key=state.masked_key, + ) + return None + + async def mark_success(self, key_id: str) -> None: + async with self._lock: + state = self._find_state(key_id) + if state is not None: + state.last_error = None + + async def mark_invalid(self, key_id: str) -> None: + await self._mark_error(key_id, "invalid", "HTTP 401") + + async def mark_cooldown(self, key_id: str, status_code: int) -> None: + async with self._lock: + state = self._find_state(key_id) + if state is None: + return + state.status = "cooldown" + state.error_count += 1 + state.cooldown_until = self._clock() + self._get_cooldown_seconds() + state.last_error = f"HTTP {status_code}" + + async def mark_transient_error(self, key_id: str, error_code: str) -> None: + await self._mark_error(key_id, "active", error_code) + + async def _mark_error(self, key_id: str, status: str, error_code: str) -> None: + async with self._lock: + state = self._find_state(key_id) + if state is None: + return + state.status = status + state.error_count += 1 + state.cooldown_until = None + state.last_error = error_code + + def _find_state(self, key_id: str) -> Optional[_ApiKeyState]: + return next((state for state in self._states if state.key_id == key_id), None) + + @staticmethod + def _format_timestamp(value: Optional[float]) -> Optional[str]: + if value is None: + return None + return datetime.fromtimestamp(value, tz=timezone.utc).isoformat() + + async def get_status(self) -> Dict[str, object]: + """Return safe status data without raw keys or file paths.""" + async with self._lock: + self._refresh_expired_cooldowns(self._clock()) + return { + "keys": [ + { + "masked_key": state.masked_key, + "status": state.status, + "request_count": state.request_count, + "error_count": state.error_count, + "last_used_at": self._format_timestamp(state.last_used_at), + "cooldown_until": self._format_timestamp(state.cooldown_until), + } + for state in self._states + ], + } + + +codebuddy_api_key_manager = CodeBuddyApiKeyManager() diff --git a/src/codebuddy_auth_router.py b/src/codebuddy_auth_router.py index 8177e65..91daafe 100644 --- a/src/codebuddy_auth_router.py +++ b/src/codebuddy_auth_router.py @@ -1,8 +1,7 @@ """ CodeBuddy Authentication Router -基于真实CodeBuddy API的认证实现 +Authentication implementation based on the real CodeBuddy API """ -import hashlib import secrets import httpx import base64 @@ -11,10 +10,9 @@ import time from typing import Dict, Any, Optional from fastapi.responses import JSONResponse -from fastapi import APIRouter, HTTPException, Depends, Body -from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials +from fastapi import APIRouter, Body, Depends -from config import get_server_password +from .auth import authenticate_admin import logging logger = logging.getLogger(__name__) @@ -27,41 +25,15 @@ # --- Router Setup --- router = APIRouter() -security = HTTPBearer() - -# --- JWT Authentication --- -import jwt - -def get_jwt_secret(): - """基于服务密码生成JWT密钥""" - password = get_server_password() - if not password: - return "fallback-secret-for-development-only" - return hashlib.sha256(password.encode()).hexdigest() - -JWT_SECRET = get_jwt_secret() -ALGORITHM = "HS256" - -def authenticate(credentials = Depends(security)) -> str: - """基于服务密码的认证""" - password = get_server_password() - if not password: - raise HTTPException(status_code=500, detail="CODEBUDDY_PASSWORD is not configured on the server.") - - token = credentials.credentials - if token != password: - raise HTTPException(status_code=403, detail="Invalid password") - return token - # --- Helper Functions --- def generate_auth_state() -> str: - """生成CodeBuddy认证的state参数""" + """Generate the state parameter for CodeBuddy authentication.""" timestamp = int(time.time()) random_part = secrets.token_hex(16) return f"{random_part}_{timestamp}" def get_auth_start_headers() -> Dict[str, str]: - """生成启动认证(/state)所需的请求头""" + """Generate headers required to start authentication through /state.""" request_id = str(uuid.uuid4()).replace('-', '') return { 'Host': 'www.codebuddy.ai', @@ -82,7 +54,7 @@ def get_auth_start_headers() -> Dict[str, str]: } def get_auth_poll_headers() -> Dict[str, str]: - """生成轮询认证(/token)所需的请求头""" + """Generate headers required to poll authentication through /token.""" request_id = str(uuid.uuid4()).replace('-', '') span_id = secrets.token_hex(8) return { @@ -108,15 +80,15 @@ def get_auth_poll_headers() -> Dict[str, str]: } async def start_codebuddy_auth() -> Dict[str, Any]: - """启动CodeBuddy认证流程""" + """Start the CodeBuddy authentication flow.""" try: - logger.info("启动CodeBuddy认证流程...") + logger.info("Starting the CodeBuddy authentication flow...") headers = get_auth_start_headers() - # 调用 /v2/plugin/auth/state 获取认证状态和URL + # Call /v2/plugin/auth/state to retrieve the authentication state and URL. async with httpx.AsyncClient(verify=False) as client: - # 为避免上游/中间层缓存，添加随机nonce参数，确保每次请求唯一 + # Add a random nonce so every request is unique and bypasses upstream/intermediary caches. nonce = secrets.token_hex(8) state_url = f"{CODEBUDDY_AUTH_STATE_ENDPOINT}?platform=CLI&nonce={nonce}" payload = {"nonce": nonce} @@ -133,7 +105,7 @@ async def start_codebuddy_auth() -> Dict[str, Any]: if auth_state and auth_url: global _last_auth_state if _last_auth_state and auth_state == _last_auth_state: - logger.warning("上游返回的state与上一次相同，尝试重新获取新的state...") + logger.warning("Upstream returned the previous state; trying to retrieve a new state...") try: nonce2 = secrets.token_hex(8) state_url2 = f"{CODEBUDDY_AUTH_STATE_ENDPOINT}?platform=CLI&nonce={nonce2}" @@ -164,27 +136,27 @@ async def start_codebuddy_auth() -> Dict[str, Any]: "expires_in": 1800, "interval": 5, "status": "awaiting_login", - "instructions": "请点击链接完成CodeBuddy登录", - "message": "请使用提供的链接登录CodeBuddy", + "instructions": "Click the link to complete CodeBuddy sign-in.", + "message": "Sign in to CodeBuddy using the provided link.", "platform": "CLI" } return { "success": False, "error": "auth_start_failed", - "message": "无法启动认证流程" + "message": "Unable to start the authentication flow." } except Exception as e: - logger.error(f"启动CodeBuddy认证失败: {e}") + logger.error(f"Failed to start CodeBuddy authentication: {e}") return { "success": False, "error": "auth_start_failed", - "message": f"认证启动失败: {str(e)}" + "message": f"Authentication startup failed: {str(e)}" } async def poll_codebuddy_auth_status(auth_state: str) -> Dict[str, Any]: - """轮询CodeBuddy认证状态""" + """Poll CodeBuddy authentication status.""" try: headers = get_auth_poll_headers() url = f"{CODEBUDDY_AUTH_TOKEN_ENDPOINT}?state={auth_state}" @@ -196,18 +168,18 @@ async def poll_codebuddy_auth_status(auth_state: str) -> Dict[str, Any]: result = response.json() if result.get('code') == 11217: - # 仍在等待登录 + # Sign-in is still pending. return { "status": "pending", "message": result.get('msg', 'login ing...'), "code": result.get('code') } elif result.get('code') == 0 and result.get('data') and result.get('data', {}).get('accessToken'): - # 认证成功，获得token + # Authentication succeeded and returned a token. data = result.get('data', {}) return { "status": "success", - "message": "认证成功！", + "message": "Authentication succeeded.", "token_data": { "access_token": data.get('accessToken'), "bearer_token": data.get('accessToken'), @@ -221,7 +193,7 @@ async def poll_codebuddy_auth_status(auth_state: str) -> Dict[str, Any]: } } else: - # 其他状态码 + # Handle other status codes. return { "status": "unknown", "message": result.get('msg', 'Unknown status'), @@ -231,54 +203,54 @@ async def poll_codebuddy_auth_status(auth_state: str) -> Dict[str, Any]: else: return { "status": "error", - "message": f"API请求失败，状态码: {response.status_code}", + "message": f"API request failed with status: {response.status_code}", "response_text": response.text } except Exception as e: - logger.error(f"轮询认证状态失败: {e}") + logger.error(f"Failed to poll authentication status: {e}") return { "status": "error", - "message": f"轮询失败: {str(e)}" + "message": f"Polling failed: {str(e)}" } async def save_codebuddy_token(token_data: Dict[str, Any]) -> bool: - """保存CodeBuddy token到文件""" + """Save the CodeBuddy token to a file.""" try: from .codebuddy_token_manager import codebuddy_token_manager - # 添加创建时间 + # Add the creation timestamp. token_data["created_at"] = int(time.time()) - # 从JWT中解析用户信息 + # Parse user information from the JWT. bearer_token = token_data.get("access_token") or token_data.get("bearer_token") user_id = "unknown" user_info = {} try: if bearer_token and '.' in bearer_token: - # 分割JWT token + # Split the JWT token. parts = bearer_token.split('.') if len(parts) >= 2: payload_part = parts[1] - # 修复Base64 padding问题 + # Repair Base64 padding. missing_padding = len(payload_part) % 4 if missing_padding: payload_part += '=' * (4 - missing_padding) - # 解码JWT payload + # Decode the JWT payload. try: payload = base64.urlsafe_b64decode(payload_part) jwt_data = json.loads(payload.decode('utf-8')) - # 提取用户信息，优先使用邮箱作为用户标识 + # Extract user information, preferring email as the identifier. user_id = (jwt_data.get('email') or jwt_data.get('preferred_username') or jwt_data.get('sub') or "unknown") - # 保存完整的用户信息 + # Preserve complete user information. user_info = { 'sub': jwt_data.get('sub'), 'email': jwt_data.get('email'), @@ -292,27 +264,27 @@ async def save_codebuddy_token(token_data: Dict[str, Any]) -> bool: 'session_state': jwt_data.get('sid') } - # 移除None值 + # Remove None values. user_info = {k: v for k, v in user_info.items() if v is not None} - logger.info(f"成功解析JWT，用户: {user_id}") - logger.debug(f"JWT用户信息: {user_info}") + logger.info(f"Successfully parsed JWT for user: {user_id}") + logger.debug(f"JWT user information: {user_info}") except (json.JSONDecodeError, UnicodeDecodeError) as decode_error: - logger.warning(f"JWT payload解码失败: {decode_error}") + logger.warning(f"Failed to decode JWT payload: {decode_error}") user_id = token_data.get('domain', 'unknown') else: - logger.warning("JWT格式无效：缺少必要的部分") + logger.warning("Invalid JWT format: required sections are missing") user_id = token_data.get('domain', 'unknown') else: - logger.warning("Bearer token为空或格式无效") + logger.warning("Bearer token is empty or malformed") user_id = token_data.get('domain', 'unknown') except Exception as e: - logger.error(f"JWT解析过程发生异常: {e}") + logger.error(f"Unexpected JWT parsing error: {e}") user_id = token_data.get('domain', 'unknown') - # 构建完整的凭证数据 + # Build the complete credential data. credential_data = { "bearer_token": bearer_token, "user_id": user_id, @@ -324,79 +296,80 @@ async def save_codebuddy_token(token_data: Dict[str, Any]) -> bool: "domain": token_data.get('domain'), "session_state": token_data.get('session_state'), "user_info": user_info, - "full_response": token_data # 保存完整的原始响应 + "full_response": token_data # Preserve the complete original response. } - # 移除None值，保持文件整洁 + # Remove None values to keep the file clean. credential_data = {k: v for k, v in credential_data.items() if v is not None} - # 生成更友好的文件名 + # Generate a readable filename. timestamp = int(time.time()) safe_user_id = "".join(c for c in user_id if c.isalnum() or c in "._-")[:20] filename = f"codebuddy_{safe_user_id}_{timestamp}.json" - # 使用token管理器保存 + # Save through the token manager. success = codebuddy_token_manager.add_credential_with_data( credential_data=credential_data, filename=filename ) if success: - logger.info(f"成功保存CodeBuddy token，用户: {user_id}，文件: {filename}") + logger.info(f"Saved CodeBuddy token for user {user_id} in file {filename}") return success except Exception as e: - logger.error(f"保存CodeBuddy token失败: {e}") + logger.error(f"Failed to save CodeBuddy token: {e}") return False # --- API Endpoints --- @router.get("/auth/start", summary="Start CodeBuddy Authentication") -async def start_device_auth(): - """启动CodeBuddy认证流程""" +async def start_device_auth(_token: str = Depends(authenticate_admin)): + """Start the CodeBuddy authentication flow.""" try: - logger.info("开始启动CodeBuddy认证流程...") + logger.info("Starting the CodeBuddy authentication flow...") - # 尝试真实的CodeBuddy认证API + # Try the real CodeBuddy authentication API. real_auth_result = await start_codebuddy_auth() if real_auth_result.get('success'): - logger.info("真实CodeBuddy认证API启动成功!") + logger.info("The real CodeBuddy authentication API started successfully.") return real_auth_result else: - logger.warning(f"真实认证API失败: {real_auth_result}") + logger.warning(f"The real authentication API failed: {real_auth_result}") return real_auth_result except Exception as e: - logger.error(f"认证启动过程发生异常: {e}") + logger.error(f"Unexpected authentication startup error: {e}") return { "success": False, "error": "Unexpected error", - "message": f"认证启动失败: {str(e)}" + "message": f"Authentication startup failed: {str(e)}" } @router.post("/auth/poll", summary="Poll for OAuth token") async def poll_for_token( device_code: str = Body(None, embed=True), code_verifier: str = Body(None, embed=True), - auth_state: str = Body(None, embed=True) + auth_state: str = Body(None, embed=True), + _token: str = Depends(authenticate_admin), ): - """轮询CodeBuddy token端点""" + """Poll the CodeBuddy token endpoint.""" from .codebuddy_token_manager import codebuddy_token_manager - # 如果有auth_state，说明是真实的CodeBuddy认证流程 + # auth_state indicates the real CodeBuddy authentication flow. if auth_state: - logger.info(f"轮询真实CodeBuddy认证状态: {auth_state}") + logger.info(f"Polling real CodeBuddy authentication state: {auth_state}") poll_result = await poll_codebuddy_auth_status(auth_state) if poll_result.get('status') == 'success': - # 认证成功，保存token + # Authentication succeeded; save the token. token_data = poll_result.get('token_data', {}) if token_data: - # 提取token信息 + # Extract token information. bearer_token = token_data.get('access_token') or token_data.get('bearer_token') if bearer_token: - # 保存token + # Save the token. token_saved = await save_codebuddy_token(token_data) return JSONResponse(content={ "access_token": bearer_token, @@ -405,47 +378,47 @@ async def poll_for_token( "refresh_token": token_data.get('refresh_token'), "scope": token_data.get('scope'), "saved": token_saved, - "message": "认证成功！🎉", + "message": "Authentication succeeded! 🎉", "user_info": token_data, "domain": token_data.get('domain') }, status_code=200) else: return JSONResponse(content={ "error": "invalid_token_response", - "error_description": "API返回的响应中没有找到token" + "error_description": "The API response did not contain a token." }, status_code=400) elif poll_result.get('status') == 'pending': - # 仍在等待 + # Authorization is still pending. return JSONResponse(content={ "error": "authorization_pending", - "error_description": poll_result.get('message', '等待用户登录...'), + "error_description": poll_result.get('message', 'Waiting for user sign-in...'), "code": poll_result.get('code') }, status_code=400) else: - # 错误状态 + # Error state. return JSONResponse(content={ "error": "auth_error", - "error_description": poll_result.get('message', '认证过程发生错误'), + "error_description": poll_result.get('message', 'The authentication flow failed.'), "details": poll_result }, status_code=400) else: return JSONResponse(content={ "error": "missing_parameters", - "error_description": "缺少必要的参数：auth_state" + "error_description": "Required parameter is missing: auth_state" }, status_code=400) @router.get("/auth/callback", summary="OAuth2 callback endpoint") async def oauth_callback(code: str = None, state: str = None, error: str = None): - """OAuth2回调端点""" + """OAuth2 callback endpoint.""" if error: return JSONResponse( - content={"error": error, "error_description": "授权被拒绝或出现错误"}, + content={"error": error, "error_description": "Authorization was denied or failed."}, status_code=400 ) return JSONResponse( content={ - "message": "授权成功！请返回应用程序。", + "message": "Authorization succeeded. Return to the application.", "code": code, "state": state } diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index 6c5acdf..ec9cfb0 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -1,33 +1,39 @@ """ -CodeBuddy API Router - 兼容CodeBuddy官方API格式 -重构版本 - 优化了代码结构、错误处理和资源管理 +CodeBuddy API Router - compatible with the official CodeBuddy API format +Refactored version - improved code structure, error handling, and resource management """ import json import time import uuid import logging import asyncio -from typing import Optional, Dict, Any, List, AsyncGenerator +from dataclasses import dataclass +from typing import Optional, Dict, Any, List, AsyncGenerator, Set import httpx from fastapi import APIRouter, HTTPException, Depends, Request, Header -from fastapi.responses import StreamingResponse +from fastapi.responses import JSONResponse, StreamingResponse -from .auth import authenticate +from .auth import ClientAuthContext, authenticate_admin, authenticate_inference from .codebuddy_api_client import codebuddy_api_client +from .codebuddy_api_key_manager import ( + ApiKeyConfigurationError, + codebuddy_api_key_manager, +) from .codebuddy_token_manager import codebuddy_token_manager from .usage_stats_manager import usage_stats_manager from .keyword_replacer import apply_keyword_replacement_to_system_message +from config import get_upstream_api_key_header logger = logging.getLogger(__name__) router = APIRouter() -# --- 延迟加载配置常量 - 避免循环导入 --- +# --- Lazily loaded configuration constants - avoids circular imports --- _codebuddy_api_url: Optional[str] = None _available_models: Optional[List[str]] = None def get_codebuddy_api_url() -> str: - """延迟加载 CodeBuddy API URL""" + """Lazily load the CodeBuddy API URL""" global _codebuddy_api_url if _codebuddy_api_url is None: from config import get_codebuddy_api_endpoint @@ -35,82 +41,92 @@ def get_codebuddy_api_url() -> str: return _codebuddy_api_url def get_available_models_list() -> List[str]: - """延迟加载可用模型列表""" + """Lazily load the list of available models""" global _available_models if _available_models is None: from config import get_available_models _available_models = get_available_models() return _available_models -# --- 配置管理 --- +# --- Configuration management --- class SecurityConfig: - """安全配置管理器""" - + """Security configuration manager""" + @staticmethod def get_ssl_verify() -> bool: - """获取SSL验证设置 - 默认关闭，可通过环境变量启用""" + """Get the SSL verification setting - disabled by default, can be enabled via environment variable""" import os - # 默认关闭SSL验证，只有明确设置为true时才启用 + # SSL verification is disabled by default; only enabled when explicitly set to true ssl_verify_env = os.getenv("CODEBUDDY_SSL_VERIFY", "false").lower() ssl_verify = ssl_verify_env == "true" - + if not ssl_verify: - logger.warning("⚠️ SSL验证已禁用 - 仅在开发环境使用！生产环境请设置 CODEBUDDY_SSL_VERIFY=true") + logger.warning("⚠️ SSL verification is disabled - for development use only! Set CODEBUDDY_SSL_VERIFY=true in production") return ssl_verify -# --- HTTP 客户端配置 --- +# --- HTTP client configuration --- HTTP_CLIENT_CONFIG = { "verify": SecurityConfig.get_ssl_verify(), "timeout": httpx.Timeout(300.0, connect=30.0, read=300.0), "limits": httpx.Limits(max_keepalive_connections=20, max_connections=100) } -# --- 异步安全的 HTTP 客户端池 --- +# --- Async-safe HTTP client pool --- _http_client_pool: Optional[httpx.AsyncClient] = None _client_lock = asyncio.Lock() async def get_http_client() -> httpx.AsyncClient: - """获取全局 HTTP 客户端池 - 异步安全""" + """Get the global HTTP client pool - async-safe""" global _http_client_pool if _http_client_pool is None: async with _client_lock: - # 双重检查锁定模式 - 异步版本 + # Double-checked locking pattern - async version if _http_client_pool is None: _http_client_pool = httpx.AsyncClient(**HTTP_CLIENT_CONFIG) return _http_client_pool async def close_http_client(): - """关闭全局 HTTP 客户端池 - 异步安全""" + """Close the global HTTP client pool - async-safe""" global _http_client_pool async with _client_lock: if _http_client_pool is not None: await _http_client_pool.aclose() _http_client_pool = None -# --- 应用生命周期管理 --- +# --- Application lifecycle management --- class AppLifecycleManager: - """应用生命周期管理器 - 处理资源清理""" - + """Application lifecycle manager - handles resource cleanup""" + @staticmethod async def startup(): - """应用启动时的初始化""" - logger.info("CodeBuddy Router 启动中...") - # 预热连接池 + """Initialization at application startup""" + logger.info("CodeBuddy Router starting up...") + # Warm up the connection pool, and start API key file reloading only in the relevant auth modes + from config import get_codebuddy_auth_mode + await get_http_client() - logger.info("HTTP 连接池已初始化") - + try: + auth_mode = get_codebuddy_auth_mode() + except ValueError: + auth_mode = "auto" + logger.error("Invalid CODEBUDDY_AUTH_MODE configuration") + if auth_mode in {"auto", "api_key_file"}: + await codebuddy_api_key_manager.start_periodic_reload() + logger.info("HTTP connection pool and API key pool initialized") + @staticmethod async def shutdown(): - """应用关闭时的清理""" - logger.info("CodeBuddy Router 关闭中...") + """Cleanup at application shutdown""" + logger.info("CodeBuddy Router shutting down...") + await codebuddy_api_key_manager.stop_periodic_reload() await close_http_client() - logger.info("资源清理完成") + logger.info("Resource cleanup complete") -# 导出生命周期管理器供主应用使用 +# Export the lifecycle manager for use by the main application lifecycle_manager = AppLifecycleManager() -# --- 标准响应头 --- +# --- Standard response headers --- SSE_HEADERS = { "Cache-Control": "no-cache", "Connection": "keep-alive", @@ -119,10 +135,10 @@ async def shutdown(): "Access-Control-Allow-Headers": "*" } -# --- 辅助函数 --- +# --- Helper functions --- def format_sse_error(message: str, error_type: str = "stream_error") -> str: - """格式化SSE错误响应""" + """Format an SSE error response""" error_data = { "error": { "message": message, @@ -132,18 +148,18 @@ def format_sse_error(message: str, error_type: str = "stream_error") -> str: return f'data: {json.dumps(error_data, ensure_ascii=False)}\n\n' class OpenAICompatibilityConverter: - """将CodeBuddy格式转换为OpenAI兼容格式""" - + """Convert the CodeBuddy format to the OpenAI-compatible format""" + @staticmethod def convert_tool_call_id(codebuddy_id: str) -> str: - """转换工具调用ID格式: tooluse_xxx -> call_xxx""" + """Convert the tool call ID format: tooluse_xxx -> call_xxx""" if codebuddy_id.startswith('tooluse_'): return f"call_{codebuddy_id[8:]}" return codebuddy_id @staticmethod def convert_sse_chunk_to_openai_format(chunk_data: Dict[str, Any], tool_call_index_map: Dict[str, int]) -> Dict[str, Any]: - """将CodeBuddy SSE块转换为OpenAI格式""" + """Convert a CodeBuddy SSE chunk to OpenAI format""" if not chunk_data.get('choices'): return chunk_data @@ -154,38 +170,38 @@ def convert_sse_chunk_to_openai_format(chunk_data: Dict[str, Any], tool_call_ind if not tool_calls: return chunk_data - # 转换工具调用格式 + # Convert tool call format converted_tool_calls = [] for tc in tool_calls: converted_tc = tc.copy() - - # 转换ID格式 + + # Convert ID format if tc.get('id'): original_id = tc['id'] converted_id = OpenAICompatibilityConverter.convert_tool_call_id(original_id) converted_tc['id'] = converted_id - - # 分配新的index + + # Assign a new index if original_id not in tool_call_index_map: tool_call_index_map[original_id] = len(tool_call_index_map) - + converted_tc['index'] = tool_call_index_map[original_id] - - # 如果没有ID，使用当前最新的index + + # If there is no ID, use the current latest index elif tool_call_index_map: - # 使用最后一个工具调用的index + # Use the index of the last tool call converted_tc['index'] = max(tool_call_index_map.values()) - + converted_tool_calls.append(converted_tc) - - # 更新chunk数据 + + # Update chunk data converted_chunk = chunk_data.copy() converted_chunk['choices'][0]['delta']['tool_calls'] = converted_tool_calls return converted_chunk def parse_sse_line(line: str) -> Optional[Dict[str, Any]]: - """解析单行SSE数据""" + """Parse a single line of SSE data""" if not line.startswith('data: '): return None @@ -199,15 +215,15 @@ def parse_sse_line(line: str) -> Optional[Dict[str, Any]]: return None def validate_and_fix_tool_call_args(args: str) -> str: - """增强版的工具调用参数验证和修复 - 专门处理多工具调用问题""" + """Enhanced validation and repair of tool call arguments - specifically handles multi-tool-call issues""" if not args: return '{}' args = args.strip() - # 检查是否是多个JSON对象连接的情况 - 这是多工具调用的主要问题 + # Check whether multiple JSON objects are concatenated - this is the main multi-tool-call problem if args.count('}{') > 0: - # 尝试分离多个JSON对象 + # Try to separate the multiple JSON objects json_objects = [] current_obj = "" brace_count = 0 @@ -219,7 +235,7 @@ def validate_and_fix_tool_call_args(args: str) -> str: elif char == '}': brace_count -= 1 if brace_count == 0 and current_obj.strip(): - # 完成了一个JSON对象 + # Completed one JSON object try: parsed = json.loads(current_obj.strip()) json_objects.append(parsed) @@ -230,14 +246,14 @@ def validate_and_fix_tool_call_args(args: str) -> str: if json_objects: return json.dumps(json_objects[0], ensure_ascii=False) - # 原有的修复逻辑 + # Original repair logic try: json.loads(args) return args except json.JSONDecodeError as e: - - - # 尝试修复常见的JSON问题 + + + # Try to fix common JSON problems original_args = args if not args.endswith('}') and args.count('{') > args.count('}'): args += '}' @@ -254,38 +270,38 @@ def validate_and_fix_tool_call_args(args: str) -> str: return '{}' class SSEConnectionManager: - """SSE 连接管理器，包含重连逻辑""" + """SSE connection manager, including reconnection logic""" def __init__(self, max_retries: int = 3, retry_delay: float = 1.0): self.max_retries = max_retries self.retry_delay = retry_delay async def stream_with_retry(self, stream_func, *args, **kwargs): - """带重连的流式处理""" + """Streaming processing with reconnection""" for attempt in range(self.max_retries + 1): try: async for chunk in stream_func(*args, **kwargs): yield chunk - break # 成功完成，退出重试循环 + break # Completed successfully, exit the retry loop except (httpx.TimeoutException, httpx.NetworkError) as e: if attempt < self.max_retries: - wait_time = self.retry_delay * (2 ** attempt) # 指数退避: 1s, 2s, 4s - logger.warning(f"连接失败，{wait_time}秒后重试 (第{attempt + 1}次): {e}") + wait_time = self.retry_delay * (2 ** attempt) # Exponential backoff: 1s, 2s, 4s + logger.warning(f"Connection failed, retrying in {wait_time}s (attempt {attempt + 1}): {e}") yield format_sse_error(f"Connection lost, retrying in {wait_time}s... (attempt {attempt + 1})", "connection_retry") await asyncio.sleep(wait_time) continue else: - logger.error(f"重连失败，已达到最大重试次数: {e}") + logger.error(f"Reconnection failed, maximum retry count reached: {e}") yield format_sse_error(f"Connection failed after {self.max_retries} retries: {str(e)}", "connection_failed") raise except Exception as e: - # 其他异常不重试，直接抛出 - logger.error(f"流式处理异常: {e}") + # Other exceptions are not retried; re-raise directly + logger.error(f"Streaming processing exception: {e}") yield format_sse_error(f"Stream error: {str(e)}", "stream_error") raise class StreamResponseAggregator: - """流式响应聚合器 - 修复多工具调用问题：使用工具调用ID作为键""" + """Streaming response aggregator - fixes multi-tool-call issues by using the tool call ID as the key""" def __init__(self): self.data = { @@ -297,14 +313,14 @@ def __init__(self): "usage": None, "system_fingerprint": None } - # 🔑 关键：使用工具调用ID作为键，因为index都是0会覆盖 + # 🔑 Key point: use the tool call ID as the key, because the index is always 0 and would overwrite self.tool_call_map = {} # key: tool_call_id, value: tool_call_data - self.tool_call_order = [] # 保持工具调用的接收顺序 - self.current_tool_id = None # 当前正在处理的工具调用ID - + self.tool_call_order = [] # Preserve the order in which tool calls are received + self.current_tool_id = None # ID of the tool call currently being processed + def process_chunk(self, obj: Dict[str, Any]): - """处理单个响应块""" - # 聚合基本信息 + """Process a single response chunk""" + # Aggregate basic information self.data["id"] = self.data["id"] or obj.get('id') self.data["model"] = self.data["model"] or obj.get('model') self.data["system_fingerprint"] = obj.get('system_fingerprint') or self.data["system_fingerprint"] @@ -322,22 +338,22 @@ def process_chunk(self, obj: Dict[str, Any]): delta = choice.get('delta', {}) - # 聚合内容 + # Aggregate content if delta.get('content'): self.data["content"] += delta.get('content') - - # 处理工具调用 + + # Handle tool calls if delta.get('tool_calls'): self._process_tool_calls(delta.get('tool_calls')) - + def _process_tool_calls(self, tool_calls: List[Dict[str, Any]]): - """处理工具调用 - 修复版：使用工具调用ID，正确处理分块传输""" + """Handle tool calls - fixed version: use the tool call ID and correctly handle chunked transfer""" for tc in tool_calls: tool_id = tc.get('id') - - # 如果有ID，这是一个新的工具调用 + + # If there is an ID, this is a new tool call if tool_id: - # 新工具调用 + # New tool call if tool_id not in self.tool_call_map: self.tool_call_map[tool_id] = { 'id': tool_id, @@ -349,12 +365,12 @@ def _process_tool_calls(self, tool_calls: List[Dict[str, Any]]): } self.tool_call_order.append(tool_id) self.current_tool_id = tool_id - logger.info(f"🔧 新工具调用: {tool_id}") + logger.info(f"🔧 New tool call: {tool_id}") else: - # 更新当前工具调用ID + # Update the current tool call ID self.current_tool_id = tool_id - - # 更新工具调用信息 + + # Update tool call information if tc.get('type'): self.tool_call_map[tool_id]['type'] = tc.get('type') @@ -364,7 +380,7 @@ def _process_tool_calls(self, tool_calls: List[Dict[str, Any]]): if func.get('arguments'): self.tool_call_map[tool_id]['function']['arguments'] += func.get('arguments') - # 如果没有ID，但有当前工具调用ID，这是增量数据 + # If there is no ID but there is a current tool call ID, this is incremental data elif self.current_tool_id and self.current_tool_id in self.tool_call_map: func = tc.get('function', {}) if func.get('name'): @@ -373,27 +389,27 @@ def _process_tool_calls(self, tool_calls: List[Dict[str, Any]]): self.tool_call_map[self.current_tool_id]['function']['arguments'] += func.get('arguments') else: - # 没有ID且没有当前工具调用，跳过 - logger.warning("⚠️ 工具调用缺少ID且无当前工具调用上下文，跳过处理") - + # No ID and no current tool call, skip + logger.warning("⚠️ Tool call is missing an ID and there is no current tool call context, skipping") + def finalize(self) -> Dict[str, Any]: - """完成聚合并返回最终响应""" - # 按接收顺序构建工具调用列表 + """Finish aggregation and return the final response""" + # Build the tool call list in the order received if self.tool_call_map: self.data["tool_calls"] = [] for tool_id in self.tool_call_order: if tool_id in self.tool_call_map: tc = self.tool_call_map[tool_id] - # 验证和修复工具调用参数 + # Validate and repair the tool call arguments tc['function']['arguments'] = validate_and_fix_tool_call_args( tc['function']['arguments'] ) self.data["tool_calls"].append(tc) - logger.info(f"📋 工具调用: {tool_id} - {tc['function']['name']}") - - logger.info(f"✅ 成功聚合 {len(self.data['tool_calls'])} 个工具调用") - - # 构建最终响应 + logger.info(f"📋 Tool call: {tool_id} - {tc['function']['name']}") + + logger.info(f"✅ Successfully aggregated {len(self.data['tool_calls'])} tool call(s)") + + # Build the final response final_message = {"role": "assistant", "content": self.data["content"]} if self.data["tool_calls"]: final_message["tool_calls"] = self.data["tool_calls"] @@ -422,106 +438,156 @@ def finalize(self) -> Dict[str, Any]: return final_response +class UpstreamAttemptError(Exception): + """Safe error that does not carry the upstream response body.""" + + def __init__(self, kind: str, status_code: int, code: str): + super().__init__(code) + self.kind = kind + self.status_code = status_code + self.code = code + + class CodeBuddyStreamService: - """CodeBuddy 流式服务类 - 职责分离，使用连接池优化""" - - def __init__(self): - self.connection_manager = SSEConnectionManager(max_retries=3, retry_delay=1.0) - - def _handle_api_error(self, status_code: int, error_msg: str) -> None: - """统一的API错误处理 - 直接抛出异常""" - logger.error(f"CodeBuddy API错误: {status_code} - {error_msg}") - + """CodeBuddy streaming service; each method performs exactly one upstream attempt.""" + + @staticmethod + def _classify_status(status_code: int) -> UpstreamAttemptError: if status_code == 401: - raise HTTPException(status_code=401, detail="CodeBuddy API authentication failed") - elif status_code == 429: - raise HTTPException(status_code=429, detail="CodeBuddy API rate limit exceeded") - elif status_code >= 500: - raise HTTPException(status_code=502, detail="CodeBuddy API server error") - else: - raise HTTPException(status_code=status_code, detail=f"CodeBuddy API error: {error_msg}") - - async def handle_stream_response(self, payload: Dict[str, Any], headers: Dict[str, str]) -> StreamingResponse: - """处理流式响应 - 使用OpenAI兼容性转换器修复格式问题""" - async def stream_core(): - client = await get_http_client() - async with client.stream("POST", get_codebuddy_api_url(), json=payload, headers=headers) as response: - if response.status_code != 200: - error_text = await response.aread() - error_msg = error_text.decode('utf-8', errors='ignore') - yield format_sse_error(f"CodeBuddy API error: {response.status_code} - {error_msg}", "api_error") - return - - buffer = "" - tool_call_index_map = {} # 用于跟踪工具调用ID到index的映射 - - - async for chunk in response.aiter_text(chunk_size=8192): - if not chunk: + return UpstreamAttemptError("invalid", 401, "upstream_authentication_failed") + if status_code in {403, 429}: + return UpstreamAttemptError("cooldown", status_code, "upstream_temporarily_unavailable") + if status_code >= 500: + return UpstreamAttemptError("transient", status_code, "upstream_server_error") + return UpstreamAttemptError("fatal", status_code, "upstream_request_rejected") + + async def open_stream_response( + self, + payload: Dict[str, Any], + headers: Dict[str, str], + key_id: Optional[str] = None, + ) -> StreamingResponse: + """Establish the connection and verify the upstream status before returning a StreamingResponse.""" + client = await get_http_client() + request = client.build_request( + "POST", get_codebuddy_api_url(), json=payload, headers=headers + ) + try: + response = await client.send(request, stream=True) + except httpx.TimeoutException as exc: + raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc + except httpx.RequestError as exc: + raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc + + if response.status_code != 200: + status_code = response.status_code + await response.aclose() + raise self._classify_status(status_code) + + # Do not expose upstream response headers; only this proxy's fixed SSE headers + # are sent downstream after the first chunk is safely available. + + async def converted_chunks(): + buffer = "" + tool_call_index_map = {} + async for chunk in response.aiter_text(chunk_size=8192): + if not chunk: + continue + buffer += chunk + + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + if not line.strip() or line.startswith(':'): continue - - buffer += chunk - - # 处理完整的SSE行 - while '\n' in buffer: - line, buffer = buffer.split('\n', 1) - - # 跳过空行和注释行 - if not line.strip() or line.startswith(':'): - continue - - # 检查是否结束 - if '[DONE]' in line: - - yield line + '\n' - return - - # 解析SSE数据 - chunk_data = parse_sse_line(line) - if chunk_data: - # 🔑 关键修改：使用OpenAI兼容性转换器 - converted_chunk = OpenAICompatibilityConverter.convert_sse_chunk_to_openai_format( - chunk_data, tool_call_index_map - ) - - # 重新格式化为SSE格式并发送 - converted_line = f"data: {json.dumps(converted_chunk, ensure_ascii=False)}" - yield converted_line + '\n' - else: - # 非数据行直接转发 - yield line + '\n' - - # 处理缓冲区中剩余的数据 - if buffer.strip(): - chunk_data = parse_sse_line(buffer.strip()) + if '[DONE]' in line: + yield line + '\n' + return + + chunk_data = parse_sse_line(line) if chunk_data: converted_chunk = OpenAICompatibilityConverter.convert_sse_chunk_to_openai_format( chunk_data, tool_call_index_map ) - converted_line = f"data: {json.dumps(converted_chunk, ensure_ascii=False)}" - yield converted_line + '\n' - else: - yield buffer + '\n' - - async def stream_with_retry(): - async for chunk in self.connection_manager.stream_with_retry(stream_core): - yield chunk - - return StreamingResponse(stream_with_retry(), media_type="text/event-stream", headers=SSE_HEADERS) - - async def handle_non_stream_response(self, payload: Dict[str, Any], headers: Dict[str, str]) -> Dict[str, Any]: - """处理非流式响应 - 使用修复后的聚合器，支持多工具调用""" + line = f"data: {json.dumps(converted_chunk, ensure_ascii=False)}" + yield line + '\n' + + if buffer.strip(): + chunk_data = parse_sse_line(buffer.strip()) + if chunk_data: + converted_chunk = OpenAICompatibilityConverter.convert_sse_chunk_to_openai_format( + chunk_data, tool_call_index_map + ) + buffer = f"data: {json.dumps(converted_chunk, ensure_ascii=False)}" + yield buffer + '\n' + + stream = converted_chunks() + try: + first_chunk = await anext(stream) + except StopAsyncIteration: + first_chunk = None + except httpx.TimeoutException as exc: + await response.aclose() + raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc + except httpx.RequestError as exc: + await response.aclose() + raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc + except Exception as exc: + await response.aclose() + raise UpstreamAttemptError("fatal", 502, "upstream_response_invalid") from exc + + async def stream_core(): + try: + if first_chunk is not None: + yield first_chunk + async for chunk in stream: + yield chunk + except httpx.RequestError: + logger.warning("CodeBuddy upstream stream interrupted") + if key_id is not None: + await codebuddy_api_key_manager.mark_transient_error( + key_id, "stream_interrupted" + ) + yield format_sse_error( + "Upstream stream interrupted", "upstream_stream_error" + ) + except Exception: + logger.error("Unexpected CodeBuddy stream processing error") + if key_id is not None: + await codebuddy_api_key_manager.mark_transient_error( + key_id, "stream_processing_error" + ) + yield format_sse_error( + "Upstream stream interrupted", "upstream_stream_error" + ) + finally: + await response.aclose() + + return StreamingResponse( + stream_core(), media_type="text/event-stream", headers=SSE_HEADERS + ) + + async def handle_non_stream_response( + self, payload: Dict[str, Any], headers: Dict[str, str] + ) -> Dict[str, Any]: + """Perform a single non-streaming upstream request and aggregate the SSE response.""" try: client = await get_http_client() - response = await client.post(get_codebuddy_api_url(), json=payload, headers=headers) - - if response.status_code != 200: - error_msg = response.text - self._handle_api_error(response.status_code, error_msg) - + response = await client.post( + get_codebuddy_api_url(), json=payload, headers=headers + ) + except httpx.TimeoutException as exc: + raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc + except httpx.RequestError as exc: + raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc + + if response.status_code != 200: + status_code = response.status_code + await response.aclose() + raise self._classify_status(status_code) + + try: aggregator = StreamResponseAggregator() buffer = "" - async for chunk in response.aiter_text(): if not chunk: continue @@ -531,42 +597,35 @@ async def handle_non_stream_response(self, payload: Dict[str, Any], headers: Dic obj = parse_sse_line(line) if obj: aggregator.process_chunk(obj) - + if buffer.strip(): obj = parse_sse_line(buffer.strip()) if obj: aggregator.process_chunk(obj) - return aggregator.finalize() - - except httpx.TimeoutException: - logger.error("CodeBuddy API 超时") - raise HTTPException(status_code=504, detail="CodeBuddy API timeout") - except httpx.NetworkError as e: - logger.error(f"网络错误: {e}") - raise HTTPException(status_code=502, detail=f"Network error: {str(e)}") - except HTTPException: + except httpx.RequestError as exc: + raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc + except UpstreamAttemptError: raise - except Exception as e: - logger.error(f"请求异常: {e}") - raise HTTPException(status_code=500, detail=f"Request error: {str(e)}") + except Exception as exc: + raise UpstreamAttemptError("fatal", 502, "upstream_response_invalid") from exc class RequestProcessor: - """请求预处理器 - 线程安全的请求处理""" - + """Request preprocessor - thread-safe request handling""" + @staticmethod def prepare_payload(request_body: Dict[str, Any]) -> Dict[str, Any]: - """准备请求载荷""" + """Prepare the request payload""" payload = request_body.copy() - payload["stream"] = True # CodeBuddy 只支持流式请求 - - # 处理消息长度要求：CodeBuddy要求至少2条消息 + payload["stream"] = True # CodeBuddy only supports streaming requests + + # Handle the message count requirement: CodeBuddy requires at least 2 messages messages = payload.get("messages", []) if len(messages) == 1 and messages[0].get("role") == "user": system_msg = {"role": "system", "content": "You are a helpful assistant."} payload["messages"] = [system_msg] + messages - # 应用关键词替换 + # Apply keyword replacement for msg in payload.get("messages", []): if msg.get("role") == "system": msg["content"] = apply_keyword_replacement_to_system_message(msg.get("content")) @@ -575,7 +634,7 @@ def prepare_payload(request_body: Dict[str, Any]) -> Dict[str, Any]: @staticmethod def validate_request(request_body: Dict[str, Any]) -> None: - """验证请求参数""" + """Validate request parameters""" if not isinstance(request_body, dict): raise HTTPException(status_code=400, detail="Request body must be a JSON object") @@ -586,32 +645,104 @@ def validate_request(request_body: Dict[str, Any]) -> None: if not messages: raise HTTPException(status_code=400, detail="At least one message is required") - # 验证消息格式 + # Validate message format for i, msg in enumerate(messages): if not isinstance(msg, dict): raise HTTPException(status_code=400, detail=f"Message {i} must be an object") if "role" not in msg or "content" not in msg: raise HTTPException(status_code=400, detail=f"Message {i} must have 'role' and 'content' fields") +@dataclass(repr=False) +class ResolvedCredential: + """Unified upstream credential; the raw token must not be logged or serialized.""" + + bearer_token: str + user_id: Optional[str] + source: str + key_id: Optional[str] = None + + class CredentialManager: - """凭证管理器 - 线程安全的凭证获取""" - + """Resolve legacy credentials or TXT API keys according to the auth mode.""" + @staticmethod - def get_valid_credential() -> Dict[str, Any]: - """获取有效凭证，包含错误处理""" + def get_legacy_credential() -> Optional[ResolvedCredential]: try: credential = codebuddy_token_manager.get_next_credential() - if not credential: - raise HTTPException(status_code=401, detail="没有可用的CodeBuddy凭证") - - bearer_token = credential.get('bearer_token') - if not bearer_token: - raise HTTPException(status_code=401, detail="无效的CodeBuddy凭证") - - return credential - except Exception as e: - logger.error(f"获取凭证失败: {e}") - raise HTTPException(status_code=401, detail="凭证获取失败") + except Exception: + logger.exception("Failed to select a legacy CodeBuddy credential") + return None + if not credential or not credential.get("bearer_token"): + return None + return ResolvedCredential( + bearer_token=credential["bearer_token"], + user_id=credential.get("user_id"), + source="credentials", + ) + + @staticmethod + async def get_api_key(excluded_ids: Set[str]) -> Optional[ResolvedCredential]: + selection = await codebuddy_api_key_manager.acquire(excluded_ids) + if selection is None: + return None + return ResolvedCredential( + bearer_token=selection.key, + user_id=None, + source="api_key_file", + key_id=selection.key_id, + ) + + @staticmethod + async def resolve_source() -> tuple[str, int]: + from config import get_codebuddy_auth_mode + + mode = get_codebuddy_auth_mode() + if mode == "credentials": + return "credentials", 1 + + eligible = await codebuddy_api_key_manager.eligible_count() + if eligible > 0: + return "api_key_file", eligible + if mode == "auto": + return "credentials", 1 + + total = await codebuddy_api_key_manager.total_count() + if total == 0: + raise ApiKeyConfigurationError( + "No API keys are configured in CODEBUDDY_API_KEYS_FILE" + ) + raise ApiKeyConfigurationError("No API keys are currently available") + + +def openai_error_response( + message: str, error_type: str, code: str, status_code: int +) -> JSONResponse: + """Construct an OpenAI-compatible error that excludes sensitive upstream information.""" + return JSONResponse( + status_code=status_code, + content={ + "error": { + "message": message, + "type": error_type, + "code": code, + } + }, + ) + + +async def record_attempt_error(credential: ResolvedCredential, error: UpstreamAttemptError) -> None: + if credential.source != "api_key_file" or credential.key_id is None: + return + if error.kind == "invalid": + await codebuddy_api_key_manager.mark_invalid(credential.key_id) + elif error.kind == "cooldown": + await codebuddy_api_key_manager.mark_cooldown( + credential.key_id, error.status_code + ) + else: + await codebuddy_api_key_manager.mark_transient_error( + credential.key_id, error.code + ) # --- API Endpoints --- @@ -622,55 +753,161 @@ async def chat_completions( x_conversation_request_id: Optional[str] = Header(None, alias="X-Conversation-Request-ID"), x_conversation_message_id: Optional[str] = Header(None, alias="X-Conversation-Message-ID"), x_request_id: Optional[str] = Header(None, alias="X-Request-ID"), - _token: str = Depends(authenticate) + auth_context: ClientAuthContext = Depends(authenticate_inference) ): - """CodeBuddy V1 聊天完成API - 重构后的简洁版本""" + """CodeBuddy V1 chat completions API, supporting relay and per-request passthrough.""" + try: + request_body = await request.json() + except Exception: + return openai_error_response( + "Invalid JSON request body", "invalid_request_error", "invalid_json", 400 + ) + try: - # 解析和验证请求体 - try: - request_body = await request.json() - except Exception as e: - logger.error(f"解析请求体失败: {e}") - raise HTTPException(status_code=400, detail=f"Invalid JSON request body: {str(e)}") - - # 验证请求参数 RequestProcessor.validate_request(request_body) - - # 获取有效凭证 - credential = CredentialManager.get_valid_credential() - - # 生成请求头 + except HTTPException as exc: + return openai_error_response( + str(exc.detail), "invalid_request_error", "invalid_request", exc.status_code + ) + + passthrough_credential: Optional[ResolvedCredential] = None + if auth_context.mode == "passthrough": + if not auth_context.passthrough_key: + return openai_error_response( + "A Bearer API key is required", + "authentication_error", + "missing_api_key", + 401, + ) + source, max_attempts = "passthrough", 1 + passthrough_credential = ResolvedCredential( + bearer_token=auth_context.passthrough_key, + user_id=None, + source="passthrough", + ) + else: + try: + source, max_attempts = await CredentialManager.resolve_source() + except (ApiKeyConfigurationError, ValueError): + logger.error("CodeBuddy API key file authentication is not configured correctly") + return openai_error_response( + "Upstream API key file is empty, unavailable, or invalid", + "configuration_error", + "api_key_file_unavailable", + 503, + ) + + payload = RequestProcessor.prepare_payload(request_body) + usage_stats_manager.record_model_usage(payload.get("model", "unknown")) + service = CodeBuddyStreamService() + client_wants_stream = request_body.get("stream", False) + excluded_ids: Set[str] = set() + last_error: Optional[UpstreamAttemptError] = None + + for _attempt in range(max_attempts): + if source == "passthrough": + credential = passthrough_credential + elif source == "api_key_file": + credential = await CredentialManager.get_api_key(excluded_ids) + else: + credential = CredentialManager.get_legacy_credential() + + if credential is None: + break + if credential.key_id is not None: + excluded_ids.add(credential.key_id) + + try: + upstream_key_header = get_upstream_api_key_header() + except ValueError: + return openai_error_response( + "Upstream API key header mode is misconfigured", + "configuration_error", + "upstream_header_mode_invalid", + 500, + ) + headers = codebuddy_api_client.generate_codebuddy_headers( - bearer_token=credential.get('bearer_token'), - user_id=credential.get('user_id'), + bearer_token=credential.bearer_token, + user_id=credential.user_id, conversation_id=x_conversation_id, conversation_request_id=x_conversation_request_id, conversation_message_id=x_conversation_message_id, - request_id=x_request_id + request_id=x_request_id, + api_key_header=upstream_key_header, + ) + + try: + if client_wants_stream: + result = await service.open_stream_response( + payload, headers, credential.key_id + ) + else: + result = await service.handle_non_stream_response(payload, headers) + if credential.key_id is not None: + await codebuddy_api_key_manager.mark_success(credential.key_id) + return result + except UpstreamAttemptError as error: + last_error = error + await record_attempt_error(credential, error) + logger.warning( + "CodeBuddy upstream attempt failed: source=%s code=%s", + credential.source, + error.code, + ) + if source != "api_key_file" or error.kind == "fatal": + break + + if source == "credentials" and last_error is None: + return openai_error_response( + "No valid CodeBuddy credentials are available", + "authentication_error", + "credentials_unavailable", + 401, ) - - # 预处理请求 - payload = RequestProcessor.prepare_payload(request_body) - usage_stats_manager.record_model_usage(payload.get("model", "unknown")) - - # 使用服务类处理请求 - service = CodeBuddyStreamService() - client_wants_stream = request_body.get("stream", False) - - if client_wants_stream: - return await service.handle_stream_response(payload, headers) - else: - return await service.handle_non_stream_response(payload, headers) - - except HTTPException: - raise - except Exception as e: - logger.error(f"CodeBuddy V1 API错误: {e}") - raise HTTPException(status_code=500, detail=f"内部服务器错误: {str(e)}") + + if source == "passthrough" and last_error is not None: + if last_error.status_code == 401: + return openai_error_response( + "Upstream CodeBuddy rejected the supplied API key", + "authentication_error", + "upstream_api_key_rejected", + 401, + ) + error_type = ( + "rate_limit_error" + if last_error.status_code == 429 + else "permission_error" + if last_error.status_code == 403 + else "upstream_error" + ) + return openai_error_response( + "Upstream CodeBuddy request failed", + error_type, + last_error.code, + last_error.status_code, + ) + + status_code = ( + 502 + if source == "credentials" and last_error and last_error.status_code >= 500 + else last_error.status_code + if source == "credentials" and last_error + else 502 + ) + code = last_error.code if source == "credentials" and last_error else "upstream_keys_exhausted" + return openai_error_response( + "Upstream CodeBuddy request failed", + "upstream_error", + code, + status_code, + ) @router.get("/v1/models") -async def list_v1_models(_token: str = Depends(authenticate)): - """获取CodeBuddy V1模型列表""" +async def list_v1_models( + _auth_context: ClientAuthContext = Depends(authenticate_inference), +): + """Get the list of CodeBuddy V1 models""" try: return { "object": "list", @@ -683,12 +920,28 @@ async def list_v1_models(_token: str = Depends(authenticate)): } except Exception as e: - logger.error(f"获取V1模型列表错误: {e}") - raise HTTPException(status_code=500, detail="获取模型列表失败") + logger.error(f"Error getting V1 model list: {e}") + raise HTTPException(status_code=500, detail="Failed to get model list") + +@router.get("/v1/api-keys/status", summary="Get upstream API key pool status") +async def get_api_keys_status(_token: str = Depends(authenticate_admin)): + """Return a safe status containing only masked keys and runtime statistics.""" + from config import get_codebuddy_auth_mode + + status = await codebuddy_api_key_manager.get_status() + return {"auth_mode": get_codebuddy_auth_mode(), **status} + + +@router.post("/v1/api-keys/reload", summary="Reload upstream API keys") +async def reload_api_keys(_token: str = Depends(authenticate_admin)): + """Force a re-read of the TXT API key file without restarting the service.""" + result = await codebuddy_api_key_manager.reload() + return {"message": "API key file reloaded", **result} + @router.get("/v1/credentials", summary="List all available credentials") -async def list_credentials(_token: str = Depends(authenticate)): - """列出所有可用凭证的详细信息，包括过期状态""" +async def list_credentials(_token: str = Depends(authenticate_admin)): + """List detailed information for all available credentials, including expiration status""" try: credentials_info = codebuddy_token_manager.get_credentials_info() safe_credentials = [] @@ -698,7 +951,7 @@ async def list_credentials(_token: str = Depends(authenticate)): for info in credentials_info: bearer_token = credentials[info['index']].get("bearer_token", "") if info['index'] < len(credentials) else "" - # 格式化时间显示 + # Format the time display if info['time_remaining'] is not None and info['time_remaining'] > 0: days, remainder = divmod(info['time_remaining'], 86400) hours, remainder = divmod(remainder, 3600) @@ -708,7 +961,7 @@ async def list_credentials(_token: str = Depends(authenticate)): time_remaining_str = "Expired" if info['time_remaining'] is not None else "Unknown" safe_credentials.append({ - **info, # 展开所有原始信息 + **info, # Expand all original info "time_remaining_str": time_remaining_str, "has_token": bool(bearer_token), "token_preview": f"{bearer_token[:10]}...{bearer_token[-4:]}" if len(bearer_token) > 14 else "Invalid Token" @@ -717,16 +970,16 @@ async def list_credentials(_token: str = Depends(authenticate)): return {"credentials": safe_credentials} except Exception as e: - logger.error(f"获取凭证列表失败: {e}") + logger.error(f"Failed to get credential list: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.post("/v1/credentials", summary="Add a new credential") async def add_credential( request: Request, - _token: str = Depends(authenticate) + _token: str = Depends(authenticate_admin) ): - """添加一个新的认证凭证""" + """Add a new authentication credential""" try: data = await request.json() if not data.get("bearer_token"): @@ -743,16 +996,16 @@ async def add_credential( return {"message": "Credential added successfully"} except Exception as e: - logger.error(f"添加凭证失败: {e}") + logger.error(f"Failed to add credential: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.post("/v1/credentials/select", summary="Manually select a credential") async def select_credential( request: Request, - _token: str = Depends(authenticate) + _token: str = Depends(authenticate_admin) ): - """手动选择指定的凭证""" + """Manually select the specified credential""" try: data = await request.json() index = data.get("index") @@ -765,25 +1018,25 @@ async def select_credential( return {"message": f"Credential #{index + 1} selected successfully"} except Exception as e: - logger.error(f"选择凭证失败: {e}") + logger.error(f"Failed to select credential: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.post("/v1/credentials/auto", summary="Resume automatic credential rotation") -async def resume_auto_rotation(_token: str = Depends(authenticate)): - """恢复自动凭证轮换""" +async def resume_auto_rotation(_token: str = Depends(authenticate_admin)): + """Resume automatic credential rotation""" try: codebuddy_token_manager.clear_manual_selection() return {"message": "Resumed automatic credential rotation"} except Exception as e: - logger.error(f"恢复自动轮换失败: {e}") + logger.error(f"Failed to resume automatic rotation: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.post("/v1/credentials/toggle-rotation", summary="Toggle automatic credential rotation") -async def toggle_auto_rotation(_token: str = Depends(authenticate)): - """切换自动轮换开关""" +async def toggle_auto_rotation(_token: str = Depends(authenticate_admin)): + """Toggle the automatic rotation switch""" try: is_enabled = codebuddy_token_manager.toggle_auto_rotation() status = "enabled" if is_enabled else "disabled" @@ -794,25 +1047,25 @@ async def toggle_auto_rotation(_token: str = Depends(authenticate)): } except Exception as e: - logger.error(f"切换自动轮换失败: {e}") + logger.error(f"Failed to toggle automatic rotation: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.get("/v1/credentials/current", summary="Get current credential info") -async def get_current_credential(_token: str = Depends(authenticate)): - """获取当前使用的凭证信息""" +async def get_current_credential(_token: str = Depends(authenticate_admin)): + """Get information about the currently used credential""" try: info = codebuddy_token_manager.get_current_credential_info() return info except Exception as e: - logger.error(f"获取当前凭证信息失败: {e}") + logger.error(f"Failed to get current credential info: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.post("/v1/credentials/delete", summary="Delete a credential by index") -async def delete_credential(request: Request, _token: str = Depends(authenticate)): - """删除一个凭证文件（通过索引）并从列表中移除""" +async def delete_credential(request: Request, _token: str = Depends(authenticate_admin)): + """Delete a credential file (by index) and remove it from the list""" try: data = await request.json() index = data.get("index") @@ -827,5 +1080,5 @@ async def delete_credential(request: Request, _token: str = Depends(authenticate except HTTPException: raise except Exception as e: - logger.error(f"删除凭证失败: {e}") + logger.error(f"Failed to delete credential: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/src/codebuddy_token_manager.py b/src/codebuddy_token_manager.py index 4de2bf0..a16c9ed 100644 --- a/src/codebuddy_token_manager.py +++ b/src/codebuddy_token_manager.py @@ -1,5 +1,5 @@ """ -CodeBuddy Token Manager - 管理CodeBuddy认证token +CodeBuddy Token Manager - Manages CodeBuddy authentication tokens """ import os import glob @@ -13,7 +13,7 @@ class CodeBuddyTokenManager: - """CodeBuddy Token管理器""" + """CodeBuddy token manager.""" def __init__(self, creds_dir=None): if creds_dir is None: @@ -25,13 +25,13 @@ def __init__(self, creds_dir=None): self.credentials = [] self.current_index = 0 # Start from the first credential self.usage_count = 0 # Counter for the current credential usage - self.manual_selected_index = None # 手动选择的凭证索引 - self.auto_rotation_enabled = True # 自动轮换开关，默认开启 + self.manual_selected_index = None # Manually selected credential index. + self.auto_rotation_enabled = True # Automatic rotation is enabled by default. self.load_all_tokens() - self.load_state() # 加载保存的状态 + self.load_state() # Load the saved state. def load_all_tokens(self): - """加载所有token文件""" + """Load all token files.""" self.credentials = [] self.current_index = -1 @@ -61,16 +61,16 @@ def load_all_tokens(self): logger.info(f"Loaded a total of {len(self.credentials)} CodeBuddy credentials.") def load_state(self): - """加载管理器状态""" + """Load manager state.""" try: if os.path.exists(self.state_file): with open(self.state_file, 'r', encoding='utf-8') as f: state = json.load(f) - # 恢复状态，但要验证索引是否还有效 + # Restore state after validating the saved index. saved_manual_index = state.get('manual_selected_index') if saved_manual_index is not None and 0 <= saved_manual_index < len(self.credentials): - # 验证凭证文件是否还存在 + # Verify that the credential file still exists. if saved_manual_index < len(self.credentials): saved_filename = state.get('manual_selected_filename') current_filename = os.path.basename(self.credentials[saved_manual_index]['file_path']) @@ -81,10 +81,10 @@ def load_state(self): else: logger.warning(f"Saved credential filename mismatch, ignoring saved selection") - # 恢复自动轮换状态 + # Restore automatic rotation state. self.auto_rotation_enabled = state.get('auto_rotation_enabled', True) - # 恢复当前索引（如果没有手动选择的话） + # Restore the current index when no manual selection exists. if self.manual_selected_index is None: saved_current_index = state.get('current_index', 0) if 0 <= saved_current_index < len(self.credentials): @@ -95,9 +95,9 @@ def load_state(self): logger.warning(f"Failed to load manager state: {e}") def save_state(self): - """保存管理器状态""" + """Save manager state.""" try: - # 确保目录存在 + # Ensure the directory exists. if not os.path.exists(self.creds_dir): os.makedirs(self.creds_dir) @@ -109,7 +109,7 @@ def save_state(self): 'saved_at': int(time.time()) } - # 如果有手动选择，保存文件名用于验证 + # Save the manually selected filename for validation. if self.manual_selected_index is not None and 0 <= self.manual_selected_index < len(self.credentials): state['manual_selected_filename'] = os.path.basename( self.credentials[self.manual_selected_index]['file_path'] @@ -123,20 +123,20 @@ def save_state(self): logger.error(f"Failed to save manager state: {e}") def is_token_expired(self, credential_data: Dict) -> bool: - """检查token是否过期""" + """Check whether a token is expired.""" try: created_at = credential_data.get('created_at') expires_in = credential_data.get('expires_in') if not created_at or not expires_in: - # 如果没有过期信息，假设未过期（向后兼容） + # Assume the token is valid when expiry data is absent for backward compatibility. return False current_time = int(time.time()) expiry_time = created_at + expires_in - # 提前5分钟认为过期，留出刷新时间 - buffer_time = 300 # 5分钟 + # Treat the token as expired five minutes early to allow refresh time. + buffer_time = 300 # Five minutes. is_expired = current_time >= (expiry_time - buffer_time) if is_expired: @@ -149,13 +149,13 @@ def is_token_expired(self, credential_data: Dict) -> bool: return False def get_next_credential(self) -> Optional[Dict]: - """获取下一个可用的凭证，根据轮换策略，并检查过期状态""" + """Return the next available credential according to rotation and expiry state.""" from config import get_rotation_count if not self.credentials: return None - # 过滤掉过期的凭证 + # Filter out expired credentials. valid_credentials = [] for i, cred in enumerate(self.credentials): if not self.is_token_expired(cred['data']): @@ -168,7 +168,7 @@ def get_next_credential(self) -> Optional[Dict]: logger.error("No valid (non-expired) credentials available") return None - # 如果当前索引无效或指向过期凭证，重置到第一个有效凭证 + # Reset to the first valid credential when the current index is invalid or expired. current_valid_indices = [i for i, _ in valid_credentials] if self.current_index not in current_valid_indices: self.current_index = current_valid_indices[0] @@ -177,7 +177,7 @@ def get_next_credential(self) -> Optional[Dict]: rotation_count = get_rotation_count() - # 如果有手动选择的凭证，优先使用（如果未过期） + # Prefer a manually selected credential when it is not expired. if self.manual_selected_index is not None and 0 <= self.manual_selected_index < len(self.credentials): manual_cred = self.credentials[self.manual_selected_index] if not self.is_token_expired(manual_cred['data']): @@ -189,7 +189,7 @@ def get_next_credential(self) -> Optional[Dict]: logger.warning("Manually selected credential is expired, falling back to automatic rotation") self.manual_selected_index = None - # 找到当前索引在有效凭证中的位置 + # Find the current index among valid credentials. try: current_valid_position = current_valid_indices.index(self.current_index) except ValueError: @@ -197,11 +197,11 @@ def get_next_credential(self) -> Optional[Dict]: self.current_index = current_valid_indices[0] self.usage_count = 0 - # 检查是否需要轮换：需要同时满足自动轮换开启 且 轮换次数大于0 + # Rotate only when automatic rotation is enabled and the rotation count is positive. should_rotate = self.auto_rotation_enabled and rotation_count > 0 if not should_rotate: - # 不轮换：固定使用当前凭证 + # Keep using the current credential when rotation is disabled. credential = self.credentials[self.current_index] credential_filename = os.path.basename(credential['file_path']) usage_stats_manager.record_credential_usage(credential_filename) @@ -211,12 +211,12 @@ def get_next_credential(self) -> Optional[Dict]: logger.info(f"Using fixed credential (auto rotation disabled): {credential_filename}") return credential['data'] - # 自动轮换逻辑：当开关开启且轮换次数>0时 + # Automatic rotation logic for a positive rotation count. if self.usage_count >= rotation_count: - # 轮换到下一个有效凭证 + # Rotate to the next valid credential. next_valid_position = (current_valid_position + 1) % len(valid_credentials) self.current_index = current_valid_indices[next_valid_position] - self.usage_count = 0 # 重置计数器 + self.usage_count = 0 # Reset the counter. logger.info("Credential rotation triggered.") credential = self.credentials[self.current_index] @@ -233,17 +233,17 @@ def get_next_credential(self) -> Optional[Dict]: return credential['data'] def get_all_credentials(self) -> List[Dict]: - """获取所有凭证""" + """Return all credentials.""" return [cred['data'] for cred in self.credentials] def get_credentials_info(self) -> List[Dict]: - """获取所有凭证的详细信息，包括过期状态""" + """Return credential details, including expiry state.""" credentials_info = [] for i, cred in enumerate(self.credentials): data = cred['data'] filename = os.path.basename(cred['file_path']) - # 计算过期信息 + # Calculate expiry information. is_expired = self.is_token_expired(data) expires_at = None time_remaining = None @@ -252,7 +252,7 @@ def get_credentials_info(self) -> List[Dict]: expires_at = data['created_at'] + data['expires_in'] time_remaining = expires_at - int(time.time()) - # 提取用户信息 + # Extract user information. user_info = data.get('user_info', {}) info = { @@ -279,7 +279,7 @@ def get_credentials_info(self) -> List[Dict]: return credentials_info def add_credential(self, bearer_token: str, user_id: str = None, filename: str = None) -> bool: - """添加新的凭证（简化版本，向后兼容）""" + """Add a credential using the simplified backward-compatible format.""" if not filename: filename = f"codebuddy_token_{len(self.credentials) + 1}.json" @@ -295,7 +295,7 @@ def add_credential(self, bearer_token: str, user_id: str = None, filename: str = return self.add_credential_with_data(credential_data, filename) def add_credential_with_data(self, credential_data: Dict[str, Any], filename: str = None) -> bool: - """添加新的凭证（完整数据版本）""" + """Add a credential using the complete data format.""" if not filename: user_id = credential_data.get('user_id', 'unknown') timestamp = credential_data.get('created_at', int(time.time())) @@ -307,12 +307,12 @@ def add_credential_with_data(self, credential_data: Dict[str, Any], filename: st file_path = os.path.join(self.creds_dir, filename) - # 确保必要字段存在 + # Ensure required fields exist. if 'created_at' not in credential_data: credential_data['created_at'] = int(time.time()) try: - # 确保目录存在 + # Ensure the directory exists. if not os.path.exists(self.creds_dir): os.makedirs(self.creds_dir) @@ -320,14 +320,14 @@ def add_credential_with_data(self, credential_data: Dict[str, Any], filename: st json.dump(credential_data, f, indent=4, ensure_ascii=False) logger.info(f"Added new credential: {filename}") - self.load_all_tokens() # 重新加载 + self.load_all_tokens() # Reload credentials. return True except Exception as e: logger.error(f"Failed to save credential: {e}") return False def delete_credential_by_index(self, index: int) -> bool: - """删除指定索引的凭证文件，并重新加载列表""" + """Delete the credential file at an index and reload the list.""" try: if not (0 <= index < len(self.credentials)): logger.error(f"Invalid credential index for deletion: {index}") @@ -342,9 +342,9 @@ def delete_credential_by_index(self, index: int) -> bool: else: logger.warning(f"Credential file already missing: {filename}") - # 重新加载凭证列表，重置索引等状态 + # Reload credentials and reset related state. self.load_all_tokens() - # 清理手动选择（若已删除的索引影响手动选择状态） + # Clear manual selection when the deleted index was selected. if self.manual_selected_index is not None and self.manual_selected_index == index: self.manual_selected_index = None logger.info("Cleared manual selection because deleted credential was selected") @@ -354,44 +354,44 @@ def delete_credential_by_index(self, index: int) -> bool: return False def set_manual_credential(self, index: int) -> bool: - """手动选择指定索引的凭证""" + """Manually select the credential at an index.""" if 0 <= index < len(self.credentials): self.manual_selected_index = index - self.current_index = index # 更新当前索引 + self.current_index = index # Update the current index. credential_filename = os.path.basename(self.credentials[index]['file_path']) logger.info(f"Manually selected credential: {credential_filename} (index: {index})") - self.save_state() # 保存状态 + self.save_state() # Save state. return True else: logger.error(f"Invalid credential index: {index}") return False def clear_manual_selection(self): - """清除手动选择，恢复自动轮换""" + """Clear manual selection and resume automatic rotation.""" self.manual_selected_index = None logger.info("Cleared manual credential selection, resumed automatic rotation") - self.save_state() # 保存状态 + self.save_state() # Save state. def enable_auto_rotation(self): - """开启自动轮换""" + """Enable automatic rotation.""" self.auto_rotation_enabled = True logger.info("Auto rotation enabled") def disable_auto_rotation(self): - """关闭自动轮换""" + """Disable automatic rotation.""" self.auto_rotation_enabled = False logger.info("Auto rotation disabled") def toggle_auto_rotation(self): - """切换自动轮换状态""" + """Toggle automatic rotation.""" self.auto_rotation_enabled = not self.auto_rotation_enabled status = "enabled" if self.auto_rotation_enabled else "disabled" logger.info(f"Auto rotation toggled: {status}") - self.save_state() # 保存状态 + self.save_state() # Save state. return self.auto_rotation_enabled def get_current_credential_info(self) -> Dict: - """获取当前使用的凭证信息""" + """Return information about the current credential.""" from config import get_rotation_count if not self.credentials: @@ -408,7 +408,7 @@ def get_current_credential_info(self) -> Dict: "user_id": credential['data'].get('user_id', 'unknown') } elif not self.auto_rotation_enabled: - # 确保current_index有效 + # Ensure current_index is valid. if not (0 <= self.current_index < len(self.credentials)): self.current_index = 0 credential = self.credentials[self.current_index] @@ -421,7 +421,7 @@ def get_current_credential_info(self) -> Dict: "auto_rotation_enabled": False } elif rotation_count == 0: - # 确保current_index有效 + # Ensure current_index is valid. if not (0 <= self.current_index < len(self.credentials)): self.current_index = 0 credential = self.credentials[self.current_index] @@ -434,7 +434,7 @@ def get_current_credential_info(self) -> Dict: "auto_rotation_enabled": True } else: - # 确保current_index有效 + # Ensure current_index is valid. if not (0 <= self.current_index < len(self.credentials)): self.current_index = 0 credential = self.credentials[self.current_index] @@ -449,5 +449,5 @@ def get_current_credential_info(self) -> Dict: } -# 全局token管理器实例 +# Global token manager instance. codebuddy_token_manager = CodeBuddyTokenManager() diff --git a/src/frontend_router.py b/src/frontend_router.py index 8632257..05a7080 100644 --- a/src/frontend_router.py +++ b/src/frontend_router.py @@ -16,7 +16,7 @@ async def serve_frontend(): if not os.path.exists(HTML_FILE_PATH): return "Frontend file not found. Please ensure frontend/admin.html exists." - # 添加缓存控制头，防止浏览器缓存 + # Add cache-control headers to prevent browser caching. headers = { "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache, no-store, must-revalidate", @@ -31,7 +31,7 @@ async def serve_admin(): if not os.path.exists(HTML_FILE_PATH): return "Frontend file not found. Please ensure frontend/admin.html exists." - # 添加缓存控制头，防止浏览器缓存 + # Add cache-control headers to prevent browser caching. headers = { "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache, no-store, must-revalidate", diff --git a/src/keyword_replacer.py b/src/keyword_replacer.py index c07f041..2647ab8 100644 --- a/src/keyword_replacer.py +++ b/src/keyword_replacer.py @@ -1,6 +1,6 @@ """ -关键词替换工具模块 - 统一处理关键词替换逻辑 -防止CodeBuddy检测到竞争对手关键词 +Keyword replacement utility with centralized replacement rules. +Prevents CodeBuddy from detecting competitor-related keywords. """ import logging @@ -9,18 +9,18 @@ def apply_keyword_replacement(text: str) -> str: """ - 统一的关键词替换函数 + Apply the standard keyword replacements. Args: - text: 需要处理的文本内容 + text: Text content to process. Returns: - str: 替换后的文本内容 + str: The processed text. """ if not isinstance(text, str): return text - # 定义替换规则 + # Define replacement rules. replacements = { "Claude Code": "CodeBuddy Code", "Anthropic's official CLI for Claude": "Tencent's official CLI for CodeBuddy", @@ -31,11 +31,11 @@ def apply_keyword_replacement(text: str) -> str: original_text = text - # 应用所有替换规则 + # Apply every replacement rule. for old_keyword, new_keyword in replacements.items(): text = text.replace(old_keyword, new_keyword) - # 记录替换日志（仅在调试模式下） + # Record replacements at debug level only. if text != original_text: logger.debug(f"[KEYWORD_REPLACE] Applied keyword replacements, original length: {len(original_text)}, new length: {len(text)}") @@ -44,22 +44,22 @@ def apply_keyword_replacement(text: str) -> str: def apply_keyword_replacement_to_system_message(content) -> str: """ - 专门用于处理系统消息的关键词替换 - 支持字符串和复杂结构的content + Apply keyword replacements to system messages. + Supports both strings and structured content. Args: - content: 消息内容，可能是字符串或列表结构 + content: Message content, either a string or a list structure. Returns: - str: 处理后的内容 + str: The processed content. """ if isinstance(content, str): return apply_keyword_replacement(content) elif isinstance(content, list): - # 处理复杂结构的系统消息 + # Process structured system-message content. for item in content: if isinstance(item, dict) and item.get("type") == "text": item["text"] = apply_keyword_replacement(item.get("text", "")) return content else: - return content \ No newline at end of file + return content diff --git a/src/logging_utils.py b/src/logging_utils.py new file mode 100644 index 0000000..cd78d0e --- /dev/null +++ b/src/logging_utils.py @@ -0,0 +1,33 @@ +"""Defense-in-depth redaction for sensitive HTTP authentication headers.""" +import logging +import re +from typing import Any + +_SENSITIVE_HEADERS = re.compile( + r"(?i)(authorization(?:['\"]?\s*[:=]\s*['\"]?\s*)bearer\s+|x-api-key['\"]?\s*[:=]\s*['\"]?\s*)([^\s,;'\"}\]]+)" +) + + +def redact_sensitive(value: Any) -> str: + """Replace bearer and X-API-Key values without retaining them.""" + return _SENSITIVE_HEADERS.sub(r"\1[REDACTED]", str(value)) + + +class SensitiveHeaderFilter(logging.Filter): + """Redact sensitive values before a record reaches a handler.""" + + def filter(self, record: logging.LogRecord) -> bool: + try: + record.msg = redact_sensitive(record.getMessage()) + record.args = () + except Exception: + record.msg = "[log message redacted after formatting failure]" + record.args = () + return True + + +def install_sensitive_header_filter() -> None: + root = logging.getLogger() + for handler in root.handlers: + if not any(isinstance(item, SensitiveHeaderFilter) for item in handler.filters): + handler.addFilter(SensitiveHeaderFilter()) diff --git a/src/models.py b/src/models.py index e2df090..e281618 100644 --- a/src/models.py +++ b/src/models.py @@ -52,7 +52,7 @@ class ModelList(BaseModel): class CredentialInfo(BaseModel): - """凭证信息""" + """Credential information.""" index: int user_id: str created_at: int diff --git a/src/settings_router.py b/src/settings_router.py index 36a8e90..5fc71a2 100644 --- a/src/settings_router.py +++ b/src/settings_router.py @@ -1,37 +1,46 @@ """ -Settings Router - For loading and saving .env configurations +Settings Router - For loading and saving environment configurations. """ -import os import logging from fastapi import APIRouter, HTTPException, Depends from pydantic import BaseModel from typing import Dict, Any -from .auth import authenticate +from .auth import authenticate_admin from config import get_active_config, update_settings from .usage_stats_manager import usage_stats_manager logger = logging.getLogger(__name__) router = APIRouter() -# 中文标签映射 +# English labels displayed by the settings interface. SETTING_LABELS = { - "CODEBUDDY_HOST": "服务主机地址", - "CODEBUDDY_PORT": "服务端口", - "CODEBUDDY_PASSWORD": "API 服务访问密码", - "CODEBUDDY_API_ENDPOINT": "CodeBuddy 官方API端点", - "CODEBUDDY_CREDS_DIR": "凭证文件目录", - "CODEBUDDY_LOG_LEVEL": "日志级别", - "CODEBUDDY_MODELS": "可用模型列表 (逗号分隔)", - "CODEBUDDY_ROTATION_COUNT": "凭证轮换频率 (N次请求/凭证，设为0关闭轮换)" + "CODEBUDDY_HOST": "Service host address", + "CODEBUDDY_PORT": "Service port", + "CODEBUDDY_PASSWORD": "API service access password", + "CODEBUDDY_API_ENDPOINT": "Official CodeBuddy API endpoint", + "CODEBUDDY_CREDS_DIR": "Credential file directory", + "CODEBUDDY_LOG_LEVEL": "Log level", + "CODEBUDDY_MODELS": "Available model list (comma-separated)", + "CODEBUDDY_ROTATION_COUNT": "Credential rotation frequency (requests per credential; 0 disables rotation)", + "CODEBUDDY_AUTH_MODE": "Upstream authentication mode (auto/api_key_file/credentials)", + "CODEBUDDY_API_KEYS_FILE": "Upstream API key TXT file", + "CODEBUDDY_API_KEY_ROTATION": "API key rotation strategy", + "CODEBUDDY_API_KEY_RELOAD_INTERVAL": "API key file reload interval (seconds)", + "CODEBUDDY_API_KEY_COOLDOWN_SECONDS": "API key cooldown period (seconds)", + "CODEBUDDY_CLIENT_AUTH_MODE": "Client authentication mode (relay/passthrough/hybrid)", + "CODEBUDDY_ADMIN_PASSWORD": "Admin dashboard password (empty falls back to service password)", + "CODEBUDDY_UPSTREAM_API_KEY_HEADER": "Upstream key header (x-api-key/bearer/both)" } + class Settings(BaseModel): settings: Dict[str, Any] + @router.get("/settings", summary="Get all current active settings and labels") -async def get_settings(_token: str = Depends(authenticate)): - """Returns the current config and their Chinese labels.""" +async def get_settings(_token: str = Depends(authenticate_admin)): + """Return the active configuration and display labels.""" try: return { "settings": get_active_config(), @@ -41,19 +50,21 @@ async def get_settings(_token: str = Depends(authenticate)): logger.error(f"Error retrieving active config: {e}") raise HTTPException(status_code=500, detail="Could not retrieve settings.") + @router.post("/settings", summary="Save and hot-reload settings") -async def save_settings(new_settings: Settings, _token: str = Depends(authenticate)): - """Saves settings to config.json and hot-reloads them into memory.""" +async def save_settings(new_settings: Settings, _token: str = Depends(authenticate_admin)): + """Save settings and hot-reload them into memory.""" try: update_settings(new_settings.settings) - return {"message": "设置已保存并成功热加载！"} + return {"message": "Settings saved and hot-reloaded successfully."} except Exception as e: logger.error(f"Error saving settings: {e}") - raise HTTPException(status_code=500, detail="无法保存设置文件。") + raise HTTPException(status_code=500, detail="Could not save the settings file.") + @router.get("/stats", summary="Get usage statistics") -async def get_usage_stats(_token: str = Depends(authenticate)): - """Returns usage statistics for models and credentials.""" +async def get_usage_stats(_token: str = Depends(authenticate_admin)): + """Return usage statistics for models and credentials.""" try: return usage_stats_manager.get_stats() except Exception as e: diff --git a/tests/test_codebuddy_api_key_failover.py b/tests/test_codebuddy_api_key_failover.py new file mode 100644 index 0000000..3a65f26 --- /dev/null +++ b/tests/test_codebuddy_api_key_failover.py @@ -0,0 +1,264 @@ +import json + +import httpx +import pytest +from fastapi import FastAPI + +import config +from src import auth, codebuddy_router +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager + + +SERVER_PASSWORD = "relay-test-password" +KEY_A = "account-alpha-secret-0001" +KEY_B = "account-beta-secret-0002" + + +class BrokenStream(httpx.AsyncByteStream): + def __init__(self, fail_before_first=False): + self.fail_before_first = fail_before_first + + async def __aiter__(self): + if self.fail_before_first: + raise httpx.ReadError("upstream disconnected") + yield b'data: {"id":"chat-1","model":"auto-chat","choices":[{"delta":{"content":"first"}}]}\n\n' + raise httpx.ReadError("upstream disconnected") + + + +def sse_response(text="ok"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":"{text}"}},"finish_reason":"stop"}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +@pytest.fixture +def app(): + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + return application + + +async def setup_pool(monkeypatch, tmp_path, keys=(KEY_A, KEY_B), mode="api_key_file"): + path = tmp_path / "keys.txt" + path.write_text("\n".join(keys), encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), cooldown_seconds=60, reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + monkeypatch.setattr(config, "get_codebuddy_auth_mode", lambda: mode) + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "relay") + monkeypatch.setattr(auth, "get_server_password", lambda: SERVER_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: SERVER_PASSWORD) + return manager + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +async def post_chat(app, stream=False): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post( + "/codebuddy/v1/chat/completions", + headers={"Authorization": f"Bearer {SERVER_PASSWORD}"}, + json={ + "model": "auto-chat", + "messages": [{"role": "user", "content": "hello"}], + "stream": stream, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first_status", "expected_state"), + [(401, "invalid"), (429, "cooldown"), (500, "active")], +) +async def test_non_stream_failover(monkeypatch, tmp_path, app, first_status, expected_state): + manager = await setup_pool(monkeypatch, tmp_path) + used = [] + + def handler(request): + token = request.headers["authorization"].removeprefix("Bearer ") + used.append(token) + return httpx.Response(first_status) if token == KEY_A else sse_response() + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app) + await upstream.aclose() + + assert response.status_code == 200 + assert response.json()["choices"][0]["message"]["content"] == "ok" + assert used == [KEY_A, KEY_B] + status = await manager.get_status() + assert status["keys"][0]["status"] == expected_state + assert status["keys"][0]["error_count"] == 1 + assert status["keys"][1]["request_count"] == 1 + + +@pytest.mark.asyncio +async def test_timeout_failover_uses_next_key_once(monkeypatch, tmp_path, app): + manager = await setup_pool(monkeypatch, tmp_path) + used = [] + + def handler(request): + token = request.headers["authorization"].removeprefix("Bearer ") + used.append(token) + if token == KEY_A: + raise httpx.ConnectTimeout("upstream timeout", request=request) + return sse_response("recovered") + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app) + await upstream.aclose() + + assert response.status_code == 200 + assert used == [KEY_A, KEY_B] + status = await manager.get_status() + assert status["keys"][0]["error_count"] == 1 + assert status["keys"][1]["request_count"] == 1 + + +@pytest.mark.asyncio +async def test_stream_failover_happens_before_first_chunk(monkeypatch, tmp_path, app): + manager = await setup_pool(monkeypatch, tmp_path) + used = [] + + def handler(request): + token = request.headers["authorization"].removeprefix("Bearer ") + used.append(token) + return httpx.Response(401) if token == KEY_A else sse_response("streamed") + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "streamed" in response.text + assert used == [KEY_A, KEY_B] + status = await manager.get_status() + assert status["keys"][0]["status"] == "invalid" + + +@pytest.mark.asyncio +async def test_stream_failover_on_error_before_first_chunk(monkeypatch, tmp_path, app): + await setup_pool(monkeypatch, tmp_path) + used = [] + + def handler(request): + token = request.headers["authorization"].removeprefix("Bearer ") + used.append(token) + if token == KEY_A: + return httpx.Response(200, stream=BrokenStream(fail_before_first=True)) + return sse_response("recovered") + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "recovered" in response.text + assert used == [KEY_A, KEY_B] + + +@pytest.mark.asyncio +async def test_stream_does_not_switch_key_after_first_chunk(monkeypatch, tmp_path, app): + manager = await setup_pool(monkeypatch, tmp_path) + used = [] + + def handler(request): + token = request.headers["authorization"].removeprefix("Bearer ") + used.append(token) + return httpx.Response(200, stream=BrokenStream()) + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "first" in response.text + assert "upstream_stream_error" in response.text + assert used == [KEY_A] + status = await manager.get_status() + assert status["keys"][0]["error_count"] == 1 + assert status["keys"][1]["request_count"] == 0 + + +@pytest.mark.asyncio +async def test_all_keys_exhausted_is_openai_error_and_sanitized(monkeypatch, tmp_path, app, caplog): + await setup_pool(monkeypatch, tmp_path) + + def handler(request): + return httpx.Response(401, text=f"invalid key {request.headers['authorization']}") + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app) + await upstream.aclose() + + assert response.status_code == 502 + assert response.json()["error"]["code"] == "upstream_keys_exhausted" + combined = response.text + caplog.text + assert KEY_A not in combined + assert KEY_B not in combined + + +@pytest.mark.asyncio +async def test_auto_mode_falls_back_to_legacy_credentials(monkeypatch, tmp_path, app): + await setup_pool(monkeypatch, tmp_path, keys=(), mode="auto") + calls = [] + + def legacy_credential(): + calls.append(True) + return {"bearer_token": "legacy-bearer", "user_id": "legacy-user"} + + monkeypatch.setattr( + codebuddy_router.codebuddy_token_manager, + "get_next_credential", + legacy_credential, + ) + + def handler(request): + assert request.headers["authorization"] == "Bearer legacy-bearer" + return sse_response("legacy") + + upstream = install_upstream(monkeypatch, handler) + response = await post_chat(app) + await upstream.aclose() + + assert response.status_code == 200 + assert calls == [True] + assert response.json()["choices"][0]["message"]["content"] == "legacy" + + +@pytest.mark.asyncio +async def test_api_key_admin_endpoints_are_protected_and_sanitized(monkeypatch, tmp_path, app): + await setup_pool(monkeypatch, tmp_path) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + denied = await client.get("/codebuddy/v1/api-keys/status") + status = await client.get( + "/codebuddy/v1/api-keys/status", + headers={"Authorization": f"Bearer {SERVER_PASSWORD}"}, + ) + reload_response = await client.post( + "/codebuddy/v1/api-keys/reload", + headers={"Authorization": f"Bearer {SERVER_PASSWORD}"}, + ) + + assert denied.status_code in {401, 403} + assert status.status_code == 200 + assert reload_response.status_code == 200 + serialized = json.dumps(status.json()) + json.dumps(reload_response.json()) + assert KEY_A not in serialized + assert KEY_B not in serialized + assert status.json()["keys"][0]["masked_key"] == "acco...0001" diff --git a/tests/test_codebuddy_api_key_manager.py b/tests/test_codebuddy_api_key_manager.py new file mode 100644 index 0000000..b9fe7a0 --- /dev/null +++ b/tests/test_codebuddy_api_key_manager.py @@ -0,0 +1,138 @@ +import asyncio +import logging + +import pytest + +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager + + +@pytest.fixture +def key_file(tmp_path): + path = tmp_path / "keys.txt" + path.write_text( + "\n# comment\n alpha-secret-0001 \n\nbeta-secret-0002\nalpha-secret-0001\n", + encoding="utf-8", + ) + return path + + +@pytest.mark.asyncio +async def test_parsing_dedup_and_round_robin(key_file): + manager = CodeBuddyApiKeyManager(str(key_file), reload_interval=0) + result = await manager.reload() + + assert result == {"loaded": 2, "added": 2, "removed": 0, "available": True} + first = await manager.acquire() + second = await manager.acquire() + third = await manager.acquire() + + assert [first.key, second.key, third.key] == [ + "alpha-secret-0001", + "beta-secret-0002", + "alpha-secret-0001", + ] + + +@pytest.mark.asyncio +async def test_concurrent_selection_is_even(key_file): + manager = CodeBuddyApiKeyManager(str(key_file), reload_interval=0) + await manager.reload() + + selections = await asyncio.gather(*(manager.acquire() for _ in range(20))) + keys = [selection.key for selection in selections] + assert keys.count("alpha-secret-0001") == 10 + assert keys.count("beta-secret-0002") == 10 + + +@pytest.mark.asyncio +async def test_status_cooldown_invalid_and_no_raw_key(key_file, caplog): + now = [1_700_000_000.0] + manager = CodeBuddyApiKeyManager( + str(key_file), cooldown_seconds=10, reload_interval=0, clock=lambda: now[0] + ) + with caplog.at_level(logging.DEBUG): + await manager.reload() + first = await manager.acquire() + second = await manager.acquire() + await manager.mark_invalid(first.key_id) + await manager.mark_cooldown(second.key_id, 429) + status = await manager.get_status() + + serialized = str(status) + caplog.text + assert "alpha-secret-0001" not in serialized + assert "beta-secret-0002" not in serialized + assert [item["status"] for item in status["keys"]] == ["invalid", "cooldown"] + assert status["keys"][0]["request_count"] == 1 + assert status["keys"][0]["error_count"] == 1 + assert set(status["keys"][1]) == { + "masked_key", + "status", + "request_count", + "error_count", + "last_used_at", + "cooldown_until", + } + assert await manager.eligible_count() == 0 + + now[0] += 11 + assert await manager.eligible_count() == 1 + + +@pytest.mark.asyncio +async def test_hot_reload_preserves_state_and_removes_key(key_file): + manager = CodeBuddyApiKeyManager(str(key_file), reload_interval=0) + await manager.reload() + first = await manager.acquire() + await manager.mark_invalid(first.key_id) + + key_file.write_text("alpha-secret-0001\ngamma-secret-0003\n", encoding="utf-8") + result = await manager.reload() + status = await manager.get_status() + + assert result == {"loaded": 2, "added": 1, "removed": 1, "available": True} + assert [item["status"] for item in status["keys"]] == ["invalid", "active"] + selections = [await manager.acquire() for _ in range(4)] + assert {selection.key for selection in selections if selection} == {"gamma-secret-0003"} + assert "beta-secret-0002" not in {selection.key for selection in selections if selection} + + +@pytest.mark.asyncio +async def test_periodic_hot_reload_adds_key_without_restart(tmp_path): + path = tmp_path / "keys.txt" + path.write_text("first-secret-0001\n", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=1) + await manager.start_periodic_reload() + try: + path.write_text("first-secret-0001\nsecond-secret-0002\n", encoding="utf-8") + await asyncio.sleep(1.1) + assert await manager.total_count() == 2 + finally: + await manager.stop_periodic_reload() + + +@pytest.mark.asyncio +async def test_missing_file_keeps_pool_safe_without_exposing_path(tmp_path): + path = tmp_path / "missing-secret-name.txt" + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + result = await manager.reload() + status = await manager.get_status() + + assert result["loaded"] == 0 + assert result["reloaded"] is False + assert status["keys"] == [] + assert "missing-secret-name" not in str(status) + + +@pytest.mark.asyncio +async def test_temporary_read_error_does_not_drop_loaded_keys(tmp_path): + path = tmp_path / "keys.txt" + path.write_text("stable-secret-0001\n", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + path.unlink() + + result = await manager.reload() + selection = await manager.acquire() + + assert result["reloaded"] is False + assert selection.key == "stable-secret-0001" diff --git a/tests/test_codebuddy_client_auth.py b/tests/test_codebuddy_client_auth.py new file mode 100644 index 0000000..815a037 --- /dev/null +++ b/tests/test_codebuddy_client_auth.py @@ -0,0 +1,422 @@ +import json +import logging + +import httpx +import pytest +from fastapi import FastAPI + +import config +from src import auth, codebuddy_auth_router, codebuddy_router, settings_router +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager +from src.codebuddy_api_client import codebuddy_api_client +from src.logging_utils import SensitiveHeaderFilter, redact_sensitive + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" +KEY_B = "passthrough-account-beta-0002" + + +def sse_response(text="ok"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":"{text}"}},"finish_reason":"stop"}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +@pytest.fixture +def app(): + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def configure(monkeypatch, client_mode="passthrough", header_mode="bearer"): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: client_mode) + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: header_mode) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, + "record_model_usage", + lambda _model: None, + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +async def request(app, path, token=None, method="GET", json_body=None): + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.request(method, path, headers=headers, json=json_body) + + +async def chat(app, token, stream=False, content="hello"): + return await request( + app, + "/codebuddy/v1/chat/completions", + token, + "POST", + { + "model": "auto-chat", + "messages": [{"role": "user", "content": content}], + "stream": stream, + }, + ) + + +@pytest.mark.asyncio +async def test_passthrough_forwards_request_key(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = [] + + def handler(upstream_request): + seen.append(upstream_request.headers.get("authorization")) + return sse_response() + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 200 + assert seen == [f"Bearer {KEY_A}"] + assert await empty_pool.total_count() == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "authorization", "x_api_key"), + [ + ("bearer", f"Bearer {KEY_A}", None), + ("x-api-key", None, KEY_A), + ("both", f"Bearer {KEY_A}", KEY_A), + ], +) +async def test_upstream_header_mapping( + monkeypatch, app, empty_pool, mode, authorization, x_api_key +): + configure(monkeypatch, header_mode=mode) + seen = [] + + def handler(upstream_request): + seen.append( + ( + upstream_request.headers.get("authorization"), + upstream_request.headers.get("x-api-key"), + ) + ) + return sse_response() + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 200 + assert seen == [(authorization, x_api_key)] + + +@pytest.mark.asyncio +async def test_concurrent_passthrough_keys_do_not_mix(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(upstream_request): + body = json.loads(upstream_request.content) + content = body["messages"][-1]["content"] + seen[content] = upstream_request.headers["authorization"] + return sse_response(content) + + upstream = install_upstream(monkeypatch, handler) + first, second = await asyncio.gather( + chat(app, KEY_A, content="a"), chat(app, KEY_B, content="b") + ) + await upstream.aclose() + + assert first.status_code == second.status_code == 200 + assert seen == {"a": f"Bearer {KEY_A}", "b": f"Bearer {KEY_B}"} + + +@pytest.mark.asyncio +async def test_passthrough_stream_uses_one_request_key(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = [] + + def handler(upstream_request): + seen.append(upstream_request.headers["authorization"]) + return sse_response("streamed") + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "streamed" in response.text + assert seen == [f"Bearer {KEY_A}"] + assert await empty_pool.total_count() == 0 + + +@pytest.mark.asyncio +async def test_passthrough_stream_does_not_restart_after_first_chunk( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + attempts = [] + + class InterruptedStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'data: {"id":"chat-1","model":"auto-chat","choices":[{"delta":{"content":"first"}}]}\n\n' + raise httpx.ReadError("interrupted") + + def handler(upstream_request): + attempts.append(upstream_request.headers["authorization"]) + return httpx.Response(200, stream=InterruptedStream()) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "first" in response.text + assert "upstream_stream_error" in response.text + assert attempts == [f"Bearer {KEY_A}"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("upstream_status", "expected_status"), + [(401, 401), (403, 403), (429, 429), (500, 500), (503, 503)], +) +async def test_passthrough_preserves_status_without_failover( + monkeypatch, app, empty_pool, upstream_status, expected_status, caplog +): + configure(monkeypatch) + attempts = [] + + def handler(upstream_request): + attempts.append(upstream_request.headers["authorization"]) + return httpx.Response(upstream_status, text=f"rejected {KEY_A}") + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == expected_status + assert attempts == [f"Bearer {KEY_A}"] + assert KEY_A not in response.text + caplog.text + if upstream_status == 401: + assert response.json()["error"]["code"] == "upstream_api_key_rejected" + + +@pytest.mark.asyncio +async def test_passthrough_timeout_is_single_504_attempt(monkeypatch, app, empty_pool): + configure(monkeypatch) + attempts = [] + + def handler(upstream_request): + attempts.append(upstream_request.headers["authorization"]) + raise httpx.ConnectTimeout("timeout", request=upstream_request) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 504 + assert attempts == [f"Bearer {KEY_A}"] + assert KEY_A not in response.text + + +@pytest.mark.asyncio +async def test_hybrid_exact_match_selects_relay_or_passthrough( + monkeypatch, app, tmp_path +): + configure(monkeypatch, client_mode="hybrid") + path = tmp_path / "keys.txt" + pool_key = "pool-key-secret-0001" + path.write_text(pool_key, encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + monkeypatch.setattr(config, "get_codebuddy_auth_mode", lambda: "api_key_file") + seen = [] + + def handler(upstream_request): + seen.append(upstream_request.headers["authorization"]) + return sse_response() + + upstream = install_upstream(monkeypatch, handler) + relay_response = await chat(app, RELAY_PASSWORD) + passthrough_response = await chat(app, KEY_A) + await upstream.aclose() + + assert relay_response.status_code == passthrough_response.status_code == 200 + assert seen == [f"Bearer {pool_key}", f"Bearer {KEY_A}"] + + +@pytest.mark.asyncio +async def test_models_and_invalid_authorization(monkeypatch, app, empty_pool): + configure(monkeypatch) + valid = await request(app, "/codebuddy/v1/models", KEY_A) + missing = await request(app, "/codebuddy/v1/models") + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + malformed = await client.get( + "/codebuddy/v1/models", headers={"Authorization": f"Basic {KEY_A}"} + ) + + assert valid.status_code == 200 + assert missing.status_code == malformed.status_code == 401 + + +@pytest.mark.asyncio +async def test_admin_password_is_separate(monkeypatch, app, empty_pool): + configure(monkeypatch) + denied_key = await request(app, "/codebuddy/v1/api-keys/status", KEY_A) + denied_relay = await request( + app, "/codebuddy/v1/api-keys/status", RELAY_PASSWORD + ) + accepted = await request(app, "/codebuddy/v1/api-keys/status", ADMIN_PASSWORD) + + assert denied_key.status_code == denied_relay.status_code == 403 + assert accepted.status_code == 200 + + +@pytest.mark.asyncio +async def test_admin_password_falls_back_to_relay(monkeypatch, app, empty_pool): + configure(monkeypatch) + monkeypatch.setattr(auth, "get_admin_password", lambda: RELAY_PASSWORD) + response = await request( + app, "/codebuddy/v1/api-keys/status", RELAY_PASSWORD + ) + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_oauth_management_requires_admin_password( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + + async def fake_start(): + return {"success": True, "verification_uri_complete": "https://example.test"} + + monkeypatch.setattr(codebuddy_auth_router, "start_codebuddy_auth", fake_start) + denied = await request(app, "/codebuddy/auth/start", KEY_A) + accepted = await request(app, "/codebuddy/auth/start", ADMIN_PASSWORD) + + assert denied.status_code == 403 + assert accepted.status_code == 200 + + +@pytest.mark.asyncio +async def test_relay_still_validates_server_password(monkeypatch, app, tmp_path): + configure(monkeypatch, client_mode="relay") + path = tmp_path / "relay-keys.txt" + path.write_text("relay-pool-key-0001", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + monkeypatch.setattr(config, "get_codebuddy_auth_mode", lambda: "api_key_file") + + upstream = install_upstream(monkeypatch, lambda _request: sse_response("relay")) + accepted = await chat(app, RELAY_PASSWORD) + denied = await chat(app, KEY_A) + await upstream.aclose() + + assert accepted.status_code == 200 + assert denied.status_code == 403 + + +@pytest.mark.asyncio +async def test_passthrough_does_not_call_global_credential_managers( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + + async def forbidden_acquire(*_args, **_kwargs): + raise AssertionError("passthrough must not acquire from TXT pool") + + def forbidden_legacy(): + raise AssertionError("passthrough must not select legacy credentials") + + monkeypatch.setattr(empty_pool, "acquire", forbidden_acquire) + monkeypatch.setattr( + codebuddy_router.codebuddy_token_manager, + "get_next_credential", + forbidden_legacy, + ) + upstream = install_upstream(monkeypatch, lambda _request: sse_response()) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 200 + + +@pytest.mark.asyncio +async def test_settings_masks_secrets_and_sentinel_does_not_overwrite( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + monkeypatch.setitem(config._config_cache, "CODEBUDDY_PASSWORD", RELAY_PASSWORD) + monkeypatch.setitem(config._config_cache, "CODEBUDDY_ADMIN_PASSWORD", ADMIN_PASSWORD) + monkeypatch.setattr(config, "save_config_to_json", lambda: None) + + response = await request(app, "/api/settings", ADMIN_PASSWORD) + assert response.status_code == 200 + settings = response.json()["settings"] + assert settings["CODEBUDDY_PASSWORD"] == "********" + assert settings["CODEBUDDY_ADMIN_PASSWORD"] == "********" + assert RELAY_PASSWORD not in response.text + assert ADMIN_PASSWORD not in response.text + + config.update_settings( + { + "CODEBUDDY_PASSWORD": "********", + "CODEBUDDY_ADMIN_PASSWORD": "********", + } + ) + assert config._config_cache["CODEBUDDY_PASSWORD"] == RELAY_PASSWORD + assert config._config_cache["CODEBUDDY_ADMIN_PASSWORD"] == ADMIN_PASSWORD + + +def test_header_builder_and_log_redaction(): + headers = codebuddy_api_client.generate_codebuddy_headers( + KEY_A, api_key_header="both" + ) + assert headers["Authorization"] == f"Bearer {KEY_A}" + assert headers["X-API-Key"] == KEY_A + + message = f"Authorization: Bearer {KEY_A}, X-API-Key: {KEY_B}" + redacted = redact_sensitive(message) + assert KEY_A not in redacted + assert KEY_B not in redacted + assert redacted.count("[REDACTED]") == 2 + + record = logging.LogRecord("test", logging.INFO, __file__, 1, message, (), None) + assert SensitiveHeaderFilter().filter(record) + assert KEY_A not in record.getMessage() diff --git a/tests/test_deploy_script.py b/tests/test_deploy_script.py new file mode 100644 index 0000000..1c39a3a --- /dev/null +++ b/tests/test_deploy_script.py @@ -0,0 +1,119 @@ +import hashlib +import os +import shutil +import stat +import subprocess +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def copy_deploy_fixture(tmp_path): + for name in ("deploy.sh", ".env.example"): + shutil.copy2(PROJECT_ROOT / name, tmp_path / name) + config_dir = tmp_path / "config" + config_dir.mkdir() + shutil.copy2( + PROJECT_ROOT / "config" / "codebuddy_api_keys.example.txt", + config_dir / "codebuddy_api_keys.example.txt", + ) + return tmp_path / "deploy.sh" + + +def run_bootstrap(script, *args): + env = os.environ.copy() + env["DEPLOY_BOOTSTRAP_ONLY"] = "1" + return subprocess.run( + ["bash", str(script), *args], + cwd=script.parent, + env=env, + capture_output=True, + text=True, + check=False, + ) + + +def env_values(path): + values = {} + for line in path.read_text(encoding="utf-8").splitlines(): + if line and not line.lstrip().startswith("#") and "=" in line: + key, value = line.split("=", 1) + values[key.strip()] = value + return values + + +def digest(path): + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def test_bootstrap_creates_secure_passthrough_environment(tmp_path): + script = copy_deploy_fixture(tmp_path) + result = run_bootstrap(script) + + assert result.returncode == 0, result.stderr + env_file = tmp_path / ".env" + values = env_values(env_file) + assert values["CODEBUDDY_CLIENT_AUTH_MODE"] == "passthrough" + assert values["CODEBUDDY_UPSTREAM_API_KEY_HEADER"] == "bearer" + assert values["CODEBUDDY_HOST"] == "0.0.0.0" + assert values["CODEBUDDY_PORT"] == "8001" + assert len(values["CODEBUDDY_PASSWORD"]) == 64 + assert len(values["CODEBUDDY_ADMIN_PASSWORD"]) == 64 + assert values["CODEBUDDY_PASSWORD"] != values["CODEBUDDY_ADMIN_PASSWORD"] + assert stat.S_IMODE(env_file.stat().st_mode) == 0o600 + assert (tmp_path / "config" / "codebuddy_api_keys.txt").exists() + assert (tmp_path / ".codebuddy_creds").is_dir() + + +def test_bootstrap_is_idempotent_and_preserves_existing_files(tmp_path): + script = copy_deploy_fixture(tmp_path) + first = run_bootstrap( + script, + "--relay-password", + "fixed-relay", + "--admin-password", + "fixed-admin", + ) + assert first.returncode == 0, first.stderr + + env_file = tmp_path / ".env" + key_file = tmp_path / "config" / "codebuddy_api_keys.txt" + key_file.write_text("real-key-must-survive\n", encoding="utf-8") + env_digest = digest(env_file) + key_digest = digest(key_file) + + second = run_bootstrap(script) + assert second.returncode == 0, second.stderr + assert digest(env_file) == env_digest + assert digest(key_file) == key_digest + values = env_values(env_file) + assert values["CODEBUDDY_PASSWORD"] == "fixed-relay" + assert values["CODEBUDDY_ADMIN_PASSWORD"] == "fixed-admin" + + +def test_bootstrap_supports_mode_port_and_header_overrides(tmp_path): + script = copy_deploy_fixture(tmp_path) + result = run_bootstrap( + script, + "--client-auth-mode", + "hybrid", + "--upstream-header", + "both", + "--port", + "18001", + ) + + assert result.returncode == 0, result.stderr + values = env_values(tmp_path / ".env") + assert values["CODEBUDDY_CLIENT_AUTH_MODE"] == "hybrid" + assert values["CODEBUDDY_UPSTREAM_API_KEY_HEADER"] == "both" + assert values["CODEBUDDY_PORT"] == "18001" + + +def test_invalid_options_fail_before_writing_environment(tmp_path): + script = copy_deploy_fixture(tmp_path) + result = run_bootstrap(script, "--client-auth-mode", "invalid") + + assert result.returncode != 0 + assert not (tmp_path / ".env").exists() diff --git a/web.py b/web.py index c9c3741..93de4ff 100644 --- a/web.py +++ b/web.py @@ -16,29 +16,32 @@ from config import get_server_host, get_server_port, get_log_level -# 配置日志 +# Configure logging. logging.basicConfig( level=getattr(logging, get_log_level().upper()), format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) +from src.logging_utils import install_sensitive_header_filter + +install_sensitive_header_filter() logger = logging.getLogger(__name__) @asynccontextmanager async def lifespan(app: FastAPI): - """应用生命周期管理""" + """Manage the application lifecycle.""" logger.info("Starting CodeBuddy2API Service") try: - # 启动时初始化资源 + # Initialize resources during startup. await lifecycle_manager.startup() yield finally: - # 关闭时清理资源 + # Clean up resources during shutdown. await lifecycle_manager.shutdown() logger.info("CodeBuddy2API Service stopped") -# 创建FastAPI应用 +# Create the FastAPI application. app = FastAPI( title="CodeBuddy2API", description="CodeBuddy API proxy with OpenAI-compatible interface", @@ -46,7 +49,7 @@ async def lifespan(app: FastAPI): lifespan=lifespan ) -# CORS中间件 +# CORS middleware. app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -55,43 +58,43 @@ async def lifespan(app: FastAPI): allow_headers=["*"], ) -# 挂载前端路由 +# Mount the frontend router. app.include_router( frontend_router, tags=["Frontend"] ) -# 挂载CodeBuddy认证路由 +# Mount the CodeBuddy authentication router. app.include_router( codebuddy_auth_router, prefix="/codebuddy", tags=["CodeBuddy OAuth2 Authentication"] ) -# 挂载CodeBuddy API路由 +# Mount the CodeBuddy API router. app.include_router( codebuddy_router, prefix="/codebuddy", tags=["CodeBuddy Compatible API"] ) -# 挂载设置路由 +# Mount the settings router. app.include_router( settings_router, prefix="/api", tags=["Settings Management"] ) -# 健康检查端点 +# Health-check endpoint. @app.get("/health") async def health_check(): - """健康检查""" + """Return service health information.""" return {"status": "healthy", "service": "codebuddy2api"} @app.get("/") async def root(): - """根路径信息""" + """Return root service information.""" return { "service": "CodeBuddy2API", "version": "1.0.0", @@ -99,6 +102,8 @@ async def root(): "endpoints": { "models": "/codebuddy/v1/models", "chat": "/codebuddy/v1/chat/completions", + "api_keys_status": "/codebuddy/v1/api-keys/status", + "api_keys_reload": "/codebuddy/v1/api-keys/reload", "credentials": "/codebuddy/v1/credentials", "auth_start": "/codebuddy/auth/start", "auth_poll": "/codebuddy/auth/poll", From 6739a0e186efaf73939ac629661c8fabf3e85064 Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 06:17:32 +0700 Subject: [PATCH 2/9] add: API Key Passthrough --- .env.example | 20 +- config.py | 32 +- src/codebuddy_api_client.py | 55 +++- src/codebuddy_message_sanitizer.py | 157 ++++++++++ src/codebuddy_router.py | 278 +++++++++++++++-- tests/test_codebuddy_client_auth.py | 5 +- tests/test_codebuddy_moderation.py | 457 ++++++++++++++++++++++++++++ 7 files changed, 971 insertions(+), 33 deletions(-) create mode 100644 src/codebuddy_message_sanitizer.py create mode 100644 tests/test_codebuddy_moderation.py diff --git a/.env.example b/.env.example index 7ddfac7..d7cd905 100644 --- a/.env.example +++ b/.env.example @@ -26,9 +26,27 @@ CODEBUDDY_PORT=8001 # ----------------- # (Optional) Official CodeBuddy API endpoint. -# This normally does not need to be changed. +# This is correct for CodeBuddy Global and normally does not need to be changed. CODEBUDDY_API_ENDPOINT=https://www.codebuddy.ai +# (Optional) Upstream request profile: web or cli. +# web -> browser-like User-Agent, sends both Authorization and X-Api-Key, +# and omits CLI/IDE identity headers. Reduces false-positive moderation. +# cli -> preserves the legacy CLI/IDE headers for backward compatibility. +# Default: web +CODEBUDDY_REQUEST_PROFILE=web + +# (Optional) Sanitize long or agent-style system prompts (Claude Code, Cursor, +# Cline, etc.) that tend to trigger CodeBuddy false-positive moderation. +# Only the system message is replaced; user/assistant/tool messages are untouched. +# Default: true +CODEBUDDY_SANITIZE_AGENT_PROMPT=true + +# (Optional) Maximum system prompt length (characters). System prompts longer +# than this are replaced with a short neutral instruction. +# Default: 2000 +CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH=2000 + # (Optional) Directory containing CodeBuddy credential JSON files. # Default: .codebuddy_creds CODEBUDDY_CREDS_DIR=.codebuddy_creds diff --git a/config.py b/config.py index d41389a..fda6d9e 100644 --- a/config.py +++ b/config.py @@ -38,7 +38,10 @@ "CODEBUDDY_API_KEY_COOLDOWN_SECONDS": 300, "CODEBUDDY_CLIENT_AUTH_MODE": "relay", "CODEBUDDY_ADMIN_PASSWORD": None, - "CODEBUDDY_UPSTREAM_API_KEY_HEADER": "bearer" + "CODEBUDDY_UPSTREAM_API_KEY_HEADER": "bearer", + "CODEBUDDY_REQUEST_PROFILE": "web", + "CODEBUDDY_SANITIZE_AGENT_PROMPT": True, + "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000 } # --- Core Functions --- @@ -152,6 +155,33 @@ def get_upstream_api_key_header() -> str: ) return mode +def get_codebuddy_request_profile() -> str: + profile = str(_get_config_value("CODEBUDDY_REQUEST_PROFILE")).strip().lower() + if profile not in {"web", "cli"}: + raise ValueError("CODEBUDDY_REQUEST_PROFILE must be web or cli") + return profile + + +def _coerce_bool(value: Any, default: bool) -> bool: + if isinstance(value, bool): + return value + if value is None: + return default + return str(value).strip().lower() in {"true", "1", "t", "y", "yes"} + + +def get_sanitize_agent_prompt() -> bool: + return _coerce_bool(_get_config_value("CODEBUDDY_SANITIZE_AGENT_PROMPT"), True) + + +def get_max_system_prompt_length() -> int: + try: + value = int(_get_config_value("CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH")) + except (TypeError, ValueError): + return 2000 + return value if value > 0 else 2000 + + def get_codebuddy_api_endpoint() -> str: return str(_get_config_value("CODEBUDDY_API_ENDPOINT")) diff --git a/src/codebuddy_api_client.py b/src/codebuddy_api_client.py index a85b05e..2ee13cc 100644 --- a/src/codebuddy_api_client.py +++ b/src/codebuddy_api_client.py @@ -164,6 +164,12 @@ def convert_openai_to_codebuddy_messages(self, openai_messages: List[Dict]) -> L return codebuddy_messages + # Browser-compatible User-Agent used by the "web" request profile. + _WEB_USER_AGENT = ( + 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 ' + '(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36' + ) + def generate_codebuddy_headers( self, bearer_token: str, @@ -172,19 +178,55 @@ def generate_codebuddy_headers( conversation_request_id: Optional[str] = None, conversation_message_id: Optional[str] = None, request_id: Optional[str] = None, - api_key_header: str = "bearer" + api_key_header: str = "bearer", + profile: str = "web" ) -> Dict[str, str]: """ Generate the complete header set required by the CodeBuddy API. + + The request profile controls which identity is presented upstream: + + * ``web`` - browser User-Agent, always sends both ``Authorization`` + and ``X-Api-Key``, and omits CLI/IDE identity headers. + This is less likely to trigger false-positive + moderation. + * ``cli`` - preserves the legacy CLI/IDE headers (X-IDE-*, + x-stainless-*, CLI User-Agent) for backward + compatibility, honoring ``api_key_header``. + Prefer supplied conversation IDs and generate missing IDs automatically. """ if api_key_header not in {"x-api-key", "bearer", "both"}: raise ValueError("api_key_header must be x-api-key, bearer, or both") + if profile not in {"web", "cli"}: + raise ValueError("profile must be web or cli") + + # Shared conversation/routing headers used by both profiles. headers = { 'Host': 'www.codebuddy.ai', - 'Accept': 'application/json', 'Content-Type': 'application/json', 'X-Requested-With': 'XMLHttpRequest', + 'X-Conversation-ID': conversation_id or str(uuid.uuid4()), + 'X-Conversation-Request-ID': conversation_request_id or secrets.token_hex(16), + 'X-Conversation-Message-ID': conversation_message_id or str(uuid.uuid4()).replace('-', ''), + 'X-Request-ID': request_id or str(uuid.uuid4()).replace('-', ''), + 'X-Domain': 'www.codebuddy.ai', + 'X-Product': 'SaaS', + 'X-User-Id': user_id or 'b5be3a67-237e-4ee6-9b9a-0b9ecd7b454b', + } + + if profile == "web": + # Browser-like identity. Always send both auth headers so the + # request matches what web/other adapters send upstream. + headers['Accept'] = 'text/event-stream' + headers['User-Agent'] = self._WEB_USER_AGENT + headers["X-API-Key"] = bearer_token + headers["Authorization"] = f"Bearer {bearer_token}" + return headers + + # Legacy CLI profile: keep IDE identity headers and honor api_key_header. + headers['Accept'] = 'application/json' + headers.update({ 'x-stainless-arch': 'x64', 'x-stainless-lang': 'js', 'x-stainless-os': 'Windows', @@ -192,19 +234,12 @@ def generate_codebuddy_headers( 'x-stainless-retry-count': '0', 'x-stainless-runtime': 'node', 'x-stainless-runtime-version': 'v22.13.1', - 'X-Conversation-ID': conversation_id or str(uuid.uuid4()), - 'X-Conversation-Request-ID': conversation_request_id or secrets.token_hex(16), - 'X-Conversation-Message-ID': conversation_message_id or str(uuid.uuid4()).replace('-', ''), - 'X-Request-ID': request_id or str(uuid.uuid4()).replace('-', ''), 'X-Agent-Intent': 'craft', 'X-IDE-Type': 'CLI', 'X-IDE-Name': 'CLI', 'X-IDE-Version': '1.0.7', - 'X-Domain': 'www.codebuddy.ai', 'User-Agent': 'CLI/1.0.7 CodeBuddy/1.0.7', - 'X-Product': 'SaaS', - 'X-User-Id': user_id or 'b5be3a67-237e-4ee6-9b9a-0b9ecd7b454b' - } + }) if api_key_header in {"x-api-key", "both"}: headers["X-API-Key"] = bearer_token if api_key_header in {"bearer", "both"}: diff --git a/src/codebuddy_message_sanitizer.py b/src/codebuddy_message_sanitizer.py new file mode 100644 index 0000000..091aaa3 --- /dev/null +++ b/src/codebuddy_message_sanitizer.py @@ -0,0 +1,157 @@ +""" +Message sanitization and moderation detection for CodeBuddy2API. + +Two concerns live here, both aimed at reducing CodeBuddy false-positive +moderation without weakening legitimate upstream moderation: + +1. Agent system-prompt sanitization. Long coding-agent identity blocks + (Claude Code, Cursor, Cline, aider, ...) trigger CodeBuddy's content + filter far more often than a plain request. We replace ONLY the system + message with a short, neutral instruction. User/assistant/tool messages + are never touched. + +2. Mandarin moderation detection. CodeBuddy returns a Mandarin refusal when + it blocks a request. We detect specific markers so the router can surface + an OpenAI-compatible ``content_filter`` error instead of leaking the + refusal as an assistant reply. +""" +import logging +from typing import Any, Dict, List, Tuple + +logger = logging.getLogger(__name__) + +# Neutral replacement used when an agent system prompt is sanitized. +NEUTRAL_SYSTEM_PROMPT = ( + "You are a helpful coding assistant. Answer accurately, follow the " + "user's request, and reply in the same language as the user." +) + +# Lower-cased substrings that identify a coding-agent system prompt. These are +# matched case-insensitively against system-message text only. +_AGENT_PROMPT_MARKERS: Tuple[str, ...] = ( + "you are claude code", + "claude code", + "official cli", + "cursor", + "windsurf", + "cline", + "aider", + "continue", + "copilot", + "coding agent", + "code agent", + "agentic coding assistant", + "cc_entrypoint", + "", + "", + "", +) + +# Specific markers found in CodeBuddy's Mandarin moderation refusal. Matching +# any one of these flags a moderation response. We deliberately do NOT treat +# arbitrary Mandarin text as moderation. +_MODERATION_MARKERS: Tuple[str, ...] = ( + "敏感内容", + "系统检测到", + "无法响应您的请求", + "请检查后重新输入", +) + + +def _extract_text(content: Any) -> str: + """Flatten string or structured content into plain text for matching.""" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: List[str] = [] + for item in content: + if isinstance(item, dict): + if item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + # Include other structured values so markers embedded in + # nested dicts are still detected. + parts.append(str(item)) + elif isinstance(item, str): + parts.append(item) + else: + parts.append(str(item)) + return "".join(parts) + if content is None: + return "" + return str(content) + + +def _looks_like_agent_prompt(text: str) -> bool: + lowered = text.lower() + return any(marker in lowered for marker in _AGENT_PROMPT_MARKERS) + + +def sanitize_messages( + messages: List[Dict[str, Any]], + *, + enabled: bool = True, + max_system_prompt_length: int = 2000, +) -> Tuple[List[Dict[str, Any]], bool]: + """ + Return a sanitized copy of ``messages`` and whether any system message was + replaced. + + Only ``system`` messages are considered. A system message is replaced with + :data:`NEUTRAL_SYSTEM_PROMPT` when either: + + * its text matches a coding-agent marker, or + * its text length exceeds ``max_system_prompt_length``. + + User, assistant, and tool messages are copied through unchanged. This never + attempts to disguise sensitive content; it only removes agent identity + boilerplate that causes false-positive moderation. + """ + if not enabled: + return messages, False + + sanitized: List[Dict[str, Any]] = [] + changed = False + + for msg in messages: + if not isinstance(msg, dict) or msg.get("role") != "system": + sanitized.append(msg) + continue + + text = _extract_text(msg.get("content")) + too_long = len(text) > max_system_prompt_length + is_agent = _looks_like_agent_prompt(text) + + if too_long or is_agent: + new_msg = dict(msg) + new_msg["content"] = NEUTRAL_SYSTEM_PROMPT + sanitized.append(new_msg) + changed = True + # Log only metadata, never the prompt contents. + logger.info( + "System prompt sanitized: agent_marker=%s over_length=%s " + "original_length=%d", + is_agent, + too_long, + len(text), + ) + else: + sanitized.append(msg) + + return sanitized, changed + + +def is_codebuddy_moderation_response(text: Any) -> bool: + """ + Return True when ``text`` contains a CodeBuddy Mandarin moderation marker. + + Accepts a string or any value coercible to text (e.g. nested JSON already + serialized). Only the specific refusal markers count; ordinary Mandarin + responses are not flagged. + """ + if text is None: + return False + haystack = text if isinstance(text, str) else str(text) + if not haystack: + return False + return any(marker in haystack for marker in _MODERATION_MARKERS) diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index ec9cfb0..d75df7b 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -5,6 +5,7 @@ import json import time import uuid +import hashlib import logging import asyncio from dataclasses import dataclass @@ -23,7 +24,16 @@ from .codebuddy_token_manager import codebuddy_token_manager from .usage_stats_manager import usage_stats_manager from .keyword_replacer import apply_keyword_replacement_to_system_message -from config import get_upstream_api_key_header +from .codebuddy_message_sanitizer import ( + is_codebuddy_moderation_response, + sanitize_messages, +) +from config import ( + get_codebuddy_request_profile, + get_max_system_prompt_length, + get_sanitize_agent_prompt, + get_upstream_api_key_header, +) logger = logging.getLogger(__name__) router = APIRouter() @@ -448,6 +458,120 @@ def __init__(self, kind: str, status_code: int, code: str): self.code = code +class CodeBuddyModerationError(Exception): + """Raised when CodeBuddy rejects a request via its content moderation. + + This is distinct from an upstream failure: the key is valid and the + request reached upstream. It must not trigger key failover or retries. + """ + + +# Human-readable message returned to the client on a moderation rejection. +MODERATION_MESSAGE = ( + "CodeBuddy rejected the request through its content moderation system. " + "This may be caused by the system prompt or conversation history." +) + +# Shorter message used inside the streaming content_filter delta. +MODERATION_STREAM_MESSAGE = ( + "CodeBuddy rejected the request through its content moderation system." +) + +# Number of assistant-content characters to buffer while deciding whether a +# streaming response is a moderation refusal. The Mandarin refusal is short and +# self-contained, so a small buffer distinguishes it from a normal reply. +MODERATION_STREAM_BUFFER_CHARS = 200 + + +def _extract_delta_content(sse_line: str) -> str: + """Return the assistant ``delta.content`` string from an OpenAI SSE line.""" + obj = parse_sse_line(sse_line.strip()) + if not obj: + return "" + try: + choices = obj.get("choices") or [] + if not choices: + return "" + delta = choices[0].get("delta") or {} + content = delta.get("content") + return content if isinstance(content, str) else "" + except (AttributeError, IndexError, TypeError): + return "" + + +async def _moderation_stream() -> AsyncGenerator[str, None]: + """Yield an OpenAI-compatible content_filter SSE stream, ending with [DONE].""" + chunk = { + "id": "chatcmpl-codebuddy-filter", + "object": "chat.completion.chunk", + "created": int(time.time()), + "choices": [ + { + "index": 0, + "delta": {"content": MODERATION_STREAM_MESSAGE}, + "finish_reason": "content_filter", + } + ], + } + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" + yield "data: [DONE]\n\n" + + +def _key_fingerprint(token: Optional[str]) -> str: + """Return a short, non-reversible SHA-256 fingerprint of an API key. + + Used only for safe diagnostics. The raw key is never logged. + """ + if not token: + return "none" + return hashlib.sha256(token.encode("utf-8")).hexdigest()[:8] + + +def _log_request_diagnostics( + *, + payload: Dict[str, Any], + client_wants_stream: bool, + system_prompt_sanitized: bool, + request_profile: str, + key_fingerprint: str, +) -> None: + """Log request metadata only. Never logs keys, prompts, or user content.""" + messages = payload.get("messages", []) or [] + roles: List[str] = [] + content_lengths: List[int] = [] + tool_count = 0 + for msg in messages: + if not isinstance(msg, dict): + continue + roles.append(str(msg.get("role", "unknown"))) + content = msg.get("content", "") + if isinstance(content, str): + content_lengths.append(len(content)) + elif isinstance(content, list): + length = 0 + for item in content: + if isinstance(item, dict): + if item.get("type") in {"tool_use", "tool_result"}: + tool_count += 1 + length += len(str(item.get("text", ""))) if item.get("type") == "text" else 0 + content_lengths.append(length) + else: + content_lengths.append(0) + tool_count += len(payload.get("tools", []) or []) + logger.info( + "CodeBuddy request model=%s stream=%s roles=[%s] content_lengths=%s " + "tool_count=%d system_prompt_sanitized=%s request_profile=%s key_fingerprint=%s", + payload.get("model", "unknown"), + client_wants_stream, + ",".join(roles), + content_lengths, + tool_count, + system_prompt_sanitized, + request_profile, + key_fingerprint, + ) + + class CodeBuddyStreamService: """CodeBuddy streaming service; each method performs exactly one upstream attempt.""" @@ -521,10 +645,25 @@ async def converted_chunks(): yield buffer + '\n' stream = converted_chunks() + + # Pre-buffer the leading chunks so a Mandarin moderation refusal can be + # detected before any assistant content is sent downstream. The refusal + # is short, so a small buffer suffices; normal replies simply get + # replayed afterwards in order. + buffered_lines: List[str] = [] + accumulated_content = "" + moderation_detected = False try: - first_chunk = await anext(stream) - except StopAsyncIteration: - first_chunk = None + async for line in stream: + buffered_lines.append(line) + if '[DONE]' in line: + break + accumulated_content += _extract_delta_content(line) + if is_codebuddy_moderation_response(accumulated_content): + moderation_detected = True + break + if len(accumulated_content) >= MODERATION_STREAM_BUFFER_CHARS: + break except httpx.TimeoutException as exc: await response.aclose() raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc @@ -535,10 +674,26 @@ async def converted_chunks(): await response.aclose() raise UpstreamAttemptError("fatal", 502, "upstream_response_invalid") from exc + if moderation_detected: + # Nothing normal was sent yet: replace the whole stream with an + # OpenAI-compatible content_filter stream. The key is valid, so do + # not mark it failed or fail over. + await response.aclose() + + async def moderation_core(): + async for item in _moderation_stream(): + yield item + + return StreamingResponse( + moderation_core(), media_type="text/event-stream", headers={ + **SSE_HEADERS, "X-CodeBuddy-Moderation": "true" + } + ) + async def stream_core(): try: - if first_chunk is not None: - yield first_chunk + for line in buffered_lines: + yield line async for chunk in stream: yield chunk except httpx.RequestError: @@ -587,10 +742,12 @@ async def handle_non_stream_response( try: aggregator = StreamResponseAggregator() + raw_text = "" buffer = "" async for chunk in response.aiter_text(): if not chunk: continue + raw_text += chunk buffer += chunk while '\n' in buffer: line, buffer = buffer.split('\n', 1) @@ -602,7 +759,7 @@ async def handle_non_stream_response( obj = parse_sse_line(buffer.strip()) if obj: aggregator.process_chunk(obj) - return aggregator.finalize() + final = aggregator.finalize() except httpx.RequestError as exc: raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc except UpstreamAttemptError: @@ -610,27 +767,68 @@ async def handle_non_stream_response( except Exception as exc: raise UpstreamAttemptError("fatal", 502, "upstream_response_invalid") from exc + # Detect a CodeBuddy moderation refusal in either the aggregated + # assistant content or the raw upstream body (covers nested/non-SSE + # bodies). This is not an upstream failure, so it must not trigger + # key failover. + aggregated_content = "" + try: + aggregated_content = final["choices"][0]["message"].get("content") or "" + except (KeyError, IndexError, TypeError): + aggregated_content = "" + if is_codebuddy_moderation_response(aggregated_content) or \ + is_codebuddy_moderation_response(raw_text): + raise CodeBuddyModerationError() + return final + class RequestProcessor: """Request preprocessor - thread-safe request handling""" @staticmethod - def prepare_payload(request_body: Dict[str, Any]) -> Dict[str, Any]: - """Prepare the request payload""" + def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], bool]: + """Prepare the request payload. + + Returns the upstream payload and a flag indicating whether an agent + system prompt was sanitized. The flag is used only for safe diagnostics + and is never sent upstream. + """ payload = request_body.copy() payload["stream"] = True # CodeBuddy only supports streaming requests - # Handle the message count requirement: CodeBuddy requires at least 2 messages messages = payload.get("messages", []) - if len(messages) == 1 and messages[0].get("role") == "user": - system_msg = {"role": "system", "content": "You are a helpful assistant."} - payload["messages"] = [system_msg] + messages - - # Apply keyword replacement - for msg in payload.get("messages", []): - if msg.get("role") == "system": + + # Sanitize only agent system prompts that tend to trigger false-positive + # moderation. User/assistant/tool messages are never modified, and + # legitimate short system prompts are left intact. + try: + sanitize_enabled = get_sanitize_agent_prompt() + max_len = get_max_system_prompt_length() + except Exception: + sanitize_enabled, max_len = True, 2000 + messages, system_prompt_sanitized = sanitize_messages( + messages, enabled=sanitize_enabled, max_system_prompt_length=max_len + ) + + # CodeBuddy requires at least two messages. Only add a default system + # prompt when there is no system message already, to avoid duplicating + # system messages or reordering the conversation. + has_system = any( + isinstance(m, dict) and m.get("role") == "system" for m in messages + ) + if not has_system and len(messages) == 1 and messages[0].get("role") == "user": + system_msg = { + "role": "system", + "content": "You are a helpful assistant. Reply in the same language as the user.", + } + messages = [system_msg] + messages + + # Apply keyword replacement to system messages only. + for msg in messages: + if isinstance(msg, dict) and msg.get("role") == "system": msg["content"] = apply_keyword_replacement_to_system_message(msg.get("content")) - - return payload + + payload["messages"] = messages + return payload, system_prompt_sanitized @staticmethod def validate_request(request_body: Dict[str, Any]) -> None: @@ -797,13 +995,23 @@ async def chat_completions( 503, ) - payload = RequestProcessor.prepare_payload(request_body) + payload, system_prompt_sanitized = RequestProcessor.prepare_payload(request_body) usage_stats_manager.record_model_usage(payload.get("model", "unknown")) service = CodeBuddyStreamService() client_wants_stream = request_body.get("stream", False) excluded_ids: Set[str] = set() last_error: Optional[UpstreamAttemptError] = None + try: + request_profile = get_codebuddy_request_profile() + except ValueError: + return openai_error_response( + "Upstream request profile is misconfigured", + "configuration_error", + "request_profile_invalid", + 500, + ) + for _attempt in range(max_attempts): if source == "passthrough": credential = passthrough_credential @@ -827,6 +1035,15 @@ async def chat_completions( 500, ) + # Safe diagnostics: metadata only, never key/prompt/user content. + _log_request_diagnostics( + payload=payload, + client_wants_stream=client_wants_stream, + system_prompt_sanitized=system_prompt_sanitized, + request_profile=request_profile, + key_fingerprint=_key_fingerprint(credential.bearer_token), + ) + headers = codebuddy_api_client.generate_codebuddy_headers( bearer_token=credential.bearer_token, user_id=credential.user_id, @@ -835,6 +1052,7 @@ async def chat_completions( conversation_message_id=x_conversation_message_id, request_id=x_request_id, api_key_header=upstream_key_header, + profile=request_profile, ) try: @@ -847,6 +1065,26 @@ async def chat_completions( if credential.key_id is not None: await codebuddy_api_key_manager.mark_success(credential.key_id) return result + except CodeBuddyModerationError: + # The key is valid and the request reached upstream; moderation is + # not a key failure. Mark success (no failover) and return a clear + # OpenAI-compatible content_filter error. Streaming moderation is + # already handled inside open_stream_response before this point. + if credential.key_id is not None: + await codebuddy_api_key_manager.mark_success(credential.key_id) + logger.info("CodeBuddy moderation rejection detected (non-stream)") + return JSONResponse( + status_code=400, + headers={"X-CodeBuddy-Moderation": "true"}, + content={ + "error": { + "message": MODERATION_MESSAGE, + "type": "content_filter", + "param": None, + "code": "codebuddy_content_filter", + } + }, + ) except UpstreamAttemptError as error: last_error = error await record_attempt_error(credential, error) diff --git a/tests/test_codebuddy_client_auth.py b/tests/test_codebuddy_client_auth.py index 815a037..1283332 100644 --- a/tests/test_codebuddy_client_auth.py +++ b/tests/test_codebuddy_client_auth.py @@ -45,11 +45,14 @@ async def empty_pool(monkeypatch, tmp_path): return manager -def configure(monkeypatch, client_mode="passthrough", header_mode="bearer"): +def configure(monkeypatch, client_mode="passthrough", header_mode="bearer", profile="cli"): monkeypatch.setattr(auth, "get_client_auth_mode", lambda: client_mode) monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: header_mode) + # Pin the request profile to cli so the upstream header-mode mapping stays + # deterministic in these tests (the web profile always sends both headers). + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: profile) monkeypatch.setattr( codebuddy_router.usage_stats_manager, "record_model_usage", diff --git a/tests/test_codebuddy_moderation.py b/tests/test_codebuddy_moderation.py new file mode 100644 index 0000000..0d20680 --- /dev/null +++ b/tests/test_codebuddy_moderation.py @@ -0,0 +1,457 @@ +""" +Tests for CodeBuddy false-positive moderation mitigation: + + * agent system-prompt sanitization (unit + end-to-end) + * web/cli request header profiles + * Mandarin moderation detection (non-stream + stream) + * safe diagnostics (no raw key in logs) + * message normalization / tool-call preservation + * upstream status propagation for 9Router fallback + +Uses a mock upstream transport; no real CodeBuddy API key is required. +""" +import json +import logging + +import httpx +import pytest + +import config +from src import auth, codebuddy_router +from src.codebuddy_api_client import codebuddy_api_client +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager +from src.codebuddy_message_sanitizer import ( + NEUTRAL_SYSTEM_PROMPT, + is_codebuddy_moderation_response, + sanitize_messages, +) + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" +KEY_B = "passthrough-account-beta-0002" + +MODERATION_TEXT = ( + "抱歉，系统检测到您当前输入的信息存在敏感内容，我无法响应您的请求，请检查后重新输入。" +) + + +# --------------------------------------------------------------------------- # +# Unit tests: sanitizer + moderation detector (no I/O) +# --------------------------------------------------------------------------- # + + +def test_simple_user_prompt_not_modified(): + messages = [{"role": "user", "content": "halo, siapa kamu"}] + result, changed = sanitize_messages(messages) + assert changed is False + assert result[0]["content"] == "halo, siapa kamu" + + +def test_user_message_never_changed_even_with_agent_markers(): + # A user message that mentions "Claude Code" must stay intact; only system + # messages are ever sanitized. + messages = [{"role": "user", "content": "explain what claude code cli does"}] + result, changed = sanitize_messages(messages) + assert changed is False + assert result[0]["content"] == "explain what claude code cli does" + + +def test_claude_code_system_prompt_detected_and_replaced(): + messages = [ + {"role": "system", "content": "You are Claude Code, Anthropic's official CLI."}, + {"role": "user", "content": "hi"}, + ] + result, changed = sanitize_messages(messages) + assert changed is True + assert result[0]["content"] == NEUTRAL_SYSTEM_PROMPT + assert result[1]["content"] == "hi" + + +def test_over_length_system_prompt_replaced(): + long_prompt = "a" * 5000 + messages = [{"role": "system", "content": long_prompt}] + result, changed = sanitize_messages(messages, max_system_prompt_length=2000) + assert changed is True + assert result[0]["content"] == NEUTRAL_SYSTEM_PROMPT + + +def test_short_normal_system_prompt_preserved(): + messages = [ + {"role": "system", "content": "You are a friendly translator."}, + {"role": "user", "content": "hola"}, + ] + result, changed = sanitize_messages(messages) + assert changed is False + assert result[0]["content"] == "You are a friendly translator." + + +def test_sanitize_disabled_is_noop(): + messages = [{"role": "system", "content": "You are Claude Code"}] + result, changed = sanitize_messages(messages, enabled=False) + assert changed is False + assert result[0]["content"] == "You are Claude Code" + + +def test_moderation_markers_detected(): + assert is_codebuddy_moderation_response(MODERATION_TEXT) is True + assert is_codebuddy_moderation_response("系统检测到问题") is True + + +def test_normal_mandarin_not_flagged(): + assert is_codebuddy_moderation_response("你好，我是一个AI助手，很高兴认识你。") is False + assert is_codebuddy_moderation_response("2 + 2 = 4") is False + assert is_codebuddy_moderation_response("") is False + assert is_codebuddy_moderation_response(None) is False + + +# --------------------------------------------------------------------------- # +# Header profile unit tests +# --------------------------------------------------------------------------- # + + +def test_web_profile_uses_browser_user_agent_and_both_headers(): + headers = codebuddy_api_client.generate_codebuddy_headers( + bearer_token=KEY_A, profile="web" + ) + assert "Mozilla/5.0" in headers["User-Agent"] + assert "CLI" not in headers["User-Agent"] + assert headers["Authorization"] == f"Bearer {KEY_A}" + assert headers["X-API-Key"] == KEY_A + # CLI/IDE identity headers must be absent on the web profile. + assert "X-IDE-Type" not in headers + assert "X-IDE-Name" not in headers + assert "x-stainless-lang" not in headers + + +def test_cli_profile_preserves_legacy_headers(): + headers = codebuddy_api_client.generate_codebuddy_headers( + bearer_token=KEY_A, profile="cli", api_key_header="bearer" + ) + assert headers["User-Agent"] == "CLI/1.0.7 CodeBuddy/1.0.7" + assert headers["X-IDE-Type"] == "CLI" + assert headers["Authorization"] == f"Bearer {KEY_A}" + # cli honors api_key_header; bearer-only means no X-API-Key. + assert "X-API-Key" not in headers + + +# --------------------------------------------------------------------------- # +# Integration fixtures / helpers +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def configure(monkeypatch, client_mode="passthrough", profile="web"): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: client_mode) + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: profile) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse(text, finish="stop"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":{json.dumps(text)}}},"finish_reason":{json.dumps(finish)}}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +async def request(app, path, token=None, method="GET", json_body=None): + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.request(method, path, headers=headers, json=json_body) + + +async def chat(app, token, stream=False, messages=None): + body = { + "model": "auto-chat", + "messages": messages or [{"role": "user", "content": "halo, siapa kamu"}], + "stream": stream, + } + return await request(app, "/codebuddy/v1/chat/completions", token, "POST", body) + + +# --------------------------------------------------------------------------- # +# Integration: header profile end-to-end +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_web_profile_sends_both_headers_upstream(monkeypatch, app, empty_pool): + configure(monkeypatch, profile="web") + seen = {} + + def handler(req): + seen["authorization"] = req.headers.get("authorization") + seen["x-api-key"] = req.headers.get("x-api-key") + seen["user-agent"] = req.headers.get("user-agent") + return sse("hi") + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["authorization"] == f"Bearer {KEY_A}" + assert seen["x-api-key"] == KEY_A + assert "Mozilla/5.0" in seen["user-agent"] + + +@pytest.mark.asyncio +async def test_cli_profile_still_works(monkeypatch, app, empty_pool): + configure(monkeypatch, profile="cli") + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "bearer") + seen = {} + + def handler(req): + seen["authorization"] = req.headers.get("authorization") + seen["user-agent"] = req.headers.get("user-agent") + seen["x-ide-type"] = req.headers.get("x-ide-type") + return sse("hi") + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["authorization"] == f"Bearer {KEY_A}" + assert seen["user-agent"] == "CLI/1.0.7 CodeBuddy/1.0.7" + assert seen["x-ide-type"] == "CLI" + + +# --------------------------------------------------------------------------- # +# Integration: sanitization end-to-end +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_agent_system_prompt_sanitized_upstream(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + body = json.loads(req.content) + seen["messages"] = body["messages"] + return sse("hi") + + upstream = install_upstream(monkeypatch, handler) + response = await chat( + app, + KEY_A, + messages=[ + {"role": "system", "content": "You are Claude Code, the official CLI."}, + {"role": "user", "content": "halo, siapa kamu"}, + ], + ) + await upstream.aclose() + + assert response.status_code == 200 + system_msg = next(m for m in seen["messages"] if m["role"] == "system") + user_msg = next(m for m in seen["messages"] if m["role"] == "user") + assert system_msg["content"] == NEUTRAL_SYSTEM_PROMPT + # The user message must be forwarded verbatim. + assert user_msg["content"] == "halo, siapa kamu" + + +@pytest.mark.asyncio +async def test_tool_calls_survive_normalization(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + tools = [ + { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object"}}, + } + ] + upstream = install_upstream(monkeypatch, handler) + response = await request( + app, + "/codebuddy/v1/chat/completions", + KEY_A, + "POST", + { + "model": "auto-chat", + "messages": [{"role": "user", "content": "weather?"}], + "tools": tools, + "stream": False, + }, + ) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["body"]["tools"] == tools + + +# --------------------------------------------------------------------------- # +# Integration: moderation detection +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_non_stream_moderation_becomes_content_filter(monkeypatch, app, empty_pool): + configure(monkeypatch) + + def handler(req): + return sse(MODERATION_TEXT) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == 400 + assert response.headers.get("X-CodeBuddy-Moderation") == "true" + payload = response.json() + assert payload["error"]["type"] == "content_filter" + assert payload["error"]["code"] == "codebuddy_content_filter" + # The raw Mandarin refusal must not be surfaced as assistant content. + assert "choices" not in payload + + +@pytest.mark.asyncio +async def test_stream_moderation_uses_content_filter_finish_reason( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + + def handler(req): + return sse(MODERATION_TEXT, finish=None) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert response.headers.get("X-CodeBuddy-Moderation") == "true" + text = response.text + assert '"finish_reason": "content_filter"' in text or '"finish_reason":"content_filter"' in text + assert "[DONE]" in text + # The Mandarin refusal text must not be streamed as assistant content. + assert "敏感内容" not in text + + +@pytest.mark.asyncio +async def test_normal_mandarin_response_not_flagged(monkeypatch, app, empty_pool): + configure(monkeypatch) + normal = "你好，我是一个AI助手。" + + def handler(req): + return sse(normal) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == 200 + assert response.headers.get("X-CodeBuddy-Moderation") is None + assert response.json()["choices"][0]["message"]["content"] == normal + + +@pytest.mark.asyncio +async def test_normal_stream_still_emits_done(monkeypatch, app, empty_pool): + configure(monkeypatch) + + def handler(req): + return sse("hello there") + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert "[DONE]" in response.text + assert "hello there" in response.text + + +# --------------------------------------------------------------------------- # +# Integration: upstream status propagation for 9Router +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "expected"), + [(401, 401), (403, 403), (429, 429), (500, 502), (503, 502)], +) +async def test_upstream_status_propagated(monkeypatch, app, empty_pool, status, expected): + configure(monkeypatch) + + def handler(req): + return httpx.Response(status, json={"error": "upstream"}) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == expected + # Never leak the key or upstream body. + assert KEY_A not in response.text + + +# --------------------------------------------------------------------------- # +# Safe diagnostics +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_raw_key_not_in_logs(monkeypatch, app, empty_pool, caplog): + configure(monkeypatch) + + def handler(req): + return sse("hi") + + upstream = install_upstream(monkeypatch, handler) + with caplog.at_level(logging.INFO, logger="src.codebuddy_router"): + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == 200 + combined = "\n".join(record.getMessage() for record in caplog.records) + assert KEY_A not in combined + # A short fingerprint should be present instead. + assert "key_fingerprint=" in combined + assert "request_profile=web" in combined From 4573feff17484f4ee5e54d809a9129b59b771746 Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 09:10:26 +0700 Subject: [PATCH 3/9] fix: Stream Tools --- tests/test_codebuddy_response_shape.py | 363 +++++++++++++++++++++++++ 1 file changed, 363 insertions(+) create mode 100644 tests/test_codebuddy_response_shape.py diff --git a/tests/test_codebuddy_response_shape.py b/tests/test_codebuddy_response_shape.py new file mode 100644 index 0000000..cd345a3 --- /dev/null +++ b/tests/test_codebuddy_response_shape.py @@ -0,0 +1,363 @@ +""" +Response-shape regression tests for /codebuddy/v1/chat/completions. + +Locks in the OpenAI-compatible contract for both transports so a future change +to the aggregator or streaming path cannot silently regress it: + + * stream=false -> a single ChatCompletion JSON object + (choices[0].message.role="assistant", non-empty content, + finish_reason="stop") + * stream=true -> text/event-stream, `data: \\n\\n` events with text in + choices[0].delta.content, terminated by exactly one + `data: [DONE]` + +Covers: sanitized system prompt still yields content, moderation maps to +content_filter (never an empty reply), tool-call-only responses return +tool_calls (never "no response"), and reasoning_content never replaces the +final assistant content. + +Uses a mock upstream transport; no real CodeBuddy API key is required. +""" +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager +from src.codebuddy_message_sanitizer import NEUTRAL_SYSTEM_PROMPT + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + +MODERATION_TEXT = ( + "抱歉，系统检测到您当前输入的信息存在敏感内容，我无法响应您的请求，请检查后重新输入。" +) + + +# --------------------------------------------------------------------------- # +# Fixtures / helpers +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def configure(monkeypatch, profile="web"): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: profile) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse_body(chunks, done=True): + """Build an SSE ``httpx.Response`` from ``(delta, finish_reason)`` pairs.""" + parts = [] + for delta, finish in chunks: + choice = {"index": 0, "delta": delta} + if finish is not None: + choice["finish_reason"] = finish + obj = { + "id": "chat-shape", + "object": "chat.completion.chunk", + "model": "auto-chat", + "choices": [choice], + } + parts.append("data: " + json.dumps(obj, ensure_ascii=False)) + text = "\n\n".join(parts) + "\n\n" + if done: + text += "data: [DONE]\n\n" + return httpx.Response(200, text=text, headers={"content-type": "text/event-stream"}) + + +async def request(app, path, token=None, method="GET", json_body=None): + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.request(method, path, headers=headers, json=json_body) + + +async def chat(app, token, stream=False, messages=None): + body = { + "model": "auto-chat", + "messages": messages or [{"role": "user", "content": "halo, siapa kamu"}], + "stream": stream, + } + return await request(app, "/codebuddy/v1/chat/completions", token, "POST", body) + + +def parse_stream_events(text): + """Return the list of parsed JSON event objects (excluding [DONE]).""" + events = [] + for raw in text.split("\n\n"): + line = raw.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:"):].strip() + if payload == "[DONE]": + continue + events.append(json.loads(payload)) + return events + + +# --------------------------------------------------------------------------- # +# 1. stream=true yields content and exactly one [DONE] +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_stream_yields_content_and_single_done(monkeypatch, app, empty_pool): + configure(monkeypatch) + + def handler(_req): + return sse_body( + [ + ({"role": "assistant"}, None), # role-only chunk (no content) + ({"content": "Hello"}, None), + ({"content": " world"}, None), + ({}, "stop"), # trailing usage/stop chunk + ] + ) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + + text = response.text + assert text.count("[DONE]") == 1 + # No HTML or extra JSON wrapper. + assert " Date: Tue, 21 Jul 2026 09:38:26 +0700 Subject: [PATCH 4/9] fix: Stream Tools --- .env.example | 12 + config.py | 32 ++- src/codebuddy_router.py | 198 +++++++++++--- tests/test_codebuddy_payload_allowlist.py | 319 ++++++++++++++++++++++ 4 files changed, 527 insertions(+), 34 deletions(-) create mode 100644 tests/test_codebuddy_payload_allowlist.py diff --git a/.env.example b/.env.example index d7cd905..02372cb 100644 --- a/.env.example +++ b/.env.example @@ -47,6 +47,18 @@ CODEBUDDY_SANITIZE_AGENT_PROMPT=true # Default: 2000 CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH=2000 +# (Optional) Fallback model used when the client sends an unknown model label +# (e.g. a UI display name). Must be a model CodeBuddy actually supports. +# Default: auto-chat +CODEBUDDY_DEFAULT_MODEL=auto-chat + +# (Optional) Model alias map for UI display labels that are not valid upstream +# model IDs. Comma-separated alias=upstream pairs, matched case-insensitively. +# Example: Claude Opus 4.7=claude-4.0,GPT-5 (UI)=gpt-5 +# Only allowlisted upstream fields (model, messages, stream, tools, tool_choice) +# are forwarded to CodeBuddy; other OpenAI fields are ignored. +CODEBUDDY_MODEL_ALIASES= + # (Optional) Directory containing CodeBuddy credential JSON files. # Default: .codebuddy_creds CODEBUDDY_CREDS_DIR=.codebuddy_creds diff --git a/config.py b/config.py index fda6d9e..7eccb7a 100644 --- a/config.py +++ b/config.py @@ -41,7 +41,9 @@ "CODEBUDDY_UPSTREAM_API_KEY_HEADER": "bearer", "CODEBUDDY_REQUEST_PROFILE": "web", "CODEBUDDY_SANITIZE_AGENT_PROMPT": True, - "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000 + "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000, + "CODEBUDDY_MODEL_ALIASES": "", + "CODEBUDDY_DEFAULT_MODEL": "auto-chat" } # --- Core Functions --- @@ -162,6 +164,34 @@ def get_codebuddy_request_profile() -> str: return profile +def get_codebuddy_default_model() -> str: + value = str(_get_config_value("CODEBUDDY_DEFAULT_MODEL")).strip() + return value or "auto-chat" + + +def get_codebuddy_model_aliases() -> Dict[str, str]: + """Parse CODEBUDDY_MODEL_ALIASES into a lower-cased alias -> upstream map. + + Format: comma-separated ``alias=upstream`` pairs, e.g. + ``Claude Opus 4.7=claude-4.0,gpt-5-ui=gpt-5``. Aliases are matched + case-insensitively; upstream IDs are preserved verbatim. + """ + raw = _get_config_value("CODEBUDDY_MODEL_ALIASES") + aliases: Dict[str, str] = {} + if not raw: + return aliases + for pair in str(raw).split(","): + pair = pair.strip() + if not pair or "=" not in pair: + continue + alias, upstream = pair.split("=", 1) + alias = alias.strip().lower() + upstream = upstream.strip() + if alias and upstream: + aliases[alias] = upstream + return aliases + + def _coerce_bool(value: Any, default: bool) -> bool: if isinstance(value, bool): return value diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index d75df7b..f2ed632 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -29,6 +29,8 @@ sanitize_messages, ) from config import ( + get_codebuddy_default_model, + get_codebuddy_model_aliases, get_codebuddy_request_profile, get_max_system_prompt_length, get_sanitize_agent_prompt, @@ -527,46 +529,74 @@ def _key_fingerprint(token: Optional[str]) -> str: return hashlib.sha256(token.encode("utf-8")).hexdigest()[:8] -def _log_request_diagnostics( - *, - payload: Dict[str, Any], - client_wants_stream: bool, - system_prompt_sanitized: bool, - request_profile: str, - key_fingerprint: str, -) -> None: - """Log request metadata only. Never logs keys, prompts, or user content.""" - messages = payload.get("messages", []) or [] +def _content_types_and_lengths(messages: List[Any]) -> tuple[List[str], List[str], List[int], int]: + """Return (roles, content_types, content_lengths, tool_content_count). + + Metadata only: describes the shape of each message without exposing text. + """ roles: List[str] = [] + content_types: List[str] = [] content_lengths: List[int] = [] - tool_count = 0 + tool_content_count = 0 for msg in messages: if not isinstance(msg, dict): continue roles.append(str(msg.get("role", "unknown"))) content = msg.get("content", "") if isinstance(content, str): + content_types.append("str") content_lengths.append(len(content)) elif isinstance(content, list): + content_types.append("list") length = 0 for item in content: if isinstance(item, dict): if item.get("type") in {"tool_use", "tool_result"}: - tool_count += 1 - length += len(str(item.get("text", ""))) if item.get("type") == "text" else 0 + tool_content_count += 1 + if item.get("type") == "text": + length += len(str(item.get("text", ""))) content_lengths.append(length) else: + content_types.append(type(content).__name__) content_lengths.append(0) - tool_count += len(payload.get("tools", []) or []) + return roles, content_types, content_lengths, tool_content_count + + +def _log_request_diagnostics( + *, + payload: Dict[str, Any], + prep_info: Dict[str, Any], + client_wants_stream: bool, + request_profile: str, + key_fingerprint: str, + outcome: str = "prepared", +) -> None: + """Log safe request metadata for successful and failed requests. + + Never logs message contents, raw keys, or the mapped-away model label as + text beyond what the client requested. ``outcome`` tags the log point + (e.g. "prepared", "success", "failed:"). + """ + messages = payload.get("messages", []) or [] + roles, content_types, content_lengths, tool_content_count = _content_types_and_lengths(messages) + tool_count = tool_content_count + len(payload.get("tools", []) or []) logger.info( - "CodeBuddy request model=%s stream=%s roles=[%s] content_lengths=%s " - "tool_count=%d system_prompt_sanitized=%s request_profile=%s key_fingerprint=%s", - payload.get("model", "unknown"), + "CodeBuddy request outcome=%s requested_model=%s mapped_model=%s stream=%s " + "roles=[%s] content_types=[%s] content_lengths=%s has_tools=%s has_tool_choice=%s " + "tool_count=%d dropped_fields=[%s] system_prompt_sanitized=%s request_profile=%s " + "key_fingerprint=%s", + outcome, + prep_info.get("requested_model", "unknown"), + prep_info.get("mapped_model", payload.get("model", "unknown")), client_wants_stream, ",".join(roles), + ",".join(content_types), content_lengths, + "tools" in payload, + "tool_choice" in payload, tool_count, - system_prompt_sanitized, + ",".join(prep_info.get("dropped_fields", []) or []), + prep_info.get("system_prompt_sanitized", False), request_profile, key_fingerprint, ) @@ -784,18 +814,76 @@ async def handle_non_stream_response( class RequestProcessor: """Request preprocessor - thread-safe request handling""" + # Top-level request fields forwarded to CodeBuddy upstream. Everything else + # (OpenAI-only fields the upstream rejects, or unknown UI fields) is dropped. + # ``tools`` and ``tool_choice`` are forwarded only when present. + _UPSTREAM_ALLOWLIST = {"model", "messages", "stream", "tools", "tool_choice"} + @staticmethod - def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], bool]: - """Prepare the request payload. + def map_model(requested_model: Any) -> str: + """Map a client/UI model label to an upstream CodeBuddy model ID. + + A requested model that already matches a configured available model is + passed through unchanged. Otherwise a configured alias + (CODEBUDDY_MODEL_ALIASES) is applied. Unknown labels fall back to the + configured default model so a UI display label (e.g. "Claude Opus 4.7") + is never sent to upstream verbatim unless it is a known model ID. + """ + try: + available = set(get_available_models_list()) + except Exception: + available = set() + try: + default_model = get_codebuddy_default_model() + except Exception: + default_model = "auto-chat" - Returns the upstream payload and a flag indicating whether an agent - system prompt was sanitized. The flag is used only for safe diagnostics - and is never sent upstream. + if not isinstance(requested_model, str) or not requested_model.strip(): + return default_model + requested = requested_model.strip() + + # Exact match against a known upstream model ID. + if requested in available: + return requested + + # Configured alias mapping (case-insensitive). + try: + aliases = get_codebuddy_model_aliases() + except Exception: + aliases = {} + mapped = aliases.get(requested.lower()) + if mapped: + return mapped + + # Unknown label: do not forward it verbatim; use the safe default. + logger.info( + "Unknown requested model mapped to default: requested_present=%s", + bool(requested), + ) + return default_model + + @staticmethod + def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[str, Any]]: + """Build a strict, upstream-safe payload from the client request body. + + Instead of forwarding the client body verbatim, only allowlisted + top-level fields are forwarded. This prevents OpenAI-only fields (e.g. + response_format, reasoning_effort, stream_options) and unknown UI fields + from reaching CodeBuddy, which is the common cause of upstream rejection. + + Returns ``(payload, prep_info)`` where ``prep_info`` carries safe + metadata for diagnostics (never message content or keys): + * system_prompt_sanitized: bool + * requested_model: str + * mapped_model: str + * dropped_fields: sorted list of ignored top-level field names """ - payload = request_body.copy() - payload["stream"] = True # CodeBuddy only supports streaming requests + source = request_body if isinstance(request_body, dict) else {} + + requested_model = source.get("model") + mapped_model = RequestProcessor.map_model(requested_model) - messages = payload.get("messages", []) + messages = source.get("messages", []) or [] # Sanitize only agent system prompts that tend to trigger false-positive # moderation. User/assistant/tool messages are never modified, and @@ -827,8 +915,32 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], bool] if isinstance(msg, dict) and msg.get("role") == "system": msg["content"] = apply_keyword_replacement_to_system_message(msg.get("content")) - payload["messages"] = messages - return payload, system_prompt_sanitized + # Build the strict upstream payload. CodeBuddy only supports streaming, + # so stream is always True upstream; the client's stream preference is + # honored separately by the router (aggregate vs pass-through). + payload: Dict[str, Any] = { + "model": mapped_model, + "messages": messages, + "stream": True, + } + # Forward tools / tool_choice only when the client actually sent them. + if source.get("tools"): + payload["tools"] = source["tools"] + if source.get("tool_choice") is not None: + payload["tool_choice"] = source["tool_choice"] + + dropped_fields = sorted( + key for key in source.keys() + if key not in RequestProcessor._UPSTREAM_ALLOWLIST + ) + + prep_info = { + "system_prompt_sanitized": system_prompt_sanitized, + "requested_model": requested_model if isinstance(requested_model, str) else "unknown", + "mapped_model": mapped_model, + "dropped_fields": dropped_fields, + } + return payload, prep_info @staticmethod def validate_request(request_body: Dict[str, Any]) -> None: @@ -995,10 +1107,10 @@ async def chat_completions( 503, ) - payload, system_prompt_sanitized = RequestProcessor.prepare_payload(request_body) + payload, prep_info = RequestProcessor.prepare_payload(request_body) usage_stats_manager.record_model_usage(payload.get("model", "unknown")) service = CodeBuddyStreamService() - client_wants_stream = request_body.get("stream", False) + client_wants_stream = bool(request_body.get("stream", False)) if isinstance(request_body, dict) else False excluded_ids: Set[str] = set() last_error: Optional[UpstreamAttemptError] = None @@ -1035,13 +1147,16 @@ async def chat_completions( 500, ) + key_fingerprint = _key_fingerprint(credential.bearer_token) + # Safe diagnostics: metadata only, never key/prompt/user content. _log_request_diagnostics( payload=payload, + prep_info=prep_info, client_wants_stream=client_wants_stream, - system_prompt_sanitized=system_prompt_sanitized, request_profile=request_profile, - key_fingerprint=_key_fingerprint(credential.bearer_token), + key_fingerprint=key_fingerprint, + outcome="prepared", ) headers = codebuddy_api_client.generate_codebuddy_headers( @@ -1064,6 +1179,14 @@ async def chat_completions( result = await service.handle_non_stream_response(payload, headers) if credential.key_id is not None: await codebuddy_api_key_manager.mark_success(credential.key_id) + _log_request_diagnostics( + payload=payload, + prep_info=prep_info, + client_wants_stream=client_wants_stream, + request_profile=request_profile, + key_fingerprint=key_fingerprint, + outcome="success", + ) return result except CodeBuddyModerationError: # The key is valid and the request reached upstream; moderation is @@ -1089,9 +1212,18 @@ async def chat_completions( last_error = error await record_attempt_error(credential, error) logger.warning( - "CodeBuddy upstream attempt failed: source=%s code=%s", + "CodeBuddy upstream attempt failed: source=%s code=%s status=%d", credential.source, error.code, + error.status_code, + ) + _log_request_diagnostics( + payload=payload, + prep_info=prep_info, + client_wants_stream=client_wants_stream, + request_profile=request_profile, + key_fingerprint=key_fingerprint, + outcome=f"failed:{error.code}", ) if source != "api_key_file" or error.kind == "fatal": break diff --git a/tests/test_codebuddy_payload_allowlist.py b/tests/test_codebuddy_payload_allowlist.py new file mode 100644 index 0000000..5e2ec3e --- /dev/null +++ b/tests/test_codebuddy_payload_allowlist.py @@ -0,0 +1,319 @@ +""" +Payload-compatibility regression tests for /codebuddy/v1/chat/completions. + +Reproduces the reported bug: Claude Code works through 9Router + CodeBuddy2API, +but the Shiteru web chat fails ("Upstream CodeBuddy request failed" / empty +reply) because its request body carries OpenAI-only fields and/or a UI model +label that CodeBuddy rejects. + +The fix builds a strict upstream allowlist (model, messages, stream, tools, +tool_choice) and maps model aliases before calling CodeBuddy. These tests lock +in that behavior with the two exact request shapes plus unit coverage of the +allowlist and model mapping. + +Uses a mock upstream transport; no real CodeBuddy API key is required. +""" +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.codebuddy_router import RequestProcessor +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + +AVAILABLE_MODELS = ["claude-4.0", "gpt-5", "auto-chat"] + +# OpenAI fields Shiteru-style clients send that CodeBuddy does not accept. +UNSUPPORTED_FIELDS = { + "response_format": {"type": "json_object"}, + "parallel_tool_calls": False, + "reasoning_effort": "high", + "stream_options": {"include_usage": True}, + "service_tier": "auto", + "store": True, + "metadata": {"session": "abc"}, + "seed": 42, + "logprobs": True, + "top_logprobs": 5, + "prediction": {"type": "content", "content": "x"}, + "modalities": ["text"], + "temperature": 0.7, + "top_p": 0.9, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "max_tokens": 1024, + "user": "shiteru-web", +} + + +# --------------------------------------------------------------------------- # +# Unit: model mapping +# --------------------------------------------------------------------------- # + + +def _patch_models(monkeypatch, aliases=None, default="auto-chat"): + monkeypatch.setattr(codebuddy_router, "get_available_models_list", lambda: list(AVAILABLE_MODELS)) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_default_model", lambda: default) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_model_aliases", lambda: dict(aliases or {})) + + +def test_known_model_passthrough(monkeypatch): + _patch_models(monkeypatch) + assert RequestProcessor.map_model("claude-4.0") == "claude-4.0" + + +def test_alias_mapped_case_insensitively(monkeypatch): + _patch_models(monkeypatch, aliases={"claude opus 4.7": "claude-4.0"}) + assert RequestProcessor.map_model("Claude Opus 4.7") == "claude-4.0" + + +def test_unknown_label_falls_back_to_default(monkeypatch): + _patch_models(monkeypatch, default="auto-chat") + # A UI display label with no alias must NOT be forwarded verbatim. + assert RequestProcessor.map_model("Claude Opus 4.7") == "auto-chat" + + +def test_missing_model_uses_default(monkeypatch): + _patch_models(monkeypatch, default="auto-chat") + assert RequestProcessor.map_model(None) == "auto-chat" + assert RequestProcessor.map_model("") == "auto-chat" + + +# --------------------------------------------------------------------------- # +# Unit: allowlist +# --------------------------------------------------------------------------- # + + +def test_prepare_payload_drops_unsupported_fields(monkeypatch): + _patch_models(monkeypatch) + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: True) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + + body = { + "model": "claude-4.0", + "messages": [{"role": "user", "content": "hi"}], + "stream": False, + "tools": [{"type": "function", "function": {"name": "f"}}], + "tool_choice": "auto", + **UNSUPPORTED_FIELDS, + } + payload, prep = RequestProcessor.prepare_payload(body) + + # Only allowlisted fields survive. + assert set(payload.keys()) == {"model", "messages", "stream", "tools", "tool_choice"} + # Upstream always receives stream=True (CodeBuddy is SSE-only). + assert payload["stream"] is True + # Every unsupported field is reported as dropped. + for field in UNSUPPORTED_FIELDS: + assert field in prep["dropped_fields"] + assert prep["mapped_model"] == "claude-4.0" + + +def test_prepare_payload_omits_absent_tools(monkeypatch): + _patch_models(monkeypatch) + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: True) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + + body = {"model": "gpt-5", "messages": [{"role": "user", "content": "hi"}]} + payload, _prep = RequestProcessor.prepare_payload(body) + assert "tools" not in payload + assert "tool_choice" not in payload + + +# --------------------------------------------------------------------------- # +# Integration fixtures +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def configure(monkeypatch, aliases=None, default="auto-chat"): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: "web") + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: True) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + _patch_models(monkeypatch, aliases=aliases, default=default) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse(text, finish="stop"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":{json.dumps(text)}}},"finish_reason":{json.dumps(finish)}}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +async def post(app, token, body): + transport = httpx.ASGITransport(app=app) + headers = {"Authorization": f"Bearer {token}"} + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post( + "/codebuddy/v1/chat/completions", headers=headers, json=body + ) + + +# --------------------------------------------------------------------------- # +# The two exact request shapes +# --------------------------------------------------------------------------- # + +# A Claude Code request: valid model, system + user messages, tools, streaming. +CLAUDE_CODE_REQUEST = { + "model": "claude-4.0", + "messages": [ + {"role": "system", "content": "You are a concise assistant."}, + {"role": "user", "content": "list two prime numbers"}, + ], + "stream": True, + "tools": [ + { + "type": "function", + "function": {"name": "noop", "parameters": {"type": "object"}}, + } + ], + "tool_choice": "auto", +} + +# A Shiteru web-chat request: UI display-label model + many OpenAI-only fields +# and stream=false. This is the shape that currently fails upstream. +SHITERU_WEB_REQUEST = { + "model": "Claude Opus 4.7", + "messages": [{"role": "user", "content": "halo, siapa kamu"}], + "stream": False, + "response_format": {"type": "text"}, + "reasoning_effort": "medium", + "parallel_tool_calls": True, + "stream_options": {"include_usage": True}, + "temperature": 0.6, + "top_p": 1.0, + "max_tokens": 800, + "metadata": {"ui": "shiteru"}, + "seed": 7, +} + + +@pytest.mark.asyncio +async def test_claude_code_request_succeeds(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("2 and 3") + + upstream = install_upstream(monkeypatch, handler) + response = await post(app, KEY_A, CLAUDE_CODE_REQUEST) + await upstream.aclose() + + assert response.status_code == 200 + # Streaming client -> SSE passthrough. + assert response.headers["content-type"].startswith("text/event-stream") + assert "[DONE]" in response.text + # Upstream received only allowlisted fields, model unchanged, stream=True. + body = seen["body"] + assert set(body.keys()) == {"model", "messages", "stream", "tools", "tool_choice"} + assert body["model"] == "claude-4.0" + assert body["stream"] is True + assert body["tools"] == CLAUDE_CODE_REQUEST["tools"] + + +@pytest.mark.asyncio +async def test_shiteru_web_request_now_succeeds(monkeypatch, app, empty_pool): + # Map the UI label so it resolves to a real upstream model. + configure(monkeypatch, aliases={"claude opus 4.7": "claude-4.0"}) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("Halo! Saya asisten AI.") + + upstream = install_upstream(monkeypatch, handler) + response = await post(app, KEY_A, SHITERU_WEB_REQUEST) + await upstream.aclose() + + # No more "Upstream CodeBuddy request failed": a proper ChatCompletion. + assert response.status_code == 200 + assert response.headers["content-type"].startswith("application/json") + data = response.json() + assert data["object"] == "chat.completion" + choice = data["choices"][0] + assert choice["message"]["role"] == "assistant" + assert choice["message"]["content"] == "Halo! Saya asisten AI." + assert choice["finish_reason"] == "stop" + + # Upstream must NOT have received any of the unsupported OpenAI fields... + body = seen["body"] + assert set(body.keys()) == {"model", "messages", "stream"} + for field in ( + "response_format", "reasoning_effort", "parallel_tool_calls", + "stream_options", "temperature", "top_p", "max_tokens", "metadata", "seed", + ): + assert field not in body + # ...and the UI label must be mapped to a real model ID. + assert body["model"] == "claude-4.0" + # CodeBuddy is SSE-only: upstream stream is always True even for stream=false. + assert body["stream"] is True + + +@pytest.mark.asyncio +async def test_shiteru_unmapped_label_uses_default_not_verbatim( + monkeypatch, app, empty_pool +): + # No alias configured: the unknown UI label must fall back to the default, + # never be sent verbatim. + configure(monkeypatch, default="auto-chat") + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + response = await post(app, KEY_A, SHITERU_WEB_REQUEST) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["body"]["model"] == "auto-chat" + assert seen["body"]["model"] != "Claude Opus 4.7" From 695ee36cfa3a6274f6ac13471c0e5588f45897e9 Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 11:40:51 +0700 Subject: [PATCH 5/9] fix: SSE Error --- config.py | 20 +- src/codebuddy_message_sanitizer.py | 154 +++++- src/codebuddy_router.py | 173 +++++-- tests/test_codebuddy_message_normalization.py | 456 ++++++++++++++++++ tests/test_codebuddy_payload_allowlist.py | 179 ++++++- 5 files changed, 934 insertions(+), 48 deletions(-) create mode 100644 tests/test_codebuddy_message_normalization.py diff --git a/config.py b/config.py index 7eccb7a..d1cab05 100644 --- a/config.py +++ b/config.py @@ -43,7 +43,8 @@ "CODEBUDDY_SANITIZE_AGENT_PROMPT": True, "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000, "CODEBUDDY_MODEL_ALIASES": "", - "CODEBUDDY_DEFAULT_MODEL": "auto-chat" + "CODEBUDDY_DEFAULT_MODEL": "auto-chat", + "CODEBUDDY_UNKNOWN_MODEL_POLICY": "passthrough" } # --- Core Functions --- @@ -169,6 +170,23 @@ def get_codebuddy_default_model() -> str: return value or "auto-chat" +def get_codebuddy_unknown_model_policy() -> str: + """How to handle a requested model that is neither a known upstream model + ID nor a configured alias. + + * passthrough (default): forward the requested model verbatim and let + CodeBuddy accept or reject it. Never silently rewrite it. + * reject: return HTTP 400 (code=unknown_model) without calling upstream. + * default: fall back to CODEBUDDY_DEFAULT_MODEL. + """ + policy = str(_get_config_value("CODEBUDDY_UNKNOWN_MODEL_POLICY")).strip().lower() + if policy not in {"passthrough", "reject", "default"}: + raise ValueError( + "CODEBUDDY_UNKNOWN_MODEL_POLICY must be passthrough, reject, or default" + ) + return policy + + def get_codebuddy_model_aliases() -> Dict[str, str]: """Parse CODEBUDDY_MODEL_ALIASES into a lower-cased alias -> upstream map. diff --git a/src/codebuddy_message_sanitizer.py b/src/codebuddy_message_sanitizer.py index 091aaa3..577ecb9 100644 --- a/src/codebuddy_message_sanitizer.py +++ b/src/codebuddy_message_sanitizer.py @@ -16,7 +16,7 @@ refusal as an assistant reply. """ import logging -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple logger = logging.getLogger(__name__) @@ -87,6 +87,158 @@ def _looks_like_agent_prompt(text: str) -> bool: return any(marker in lowered for marker in _AGENT_PROMPT_MARKERS) +# Roles CodeBuddy accepts on an upstream message. +_VALID_UPSTREAM_ROLES: Tuple[str, ...] = ("system", "user", "assistant", "tool") + + +class MessageNormalizationError(Exception): + """Raised when an upstream message cannot be given a valid role. + + Carries the offending message index so the router can return a clear local + HTTP 400 identifying which message is malformed, instead of forwarding data + that CodeBuddy rejects with a generic "Message N must have 'role' and + 'content' fields". + """ + + def __init__(self, index: int, reason: str): + super().__init__(f"message {index}: {reason}") + self.index = index + self.reason = reason + + +def _log_structural(level: int, prefix: str, index: int, msg: Any) -> None: + """Log safe structural metadata for a message. + + Logs the message index, role, its field-name set, the content field's type, + and whether it carries tool_calls / tool_call_id. It never logs message + content values or any credential: the field names logged (role, content, + tool_calls, ...) are fixed OpenAI-schema identifiers, not user data. + """ + if isinstance(msg, dict): + role = msg.get("role") + field_names = ",".join(sorted(str(k) for k in msg.keys())) + content_type = type(msg.get("content")).__name__ + has_tool_calls = bool(msg.get("tool_calls")) + has_tool_call_id = bool(msg.get("tool_call_id")) + else: + role = None + field_names = "" + content_type = type(msg).__name__ + has_tool_calls = False + has_tool_call_id = False + + logger.log( + level, + "%s index=%d role=%s keys=[%s] content_type=%s has_tool_calls=%s " + "has_tool_call_id=%s", + prefix, + index, + role, + field_names, + content_type, + has_tool_calls, + has_tool_call_id, + ) + + +def _infer_role(msg: Dict[str, Any]) -> Any: + """Infer a message role from structure when it is missing or blank. + + Returns a valid role string, or ``None`` when it genuinely cannot be + determined. + """ + role = msg.get("role") + if isinstance(role, str) and role.strip(): + return role.strip() + # A message carrying a tool_call_id is a tool result. + if msg.get("tool_call_id"): + return "tool" + # A message carrying tool_calls is an assistant turn. + if msg.get("tool_calls"): + return "assistant" + return None + + +def normalize_messages_for_upstream( + messages: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Return a copy of ``messages`` guaranteed valid for the CodeBuddy upstream. + + This is the final step before the upstream request, run after any + Anthropic/OpenAI/tool conversion and system-prompt sanitization, so it + inspects the exact messages CodeBuddy will receive rather than the original + client messages. It guarantees every message has: + + * a valid, non-empty ``role`` + * a ``content`` field that is never missing or ``null`` + + Rules: + * A message with an explicit non-empty ``role`` keeps it verbatim; a role + that is present is considered determined even if it is not one of the + common OpenAI roles, so legitimate conversations are never rejected on a + role-value technicality. + * Assistant messages carrying ``tool_calls`` keep ``content: ""`` when + content is missing or null, so an assistant tool-call turn is never sent + without a content field. + * Tool-result messages resolve to role ``tool`` (inferred from + ``tool_call_id`` when the role is absent) and keep their + ``tool_call_id`` and ``content``. + * Existing non-null content is preserved verbatim, including multimodal + content arrays CodeBuddy already accepts; an empty string is only added + when content is missing or ``null``. + * A message whose role genuinely cannot be determined (absent/blank role + with no ``tool_call_id`` or ``tool_calls`` to infer from) raises + :class:`MessageNormalizationError` with its index, so the caller can + return a local HTTP 400 instead of forwarding malformed data. + """ + normalized: List[Dict[str, Any]] = [] + repaired = 0 + + for index, msg in enumerate(messages): + # Structural, content-free trace of every message inspected upstream. + _log_structural(logging.DEBUG, "upstream message", index, msg) + + if not isinstance(msg, dict): + _log_structural( + logging.WARNING, "upstream message not an object", index, msg + ) + raise MessageNormalizationError(index, "message is not an object") + + role = _infer_role(msg) + if not role: + _log_structural( + logging.WARNING, "upstream message role undeterminable", index, msg + ) + raise MessageNormalizationError( + index, + "role is missing and cannot be inferred from message structure", + ) + + new_msg = dict(msg) + new_msg["role"] = role + + # Add an empty content string only when the field is missing or null; + # never overwrite existing content (strings, empty strings, or + # multimodal arrays CodeBuddy already accepts). + if "content" not in new_msg or new_msg["content"] is None: + new_msg["content"] = "" + repaired += 1 + _log_structural( + logging.INFO, "upstream message content defaulted to empty", index, msg + ) + + normalized.append(new_msg) + + if repaired: + logger.info( + "Normalized upstream messages: total=%d content_defaulted=%d", + len(messages), + repaired, + ) + + return normalized + + def sanitize_messages( messages: List[Dict[str, Any]], *, diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index f2ed632..d8944ff 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -25,13 +25,16 @@ from .usage_stats_manager import usage_stats_manager from .keyword_replacer import apply_keyword_replacement_to_system_message from .codebuddy_message_sanitizer import ( + MessageNormalizationError, is_codebuddy_moderation_response, + normalize_messages_for_upstream, sanitize_messages, ) from config import ( get_codebuddy_default_model, get_codebuddy_model_aliases, get_codebuddy_request_profile, + get_codebuddy_unknown_model_policy, get_max_system_prompt_length, get_sanitize_agent_prompt, get_upstream_api_key_header, @@ -468,6 +471,19 @@ class CodeBuddyModerationError(Exception): """ +class UnknownModelError(Exception): + """Raised when a requested model is unknown and the policy is 'reject'. + + Signals that the request must be rejected with HTTP 400 (code + ``unknown_model``) before any upstream call, instead of silently rewriting + the requested model to a fallback. + """ + + def __init__(self, requested_model: str): + super().__init__(requested_model) + self.requested_model = requested_model + + # Human-readable message returned to the client on a moderation rejection. MODERATION_MESSAGE = ( "CodeBuddy rejected the request through its content moderation system. " @@ -570,24 +586,34 @@ def _log_request_diagnostics( request_profile: str, key_fingerprint: str, outcome: str = "prepared", + upstream_response_model: Optional[str] = None, ) -> None: """Log safe request metadata for successful and failed requests. Never logs message contents, raw keys, or the mapped-away model label as text beyond what the client requested. ``outcome`` tags the log point (e.g. "prepared", "success", "failed:"). + + ``requested_model`` (what the client asked for), ``mapped_model`` (what we + send upstream), ``mapping_source`` (how it was resolved), and + ``upstream_response_model`` (what CodeBuddy actually reported) are logged + separately so a divergence between the selected UI model and the upstream + model is visible rather than hidden. """ messages = payload.get("messages", []) or [] roles, content_types, content_lengths, tool_content_count = _content_types_and_lengths(messages) tool_count = tool_content_count + len(payload.get("tools", []) or []) logger.info( - "CodeBuddy request outcome=%s requested_model=%s mapped_model=%s stream=%s " + "CodeBuddy request outcome=%s requested_model=%s mapped_model=%s mapping_source=%s " + "upstream_response_model=%s stream=%s " "roles=[%s] content_types=[%s] content_lengths=%s has_tools=%s has_tool_choice=%s " "tool_count=%d dropped_fields=[%s] system_prompt_sanitized=%s request_profile=%s " "key_fingerprint=%s", outcome, prep_info.get("requested_model", "unknown"), prep_info.get("mapped_model", payload.get("model", "unknown")), + prep_info.get("mapping_source", "unknown"), + upstream_response_model or "unknown", client_wants_stream, ",".join(roles), ",".join(content_types), @@ -820,14 +846,22 @@ class RequestProcessor: _UPSTREAM_ALLOWLIST = {"model", "messages", "stream", "tools", "tool_choice"} @staticmethod - def map_model(requested_model: Any) -> str: - """Map a client/UI model label to an upstream CodeBuddy model ID. - - A requested model that already matches a configured available model is - passed through unchanged. Otherwise a configured alias - (CODEBUDDY_MODEL_ALIASES) is applied. Unknown labels fall back to the - configured default model so a UI display label (e.g. "Claude Opus 4.7") - is never sent to upstream verbatim unless it is a known model ID. + def resolve_model(requested_model: Any) -> tuple[str, str]: + """Resolve a client/UI model label to an upstream CodeBuddy model ID. + + Returns ``(mapped_model, mapping_source)`` where ``mapping_source`` is: + * ``"exact"`` - matched a configured available model ID + * ``"alias"`` - matched a configured alias (case-insensitive) + * ``"default_empty"`` - no model supplied; used the default + * ``"passthrough"`` - unknown label forwarded verbatim + * ``"default"`` - unknown label mapped to the default model + + An unknown label (neither a known model ID nor a configured alias) is + handled per CODEBUDDY_UNKNOWN_MODEL_POLICY: + * ``passthrough`` (default): forward it verbatim and let CodeBuddy + accept or reject it. The requested model is never silently rewritten. + * ``reject``: raise :class:`UnknownModelError` (HTTP 400 upstream). + * ``default``: fall back to CODEBUDDY_DEFAULT_MODEL. """ try: available = set(get_available_models_list()) @@ -838,13 +872,16 @@ def map_model(requested_model: Any) -> str: except Exception: default_model = "auto-chat" + # No usable model supplied: there is nothing to pass through, so use + # the configured default regardless of policy. if not isinstance(requested_model, str) or not requested_model.strip(): - return default_model + return default_model, "default_empty" requested = requested_model.strip() - # Exact match against a known upstream model ID. + # Exact match against a known upstream model ID (this covers a client + # that explicitly requests "auto-chat"). if requested in available: - return requested + return requested, "exact" # Configured alias mapping (case-insensitive). try: @@ -853,14 +890,34 @@ def map_model(requested_model: Any) -> str: aliases = {} mapped = aliases.get(requested.lower()) if mapped: - return mapped + return mapped, "alias" - # Unknown label: do not forward it verbatim; use the safe default. - logger.info( - "Unknown requested model mapped to default: requested_present=%s", - bool(requested), - ) - return default_model + # Unknown label: behavior is governed by the configured policy. Never + # silently rewrite an unknown model to the default. + try: + policy = get_codebuddy_unknown_model_policy() + except Exception: + policy = "passthrough" + + if policy == "reject": + logger.info("Unknown requested model rejected (policy=reject)") + raise UnknownModelError(requested) + if policy == "default": + logger.info("Unknown requested model mapped to default (policy=default)") + return default_model, "default" + # passthrough (default): forward the requested model verbatim so the + # real upstream model is preserved and CodeBuddy decides if it is valid. + return requested, "passthrough" + + @staticmethod + def map_model(requested_model: Any) -> str: + """Backward-compatible wrapper returning only the resolved model ID. + + May raise :class:`UnknownModelError` when the unknown-model policy is + ``reject``. + """ + mapped, _source = RequestProcessor.resolve_model(requested_model) + return mapped @staticmethod def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[str, Any]]: @@ -876,12 +933,16 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[ * system_prompt_sanitized: bool * requested_model: str * mapped_model: str + * mapping_source: str (how the model was resolved) * dropped_fields: sorted list of ignored top-level field names + + May raise :class:`UnknownModelError` when the requested model is unknown + and CODEBUDDY_UNKNOWN_MODEL_POLICY is ``reject``. """ source = request_body if isinstance(request_body, dict) else {} requested_model = source.get("model") - mapped_model = RequestProcessor.map_model(requested_model) + mapped_model, mapping_source = RequestProcessor.resolve_model(requested_model) messages = source.get("messages", []) or [] @@ -915,6 +976,16 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[ if isinstance(msg, dict) and msg.get("role") == "system": msg["content"] = apply_keyword_replacement_to_system_message(msg.get("content")) + # Final transform before upstream: after all Anthropic/OpenAI/tool + # conversion and system-prompt sanitization, guarantee every message + # has a valid role and a content field. This inspects the exact messages + # CodeBuddy will receive (e.g. Claude Code assistant tool_calls turns + # with null content, or tool-result messages), preventing the upstream + # "Message N must have 'role' and 'content' fields" rejection. Raises + # MessageNormalizationError (handled by the caller) when a role cannot + # be determined. + messages = normalize_messages_for_upstream(messages) + # Build the strict upstream payload. CodeBuddy only supports streaming, # so stream is always True upstream; the client's stream preference is # honored separately by the router (aggregate vs pass-through). @@ -938,29 +1009,36 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[ "system_prompt_sanitized": system_prompt_sanitized, "requested_model": requested_model if isinstance(requested_model, str) else "unknown", "mapped_model": mapped_model, + "mapping_source": mapping_source, "dropped_fields": dropped_fields, } return payload, prep_info @staticmethod def validate_request(request_body: Dict[str, Any]) -> None: - """Validate request parameters""" + """Validate request structure. + + Only structural invariants are enforced here: the body must be an + object with a non-empty ``messages`` array of objects. Role/content + completeness is intentionally NOT enforced at this point: valid + OpenAI/Anthropic tool-use turns legitimately omit ``content`` (assistant + messages carrying ``tool_calls``) or ``role`` (tool results identified + by ``tool_call_id``). Those fields are guaranteed later by + ``normalize_messages_for_upstream`` immediately before the upstream call, + which repairs what it can and returns an indexed 400 for what it cannot. + """ if not isinstance(request_body, dict): raise HTTPException(status_code=400, detail="Request body must be a JSON object") - + messages = request_body.get("messages") if not messages or not isinstance(messages, list): raise HTTPException(status_code=400, detail="Messages field is required and must be an array") - - if not messages: - raise HTTPException(status_code=400, detail="At least one message is required") - - # Validate message format + + # Validate message container type only; field-level normalization runs + # later against the final transformed messages. for i, msg in enumerate(messages): if not isinstance(msg, dict): raise HTTPException(status_code=400, detail=f"Message {i} must be an object") - if "role" not in msg or "content" not in msg: - raise HTTPException(status_code=400, detail=f"Message {i} must have 'role' and 'content' fields") @dataclass(repr=False) class ResolvedCredential: @@ -1107,7 +1185,35 @@ async def chat_completions( 503, ) - payload, prep_info = RequestProcessor.prepare_payload(request_body) + try: + payload, prep_info = RequestProcessor.prepare_payload(request_body) + except UnknownModelError as exc: + logger.info( + "Rejecting unknown requested model (policy=reject): requested_present=%s", + bool(exc.requested_model), + ) + return openai_error_response( + f"Unknown model: {exc.requested_model}", + "invalid_request_error", + "unknown_model", + 400, + ) + except MessageNormalizationError as exc: + # A message reached the upstream boundary without a determinable role. + # Reject locally with the offending index instead of forwarding malformed + # data that CodeBuddy rejects with a generic "must have 'role' and + # 'content'" error. + logger.info( + "Rejecting malformed upstream message: index=%d reason=%s", + exc.index, + exc.reason, + ) + return openai_error_response( + f"Message {exc.index} is malformed: {exc.reason}", + "invalid_request_error", + "invalid_message", + 400, + ) usage_stats_manager.record_model_usage(payload.get("model", "unknown")) service = CodeBuddyStreamService() client_wants_stream = bool(request_body.get("stream", False)) if isinstance(request_body, dict) else False @@ -1179,6 +1285,12 @@ async def chat_completions( result = await service.handle_non_stream_response(payload, headers) if credential.key_id is not None: await codebuddy_api_key_manager.mark_success(credential.key_id) + # For the non-stream path we have the aggregated upstream response + # and can log the model CodeBuddy actually reported. It is logged + # separately and never used to rewrite the response model. + upstream_response_model = ( + result.get("model") if isinstance(result, dict) else None + ) _log_request_diagnostics( payload=payload, prep_info=prep_info, @@ -1186,6 +1298,7 @@ async def chat_completions( request_profile=request_profile, key_fingerprint=key_fingerprint, outcome="success", + upstream_response_model=upstream_response_model, ) return result except CodeBuddyModerationError: diff --git a/tests/test_codebuddy_message_normalization.py b/tests/test_codebuddy_message_normalization.py new file mode 100644 index 0000000..7d0728d --- /dev/null +++ b/tests/test_codebuddy_message_normalization.py @@ -0,0 +1,456 @@ +""" +Upstream message-normalization regression tests. + +Reproduces and locks the fix for the reported upstream failure during Claude +Code tool-use conversations: + + Message 13 must have 'role' and 'content' fields + +CodeBuddy rejects any message missing a role or content. Valid OpenAI/Anthropic +tool-use turns legitimately omit those fields (an assistant turn carrying +``tool_calls`` has null/absent content; a tool result is identified by +``tool_call_id`` and may omit ``role``). ``normalize_messages_for_upstream`` +runs as the final transform before the upstream request and guarantees every +message has a valid role and a content field, while preserving content +CodeBuddy already accepts (including multimodal arrays). + +These tests assert both the unit behavior of the normalizer and the exact +messages that reach the (mocked) upstream through the full request path. + +Uses a mock upstream transport; no real CodeBuddy API key is required. +""" +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.codebuddy_router import RequestProcessor +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager +from src.codebuddy_message_sanitizer import ( + MessageNormalizationError, + normalize_messages_for_upstream, +) + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + +AVAILABLE_MODELS = ["claude-4.0", "gpt-5", "auto-chat"] + + +# --------------------------------------------------------------------------- # +# Fixtures / helpers (mirror test_codebuddy_payload_allowlist.py) +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def _patch_models(monkeypatch): + monkeypatch.setattr(codebuddy_router, "get_available_models_list", lambda: list(AVAILABLE_MODELS)) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_default_model", lambda: "auto-chat") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_model_aliases", lambda: {}) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_unknown_model_policy", lambda: "passthrough") + + +def configure(monkeypatch): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: "web") + # Disable system-prompt sanitization so it does not perturb the message + # sequences under test; normalization is independent of it. + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: False) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + _patch_models(monkeypatch) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse(text, finish="stop"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":{json.dumps(text)}}},"finish_reason":{json.dumps(finish)}}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +async def post(app, token, body): + transport = httpx.ASGITransport(app=app) + headers = {"Authorization": f"Bearer {token}"} + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post( + "/codebuddy/v1/chat/completions", headers=headers, json=body + ) + + +def _assert_all_valid_upstream(messages): + """Every upstream message must have a non-empty role and a content field.""" + for i, msg in enumerate(messages): + assert isinstance(msg, dict), f"message {i} is not an object" + role = msg.get("role") + assert isinstance(role, str) and role.strip(), f"message {i} has no valid role" + assert "content" in msg, f"message {i} is missing content" + assert msg["content"] is not None, f"message {i} has null content" + + +# --------------------------------------------------------------------------- # +# Unit: normalize_messages_for_upstream +# --------------------------------------------------------------------------- # + + +def test_assistant_tool_calls_missing_content_gets_empty_string(): + tool_calls = [ + { + "id": "call_abc", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Jakarta"}'}, + } + ] + messages = [ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "tool_calls": tool_calls}, # no content field + ] + out = normalize_messages_for_upstream(messages) + _assert_all_valid_upstream(out) + assert out[1]["role"] == "assistant" + assert out[1]["content"] == "" + # tool_calls preserved verbatim. + assert out[1]["tool_calls"] == tool_calls + + +def test_assistant_tool_calls_null_content_gets_empty_string(): + messages = [ + {"role": "assistant", "content": None, "tool_calls": [{"id": "c1"}]}, + ] + out = normalize_messages_for_upstream(messages) + assert out[0]["content"] == "" + assert out[0]["role"] == "assistant" + + +def test_tool_result_missing_role_inferred_from_tool_call_id(): + messages = [ + # No role, but tool_call_id identifies it as a tool result. + {"tool_call_id": "call_abc", "content": "23C and sunny"}, + ] + out = normalize_messages_for_upstream(messages) + _assert_all_valid_upstream(out) + assert out[0]["role"] == "tool" + assert out[0]["tool_call_id"] == "call_abc" + assert out[0]["content"] == "23C and sunny" + + +def test_tool_result_shape_is_preserved(): + messages = [ + {"role": "tool", "tool_call_id": "call_xyz", "content": "result text"}, + ] + out = normalize_messages_for_upstream(messages) + assert out[0] == { + "role": "tool", + "tool_call_id": "call_xyz", + "content": "result text", + } + + +def test_existing_empty_string_content_is_not_touched(): + messages = [{"role": "assistant", "content": "", "tool_calls": [{"id": "c1"}]}] + out = normalize_messages_for_upstream(messages) + assert out[0]["content"] == "" + + +def test_multimodal_content_array_is_preserved(): + multimodal = [ + {"type": "text", "text": "look at this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}, + ] + messages = [{"role": "user", "content": multimodal}] + out = normalize_messages_for_upstream(messages) + _assert_all_valid_upstream(out) + # The array must be preserved verbatim, not flattened or replaced. + assert out[0]["content"] == multimodal + assert isinstance(out[0]["content"], list) + + +def test_undeterminable_role_raises_with_index(): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"content": "orphan with no role and no tool markers"}, # index 2 + ] + with pytest.raises(MessageNormalizationError) as excinfo: + normalize_messages_for_upstream(messages) + assert excinfo.value.index == 2 + + +def test_blank_role_string_is_treated_as_missing(): + messages = [{"role": " ", "content": "x"}] + with pytest.raises(MessageNormalizationError) as excinfo: + normalize_messages_for_upstream(messages) + assert excinfo.value.index == 0 + + +def test_normalizer_does_not_mutate_input(): + messages = [{"role": "assistant", "tool_calls": [{"id": "c1"}]}] + normalize_messages_for_upstream(messages) + # Original list/message untouched (copy-on-write). + assert "content" not in messages[0] + + +# --------------------------------------------------------------------------- # +# Integration: full request path, asserting what reaches upstream +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_claude_code_tool_loop_reaches_upstream_valid(monkeypatch, app, empty_pool): + """Assistant tool_calls (no content) + tool result (no role) both survive.""" + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("done") + + upstream = install_upstream(monkeypatch, handler) + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": "weather in Jakarta?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Jakarta"}'}, + } + ], + }, + {"tool_call_id": "call_1", "content": "31C, humid"}, + {"role": "user", "content": "and tomorrow?"}, + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + _assert_all_valid_upstream(sent) + # Assistant tool-call turn keeps its tool_calls and gains content: "". + assert sent[1]["role"] == "assistant" + assert sent[1]["content"] == "" + assert sent[1]["tool_calls"][0]["id"] == "call_1" + # Tool result got role tool inferred from tool_call_id. + assert sent[2]["role"] == "tool" + assert sent[2]["tool_call_id"] == "call_1" + + +@pytest.mark.asyncio +async def test_multiple_claude_code_tool_loops(monkeypatch, app, empty_pool): + """Several back-to-back tool loops all normalize correctly.""" + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + + messages = [{"role": "user", "content": "start a multi-step task"}] + for n in range(3): + messages.append( + { + "role": "assistant", + "tool_calls": [ + { + "id": f"call_{n}", + "type": "function", + "function": {"name": "step", "arguments": f'{{"n":{n}}}'}, + } + ], + } + ) + # Tool result with no role (identified by tool_call_id). + messages.append({"tool_call_id": f"call_{n}", "content": f"step {n} done"}) + messages.append({"role": "assistant", "content": "all steps complete"}) + + body = {"model": "claude-4.0", "stream": False, "messages": messages} + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + _assert_all_valid_upstream(sent) + # Each tool-call assistant turn has content "" and each tool result role tool. + assistants = [m for m in sent if m["role"] == "assistant" and m.get("tool_calls")] + tools = [m for m in sent if m["role"] == "tool"] + assert len(assistants) == 3 + assert len(tools) == 3 + assert all(m["content"] == "" for m in assistants) + assert all(m["tool_call_id"].startswith("call_") for m in tools) + + +@pytest.mark.asyncio +async def test_long_conversation_at_least_15_messages(monkeypatch, app, empty_pool): + """A >=15-message Claude Code conversation reaches upstream fully valid. + + Regression for "Message 13 must have 'role' and 'content' fields": the + message at index 13 is an assistant tool-call turn with no content. + """ + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("final answer") + + upstream = install_upstream(monkeypatch, handler) + + messages = [{"role": "system", "content": "You are helpful."}] + # Build alternating user / assistant-tool-call / tool-result turns until + # we have well over 15 messages, including one at index 13. + for n in range(6): + messages.append({"role": "user", "content": f"question {n}"}) + messages.append( + { + "role": "assistant", + "tool_calls": [ + { + "id": f"call_{n}", + "type": "function", + "function": {"name": "lookup", "arguments": f'{{"q":{n}}}'}, + } + ], + } + ) + messages.append({"tool_call_id": f"call_{n}", "content": f"answer {n}"}) + messages.append({"role": "user", "content": "summarize"}) + + assert len(messages) >= 15 + # Sanity: index 13 is the kind of message that used to fail upstream. + assert "content" not in messages[13] or messages[13].get("content") is None + + body = {"model": "claude-4.0", "stream": False, "messages": messages} + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + assert len(sent) >= 15 + _assert_all_valid_upstream(sent) + + +@pytest.mark.asyncio +async def test_multimodal_array_survives_full_path(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + multimodal = [ + {"type": "text", "text": "describe this image"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + ] + body = { + "model": "claude-4.0", + "stream": False, + "messages": [{"role": "user", "content": multimodal}], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + # The multimodal content array reaches upstream unchanged. + user_msg = next(m for m in sent if m["role"] == "user") + assert user_msg["content"] == multimodal + + +@pytest.mark.asyncio +async def test_undeterminable_role_returns_local_400_with_index( + monkeypatch, app, empty_pool +): + """A message with no determinable role is rejected locally, not forwarded.""" + configure(monkeypatch) + called = {"upstream": False} + + def handler(req): + called["upstream"] = True + return sse("should not be called") + + upstream = install_upstream(monkeypatch, handler) + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": "hi"}, + {"content": "orphan with no role and no tool markers"}, # index 1 + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 400 + assert called["upstream"] is False + error = response.json()["error"] + assert error["code"] == "invalid_message" + assert error["type"] == "invalid_request_error" + # The message index is identified in the error message. + assert "1" in error["message"] + + +@pytest.mark.asyncio +async def test_prepare_payload_normalizes_missing_content(monkeypatch): + _patch_models(monkeypatch) + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: False) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + + body = { + "model": "claude-4.0", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "tool_calls": [{"id": "c1"}]}, + {"tool_call_id": "c1", "content": "tool output"}, + ], + } + payload, _prep = RequestProcessor.prepare_payload(body) + _assert_all_valid_upstream(payload["messages"]) diff --git a/tests/test_codebuddy_payload_allowlist.py b/tests/test_codebuddy_payload_allowlist.py index 5e2ec3e..6d1ac3c 100644 --- a/tests/test_codebuddy_payload_allowlist.py +++ b/tests/test_codebuddy_payload_allowlist.py @@ -56,32 +56,75 @@ # --------------------------------------------------------------------------- # -def _patch_models(monkeypatch, aliases=None, default="auto-chat"): +def _patch_models(monkeypatch, aliases=None, default="auto-chat", policy="passthrough"): monkeypatch.setattr(codebuddy_router, "get_available_models_list", lambda: list(AVAILABLE_MODELS)) monkeypatch.setattr(codebuddy_router, "get_codebuddy_default_model", lambda: default) monkeypatch.setattr(codebuddy_router, "get_codebuddy_model_aliases", lambda: dict(aliases or {})) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_unknown_model_policy", lambda: policy) def test_known_model_passthrough(monkeypatch): _patch_models(monkeypatch) assert RequestProcessor.map_model("claude-4.0") == "claude-4.0" + assert RequestProcessor.resolve_model("claude-4.0") == ("claude-4.0", "exact") def test_alias_mapped_case_insensitively(monkeypatch): _patch_models(monkeypatch, aliases={"claude opus 4.7": "claude-4.0"}) assert RequestProcessor.map_model("Claude Opus 4.7") == "claude-4.0" + assert RequestProcessor.resolve_model("Claude Opus 4.7") == ("claude-4.0", "alias") -def test_unknown_label_falls_back_to_default(monkeypatch): - _patch_models(monkeypatch, default="auto-chat") - # A UI display label with no alias must NOT be forwarded verbatim. - assert RequestProcessor.map_model("Claude Opus 4.7") == "auto-chat" +def test_unknown_label_passthrough_by_default(monkeypatch): + # Default policy is passthrough: an unknown label is forwarded verbatim, + # never silently rewritten to the default model. This is the core bug fix. + _patch_models(monkeypatch, default="auto-chat", policy="passthrough") + assert RequestProcessor.resolve_model("Claude Opus 4.7") == ( + "Claude Opus 4.7", + "passthrough", + ) + assert RequestProcessor.map_model("Claude Opus 4.7") == "Claude Opus 4.7" + + +def test_unknown_label_default_policy_falls_back(monkeypatch): + # Only when policy=default does an unknown label map to the default model. + _patch_models(monkeypatch, default="auto-chat", policy="default") + assert RequestProcessor.resolve_model("Claude Opus 4.7") == ("auto-chat", "default") + + +def test_unknown_label_reject_policy_raises(monkeypatch): + _patch_models(monkeypatch, default="auto-chat", policy="reject") + with pytest.raises(codebuddy_router.UnknownModelError) as excinfo: + RequestProcessor.resolve_model("Claude Opus 4.7") + assert excinfo.value.requested_model == "Claude Opus 4.7" + + +def test_missing_model_uses_default_under_every_policy(monkeypatch): + # A missing/blank model has nothing to pass through, so it uses the default + # regardless of policy (and never raises under reject). + for policy in ("passthrough", "default", "reject"): + _patch_models(monkeypatch, default="auto-chat", policy=policy) + assert RequestProcessor.resolve_model(None) == ("auto-chat", "default_empty") + assert RequestProcessor.resolve_model("") == ("auto-chat", "default_empty") + assert RequestProcessor.map_model(None) == "auto-chat" + + +def test_explicit_auto_chat_request_is_exact_not_default(monkeypatch): + # A client explicitly asking for auto-chat matches exactly; it is not the + # unknown-model fallback path. + _patch_models(monkeypatch, default="auto-chat", policy="reject") + assert RequestProcessor.resolve_model("auto-chat") == ("auto-chat", "exact") -def test_missing_model_uses_default(monkeypatch): - _patch_models(monkeypatch, default="auto-chat") - assert RequestProcessor.map_model(None) == "auto-chat" - assert RequestProcessor.map_model("") == "auto-chat" +def test_opus_4_7_1m_is_not_converted_to_auto_chat(monkeypatch): + # Regression for the reported log: requested_model=claude-opus-4.7-1m must + # NOT be silently mapped to auto-chat. Under the default passthrough policy + # it is forwarded verbatim so the real upstream model is preserved. + _patch_models(monkeypatch, default="auto-chat", policy="passthrough") + mapped, source = RequestProcessor.resolve_model("claude-opus-4.7-1m") + assert mapped == "claude-opus-4.7-1m" + assert mapped != "auto-chat" + assert source == "passthrough" # --------------------------------------------------------------------------- # @@ -153,7 +196,7 @@ async def empty_pool(monkeypatch, tmp_path): return manager -def configure(monkeypatch, aliases=None, default="auto-chat"): +def configure(monkeypatch, aliases=None, default="auto-chat", policy="passthrough"): monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) @@ -161,7 +204,7 @@ def configure(monkeypatch, aliases=None, default="auto-chat"): monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: "web") monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: True) monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) - _patch_models(monkeypatch, aliases=aliases, default=default) + _patch_models(monkeypatch, aliases=aliases, default=default, policy=policy) monkeypatch.setattr( codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None ) @@ -298,12 +341,61 @@ def handler(req): @pytest.mark.asyncio -async def test_shiteru_unmapped_label_uses_default_not_verbatim( +async def test_shiteru_unmapped_label_passthrough_by_default( + monkeypatch, app, empty_pool +): + # No alias configured and default policy (passthrough): the unknown UI label + # is forwarded verbatim, NOT silently rewritten to auto-chat. + configure(monkeypatch, default="auto-chat", policy="passthrough") + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + response = await post(app, KEY_A, SHITERU_WEB_REQUEST) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["body"]["model"] == "Claude Opus 4.7" + assert seen["body"]["model"] != "auto-chat" + + +@pytest.mark.asyncio +async def test_opus_4_7_1m_not_rewritten_to_auto_chat_integration( + monkeypatch, app, empty_pool +): + # End-to-end regression for the reported log line + # (requested_model=claude-opus-4.7-1m mapped_model=auto-chat): the real + # model ID must reach upstream unchanged under the default policy. + configure(monkeypatch, default="auto-chat", policy="passthrough") + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + body = { + "model": "claude-opus-4.7-1m", + "messages": [{"role": "user", "content": "halo"}], + "stream": False, + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + assert seen["body"]["model"] == "claude-opus-4.7-1m" + assert seen["body"]["model"] != "auto-chat" + + +@pytest.mark.asyncio +async def test_shiteru_unmapped_label_default_policy_uses_default( monkeypatch, app, empty_pool ): - # No alias configured: the unknown UI label must fall back to the default, - # never be sent verbatim. - configure(monkeypatch, default="auto-chat") + # Only with policy=default does the unknown label fall back to the default. + configure(monkeypatch, default="auto-chat", policy="default") seen = {} def handler(req): @@ -316,4 +408,59 @@ def handler(req): assert response.status_code == 200 assert seen["body"]["model"] == "auto-chat" - assert seen["body"]["model"] != "Claude Opus 4.7" + + +@pytest.mark.asyncio +async def test_unknown_model_reject_policy_returns_400(monkeypatch, app, empty_pool): + # policy=reject: an unknown model is rejected with HTTP 400 code=unknown_model + # before any upstream call is made. + configure(monkeypatch, default="auto-chat", policy="reject") + called = {"upstream": False} + + def handler(req): + called["upstream"] = True + return sse("should not be called") + + upstream = install_upstream(monkeypatch, handler) + response = await post(app, KEY_A, SHITERU_WEB_REQUEST) + await upstream.aclose() + + assert response.status_code == 400 + assert called["upstream"] is False + error = response.json()["error"] + assert error["code"] == "unknown_model" + assert error["type"] == "invalid_request_error" + + +@pytest.mark.asyncio +async def test_passthrough_preserves_upstream_400_no_fallback( + monkeypatch, app, empty_pool +): + # Requirement 10: if CodeBuddy rejects the passed-through model with a 400, + # that rejection is surfaced rather than retried with a different model. + configure(monkeypatch, default="auto-chat", policy="passthrough") + seen = {"count": 0, "models": []} + + def handler(req): + seen["count"] += 1 + seen["models"].append(json.loads(req.content)["model"]) + return httpx.Response( + 400, + json={"error": {"message": "unknown model", "code": "model_not_found"}}, + headers={"content-type": "application/json"}, + ) + + upstream = install_upstream(monkeypatch, handler) + body = { + "model": "claude-opus-4.7-1m", + "messages": [{"role": "user", "content": "halo"}], + "stream": False, + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + # Upstream saw exactly one attempt with the verbatim model; no fallback to + # auto-chat was attempted. + assert seen["count"] == 1 + assert seen["models"] == ["claude-opus-4.7-1m"] + assert response.status_code == 400 From 5d94ffc5d3c02fd3fb9232365746bff045196d99 Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 11:42:01 +0700 Subject: [PATCH 6/9] update: Docker Network --- docker-compose.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docker-compose.yml b/docker-compose.yml index a307dd5..7cf895b 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -25,3 +25,12 @@ services: CODEBUDDY_CLIENT_AUTH_MODE: ${CODEBUDDY_CLIENT_AUTH_MODE:-relay} CODEBUDDY_ADMIN_PASSWORD: ${CODEBUDDY_ADMIN_PASSWORD:-} CODEBUDDY_UPSTREAM_API_KEY_HEADER: ${CODEBUDDY_UPSTREAM_API_KEY_HEADER:-bearer} + networks: + npm_network: + aliases: + - cb2api + +networks: + npm_network: + external: true + name: lizeu_default From 6bd6ee418abd39e4fb7c752159286d323833596c Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 13:27:48 +0700 Subject: [PATCH 7/9] fix: SSE --- src/codebuddy_message_sanitizer.py | 328 +++++++++- src/codebuddy_router.py | 25 + ...est_codebuddy_anthropic_tool_conversion.py | 608 ++++++++++++++++++ 3 files changed, 946 insertions(+), 15 deletions(-) create mode 100644 tests/test_codebuddy_anthropic_tool_conversion.py diff --git a/src/codebuddy_message_sanitizer.py b/src/codebuddy_message_sanitizer.py index 577ecb9..4f57520 100644 --- a/src/codebuddy_message_sanitizer.py +++ b/src/codebuddy_message_sanitizer.py @@ -15,6 +15,7 @@ an OpenAI-compatible ``content_filter`` error instead of leaking the refusal as an assistant reply. """ +import json import logging from typing import Any, Dict, List, Optional, Tuple @@ -106,41 +107,101 @@ def __init__(self, index: int, reason: str): self.reason = reason -def _log_structural(level: int, prefix: str, index: int, msg: Any) -> None: - """Log safe structural metadata for a message. +def _content_block_types(content: Any) -> List[str]: + """Return the ``type`` of each block in an array content, else empty list.""" + if not isinstance(content, list): + return [] + types: List[str] = [] + for block in content: + if isinstance(block, dict): + types.append(str(block.get("type", "unknown"))) + else: + types.append(type(block).__name__) + return types + - Logs the message index, role, its field-name set, the content field's type, - and whether it carries tool_calls / tool_call_id. It never logs message - content values or any credential: the field names logged (role, content, - tool_calls, ...) are fixed OpenAI-schema identifiers, not user data. +def _collect_tool_ids(msg: Dict[str, Any]) -> Tuple[List[str], List[str]]: + """Return ``(tool_call_ids, tool_result_ids)`` for a message. + + Handles both OpenAI shape (``tool_calls[*].id`` / ``tool_call_id``) and + Anthropic array shape (``tool_use`` / ``tool_result`` blocks). IDs are + opaque routing tokens (e.g. ``toolu_...`` / ``call_...``), never secrets or + content, so logging them is safe and is required for diagnosing lost tool + results. + """ + tool_call_ids: List[str] = [] + tool_result_ids: List[str] = [] + + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list): + for tc in tool_calls: + if isinstance(tc, dict) and tc.get("id"): + tool_call_ids.append(str(tc.get("id"))) + + if msg.get("tool_call_id"): + tool_result_ids.append(str(msg.get("tool_call_id"))) + + content = msg.get("content") + if isinstance(content, list): + for block in content: + if not isinstance(block, dict): + continue + if block.get("type") == "tool_use" and block.get("id"): + tool_call_ids.append(str(block.get("id"))) + elif block.get("type") == "tool_result" and block.get("tool_use_id"): + tool_result_ids.append(str(block.get("tool_use_id"))) + + return tool_call_ids, tool_result_ids + + +def _log_structural(level: int, prefix: str, index: int, msg: Any) -> None: + """Log safe structural metadata for a single message. + + Emits, per requirement, the fields needed to diagnose lost/malformed + tool-use conversion: index, role, content_type, content_block_types, + has_tool_calls, tool_call_ids, tool_result_ids, content_length. It never + logs message content values, prompts, or credentials: block *types* and + opaque tool IDs are structural routing metadata, not user data. content_length + is a character count only. """ if isinstance(msg, dict): role = msg.get("role") - field_names = ",".join(sorted(str(k) for k in msg.keys())) - content_type = type(msg.get("content")).__name__ + content = msg.get("content") + content_type = type(content).__name__ + block_types = _content_block_types(content) has_tool_calls = bool(msg.get("tool_calls")) - has_tool_call_id = bool(msg.get("tool_call_id")) + tool_call_ids, tool_result_ids = _collect_tool_ids(msg) + content_length = len(_extract_text(content)) else: role = None - field_names = "" content_type = type(msg).__name__ + block_types = [] has_tool_calls = False - has_tool_call_id = False + tool_call_ids, tool_result_ids = [], [] + content_length = 0 logger.log( level, - "%s index=%d role=%s keys=[%s] content_type=%s has_tool_calls=%s " - "has_tool_call_id=%s", + "%s index=%d role=%s content_type=%s content_block_types=[%s] " + "has_tool_calls=%s tool_call_ids=[%s] tool_result_ids=[%s] content_length=%d", prefix, index, role, - field_names, content_type, + ",".join(block_types), has_tool_calls, - has_tool_call_id, + ",".join(tool_call_ids), + ",".join(tool_result_ids), + content_length, ) +def log_messages_structural(prefix: str, messages: List[Dict[str, Any]], level: int = logging.DEBUG) -> None: + """Log structural diagnostics for an entire message list (content-free).""" + for index, msg in enumerate(messages): + _log_structural(level, prefix, index, msg) + + def _infer_role(msg: Dict[str, Any]) -> Any: """Infer a message role from structure when it is missing or blank. @@ -239,6 +300,243 @@ def normalize_messages_for_upstream( return normalized +# Anthropic content-block types that must be converted away before upstream. +_ANTHROPIC_TOOL_BLOCK_TYPES: Tuple[str, ...] = ("tool_use", "tool_result") + + +def _flatten_tool_result_content(content: Any) -> str: + """Flatten an Anthropic ``tool_result.content`` into a complete text string. + + ``tool_result.content`` may be a plain string or a list of blocks + (``text``/``image``/...). Every part is preserved so a tool result (e.g. the + full text of a file read) is never truncated or dropped. Non-text blocks are + serialized to JSON so their information survives rather than being discarded. + """ + if isinstance(content, str): + return content + if content is None: + return "" + if isinstance(content, list): + parts: List[str] = [] + for item in content: + if isinstance(item, dict): + if item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + parts.append(json.dumps(item, ensure_ascii=False)) + elif isinstance(item, str): + parts.append(item) + else: + parts.append(str(item)) + return "".join(parts) + return str(content) + + +def _assistant_text_from_blocks(blocks: List[Any]) -> str: + """Concatenate the text of the ``text`` blocks in an assistant content array.""" + parts: List[str] = [] + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + return "".join(parts) + + +def _has_block_type(content: Any, block_type: str) -> bool: + if not isinstance(content, list): + return False + return any( + isinstance(b, dict) and b.get("type") == block_type for b in content + ) + + +def convert_anthropic_messages_to_openai( + messages: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Convert Anthropic block-array messages into OpenAI-shaped messages. + + Claude Code sends tool interactions as Anthropic content blocks: + + * an assistant turn whose ``content`` array contains ``tool_use`` blocks + * a following user turn whose ``content`` array contains ``tool_result`` + blocks carrying the tool output (e.g. the full text of a file read) + + CodeBuddy speaks the OpenAI schema, so without conversion these blocks are + forwarded verbatim and the tool output never reaches the model — the + reported "file contents are not in context" failure. This function rewrites + them into the OpenAI shape while preserving every relationship: + + * assistant ``tool_use`` -> ``role: assistant`` message with ``content`` + (the text blocks, or ``""``) and a ``tool_calls`` list. Each tool call + keeps the Anthropic ``tool_use.id`` verbatim and encodes ``input`` as a + JSON-string ``function.arguments``. + * each user ``tool_result`` -> its own ``role: tool`` message whose + ``tool_call_id`` is the verbatim ``tool_use_id`` and whose ``content`` is + the complete flattened tool output. Tool results become standalone + messages and are never merged into an unrelated user message. + * messages whose content is a plain string, or an array with no tool + blocks (plain text or multimodal image arrays), are passed through + unchanged — array content is never replaced with an empty string. + + The verbatim ID reuse guarantees ``tool_use.id`` == the emitted + ``tool_call.id`` == the matching tool message's ``tool_call_id``, so the + OpenAI tool-call/tool-result pairing mirrors the Anthropic one exactly. Each + block is emitted exactly once, so no tool call or result is duplicated. + """ + converted: List[Dict[str, Any]] = [] + + for msg in messages: + if not isinstance(msg, dict): + converted.append(msg) + continue + + role = msg.get("role") + content = msg.get("content") + + # Non-array content (string / None) and arrays without tool blocks are + # preserved verbatim. This keeps plain text, and multimodal image + # arrays, exactly as the client sent them. + if not isinstance(content, list): + converted.append(msg) + continue + + has_tool_use = _has_block_type(content, "tool_use") + has_tool_result = _has_block_type(content, "tool_result") + + if not has_tool_use and not has_tool_result: + # Plain text or multimodal (image) array: preserve as-is. + converted.append(msg) + continue + + if has_tool_use: + # Assistant turn issuing one or more tool calls. + tool_calls: List[Dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict) or block.get("type") != "tool_use": + continue + tool_calls.append( + { + "id": block.get("id", ""), + "type": "function", + "function": { + "name": block.get("name", ""), + # OpenAI arguments must be a JSON *string*. + "arguments": json.dumps( + block.get("input", {}) or {}, ensure_ascii=False + ), + }, + } + ) + new_msg = dict(msg) + new_msg["role"] = role or "assistant" + # Assistant tool-call turns carry text content when present, else "". + new_msg["content"] = _assistant_text_from_blocks(content) + new_msg["tool_calls"] = tool_calls + converted.append(new_msg) + continue + + # has_tool_result: split the array so each tool_result becomes its own + # tool message; any remaining non-tool blocks become a trailing user + # message so tool output is never merged into unrelated user text. + leftover_blocks: List[Any] = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + converted.append( + { + "role": "tool", + "tool_call_id": block.get("tool_use_id", ""), + "content": _flatten_tool_result_content(block.get("content")), + } + ) + else: + leftover_blocks.append(block) + + if leftover_blocks: + # Preserve any accompanying user content (text/image) as a separate + # message, keeping array shape so multimodal content survives. + converted.append({"role": role or "user", "content": leftover_blocks}) + + return converted + + +def validate_upstream_messages(messages: List[Dict[str, Any]]) -> None: + """Validate the final OpenAI-shaped messages just before the upstream call. + + Raises :class:`MessageNormalizationError` (carrying the offending message + index and a structural description) when any of these hold, so the router + can return a local HTTP 400 instead of forwarding a malformed payload: + + * a message is missing ``role`` or ``content`` + * a ``tool`` message's ``tool_call_id`` has no preceding matching + ``tool_calls`` id (an orphaned tool result) + * a ``tool_calls`` entry has ``function.arguments`` that is not a valid + JSON string + * any unconverted Anthropic ``tool_use`` / ``tool_result`` block remains in + a message's content array + """ + seen_tool_call_ids: set = set() + + for index, msg in enumerate(messages): + _log_structural(logging.DEBUG, "validate upstream message", index, msg) + + if not isinstance(msg, dict): + raise MessageNormalizationError(index, "message is not an object") + + role = msg.get("role") + if not isinstance(role, str) or not role.strip(): + raise MessageNormalizationError(index, "message is missing a valid role") + if "content" not in msg: + raise MessageNormalizationError(index, "message is missing content") + + # No unconverted Anthropic tool blocks may remain in content arrays. + content = msg.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") in _ANTHROPIC_TOOL_BLOCK_TYPES: + raise MessageNormalizationError( + index, + f"unconverted Anthropic '{block.get('type')}' block remains " + f"in content", + ) + + # Register tool_call ids and validate their arguments are JSON strings. + tool_calls = msg.get("tool_calls") + if tool_calls is not None: + if not isinstance(tool_calls, list): + raise MessageNormalizationError(index, "tool_calls must be a list") + for tc in tool_calls: + if not isinstance(tc, dict): + raise MessageNormalizationError(index, "tool_call must be an object") + tc_id = tc.get("id") + if tc_id: + seen_tool_call_ids.add(str(tc_id)) + func = tc.get("function", {}) + arguments = func.get("arguments") if isinstance(func, dict) else None + if not isinstance(arguments, str): + raise MessageNormalizationError( + index, "tool_call function.arguments must be a JSON string" + ) + try: + json.loads(arguments) + except (ValueError, TypeError): + raise MessageNormalizationError( + index, "tool_call function.arguments is not valid JSON" + ) + + # A tool result must reference a tool_call id already seen upstream. + if role == "tool": + tool_call_id = msg.get("tool_call_id") + if not tool_call_id: + raise MessageNormalizationError( + index, "tool message is missing tool_call_id" + ) + if str(tool_call_id) not in seen_tool_call_ids: + raise MessageNormalizationError( + index, + "tool result references tool_call_id with no preceding " + "matching tool call", + ) + + def sanitize_messages( messages: List[Dict[str, Any]], *, diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index d8944ff..d0e6cff 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -26,9 +26,12 @@ from .keyword_replacer import apply_keyword_replacement_to_system_message from .codebuddy_message_sanitizer import ( MessageNormalizationError, + convert_anthropic_messages_to_openai, is_codebuddy_moderation_response, + log_messages_structural, normalize_messages_for_upstream, sanitize_messages, + validate_upstream_messages, ) from config import ( get_codebuddy_default_model, @@ -946,6 +949,19 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[ messages = source.get("messages", []) or [] + # Requirement: inspect the ORIGINAL request messages (content-free + # structural trace) before any transformation. + log_messages_structural("original request message", messages) + + # Convert Anthropic content-block messages (tool_use / tool_result + # arrays that Claude Code sends) into OpenAI-shaped messages FIRST, so + # tool calls become `tool_calls` and tool results become standalone + # `role: tool` messages. Without this, tool output (e.g. the full text of + # a file read) is forwarded as opaque Anthropic blocks and never reaches + # the model, which is the reported tool-result context loss. Plain and + # multimodal content arrays are preserved unchanged. + messages = convert_anthropic_messages_to_openai(messages) + # Sanitize only agent system prompts that tend to trigger false-positive # moderation. User/assistant/tool messages are never modified, and # legitimate short system prompts are left intact. @@ -986,6 +1002,15 @@ def prepare_payload(request_body: Dict[str, Any]) -> tuple[Dict[str, Any], Dict[ # be determined. messages = normalize_messages_for_upstream(messages) + # Final validation of the exact messages CodeBuddy will receive: every + # message has role+content, every tool result references a preceding + # matching tool call, function arguments are valid JSON strings, and no + # unconverted Anthropic blocks remain. Raises MessageNormalizationError + # (handled by the caller as a local HTTP 400 with the message index) + # rather than forwarding a malformed payload. + log_messages_structural("final upstream message", messages, level=logging.INFO) + validate_upstream_messages(messages) + # Build the strict upstream payload. CodeBuddy only supports streaming, # so stream is always True upstream; the client's stream preference is # honored separately by the router (aggregate vs pass-through). diff --git a/tests/test_codebuddy_anthropic_tool_conversion.py b/tests/test_codebuddy_anthropic_tool_conversion.py new file mode 100644 index 0000000..32e4709 --- /dev/null +++ b/tests/test_codebuddy_anthropic_tool_conversion.py @@ -0,0 +1,608 @@ +""" +Anthropic tool-block -> OpenAI conversion regression tests. + +Reproduces and locks the fix for Claude Code tool-result context loss: + + - Claude Code runs Read on portfolio.html and the client sends the file + back as an Anthropic ``tool_result`` block. + - Without conversion, the block is forwarded to CodeBuddy verbatim, the + OpenAI-speaking upstream never sees the file text, and the model reports + "the file contents are not in context" and falls back to shell commands. + +``convert_anthropic_messages_to_openai`` rewrites Anthropic content-block +messages into OpenAI-shaped messages before the upstream call: + + * assistant ``tool_use`` blocks -> a ``tool_calls`` list (arguments as a JSON + string), keeping the ``tool_use.id`` verbatim; + * user ``tool_result`` blocks -> standalone ``role: "tool"`` messages whose + ``tool_call_id`` is the verbatim ``tool_use_id`` and whose ``content`` is + the complete tool output. + +These tests cover the exact portfolio.html workflow, a multi-step +Read -> Grep -> Edit -> Read loop, multimodal preservation, that system-prompt +sanitization never touches tool results, and that payload allowlisting keeps +content blocks intact. + +Uses a mock upstream transport; no real CodeBuddy API key is required. +""" +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.codebuddy_router import RequestProcessor +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager +from src.codebuddy_message_sanitizer import ( + MessageNormalizationError, + convert_anthropic_messages_to_openai, + sanitize_messages, + validate_upstream_messages, +) + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + +AVAILABLE_MODELS = ["claude-4.0", "gpt-5", "auto-chat"] + +# A distinctive file body so we can assert it survives the pipeline exactly. +PORTFOLIO_HTML = ( + "\n\nMy Portfolio\n" + "\n

Jane Developer

\n

Full-stack engineer.

\n" + " \n\n\n" +) + + +# --------------------------------------------------------------------------- # +# Fixtures / helpers (mirror the other integration suites) +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +def _patch_models(monkeypatch): + monkeypatch.setattr(codebuddy_router, "get_available_models_list", lambda: list(AVAILABLE_MODELS)) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_default_model", lambda: "auto-chat") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_model_aliases", lambda: {}) + monkeypatch.setattr(codebuddy_router, "get_codebuddy_unknown_model_policy", lambda: "passthrough") + + +def configure(monkeypatch, sanitize=False): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: "web") + monkeypatch.setattr(codebuddy_router, "get_sanitize_agent_prompt", lambda: sanitize) + monkeypatch.setattr(codebuddy_router, "get_max_system_prompt_length", lambda: 2000) + _patch_models(monkeypatch) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse(text, finish="stop"): + body = ( + 'data: {"id":"chat-1","model":"auto-chat","choices":' + f'[{{"delta":{{"content":{json.dumps(text)}}},"finish_reason":{json.dumps(finish)}}}]}}\n\n' + "data: [DONE]\n\n" + ) + return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"}) + + +async def post(app, token, body): + transport = httpx.ASGITransport(app=app) + headers = {"Authorization": f"Bearer {token}"} + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post( + "/codebuddy/v1/chat/completions", headers=headers, json=body + ) + + +def _tool_use_msg(text, tool_id, name, tool_input): + """An Anthropic assistant turn: optional text + a tool_use block.""" + content = [] + if text: + content.append({"type": "text", "text": text}) + content.append( + {"type": "tool_use", "id": tool_id, "name": name, "input": tool_input} + ) + return {"role": "assistant", "content": content} + + +def _tool_result_msg(tool_use_id, result): + """An Anthropic user turn carrying a tool_result block.""" + return { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": tool_use_id, "content": result} + ], + } + + +def _find_tool_messages(messages): + return [m for m in messages if m.get("role") == "tool"] + + +# --------------------------------------------------------------------------- # +# Unit: converter shape +# --------------------------------------------------------------------------- # + + +def test_assistant_tool_use_becomes_openai_tool_calls(): + messages = [ + _tool_use_msg("I'll read it.", "toolu_01", "Read", {"file_path": "portfolio.html"}), + ] + out = convert_anthropic_messages_to_openai(messages) + assert len(out) == 1 + msg = out[0] + assert msg["role"] == "assistant" + assert msg["content"] == "I'll read it." + assert isinstance(msg["tool_calls"], list) and len(msg["tool_calls"]) == 1 + tc = msg["tool_calls"][0] + assert tc["id"] == "toolu_01" + assert tc["type"] == "function" + assert tc["function"]["name"] == "Read" + # arguments is a JSON *string* and round-trips. + assert json.loads(tc["function"]["arguments"]) == {"file_path": "portfolio.html"} + + +def test_assistant_tool_use_without_text_gets_empty_content(): + messages = [_tool_use_msg("", "toolu_x", "Read", {"file_path": "a"})] + out = convert_anthropic_messages_to_openai(messages) + assert out[0]["content"] == "" + assert out[0]["tool_calls"][0]["id"] == "toolu_x" + + +def test_tool_result_becomes_standalone_tool_message(): + messages = [_tool_result_msg("toolu_01", PORTFOLIO_HTML)] + out = convert_anthropic_messages_to_openai(messages) + assert len(out) == 1 + assert out[0] == { + "role": "tool", + "tool_call_id": "toolu_01", + "content": PORTFOLIO_HTML, + } + + +def test_tool_result_with_block_list_content_is_flattened(): + messages = [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_2", + "content": [{"type": "text", "text": PORTFOLIO_HTML}], + } + ], + } + ] + out = convert_anthropic_messages_to_openai(messages) + assert out[0]["role"] == "tool" + assert out[0]["content"] == PORTFOLIO_HTML + + +def test_id_relationship_preserved_end_to_end(): + messages = [ + {"role": "user", "content": "read portfolio.html"}, + _tool_use_msg("Reading.", "toolu_ABC", "Read", {"file_path": "portfolio.html"}), + _tool_result_msg("toolu_ABC", PORTFOLIO_HTML), + ] + out = convert_anthropic_messages_to_openai(messages) + assistant = next(m for m in out if m.get("tool_calls")) + tool = next(m for m in out if m.get("role") == "tool") + # tool_use.id == tool_call.id == tool_result tool_call_id. + assert assistant["tool_calls"][0]["id"] == "toolu_ABC" + assert tool["tool_call_id"] == "toolu_ABC" + + +def test_tool_result_not_merged_into_unrelated_user_text(): + # A user turn that contains BOTH a tool_result and trailing text. + messages = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_9", "content": "RESULT"}, + {"type": "text", "text": "now edit it"}, + ], + } + ] + out = convert_anthropic_messages_to_openai(messages) + # The tool result is its own message; the text is a separate user message. + tool_msgs = [m for m in out if m.get("role") == "tool"] + user_msgs = [m for m in out if m.get("role") == "user"] + assert len(tool_msgs) == 1 + assert tool_msgs[0]["content"] == "RESULT" + assert len(user_msgs) == 1 + assert user_msgs[0]["content"] == [{"type": "text", "text": "now edit it"}] + + +def test_plain_string_content_passed_through(): + messages = [{"role": "user", "content": "hello"}] + out = convert_anthropic_messages_to_openai(messages) + assert out == messages + + +def test_multimodal_image_array_preserved_by_converter(): + multimodal = [ + {"type": "text", "text": "describe"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}}, + ] + messages = [{"role": "user", "content": multimodal}] + out = convert_anthropic_messages_to_openai(messages) + # No tool blocks -> preserved verbatim, never replaced with "". + assert out[0]["content"] == multimodal + + +def test_no_duplication_of_tool_calls_or_results(): + messages = [ + _tool_use_msg("", "toolu_1", "Read", {"file_path": "a"}), + _tool_result_msg("toolu_1", "A"), + _tool_use_msg("", "toolu_2", "Read", {"file_path": "b"}), + _tool_result_msg("toolu_2", "B"), + ] + out = convert_anthropic_messages_to_openai(messages) + all_call_ids = [ + tc["id"] for m in out for tc in (m.get("tool_calls") or []) + ] + all_result_ids = [m["tool_call_id"] for m in out if m.get("role") == "tool"] + assert all_call_ids == ["toolu_1", "toolu_2"] + assert all_result_ids == ["toolu_1", "toolu_2"] + + +# --------------------------------------------------------------------------- # +# Unit: validation (requirement 12/13) +# --------------------------------------------------------------------------- # + + +def test_validate_rejects_orphan_tool_result(): + messages = [ + {"role": "user", "content": "hi"}, + {"role": "tool", "tool_call_id": "toolu_missing", "content": "x"}, # index 1 + ] + with pytest.raises(MessageNormalizationError) as excinfo: + validate_upstream_messages(messages) + assert excinfo.value.index == 1 + + +def test_validate_rejects_unconverted_anthropic_block(): + messages = [ + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t", "content": "x"}], + } + ] + with pytest.raises(MessageNormalizationError) as excinfo: + validate_upstream_messages(messages) + assert excinfo.value.index == 0 + + +def test_validate_rejects_non_json_arguments(): + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "f", "arguments": "{not json"}} + ], + } + ] + with pytest.raises(MessageNormalizationError) as excinfo: + validate_upstream_messages(messages) + assert excinfo.value.index == 0 + + +def test_validate_accepts_well_formed_pairing(): + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "Read", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "ok"}, + ] + # Should not raise. + validate_upstream_messages(messages) + + +# --------------------------------------------------------------------------- # +# Requirement 17: sanitization only touches system messages +# --------------------------------------------------------------------------- # + + +def test_sanitization_never_removes_tool_result_content(): + messages = [ + {"role": "system", "content": "You are Claude Code, the official CLI. " + "x" * 3000}, + {"role": "tool", "tool_call_id": "c1", "content": PORTFOLIO_HTML}, + {"role": "assistant", "content": "", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "Read", "arguments": "{}"}} + ]}, + ] + out, changed = sanitize_messages(messages, enabled=True, max_system_prompt_length=2000) + assert changed is True # the system prompt was replaced + # The tool result content is untouched. + tool_msg = next(m for m in out if m.get("role") == "tool") + assert tool_msg["content"] == PORTFOLIO_HTML + # The assistant tool call is untouched. + asst = next(m for m in out if m.get("tool_calls")) + assert asst["tool_calls"][0]["id"] == "c1" + + +# --------------------------------------------------------------------------- # +# Requirement 14: the exact portfolio.html workflow (integration) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_portfolio_read_workflow_delivers_full_file(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("I can see the portfolio file.") + + upstream = install_upstream(monkeypatch, handler) + + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": "please read portfolio.html"}, + _tool_use_msg("I'll read it.", "toolu_read1", "Read", {"file_path": "portfolio.html"}), + _tool_result_msg("toolu_read1", PORTFOLIO_HTML), + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + + # The assistant Read tool call reached upstream as an OpenAI tool call. + asst = next(m for m in sent if m.get("tool_calls")) + assert asst["tool_calls"][0]["id"] == "toolu_read1" + assert asst["tool_calls"][0]["function"]["name"] == "Read" + + # The file contents are present, complete, and unmodified in a tool message. + tool_msgs = _find_tool_messages(sent) + assert len(tool_msgs) == 1 + assert tool_msgs[0]["tool_call_id"] == "toolu_read1" + assert tool_msgs[0]["content"] == PORTFOLIO_HTML + assert "SENTINEL_MARKER_9F3A" in tool_msgs[0]["content"] + + +@pytest.mark.asyncio +async def test_portfolio_workflow_then_edit_call_converts(monkeypatch, app, empty_pool): + # After the file is read, the model proposes an Edit call: that Anthropic + # tool_use must also convert cleanly. + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("Editing now.") + + upstream = install_upstream(monkeypatch, handler) + + edit_input = {"file_path": "portfolio.html", "old_string": "Jane Developer", "new_string": "Jane D."} + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": "read then rename the heading"}, + _tool_use_msg("Reading.", "toolu_read1", "Read", {"file_path": "portfolio.html"}), + _tool_result_msg("toolu_read1", PORTFOLIO_HTML), + _tool_use_msg("Now editing.", "toolu_edit1", "Edit", edit_input), + _tool_result_msg("toolu_edit1", "Edit applied."), + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + call_ids = [tc["id"] for m in sent for tc in (m.get("tool_calls") or [])] + assert call_ids == ["toolu_read1", "toolu_edit1"] + # The Edit arguments survived as valid JSON. + edit_call = next( + tc for m in sent for tc in (m.get("tool_calls") or []) if tc["id"] == "toolu_edit1" + ) + assert json.loads(edit_call["function"]["arguments"]) == edit_input + # Both tool results are present and none dropped. + tool_msgs = _find_tool_messages(sent) + assert [t["tool_call_id"] for t in tool_msgs] == ["toolu_read1", "toolu_edit1"] + assert tool_msgs[0]["content"] == PORTFOLIO_HTML + + +# --------------------------------------------------------------------------- # +# Requirement 15: multi-step Read -> Grep -> Edit -> Read, nothing dropped +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_read_grep_edit_read_loop_preserves_every_result(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("done") + + upstream = install_upstream(monkeypatch, handler) + + steps = [ + ("toolu_read1", "Read", {"file_path": "portfolio.html"}, PORTFOLIO_HTML), + ("toolu_grep1", "Grep", {"pattern": "Jane", "path": "portfolio.html"}, "portfolio.html:4:

Jane Developer

"), + ("toolu_edit1", "Edit", {"file_path": "portfolio.html", "old_string": "Jane", "new_string": "Janet"}, "Edit applied."), + ("toolu_read2", "Read", {"file_path": "portfolio.html"}, PORTFOLIO_HTML.replace("Jane", "Janet")), + ] + + messages = [{"role": "user", "content": "read, grep, edit, then re-read"}] + for tool_id, name, tool_input, result in steps: + messages.append(_tool_use_msg("", tool_id, name, tool_input)) + messages.append(_tool_result_msg(tool_id, result)) + + body = {"model": "claude-4.0", "stream": False, "messages": messages} + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + + # Every tool call id appears exactly once, in order. + call_ids = [tc["id"] for m in sent for tc in (m.get("tool_calls") or [])] + assert call_ids == [s[0] for s in steps] + + # Every tool result is present, in order, with its exact content. + tool_msgs = _find_tool_messages(sent) + assert [t["tool_call_id"] for t in tool_msgs] == [s[0] for s in steps] + for tmsg, (_id, _name, _input, result) in zip(tool_msgs, steps): + assert tmsg["content"] == result + + # No duplication: as many tool results as tool calls. + assert len(tool_msgs) == len(call_ids) == 4 + + +# --------------------------------------------------------------------------- # +# Requirement 16/18: multimodal + allowlist preserve content blocks +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_multimodal_array_survives_full_path(monkeypatch, app, empty_pool): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + multimodal = [ + {"type": "text", "text": "what is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/x.png"}}, + ] + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": multimodal}, + {"role": "assistant", "content": "A cat."}, + {"role": "user", "content": "thanks"}, + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + sent = seen["body"]["messages"] + first_user = next(m for m in sent if m.get("role") == "user") + # The multimodal content array reaches upstream unchanged (req 16, 18). + assert first_user["content"] == multimodal + + +@pytest.mark.asyncio +async def test_allowlist_keeps_tool_calls_inside_messages(monkeypatch, app, empty_pool): + # Top-level allowlist drops unsupported fields but must not strip the + # content/tool_calls INSIDE the message objects. + configure(monkeypatch) + seen = {} + + def handler(req): + seen["body"] = json.loads(req.content) + return sse("ok") + + upstream = install_upstream(monkeypatch, handler) + body = { + "model": "claude-4.0", + "stream": False, + "temperature": 0.5, # unsupported top-level field, must be dropped + "metadata": {"ui": "shiteru"}, + "messages": [ + {"role": "user", "content": "read it"}, + _tool_use_msg("Reading.", "toolu_1", "Read", {"file_path": "portfolio.html"}), + _tool_result_msg("toolu_1", PORTFOLIO_HTML), + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 200 + body_sent = seen["body"] + # Unsupported top-level fields were dropped... + assert "temperature" not in body_sent + assert "metadata" not in body_sent + # ...but the tool call and tool result inside the messages survived. + assert any(m.get("tool_calls") for m in body_sent["messages"]) + tool_msgs = _find_tool_messages(body_sent["messages"]) + assert tool_msgs and tool_msgs[0]["content"] == PORTFOLIO_HTML + + +# --------------------------------------------------------------------------- # +# Requirement 13: orphan tool result -> local 400 with index, not sent upstream +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_orphan_tool_result_returns_local_400(monkeypatch, app, empty_pool): + configure(monkeypatch) + called = {"upstream": False} + + def handler(req): + called["upstream"] = True + return sse("should not be called") + + upstream = install_upstream(monkeypatch, handler) + # A tool_result whose tool_use_id has no preceding tool_use. + body = { + "model": "claude-4.0", + "stream": False, + "messages": [ + {"role": "user", "content": "hi"}, + _tool_result_msg("toolu_orphan", "dangling result"), + ], + } + response = await post(app, KEY_A, body) + await upstream.aclose() + + assert response.status_code == 400 + assert called["upstream"] is False + error = response.json()["error"] + assert error["code"] == "invalid_message" + assert error["type"] == "invalid_request_error" From 2fafeb1113e2ae7204de1e1a766b49a671d0d143 Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 15:46:19 +0700 Subject: [PATCH 8/9] fix: Stream & Timeout Error --- .env.example | 62 ++ codebuddy2api-fix.md | 904 ++++++++++++++++++++ config.py | 110 ++- src/codebuddy_router.py | 622 ++++++++++++-- tests/test_codebuddy_streaming_stability.py | 537 ++++++++++++ 5 files changed, 2154 insertions(+), 81 deletions(-) create mode 100644 codebuddy2api-fix.md create mode 100644 tests/test_codebuddy_streaming_stability.py diff --git a/.env.example b/.env.example index 02372cb..45b2e42 100644 --- a/.env.example +++ b/.env.example @@ -100,6 +100,68 @@ CODEBUDDY_UPSTREAM_API_KEY_HEADER=bearer # Default: INFO CODEBUDDY_LOG_LEVEL=INFO + +# ----------------- +# Upstream connection / timeout tuning +# ----------------- +# These timeouts are intentionally separate so a slow-but-alive agentic stream +# (many tools, long tool results) is not killed by the same short deadline used +# to open a connection. Do NOT use a single global 15s timeout for agentic +# requests. + +# (Optional) Time allowed to open a TCP/TLS connection to CodeBuddy. Default: 30 +CODEBUDDY_CONNECT_TIMEOUT_SECONDS=30 + +# (Optional) Time allowed to wait for a free slot in the connection pool. Default: 30 +CODEBUDDY_POOL_TIMEOUT_SECONDS=30 + +# (Optional) Time allowed to write the request body upstream. Default: 60 +CODEBUDDY_WRITE_TIMEOUT_SECONDS=60 + +# (Optional) Time allowed for upstream to return headers + the first chunk. +# Sized for tool-heavy agentic requests. Default: 180 +CODEBUDDY_FIRST_BYTE_TIMEOUT_SECONDS=180 + +# (Optional) Per-read timeout once the stream is flowing. 0 = unlimited. Default: 0 +CODEBUDDY_STREAM_READ_TIMEOUT_SECONDS=0 + +# (Optional) Connection pool sizing. Defaults: 30 keepalive / 100 max / 60s expiry +CODEBUDDY_MAX_KEEPALIVE_CONNECTIONS=30 +CODEBUDDY_MAX_CONNECTIONS=100 +CODEBUDDY_KEEPALIVE_EXPIRY_SECONDS=60 + + +# ----------------- +# Upstream concurrency limiter +# ----------------- +# Bounds how many upstream requests run at once so tool-heavy bursts do not pile +# up sockets/tasks without limit. Requests beyond the limit queue up to the +# queue timeout, then receive HTTP 503 (code=upstream_queue_timeout). + +# (Optional) Max simultaneous in-flight upstream requests. Default: 20 +CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS=20 + +# (Optional) Max time a request waits for a concurrency slot before 503. Default: 60 +CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS=60 + + +# ----------------- +# SSE heartbeat +# ----------------- +# While waiting for the first upstream chunk (after upstream status is known to +# be 200), a heartbeat comment (": ping") is sent to keep the downstream +# connection alive. Heartbeats never become assistant content and stop once the +# real stream starts or the client disconnects. 0 disables heartbeats. +CODEBUDDY_HEARTBEAT_INTERVAL_SECONDS=15 + + +# ----------------- +# Large-request warnings (log only; never truncates the request) +# ----------------- +CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH=50000 +CODEBUDDY_WARN_MESSAGE_COUNT=40 +CODEBUDDY_WARN_TOOL_COUNT=30 + # (Optional) Comma-separated model list reported to clients. # Adjust this list to match the models supported by your CodeBuddy account. CODEBUDDY_MODELS=claude-4.0,claude-3.7,gpt-5,gpt-5-mini,gpt-5-nano,o4-mini,gemini-2.5-flash,gemini-2.5-pro,auto-chat diff --git a/codebuddy2api-fix.md b/codebuddy2api-fix.md new file mode 100644 index 0000000..ddd9c60 --- /dev/null +++ b/codebuddy2api-fix.md @@ -0,0 +1,904 @@ +# PROMPT 1 — CodeBuddy2API: Audit Latency, Streaming, Cancellation, dan Connection Pool + +Repository: + +```text +https://github.com/xueyue33/codebuddy2api +``` + +## Konteks Masalah + +CodeBuddy2API digunakan dengan alur: + +```text +Claude Code / Shiteru +→ 9Router +→ CodeBuddy2API +→ https://www.codebuddy.ai/v2/chat/completions +``` + +Chat biasa relatif cepat. Namun request Claude Code yang memiliki banyak messages, tool calls, tool results, dan sekitar 30 tools terkadang: + +- lama setelah log `outcome=prepared`; +- menunggu lama sebelum respons pertama; +- menyebabkan 9Router menampilkan `502 fetch connect timeout`; +- meninggalkan request yang masih berjalan setelah client berhenti; +- menghasilkan beberapa request upstream berdekatan; +- terasa stuck saat melanjutkan percakapan lama. + +Contoh request berat: + +```text +requested_model=claude-opus-4.7-1m +mapped_model=claude-opus-4.7-1m +stream=True +message_count=20+ +tool_count=30 +has_tools=True +system_prompt_sanitized=False +request_profile=cli +``` + +Model mapping, tool-result conversion, dan basic streaming sudah bekerja. Jangan merusak perbaikan yang sudah ada. + +## Tujuan + +Perbaiki observability dan stabilitas CodeBuddy2API agar: + +1. Tahap yang lambat dapat diketahui secara pasti. +2. Request tidak diam setelah `outcome=prepared`. +3. Streaming benar-benar diteruskan secara incremental. +4. Request upstream dibatalkan ketika downstream disconnect. +5. Tidak ada ghost request setelah Claude Code dibatalkan. +6. Connection pool dipakai ulang dengan benar. +7. Semaphore atau concurrency limiter tidak bocor. +8. Request besar tidak membuat task/socket menumpuk tanpa batas. +9. Error upstream tetap dikembalikan dengan status yang tepat. +10. API key, prompt, dan isi file tidak bocor ke log. + +## 1. Audit Sebelum Mengubah Kode + +Temukan dan dokumentasikan: + +- endpoint `chat/completions`; +- lokasi log `outcome=prepared`; +- pembentukan payload upstream; +- pembuatan `httpx.AsyncClient`; +- semaphore/concurrency limiter; +- pemanggilan CodeBuddy upstream; +- pembuatan `StreamingResponse`; +- pembacaan SSE upstream; +- cancellation handling; +- error propagation; +- lifecycle startup dan shutdown. + +Jangan langsung menulis ulang streaming. Audit terlebih dahulu apakah body upstream masih dibuffer atau sudah diteruskan incremental. + +## 2. Request ID dan Stage Telemetry + +Setiap inbound request harus memiliki `request_id`. + +Tambahkan log pada tahap: + +```text +request_received +request_normalized +request_prepared +upstream_slot_wait_start +upstream_slot_acquired +upstream_send_start +upstream_headers_received +upstream_first_chunk +upstream_stream_finished +downstream_disconnected +upstream_cancelled +request_failed +``` + +Metadata aman: + +```text +request_id +requested_model +mapped_model +stream +message_count +tool_count +total_content_length +key_fingerprint +queue_wait_ms +pool_wait_ms +time_to_upstream_headers_ms +time_to_first_chunk_ms +stream_duration_ms +chunk_count +upstream_status +finish_reason +``` + +Contoh: + +```text +request_id=abc123 stage=upstream_slot_wait_start +request_id=abc123 stage=upstream_slot_acquired queue_wait_ms=125 +request_id=abc123 stage=upstream_send_start +request_id=abc123 stage=upstream_headers_received status=200 elapsed_ms=2180 +request_id=abc123 stage=upstream_first_chunk elapsed_ms=3475 +request_id=abc123 stage=upstream_stream_finished chunks=82 duration_ms=18230 +``` + +Jangan log: + +- raw API key; +- Authorization; +- X-Api-Key; +- isi prompt; +- isi file; +- isi tool result; +- cookie atau credential. + +## 3. Shared HTTP Client + +Pastikan aplikasi menggunakan satu reusable `httpx.AsyncClient`. + +Contoh konfigurasi awal: + +```python +httpx.AsyncClient( + timeout=httpx.Timeout( + connect=30.0, + read=None, + write=60.0, + pool=30.0, + ), + limits=httpx.Limits( + max_connections=100, + max_keepalive_connections=30, + keepalive_expiry=60.0, + ), +) +``` + +Ketentuan: + +- dibuat saat startup; +- ditutup saat shutdown; +- tidak membuat client baru pada setiap request; +- tidak menutup shared client setelah satu request; +- connection pool memiliki limit yang jelas; +- waktu menunggu pool dicatat. + +## 4. Audit True Streaming + +Untuk `stream=true`, pastikan jalur upstream menggunakan: + +```python +client.stream(...) +``` + +atau: + +```python +client.send(request, stream=True) +``` + +Dilarang pada jalur streaming: + +```python +await response.aread() +response.text +response.json() +list(response.aiter_lines()) +``` + +Jangan mengumpulkan seluruh chunk sebelum mengirim ke downstream. + +Koneksi upstream harus tetap terbuka selama downstream membaca stream. + +Response SSE harus memakai: + +```http +Content-Type: text/event-stream +Cache-Control: no-cache, no-transform +X-Accel-Buffering: no +``` + +Pastikan GZip tidak membuffer endpoint SSE. + +## 5. Jangan Mengirim Heartbeat Sebelum Status Upstream Diketahui Secara Aman + +Audit kemungkinan penggunaan SSE heartbeat. + +Heartbeat boleh dipakai untuk menjaga koneksi: + +```text +: ping + +``` + +tetapi jangan langsung mengirim HTTP 200 ke downstream sebelum diketahui bahwa upstream menerima request. + +Gunakan strategi aman: + +1. buka koneksi upstream; +2. tunggu status/header upstream; +3. jika status error, kembalikan error HTTP yang benar; +4. jika status berhasil, mulai StreamingResponse; +5. selama menunggu chunk pertama, heartbeat SSE boleh dikirim setiap 15–20 detik. + +Heartbeat: + +- tidak boleh masuk ke assistant content; +- tidak boleh merusak `[DONE]`; +- harus berhenti saat client disconnect; +- tidak boleh ditulis paralel tanpa sinkronisasi. + +## 6. Cancellation dan Disconnect + +Jika downstream disconnect: + +- batalkan request upstream; +- tutup response upstream; +- lepaskan semaphore; +- hentikan heartbeat; +- jangan lanjutkan retry; +- jangan meninggalkan background task. + +Tangani: + +```python +asyncio.CancelledError +``` + +Pastikan cleanup menggunakan `finally`. + +Contoh kebutuhan: + +```python +try: + ... +except asyncio.CancelledError: + logger.info( + "request_id=%s stage=downstream_disconnected", + request_id, + ) + raise +finally: + if upstream_response is not None: + await upstream_response.aclose() +``` + +Tambahkan test bahwa setelah client disconnect, tidak ada request upstream yang terus berjalan. + +## 7. Semaphore dan Concurrency + +Audit semua semaphore. + +Pastikan: + +- semaphore selalu dilepas; +- cancellation tidak menyebabkan slot bocor; +- tidak ada semaphore ganda tanpa alasan; +- queue wait dicatat; +- request tidak mengantre tanpa batas. + +Tambahkan konfigurasi: + +```env +CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS=20 +CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS=60 +``` + +Jika queue timeout: + +```json +{ + "error": { + "message": "CodeBuddy relay is temporarily at capacity", + "type": "server_overloaded", + "code": "upstream_queue_timeout" + } +} +``` + +Gunakan HTTP `503`. + +## 8. Timeout Terpisah + +Tambahkan konfigurasi terpisah: + +```env +CODEBUDDY_CONNECT_TIMEOUT_SECONDS=30 +CODEBUDDY_POOL_TIMEOUT_SECONDS=30 +CODEBUDDY_WRITE_TIMEOUT_SECONDS=60 +CODEBUDDY_FIRST_BYTE_TIMEOUT_SECONDS=180 +CODEBUDDY_STREAM_READ_TIMEOUT_SECONDS=0 +``` + +Interpretasi: + +```text +connect timeout +→ gagal membuka koneksi + +pool timeout +→ menunggu slot connection pool terlalu lama + +first-byte timeout +→ upstream terlalu lama mengirim header/chunk awal + +stream read timeout +→ timeout setelah stream berjalan +``` + +Nilai `0` untuk stream read berarti unlimited jika implementasi mendukung. + +Jangan memakai satu timeout global 15 detik untuk seluruh agentic request. + +## 9. Jangan Retry Internal pada Passthrough + +Pada: + +```env +CODEBUDDY_CLIENT_AUTH_MODE=passthrough +``` + +CodeBuddy2API tidak boleh berpindah ke key lain. + +Tidak boleh retry untuk: + +```text +400 invalid_request +400 content_filter +401 invalid API key +403 permission denied +``` + +Retry terbatas hanya boleh terjadi sebelum first chunk untuk: + +```text +408 +429 +502 +503 +504 +connect timeout +connection reset +``` + +Ketentuan: + +- maksimal 2 attempt; +- tidak retry setelah first chunk; +- tidak retry setelah tool call diteruskan; +- tidak retry setelah client disconnect; +- tidak retry payload malformed. + +## 10. Pertahankan Tool Messages + +Jangan menghapus atau merangkum tool results secara diam-diam. + +Pertahankan perbaikan yang memastikan: + +```json +{ + "role": "assistant", + "content": "", + "tool_calls": [...] +} +``` + +dan: + +```json +{ + "role": "tool", + "tool_call_id": "call_xxx", + "content": "..." +} +``` + +Sebelum upstream request, validasi: + +- semua message memiliki `role`; +- semua message memiliki `content`; +- setiap tool result memiliki pasangan tool call; +- tool call arguments merupakan JSON string valid; +- content array yang valid tidak dibuang. + +## 11. Large Request Warning + +Tambahkan warning metadata bila request besar: + +```env +CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH=50000 +CODEBUDDY_WARN_MESSAGE_COUNT=40 +CODEBUDDY_WARN_TOOL_COUNT=30 +``` + +Log hanya warning, jangan memotong request otomatis. + +Contoh: + +```text +request_id=abc123 large_agentic_request=true message_count=42 tool_count=30 total_content_length=68420 +``` + +## 12. Tests + +Tambahkan automated tests: + +1. Upstream memberi tiga chunk dengan delay dan downstream menerima incremental. +2. Header SSE benar. +3. Tidak ada buffering body pada stream path. +4. Shared HTTP client digunakan ulang. +5. Queue wait dicatat. +6. Queue timeout menghasilkan 503. +7. Client disconnect membatalkan upstream. +8. Semaphore dilepas setelah cancellation. +9. First-byte timeout berbeda dari stream read timeout. +10. Heartbeat tidak menjadi assistant content. +11. Tidak ada retry setelah first chunk. +12. HTTP 400 tidak di-retry. +13. HTTP 429 hanya di-retry secara terbatas. +14. Request dengan 30 tools tetap dapat streaming. +15. Request dengan 25+ messages tidak kehilangan tool result. +16. Non-streaming tetap menghasilkan `message.content`. +17. Raw key, prompt, dan isi file tidak muncul di log. + +## Acceptance Criteria + +Pekerjaan selesai jika: + +- tahap lambat dapat diketahui dari log; +- log tidak berhenti tanpa penjelasan setelah `outcome=prepared`; +- streaming tetap incremental; +- client disconnect menghentikan upstream; +- tidak ada ghost request; +- semaphore dan connection pool tidak bocor; +- request tool-heavy tidak menumpuk tanpa batas; +- error HTTP tetap dipertahankan; +- mode passthrough tidak merotasi key internal; +- semua test lulus. + +## Output Agent + +Sebelum implementasi: + +- jelaskan alur request saat ini; +- tunjukkan kemungkinan lokasi delay; +- sebutkan file yang akan diubah. + +Setelah implementasi: + +- daftar file yang diubah; +- arsitektur shared client; +- lifecycle streaming; +- konfigurasi environment baru; +- contoh telemetry; +- hasil test; +- perintah deployment. + +Implementasikan perubahan secara langsung. Jangan hanya memberikan saran. +``` + +--- + +# PROMPT 2 — 9Router: Retry Policy, Timeout, Fallback, dan Streaming Cancellation + +Repository: + +```text +https://github.com/decolua/9router +``` + +## Konteks Masalah + +9Router digunakan dengan alur: + +```text +Claude Code / Shiteru +→ 9Router +→ http://cb2api:8001/codebuddy/v1 +→ CodeBuddy2API +→ CodeBuddy +``` + +Chat biasa bekerja. Namun workflow Claude Code dengan tool calls kadang menampilkan: + +```text +API error · Retrying · attempt 1/10 +``` + +Dashboard 9Router kadang menampilkan: + +```text +[502]: fetch connect timeout +13s–15s +``` + +atau menandai banyak key sebagai: + +```text +unavailable +``` + +Beberapa error yang pernah terjadi: + +```text +400 invalid_request +Message must have role and content + +502 fetch connect timeout + +upstream response lambat saat request tool-heavy +``` + +Saat ini retry sampai 10 kali membuat kegagalan terasa sangat lama dan dapat menghasilkan request berulang. + +## Tujuan + +Perbaiki 9Router agar: + +1. HTTP 400 tidak di-retry. +2. Retry dibatasi berdasarkan jenis error. +3. Connect timeout terpisah dari first-byte dan stream timeout. +4. Tidak melakukan fallback setelah stream dimulai. +5. Tidak menggandakan tool call. +6. Downstream disconnect membatalkan upstream. +7. Provider tidak ditandai unavailable terlalu agresif. +8. Semua attempt memiliki telemetry. +9. Base URL internal Docker digunakan secara stabil. +10. API key tidak bocor ke log. + +## 1. Audit Retry dan Timeout + +Temukan: + +- fungsi request upstream; +- default connect timeout; +- retry loop; +- fallback loop; +- provider health state; +- unavailable/cooldown handling; +- streaming parser; +- cancellation handling; +- logika kapan attempt dianggap gagal. + +Jelaskan alasan angka timeout sekitar 13–15 detik yang terlihat di dashboard. + +## 2. Pisahkan Timeout + +Tambahkan konfigurasi: + +```env +ROUTER_CONNECT_TIMEOUT_SECONDS=30 +ROUTER_FIRST_BYTE_TIMEOUT_SECONDS=180 +ROUTER_STREAM_IDLE_TIMEOUT_SECONDS=600 +ROUTER_TOTAL_REQUEST_TIMEOUT_SECONDS=0 +``` + +Makna: + +```text +connect timeout +→ waktu membuka koneksi ke CodeBuddy2API + +first-byte timeout +→ waktu menunggu header/chunk pertama + +stream idle timeout +→ maksimum waktu tanpa event setelah stream dimulai + +total timeout +→ 0 berarti tidak membatasi seluruh agentic request +``` + +Jangan memakai connect timeout sebagai total timeout. + +## 3. Retry Policy Berdasarkan Status + +Jangan retry: + +```text +400 invalid_request +400 content_filter +401 invalid credentials +403 permission denied +404 model/endpoint not found +422 validation error +``` + +Boleh retry/fallback secara terbatas: + +```text +408 request timeout +429 rate limited +502 bad gateway +503 service unavailable +504 gateway timeout +connect timeout +connection reset +DNS failure sementara +``` + +Konfigurasi: + +```env +ROUTER_MAX_RETRY_ATTEMPTS=2 +ROUTER_RETRY_BASE_DELAY_MS=500 +ROUTER_RETRY_MAX_DELAY_MS=3000 +``` + +Jangan gunakan 10 attempt sebagai default. + +## 4. Jangan Retry Setelah Stream Dimulai + +Setelah salah satu dari berikut diterima: + +- assistant content; +- reasoning chunk; +- tool call; +- role chunk; +- event SSE valid; + +request dianggap sudah dimulai. + +Setelah itu: + +- jangan fallback ke provider lain; +- jangan retry request dari awal; +- jangan mengirim tool call dua kali; +- jangan menandai key lain sebagai gagal karena stream yang sama; +- jika stream terputus, return error stream ke client. + +Tambahkan state: + +```text +stream_started=true +``` + +dan log: + +```text +retry_skipped_reason=stream_already_started +``` + +## 5. Cancellation + +Jika Claude Code atau Shiteru membatalkan request: + +- abort fetch upstream; +- hentikan retry; +- hentikan fallback; +- jangan lanjutkan request pada background; +- jangan menandai key unavailable karena downstream cancel. + +Gunakan `AbortController` atau mekanisme cancellation runtime yang sesuai. + +Log: + +```text +request_id=... +downstream_disconnected=true +upstream_aborted=true +retry_cancelled=true +``` + +## 6. Provider Health dan Unavailable State + +Jangan menandai key unavailable permanen hanya karena satu connect timeout ke relay lokal. + +Bedakan: + +```text +invalid_key +rate_limited +provider_unavailable +relay_connect_timeout +client_cancelled +payload_invalid +``` + +Aturan: + +```text +400 payload_invalid +→ jangan mengubah status key + +client cancelled +→ jangan mengubah status key + +connect timeout lokal +→ cooldown pendek, bukan invalid + +401 +→ invalid key + +429 +→ cooldown sesuai retry-after + +5xx +→ temporary unavailable +``` + +Tambahkan konfigurasi: + +```env +ROUTER_CONNECT_FAILURE_COOLDOWN_SECONDS=15 +ROUTER_5XX_COOLDOWN_SECONDS=30 +ROUTER_429_DEFAULT_COOLDOWN_SECONDS=60 +``` + +## 7. Internal Docker Endpoint + +Untuk provider CodeBuddy2API pada VPS yang sama, gunakan: + +```text +http://cb2api:8001/codebuddy/v1 +``` + +Jangan menggunakan: + +```text +https://cb2api.heracles.id/codebuddy/v1 +http://127.0.0.1:8001/codebuddy/v1 +``` + +Tambahkan startup validation atau provider test yang memeriksa: + +```text +GET http://cb2api:8001/health +``` + +Namun health check tidak boleh terlalu sering atau memicu load besar. + +## 8. Safe Telemetry + +Setiap inbound request memiliki `request_id`. + +Catat: + +```text +request_id +provider +key_fingerprint +attempt +max_attempts +base_url_host +connect_ms +time_to_headers_ms +time_to_first_chunk_ms +stream_duration_ms +stream_started +upstream_status +error_type +fallback_reason +cooldown_seconds +downstream_disconnected +``` + +Jangan log: + +- raw API key; +- Authorization; +- prompt; +- isi file; +- tool result; +- cookie. + +Contoh: + +```text +request_id=abc123 provider=codebuddy attempt=1 stage=connect_start +request_id=abc123 provider=codebuddy attempt=1 connected_ms=12 +request_id=abc123 provider=codebuddy first_chunk_ms=2450 +request_id=abc123 stream_started=true +request_id=abc123 retry_skipped_reason=stream_already_started +``` + +## 9. Preserve Error Detail + +Jangan mengubah seluruh error menjadi: + +```text +fetch connect timeout +``` + +Jika upstream memberi status HTTP, pertahankan: + +```json +{ + "error": { + "message": "...", + "type": "invalid_request_error", + "code": "invalid_request" + } +} +``` + +Bedakan error: + +```text +connect_timeout +first_byte_timeout +stream_idle_timeout +upstream_http_error +client_cancelled +payload_invalid +``` + +## 10. Concurrency + +Audit apakah fallback atau retry dijalankan paralel. + +Jangan menjalankan banyak provider secara paralel untuk satu request kecuali memang mode hedging diaktifkan secara eksplisit. + +Tambahkan konfigurasi: + +```env +ROUTER_MAX_CONCURRENT_REQUESTS_PER_PROVIDER=4 +ROUTER_MAX_CONCURRENT_REQUESTS_PER_KEY=2 +``` + +Pastikan tool-heavy request tidak menggandakan attempt secara paralel. + +## 11. Tests + +Tambahkan automated tests: + +1. HTTP 400 tidak di-retry. +2. HTTP 422 tidak di-retry. +3. HTTP 429 di-retry maksimal sesuai konfigurasi. +4. Connect timeout dibedakan dari first-byte timeout. +5. Setelah first chunk, tidak ada retry. +6. Tool call pertama menandai stream sudah dimulai. +7. Client disconnect membatalkan upstream. +8. Client disconnect tidak menandai key unavailable. +9. Connect timeout hanya memberi cooldown pendek. +10. 401 menandai key invalid. +11. 429 memakai `Retry-After` bila tersedia. +12. Maksimal retry default adalah 2. +13. Tidak ada retry paralel. +14. Error body OpenAI-compatible dipertahankan. +15. Raw API key tidak muncul di log. +16. Provider internal `http://cb2api:8001` dapat divalidasi. +17. Stream SSE normal diteruskan incremental. +18. `[DONE]` hanya dikirim sekali. + +## Acceptance Criteria + +Pekerjaan selesai jika: + +- HTTP 400 langsung dikembalikan tanpa retry; +- `attempt 1/10` tidak lagi menjadi default; +- connect timeout tidak membatasi seluruh request; +- first-byte timeout cocok untuk request agentic; +- tidak ada fallback setelah stream dimulai; +- tidak ada tool call ganda; +- client disconnect menghentikan seluruh attempt; +- key tidak ditandai unavailable karena payload invalid; +- retry/fallback dapat dijelaskan dari log; +- semua test lulus. + +## Output Agent + +Sebelum implementasi: + +- analisis retry saat ini; +- analisis timeout saat ini; +- jelaskan kapan key ditandai unavailable; +- sebutkan file yang akan diubah. + +Setelah implementasi: + +- daftar file yang diubah; +- tabel retry policy; +- konfigurasi environment baru; +- contoh telemetry; +- hasil test; +- perintah deployment; +- catatan backward compatibility. + +Implementasikan perubahan secara langsung. Jangan hanya memberikan saran. +``` \ No newline at end of file diff --git a/config.py b/config.py index d1cab05..3f26fa8 100644 --- a/config.py +++ b/config.py @@ -44,7 +44,28 @@ "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000, "CODEBUDDY_MODEL_ALIASES": "", "CODEBUDDY_DEFAULT_MODEL": "auto-chat", - "CODEBUDDY_UNKNOWN_MODEL_POLICY": "passthrough" + "CODEBUDDY_UNKNOWN_MODEL_POLICY": "passthrough", + # --- Upstream connection / timeout tuning --- + # Separate timeouts so a slow-but-alive agentic stream is not killed by the + # same deadline that guards opening a connection. See get_codebuddy_* below. + "CODEBUDDY_CONNECT_TIMEOUT_SECONDS": 30, + "CODEBUDDY_POOL_TIMEOUT_SECONDS": 30, + "CODEBUDDY_WRITE_TIMEOUT_SECONDS": 60, + "CODEBUDDY_FIRST_BYTE_TIMEOUT_SECONDS": 180, + # 0 means unlimited read timeout once the stream is flowing. + "CODEBUDDY_STREAM_READ_TIMEOUT_SECONDS": 0, + "CODEBUDDY_MAX_KEEPALIVE_CONNECTIONS": 30, + "CODEBUDDY_MAX_CONNECTIONS": 100, + "CODEBUDDY_KEEPALIVE_EXPIRY_SECONDS": 60, + # --- Upstream concurrency limiter --- + "CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS": 20, + "CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS": 60, + # --- SSE heartbeat while waiting for the first upstream chunk --- + "CODEBUDDY_HEARTBEAT_INTERVAL_SECONDS": 15, + # --- Large-request warning thresholds (log only; never truncates) --- + "CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH": 50000, + "CODEBUDDY_WARN_MESSAGE_COUNT": 40, + "CODEBUDDY_WARN_TOOL_COUNT": 30, } # --- Core Functions --- @@ -278,6 +299,93 @@ def get_codebuddy_api_key_cooldown_seconds() -> int: raise ValueError("CODEBUDDY_API_KEY_COOLDOWN_SECONDS must be zero or greater") return cooldown + +# --- Upstream connection / timeout tuning --- + +def _get_non_negative_float(key: str, default: float) -> float: + """Return a validated non-negative float config value (0 allowed).""" + try: + value = float(_get_config_value(key)) + except (TypeError, ValueError): + return default + return value if value >= 0 else default + + +def _get_positive_int(key: str, default: int) -> int: + """Return a validated strictly-positive int config value.""" + try: + value = int(_get_config_value(key)) + except (TypeError, ValueError): + return default + return value if value > 0 else default + + +def get_codebuddy_connect_timeout_seconds() -> float: + return _get_non_negative_float("CODEBUDDY_CONNECT_TIMEOUT_SECONDS", 30.0) + + +def get_codebuddy_pool_timeout_seconds() -> float: + return _get_non_negative_float("CODEBUDDY_POOL_TIMEOUT_SECONDS", 30.0) + + +def get_codebuddy_write_timeout_seconds() -> float: + return _get_non_negative_float("CODEBUDDY_WRITE_TIMEOUT_SECONDS", 60.0) + + +def get_codebuddy_first_byte_timeout_seconds() -> float: + """Deadline for the upstream to return status/headers and the first chunk. + + Distinct from the stream read timeout: this guards the wait *before* the + stream starts flowing, so a dead upstream fails fast without capping a + long-running agentic stream. + """ + return _get_non_negative_float("CODEBUDDY_FIRST_BYTE_TIMEOUT_SECONDS", 180.0) + + +def get_codebuddy_stream_read_timeout_seconds() -> float: + """Per-read timeout once the stream is flowing. 0 means unlimited.""" + return _get_non_negative_float("CODEBUDDY_STREAM_READ_TIMEOUT_SECONDS", 0.0) + + +def get_codebuddy_max_keepalive_connections() -> int: + return _get_positive_int("CODEBUDDY_MAX_KEEPALIVE_CONNECTIONS", 30) + + +def get_codebuddy_max_connections() -> int: + return _get_positive_int("CODEBUDDY_MAX_CONNECTIONS", 100) + + +def get_codebuddy_keepalive_expiry_seconds() -> float: + return _get_non_negative_float("CODEBUDDY_KEEPALIVE_EXPIRY_SECONDS", 60.0) + + +def get_codebuddy_max_concurrent_upstream_requests() -> int: + return _get_positive_int("CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS", 20) + + +def get_codebuddy_upstream_queue_timeout_seconds() -> float: + return _get_non_negative_float("CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS", 60.0) + + +def get_codebuddy_heartbeat_interval_seconds() -> float: + """SSE heartbeat cadence while waiting for the first upstream chunk. + + 0 disables heartbeats entirely. + """ + return _get_non_negative_float("CODEBUDDY_HEARTBEAT_INTERVAL_SECONDS", 15.0) + + +def get_codebuddy_warn_total_content_length() -> int: + return _get_positive_int("CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH", 50000) + + +def get_codebuddy_warn_message_count() -> int: + return _get_positive_int("CODEBUDDY_WARN_MESSAGE_COUNT", 40) + + +def get_codebuddy_warn_tool_count() -> int: + return _get_positive_int("CODEBUDDY_WARN_TOOL_COUNT", 30) + # --- Public Setter for Hot-Reload --- def update_settings(new_settings: Dict[str, Any]): diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index d0e6cff..874bc6b 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -8,6 +8,8 @@ import hashlib import logging import asyncio +import contextlib +import contextvars from dataclasses import dataclass from typing import Optional, Dict, Any, List, AsyncGenerator, Set @@ -84,24 +86,75 @@ def get_ssl_verify() -> bool: return ssl_verify # --- HTTP client configuration --- -HTTP_CLIENT_CONFIG = { - "verify": SecurityConfig.get_ssl_verify(), - "timeout": httpx.Timeout(300.0, connect=30.0, read=300.0), - "limits": httpx.Limits(max_keepalive_connections=20, max_connections=100) -} + +def _build_http_client_config() -> Dict[str, Any]: + """Build the shared client config from separated upstream timeout settings. + + The read timeout intentionally covers only the wait for the first + status/headers response; per-chunk stream reads are governed separately by + ``CODEBUDDY_STREAM_READ_TIMEOUT_SECONDS`` (0 = unlimited) enforced in the + streaming loop. A single global timeout must never cap an agentic stream. + """ + from config import ( + get_codebuddy_connect_timeout_seconds, + get_codebuddy_pool_timeout_seconds, + get_codebuddy_write_timeout_seconds, + get_codebuddy_first_byte_timeout_seconds, + get_codebuddy_stream_read_timeout_seconds, + get_codebuddy_max_keepalive_connections, + get_codebuddy_max_connections, + get_codebuddy_keepalive_expiry_seconds, + ) + + connect = get_codebuddy_connect_timeout_seconds() + pool = get_codebuddy_pool_timeout_seconds() + write = get_codebuddy_write_timeout_seconds() + first_byte = get_codebuddy_first_byte_timeout_seconds() + stream_read = get_codebuddy_stream_read_timeout_seconds() + + # httpx uses one `read` timeout per read() call. Set it to the larger of the + # first-byte deadline and the per-chunk stream read timeout so neither the + # initial header wait nor a long inter-chunk gap trips it prematurely. A + # value of 0 (unlimited stream read) maps to None (no httpx read timeout); + # the first-byte deadline is then enforced explicitly by the caller. + read_timeout: Optional[float] + if stream_read and stream_read > 0: + read_timeout = max(first_byte, stream_read) + else: + read_timeout = None + + return { + "verify": SecurityConfig.get_ssl_verify(), + "timeout": httpx.Timeout( + connect=connect or None, + read=read_timeout, + write=write or None, + pool=pool or None, + ), + "limits": httpx.Limits( + max_keepalive_connections=get_codebuddy_max_keepalive_connections(), + max_connections=get_codebuddy_max_connections(), + keepalive_expiry=get_codebuddy_keepalive_expiry_seconds(), + ), + } # --- Async-safe HTTP client pool --- _http_client_pool: Optional[httpx.AsyncClient] = None _client_lock = asyncio.Lock() async def get_http_client() -> httpx.AsyncClient: - """Get the global HTTP client pool - async-safe""" + """Get the global HTTP client pool - async-safe. + + A single reusable AsyncClient is created lazily and shared by every request + so the connection pool is reused. It is never created per-request and never + closed after a single request; only ``close_http_client`` (shutdown) closes it. + """ global _http_client_pool if _http_client_pool is None: async with _client_lock: # Double-checked locking pattern - async version if _http_client_pool is None: - _http_client_pool = httpx.AsyncClient(**HTTP_CLIENT_CONFIG) + _http_client_pool = httpx.AsyncClient(**_build_http_client_config()) return _http_client_pool async def close_http_client(): @@ -112,6 +165,130 @@ async def close_http_client(): await _http_client_pool.aclose() _http_client_pool = None + +# --- Request-scoped telemetry --- +# A contextvar carries the current request_id so stage logs can be correlated +# without threading it through every call. Never carries prompt/key material. +_request_id_var: contextvars.ContextVar[str] = contextvars.ContextVar( + "codebuddy_request_id", default="-" +) + + +def log_stage(stage: str, **fields: Any) -> None: + """Emit a single safe stage-telemetry line for the current request. + + Only metadata is logged (durations, counts, statuses, fingerprints). Never + logs prompts, keys, tool results, or headers. Field values are rendered as + ``key=value`` pairs in a stable order. + """ + request_id = _request_id_var.get() + if fields: + extra = " " + " ".join(f"{k}={v}" for k, v in fields.items()) + else: + extra = "" + logger.info("request_id=%s stage=%s%s", request_id, stage, extra) + + +def _now_ms() -> float: + """Monotonic clock in milliseconds for elapsed-time telemetry.""" + return time.monotonic() * 1000.0 + + +# --- Upstream concurrency limiter --- +# Bounds simultaneous in-flight upstream requests so tool-heavy bursts do not +# pile up sockets/tasks without limit. Requests beyond the limit wait for a slot +# up to the queue timeout, then receive HTTP 503. The semaphore is rebuilt when +# the configured size changes (hot-reload safe). +_upstream_semaphore: Optional[asyncio.Semaphore] = None +_upstream_semaphore_size: int = 0 +_semaphore_lock = asyncio.Lock() + + +class UpstreamQueueTimeout(Exception): + """Raised when a request cannot acquire an upstream slot within the timeout.""" + + +async def _get_upstream_semaphore() -> asyncio.Semaphore: + global _upstream_semaphore, _upstream_semaphore_size + from config import get_codebuddy_max_concurrent_upstream_requests + + desired = get_codebuddy_max_concurrent_upstream_requests() + if _upstream_semaphore is None or desired != _upstream_semaphore_size: + async with _semaphore_lock: + if _upstream_semaphore is None or desired != _upstream_semaphore_size: + _upstream_semaphore = asyncio.Semaphore(desired) + _upstream_semaphore_size = desired + return _upstream_semaphore + + +class UpstreamSlot: + """Async context manager that acquires an upstream concurrency slot. + + Records queue wait time and guarantees the slot is released exactly once. + ``release()`` is idempotent and safe to call from either ``__aexit__`` (the + non-stream path) or a streaming generator's ``finally`` (the stream path). + + For streaming, ownership is transferred to the generator via ``detach()`` + so the slot is held for the full stream lifetime and released when the + stream ends, errors, or the client disconnects — not when the endpoint's + ``async with`` block exits. + """ + + def __init__(self) -> None: + self._semaphore: Optional[asyncio.Semaphore] = None + self._acquired = False + self._detached = False + self.queue_wait_ms = 0.0 + + async def __aenter__(self) -> "UpstreamSlot": + from config import get_codebuddy_upstream_queue_timeout_seconds + + self._semaphore = await _get_upstream_semaphore() + timeout = get_codebuddy_upstream_queue_timeout_seconds() + start = _now_ms() + try: + if timeout and timeout > 0: + await asyncio.wait_for(self._semaphore.acquire(), timeout=timeout) + else: + await self._semaphore.acquire() + except asyncio.TimeoutError as exc: + self.queue_wait_ms = _now_ms() - start + raise UpstreamQueueTimeout() from exc + self._acquired = True + self.queue_wait_ms = _now_ms() - start + return self + + def release(self) -> None: + """Release the slot exactly once; safe to call multiple times. + + Unconditional: used by the streaming generator's ``finally`` (which owns + the slot after ``detach()``) and by the non-stream path. + """ + if self._acquired and self._semaphore is not None: + self._semaphore.release() + self._acquired = False + + def release_unless_detached(self) -> None: + """Release only if ownership was NOT handed off to a streaming generator. + + The endpoint's ``finally`` calls this so a live streaming response keeps + its slot for the full stream lifetime; the generator releases it later. + """ + if not self._detached: + self.release() + + def detach(self) -> None: + """Transfer release responsibility to the streaming generator. + + After ``detach()`` neither ``release_unless_detached()`` nor + ``__aexit__`` will free the slot; the streaming generator's ``finally`` + must call ``release()``. + """ + self._detached = True + + async def __aexit__(self, exc_type, exc, tb) -> None: + self.release_unless_detached() + # --- Application lifecycle management --- class AppLifecycleManager: """Application lifecycle manager - handles resource cleanup""" @@ -145,9 +322,13 @@ async def shutdown(): lifecycle_manager = AppLifecycleManager() # --- Standard response headers --- +# no-transform prevents proxies from buffering/altering the SSE body; +# X-Accel-Buffering: no disables nginx/proxy response buffering so chunks are +# flushed to the client incrementally rather than accumulated. SSE_HEADERS = { - "Cache-Control": "no-cache", + "Cache-Control": "no-cache, no-transform", "Connection": "keep-alive", + "X-Accel-Buffering": "no", "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "GET, POST, OPTIONS", "Access-Control-Allow-Headers": "*" @@ -631,6 +812,49 @@ def _log_request_diagnostics( ) +def _request_size_metadata(payload: Dict[str, Any]) -> tuple[int, int, int]: + """Return (message_count, tool_count, total_content_length) as safe metadata. + + ``total_content_length`` counts text characters only; it never includes key + or header material and is used for the large-request warning. + """ + messages = payload.get("messages", []) or [] + _roles, _types, content_lengths, tool_content_count = _content_types_and_lengths(messages) + message_count = len(messages) + tool_count = tool_content_count + len(payload.get("tools", []) or []) + total_content_length = sum(content_lengths) + return message_count, tool_count, total_content_length + + +def _log_large_request_warning(payload: Dict[str, Any]) -> None: + """Emit a warning (never truncate) when a request exceeds configured thresholds.""" + try: + from config import ( + get_codebuddy_warn_total_content_length, + get_codebuddy_warn_message_count, + get_codebuddy_warn_tool_count, + ) + warn_len = get_codebuddy_warn_total_content_length() + warn_msgs = get_codebuddy_warn_message_count() + warn_tools = get_codebuddy_warn_tool_count() + except Exception: + return + + message_count, tool_count, total_content_length = _request_size_metadata(payload) + if ( + total_content_length >= warn_len + or message_count >= warn_msgs + or tool_count >= warn_tools + ): + log_stage( + "large_agentic_request", + large_agentic_request="true", + message_count=message_count, + tool_count=tool_count, + total_content_length=total_content_length, + ) + + class CodeBuddyStreamService: """CodeBuddy streaming service; each method performs exactly one upstream attempt.""" @@ -649,14 +873,50 @@ async def open_stream_response( payload: Dict[str, Any], headers: Dict[str, str], key_id: Optional[str] = None, + slot: Optional["UpstreamSlot"] = None, ) -> StreamingResponse: - """Establish the connection and verify the upstream status before returning a StreamingResponse.""" + """Open the upstream connection, verify status, then return a StreamingResponse. + + The connection is opened and its status/headers awaited under the + first-byte timeout. Only a non-200 status blocks here (so failover can + pick another key). On a 200 the StreamingResponse is returned + immediately; moderation detection, heartbeats, and per-chunk timeouts + all happen inside the streaming generator so the HTTP 200 reaches the + downstream client without waiting for the first content chunk. This is + what stops requests from stalling silently after ``outcome=prepared``. + + ``slot``, when provided, is released in the generator's ``finally`` so + the concurrency slot is held for the full stream lifetime and freed on + completion, error, or downstream disconnect. + """ + from config import ( + get_codebuddy_first_byte_timeout_seconds, + get_codebuddy_stream_read_timeout_seconds, + get_codebuddy_heartbeat_interval_seconds, + ) + + first_byte_timeout = get_codebuddy_first_byte_timeout_seconds() + stream_read_timeout = get_codebuddy_stream_read_timeout_seconds() + heartbeat_interval = get_codebuddy_heartbeat_interval_seconds() + client = await get_http_client() request = client.build_request( "POST", get_codebuddy_api_url(), json=payload, headers=headers ) + log_stage("upstream_send_start") + send_start = _now_ms() try: - response = await client.send(request, stream=True) + # Enforce the first-byte/header deadline explicitly: with an + # unlimited stream read timeout the httpx read timeout is None, so + # the header wait must be bounded here instead. + if first_byte_timeout and first_byte_timeout > 0: + response = await asyncio.wait_for( + client.send(request, stream=True), timeout=first_byte_timeout + ) + else: + response = await client.send(request, stream=True) + except asyncio.TimeoutError as exc: + raise UpstreamAttemptError("transient", 504, "upstream_first_byte_timeout") from exc except httpx.TimeoutException as exc: raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc except httpx.RequestError as exc: @@ -667,8 +927,14 @@ async def open_stream_response( await response.aclose() raise self._classify_status(status_code) - # Do not expose upstream response headers; only this proxy's fixed SSE headers - # are sent downstream after the first chunk is safely available. + log_stage( + "upstream_headers_received", + status=response.status_code, + elapsed_ms=round(_now_ms() - send_start, 1), + ) + + # Do not expose upstream response headers; only this proxy's fixed SSE + # headers are sent downstream. async def converted_chunks(): buffer = "" @@ -705,76 +971,161 @@ async def converted_chunks(): stream = converted_chunks() - # Pre-buffer the leading chunks so a Mandarin moderation refusal can be - # detected before any assistant content is sent downstream. The refusal - # is short, so a small buffer suffices; normal replies simply get - # replayed afterwards in order. - buffered_lines: List[str] = [] - accumulated_content = "" - moderation_detected = False - try: - async for line in stream: - buffered_lines.append(line) - if '[DONE]' in line: - break - accumulated_content += _extract_delta_content(line) - if is_codebuddy_moderation_response(accumulated_content): - moderation_detected = True - break - if len(accumulated_content) >= MODERATION_STREAM_BUFFER_CHARS: - break - except httpx.TimeoutException as exc: - await response.aclose() - raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc - except httpx.RequestError as exc: - await response.aclose() - raise UpstreamAttemptError("transient", 502, "upstream_network_error") from exc - except Exception as exc: - await response.aclose() - raise UpstreamAttemptError("fatal", 502, "upstream_response_invalid") from exc - - if moderation_detected: - # Nothing normal was sent yet: replace the whole stream with an - # OpenAI-compatible content_filter stream. The key is valid, so do - # not mark it failed or fail over. - await response.aclose() - - async def moderation_core(): - async for item in _moderation_stream(): - yield item + def _release_slot() -> None: + if slot is not None: + slot.release() - return StreamingResponse( - moderation_core(), media_type="text/event-stream", headers={ - **SSE_HEADERS, "X-CodeBuddy-Moderation": "true" - } - ) + # Capture the request_id now: the StreamingResponse generator runs in a + # separate task after this function returns, so the contextvar must be + # re-bound inside it for stage logs to carry the correct request_id. + stream_request_id = _request_id_var.get() async def stream_core(): + _request_id_var.set(stream_request_id) + # A background producer isolates the upstream read from the + # heartbeat/timeout timer: timeouts never cancel a read mid-flight + # (which would corrupt the httpx stream); they only decide whether + # to emit a heartbeat or give up. + queue: asyncio.Queue = asyncio.Queue(maxsize=64) + + async def producer(): + try: + async for line in stream: + await queue.put(("line", line)) + await queue.put(("end", None)) + except httpx.RequestError: + await queue.put(("error", "upstream_stream_error")) + except Exception: + await queue.put(("error", "upstream_stream_processing_error")) + + producer_task = asyncio.create_task(producer()) + + async def next_event(timeout: Optional[float]): + if timeout and timeout > 0: + return await asyncio.wait_for(queue.get(), timeout=timeout) + return await queue.get() + + first_chunk_ms: Optional[float] = None + chunk_count = 0 + stream_start = _now_ms() try: + # --- Phase 1: buffer leading lines to detect a moderation + # refusal before any assistant content is emitted. Heartbeats + # keep the connection alive while we wait; the first-byte + # deadline bounds the wait for the very first line. + buffered_lines: List[str] = [] + accumulated_content = "" + moderation_detected = False + stream_ended = False + first_byte_start = _now_ms() + + while True: + remaining: Optional[float] = None + if first_byte_timeout and first_byte_timeout > 0: + elapsed = (_now_ms() - first_byte_start) / 1000.0 + remaining = first_byte_timeout - elapsed + if remaining <= 0: + log_stage("request_failed", reason="first_byte_timeout") + yield format_sse_error( + "Upstream timed out before the first chunk", + "upstream_first_byte_timeout", + ) + return + + wait = None + if heartbeat_interval and heartbeat_interval > 0: + wait = heartbeat_interval + if remaining is not None: + wait = remaining if wait is None else min(wait, remaining) + + try: + kind, value = await next_event(wait) + except asyncio.TimeoutError: + # Heartbeat only before the first content chunk. It is + # an SSE comment, never assistant content, and never a + # data event, so it cannot corrupt [DONE]. + yield ": ping\n\n" + continue + + if kind == "error": + if key_id is not None: + await codebuddy_api_key_manager.mark_transient_error( + key_id, value + ) + yield format_sse_error("Upstream stream interrupted", value) + return + if kind == "end": + stream_ended = True + break + + line = value + buffered_lines.append(line) + if '[DONE]' in line: + break + accumulated_content += _extract_delta_content(line) + if is_codebuddy_moderation_response(accumulated_content): + moderation_detected = True + break + if len(accumulated_content) >= MODERATION_STREAM_BUFFER_CHARS: + break + + if moderation_detected: + # Nothing normal was sent yet: emit an OpenAI-compatible + # content_filter stream instead. The key is valid, so it is + # not marked failed and no failover occurs. + log_stage("upstream_moderation_detected") + async for item in _moderation_stream(): + yield item + return + + # --- Phase 2: flush buffered lines then pass through the rest, + # applying the stream-read timeout (0 = unlimited). + first_chunk_ms = _now_ms() - stream_start + log_stage("upstream_first_chunk", elapsed_ms=round(first_chunk_ms, 1)) for line in buffered_lines: + chunk_count += 1 yield line - async for chunk in stream: - yield chunk - except httpx.RequestError: - logger.warning("CodeBuddy upstream stream interrupted") - if key_id is not None: - await codebuddy_api_key_manager.mark_transient_error( - key_id, "stream_interrupted" - ) - yield format_sse_error( - "Upstream stream interrupted", "upstream_stream_error" - ) - except Exception: - logger.error("Unexpected CodeBuddy stream processing error") - if key_id is not None: - await codebuddy_api_key_manager.mark_transient_error( - key_id, "stream_processing_error" - ) - yield format_sse_error( - "Upstream stream interrupted", "upstream_stream_error" + + if not stream_ended: + while True: + try: + kind, value = await next_event(stream_read_timeout) + except asyncio.TimeoutError: + log_stage("request_failed", reason="stream_read_timeout") + yield format_sse_error( + "Upstream stream idle timeout", + "upstream_stream_idle_timeout", + ) + return + if kind == "error": + if key_id is not None: + await codebuddy_api_key_manager.mark_transient_error( + key_id, value + ) + yield format_sse_error("Upstream stream interrupted", value) + return + if kind == "end": + break + chunk_count += 1 + yield value + + log_stage( + "upstream_stream_finished", + chunks=chunk_count, + duration_ms=round(_now_ms() - stream_start, 1), ) + except asyncio.CancelledError: + # Downstream disconnected: stop everything, do not retry, do not + # leave a background upstream request running. + log_stage("downstream_disconnected") + raise finally: - await response.aclose() + producer_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await producer_task + with contextlib.suppress(Exception): + await response.aclose() + _release_slot() return StreamingResponse( stream_core(), media_type="text/event-stream", headers=SSE_HEADERS @@ -783,12 +1134,32 @@ async def stream_core(): async def handle_non_stream_response( self, payload: Dict[str, Any], headers: Dict[str, str] ) -> Dict[str, Any]: - """Perform a single non-streaming upstream request and aggregate the SSE response.""" + """Perform a single non-streaming upstream request and aggregate the SSE response. + + Upstream is always streamed (CodeBuddy only supports streaming); this + method aggregates the SSE stream into a single OpenAI-compatible + completion. The first-byte deadline bounds the wait for status/headers + so a dead upstream fails fast, while the aggregation loop reads chunks + as they arrive rather than buffering the whole body up front. + """ + from config import get_codebuddy_first_byte_timeout_seconds + + first_byte_timeout = get_codebuddy_first_byte_timeout_seconds() + client = await get_http_client() + request = client.build_request( + "POST", get_codebuddy_api_url(), json=payload, headers=headers + ) + log_stage("upstream_send_start") + send_start = _now_ms() try: - client = await get_http_client() - response = await client.post( - get_codebuddy_api_url(), json=payload, headers=headers - ) + if first_byte_timeout and first_byte_timeout > 0: + response = await asyncio.wait_for( + client.send(request, stream=True), timeout=first_byte_timeout + ) + else: + response = await client.send(request, stream=True) + except asyncio.TimeoutError as exc: + raise UpstreamAttemptError("transient", 504, "upstream_first_byte_timeout") from exc except httpx.TimeoutException as exc: raise UpstreamAttemptError("transient", 504, "upstream_timeout") from exc except httpx.RequestError as exc: @@ -799,6 +1170,12 @@ async def handle_non_stream_response( await response.aclose() raise self._classify_status(status_code) + log_stage( + "upstream_headers_received", + status=response.status_code, + elapsed_ms=round(_now_ms() - send_start, 1), + ) + try: aggregator = StreamResponseAggregator() raw_text = "" @@ -1169,6 +1546,13 @@ async def chat_completions( auth_context: ClientAuthContext = Depends(authenticate_inference) ): """CodeBuddy V1 chat completions API, supporting relay and per-request passthrough.""" + # Assign a request_id for stage telemetry. Prefer the client-supplied + # X-Request-ID for cross-service correlation (9Router → cb2api) and fall + # back to a generated id. + request_id = x_request_id or uuid.uuid4().hex + _request_id_var.set(request_id) + log_stage("request_received") + try: request_body = await request.json() except Exception: @@ -1182,6 +1566,7 @@ async def chat_completions( return openai_error_response( str(exc.detail), "invalid_request_error", "invalid_request", exc.status_code ) + log_stage("request_normalized") passthrough_credential: Optional[ResolvedCredential] = None if auth_context.mode == "passthrough": @@ -1240,6 +1625,8 @@ async def chat_completions( 400, ) usage_stats_manager.record_model_usage(payload.get("model", "unknown")) + log_stage("request_prepared") + _log_large_request_warning(payload) service = CodeBuddyStreamService() client_wants_stream = bool(request_body.get("stream", False)) if isinstance(request_body, dict) else False excluded_ids: Set[str] = set() @@ -1255,6 +1642,78 @@ async def chat_completions( 500, ) + # Acquire an upstream concurrency slot before making any upstream call so + # tool-heavy bursts cannot pile up sockets/tasks without bound. The slot is + # released here for the non-stream and error paths; for a successful stream + # it is detached and released by the streaming generator instead (held for + # the stream's full lifetime). + log_stage("upstream_slot_wait_start") + slot = UpstreamSlot() + try: + await slot.__aenter__() + except UpstreamQueueTimeout: + log_stage( + "request_failed", + reason="upstream_queue_timeout", + queue_wait_ms=round(slot.queue_wait_ms, 1), + ) + return openai_error_response( + "CodeBuddy relay is temporarily at capacity", + "server_overloaded", + "upstream_queue_timeout", + 503, + ) + log_stage("upstream_slot_acquired", queue_wait_ms=round(slot.queue_wait_ms, 1)) + + try: + return await _run_attempts( + request=request, + service=service, + slot=slot, + payload=payload, + prep_info=prep_info, + source=source, + max_attempts=max_attempts, + passthrough_credential=passthrough_credential, + client_wants_stream=client_wants_stream, + request_profile=request_profile, + x_conversation_id=x_conversation_id, + x_conversation_request_id=x_conversation_request_id, + x_conversation_message_id=x_conversation_message_id, + x_request_id=x_request_id, + ) + finally: + # Release the slot unless a streaming response took ownership of it + # (a live stream holds its slot until its generator's finally runs). + slot.release_unless_detached() + + +async def _run_attempts( + *, + request: Request, + service: "CodeBuddyStreamService", + slot: "UpstreamSlot", + payload: Dict[str, Any], + prep_info: Dict[str, Any], + source: str, + max_attempts: int, + passthrough_credential: Optional["ResolvedCredential"], + client_wants_stream: bool, + request_profile: str, + x_conversation_id: Optional[str], + x_conversation_request_id: Optional[str], + x_conversation_message_id: Optional[str], + x_request_id: Optional[str], +): + """Run the credential/failover attempt loop for a prepared request. + + A concurrency ``slot`` is already held. On a successful streaming response + the slot is detached so the streaming generator owns its release; on every + other path the caller's ``finally`` releases it. + """ + excluded_ids: Set[str] = set() + last_error: Optional[UpstreamAttemptError] = None + for _attempt in range(max_attempts): if source == "passthrough": credential = passthrough_credential @@ -1304,8 +1763,11 @@ async def chat_completions( try: if client_wants_stream: result = await service.open_stream_response( - payload, headers, credential.key_id + payload, headers, credential.key_id, slot=slot ) + # The stream is live and owns the concurrency slot for its full + # lifetime; hand off release responsibility to its generator. + slot.detach() else: result = await service.handle_non_stream_response(payload, headers) if credential.key_id is not None: diff --git a/tests/test_codebuddy_streaming_stability.py b/tests/test_codebuddy_streaming_stability.py new file mode 100644 index 0000000..1b2b0e1 --- /dev/null +++ b/tests/test_codebuddy_streaming_stability.py @@ -0,0 +1,537 @@ +""" +Stability tests for the CodeBuddy relay: true streaming, SSE headers, shared +HTTP client reuse, upstream concurrency limiting (queue wait + 503), separated +timeouts (first-byte vs stream read), heartbeats, cancellation/slot release, and +no-retry-after-stream-start. + +These lock in the latency/streaming/cancellation fixes so a future change cannot +silently reintroduce full-body buffering, a leaked concurrency slot, a single +global timeout, or a retry after the stream has started. + +Uses mock upstream transports; no real CodeBuddy API key is required. +""" +import asyncio +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.codebuddy_api_key_manager import CodeBuddyApiKeyManager + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + + +# --------------------------------------------------------------------------- # +# Fixtures / helpers +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + from src import codebuddy_auth_router, settings_router + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + application.include_router(codebuddy_auth_router.router, prefix="/codebuddy") + application.include_router(settings_router.router, prefix="/api") + return application + + +@pytest.fixture +async def empty_pool(monkeypatch, tmp_path): + path = tmp_path / "keys.txt" + path.write_text("", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(path), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + return manager + + +@pytest.fixture(autouse=True) +def reset_semaphore(monkeypatch): + """Reset the global upstream concurrency limiter between tests.""" + codebuddy_router._upstream_semaphore = None + codebuddy_router._upstream_semaphore_size = 0 + yield + codebuddy_router._upstream_semaphore = None + codebuddy_router._upstream_semaphore_size = 0 + + +def configure(monkeypatch, profile="web"): + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + monkeypatch.setattr(codebuddy_router, "get_upstream_api_key_header", lambda: "both") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: profile) + monkeypatch.setattr( + codebuddy_router.usage_stats_manager, "record_model_usage", lambda _m: None + ) + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def install_transport(monkeypatch, transport): + """Install a shared client backed by a custom async transport.""" + client = httpx.AsyncClient(transport=transport) + + async def get_client(): + return client + + monkeypatch.setattr(codebuddy_router, "get_http_client", get_client) + return client + + +def sse_body(chunks, done=True): + """Build an SSE ``httpx.Response`` from ``(delta, finish_reason)`` pairs.""" + parts = [] + for delta, finish in chunks: + choice = {"index": 0, "delta": delta} + if finish is not None: + choice["finish_reason"] = finish + obj = { + "id": "chat-stab", + "object": "chat.completion.chunk", + "model": "auto-chat", + "choices": [choice], + } + parts.append("data: " + json.dumps(obj, ensure_ascii=False)) + text = "\n\n".join(parts) + "\n\n" + if done: + text += "data: [DONE]\n\n" + return httpx.Response(200, text=text, headers={"content-type": "text/event-stream"}) + + +async def request(app, path, token=None, method="GET", json_body=None): + headers = {} + if token is not None: + headers["Authorization"] = f"Bearer {token}" + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.request(method, path, headers=headers, json=json_body) + + +async def chat(app, token, stream=False, messages=None, tools=None): + body = { + "model": "auto-chat", + "messages": messages or [{"role": "user", "content": "halo, siapa kamu"}], + "stream": stream, + } + if tools is not None: + body["tools"] = tools + return await request(app, "/codebuddy/v1/chat/completions", token, "POST", body) + + +def parse_stream_events(text): + events = [] + for raw in text.split("\n\n"): + line = raw.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:"):].strip() + if payload == "[DONE]": + continue + events.append(json.loads(payload)) + return events + + +# --------------------------------------------------------------------------- # +# Custom transports for delay-based tests +# --------------------------------------------------------------------------- # + + +class _ChunkStream(httpx.AsyncByteStream): + """Async byte stream that yields chunks with an inter-chunk delay.""" + + def __init__(self, chunks, delay=0.0): + self._chunks = chunks + self._delay = delay + + async def __aiter__(self): + for chunk in self._chunks: + if self._delay: + await asyncio.sleep(self._delay) + yield chunk + + async def aclose(self): + return None + + +class SlowFirstByteTransport(httpx.AsyncBaseTransport): + """Delays returning the response object itself (simulates slow headers).""" + + def __init__(self, header_delay): + self._header_delay = header_delay + self.calls = 0 + + async def handle_async_request(self, request): + self.calls += 1 + await asyncio.sleep(self._header_delay) + body = b"data: " + json.dumps( + {"choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": "stop"}]} + ).encode() + b"\n\ndata: [DONE]\n\n" + return httpx.Response( + 200, stream=_ChunkStream([body]), headers={"content-type": "text/event-stream"} + ) + + +class SlowChunkTransport(httpx.AsyncBaseTransport): + """Returns headers immediately but delays the first body chunk.""" + + def __init__(self, first_chunk_delay): + self._delay = first_chunk_delay + self.calls = 0 + + async def handle_async_request(self, request): + self.calls += 1 + first = "data: " + json.dumps( + {"choices": [{"index": 0, "delta": {"content": "hello"}}]} + ) + "\n\n" + second = "data: [DONE]\n\n" + stream = _ChunkStream( + [first.encode(), second.encode()], delay=self._delay + ) + return httpx.Response( + 200, stream=stream, headers={"content-type": "text/event-stream"} + ) + + +# --------------------------------------------------------------------------- # +# 1. SSE headers: no-transform + X-Accel-Buffering: no +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_sse_headers_disable_proxy_buffering(monkeypatch, app, empty_pool): + configure(monkeypatch) + + def handler(_req): + return sse_body([({"content": "hi"}, "stop")]) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + assert response.headers["content-type"].startswith("text/event-stream") + assert "no-transform" in response.headers.get("cache-control", "") + assert response.headers.get("x-accel-buffering") == "no" + + +# --------------------------------------------------------------------------- # +# 2. Streaming preserves multiple incremental chunks in order +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_stream_preserves_incremental_chunks(monkeypatch, app, empty_pool): + configure(monkeypatch) + + def handler(_req): + return sse_body( + [ + ({"role": "assistant"}, None), + ({"content": "A"}, None), + ({"content": "B"}, None), + ({"content": "C"}, None), + ({}, "stop"), + ] + ) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + await upstream.aclose() + + assert response.status_code == 200 + text = response.text + assert text.count("[DONE]") == 1 + events = parse_stream_events(text) + streamed = "".join( + (e["choices"][0].get("delta", {}) or {}).get("content", "") for e in events + ) + assert streamed == "ABC" + + +# --------------------------------------------------------------------------- # +# 3. Shared HTTP client is built once and reused +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_shared_http_client_is_reused(monkeypatch): + # Force a fresh pool, then confirm repeated calls return the same instance. + await codebuddy_router.close_http_client() + try: + first = await codebuddy_router.get_http_client() + second = await codebuddy_router.get_http_client() + assert first is second + finally: + await codebuddy_router.close_http_client() + + +def test_http_client_config_uses_separated_timeouts(monkeypatch): + monkeypatch.setattr(codebuddy_router, "get_codebuddy_stream_read_timeout_seconds", None, raising=False) + from config import ( + get_codebuddy_connect_timeout_seconds, + get_codebuddy_first_byte_timeout_seconds, + ) + + cfg = codebuddy_router._build_http_client_config() + timeout = cfg["timeout"] + # connect uses the connect timeout; unlimited stream read maps to no read + # timeout so a long agentic stream is never capped by a global deadline. + assert timeout.connect == get_codebuddy_connect_timeout_seconds() + assert timeout.read is None # default stream read timeout is 0 (unlimited) + + +# --------------------------------------------------------------------------- # +# 4. Concurrency limiter: queue wait recorded, timeout -> 503 +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_queue_timeout_returns_503(monkeypatch, app, empty_pool): + configure(monkeypatch) + + # Exhaust the single slot and force a near-instant queue timeout. + exhausted = asyncio.Semaphore(1) + await exhausted.acquire() + + async def get_sema(): + return exhausted + + monkeypatch.setattr(codebuddy_router, "_get_upstream_semaphore", get_sema) + monkeypatch.setattr( + codebuddy_router, + "get_codebuddy_upstream_queue_timeout_seconds", + lambda: 0.05, + raising=False, + ) + monkeypatch.setattr( + "config.get_codebuddy_upstream_queue_timeout_seconds", lambda: 0.05 + ) + + def handler(_req): + return sse_body([({"content": "unreachable"}, "stop")]) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == 503 + error = response.json()["error"] + assert error["code"] == "upstream_queue_timeout" + assert error["type"] == "server_overloaded" + + +@pytest.mark.asyncio +async def test_slot_released_after_non_stream(monkeypatch, app, empty_pool): + configure(monkeypatch) + monkeypatch.setattr( + "config.get_codebuddy_max_concurrent_upstream_requests", lambda: 2 + ) + + def handler(_req): + return sse_body([({"content": "ok"}, "stop")]) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + assert response.status_code == 200 + sema = await codebuddy_router._get_upstream_semaphore() + # Both slots must be free again once the request completed. + assert sema._value == 2 + + +@pytest.mark.asyncio +async def test_slot_released_after_stream_completes(monkeypatch, app, empty_pool): + configure(monkeypatch) + monkeypatch.setattr( + "config.get_codebuddy_max_concurrent_upstream_requests", lambda: 2 + ) + + def handler(_req): + return sse_body([({"content": "hi"}, "stop")]) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True) + # Fully drain the streaming body so the generator's finally runs. + _ = response.text + await upstream.aclose() + + assert response.status_code == 200 + sema = await codebuddy_router._get_upstream_semaphore() + assert sema._value == 2 + + +# --------------------------------------------------------------------------- # +# 5. Separated timeouts: first-byte deadline fails fast; stream read is separate +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_first_byte_timeout_is_enforced(monkeypatch, app, empty_pool): + configure(monkeypatch) + monkeypatch.setattr( + "config.get_codebuddy_first_byte_timeout_seconds", lambda: 0.1 + ) + monkeypatch.setattr( + "config.get_codebuddy_stream_read_timeout_seconds", lambda: 0.0 + ) + monkeypatch.setattr( + "config.get_codebuddy_heartbeat_interval_seconds", lambda: 0.0 + ) + + transport = SlowFirstByteTransport(header_delay=0.5) + upstream = install_transport(monkeypatch, transport) + response = await chat(app, KEY_A, stream=False) + await upstream.aclose() + + # Slow headers trip the first-byte deadline -> upstream failure, not a hang. + assert response.status_code in (502, 504) + assert transport.calls == 1 + + +@pytest.mark.asyncio +async def test_heartbeat_emitted_while_waiting_and_not_content( + monkeypatch, app, empty_pool +): + configure(monkeypatch) + # Generous first-byte deadline, frequent heartbeats, slow first chunk. + monkeypatch.setattr( + "config.get_codebuddy_first_byte_timeout_seconds", lambda: 5.0 + ) + monkeypatch.setattr( + "config.get_codebuddy_stream_read_timeout_seconds", lambda: 0.0 + ) + monkeypatch.setattr( + "config.get_codebuddy_heartbeat_interval_seconds", lambda: 0.05 + ) + + transport = SlowChunkTransport(first_chunk_delay=0.18) + upstream = install_transport(monkeypatch, transport) + response = await chat(app, KEY_A, stream=True) + text = response.text + await upstream.aclose() + + assert response.status_code == 200 + # Heartbeat is an SSE comment, never a data event / assistant content. + assert ": ping" in text + events = parse_stream_events(text) + streamed = "".join( + (e["choices"][0].get("delta", {}) or {}).get("content", "") for e in events + ) + assert streamed == "hello" + # The ping comment must not have been parsed as a data event. + for e in events: + assert (e["choices"][0].get("delta", {}) or {}).get("content") != ": ping" + + +# --------------------------------------------------------------------------- # +# 6. No retry on fatal (HTTP 400) and single upstream call per stream +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_http_400_is_not_retried(monkeypatch, app, tmp_path): + # Multiple keys available so retry, if it happened, would call upstream again. + configure(monkeypatch) + keys = tmp_path / "keys.txt" + keys.write_text("key-one-000000000000\nkey-two-000000000000\n", encoding="utf-8") + manager = CodeBuddyApiKeyManager(str(keys), reload_interval=0) + await manager.reload() + monkeypatch.setattr(codebuddy_router, "codebuddy_api_key_manager", manager) + # Use api_key_file source (relay), not passthrough, to exercise failover. + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "relay") + monkeypatch.setattr(codebuddy_router, "get_codebuddy_request_profile", lambda: "web") + + calls = {"n": 0} + + def handler(_req): + calls["n"] += 1 + return httpx.Response( + 400, json={"error": {"message": "bad", "code": "invalid_request"}} + ) + + upstream = install_upstream(monkeypatch, handler) + response = await request( + app, + "/codebuddy/v1/chat/completions", + RELAY_PASSWORD, + "POST", + {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}], "stream": False}, + ) + await upstream.aclose() + + # A 400 is fatal: exactly one upstream call, no failover to the second key. + assert calls["n"] == 1 + assert response.status_code in (400, 502) + + +# --------------------------------------------------------------------------- # +# 7. Tool-heavy request (30 tools) still streams +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_thirty_tools_still_streams(monkeypatch, app, empty_pool): + configure(monkeypatch) + tools = [ + { + "type": "function", + "function": { + "name": f"tool_{i}", + "description": f"tool number {i}", + "parameters": {"type": "object", "properties": {}}, + }, + } + for i in range(30) + ] + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "done"}, "stop")]) + + upstream = install_upstream(monkeypatch, handler) + response = await chat(app, KEY_A, stream=True, tools=tools) + text = response.text + await upstream.aclose() + + assert response.status_code == 200 + assert len(seen["payload"]["tools"]) == 30 + events = parse_stream_events(text) + streamed = "".join( + (e["choices"][0].get("delta", {}) or {}).get("content", "") for e in events + ) + assert streamed == "done" + + +# --------------------------------------------------------------------------- # +# 8. Config getters validate and clamp +# --------------------------------------------------------------------------- # + + +def test_config_getters_defaults_and_validation(): + import config + + assert config.get_codebuddy_connect_timeout_seconds() == 30.0 + assert config.get_codebuddy_first_byte_timeout_seconds() == 180.0 + # 0 is a valid (unlimited) stream read timeout. + assert config.get_codebuddy_stream_read_timeout_seconds() == 0.0 + assert config.get_codebuddy_max_concurrent_upstream_requests() == 20 + assert config.get_codebuddy_upstream_queue_timeout_seconds() == 60.0 + assert config.get_codebuddy_heartbeat_interval_seconds() == 15.0 From 047a92ca3f4ba8433a9ad495fc9a419be1ed72df Mon Sep 17 00:00:00 2001 From: ranggaalk Date: Tue, 21 Jul 2026 21:19:39 +0700 Subject: [PATCH 9/9] add: Tools Adapter --- .env.example | 61 + config.py | 87 +- src/adapters/__init__.py | 1 + src/adapters/codebuddy/__init__.py | 29 + src/adapters/codebuddy/adapter.py | 702 ++++++++++ src/adapters/codebuddy/config.py | 77 ++ src/adapters/codebuddy/errors.py | 185 +++ src/adapters/codebuddy/message_normalizer.py | 278 ++++ src/adapters/codebuddy/models.py | 97 ++ src/adapters/codebuddy/request_mapper.py | 124 ++ src/adapters/codebuddy/response_mapper.py | 143 ++ src/adapters/codebuddy/stream_decoder.py | 138 ++ src/adapters/codebuddy/tool_call_state.py | 169 +++ src/adapters/codebuddy/tool_schema_adapter.py | 190 +++ src/adapters/codebuddy/transport.py | 212 +++ src/adapters/codebuddy/validation.py | 108 ++ src/codebuddy_router.py | 41 + tests/test_adapter_v2_integration.py | 443 ++++++ tests/test_adapter_v2_units.py | 661 +++++++++ tools-refactor-adapter.md | 1204 +++++++++++++++++ 20 files changed, 4949 insertions(+), 1 deletion(-) create mode 100644 src/adapters/__init__.py create mode 100644 src/adapters/codebuddy/__init__.py create mode 100644 src/adapters/codebuddy/adapter.py create mode 100644 src/adapters/codebuddy/config.py create mode 100644 src/adapters/codebuddy/errors.py create mode 100644 src/adapters/codebuddy/message_normalizer.py create mode 100644 src/adapters/codebuddy/models.py create mode 100644 src/adapters/codebuddy/request_mapper.py create mode 100644 src/adapters/codebuddy/response_mapper.py create mode 100644 src/adapters/codebuddy/stream_decoder.py create mode 100644 src/adapters/codebuddy/tool_call_state.py create mode 100644 src/adapters/codebuddy/tool_schema_adapter.py create mode 100644 src/adapters/codebuddy/transport.py create mode 100644 src/adapters/codebuddy/validation.py create mode 100644 tests/test_adapter_v2_integration.py create mode 100644 tests/test_adapter_v2_units.py create mode 100644 tools-refactor-adapter.md diff --git a/.env.example b/.env.example index 02372cb..016c3fe 100644 --- a/.env.example +++ b/.env.example @@ -21,6 +21,19 @@ CODEBUDDY_HOST=127.0.0.1 CODEBUDDY_PORT=8001 +# ----------------- +# Adapter selection (spec §1) +# ----------------- + +# (Optional) Which request pipeline handles /v1/chat/completions. +# v2 -> the isolated CodeBuddyAdapterV2 (default): strict model resolver, +# deterministic Anthropic/OpenAI message normalization, streaming +# tool-call state machine, granular timeouts, and cancellation. +# legacy -> the original monolithic handler, kept as a fallback. +# Default: v2 +CODEBUDDY_ADAPTER_VERSION=v2 + + # ----------------- # CodeBuddy API configuration # ----------------- @@ -52,6 +65,14 @@ CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH=2000 # Default: auto-chat CODEBUDDY_DEFAULT_MODEL=auto-chat +# (Optional) How an unknown model label (neither a known upstream model ID nor a +# configured alias) is handled. An unknown model is NEVER silently rewritten. +# passthrough (default) -> forward the requested model verbatim; CodeBuddy +# accepts or rejects it. +# reject -> return HTTP 400 (code=unknown_model) locally. +# default -> fall back to CODEBUDDY_DEFAULT_MODEL. +CODEBUDDY_UNKNOWN_MODEL_POLICY=passthrough + # (Optional) Model alias map for UI display labels that are not valid upstream # model IDs. Comma-separated alias=upstream pairs, matched case-insensitively. # Example: Claude Opus 4.7=claude-4.0,GPT-5 (UI)=gpt-5 @@ -100,6 +121,46 @@ CODEBUDDY_UPSTREAM_API_KEY_HEADER=bearer # Default: INFO CODEBUDDY_LOG_LEVEL=INFO + +# ----------------- +# Adapter V2 timeouts and concurrency (spec §12, §16) +# ----------------- +# Each upstream stage has its own timeout so a failure is reported precisely +# (e.g. upstream_first_chunk_timeout) instead of a generic connect timeout. + +# (Optional) TCP connection establishment timeout (seconds). Default: 30 +CODEBUDDY_CONNECT_TIMEOUT_SECONDS=30 + +# (Optional) Connection-pool acquisition timeout (seconds). Default: 30 +CODEBUDDY_POOL_TIMEOUT_SECONDS=30 + +# (Optional) Request-body write timeout (seconds). Default: 60 +CODEBUDDY_WRITE_TIMEOUT_SECONDS=60 + +# (Optional) Time to wait for upstream response headers (seconds). Default: 300 +CODEBUDDY_HEADERS_TIMEOUT_SECONDS=300 + +# (Optional) Time to wait for the first SSE chunk after headers (seconds). +# Default: 300 +CODEBUDDY_FIRST_CHUNK_TIMEOUT_SECONDS=300 + +# (Optional) Maximum idle gap between SSE chunks before aborting (seconds). +# Default: 600 +CODEBUDDY_STREAM_IDLE_TIMEOUT_SECONDS=600 + +# (Optional) Maximum concurrent in-flight upstream requests. Default: 20 +CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS=20 + +# (Optional) Max time to wait for a concurrency slot before returning +# upstream_queue_timeout (seconds). Default: 60 +CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS=60 + +# (Optional) Large-agentic-request warning thresholds. Exceeding any of these +# only logs large_agentic_request=true; context and tools are NEVER truncated. +CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH=50000 +CODEBUDDY_WARN_MESSAGE_COUNT=40 +CODEBUDDY_WARN_TOOL_COUNT=30 + # (Optional) Comma-separated model list reported to clients. # Adjust this list to match the models supported by your CodeBuddy account. CODEBUDDY_MODELS=claude-4.0,claude-3.7,gpt-5,gpt-5-mini,gpt-5-nano,o4-mini,gemini-2.5-flash,gemini-2.5-pro,auto-chat diff --git a/config.py b/config.py index d1cab05..bdb466d 100644 --- a/config.py +++ b/config.py @@ -44,7 +44,26 @@ "CODEBUDDY_MAX_SYSTEM_PROMPT_LENGTH": 2000, "CODEBUDDY_MODEL_ALIASES": "", "CODEBUDDY_DEFAULT_MODEL": "auto-chat", - "CODEBUDDY_UNKNOWN_MODEL_POLICY": "passthrough" + "CODEBUDDY_UNKNOWN_MODEL_POLICY": "passthrough", + # --- Adapter V2 --- + # Which request pipeline handles /v1/chat/completions: "v2" (the isolated + # CodeBuddyAdapterV2) or "legacy" (the original monolithic handler). + "CODEBUDDY_ADAPTER_VERSION": "v2", + # Granular upstream timeouts (seconds). Each stage is distinguished so a + # failure can be reported precisely instead of a generic "connect timeout". + "CODEBUDDY_CONNECT_TIMEOUT_SECONDS": 30, + "CODEBUDDY_POOL_TIMEOUT_SECONDS": 30, + "CODEBUDDY_WRITE_TIMEOUT_SECONDS": 60, + "CODEBUDDY_HEADERS_TIMEOUT_SECONDS": 300, + "CODEBUDDY_FIRST_CHUNK_TIMEOUT_SECONDS": 300, + "CODEBUDDY_STREAM_IDLE_TIMEOUT_SECONDS": 600, + # Upstream concurrency control. + "CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS": 20, + "CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS": 60, + # Large-agentic-request warning thresholds (never used to truncate). + "CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH": 50000, + "CODEBUDDY_WARN_MESSAGE_COUNT": 40, + "CODEBUDDY_WARN_TOOL_COUNT": 30, } # --- Core Functions --- @@ -210,6 +229,72 @@ def get_codebuddy_model_aliases() -> Dict[str, str]: return aliases +def get_codebuddy_adapter_version() -> str: + """Which chat-completions pipeline handles the request. + + * ``v2`` (default): the isolated :class:`CodeBuddyAdapterV2`. + * ``legacy``: the original monolithic handler in ``codebuddy_router``. + + An unrecognized value falls back to ``v2`` rather than raising, so a typo + never takes the service down; the effective value is logged by the router. + """ + version = str(_get_config_value("CODEBUDDY_ADAPTER_VERSION")).strip().lower() + return version if version in {"v2", "legacy"} else "v2" + + +def _get_positive_int(key: str, default: int) -> int: + """Return a strictly-positive int config value, else ``default``.""" + try: + value = int(_get_config_value(key)) + except (TypeError, ValueError): + return default + return value if value > 0 else default + + +def get_codebuddy_connect_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_CONNECT_TIMEOUT_SECONDS", 30)) + + +def get_codebuddy_pool_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_POOL_TIMEOUT_SECONDS", 30)) + + +def get_codebuddy_write_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_WRITE_TIMEOUT_SECONDS", 60)) + + +def get_codebuddy_headers_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_HEADERS_TIMEOUT_SECONDS", 300)) + + +def get_codebuddy_first_chunk_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_FIRST_CHUNK_TIMEOUT_SECONDS", 300)) + + +def get_codebuddy_stream_idle_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_STREAM_IDLE_TIMEOUT_SECONDS", 600)) + + +def get_codebuddy_max_concurrent_upstream_requests() -> int: + return _get_positive_int("CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS", 20) + + +def get_codebuddy_upstream_queue_timeout_seconds() -> float: + return float(_get_positive_int("CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS", 60)) + + +def get_codebuddy_warn_total_content_length() -> int: + return _get_positive_int("CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH", 50000) + + +def get_codebuddy_warn_message_count() -> int: + return _get_positive_int("CODEBUDDY_WARN_MESSAGE_COUNT", 40) + + +def get_codebuddy_warn_tool_count() -> int: + return _get_positive_int("CODEBUDDY_WARN_TOOL_COUNT", 30) + + def _coerce_bool(value: Any, default: bool) -> bool: if isinstance(value, bool): return value diff --git a/src/adapters/__init__.py b/src/adapters/__init__.py new file mode 100644 index 0000000..27fa7a2 --- /dev/null +++ b/src/adapters/__init__.py @@ -0,0 +1 @@ +"""Provider adapters for CodeBuddy2API.""" diff --git a/src/adapters/codebuddy/__init__.py b/src/adapters/codebuddy/__init__.py new file mode 100644 index 0000000..b771864 --- /dev/null +++ b/src/adapters/codebuddy/__init__.py @@ -0,0 +1,29 @@ +"""CodeBuddyAdapterV2 — an isolated, testable CodeBuddy provider adapter. + +This package converts OpenAI/Claude-compatible chat-completion requests into +CodeBuddy-compatible requests, streams the SSE response back, and maps it to the +OpenAI schema. It is stateless per request: the upstream API key is supplied by +the caller (9Router) on each request and is never stored or rotated internally. + +The router selects this adapter when ``CODEBUDDY_ADAPTER_VERSION=v2``. The +public entry point is :func:`get_adapter`, which returns a process-wide adapter +instance sharing a single pooled ``httpx.AsyncClient``. +""" +from __future__ import annotations + +from typing import Optional + +from .adapter import CodeBuddyAdapterV2 + +_adapter_instance: Optional[CodeBuddyAdapterV2] = None + + +def get_adapter() -> CodeBuddyAdapterV2: + """Return the process-wide :class:`CodeBuddyAdapterV2` singleton.""" + global _adapter_instance + if _adapter_instance is None: + _adapter_instance = CodeBuddyAdapterV2() + return _adapter_instance + + +__all__ = ["CodeBuddyAdapterV2", "get_adapter"] diff --git a/src/adapters/codebuddy/adapter.py b/src/adapters/codebuddy/adapter.py new file mode 100644 index 0000000..39ec5a2 --- /dev/null +++ b/src/adapters/codebuddy/adapter.py @@ -0,0 +1,702 @@ +"""CodeBuddyAdapterV2 orchestrator (spec §11, §13, §14, §15, §16). + +Owns the full request lifecycle for ``/v1/chat/completions`` when +``CODEBUDDY_ADAPTER_VERSION=v2``: + + 1. resolve the request-local upstream credential (passthrough key, or a legacy + credential for relay mode); + 2. normalize + validate messages, resolve the model, sanitize tool schemas, + build the strict upstream payload; + 3. acquire a concurrency slot (bounded, with a queue timeout); + 4. open the upstream stream with at most one pre-first-event retry; + 5. stream through incrementally, or aggregate for a non-streaming client; + 6. on downstream disconnect, cancel upstream and release the slot — no ghost + requests, no retry after the stream has started. + +Telemetry is emitted per stage with safe metadata only (spec §15). Nothing here +logs prompts, file contents, tool argument values, or raw keys. +""" +from __future__ import annotations + +import asyncio +import logging +import time +import uuid +from typing import Any, AsyncIterator, Dict, List, Optional + +import httpx +from fastapi import Request +from fastapi.responses import JSONResponse, StreamingResponse + +from src.auth import ClientAuthContext +from src.codebuddy_message_sanitizer import ( + is_codebuddy_moderation_response, + sanitize_messages, +) +from src.keyword_replacer import apply_keyword_replacement_to_system_message + +from .config import AdapterSettings, load_adapter_settings +from .errors import ( + AdapterError, + ModerationError, + QueueTimeoutError, + UnknownModelError, + UpstreamAuthError, + UpstreamNetworkError, + UpstreamRetryableError, +) +from .errors import ConnectTimeoutError +from .message_normalizer import ( + convert_anthropic_messages_to_openai, + normalize_messages_for_upstream, +) +from .models import TelemetryContext, UpstreamCredential +from .request_mapper import build_payload, dropped_fields, resolve_model +from . import response_mapper +from .response_mapper import sse_chunk, sse_done +from .stream_decoder import ( + SSELineBuffer, + StreamAggregator, + _extract_delta_content, + parse_sse_data_line, +) +from .tool_schema_adapter import sanitize_tools +from .transport import UpstreamTransport, key_fingerprint +from .validation import validate_conversation + +logger = logging.getLogger(__name__) + +# SSE response headers (spec §11): disable caching and proxy buffering so the +# stream is delivered incrementally. +SSE_HEADERS = { + "Cache-Control": "no-cache, no-transform", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + "Access-Control-Allow-Origin": "*", + "Access-Control-Allow-Methods": "GET, POST, OPTIONS", + "Access-Control-Allow-Headers": "*", +} + +# How many characters of assistant content to buffer while deciding whether a +# streaming reply is a moderation refusal (the refusal is short). +_MODERATION_BUFFER_CHARS = 200 + +_MODERATION_MESSAGE = ( + "CodeBuddy rejected the request through its content moderation system." +) + + +class CodeBuddyAdapterV2: + """Stateless-per-request CodeBuddy provider adapter.""" + + def __init__(self, transport: Optional[UpstreamTransport] = None) -> None: + self.transport = transport or UpstreamTransport() + self._semaphore: Optional[asyncio.Semaphore] = None + self._semaphore_limit: Optional[int] = None + + # --- Lifecycle ------------------------------------------------------- + + async def startup(self) -> None: + """Warm the shared HTTP client at application startup.""" + try: + settings = load_adapter_settings() + except Exception: + logger.exception("CodeBuddyAdapterV2 startup: settings unavailable") + return + await self.transport.get_client(settings) + + async def shutdown(self) -> None: + """Close the shared HTTP client at application shutdown.""" + await self.transport.aclose() + + def _get_semaphore(self, settings: AdapterSettings) -> asyncio.Semaphore: + """Return the concurrency semaphore, (re)building it if the limit changed.""" + limit = settings.max_concurrent_upstream_requests + if self._semaphore is None or self._semaphore_limit != limit: + self._semaphore = asyncio.Semaphore(limit) + self._semaphore_limit = limit + return self._semaphore + + # --- Telemetry ------------------------------------------------------- + + @staticmethod + def _log_stage(stage: str, telemetry: TelemetryContext, **extra: Any) -> None: + fields = telemetry.as_log_fields() + fields.update(extra) + logger.info( + "codebuddy_v2 stage=%s %s", + stage, + " ".join(f"{k}={v}" for k, v in fields.items()), + ) + + # --- Entry point ----------------------------------------------------- + + async def chat_completion( + self, + *, + request: Request, + request_body: Any, + auth_context: ClientAuthContext, + conversation_id: Optional[str] = None, + conversation_request_id: Optional[str] = None, + conversation_message_id: Optional[str] = None, + request_id: Optional[str] = None, + ): + """Handle one chat-completions request end to end.""" + telemetry = TelemetryContext(request_id=request_id or uuid.uuid4().hex) + self._log_stage("request_received", telemetry) + + try: + settings = load_adapter_settings() + except ValueError: + return self._error_json( + "Upstream configuration is invalid", + "configuration_error", + "configuration_invalid", + 500, + ) + + # Structural request validation (object with non-empty messages array). + if not isinstance(request_body, dict): + return self._error_json( + "Request body must be a JSON object", + "invalid_request_error", + "invalid_request", + 400, + ) + messages_in = request_body.get("messages") + if not messages_in or not isinstance(messages_in, list): + return self._error_json( + "Messages field is required and must be an array", + "invalid_request_error", + "invalid_request", + 400, + ) + + # Resolve the request-local upstream credential. + credential, cred_error = await self._resolve_credential(auth_context, settings) + if cred_error is not None: + return cred_error + telemetry.key_fingerprint = credential.fingerprint + + client_wants_stream = bool(request_body.get("stream", False)) + telemetry.client_wants_stream = client_wants_stream + + # Prepare the payload (model resolve → normalize → validate → tools). + try: + payload, prepared_telemetry = self._prepare_payload( + request_body, settings, telemetry + ) + except AdapterError as exc: + self._log_stage("request_failed", telemetry, code=exc.code) + return self._error_from_adapter(exc) + + self._log_stage("request_prepared", telemetry) + + # Acquire a concurrency slot (bounded), honoring the queue timeout. + semaphore = self._get_semaphore(settings) + self._log_stage("upstream_slot_wait_start", telemetry) + try: + await asyncio.wait_for( + semaphore.acquire(), timeout=settings.upstream_queue_timeout + ) + except asyncio.TimeoutError: + self._log_stage("request_failed", telemetry, code="upstream_queue_timeout") + return self._error_from_adapter(QueueTimeoutError()) + self._log_stage("upstream_slot_acquired", telemetry) + + slot_owned_by_caller = True + try: + headers = self.transport.build_headers( + bearer_token=credential.bearer_token, + settings=settings, + user_id=credential.user_id, + conversation_id=conversation_id, + conversation_request_id=conversation_request_id, + conversation_message_id=conversation_message_id, + request_id=request_id, + ) + + self._log_stage("upstream_send_start", telemetry) + response = await self._open_with_retry(settings, payload, headers, telemetry) + self._log_stage("upstream_headers_received", telemetry) + + if client_wants_stream: + # Hand slot + response ownership to the streaming generator, + # which releases the slot and closes the response in its finally. + slot_owned_by_caller = False + generator = self._stream_response( + request=request, + response=response, + settings=settings, + telemetry=telemetry, + semaphore=semaphore, + ) + return StreamingResponse( + generator, media_type="text/event-stream", headers=SSE_HEADERS + ) + + # Non-streaming: aggregate, then release the slot in finally. + result = await self._aggregate_response( + response=response, settings=settings, telemetry=telemetry + ) + if isinstance(result, JSONResponse): + return result + self._log_stage("stream_finished", telemetry, mode="aggregate") + return JSONResponse(content=result) + except ModerationError: + self._log_stage("request_failed", telemetry, code="content_filter") + return JSONResponse( + status_code=400, + headers={"X-CodeBuddy-Moderation": "true"}, + content={ + "error": { + "message": _MODERATION_MESSAGE, + "type": "content_filter", + "param": None, + "code": "codebuddy_content_filter", + } + }, + ) + except AdapterError as exc: + self._log_stage("request_failed", telemetry, code=exc.code) + return self._error_from_adapter(exc) + finally: + if slot_owned_by_caller: + semaphore.release() + + # --- Credential resolution ------------------------------------------ + + async def _resolve_credential( + self, auth_context: ClientAuthContext, settings: AdapterSettings + ): + """Resolve the upstream credential for this request. + + Passthrough mode uses the caller's Bearer token verbatim, only for this + request (spec §2) — no storage, no rotation, no failover. Relay/legacy + modes reuse the existing credential managers via a lazy import so this + adapter stays isolated while remaining backward compatible. + + Returns ``(credential, None)`` on success or ``(None, JSONResponse)`` on + failure. + """ + if auth_context.mode == "passthrough": + if not auth_context.passthrough_key: + return None, self._error_json( + "A Bearer API key is required", + "authentication_error", + "missing_api_key", + 401, + ) + token = auth_context.passthrough_key + return ( + UpstreamCredential( + bearer_token=token, + fingerprint=key_fingerprint(token), + source="passthrough", + ), + None, + ) + + # Relay / legacy credential resolution (backward compatibility). + try: + from src.codebuddy_router import CredentialManager + from src.codebuddy_api_key_manager import ApiKeyConfigurationError + + source, _max_attempts = await CredentialManager.resolve_source() + if source == "api_key_file": + resolved = await CredentialManager.get_api_key(set()) + else: + resolved = CredentialManager.get_legacy_credential() + except Exception: + logger.exception("CodeBuddyAdapterV2 credential resolution failed") + return None, self._error_json( + "Upstream API key is unavailable", + "configuration_error", + "api_key_file_unavailable", + 503, + ) + + if resolved is None or not resolved.bearer_token: + return None, self._error_json( + "No valid CodeBuddy credentials are available", + "authentication_error", + "credentials_unavailable", + 401, + ) + return ( + UpstreamCredential( + bearer_token=resolved.bearer_token, + fingerprint=key_fingerprint(resolved.bearer_token), + source=resolved.source, + user_id=resolved.user_id, + key_id=resolved.key_id, + ), + None, + ) + + # --- Payload preparation -------------------------------------------- + + def _prepare_payload( + self, + request_body: Dict[str, Any], + settings: AdapterSettings, + telemetry: TelemetryContext, + ): + """Resolve model, normalize + validate messages, and build the payload.""" + # Model resolution (may raise UnknownModelError → local 400). + resolved = resolve_model(request_body.get("model"), settings) + telemetry.requested_model = resolved.requested or "unknown" + telemetry.mapped_model = resolved.mapped + telemetry.mapping_source = resolved.source + + raw_messages = request_body.get("messages", []) or [] + telemetry.message_count = len(raw_messages) + + # 1) Convert Anthropic tool blocks → OpenAI shape. + messages = convert_anthropic_messages_to_openai(raw_messages) + + # 2) Sanitize agent system prompts (system messages only). + messages, _sanitized = sanitize_messages( + messages, + enabled=settings.sanitize_agent_prompt, + max_system_prompt_length=settings.max_system_prompt_length, + ) + + # 3) Keyword replacement on system messages only. + for msg in messages: + if isinstance(msg, dict) and msg.get("role") == "system": + msg["content"] = apply_keyword_replacement_to_system_message( + msg.get("content") + ) + + # 4) Guarantee role + content on every message. + messages = normalize_messages_for_upstream(messages) + self._log_stage("request_normalized", telemetry) + + # 5) Validate the fully-converted conversation (→ local 400). + validate_conversation(messages) + self._log_stage("request_validated", telemetry) + + # 6) Sanitize tool schemas; forward tools only when present. + tools = sanitize_tools(request_body.get("tools")) if request_body.get("tools") else None + + payload = build_payload(request_body, resolved.mapped, messages) + if tools: + payload["tools"] = tools + if request_body.get("tool_choice") is not None: + payload["tool_choice"] = request_body["tool_choice"] + + # Telemetry: tool + content-size metrics and large-request warning. + telemetry.tool_count = len(tools or []) + self._count_tool_messages(messages) + telemetry.total_content_length = self._total_content_length(messages) + self._maybe_warn_large_request(settings, telemetry, dropped_fields(request_body)) + + return payload, telemetry + + @staticmethod + def _count_tool_messages(messages: List[Dict[str, Any]]) -> int: + return sum( + 1 + for m in messages + if isinstance(m, dict) and (m.get("role") == "tool" or m.get("tool_calls")) + ) + + @staticmethod + def _total_content_length(messages: List[Dict[str, Any]]) -> int: + total = 0 + for m in messages: + if not isinstance(m, dict): + continue + content = m.get("content") + if isinstance(content, str): + total += len(content) + elif isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("type") == "text": + total += len(str(block.get("text", ""))) + return total + + def _maybe_warn_large_request( + self, + settings: AdapterSettings, + telemetry: TelemetryContext, + dropped: List[str], + ) -> None: + """Log a warning for large agentic requests without ever truncating.""" + if ( + telemetry.total_content_length > settings.warn_total_content_length + or telemetry.message_count > settings.warn_message_count + or telemetry.tool_count > settings.warn_tool_count + ): + telemetry.large_agentic_request = True + logger.warning( + "codebuddy_v2 large_agentic_request=true request_id=%s message_count=%d " + "tool_count=%d total_content_length=%d dropped_fields=[%s]", + telemetry.request_id, + telemetry.message_count, + telemetry.tool_count, + telemetry.total_content_length, + ",".join(dropped), + ) + + # --- Upstream open with single pre-first-event retry ---------------- + + async def _open_with_retry( + self, + settings: AdapterSettings, + payload: Dict[str, Any], + headers: Dict[str, str], + telemetry: TelemetryContext, + ) -> httpx.Response: + """Open the upstream stream, retrying at most once before any event. + + Only transient pre-event failures are retried (connect timeout, connect + reset/network error, 408/429/5xx). The same credential is reused — no key + rotation/fallback (9Router owns fallback, spec §14). Fatal statuses + (400/401/403/404/422) are never retried. + """ + attempts = 0 + while True: + try: + return await self.transport.open_stream( + settings=settings, payload=payload, headers=headers + ) + except (ConnectTimeoutError, UpstreamNetworkError, UpstreamRetryableError) as exc: + attempts += 1 + if attempts > 1: + raise + logger.warning( + "codebuddy_v2 retrying upstream once before first event: code=%s", + exc.code, + ) + continue + + # --- Streaming path -------------------------------------------------- + + async def _stream_response( + self, + *, + request: Request, + response: httpx.Response, + settings: AdapterSettings, + telemetry: TelemetryContext, + semaphore: asyncio.Semaphore, + ) -> AsyncIterator[str]: + """Yield OpenAI SSE frames, with moderation pre-buffering and cancellation. + + A downstream disconnect (client cancel) raises ``CancelledError``/ + ``GeneratorExit`` into this generator; the ``finally`` closes the upstream + response and releases the concurrency slot so no ghost request continues. + Once the first content/tool event is forwarded, the stream is never + retried. + """ + line_buffer = SSELineBuffer() + accumulated_content = "" + moderation_window_open = True + buffered_frames: List[str] = [] + first_event_seen = False + released = False + + def release() -> None: + nonlocal released + if not released: + released = True + semaphore.release() + + try: + async for chunk in self.transport.iter_lines_with_idle_timeout( + response, + first_chunk_timeout=settings.first_chunk_timeout, + idle_timeout=settings.stream_idle_timeout, + ): + if await request.is_disconnected(): + self._log_stage("downstream_disconnected", telemetry) + break + + for line in line_buffer.feed(chunk): + frame = self._process_stream_line(line) + if frame is None: + continue + + if not first_event_seen: + first_event_seen = True + self._log_stage("upstream_first_chunk", telemetry) + + # Moderation detection window: buffer initial frames until we + # can tell a normal reply from a Mandarin refusal. + if moderation_window_open: + obj = parse_sse_data_line(line) + if obj is not None: + accumulated_content += _extract_delta_content(obj) + buffered_frames.append(frame) + if is_codebuddy_moderation_response(accumulated_content): + for item in self._moderation_frames(telemetry): + yield item + self._log_stage("stream_finished", telemetry, moderation=True) + return + if ( + len(accumulated_content) >= _MODERATION_BUFFER_CHARS + or "[DONE]" in line + ): + moderation_window_open = False + for item in buffered_frames: + yield item + buffered_frames.clear() + continue + + yield frame + + # Flush any buffered frames that never crossed the window threshold + # (short, non-moderation responses) and any trailing partial line. + if buffered_frames: + for item in buffered_frames: + yield item + tail = line_buffer.flush() + if tail is not None: + frame = self._process_stream_line(tail) + if frame is not None: + yield frame + # Ensure the stream is terminated exactly once. + yield sse_done() + self._log_stage("stream_finished", telemetry, mode="stream") + except (asyncio.CancelledError, GeneratorExit): + self._log_stage("upstream_cancelled", telemetry) + raise + except AdapterError as exc: + # A timeout mid-stream: surface a safe SSE error frame. + logger.warning("codebuddy_v2 stream error code=%s", exc.code) + yield self._sse_error(exc.message, exc.code) + except httpx.RequestError: + logger.warning("codebuddy_v2 upstream stream interrupted") + yield self._sse_error("Upstream stream interrupted", "upstream_stream_error") + finally: + await response.aclose() + release() + + def _process_stream_line(self, line: str) -> Optional[str]: + """Convert one upstream SSE line into a downstream OpenAI SSE frame. + + Comments/blank lines are dropped. ``[DONE]`` is suppressed here so the + generator can emit exactly one terminator itself. Data frames have their + tool-call IDs normalized to the OpenAI ``call_`` convention for + consistency and are re-serialized. + """ + stripped = line.strip() + if not stripped or stripped.startswith(":"): + return None + if "[DONE]" in stripped: + return None + obj = parse_sse_data_line(line) + if obj is None: + return None + self._normalize_chunk_tool_ids(obj) + return sse_chunk(obj) + + @staticmethod + def _normalize_chunk_tool_ids(obj: Dict[str, Any]) -> None: + """Rewrite ``tooluse_`` tool-call IDs to ``call_`` in place (consistent).""" + try: + choices = obj.get("choices") or [] + if not choices: + return + delta = choices[0].get("delta") or {} + tool_calls = delta.get("tool_calls") + if not isinstance(tool_calls, list): + return + for tc in tool_calls: + if isinstance(tc, dict) and tc.get("id"): + tc["id"] = response_mapper._tool_call_id_openai(tc["id"]) + except (AttributeError, IndexError, TypeError): + return + + def _moderation_frames(self, telemetry: TelemetryContext) -> List[str]: + created = self._now() + chunk = response_mapper.moderation_chunk( + response_id="chatcmpl-codebuddy-filter", + created=created, + message=_MODERATION_MESSAGE, + ) + return [sse_chunk(chunk), sse_done()] + + @staticmethod + def _sse_error(message: str, code: str) -> str: + import json as _json + + return f'data: {_json.dumps({"error": {"message": message, "type": "stream_error", "code": code}}, ensure_ascii=False)}\n\n' + + # --- Non-streaming path --------------------------------------------- + + async def _aggregate_response( + self, + *, + response: httpx.Response, + settings: AdapterSettings, + telemetry: TelemetryContext, + ): + """Consume the whole SSE stream and build one chat.completion object. + + Detects a moderation refusal in the aggregated content or the raw body + and raises :class:`ModerationError` (mapped to a content_filter 400). + """ + aggregator = StreamAggregator() + line_buffer = SSELineBuffer() + raw_text = "" + first = True + try: + async for chunk in self.transport.iter_lines_with_idle_timeout( + response, + first_chunk_timeout=settings.first_chunk_timeout, + idle_timeout=settings.stream_idle_timeout, + ): + if first: + first = False + self._log_stage("upstream_first_chunk", telemetry) + raw_text += chunk + for line in line_buffer.feed(chunk): + obj = parse_sse_data_line(line) + if obj is not None: + aggregator.process_chunk(obj) + tail = line_buffer.flush() + if tail is not None: + obj = parse_sse_data_line(tail) + if obj is not None: + aggregator.process_chunk(obj) + content, tool_calls, finish_reason = aggregator.finalize() + finally: + await response.aclose() + + if is_codebuddy_moderation_response(content) or is_codebuddy_moderation_response( + raw_text + ): + raise ModerationError() + + return response_mapper.build_non_stream_response( + response_id=aggregator.id or f"chatcmpl-{uuid.uuid4().hex}", + model=aggregator.model or telemetry.mapped_model, + created=self._now(), + content=content, + tool_calls=tool_calls, + finish_reason=finish_reason, + usage=aggregator.usage, + system_fingerprint=aggregator.system_fingerprint, + reasoning_content=aggregator.reasoning_content, + ) + + # --- Error mapping --------------------------------------------------- + + @staticmethod + def _now() -> int: + return int(time.time()) + + @staticmethod + def _error_json( + message: str, error_type: str, code: str, status_code: int + ) -> JSONResponse: + return JSONResponse( + status_code=status_code, + content={"error": {"message": message, "type": error_type, "code": code}}, + ) + + def _error_from_adapter(self, exc: AdapterError) -> JSONResponse: + return self._error_json(exc.message, exc.error_type, exc.code, exc.status_code) diff --git a/src/adapters/codebuddy/config.py b/src/adapters/codebuddy/config.py new file mode 100644 index 0000000..7443b71 --- /dev/null +++ b/src/adapters/codebuddy/config.py @@ -0,0 +1,77 @@ +"""Settings snapshot for CodeBuddyAdapterV2. + +The adapter reads configuration through :func:`load_adapter_settings`, which +pulls from the root layered config system (in-memory → config.json → env → +defaults). Bundling the values into one immutable dataclass keeps request +handling free of scattered ``get_*`` calls and makes tests trivial: build an +``AdapterSettings`` directly instead of monkeypatching many getters. +""" +from __future__ import annotations + +from dataclasses import dataclass + +import config as root_config + + +@dataclass(frozen=True) +class AdapterSettings: + """Immutable per-request snapshot of adapter configuration.""" + + adapter_version: str + request_profile: str + upstream_api_key_header: str + sanitize_agent_prompt: bool + max_system_prompt_length: int + + default_model: str + unknown_model_policy: str + + connect_timeout: float + pool_timeout: float + write_timeout: float + headers_timeout: float + first_chunk_timeout: float + stream_idle_timeout: float + + max_concurrent_upstream_requests: int + upstream_queue_timeout: float + + warn_total_content_length: int + warn_message_count: int + warn_tool_count: int + + api_endpoint: str + + @property + def chat_completions_url(self) -> str: + return f"{self.api_endpoint}/v2/chat/completions" + + +def load_adapter_settings() -> AdapterSettings: + """Build an :class:`AdapterSettings` from the current root configuration. + + Config-validation errors (e.g. an invalid profile or header mode) propagate + to the caller, which surfaces them as a 500 misconfiguration error, exactly + like the legacy path. + """ + return AdapterSettings( + adapter_version=root_config.get_codebuddy_adapter_version(), + request_profile=root_config.get_codebuddy_request_profile(), + upstream_api_key_header=root_config.get_upstream_api_key_header(), + sanitize_agent_prompt=root_config.get_sanitize_agent_prompt(), + max_system_prompt_length=root_config.get_max_system_prompt_length(), + default_model=root_config.get_codebuddy_default_model(), + unknown_model_policy=root_config.get_codebuddy_unknown_model_policy(), + connect_timeout=root_config.get_codebuddy_connect_timeout_seconds(), + pool_timeout=root_config.get_codebuddy_pool_timeout_seconds(), + write_timeout=root_config.get_codebuddy_write_timeout_seconds(), + headers_timeout=root_config.get_codebuddy_headers_timeout_seconds(), + first_chunk_timeout=root_config.get_codebuddy_first_chunk_timeout_seconds(), + stream_idle_timeout=root_config.get_codebuddy_stream_idle_timeout_seconds(), + max_concurrent_upstream_requests=root_config.get_codebuddy_max_concurrent_upstream_requests(), + upstream_queue_timeout=root_config.get_codebuddy_upstream_queue_timeout_seconds(), + warn_total_content_length=root_config.get_codebuddy_warn_total_content_length(), + warn_message_count=root_config.get_codebuddy_warn_message_count(), + warn_tool_count=root_config.get_codebuddy_warn_tool_count(), + api_endpoint=root_config.get_codebuddy_api_endpoint(), + ) diff --git a/src/adapters/codebuddy/errors.py b/src/adapters/codebuddy/errors.py new file mode 100644 index 0000000..e0457bb --- /dev/null +++ b/src/adapters/codebuddy/errors.py @@ -0,0 +1,185 @@ +"""Typed error hierarchy for CodeBuddyAdapterV2. + +Every adapter failure maps to one of these so the router can translate it into +a precise OpenAI-compatible HTTP error. Crucially, the six upstream timeout +stages are distinct types (spec §12) so a failure is never collapsed into a +generic "fetch connect timeout". None of these carry raw upstream bodies, keys, +or message content — only safe codes and status numbers. +""" +from __future__ import annotations + +from typing import Optional + + +class AdapterError(Exception): + """Base class for every CodeBuddyAdapterV2 error. + + ``code`` is a stable machine-readable slug surfaced to the client; + ``status_code`` is the HTTP status the router should return; ``message`` is a + safe human-readable summary that never contains upstream bodies or secrets. + """ + + code: str = "adapter_error" + status_code: int = 502 + error_type: str = "upstream_error" + + def __init__(self, message: Optional[str] = None): + super().__init__(message or self.code) + self.message = message or self.code + + +# --- Request-shaping errors (local HTTP 400, never reach upstream) --- + + +class UnknownModelError(AdapterError): + """Requested model is unknown and CODEBUDDY_UNKNOWN_MODEL_POLICY=reject.""" + + code = "unknown_model" + status_code = 400 + error_type = "invalid_request_error" + + def __init__(self, requested_model: str): + super().__init__(f"Unknown model: {requested_model}") + self.requested_model = requested_model + + +class MessageNormalizationError(AdapterError): + """A message cannot be given a valid role/content before upstream. + + Carries the offending message ``index`` so the router returns a local 400 + identifying exactly which message is malformed. + """ + + code = "invalid_message" + status_code = 400 + error_type = "invalid_request_error" + + def __init__(self, index: int, reason: str): + super().__init__(f"Message {index} is malformed: {reason}") + self.index = index + self.reason = reason + + +class InvalidToolConversationError(AdapterError): + """The tool-call/tool-result structure is invalid (spec §7). + + Raised when validation of the fully-converted conversation fails: an orphan + tool result, non-JSON function arguments, a leftover Anthropic block, etc. + """ + + code = "invalid_tool_conversation" + status_code = 400 + error_type = "invalid_request_error" + + def __init__(self, index: int, reason: str): + super().__init__( + f"Invalid tool conversation structure at message {index}: {reason}" + ) + self.index = index + self.reason = reason + + +class InvalidToolArgumentsError(AdapterError): + """A streamed tool call finished with arguments that are not valid JSON. + + Per spec §9 the adapter must NOT silently coerce invalid arguments to ``{}``; + it raises this structured error instead. + """ + + code = "invalid_tool_arguments" + status_code = 502 + error_type = "upstream_error" + + def __init__(self, tool_name: Optional[str] = None): + super().__init__("Upstream tool call produced invalid JSON arguments") + self.tool_name = tool_name + + +# --- Moderation (not an upstream failure; must not trigger key failover) --- + + +class ModerationError(AdapterError): + """CodeBuddy rejected the request via its content moderation system.""" + + code = "codebuddy_content_filter" + status_code = 400 + error_type = "content_filter" + + +# --- Upstream transport errors --- + + +class UpstreamRejectedError(AdapterError): + """Upstream returned a non-200 status that is not retryable.""" + + def __init__(self, status_code: int, code: str = "upstream_request_rejected"): + super().__init__("Upstream CodeBuddy request was rejected") + self.status_code = status_code + self.code = code + + +class UpstreamRetryableError(AdapterError): + """Upstream returned a status eligible for a single pre-stream retry.""" + + error_type = "upstream_error" + + def __init__(self, status_code: int, code: str = "upstream_server_error"): + super().__init__("Upstream CodeBuddy request failed") + self.status_code = status_code + self.code = code + + +class UpstreamAuthError(AdapterError): + """Upstream rejected the supplied API key (401).""" + + code = "upstream_api_key_rejected" + status_code = 401 + error_type = "authentication_error" + + +# --- Distinct timeout stages (spec §12) --- + + +class UpstreamTimeoutError(AdapterError): + """Base for the distinct upstream timeout stages.""" + + error_type = "upstream_error" + status_code = 504 + + +class ConnectTimeoutError(UpstreamTimeoutError): + code = "upstream_connect_timeout" + + +class PoolTimeoutError(UpstreamTimeoutError): + code = "upstream_pool_timeout" + + +class WriteTimeoutError(UpstreamTimeoutError): + code = "upstream_write_timeout" + + +class HeadersTimeoutError(UpstreamTimeoutError): + code = "upstream_headers_timeout" + + +class FirstChunkTimeoutError(UpstreamTimeoutError): + code = "upstream_first_chunk_timeout" + + +class StreamIdleTimeoutError(UpstreamTimeoutError): + code = "upstream_stream_idle_timeout" + + +class QueueTimeoutError(UpstreamTimeoutError): + """Timed out waiting for a concurrency slot before contacting upstream.""" + + code = "upstream_queue_timeout" + status_code = 503 + + +class UpstreamNetworkError(UpstreamRetryableError): + """A connection reset or other network error before the first event.""" + + def __init__(self): + super().__init__(status_code=502, code="upstream_network_error") diff --git a/src/adapters/codebuddy/message_normalizer.py b/src/adapters/codebuddy/message_normalizer.py new file mode 100644 index 0000000..f967cab --- /dev/null +++ b/src/adapters/codebuddy/message_normalizer.py @@ -0,0 +1,278 @@ +"""Deterministic message normalizer for CodeBuddyAdapterV2 (spec §4, §5, §6). + +Claude Code and 9Router send messages in a mix of two shapes: + + * OpenAI shape — ``content`` is a string (or a multimodal array), assistant + tool calls live in ``tool_calls``, and tool results are ``role: "tool"`` + messages keyed by ``tool_call_id``. + * Anthropic shape — ``content`` is an array of typed blocks: ``text``, + ``image``, ``tool_use`` (an assistant issuing a call), and ``tool_result`` + (a user turn carrying the tool output). + +CodeBuddy speaks the OpenAI schema. This module rewrites Anthropic tool blocks +into OpenAI messages while preserving every relationship, and guarantees that +every final message has both ``role`` and ``content``. The cardinal rule +(spec §5) is that array content is NEVER replaced with an empty string — doing +so is what silently dropped ``tool_result`` file contents in the legacy path. + +The transformation is a pure function of its input: no I/O, no globals, fully +unit-testable. +""" +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, List, Tuple + +from .errors import MessageNormalizationError + +logger = logging.getLogger(__name__) + +# Anthropic block types that must be converted away before upstream. +_ANTHROPIC_TOOL_BLOCK_TYPES: Tuple[str, ...] = ("tool_use", "tool_result") + + +def _text_from_content(content: Any) -> str: + """Flatten any content shape into plain text (for length metrics / matching). + + Never used to REPLACE content — only to measure or to concatenate assistant + text blocks. Non-text blocks are serialized so their information is counted + rather than silently dropped. + """ + if isinstance(content, str): + return content + if content is None: + return "" + if isinstance(content, list): + parts: List[str] = [] + for item in content: + if isinstance(item, dict): + if item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + parts.append(json.dumps(item, ensure_ascii=False)) + elif isinstance(item, str): + parts.append(item) + else: + parts.append(str(item)) + return "".join(parts) + return str(content) + + +def _assistant_text_from_blocks(blocks: List[Any]) -> str: + """Concatenate only the ``text`` blocks of an assistant content array.""" + parts: List[str] = [] + for block in blocks: + if isinstance(block, dict) and block.get("type") == "text": + parts.append(str(block.get("text", ""))) + return "".join(parts) + + +def _flatten_tool_result_content(content: Any) -> str: + """Flatten an Anthropic ``tool_result.content`` into a complete string. + + The content may be a plain string or a list of blocks (``text``/``image``/ + ...). Every part is preserved so a tool result — e.g. the full text of a + file read — is never truncated or dropped. Non-text blocks are serialized to + JSON so their information survives. + """ + if isinstance(content, str): + return content + if content is None: + return "" + if isinstance(content, list): + parts: List[str] = [] + for item in content: + if isinstance(item, dict): + if item.get("type") == "text": + parts.append(str(item.get("text", ""))) + else: + parts.append(json.dumps(item, ensure_ascii=False)) + elif isinstance(item, str): + parts.append(item) + else: + parts.append(str(item)) + return "".join(parts) + return str(content) + + +def _has_block_type(content: Any, block_type: str) -> bool: + if not isinstance(content, list): + return False + return any(isinstance(b, dict) and b.get("type") == block_type for b in content) + + +def convert_anthropic_messages_to_openai( + messages: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Rewrite Anthropic block-array messages into OpenAI-shaped messages. + + Rules (spec §6): + + * An assistant turn whose content array contains ``tool_use`` blocks + becomes a ``role: assistant`` message with ``content`` = the text blocks + (or ``""``) and a ``tool_calls`` list. Each tool call keeps the + ``tool_use.id`` verbatim and encodes ``input`` as a JSON-string + ``function.arguments``. Multiple ``tool_use`` blocks → multiple entries. + * A user turn whose content array contains ``tool_result`` blocks is split + so each ``tool_result`` becomes its own ``role: tool`` message whose + ``tool_call_id`` is the verbatim ``tool_use_id`` and whose ``content`` is + the complete flattened tool output. Multiple results → multiple + messages, each standalone — never merged into unrelated user text. + * When a user content array mixes ``text``/``image`` with ``tool_result``, + the order is preserved by emitting each contiguous run of non-tool blocks + as its own user message positioned exactly where it appeared. + * Messages whose content is a plain string, or an array with no tool + blocks (plain text or multimodal image arrays), are passed through + unchanged — array content is never replaced with an empty string. + + The verbatim ID reuse guarantees ``tool_use.id`` == emitted + ``tool_call.id`` == matching tool message's ``tool_call_id``. Each block is + emitted exactly once, so no tool call or result is duplicated. + """ + converted: List[Dict[str, Any]] = [] + + for msg in messages: + if not isinstance(msg, dict): + converted.append(msg) + continue + + role = msg.get("role") + content = msg.get("content") + + # Non-array content (string / None) passes through untouched. + if not isinstance(content, list): + converted.append(msg) + continue + + has_tool_use = _has_block_type(content, "tool_use") + has_tool_result = _has_block_type(content, "tool_result") + + # Plain text or multimodal (image) array with no tool blocks: preserve. + if not has_tool_use and not has_tool_result: + converted.append(msg) + continue + + if has_tool_use: + # Assistant turn issuing one or more tool calls. Preserve any + # accompanying assistant text as the message content. + tool_calls: List[Dict[str, Any]] = [] + for block in content: + if not isinstance(block, dict) or block.get("type") != "tool_use": + continue + tool_calls.append( + { + "id": block.get("id", ""), + "type": "function", + "function": { + "name": block.get("name", ""), + # OpenAI requires arguments to be a JSON *string*. + "arguments": json.dumps( + block.get("input", {}) or {}, ensure_ascii=False + ), + }, + } + ) + new_msg: Dict[str, Any] = { + "role": role or "assistant", + "content": _assistant_text_from_blocks(content), + "tool_calls": tool_calls, + } + converted.append(new_msg) + continue + + # has_tool_result: walk the array in order. Each ``tool_result`` becomes + # its own tool message; contiguous non-tool blocks (text/image) are + # flushed as a separate user message in their original position so + # ordering and multimodal content are preserved and nothing is merged + # into an unrelated turn. + pending_blocks: List[Any] = [] + + def _flush_pending() -> None: + if pending_blocks: + converted.append({"role": role or "user", "content": list(pending_blocks)}) + pending_blocks.clear() + + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + _flush_pending() + converted.append( + { + "role": "tool", + "tool_call_id": block.get("tool_use_id", ""), + "content": _flatten_tool_result_content(block.get("content")), + } + ) + else: + pending_blocks.append(block) + _flush_pending() + + return converted + + +def normalize_messages_for_upstream( + messages: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Guarantee every message has a valid ``role`` and a ``content`` field. + + Run as the final shaping step before validation. Rules (spec §4, §5): + + * A message with an explicit non-empty ``role`` keeps it. A role that is + present but nonstandard is accepted (never rejected on a technicality). + * A missing role is inferred: ``tool_call_id`` → ``tool``; ``tool_calls`` + → ``assistant``. If still undeterminable, raise + :class:`MessageNormalizationError` with the index (local HTTP 400). + * ``content`` is added as ``""`` ONLY when missing or ``null``. Existing + content — including empty strings and multimodal arrays — is preserved + verbatim. Assistant tool-call turns therefore get ``content: ""`` only + when they truly lack content. + """ + normalized: List[Dict[str, Any]] = [] + repaired = 0 + + for index, msg in enumerate(messages): + if not isinstance(msg, dict): + raise MessageNormalizationError(index, "message is not an object") + + role = msg.get("role") + if not (isinstance(role, str) and role.strip()): + if msg.get("tool_call_id"): + role = "tool" + elif msg.get("tool_calls"): + role = "assistant" + else: + raise MessageNormalizationError( + index, + "role is missing and cannot be inferred from message structure", + ) + else: + role = role.strip() + + new_msg = dict(msg) + new_msg["role"] = role + + if "content" not in new_msg or new_msg["content"] is None: + new_msg["content"] = "" + repaired += 1 + + normalized.append(new_msg) + + if repaired: + logger.info( + "Normalized upstream messages: total=%d content_defaulted=%d", + len(messages), + repaired, + ) + + return normalized + + +def normalize_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Full normalization pipeline: Anthropic conversion then role/content fill. + + Returns OpenAI-shaped messages ready for validation and the upstream + payload. Sanitization of system prompts is applied separately by the + orchestrator so this stays a pure structural transform. + """ + converted = convert_anthropic_messages_to_openai(messages) + return normalize_messages_for_upstream(converted) diff --git a/src/adapters/codebuddy/models.py b/src/adapters/codebuddy/models.py new file mode 100644 index 0000000..d4662a8 --- /dev/null +++ b/src/adapters/codebuddy/models.py @@ -0,0 +1,97 @@ +"""Dataclasses shared across CodeBuddyAdapterV2. + +These are plain data holders — no I/O, no upstream calls — so they are trivially +unit-testable and safe to import from anywhere in the package. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +@dataclass +class ResolvedModel: + """Outcome of resolving a client model label to an upstream model ID. + + ``source`` records how it was resolved (exact / alias / default_empty / + passthrough / default) for safe diagnostics. + """ + + requested: str + mapped: str + source: str + + +@dataclass(repr=False) +class UpstreamCredential: + """A request-local upstream credential. + + The raw ``bearer_token`` is never logged or serialized; ``repr`` is disabled + so it cannot leak through incidental logging. ``fingerprint`` is a short, + non-reversible SHA-256 prefix suitable for correlating logs. + """ + + bearer_token: str = field(repr=False) + fingerprint: str = "none" + source: str = "passthrough" + user_id: Optional[str] = None + key_id: Optional[str] = None + + +@dataclass +class ToolCallState: + """Accumulates one streaming tool call across many SSE chunks (spec §9). + + Upstream may split ``function.arguments`` across arbitrary chunk boundaries + (e.g. ``{"file_`` then ``path":"a.txt"}``). Fragments are buffered verbatim + and only joined + parsed once the tool call is complete, so a partial JSON + fragment is never parsed or emitted as tool input. + """ + + index: int + id: Optional[str] = None + name: Optional[str] = None + argument_fragments: List[str] = field(default_factory=list) + emitted: bool = False + + def append_arguments(self, fragment: str) -> None: + if fragment: + self.argument_fragments.append(fragment) + + def raw_arguments(self) -> str: + """The concatenated argument fragments (may be incomplete mid-stream).""" + return "".join(self.argument_fragments) + + +@dataclass +class TelemetryContext: + """Per-request safe telemetry accumulator (spec §15). + + Holds only structural metadata — never prompts, file contents, tool + argument values, or credentials. ``request_id`` correlates all stage logs. + """ + + request_id: str + requested_model: str = "unknown" + mapped_model: str = "unknown" + mapping_source: str = "unknown" + message_count: int = 0 + tool_count: int = 0 + total_content_length: int = 0 + key_fingerprint: str = "none" + client_wants_stream: bool = False + large_agentic_request: bool = False + + def as_log_fields(self) -> Dict[str, Any]: + return { + "request_id": self.request_id, + "requested_model": self.requested_model, + "mapped_model": self.mapped_model, + "mapping_source": self.mapping_source, + "message_count": self.message_count, + "tool_count": self.tool_count, + "total_content_length": self.total_content_length, + "key_fingerprint": self.key_fingerprint, + "stream": self.client_wants_stream, + "large_agentic_request": self.large_agentic_request, + } diff --git a/src/adapters/codebuddy/request_mapper.py b/src/adapters/codebuddy/request_mapper.py new file mode 100644 index 0000000..6c5e02c --- /dev/null +++ b/src/adapters/codebuddy/request_mapper.py @@ -0,0 +1,124 @@ +"""Model resolution and upstream payload construction for CodeBuddyAdapterV2. + +Two responsibilities (spec §3): + + * :func:`resolve_model` — map a client/UI model label to an upstream CodeBuddy + model ID WITHOUT ever silently rewriting an unknown label to ``auto-chat``. + An unknown label is handled per ``CODEBUDDY_UNKNOWN_MODEL_POLICY``. + * :func:`build_payload` — assemble the strict, allowlisted upstream payload. + Only fields CodeBuddy understands are forwarded; OpenAI-only and unknown UI + fields are dropped so they cannot trigger an upstream rejection. + +Both are pure functions of their inputs plus the injected settings/model list, +so they are directly unit-testable without touching global config. +""" +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional, Tuple + +import config as root_config + +from .config import AdapterSettings +from .errors import UnknownModelError +from .models import ResolvedModel + +logger = logging.getLogger(__name__) + +# Top-level request fields forwarded to CodeBuddy. Everything else (OpenAI-only +# fields like response_format / reasoning_effort / stream_options, and unknown +# UI fields) is dropped. ``tools`` / ``tool_choice`` are forwarded only when the +# client actually sent them. +_UPSTREAM_ALLOWLIST = {"model", "messages", "stream", "tools", "tool_choice"} + + +def resolve_model( + requested_model: Any, + settings: AdapterSettings, + available_models: Optional[List[str]] = None, + aliases: Optional[Dict[str, str]] = None, +) -> ResolvedModel: + """Resolve a model label to an upstream model ID and record how. + + ``source`` is one of: + * ``exact`` — matched a configured available model ID + * ``alias`` — matched a configured alias (case-insensitive) + * ``default_empty`` — no model supplied; used the default + * ``passthrough`` — unknown label forwarded verbatim + * ``default`` — unknown label mapped to the default model + + An unknown label follows ``settings.unknown_model_policy``: + * ``passthrough`` (default): forward it verbatim. + * ``reject``: raise :class:`UnknownModelError` (local HTTP 400). + * ``default``: fall back to ``settings.default_model``. + """ + if available_models is None: + try: + available_models = root_config.get_available_models() + except Exception: + available_models = [] + if aliases is None: + try: + aliases = root_config.get_codebuddy_model_aliases() + except Exception: + aliases = {} + + available = set(available_models) + default_model = settings.default_model or "auto-chat" + + # No usable model supplied: nothing to pass through, so use the default + # regardless of policy. + if not isinstance(requested_model, str) or not requested_model.strip(): + return ResolvedModel(requested="", mapped=default_model, source="default_empty") + requested = requested_model.strip() + + # Exact match against a known upstream model ID (covers an explicit + # "auto-chat" request, which is therefore "exact", not "default"). + if requested in available: + return ResolvedModel(requested=requested, mapped=requested, source="exact") + + # Configured alias mapping (case-insensitive). + mapped = aliases.get(requested.lower()) + if mapped: + return ResolvedModel(requested=requested, mapped=mapped, source="alias") + + # Unknown label: governed by policy. Never silently rewrite to the default. + policy = settings.unknown_model_policy + if policy == "reject": + logger.info("Unknown requested model rejected (policy=reject)") + raise UnknownModelError(requested) + if policy == "default": + logger.info("Unknown requested model mapped to default (policy=default)") + return ResolvedModel(requested=requested, mapped=default_model, source="default") + # passthrough (default): forward verbatim so the real upstream model is + # preserved and CodeBuddy decides whether it is valid. + return ResolvedModel(requested=requested, mapped=requested, source="passthrough") + + +def dropped_fields(request_body: Dict[str, Any]) -> List[str]: + """Return the sorted top-level field names that will NOT be forwarded.""" + return sorted(k for k in request_body.keys() if k not in _UPSTREAM_ALLOWLIST) + + +def build_payload( + request_body: Dict[str, Any], + mapped_model: str, + messages: List[Dict[str, Any]], +) -> Dict[str, Any]: + """Assemble the strict upstream payload. + + CodeBuddy is SSE-only, so ``stream`` is always ``True`` upstream regardless + of the client's preference; the adapter decides separately whether to + aggregate (non-stream client) or pass the stream through. ``tools`` and + ``tool_choice`` are forwarded only when present. + """ + payload: Dict[str, Any] = { + "model": mapped_model, + "messages": messages, + "stream": True, + } + if request_body.get("tools"): + payload["tools"] = request_body["tools"] + if request_body.get("tool_choice") is not None: + payload["tool_choice"] = request_body["tool_choice"] + return payload diff --git a/src/adapters/codebuddy/response_mapper.py b/src/adapters/codebuddy/response_mapper.py new file mode 100644 index 0000000..a68f516 --- /dev/null +++ b/src/adapters/codebuddy/response_mapper.py @@ -0,0 +1,143 @@ +"""OpenAI response shaping for CodeBuddyAdapterV2 (spec §10). + +Two output shapes: + + * non-streaming — a single ``chat.completion`` object with a ``message`` that + may carry ``content`` and/or ``tool_calls``. Never emit a + ``chat.completion.chunk`` here. + * streaming — ``chat.completion.chunk`` frames, ending with a single + ``data: [DONE]``. + +``reasoning_content`` is surfaced as a distinct field, never merged into the +final answer content (spec §10). + +Every builder is pure: it takes decoded values and returns dict/str output, so +these are trivially unit-testable without any upstream call. +""" +from __future__ import annotations + +import json +from typing import Any, Dict, List, Optional + + +def _tool_call_id_openai(codebuddy_id: str) -> str: + """Normalize a CodeBuddy tool-call id to the OpenAI ``call_`` convention. + + CodeBuddy sometimes emits ``tooluse_``; OpenAI clients expect ``call_``. + A verbatim id (already ``call_`` / ``toolu_``) is returned unchanged so the + tool_call/tool_result pairing is preserved. + """ + if isinstance(codebuddy_id, str) and codebuddy_id.startswith("tooluse_"): + return f"call_{codebuddy_id[len('tooluse_'):]}" + return codebuddy_id + + +def build_non_stream_response( + *, + response_id: str, + model: str, + created: int, + content: str, + tool_calls: List[Dict[str, Any]], + finish_reason: str, + usage: Optional[Dict[str, Any]] = None, + system_fingerprint: Optional[str] = None, + reasoning_content: str = "", +) -> Dict[str, Any]: + """Build a complete OpenAI ``chat.completion`` object. + + ``model`` is the model CodeBuddy reported (or the mapped model when upstream + omitted it); the adapter never rewrites it to hide the upstream model. + """ + message: Dict[str, Any] = {"role": "assistant", "content": content or ""} + if reasoning_content: + # Surface reasoning separately; do NOT fold it into content. + message["reasoning_content"] = reasoning_content + if tool_calls: + message["tool_calls"] = [ + { + "id": _tool_call_id_openai(tc.get("id", "")), + "type": "function", + "function": { + "name": tc.get("function", {}).get("name", ""), + "arguments": tc.get("function", {}).get("arguments", "{}"), + }, + } + for tc in tool_calls + ] + + response: Dict[str, Any] = { + "id": response_id, + "object": "chat.completion", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + "logprobs": None, + } + ], + } + if usage: + response["usage"] = usage + if system_fingerprint: + response["system_fingerprint"] = system_fingerprint + return response + + +def sse_chunk(obj: Dict[str, Any]) -> str: + """Serialize a chunk object as an SSE ``data:`` frame.""" + return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n" + + +def sse_done() -> str: + """The single terminating SSE frame.""" + return "data: [DONE]\n\n" + + +def content_chunk( + *, response_id: str, model: str, created: int, content: str +) -> Dict[str, Any]: + """Build a streaming ``chat.completion.chunk`` carrying a content delta.""" + return { + "id": response_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + {"index": 0, "delta": {"content": content}, "finish_reason": None} + ], + } + + +def final_chunk( + *, response_id: str, model: str, created: int, finish_reason: str +) -> Dict[str, Any]: + """Build the terminal streaming chunk carrying only ``finish_reason``.""" + return { + "id": response_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], + } + + +def moderation_chunk( + *, response_id: str, created: int, message: str +) -> Dict[str, Any]: + """Build a content_filter streaming chunk for a moderation refusal.""" + return { + "id": response_id, + "object": "chat.completion.chunk", + "created": created, + "choices": [ + { + "index": 0, + "delta": {"content": message}, + "finish_reason": "content_filter", + } + ], + } diff --git a/src/adapters/codebuddy/stream_decoder.py b/src/adapters/codebuddy/stream_decoder.py new file mode 100644 index 0000000..11451d7 --- /dev/null +++ b/src/adapters/codebuddy/stream_decoder.py @@ -0,0 +1,138 @@ +"""SSE decoding for CodeBuddyAdapterV2 (spec §9, §10). + +CodeBuddy responds only with an OpenAI-style SSE stream. :class:`StreamDecoder` +turns a raw byte/text stream into discrete parsed chunk objects and offers two +consumption modes: + + * :meth:`iter_events` — line-oriented parsing that yields each decoded SSE + ``data:`` object (skipping comments, blanks, and ``[DONE]``). The caller + handles pass-through vs aggregation. + * :meth:`aggregate` — consume the whole stream into a single non-streaming + ``chat.completion`` message, reconstructing split tool-call arguments via + :class:`ToolCallAccumulator`. + +Parsing is incremental and buffered so a chunk boundary in the middle of a line +never corrupts a frame. +""" +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple + +from .tool_call_state import ToolCallAccumulator + +logger = logging.getLogger(__name__) + + +def parse_sse_data_line(line: str) -> Optional[Dict[str, Any]]: + """Parse a single ``data:`` SSE line into an object, or ``None``. + + Returns ``None`` for comments (``:`` prefix), blank lines, non-data lines, + the ``[DONE]`` sentinel, and any line whose payload is not valid JSON. + """ + stripped = line.strip() + if not stripped or stripped.startswith(":"): + return None + if not stripped.startswith("data:"): + return None + data = stripped[len("data:"):].strip() + if not data or data == "[DONE]": + return None + try: + return json.loads(data) + except (ValueError, TypeError): + return None + + +class SSELineBuffer: + """Incrementally split a chunked text stream into complete lines.""" + + def __init__(self) -> None: + self._buffer = "" + + def feed(self, chunk: str) -> List[str]: + """Add a chunk and return any complete lines it produced.""" + if not chunk: + return [] + self._buffer += chunk + lines: List[str] = [] + while "\n" in self._buffer: + line, self._buffer = self._buffer.split("\n", 1) + lines.append(line) + return lines + + def flush(self) -> Optional[str]: + """Return any trailing partial line and clear the buffer.""" + if self._buffer.strip(): + line = self._buffer + self._buffer = "" + return line + self._buffer = "" + return None + + +def _extract_delta_content(obj: Dict[str, Any]) -> str: + """Return the assistant ``delta.content`` string from a chunk object.""" + try: + choices = obj.get("choices") or [] + if not choices: + return "" + delta = choices[0].get("delta") or {} + content = delta.get("content") + return content if isinstance(content, str) else "" + except (AttributeError, IndexError, TypeError): + return "" + + +class StreamAggregator: + """Aggregate decoded SSE chunk objects into one ``chat.completion`` message. + + Content deltas are concatenated; tool-call fragments are routed to a + :class:`ToolCallAccumulator` so split ``function.arguments`` are only parsed + once complete. ``finalize`` returns the OpenAI non-streaming message body. + """ + + def __init__(self) -> None: + self.id: Optional[str] = None + self.model: Optional[str] = None + self.system_fingerprint: Optional[str] = None + self.content = "" + self.finish_reason: Optional[str] = None + self.usage: Optional[Dict[str, Any]] = None + self.reasoning_content = "" + self._tools = ToolCallAccumulator() + + def process_chunk(self, obj: Dict[str, Any]) -> None: + if not isinstance(obj, dict): + return + self.id = self.id or obj.get("id") + self.model = self.model or obj.get("model") + self.system_fingerprint = obj.get("system_fingerprint") or self.system_fingerprint + if obj.get("usage"): + self.usage = obj.get("usage") + + choices = obj.get("choices") or [] + if not choices: + return + choice = choices[0] + if choice.get("finish_reason"): + self.finish_reason = choice.get("finish_reason") + + delta = choice.get("delta") or {} + if isinstance(delta.get("content"), str): + self.content += delta["content"] + if isinstance(delta.get("reasoning_content"), str): + self.reasoning_content += delta["reasoning_content"] + if delta.get("tool_calls"): + self._tools.process_delta_tool_calls(delta["tool_calls"]) + + def finalize(self) -> Tuple[str, List[Dict[str, Any]], Optional[str]]: + """Return ``(content, tool_calls, finish_reason)``. + + May raise :class:`InvalidToolArgumentsError` from the accumulator when a + completed tool call has malformed JSON arguments. + """ + tool_calls = self._tools.finalize() if self._tools.has_pending() else [] + finish_reason = "tool_calls" if tool_calls else (self.finish_reason or "stop") + return self.content, tool_calls, finish_reason diff --git a/src/adapters/codebuddy/tool_call_state.py b/src/adapters/codebuddy/tool_call_state.py new file mode 100644 index 0000000..5108c7c --- /dev/null +++ b/src/adapters/codebuddy/tool_call_state.py @@ -0,0 +1,169 @@ +"""Streaming tool-call reconstruction state machine (spec §9). + +Upstream delivers a tool call across many SSE chunks, and ``function.arguments`` +may be split at arbitrary byte boundaries: + + chunk 1: {"file_ + chunk 2: path":"port + chunk 3: folio.html"} + +Parsing any single fragment as JSON would fail or, worse, silently produce ``{}`` +— the exact bug the legacy ``validate_and_fix_tool_call_args`` had. This state +machine instead: + + * groups fragments by tool-call ``index`` AND ``id`` (upstream reuses index 0 + for every call, so the ID is the real key once known); + * appends every ``arguments`` fragment verbatim without parsing; + * parses the joined JSON exactly once, when the call is finalized; + * raises :class:`InvalidToolArgumentsError` if the *complete* JSON is invalid, + rather than coercing it to ``{}``; + * emits each tool call exactly once and never mutates its ID. + +The machine is transport-agnostic: :class:`StreamDecoder` feeds it delta +fragments and asks it to finalize when ``finish_reason == "tool_calls"`` or the +stream ends. +""" +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, List, Optional + +from .errors import InvalidToolArgumentsError +from .models import ToolCallState + +logger = logging.getLogger(__name__) + + +class ToolCallAccumulator: + """Accumulates streaming tool-call fragments into complete tool calls.""" + + def __init__(self) -> None: + # Keyed by a stable slot key so index-0-repeated calls stay separate + # once their IDs are known. + self._states: Dict[str, ToolCallState] = {} + # Preserve the order tool calls were first seen. + self._order: List[str] = [] + # The slot key of the tool call currently receiving fragments; used for + # incremental deltas that arrive without an id/index. + self._current_key: Optional[str] = None + + @staticmethod + def _slot_key(index: Optional[int], tool_id: Optional[str]) -> str: + """Build a stable key. Prefer the ID; fall back to the index.""" + if tool_id: + return f"id:{tool_id}" + if index is not None: + return f"idx:{index}" + return "idx:0" + + def has_pending(self) -> bool: + return bool(self._states) + + def process_delta_tool_calls(self, tool_calls: List[Dict[str, Any]]) -> None: + """Consume the ``delta.tool_calls`` array from one SSE chunk. + + Each entry may carry an ``id``/``index``/``function.name`` (a new call or + the first fragment) and/or a ``function.arguments`` fragment (an + increment). Nothing is parsed here — fragments are only buffered. + """ + for tc in tool_calls: + if not isinstance(tc, dict): + continue + index = tc.get("index") + tool_id = tc.get("id") + func = tc.get("function") or {} + name = func.get("name") if isinstance(func, dict) else None + arguments = func.get("arguments") if isinstance(func, dict) else None + + key = self._resolve_key(index, tool_id) + state = self._states.get(key) + if state is None: + state = ToolCallState(index=index if index is not None else 0) + self._states[key] = state + self._order.append(key) + logger.debug("New streaming tool call slot=%s", key) + + # Fill in identity fields as they arrive. + if tool_id and not state.id: + state.id = tool_id + if name: + state.name = name + if arguments: + state.append_arguments(arguments) + + self._current_key = key + + def _resolve_key(self, index: Optional[int], tool_id: Optional[str]) -> str: + """Map an incoming delta entry to an existing slot or a new one. + + When an id is present we key by id, merging any earlier index-only slot + that was opened before the id was known. When neither id nor index is + present, the fragment belongs to the call currently in progress. + """ + if tool_id: + id_key = f"id:{tool_id}" + if id_key in self._states: + return id_key + # Promote a prior index-only slot to this id, if one exists and has + # not yet been assigned an id. + if index is not None: + idx_key = f"idx:{index}" + existing = self._states.get(idx_key) + if existing is not None and not existing.id: + self._states[id_key] = existing + self._states.pop(idx_key, None) + self._order[:] = [ + id_key if k == idx_key else k for k in self._order + ] + return id_key + return id_key + if index is not None: + return f"idx:{index}" + if self._current_key is not None: + return self._current_key + return "idx:0" + + def finalize(self) -> List[Dict[str, Any]]: + """Join fragments, parse arguments once, and return complete tool calls. + + Raises :class:`InvalidToolArgumentsError` when a completed tool call's + joined arguments are not valid JSON. Empty arguments are normalized to + ``"{}"`` (a call with no arguments is valid); a NON-empty but malformed + argument string is an error, never silently replaced. + """ + result: List[Dict[str, Any]] = [] + for key in self._order: + state = self._states.get(key) + if state is None or state.emitted: + continue + + raw = state.raw_arguments().strip() + if not raw: + arguments = "{}" + else: + try: + parsed = json.loads(raw) + except (ValueError, TypeError) as exc: + logger.warning( + "Streaming tool call produced invalid JSON arguments " + "(slot=%s, name=%s)", + key, + state.name, + ) + raise InvalidToolArgumentsError(state.name) from exc + # Re-serialize compactly to guarantee a canonical JSON string. + arguments = json.dumps(parsed, ensure_ascii=False) + + state.emitted = True + result.append( + { + "id": state.id or "", + "type": "function", + "function": { + "name": state.name or "", + "arguments": arguments, + }, + } + ) + return result diff --git a/src/adapters/codebuddy/tool_schema_adapter.py b/src/adapters/codebuddy/tool_schema_adapter.py new file mode 100644 index 0000000..4875ced --- /dev/null +++ b/src/adapters/codebuddy/tool_schema_adapter.py @@ -0,0 +1,190 @@ +"""Tool JSON-Schema sanitization for CodeBuddyAdapterV2 (spec §8). + +Claude Code sends tool definitions whose ``function.parameters`` are JSON Schema +documents that may use ``$ref``/``$defs``/``definitions`` and a range of +keywords. This module produces a deterministic, upstream-safe copy of each tool +schema: + + * local ``$ref`` (``#/$defs/...`` / ``#/definitions/...``) are resolved inline + so the upstream never has to follow references it may not support; + * supported structural keywords are preserved exactly (``type``, + ``properties``, ``required``, ``items``, ``enum``, ``additionalProperties``, + ``anyOf``/``oneOf``/``allOf``); + * tool names and parameter names are preserved verbatim — nothing is renamed, + no required field is dropped, no nested object is flattened; + * unknown/unsupported keywords are dropped from the copy rather than forwarded + raw. + +Everything here is pure: it takes a tools list and returns a new one, so it is +fully unit-testable and never mutates the caller's data. +""" +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + +# JSON Schema keywords preserved verbatim (structure-bearing). +_SUPPORTED_SCHEMA_KEYS = { + "type", + "properties", + "required", + "items", + "enum", + "const", + "additionalProperties", + "anyOf", + "oneOf", + "allOf", + "not", + "description", + "title", + "default", + "format", + "minimum", + "maximum", + "exclusiveMinimum", + "exclusiveMaximum", + "minLength", + "maxLength", + "pattern", + "minItems", + "maxItems", + "uniqueItems", + "minProperties", + "maxProperties", + "nullable", +} + +# Maximum depth guard so a pathological or cyclic schema cannot recurse forever. +_MAX_DEPTH = 64 + + +def _resolve_ref(ref: str, defs: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Resolve a local ``#/$defs/Name`` or ``#/definitions/Name`` reference. + + Returns the referenced schema dict, or ``None`` for non-local / unknown refs + (which are then simply dropped, leaving a permissive empty schema). + """ + if not isinstance(ref, str) or not ref.startswith("#/"): + return None + parts = ref[2:].split("/") + if len(parts) != 2 or parts[0] not in ("$defs", "definitions"): + return None + target = defs.get(parts[1]) + return target if isinstance(target, dict) else None + + +def _sanitize_schema( + schema: Any, + defs: Dict[str, Any], + depth: int = 0, + seen_refs: Optional[frozenset] = None, +) -> Any: + """Return an upstream-safe deep copy of a JSON Schema node. + + ``defs`` is the pooled ``$defs``/``definitions`` map from the root schema so + references can be resolved at any depth. ``seen_refs`` breaks reference + cycles: a ``$ref`` already being expanded on the current path resolves to a + permissive ``{}`` instead of recursing forever. + """ + if seen_refs is None: + seen_refs = frozenset() + + if depth > _MAX_DEPTH: + return {} + + if isinstance(schema, list): + return [_sanitize_schema(item, defs, depth + 1, seen_refs) for item in schema] + + if not isinstance(schema, dict): + # Primitive (already-resolved default/enum value etc.) — copy verbatim. + return schema + + # Resolve a local $ref by inlining the referenced schema. + ref = schema.get("$ref") + if isinstance(ref, str): + if ref in seen_refs: + # Cycle: stop expanding and emit a permissive schema. + return {} + resolved = _resolve_ref(ref, defs) + if resolved is not None: + return _sanitize_schema(resolved, defs, depth + 1, seen_refs | {ref}) + # Unknown/non-local ref: drop it, leaving a permissive schema. + return {} + + out: Dict[str, Any] = {} + for key, value in schema.items(): + if key in ("$defs", "definitions", "$ref", "$schema", "$id"): + # Definitions are inlined at reference sites; meta-keywords dropped. + continue + if key not in _SUPPORTED_SCHEMA_KEYS: + # Unsupported keyword: omit from the upstream copy. + continue + + if key == "properties" and isinstance(value, dict): + out["properties"] = { + prop_name: _sanitize_schema(prop_schema, defs, depth + 1, seen_refs) + for prop_name, prop_schema in value.items() + } + elif key in ("items", "not") and isinstance(value, (dict, list)): + out[key] = _sanitize_schema(value, defs, depth + 1, seen_refs) + elif key in ("anyOf", "oneOf", "allOf") and isinstance(value, list): + out[key] = [ + _sanitize_schema(sub, defs, depth + 1, seen_refs) for sub in value + ] + elif key == "additionalProperties" and isinstance(value, dict): + out[key] = _sanitize_schema(value, defs, depth + 1, seen_refs) + elif key == "required" and isinstance(value, list): + # Preserve required field names verbatim — never drop them. + out["required"] = list(value) + else: + out[key] = value + + return out + + +def sanitize_tool(tool: Dict[str, Any]) -> Dict[str, Any]: + """Return an upstream-safe copy of a single OpenAI tool definition. + + The tool ``name`` is preserved verbatim. ``function.parameters`` is + sanitized (refs resolved, unsupported keywords dropped). A tool without a + parameters schema is returned structurally unchanged. + """ + if not isinstance(tool, dict): + return tool + + func = tool.get("function") + if not isinstance(func, dict): + return tool + + params = func.get("parameters") + if not isinstance(params, dict): + return tool + + # Pool both $defs and definitions so references anywhere resolve. + defs: Dict[str, Any] = {} + for defs_key in ("$defs", "definitions"): + block = params.get(defs_key) + if isinstance(block, dict): + defs.update(block) + + sanitized_params = _sanitize_schema(params, defs) + + new_func = dict(func) + new_func["parameters"] = sanitized_params + new_tool = dict(tool) + new_tool["function"] = new_func + return new_tool + + +def sanitize_tools(tools: Any) -> Any: + """Return an upstream-safe copy of a tools list. + + Non-list input is returned unchanged so callers can forward it as-is. Each + tool is sanitized independently; tool ordering and names are preserved. + """ + if not isinstance(tools, list): + return tools + return [sanitize_tool(tool) for tool in tools] diff --git a/src/adapters/codebuddy/transport.py b/src/adapters/codebuddy/transport.py new file mode 100644 index 0000000..e3e9975 --- /dev/null +++ b/src/adapters/codebuddy/transport.py @@ -0,0 +1,212 @@ +"""Upstream transport for CodeBuddyAdapterV2 (spec §2, §11, §12, §14). + +Owns the single shared ``httpx.AsyncClient`` and everything about talking to +CodeBuddy over the wire: + + * a process-wide pooled client created at startup and closed at shutdown — + never one client per request (spec §11); + * granular per-stage timeouts mapped to distinct error types so a failure is + reported precisely (connect / pool / write / headers / first-chunk / idle), + never a generic "connect timeout" (spec §12); + * request-local upstream credential handling: the caller's key is used only + for that request, never stored, and only its SHA-256 fingerprint is logged + (spec §2); + * true streaming via ``client.send(request, stream=True)`` — the body is never + pre-read with ``aread()``/``.text``/``.json()`` before the caller consumes + it (spec §11). + +Header construction reuses the existing ``codebuddy_api_client`` so the web/cli +profile behavior stays identical to the legacy path. +""" +from __future__ import annotations + +import asyncio +import hashlib +import logging +from typing import AsyncIterator, Dict, Optional + +import httpx + +from src.codebuddy_api_client import codebuddy_api_client + +from .config import AdapterSettings +from .errors import ( + ConnectTimeoutError, + FirstChunkTimeoutError, + HeadersTimeoutError, + PoolTimeoutError, + StreamIdleTimeoutError, + UpstreamAuthError, + UpstreamNetworkError, + UpstreamRejectedError, + UpstreamRetryableError, + WriteTimeoutError, +) + +logger = logging.getLogger(__name__) + + +def key_fingerprint(token: Optional[str]) -> str: + """Return a short, non-reversible SHA-256 fingerprint of an API key. + + Used only for safe diagnostics; the raw key is never logged. + """ + if not token: + return "none" + return hashlib.sha256(token.encode("utf-8")).hexdigest()[:8] + + +class UpstreamTransport: + """Shared HTTP transport to the CodeBuddy upstream. + + A single instance is held by the adapter for the process lifetime. Timeouts + are read from :class:`AdapterSettings` when the client is first built. + """ + + def __init__(self, ssl_verify: bool = False) -> None: + self._client: Optional[httpx.AsyncClient] = None + self._lock = asyncio.Lock() + self._ssl_verify = ssl_verify + + async def get_client(self, settings: AdapterSettings) -> httpx.AsyncClient: + """Return the shared client, creating it once under a lock.""" + if self._client is None: + async with self._lock: + if self._client is None: + self._client = httpx.AsyncClient( + verify=self._ssl_verify, + timeout=httpx.Timeout( + connect=settings.connect_timeout, + read=settings.headers_timeout, + write=settings.write_timeout, + pool=settings.pool_timeout, + ), + limits=httpx.Limits( + max_connections=100, + max_keepalive_connections=30, + keepalive_expiry=60.0, + ), + ) + logger.info("CodeBuddyAdapterV2 shared HTTP client initialized") + return self._client + + async def aclose(self) -> None: + """Close the shared client (called at application shutdown).""" + async with self._lock: + if self._client is not None: + await self._client.aclose() + self._client = None + logger.info("CodeBuddyAdapterV2 shared HTTP client closed") + + def build_headers( + self, + *, + bearer_token: str, + settings: AdapterSettings, + user_id: Optional[str] = None, + conversation_id: Optional[str] = None, + conversation_request_id: Optional[str] = None, + conversation_message_id: Optional[str] = None, + request_id: Optional[str] = None, + ) -> Dict[str, str]: + """Build upstream headers, honoring the configured key header + profile.""" + return codebuddy_api_client.generate_codebuddy_headers( + bearer_token=bearer_token, + user_id=user_id, + conversation_id=conversation_id, + conversation_request_id=conversation_request_id, + conversation_message_id=conversation_message_id, + request_id=request_id, + api_key_header=settings.upstream_api_key_header, + profile=settings.request_profile, + ) + + @staticmethod + def _classify_timeout(exc: httpx.TimeoutException): + """Map an httpx timeout to a distinct adapter timeout error.""" + if isinstance(exc, httpx.ConnectTimeout): + return ConnectTimeoutError() + if isinstance(exc, httpx.PoolTimeout): + return PoolTimeoutError() + if isinstance(exc, httpx.WriteTimeout): + return WriteTimeoutError() + if isinstance(exc, httpx.ReadTimeout): + # A read timeout while awaiting headers is a headers timeout; the + # first-chunk/idle stages are enforced separately during streaming. + return HeadersTimeoutError() + return HeadersTimeoutError() + + @staticmethod + def _classify_status(status_code: int): + """Map a non-200 upstream status to a retryable/fatal adapter error. + + Retryable (a single pre-stream retry is allowed): 408, 429, 502, 503, + 504. Fatal (never retried): 400, 401, 403, 404, 422, and other 4xx. + """ + if status_code == 401: + return UpstreamAuthError() + if status_code in (408, 429, 502, 503, 504): + return UpstreamRetryableError(status_code, "upstream_temporarily_unavailable") + if status_code >= 500: + return UpstreamRetryableError(status_code, "upstream_server_error") + return UpstreamRejectedError(status_code, "upstream_request_rejected") + + async def open_stream( + self, + *, + settings: AdapterSettings, + payload: dict, + headers: Dict[str, str], + ) -> httpx.Response: + """Open a streaming upstream response and verify its status. + + Returns an open ``httpx.Response`` whose body has NOT been read; the + caller must iterate it and ensure it is closed. Raises a distinct + timeout/te transport error, or a status-classified error on non-200. + """ + client = await self.get_client(settings) + request = client.build_request( + "POST", settings.chat_completions_url, json=payload, headers=headers + ) + try: + response = await client.send(request, stream=True) + except httpx.TimeoutException as exc: + raise self._classify_timeout(exc) from exc + except httpx.RequestError as exc: + raise UpstreamNetworkError() from exc + + if response.status_code != 200: + status_code = response.status_code + await response.aclose() + raise self._classify_status(status_code) + return response + + async def iter_lines_with_idle_timeout( + self, + response: httpx.Response, + *, + first_chunk_timeout: float, + idle_timeout: float, + ) -> AsyncIterator[str]: + """Yield decoded text chunks, enforcing first-chunk and idle timeouts. + + The wait for the FIRST chunk uses ``first_chunk_timeout``; each + subsequent gap uses ``idle_timeout``. These are distinct from the + connect/headers timeouts so a stall after headers is reported as + ``upstream_first_chunk_timeout`` / ``upstream_stream_idle_timeout``. + """ + iterator = response.aiter_text(chunk_size=8192).__aiter__() + first = True + while True: + timeout = first_chunk_timeout if first else idle_timeout + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout=timeout) + except StopAsyncIteration: + return + except asyncio.TimeoutError as exc: + if first: + raise FirstChunkTimeoutError() from exc + raise StreamIdleTimeoutError() from exc + first = False + if chunk: + yield chunk diff --git a/src/adapters/codebuddy/validation.py b/src/adapters/codebuddy/validation.py new file mode 100644 index 0000000..780dfd9 --- /dev/null +++ b/src/adapters/codebuddy/validation.py @@ -0,0 +1,108 @@ +"""Pre-upstream conversation validation for CodeBuddyAdapterV2 (spec §7). + +After normalization, and immediately before building the upstream payload, the +entire conversation is validated. A malformed tool conversation is rejected +LOCALLY with HTTP 400 (:class:`InvalidToolConversationError` carrying the +offending message index) instead of being forwarded to CodeBuddy, which would +otherwise fail with an opaque "Message N must have 'role' and 'content'" error. + +Rules enforced (spec §7): + 1. Every tool result references a preceding assistant tool call with the same + ID. + 2. Every assistant tool call carries ``content`` (at least ``""``). + 3. Every tool result has ``role == "tool"``. + 4. Every tool result has a ``tool_call_id``. + 5. Tool call IDs are unchanged (validated structurally — the normalizer copies + them verbatim). + 6. ``function.arguments`` is a valid JSON string. + 7. No unsupported Anthropic blocks remain in the final OpenAI payload. + +This is a pure function: no I/O, no globals. +""" +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, List + +from .errors import InvalidToolConversationError, MessageNormalizationError + +logger = logging.getLogger(__name__) + +_ANTHROPIC_TOOL_BLOCK_TYPES = ("tool_use", "tool_result") + + +def validate_conversation(messages: List[Dict[str, Any]]) -> None: + """Validate the final OpenAI-shaped conversation just before upstream. + + Raises :class:`MessageNormalizationError` for basic role/content problems and + :class:`InvalidToolConversationError` for tool-structure problems. Both carry + the offending message index so the router can return a precise local 400. + """ + seen_tool_call_ids: set = set() + + for index, msg in enumerate(messages): + if not isinstance(msg, dict): + raise MessageNormalizationError(index, "message is not an object") + + role = msg.get("role") + if not (isinstance(role, str) and role.strip()): + raise MessageNormalizationError(index, "message is missing a valid role") + + # Rule 2: content must be present (assistant tool-call turns included). + if "content" not in msg: + raise MessageNormalizationError(index, "message is missing content") + + # Rule 7: no unconverted Anthropic blocks may remain in content arrays. + content = msg.get("content") + if isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") in _ANTHROPIC_TOOL_BLOCK_TYPES + ): + raise InvalidToolConversationError( + index, + f"unconverted Anthropic '{block.get('type')}' block remains " + f"in content", + ) + + # Rules 5 & 6: register tool-call IDs and validate their JSON arguments. + tool_calls = msg.get("tool_calls") + if tool_calls is not None: + if not isinstance(tool_calls, list): + raise InvalidToolConversationError(index, "tool_calls must be a list") + for tc in tool_calls: + if not isinstance(tc, dict): + raise InvalidToolConversationError( + index, "tool_call must be an object" + ) + tc_id = tc.get("id") + if tc_id: + seen_tool_call_ids.add(str(tc_id)) + func = tc.get("function", {}) + arguments = func.get("arguments") if isinstance(func, dict) else None + if not isinstance(arguments, str): + raise InvalidToolConversationError( + index, "tool_call function.arguments must be a JSON string" + ) + try: + json.loads(arguments) + except (ValueError, TypeError): + raise InvalidToolConversationError( + index, "tool_call function.arguments is not valid JSON" + ) + + # Rules 1, 3 & 4: a tool result must reference a preceding tool call. + if role == "tool": + tool_call_id = msg.get("tool_call_id") + if not tool_call_id: + raise InvalidToolConversationError( + index, "tool message is missing tool_call_id" + ) + if str(tool_call_id) not in seen_tool_call_ids: + raise InvalidToolConversationError( + index, + "tool result references tool_call_id with no preceding " + "matching tool call", + ) diff --git a/src/codebuddy_router.py b/src/codebuddy_router.py index d0e6cff..2fa4cda 100644 --- a/src/codebuddy_router.py +++ b/src/codebuddy_router.py @@ -34,6 +34,7 @@ validate_upstream_messages, ) from config import ( + get_codebuddy_adapter_version, get_codebuddy_default_model, get_codebuddy_model_aliases, get_codebuddy_request_profile, @@ -131,6 +132,16 @@ async def startup(): logger.error("Invalid CODEBUDDY_AUTH_MODE configuration") if auth_mode in {"auto", "api_key_file"}: await codebuddy_api_key_manager.start_periodic_reload() + + # Warm the CodeBuddyAdapterV2 shared HTTP client too, so its pooled + # connections are ready before the first request. Failures here are + # non-fatal: the adapter lazily builds the client on first use, and the + # legacy path is unaffected. + try: + from .adapters.codebuddy import get_adapter + await get_adapter().startup() + except Exception: + logger.exception("CodeBuddyAdapterV2 startup skipped") logger.info("HTTP connection pool and API key pool initialized") @staticmethod @@ -139,6 +150,12 @@ async def shutdown(): logger.info("CodeBuddy Router shutting down...") await codebuddy_api_key_manager.stop_periodic_reload() await close_http_client() + # Close the CodeBuddyAdapterV2 shared HTTP client as well. + try: + from .adapters.codebuddy import get_adapter + await get_adapter().shutdown() + except Exception: + logger.exception("CodeBuddyAdapterV2 shutdown skipped") logger.info("Resource cleanup complete") # Export the lifecycle manager for use by the main application @@ -1176,6 +1193,30 @@ async def chat_completions( "Invalid JSON request body", "invalid_request_error", "invalid_json", 400 ) + # Feature-flag dispatch (spec §1). When CODEBUDDY_ADAPTER_VERSION=v2 the + # isolated CodeBuddyAdapterV2 owns the entire request lifecycle. The import + # is lazy and guarded so a partially-built or unavailable adapter falls back + # to the legacy handler below rather than taking the endpoint down. The + # legacy path is preserved verbatim and selectable via + # CODEBUDDY_ADAPTER_VERSION=legacy. + if get_codebuddy_adapter_version() == "v2": + try: + from .adapters.codebuddy import get_adapter + except Exception: + logger.exception( + "CodeBuddyAdapterV2 is unavailable; falling back to the legacy handler" + ) + else: + return await get_adapter().chat_completion( + request=request, + request_body=request_body, + auth_context=auth_context, + conversation_id=x_conversation_id, + conversation_request_id=x_conversation_request_id, + conversation_message_id=x_conversation_message_id, + request_id=x_request_id, + ) + try: RequestProcessor.validate_request(request_body) except HTTPException as exc: diff --git a/tests/test_adapter_v2_integration.py b/tests/test_adapter_v2_integration.py new file mode 100644 index 0000000..34d9ba6 --- /dev/null +++ b/tests/test_adapter_v2_integration.py @@ -0,0 +1,443 @@ +"""Integration tests for CodeBuddyAdapterV2 through the FastAPI app (spec §18). + +These drive real HTTP requests through the router (which dispatches to the v2 +adapter) against a mocked CodeBuddy upstream, so the full lifecycle is covered: +auth → normalize → validate → payload → upstream → stream/aggregate → response. +No real CodeBuddy API key is required. + +The upstream is mocked by replacing the adapter transport's ``get_client`` with +an ``httpx.AsyncClient`` backed by ``httpx.MockTransport``. The handler can +capture the exact payload the adapter sends upstream for assertions. +""" +import json + +import httpx +import pytest + +from src import auth, codebuddy_router +from src.adapters.codebuddy import get_adapter + +RELAY_PASSWORD = "relay-password" +ADMIN_PASSWORD = "admin-password" +KEY_A = "passthrough-account-alpha-0001" + +MODERATION_TEXT = ( + "抱歉，系统检测到您当前输入的信息存在敏感内容，我无法响应您的请求，请检查后重新输入。" +) + +SENTINEL = "SENTINEL_MARKER_9F3A" +PORTFOLIO_HTML = f"

{SENTINEL}

" + + +# --------------------------------------------------------------------------- # +# Fixtures / helpers +# --------------------------------------------------------------------------- # + + +@pytest.fixture +def app(): + from fastapi import FastAPI + + application = FastAPI() + application.include_router(codebuddy_router.router, prefix="/codebuddy") + return application + + +def configure(monkeypatch, unknown_policy="passthrough", sanitize=False): + # Route to the v2 adapter. + monkeypatch.setattr(codebuddy_router, "get_codebuddy_adapter_version", lambda: "v2") + # Passthrough auth: the client Bearer token is the upstream key. + monkeypatch.setattr(auth, "get_client_auth_mode", lambda: "passthrough") + monkeypatch.setattr(auth, "get_server_password", lambda: RELAY_PASSWORD) + monkeypatch.setattr(auth, "get_admin_password", lambda: ADMIN_PASSWORD) + + import config as root_config + from src.adapters.codebuddy import config as adapter_config + + monkeypatch.setattr(root_config, "get_codebuddy_unknown_model_policy", lambda: unknown_policy) + monkeypatch.setattr(root_config, "get_sanitize_agent_prompt", lambda: sanitize) + monkeypatch.setattr(root_config, "get_codebuddy_request_profile", lambda: "web") + monkeypatch.setattr(root_config, "get_upstream_api_key_header", lambda: "bearer") + # Ensure the adapter's settings loader sees the patched getters. + assert adapter_config.load_adapter_settings().unknown_model_policy == unknown_policy + + +def install_upstream(monkeypatch, handler): + client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + adapter = get_adapter() + + async def get_client(_settings): + return client + + monkeypatch.setattr(adapter.transport, "get_client", get_client) + return client + + +def sse_body(chunks, done=True, model="auto-chat"): + parts = [] + for delta, finish in chunks: + choice = {"index": 0, "delta": delta} + if finish is not None: + choice["finish_reason"] = finish + obj = { + "id": "chat-v2", + "object": "chat.completion.chunk", + "model": model, + "choices": [choice], + } + parts.append("data: " + json.dumps(obj, ensure_ascii=False)) + text = "\n\n".join(parts) + "\n\n" + if done: + text += "data: [DONE]\n\n" + return httpx.Response(200, text=text, headers={"content-type": "text/event-stream"}) + + +async def post_chat(app, token, body): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.post( + "/codebuddy/v1/chat/completions", + headers={"Authorization": f"Bearer {token}"}, + json=body, + ) + + +def parse_stream_events(text): + events = [] + for raw in text.split("\n\n"): + line = raw.strip() + if not line.startswith("data:"): + continue + payload = line[len("data:"):].strip() + if payload == "[DONE]": + continue + events.append(json.loads(payload)) + return events + + +# --------------------------------------------------------------------------- # +# Basic chat (spec §18.1-4) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_non_stream_basic(monkeypatch, app): + configure(monkeypatch) + install_upstream(monkeypatch, lambda _r: sse_body([({"content": "Halo!"}, "stop")])) + + resp = await post_chat( + app, KEY_A, {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}]} + ) + assert resp.status_code == 200 + data = resp.json() + assert data["object"] == "chat.completion" + assert data["choices"][0]["message"]["content"] == "Halo!" + assert data["choices"][0]["finish_reason"] == "stop" + + +@pytest.mark.asyncio +async def test_stream_basic_single_done(monkeypatch, app): + configure(monkeypatch) + install_upstream( + monkeypatch, + lambda _r: sse_body( + [({"role": "assistant"}, None), ({"content": "Hello"}, None), ({"content": " world"}, None), ({}, "stop")] + ), + ) + + resp = await post_chat( + app, + KEY_A, + {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + text = resp.text + assert text.count("[DONE]") == 1 + events = parse_stream_events(text) + streamed = "".join( + (e["choices"][0].get("delta", {}) or {}).get("content", "") for e in events + ) + assert streamed == "Hello world" + + +@pytest.mark.asyncio +async def test_sse_headers_present(monkeypatch, app): + configure(monkeypatch) + install_upstream(monkeypatch, lambda _r: sse_body([({"content": "hi"}, "stop")])) + resp = await post_chat( + app, + KEY_A, + {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + ) + assert resp.headers.get("cache-control") == "no-cache, no-transform" + assert resp.headers.get("x-accel-buffering") == "no" + + +# --------------------------------------------------------------------------- # +# Model mapping (spec §18.5-8) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_opus_label_not_rewritten_to_auto_chat(monkeypatch, app): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "ok"}, "stop")]) + + install_upstream(monkeypatch, handler) + await post_chat( + app, + KEY_A, + {"model": "claude-opus-4.7-1m", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert seen["payload"]["model"] == "claude-opus-4.7-1m" + assert seen["payload"]["model"] != "auto-chat" + + +@pytest.mark.asyncio +async def test_unknown_model_reject_returns_400(monkeypatch, app): + configure(monkeypatch, unknown_policy="reject") + install_upstream(monkeypatch, lambda _r: sse_body([({"content": "x"}, "stop")])) + resp = await post_chat( + app, KEY_A, {"model": "totally-unknown", "messages": [{"role": "user", "content": "hi"}]} + ) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "unknown_model" + + +# --------------------------------------------------------------------------- # +# Tool workflow (spec §18.14-21, §18.35-40) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_read_workflow_delivers_full_file_upstream(monkeypatch, app): + """Read → tool_result(portfolio.html) → the NEXT request must carry the + complete file to the model as a role:tool message.""" + configure(monkeypatch) + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "I can see the file."}, "stop")]) + + install_upstream(monkeypatch, handler) + + body = { + "model": "auto-chat", + "messages": [ + {"role": "user", "content": "read portfolio.html"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Reading."}, + {"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "portfolio.html"}}, + ], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": PORTFOLIO_HTML} + ], + }, + ], + } + resp = await post_chat(app, KEY_A, body) + assert resp.status_code == 200 + + upstream_msgs = seen["payload"]["messages"] + # The assistant tool_use became an OpenAI tool_calls message. + asst = next(m for m in upstream_msgs if m.get("tool_calls")) + assert asst["tool_calls"][0]["id"] == "toolu_1" + assert asst["tool_calls"][0]["function"]["name"] == "Read" + # The tool_result became a role:tool message carrying the FULL file. + tool_msg = next(m for m in upstream_msgs if m.get("role") == "tool") + assert tool_msg["tool_call_id"] == "toolu_1" + assert SENTINEL in tool_msg["content"] + assert tool_msg["content"] == PORTFOLIO_HTML + + +@pytest.mark.asyncio +async def test_read_grep_edit_read_loop_preserves_every_result(monkeypatch, app): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "done"}, "stop")]) + + install_upstream(monkeypatch, handler) + + def tu(i, name): + return { + "role": "assistant", + "content": [{"type": "tool_use", "id": f"toolu_{i}", "name": name, "input": {"n": i}}], + } + + def tr(i, payload): + return { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": f"toolu_{i}", "content": payload}], + } + + body = { + "model": "auto-chat", + "messages": [ + {"role": "user", "content": "improve portfolio"}, + tu(1, "Read"), tr(1, "READ_RESULT_1"), + tu(2, "Grep"), tr(2, "GREP_RESULT_2"), + tu(3, "Edit"), tr(3, "EDIT_RESULT_3"), + tu(4, "Read"), tr(4, "READ_RESULT_4"), + ], + } + resp = await post_chat(app, KEY_A, body) + assert resp.status_code == 200 + + tool_msgs = [m for m in seen["payload"]["messages"] if m.get("role") == "tool"] + contents = [m["content"] for m in tool_msgs] + assert contents == ["READ_RESULT_1", "GREP_RESULT_2", "EDIT_RESULT_3", "READ_RESULT_4"] + # Each tool result is paired to its originating tool call id. + assert [m["tool_call_id"] for m in tool_msgs] == ["toolu_1", "toolu_2", "toolu_3", "toolu_4"] + + +@pytest.mark.asyncio +async def test_multimodal_array_survives(monkeypatch, app): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "ok"}, "stop")]) + + install_upstream(monkeypatch, handler) + + image_block = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}} + body = { + "model": "auto-chat", + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "what is this"}, image_block]} + ], + } + resp = await post_chat(app, KEY_A, body) + assert resp.status_code == 200 + user_msg = seen["payload"]["messages"][-1] + assert isinstance(user_msg["content"], list) + assert image_block in user_msg["content"] + + +@pytest.mark.asyncio +async def test_orphan_tool_result_returns_local_400(monkeypatch, app): + configure(monkeypatch) + # Upstream must never be called; make it explode if it is. + def handler(_req): + raise AssertionError("upstream must not be called for a local 400") + + install_upstream(monkeypatch, handler) + + body = { + "model": "auto-chat", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "ghost", "content": "x"}]}, + ], + } + resp = await post_chat(app, KEY_A, body) + assert resp.status_code == 400 + assert resp.json()["error"]["code"] == "invalid_tool_conversation" + + +# --------------------------------------------------------------------------- # +# Tools passthrough + schema (spec §8, §18) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_tools_are_forwarded_with_names_preserved(monkeypatch, app): + configure(monkeypatch) + seen = {} + + def handler(req): + seen["payload"] = json.loads(req.content) + return sse_body([({"content": "ok"}, "stop")]) + + install_upstream(monkeypatch, handler) + + tools = [ + { + "type": "function", + "function": { + "name": "Read", + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + }, + } + ] + body = {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}], "tools": tools} + resp = await post_chat(app, KEY_A, body) + assert resp.status_code == 200 + fwd = seen["payload"]["tools"] + assert fwd[0]["function"]["name"] == "Read" + assert fwd[0]["function"]["parameters"]["required"] == ["file_path"] + + +# --------------------------------------------------------------------------- # +# Moderation (spec §18) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_non_stream_moderation_is_content_filter(monkeypatch, app): + configure(monkeypatch) + install_upstream(monkeypatch, lambda _r: sse_body([({"content": MODERATION_TEXT}, "stop")])) + resp = await post_chat( + app, KEY_A, {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}]} + ) + assert resp.status_code == 400 + assert resp.headers.get("X-CodeBuddy-Moderation") == "true" + assert resp.json()["error"]["type"] == "content_filter" + + +# --------------------------------------------------------------------------- # +# Streaming tool calls end-to-end (spec §9) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_non_stream_fragmented_tool_args_reconstructed(monkeypatch, app): + configure(monkeypatch) + install_upstream( + monkeypatch, + lambda _r: sse_body( + [ + ({"tool_calls": [{"index": 0, "id": "tooluse_1", "function": {"name": "Read", "arguments": '{"file_'}}]}, None), + ({"tool_calls": [{"index": 0, "function": {"arguments": 'path":"a.html"}'}}]}, None), + ({}, "tool_calls"), + ] + ), + ) + resp = await post_chat( + app, KEY_A, {"model": "auto-chat", "messages": [{"role": "user", "content": "read a"}]} + ) + assert resp.status_code == 200 + msg = resp.json()["choices"][0]["message"] + tc = msg["tool_calls"][0] + assert tc["id"] == "call_1" # tooluse_ normalized to call_ + assert json.loads(tc["function"]["arguments"]) == {"file_path": "a.html"} + assert resp.json()["choices"][0]["finish_reason"] == "tool_calls" + + +@pytest.mark.asyncio +async def test_security_key_never_appears_in_response(monkeypatch, app): + configure(monkeypatch) + install_upstream(monkeypatch, lambda _r: sse_body([({"content": "ok"}, "stop")])) + resp = await post_chat( + app, KEY_A, {"model": "auto-chat", "messages": [{"role": "user", "content": "hi"}]} + ) + assert KEY_A not in resp.text diff --git a/tests/test_adapter_v2_units.py b/tests/test_adapter_v2_units.py new file mode 100644 index 0000000..e27cf83 --- /dev/null +++ b/tests/test_adapter_v2_units.py @@ -0,0 +1,661 @@ +"""Unit tests for CodeBuddyAdapterV2 pure modules (spec §18). + +These exercise the deterministic building blocks with no HTTP: model +resolution, Anthropic→OpenAI message conversion, normalization, tool-conversation +validation, tool-schema sanitization, the streaming tool-call state machine, +SSE aggregation, and response shaping. They require no CodeBuddy API key. +""" +import json + +import pytest + +from src.adapters.codebuddy import response_mapper +from src.adapters.codebuddy.config import AdapterSettings +from src.adapters.codebuddy.errors import ( + InvalidToolArgumentsError, + InvalidToolConversationError, + MessageNormalizationError, + UnknownModelError, +) +from src.adapters.codebuddy.message_normalizer import ( + convert_anthropic_messages_to_openai, + normalize_messages_for_upstream, +) +from src.adapters.codebuddy.request_mapper import build_payload, dropped_fields, resolve_model +from src.adapters.codebuddy.stream_decoder import ( + SSELineBuffer, + StreamAggregator, + parse_sse_data_line, +) +from src.adapters.codebuddy.tool_call_state import ToolCallAccumulator +from src.adapters.codebuddy.tool_schema_adapter import sanitize_tools +from src.adapters.codebuddy.validation import validate_conversation + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # + + +def make_settings(**overrides) -> AdapterSettings: + base = dict( + adapter_version="v2", + request_profile="web", + upstream_api_key_header="bearer", + sanitize_agent_prompt=False, + max_system_prompt_length=2000, + default_model="auto-chat", + unknown_model_policy="passthrough", + connect_timeout=30.0, + pool_timeout=30.0, + write_timeout=60.0, + headers_timeout=300.0, + first_chunk_timeout=300.0, + stream_idle_timeout=600.0, + max_concurrent_upstream_requests=20, + upstream_queue_timeout=60.0, + warn_total_content_length=50000, + warn_message_count=40, + warn_tool_count=30, + api_endpoint="https://www.codebuddy.ai", + ) + base.update(overrides) + return AdapterSettings(**base) + + +AVAILABLE = ["claude-4.0", "gpt-5", "auto-chat"] + + +def _tool_use_msg(text, tool_id, name, tool_input): + content = [] + if text: + content.append({"type": "text", "text": text}) + content.append({"type": "tool_use", "id": tool_id, "name": name, "input": tool_input}) + return {"role": "assistant", "content": content} + + +def _tool_result_msg(tool_use_id, result, extra_text=None): + content = [{"type": "tool_result", "tool_use_id": tool_use_id, "content": result}] + if extra_text: + content.append({"type": "text", "text": extra_text}) + return {"role": "user", "content": content} + + +# --------------------------------------------------------------------------- # +# Model mapping (spec §18.5-8) +# --------------------------------------------------------------------------- # + + +def test_opus_label_is_not_rewritten_to_auto_chat(): + settings = make_settings() + resolved = resolve_model("claude-opus-4.7-1m", settings, AVAILABLE, {}) + assert resolved.mapped == "claude-opus-4.7-1m" + assert resolved.source == "passthrough" + assert resolved.mapped != "auto-chat" + + +def test_unknown_model_passthrough_by_default(): + settings = make_settings(unknown_model_policy="passthrough") + resolved = resolve_model("mystery-model", settings, AVAILABLE, {}) + assert resolved.mapped == "mystery-model" + assert resolved.source == "passthrough" + + +def test_unknown_model_reject_policy_raises(): + settings = make_settings(unknown_model_policy="reject") + with pytest.raises(UnknownModelError): + resolve_model("mystery-model", settings, AVAILABLE, {}) + + +def test_unknown_model_default_policy_falls_back(): + settings = make_settings(unknown_model_policy="default", default_model="auto-chat") + resolved = resolve_model("mystery-model", settings, AVAILABLE, {}) + assert resolved.mapped == "auto-chat" + assert resolved.source == "default" + + +def test_alias_is_mapped_case_insensitively(): + settings = make_settings() + resolved = resolve_model( + "Claude Opus 4.7", settings, AVAILABLE, {"claude opus 4.7": "claude-4.0"} + ) + assert resolved.mapped == "claude-4.0" + assert resolved.source == "alias" + + +def test_exact_auto_chat_is_exact_not_default(): + settings = make_settings() + resolved = resolve_model("auto-chat", settings, AVAILABLE, {}) + assert resolved.source == "exact" + + +def test_missing_model_uses_default_under_every_policy(): + for policy in ("passthrough", "reject", "default"): + settings = make_settings(unknown_model_policy=policy, default_model="auto-chat") + resolved = resolve_model(None, settings, AVAILABLE, {}) + assert resolved.mapped == "auto-chat" + assert resolved.source == "default_empty" + + +def test_dropped_fields_excludes_allowlisted(): + body = { + "model": "x", + "messages": [], + "stream": True, + "tools": [], + "tool_choice": "auto", + "response_format": {"type": "json"}, + "reasoning_effort": "high", + } + assert dropped_fields(body) == ["reasoning_effort", "response_format"] + + +def test_build_payload_always_streams_upstream(): + body = {"stream": False, "tools": [{"x": 1}], "tool_choice": "auto"} + payload = build_payload(body, "auto-chat", [{"role": "user", "content": "hi"}]) + assert payload["stream"] is True + assert payload["model"] == "auto-chat" + assert payload["tools"] == [{"x": 1}] + assert payload["tool_choice"] == "auto" + + +# --------------------------------------------------------------------------- # +# Messages: role/content guarantees (spec §18.9-13) +# --------------------------------------------------------------------------- # + + +def test_string_content_survives_normalization(): + out = normalize_messages_for_upstream([{"role": "user", "content": "hello"}]) + assert out[0]["content"] == "hello" + + +def test_array_content_is_not_replaced_with_empty_string(): + multimodal = [ + {"type": "text", "text": "look"}, + {"type": "image", "source": {"type": "base64", "data": "AAAA"}}, + ] + out = convert_anthropic_messages_to_openai( + [{"role": "user", "content": multimodal}] + ) + assert out[0]["content"] == multimodal # preserved verbatim, not "" + + +def test_assistant_tool_calls_get_empty_content(): + out = normalize_messages_for_upstream( + [{"role": "assistant", "tool_calls": [{"id": "c1"}]}] + ) + assert out[0]["content"] == "" + assert out[0]["role"] == "assistant" + + +def test_missing_role_inferred_from_tool_call_id(): + out = normalize_messages_for_upstream( + [ + {"role": "assistant", "tool_calls": [{"id": "c1"}], "content": ""}, + {"tool_call_id": "c1", "content": "result"}, + ] + ) + assert out[1]["role"] == "tool" + + +def test_undeterminable_role_raises_indexed_error(): + with pytest.raises(MessageNormalizationError) as exc: + normalize_messages_for_upstream( + [{"role": "user", "content": "ok"}, {"content": "who am i"}] + ) + assert exc.value.index == 1 + + +def test_every_normalized_message_has_role_and_content(): + out = normalize_messages_for_upstream( + [ + {"role": "system", "content": "s"}, + {"role": "user", "content": "u"}, + {"role": "assistant", "tool_calls": [{"id": "c1"}]}, + {"tool_call_id": "c1", "content": "r"}, + ] + ) + for msg in out: + assert msg.get("role") + assert "content" in msg + + +# --------------------------------------------------------------------------- # +# Tool conversion (spec §18.14-21) +# --------------------------------------------------------------------------- # + + +def test_tool_use_becomes_openai_tool_calls(): + out = convert_anthropic_messages_to_openai( + [_tool_use_msg("Reading", "toolu_1", "Read", {"file_path": "a.html"})] + ) + assert len(out) == 1 + msg = out[0] + assert msg["role"] == "assistant" + assert msg["content"] == "Reading" + tc = msg["tool_calls"][0] + assert tc["id"] == "toolu_1" + assert tc["type"] == "function" + assert tc["function"]["name"] == "Read" + assert json.loads(tc["function"]["arguments"]) == {"file_path": "a.html"} + + +def test_tool_result_becomes_role_tool(): + out = convert_anthropic_messages_to_openai( + [_tool_result_msg("toolu_1", "full file")] + ) + assert len(out) == 1 + assert out[0]["role"] == "tool" + assert out[0]["tool_call_id"] == "toolu_1" + assert out[0]["content"] == "full file" + + +def test_tool_use_id_matches_tool_result_id(): + msgs = [ + _tool_use_msg("", "toolu_ABC", "Read", {"file_path": "a"}), + _tool_result_msg("toolu_ABC", "data"), + ] + out = convert_anthropic_messages_to_openai(msgs) + assert out[0]["tool_calls"][0]["id"] == out[1]["tool_call_id"] == "toolu_ABC" + + +def test_assistant_tool_call_content_empty_when_no_text(): + out = convert_anthropic_messages_to_openai( + [_tool_use_msg("", "toolu_1", "Read", {"file_path": "a"})] + ) + assert out[0]["content"] == "" + + +def test_tool_result_preserves_full_file_contents(): + big = "" + ("x" * 10000) + "" + out = convert_anthropic_messages_to_openai([_tool_result_msg("toolu_1", big)]) + assert out[0]["content"] == big + + +def test_tool_result_not_merged_into_user_message(): + # A user turn mixing text + tool_result must split, never merge the result + # into the text message. + out = convert_anthropic_messages_to_openai( + [_tool_result_msg("toolu_1", "RESULT", extra_text="and please continue")] + ) + tool_msgs = [m for m in out if m.get("role") == "tool"] + user_msgs = [m for m in out if m.get("role") == "user"] + assert len(tool_msgs) == 1 + assert tool_msgs[0]["content"] == "RESULT" + # The trailing user text is a separate message, not merged with the result. + assert user_msgs and all("RESULT" not in str(m["content"]) for m in user_msgs) + + +def test_multiple_tool_calls_in_one_assistant_message(): + content = [ + {"type": "tool_use", "id": "t1", "name": "Read", "input": {"p": 1}}, + {"type": "tool_use", "id": "t2", "name": "Grep", "input": {"q": "x"}}, + ] + out = convert_anthropic_messages_to_openai( + [{"role": "assistant", "content": content}] + ) + assert len(out) == 1 + ids = [tc["id"] for tc in out[0]["tool_calls"]] + assert ids == ["t1", "t2"] + + +def test_multiple_tool_results_in_one_user_array(): + content = [ + {"type": "tool_result", "tool_use_id": "t1", "content": "r1"}, + {"type": "tool_result", "tool_use_id": "t2", "content": "r2"}, + ] + out = convert_anthropic_messages_to_openai([{"role": "user", "content": content}]) + assert [m["tool_call_id"] for m in out] == ["t1", "t2"] + assert [m["content"] for m in out] == ["r1", "r2"] + + +def test_no_duplication_of_tool_calls_or_results(): + msgs = [ + _tool_use_msg("", "t1", "Read", {"p": 1}), + _tool_result_msg("t1", "r1"), + ] + out = convert_anthropic_messages_to_openai(msgs) + all_tc_ids = [tc["id"] for m in out for tc in m.get("tool_calls", [])] + all_result_ids = [m["tool_call_id"] for m in out if m.get("role") == "tool"] + assert all_tc_ids == ["t1"] + assert all_result_ids == ["t1"] + + +def test_tool_result_content_list_is_flattened_completely(): + content = [ + { + "type": "tool_result", + "tool_use_id": "t1", + "content": [ + {"type": "text", "text": "part1 "}, + {"type": "text", "text": "part2"}, + ], + } + ] + out = convert_anthropic_messages_to_openai([{"role": "user", "content": content}]) + assert out[0]["content"] == "part1 part2" + + +# --------------------------------------------------------------------------- # +# Validation (spec §7, §18) +# --------------------------------------------------------------------------- # + + +def test_validate_accepts_well_formed_tool_conversation(): + msgs = [ + {"role": "user", "content": "read a"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "t1", + "type": "function", + "function": {"name": "Read", "arguments": '{"p":1}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "t1", "content": "data"}, + ] + validate_conversation(msgs) # must not raise + + +def test_validate_rejects_orphan_tool_result(): + msgs = [{"role": "tool", "tool_call_id": "ghost", "content": "x"}] + with pytest.raises(InvalidToolConversationError) as exc: + validate_conversation(msgs) + assert exc.value.index == 0 + + +def test_validate_rejects_non_json_arguments(): + msgs = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "t1", + "type": "function", + "function": {"name": "Read", "arguments": "{not json"}, + } + ], + } + ] + with pytest.raises(InvalidToolConversationError): + validate_conversation(msgs) + + +def test_validate_rejects_leftover_anthropic_block(): + msgs = [ + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1"}]} + ] + with pytest.raises(InvalidToolConversationError): + validate_conversation(msgs) + + +# --------------------------------------------------------------------------- # +# Tool schema adapter (spec §8, §18.30-34) +# --------------------------------------------------------------------------- # + + +def test_schema_preserves_name_and_required(): + tools = [ + { + "type": "function", + "function": { + "name": "Read", + "parameters": { + "type": "object", + "properties": {"file_path": {"type": "string"}}, + "required": ["file_path"], + }, + }, + } + ] + out = sanitize_tools(tools) + assert out[0]["function"]["name"] == "Read" + assert out[0]["function"]["parameters"]["required"] == ["file_path"] + + +def test_schema_resolves_local_ref(): + tools = [ + { + "type": "function", + "function": { + "name": "AskUserQuestion", + "parameters": { + "type": "object", + "properties": {"q": {"$ref": "#/$defs/Question"}}, + "required": ["q"], + "$defs": { + "Question": { + "type": "object", + "properties": {"text": {"type": "string"}}, + "required": ["text"], + } + }, + }, + }, + } + ] + out = sanitize_tools(tools) + q = out[0]["function"]["parameters"]["properties"]["q"] + # $ref inlined; $defs stripped from output. + assert q["type"] == "object" + assert q["properties"]["text"]["type"] == "string" + assert "$defs" not in out[0]["function"]["parameters"] + + +def test_schema_drops_unsupported_keyword_but_keeps_structure(): + tools = [ + { + "type": "function", + "function": { + "name": "Write", + "parameters": { + "type": "object", + "properties": {"content": {"type": "string"}}, + "required": ["content"], + "$comment": "internal", + "unevaluatedProperties": False, + }, + }, + } + ] + out = sanitize_tools(tools) + params = out[0]["function"]["parameters"] + assert "$comment" not in params + assert "unevaluatedProperties" not in params + assert params["properties"]["content"]["type"] == "string" + + +def test_schema_handles_cyclic_ref_without_infinite_loop(): + tools = [ + { + "type": "function", + "function": { + "name": "Tree", + "parameters": { + "type": "object", + "properties": {"child": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": {"child": {"$ref": "#/$defs/Node"}}, + } + }, + }, + }, + } + ] + out = sanitize_tools(tools) # must terminate + assert out[0]["function"]["name"] == "Tree" + + +# --------------------------------------------------------------------------- # +# Streaming tool-call state machine (spec §9, §18.22-29) +# --------------------------------------------------------------------------- # + + +def test_fragmented_arguments_are_joined(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [{"index": 0, "id": "t1", "function": {"name": "Read", "arguments": '{"file_'}}] + ) + acc.process_delta_tool_calls([{"index": 0, "function": {"arguments": 'path":"a'}}]) + acc.process_delta_tool_calls([{"index": 0, "function": {"arguments": '.html"}'}}]) + calls = acc.finalize() + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == {"file_path": "a.html"} + + +def test_arguments_parsed_only_after_completion(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [{"index": 0, "id": "t1", "function": {"name": "Read", "arguments": '{"a":'}}] + ) + # Mid-stream: nothing is parsed/validated yet, so no error despite bad JSON. + assert acc.has_pending() + + +def test_nested_and_escaped_arguments(): + acc = ToolCallAccumulator() + payload = '{"query":"a\\"b","opts":{"deep":[1,2,{"k":"v"}]}}' + for frag in [payload[:5], payload[5:15], payload[15:]]: + acc.process_delta_tool_calls( + [{"index": 0, "id": "t1", "function": {"name": "Grep", "arguments": frag}}] + ) + calls = acc.finalize() + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": 'a"b', + "opts": {"deep": [1, 2, {"k": "v"}]}, + } + + +def test_two_concurrent_tool_calls_are_separated(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [ + {"index": 0, "id": "t1", "function": {"name": "Read", "arguments": '{"p":1}'}}, + {"index": 1, "id": "t2", "function": {"name": "Grep", "arguments": '{"q":"x"}'}}, + ] + ) + calls = acc.finalize() + assert [c["id"] for c in calls] == ["t1", "t2"] + + +def test_tool_call_id_preserved_and_not_emitted_twice(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [{"index": 0, "id": "toolu_keep", "function": {"name": "Read", "arguments": "{}"}}] + ) + first = acc.finalize() + second = acc.finalize() # already emitted → nothing new + assert first[0]["id"] == "toolu_keep" + assert second == [] + + +def test_invalid_complete_json_raises_not_empty_object(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [{"index": 0, "id": "t1", "function": {"name": "Read", "arguments": "{broken"}}] + ) + with pytest.raises(InvalidToolArgumentsError): + acc.finalize() + + +def test_empty_arguments_normalize_to_empty_object(): + acc = ToolCallAccumulator() + acc.process_delta_tool_calls( + [{"index": 0, "id": "t1", "function": {"name": "NoArgs", "arguments": ""}}] + ) + calls = acc.finalize() + assert calls[0]["function"]["arguments"] == "{}" + + +# --------------------------------------------------------------------------- # +# SSE parsing + aggregation (spec §10) +# --------------------------------------------------------------------------- # + + +def test_parse_sse_line_ignores_done_and_comments(): + assert parse_sse_data_line("data: [DONE]") is None + assert parse_sse_data_line(": keep-alive") is None + assert parse_sse_data_line("") is None + assert parse_sse_data_line('data: {"a":1}') == {"a": 1} + + +def test_line_buffer_splits_across_chunks(): + buf = SSELineBuffer() + assert buf.feed("data: {") == [] + lines = buf.feed('"a":1}\n\n') + assert 'data: {"a":1}' in lines + + +def test_aggregator_concatenates_content_and_reasoning_separately(): + agg = StreamAggregator() + agg.process_chunk({"choices": [{"delta": {"reasoning_content": "think"}}]}) + agg.process_chunk({"choices": [{"delta": {"content": "answer"}}]}) + content, tool_calls, finish = agg.finalize() + assert content == "answer" + assert agg.reasoning_content == "think" + assert tool_calls == [] + + +def test_aggregator_reconstructs_split_tool_call(): + agg = StreamAggregator() + agg.process_chunk( + {"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "t1", "function": {"name": "Read", "arguments": '{"p"'}}]}}]} + ) + agg.process_chunk( + {"choices": [{"delta": {"tool_calls": [{"index": 0, "function": {"arguments": ':1}'}}]}, "finish_reason": "tool_calls"}]} + ) + content, tool_calls, finish = agg.finalize() + assert finish == "tool_calls" + assert json.loads(tool_calls[0]["function"]["arguments"]) == {"p": 1} + + +# --------------------------------------------------------------------------- # +# Response mapper (spec §10) +# --------------------------------------------------------------------------- # + + +def test_non_stream_response_shape(): + resp = response_mapper.build_non_stream_response( + response_id="id1", + model="auto-chat", + created=123, + content="hi", + tool_calls=[], + finish_reason="stop", + ) + assert resp["object"] == "chat.completion" + assert resp["choices"][0]["message"]["content"] == "hi" + assert resp["choices"][0]["finish_reason"] == "stop" + + +def test_non_stream_reasoning_not_folded_into_content(): + resp = response_mapper.build_non_stream_response( + response_id="id1", + model="m", + created=1, + content="final", + tool_calls=[], + finish_reason="stop", + reasoning_content="secret thoughts", + ) + msg = resp["choices"][0]["message"] + assert msg["content"] == "final" + assert msg["reasoning_content"] == "secret thoughts" + + +def test_tooluse_id_normalized_to_call_prefix(): + resp = response_mapper.build_non_stream_response( + response_id="id1", + model="m", + created=1, + content="", + tool_calls=[{"id": "tooluse_abc", "function": {"name": "Read", "arguments": "{}"}}], + finish_reason="tool_calls", + ) + assert resp["choices"][0]["message"]["tool_calls"][0]["id"] == "call_abc" diff --git a/tools-refactor-adapter.md b/tools-refactor-adapter.md new file mode 100644 index 0000000..597cdbe --- /dev/null +++ b/tools-refactor-adapter.md @@ -0,0 +1,1204 @@ +# Refactor CodeBuddy2API Menjadi Adapter V2 yang Stabil untuk Claude Code dan 9Router + +Kerjakan langsung pada repository: + +```text +https://github.com/xueyue33/codebuddy2api +``` + +Gunakan implementasi provider berikut hanya sebagai referensi teknis untuk normalisasi messages, tool schema, tool calls, streaming, dan respons CodeBuddy: + +```text +https://github.com/priyo000/etteum-pool +``` + +Fokus referensi: + +```text +src/proxy/providers/codebuddy.ts +``` + +Jangan menyalin dashboard, database, Camoufox, sistem login, account pool, atau arsitektur Etteum secara keseluruhan. + +## Latar Belakang + +CodeBuddy2API saat ini digunakan dengan alur: + +```text +Claude Code +→ 9Router +→ CodeBuddy2API +→ https://www.codebuddy.ai/v2/chat/completions +``` + +9Router tetap bertanggung jawab atas: + +- round-robin API key; +- fallback provider; +- pemilihan model; +- API key per provider. + +CodeBuddy2API hanya bertindak sebagai adapter stateless: + +```text +OpenAI/Claude-compatible request +→ CodeBuddy-compatible request +→ CodeBuddy response +→ OpenAI-compatible response +``` + +API key CodeBuddy dikirim oleh 9Router pada setiap request: + +```http +Authorization: Bearer +``` + +Adapter harus menggunakan key tersebut hanya untuk request terkait. + +## Masalah yang Pernah Terjadi + +Implementasi lama mengalami masalah berikut: + +```text +- Message N must have role and content fields +- Invalid tool parameters +- tool_result hilang dari context +- model tidak dapat melihat isi file yang sudah dibaca +- tool call dan tool result tidak cocok +- function.arguments diproses ketika JSON masih berupa fragmen +- request tool-heavy terus berulang +- retry setelah stream mulai +- beberapa request upstream berjalan setelah client disconnect +- request berhenti setelah outcome=prepared +- timeout sekitar 60 detik sebelum first byte +- 9Router menampilkan fetch connect timeout +- conversation dengan 30 tools menjadi sangat lambat +- unknown model pernah diubah diam-diam menjadi auto-chat +- system prompt Claude Code pernah disanitasi sehingga instruksi tool hilang +``` + +Jangan menambahkan patch baru ke handler lama tanpa struktur yang jelas. + +## Tujuan Utama + +Bangun `CodeBuddyAdapterV2` yang terisolasi, dapat diuji, dan tetap kompatibel dengan endpoint serta konfigurasi deployment yang sudah ada. + +Target arsitektur: + +```text +src/ +├── codebuddy_router.py +└── adapters/ + └── codebuddy/ + ├── __init__.py + ├── adapter.py + ├── config.py + ├── models.py + ├── request_mapper.py + ├── message_normalizer.py + ├── tool_schema_adapter.py + ├── tool_call_state.py + ├── transport.py + ├── stream_decoder.py + ├── response_mapper.py + ├── validation.py + └── errors.py +``` + +Struktur boleh disesuaikan dengan repository, tetapi pemisahan tanggung jawab harus tetap jelas. + +--- + +# 1. Feature Flag dan Backward Compatibility + +Tambahkan: + +```env +CODEBUDDY_ADAPTER_VERSION=v2 +``` + +Nilai yang didukung: + +```text +legacy +v2 +``` + +Default sementara: + +```env +CODEBUDDY_ADAPTER_VERSION=v2 +``` + +Routing: + +```python +if settings.codebuddy_adapter_version == "v2": + return await codebuddy_adapter_v2.chat_completion(...) +return await legacy_handler(...) +``` + +Jangan menghapus implementasi lama sampai seluruh regression test v2 lulus. + +Endpoint lama harus tetap tersedia: + +```text +POST /codebuddy/v1/chat/completions +GET /codebuddy/v1/models +GET /health +``` + +--- + +# 2. API Key Passthrough Harus Request-Local + +Mode utama: + +```env +CODEBUDDY_CLIENT_AUTH_MODE=passthrough +``` + +Perilaku: + +```text +Authorization Bearer dari request 9Router +→ dipakai sebagai API key CodeBuddy upstream +→ hanya untuk request tersebut +``` + +Ketentuan wajib: + +- jangan menyimpan key ke global state; +- jangan memasukkan key ke TXT pool; +- jangan memasukkan key ke database; +- jangan melakukan rotasi internal; +- jangan mencampur key antar-request; +- streaming harus menggunakan key yang sama sampai selesai; +- jangan fallback ke key lain; +- jangan mencetak raw key ke log; +- jangan mencetak header Authorization; +- jangan mencetak `X-Api-Key`. + +Gunakan fingerprint aman: + +```python +sha256(api_key.encode()).hexdigest()[:8] +``` + +Header upstream harus configurable: + +```env +CODEBUDDY_UPSTREAM_API_KEY_HEADER=bearer +``` + +Dukungan: + +```text +bearer +x-api-key +both +``` + +Untuk mode `bearer`: + +```http +Authorization: Bearer +``` + +Untuk mode `both`: + +```http +Authorization: Bearer +X-Api-Key: +``` + +--- + +# 3. Model Resolver yang Ketat + +Jangan mengubah unknown model menjadi `auto-chat` secara diam-diam. + +Tambahkan: + +```env +CODEBUDDY_UNKNOWN_MODEL_POLICY=passthrough +CODEBUDDY_DEFAULT_MODEL=claude-opus-4.7-1m +CODEBUDDY_MODEL_ALIASES=Claude Opus 4.7=claude-opus-4.7-1m +``` + +Nilai policy: + +```text +passthrough +reject +default +``` + +Perilaku: + +```text +passthrough +→ teruskan requested model tanpa perubahan + +reject +→ return HTTP 400 unknown_model + +default +→ gunakan CODEBUDDY_DEFAULT_MODEL +``` + +Default harus: + +```env +CODEBUDDY_UNKNOWN_MODEL_POLICY=passthrough +``` + +Log aman: + +```text +requested_model +mapped_model +mapping_source +upstream_response_model +``` + +Jangan mengubah metadata respons untuk menyembunyikan model upstream. + +--- + +# 4. Jangan Sanitasi System Prompt Claude Code + +Untuk request agentic dengan tools, system prompt Claude Code harus dipertahankan. + +Default: + +```env +CODEBUDDY_SANITIZE_AGENT_PROMPT=false +CODEBUDDY_REQUEST_PROFILE=cli +``` + +Ketentuan: + +- jangan mengganti system prompt Claude Code dengan prompt generik; +- jangan menghapus instruksi tool; +- jangan menghapus permission instructions; +- jangan menghapus aturan agent loop; +- jangan menghapus context yang diperlukan model. + +Jika sanitasi masih ingin dipertahankan untuk web chat sederhana: + +```text +has_tools=true atau Claude Code terdeteksi +→ sanitasi otomatis dilewati +``` + +Sanitasi hanya boleh menyentuh system message, tidak pernah: + +```text +user +assistant +tool +tool_result +``` + +--- + +# 5. Message Normalizer + +Bangun normalizer deterministik untuk format OpenAI dan Anthropic content blocks. + +Dukung: + +```text +string content +array content +text blocks +image blocks jika sudah didukung +tool_use +tool_result +assistant tool_calls +OpenAI role=tool +``` + +Jangan pernah melakukan pola berikut: + +```python +if not isinstance(content, str): + content = "" +``` + +Karena itu akan membuang `tool_result`, file content, dan array multimodal. + +Setiap final upstream message wajib memiliki: + +```text +role +content +``` + +Assistant yang hanya berisi tool calls: + +```json +{ + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "toolu_xxx", + "type": "function", + "function": { + "name": "Read", + "arguments": "{\"file_path\":\"portfolio.html\"}" + } + } + ] +} +``` + +Tool result: + +```json +{ + "role": "tool", + "tool_call_id": "toolu_xxx", + "content": "isi hasil tool" +} +``` + +Ketentuan: + +- pertahankan urutan message; +- jangan menduplikasi user message; +- jangan menduplikasi tool call; +- jangan menduplikasi tool result; +- jangan menggabungkan tool result ke user message yang tidak terkait; +- jangan membuang content array; +- jangan membuat message hanya berisi `tool_calls` tanpa `role` dan `content`; +- jangan membuat message hanya berisi `tool_call_id`. + +--- + +# 6. Konversi Anthropic Tool Blocks + +Input seperti: + +```json +{ + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Saya akan membaca file." + }, + { + "type": "tool_use", + "id": "toolu_123", + "name": "Read", + "input": { + "file_path": "portfolio.html" + } + } + ] +} +``` + +Harus dikonversi menjadi: + +```json +{ + "role": "assistant", + "content": "Saya akan membaca file.", + "tool_calls": [ + { + "id": "toolu_123", + "type": "function", + "function": { + "name": "Read", + "arguments": "{\"file_path\":\"portfolio.html\"}" + } + } + ] +} +``` + +Input tool result: + +```json +{ + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_123", + "content": "..." + } + ] +} +``` + +Harus menjadi: + +```json +{ + "role": "tool", + "tool_call_id": "toolu_123", + "content": "..." +} +``` + +Jika satu user content array berisi kombinasi: + +```text +text +tool_result +text +``` + +pecah menjadi beberapa message yang urut dan valid tanpa kehilangan isi. + +Pertahankan hubungan: + +```text +tool_use.id +↔ tool_result.tool_use_id +``` + +--- + +# 7. Validasi Urutan Tool Call + +Sebelum request dikirim upstream, validasi seluruh conversation. + +Aturan: + +1. Setiap tool result harus memiliki preceding assistant tool call dengan ID yang sama. +2. Setiap assistant tool call harus memiliki `content`, minimal string kosong. +3. Setiap tool result harus memiliki `role=tool`. +4. Setiap tool result harus memiliki `tool_call_id`. +5. Tool call ID tidak boleh berubah selama conversion. +6. Function arguments harus berupa JSON string valid. +7. Jangan mengirim unsupported Anthropic blocks ke upstream. + +Jika invalid, return HTTP `400` lokal: + +```json +{ + "error": { + "message": "Invalid tool conversation structure at message 13", + "type": "invalid_request_error", + "code": "invalid_tool_conversation" + } +} +``` + +Jangan mengirim payload malformed ke CodeBuddy. + +--- + +# 8. Tool Schema Adapter + +Port konsep sanitasi tool schema dari Etteum. + +Dukung: + +```text +type +properties +required +items +enum +additionalProperties +anyOf +oneOf +allOf +$ref +$defs +definitions +``` + +Ketentuan: + +- resolve local `$ref` bila diperlukan; +- jangan membuang required fields; +- jangan mengganti nama parameter; +- jangan mengubah struktur nested object secara sembarangan; +- jangan menghapus array item schema; +- jangan mengirim field JSON Schema yang tidak didukung upstream tanpa sanitasi; +- pertahankan nama tool persis seperti yang diterima dari Claude Code. + +Buat mapping schema yang deterministik dan memiliki unit test. + +--- + +# 9. State Machine untuk Streaming Tool Calls + +Jangan memproses `function.arguments` per chunk. + +Arguments dapat datang seperti: + +```text +chunk 1: {"file_ +chunk 2: path":"port +chunk 3: folio.html"} +``` + +Buat state: + +```python +@dataclass +class ToolCallState: + index: int + id: str | None + name: str | None + argument_fragments: list[str] + emitted: bool = False +``` + +Kelompokkan fragmen berdasarkan: + +```text +tool_call index +tool_call ID +``` + +Alur: + +```text +chunk pertama +→ buat state + +chunk berikutnya +→ append argument fragment + +finish_reason=tool_calls atau stream selesai +→ gabungkan fragments +→ parse JSON +→ validasi terhadap schema +→ emit satu tool call lengkap +``` + +Ketentuan: + +- jangan emit partial JSON sebagai tool input; +- jangan parse arguments sebelum lengkap; +- jangan mengganti invalid JSON menjadi `{}`; +- jangan mengambil substring JSON pertama; +- jangan menghapus parameter yang dianggap tidak dikenal; +- jangan mengirim tool call dua kali; +- jangan mengubah tool call ID; +- jangan mengirim `input` sebagai string pada format Anthropic; +- `input` harus object setelah JSON selesai di-parse. + +Jika JSON tetap invalid setelah tool call selesai, return error terstruktur. Jangan memperbaikinya secara heuristik menjadi `{}`. + +--- + +# 10. Response Mapper + +Dukung respons: + +```text +assistant text +reasoning_content bila ada +tool_calls +finish_reason +usage +moderation response +normal completion +streaming completion +``` + +Untuk OpenAI streaming: + +```json +{ + "choices": [ + { + "index": 0, + "delta": { + "content": "Halo" + }, + "finish_reason": null + } + ] +} +``` + +Untuk tool calls, emit fragmen yang valid dan konsisten. + +Untuk non-streaming: + +```json +{ + "id": "...", + "object": "chat.completion", + "model": "claude-opus-4.7-1m", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "...", + "tool_calls": [] + }, + "finish_reason": "stop" + } + ] +} +``` + +Jangan return objek `chat.completion.chunk` pada `stream=false`. + +Jangan menganggap `reasoning_content` sebagai final answer. + +--- + +# 11. True End-to-End Streaming + +Gunakan shared `httpx.AsyncClient`. + +Contoh konfigurasi: + +```python +httpx.AsyncClient( + timeout=httpx.Timeout( + connect=30.0, + read=None, + write=60.0, + pool=30.0, + ), + limits=httpx.Limits( + max_connections=100, + max_keepalive_connections=30, + keepalive_expiry=60.0, + ), +) +``` + +Lifecycle: + +```text +startup → buat client +shutdown → tutup client +``` + +Jangan membuat client baru untuk setiap request. + +Pada stream path gunakan: + +```python +client.stream(...) +``` + +atau: + +```python +client.send(request, stream=True) +``` + +Dilarang: + +```python +await response.aread() +response.text +response.json() +list(response.aiter_lines()) +``` + +sebelum downstream menerima stream. + +Response SSE harus memiliki: + +```http +Content-Type: text/event-stream +Cache-Control: no-cache, no-transform +X-Accel-Buffering: no +``` + +Pastikan GZip tidak membuffer SSE. + +Akhiri normal stream tepat satu kali: + +```text +data: [DONE] +``` + +--- + +# 12. Timeout Terpisah + +Tambahkan: + +```env +CODEBUDDY_CONNECT_TIMEOUT_SECONDS=30 +CODEBUDDY_POOL_TIMEOUT_SECONDS=30 +CODEBUDDY_WRITE_TIMEOUT_SECONDS=60 +CODEBUDDY_HEADERS_TIMEOUT_SECONDS=300 +CODEBUDDY_FIRST_CHUNK_TIMEOUT_SECONDS=300 +CODEBUDDY_STREAM_IDLE_TIMEOUT_SECONDS=600 +CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS=20 +CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS=60 +``` + +Bedakan error: + +```text +upstream_connect_timeout +upstream_pool_timeout +upstream_headers_timeout +upstream_first_chunk_timeout +upstream_stream_idle_timeout +upstream_queue_timeout +``` + +Jangan menyebut semua timeout sebagai: + +```text +fetch connect timeout +``` + +--- + +# 13. Cancellation dan Disconnect + +Jika downstream disconnect: + +- hentikan stream; +- batalkan request upstream; +- tutup response upstream; +- lepaskan semaphore; +- jangan menyimpan task di background; +- jangan retry; +- jangan lanjutkan request setelah Claude Code dihentikan. + +Tangani: + +```python +asyncio.CancelledError +``` + +Gunakan `finally` untuk cleanup. + +Log: + +```text +request_id=... +stage=downstream_disconnected +stage=upstream_cancelled +``` + +Tidak boleh ada ghost request setelah client berhenti. + +--- + +# 14. Retry Policy Internal + +Dalam mode passthrough, adapter idealnya tidak melakukan rotasi atau fallback. + +Tidak boleh retry: + +```text +400 invalid_request +400 content_filter +401 +403 +404 +422 +``` + +Retry maksimal satu kali hanya sebelum event pertama untuk: + +```text +408 +429 +502 +503 +504 +connect reset +connect timeout +``` + +Setelah salah satu event berikut diteruskan: + +```text +assistant content +reasoning delta +tool call +role delta +SSE event valid +``` + +maka: + +```text +stream_started=true +``` + +Setelah itu: + +- jangan retry; +- jangan fallback; +- jangan mengulang tool call; +- jangan mengirim duplicate content. + +9Router tetap menjadi pihak utama yang menentukan fallback. + +--- + +# 15. Telemetry Bertahap + +Setiap request memiliki `request_id`. + +Log stage: + +```text +request_received +request_normalized +request_validated +request_prepared +upstream_slot_wait_start +upstream_slot_acquired +upstream_send_start +upstream_headers_received +upstream_first_chunk +stream_finished +downstream_disconnected +request_failed +``` + +Metadata aman: + +```text +request_id +requested_model +mapped_model +mapping_source +message_count +tool_count +total_content_length +key_fingerprint +queue_wait_ms +time_to_headers_ms +time_to_first_chunk_ms +stream_duration_ms +chunk_count +finish_reason +upstream_status +``` + +Untuk debug tool structure, log hanya: + +```text +message_index +role +content_type +content_block_types +has_tool_calls +tool_call_ids +tool_result_ids +content_length +``` + +Jangan log: + +- prompt; +- isi file; +- tool argument values; +- raw API key; +- Authorization; +- X-Api-Key; +- cookies; +- complete tool result. + +--- + +# 16. Large Agentic Request Warning + +Tambahkan: + +```env +CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH=50000 +CODEBUDDY_WARN_MESSAGE_COUNT=40 +CODEBUDDY_WARN_TOOL_COUNT=30 +``` + +Jika melewati batas, log warning: + +```text +large_agentic_request=true +``` + +Jangan memotong context secara diam-diam. + +Jangan menghapus tools secara otomatis. + +Jangan meringkas tool result tanpa permintaan eksplisit. + +--- + +# 17. `.env.example` + +Tambahkan atau perbarui: + +```env +CODEBUDDY_ADAPTER_VERSION=v2 + +CODEBUDDY_HOST=0.0.0.0 +CODEBUDDY_PORT=8001 + +CODEBUDDY_API_ENDPOINT=https://www.codebuddy.ai +CODEBUDDY_CLIENT_AUTH_MODE=passthrough +CODEBUDDY_REQUEST_PROFILE=cli +CODEBUDDY_SANITIZE_AGENT_PROMPT=false + +CODEBUDDY_UNKNOWN_MODEL_POLICY=passthrough +CODEBUDDY_DEFAULT_MODEL=claude-opus-4.7-1m +CODEBUDDY_MODEL_ALIASES=Claude Opus 4.7=claude-opus-4.7-1m + +CODEBUDDY_UPSTREAM_API_KEY_HEADER=bearer + +CODEBUDDY_CONNECT_TIMEOUT_SECONDS=30 +CODEBUDDY_POOL_TIMEOUT_SECONDS=30 +CODEBUDDY_WRITE_TIMEOUT_SECONDS=60 +CODEBUDDY_HEADERS_TIMEOUT_SECONDS=300 +CODEBUDDY_FIRST_CHUNK_TIMEOUT_SECONDS=300 +CODEBUDDY_STREAM_IDLE_TIMEOUT_SECONDS=600 + +CODEBUDDY_MAX_CONCURRENT_UPSTREAM_REQUESTS=20 +CODEBUDDY_UPSTREAM_QUEUE_TIMEOUT_SECONDS=60 + +CODEBUDDY_WARN_TOTAL_CONTENT_LENGTH=50000 +CODEBUDDY_WARN_MESSAGE_COUNT=40 +CODEBUDDY_WARN_TOOL_COUNT=30 + +CODEBUDDY_LOG_LEVEL=INFO +``` + +--- + +# 18. Automated Tests + +Gunakan mock upstream. Jangan membutuhkan API key CodeBuddy asli. + +Test wajib: + +## Basic chat + +1. Chat sederhana non-streaming. +2. Chat sederhana streaming. +3. Stream selesai dengan satu `[DONE]`. +4. `stream=false` menghasilkan `message.content`. + +## Model mapping + +5. `claude-opus-4.7-1m` tidak berubah menjadi `auto-chat`. +6. Unknown model passthrough. +7. Policy reject menghasilkan HTTP 400. +8. Alias UI dipetakan dengan benar. + +## Messages + +9. Semua final message memiliki `role`. +10. Semua final message memiliki `content`. +11. Content string tetap utuh. +12. Content array tidak dibuang. +13. Multimodal content tetap valid bila didukung. + +## Tool conversion + +14. Anthropic `tool_use` menjadi OpenAI `tool_calls`. +15. Anthropic `tool_result` menjadi role `tool`. +16. `tool_use.id` cocok dengan `tool_result.tool_use_id`. +17. Assistant tool call memiliki `content=""`. +18. Tool result mempertahankan isi file. +19. Tool result tidak digabung ke user message lain. +20. Multiple tool calls pada satu assistant response. +21. Multiple tool results pada satu user content array. + +## Streaming tool calls + +22. Arguments yang terpecah dalam banyak chunk digabung. +23. Arguments baru di-parse setelah selesai. +24. Nested JSON arguments. +25. Escaped string arguments. +26. Dua tool calls bersamaan. +27. Tool call ID dipertahankan. +28. Tool call tidak di-emit dua kali. +29. Invalid complete JSON menghasilkan error, bukan `{}`. +30. `AskUserQuestion` schema tetap valid. +31. `Read` schema tetap valid. +32. `Write` schema tetap valid. +33. `Edit` schema tetap valid. +34. `Bash` schema tetap valid. + +## Tool workflow + +35. `Read → tool_result → assistant` melihat isi file. +36. `Read → AskUserQuestion → user answer → Edit`. +37. `Read → Grep → Edit → Read`. +38. Conversation dengan 25+ messages. +39. Request dengan 30 tool definitions. +40. Tool result tidak hilang setelah banyak iterasi. + +## Streaming transport + +41. Mock upstream mengirim tiga chunk dengan delay. +42. Downstream menerima chunk secara incremental. +43. Response tidak dibuffer sampai selesai. +44. Shared HTTP client digunakan ulang. +45. SSE headers benar. +46. GZip tidak membuffer SSE. +47. First chunk timeout berbeda dari connect timeout. +48. Stream idle timeout berbeda dari headers timeout. + +## Cancellation + +49. Client disconnect membatalkan upstream. +50. Semaphore dilepas setelah cancellation. +51. Tidak ada ghost request. +52. Tidak ada retry setelah disconnect. +53. Tidak ada retry setelah first chunk. + +## Errors + +54. HTTP 400 tidak di-retry. +55. HTTP 401 tidak di-retry. +56. HTTP 403 tidak di-retry. +57. HTTP 422 tidak di-retry. +58. HTTP 429 boleh di-retry maksimal sekali sebelum stream. +59. HTTP 5xx boleh di-retry maksimal sekali sebelum stream. +60. Error upstream mempertahankan status dan body yang sudah disanitasi. + +## Security + +61. Concurrent request dengan dua key berbeda tidak tertukar. +62. Raw API key tidak muncul di log. +63. Authorization tidak muncul di log. +64. Isi file tidak muncul di log. +65. Tool argument values tidak muncul di log. + +--- + +# 19. Skenario Acceptance Manual + +Uji melalui: + +```text +Claude Code +→ 9Router +→ CodeBuddyAdapterV2 +``` + +Base URL 9Router: + +```text +http://cb2api:8001/codebuddy/v1 +``` + +Skenario wajib: + +```text +1. "Halo, siapa kamu?" +2. "Baca portfolio.html" +3. "Buat portfolio ini lebih bagus" +4. Model mengirim AskUserQuestion +5. User menjawab pertanyaan +6. Model membaca file +7. Model mengedit atau overwrite file +8. Model membaca ulang hasil +9. Model selesai tanpa retry +``` + +Uji juga: + +```text +- cancel request saat upstream belum memberi first token; +- cancel saat tool loop berjalan; +- lanjutkan conversation dengan 25+ messages; +- dua Claude Code session paralel dengan API key berbeda; +- prompt dengan 30 tools; +- upstream menunggu 90 detik sebelum first chunk. +``` + +Tidak boleh terjadi: + +```text +Message must have role and content +Invalid tool parameters karena fragment JSON +tool result hilang +file sudah dibaca tetapi isi tidak terlihat model +duplicate tool call +ghost request +retry setelah stream dimulai +model otomatis berubah ke auto-chat +system prompt Claude Code dihapus +``` + +--- + +# 20. Acceptance Criteria + +Implementasi selesai jika: + +- CodeBuddy2API memiliki adapter v2 terpisah; +- router tidak lagi berisi seluruh logika provider; +- API key passthrough aman untuk concurrent request; +- model tidak diubah diam-diam; +- tool calls dan tool results selalu valid; +- file content tidak hilang dari context; +- streamed arguments direkonstruksi dengan state machine; +- invalid arguments tidak diubah menjadi `{}`; +- streaming diteruskan incremental; +- cancellation menghentikan upstream; +- tidak ada request berjalan setelah client disconnect; +- timeout dibedakan berdasarkan tahap; +- tidak ada retry setelah stream mulai; +- request tool-heavy dengan 30 tools berhasil; +- seluruh automated test lulus; +- implementasi legacy masih bisa dipilih melalui feature flag. + +--- + +# Output Wajib dari Agent + +Sebelum implementasi, tampilkan: + +```text +1. Analisis arsitektur lama. +2. Root cause utama. +3. Daftar kode yang akan dipertahankan. +4. Daftar kode yang akan dipindahkan ke adapter. +5. Referensi logika yang akan di-port dari Etteum. +6. Daftar file yang akan dibuat atau diubah. +``` + +Setelah implementasi, tampilkan: + +```text +1. Daftar file yang dibuat atau diubah. +2. Diagram alur request v2. +3. Penjelasan message normalizer. +4. Penjelasan tool-call state machine. +5. Penjelasan shared HTTP client. +6. Penjelasan cancellation. +7. Contoh `.env`. +8. Hasil lint. +9. Hasil type checking. +10. Hasil automated tests. +11. Curl streaming. +12. Curl non-streaming. +13. Contoh konfigurasi 9Router. +14. Perintah deployment Docker. +15. Catatan backward compatibility. +16. Risiko yang masih tersisa. +``` + +Implementasikan perubahan secara langsung pada repository. + +Jangan hanya memberikan analisis, pseudocode, atau potongan kode. + +Jangan memodifikasi 9Router dalam pekerjaan ini. + +Jangan menyalin seluruh Etteum. + +Port hanya logika adapter/provider yang relevan dan tulis ulang agar sesuai dengan struktur Python CodeBuddy2API. \ No newline at end of file
凭证文件使用次数Credential fileUsage count