Skip to content

A
aisbf
Commit 9bf96583 authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

feat: expand CoderAI broker telemetry

parent 68b2e467
Show whitespace changes
Inline Side-by-side
Showing with 555 additions and 78 deletions
aisbf/app/model_cache.py
View file @ 9bf96583
......@@ -16,6 +16,16 @@ _endpoint_model_cache: dict = {}
_background_tasks: set = set()
def _is_broker_only_coderai(provider_config) -> bool:
provider_type = getattr(provider_config, 'type', '')
if provider_type != 'coderai':
return False
coderai_config = getattr(provider_config, 'coderai_config', None) or {}
if not isinstance(coderai_config, dict):
return False
return bool(coderai_config.get('broker_mode', False))
async def fetch_provider_models(provider_id: str, config, user_id: Optional[int] = None) -> list:
global _model_cache, _model_cache_timestamps, _endpoint_model_cache
......@@ -172,6 +182,9 @@ async def prefetch_global_provider_models(config):
total += 1
if hasattr(provider_config, 'models') and provider_config.models:
continue
if _is_broker_only_coderai(provider_config):
logger.info(f"Skipping model prefetch for broker-only CoderAI provider '{provider_id}' until a broker session connects")
continue
provider_type = getattr(provider_config, 'type', '')
if provider_type in ('kilo', 'kilocode', 'coderai'):
......
aisbf/coderai_broker.py
View file @ 9bf96583
......@@ -5,6 +5,7 @@ CoderAI broker and session registry for NAT-friendly outbound connections.
from __future__ import annotations
import asyncio
from collections import deque
import json
import logging
import time
......@@ -24,6 +25,7 @@ SESSION_TTL_SECONDS = 120
REQUEST_TTL_SECONDS = 300
REPLY_TTL_SECONDS = 300
HEARTBEAT_POLL_SECONDS = 1
PERFORMANCE_WINDOW_SIZE = 100
class BrokerPlaceholderWebSocket:
......@@ -37,6 +39,7 @@ class PendingCoderAIRequest:
created_at: float
stream_queue: asyncio.Queue | None = None
event_log: list[Dict[str, Any]] = field(default_factory=list)
request_snapshot: Dict[str, Any] = field(default_factory=dict)
@dataclass
......@@ -50,6 +53,7 @@ class CoderAISession:
capabilities: Dict[str, Any] = field(default_factory=dict)
last_seen: float = field(default_factory=time.time)
closed: bool = False
recent_requests: deque[Dict[str, Any]] = field(default_factory=lambda: deque(maxlen=PERFORMANCE_WINDOW_SIZE))
class CoderAIBroker:
......@@ -94,8 +98,135 @@ class CoderAIBroker:
"closed": session.closed,
"metadata": dict(session.metadata),
"capabilities": dict(session.capabilities),
"performance": self._build_performance_snapshot(session),
}
@staticmethod
def _normalize_hardware_metadata(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]:
result = dict(metadata or {})
gpus = result.get("gpus") or result.get("gpu") or []
if isinstance(gpus, dict):
gpus = [gpus]
normalized_gpus = []
total_vram_mb = 0.0
available_vram_mb = 0.0
for gpu in gpus if isinstance(gpus, list) else []:
if not isinstance(gpu, dict):
continue
normalized_gpu = dict(gpu)
total_mb = normalized_gpu.get("total_vram_mb")
available_mb = normalized_gpu.get("available_vram_mb")
used_mb = normalized_gpu.get("used_vram_mb")
if total_mb is None and normalized_gpu.get("total_vram_gb") is not None:
total_mb = float(normalized_gpu.get("total_vram_gb")) * 1024.0
if available_mb is None and normalized_gpu.get("available_vram_gb") is not None:
available_mb = float(normalized_gpu.get("available_vram_gb")) * 1024.0
if used_mb is None and normalized_gpu.get("used_vram_gb") is not None:
used_mb = float(normalized_gpu.get("used_vram_gb")) * 1024.0
if total_mb is None and available_mb is not None and used_mb is not None:
total_mb = float(available_mb) + float(used_mb)
if available_mb is None and total_mb is not None and used_mb is not None:
available_mb = max(float(total_mb) - float(used_mb), 0.0)
if used_mb is None and total_mb is not None and available_mb is not None:
used_mb = max(float(total_mb) - float(available_mb), 0.0)
if total_mb is not None:
normalized_gpu["total_vram_mb"] = float(total_mb)
total_vram_mb += float(total_mb)
if available_mb is not None:
normalized_gpu["available_vram_mb"] = float(available_mb)
available_vram_mb += float(available_mb)
if used_mb is not None:
normalized_gpu["used_vram_mb"] = float(used_mb)
normalized_gpus.append(normalized_gpu)
if normalized_gpus:
result["gpus"] = normalized_gpus
result["gpu_count"] = len(normalized_gpus)
result["total_vram_mb"] = total_vram_mb
result["available_vram_mb"] = available_vram_mb
result["used_vram_mb"] = max(total_vram_mb - available_vram_mb, 0.0)
return result
@staticmethod
def _estimate_performance_sample(message: Dict[str, Any], started_at: float) -> Dict[str, Any]:
payload = message.get("payload") or {}
latency_ms = payload.get("latency_ms")
if latency_ms is None:
latency_ms = max((time.time() - started_at) * 1000.0, 0.0)
usage = payload.get("usage") or {}
prompt_tokens = payload.get("prompt_tokens")
completion_tokens = payload.get("completion_tokens")
total_tokens = payload.get("total_tokens")
if prompt_tokens is None:
prompt_tokens = usage.get("prompt_tokens")
if completion_tokens is None:
completion_tokens = usage.get("completion_tokens")
if total_tokens is None:
total_tokens = usage.get("total_tokens")
if total_tokens is None:
values = [value for value in (prompt_tokens, completion_tokens) if isinstance(value, (int, float))]
total_tokens = sum(values) if values else None
tokens_per_second = payload.get("tokens_per_second") or payload.get("tok_per_s") or payload.get("throughput_tps")
if tokens_per_second is None and isinstance(total_tokens, (int, float)) and latency_ms:
seconds = float(latency_ms) / 1000.0
if seconds > 0:
tokens_per_second = float(total_tokens) / seconds
return {
"latency_ms": float(latency_ms),
"prompt_tokens": int(prompt_tokens) if isinstance(prompt_tokens, (int, float)) else None,
"completion_tokens": int(completion_tokens) if isinstance(completion_tokens, (int, float)) else None,
"total_tokens": int(total_tokens) if isinstance(total_tokens, (int, float)) else None,
"tokens_per_second": float(tokens_per_second) if isinstance(tokens_per_second, (int, float)) else None,
"success": (message.get("status") or "ok") != "error",
"recorded_at": time.time(),
}
@staticmethod
def _build_performance_snapshot(session: CoderAISession) -> Dict[str, Any]:
samples = list(session.recent_requests)
if not samples:
return {
"window_size": PERFORMANCE_WINDOW_SIZE,
"sample_count": 0,
"estimated": True,
"avg_latency_ms": 0.0,
"avg_tokens_per_second": 0.0,
"avg_total_tokens": 0.0,
"success_rate": 0.0,
"last_latency_ms": None,
"last_tokens_per_second": None,
"last_total_tokens": None,
}
latencies = [sample["latency_ms"] for sample in samples if isinstance(sample.get("latency_ms"), (int, float))]
tps_values = [sample["tokens_per_second"] for sample in samples if isinstance(sample.get("tokens_per_second"), (int, float))]
total_tokens = [sample["total_tokens"] for sample in samples if isinstance(sample.get("total_tokens"), (int, float))]
success_count = sum(1 for sample in samples if sample.get("success"))
last = samples[-1]
return {
"window_size": PERFORMANCE_WINDOW_SIZE,
"sample_count": len(samples),
"estimated": True,
"avg_latency_ms": (sum(latencies) / len(latencies)) if latencies else 0.0,
"avg_tokens_per_second": (sum(tps_values) / len(tps_values)) if tps_values else 0.0,
"avg_total_tokens": (sum(total_tokens) / len(total_tokens)) if total_tokens else 0.0,
"success_rate": success_count / len(samples),
"last_latency_ms": last.get("latency_ms"),
"last_tokens_per_second": last.get("tokens_per_second"),
"last_total_tokens": last.get("total_tokens"),
}
async def _record_request_metric(self, request_id: str, message: Dict[str, Any]) -> None:
async with self._lock:
pending = self._pending.get(request_id)
if not pending:
return
session_id = (pending.request_snapshot or {}).get("session_id")
session = self._sessions_by_id.get(session_id) if session_id else None
if not session:
return
session.recent_requests.append(self._estimate_performance_sample(message, pending.created_at))
self._store_session_cache(session)
self._persist_sessions_locked()
def _persist_sessions_locked(self) -> None:
payload = {
"sessions": [
......@@ -136,6 +267,17 @@ class CoderAIBroker:
last_seen=float(raw_session.get("last_seen") or time.time()),
closed=bool(raw_session.get("closed", True)),
)
perf = raw_session.get("performance") or {}
if isinstance(perf, dict) and int(perf.get("sample_count") or 0) > 0:
sample_count = min(int(perf.get("sample_count") or 0), PERFORMANCE_WINDOW_SIZE)
for _ in range(sample_count):
session.recent_requests.append({
"latency_ms": float(perf.get("avg_latency_ms") or 0.0),
"tokens_per_second": float(perf.get("avg_tokens_per_second")) if isinstance(perf.get("avg_tokens_per_second"), (int, float)) else None,
"total_tokens": float(perf.get("avg_total_tokens")) if isinstance(perf.get("avg_total_tokens"), (int, float)) else None,
"success": True,
"recorded_at": session.last_seen,
})
session.metadata.setdefault("persisted", True)
session.metadata.setdefault("connection_state", "disconnected")
key = self._session_key(provider_id, client_id)
......@@ -178,7 +320,7 @@ class CoderAIBroker:
provider_id=provider_id,
client_id=client_id,
websocket=websocket,
metadata=dict(metadata or {}),
metadata=self._normalize_hardware_metadata(metadata),
capabilities=dict(capabilities or {}),
)
session.closed = False
......@@ -213,7 +355,7 @@ class CoderAIBroker:
return None
session.last_seen = time.time()
if metadata:
session.metadata.update(metadata)
session.metadata.update(self._normalize_hardware_metadata(metadata))
if capabilities:
session.capabilities = dict(capabilities)
self._store_session_cache(session)
......@@ -336,6 +478,12 @@ class CoderAIBroker:
future=future,
created_at=time.time(),
stream_queue=stream_queue,
request_snapshot={
"session_id": snapshot.get("session_id"),
"provider_id": provider_id,
"client_id": snapshot.get("client_id") or client_id,
"op": op,
},
)
local_session = None
......@@ -395,6 +543,7 @@ class CoderAIBroker:
if event in {"chunk", "progress", "output", "log", "data", "done", "completed"}:
await self._publish_stream_response(message)
return
await self._record_request_metric(request_id, message)
self._cache.broker_set(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
await self.resolve_response(message)
......@@ -406,6 +555,7 @@ class CoderAIBroker:
if queue is not None:
pending.event_log.append(message)
if message.get("event") in {"done", "completed"} and not pending.future.done():
await self._record_request_metric(request_id, message)
pending.future.set_result(message)
else:
await queue.put(message)
......
aisbf/providers/coderai.py
View file @ 9bf96583
......@@ -68,6 +68,8 @@ class CoderAIProviderHandler(BaseProviderHandler):
self._broker_enabled = bool(self._coderai_config.get("broker_enabled", True))
self._broker_mode = bool(self._coderai_config.get("broker_mode", False))
self._broker_preferred = bool(self._coderai_config.get("broker_preferred", False))
if self._broker_mode:
self._transport = "broker"
self._base_endpoint = self._normalize_http_base(self._raw_endpoint)
self._ws_endpoint = self._normalize_ws_endpoint(self._raw_endpoint)
self._apply_provider_defaults()
......@@ -184,6 +186,12 @@ class CoderAIProviderHandler(BaseProviderHandler):
headers["Authorization"] = f"Bearer {token}"
return headers
def _is_direct_http_mode(self) -> bool:
return self._transport == "http" and not self._broker_mode
def _is_direct_websocket_mode(self) -> bool:
return self._transport == "websocket" and not self._broker_mode
async def _use_broker(self) -> bool:
if not self._broker_enabled:
return False
......@@ -264,6 +272,9 @@ class CoderAIProviderHandler(BaseProviderHandler):
message = await coderai_broker.wait_for_stream_event(next_request_id, timeout=timeout)
def validate_credentials(self) -> bool:
if self._broker_mode:
logger.info(f"[{self.provider_id}] CoderAI broker mode enabled; waiting for inbound broker session")
return True
if self._transport == "websocket" and not self._websocket_enabled:
logger.error(f"[{self.provider_id}] WebSocket transport selected but disabled")
return False
......@@ -444,7 +455,7 @@ class CoderAIProviderHandler(BaseProviderHandler):
raise Exception(message.get("error") or "CoderAI broker request failed")
self.record_success()
return message.get("payload") or {}
if self._transport == "websocket":
if self._is_direct_websocket_mode():
if stream:
return self._ws_stream("chat.completions", payload, timeout=self._request_timeout)
message = await self._ws_roundtrip("chat.completions", payload, timeout=self._request_timeout)
......@@ -468,7 +479,7 @@ class CoderAIProviderHandler(BaseProviderHandler):
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI broker model discovery failed")
return self._extract_models(message.get("payload") or {})
if self._transport == "websocket":
if self._is_direct_websocket_mode():
message = await self._ws_roundtrip("models.list", {}, timeout=self._model_timeout)
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI model discovery failed")
......@@ -488,11 +499,13 @@ class CoderAIProviderHandler(BaseProviderHandler):
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI broker capability discovery failed")
return message.get("payload") or {}
if self._transport == "websocket":
if self._is_direct_websocket_mode():
message = await self._ws_roundtrip("capabilities", {}, timeout=self._model_timeout)
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI capability discovery failed")
return message.get("payload") or {}
if self._broker_mode:
raise Exception("CoderAI broker mode requires an active broker session")
return await self._http_json("GET", "/coderai/capabilities", timeout=self._model_timeout)
async def register_client(self) -> Dict[str, Any]:
......@@ -507,11 +520,13 @@ class CoderAIProviderHandler(BaseProviderHandler):
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI broker registration failed")
return message.get("payload") or {}
if self._transport == "websocket":
if self._is_direct_websocket_mode():
message = await self._ws_roundtrip("register", payload, timeout=self._model_timeout)
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI registration failed")
return message.get("payload") or {}
if self._broker_mode:
raise Exception("CoderAI broker mode does not support outbound registration")
return await self._http_json("POST", self._registration_path, payload, timeout=self._model_timeout)
async def proxy_native_request(
......@@ -549,11 +564,13 @@ class CoderAIProviderHandler(BaseProviderHandler):
raise Exception(message.get("error") or "CoderAI broker proxy request failed")
envelope = message.get("payload") or {}
return int(envelope.get("status_code") or 200), envelope
if self._transport == "websocket":
if self._is_direct_websocket_mode():
message = await self._ws_roundtrip("proxy", payload, timeout=self._request_timeout)
if (message.get("status") or "ok") == "error":
raise Exception(message.get("error") or "CoderAI proxy request failed")
envelope = message.get("payload") or {}
return int(envelope.get("status_code") or 200), envelope
if self._broker_mode:
raise Exception("CoderAI broker mode requires an active broker session")
response = await self._http_json(method.upper(), endpoint_path, body or {}, timeout=self._request_timeout)
return 200, response
aisbf/routes/coderai_broker.py
View file @ 9bf96583
......@@ -62,6 +62,11 @@ async def _coderai_broker_websocket_impl(websocket: WebSocket, scope_name: str):
"endpoint": payload.get("endpoint"),
"transport": payload.get("transport"),
"studio_endpoints": payload.get("studio_endpoints") or [],
"hardware": payload.get("hardware") or {},
"gpus": payload.get("gpus") or ((payload.get("hardware") or {}).get("gpus")) or [],
"gpu_count": payload.get("gpu_count") or ((payload.get("hardware") or {}).get("gpu_count")),
"total_vram_mb": payload.get("total_vram_mb") or ((payload.get("hardware") or {}).get("total_vram_mb")),
"available_vram_mb": payload.get("available_vram_mb") or ((payload.get("hardware") or {}).get("available_vram_mb")),
"owner_user_id": owner_user_id,
"username": username,
"scope_name": expected_scope,
......
config/providers.json
View file @ 9bf96583
......@@ -65,17 +65,17 @@
},
"coderai": {
"id": "coderai",
"name": "CoderAI Local Bridge",
"name": "CoderAI Broker",
"endpoint": "http://127.0.0.1:11437",
"type": "coderai",
"api_key_required": false,
"rate_limit": 0,
"coderai_config": {
"transport": "http",
"http_enabled": true,
"websocket_enabled": true,
"transport": "broker",
"http_enabled": false,
"websocket_enabled": false,
"broker_enabled": true,
"broker_mode": false,
"broker_mode": true,
"broker_preferred": true,
"discovery_enabled": true,
"client_id": "aisbf-default",
......
docs/coderai-broker-implementation-reference.md
View file @ 9bf96583
......@@ -25,12 +25,13 @@ The broker client should:
1. Dial AISBF over `ws://` or `wss://`
2. Authenticate using the provider-scoped `registration_token`
3. Register metadata and capabilities after connect
3. Register metadata, hardware inventory, and capabilities after connect
4. Stay connected with heartbeat support
5. Receive queued or direct broker requests from AISBF
6. Execute supported operations locally
7. Send back success, error, binary, and streaming envelopes with the same `request_id`
8. Automatically reconnect if the connection drops
8. Include performance metrics for completed requests whenever available
9. Automatically reconnect if the connection drops
## AISBF Concepts You Must Match
......@@ -137,7 +138,7 @@ Store:
### 3. Send explicit `register` operation
After the `registered` event, CoderAI must send a `register` message describing its capabilities and advertised endpoints.
After the `registered` event, CoderAI must send a `register` message describing its capabilities, hardware inventory, and advertised endpoints.
### 4. Enter long-lived receive loop
......@@ -164,6 +165,23 @@ CoderAI should send this after receiving the initial AISBF `registered` event.
"endpoint": "ws://local-coderai-or-descriptive-endpoint",
"transport": "websocket",
"registration_token": "<same_registration_token>",
"hardware": {
"hostname": "workstation-01",
"platform": "linux",
"gpus": [
{
"index": 0,
"name": "NVIDIA RTX 4090",
"vendor": "nvidia",
"total_vram_mb": 24576,
"available_vram_mb": 20480,
"used_vram_mb": 4096
}
],
"gpu_count": 1,
"total_vram_mb": 24576,
"available_vram_mb": 20480
},
"studio_endpoints": [
"v1/images/generate",
"v1/audio/tts",
......@@ -215,6 +233,31 @@ CoderAI should send this after receiving the initial AISBF `registered` event.
AISBF replies with a success envelope.
### Hardware Reporting Requirements
The `register` payload should include the best hardware view available to the running CoderAI process.
Required if detectable:
- `hardware.gpus`: array of GPU objects visible to the process
- `hardware.gpu_count`: integer count
- `hardware.total_vram_mb`: total usable VRAM across advertised GPUs
- `hardware.available_vram_mb`: currently free or available VRAM across advertised GPUs
Recommended per GPU fields:
- `index`
- `name`
- `vendor`
- `total_vram_mb`
- `available_vram_mb`
- `used_vram_mb`
- backend-specific extras if trivial to expose
If exact values are unavailable, estimate conservatively and include any supporting marker such as `estimated: true`.
AISBF stores this in broker session metadata so dashboards and future routing logic can reason about available hardware.
## Required Heartbeat Support
AISBF may send heartbeat requests, and CoderAI may also proactively keep the socket alive.
......@@ -261,6 +304,30 @@ CoderAI may also periodically send:
}
```
Heartbeat payloads may also refresh dynamic hardware state such as changing free VRAM:
```json
{
"v": 1,
"op": "heartbeat",
"request_id": "hb-self-2",
"payload": {
"hardware": {
"available_vram_mb": 18432,
"gpus": [
{
"index": 0,
"available_vram_mb": 18432,
"used_vram_mb": 6144
}
]
}
}
}
```
AISBF merges those updates into the broker session metadata.
## Local HTTP Endpoints CoderAI Should Expose
### OpenAI-compatible endpoints
......@@ -427,6 +494,16 @@ Response payload should match `GET /v1/models`.
Request payload matches OpenAI `POST /v1/chat/completions` body.
Completed responses should include performance metrics whenever practical:
- `latency_ms`
- `prompt_tokens`
- `completion_tokens`
- `total_tokens`
- `tokens_per_second`
If exact values are unavailable, AISBF estimates latency from broker timing and may estimate throughput from tokens plus latency.
### `op = "capabilities"`
Response payload should match `GET /coderai/capabilities`.
......@@ -505,6 +582,8 @@ Semantics:
"finish_reason": "stop"
}
],
"latency_ms": 842,
"tokens_per_second": 17.8,
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
......@@ -514,6 +593,20 @@ Semantics:
}
```
AISBF keeps a rolling performance window for the latest 100 completed broker requests per connected session.
When exact metrics are present, AISBF stores them directly. Otherwise it estimates:
- latency from broker request start to final reply time
- throughput from `total_tokens / latency_seconds` when token counts are present
The resulting session snapshot tracks:
- average latency
- average throughput
- average total tokens
- success rate
### Error
```json
......
templates/dashboard/providers.html
View file @ 9bf96583
......@@ -739,9 +739,9 @@ function renderProviderDetails(key) {
}
if (isCoderAIProvider && !provider.coderai_config) {
provider.coderai_config = {
transport: 'http',
transport: 'broker',
broker_enabled: true,
broker_mode: false,
broker_mode: true,
broker_preferred: true,
discovery_enabled: true,
client_id: key,
......@@ -1016,9 +1016,28 @@ function renderProviderDetails(key) {
} else if (isCoderAIProvider) {
authFieldsHtml = `
<div style="background: var(--bg-panel); padding: 15px; border-radius: 5px; margin-bottom: 15px; border-left: 3px solid #4a9eff;">
<h4 style="margin: 0 0 15px 0; color: var(--color-link);">CoderAI Broker Registration</h4>
<h4 style="margin: 0 0 15px 0; color: var(--color-link);">CoderAI Connection Mode</h4>
<small style="color: var(--color-muted); display: block; margin-bottom: 15px;">
This token belongs to this provider configuration. Share it only with the specific CoderAI instance that should register for this provider.
In broker mode, CoderAI connects inbound to AISBF and all traffic stays on that registered broker WebSocket. In direct mode, AISBF uses the configured endpoint as a normal OpenAI-compatible API.
</small>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_mode !== false ? 'checked' : ''} onchange="updateCoderAIBrokerMode('${key}', this.checked)">
<span>Use broker mode</span>
</label>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Recommended for NATed CoderAI instances. Disable this only for directly reachable OpenAI-compatible HTTP endpoints.</small>
</div>
<div class="form-group" style="display:${coderaiConfig.broker_mode === false ? 'block' : 'none'};" id="coderai-api-key-group-${key}">
<label>API Token</label>
<input type="password" value="${provider.api_key || ''}" onchange="updateProvider('${key}', 'api_key', this.value)" placeholder="Enter CoderAI API token for direct HTTP mode">
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Used only when broker mode is disabled, so AISBF can call the endpoint as a normal OpenAI-compatible API.</small>
</div>
<div id="coderai-broker-group-${key}" style="display:${coderaiConfig.broker_mode === false ? 'none' : 'block'};">
<small style="color: var(--color-muted); display: block; margin-bottom: 15px;">
Generate or rotate this registration token here, then copy it into the matching CoderAI configuration so it can authenticate when connecting to AISBF.
</small>
<div class="form-group">
......@@ -1031,34 +1050,13 @@ function renderProviderDetails(key) {
<label>Registration Token</label>
<div style="display:flex; gap:10px; align-items:center; flex-wrap:wrap;">
<input type="password" value="${coderaiConfig.registration_token || ''}" onchange="updateCoderAIConfig('${key}', 'registration_token', this.value)" placeholder="Provider-scoped token" style="flex:1; min-width:320px;">
<button type="button" class="btn btn-secondary" onclick="copyCoderAIToken('${safeKey}')">Copy Token</button>
<button type="button" class="btn btn-secondary" onclick="rotateCoderAIToken('${safeKey}')">Rotate Token</button>
</div>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Generated and stored directly in this provider configuration for this owner. Rotating immediately invalidates the old token for future broker registrations.</small>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Generated and stored directly in this provider configuration for this owner. Rotating invalidates the old token for future CoderAI broker registrations.</small>
<div id="coderai-token-status-${key}" style="margin-top:10px; font-size:13px; color: var(--color-muted);"></div>
</div>
<div class="form-group">
<label>Transport</label>
<select onchange="updateCoderAIConfig('${key}', 'transport', this.value)">
<option value="http" ${coderaiConfig.transport === 'http' || !coderaiConfig.transport ? 'selected' : ''}>HTTP</option>
<option value="websocket" ${coderaiConfig.transport === 'websocket' ? 'selected' : ''}>WebSocket</option>
</select>
</div>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_enabled !== false ? 'checked' : ''} onchange="updateCoderAIConfig('${key}', 'broker_enabled', this.checked)">
<span>Enable AISBF broker sessions</span>
</label>
</div>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_preferred !== false ? 'checked' : ''} onchange="updateCoderAIConfig('${key}', 'broker_preferred', this.checked)">
<span>Prefer outbound broker session over direct HTTP/WebSocket</span>
</label>
</div>
<div class="form-group" style="margin-top:20px; padding:14px; border-radius:6px; background:${brokerConnected ? 'rgba(34,197,94,0.08)' : 'rgba(248,113,113,0.08)'}; border:1px solid ${brokerConnected ? 'rgba(34,197,94,0.35)' : 'rgba(248,113,113,0.35)'};">
<label>Broker Session Status</label>
<div style="font-size:14px; color:var(--color-text); line-height:1.6;">
......@@ -1069,11 +1067,23 @@ function renderProviderDetails(key) {
<div><strong>Last Seen:</strong> ${escHtmlAttr(formatBrokerTimestamp(brokerSession.last_seen))}</div>
<div><strong>Connected At:</strong> ${escHtmlAttr(formatBrokerTimestamp(brokerSession.connected_at))}</div>
<div><strong>Remote Endpoint:</strong> ${escHtmlAttr(brokerSession.endpoint || 'Unknown')}</div>
<div><strong>Transport:</strong> ${escHtmlAttr(brokerSession.transport || coderaiConfig.transport || 'http')}</div>
<div><strong>Transport:</strong> ${escHtmlAttr(brokerSession.transport || 'broker')}</div>
<div><strong>Studio Endpoints:</strong> ${escHtmlAttr((brokerSession.studio_endpoints || []).join(', ') || 'None advertised')}</div>
</div>
</div>
</div>
<div id="coderai-direct-group-${key}" style="display:${coderaiConfig.broker_mode === false ? 'block' : 'none'}; margin-top:20px;">
<div class="form-group">
<label>Direct Transport</label>
<select onchange="updateCoderAITransportMode('${key}', this.value)">
<option value="http" ${coderaiConfig.transport === 'http' || coderaiConfig.broker_mode !== false ? 'selected' : ''}>HTTP</option>
<option value="websocket" ${coderaiConfig.transport === 'websocket' ? 'selected' : ''}>WebSocket</option>
</select>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Use direct mode only when AISBF can reach the CoderAI endpoint itself.</small>
</div>
</div>
</div>
`;
} else if (isQwenProvider) {
// Qwen authentication fields - supports both API key and OAuth2
......@@ -1597,7 +1607,7 @@ function updateNewProviderDefaults() {
'kilocode': 'Kilocode provider. Uses OAuth2 Device Authorization Grant. Endpoint: https://api.kilo.ai/api/gateway',
'qwen': 'Qwen provider. Uses OAuth2 Device Authorization Grant or API key. Endpoint: https://dashscope.aliyuncs.com/compatible-mode/v1',
'codex': 'Codex provider. Uses OAuth2 Device Authorization Grant (same protocol as OpenAI). Endpoint: https://api.openai.com/v1',
'coderai': 'CoderAI provider. Uses local OpenAI-compatible HTTP or AISBF broker/WebSocket registration with a provider-scoped token. Endpoint: http://127.0.0.1:11437'
'coderai': 'CoderAI provider. In broker mode, CoderAI connects inbound to AISBF using a provider-scoped registration token. In direct mode, AISBF calls it as an OpenAI-compatible endpoint. Default endpoint: http://127.0.0.1:11437'
};
descriptionEl.textContent = descriptions[providerType] || window.i18n.t('providers.standard_config');
......@@ -2323,6 +2333,49 @@ function updateCoderAIConfig(key, field, value) {
providersData[key].coderai_config[field] = value;
}
function updateCoderAIBrokerMode(key, enabled) {
if (!providersData[key].coderai_config) {
providersData[key].coderai_config = {};
}
providersData[key].coderai_config.broker_mode = enabled;
providersData[key].coderai_config.broker_enabled = enabled;
providersData[key].coderai_config.broker_preferred = enabled;
providersData[key].coderai_config.transport = enabled ? 'broker' : ((providersData[key].coderai_config.transport === 'websocket') ? 'websocket' : 'http');
renderProviderDetails(key);
}
function updateCoderAITransportMode(key, value) {
if (!providersData[key].coderai_config) {
providersData[key].coderai_config = {};
}
providersData[key].coderai_config.transport = value;
renderProviderDetails(key);
}
async function copyCoderAIToken(key) {
const statusEl = document.getElementById(`coderai-token-status-${key}`);
const token = (((providersData[key] || {}).coderai_config || {}).registration_token || '').trim();
if (!token) {
if (statusEl) {
statusEl.textContent = 'No token to copy';
statusEl.style.color = '#f87171';
}
return;
}
try {
await navigator.clipboard.writeText(token);
if (statusEl) {
statusEl.textContent = 'Token copied to clipboard';
statusEl.style.color = '#4ade80';
}
} catch (e) {
if (statusEl) {
statusEl.textContent = 'Copy failed: ' + e.message;
statusEl.style.color = '#f87171';
}
}
}
async function rotateCoderAIToken(key) {
const statusEl = document.getElementById(`coderai-token-status-${key}`);
if (statusEl) {
......@@ -2452,9 +2505,9 @@ async function confirmAddProvider() {
};
} else if (type === 'coderai') {
providersData[key].coderai_config = {
transport: 'http',
transport: 'broker',
broker_enabled: true,
broker_mode: false,
broker_mode: true,
broker_preferred: true,
discovery_enabled: true,
client_id: key,
......
templates/dashboard/user_providers.html
View file @ 9bf96583
......@@ -784,9 +784,9 @@ function renderProviderDetails(key) {
}
if (isCoderAIProvider && !provider.coderai_config) {
provider.coderai_config = {
transport: 'http',
transport: 'broker',
broker_enabled: true,
broker_mode: false,
broker_mode: true,
broker_preferred: true,
discovery_enabled: true,
client_id: key,
......@@ -1025,9 +1025,28 @@ function renderProviderDetails(key) {
} else if (isCoderAIProvider) {
authFieldsHtml = `
<div style="background: var(--bg-panel); padding: 15px; border-radius: 5px; margin-bottom: 15px; border-left: 3px solid #4a9eff;">
<h4 style="margin: 0 0 15px 0; color: var(--color-link);">CoderAI Broker Registration</h4>
<h4 style="margin: 0 0 15px 0; color: var(--color-link);">CoderAI Connection Mode</h4>
<small style="color: var(--color-muted); display: block; margin-bottom: 15px;">
This token belongs to your provider configuration. Share it only with the matching CoderAI instance for this provider.
In broker mode, CoderAI connects inbound to AISBF and all traffic stays on that registered broker WebSocket. In direct mode, AISBF uses the configured endpoint as a normal OpenAI-compatible API.
</small>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_mode !== false ? 'checked' : ''} onchange="updateCoderAIBrokerMode('${key}', this.checked)">
<span>Use broker mode</span>
</label>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Recommended for NATed CoderAI instances. Disable this only for directly reachable OpenAI-compatible HTTP endpoints.</small>
</div>
<div class="form-group" style="display:${coderaiConfig.broker_mode === false ? 'block' : 'none'};" id="coderai-api-key-group-${key}">
<label>API Token</label>
<input type="password" value="${provider.api_key || ''}" onchange="updateProvider('${key}', 'api_key', this.value)" placeholder="Enter CoderAI API token for direct HTTP mode">
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Used only when broker mode is disabled, so AISBF can call the endpoint as a normal OpenAI-compatible API.</small>
</div>
<div id="coderai-broker-group-${key}" style="display:${coderaiConfig.broker_mode === false ? 'none' : 'block'};">
<small style="color: var(--color-muted); display: block; margin-bottom: 15px;">
Generate or rotate this registration token here, then copy it into the matching CoderAI configuration so it can authenticate when connecting to AISBF.
</small>
<div class="form-group">
......@@ -1040,34 +1059,13 @@ function renderProviderDetails(key) {
<label>Registration Token</label>
<div style="display:flex; gap:10px; align-items:center; flex-wrap:wrap;">
<input type="password" value="${coderaiConfig.registration_token || ''}" onchange="updateCoderAIConfig('${key}', 'registration_token', this.value)" placeholder="Provider-scoped token" style="flex:1; min-width:320px;">
<button type="button" class="btn btn-secondary" onclick="copyCoderAIToken('${safeKey}')">Copy Token</button>
<button type="button" class="btn btn-secondary" onclick="rotateCoderAIToken('${safeKey}')">Rotate Token</button>
</div>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Stored directly in this provider configuration for your account. Rotating immediately invalidates the old token for future broker registrations.</small>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Stored directly in this provider configuration for your account. Rotating invalidates the old token for future CoderAI broker registrations.</small>
<div id="coderai-token-status-${key}" style="margin-top:10px; font-size:13px; color: var(--color-muted);"></div>
</div>
<div class="form-group">
<label>Transport</label>
<select onchange="updateCoderAIConfig('${key}', 'transport', this.value)">
<option value="http" ${coderaiConfig.transport === 'http' || !coderaiConfig.transport ? 'selected' : ''}>HTTP</option>
<option value="websocket" ${coderaiConfig.transport === 'websocket' ? 'selected' : ''}>WebSocket</option>
</select>
</div>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_enabled !== false ? 'checked' : ''} onchange="updateCoderAIConfig('${key}', 'broker_enabled', this.checked)">
<span>Enable AISBF broker sessions</span>
</label>
</div>
<div class="form-group">
<label style="display:flex; align-items:center; gap:10px; cursor:pointer;">
<input type="checkbox" ${coderaiConfig.broker_preferred !== false ? 'checked' : ''} onchange="updateCoderAIConfig('${key}', 'broker_preferred', this.checked)">
<span>Prefer outbound broker session when available</span>
</label>
</div>
<div class="form-group" style="margin-top:20px; padding:14px; border-radius:6px; background:${brokerConnected ? 'rgba(34,197,94,0.08)' : 'rgba(248,113,113,0.08)'}; border:1px solid ${brokerConnected ? 'rgba(34,197,94,0.35)' : 'rgba(248,113,113,0.35)'};">
<label>Broker Session Status</label>
<div style="font-size:14px; color:var(--color-text); line-height:1.6;">
......@@ -1078,11 +1076,23 @@ function renderProviderDetails(key) {
<div><strong>Last Seen:</strong> ${escHtmlAttr(formatBrokerTimestamp(brokerSession.last_seen))}</div>
<div><strong>Connected At:</strong> ${escHtmlAttr(formatBrokerTimestamp(brokerSession.connected_at))}</div>
<div><strong>Remote Endpoint:</strong> ${escHtmlAttr(brokerSession.endpoint || 'Unknown')}</div>
<div><strong>Transport:</strong> ${escHtmlAttr(brokerSession.transport || coderaiConfig.transport || 'http')}</div>
<div><strong>Transport:</strong> ${escHtmlAttr(brokerSession.transport || 'broker')}</div>
<div><strong>Studio Endpoints:</strong> ${escHtmlAttr((brokerSession.studio_endpoints || []).join(', ') || 'None advertised')}</div>
</div>
</div>
</div>
<div id="coderai-direct-group-${key}" style="display:${coderaiConfig.broker_mode === false ? 'block' : 'none'}; margin-top:20px;">
<div class="form-group">
<label>Direct Transport</label>
<select onchange="updateCoderAITransportMode('${key}', this.value)">
<option value="http" ${coderaiConfig.transport === 'http' || coderaiConfig.broker_mode !== false ? 'selected' : ''}>HTTP</option>
<option value="websocket" ${coderaiConfig.transport === 'websocket' ? 'selected' : ''}>WebSocket</option>
</select>
<small style="color: var(--color-muted); display: block; margin-top: 5px;">Use direct mode only when AISBF can reach the CoderAI endpoint itself.</small>
</div>
</div>
</div>
`;
} else if (isQwenProvider) {
// Qwen authentication fields - supports both API key and OAuth2
......@@ -1617,7 +1627,7 @@ function updateNewProviderDefaults() {
'kilocode': 'Kilocode provider. Uses OAuth2 Device Authorization Grant. Endpoint: https://api.kilo.ai/api/gateway',
'qwen': 'Qwen provider. Uses OAuth2 Device Authorization Grant or API key. Endpoint: https://dashscope.aliyuncs.com/compatible-mode/v1',
'codex': 'Codex provider. Uses OAuth2 Device Authorization Grant (same protocol as OpenAI). Endpoint: https://api.openai.com/v1',
'coderai': 'CoderAI provider. Uses local OpenAI-compatible HTTP or AISBF broker/WebSocket registration with a provider-scoped token. Endpoint: http://127.0.0.1:11437'
'coderai': 'CoderAI provider. In broker mode, CoderAI connects inbound to AISBF using a provider-scoped registration token. In direct mode, AISBF calls it as an OpenAI-compatible endpoint. Default endpoint: http://127.0.0.1:11437'
};
descriptionEl.textContent = descriptions[providerType] || window.i18n.t('providers.standard_config');
......@@ -2339,6 +2349,49 @@ function updateCoderAIConfig(key, field, value) {
providersData[key].coderai_config[field] = value;
}
function updateCoderAIBrokerMode(key, enabled) {
if (!providersData[key].coderai_config) {
providersData[key].coderai_config = {};
}
providersData[key].coderai_config.broker_mode = enabled;
providersData[key].coderai_config.broker_enabled = enabled;
providersData[key].coderai_config.broker_preferred = enabled;
providersData[key].coderai_config.transport = enabled ? 'broker' : ((providersData[key].coderai_config.transport === 'websocket') ? 'websocket' : 'http');
renderProviderDetails(key);
}
function updateCoderAITransportMode(key, value) {
if (!providersData[key].coderai_config) {
providersData[key].coderai_config = {};
}
providersData[key].coderai_config.transport = value;
renderProviderDetails(key);
}
async function copyCoderAIToken(key) {
const statusEl = document.getElementById(`coderai-token-status-${key}`);
const token = (((providersData[key] || {}).coderai_config || {}).registration_token || '').trim();
if (!token) {
if (statusEl) {
statusEl.textContent = 'No token to copy';
statusEl.style.color = '#f87171';
}
return;
}
try {
await navigator.clipboard.writeText(token);
if (statusEl) {
statusEl.textContent = 'Token copied to clipboard';
statusEl.style.color = '#4ade80';
}
} catch (e) {
if (statusEl) {
statusEl.textContent = 'Copy failed: ' + e.message;
statusEl.style.color = '#f87171';
}
}
}
async function rotateCoderAIToken(key) {
const statusEl = document.getElementById(`coderai-token-status-${key}`);
if (statusEl) {
......@@ -2461,9 +2514,9 @@ async function confirmAddProvider() {
};
} else if (type === 'coderai') {
providersData[key].coderai_config = {
transport: 'http',
transport: 'broker',
broker_enabled: true,
broker_mode: false,
broker_mode: true,
broker_preferred: true,
discovery_enabled: true,
client_id: key,
......
tests/providers/test_coderai_provider.py
View file @ 9bf96583
import json
import base64
from unittest.mock import Mock
from types import SimpleNamespace
import pytest
from aisbf.providers.coderai import CoderAIProviderHandler
from aisbf.config import config
from aisbf.config import ProviderConfig
from aisbf.app.model_cache import _is_broker_only_coderai
@pytest.fixture(autouse=True)
......@@ -190,3 +192,56 @@ async def test_coderai_broker_stream_supports_progress_and_binary_chunks(monkeyp
assert status_code == 200
assert payload["stream_encoding"] == "base64"
assert base64.b64decode(payload["stream_chunks"][0]).startswith(b"event: progress")
def test_coderai_broker_mode_forces_broker_transport_and_skips_outbound_validation():
provider_config = {
"id": "coderai_nat",
"name": "CoderAI NAT",
"endpoint": "http://127.0.0.1:11437",
"type": "coderai",
"api_key_required": False,
"coderai_config": {"broker_mode": True, "registration_token": "nat-token"},
}
handler = CoderAIProviderHandler("coderai_nat", provider_config=provider_config)
assert handler._transport == "broker"
assert handler.validate_credentials() is True
@pytest.mark.asyncio
async def test_coderai_non_broker_mode_uses_openai_compatible_http_api():
provider_config = {
"id": "coderai_local",
"name": "CoderAI",
"endpoint": "http://127.0.0.1:11437",
"type": "coderai",
"api_key_required": False,
"api_key": "local-api-token",
"coderai_config": {"broker_mode": False, "broker_enabled": False, "transport": "http"},
}
handler = CoderAIProviderHandler("coderai_local", provider_config=provider_config)
assert handler._transport == "http"
assert handler._effective_api_key() == "local-api-token"
assert handler.validate_credentials() is True
def test_model_cache_detects_broker_only_coderai_provider():
provider_config = SimpleNamespace(
type="coderai",
coderai_config={"broker_mode": True},
)
assert _is_broker_only_coderai(provider_config) is True
def test_model_cache_allows_direct_http_coderai_provider():
provider_config = SimpleNamespace(
type="coderai",
coderai_config={"broker_mode": False, "broker_enabled": False, "transport": "http"},
)
assert _is_broker_only_coderai(provider_config) is False
tests/routes/test_dashboard_coderai.py
View file @ 9bf96583
......@@ -149,6 +149,27 @@ def test_api_provider_save_rejects_coderai_broker_without_registration_token(mon
assert "requires a registration token" in str(exc)
def test_api_provider_save_allows_non_broker_coderai_without_registration_token():
provider_config = dashboard_providers._ensure_coderai_token(
{
"id": "coderai-http",
"name": "CoderAI HTTP",
"endpoint": "http://127.0.0.1:11437",
"api_key_required": False,
"api_key": "local-token",
"type": "coderai",
"coderai_config": {
"broker_enabled": False,
"broker_mode": False,
"registration_token": "",
},
}
)
dashboard_providers._validate_coderai_provider_config("coderai-http", provider_config)
assert provider_config["coderai_config"]["registration_token"] == ""
def test_api_provider_save_persists_coderai_token_for_user(monkeypatch):
db = DbStub()
......
tests/test_coderai_broker.py
View file @ 9bf96583
......@@ -54,6 +54,15 @@ def test_broker_registers_websocket_session_and_reports_status():
"endpoint": "ws://local-tunnel",
"transport": "websocket",
"registration_token": "global-token",
"hardware": {
"gpus": [
{
"name": "RTX 4090",
"total_vram_mb": 24576,
"available_vram_mb": 20480
}
]
},
"capabilities": {"studio": {"enabled": True}},
},
})
......@@ -65,6 +74,9 @@ def test_broker_registers_websocket_session_and_reports_status():
payload = response.json()
assert payload["connected"] is True
assert payload["client_id"] == "nat-client"
assert payload["metadata"]["gpu_count"] == 1
assert payload["metadata"]["total_vram_mb"] == 24576
assert payload["metadata"]["available_vram_mb"] == 20480
def test_broker_rejects_missing_registration_token():
......@@ -83,7 +95,7 @@ def test_broker_routes_request_to_registered_session():
async def send_text(self, payload: str):
self.sent.append(payload)
message = json.loads(payload)
await broker.resolve_response({
await broker.publish_response({
"request_id": message["request_id"],
"status": "ok",
"payload": {"data": [{"id": "llama3.1:8b"}]},
......@@ -97,6 +109,9 @@ def test_broker_routes_request_to_registered_session():
assert sent_message["op"] == "models.list"
assert response["status"] == "ok"
assert response["payload"]["data"][0]["id"] == "llama3.1:8b"
snapshot = await broker.get_session_snapshot("coderai", "bridge-client")
assert snapshot["performance"]["sample_count"] == 1
assert snapshot["performance"]["avg_latency_ms"] >= 0
queued = await broker.consume_request(session.session_id, timeout=0)
assert queued is None
finally:
......@@ -163,6 +178,8 @@ def test_broker_stream_events_are_delivered_to_waiter():
try:
response = await broker.send_request("coderai", "proxy", {"stream": True}, client_id="stream-client", owner_user_id=None, timeout=3.0)
assert response["event"] == "done"
snapshot = await broker.get_session_snapshot("coderai", "stream-client")
assert snapshot["performance"]["sample_count"] == 1
finally:
await broker.unregister(session.session_id)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment