Merge branch 'feature/aisbf-broker-client'

codai/api/app.py

@@ -27,11 +27,15 @@ from typing import List
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import FileResponse, JSONResponse
 
+from codai.broker.client import BrokerClient
+from codai.broker.service import BrokerService
+
 logger = logging.getLogger(__name__)
 
 # Import from codai modules
 from codai.pydantic.textrequest import ModelList
 from codai.models.manager import model_manager, multi_model_manager
+from codai.broker import build_capabilities_document, build_hardware_summary
 
 # Import global state functions - re-export for backward compatibility
 from codai.api.state import (
@@ -78,6 +82,12 @@ async def lifespan(app: FastAPI):
                 pass
 
     cleanup_task = _asyncio.create_task(_archive_cleanup_loop())
+    broker_service = None
+    broker_runtime = getattr(app.state, "broker_runtime", None)
+    if broker_runtime is not None:
+        broker_service = BrokerService(BrokerClient(broker_runtime), app)
+        app.state.broker_service = broker_service
+        broker_service.start()
     # Run initial cleanup on startup
     try:
         from codai.api.archive import archive_manager
@@ -87,6 +97,9 @@ async def lifespan(app: FastAPI):
 
     yield
 
+    if broker_service is not None:
+        await broker_service.stop()
+
     cleanup_task.cancel()
     try:
         await cleanup_task
@@ -214,6 +227,12 @@ async def list_models():
     return ModelList(data=models)
 
 
+@app.get("/coderai/capabilities")
+async def get_broker_capabilities():
+    """Return broker capability metadata."""
+    return build_capabilities_document(hardware=build_hardware_summary())
+
+
 @app.get("/v1/files/{filename}")
 async def get_file(filename: str):
     """Serve uploaded/generated files."""
@@ -283,4 +302,4 @@ async def delete_archive_file(filename: str):
     if not os.path.isfile(candidate):
         raise HTTPException(status_code=404, detail="File not found")
     os.remove(candidate)
-    return {"deleted": filename}
+    return {"deleted": filename}
\ No newline at end of file

codai/broker/__init__.py

+"""Public broker helpers."""
+
+from codai.broker.capabilities import (
+    DEFAULT_STUDIO_ENDPOINTS,
+    build_capabilities_document,
+    build_hardware_summary,
+    build_register_message,
+)
+from codai.broker.config import (
+    BrokerConfig,
+    BrokerConfigError,
+    BrokerRuntimeConfig,
+    build_broker_runtime_config,
+)
+
+__all__ = [
+    "DEFAULT_STUDIO_ENDPOINTS",
+    "BrokerConfig",
+    "BrokerConfigError",
+    "BrokerRuntimeConfig",
+    "build_broker_runtime_config",
+    "build_capabilities_document",
+    "build_hardware_summary",
+    "build_register_message",
+]

codai/broker/asgi_bridge.py

+"""Helpers for executing in-process ASGI HTTP requests."""
+
+from urllib.parse import urlencode
+
+
+async def execute_internal_request(app, *, method, path, headers=None, query=None, body=b""):
+    request_headers = []
+    for key, value in (headers or {}).items():
+        request_headers.append((key.lower().encode("latin-1"), str(value).encode("latin-1")))
+
+    query_string = urlencode(query or {}, doseq=True).encode("ascii")
+
+    scope = {
+        "type": "http",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "http_version": "1.1",
+        "method": method.upper(),
+        "scheme": "http",
+        "path": path,
+        "raw_path": path.encode("ascii"),
+        "query_string": query_string,
+        "headers": request_headers,
+        "client": ("127.0.0.1", 0),
+        "server": ("internal", 80),
+    }
+
+    response = {"status_code": 500, "headers": {}, "body": b""}
+    request_sent = False
+
+    async def receive():
+        nonlocal request_sent
+        if request_sent:
+            return {"type": "http.disconnect"}
+        request_sent = True
+        return {
+            "type": "http.request",
+            "body": body,
+            "more_body": False,
+        }
+
+    async def send(message):
+        if message["type"] == "http.response.start":
+            response["status_code"] = message["status"]
+            response["headers"] = {
+                key.decode("latin-1").lower(): value.decode("latin-1")
+                for key, value in message.get("headers", [])
+            }
+        elif message["type"] == "http.response.body":
+            response["body"] += message.get("body", b"")
+
+    await app(scope, receive, send)
+    return response

codai/broker/capabilities.py

+"""Broker capability documents and registration payloads."""
+
+import platform
+import socket
+from typing import Any, Dict, Sequence
+
+from codai.broker.config import BrokerRuntimeConfig
+
+DEFAULT_STUDIO_ENDPOINTS = [
+    "v1/audio/speech",
+    "v1/audio/transcriptions",
+    "v1/audio/progress",
+    "v1/audio/generate",
+    "v1/audio/stems",
+    "v1/audio/cleanup",
+    "v1/audio/voices",
+    "v1/audio/voices/{name}",
+    "v1/audio/voices/extract",
+    "v1/audio/clone",
+    "v1/images/progress",
+    "v1/images/generations",
+    "v1/images/edits",
+    "v1/images/inpaint",
+    "v1/images/upscale",
+    "v1/images/depth",
+    "v1/images/segment",
+    "v1/images/deblur",
+    "v1/images/unpixelate",
+    "v1/images/outfit",
+    "v1/images/to3d",
+    "v1/images/from3d",
+    "v1/video/progress",
+    "v1/video/generations",
+    "v1/video/upscale",
+    "v1/video/subtitle",
+    "v1/video/interpolate",
+    "v1/video/dub",
+    "v1/video/to3d",
+    "v1/video/from3d",
+]
+
+
+def build_hardware_summary() -> Dict[str, Any]:
+    """Build a conservative default hardware summary."""
+
+    return {
+        "hostname": socket.gethostname(),
+        "platform": platform.platform(),
+        "gpus": [],
+        "gpu_count": 0,
+        "total_vram_mb": 0,
+        "available_vram_mb": 0,
+    }
+
+
+def build_capabilities_document(
+    version: str = "2.0.0",
+    studio_endpoints: Sequence[str] | None = None,
+    hardware: Dict[str, Any] | None = None,
+) -> Dict[str, Any]:
+    """Build broker capability metadata for registration."""
+
+    document = {
+        "server": "codai",
+        "version": version,
+        "transports": ["websocket"],
+        "openai_compat": {
+            "chat_completions": True,
+            "responses": False,
+            "models": True,
+        },
+        "studio": {
+            "supported": True,
+            "endpoints": list(studio_endpoints or DEFAULT_STUDIO_ENDPOINTS),
+        },
+    }
+    if hardware is not None:
+        document["hardware"] = hardware
+    return document
+
+
+def build_register_message(
+    runtime: BrokerRuntimeConfig,
+    request_id: str,
+    hardware: Dict[str, Any] | None,
+    capabilities: Dict[str, Any],
+    studio_endpoints: Sequence[str] | None,
+) -> Dict[str, Any]:
+    """Build broker registration frame."""
+
+    return {
+        "v": 1,
+        "op": "register",
+        "request_id": request_id,
+        "payload": {
+            "endpoint": runtime.advertised_endpoint,
+            "transport": runtime.transport,
+            "registration_token": runtime.headers.get("Authorization", "").removeprefix("Bearer "),
+            "hardware": hardware,
+            "studio_endpoints": list(studio_endpoints or DEFAULT_STUDIO_ENDPOINTS),
+            "capabilities": capabilities,
+        },
+    }

codai/broker/client.py

+"""Minimal broker websocket client for registration and heartbeats."""
+
+from __future__ import annotations
+
+import json
+import asyncio
+import time
+import uuid
+from typing import Any, Awaitable, Callable
+
+import websockets
+
+from codai.broker.capabilities import (
+    DEFAULT_STUDIO_ENDPOINTS,
+    build_capabilities_document,
+    build_hardware_summary,
+    build_register_message,
+)
+from codai.broker.models import BrokerRequestEnvelope, success_envelope
+
+Dispatcher = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
+
+
+class BrokerClient:
+    def __init__(self, runtime, dispatcher: Dispatcher | None = None):
+        self.runtime = runtime
+        self.dispatcher = dispatcher
+        self.websocket = None
+        self.session_id = None
+
+    async def connect_and_register(self):
+        self.websocket = await websockets.connect(
+            self.runtime.websocket_url,
+            additional_headers=self.runtime.headers,
+            open_timeout=self.runtime.connect_timeout_seconds,
+        )
+
+        registered_message = json.loads(
+            await asyncio.wait_for(
+                self.websocket.recv(),
+                timeout=self.runtime.request_timeout_seconds,
+            )
+        )
+        session_id = registered_message.get("session_id")
+        if (
+            registered_message.get("event") != "registered"
+            or registered_message.get("accepted") is not True
+            or not session_id
+        ):
+            raise ValueError("broker did not accept registration")
+
+        self.session_id = session_id
+        hardware = build_hardware_summary()
+        capabilities = build_capabilities_document(
+            studio_endpoints=DEFAULT_STUDIO_ENDPOINTS,
+            hardware=hardware,
+        )
+        register_message = build_register_message(
+            runtime=self.runtime,
+            request_id=str(uuid.uuid4()),
+            hardware=hardware,
+            capabilities=capabilities,
+            studio_endpoints=DEFAULT_STUDIO_ENDPOINTS,
+        )
+        await self.websocket.send(json.dumps(register_message))
+        return register_message
+
+    def message_to_envelope(self, message: dict[str, Any]) -> BrokerRequestEnvelope:
+        return BrokerRequestEnvelope(
+            request_id=message["request_id"],
+            method=message["method"],
+            path=message["path"],
+            headers=message.get("headers", {}),
+            query=message.get("query", {}),
+            payload=message.get("payload"),
+            stream=message.get("stream", False),
+            content_type=message.get("content_type", "application/json"),
+        )
+
+    def next_reconnect_delay(self, current_delay):
+        return min(current_delay * 2, self.runtime.reconnect_max_delay_seconds)
+
+    async def run_forever(self):
+        delay = self.runtime.reconnect_initial_delay_seconds
+        while True:
+            try:
+                await self.connect_and_register()
+                delay = self.runtime.reconnect_initial_delay_seconds
+                while True:
+                    message = await asyncio.wait_for(
+                        self.websocket.recv(),
+                        timeout=self.runtime.request_timeout_seconds,
+                    )
+                    await self.handle_message(message)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                await asyncio.sleep(delay)
+                delay = self.next_reconnect_delay(delay)
+
+    async def handle_message(self, raw_message):
+        message = json.loads(raw_message)
+        if message.get("op") == "heartbeat":
+            response = success_envelope(
+                message["request_id"],
+                payload={"ts": int(time.time())},
+                event="heartbeat",
+            )
+            await self.websocket.send(json.dumps(response))
+            return response
+
+        if self.dispatcher is not None:
+            response = await self.dispatcher(message)
+            await self.websocket.send(json.dumps(response))
+            return response
+
+        return None

codai/broker/config.py

+"""Broker configuration and runtime validation."""
+
+from dataclasses import dataclass, field
+from typing import Dict
+from urllib.parse import quote, urlencode, urlsplit, urlunsplit
+
+
+class BrokerConfigError(ValueError):
+    """Raised when broker configuration is invalid."""
+
+
+@dataclass
+class BrokerConfig:
+    """Persistent broker configuration."""
+
+    enabled: bool = False
+    base_url: str = ""
+    scope: str = "user"
+    username: str = ""
+    provider_id: str = ""
+    client_id: str = ""
+    registration_token: str = ""
+    advertised_endpoint: str = ""
+    transport: str = "websocket"
+    heartbeat_interval_seconds: int = 30
+    connect_timeout_seconds: int = 10
+    request_timeout_seconds: int = 30
+    reconnect_initial_delay_seconds: int = 1
+    reconnect_max_delay_seconds: int = 60
+
+
+@dataclass
+class BrokerRuntimeConfig:
+    """Validated broker runtime configuration."""
+
+    enabled: bool
+    websocket_url: str = ""
+    headers: Dict[str, str] = field(default_factory=dict)
+    advertised_endpoint: str = ""
+    transport: str = "websocket"
+    heartbeat_interval_seconds: int = 30
+    connect_timeout_seconds: int = 10
+    request_timeout_seconds: int = 30
+    reconnect_initial_delay_seconds: int = 1
+    reconnect_max_delay_seconds: int = 60
+
+
+def build_broker_runtime_config(config: BrokerConfig) -> BrokerRuntimeConfig:
+    """Validate broker configuration and build runtime websocket settings."""
+
+    runtime = BrokerRuntimeConfig(
+        enabled=config.enabled,
+        advertised_endpoint=config.advertised_endpoint,
+        transport=config.transport,
+        heartbeat_interval_seconds=config.heartbeat_interval_seconds,
+        connect_timeout_seconds=config.connect_timeout_seconds,
+        request_timeout_seconds=config.request_timeout_seconds,
+        reconnect_initial_delay_seconds=config.reconnect_initial_delay_seconds,
+        reconnect_max_delay_seconds=config.reconnect_max_delay_seconds,
+    )
+    if not config.enabled:
+        return runtime
+
+    if config.scope == "global":
+        if config.username != "global":
+            raise BrokerConfigError("global broker scope requires username 'global'")
+        suffix = "/api/coderai/wss"
+    elif config.scope == "user":
+        if not config.username or config.username == "global":
+            raise BrokerConfigError("user broker scope requires a non-global username")
+        suffix = f"/api/u/{quote(config.username, safe='')}/coderai/wss"
+    else:
+        raise BrokerConfigError(f"unsupported broker scope: {config.scope}")
+
+    for field_name, value in (
+        ("provider_id", config.provider_id),
+        ("client_id", config.client_id),
+        ("registration_token", config.registration_token),
+    ):
+        if not value:
+            raise BrokerConfigError(f"enabled broker config requires {field_name}")
+
+    split_url = urlsplit(config.base_url)
+    if split_url.scheme not in {"http", "https", "ws", "wss"}:
+        raise BrokerConfigError("broker base_url must use http, https, ws, or wss")
+    if not split_url.netloc:
+        raise BrokerConfigError("broker base_url must include a host")
+
+    scheme = {"http": "ws", "https": "wss"}.get(split_url.scheme, split_url.scheme)
+    base_path = split_url.path.rstrip("/")
+    path = f"{base_path}{suffix}" if base_path else suffix
+    query = urlencode(
+        {
+            "provider_id": config.provider_id,
+            "client_id": config.client_id,
+            "username": config.username,
+            "registration_token": config.registration_token,
+        }
+    )
+    runtime.websocket_url = urlunsplit((scheme, split_url.netloc, path, query, ""))
+    runtime.headers = {
+        "Authorization": f"Bearer {config.registration_token}",
+        "x-coderai-provider-id": config.provider_id,
+        "x-coderai-client-id": config.client_id,
+        "x-coderai-username": config.username,
+        "x-coderai-advertised-endpoint": config.advertised_endpoint,
+    }
+    return runtime

codai/broker/dispatcher.py

+"""Broker request dispatcher helpers."""
+
+from __future__ import annotations
+
+import json
+from base64 import b64encode
+from time import perf_counter
+from typing import Any
+
+from codai.broker.asgi_bridge import execute_internal_request
+from codai.broker.models import error_envelope, success_envelope
+
+SUPPORTED_PREFIXES = (
+    "/v1/models",
+    "/v1/chat/completions",
+    "/v1/images",
+    "/v1/audio",
+    "/v1/video",
+    "/v1/pipelines",
+    "/coderai/capabilities",
+)
+
+TEXT_CONTENT_TYPES = (
+    "application/json",
+    "application/ld+json",
+    "application/problem+json",
+    "application/xml",
+    "application/x-www-form-urlencoded",
+    "text/",
+)
+
+
+def is_supported_path(path: str) -> bool:
+    """Return whether a broker path is supported."""
+
+    return any(path == prefix or path.startswith(f"{prefix}/") for prefix in SUPPORTED_PREFIXES)
+
+
+def _is_text_response(content_type: str | None) -> bool:
+    if not content_type:
+        return True
+
+    normalized = content_type.split(";", 1)[0].strip().lower()
+    return any(
+        normalized == candidate or normalized.startswith(candidate)
+        for candidate in TEXT_CONTENT_TYPES
+    )
+
+
+async def execute_broker_request(app, envelope):
+    """Validate and execute a broker request envelope."""
+
+    envelope.validate()
+
+    if not is_supported_path(envelope.path):
+        return error_envelope(
+            envelope.request_id,
+            code="unsupported_endpoint",
+            message=f"Unsupported endpoint: {envelope.path}",
+        )
+
+    body: bytes
+    if isinstance(envelope.payload, (dict, list)):
+        body = json.dumps(envelope.payload, separators=(",", ":")).encode("utf-8")
+    elif isinstance(envelope.payload, str):
+        body = envelope.payload.encode("utf-8")
+    elif isinstance(envelope.payload, bytes):
+        body = envelope.payload
+    elif envelope.payload is None:
+        body = b""
+    else:
+        body = json.dumps(envelope.payload, separators=(",", ":")).encode("utf-8")
+
+    headers = dict(envelope.headers)
+    if body and "content-type" not in {key.lower() for key in headers}:
+        headers["content-type"] = envelope.content_type
+
+    started_at = perf_counter()
+    response = await execute_internal_request(
+        app,
+        method=envelope.method,
+        path=envelope.path,
+        headers=headers,
+        query=envelope.query,
+        body=body,
+    )
+    elapsed_ms = round((perf_counter() - started_at) * 1000, 3)
+
+    response_headers = response["headers"]
+    if envelope.stream:
+        response_headers = {
+            key: value
+            for key, value in response_headers.items()
+            if key.lower() != "content-length"
+        }
+
+    payload: dict[str, Any] = {
+        "status_code": response["status_code"],
+        "headers": response_headers,
+    }
+    content_type = response_headers.get("content-type")
+    if content_type:
+        payload["content_type"] = content_type
+
+    if _is_text_response(content_type):
+        payload["body"] = response["body"].decode("utf-8")
+    else:
+        payload["body_base64"] = b64encode(response["body"]).decode("ascii")
+        filename = response_headers.get("x-filename")
+        if filename:
+            payload["filename"] = filename
+
+    if envelope.stream:
+        payload["stream"] = True
+    return success_envelope(
+        envelope.request_id,
+        payload=payload,
+        metrics={"elapsed_ms": elapsed_ms},
+    )

codai/broker/models.py

+"""Broker protocol models."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict
+
+
+@dataclass
+class BrokerRequestEnvelope:
+    """Normalized broker request payload."""
+
+    request_id: str
+    method: str
+    path: str
+    headers: Dict[str, str] = field(default_factory=dict)
+    query: Dict[str, Any] = field(default_factory=dict)
+    payload: Any = None
+    stream: bool = False
+    content_type: str = "application/json"
+
+    def validate(self) -> None:
+        """Validate required request envelope fields."""
+
+        if not self.request_id or not isinstance(self.request_id, str):
+            raise ValueError("request_id is required")
+        if not self.method or not isinstance(self.method, str):
+            raise ValueError("method is required")
+        if not self.path or not isinstance(self.path, str):
+            raise ValueError("path is required")
+
+
+def success_envelope(request_id: str, payload: Any, event: str | None = None, metrics: Dict[str, Any] | None = None) -> Dict[str, Any]:
+    """Build a success response envelope."""
+
+    envelope = {
+        "request_id": request_id,
+        "ok": True,
+        "payload": payload,
+    }
+    if event is not None:
+        envelope["event"] = event
+    if metrics is not None:
+        envelope["metrics"] = metrics
+    return envelope
+
+
+def error_envelope(request_id: str, code: str, message: str, details: Dict[str, Any] | None = None) -> Dict[str, Any]:
+    """Build an error response envelope."""
+
+    error = {
+        "code": code,
+        "message": message,
+    }
+    if details is not None:
+        error["details"] = details
+    return {
+        "request_id": request_id,
+        "ok": False,
+        "error": error,
+    }

codai/broker/service.py

+"""Broker client lifecycle wrapper."""
+
+from __future__ import annotations
+
+import asyncio
+
+from codai.broker.client import BrokerClient
+from codai.broker.dispatcher import execute_broker_request
+
+
+class BrokerService:
+    def __init__(self, client: BrokerClient, app=None):
+        self.client = client
+        if app is not None:
+            async def dispatch(message):
+                envelope = self.client.message_to_envelope(message)
+                return await execute_broker_request(app, envelope)
+
+            self.client.dispatcher = dispatch
+        self.task: asyncio.Task | None = None
+
+    def start(self):
+        if not self.client.runtime.enabled or self.task is not None:
+            return
+        self.task = asyncio.create_task(self.client.run_forever())
+
+    async def stop(self):
+        if self.task is None:
+            return
+        task = self.task
+        self.task = None
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass

codai/broker/streaming.py

+"""Broker streaming response helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from codai.broker.models import success_envelope
+
+
+def stream_chunk_envelope(request_id: str, sequence: int, data: Any) -> dict[str, Any]:
+    """Build a normalized streaming chunk envelope."""
+
+    return success_envelope(
+        request_id,
+        event="stream",
+        payload={"sequence": sequence, "data": data},
+    )
+
+
+def finalize_stream(request_id: str, total_chunks: int, elapsed_ms: float) -> dict[str, Any]:
+    """Build the terminal streaming metadata envelope."""
+
+    return success_envelope(
+        request_id,
+        event="stream_end",
+        payload={"total_chunks": total_chunks},
+        metrics={"elapsed_ms": elapsed_ms},
+    )

codai/config.py

@@ -21,6 +21,8 @@ from pathlib import Path
 from typing import Any, Dict, Optional
 from dataclasses import dataclass, field
 
+from codai.broker.config import BrokerConfig
+
 
 @dataclass
 class ServerConfig:
@@ -118,6 +120,7 @@ class Config:
     image: ImageConfig = field(default_factory=ImageConfig)
     whisper: WhisperConfig = field(default_factory=WhisperConfig)
     archive: ArchiveConfig = field(default_factory=ArchiveConfig)
+    broker: BrokerConfig = field(default_factory=BrokerConfig)
     system_prompt: Optional[str] = None
     tools_closer_prompt: bool = False
     grammar_guided: bool = False
@@ -178,6 +181,22 @@ class ConfigManager:
                 "offload": {
                     "directory": "./offload"
                 },
+                "broker": {
+                    "enabled": False,
+                    "base_url": "",
+                    "scope": "user",
+                    "username": "",
+                    "provider_id": "",
+                    "client_id": "",
+                    "registration_token": "",
+                    "advertised_endpoint": "",
+                    "transport": "websocket",
+                    "heartbeat_interval_seconds": 30,
+                    "connect_timeout_seconds": 10,
+                    "request_timeout_seconds": 30,
+                    "reconnect_initial_delay_seconds": 1,
+                    "reconnect_max_delay_seconds": 60
+                },
                 "system_prompt": None,
                 "tools_closer_prompt": False,
                 "grammar_guided": False,
@@ -254,6 +273,7 @@ class ConfigManager:
                 image=ImageConfig(**config_data.get("image", {})),
                 whisper=WhisperConfig(**config_data.get("whisper", {})),
                 archive=ArchiveConfig(**config_data.get("archive", {})),
+                broker=BrokerConfig(**config_data.get("broker", {})),
                 system_prompt=config_data.get("system_prompt"),
                 tools_closer_prompt=config_data.get("tools_closer_prompt", False),
                 grammar_guided=config_data.get("grammar_guided", False),
@@ -362,6 +382,22 @@ class ConfigManager:
                 "directory": self.config.archive.directory,
                 "retention": self.config.archive.retention,
             },
+            "broker": {
+                "enabled": self.config.broker.enabled,
+                "base_url": self.config.broker.base_url,
+                "scope": self.config.broker.scope,
+                "username": self.config.broker.username,
+                "provider_id": self.config.broker.provider_id,
+                "client_id": self.config.broker.client_id,
+                "registration_token": self.config.broker.registration_token,
+                "advertised_endpoint": self.config.broker.advertised_endpoint,
+                "transport": self.config.broker.transport,
+                "heartbeat_interval_seconds": self.config.broker.heartbeat_interval_seconds,
+                "connect_timeout_seconds": self.config.broker.connect_timeout_seconds,
+                "request_timeout_seconds": self.config.broker.request_timeout_seconds,
+                "reconnect_initial_delay_seconds": self.config.broker.reconnect_initial_delay_seconds,
+                "reconnect_max_delay_seconds": self.config.broker.reconnect_max_delay_seconds,
+            },
             "system_prompt": self.config.system_prompt,
             "tools_closer_prompt": self.config.tools_closer_prompt,
             "grammar_guided": self.config.grammar_guided,

codai/main.py

@@ -17,12 +17,17 @@
 """Main entry point for codai server."""
 import sys
 import os
+import logging
 
 # Import configuration from codai modules
 from codai.cli import parse_args
 from codai.config import ConfigManager
 from codai.admin.routes import init_session_manager, set_config_manager
 from codai.api.archive import archive_manager
+from codai.broker import BrokerConfigError, build_broker_runtime_config
+
+
+logger = logging.getLogger(__name__)
 
 
 def main():
@@ -194,6 +199,11 @@ def main():
     fastapi_app = app.app
     fastapi_app.state.config_mgr = config_mgr
     fastapi_app.state.config = config
+    try:
+        fastapi_app.state.broker_runtime = build_broker_runtime_config(config.broker)
+    except BrokerConfigError as error:
+        logger.warning("Broker disabled due to invalid configuration: %s", error)
+        fastapi_app.state.broker_runtime = build_broker_runtime_config(config.broker.__class__())
     
     # Set global variables from config and args (args override config for now)
     global global_system_prompt, global_tools_closer_prompt, global_debug, global_dump, global_file_path, grammar_guided_gen

tests/test_broker_dispatch.py

+import json
+from pathlib import Path
+import sys
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import pytest
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse
+from fastapi.testclient import TestClient
+from starlette.responses import Response
+
+from codai.api.app import app as real_app
+from codai.broker.asgi_bridge import execute_internal_request
+from codai.broker.dispatcher import execute_broker_request
+from codai.broker.models import BrokerRequestEnvelope
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_internal_request_returns_json_response():
+    app = FastAPI()
+
+    @app.get("/internal")
+    async def internal_route():
+        return JSONResponse(
+            status_code=202,
+            content={"message": "ok"},
+            headers={"x-broker": "ready"},
+        )
+
+    response = await execute_internal_request(
+        app,
+        method="GET",
+        path="/internal",
+        headers={"accept": "application/json"},
+        query={"source": "broker", "limit": "1"},
+    )
+
+    assert response == {
+        "status_code": 202,
+        "headers": {
+            "content-length": "16",
+            "content-type": "application/json",
+            "x-broker": "ready",
+        },
+        "body": b'{"message":"ok"}',
+    }
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_internal_request_supports_json_body():
+    app = FastAPI()
+
+    @app.post("/internal")
+    async def internal_route(payload: dict):
+        return {"received": payload}
+
+    response = await execute_internal_request(
+        app,
+        method="POST",
+        path="/internal",
+        headers={"content-type": "application/json"},
+        body=b'{"message":"hello","count":2}',
+    )
+
+    assert response == {
+        "status_code": 200,
+        "headers": {
+            "content-length": "42",
+            "content-type": "application/json",
+        },
+        "body": b'{"received":{"message":"hello","count":2}}',
+    }
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_broker_request_returns_success_envelope_for_json_route():
+    app = FastAPI()
+
+    @app.post("/v1/chat/completions")
+    async def chat_route(payload: dict):
+        return JSONResponse(
+            status_code=201,
+            content={"received": payload, "route": "chat"},
+            headers={"x-broker": "executed"},
+        )
+
+    envelope = BrokerRequestEnvelope(
+        request_id="req-123",
+        method="POST",
+        path="/v1/chat/completions",
+        headers={"accept": "application/json"},
+        payload={"message": "hello", "count": 2},
+    )
+
+    response = await execute_broker_request(app, envelope)
+
+    assert response["request_id"] == "req-123"
+    assert response["ok"] is True
+    assert response["payload"] == {
+        "status_code": 201,
+        "headers": {
+            "content-length": "57",
+            "content-type": "application/json",
+            "x-broker": "executed",
+        },
+        "content_type": "application/json",
+        "body": '{"received":{"message":"hello","count":2},"route":"chat"}',
+    }
+    assert isinstance(response["metrics"]["elapsed_ms"], int | float)
+    assert response["metrics"]["elapsed_ms"] >= 0
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_broker_request_preserves_binary_payload_metadata():
+    app = FastAPI()
+
+    @app.get("/v1/images/render")
+    async def image_route():
+        return Response(
+            content=b"\x89PNG\r\n",
+            media_type="image/png",
+            headers={"x-filename": "image.png"},
+        )
+
+    envelope = BrokerRequestEnvelope(
+        request_id="req-binary",
+        method="GET",
+        path="/v1/images/render",
+    )
+
+    response = await execute_broker_request(app, envelope)
+
+    assert response["request_id"] == "req-binary"
+    assert response["ok"] is True
+    assert response["payload"] == {
+        "status_code": 200,
+        "headers": {
+            "content-length": "6",
+            "content-type": "image/png",
+            "x-filename": "image.png",
+        },
+        "content_type": "image/png",
+        "filename": "image.png",
+        "body_base64": "iVBORw0K",
+    }
+    assert isinstance(response["metrics"]["elapsed_ms"], int | float)
+    assert response["metrics"]["elapsed_ms"] >= 0
+
+
+@pytest.mark.anyio("asyncio")
+async def test_brokered_models_match_direct_http_response_shape():
+    direct_response = TestClient(real_app).get("/v1/models")
+
+    envelope = BrokerRequestEnvelope(
+        request_id="req-models-shape",
+        method="GET",
+        path="/v1/models",
+        headers={"accept": "application/json"},
+    )
+
+    brokered_response = await execute_broker_request(real_app, envelope)
+    brokered_body = json.loads(brokered_response["payload"]["body"])
+    direct_body = direct_response.json()
+
+    assert brokered_response["ok"] is True
+    assert brokered_response["payload"]["status_code"] == direct_response.status_code
+    assert brokered_response["payload"]["content_type"] == direct_response.headers["content-type"]
+    assert brokered_response["payload"]["headers"]["content-type"] == direct_response.headers["content-type"]
+    assert brokered_response["payload"]["headers"]["content-length"] == direct_response.headers["content-length"]
+    assert isinstance(brokered_body, dict)
+    assert isinstance(direct_body, dict)
+    assert set(brokered_body) == {"data", "object"}
+    assert set(direct_body) == {"data", "object"}
+    assert isinstance(brokered_body["object"], str)
+    assert isinstance(direct_body["object"], str)
+    assert isinstance(brokered_body["data"], list)
+    assert isinstance(direct_body["data"], list)
+    for brokered_model, direct_model in zip(brokered_body["data"], direct_body["data"], strict=True):
+        assert isinstance(brokered_model, dict)
+        assert isinstance(direct_model, dict)
+        assert brokered_model.keys() == direct_model.keys()
+        assert {"id", "object", "owned_by"}.issubset(brokered_model)
+        assert isinstance(brokered_model["id"], str)
+        assert isinstance(brokered_model["object"], str)
+        assert isinstance(brokered_model["owned_by"], str)
+    assert brokered_body == direct_body
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_broker_request_rejects_unsupported_endpoint():
+    app = FastAPI()
+    envelope = BrokerRequestEnvelope(
+        request_id="req-unsupported",
+        method="GET",
+        path="/internal",
+    )
+
+    response = await execute_broker_request(app, envelope)
+
+    assert response == {
+        "request_id": "req-unsupported",
+        "ok": False,
+        "error": {
+            "code": "unsupported_endpoint",
+            "message": "Unsupported endpoint: /internal",
+        },
+    }

