agentbridge-cli 0.1.11__py3-none-any.whl

agentbridge/models.py

@@ -0,0 +1,334 @@
+"""OpenAI-compatible request/response models and provider model mapping."""
+
+import re
+from dataclasses import dataclass
+from typing import Annotated, Any, Literal, Union
+
+from pydantic import BaseModel, Field
+
+
+# Tool-related types (OpenAI format)
+class FunctionDefinition(BaseModel):
+    name: str
+    description: str | None = None
+    parameters: dict | None = None
+
+
+class Tool(BaseModel):
+    type: Literal["function"]
+    function: FunctionDefinition
+
+
+class ToolChoiceFunction(BaseModel):
+    name: str
+
+
+class ToolChoiceObject(BaseModel):
+    type: Literal["function"]
+    function: ToolChoiceFunction
+
+
+# Tool call types for responses
+class FunctionCall(BaseModel):
+    name: str
+    arguments: str  # JSON string
+
+
+class ToolCall(BaseModel):
+    id: str
+    type: Literal["function"] = "function"
+    function: FunctionCall
+
+
+# Multimodal content types (OpenAI format)
+class ImageUrl(BaseModel):
+    url: str  # data:image/xxx;base64,... or https://...
+    detail: Literal["auto", "low", "high"] | None = None  # For future SDK support
+
+
+class ImageUrlContent(BaseModel):
+    type: Literal["image_url"]
+    image_url: ImageUrl
+
+
+class TextContent(BaseModel):
+    type: Literal["text"]
+    text: str
+
+
+ContentPart = Annotated[Union[TextContent, ImageUrlContent], Field(discriminator='type')]
+
+ReasoningEffort = Literal["minimal", "low", "medium", "high", "xhigh"]
+
+
+class Message(BaseModel):
+    role: Literal["system", "user", "assistant", "tool"]
+    content: str | list[ContentPart] | None = None  # None when tool_calls present
+    tool_calls: list[ToolCall] | None = None
+    tool_call_id: str | None = None  # Present when role="tool"
+    name: str | None = None  # Tool name when role="tool"
+
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: list[Message] = Field(min_length=1)
+    temperature: float | None = None
+    max_tokens: int | None = None
+    reasoning_effort: ReasoningEffort | None = None
+    reasoning: dict[str, Any] | None = None
+    stream: bool = False
+    # Tool calling support
+    tools: list[Tool] | None = None
+    tool_choice: ToolChoiceObject | str | None = None  # "auto", "none", or specific
+    # Additional fields accepted for OpenRouter/OpenAI compatibility
+    top_p: float | None = None
+    frequency_penalty: float | None = None
+    presence_penalty: float | None = None
+    stop: str | list[str] | None = None
+    n: int | None = None
+    seed: int | None = None
+    user: str | None = None
+    response_format: dict | None = None
+    logit_bias: dict | None = None
+    logprobs: bool | None = None
+    top_logprobs: int | None = None
+    parallel_tool_calls: bool | None = None
+    stream_options: dict | None = None
+    store: bool | None = None
+
+
+class ImageReference(BaseModel):
+    type: Literal["image_url"]
+    image_url: ImageUrl
+
+
+class ImageGenerationRequest(BaseModel):
+    model: str
+    prompt: str = Field(min_length=1, max_length=20_000)
+    input_references: list[ImageReference] = Field(min_length=1, max_length=1)
+    n: Literal[1] = 1
+    store: Literal[False] = False
+
+
+class ImageData(BaseModel):
+    b64_json: str
+
+
+class Choice(BaseModel):
+    index: int = 0
+    message: Message
+    finish_reason: str = "stop"  # "stop" or "tool_calls"
+
+
+class Usage(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+class ImageGenerationResponse(BaseModel):
+    id: str
+    created: int
+    model: str
+    data: list[ImageData]
+    usage: Usage = Field(default_factory=Usage)
+    usage_scope: Literal["codex_orchestration_only"] = "codex_orchestration_only"
+
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: list[Choice]
+    usage: Usage = Field(default_factory=Usage)
+    system_fingerprint: str | None = None
+
+
+class DeltaMessage(BaseModel):
+    role: str | None = None
+    content: str | None = None
+    tool_calls: list[ToolCall] | None = None
+
+
+class StreamChoice(BaseModel):
+    index: int = 0
+    delta: DeltaMessage
+    finish_reason: str | None = None
+
+
+class ChatCompletionChunk(BaseModel):
+    id: str
+    object: str = "chat.completion.chunk"
+    created: int
+    model: str
+    choices: list[StreamChoice]
+    usage: Usage | None = None
+    system_fingerprint: str | None = None
+
+
+class ModelInfo(BaseModel):
+    id: str
+    object: str = "model"
+    created: int = 0
+    owned_by: str = "agentbridge"
+
+
+class ModelList(BaseModel):
+    object: str = "list"
+    data: list[ModelInfo]
+
+
+# Error response models (OpenAI format)
+class ErrorDetail(BaseModel):
+    message: str
+    type: str  # "invalid_request_error", "server_error", etc.
+    param: str | None = None
+    code: str | None = None
+
+
+class ErrorResponse(BaseModel):
+    error: ErrorDetail
+
+
+# ---------------------------------------------------------------------------
+# Provider model mapping
+# ---------------------------------------------------------------------------
+
+# Simple names that map directly to Claude Code model identifiers
+CLAUDE_SIMPLE_NAMES: set[str] = {"opus", "sonnet", "haiku"}
+
+# Backwards-compatible export used by tests and callers.
+SIMPLE_NAMES = CLAUDE_SIMPLE_NAMES
+
+PROVIDER_NAMES: set[str] = {"claudecode", "codex", "openrouter"}
+
+CODEX_MODELS: dict[str, ReasoningEffort | None] = {
+    "gpt-5.6-sol": "high",
+    "gpt-5.5": "high",
+    "gpt-5.4": None,
+    "gpt-5.4-mini": None,
+    "gpt-5.3-codex": None,
+    "gpt-5.3-codex-spark": None,
+    "gpt-5.2": None,
+}
+
+# Backwards-compatible catalog views used by tests and callers.
+CODEX_MODEL_SLUGS: set[str] = set(CODEX_MODELS)
+CODEX_DEFAULT_REASONING_EFFORT_BY_MODEL: dict[str, ReasoningEffort] = {
+    model: effort for model, effort in CODEX_MODELS.items() if effort is not None
+}
+
+OPENROUTER_EXAMPLE_SLUGS: set[str] = {
+    "anthropic/claude-opus-4",
+    "anthropic/claude-sonnet-4",
+    "openai/gpt-5",
+}
+
+
+@dataclass(frozen=True)
+class ModelResolution:
+    """Resolved provider and backend model identifier."""
+
+    provider: Literal["claudecode", "codex", "openrouter"]
+    model: str
+
+# Word-boundary pattern for matching model names in slugs
+_MODEL_PATTERN = re.compile(
+    r'(?:^|[^a-zA-Z])(' + '|'.join(sorted(CLAUDE_SIMPLE_NAMES)) + r')(?:[^a-zA-Z]|$)',
+    re.IGNORECASE,
+)
+# Available models for /api/v1/models endpoint.
+AVAILABLE_MODELS: list[dict[str, str]] = [
+    *[
+        {
+            "slug": f"claudecode/{name}",
+            "name": f"Claude {name.capitalize()}",
+            "owned_by": "claude-code",
+        }
+        for name in sorted(CLAUDE_SIMPLE_NAMES)
+    ],
+    *[
+        {"slug": f"codex/{name}", "name": name.upper(), "owned_by": "codex-cli"}
+        for name in sorted(CODEX_MODEL_SLUGS)
+    ],
+    *[
+        {
+            "slug": f"openrouter/{name}",
+            "name": name,
+            "owned_by": "openrouter",
+        }
+        for name in sorted(OPENROUTER_EXAMPLE_SLUGS)
+    ],
+]
+
+
+class UnsupportedModelError(ValueError):
+    """Raised when an unsupported model identifier is provided."""
+
+    def __init__(self, model: str):
+        self.model = model
+        super().__init__(
+            f"Unsupported model: '{model}'. "
+            "Model IDs must start with a provider namespace. "
+            "Use claudecode/<opus|sonnet|haiku>, codex/<model>, "
+            "or openrouter/<provider>/<model>."
+        )
+
+
+def resolve_model_request(model: str) -> ModelResolution:
+    """Resolve a namespaced model ID to a backend provider.
+
+    Model IDs must begin with an AgentBridge provider namespace:
+    claudecode/<model>, codex/<model>, or openrouter/<provider>/<model>.
+
+    Args:
+        model: Namespaced model identifier.
+
+    Returns:
+        Provider and backend model identifier.
+
+    Raises:
+        UnsupportedModelError: If model is not recognized
+    """
+    model_stripped = model.strip()
+
+    provider, sep, provider_model = model_stripped.partition("/")
+    if not sep:
+        raise UnsupportedModelError(model)
+
+    provider_lower = provider.lower()
+    provider_model = provider_model.strip()
+    provider_model_lower = provider_model.lower()
+
+    if provider_lower not in PROVIDER_NAMES or not provider_model:
+        raise UnsupportedModelError(model)
+
+    if provider_lower == "codex":
+        return ModelResolution(provider="codex", model=provider_model)
+
+    if provider_lower == "openrouter":
+        if "/" not in provider_model:
+            raise UnsupportedModelError(model)
+        return ModelResolution(provider="openrouter", model=provider_model)
+
+    if provider_model_lower in CLAUDE_SIMPLE_NAMES:
+        return ModelResolution(provider="claudecode", model=provider_model_lower)
+
+    # Word-boundary match: find model name as a distinct segment in the slug
+    match = _MODEL_PATTERN.search(provider_model_lower)
+    if match:
+        return ModelResolution(provider="claudecode", model=match.group(1).lower())
+
+    # Unknown model - raise error
+    raise UnsupportedModelError(model)
+
+
+def resolve_model(model: str) -> str:
+    """Resolve a model identifier to its backend model name.
+
+    Kept for backwards compatibility. Use resolve_model_request() when the
+    provider is needed.
+    """
+    resolution = resolve_model_request(model)
+    return resolution.model

agentbridge/pool.py

@@ -0,0 +1,338 @@
+"""Lazy reusable Claude SDK client pool."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, AsyncIterator, Callable
+
+from .config import DEFAULT_POOL_SIZE
+
+if TYPE_CHECKING:
+    from claude_agent_sdk import ClaudeSDKClient
+
+logger = logging.getLogger(__name__)
+
+# Claude palette (24-bit true color)
+_CLAUDE = "\033[38;2;218;119;86m"  # Terracotta — Claude's signature orange
+_GREEN = "\033[32m"
+_RESET = "\033[0m"
+
+# Health check interval (seconds)
+HEALTH_CHECK_INTERVAL = 60
+
+
+def make_options(model: str):
+    """Create ClaudeAgentOptions with model."""
+    from claude_agent_sdk import ClaudeAgentOptions
+
+    return ClaudeAgentOptions(
+        max_turns=1,
+        setting_sources=None,  # Don't load user filesystem settings
+        system_prompt={"type": "preset", "preset": "claude_code"},
+        model=model,
+        env={"CLAUDE_CODE_BRIDGE": "1"},
+        tools=[],  # No built-in tools - pure chat mode
+    )
+
+
+class ClientPool:
+    """Lazy reusable client pool with a hard concurrency limit.
+
+    Clients are created on demand for the requested model and returned to the
+    idle pool after successful use. The pool never creates model-specific
+    clients at process startup.
+    """
+
+    def __init__(
+        self,
+        size: int = DEFAULT_POOL_SIZE,
+        on_change: Callable[[], None] | None = None,
+    ):
+        """Initialize client pool.
+
+        Args:
+            size: Maximum concurrent clients / idle slots.
+            on_change: Optional callback invoked on every state mutation.
+        """
+        self.size = size
+        self._on_change = on_change
+        self._client_models: dict[ClaudeSDKClient, str] = {}  # client -> model
+        self._available: list[ClaudeSDKClient] = []  # idle reusable clients
+        self._lock = asyncio.Lock()
+        self._semaphore = asyncio.Semaphore(size)  # limits concurrent usage
+        self._initialized = False
+        self._in_use = 0
+        self._acquire_timeout = int(os.environ.get("CLAUDE_TIMEOUT", 120))
+        self._health_check_task: asyncio.Task | None = None
+
+    def _fire_change(self) -> None:
+        """Invoke on_change callback if set."""
+        if self._on_change is not None:
+            self._on_change()
+
+    async def initialize(self) -> None:
+        """Mark the pool ready without creating any model-specific clients."""
+        if self._initialized:
+            return
+
+        self._initialized = True
+        logger.info(
+            f"{_CLAUDE}[pool]{_RESET} {_GREEN}Ready{_RESET} | max={self.size} available=0"
+        )
+        self._fire_change()
+
+        # Start periodic health check
+        self._health_check_task = asyncio.create_task(self._periodic_health_check())
+
+    async def _create_client(self, model: str) -> ClaudeSDKClient:
+        """Create and connect a new client with specified model."""
+        from claude_agent_sdk import ClaudeSDKClient
+
+        client = ClaudeSDKClient(make_options(model))
+        try:
+            await asyncio.wait_for(client.connect(), timeout=30)
+        except BaseException:
+            try:
+                await client.disconnect()
+            except BaseException:
+                pass
+            raise
+        self._client_models[client] = model
+        return client
+
+    async def _disconnect_client(self, client: ClaudeSDKClient) -> None:
+        """Disconnect client and remove from tracking."""
+        try:
+            await client.disconnect()
+        except BaseException:
+            pass
+        self._client_models.pop(client, None)
+
+    def _log_tag(self, request_id: str | None = None) -> str:
+        """Return log prefix with optional request ID."""
+        if request_id:
+            return f"[pool] [{request_id}]"
+        return "[pool]"
+
+    def _log_status(self, action: str, request_id: str | None = None) -> None:
+        """Log current pool status with model breakdown."""
+        available_models = [self._client_models[c] for c in self._available]
+        available_str = f"[{', '.join(available_models)}]"
+        tag = self._log_tag(request_id)
+
+        logger.info(
+            f"{tag} {action} | in_use={self._in_use} available={len(self._available)} "
+            f"models={available_str}"
+        )
+
+    def status(self) -> dict:
+        """Return current pool status metrics."""
+        # Include in-use clients too
+        all_models = list(self._client_models.values())
+        return {
+            "size": self.size,
+            "available": len(self._available),
+            "in_use": self._in_use,
+            "models": all_models,
+        }
+
+    def snapshot(self) -> dict:
+        """Return pool state snapshot for error diagnostics."""
+        available_models = [self._client_models[c] for c in self._available]
+        all_models = list(self._client_models.values())
+        return {
+            "size": self.size,
+            "in_use": self._in_use,
+            "available": len(self._available),
+            "available_models": available_models,
+            "all_models": all_models,
+        }
+
+    @asynccontextmanager
+    async def acquire(
+        self, model: str, request_id: str | None = None
+    ) -> AsyncIterator[ClaudeSDKClient]:
+        """Get a client for the specified model.
+
+        Takes an idle client with matching model if available, otherwise creates
+        a new one. If the pool is full of idle clients for other models, the
+        oldest idle client is discarded to make room for the requested model.
+
+        Args:
+            model: Model name (opus, sonnet, haiku).
+            request_id: Optional request ID for log correlation.
+
+        Yields:
+            A ClaudeSDKClient configured for the specified model.
+
+        Raises:
+            HTTPException: 503 if acquire times out waiting for a slot.
+        """
+        tag = self._log_tag(request_id)
+        try:
+            await asyncio.wait_for(
+                self._semaphore.acquire(), timeout=self._acquire_timeout
+            )
+        except asyncio.TimeoutError:
+            logger.error(
+                f"{tag} Semaphore timeout after {self._acquire_timeout}s | pool={self.snapshot()}"
+            )
+            from fastapi import HTTPException
+
+            raise HTTPException(
+                status_code=503,
+                detail=(
+                    f"All pool clients busy. Timed out after {self._acquire_timeout}s "
+                    "waiting for an available client."
+                ),
+            )
+
+        try:
+            client: ClaudeSDKClient | None = None
+            old_to_discard: ClaudeSDKClient | None = None
+            completed = False
+
+            async with self._lock:
+                # Take an idle client with matching model if available.
+                matching = [
+                    c for c in self._available if self._client_models[c] == model
+                ]
+                if matching:
+                    client = matching[0]
+                    self._available.remove(client)
+                    logger.info(f"{tag} Reusing idle {model} client")
+                elif len(self._client_models) >= self.size and self._available:
+                    # Pool is full, but only with idle clients for other models.
+                    old_to_discard = self._available.pop(0)
+
+                self._in_use += 1
+                self._log_status("Acquired", request_id)
+                self._fire_change()
+
+            # Disconnect non-matching idle client outside the lock.
+            if old_to_discard:
+                old_model = self._client_models.get(old_to_discard, "unknown")
+                logger.info(
+                    f"{tag} Discarding idle {old_model} client, need {model}"
+                )
+                try:
+                    await self._disconnect_client(old_to_discard)
+                except asyncio.CancelledError:
+                    logger.warning(
+                        f"{tag} Cancelled while discarding {old_model} client"
+                    )
+                    raise
+
+            # Create fresh if no reusable client was available.
+            if client is None:
+                try:
+                    client = await self._create_client(model)
+                    logger.info(f"{tag} Created fresh {model} client")
+                except asyncio.CancelledError:
+                    logger.warning(f"{tag} Cancelled while creating {model} client")
+                    raise
+
+            yield client
+            completed = True
+        except Exception:
+            raise
+        finally:
+            disconnect_client: ClaudeSDKClient | None = None
+            returned = False
+            async with self._lock:
+                self._in_use -= 1
+                if (
+                    client is not None
+                    and completed
+                    and self._initialized
+                    and len(self._available) + self._in_use < self.size
+                ):
+                    self._available.append(client)
+                    returned = True
+                else:
+                    disconnect_client = client
+                self._fire_change()
+            if returned:
+                logger.info(f"{tag} Returned {model} client to idle pool")
+            elif disconnect_client is not None:
+                status = "after cancellation" if not completed else "during shutdown"
+                logger.info(f"{tag} Destroyed {model} client {status}")
+                await self._disconnect_client(disconnect_client)
+            self._semaphore.release()
+            self._log_status("Released", request_id)
+
+    async def _periodic_health_check(self) -> None:
+        """Background task that checks idle client health every HEALTH_CHECK_INTERVAL seconds."""
+        while True:
+            await asyncio.sleep(HEALTH_CHECK_INTERVAL)
+            try:
+                await self._check_idle_clients()
+            except Exception as e:
+                logger.warning(f"[pool] Health check error: {e}")
+
+    async def _check_idle_clients(self) -> None:
+        """Check each idle client's process is alive, removing dead clients."""
+        async with self._lock:
+            clients_to_check = list(self._available)
+
+        for client in clients_to_check:
+            try:
+                # Check if the underlying process is still alive
+                transport = getattr(client, "_transport", None)
+                if transport is not None:
+                    process = getattr(transport, "_process", None)
+                    if process is not None and process.returncode is not None:
+                        raise RuntimeError("Process exited")
+            except Exception:
+                model = self._client_models.get(client, "unknown")
+                logger.warning(f"[pool] Health check: {model} client unresponsive")
+                removed = False
+                async with self._lock:
+                    if client in self._available:
+                        self._available.remove(client)
+                        removed = True
+                if removed:
+                    await self._disconnect_client(client)
+                    self._fire_change()
+
+    async def shutdown(self) -> None:
+        """Disconnect all clients and clean up."""
+        logger.info(f"[pool] Shutting down {len(self._client_models)} clients...")
+
+        # Prevent in-flight requests from returning clients to the idle pool.
+        self._initialized = False
+
+        # Cancel health check task
+        if self._health_check_task is not None:
+            self._health_check_task.cancel()
+            try:
+                await self._health_check_task
+            except asyncio.CancelledError:
+                pass
+            self._health_check_task = None
+
+        # Drain in-flight requests before disconnecting clients
+        drain_timeout = int(os.environ.get("SHUTDOWN_DRAIN_TIMEOUT", 30))
+        deadline = time.monotonic() + drain_timeout
+        while self._in_use > 0:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                logger.warning(
+                    f"[pool] Shutdown drain timeout after {drain_timeout}s "
+                    f"({self._in_use} request(s) still in flight — forcing cleanup)"
+                )
+                break
+            await asyncio.sleep(0.1)
+
+        for client in list(self._client_models.keys()):
+            await self._disconnect_client(client)
+
+        self._available.clear()
+        self._in_use = 0
+        self._fire_change()
+
+        logger.info("[pool] Shutdown complete")