mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-26 15:40:18 +00:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
17434d6499 | ||
|
|
3f58ed6232 | ||
|
|
097a26ace6 | ||
|
|
e87d7bd066 | ||
|
|
e42b1e5ec1 | ||
|
|
ccf60055f4 |
@@ -7,6 +7,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.7.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.6.0...ironclaw-v0.7.0) - 2026-02-19
|
||||
|
||||
### Added
|
||||
|
||||
- extend lifecycle hooks with declarative bundles ([#176](https://github.com/nearai/ironclaw/pull/176))
|
||||
- support per-request model override in /v1/chat/completions ([#103](https://github.com/nearai/ironclaw/pull/103))
|
||||
|
||||
### Fixed
|
||||
|
||||
- harden openai-compatible provider, approval replay, and embeddings defaults ([#237](https://github.com/nearai/ironclaw/pull/237))
|
||||
- Network Security Findings ([#201](https://github.com/nearai/ironclaw/pull/201))
|
||||
|
||||
### Added
|
||||
|
||||
- Refactored OpenAI-compatible chat completion routing to use the rig adapter and `RetryProvider` composition for custom base URL usage.
|
||||
- Added Ollama embeddings provider support (`EMBEDDING_PROVIDER=ollama`, `OLLAMA_BASE_URL`) in workspace embeddings.
|
||||
- Added migration `V9__flexible_embedding_dimension.sql` for flexible embedding vector dimensions.
|
||||
|
||||
### Changed
|
||||
|
||||
- Changed default sandbox image to `ironclaw-worker:latest` in config/settings/sandbox defaults.
|
||||
- Improved tool-message sanitization and provider compatibility handling across NEAR AI, rig adapter, and shared LLM provider code.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed approval-input aliases (`a`, `/approve`, `/always`, `/deny`, etc.) in submission parsing.
|
||||
- Fixed multi-tool approval resume flow by preserving and replaying deferred tool calls so all prior `tool_use` IDs receive matching `tool_result` messages.
|
||||
- Fixed REPL quit/exit handling to route shutdown through the agent loop for graceful termination.
|
||||
|
||||
## [0.6.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.5.0...ironclaw-v0.6.0) - 2026-02-19
|
||||
|
||||
### Added
|
||||
@@ -94,6 +123,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
- Bump MSRV to 1.92, add GCP deployment files ([#40](https://github.com/nearai/ironclaw/pull/40))
|
||||
- Add OpenAI-compatible HTTP API (/v1/chat/completions, /v1/models) ([#31](https://github.com/nearai/ironclaw/pull/31))
|
||||
|
||||
|
||||
## [0.1.3](https://github.com/nearai/ironclaw/compare/v0.1.2...v0.1.3) - 2026-02-12
|
||||
|
||||
### Other
|
||||
|
||||
Generated
+1
-1
@@ -2490,7 +2490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "ironclaw"
|
||||
version = "0.6.0"
|
||||
version = "0.7.0"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"aho-corasick",
|
||||
|
||||
+2
-2
@@ -9,7 +9,7 @@ exclude = [
|
||||
|
||||
[package]
|
||||
name = "ironclaw"
|
||||
version = "0.6.0"
|
||||
version = "0.7.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.92"
|
||||
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
||||
@@ -78,7 +78,7 @@ termimad = "0.34"
|
||||
# Channel integrations
|
||||
axum = { version = "0.8", features = ["ws"] }
|
||||
tower = "0.5"
|
||||
tower-http = { version = "0.6", features = ["trace", "cors"] }
|
||||
tower-http = { version = "0.6", features = ["trace", "cors", "set-header"] }
|
||||
|
||||
# Cron scheduling for routines
|
||||
cron = "0.13"
|
||||
|
||||
+7
-7
@@ -37,7 +37,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Session management/routing | ✅ | ✅ | SessionManager exists |
|
||||
| Configuration hot-reload | ✅ | ❌ | |
|
||||
| Network modes (loopback/LAN/remote) | ✅ | 🚧 | HTTP only |
|
||||
| OpenAI-compatible HTTP API | ✅ | ✅ | /v1/chat/completions |
|
||||
| OpenAI-compatible HTTP API | ✅ | ✅ | /v1/chat/completions, per-request `model` override |
|
||||
| Canvas hosting | ✅ | ❌ | Agent-driven UI |
|
||||
| Gateway lock (PID-based) | ✅ | ❌ | |
|
||||
| launchd/systemd integration | ✅ | ❌ | |
|
||||
@@ -278,7 +278,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| Auth plugins | ✅ | ❌ | |
|
||||
| Memory plugins | ✅ | ❌ | Custom backends |
|
||||
| Tool plugins | ✅ | ✅ | WASM tools |
|
||||
| Hook plugins | ✅ | ❌ | |
|
||||
| Hook plugins | ✅ | ✅ | Declarative hooks from extension capabilities |
|
||||
| Provider plugins | ✅ | ❌ | |
|
||||
| Plugin CLI (`install`, `list`) | ✅ | ✅ | `tool` subcommand |
|
||||
| ClawHub registry | ✅ | ❌ | Discovery |
|
||||
@@ -421,10 +421,10 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
| `transcribeAudio` hook | ✅ | ❌ | P3 | |
|
||||
| `transformResponse` hook | ✅ | ✅ | P2 | |
|
||||
| `llm_input`/`llm_output` hooks | ✅ | ❌ | P3 | LLM payload inspection |
|
||||
| Bundled hooks | ✅ | ❌ | P2 | |
|
||||
| Plugin hooks | ✅ | ❌ | P3 | |
|
||||
| Workspace hooks | ✅ | ❌ | P2 | Inline code |
|
||||
| Outbound webhooks | ✅ | ❌ | P2 | |
|
||||
| Bundled hooks | ✅ | ✅ | P2 | Audit + declarative rule/webhook hooks |
|
||||
| Plugin hooks | ✅ | ✅ | P3 | Registered from WASM `capabilities.json` |
|
||||
| Workspace hooks | ✅ | ✅ | P2 | `hooks/hooks.json` and `hooks/*.hook.json` |
|
||||
| Outbound webhooks | ✅ | ✅ | P2 | Fire-and-forget lifecycle event delivery |
|
||||
| Heartbeat system | ✅ | ✅ | - | Periodic execution |
|
||||
| Gmail pub/sub | ✅ | ❌ | P3 | |
|
||||
|
||||
@@ -528,7 +528,7 @@ This document tracks feature parity between IronClaw (Rust implementation) and O
|
||||
- ✅ Telegram channel (WASM, DM pairing, caption, /start)
|
||||
- ❌ WhatsApp channel
|
||||
- ✅ Multi-provider failover (`FailoverProvider` with retryable error classification)
|
||||
- ✅ Hooks system (beforeInbound, beforeToolCall, beforeOutbound, onSessionStart, onSessionEnd, transformResponse)
|
||||
- ✅ Hooks system (core lifecycle hooks + bundled/plugin/workspace hooks + outbound webhooks)
|
||||
|
||||
### P2 - Medium Priority
|
||||
- ❌ Media handling (images, PDFs)
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
-- Allow embedding vectors of any dimension (not just 1536).
|
||||
-- This supports Ollama models (768-dim nomic-embed-text, 1024-dim mxbai-embed-large)
|
||||
-- alongside OpenAI models (1536-dim text-embedding-3-small, 3072-dim text-embedding-3-large).
|
||||
--
|
||||
-- NOTE: HNSW indexes require a fixed dimension, so we drop the index.
|
||||
-- Exact (sequential) cosine distance search still works without the index.
|
||||
-- For a personal assistant workspace the dataset is small enough that this
|
||||
-- has negligible impact on query latency.
|
||||
|
||||
-- Drop dependent views first
|
||||
DROP VIEW IF EXISTS chunks_pending_embedding;
|
||||
DROP VIEW IF EXISTS memory_documents_summary;
|
||||
|
||||
DROP INDEX IF EXISTS idx_memory_chunks_embedding;
|
||||
|
||||
ALTER TABLE memory_chunks
|
||||
ALTER COLUMN embedding TYPE vector
|
||||
USING embedding::vector;
|
||||
|
||||
-- Recreate the views
|
||||
CREATE VIEW memory_documents_summary AS
|
||||
SELECT
|
||||
d.id,
|
||||
d.user_id,
|
||||
d.path,
|
||||
d.created_at,
|
||||
d.updated_at,
|
||||
COUNT(c.id) as chunk_count,
|
||||
COUNT(c.embedding) as embedded_chunk_count
|
||||
FROM memory_documents d
|
||||
LEFT JOIN memory_chunks c ON c.document_id = d.id
|
||||
GROUP BY d.id;
|
||||
|
||||
CREATE VIEW chunks_pending_embedding AS
|
||||
SELECT
|
||||
c.id as chunk_id,
|
||||
c.document_id,
|
||||
d.user_id,
|
||||
d.path,
|
||||
LENGTH(c.content) as content_length
|
||||
FROM memory_chunks c
|
||||
JOIN memory_documents d ON d.id = c.document_id
|
||||
WHERE c.embedding IS NULL;
|
||||
@@ -0,0 +1,566 @@
|
||||
# IronClaw Network Security Reference
|
||||
|
||||
This document catalogs every network-facing surface in IronClaw, its authentication mechanism, bind address, security controls, and known findings. Use this as the authoritative reference during code reviews that touch network-facing code.
|
||||
|
||||
**Last updated:** 2026-02-18
|
||||
|
||||
---
|
||||
|
||||
## Threat Model
|
||||
|
||||
IronClaw operates across four trust boundaries:
|
||||
|
||||
| Boundary | Trust Level | Examples |
|
||||
|----------|------------|---------|
|
||||
| **Local user** | Fully trusted | TUI, web gateway (loopback), CLI commands |
|
||||
| **Browser client** | Authenticated | Web UI connected via bearer token; subject to CORS, Origin validation, CSRF protections |
|
||||
| **Docker containers** | Untrusted (sandboxed) | Worker containers executing user jobs; isolated via per-job tokens, allowlisted egress, dropped capabilities |
|
||||
| **External services** | Untrusted | Webhook senders (Telegram, Slack); authenticated via shared secret |
|
||||
|
||||
**Key assumptions:**
|
||||
|
||||
- The local machine is single-user. The web gateway and OAuth listener bind to loopback and do not defend against other local users.
|
||||
- Docker containers are adversarial. A compromised container should not be able to access other jobs, exfiltrate secrets, or reach the host network beyond the orchestrator API.
|
||||
- Webhook senders must prove knowledge of the shared secret. The secret is never transmitted in the clear by IronClaw itself.
|
||||
- MCP server URLs are operator-configured and treated as trusted destinations (see [MCP Client](#mcp-client)).
|
||||
|
||||
---
|
||||
|
||||
## Network Surface Inventory
|
||||
|
||||
| Listener | Default Port | Default Bind | Auth Mechanism | Config Env Var | Source |
|
||||
|----------|-------------|-------------|----------------|----------------|--------|
|
||||
| Web Gateway | 3000 | `127.0.0.1` | Bearer token (constant-time) | `GATEWAY_HOST`, `GATEWAY_PORT`, `GATEWAY_AUTH_TOKEN` | `server.rs` — `start_server()` |
|
||||
| HTTP Webhook Server | 8080 | `0.0.0.0` | Shared secret (body field) | `HTTP_HOST`, `HTTP_PORT`, `HTTP_WEBHOOK_SECRET` | `webhook_server.rs` — `start()` |
|
||||
| Orchestrator Internal API | 50051 | `127.0.0.1` (macOS/Win) / `0.0.0.0` (Linux) | Per-job bearer token (constant-time) | `ORCHESTRATOR_PORT` | `api.rs` — `OrchestratorApi::start()` |
|
||||
| OAuth Callback Listener | 9876 | `127.0.0.1` | None (ephemeral, 5-min timeout) | N/A (hardcoded) | `oauth_defaults.rs` — `bind_callback_listener()` |
|
||||
| Sandbox HTTP Proxy | OS-assigned (ephemeral) | `127.0.0.1` | None (loopback only) | N/A (auto-assigned) | `proxy/http.rs` — `SandboxProxy::start()` |
|
||||
|
||||
---
|
||||
|
||||
## 1. Web Gateway
|
||||
|
||||
**Source:** `src/channels/web/server.rs`, `src/channels/web/auth.rs`
|
||||
|
||||
### Bind Address
|
||||
|
||||
Configurable via `GATEWAY_HOST` (default `127.0.0.1`) and `GATEWAY_PORT` (default `3000`). The gateway is designed as a local-first, single-user service.
|
||||
|
||||
**Reference:** `src/config.rs` — `gateway_host` default (`"127.0.0.1"`), `gateway_port` default (`3000`)
|
||||
|
||||
### Authentication
|
||||
|
||||
Bearer token middleware applied to all `/api/*` routes via `route_layer`. Token checked in two locations:
|
||||
|
||||
1. `Authorization: Bearer <token>` header (primary)
|
||||
2. `?token=<token>` query parameter (fallback for SSE `EventSource` which cannot set headers)
|
||||
|
||||
Both paths use **constant-time comparison** via `subtle::ConstantTimeEq` (`ct_eq`).
|
||||
|
||||
**Reference:** `src/channels/web/auth.rs` — `auth_middleware()`, header check and query-param fallback both use `ct_eq`
|
||||
|
||||
If `GATEWAY_AUTH_TOKEN` is not set, a random hex token is generated at startup.
|
||||
|
||||
### Unauthenticated Routes
|
||||
|
||||
| Route | Purpose | Response |
|
||||
|-------|---------|----------|
|
||||
| `/api/health` | Health check endpoint | `{"status":"healthy","channel":"gateway"}` — no version, uptime, or fingerprinting data |
|
||||
| `/` | Static HTML (embedded) | Single-page app shell |
|
||||
| `/style.css` | Static CSS (embedded) | Stylesheet |
|
||||
| `/app.js` | Static JS (embedded) | Client-side app |
|
||||
|
||||
### CORS Policy
|
||||
|
||||
Restricted to a two-origin allowlist (not browser same-origin policy, but a CORS allowlist that achieves equivalent protection):
|
||||
|
||||
- `http://<bind_ip>:<bind_port>`
|
||||
- `http://localhost:<bind_port>`
|
||||
|
||||
Allowed methods: `GET`, `POST`, `PUT`, `DELETE`. Allowed headers: `Content-Type`, `Authorization`. Credentials allowed.
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — `CorsLayer::new()` block
|
||||
|
||||
### WebSocket Origin Validation
|
||||
|
||||
The `/api/chat/ws` endpoint has two layers of protection:
|
||||
|
||||
1. **Bearer token auth** — the route is inside the `protected` router with `route_layer`, so `auth_middleware` runs before the handler. The token is passed via the `Authorization: Bearer` header on the HTTP upgrade request (not via query parameter).
|
||||
|
||||
2. **Origin header validation** (inside the handler) as a defense-in-depth guard against cross-site WebSocket hijacking (CSWSH):
|
||||
- Origin header is **required** — missing Origin returns 403 (browsers always send it for WS upgrades; absence implies a non-browser client)
|
||||
- Origin host is extracted by stripping scheme and port, then compared **exactly** against `localhost`, `127.0.0.1`, and `[::1]`
|
||||
- Partial matches like `localhost.evil.com` are rejected because the check extracts the host portion before the first `:` or `/`
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — `chat_ws_handler()` (origin validation block)
|
||||
|
||||
### Rate Limiting
|
||||
|
||||
Chat endpoint (`/api/chat/send`) enforces a sliding-window rate limit: **30 requests per 60 seconds** (global, not per-IP — single-user gateway).
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — `RateLimiter` struct, `chat_rate_limiter` field
|
||||
|
||||
### Body Limits
|
||||
|
||||
- Global: **1 MB** max request body (`DefaultBodyLimit::max(1024 * 1024)`)
|
||||
- **Reference:** `src/channels/web/server.rs` — `.layer(DefaultBodyLimit::max(...))`
|
||||
|
||||
### Project File Serving
|
||||
|
||||
The `/projects/{project_id}/*` routes serve files from project directories. These are **behind auth middleware** to prevent unauthorized file access.
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — project file routes in `protected` router
|
||||
|
||||
### Security Headers
|
||||
|
||||
The gateway sets the following security headers on all responses (via `SetResponseHeaderLayer::if_not_present`, so handlers can override):
|
||||
|
||||
- `X-Content-Type-Options: nosniff` — prevents MIME-sniffing
|
||||
- `X-Frame-Options: DENY` — prevents clickjacking via iframes
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — `SetResponseHeaderLayer` calls
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
Shutdown is triggered via a `oneshot::Sender` stored in `GatewayState::shutdown_tx`. The server uses `axum::serve(...).with_graceful_shutdown(...)` to drain in-flight requests before closing the listener.
|
||||
|
||||
**Reference:** `src/channels/web/server.rs` — `shutdown_tx` / `shutdown_rx` setup
|
||||
|
||||
---
|
||||
|
||||
## 2. HTTP Webhook Server
|
||||
|
||||
**Source:** `src/channels/webhook_server.rs`, `src/channels/http.rs`
|
||||
|
||||
### Bind Address
|
||||
|
||||
Configurable via `HTTP_HOST` (default `0.0.0.0`) and `HTTP_PORT` (default `8080`).
|
||||
|
||||
**WARNING:** The default bind address is `0.0.0.0`, meaning the webhook server listens on **all interfaces** by default. This is intentional (webhooks must be reachable from external services like Telegram/Slack), but operators should be aware of the exposure.
|
||||
|
||||
**Reference:** `src/config.rs` — `http_host` default (`"0.0.0.0"`), `http_port` default (`8080`)
|
||||
|
||||
### Authentication
|
||||
|
||||
Webhook secret is passed **in the JSON request body** (`secret` field), not as a header. The secret is compared using **constant-time** `subtle::ConstantTimeEq` (`ct_eq`).
|
||||
|
||||
The secret is required to start the channel — if `HTTP_WEBHOOK_SECRET` is not set, `start()` returns an error.
|
||||
|
||||
**CSRF note:** Because the secret is in the JSON body (not a cookie or header that browsers auto-attach), a cross-origin form POST cannot forge a valid request. Browsers would send `application/x-www-form-urlencoded`, which the `Json<T>` extractor rejects with HTTP 415. Even if `Content-Type` were spoofed via CORS preflight, the attacker would need the secret value, which is never stored in the browser.
|
||||
|
||||
**Reference:** `src/channels/http.rs` — `webhook_handler()` (secret validation with `ct_eq`), `start()` (required-secret check)
|
||||
|
||||
### Content-Type Validation
|
||||
|
||||
The webhook endpoint uses axum's `Json<WebhookRequest>` extractor, which enforces `Content-Type: application/json`. Requests with missing or incorrect Content-Type are rejected with **HTTP 415 Unsupported Media Type** before the handler body executes. Malformed JSON bodies are rejected with **HTTP 422 Unprocessable Entity**.
|
||||
|
||||
**Reference:** `src/channels/http.rs` — `webhook_handler()` function signature (`Json(req): Json<WebhookRequest>`)
|
||||
|
||||
### Rate Limiting
|
||||
|
||||
**60 requests per minute**, enforced via a mutex-protected sliding window.
|
||||
|
||||
**Reference:** `src/channels/http.rs` — `MAX_REQUESTS_PER_MINUTE` constant, rate-limit check in `webhook_handler()`
|
||||
|
||||
### Body Limits
|
||||
|
||||
- JSON body: **64 KB** max (`MAX_BODY_BYTES`)
|
||||
- Message content: **32 KB** max (`MAX_CONTENT_BYTES`)
|
||||
- Pending synchronous responses: **100 max** (`MAX_PENDING_RESPONSES`)
|
||||
- Synchronous response timeout: **60 seconds**
|
||||
|
||||
**Reference:** `src/channels/http.rs` — constants block (`MAX_BODY_BYTES`, `MAX_CONTENT_BYTES`, `MAX_PENDING_RESPONSES`, `MAX_REQUESTS_PER_MINUTE`)
|
||||
|
||||
### Routes
|
||||
|
||||
| Route | Auth | Purpose | Response |
|
||||
|-------|------|---------|----------|
|
||||
| `/health` | None | Health check | `{"status":"healthy","channel":"http"}` — no fingerprinting data |
|
||||
| `/webhook` | Webhook secret | Receive messages | Webhook response |
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
Shutdown is triggered via a `oneshot::Sender` stored on the `WebhookServer` struct. The server uses `axum::serve(...).with_graceful_shutdown(...)`. The public `shutdown()` method sends the signal and awaits the task join handle, ensuring a clean drain-and-wait.
|
||||
|
||||
**Reference:** `src/channels/webhook_server.rs` — `shutdown()` method
|
||||
|
||||
---
|
||||
|
||||
## 3. Orchestrator Internal API
|
||||
|
||||
**Source:** `src/orchestrator/api.rs`, `src/orchestrator/auth.rs`
|
||||
|
||||
### Bind Address
|
||||
|
||||
Platform-dependent:
|
||||
|
||||
- **macOS / Windows**: `127.0.0.1:<port>` — Docker Desktop routes `host.docker.internal` through its VM to `127.0.0.1`
|
||||
- **Linux**: `0.0.0.0:<port>` — containers reach the host via the Docker bridge gateway (`172.17.0.1`), which is not loopback
|
||||
|
||||
Default port: `50051`.
|
||||
|
||||
**Reference:** `src/orchestrator/api.rs` — `OrchestratorApi::start()`, platform-conditional bind address block
|
||||
|
||||
### Authentication
|
||||
|
||||
Per-job bearer tokens validated by `worker_auth_middleware`:
|
||||
|
||||
1. Tokens are **cryptographically random** (32 bytes, hex-encoded = 64 chars)
|
||||
2. Tokens are **scoped to a specific job_id** — a token for job A cannot access endpoints for job B
|
||||
3. Comparison uses **constant-time** `subtle::ConstantTimeEq`
|
||||
4. Tokens are **ephemeral** (in-memory only, never persisted to disk or DB)
|
||||
5. Tokens and associated credential grants are **revoked** when the container is cleaned up
|
||||
|
||||
**Reference:** `src/orchestrator/auth.rs` — `TokenStore::create_token()`, `TokenStore::validate()`, `generate_token()`
|
||||
|
||||
### Token Extraction
|
||||
|
||||
The middleware extracts the job UUID from the URL path (`/worker/{job_id}/...`) and validates the `Authorization: Bearer` header against the stored token for that specific job.
|
||||
|
||||
**Reference:** `src/orchestrator/auth.rs` — `worker_auth_middleware()`, `extract_job_id_from_path()`
|
||||
|
||||
### Credential Grants
|
||||
|
||||
The orchestrator can grant per-job access to specific secrets from the encrypted secrets store. Grants are:
|
||||
|
||||
- Stored alongside the token in the `TokenStore`
|
||||
- Scoped to specific `(secret_name, env_var)` pairs
|
||||
- Revoked when the job token is revoked
|
||||
- Decrypted on-demand when the worker requests `/worker/{job_id}/credentials`
|
||||
|
||||
**Reference:** `src/orchestrator/auth.rs` — `CredentialGrant` struct, `src/orchestrator/api.rs` — `get_credentials_handler()`
|
||||
|
||||
### Rate Limiting
|
||||
|
||||
**None.** The orchestrator API has no rate limiting. All `/worker/*` endpoints are authenticated via per-job bearer tokens, but a compromised container could spam authenticated endpoints without throttling.
|
||||
|
||||
**Mitigation:** Tokens are scoped per-job so a compromised container can only abuse its own job's endpoints. Container execution is time-bounded (see [Docker Container Security](#docker-container-security)), which limits the window for abuse.
|
||||
|
||||
### Routes
|
||||
|
||||
| Route | Auth | Purpose | Response |
|
||||
|-------|------|---------|----------|
|
||||
| `/health` | None | Health check | `"ok"` (plain text) — no fingerprinting data |
|
||||
| `/worker/{job_id}/job` | Per-job token | Get job description | Job JSON |
|
||||
| `/worker/{job_id}/llm/complete` | Per-job token | Proxy LLM completion | LLM response |
|
||||
| `/worker/{job_id}/llm/complete_with_tools` | Per-job token | Proxy LLM tool completion | LLM response |
|
||||
| `/worker/{job_id}/status` | Per-job token | Report worker status | Ack |
|
||||
| `/worker/{job_id}/complete` | Per-job token | Report job completion | Ack |
|
||||
| `/worker/{job_id}/event` | Per-job token | Send job events (SSE broadcast) | Ack |
|
||||
| `/worker/{job_id}/prompt` | Per-job token | Poll for follow-up prompts | Prompt or empty |
|
||||
| `/worker/{job_id}/credentials` | Per-job token | Retrieve decrypted credentials | Credentials JSON |
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
**None.** The orchestrator calls `axum::serve(listener, router).await?` without `.with_graceful_shutdown()`. The server stops only when the task is dropped (process exit or tokio task cancellation). In-flight requests may be interrupted.
|
||||
|
||||
**Reference:** `src/orchestrator/api.rs` — `OrchestratorApi::start()`
|
||||
|
||||
---
|
||||
|
||||
## 4. OAuth Callback Listener
|
||||
|
||||
**Source:** `src/cli/oauth_defaults.rs`
|
||||
|
||||
### Bind Address
|
||||
|
||||
Always binds to **loopback only**: `127.0.0.1:9876`. Falls back to `[::1]:9876` (IPv6 loopback) if IPv4 binding fails for reasons other than `AddrInUse`. If the port is already in use, the error is returned immediately (fail-fast).
|
||||
|
||||
Both IPv4 and IPv6 loopback addresses are security-equivalent — they are only reachable from the local machine.
|
||||
|
||||
**Reference:** `src/cli/oauth_defaults.rs` — `OAUTH_CALLBACK_PORT` constant, `bind_callback_listener()`
|
||||
|
||||
### Lifecycle
|
||||
|
||||
The listener is **ephemeral** — it is started only when an OAuth flow is initiated (e.g., `ironclaw tool auth <name>`) and shut down after the callback is received or the timeout expires.
|
||||
|
||||
### Timeout
|
||||
|
||||
**5-minute timeout** (`Duration::from_secs(300)`). If the user does not complete the OAuth flow in the browser within 5 minutes, the listener shuts down.
|
||||
|
||||
**Reference:** `src/cli/oauth_defaults.rs` — `tokio::time::timeout(Duration::from_secs(300), ...)`
|
||||
|
||||
### Security Controls
|
||||
|
||||
- **HTML escaping**: Provider names displayed in the landing page are HTML-escaped to prevent XSS (escapes `&`, `<`, `>`, `"`, `'`)
|
||||
- **Error parameter checking**: The handler checks for `error=` in the callback query string before extracting the auth code
|
||||
- **URL decoding**: Callback parameters are URL-decoded safely
|
||||
|
||||
**Reference:** `src/cli/oauth_defaults.rs` — `html_escape()`
|
||||
|
||||
### Built-in OAuth Credentials
|
||||
|
||||
Google OAuth client ID and secret are compiled into the binary (with compile-time override via `IRONCLAW_GOOGLE_CLIENT_ID` / `IRONCLAW_GOOGLE_CLIENT_SECRET`). As noted in the source, Google Desktop App client secrets are [not actually secret](https://developers.google.com/identity/protocols/oauth2/native-app) per Google's documentation.
|
||||
|
||||
**Reference:** `src/cli/oauth_defaults.rs` — `GOOGLE_CLIENT_ID` / `GOOGLE_CLIENT_SECRET` constants
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
Implicit. The listener is a raw `TcpListener` (not axum) inside a `tokio::time::timeout` future. Once the authorization code or error is received, the future returns and the `TcpListener` is dropped, closing the port. No explicit shutdown signal is needed.
|
||||
|
||||
**Reference:** `src/cli/oauth_defaults.rs` — `wait_for_callback()`
|
||||
|
||||
---
|
||||
|
||||
## 5. Sandbox HTTP Proxy
|
||||
|
||||
**Source:** `src/sandbox/proxy/http.rs`, `src/sandbox/proxy/allowlist.rs`, `src/sandbox/proxy/policy.rs`
|
||||
|
||||
### Bind Address
|
||||
|
||||
Always binds to **`127.0.0.1`** (localhost only). Port is OS-assigned (port `0`, ephemeral). Falls back to `[::1]` (IPv6 loopback) if IPv4 is unavailable.
|
||||
|
||||
Both IPv4 and IPv6 loopback addresses are security-equivalent — they are only reachable from the local machine.
|
||||
|
||||
**Reference:** `src/sandbox/proxy/http.rs` — `SandboxProxy::start()`, `TcpListener::bind("127.0.0.1:0")`
|
||||
|
||||
### Purpose
|
||||
|
||||
Acts as an HTTP/HTTPS proxy for Docker sandbox containers. Containers are configured with `http_proxy` / `https_proxy` environment variables pointing to this proxy, so all outbound HTTP traffic is routed through it.
|
||||
|
||||
### Domain Allowlisting
|
||||
|
||||
All requests are validated against a domain allowlist before being forwarded:
|
||||
|
||||
- **Empty allowlist = deny all** (fail-closed default)
|
||||
- Supports exact matches and wildcard patterns (`*.example.com`)
|
||||
- Validates URL scheme (HTTP/HTTPS only, rejects `ftp://`, `file://`, etc.)
|
||||
|
||||
**Reference:** `src/sandbox/proxy/allowlist.rs` — `DomainAllowlist` struct, `is_allowed()` method
|
||||
|
||||
### HTTPS Tunneling (CONNECT)
|
||||
|
||||
- CONNECT requests for HTTPS tunneling are subject to the same allowlist
|
||||
- **30-minute timeout** on established tunnels to prevent indefinite holds
|
||||
- **No MITM**: the proxy cannot inspect or inject credentials into HTTPS traffic (by design — containers that need credentials must use the orchestrator's `/worker/{job_id}/credentials` endpoint)
|
||||
|
||||
**Reference:** `src/sandbox/proxy/http.rs` — `handle_connect()` function
|
||||
|
||||
### Credential Injection (HTTP only)
|
||||
|
||||
For plain HTTP requests to allowed hosts, the proxy can inject credentials:
|
||||
|
||||
- Bearer tokens in `Authorization` header
|
||||
- Custom headers (e.g., `X-API-Key`)
|
||||
- Query parameters
|
||||
- Credentials are resolved at request time from the encrypted secrets store
|
||||
- Credentials never enter the container's environment or filesystem
|
||||
|
||||
**Reference:** `src/sandbox/proxy/http.rs` — credential injection block in `handle_request()`
|
||||
|
||||
### Hop-by-Hop Header Filtering
|
||||
|
||||
The proxy strips hop-by-hop headers to prevent header-based attacks: `connection`, `keep-alive`, `proxy-authenticate`, `proxy-authorization`, `te`, `trailers`, `transfer-encoding`, `upgrade`.
|
||||
|
||||
**Reference:** `src/sandbox/proxy/http.rs` — `is_hop_by_hop_header()`
|
||||
|
||||
### Docker Container Security
|
||||
|
||||
Containers that use the proxy are configured with defense-in-depth:
|
||||
|
||||
| Control | Setting | Reference |
|
||||
|---------|---------|-----------|
|
||||
| Capabilities | Drop ALL, add only CHOWN | `src/sandbox/container.rs` — `cap_drop` / `cap_add` |
|
||||
| Privilege escalation | `no-new-privileges:true` | `src/sandbox/container.rs` — `security_opt` |
|
||||
| Root filesystem | Read-only (except FullAccess policy) | `src/sandbox/container.rs` — `readonly_rootfs` |
|
||||
| User | Non-root (UID 1000:1000) | `src/sandbox/container.rs` — `user` field |
|
||||
| Network | Bridge mode (isolated) | `src/sandbox/container.rs` — `network_mode` |
|
||||
| Tmpfs | `/tmp` (512 MB), `/home/sandbox/.cargo/registry` (1 GB) | `src/sandbox/container.rs` — `tmpfs` block |
|
||||
| Auto-remove | Enabled | `src/sandbox/container.rs` — `auto_remove` |
|
||||
| Output limits | Configurable max stdout/stderr | `src/sandbox/container.rs` — `collect_logs()` |
|
||||
| Timeout | Enforced with forced container removal | `src/sandbox/container.rs` — `tokio::time::timeout` in `run()` |
|
||||
|
||||
### Graceful Shutdown
|
||||
|
||||
Shutdown is triggered via a `oneshot::Sender` stored on the proxy. The accept loop uses `tokio::select!` to race `listener.accept()` against the shutdown signal. The `stop()` method fires the signal; the loop breaks on the next iteration. Note: `stop()` does not await a join handle, so there is no drain-and-wait for in-flight connections.
|
||||
|
||||
**Reference:** `src/sandbox/proxy/http.rs` — `stop()` method, `tokio::select!` loop
|
||||
|
||||
---
|
||||
|
||||
## Egress Controls
|
||||
|
||||
### WASM Tool HTTP Requests
|
||||
|
||||
WASM tools execute HTTP requests through the host runtime, subject to:
|
||||
|
||||
1. **Endpoint allowlist** — declared in `<tool>.capabilities.json`, validated by `AllowlistValidator`
|
||||
- Host matching (exact or wildcard)
|
||||
- Path prefix matching
|
||||
- HTTP method restriction
|
||||
- HTTPS required by default
|
||||
- Userinfo in URLs (`user:pass@host`) rejected to prevent allowlist bypass
|
||||
- Path traversal (`../`, `%2e%2e/`) normalized and blocked
|
||||
- Invalid percent-encoding rejected
|
||||
- **Reference:** `src/tools/wasm/allowlist.rs`
|
||||
|
||||
2. **Credential injection** — secrets injected at the host boundary by `CredentialInjector`
|
||||
- WASM code never sees actual credential values
|
||||
- Secrets must be in the tool's `allowed_secrets` list
|
||||
- Injection supports: Bearer header, Basic auth, custom header, query parameter
|
||||
- **Reference:** `src/tools/wasm/credential_injector.rs`
|
||||
|
||||
3. **Leak detection** — `LeakDetector` scans both outbound requests and inbound responses for secret patterns
|
||||
- Runs at two points: before sending and after receiving
|
||||
- Uses Aho-Corasick for fast multi-pattern matching
|
||||
- **Reference:** `src/safety/leak_detector.rs`
|
||||
|
||||
### Built-in HTTP Tool
|
||||
|
||||
The `http` tool (`src/tools/builtin/http.rs`) has its own SSRF protections:
|
||||
|
||||
| Protection | Details | Reference |
|
||||
|-----------|---------|-----------|
|
||||
| HTTPS only | Rejects `http://` URLs | `http.rs` — scheme check |
|
||||
| Localhost blocked | Rejects `localhost` and `*.localhost` | `http.rs` — host check |
|
||||
| Private IP blocked | Rejects RFC 1918, loopback, link-local, multicast, unspecified | `http.rs` — `is_disallowed_ip()` |
|
||||
| DNS rebinding | Resolves hostname and checks all resolved IPs against blocklist | `http.rs` — DNS resolution block |
|
||||
| Cloud metadata | Blocks `169.254.169.254` (AWS/GCP metadata endpoint) | `http.rs` — `is_disallowed_ip()` |
|
||||
| Redirect blocking | Returns error on 3xx responses (prevents SSRF via redirect) | `http.rs` — status code check |
|
||||
| Response size limit | **5 MB** max, enforced both via Content-Length header and streaming | `http.rs` — `MAX_RESPONSE_SIZE` constant, streaming cap |
|
||||
| Outbound leak scan | Scans URL, headers, and body for secrets before sending | `http.rs` — `LeakDetector::scan_http_request()` |
|
||||
| Approval required | Requires user approval before execution | `http.rs` — `requires_approval()` returns `true` |
|
||||
| Timeout | 30 seconds default | `http.rs` — `reqwest::Client` builder |
|
||||
| No redirects | `redirect::Policy::none()` — redirects are not followed | `http.rs` — `reqwest::Client` builder |
|
||||
|
||||
### MCP Client
|
||||
|
||||
MCP servers are external processes accessed via HTTP. The MCP client (`src/tools/mcp/client.rs`) uses `reqwest` with a 30-second timeout but has **no SSRF protections** — it connects to whatever URL is configured for the MCP server.
|
||||
|
||||
This is by design: MCP server URLs come from **operator-controlled configuration** (config files, environment variables, or the CLI `tool install` command), not from user input or LLM output. A compromised config file is outside IronClaw's threat model — it would imply the operator's machine is already compromised.
|
||||
|
||||
**Reference:** `src/tools/mcp/client.rs` — `reqwest::Client` builder
|
||||
|
||||
### Sandbox Domain Allowlists
|
||||
|
||||
Sandbox containers route all HTTP traffic through the proxy, which enforces a domain allowlist. The allowlist is built from:
|
||||
|
||||
1. A default set of domains (`src/sandbox/config.rs` — `default_allowlist()`)
|
||||
2. Additional domains from `SANDBOX_EXTRA_DOMAINS` env var (comma-separated)
|
||||
|
||||
**Reference:** `src/config.rs` — sandbox allowlist assembly
|
||||
|
||||
---
|
||||
|
||||
## Authentication Mechanisms Summary
|
||||
|
||||
| Mechanism | Constant-Time | Used By | Reference |
|
||||
|-----------|:------------:|---------|-----------|
|
||||
| Gateway bearer token | Yes | Web gateway (header + query) | `src/channels/web/auth.rs` — `auth_middleware()` |
|
||||
| Webhook shared secret | Yes | HTTP webhook (`ct_eq` comparison) | `src/channels/http.rs` — `webhook_handler()` |
|
||||
| Per-job bearer token | Yes | Orchestrator worker API | `src/orchestrator/auth.rs` — `TokenStore::validate()` |
|
||||
| OAuth callback | N/A | CLI OAuth flow (no auth, loopback-only) | `src/cli/oauth_defaults.rs` — `bind_callback_listener()` |
|
||||
| Sandbox proxy | N/A | No auth (loopback-only, ephemeral) | `src/sandbox/proxy/http.rs` — `SandboxProxy::start()` |
|
||||
|
||||
---
|
||||
|
||||
## Known Security Findings
|
||||
|
||||
### Open
|
||||
|
||||
#### F-2. No TLS at the application layer
|
||||
|
||||
**Severity:** Low (for local deployment)
|
||||
**Details:** None of the listeners terminate TLS. All communication is plain HTTP.
|
||||
**Mitigation:** The web gateway and OAuth callback bind to loopback by default. For production, users are expected to front the gateway with a reverse proxy (nginx, Caddy) or tunnel (Cloudflare, ngrok) that provides TLS.
|
||||
**Recommendation:** Document the requirement for a TLS-terminating reverse proxy in deployment guides.
|
||||
|
||||
#### F-3. Orchestrator binds to `0.0.0.0` on Linux
|
||||
|
||||
**Severity:** Medium
|
||||
**Location:** `src/orchestrator/api.rs` — platform-conditional bind in `OrchestratorApi::start()`
|
||||
**Details:** On Linux, the orchestrator API binds to all interfaces because Docker containers reach the host via the bridge gateway (`172.17.0.1`), not loopback. This means the API is reachable from any network interface on the host.
|
||||
**Mitigation:** All `/worker/*` endpoints require per-job bearer tokens (constant-time, cryptographically random). The `/health` endpoint is the only unauthenticated route and returns only `"ok"`. Firewall rules should block external access to port 50051.
|
||||
**Recommendation:** Document firewall requirements for Linux deployments. Consider binding to the Docker bridge IP (`172.17.0.1`) instead of `0.0.0.0`.
|
||||
|
||||
#### F-6. WebSocket/SSE connection limit
|
||||
|
||||
**Severity:** Info
|
||||
**Details:** The `SseManager` enforces a hard limit of **100 concurrent connections** (`MAX_CONNECTIONS` constant in `src/channels/web/sse.rs`). Both SSE subscribers and WebSocket connections share this counter. When exceeded, new WebSocket upgrades are rejected with a warning log and the connection is immediately closed.
|
||||
**Reference:** `src/channels/web/sse.rs` — `MAX_CONNECTIONS`, `src/channels/web/ws.rs` — `handle_ws_connection()` early return
|
||||
|
||||
#### F-7. Orchestrator API has no rate limiting
|
||||
|
||||
**Severity:** Low
|
||||
**Details:** The orchestrator API has no request-rate throttling. A compromised container could spam authenticated endpoints (e.g., `/worker/{job_id}/llm/complete`) to drive up LLM costs or degrade service for other jobs.
|
||||
**Mitigation:** Tokens are scoped per-job, limiting blast radius. Container execution is time-bounded by the sandbox timeout, which caps the abuse window.
|
||||
**Recommendation:** Consider adding per-token rate limiting on the LLM proxy endpoints.
|
||||
|
||||
#### F-8. Orchestrator API has no graceful shutdown
|
||||
|
||||
**Severity:** Info
|
||||
**Details:** The orchestrator calls `axum::serve(listener, router).await?` without `.with_graceful_shutdown()`. In-flight requests (including LLM proxy calls) may be interrupted during process shutdown.
|
||||
**Reference:** `src/orchestrator/api.rs` — `OrchestratorApi::start()`
|
||||
|
||||
### Resolved / Mitigated
|
||||
|
||||
<details>
|
||||
<summary>Resolved and mitigated findings (click to expand)</summary>
|
||||
|
||||
#### F-1. ~~Webhook secret comparison is not constant-time~~ (Resolved)
|
||||
|
||||
**Severity:** Low
|
||||
**Location:** `src/channels/http.rs` — `webhook_handler()`
|
||||
**Status:** Resolved — webhook secret now uses `subtle::ConstantTimeEq` (`ct_eq`), consistent with web gateway and orchestrator auth.
|
||||
|
||||
#### F-4. ~~HTTP webhook server binds to `0.0.0.0` by default~~ (Mitigated)
|
||||
|
||||
**Severity:** Low
|
||||
**Location:** `src/config.rs`, `src/main.rs`
|
||||
**Status:** Mitigated — a `tracing::warn!` is now emitted at startup when the webhook server binds to an unspecified address (`0.0.0.0` or `::`), advising operators to set `HTTP_HOST=127.0.0.1` to restrict to localhost. The default bind address remains `0.0.0.0`, so webhook exposure is still controlled by operator configuration and external network controls (firewalls, ingress rules).
|
||||
|
||||
#### F-5. ~~Missing security headers on web gateway~~ (Mitigated)
|
||||
|
||||
**Severity:** Low
|
||||
**Status:** Mitigated — `X-Content-Type-Options: nosniff` and `X-Frame-Options: DENY` are now set on all gateway responses via `SetResponseHeaderLayer::if_not_present`. Layer ordering ensures these headers are applied even to error responses generated by inner layers (e.g., `DefaultBodyLimit` 413 rejections).
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
## Review Checklist for Network Changes
|
||||
|
||||
Use this checklist for any PR that adds or modifies network-facing code.
|
||||
|
||||
### New Listener
|
||||
|
||||
- [ ] **Bind address**: Does it bind to loopback (`127.0.0.1`) or all interfaces (`0.0.0.0`)? Justify if `0.0.0.0`.
|
||||
- [ ] **Port configuration**: Is the port configurable via env var? Is a sensible default set?
|
||||
- [ ] **Authentication**: Is auth required? If yes, is it constant-time? If no, why not?
|
||||
- [ ] **Rate limiting**: Is there a rate limiter? What are the limits?
|
||||
- [ ] **Body size limit**: Is `DefaultBodyLimit` (or equivalent) set?
|
||||
- [ ] **Content-Type validation**: Does the handler validate Content-Type (e.g., via axum `Json<T>` extractor)?
|
||||
- [ ] **Graceful shutdown**: Does the listener support graceful shutdown via oneshot or similar?
|
||||
- [ ] **Inventory update**: Is this document updated with the new listener?
|
||||
|
||||
### New Route on Existing Listener
|
||||
|
||||
- [ ] **Auth layer**: Is the route behind the auth middleware? If public, why?
|
||||
- [ ] **Input validation**: Are path parameters, query parameters, and body fields validated?
|
||||
- [ ] **Error responses**: Do error responses avoid leaking internal details?
|
||||
|
||||
### Egress (Outbound HTTP)
|
||||
|
||||
- [ ] **SSRF protection**: Does the code block private IPs, localhost, and cloud metadata endpoints?
|
||||
- [ ] **DNS rebinding**: Are resolved IPs checked (not just the hostname)?
|
||||
- [ ] **Redirect handling**: Are redirects blocked or validated?
|
||||
- [ ] **Response size**: Is there a max response size?
|
||||
- [ ] **Timeout**: Is a request timeout set?
|
||||
- [ ] **Leak detection**: Is the outbound request scanned for secrets?
|
||||
|
||||
### Credential Handling
|
||||
|
||||
- [ ] **Constant-time comparison**: Are secrets compared with `subtle::ConstantTimeEq`?
|
||||
- [ ] **No logging**: Are credentials excluded from log messages?
|
||||
- [ ] **Ephemeral storage**: Are tokens stored in memory only (not persisted)?
|
||||
- [ ] **Scope**: Are credentials scoped to the minimum necessary (per-job, per-tool)?
|
||||
- [ ] **Revocation**: Are credentials revoked when no longer needed?
|
||||
|
||||
### Container / Sandbox
|
||||
|
||||
- [ ] **Capabilities**: Are all capabilities dropped except what's needed?
|
||||
- [ ] **Filesystem**: Is the root filesystem read-only?
|
||||
- [ ] **User**: Does the container run as non-root?
|
||||
- [ ] **Network**: Is network access routed through the proxy?
|
||||
- [ ] **Timeout**: Is there an execution timeout with forced cleanup?
|
||||
- [ ] **Output limits**: Are stdout/stderr capped?
|
||||
+10
-2
@@ -255,7 +255,10 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Execute each tool (with approval checking and hook interception)
|
||||
for mut tc in tool_calls {
|
||||
let mut idx = 0usize;
|
||||
while idx < tool_calls.len() {
|
||||
let mut tc = tool_calls[idx].clone();
|
||||
|
||||
// Check if tool requires approval
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
@@ -277,7 +280,9 @@ impl Agent {
|
||||
}
|
||||
|
||||
if !is_auto_approved {
|
||||
// Need approval - store pending request and return
|
||||
// Need approval - store pending request and return.
|
||||
// Preserve remaining tool calls so they can be replayed
|
||||
// after approval.
|
||||
let pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
@@ -285,6 +290,7 @@ impl Agent {
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: tool_calls[idx + 1..].to_vec(),
|
||||
};
|
||||
|
||||
return Ok(AgenticLoopResult::NeedApproval { pending });
|
||||
@@ -441,6 +447,8 @@ impl Agent {
|
||||
&tc.name,
|
||||
result_content,
|
||||
));
|
||||
|
||||
idx += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::llm::ChatMessage;
|
||||
use crate::llm::{ChatMessage, ToolCall};
|
||||
|
||||
/// A session containing one or more threads.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
@@ -148,6 +148,10 @@ pub struct PendingApproval {
|
||||
pub tool_call_id: String,
|
||||
/// Context messages at the time of the request (to resume from).
|
||||
pub context_messages: Vec<ChatMessage>,
|
||||
/// Remaining tool calls from the same assistant message that were not
|
||||
/// executed yet when approval was requested.
|
||||
#[serde(default)]
|
||||
pub deferred_tool_calls: Vec<ToolCall>,
|
||||
}
|
||||
|
||||
/// A conversation thread within a session.
|
||||
@@ -946,6 +950,7 @@ mod tests {
|
||||
description: "dangerous command".to_string(),
|
||||
tool_call_id: "call_123".to_string(),
|
||||
context_messages: vec![ChatMessage::user("do it")],
|
||||
deferred_tool_calls: vec![],
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
@@ -969,6 +974,7 @@ mod tests {
|
||||
description: "test".to_string(),
|
||||
tool_call_id: "call_456".to_string(),
|
||||
context_messages: vec![],
|
||||
deferred_tool_calls: vec![],
|
||||
};
|
||||
|
||||
thread.await_approval(approval);
|
||||
|
||||
+54
-3
@@ -118,19 +118,19 @@ impl SubmissionParser {
|
||||
// Approval responses (simple yes/no/always for pending approvals)
|
||||
// These are short enough to check explicitly
|
||||
match lower.as_str() {
|
||||
"yes" | "y" | "approve" | "ok" => {
|
||||
"yes" | "y" | "approve" | "ok" | "/approve" | "/yes" | "/y" => {
|
||||
return Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: false,
|
||||
};
|
||||
}
|
||||
"always" | "yes always" | "approve always" => {
|
||||
"always" | "a" | "yes always" | "approve always" | "/always" | "/a" => {
|
||||
return Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: true,
|
||||
};
|
||||
}
|
||||
"no" | "n" | "deny" | "reject" | "cancel" => {
|
||||
"no" | "n" | "deny" | "reject" | "cancel" | "/deny" | "/no" | "/n" => {
|
||||
return Submission::ApprovalResponse {
|
||||
approved: false,
|
||||
always: false,
|
||||
@@ -475,6 +475,57 @@ mod tests {
|
||||
assert!(matches!(submission, Submission::UserInput { content } if content == "/unknown"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_approval_response_aliases() {
|
||||
// approve once
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("y"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: false
|
||||
}
|
||||
));
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("/approve"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: false
|
||||
}
|
||||
));
|
||||
|
||||
// approve always
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("a"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: true
|
||||
}
|
||||
));
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("/always"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: true,
|
||||
always: true
|
||||
}
|
||||
));
|
||||
|
||||
// deny
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("n"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: false,
|
||||
always: false
|
||||
}
|
||||
));
|
||||
assert!(matches!(
|
||||
SubmissionParser::parse("/deny"),
|
||||
Submission::ApprovalResponse {
|
||||
approved: false,
|
||||
always: false
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parser_json_exec_approval() {
|
||||
let req_id = Uuid::new_v4();
|
||||
|
||||
+174
-1
@@ -11,7 +11,7 @@ use uuid::Uuid;
|
||||
use crate::agent::Agent;
|
||||
use crate::agent::compaction::ContextCompactor;
|
||||
use crate::agent::dispatcher::{AgenticLoopResult, detect_auth_awaiting, parse_auth_result};
|
||||
use crate::agent::session::{Session, ThreadState};
|
||||
use crate::agent::session::{PendingApproval, Session, ThreadState};
|
||||
use crate::agent::submission::SubmissionResult;
|
||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||
use crate::context::JobContext;
|
||||
@@ -712,6 +712,7 @@ impl Agent {
|
||||
|
||||
// Build context including the tool result
|
||||
let mut context_messages = pending.context_messages;
|
||||
let deferred_tool_calls = pending.deferred_tool_calls;
|
||||
|
||||
// Record result in thread
|
||||
{
|
||||
@@ -780,6 +781,178 @@ impl Agent {
|
||||
result_content,
|
||||
));
|
||||
|
||||
// Replay deferred tool calls from the same assistant message so
|
||||
// every tool_use ID gets a matching tool_result before the next
|
||||
// LLM call.
|
||||
if !deferred_tool_calls.is_empty() {
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Thinking(format!(
|
||||
"Executing {} deferred tool(s)...",
|
||||
deferred_tool_calls.len()
|
||||
)),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let mut deferred_queue = std::collections::VecDeque::from(deferred_tool_calls);
|
||||
while let Some(tc) = deferred_queue.pop_front() {
|
||||
// Re-check approval for each deferred tool call
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
let is_auto_approved = {
|
||||
let sess = session.lock().await;
|
||||
let mut approved = sess.is_tool_auto_approved(&tc.name);
|
||||
if approved && tool.requires_approval_for(&tc.arguments) {
|
||||
approved = false;
|
||||
}
|
||||
approved
|
||||
};
|
||||
|
||||
if !is_auto_approved {
|
||||
let new_pending = PendingApproval {
|
||||
request_id: Uuid::new_v4(),
|
||||
tool_name: tc.name.clone(),
|
||||
parameters: tc.arguments.clone(),
|
||||
description: tool.description().to_string(),
|
||||
tool_call_id: tc.id.clone(),
|
||||
context_messages: context_messages.clone(),
|
||||
deferred_tool_calls: deferred_queue.iter().cloned().collect(),
|
||||
};
|
||||
|
||||
let request_id = new_pending.request_id;
|
||||
let tool_name = new_pending.tool_name.clone();
|
||||
let description = new_pending.description.clone();
|
||||
let parameters = new_pending.parameters.clone();
|
||||
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.await_approval(new_pending);
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::Status("Awaiting approval".into()),
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
return Ok(SubmissionResult::NeedApproval {
|
||||
request_id,
|
||||
tool_name,
|
||||
description,
|
||||
parameters,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolStarted {
|
||||
name: tc.name.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
let deferred_result = self
|
||||
.execute_chat_tool(&tc.name, &tc.arguments, &job_ctx)
|
||||
.await;
|
||||
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolCompleted {
|
||||
name: tc.name.clone(),
|
||||
success: deferred_result.is_ok(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Ok(ref output) = deferred_result
|
||||
&& !output.is_empty()
|
||||
{
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::ToolResult {
|
||||
name: tc.name.clone(),
|
||||
preview: output.clone(),
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Record in thread
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id)
|
||||
&& let Some(turn) = thread.last_turn_mut()
|
||||
{
|
||||
match &deferred_result {
|
||||
Ok(output) => turn.record_tool_result(serde_json::json!(output)),
|
||||
Err(e) => turn.record_tool_error(e.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Auth detection for deferred tools
|
||||
if let Some((ext_name, instructions)) =
|
||||
detect_auth_awaiting(&tc.name, &deferred_result)
|
||||
{
|
||||
let auth_data = parse_auth_result(&deferred_result);
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
if let Some(thread) = sess.threads.get_mut(&thread_id) {
|
||||
thread.enter_auth_mode(ext_name.clone());
|
||||
thread.complete_turn(&instructions);
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
&message.channel,
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name: ext_name,
|
||||
instructions: Some(instructions.clone()),
|
||||
auth_url: auth_data.auth_url,
|
||||
setup_url: auth_data.setup_url,
|
||||
},
|
||||
&message.metadata,
|
||||
)
|
||||
.await;
|
||||
return Ok(SubmissionResult::response(instructions));
|
||||
}
|
||||
|
||||
let deferred_content = match deferred_result {
|
||||
Ok(output) => {
|
||||
let sanitized = self.safety().sanitize_tool_output(&tc.name, &output);
|
||||
self.safety().wrap_for_llm(
|
||||
&tc.name,
|
||||
&sanitized.content,
|
||||
sanitized.was_modified,
|
||||
)
|
||||
}
|
||||
Err(e) => format!("Error: {}", e),
|
||||
};
|
||||
|
||||
context_messages.push(ChatMessage::tool_result(&tc.id, &tc.name, deferred_content));
|
||||
}
|
||||
|
||||
// Continue the agentic loop (a tool was already executed this turn)
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages, true)
|
||||
|
||||
+7
-2
@@ -473,6 +473,7 @@ impl AppBuilder {
|
||||
pub async fn init_extensions(
|
||||
&self,
|
||||
tools: &Arc<ToolRegistry>,
|
||||
hooks: &Arc<HookRegistry>,
|
||||
) -> Result<
|
||||
(
|
||||
Arc<McpSessionManager>,
|
||||
@@ -661,6 +662,7 @@ impl AppBuilder {
|
||||
Arc::clone(&mcp_session_manager),
|
||||
Arc::clone(secrets),
|
||||
Arc::clone(tools),
|
||||
Some(Arc::clone(hooks)),
|
||||
wasm_tool_runtime.clone(),
|
||||
self.config.wasm.tools_dir.clone(),
|
||||
self.config.channels.wasm_channels_dir.clone(),
|
||||
@@ -697,8 +699,12 @@ impl AppBuilder {
|
||||
|
||||
let (llm, cheap_llm) = self.init_llm()?;
|
||||
let (safety, tools, embeddings, workspace) = self.init_tools(&llm).await?;
|
||||
|
||||
// Create hook registry early so runtime extension activation can register hooks.
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
|
||||
let (mcp_session_manager, wasm_tool_runtime, extension_manager) =
|
||||
self.init_extensions(&tools).await?;
|
||||
self.init_extensions(&tools, &hooks).await?;
|
||||
|
||||
// Seed workspace and backfill embeddings
|
||||
if let Some(ref ws) = workspace {
|
||||
@@ -741,7 +747,6 @@ impl AppBuilder {
|
||||
};
|
||||
|
||||
let context_manager = Arc::new(ContextManager::new(self.config.agent.max_parallel_jobs));
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
let cost_guard = Arc::new(crate::agent::cost_guard::CostGuard::new(
|
||||
crate::agent::cost_guard::CostGuardConfig {
|
||||
max_cost_per_day_cents: self.config.agent.max_cost_per_day_cents,
|
||||
|
||||
@@ -492,4 +492,32 @@ INJECTED="pwned"#;
|
||||
assert_eq!(parsed.len(), 2);
|
||||
assert!(parsed.iter().all(|(k, _)| k != "DATABASE_URL"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_onboard_completed_round_trips_through_env() {
|
||||
let dir = tempdir().unwrap();
|
||||
let env_path = dir.path().join(".env");
|
||||
|
||||
// Simulate what the wizard writes: bootstrap vars + ONBOARD_COMPLETED
|
||||
let vars = [
|
||||
("DATABASE_BACKEND", "libsql"),
|
||||
("ONBOARD_COMPLETED", "true"),
|
||||
];
|
||||
let mut content = String::new();
|
||||
for (key, value) in &vars {
|
||||
let escaped = value.replace('\\', "\\\\").replace('"', "\\\"");
|
||||
content.push_str(&format!("{}=\"{}\"\n", key, escaped));
|
||||
}
|
||||
std::fs::write(&env_path, &content).unwrap();
|
||||
|
||||
// Verify dotenvy parses ONBOARD_COMPLETED correctly
|
||||
let parsed: Vec<(String, String)> = dotenvy::from_path_iter(&env_path)
|
||||
.unwrap()
|
||||
.filter_map(|r| r.ok())
|
||||
.collect();
|
||||
assert_eq!(parsed.len(), 2);
|
||||
let onboard = parsed.iter().find(|(k, _)| k == "ONBOARD_COMPLETED");
|
||||
assert!(onboard.is_some(), "ONBOARD_COMPLETED must be present");
|
||||
assert_eq!(onboard.unwrap().1, "true");
|
||||
}
|
||||
}
|
||||
|
||||
+80
-9
@@ -12,6 +12,7 @@ use axum::{
|
||||
};
|
||||
use secrecy::ExposeSecret;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use subtle::ConstantTimeEq;
|
||||
use tokio::sync::{RwLock, mpsc, oneshot};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use uuid::Uuid;
|
||||
@@ -173,7 +174,7 @@ async fn webhook_handler(
|
||||
// Validate secret if configured
|
||||
if let Some(ref expected_secret) = state.webhook_secret {
|
||||
match &req.secret {
|
||||
Some(provided) if provided == expected_secret => {
|
||||
Some(provided) if bool::from(provided.as_bytes().ct_eq(expected_secret.as_bytes())) => {
|
||||
// Secret matches, continue
|
||||
}
|
||||
Some(_) => {
|
||||
@@ -356,19 +357,89 @@ impl Channel for HttpChannel {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use secrecy::SecretString;
|
||||
use tower::ServiceExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn test_channel(secret: Option<&str>) -> HttpChannel {
|
||||
HttpChannel::new(HttpConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
webhook_secret: secret.map(|s| SecretString::from(s.to_string())),
|
||||
user_id: "http".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_http_channel_requires_secret() {
|
||||
let config = HttpConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
webhook_secret: None,
|
||||
user_id: "http".to_string(),
|
||||
};
|
||||
|
||||
let channel = HttpChannel::new(config);
|
||||
let channel = test_channel(None);
|
||||
let result = channel.start().await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn webhook_correct_secret_returns_ok() {
|
||||
let channel = test_channel(Some("test-secret-123"));
|
||||
// Start the channel so the tx sender is populated (otherwise 503).
|
||||
let _stream = channel.start().await.unwrap();
|
||||
let app = channel.routes();
|
||||
|
||||
let body = serde_json::json!({
|
||||
"content": "hello",
|
||||
"secret": "test-secret-123"
|
||||
});
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
|
||||
let resp = app.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn webhook_wrong_secret_returns_unauthorized() {
|
||||
let channel = test_channel(Some("correct-secret"));
|
||||
let _stream = channel.start().await.unwrap();
|
||||
let app = channel.routes();
|
||||
|
||||
let body = serde_json::json!({
|
||||
"content": "hello",
|
||||
"secret": "wrong-secret"
|
||||
});
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
|
||||
let resp = app.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn webhook_missing_secret_returns_unauthorized() {
|
||||
let channel = test_channel(Some("correct-secret"));
|
||||
let _stream = channel.start().await.unwrap();
|
||||
let app = channel.routes();
|
||||
|
||||
let body = serde_json::json!({
|
||||
"content": "hello"
|
||||
});
|
||||
let req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/webhook")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::to_vec(&body).unwrap()))
|
||||
.unwrap();
|
||||
|
||||
let resp = app.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -330,7 +330,13 @@ impl Channel for ReplChannel {
|
||||
// Handle local REPL commands (only commands that need
|
||||
// immediate local handling stay here)
|
||||
match line.to_lowercase().as_str() {
|
||||
"/quit" | "/exit" => break,
|
||||
"/quit" | "/exit" => {
|
||||
// Forward shutdown command so the agent loop exits even
|
||||
// when other channels (e.g. web gateway) are still active.
|
||||
let msg = IncomingMessage::new("repl", "default", "/quit");
|
||||
let _ = tx.blocking_send(msg);
|
||||
break;
|
||||
}
|
||||
"/help" => {
|
||||
print_help();
|
||||
continue;
|
||||
|
||||
@@ -305,6 +305,7 @@ impl Channel for GatewayChannel {
|
||||
description,
|
||||
parameters: serde_json::to_string_pretty(¶meters)
|
||||
.unwrap_or_else(|_| parameters.to_string()),
|
||||
thread_id,
|
||||
},
|
||||
StatusUpdate::AuthRequired {
|
||||
extension_name,
|
||||
|
||||
@@ -24,6 +24,8 @@ use crate::llm::{
|
||||
|
||||
use super::server::GatewayState;
|
||||
|
||||
const MAX_MODEL_NAME_BYTES: usize = 256;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// OpenAI request types
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -380,6 +382,27 @@ fn unix_timestamp() -> u64 {
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
fn validate_model_name(model: &str) -> Result<(), String> {
|
||||
let trimmed = model.trim();
|
||||
|
||||
if trimmed.is_empty() {
|
||||
return Err("model must not be empty".to_string());
|
||||
}
|
||||
if trimmed != model {
|
||||
return Err("model must not have leading or trailing whitespace".to_string());
|
||||
}
|
||||
if model.len() > MAX_MODEL_NAME_BYTES {
|
||||
return Err(format!(
|
||||
"model must be at most {} bytes",
|
||||
MAX_MODEL_NAME_BYTES
|
||||
));
|
||||
}
|
||||
if model.chars().any(char::is_control) {
|
||||
return Err("model contains control characters".to_string());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Extract stop sequences from the flexible `stop` field.
|
||||
fn parse_stop(val: &serde_json::Value) -> Option<Vec<String>> {
|
||||
match val {
|
||||
@@ -426,29 +449,17 @@ pub async fn chat_completions_handler(
|
||||
"invalid_request_error",
|
||||
));
|
||||
}
|
||||
|
||||
// Validate the requested model matches the active model.
|
||||
// Per-request model switching is not yet supported (see GH issue).
|
||||
let active_model = llm.active_model_name();
|
||||
if req.model != active_model {
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
Json(OpenAiErrorResponse {
|
||||
error: OpenAiErrorDetail {
|
||||
message: format!(
|
||||
"Model '{}' not found. The active model is '{}'.",
|
||||
req.model, active_model
|
||||
),
|
||||
error_type: "invalid_request_error".to_string(),
|
||||
param: Some("model".to_string()),
|
||||
code: Some("model_not_found".to_string()),
|
||||
},
|
||||
}),
|
||||
if let Err(e) = validate_model_name(&req.model) {
|
||||
return Err(openai_error(
|
||||
StatusCode::BAD_REQUEST,
|
||||
e,
|
||||
"invalid_request_error",
|
||||
));
|
||||
}
|
||||
|
||||
let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty());
|
||||
let stream = req.stream.unwrap_or(false);
|
||||
let requested_model = req.model.clone();
|
||||
|
||||
if stream {
|
||||
return handle_streaming(llm.clone(), req, has_tools)
|
||||
@@ -460,13 +471,12 @@ pub async fn chat_completions_handler(
|
||||
|
||||
let messages = convert_messages(&req.messages)
|
||||
.map_err(|e| openai_error(StatusCode::BAD_REQUEST, e, "invalid_request_error"))?;
|
||||
let model_name = llm.active_model_name();
|
||||
let id = chat_completion_id();
|
||||
let created = unix_timestamp();
|
||||
|
||||
if has_tools {
|
||||
let tools = convert_tools(req.tools.as_deref().unwrap_or(&[]));
|
||||
let mut tool_req = ToolCompletionRequest::new(messages, tools);
|
||||
let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model);
|
||||
if let Some(t) = req.temperature {
|
||||
tool_req = tool_req.with_temperature(t);
|
||||
}
|
||||
@@ -483,6 +493,7 @@ pub async fn chat_completions_handler(
|
||||
.complete_with_tools(tool_req)
|
||||
.await
|
||||
.map_err(map_llm_error)?;
|
||||
let model_name = llm.effective_model_name(Some(requested_model.as_str()));
|
||||
|
||||
let tool_calls_openai = if resp.tool_calls.is_empty() {
|
||||
None
|
||||
@@ -515,7 +526,7 @@ pub async fn chat_completions_handler(
|
||||
|
||||
Ok(Json(response).into_response())
|
||||
} else {
|
||||
let mut comp_req = CompletionRequest::new(messages);
|
||||
let mut comp_req = CompletionRequest::new(messages).with_model(req.model);
|
||||
if let Some(t) = req.temperature {
|
||||
comp_req = comp_req.with_temperature(t);
|
||||
}
|
||||
@@ -527,6 +538,7 @@ pub async fn chat_completions_handler(
|
||||
}
|
||||
|
||||
let resp = llm.complete(comp_req).await.map_err(map_llm_error)?;
|
||||
let model_name = llm.effective_model_name(Some(requested_model.as_str()));
|
||||
|
||||
let response = OpenAiChatResponse {
|
||||
id,
|
||||
@@ -570,7 +582,7 @@ async fn handle_streaming(
|
||||
let messages = convert_messages(&req.messages)
|
||||
.map_err(|e| openai_error(StatusCode::BAD_REQUEST, e, "invalid_request_error"))?;
|
||||
|
||||
let model_name = llm.active_model_name();
|
||||
let requested_model = req.model.clone();
|
||||
let id = chat_completion_id();
|
||||
let created = unix_timestamp();
|
||||
|
||||
@@ -584,7 +596,7 @@ async fn handle_streaming(
|
||||
|
||||
let llm_result = if has_tools {
|
||||
let tools = convert_tools(req.tools.as_deref().unwrap_or(&[]));
|
||||
let mut tool_req = ToolCompletionRequest::new(messages, tools);
|
||||
let mut tool_req = ToolCompletionRequest::new(messages, tools).with_model(req.model);
|
||||
if let Some(t) = req.temperature {
|
||||
tool_req = tool_req.with_temperature(t);
|
||||
}
|
||||
@@ -602,7 +614,7 @@ async fn handle_streaming(
|
||||
.map_err(map_llm_error)?,
|
||||
)
|
||||
} else {
|
||||
let mut comp_req = CompletionRequest::new(messages);
|
||||
let mut comp_req = CompletionRequest::new(messages).with_model(req.model);
|
||||
if let Some(t) = req.temperature {
|
||||
comp_req = comp_req.with_temperature(t);
|
||||
}
|
||||
@@ -614,6 +626,7 @@ async fn handle_streaming(
|
||||
}
|
||||
LlmResult::Simple(llm.complete(comp_req).await.map_err(map_llm_error)?)
|
||||
};
|
||||
let model_name = llm.effective_model_name(Some(requested_model.as_str()));
|
||||
|
||||
// LLM succeeded — emit the response as SSE chunks
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, std::convert::Infallible>>(64);
|
||||
@@ -1091,4 +1104,18 @@ mod tests {
|
||||
let v = serde_json::Value::Null;
|
||||
assert_eq!(parse_stop(&v), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_model_name_rejects_leading_or_trailing_whitespace() {
|
||||
let err = validate_model_name(" gpt-4").unwrap_err();
|
||||
assert!(err.contains("leading or trailing whitespace"));
|
||||
|
||||
let err = validate_model_name("gpt-4 ").unwrap_err();
|
||||
assert!(err.contains("leading or trailing whitespace"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_model_name_accepts_normal_name() {
|
||||
assert!(validate_model_name("gpt-4").is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -22,6 +22,7 @@ use serde::Deserialize;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio_stream::StreamExt;
|
||||
use tower_http::cors::{AllowHeaders, CorsLayer};
|
||||
use tower_http::set_header::SetResponseHeaderLayer;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agent::SessionManager;
|
||||
@@ -304,8 +305,16 @@ pub async fn start_server(
|
||||
.merge(statics)
|
||||
.merge(projects)
|
||||
.merge(protected)
|
||||
.layer(cors)
|
||||
.layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body
|
||||
.layer(cors)
|
||||
.layer(SetResponseHeaderLayer::if_not_present(
|
||||
header::X_CONTENT_TYPE_OPTIONS,
|
||||
header::HeaderValue::from_static("nosniff"),
|
||||
))
|
||||
.layer(SetResponseHeaderLayer::if_not_present(
|
||||
header::X_FRAME_OPTIONS,
|
||||
header::HeaderValue::from_static("DENY"),
|
||||
))
|
||||
.with_state(state.clone());
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel();
|
||||
|
||||
@@ -159,6 +159,7 @@ function connectSSE() {
|
||||
|
||||
eventSource.addEventListener('approval_needed', (e) => {
|
||||
const data = JSON.parse(e.data);
|
||||
if (!isCurrentThread(data.thread_id)) return;
|
||||
showApproval(data);
|
||||
});
|
||||
|
||||
|
||||
@@ -137,6 +137,8 @@ pub enum SseEvent {
|
||||
tool_name: String,
|
||||
description: String,
|
||||
parameters: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thread_id: Option<String>,
|
||||
},
|
||||
#[serde(rename = "auth_required")]
|
||||
AuthRequired {
|
||||
@@ -785,12 +787,14 @@ mod tests {
|
||||
tool_name: "shell".to_string(),
|
||||
description: "Run ls".to_string(),
|
||||
parameters: "{}".to_string(),
|
||||
thread_id: Some("t1".to_string()),
|
||||
};
|
||||
let ws = WsServerMessage::from_sse_event(&sse);
|
||||
match ws {
|
||||
WsServerMessage::Event { event_type, data } => {
|
||||
assert_eq!(event_type, "approval_needed");
|
||||
assert_eq!(data["tool_name"], "shell");
|
||||
assert_eq!(data["thread_id"], "t1");
|
||||
}
|
||||
_ => panic!("Expected Event variant"),
|
||||
}
|
||||
|
||||
@@ -9,25 +9,48 @@ use crate::settings::Settings;
|
||||
pub struct EmbeddingsConfig {
|
||||
/// Whether embeddings are enabled.
|
||||
pub enabled: bool,
|
||||
/// Provider to use: "openai" or "nearai"
|
||||
/// Provider to use: "openai", "nearai", or "ollama"
|
||||
pub provider: String,
|
||||
/// OpenAI API key (for OpenAI provider).
|
||||
pub openai_api_key: Option<SecretString>,
|
||||
/// Model to use for embeddings.
|
||||
pub model: String,
|
||||
/// Ollama base URL (for Ollama provider). Defaults to http://localhost:11434.
|
||||
pub ollama_base_url: String,
|
||||
/// Embedding vector dimension. Inferred from the model name when not set explicitly.
|
||||
pub dimension: usize,
|
||||
}
|
||||
|
||||
impl Default for EmbeddingsConfig {
|
||||
fn default() -> Self {
|
||||
let model = "text-embedding-3-small".to_string();
|
||||
let dimension = default_dimension_for_model(&model);
|
||||
Self {
|
||||
enabled: false,
|
||||
provider: "openai".to_string(),
|
||||
openai_api_key: None,
|
||||
model: "text-embedding-3-small".to_string(),
|
||||
model,
|
||||
ollama_base_url: "http://localhost:11434".to_string(),
|
||||
dimension,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Infer the embedding dimension from a well-known model name.
|
||||
///
|
||||
/// Falls back to 1536 (OpenAI text-embedding-3-small default) for unknown models.
|
||||
fn default_dimension_for_model(model: &str) -> usize {
|
||||
match model {
|
||||
"text-embedding-3-small" => 1536,
|
||||
"text-embedding-3-large" => 3072,
|
||||
"text-embedding-ada-002" => 1536,
|
||||
"nomic-embed-text" => 768,
|
||||
"mxbai-embed-large" => 1024,
|
||||
"all-minilm" => 384,
|
||||
_ => 1536,
|
||||
}
|
||||
}
|
||||
|
||||
impl EmbeddingsConfig {
|
||||
pub(crate) fn resolve(settings: &Settings) -> Result<Self, ConfigError> {
|
||||
let openai_api_key = optional_env("OPENAI_API_KEY")?.map(SecretString::from);
|
||||
@@ -38,6 +61,19 @@ impl EmbeddingsConfig {
|
||||
let model =
|
||||
optional_env("EMBEDDING_MODEL")?.unwrap_or_else(|| settings.embeddings.model.clone());
|
||||
|
||||
let ollama_base_url = optional_env("OLLAMA_BASE_URL")?
|
||||
.or_else(|| settings.ollama_base_url.clone())
|
||||
.unwrap_or_else(|| "http://localhost:11434".to_string());
|
||||
|
||||
let dimension = optional_env("EMBEDDING_DIMENSION")?
|
||||
.map(|s| s.parse::<usize>())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "EMBEDDING_DIMENSION".to_string(),
|
||||
message: format!("must be a positive integer: {e}"),
|
||||
})?
|
||||
.unwrap_or_else(|| default_dimension_for_model(&model));
|
||||
|
||||
let enabled = optional_env("EMBEDDING_ENABLED")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
@@ -52,6 +88,8 @@ impl EmbeddingsConfig {
|
||||
provider,
|
||||
openai_api_key,
|
||||
model,
|
||||
ollama_base_url,
|
||||
dimension,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+16
-2
@@ -64,6 +64,8 @@ impl std::fmt::Display for LlmBackend {
|
||||
pub struct OpenAiDirectConfig {
|
||||
pub api_key: SecretString,
|
||||
pub model: String,
|
||||
/// Optional base URL override (e.g. for proxies like VibeProxy).
|
||||
pub base_url: Option<String>,
|
||||
}
|
||||
|
||||
/// Configuration for direct Anthropic API access.
|
||||
@@ -71,6 +73,8 @@ pub struct OpenAiDirectConfig {
|
||||
pub struct AnthropicDirectConfig {
|
||||
pub api_key: SecretString,
|
||||
pub model: String,
|
||||
/// Optional base URL override (e.g. for proxies like VibeProxy).
|
||||
pub base_url: Option<String>,
|
||||
}
|
||||
|
||||
/// Configuration for local Ollama.
|
||||
@@ -274,7 +278,12 @@ impl LlmConfig {
|
||||
hint: "Set OPENAI_API_KEY when LLM_BACKEND=openai".to_string(),
|
||||
})?;
|
||||
let model = optional_env("OPENAI_MODEL")?.unwrap_or_else(|| "gpt-4o".to_string());
|
||||
Some(OpenAiDirectConfig { api_key, model })
|
||||
let base_url = optional_env("OPENAI_BASE_URL")?;
|
||||
Some(OpenAiDirectConfig {
|
||||
api_key,
|
||||
model,
|
||||
base_url,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
@@ -288,7 +297,12 @@ impl LlmConfig {
|
||||
})?;
|
||||
let model = optional_env("ANTHROPIC_MODEL")?
|
||||
.unwrap_or_else(|| "claude-sonnet-4-20250514".to_string());
|
||||
Some(AnthropicDirectConfig { api_key, model })
|
||||
let base_url = optional_env("ANTHROPIC_BASE_URL")?;
|
||||
Some(AnthropicDirectConfig {
|
||||
api_key,
|
||||
model,
|
||||
base_url,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
@@ -30,7 +30,7 @@ impl Default for SandboxModeConfig {
|
||||
timeout_secs: 120,
|
||||
memory_limit_mb: 2048,
|
||||
cpu_shares: 1024,
|
||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
||||
image: "ironclaw-worker:latest".to_string(),
|
||||
auto_pull_image: true,
|
||||
extra_allowed_domains: Vec::new(),
|
||||
}
|
||||
@@ -57,7 +57,7 @@ impl SandboxModeConfig {
|
||||
memory_limit_mb: parse_optional_env("SANDBOX_MEMORY_LIMIT_MB", 2048)?,
|
||||
cpu_shares: parse_optional_env("SANDBOX_CPU_SHARES", 1024)?,
|
||||
image: optional_env("SANDBOX_IMAGE")?
|
||||
.unwrap_or_else(|| "ghcr.io/nearai/sandbox:latest".to_string()),
|
||||
.unwrap_or_else(|| "ironclaw-worker:latest".to_string()),
|
||||
auto_pull_image: optional_env("SANDBOX_AUTO_PULL")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
|
||||
@@ -16,6 +16,7 @@ use crate::extensions::{
|
||||
ActivateResult, AuthResult, ExtensionError, ExtensionKind, ExtensionSource, InstallResult,
|
||||
InstalledExtension, RegistryEntry, ResultSource, SearchResult,
|
||||
};
|
||||
use crate::hooks::HookRegistry;
|
||||
use crate::secrets::{CreateSecretParams, SecretsStore};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::tools::mcp::McpClient;
|
||||
@@ -52,6 +53,7 @@ pub struct ExtensionManager {
|
||||
// Shared
|
||||
secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
hooks: Option<Arc<HookRegistry>>,
|
||||
pending_auth: RwLock<HashMap<String, PendingAuth>>,
|
||||
/// Tunnel URL for remote OAuth callbacks (used in future iterations).
|
||||
_tunnel_url: Option<String>,
|
||||
@@ -66,6 +68,7 @@ impl ExtensionManager {
|
||||
mcp_session_manager: Arc<McpSessionManager>,
|
||||
secrets: Arc<dyn SecretsStore + Send + Sync>,
|
||||
tool_registry: Arc<ToolRegistry>,
|
||||
hooks: Option<Arc<HookRegistry>>,
|
||||
wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
|
||||
wasm_tools_dir: PathBuf,
|
||||
wasm_channels_dir: PathBuf,
|
||||
@@ -83,6 +86,7 @@ impl ExtensionManager {
|
||||
wasm_channels_dir,
|
||||
secrets,
|
||||
tool_registry,
|
||||
hooks,
|
||||
pending_auth: RwLock::new(HashMap::new()),
|
||||
_tunnel_url: tunnel_url,
|
||||
user_id,
|
||||
@@ -320,6 +324,21 @@ impl ExtensionManager {
|
||||
// Unregister from tool registry
|
||||
self.tool_registry.unregister(name).await;
|
||||
|
||||
// Unregister hooks registered from this plugin source.
|
||||
let removed_hooks = self
|
||||
.unregister_hook_prefix(&format!("plugin.tool:{}::", name))
|
||||
.await
|
||||
+ self
|
||||
.unregister_hook_prefix(&format!("plugin.dev_tool:{}::", name))
|
||||
.await;
|
||||
if removed_hooks > 0 {
|
||||
tracing::info!(
|
||||
extension = name,
|
||||
removed_hooks = removed_hooks,
|
||||
"Removed plugin hooks for WASM tool"
|
||||
);
|
||||
}
|
||||
|
||||
// Delete files
|
||||
let wasm_path = self.wasm_tools_dir.join(format!("{}.wasm", name));
|
||||
let cap_path = self
|
||||
@@ -969,6 +988,34 @@ impl ExtensionManager {
|
||||
.await
|
||||
.map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?;
|
||||
|
||||
if let Some(ref hooks) = self.hooks
|
||||
&& let Some(cap_path) = cap_path_option
|
||||
{
|
||||
let source = format!("plugin.tool:{}", name);
|
||||
let registration =
|
||||
crate::hooks::bootstrap::register_plugin_bundle_from_capabilities_file(
|
||||
hooks, &source, cap_path,
|
||||
)
|
||||
.await;
|
||||
|
||||
if registration.total_registered() > 0 {
|
||||
tracing::info!(
|
||||
extension = name,
|
||||
hooks = registration.hooks,
|
||||
outbound_webhooks = registration.outbound_webhooks,
|
||||
"Registered plugin hooks for activated WASM tool"
|
||||
);
|
||||
}
|
||||
|
||||
if registration.errors > 0 {
|
||||
tracing::warn!(
|
||||
extension = name,
|
||||
errors = registration.errors,
|
||||
"Some plugin hooks failed to register"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Activated WASM tool '{}'", name);
|
||||
|
||||
Ok(ActivateResult {
|
||||
@@ -1008,6 +1055,21 @@ impl ExtensionManager {
|
||||
let mut pending = self.pending_auth.write().await;
|
||||
pending.retain(|_, auth| auth.created_at.elapsed() < std::time::Duration::from_secs(300));
|
||||
}
|
||||
|
||||
async fn unregister_hook_prefix(&self, prefix: &str) -> usize {
|
||||
let Some(ref hooks) = self.hooks else {
|
||||
return 0;
|
||||
};
|
||||
|
||||
let names = hooks.list().await;
|
||||
let mut removed = 0;
|
||||
for hook_name in names {
|
||||
if hook_name.starts_with(prefix) && hooks.unregister(&hook_name).await {
|
||||
removed += 1;
|
||||
}
|
||||
}
|
||||
removed
|
||||
}
|
||||
}
|
||||
|
||||
/// Infer the extension kind from a URL.
|
||||
|
||||
@@ -0,0 +1,378 @@
|
||||
//! Hook bootstrap helpers for loading bundled, plugin, and workspace hooks.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::channels::wasm::discover_channels;
|
||||
use crate::hooks::bundled::{
|
||||
HookBundleConfig, HookRegistrationSummary, register_bundle, register_bundled_hooks,
|
||||
};
|
||||
use crate::hooks::registry::HookRegistry;
|
||||
use crate::tools::wasm::{discover_dev_tools, discover_tools};
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// Summary of hook bootstrap work done at startup.
|
||||
#[derive(Debug, Default, Clone, Copy)]
|
||||
pub struct HookBootstrapSummary {
|
||||
/// Number of bundled built-in hooks registered.
|
||||
pub bundled_hooks: usize,
|
||||
/// Number of plugin-provided rule hooks registered.
|
||||
pub plugin_hooks: usize,
|
||||
/// Number of workspace-provided rule hooks registered.
|
||||
pub workspace_hooks: usize,
|
||||
/// Number of outbound webhook hooks registered.
|
||||
pub outbound_webhooks: usize,
|
||||
/// Number of invalid hook configs skipped.
|
||||
pub errors: usize,
|
||||
}
|
||||
|
||||
impl HookBootstrapSummary {
|
||||
/// Total number of hooks registered across all categories.
|
||||
pub fn total_hooks(&self) -> usize {
|
||||
self.bundled_hooks + self.plugin_hooks + self.workspace_hooks + self.outbound_webhooks
|
||||
}
|
||||
}
|
||||
|
||||
/// Register bundled hooks, then load plugin and workspace hook bundles.
|
||||
pub async fn bootstrap_hooks(
|
||||
registry: &Arc<HookRegistry>,
|
||||
workspace: Option<&Arc<Workspace>>,
|
||||
wasm_tools_dir: &Path,
|
||||
wasm_channels_dir: &Path,
|
||||
active_tool_names: &[String],
|
||||
active_channel_names: &[String],
|
||||
dev_loaded_tool_names: &[String],
|
||||
) -> HookBootstrapSummary {
|
||||
let mut summary = HookBootstrapSummary::default();
|
||||
|
||||
let bundled = register_bundled_hooks(registry).await;
|
||||
summary.bundled_hooks += bundled.hooks;
|
||||
summary.outbound_webhooks += bundled.outbound_webhooks;
|
||||
summary.errors += bundled.errors;
|
||||
|
||||
let plugin = register_plugin_bundles(
|
||||
registry,
|
||||
wasm_tools_dir,
|
||||
wasm_channels_dir,
|
||||
active_tool_names,
|
||||
active_channel_names,
|
||||
dev_loaded_tool_names,
|
||||
)
|
||||
.await;
|
||||
summary.plugin_hooks += plugin.hooks;
|
||||
summary.outbound_webhooks += plugin.outbound_webhooks;
|
||||
summary.errors += plugin.errors;
|
||||
|
||||
if let Some(workspace) = workspace {
|
||||
let workspace_loaded = register_workspace_bundles(registry, workspace).await;
|
||||
summary.workspace_hooks += workspace_loaded.hooks;
|
||||
summary.outbound_webhooks += workspace_loaded.outbound_webhooks;
|
||||
summary.errors += workspace_loaded.errors;
|
||||
}
|
||||
|
||||
summary
|
||||
}
|
||||
|
||||
async fn register_plugin_bundles(
|
||||
registry: &Arc<HookRegistry>,
|
||||
wasm_tools_dir: &Path,
|
||||
wasm_channels_dir: &Path,
|
||||
active_tool_names: &[String],
|
||||
active_channel_names: &[String],
|
||||
dev_loaded_tool_names: &[String],
|
||||
) -> HookRegistrationSummary {
|
||||
let mut summary = HookRegistrationSummary::default();
|
||||
let files = collect_plugin_capability_files(
|
||||
wasm_tools_dir,
|
||||
wasm_channels_dir,
|
||||
active_tool_names,
|
||||
active_channel_names,
|
||||
dev_loaded_tool_names,
|
||||
)
|
||||
.await;
|
||||
|
||||
for (source, path) in files {
|
||||
let registered =
|
||||
register_plugin_bundle_from_capabilities_file(registry, &source, &path).await;
|
||||
summary.merge(registered);
|
||||
}
|
||||
|
||||
summary
|
||||
}
|
||||
|
||||
/// Register a plugin hook bundle from a single capabilities file.
|
||||
///
|
||||
/// This is used by startup bootstrap and by runtime extension activation.
|
||||
pub async fn register_plugin_bundle_from_capabilities_file(
|
||||
registry: &Arc<HookRegistry>,
|
||||
source: &str,
|
||||
path: &Path,
|
||||
) -> HookRegistrationSummary {
|
||||
match load_plugin_bundle_from_capabilities_file(path).await {
|
||||
Ok(Some(bundle)) => register_bundle(registry, source, bundle).await,
|
||||
Ok(None) => HookRegistrationSummary::default(),
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
source = source,
|
||||
path = %path.display(),
|
||||
error = %err,
|
||||
"Skipping plugin hook bundle"
|
||||
);
|
||||
HookRegistrationSummary {
|
||||
hooks: 0,
|
||||
outbound_webhooks: 0,
|
||||
errors: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn collect_plugin_capability_files(
|
||||
wasm_tools_dir: &Path,
|
||||
wasm_channels_dir: &Path,
|
||||
active_tool_names: &[String],
|
||||
active_channel_names: &[String],
|
||||
dev_loaded_tool_names: &[String],
|
||||
) -> Vec<(String, PathBuf)> {
|
||||
let mut files: Vec<(String, PathBuf)> = Vec::new();
|
||||
let mut seen: HashSet<String> = HashSet::new();
|
||||
let active_tools: HashSet<&str> = active_tool_names.iter().map(String::as_str).collect();
|
||||
let active_channels: HashSet<&str> = active_channel_names.iter().map(String::as_str).collect();
|
||||
let dev_loaded_tools: HashSet<&str> =
|
||||
dev_loaded_tool_names.iter().map(String::as_str).collect();
|
||||
|
||||
if wasm_tools_dir.exists() {
|
||||
match discover_tools(wasm_tools_dir).await {
|
||||
Ok(tools) => {
|
||||
for (name, tool) in tools {
|
||||
if let Some(path) = tool.capabilities_path
|
||||
&& active_tools.contains(name.as_str())
|
||||
&& !dev_loaded_tools.contains(name.as_str())
|
||||
{
|
||||
insert_unique(&mut files, &mut seen, format!("plugin.tool:{}", name), path);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
path = %wasm_tools_dir.display(),
|
||||
error = %err,
|
||||
"Failed to discover WASM tool capabilities for plugin hooks"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match discover_dev_tools().await {
|
||||
Ok(dev_tools) => {
|
||||
for (name, tool) in dev_tools {
|
||||
if let Some(path) = tool.capabilities_path
|
||||
&& active_tools.contains(name.as_str())
|
||||
&& dev_loaded_tools.contains(name.as_str())
|
||||
{
|
||||
insert_unique(
|
||||
&mut files,
|
||||
&mut seen,
|
||||
format!("plugin.dev_tool:{}", name),
|
||||
path,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::debug!(error = %err, "No dev tool capabilities discovered for plugin hooks");
|
||||
}
|
||||
}
|
||||
|
||||
if wasm_channels_dir.exists() {
|
||||
match discover_channels(wasm_channels_dir).await {
|
||||
Ok(channels) => {
|
||||
for (name, channel) in channels {
|
||||
if let Some(path) = channel.capabilities_path
|
||||
&& active_channels.contains(name.as_str())
|
||||
{
|
||||
insert_unique(
|
||||
&mut files,
|
||||
&mut seen,
|
||||
format!("plugin.channel:{}", name),
|
||||
path,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
path = %wasm_channels_dir.display(),
|
||||
error = %err,
|
||||
"Failed to discover WASM channel capabilities for plugin hooks"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
files.sort_by(|a, b| a.0.cmp(&b.0));
|
||||
files
|
||||
}
|
||||
|
||||
fn insert_unique(
|
||||
files: &mut Vec<(String, PathBuf)>,
|
||||
seen: &mut HashSet<String>,
|
||||
source: String,
|
||||
path: PathBuf,
|
||||
) {
|
||||
let key = path.to_string_lossy().to_string();
|
||||
if seen.insert(key) {
|
||||
files.push((source, path));
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_plugin_bundle_from_capabilities_file(
|
||||
path: &Path,
|
||||
) -> Result<Option<HookBundleConfig>, String> {
|
||||
let bytes = tokio::fs::read(path)
|
||||
.await
|
||||
.map_err(|e| format!("read failed: {e}"))?;
|
||||
|
||||
let value: serde_json::Value =
|
||||
serde_json::from_slice(&bytes).map_err(|e| format!("invalid JSON: {e}"))?;
|
||||
|
||||
let Some(hooks_value) = extract_hooks_section(&value) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
HookBundleConfig::from_value(hooks_value)
|
||||
.map(Some)
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
fn extract_hooks_section(root: &serde_json::Value) -> Option<&serde_json::Value> {
|
||||
root.get("hooks")
|
||||
.or_else(|| root.get("capabilities").and_then(|c| c.get("hooks")))
|
||||
}
|
||||
|
||||
async fn register_workspace_bundles(
|
||||
registry: &Arc<HookRegistry>,
|
||||
workspace: &Arc<Workspace>,
|
||||
) -> HookRegistrationSummary {
|
||||
let mut summary = HookRegistrationSummary::default();
|
||||
|
||||
let paths = match workspace.list_all().await {
|
||||
Ok(paths) => paths,
|
||||
Err(err) => {
|
||||
summary.errors += 1;
|
||||
tracing::warn!(error = %err, "Failed to list workspace paths for hooks");
|
||||
return summary;
|
||||
}
|
||||
};
|
||||
|
||||
let mut hook_paths: Vec<String> = paths
|
||||
.into_iter()
|
||||
.filter(|path| is_workspace_hook_file(path))
|
||||
.collect();
|
||||
hook_paths.sort();
|
||||
|
||||
for path in hook_paths {
|
||||
let doc = match workspace.read(&path).await {
|
||||
Ok(doc) => doc,
|
||||
Err(err) => {
|
||||
summary.errors += 1;
|
||||
tracing::warn!(path = %path, error = %err, "Skipping unreadable workspace hook file");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let parsed: serde_json::Value = match serde_json::from_str(&doc.content) {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
summary.errors += 1;
|
||||
tracing::warn!(path = %path, error = %err, "Workspace hook file is not valid JSON");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let bundle = match parse_workspace_bundle(&parsed) {
|
||||
Ok(bundle) => bundle,
|
||||
Err(err) => {
|
||||
summary.errors += 1;
|
||||
tracing::warn!(path = %path, error = %err, "Skipping invalid workspace hook bundle");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let source = format!("workspace:{}", path);
|
||||
let registered = register_bundle(registry, &source, bundle).await;
|
||||
summary.merge(registered);
|
||||
}
|
||||
|
||||
summary
|
||||
}
|
||||
|
||||
fn parse_workspace_bundle(value: &serde_json::Value) -> Result<HookBundleConfig, String> {
|
||||
if let Some(nested) = value.get("hooks") {
|
||||
HookBundleConfig::from_value(nested).map_err(|e| e.to_string())
|
||||
} else {
|
||||
HookBundleConfig::from_value(value).map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn is_workspace_hook_file(path: &str) -> bool {
|
||||
path == "hooks/hooks.json" || (path.starts_with("hooks/") && path.ends_with(".hook.json"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_extract_hooks_section_from_tool_caps() {
|
||||
let value = serde_json::json!({
|
||||
"http": {"allowlist": []},
|
||||
"hooks": {"rules": []}
|
||||
});
|
||||
|
||||
let extracted = extract_hooks_section(&value).unwrap();
|
||||
assert!(extracted.get("rules").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_hooks_section_from_channel_caps() {
|
||||
let value = serde_json::json!({
|
||||
"type": "channel",
|
||||
"capabilities": {
|
||||
"hooks": {
|
||||
"rules": []
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let extracted = extract_hooks_section(&value).unwrap();
|
||||
assert!(extracted.get("rules").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_workspace_hook_file_filter() {
|
||||
assert!(is_workspace_hook_file("hooks/hooks.json"));
|
||||
assert!(is_workspace_hook_file("hooks/redact.hook.json"));
|
||||
assert!(!is_workspace_hook_file("hooks/readme.md"));
|
||||
assert!(!is_workspace_hook_file("MEMORY.md"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_workspace_bundle_wrapped_hooks() {
|
||||
let value = serde_json::json!({
|
||||
"hooks": {
|
||||
"rules": [
|
||||
{
|
||||
"name": "append-bang",
|
||||
"points": ["beforeInbound"],
|
||||
"append": "!"
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
|
||||
let bundle = parse_workspace_bundle(&value).unwrap();
|
||||
assert_eq!(bundle.rules.len(), 1);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
+20
-3
@@ -3,9 +3,11 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Points in the agent lifecycle where hooks can be attached.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum HookPoint {
|
||||
/// Before processing an inbound user message.
|
||||
BeforeInbound,
|
||||
@@ -21,8 +23,22 @@ pub enum HookPoint {
|
||||
TransformResponse,
|
||||
}
|
||||
|
||||
impl HookPoint {
|
||||
/// Human-readable hook point identifier.
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
HookPoint::BeforeInbound => "beforeInbound",
|
||||
HookPoint::BeforeToolCall => "beforeToolCall",
|
||||
HookPoint::BeforeOutbound => "beforeOutbound",
|
||||
HookPoint::OnSessionStart => "onSessionStart",
|
||||
HookPoint::OnSessionEnd => "onSessionEnd",
|
||||
HookPoint::TransformResponse => "transformResponse",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Contextual data carried with each hook invocation.
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum HookEvent {
|
||||
/// An inbound user message about to be processed.
|
||||
Inbound {
|
||||
@@ -133,7 +149,8 @@ impl HookOutcome {
|
||||
}
|
||||
|
||||
/// How to handle hook execution failures.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum HookFailureMode {
|
||||
/// On error/timeout, continue processing as if the hook returned `ok()`.
|
||||
FailOpen,
|
||||
|
||||
@@ -12,8 +12,14 @@
|
||||
//! Hooks are executed in priority order (lower number = higher priority).
|
||||
//! Each hook can pass through, modify content, or reject the event.
|
||||
|
||||
pub mod bootstrap;
|
||||
pub mod bundled;
|
||||
pub mod hook;
|
||||
pub mod registry;
|
||||
|
||||
pub use bootstrap::{HookBootstrapSummary, bootstrap_hooks};
|
||||
pub use bundled::{
|
||||
HookBundleConfig, HookRegistrationSummary, register_bundle, register_bundled_hooks,
|
||||
};
|
||||
pub use hook::{Hook, HookContext, HookError, HookEvent, HookFailureMode, HookOutcome, HookPoint};
|
||||
pub use registry::HookRegistry;
|
||||
|
||||
+54
-1
@@ -39,7 +39,22 @@ impl HookRegistry {
|
||||
/// Lower priority number = runs first.
|
||||
pub async fn register_with_priority(&self, hook: Arc<dyn Hook>, priority: u32) {
|
||||
let mut hooks = self.hooks.write().await;
|
||||
hooks.push(HookEntry { hook, priority });
|
||||
let hook_name = hook.name().to_string();
|
||||
|
||||
if let Some(existing) = hooks
|
||||
.iter_mut()
|
||||
.find(|entry| entry.hook.name() == hook_name)
|
||||
{
|
||||
tracing::warn!(
|
||||
hook = %hook_name,
|
||||
"Replacing existing hook registration with same name"
|
||||
);
|
||||
existing.hook = hook;
|
||||
existing.priority = priority;
|
||||
} else {
|
||||
hooks.push(HookEntry { hook, priority });
|
||||
}
|
||||
|
||||
hooks.sort_by_key(|e| e.priority);
|
||||
}
|
||||
|
||||
@@ -346,6 +361,44 @@ mod tests {
|
||||
assert_eq!(names, vec!["hook-a", "hook-b"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_register_duplicate_name_replaces_existing() {
|
||||
let registry = HookRegistry::new();
|
||||
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "dup".into(),
|
||||
suffix: "-A".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
100,
|
||||
)
|
||||
.await;
|
||||
|
||||
registry
|
||||
.register_with_priority(
|
||||
Arc::new(ModifyHook {
|
||||
name: "dup".into(),
|
||||
suffix: "-B".into(),
|
||||
points: vec![HookPoint::BeforeInbound],
|
||||
}),
|
||||
10,
|
||||
)
|
||||
.await;
|
||||
|
||||
let names = registry.list().await;
|
||||
assert_eq!(names, vec!["dup"]);
|
||||
|
||||
let result = registry.run(&test_event()).await.unwrap();
|
||||
match result {
|
||||
HookOutcome::Continue {
|
||||
modified: Some(value),
|
||||
} => assert_eq!(value, "hello-B"),
|
||||
other => panic!("expected modified output, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_priority_ordering() {
|
||||
let registry = HookRegistry::new();
|
||||
|
||||
@@ -123,7 +123,11 @@ impl CircuitBreakerProvider {
|
||||
);
|
||||
Ok(())
|
||||
} else {
|
||||
let remaining = self.config.recovery_timeout - opened_at.elapsed();
|
||||
let remaining = self
|
||||
.config
|
||||
.recovery_timeout
|
||||
.checked_sub(opened_at.elapsed())
|
||||
.unwrap_or(Duration::ZERO);
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: self.inner.model_name().to_string(),
|
||||
reason: format!(
|
||||
@@ -208,8 +212,16 @@ impl CircuitBreakerProvider {
|
||||
/// Returns `true` for errors that indicate the provider is degraded
|
||||
/// (server errors, rate limits, network failures, auth infrastructure down).
|
||||
///
|
||||
/// Client errors (wrong model, bad credentials, context overflow) are NOT
|
||||
/// transient: they are the caller's problem, not a sign of backend trouble.
|
||||
/// This answers: "should this error count toward tripping the circuit breaker?"
|
||||
///
|
||||
/// Includes `SessionExpired` because repeated session failures signal backend
|
||||
/// auth infrastructure trouble.
|
||||
///
|
||||
/// Excludes client errors that are the caller's problem, not backend trouble:
|
||||
/// `AuthFailed`, `ContextLengthExceeded`, `ModelNotAvailable`, `Json`.
|
||||
///
|
||||
/// See also `retry::is_retryable()` which answers a different question:
|
||||
/// "could retrying this exact request succeed?"
|
||||
fn is_transient(err: &LlmError) -> bool {
|
||||
matches!(
|
||||
err,
|
||||
@@ -219,7 +231,6 @@ fn is_transient(err: &LlmError) -> bool {
|
||||
| LlmError::SessionExpired { .. }
|
||||
| LlmError::SessionRenewalFailed { .. }
|
||||
| LlmError::Http(_)
|
||||
| LlmError::Json(_)
|
||||
| LlmError::Io(_)
|
||||
)
|
||||
}
|
||||
@@ -273,6 +284,10 @@ impl LlmProvider for CircuitBreakerProvider {
|
||||
self.inner.model_metadata().await
|
||||
}
|
||||
|
||||
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
||||
self.inner.effective_model_name(requested_model)
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.inner.active_model_name()
|
||||
}
|
||||
@@ -543,6 +558,9 @@ mod tests {
|
||||
provider: "p".into(),
|
||||
model: "m".into(),
|
||||
}));
|
||||
assert!(!is_transient(&LlmError::Json(
|
||||
serde_json::from_str::<String>("bad").unwrap_err()
|
||||
)));
|
||||
}
|
||||
|
||||
// -- Passthrough delegation tests --
|
||||
|
||||
+128
-41
@@ -7,8 +7,10 @@
|
||||
//! so subsequent requests skip them, reducing latency when a provider
|
||||
//! is known to be down. Cooldown state is lock-free (atomics only).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -17,34 +19,11 @@ use rust_decimal::Decimal;
|
||||
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::provider::{
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
|
||||
ToolCompletionResponse,
|
||||
};
|
||||
|
||||
/// Returns `true` if the error is transient and the request should be retried
|
||||
/// on the next provider in the failover chain.
|
||||
///
|
||||
/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`,
|
||||
/// `SessionRenewalFailed`, `ModelNotAvailable`, `Http`, `Io`.
|
||||
///
|
||||
/// `ModelNotAvailable` is retryable because the next provider in the chain may
|
||||
/// offer a different model, so it's worth trying.
|
||||
///
|
||||
/// Non-retryable errors (`AuthFailed`, `SessionExpired`, `ContextLengthExceeded`)
|
||||
/// propagate immediately because a different provider won't fix them.
|
||||
fn is_retryable(err: &LlmError) -> bool {
|
||||
matches!(
|
||||
err,
|
||||
LlmError::RequestFailed { .. }
|
||||
| LlmError::RateLimited { .. }
|
||||
| LlmError::InvalidResponse { .. }
|
||||
| LlmError::SessionRenewalFailed { .. }
|
||||
// ModelNotAvailable is retryable: the next provider may offer a different model.
|
||||
| LlmError::ModelNotAvailable { .. }
|
||||
| LlmError::Http(_)
|
||||
| LlmError::Io(_)
|
||||
)
|
||||
}
|
||||
use crate::llm::retry::is_retryable;
|
||||
|
||||
/// Configuration for per-provider cooldown behavior.
|
||||
///
|
||||
@@ -139,6 +118,12 @@ pub struct FailoverProvider {
|
||||
epoch: Instant,
|
||||
/// Cooldown configuration.
|
||||
cooldown_config: CooldownConfig,
|
||||
/// Request-scoped provider index keyed by Tokio task ID.
|
||||
///
|
||||
/// This allows `effective_model_name()` to report the provider that handled
|
||||
/// the *current* request, even when other concurrent requests update
|
||||
/// `last_used`.
|
||||
provider_for_task: Mutex<HashMap<tokio::task::Id, usize>>,
|
||||
}
|
||||
|
||||
impl FailoverProvider {
|
||||
@@ -171,6 +156,7 @@ impl FailoverProvider {
|
||||
cooldowns,
|
||||
epoch: Instant::now(),
|
||||
cooldown_config,
|
||||
provider_for_task: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -182,12 +168,36 @@ impl FailoverProvider {
|
||||
self.epoch.elapsed().as_nanos() as u64
|
||||
}
|
||||
|
||||
/// Current Tokio task ID if available.
|
||||
fn current_task_id() -> Option<tokio::task::Id> {
|
||||
tokio::task::try_id()
|
||||
}
|
||||
|
||||
/// Bind the selected provider index to the current task.
|
||||
fn bind_provider_to_current_task(&self, provider_idx: usize) {
|
||||
let Some(task_id) = Self::current_task_id() else {
|
||||
return;
|
||||
};
|
||||
if let Ok(mut guard) = self.provider_for_task.lock() {
|
||||
guard.insert(task_id, provider_idx);
|
||||
}
|
||||
}
|
||||
|
||||
/// Take and remove the provider index bound to the current task.
|
||||
fn take_bound_provider_for_current_task(&self) -> Option<usize> {
|
||||
let task_id = Self::current_task_id()?;
|
||||
self.provider_for_task
|
||||
.lock()
|
||||
.ok()
|
||||
.and_then(|mut guard| guard.remove(&task_id))
|
||||
}
|
||||
|
||||
/// Try each provider in sequence until one succeeds or all fail.
|
||||
///
|
||||
/// Providers in cooldown are skipped unless *all* providers are in
|
||||
/// cooldown, in which case the one with the oldest cooldown timestamp
|
||||
/// (most likely to have recovered) is tried.
|
||||
async fn try_providers<T, F, Fut>(&self, mut call: F) -> Result<T, LlmError>
|
||||
async fn try_providers<T, F, Fut>(&self, mut call: F) -> Result<(usize, T), LlmError>
|
||||
where
|
||||
F: FnMut(Arc<dyn LlmProvider>) -> Fut,
|
||||
Fut: Future<Output = Result<T, LlmError>>,
|
||||
@@ -236,7 +246,7 @@ impl FailoverProvider {
|
||||
Ok(response) => {
|
||||
self.last_used.store(i, Ordering::Relaxed);
|
||||
self.cooldowns[i].reset();
|
||||
return Ok(response);
|
||||
return Ok((i, response));
|
||||
}
|
||||
Err(err) => {
|
||||
if !is_retryable(&err) {
|
||||
@@ -287,22 +297,28 @@ impl LlmProvider for FailoverProvider {
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
self.try_providers(|provider| {
|
||||
let req = request.clone();
|
||||
async move { provider.complete(req).await }
|
||||
})
|
||||
.await
|
||||
let (provider_idx, response) = self
|
||||
.try_providers(|provider| {
|
||||
let req = request.clone();
|
||||
async move { provider.complete(req).await }
|
||||
})
|
||||
.await?;
|
||||
self.bind_provider_to_current_task(provider_idx);
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
self.try_providers(|provider| {
|
||||
let req = request.clone();
|
||||
async move { provider.complete_with_tools(req).await }
|
||||
})
|
||||
.await
|
||||
let (provider_idx, response) = self
|
||||
.try_providers(|provider| {
|
||||
let req = request.clone();
|
||||
async move { provider.complete_with_tools(req).await }
|
||||
})
|
||||
.await?;
|
||||
self.bind_provider_to_current_task(provider_idx);
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
@@ -336,6 +352,34 @@ impl LlmProvider for FailoverProvider {
|
||||
all_models.dedup();
|
||||
Ok(all_models)
|
||||
}
|
||||
|
||||
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||
.model_metadata()
|
||||
.await
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||
.seed_response_chain(thread_id, response_id);
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)].get_response_chain_id(thread_id)
|
||||
}
|
||||
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||
.calculate_cost(input_tokens, output_tokens)
|
||||
}
|
||||
|
||||
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
||||
if let Some(provider_idx) = self.take_bound_provider_for_current_task() {
|
||||
return self.providers[provider_idx].effective_model_name(requested_model);
|
||||
}
|
||||
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)].effective_model_name(requested_model)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -610,6 +654,49 @@ mod tests {
|
||||
assert_eq!(failover.cost_per_token(), (fallback_cost, fallback_cost));
|
||||
}
|
||||
|
||||
// Test: model reporting is request-scoped under concurrent requests.
|
||||
#[tokio::test]
|
||||
async fn effective_model_name_is_request_scoped_under_concurrency() {
|
||||
let config = CooldownConfig {
|
||||
cooldown_duration: Duration::from_secs(60),
|
||||
failure_threshold: 3,
|
||||
};
|
||||
let primary = Arc::new(MultiCallMockProvider::fail_then_ok("primary", 1));
|
||||
let fallback = Arc::new(MultiCallMockProvider::always_ok("fallback"));
|
||||
let failover =
|
||||
Arc::new(FailoverProvider::with_cooldown(vec![primary, fallback], config).unwrap());
|
||||
|
||||
let (first_done_tx, first_done_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
let (second_done_tx, second_done_rx) = tokio::sync::oneshot::channel::<()>();
|
||||
|
||||
let failover_a = Arc::clone(&failover);
|
||||
let task_a = tokio::spawn(async move {
|
||||
// First request: primary fails once, fallback serves.
|
||||
let _ = failover_a.complete(make_request()).await.unwrap();
|
||||
let _ = first_done_tx.send(());
|
||||
|
||||
// Wait until the second request finishes and updates global state.
|
||||
let _ = second_done_rx.await;
|
||||
failover_a.effective_model_name(None)
|
||||
});
|
||||
|
||||
let failover_b = Arc::clone(&failover);
|
||||
let task_b = tokio::spawn(async move {
|
||||
let _ = first_done_rx.await;
|
||||
// Second request: primary now succeeds.
|
||||
let _ = failover_b.complete(make_request()).await.unwrap();
|
||||
let model = failover_b.effective_model_name(None);
|
||||
let _ = second_done_tx.send(());
|
||||
model
|
||||
});
|
||||
|
||||
let model_b = task_b.await.unwrap();
|
||||
let model_a = task_a.await.unwrap();
|
||||
|
||||
assert_eq!(model_a, "fallback");
|
||||
assert_eq!(model_b, "primary");
|
||||
}
|
||||
|
||||
// Test: list_models aggregates from all providers.
|
||||
#[tokio::test]
|
||||
async fn list_models_aggregates_all() {
|
||||
@@ -1021,10 +1108,6 @@ mod tests {
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset"
|
||||
))));
|
||||
assert!(is_retryable(&LlmError::ModelNotAvailable {
|
||||
provider: "p".into(),
|
||||
model: "m".into(),
|
||||
}));
|
||||
|
||||
// Non-retryable
|
||||
assert!(!is_retryable(&LlmError::AuthFailed {
|
||||
@@ -1037,6 +1120,10 @@ mod tests {
|
||||
used: 100_000,
|
||||
limit: 50_000,
|
||||
}));
|
||||
assert!(!is_retryable(&LlmError::ModelNotAvailable {
|
||||
provider: "p".into(),
|
||||
model: "m".into(),
|
||||
}));
|
||||
}
|
||||
|
||||
// Test: empty providers list returns error (not panic).
|
||||
|
||||
+55
-31
@@ -15,7 +15,7 @@ mod nearai_chat;
|
||||
mod provider;
|
||||
mod reasoning;
|
||||
pub mod response_cache;
|
||||
mod retry;
|
||||
pub mod retry;
|
||||
mod rig_adapter;
|
||||
pub mod session;
|
||||
|
||||
@@ -32,6 +32,7 @@ pub use reasoning::{
|
||||
ToolSelection,
|
||||
};
|
||||
pub use response_cache::{CachedProvider, ResponseCacheConfig};
|
||||
pub use retry::{RetryConfig, RetryProvider};
|
||||
pub use rig_adapter::RigAdapter;
|
||||
pub use session::{SessionConfig, SessionManager, create_session_manager};
|
||||
|
||||
@@ -76,7 +77,7 @@ pub fn create_llm_provider_with_config(
|
||||
model = %config.model,
|
||||
"Using Responses API (chat-api) with session auth"
|
||||
);
|
||||
Ok(Arc::new(NearAiProvider::new(config.clone(), session)))
|
||||
Ok(Arc::new(NearAiProvider::new(config.clone(), session)?))
|
||||
}
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
tracing::info!(
|
||||
@@ -99,15 +100,30 @@ fn create_openai_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, Ll
|
||||
// (Responses API). The Responses API path in rig-core panics when tool results
|
||||
// are sent back because ironclaw doesn't thread `call_id` through its ToolCall
|
||||
// type. The Chat Completions API works correctly with the existing code.
|
||||
let client: openai::CompletionsClient = openai::Client::new(oai.api_key.expose_secret())
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "openai".to_string(),
|
||||
reason: format!("Failed to create OpenAI client: {}", e),
|
||||
})?
|
||||
.completions_api();
|
||||
let client: openai::CompletionsClient = if let Some(ref base_url) = oai.base_url {
|
||||
tracing::info!(
|
||||
"Using OpenAI direct API (chat completions, model: {}, base_url: {})",
|
||||
oai.model,
|
||||
base_url,
|
||||
);
|
||||
openai::Client::builder()
|
||||
.base_url(base_url)
|
||||
.api_key(oai.api_key.expose_secret())
|
||||
.build()
|
||||
} else {
|
||||
tracing::info!(
|
||||
"Using OpenAI direct API (chat completions, model: {}, base_url: default)",
|
||||
oai.model,
|
||||
);
|
||||
openai::Client::new(oai.api_key.expose_secret())
|
||||
}
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "openai".to_string(),
|
||||
reason: format!("Failed to create OpenAI client: {}", e),
|
||||
})?
|
||||
.completions_api();
|
||||
|
||||
let model = client.completion_model(&oai.model);
|
||||
tracing::info!("Using OpenAI direct API (model: {})", oai.model);
|
||||
Ok(Arc::new(RigAdapter::new(model, &oai.model)))
|
||||
}
|
||||
|
||||
@@ -121,16 +137,25 @@ fn create_anthropic_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>,
|
||||
|
||||
use rig::providers::anthropic;
|
||||
|
||||
let client: anthropic::Client =
|
||||
anthropic::Client::new(anth.api_key.expose_secret()).map_err(|e| {
|
||||
LlmError::RequestFailed {
|
||||
provider: "anthropic".to_string(),
|
||||
reason: format!("Failed to create Anthropic client: {}", e),
|
||||
}
|
||||
})?;
|
||||
let client: anthropic::Client = if let Some(ref base_url) = anth.base_url {
|
||||
anthropic::Client::builder()
|
||||
.api_key(anth.api_key.expose_secret())
|
||||
.base_url(base_url)
|
||||
.build()
|
||||
} else {
|
||||
anthropic::Client::new(anth.api_key.expose_secret())
|
||||
}
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "anthropic".to_string(),
|
||||
reason: format!("Failed to create Anthropic client: {}", e),
|
||||
})?;
|
||||
|
||||
let model = client.completion_model(&anth.model);
|
||||
tracing::info!("Using Anthropic direct API (model: {})", anth.model);
|
||||
tracing::info!(
|
||||
"Using Anthropic direct API (model: {}, base_url: {})",
|
||||
anth.model,
|
||||
anth.base_url.as_deref().unwrap_or("default"),
|
||||
);
|
||||
Ok(Arc::new(RigAdapter::new(model, &anth.model)))
|
||||
}
|
||||
|
||||
@@ -199,26 +224,25 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
|
||||
|
||||
use rig::providers::openai;
|
||||
|
||||
let api_key = compat
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|k| k.expose_secret().to_string())
|
||||
.unwrap_or_else(|| "no-key".to_string());
|
||||
|
||||
let client: openai::Client = openai::Client::builder()
|
||||
let client: openai::CompletionsClient = openai::Client::builder()
|
||||
.base_url(&compat.base_url)
|
||||
.api_key(api_key)
|
||||
.api_key(
|
||||
compat
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|k| k.expose_secret().to_string())
|
||||
.unwrap_or_else(|| "no-key".to_string()),
|
||||
)
|
||||
.build()
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "openai_compatible".to_string(),
|
||||
reason: format!("Failed to create OpenAI-compatible client: {}", e),
|
||||
})?;
|
||||
})?
|
||||
.completions_api();
|
||||
|
||||
// OpenAI-compatible providers (e.g. OpenRouter) are most reliable on Chat Completions.
|
||||
// This avoids Responses-API-specific assumptions such as required tool call IDs.
|
||||
let model = client.completions_api().completion_model(&compat.model);
|
||||
let model = client.completion_model(&compat.model);
|
||||
tracing::info!(
|
||||
"Using OpenAI-compatible endpoint via Chat Completions API (base_url: {}, model: {})",
|
||||
"Using OpenAI-compatible endpoint (chat completions, base_url: {}, model: {})",
|
||||
compat.base_url,
|
||||
compat.model
|
||||
);
|
||||
@@ -252,7 +276,7 @@ pub fn create_cheap_llm_provider(
|
||||
tracing::info!("Cheap LLM provider: {}", cheap_model);
|
||||
|
||||
match cheap_config.api_mode {
|
||||
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)))),
|
||||
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)?))),
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
|
||||
}
|
||||
|
||||
+140
-150
@@ -19,7 +19,6 @@ use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
||||
use crate::llm::session::SessionManager;
|
||||
|
||||
/// Information about an available model from NEAR AI API.
|
||||
@@ -54,28 +53,34 @@ pub struct NearAiProvider {
|
||||
|
||||
impl NearAiProvider {
|
||||
/// Create a new NEAR AI provider with a session manager.
|
||||
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Self {
|
||||
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Result<Self, LlmError> {
|
||||
let client = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new());
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to build HTTP client: {}", e),
|
||||
})?;
|
||||
|
||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||
Self {
|
||||
Ok(Self {
|
||||
client,
|
||||
config,
|
||||
session,
|
||||
active_model,
|
||||
response_chains: std::sync::RwLock::new(HashMap::new()),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Seed a response chain for a thread (e.g. when restoring from DB).
|
||||
pub fn seed_response_id(&self, thread_id: &str, response_id: String) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
let mut chains = match self.response_chains.write() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("response_chains lock poisoned in seed; recovering");
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
chains.insert(
|
||||
thread_id.to_string(),
|
||||
ChainState {
|
||||
@@ -87,19 +92,25 @@ impl NearAiProvider {
|
||||
|
||||
/// Get the last response ID for a thread (for persistence).
|
||||
pub fn get_response_id(&self, thread_id: &str) -> Option<String> {
|
||||
let chains = self
|
||||
.response_chains
|
||||
.read()
|
||||
.expect("response_chains lock poisoned");
|
||||
let chains = match self.response_chains.read() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("response_chains lock poisoned in get; recovering");
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
chains.get(thread_id).map(|c| c.response_id.clone())
|
||||
}
|
||||
|
||||
/// Store a response chain state after a successful call.
|
||||
fn store_chain(&self, thread_id: &str, response_id: String, input_count: usize) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
let mut chains = match self.response_chains.write() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("response_chains lock poisoned in store; recovering");
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
chains.insert(
|
||||
thread_id.to_string(),
|
||||
ChainState {
|
||||
@@ -111,10 +122,13 @@ impl NearAiProvider {
|
||||
|
||||
/// Clear the chain for a thread (on error / fallback).
|
||||
fn clear_chain(&self, thread_id: &str) {
|
||||
let mut chains = self
|
||||
.response_chains
|
||||
.write()
|
||||
.expect("response_chains lock poisoned");
|
||||
let mut chains = match self.response_chains.write() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("response_chains lock poisoned in clear; recovering");
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
chains.remove(thread_id);
|
||||
}
|
||||
|
||||
@@ -160,7 +174,10 @@ impl NearAiProvider {
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to read response body: {}", e),
|
||||
})?;
|
||||
|
||||
if !status.is_success() {
|
||||
if status.as_u16() == 401 {
|
||||
@@ -283,139 +300,95 @@ impl NearAiProvider {
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner request implementation with retry logic for transient errors.
|
||||
/// Inner request implementation (single attempt).
|
||||
///
|
||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
||||
/// Does not retry internally — retries are handled by the external
|
||||
/// `RetryProvider` wrapper in the composition chain.
|
||||
async fn send_request_inner<T: Serialize + std::fmt::Debug, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
path: &str,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
let url = self.api_url(path);
|
||||
let max_retries = self.config.max_retries;
|
||||
let token = self.session.get_token().await?;
|
||||
|
||||
for attempt in 0..=max_retries {
|
||||
let token = self.session.get_token().await?;
|
||||
tracing::debug!("Sending request to NEAR AI: {}", url);
|
||||
tracing::debug!("Request body: {:?}", body);
|
||||
|
||||
tracing::debug!(
|
||||
"Sending request to NEAR AI: {} (attempt {})",
|
||||
url,
|
||||
attempt + 1
|
||||
);
|
||||
tracing::debug!("Request body: {:?}", body);
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!("NEAR AI request failed: {}", e);
|
||||
LlmError::Http(e)
|
||||
})?;
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", token.expose_secret()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await;
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Failed to read response body: {}", e),
|
||||
})?;
|
||||
|
||||
let response = match response {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::error!("NEAR AI request failed: {}", e);
|
||||
// Network errors (timeout, connection refused) are transient
|
||||
if attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI request error (attempt {}/{}), retrying in {:?}: {}",
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
e,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
return Err(e.into());
|
||||
}
|
||||
};
|
||||
tracing::debug!("NEAR AI response status: {}", status);
|
||||
tracing::debug!("NEAR AI response body: {}", response_text);
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
tracing::debug!("NEAR AI response status: {}", status);
|
||||
tracing::debug!("NEAR AI response body: {}", response_text);
|
||||
// Check for session expiration (401 with specific message patterns)
|
||||
if status_code == 401 {
|
||||
let lower = response_text.to_lowercase();
|
||||
let is_session_expired = lower.contains("session")
|
||||
&& (lower.contains("expired") || lower.contains("invalid"));
|
||||
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
// Check for session expiration (401 with specific message patterns)
|
||||
if status_code == 401 {
|
||||
let lower = response_text.to_lowercase();
|
||||
let is_session_expired = lower.contains("session")
|
||||
&& (lower.contains("expired") || lower.contains("invalid"));
|
||||
|
||||
if is_session_expired {
|
||||
return Err(LlmError::SessionExpired {
|
||||
provider: "nearai".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Generic 401 -- not retryable
|
||||
return Err(LlmError::AuthFailed {
|
||||
if is_session_expired {
|
||||
return Err(LlmError::SessionExpired {
|
||||
provider: "nearai".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check if this is a transient error worth retrying
|
||||
if is_retryable_status(status_code) && attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
||||
status_code,
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Non-retryable error or exhausted retries
|
||||
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: error.error,
|
||||
});
|
||||
}
|
||||
|
||||
return Err(LlmError::RequestFailed {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
}
|
||||
|
||||
// Success -- parse the response
|
||||
return match serde_json::from_str::<R>(&response_text) {
|
||||
Ok(parsed) => Ok(parsed),
|
||||
Err(e) => {
|
||||
tracing::debug!("Response is not expected JSON format: {}", e);
|
||||
tracing::debug!("Will try alternative parsing in caller");
|
||||
Err(LlmError::InvalidResponse {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
||||
})
|
||||
}
|
||||
};
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
|
||||
if let Ok(error) = serde_json::from_str::<NearAiErrorResponse>(&response_text) {
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: error.error,
|
||||
});
|
||||
}
|
||||
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
}
|
||||
|
||||
// This is unreachable because the loop always returns, but the compiler
|
||||
// cannot prove that. Return a generic error as a safety net.
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
})
|
||||
// Success -- parse the response
|
||||
match serde_json::from_str::<R>(&response_text) {
|
||||
Ok(parsed) => Ok(parsed),
|
||||
Err(e) => {
|
||||
tracing::debug!("Response is not expected JSON format: {}", e);
|
||||
tracing::debug!("Will try alternative parsing in caller");
|
||||
Err(LlmError::InvalidResponse {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("Parse error: {}. Raw: {}", e, response_text),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -462,11 +435,14 @@ fn split_messages(
|
||||
#[async_trait]
|
||||
impl LlmProvider for NearAiProvider {
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||
let thread_id = req.metadata.get("thread_id").cloned();
|
||||
let (instructions, input) = split_messages(req.messages, false);
|
||||
let mut messages = req.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||
let (instructions, input) = split_messages(messages, false);
|
||||
|
||||
let request = NearAiRequest {
|
||||
model: self.active_model_name(),
|
||||
model,
|
||||
instructions,
|
||||
input,
|
||||
previous_response_id: None,
|
||||
@@ -579,14 +555,22 @@ impl LlmProvider for NearAiProvider {
|
||||
&self,
|
||||
req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||
let thread_id = req.metadata.get("thread_id").cloned();
|
||||
let mut messages = req.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||
|
||||
// Look up chaining state for this thread
|
||||
let chain_state = thread_id.as_ref().and_then(|tid| {
|
||||
let chains = self
|
||||
.response_chains
|
||||
.read()
|
||||
.expect("response_chains lock poisoned");
|
||||
let chains = match self.response_chains.read() {
|
||||
Ok(guard) => guard,
|
||||
Err(poisoned) => {
|
||||
tracing::warn!(
|
||||
"response_chains lock poisoned in complete_with_tools; recovering"
|
||||
);
|
||||
poisoned.into_inner()
|
||||
}
|
||||
};
|
||||
chains
|
||||
.get(tid)
|
||||
.map(|c| (c.response_id.clone(), c.input_count))
|
||||
@@ -599,7 +583,7 @@ impl LlmProvider for NearAiProvider {
|
||||
|
||||
// When chaining, only send new messages (the delta since last call).
|
||||
// Tool results are converted to function_call_output items.
|
||||
let (instructions, all_input) = split_messages(req.messages, chaining);
|
||||
let (instructions, all_input) = split_messages(messages, chaining);
|
||||
let input = if chaining && all_input.len() > prev_input_count {
|
||||
all_input[prev_input_count..].to_vec()
|
||||
} else {
|
||||
@@ -619,7 +603,7 @@ impl LlmProvider for NearAiProvider {
|
||||
.collect();
|
||||
|
||||
let request = NearAiRequest {
|
||||
model: self.active_model_name(),
|
||||
model: model.clone(),
|
||||
instructions: if chaining { None } else { instructions.clone() },
|
||||
input,
|
||||
previous_response_id: previous_response_id.clone(),
|
||||
@@ -660,7 +644,7 @@ impl LlmProvider for NearAiProvider {
|
||||
false,
|
||||
);
|
||||
let retry_request = NearAiRequest {
|
||||
model: self.active_model_name(),
|
||||
model,
|
||||
instructions: instructions_full,
|
||||
input: input_full,
|
||||
previous_response_id: None,
|
||||
@@ -804,18 +788,25 @@ impl LlmProvider for NearAiProvider {
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.active_model
|
||||
.read()
|
||||
.expect("active_model lock poisoned")
|
||||
.clone()
|
||||
match self.active_model.read() {
|
||||
Ok(guard) => guard.clone(),
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("active_model lock poisoned while reading; continuing");
|
||||
poisoned.into_inner().clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||
let mut guard = self
|
||||
.active_model
|
||||
.write()
|
||||
.expect("active_model lock poisoned");
|
||||
*guard = model.to_string();
|
||||
match self.active_model.write() {
|
||||
Ok(mut guard) => {
|
||||
*guard = model.to_string();
|
||||
}
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("active_model lock poisoned while writing; continuing");
|
||||
*poisoned.into_inner() = model.to_string();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -950,7 +941,6 @@ struct NearAiTool {
|
||||
/// Primary response format (output array style)
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct NearAiResponse {
|
||||
#[allow(dead_code)]
|
||||
id: String,
|
||||
output: Vec<NearAiOutputItem>,
|
||||
usage: NearAiUsage,
|
||||
|
||||
+201
-127
@@ -16,18 +16,29 @@ use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::llm::retry::{is_retryable_status, retry_backoff_delay};
|
||||
|
||||
/// NEAR AI Chat Completions API provider.
|
||||
pub struct NearAiChatProvider {
|
||||
client: Client,
|
||||
config: NearAiConfig,
|
||||
active_model: std::sync::RwLock<String>,
|
||||
flatten_tool_messages: bool,
|
||||
}
|
||||
|
||||
impl NearAiChatProvider {
|
||||
/// Create a new NEAR AI chat completions provider with API key auth.
|
||||
///
|
||||
/// By default this enables tool-message flattening for compatibility with
|
||||
/// providers that reject `role: "tool"` messages (e.g. NEAR cloud-api).
|
||||
pub fn new(config: NearAiConfig) -> Result<Self, LlmError> {
|
||||
Self::new_with_flatten(config, true)
|
||||
}
|
||||
|
||||
/// Create a chat completions provider with configurable tool-message flattening.
|
||||
pub fn new_with_flatten(
|
||||
config: NearAiConfig,
|
||||
flatten_tool_messages: bool,
|
||||
) -> Result<Self, LlmError> {
|
||||
if config.api_key.is_none() {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
@@ -37,22 +48,29 @@ impl NearAiChatProvider {
|
||||
let client = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new());
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("Failed to build HTTP client: {}", e),
|
||||
})?;
|
||||
|
||||
let active_model = std::sync::RwLock::new(config.model.clone());
|
||||
Ok(Self {
|
||||
client,
|
||||
config,
|
||||
active_model,
|
||||
flatten_tool_messages,
|
||||
})
|
||||
}
|
||||
|
||||
fn api_url(&self, path: &str) -> String {
|
||||
format!(
|
||||
"{}/v1/{}",
|
||||
self.config.base_url,
|
||||
path.trim_start_matches('/')
|
||||
)
|
||||
let base = self.config.base_url.trim_end_matches('/');
|
||||
let path = path.trim_start_matches('/');
|
||||
|
||||
if base.ends_with("/v1") {
|
||||
format!("{}/{}", base, path)
|
||||
} else {
|
||||
format!("{}/v1/{}", base, path)
|
||||
}
|
||||
}
|
||||
|
||||
fn api_key(&self) -> String {
|
||||
@@ -63,116 +81,75 @@ impl NearAiChatProvider {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Send a request to the chat completions API with retry on transient errors.
|
||||
/// Send a single request to the chat completions API.
|
||||
///
|
||||
/// Retries on HTTP 429, 500, 502, 503, 504 with exponential backoff.
|
||||
/// Does not retry on client errors (400, 401, 403, 404) or parse errors.
|
||||
/// Does not retry internally — retries are handled by the external
|
||||
/// `RetryProvider` wrapper in the composition chain.
|
||||
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
let url = self.api_url("chat/completions");
|
||||
let max_retries = self.config.max_retries;
|
||||
|
||||
for attempt in 0..=max_retries {
|
||||
tracing::debug!(
|
||||
"Sending request to NEAR AI Chat: {} (attempt {})",
|
||||
url,
|
||||
attempt + 1,
|
||||
);
|
||||
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
|
||||
|
||||
if tracing::enabled!(tracing::Level::DEBUG)
|
||||
&& let Ok(json) = serde_json::to_string(body)
|
||||
{
|
||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||
}
|
||||
if tracing::enabled!(tracing::Level::DEBUG)
|
||||
&& let Ok(json) = serde_json::to_string(body)
|
||||
{
|
||||
tracing::debug!("NEAR AI Chat request body: {}", json);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await;
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let response = match response {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::error!("NEAR AI Chat request failed: {}", e);
|
||||
if attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI Chat request error (attempt {}/{}), retrying in {:?}: {}",
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
e,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: e.to_string(),
|
||||
});
|
||||
}
|
||||
};
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("Failed to read response body: {}", e),
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
tracing::debug!("NEAR AI Chat response status: {}", status);
|
||||
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
||||
|
||||
tracing::debug!("NEAR AI Chat response status: {}", status);
|
||||
tracing::debug!("NEAR AI Chat response body: {}", response_text);
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
if !status.is_success() {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
// Auth errors are not retryable
|
||||
if status_code == 401 {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Transient errors: retry with backoff
|
||||
if is_retryable_status(status_code) && attempt < max_retries {
|
||||
let delay = retry_backoff_delay(attempt);
|
||||
tracing::warn!(
|
||||
"NEAR AI Chat returned HTTP {} (attempt {}/{}), retrying in {:?}",
|
||||
status_code,
|
||||
attempt + 1,
|
||||
max_retries + 1,
|
||||
delay,
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Non-retryable or exhausted retries
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai_chat".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
return Err(LlmError::RequestFailed {
|
||||
if status_code == 401 {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
});
|
||||
}
|
||||
|
||||
// Success — parse the response
|
||||
return serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
|
||||
if status_code == 429 {
|
||||
return Err(LlmError::RateLimited {
|
||||
provider: "nearai_chat".to_string(),
|
||||
retry_after: None,
|
||||
});
|
||||
}
|
||||
|
||||
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("JSON parse error: {}. Raw: {}", e, response_text),
|
||||
reason: format!("HTTP {}: {}", status, truncated),
|
||||
});
|
||||
}
|
||||
|
||||
// Safety net: unreachable because the loop always returns
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
serde_json::from_str(&response_text).map_err(|e| {
|
||||
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||
LlmError::InvalidResponse {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("JSON parse error: {}. Raw: {}", e, truncated),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -192,12 +169,16 @@ impl NearAiChatProvider {
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let response_text = response.text().await.unwrap_or_default();
|
||||
let response_text = response.text().await.map_err(|e| LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("Failed to read response body: {}", e),
|
||||
})?;
|
||||
|
||||
if !status.is_success() {
|
||||
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("HTTP {}: {}", status, response_text),
|
||||
reason: format!("HTTP {}: {}", status, truncated),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -227,11 +208,14 @@ struct ApiModelEntry {
|
||||
#[async_trait]
|
||||
impl LlmProvider for NearAiChatProvider {
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||
let mut raw_messages = req.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
||||
let messages: Vec<ChatCompletionMessage> =
|
||||
req.messages.into_iter().map(|m| m.into()).collect();
|
||||
raw_messages.into_iter().map(|m| m.into()).collect();
|
||||
|
||||
let request = ChatCompletionRequest {
|
||||
model: self.active_model_name(),
|
||||
model,
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
max_tokens: req.max_tokens,
|
||||
@@ -260,11 +244,13 @@ impl LlmProvider for NearAiChatProvider {
|
||||
_ => FinishReason::Unknown,
|
||||
};
|
||||
|
||||
let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref());
|
||||
|
||||
Ok(CompletionResponse {
|
||||
content,
|
||||
finish_reason,
|
||||
input_tokens: response.usage.prompt_tokens,
|
||||
output_tokens: response.usage.completion_tokens,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
@@ -273,14 +259,19 @@ impl LlmProvider for NearAiChatProvider {
|
||||
&self,
|
||||
req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let model = req.model.unwrap_or_else(|| self.active_model_name());
|
||||
let mut raw_messages = req.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut raw_messages);
|
||||
let messages: Vec<ChatCompletionMessage> =
|
||||
req.messages.into_iter().map(|m| m.into()).collect();
|
||||
raw_messages.into_iter().map(|m| m.into()).collect();
|
||||
|
||||
// NEAR AI cloud-api does not support multi-turn tool calling (rejects
|
||||
// any request containing role:"tool" messages with HTTP 400). Rewrite
|
||||
// tool-call / tool-result pairs into plain text so the conversation
|
||||
// history is preserved without using unsupported message roles.
|
||||
let messages = flatten_tool_messages(messages);
|
||||
// Some OpenAI-compatible providers reject `role:"tool"` messages.
|
||||
// When enabled, rewrite tool-call / tool-result pairs into plain text.
|
||||
let messages = if self.flatten_tool_messages {
|
||||
flatten_tool_messages(messages)
|
||||
} else {
|
||||
messages
|
||||
};
|
||||
|
||||
let tools: Vec<ChatCompletionTool> = req
|
||||
.tools
|
||||
@@ -296,7 +287,7 @@ impl LlmProvider for NearAiChatProvider {
|
||||
.collect();
|
||||
|
||||
let request = ChatCompletionRequest {
|
||||
model: self.active_model_name(),
|
||||
model,
|
||||
messages,
|
||||
temperature: req.temperature,
|
||||
max_tokens: req.max_tokens,
|
||||
@@ -347,12 +338,14 @@ impl LlmProvider for NearAiChatProvider {
|
||||
}
|
||||
};
|
||||
|
||||
let (input_tokens, output_tokens) = parse_usage(response.usage.as_ref());
|
||||
|
||||
Ok(ToolCompletionResponse {
|
||||
content,
|
||||
tool_calls,
|
||||
finish_reason,
|
||||
input_tokens: response.usage.prompt_tokens,
|
||||
output_tokens: response.usage.completion_tokens,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
@@ -382,18 +375,25 @@ impl LlmProvider for NearAiChatProvider {
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.active_model
|
||||
.read()
|
||||
.expect("active_model lock poisoned")
|
||||
.clone()
|
||||
match self.active_model.read() {
|
||||
Ok(guard) => guard.clone(),
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("active_model lock poisoned while reading; continuing");
|
||||
poisoned.into_inner().clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn set_model(&self, model: &str) -> Result<(), crate::error::LlmError> {
|
||||
let mut guard = self
|
||||
.active_model
|
||||
.write()
|
||||
.expect("active_model lock poisoned");
|
||||
*guard = model.to_string();
|
||||
match self.active_model.write() {
|
||||
Ok(mut guard) => {
|
||||
*guard = model.to_string();
|
||||
}
|
||||
Err(poisoned) => {
|
||||
tracing::warn!("active_model lock poisoned while writing; continuing");
|
||||
*poisoned.into_inner() = model.to_string();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -543,9 +543,11 @@ struct ChatCompletionFunction {
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ChatCompletionResponse {
|
||||
#[allow(dead_code)]
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
choices: Vec<ChatCompletionChoice>,
|
||||
usage: ChatCompletionUsage,
|
||||
#[serde(default)]
|
||||
usage: Option<ChatCompletionUsage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -577,18 +579,90 @@ struct ChatCompletionToolCallFunction {
|
||||
arguments: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Deserialize, Default)]
|
||||
struct ChatCompletionUsage {
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
#[allow(dead_code)]
|
||||
total_tokens: u32,
|
||||
#[serde(default)]
|
||||
prompt_tokens: Option<u64>,
|
||||
#[serde(default)]
|
||||
completion_tokens: Option<u64>,
|
||||
#[serde(default)]
|
||||
total_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
fn saturate_u32(val: u64) -> u32 {
|
||||
val.min(u32::MAX as u64) as u32
|
||||
}
|
||||
|
||||
fn parse_usage(usage: Option<&ChatCompletionUsage>) -> (u32, u32) {
|
||||
let Some(u) = usage else {
|
||||
return (0, 0);
|
||||
};
|
||||
let input = u.prompt_tokens.map(saturate_u32).unwrap_or(0);
|
||||
let output = u.completion_tokens.map(saturate_u32).unwrap_or_else(|| {
|
||||
// Fall back to total - prompt if completion is missing.
|
||||
match (u.total_tokens, u.prompt_tokens) {
|
||||
(Some(total), Some(prompt)) => saturate_u32(total.saturating_sub(prompt)),
|
||||
(Some(total), None) => saturate_u32(total),
|
||||
_ => 0,
|
||||
}
|
||||
});
|
||||
(input, output)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_nearai_config(base_url: &str) -> NearAiConfig {
|
||||
NearAiConfig {
|
||||
model: "test-model".to_string(),
|
||||
base_url: base_url.to_string(),
|
||||
auth_base_url: "https://private.near.ai".to_string(),
|
||||
session_path: std::path::PathBuf::from("/tmp/session.json"),
|
||||
api_mode: crate::config::NearAiApiMode::ChatCompletions,
|
||||
api_key: Some(secrecy::SecretString::from("test-key".to_string())),
|
||||
cheap_model: None,
|
||||
fallback_model: None,
|
||||
max_retries: 0,
|
||||
circuit_breaker_threshold: None,
|
||||
circuit_breaker_recovery_secs: 30,
|
||||
response_cache_enabled: false,
|
||||
response_cache_ttl_secs: 3600,
|
||||
response_cache_max_entries: 1000,
|
||||
failover_cooldown_secs: 300,
|
||||
failover_cooldown_threshold: 3,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_api_url_with_base_without_v1() {
|
||||
let mut cfg = test_nearai_config("http://127.0.0.1:8318");
|
||||
|
||||
let provider = NearAiChatProvider::new(cfg.clone()).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
);
|
||||
|
||||
cfg.base_url = "http://127.0.0.1:8318/".to_string();
|
||||
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("/chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_api_url_with_base_already_v1() {
|
||||
let cfg = test_nearai_config("http://127.0.0.1:8318/v1");
|
||||
|
||||
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_conversion() {
|
||||
let msg = ChatMessage::user("Hello");
|
||||
|
||||
@@ -105,6 +105,8 @@ impl ChatMessage {
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
/// Optional per-request model override.
|
||||
pub model: Option<String>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
@@ -117,6 +119,7 @@ impl CompletionRequest {
|
||||
pub fn new(messages: Vec<ChatMessage>) -> Self {
|
||||
Self {
|
||||
messages,
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
@@ -124,6 +127,12 @@ impl CompletionRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// Set model override.
|
||||
pub fn with_model(mut self, model: impl Into<String>) -> Self {
|
||||
self.model = Some(model.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set max tokens.
|
||||
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
|
||||
self.max_tokens = Some(max_tokens);
|
||||
@@ -188,6 +197,8 @@ pub struct ToolResult {
|
||||
pub struct ToolCompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub tools: Vec<ToolDefinition>,
|
||||
/// Optional per-request model override.
|
||||
pub model: Option<String>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
/// How to handle tool use: "auto", "required", or "none".
|
||||
@@ -202,6 +213,7 @@ impl ToolCompletionRequest {
|
||||
Self {
|
||||
messages,
|
||||
tools,
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
tool_choice: None,
|
||||
@@ -209,6 +221,12 @@ impl ToolCompletionRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// Set model override.
|
||||
pub fn with_model(mut self, model: impl Into<String>) -> Self {
|
||||
self.model = Some(model.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set max tokens.
|
||||
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
|
||||
self.max_tokens = Some(max_tokens);
|
||||
@@ -283,6 +301,16 @@ pub trait LlmProvider: Send + Sync {
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve which model should be reported for a given request.
|
||||
///
|
||||
/// Providers that ignore per-request model overrides should override this
|
||||
/// and return `active_model_name()`.
|
||||
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
||||
requested_model
|
||||
.map(std::borrow::ToOwned::to_owned)
|
||||
.unwrap_or_else(|| self.active_model_name())
|
||||
}
|
||||
|
||||
/// Get the currently active model name.
|
||||
///
|
||||
/// May differ from `model_name()` if the model was switched at runtime
|
||||
@@ -319,3 +347,124 @@ pub trait LlmProvider: Send + Sync {
|
||||
input_cost * Decimal::from(input_tokens) + output_cost * Decimal::from(output_tokens)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sanitize a message list to ensure tool_use / tool_result integrity.
|
||||
///
|
||||
/// LLM APIs (especially Anthropic) require every tool_result to reference a
|
||||
/// tool_call_id that exists in an immediately preceding assistant message's
|
||||
/// tool_calls. Orphaned tool_results cause HTTP 400 errors.
|
||||
///
|
||||
/// This function:
|
||||
/// 1. Tracks all tool_call_ids emitted by assistant messages.
|
||||
/// 2. Rewrites orphaned tool_result messages (whose tool_call_id has no
|
||||
/// matching assistant tool_call) as user messages so the content is
|
||||
/// preserved without violating the protocol.
|
||||
///
|
||||
/// Call this before sending messages to any LLM provider.
|
||||
pub fn sanitize_tool_messages(messages: &mut [ChatMessage]) {
|
||||
use std::collections::HashSet;
|
||||
|
||||
// Collect all tool_call_ids from assistant messages with tool_calls.
|
||||
let mut known_ids: HashSet<String> = HashSet::new();
|
||||
for msg in messages.iter() {
|
||||
if msg.role == Role::Assistant
|
||||
&& let Some(ref calls) = msg.tool_calls
|
||||
{
|
||||
for tc in calls {
|
||||
known_ids.insert(tc.id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rewrite orphaned tool_result messages as user messages.
|
||||
for msg in messages.iter_mut() {
|
||||
if msg.role != Role::Tool {
|
||||
continue;
|
||||
}
|
||||
let is_orphaned = match &msg.tool_call_id {
|
||||
Some(id) => !known_ids.contains(id),
|
||||
None => true,
|
||||
};
|
||||
if is_orphaned {
|
||||
let tool_name = msg.name.as_deref().unwrap_or("unknown");
|
||||
tracing::debug!(
|
||||
tool_call_id = ?msg.tool_call_id,
|
||||
tool_name,
|
||||
"Rewriting orphaned tool_result as user message",
|
||||
);
|
||||
msg.role = Role::User;
|
||||
msg.content = format!("[Tool `{}` returned: {}]", tool_name, msg.content);
|
||||
msg.tool_call_id = None;
|
||||
msg.name = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_preserves_valid_pairs() {
|
||||
let tc = ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "echo".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
};
|
||||
let mut messages = vec![
|
||||
ChatMessage::user("hello"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||
ChatMessage::tool_result("call_1", "echo", "result"),
|
||||
];
|
||||
sanitize_tool_messages(&mut messages);
|
||||
assert_eq!(messages[2].role, Role::Tool);
|
||||
assert_eq!(messages[2].tool_call_id, Some("call_1".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_rewrites_orphaned_tool_result() {
|
||||
let mut messages = vec![
|
||||
ChatMessage::user("hello"),
|
||||
ChatMessage::assistant("I'll use a tool"),
|
||||
ChatMessage::tool_result("call_missing", "search", "some result"),
|
||||
];
|
||||
sanitize_tool_messages(&mut messages);
|
||||
assert_eq!(messages[2].role, Role::User);
|
||||
assert!(messages[2].content.contains("[Tool `search` returned:"));
|
||||
assert!(messages[2].tool_call_id.is_none());
|
||||
assert!(messages[2].name.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_handles_no_tool_messages() {
|
||||
let mut messages = vec![
|
||||
ChatMessage::system("prompt"),
|
||||
ChatMessage::user("hello"),
|
||||
ChatMessage::assistant("hi"),
|
||||
];
|
||||
let original_len = messages.len();
|
||||
sanitize_tool_messages(&mut messages);
|
||||
assert_eq!(messages.len(), original_len);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sanitize_multiple_orphaned() {
|
||||
let tc = ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "echo".to_string(),
|
||||
arguments: serde_json::json!({}),
|
||||
};
|
||||
let mut messages = vec![
|
||||
ChatMessage::user("test"),
|
||||
ChatMessage::assistant_with_tool_calls(None, vec![tc]),
|
||||
ChatMessage::tool_result("call_1", "echo", "ok"),
|
||||
// These are orphaned (call_2 and call_3 have no matching assistant message)
|
||||
ChatMessage::tool_result("call_2", "search", "orphan 1"),
|
||||
ChatMessage::tool_result("call_3", "http", "orphan 2"),
|
||||
];
|
||||
sanitize_tool_messages(&mut messages);
|
||||
assert_eq!(messages[2].role, Role::Tool); // call_1 is valid
|
||||
assert_eq!(messages[3].role, Role::User); // call_2 orphaned
|
||||
assert_eq!(messages[4].role, Role::User); // call_3 orphaned
|
||||
}
|
||||
}
|
||||
|
||||
@@ -144,7 +144,8 @@ impl LlmProvider for CachedProvider {
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let key = cache_key(self.inner.model_name(), &request);
|
||||
let effective_model = self.inner.effective_model_name(request.model.as_deref());
|
||||
let key = cache_key(&effective_model, &request);
|
||||
let now = Instant::now();
|
||||
|
||||
// Check cache
|
||||
@@ -216,6 +217,10 @@ impl LlmProvider for CachedProvider {
|
||||
self.inner.model_metadata().await
|
||||
}
|
||||
|
||||
fn effective_model_name(&self, requested_model: Option<&str>) -> String {
|
||||
self.inner.effective_model_name(requested_model)
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.inner.active_model_name()
|
||||
}
|
||||
@@ -242,6 +247,7 @@ mod tests {
|
||||
fn simple_request() -> CompletionRequest {
|
||||
CompletionRequest {
|
||||
messages: vec![ChatMessage::user("hello")],
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
@@ -252,6 +258,7 @@ mod tests {
|
||||
fn different_request() -> CompletionRequest {
|
||||
CompletionRequest {
|
||||
messages: vec![ChatMessage::user("goodbye")],
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
@@ -378,6 +385,7 @@ mod tests {
|
||||
// Add a third: should evict the oldest
|
||||
let third = CompletionRequest {
|
||||
messages: vec![ChatMessage::user("third")],
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
stop_sequences: None,
|
||||
@@ -396,6 +404,7 @@ mod tests {
|
||||
let req = ToolCompletionRequest {
|
||||
messages: vec![ChatMessage::user("use tool")],
|
||||
tools: vec![],
|
||||
model: None,
|
||||
max_tokens: None,
|
||||
temperature: None,
|
||||
tool_choice: None,
|
||||
@@ -444,6 +453,23 @@ mod tests {
|
||||
assert!(cached.is_empty().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_override_gets_distinct_cache_entries() {
|
||||
let stub = Arc::new(StubLlm::new("cached response"));
|
||||
let cached = CachedProvider::new(stub.clone(), ResponseCacheConfig::default());
|
||||
|
||||
let mut req_a = simple_request();
|
||||
req_a.model = Some("model-a".to_string());
|
||||
let mut req_b = simple_request();
|
||||
req_b.model = Some("model-b".to_string());
|
||||
|
||||
cached.complete(req_a).await.unwrap();
|
||||
cached.complete(req_b).await.unwrap();
|
||||
|
||||
assert_eq!(stub.calls(), 2);
|
||||
assert_eq!(cached.len().await, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_config_is_reasonable() {
|
||||
let cfg = ResponseCacheConfig::default();
|
||||
|
||||
+334
-24
@@ -1,15 +1,50 @@
|
||||
//! Shared retry helpers for LLM providers.
|
||||
//! Shared retry helpers and composable `RetryProvider` decorator for LLM providers.
|
||||
//!
|
||||
//! Provides exponential backoff with jitter and retryable status classification
|
||||
//! used by both `NearAiProvider` and `NearAiChatProvider`.
|
||||
//! Provides:
|
||||
//! - `is_retryable()` — `LlmError`-level retryability classification (shared with `failover.rs`)
|
||||
//! - `retry_backoff_delay()` — exponential backoff with jitter
|
||||
//! - `RetryProvider` — decorator that wraps any `LlmProvider` with automatic retries
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rand::Rng;
|
||||
use rust_decimal::Decimal;
|
||||
|
||||
/// Returns `true` if the HTTP status code is transient and worth retrying.
|
||||
pub(crate) fn is_retryable_status(status: u16) -> bool {
|
||||
matches!(status, 429 | 500 | 502 | 503 | 504)
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::provider::{
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ModelMetadata, ToolCompletionRequest,
|
||||
ToolCompletionResponse,
|
||||
};
|
||||
|
||||
/// Returns `true` if the `LlmError` is transient and the request should be retried.
|
||||
///
|
||||
/// Used by `RetryProvider` (retry the same provider) and `FailoverProvider`
|
||||
/// (try the next provider). The question is: "could this exact same request
|
||||
/// succeed if we try again?"
|
||||
///
|
||||
/// Retryable: `RequestFailed`, `RateLimited`, `InvalidResponse`,
|
||||
/// `SessionRenewalFailed`, `Http`, `Io`.
|
||||
///
|
||||
/// Non-retryable: `AuthFailed`, `SessionExpired`, `ContextLengthExceeded`,
|
||||
/// `ModelNotAvailable`, `Json`.
|
||||
/// - `SessionExpired` — handled by session renewal layer, not by retry
|
||||
/// - `ModelNotAvailable` — the model won't appear between attempts
|
||||
/// - `Json` — a serde parse bug, not a transient failure
|
||||
///
|
||||
/// See also `circuit_breaker::is_transient()` which answers a different
|
||||
/// question: "does this error indicate the backend is degraded?"
|
||||
pub(crate) fn is_retryable(err: &LlmError) -> bool {
|
||||
matches!(
|
||||
err,
|
||||
LlmError::RequestFailed { .. }
|
||||
| LlmError::RateLimited { .. }
|
||||
| LlmError::InvalidResponse { .. }
|
||||
| LlmError::SessionRenewalFailed { .. }
|
||||
| LlmError::Http(_)
|
||||
| LlmError::Io(_)
|
||||
)
|
||||
}
|
||||
|
||||
/// Calculate exponential backoff delay with random jitter.
|
||||
@@ -31,31 +66,183 @@ pub(crate) fn retry_backoff_delay(attempt: u32) -> Duration {
|
||||
Duration::from_millis(delay_ms)
|
||||
}
|
||||
|
||||
/// Configuration for the retry decorator.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RetryConfig {
|
||||
/// Maximum number of retry attempts (not counting the initial attempt).
|
||||
/// Default: 3.
|
||||
pub max_retries: u32,
|
||||
}
|
||||
|
||||
impl Default for RetryConfig {
|
||||
fn default() -> Self {
|
||||
Self { max_retries: 3 }
|
||||
}
|
||||
}
|
||||
|
||||
/// Composable decorator that wraps any `LlmProvider` with automatic retries.
|
||||
///
|
||||
/// On transient errors, sleeps using exponential backoff and retries.
|
||||
/// On non-transient errors (`AuthFailed`, `ContextLengthExceeded`, `SessionExpired`),
|
||||
/// returns immediately.
|
||||
///
|
||||
/// Special handling for `RateLimited { retry_after }`: uses the provider-suggested
|
||||
/// duration if available, otherwise falls back to standard backoff.
|
||||
pub struct RetryProvider {
|
||||
inner: Arc<dyn LlmProvider>,
|
||||
config: RetryConfig,
|
||||
}
|
||||
|
||||
impl RetryProvider {
|
||||
pub fn new(inner: Arc<dyn LlmProvider>, config: RetryConfig) -> Self {
|
||||
Self { inner, config }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for RetryProvider {
|
||||
fn model_name(&self) -> &str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
self.inner.cost_per_token()
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let mut last_error: Option<LlmError> = None;
|
||||
|
||||
for attempt in 0..=self.config.max_retries {
|
||||
let req = request.clone();
|
||||
match self.inner.complete(req).await {
|
||||
Ok(resp) => return Ok(resp),
|
||||
Err(err) => {
|
||||
if !is_retryable(&err) || attempt == self.config.max_retries {
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
let delay = match &err {
|
||||
LlmError::RateLimited {
|
||||
retry_after: Some(duration),
|
||||
..
|
||||
} => *duration,
|
||||
_ => retry_backoff_delay(attempt),
|
||||
};
|
||||
|
||||
tracing::warn!(
|
||||
provider = %self.inner.model_name(),
|
||||
attempt = attempt + 1,
|
||||
max_retries = self.config.max_retries,
|
||||
delay_ms = delay.as_millis() as u64,
|
||||
error = %err,
|
||||
"Retrying after transient error"
|
||||
);
|
||||
|
||||
last_error = Some(err);
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| LlmError::RequestFailed {
|
||||
provider: self.inner.model_name().to_string(),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let mut last_error: Option<LlmError> = None;
|
||||
|
||||
for attempt in 0..=self.config.max_retries {
|
||||
let req = request.clone();
|
||||
match self.inner.complete_with_tools(req).await {
|
||||
Ok(resp) => return Ok(resp),
|
||||
Err(err) => {
|
||||
if !is_retryable(&err) || attempt == self.config.max_retries {
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
let delay = match &err {
|
||||
LlmError::RateLimited {
|
||||
retry_after: Some(duration),
|
||||
..
|
||||
} => *duration,
|
||||
_ => retry_backoff_delay(attempt),
|
||||
};
|
||||
|
||||
tracing::warn!(
|
||||
provider = %self.inner.model_name(),
|
||||
attempt = attempt + 1,
|
||||
max_retries = self.config.max_retries,
|
||||
delay_ms = delay.as_millis() as u64,
|
||||
error = %err,
|
||||
"Retrying after transient error (tools)"
|
||||
);
|
||||
|
||||
last_error = Some(err);
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or_else(|| LlmError::RequestFailed {
|
||||
provider: self.inner.model_name().to_string(),
|
||||
reason: "retry loop exited unexpectedly".to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
self.inner.list_models().await
|
||||
}
|
||||
|
||||
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||
self.inner.model_metadata().await
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
self.inner.active_model_name()
|
||||
}
|
||||
|
||||
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||
self.inner.set_model(model)
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.inner.seed_response_chain(thread_id, response_id)
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.inner.get_response_chain_id(thread_id)
|
||||
}
|
||||
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
self.inner.calculate_cost(input_tokens, output_tokens)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_is_retryable_status() {
|
||||
// Transient errors should be retryable
|
||||
assert!(is_retryable_status(429));
|
||||
assert!(is_retryable_status(500));
|
||||
assert!(is_retryable_status(502));
|
||||
assert!(is_retryable_status(503));
|
||||
assert!(is_retryable_status(504));
|
||||
use crate::testing::StubLlm;
|
||||
|
||||
// Client errors should not be retryable
|
||||
assert!(!is_retryable_status(400));
|
||||
assert!(!is_retryable_status(401));
|
||||
assert!(!is_retryable_status(403));
|
||||
assert!(!is_retryable_status(404));
|
||||
assert!(!is_retryable_status(422));
|
||||
|
||||
// Success codes should not be retryable
|
||||
assert!(!is_retryable_status(200));
|
||||
assert!(!is_retryable_status(201));
|
||||
fn make_request() -> CompletionRequest {
|
||||
CompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")])
|
||||
}
|
||||
|
||||
fn make_tool_request() -> ToolCompletionRequest {
|
||||
ToolCompletionRequest::new(vec![crate::llm::ChatMessage::user("hello")], vec![])
|
||||
}
|
||||
|
||||
fn fast_config(max_retries: u32) -> RetryConfig {
|
||||
RetryConfig { max_retries }
|
||||
}
|
||||
|
||||
// -- Backoff delay tests --
|
||||
|
||||
#[test]
|
||||
fn test_retry_backoff_delay_exponential_growth() {
|
||||
// Run multiple samples to verify the range, accounting for jitter
|
||||
@@ -93,4 +280,127 @@ mod tests {
|
||||
let delay = retry_backoff_delay(30);
|
||||
assert!(delay.as_millis() >= 100);
|
||||
}
|
||||
|
||||
// -- is_retryable() classification tests --
|
||||
|
||||
#[test]
|
||||
fn test_is_retryable_classification() {
|
||||
// Retryable
|
||||
assert!(is_retryable(&LlmError::RequestFailed {
|
||||
provider: "p".into(),
|
||||
reason: "err".into(),
|
||||
}));
|
||||
assert!(is_retryable(&LlmError::RateLimited {
|
||||
provider: "p".into(),
|
||||
retry_after: None,
|
||||
}));
|
||||
assert!(is_retryable(&LlmError::InvalidResponse {
|
||||
provider: "p".into(),
|
||||
reason: "bad".into(),
|
||||
}));
|
||||
assert!(is_retryable(&LlmError::SessionRenewalFailed {
|
||||
provider: "p".into(),
|
||||
reason: "timeout".into(),
|
||||
}));
|
||||
assert!(is_retryable(&LlmError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::ConnectionReset,
|
||||
"reset"
|
||||
))));
|
||||
|
||||
// NOT retryable
|
||||
assert!(!is_retryable(&LlmError::AuthFailed {
|
||||
provider: "p".into(),
|
||||
}));
|
||||
assert!(!is_retryable(&LlmError::SessionExpired {
|
||||
provider: "p".into(),
|
||||
}));
|
||||
assert!(!is_retryable(&LlmError::ContextLengthExceeded {
|
||||
used: 100_000,
|
||||
limit: 50_000,
|
||||
}));
|
||||
assert!(!is_retryable(&LlmError::ModelNotAvailable {
|
||||
provider: "p".into(),
|
||||
model: "m".into(),
|
||||
}));
|
||||
}
|
||||
|
||||
// -- RetryProvider tests --
|
||||
|
||||
#[tokio::test]
|
||||
async fn success_on_first_attempt() {
|
||||
let stub = Arc::new(StubLlm::new("ok").with_model_name("test"));
|
||||
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||
|
||||
let resp = retry.complete(make_request()).await;
|
||||
assert!(resp.is_ok());
|
||||
assert_eq!(resp.unwrap().content, "ok");
|
||||
assert_eq!(stub.calls(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retries_transient_errors_then_succeeds() {
|
||||
// StubLlm starts failing, then we flip it to succeed.
|
||||
// With max_retries=2, it will try 3 times total.
|
||||
let stub = Arc::new(StubLlm::failing("test"));
|
||||
let retry = RetryProvider::new(stub.clone(), fast_config(2));
|
||||
|
||||
// Spawn a task that flips the stub to succeed after a short delay
|
||||
let stub_clone = stub.clone();
|
||||
tokio::spawn(async move {
|
||||
// Wait for at least 1 retry attempt (backoff is ~1s, so 1.5s should be enough)
|
||||
tokio::time::sleep(Duration::from_millis(1500)).await;
|
||||
stub_clone.set_failing(false);
|
||||
});
|
||||
|
||||
let resp = retry.complete(make_request()).await;
|
||||
assert!(resp.is_ok());
|
||||
// Should have called at least twice (first fail, then succeed after flip)
|
||||
assert!(stub.calls() >= 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_transient_error_fails_immediately() {
|
||||
let stub = Arc::new(StubLlm::failing_non_transient("test"));
|
||||
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||
|
||||
let err = retry.complete(make_request()).await.unwrap_err();
|
||||
assert!(matches!(err, LlmError::ContextLengthExceeded { .. }));
|
||||
// Should only be called once — no retries for non-transient errors
|
||||
assert_eq!(stub.calls(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exhausts_retries_then_returns_error() {
|
||||
let stub = Arc::new(StubLlm::failing("test"));
|
||||
// max_retries=0 means only the initial attempt, no retries
|
||||
let retry = RetryProvider::new(stub.clone(), fast_config(0));
|
||||
|
||||
let err = retry.complete(make_request()).await.unwrap_err();
|
||||
assert!(matches!(err, LlmError::RequestFailed { .. }));
|
||||
assert_eq!(stub.calls(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn complete_with_tools_retries_same_as_complete() {
|
||||
let stub = Arc::new(StubLlm::failing_non_transient("test"));
|
||||
let retry = RetryProvider::new(stub.clone(), fast_config(3));
|
||||
|
||||
let err = retry
|
||||
.complete_with_tools(make_tool_request())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, LlmError::ContextLengthExceeded { .. }));
|
||||
assert_eq!(stub.calls(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn passthrough_methods_delegate_to_inner() {
|
||||
let stub = Arc::new(StubLlm::new("ok").with_model_name("my-model"));
|
||||
let retry = RetryProvider::new(stub, fast_config(3));
|
||||
|
||||
assert_eq!(retry.model_name(), "my-model");
|
||||
assert_eq!(retry.active_model_name(), "my-model");
|
||||
assert_eq!(retry.cost_per_token(), (Decimal::ZERO, Decimal::ZERO));
|
||||
assert_eq!(retry.calculate_cost(100, 50), Decimal::ZERO);
|
||||
}
|
||||
}
|
||||
|
||||
+97
-3
@@ -18,6 +18,8 @@ use serde::Serialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::costs;
|
||||
use crate::llm::provider::{
|
||||
@@ -404,7 +406,19 @@ where
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let (preamble, history) = convert_messages(&request.messages);
|
||||
if let Some(requested_model) = request.model.as_deref()
|
||||
&& requested_model != self.model_name.as_str()
|
||||
{
|
||||
tracing::warn!(
|
||||
requested_model = requested_model,
|
||||
active_model = %self.model_name,
|
||||
"Per-request model override is not supported for this provider; using configured model"
|
||||
);
|
||||
}
|
||||
|
||||
let mut messages = request.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
|
||||
let rig_req = build_rig_request(
|
||||
preamble,
|
||||
@@ -439,7 +453,22 @@ where
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let (preamble, history) = convert_messages(&request.messages);
|
||||
if let Some(requested_model) = request.model.as_deref()
|
||||
&& requested_model != self.model_name.as_str()
|
||||
{
|
||||
tracing::warn!(
|
||||
requested_model = requested_model,
|
||||
active_model = %self.model_name,
|
||||
"Per-request model override is not supported for this provider; using configured model"
|
||||
);
|
||||
}
|
||||
|
||||
let known_tool_names: HashSet<String> =
|
||||
request.tools.iter().map(|t| t.name.clone()).collect();
|
||||
|
||||
let mut messages = request.messages;
|
||||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||||
let (preamble, history) = convert_messages(&messages);
|
||||
let tools = convert_tools(&request.tools);
|
||||
let tool_choice = convert_tool_choice(request.tool_choice.as_deref());
|
||||
|
||||
@@ -461,7 +490,20 @@ where
|
||||
reason: e.to_string(),
|
||||
})?;
|
||||
|
||||
let (text, tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||
let (text, mut tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||||
|
||||
// Normalize tool call names: some proxies prepend "proxy_" prefixes.
|
||||
for tc in &mut tool_calls {
|
||||
let normalized = normalize_tool_name(&tc.name, &known_tool_names);
|
||||
if normalized != tc.name {
|
||||
tracing::debug!(
|
||||
original = %tc.name,
|
||||
normalized = %normalized,
|
||||
"Normalized tool call name from provider",
|
||||
);
|
||||
tc.name = normalized;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ToolCompletionResponse {
|
||||
content: text,
|
||||
@@ -477,6 +519,10 @@ where
|
||||
self.model_name.clone()
|
||||
}
|
||||
|
||||
fn effective_model_name(&self, _requested_model: Option<&str>) -> String {
|
||||
self.active_model_name()
|
||||
}
|
||||
|
||||
fn set_model(&self, _model: &str) -> Result<(), LlmError> {
|
||||
// rig-core models are baked at construction time.
|
||||
// Switching requires creating a new adapter.
|
||||
@@ -489,6 +535,25 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// Normalize a tool call name returned by an OpenAI-compatible provider.
|
||||
///
|
||||
/// Some proxies (e.g. VibeProxy) prepend `proxy_` to tool names.
|
||||
/// If the returned name doesn't match any known tool but stripping a
|
||||
/// `proxy_` prefix yields a match, use the stripped version.
|
||||
fn normalize_tool_name(name: &str, known_tools: &HashSet<String>) -> String {
|
||||
if known_tools.contains(name) {
|
||||
return name.to_string();
|
||||
}
|
||||
|
||||
if let Some(stripped) = name.strip_prefix("proxy_")
|
||||
&& known_tools.contains(stripped)
|
||||
{
|
||||
return stripped.to_string();
|
||||
}
|
||||
|
||||
name.to_string()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -777,4 +842,33 @@ mod tests {
|
||||
assert_eq!(saturate_u32(u64::MAX), u32::MAX);
|
||||
assert_eq!(saturate_u32(u32::MAX as u64), u32::MAX);
|
||||
}
|
||||
|
||||
// -- normalize_tool_name tests --
|
||||
|
||||
#[test]
|
||||
fn test_normalize_tool_name_exact_match() {
|
||||
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||
assert_eq!(normalize_tool_name("echo", &known), "echo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_tool_name_proxy_prefix_match() {
|
||||
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||
assert_eq!(normalize_tool_name("proxy_echo", &known), "echo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_tool_name_proxy_prefix_no_match_kept() {
|
||||
let known = HashSet::from(["echo".to_string(), "list_jobs".to_string()]);
|
||||
assert_eq!(
|
||||
normalize_tool_name("proxy_unknown", &known),
|
||||
"proxy_unknown"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_tool_name_unknown_passthrough() {
|
||||
let known = HashSet::from(["echo".to_string()]);
|
||||
assert_eq!(normalize_tool_name("other_tool", &known), "other_tool");
|
||||
}
|
||||
}
|
||||
|
||||
+142
-29
@@ -23,12 +23,12 @@ use ironclaw::{
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
extensions::ExtensionManager,
|
||||
hooks::HookRegistry,
|
||||
hooks::{HookRegistry, bootstrap_hooks},
|
||||
llm::{
|
||||
CachedProvider, CircuitBreakerConfig, CircuitBreakerProvider, CooldownConfig,
|
||||
FailoverProvider, LlmProvider, ResponseCacheConfig, SessionConfig,
|
||||
create_cheap_llm_provider, create_llm_provider, create_llm_provider_with_config,
|
||||
create_session_manager,
|
||||
FailoverProvider, LlmProvider, ResponseCacheConfig, RetryConfig, RetryProvider,
|
||||
SessionConfig, create_cheap_llm_provider, create_llm_provider,
|
||||
create_llm_provider_with_config, create_session_manager,
|
||||
},
|
||||
orchestrator::{
|
||||
ContainerJobConfig, ContainerJobManager, OrchestratorApi, TokenStore,
|
||||
@@ -42,7 +42,9 @@ use ironclaw::{
|
||||
mcp::{McpClient, McpSessionManager, config::load_mcp_servers_from_db, is_authenticated},
|
||||
wasm::{WasmToolLoader, WasmToolRuntime, load_dev_tools},
|
||||
},
|
||||
workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
|
||||
workspace::{
|
||||
EmbeddingProvider, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings, Workspace,
|
||||
},
|
||||
};
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
@@ -115,18 +117,20 @@ async fn main() -> anyhow::Result<()> {
|
||||
&config.llm.nearai.base_url,
|
||||
session,
|
||||
)
|
||||
.with_model(&config.embeddings.model, 1536),
|
||||
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||
)),
|
||||
"ollama" => Some(Arc::new(
|
||||
ironclaw::workspace::OllamaEmbeddings::new(
|
||||
&config.embeddings.ollama_base_url,
|
||||
)
|
||||
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||
)),
|
||||
_ => {
|
||||
if let Some(api_key) = config.embeddings.openai_api_key() {
|
||||
let dim = match config.embeddings.model.as_str() {
|
||||
"text-embedding-3-large" => 3072,
|
||||
_ => 1536,
|
||||
};
|
||||
Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model(
|
||||
api_key,
|
||||
&config.embeddings.model,
|
||||
dim,
|
||||
config.embeddings.dimension,
|
||||
)))
|
||||
} else {
|
||||
None
|
||||
@@ -137,6 +141,23 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
// Warn if libSQL backend is used with non-1536 embedding dimension.
|
||||
// libSQL schema uses F32_BLOB(1536) which cannot be altered without a
|
||||
// table rebuild, so non-1536 embeddings will cause storage failures.
|
||||
if config.database.backend == ironclaw::config::DatabaseBackend::LibSql
|
||||
&& config.embeddings.enabled
|
||||
&& config.embeddings.dimension != 1536
|
||||
{
|
||||
tracing::warn!(
|
||||
configured_dimension = config.embeddings.dimension,
|
||||
"Embedding dimension {} is not 1536. The libSQL schema uses \
|
||||
F32_BLOB(1536) which requires exactly 1536 dimensions. \
|
||||
Embedding storage will fail. Use PostgreSQL or set \
|
||||
EMBEDDING_DIMENSION=1536.",
|
||||
config.embeddings.dimension
|
||||
);
|
||||
}
|
||||
|
||||
// Create a Database-trait-backed workspace for the memory command
|
||||
let db: Arc<dyn ironclaw::db::Database> =
|
||||
ironclaw::db::connect_from_config(&config.database)
|
||||
@@ -599,6 +620,22 @@ async fn main() -> anyhow::Result<()> {
|
||||
let llm = create_llm_provider(&config.llm, session.clone())?;
|
||||
tracing::info!("LLM provider initialized: {}", llm.model_name());
|
||||
|
||||
// Wrap each provider with RetryProvider for automatic retries on transient errors.
|
||||
// RetryProvider sits inside FailoverProvider so each provider in the failover chain
|
||||
// gets its own retry attempts before the failover moves to the next provider.
|
||||
let retry_config = RetryConfig {
|
||||
max_retries: config.llm.nearai.max_retries,
|
||||
};
|
||||
let llm: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
|
||||
tracing::info!(
|
||||
max_retries = retry_config.max_retries,
|
||||
"LLM retry wrapper enabled"
|
||||
);
|
||||
Arc::new(RetryProvider::new(llm, retry_config.clone()))
|
||||
} else {
|
||||
llm
|
||||
};
|
||||
|
||||
// Wrap in failover if a fallback model is configured
|
||||
let llm: Arc<dyn LlmProvider> =
|
||||
if let Some(fallback_model) = config.llm.nearai.fallback_model.as_ref() {
|
||||
@@ -615,6 +652,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
fallback = %fallback.model_name(),
|
||||
"LLM failover enabled"
|
||||
);
|
||||
// Wrap fallback with retry too
|
||||
let fallback: Arc<dyn LlmProvider> = if retry_config.max_retries > 0 {
|
||||
Arc::new(RetryProvider::new(fallback, retry_config.clone()))
|
||||
} else {
|
||||
fallback
|
||||
};
|
||||
let cooldown_config = CooldownConfig {
|
||||
cooldown_duration: std::time::Duration::from_secs(
|
||||
config.llm.nearai.failover_cooldown_secs,
|
||||
@@ -685,28 +728,39 @@ async fn main() -> anyhow::Result<()> {
|
||||
match config.embeddings.provider.as_str() {
|
||||
"nearai" => {
|
||||
tracing::info!(
|
||||
"Embeddings enabled via NEAR AI (model: {})",
|
||||
config.embeddings.model
|
||||
"Embeddings enabled via NEAR AI (model: {}, dim: {})",
|
||||
config.embeddings.model,
|
||||
config.embeddings.dimension,
|
||||
);
|
||||
Some(Arc::new(
|
||||
NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone())
|
||||
.with_model(&config.embeddings.model, 1536),
|
||||
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||
))
|
||||
}
|
||||
"ollama" => {
|
||||
tracing::info!(
|
||||
"Embeddings enabled via Ollama (model: {}, url: {}, dim: {})",
|
||||
config.embeddings.model,
|
||||
config.embeddings.ollama_base_url,
|
||||
config.embeddings.dimension,
|
||||
);
|
||||
Some(Arc::new(
|
||||
OllamaEmbeddings::new(&config.embeddings.ollama_base_url)
|
||||
.with_model(&config.embeddings.model, config.embeddings.dimension),
|
||||
))
|
||||
}
|
||||
_ => {
|
||||
// Default to OpenAI for unknown providers
|
||||
if let Some(api_key) = config.embeddings.openai_api_key() {
|
||||
tracing::info!(
|
||||
"Embeddings enabled via OpenAI (model: {})",
|
||||
config.embeddings.model
|
||||
"Embeddings enabled via OpenAI (model: {}, dim: {})",
|
||||
config.embeddings.model,
|
||||
config.embeddings.dimension,
|
||||
);
|
||||
Some(Arc::new(OpenAiEmbeddings::with_model(
|
||||
api_key,
|
||||
&config.embeddings.model,
|
||||
match config.embeddings.model.as_str() {
|
||||
"text-embedding-3-large" => 3072,
|
||||
_ => 1536, // text-embedding-3-small and ada-002
|
||||
},
|
||||
config.embeddings.dimension,
|
||||
)))
|
||||
} else {
|
||||
tracing::warn!("Embeddings configured but OPENAI_API_KEY not set");
|
||||
@@ -719,6 +773,21 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
// Warn if libSQL backend is used with non-1536 embedding dimension.
|
||||
if config.database.backend == ironclaw::config::DatabaseBackend::LibSql
|
||||
&& config.embeddings.enabled
|
||||
&& config.embeddings.dimension != 1536
|
||||
{
|
||||
tracing::warn!(
|
||||
configured_dimension = config.embeddings.dimension,
|
||||
"Embedding dimension {} is not 1536. The libSQL schema uses \
|
||||
F32_BLOB(1536) which requires exactly 1536 dimensions. \
|
||||
Embedding storage will fail. Use PostgreSQL or set \
|
||||
EMBEDDING_DIMENSION=1536.",
|
||||
config.embeddings.dimension
|
||||
);
|
||||
}
|
||||
|
||||
// Register memory tools if database is available
|
||||
if let Some(ref db) = db {
|
||||
let mut workspace = Workspace::new_with_db("default", Arc::clone(db));
|
||||
@@ -746,6 +815,9 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
|
||||
// Create hook registry early so runtime extension activation can register hooks.
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
|
||||
// Create WASM tool runtime (sync, just builds the wasmtime engine)
|
||||
let wasm_tool_runtime: Option<Arc<WasmToolRuntime>> =
|
||||
if config.wasm.enabled && config.wasm.tools_dir.exists() {
|
||||
@@ -763,6 +835,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Load WASM tools and MCP servers concurrently.
|
||||
// Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe.
|
||||
let wasm_tools_future = async {
|
||||
let mut dev_loaded_tool_names: Vec<String> = Vec::new();
|
||||
|
||||
if let Some(ref runtime) = wasm_tool_runtime {
|
||||
let mut loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&tools));
|
||||
if let Some(ref secrets) = secrets_store {
|
||||
@@ -791,6 +865,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Load dev tools from build artifacts (overrides installed if newer)
|
||||
match load_dev_tools(&loader, &config.wasm.tools_dir).await {
|
||||
Ok(results) => {
|
||||
dev_loaded_tool_names.extend(results.loaded.iter().cloned());
|
||||
if !results.loaded.is_empty() {
|
||||
tracing::info!(
|
||||
"Loaded {} dev WASM tools from build artifacts",
|
||||
@@ -803,6 +878,8 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dev_loaded_tool_names
|
||||
};
|
||||
|
||||
let mcp_servers_future = async {
|
||||
@@ -908,7 +985,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
};
|
||||
|
||||
tokio::join!(wasm_tools_future, mcp_servers_future);
|
||||
let (dev_loaded_tool_names, _) = tokio::join!(wasm_tools_future, mcp_servers_future);
|
||||
|
||||
// Create extension manager for in-chat discovery/install/auth/activate
|
||||
let extension_manager = if let Some(ref secrets) = secrets_store {
|
||||
@@ -916,6 +993,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
Arc::clone(&mcp_session_manager),
|
||||
Arc::clone(secrets),
|
||||
Arc::clone(&tools),
|
||||
Some(Arc::clone(&hooks)),
|
||||
wasm_tool_runtime.clone(),
|
||||
config.wasm.tools_dir.clone(),
|
||||
config.channels.wasm_channels_dir.clone(),
|
||||
@@ -1013,6 +1091,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Initialize channel manager
|
||||
let mut channels = ChannelManager::new();
|
||||
let mut channel_names: Vec<String> = Vec::new();
|
||||
let mut loaded_wasm_channel_names: Vec<String> = Vec::new();
|
||||
|
||||
if let Some(repl) = repl_channel {
|
||||
channels.add(Box::new(repl));
|
||||
@@ -1045,6 +1124,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
for loaded in results.loaded {
|
||||
let channel_name = loaded.name().to_string();
|
||||
loaded_wasm_channel_names.push(channel_name.clone());
|
||||
tracing::info!("Loaded WASM channel: {}", channel_name);
|
||||
|
||||
let secret_name = loaded.webhook_secret_name();
|
||||
@@ -1213,6 +1293,13 @@ async fn main() -> anyhow::Result<()> {
|
||||
let mut webhook_server = if !webhook_routes.is_empty() {
|
||||
let addr =
|
||||
webhook_server_addr.unwrap_or_else(|| std::net::SocketAddr::from(([0, 0, 0, 0], 8080)));
|
||||
if addr.ip().is_unspecified() {
|
||||
tracing::warn!(
|
||||
"Webhook server is binding to {} — it will be reachable from all network interfaces. \
|
||||
Set HTTP_HOST=127.0.0.1 to restrict to localhost.",
|
||||
addr.ip()
|
||||
);
|
||||
}
|
||||
let mut server = WebhookServer::new(WebhookServerConfig { addr });
|
||||
for routes in webhook_routes {
|
||||
server.add_routes(routes);
|
||||
@@ -1263,8 +1350,27 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Create context manager (shared between job tools and agent)
|
||||
let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs));
|
||||
|
||||
// Create hook registry
|
||||
let hooks = Arc::new(HookRegistry::new());
|
||||
// Register bundled/plugin/workspace hooks.
|
||||
let active_tool_names = tools.list().await;
|
||||
|
||||
let hook_bootstrap = bootstrap_hooks(
|
||||
&hooks,
|
||||
workspace.as_ref(),
|
||||
&config.wasm.tools_dir,
|
||||
&config.channels.wasm_channels_dir,
|
||||
&active_tool_names,
|
||||
&loaded_wasm_channel_names,
|
||||
&dev_loaded_tool_names,
|
||||
)
|
||||
.await;
|
||||
tracing::info!(
|
||||
bundled = hook_bootstrap.bundled_hooks,
|
||||
plugin = hook_bootstrap.plugin_hooks,
|
||||
workspace = hook_bootstrap.workspace_hooks,
|
||||
outbound_webhooks = hook_bootstrap.outbound_webhooks,
|
||||
errors = hook_bootstrap.errors,
|
||||
"Lifecycle hooks initialized"
|
||||
);
|
||||
|
||||
// Create session manager (shared between agent and web gateway)
|
||||
let session_manager = Arc::new(SessionManager::new().with_hooks(hooks.clone()));
|
||||
@@ -1305,7 +1411,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Add web gateway channel if configured
|
||||
let mut gateway_url: Option<String> = None;
|
||||
if let Some(ref gw_config) = config.channels.gateway {
|
||||
let mut gw = GatewayChannel::new(gw_config.clone());
|
||||
let mut gw = GatewayChannel::new(gw_config.clone()).with_llm_provider(Arc::clone(&llm));
|
||||
if let Some(ref ws) = workspace {
|
||||
gw = gw.with_workspace(Arc::clone(ws));
|
||||
}
|
||||
@@ -1469,14 +1575,21 @@ fn check_onboard_needed() -> Option<&'static str> {
|
||||
return Some("Database not configured");
|
||||
}
|
||||
|
||||
// The wizard writes ONBOARD_COMPLETED=true to ~/.ironclaw/.env,
|
||||
// which load_ironclaw_env() loads before this function runs.
|
||||
if std::env::var("ONBOARD_COMPLETED")
|
||||
.map(|v| v == "true")
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
// First run (onboarding never completed and no session).
|
||||
// Reads NEARAI_API_KEY env var directly because this function runs
|
||||
// before Config is loaded -- Config::from_env() may fail without a
|
||||
// database URL, which is what triggers onboarding in the first place.
|
||||
// Check for a NEAR AI API key or session file as a fallback
|
||||
// for users who configured credentials manually (no wizard).
|
||||
if std::env::var("NEARAI_API_KEY").is_err() {
|
||||
let settings = ironclaw::settings::Settings::load();
|
||||
let session_path = ironclaw::llm::session::default_session_path();
|
||||
if !settings.onboard_completed && !session_path.exists() {
|
||||
if !session_path.exists() {
|
||||
return Some("First run");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -143,6 +143,7 @@ async fn llm_complete(
|
||||
) -> Result<Json<ProxyCompletionResponse>, StatusCode> {
|
||||
let completion_req = CompletionRequest {
|
||||
messages: req.messages,
|
||||
model: req.model,
|
||||
max_tokens: req.max_tokens,
|
||||
temperature: req.temperature,
|
||||
stop_sequences: req.stop_sequences,
|
||||
@@ -170,6 +171,7 @@ async fn llm_complete_with_tools(
|
||||
let tool_req = ToolCompletionRequest {
|
||||
messages: req.messages,
|
||||
tools: req.tools,
|
||||
model: req.model,
|
||||
max_tokens: req.max_tokens,
|
||||
temperature: req.temperature,
|
||||
tool_choice: req.tool_choice,
|
||||
|
||||
@@ -34,7 +34,7 @@ impl Default for SandboxConfig {
|
||||
memory_limit_mb: 2048,
|
||||
cpu_shares: 1024,
|
||||
network_allowlist: default_allowlist(),
|
||||
image: "ghcr.io/nearai/sandbox:latest".to_string(),
|
||||
image: "ironclaw-worker:latest".to_string(),
|
||||
auto_pull_image: true,
|
||||
proxy_port: 0,
|
||||
}
|
||||
|
||||
+1
-1
@@ -462,7 +462,7 @@ fn default_sandbox_cpu_shares() -> u32 {
|
||||
}
|
||||
|
||||
fn default_sandbox_image() -> String {
|
||||
"ghcr.io/nearai/sandbox:latest".to_string()
|
||||
"ironclaw-worker:latest".to_string()
|
||||
}
|
||||
|
||||
impl Default for SandboxSettings {
|
||||
|
||||
+11
-9
@@ -19,8 +19,9 @@ Explicit invocation. Loads `.env` files, runs the wizard, exits.
|
||||
ironclaw (first run, no database configured)
|
||||
```
|
||||
|
||||
Auto-detection via `check_onboard_needed()` in `main.rs`. Triggers when
|
||||
none of these are true:
|
||||
Auto-detection via `check_onboard_needed()` in `main.rs`. Skips onboarding
|
||||
when `ONBOARD_COMPLETED` env var is set (written to `~/.ironclaw/.env` by
|
||||
the wizard). Otherwise triggers when no database is configured:
|
||||
- `DATABASE_URL` env var is set
|
||||
- `LIBSQL_PATH` env var is set
|
||||
- `~/.ironclaw/ironclaw.db` exists on disk
|
||||
@@ -345,13 +346,14 @@ Final step of the wizard:
|
||||
1. Mark onboard_completed = true
|
||||
2. Write ALL settings to database (try postgres pool, then libSQL backend)
|
||||
3. Write bootstrap vars to ~/.ironclaw/.env:
|
||||
- DATABASE_BACKEND (always)
|
||||
- DATABASE_URL (if postgres)
|
||||
- LIBSQL_PATH (if libsql)
|
||||
- LIBSQL_URL (if turso sync)
|
||||
- LLM_BACKEND (always, when set)
|
||||
- LLM_BASE_URL (if openai_compatible)
|
||||
- OLLAMA_BASE_URL (if ollama)
|
||||
- DATABASE_BACKEND (always)
|
||||
- DATABASE_URL (if postgres)
|
||||
- LIBSQL_PATH (if libsql)
|
||||
- LIBSQL_URL (if turso sync)
|
||||
- LLM_BACKEND (always, when set)
|
||||
- LLM_BASE_URL (if openai_compatible)
|
||||
- OLLAMA_BASE_URL (if ollama)
|
||||
- ONBOARD_COMPLETED (always, "true")
|
||||
4. Print configuration summary
|
||||
```
|
||||
|
||||
|
||||
@@ -1542,6 +1542,10 @@ impl SetupWizard {
|
||||
env_vars.push(("OLLAMA_BASE_URL", url.clone()));
|
||||
}
|
||||
|
||||
// Always write ONBOARD_COMPLETED so that check_onboard_needed()
|
||||
// (which runs before the DB is connected) knows to skip re-onboarding.
|
||||
env_vars.push(("ONBOARD_COMPLETED", "true".to_string()));
|
||||
|
||||
if !env_vars.is_empty() {
|
||||
let pairs: Vec<(&str, &str)> =
|
||||
env_vars.iter().map(|(k, v)| (*k, v.as_str())).collect();
|
||||
|
||||
@@ -570,6 +570,7 @@ mod tests {
|
||||
Arc::new(InMemorySecretsStore::new(crypto)),
|
||||
Arc::new(ToolRegistry::new()),
|
||||
None,
|
||||
None,
|
||||
std::path::PathBuf::from("/tmp/ironclaw-test-tools"),
|
||||
std::path::PathBuf::from("/tmp/ironclaw-test-channels"),
|
||||
None,
|
||||
|
||||
@@ -40,6 +40,7 @@ pub struct JobDescription {
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct ProxyCompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub model: Option<String>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
@@ -57,6 +58,7 @@ pub struct ProxyCompletionResponse {
|
||||
pub struct ProxyToolCompletionRequest {
|
||||
pub messages: Vec<ChatMessage>,
|
||||
pub tools: Vec<ToolDefinition>,
|
||||
pub model: Option<String>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub temperature: Option<f32>,
|
||||
pub tool_choice: Option<String>,
|
||||
@@ -210,6 +212,7 @@ impl WorkerHttpClient {
|
||||
) -> Result<CompletionResponse, WorkerError> {
|
||||
let proxy_req = ProxyCompletionRequest {
|
||||
messages: request.messages.clone(),
|
||||
model: request.model.clone(),
|
||||
max_tokens: request.max_tokens,
|
||||
temperature: request.temperature,
|
||||
stop_sequences: request.stop_sequences.clone(),
|
||||
@@ -236,6 +239,7 @@ impl WorkerHttpClient {
|
||||
let proxy_req = ProxyToolCompletionRequest {
|
||||
messages: request.messages.clone(),
|
||||
tools: request.tools.clone(),
|
||||
model: request.model.clone(),
|
||||
max_tokens: request.max_tokens,
|
||||
temperature: request.temperature,
|
||||
tool_choice: request.tool_choice.clone(),
|
||||
|
||||
@@ -354,6 +354,123 @@ impl EmbeddingProvider for NearAiEmbeddings {
|
||||
}
|
||||
}
|
||||
|
||||
/// Ollama embedding provider using a local Ollama instance.
|
||||
///
|
||||
/// Ollama serves embedding models (e.g. `nomic-embed-text`, `mxbai-embed-large`)
|
||||
/// via a REST API, typically at `http://localhost:11434`.
|
||||
pub struct OllamaEmbeddings {
|
||||
client: reqwest::Client,
|
||||
base_url: String,
|
||||
model: String,
|
||||
dimension: usize,
|
||||
}
|
||||
|
||||
impl OllamaEmbeddings {
|
||||
/// Create a new Ollama embedding provider.
|
||||
///
|
||||
/// Defaults to `nomic-embed-text` (768 dimensions).
|
||||
pub fn new(base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
base_url: base_url.into(),
|
||||
model: "nomic-embed-text".to_string(),
|
||||
dimension: 768,
|
||||
}
|
||||
}
|
||||
|
||||
/// Use a specific model with a given dimension.
|
||||
pub fn with_model(mut self, model: impl Into<String>, dimension: usize) -> Self {
|
||||
self.model = model.into();
|
||||
self.dimension = dimension;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OllamaEmbedRequest<'a> {
|
||||
model: &'a str,
|
||||
input: &'a [String],
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OllamaEmbedResponse {
|
||||
embeddings: Vec<Vec<f32>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EmbeddingProvider for OllamaEmbeddings {
|
||||
fn dimension(&self) -> usize {
|
||||
self.dimension
|
||||
}
|
||||
|
||||
fn model_name(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
fn max_input_length(&self) -> usize {
|
||||
// Most Ollama embedding models support 8192 tokens (~32k chars)
|
||||
32_000
|
||||
}
|
||||
|
||||
async fn embed(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
|
||||
if text.len() > self.max_input_length() {
|
||||
return Err(EmbeddingError::TextTooLong {
|
||||
length: text.len(),
|
||||
max: self.max_input_length(),
|
||||
});
|
||||
}
|
||||
|
||||
let embeddings = self.embed_batch(&[text.to_string()]).await?;
|
||||
embeddings
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| EmbeddingError::InvalidResponse("No embedding returned".to_string()))
|
||||
}
|
||||
|
||||
async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let request = OllamaEmbedRequest {
|
||||
model: &self.model,
|
||||
input: texts,
|
||||
};
|
||||
|
||||
let url = format!("{}/api/embed", self.base_url);
|
||||
|
||||
let response = self.client.post(&url).json(&request).send().await?;
|
||||
|
||||
let status = response.status();
|
||||
|
||||
if !status.is_success() {
|
||||
let error_text = response.text().await.unwrap_or_default();
|
||||
return Err(EmbeddingError::HttpError(format!(
|
||||
"Ollama returned HTTP {}: {}",
|
||||
status, error_text
|
||||
)));
|
||||
}
|
||||
|
||||
let result: OllamaEmbedResponse = response.json().await.map_err(|e| {
|
||||
EmbeddingError::InvalidResponse(format!("Failed to parse Ollama response: {}", e))
|
||||
})?;
|
||||
|
||||
// Validate that returned embeddings match the configured dimension.
|
||||
for (i, emb) in result.embeddings.iter().enumerate() {
|
||||
if emb.len() != self.dimension {
|
||||
return Err(EmbeddingError::InvalidResponse(format!(
|
||||
"Ollama returned embedding of dimension {}, expected {} at index {}",
|
||||
emb.len(),
|
||||
self.dimension,
|
||||
i
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result.embeddings)
|
||||
}
|
||||
}
|
||||
|
||||
/// A mock embedding provider for testing.
|
||||
///
|
||||
/// Generates deterministic embeddings based on text hash.
|
||||
|
||||
@@ -50,7 +50,9 @@ mod search;
|
||||
|
||||
pub use chunker::{ChunkConfig, chunk_document};
|
||||
pub use document::{MemoryChunk, MemoryDocument, WorkspaceEntry, paths};
|
||||
pub use embeddings::{EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OpenAiEmbeddings};
|
||||
pub use embeddings::{
|
||||
EmbeddingProvider, MockEmbeddings, NearAiEmbeddings, OllamaEmbeddings, OpenAiEmbeddings,
|
||||
};
|
||||
#[cfg(feature = "postgres")]
|
||||
pub use repository::Repository;
|
||||
pub use search::{RankedResult, SearchConfig, SearchResult, reciprocal_rank_fusion};
|
||||
|
||||
@@ -24,7 +24,21 @@ const AUTH_TOKEN: &str = "test-openai-token";
|
||||
// Mock LLM provider
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct MockLlmProvider;
|
||||
#[derive(Default)]
|
||||
struct MockLlmState {
|
||||
completion_models: tokio::sync::Mutex<Vec<Option<String>>>,
|
||||
tool_completion_models: tokio::sync::Mutex<Vec<Option<String>>>,
|
||||
}
|
||||
|
||||
struct MockLlmProvider {
|
||||
state: Arc<MockLlmState>,
|
||||
}
|
||||
|
||||
impl MockLlmProvider {
|
||||
fn new(state: Arc<MockLlmState>) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for MockLlmProvider {
|
||||
@@ -37,6 +51,12 @@ impl LlmProvider for MockLlmProvider {
|
||||
}
|
||||
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
self.state
|
||||
.completion_models
|
||||
.lock()
|
||||
.await
|
||||
.push(req.model.clone());
|
||||
|
||||
// Echo the last user message back
|
||||
let user_msg = req
|
||||
.messages
|
||||
@@ -59,6 +79,12 @@ impl LlmProvider for MockLlmProvider {
|
||||
&self,
|
||||
req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
self.state
|
||||
.tool_completion_models
|
||||
.lock()
|
||||
.await
|
||||
.push(req.model.clone());
|
||||
|
||||
// If tools are provided, return a tool call
|
||||
if let Some(tool) = req.tools.first() {
|
||||
Ok(ToolCompletionResponse {
|
||||
@@ -93,11 +119,71 @@ impl LlmProvider for MockLlmProvider {
|
||||
}
|
||||
}
|
||||
|
||||
struct FixedModelProvider {
|
||||
model: &'static str,
|
||||
}
|
||||
|
||||
impl FixedModelProvider {
|
||||
fn new(model: &'static str) -> Self {
|
||||
Self { model }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for FixedModelProvider {
|
||||
fn model_name(&self) -> &str {
|
||||
self.model
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
(Decimal::ZERO, Decimal::ZERO)
|
||||
}
|
||||
|
||||
async fn complete(&self, _req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
Ok(CompletionResponse {
|
||||
content: "fixed response".to_string(),
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
Ok(ToolCompletionResponse {
|
||||
content: Some("fixed response".to_string()),
|
||||
tool_calls: vec![],
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn effective_model_name(&self, _requested_model: Option<&str>) -> String {
|
||||
self.model.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Test helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn start_test_server() -> (SocketAddr, Arc<GatewayState>) {
|
||||
async fn start_test_server() -> (SocketAddr, Arc<GatewayState>, Arc<MockLlmState>) {
|
||||
let mock_state = Arc::new(MockLlmState::default());
|
||||
|
||||
let llm_provider: Arc<dyn LlmProvider> = Arc::new(MockLlmProvider::new(mock_state.clone()));
|
||||
let (bound_addr, state) = start_test_server_with_provider(llm_provider).await;
|
||||
|
||||
(bound_addr, state, mock_state)
|
||||
}
|
||||
|
||||
async fn start_test_server_with_provider(
|
||||
llm_provider: Arc<dyn LlmProvider>,
|
||||
) -> (SocketAddr, Arc<GatewayState>) {
|
||||
let state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(None),
|
||||
sse: SseManager::new(),
|
||||
@@ -112,7 +198,7 @@ async fn start_test_server() -> (SocketAddr, Arc<GatewayState>) {
|
||||
user_id: "test-user".to_string(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(Arc::new(MockLlmProvider)),
|
||||
llm_provider: Some(llm_provider),
|
||||
skill_registry: None,
|
||||
skill_catalog: None,
|
||||
chat_rate_limiter: ironclaw::channels::web::server::RateLimiter::new(30, 60),
|
||||
@@ -139,7 +225,7 @@ fn client() -> reqwest::Client {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_basic() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -173,11 +259,14 @@ async fn test_chat_completions_basic() {
|
||||
assert_eq!(body["usage"]["prompt_tokens"], 10);
|
||||
assert_eq!(body["usage"]["completion_tokens"], 5);
|
||||
assert_eq!(body["usage"]["total_tokens"], 15);
|
||||
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_with_system_message() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -204,7 +293,7 @@ async fn test_chat_completions_with_system_message() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_with_tools() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -243,11 +332,14 @@ async fn test_chat_completions_with_tools() {
|
||||
assert_eq!(tool_calls[0]["id"], "call_mock_001");
|
||||
assert_eq!(tool_calls[0]["type"], "function");
|
||||
assert_eq!(tool_calls[0]["function"]["name"], "get_weather");
|
||||
|
||||
let models = mock_state.tool_completion_models.lock().await;
|
||||
assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_streaming() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -316,11 +408,14 @@ async fn test_chat_completions_streaming() {
|
||||
"Expected reassembled content to contain 'Stream test', got: '{}'",
|
||||
full_content
|
||||
);
|
||||
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert_eq!(*models, vec![Some("mock-model-v1".to_string())]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_empty_messages() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -340,8 +435,8 @@ async fn test_chat_completions_empty_messages() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_model_mismatch() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
async fn test_chat_completions_model_override() {
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -355,20 +450,173 @@ async fn test_chat_completions_model_mismatch() {
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 404);
|
||||
assert_eq!(resp.status(), 200);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["model"], "gpt-4");
|
||||
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert_eq!(*models, vec![Some("gpt-4".to_string())]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_uses_effective_model_when_override_ignored() {
|
||||
let provider: Arc<dyn LlmProvider> = Arc::new(FixedModelProvider::new("configured-model"));
|
||||
let (addr, _state) = start_test_server_with_provider(provider).await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
.post(&url)
|
||||
.bearer_auth(AUTH_TOKEN)
|
||||
.json(&serde_json::json!({
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hi"}]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 200);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["model"], "configured-model");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_streaming_uses_effective_model_when_override_ignored() {
|
||||
let provider: Arc<dyn LlmProvider> = Arc::new(FixedModelProvider::new("configured-model"));
|
||||
let (addr, _state) = start_test_server_with_provider(provider).await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
.post(&url)
|
||||
.bearer_auth(AUTH_TOKEN)
|
||||
.json(&serde_json::json!({
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hi"}],
|
||||
"stream": true
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 200);
|
||||
let text = resp.text().await.unwrap();
|
||||
assert!(
|
||||
text.contains("\"model\":\"configured-model\""),
|
||||
"Expected streaming chunks to report configured model, got: {}",
|
||||
text
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_model_too_long() {
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
.post(&url)
|
||||
.bearer_auth(AUTH_TOKEN)
|
||||
.json(&serde_json::json!({
|
||||
"model": "m".repeat(300),
|
||||
"messages": [{"role": "user", "content": "Hi"}]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 400);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
assert!(
|
||||
body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("mock-model-v1")
|
||||
.unwrap_or("")
|
||||
.contains("model"),
|
||||
"Expected model validation error, got: {}",
|
||||
body
|
||||
);
|
||||
|
||||
// Validation should fail before provider invocation.
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert!(
|
||||
models.is_empty(),
|
||||
"provider should not be called: {:?}",
|
||||
*models
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_model_with_control_chars() {
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
.post(&url)
|
||||
.bearer_auth(AUTH_TOKEN)
|
||||
.json(&serde_json::json!({
|
||||
"model": "gpt-4\noops",
|
||||
"messages": [{"role": "user", "content": "Hi"}]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 400);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert!(
|
||||
body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap_or("")
|
||||
.contains("control"),
|
||||
"Expected model validation error, got: {}",
|
||||
body
|
||||
);
|
||||
|
||||
// Validation should fail before provider invocation.
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert!(
|
||||
models.is_empty(),
|
||||
"provider should not be called: {:?}",
|
||||
*models
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_model_with_surrounding_whitespace() {
|
||||
let (addr, _state, mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
.post(&url)
|
||||
.bearer_auth(AUTH_TOKEN)
|
||||
.json(&serde_json::json!({
|
||||
"model": " gpt-4 ",
|
||||
"messages": [{"role": "user", "content": "Hi"}]
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(resp.status(), 400);
|
||||
let body: serde_json::Value = resp.json().await.unwrap();
|
||||
assert!(
|
||||
body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap_or("")
|
||||
.contains("leading or trailing whitespace"),
|
||||
"Expected model validation error, got: {}",
|
||||
body
|
||||
);
|
||||
|
||||
let models = mock_state.completion_models.lock().await;
|
||||
assert!(
|
||||
models.is_empty(),
|
||||
"provider should not be called: {:?}",
|
||||
*models
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_no_auth() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -387,7 +635,7 @@ async fn test_chat_completions_no_auth() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_models_endpoint() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/models", addr);
|
||||
|
||||
let resp = client()
|
||||
@@ -410,7 +658,7 @@ async fn test_models_endpoint() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_models_no_auth() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/models", addr);
|
||||
|
||||
let resp = client().get(&url).send().await.unwrap();
|
||||
@@ -462,7 +710,7 @@ async fn test_no_llm_provider_returns_503() {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_chat_completions_body_too_large() {
|
||||
let (addr, _state) = start_test_server().await;
|
||||
let (addr, _state, _mock_state) = start_test_server().await;
|
||||
let url = format!("http://{}/v1/chat/completions", addr);
|
||||
|
||||
// Build a payload over 1 MB (the gateway's DefaultBodyLimit)
|
||||
|
||||
Reference in New Issue
Block a user