tests/test_broker_streaming.py

+from pathlib import Path
+import sys
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+import pytest
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from fastapi.testclient import TestClient
+
+from codai.api.app import app as real_app
+from codai.broker.dispatcher import execute_broker_request
+from codai.broker.models import BrokerRequestEnvelope
+from codai.broker.streaming import finalize_stream, stream_chunk_envelope
+
+
+@pytest.mark.anyio("asyncio")
+async def test_stream_chunk_envelope_preserves_request_id_and_order():
+    chunk_one = stream_chunk_envelope(
+        request_id="req-stream",
+        sequence=0,
+        data={"delta": "hel"},
+    )
+    chunk_two = stream_chunk_envelope(
+        request_id="req-stream",
+        sequence=1,
+        data={"delta": "lo"},
+    )
+
+    assert chunk_one == {
+        "request_id": "req-stream",
+        "ok": True,
+        "event": "stream",
+        "payload": {
+            "sequence": 0,
+            "data": {"delta": "hel"},
+        },
+    }
+    assert chunk_two == {
+        "request_id": "req-stream",
+        "ok": True,
+        "event": "stream",
+        "payload": {
+            "sequence": 1,
+            "data": {"delta": "lo"},
+        },
+    }
+
+
+@pytest.mark.anyio("asyncio")
+async def test_finalize_stream_attaches_metrics():
+    response = finalize_stream(
+        request_id="req-stream",
+        total_chunks=2,
+        elapsed_ms=15.25,
+    )
+
+    assert response == {
+        "request_id": "req-stream",
+        "ok": True,
+        "event": "stream_end",
+        "payload": {
+            "total_chunks": 2,
+        },
+        "metrics": {
+            "elapsed_ms": 15.25,
+        },
+    }
+
+
+@pytest.mark.anyio("asyncio")
+async def test_execute_broker_request_wraps_streaming_response_metadata():
+    app = FastAPI()
+
+    @app.get("/v1/chat/completions")
+    async def stream_route():
+        async def generate():
+            yield b"data: first\n\n"
+            yield b"data: second\n\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+
+    envelope = BrokerRequestEnvelope(
+        request_id="req-stream-route",
+        method="GET",
+        path="/v1/chat/completions",
+        stream=True,
+    )
+
+    response = await execute_broker_request(app, envelope)
+
+    assert set(response) == {"request_id", "ok", "payload", "metrics"}
+    assert response["request_id"] == "req-stream-route"
+    assert response["ok"] is True
+    assert response["payload"] == {
+        "status_code": 200,
+        "headers": {
+            "content-type": "text/event-stream; charset=utf-8",
+        },
+        "content_type": "text/event-stream; charset=utf-8",
+        "body": "data: first\n\ndata: second\n\n",
+        "stream": True,
+    }
+    assert isinstance(response["metrics"], dict)
+    assert set(response["metrics"]) == {"elapsed_ms"}
+    assert isinstance(response["metrics"]["elapsed_ms"], int | float)
+    assert response["metrics"]["elapsed_ms"] >= 0
+
+
+@pytest.mark.anyio("asyncio")
+async def test_brokered_streaming_route_preserves_event_stream_body():
+    app = FastAPI()
+
+    @app.get("/v1/chat/completions")
+    async def stream_route():
+        async def generate():
+            yield b'data: {"chunk":1}\n\n'
+            yield b"data: [DONE]\n\n"
+
+        return StreamingResponse(generate(), media_type="text/event-stream")
+
+    with TestClient(app) as client:
+        with client.stream("GET", "/v1/chat/completions") as direct_response:
+            direct_body = b"".join(direct_response.iter_bytes())
+            direct_content_type = direct_response.headers["content-type"]
+            direct_status_code = direct_response.status_code
+
+    envelope = BrokerRequestEnvelope(
+        request_id="req-stream-equivalence",
+        method="GET",
+        path="/v1/chat/completions",
+        headers={"accept": "text/event-stream"},
+        stream=True,
+    )
+
+    brokered_response = await execute_broker_request(app, envelope)
+
+    assert brokered_response["ok"] is True
+    assert brokered_response["payload"]["status_code"] == direct_status_code
+    assert brokered_response["payload"]["content_type"] == direct_content_type
+    assert brokered_response["payload"]["headers"]["content-type"] == direct_content_type
+    assert brokered_response["payload"]["stream"] is True
+    assert brokered_response["payload"]["body"].encode("utf-8") == direct_body
+    assert "data: [DONE]" in brokered_response["payload"]["body"]