mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-09-01 09:09:19 +00:00
* feat: complete multi-tenant isolation — per-user budgets, model selection, heartbeat cycling Finishes the remaining isolation work from phases 2–4 of #59: Phase 2 (DB scoping): Fix /status and /list commands to use _for_user DB variants instead of global queries that leaked cross-user job data. Phase 3 (Runtime isolation): Per-user workspace in routine engine's spawn_fire so lightweight routines run in the correct user context. Per-user daily cost tracking in CostGuard with configurable budget via MAX_COST_PER_USER_PER_DAY_CENTS. Multi-user heartbeat that cycles through all users with routines, auto-detected from GATEWAY_USER_TOKENS. Phase 4 (Provider/tools): Per-user model selection via preferred_model setting — looked up from SettingsStore on first iteration, threaded through ReasoningContext.model_override to CompletionRequest. Works with providers that support per-request model overrides (NearAI). Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: use selected_model setting key to match /model command persistence The dispatcher was reading "preferred_model" but the /model command (merged from staging) persists to "selected_model". Since set_setting is already per-user scoped, using the same key makes /model work as the per-user model override in multi-tenant mode. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: heartbeat hygiene, /model multi-tenant guard, RigAdapter model override Three follow-up fixes for multi-tenant isolation: 1. Multi-user heartbeat now runs memory hygiene per user before each heartbeat check, matching single-user heartbeat behavior. 2. /model command in multi-tenant mode only persists to per-user settings (selected_model) without calling set_model() on the shared LlmProvider. The per-request model_override in the dispatcher reads from the same setting. Added multi_tenant flag to AgentConfig (auto-detected from GATEWAY_USER_TOKENS). 3. RigAdapter now supports per-request model overrides by injecting the model name into rig-core's additional_params. OpenAI/Anthropic/Ollama API servers use last-key-wins for duplicate JSON keys, so the override takes effect via serde's flatten serialization order. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address PR review — cost model attribution, heartbeat concurrency, pruning Fixes from review comments on #1614: - Cost tracking now uses the override model name (not active_model_name) when a per-user model override is active, for accurate attribution. - Multi-user heartbeat runs per-user checks concurrently via JoinSet instead of sequentially, preventing one slow user from blocking others. - Per-user failure counts tracked independently; users exceeding max_failures are skipped (matching single-user semantics). - per_user_daily_cost HashMap pruned on day rollover to prevent unbounded growth in long-lived deployments. - Doc comment fixed: says "routines" not "active routines". Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: /status ownership, model persistence scoping, heartbeat robustness Addresses second round of PR review on #1614: - /status <job_id> DB path now validates job.user_id == requesting user before returning data (was missing ownership check, security fix). - persist_selected_model takes user_id param instead of owner_id, and skips .env/TOML writes in multi-tenant mode (these are shared global files). handle_system_command now receives user_id from caller. - JoinSet collection handles Err(JoinError) explicitly instead of silently dropping panicked tasks. - Notification forwarder extracts owner_id from response metadata in multi-tenant mode for per-user routing instead of broadcasting to the agent owner. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: cost pricing, fire_manual workspace, heartbeat concurrency cap Round 3 review fixes: - Cost tracking passes None for cost_per_token when model override is active, letting CostGuard look up pricing by model name instead of using the default provider's rates (serrrfirat). - fire_manual() now uses per-user workspace, matching spawn_fire() pattern (serrrfirat). - Removed MULTI_TENANT env var — multi-tenant mode is auto-detected solely from GATEWAY_USER_TOKENS presence (serrrfirat + Copilot). - Multi-user heartbeat capped at 8 concurrent tasks to avoid flooding the LLM provider (serrrfirat + Copilot). - Fixed inject_model_override doc comment accuracy (Copilot). - Added comment explaining multi-tenant notification routing priority (Copilot). Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: user-scoped webhook endpoint for multi-tenant isolation Adds POST /api/webhooks/u/{user_id}/{path} — a user-scoped webhook endpoint that filters the routine lookup by user_id, preventing cross-user webhook triggering when paths collide. The existing /api/webhooks/{path} endpoint remains unchanged for backward compatibility in single-user deployments. Changes: - get_webhook_routine_by_path gains user_id: Option<&str> param - Both postgres and libsql implementations add AND user_id = ? filter when user_id is provided - New webhook_trigger_user_scoped_handler extracts (user_id, path) from URL and passes to shared fire_webhook_inner logic - Route registered on public router (webhooks are called by external services that can't send bearer tokens) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat(db): add UserStore trait with users, api_tokens, invitations tables Foundation for DB-backed user management (#1605): - UserRecord, ApiTokenRecord, InvitationRecord types in db/mod.rs - UserStore sub-trait (17 methods) added to Database supertrait - PostgreSQL migration V14__users.sql (users, api_tokens, invitations) - libSQL schema + incremental migration V14 - Full implementations for both PgBackend (via Store delegation) and LibSqlBackend (direct SQL in libsql/users.rs) - authenticate_token JOINs api_tokens+users with active/non-revoked checks; has_any_users for bootstrap detection Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat(web): DB-backed auth, user/token/invitation API handlers Adds the web gateway layer for DB-backed user management (#1605): Auth refactor: - CombinedAuthState wraps env-var tokens (MultiAuthState) + optional DbAuthenticator for DB-backed token lookup with LRU cache (60s TTL, 1024 max entries) - auth_middleware tries env-var tokens first, then DB fallback - From<MultiAuthState> impl for backward compatibility - main.rs wires with_db_auth when database is available API handlers (12 new endpoints): - /api/admin/users — CRUD: create, list, detail, update, suspend, activate - /api/tokens — create (returns plaintext once), list, revoke - /api/invitations — create, list, accept (creates user + first token) Token creation: 32 random bytes → hex plaintext, SHA-256 hash stored. Invitation accept: validates hash + pending + not expired, creates user record and first API token atomically. All test files updated for CombinedAuthState type change. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: startup env-var user migration + UserStore integration tests Completes the DB-backed user management feature (#1605): - Startup migration: when GATEWAY_USER_TOKENS is set and the users table is empty, inserts env-var users + hashed tokens into DB. Logs deprecation notice when DB already has users. - hash_token made pub for reuse in migration code. - 10 integration tests for UserStore (libsql file-backed): - has_any_users bootstrap detection - create/get/get_by_email/list/update user lifecycle - token create → authenticate → revoke → reject cycle - suspended user tokens rejected - wrong-user token revoke returns false - invitation create → accept → user created - record_login and record_token_usage timestamps - libSQL migration: removed FK constraints from V14 (incompatible with execute_batch inside transactions). Tables in both base SCHEMA and incremental migration for fresh and existing databases. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * refactor: remove GATEWAY_USER_TOKENS, fix review feedback GATEWAY_USER_TOKENS never went to production — replaced entirely by DB-backed user management via /api/admin/users and /api/tokens. Removed: - UserTokenConfig struct and GATEWAY_USER_TOKENS env var parsing - user_tokens field from GatewayConfig - GatewayChannel::new_multi_auth() constructor - Env-var user migration block in main.rs (~90 lines) - multi_tenant auto-detection from GATEWAY_USER_TOKENS (now runtime via db.has_any_users() in app.rs) Review fixes (zmanian): - User ID generation: UUID instead of display-name derivation (#1) - Invitation accept moved to public router (no auth needed) (#3) - libSQL get_invitation_by_hash aligned with postgres: filters status='pending' AND expires_at > now (#4) - UUID parse: returns DatabaseError::Serialization instead of unwrap_or_default (#7) - PostgreSQL SELECT * replaced with explicit column lists (#8) - Sort order aligned (both backends use DESC) (#6) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: add role-based access control (admin/member) Adds a `role` field (admin|member) to user management: Schema: - `role TEXT NOT NULL DEFAULT 'member'` added to users table in both PostgreSQL V14 migration and libSQL schema/incremental migration - UserRecord gains `role: String` field - UserIdentity gains `role: String` field, populated from DB in DbAuthenticator and defaulting to "admin" for single-user mode Access control: - AdminUser extractor: returns 403 Forbidden if role != "admin" - /api/admin/users/* handlers: require AdminUser (create, list, detail, update, suspend, activate) - POST /api/invitations: requires AdminUser (only admins can invite) - User creation accepts optional "role" param (defaults to "member") - Invitation acceptance creates users with "member" role Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat(web): add Users admin tab to web UI Adds a Users tab to the web gateway UI for managing users, tokens, and roles without needing direct API calls. Features: - User list table with ID, name, email, role, status, created date - Create user form with display name, email, role selector - Suspend/activate actions per user - Create API token for any user (shows plaintext once with copy button) - Role badges (admin highlighted, member muted) - Non-admin users see "Admin access required" message - Keyboard shortcut: Cmd/Ctrl+5 switches to Users tab CSS: - Reuses routines-table styles for the user list - Badge, token-display, btn-small, btn-danger, btn-primary components Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: move Users to Settings subtab, bootstrap admin user on first run - Moved Users from top-level tab to Settings sidebar subtab (under Skills, before Theme toggle) - On first startup with empty users table, automatically creates an admin user from GATEWAY_USER_ID config with a corresponding API token from GATEWAY_AUTH_TOKEN. This ensures the owner appears in the Users panel immediately. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: user creation shows token, + Token works, no password save popup Three UI/UX fixes: 1. Create user now generates an initial API token and shows it in a copy-able banner instead of triggering the browser's password save dialog. Uses autocomplete="off" and type="text" for email field. 2. "+ Token" button works: exposed createTokenForUser/suspendUser/ activateUser on window for inline onclick handlers in dynamically generated table rows. Token creation uses showTokenBanner helper. 3. Admin token creation: POST /api/tokens now accepts optional "user_id" field when the requesting user is admin, allowing token creation for other users from the Users panel. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: use event delegation for user action buttons (CSP compliance) Inline onclick handlers are blocked by the Content-Security-Policy (script-src 'self' without 'unsafe-inline'). Switched to data-action attributes with a delegated click listener on the users table. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: add i18n for Users subtab, show login link on user creation - Added 'settings.users' i18n key for English and Chinese - Token banner now shows a full login link (domain/?token=xxx) with a Copy Link button, plus the raw token below - Login link works automatically via existing ?token= auto-auth Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: token hash mismatch — hash hex string, not raw bytes Critical auth bug: token creation hashed the raw 32 bytes (hasher.update(token_bytes)) but authentication hashed the hex-encoded string (hash_token(candidate) where candidate is the hex string the user sends). This meant newly created tokens could never authenticate. Fixed all 4 token creation sites (users, tokens, invitations create, invitations accept) to use hash_token(&plaintext_token) which hashes the hex string consistently with the auth lookup path. Removed now-unused sha2::Digest imports from handlers. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * refactor: remove invitation system The invitation flow is redundant — admin create user already generates a token and shows a login link. Invitations add complexity without value until email integration exists. Removed: - InvitationRecord struct and 4 UserStore trait methods - invitations table from V14 migration (postgres + both libsql schemas) - PostgreSQL Store methods (create/get/accept/list invitations) - libSQL UserStore invitation methods + row_to_invitation helper - invitations.rs handler file (212 lines) - /api/invitations routes (create, list, accept) - test_invitation_lifecycle test Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: user deletion, self-service profile, per-user job limits, usage API Four multi-tenancy improvements: 1. User deletion cascade (DELETE /api/admin/users/{id}): Deletes user and all data across 11 user-scoped tables (settings, secrets, routines, memory, jobs, conversations, etc.). Admin only. 2. Self-service profile (GET/PATCH /api/profile): Users can read and update their own display_name and metadata without admin privileges. 3. Per-user job concurrency (MAX_JOBS_PER_USER env var): Scheduler checks active_jobs_for(user_id) before dispatch. Prevents one user from exhausting all job slots. 4. Usage reporting (GET /api/admin/usage?user_id=X&period=day|week|month): Aggregates LLM costs from llm_calls via agent_jobs.user_id. Returns per-user, per-model breakdown of calls, tokens, and cost. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: add TenantCtx for compile-time tenant isolation Implements zmanian's architectural proposal from #1614 review: two-tier scoped database access (TenantScope/AdminScope) so handler code cannot accidentally bypass tenant scoping. TenantScope (default): wraps user_id + Arc<dyn Database>, auto-binds user_id on every operation. ID-based lookups return None for cross- tenant resources. No escape hatch — forgetting to scope is a compile error. AdminScope (explicit opt-in): cross-tenant access for system-level components (heartbeat, routine engine, self-repair, scheduler, worker). TenantCtx bundles TenantScope + workspace + cost guard + per-user rate limiting. Constructed once per request in handle_message, threaded through all command handlers and ChatDelegate. Key changes: - New src/tenant.rs (~920 lines): TenantScope, AdminScope, TenantCtx, TenantRateState, TenantRateRegistry - All command handlers: user_id: &str → ctx: &TenantCtx - ChatDelegate: cost check/record/settings via self.tenant - System components: store field changed to AdminScope - Config: TENANT_MAX_LLM_CONCURRENT, TENANT_MAX_JOBS_CONCURRENT env vars - Fixes bug: /status <job_id> cross-tenant leak (now auto-filtered) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address PR #1626 review feedback — bounded LRU cache, admin auth, FK cleanup - Replace HashMap with lru::LruCache in DbAuthenticator so the token cache is hard-bounded at 1024 entries (evicts LRU, not just expired) - Gate admin user endpoints (list/detail/update/suspend/activate) with AdminUser extractor so members get 403 instead of full access - Add api_tokens to libSQL delete_user cleanup list to prevent orphaned tokens (libSQL has no FK cascade) - Add regression tests for all three fixes Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: update CA certificates in runtime Docker image Ensures the root certificate bundle is current so TLS handshakes to services like Supabase succeed on Railway. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: resolve CI failures — formatting, no-panics check - Run cargo fmt on test code - Replace .expect() with const NonZeroUsize in DbAuthenticator - Add // safety: comments for test-only code in multi_tenant.rs Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: switch PostgreSQL TLS from rustls to native-tls rustls with rustls-native-certs fails TLS handshake on Railway's slim container (empty or stale root cert store). native-tls delegates to OpenSSL on Linux which handles system certs more reliably. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * Adding user management api * feat: admin secrets provisioning API + API documentation - Add PUT/GET/DELETE /api/admin/users/{id}/secrets/{name} endpoints for application backends to provision per-user secrets (AES-256-GCM encrypted) - Add secrets_store field to GatewayState with builder wiring - Create docs/USER_MANAGEMENT_API.md with full API spec covering users, secrets, tokens, profile, and usage endpoints - Update web gateway CLAUDE.md route table Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: add CatchPanicLayer to capture handler panics Without this, panics in async handlers silently drop the connection and the edge proxy returns a generic 503. Now panics are caught, logged, and returned as 500 with the panic message. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address second-round review — transactional delete, overflow, error logging - C1: Wrap PostgreSQL delete_user() in a transaction so partial cleanup can't leave users in a half-deleted state - M2: Add job_events to delete cleanup (both backends) — FK to agent_jobs without CASCADE would cause FK violation - H1/M4: Cap expires_in_days to 36500 before i64 cast (tokens + secrets) - H2: Validate target user exists before creating admin token to prevent orphan tokens on libSQL - H3: Log DB errors in DbAuthenticator::authenticate() instead of silently swallowing them as 401 Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: revert to rustls with webpki-roots fallback for PostgreSQL TLS native-tls/OpenSSL caused silent crashes (segfaults in C code) during DB writes on Railway containers. Switch back to rustls but add webpki-roots as a fallback when system certs are missing, which was the original TLS handshake failure on slim container images. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * chore: update Cargo.lock for rustls + webpki-roots Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * debug: add /api/debug/db-write endpoint to diagnose user insert failure Temporary diagnostic endpoint that tests DB INSERT to users table with full error logging. No auth required. Will be removed after debugging. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * perf: use cargo-chef in Dockerfile for dependency caching Splits the build into planner/deps/builder stages. Dependencies are only recompiled when Cargo.toml or Cargo.lock change. Source-only changes skip straight to the final build stage. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * debug: add tracing to users_create_handler Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: guard created_by FK in user creation handler The auth identity user_id (from owner_id scope) may not match any user row in the DB, causing a FK violation on the created_by column. Check that the referenced user exists before setting created_by. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * refactor: collapse GATEWAY_USER_ID into IRONCLAW_OWNER_ID Remove the separate GATEWAY_USER_ID config. The gateway now uses IRONCLAW_OWNER_ID (config.owner_id) directly for auth identity, bootstrap user creation, and workspace scoping. Previously, with_owner_scope() rebinds the auth identity to owner_id while keeping default_sender_id as the gateway user_id. This caused a FK constraint violation when creating users because the auth identity ("default") didn't match any user in the DB ("nearai"). Changes: - Remove GATEWAY_USER_ID env var and gateway_user_id from settings - Remove user_id field from GatewayConfig - Add owner_id parameter to GatewayChannel::new() - Remove with_owner_scope() method - Remove default_sender_id from GatewayState - Remove sender override logic in chat/approval handlers - Remove debug endpoint and tracing from prior debugging - Update all tests and E2E fixtures Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: hide Users tab for non-admins, remove auth hint text - Fetch /api/profile after login and hide the Users settings tab when the user's role is not admin - Remove the "Enter the GATEWAY_AUTH_TOKEN" hint from the login page since tokens are now managed via the admin panel, not .env files Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address review feedback (auth 503, token expiry, CORS PATCH) - DB auth errors now return 503 instead of 401 so outages are distinguishable from invalid tokens (serrrfirat H3) - Cap expires_in_days to 36500 before i64 cast to prevent negative duration from u64 overflow (serrrfirat H1) - Add PATCH to CORS allowed methods for profile/user update endpoints (Copilot) - Stop leaking panic details in CatchPanicLayer response body Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: harden multi-tenant isolation — review fixes from #1614 - Add conversation ownership checks in TenantScope: add_conversation_message, touch_conversation, list_conversation_messages (+ paginated), update_conversation_metadata_field, get_conversation_metadata now return NotFound for conversations not owned by the tenant (cross-tenant data leak) - Fix multi-user heartbeat: clear notify_user_id per runner so notifications persist to the correct user, not the shared config target - Move hygiene tasks into bounded JoinSet instead of unbounded tokio::spawn - Revert send_notification to private visibility (only used within module) - Use effective_model_name() for cost attribution in dispatcher so providers that ignore per-request model overrides report the actual model used - Fix inject_model_override doc comment; add 3 unit tests - Fix heartbeat doc comment ("routines" not "active routines") Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: add Jobs, Cost, Last Active columns to admin Users table Add UserSummaryStats struct and user_summary_stats() batch query to the UserStore trait (both PostgreSQL and libSQL backends). The admin users list endpoint now fetches per-user aggregates (job count, total LLM spend, most recent activity) in a single query and includes them inline in the response. The frontend Users table displays three new columns. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address review comments and CI formatting failures CI fixes: - cargo fmt fixes in cli/mod.rs and db/tls.rs Security/correctness (from Copilot + serrrfirat + pranavraja99 reviews): - Token create: reject expires_in_days > 36500 with 400 instead of silent clamp - Token create: return 404 when admin targets non-existent user - User create: map duplicate email constraint violations to 409 Conflict - User create: remove unnecessary DB roundtrip for created_by (use AdminUser directly) - DB auth: log warn on DB lookup failures instead of silently swallowing errors - libSQL: add FK constraints on users.created_by and api_tokens.user_id Config fixes: - agent.multi_tenant: resolve from AGENT_MULTI_TENANT env var instead of hardcoding false - heartbeat.multi_tenant: fix doc comment to match actual env-var-based behavior UI fix: - showTokenBanner: pass correct title ("Token created!" vs "User created!") Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address remaining review comments (round 2) - Secrets handlers: normalize name to lowercase before store operations, validate target user_id exists (returns 404 if not found) - libSQL: propagate cost parsing errors instead of unwrap_or_default() in both user_usage_stats and user_summary_stats - users_list_handler: propagate user_summary_stats DB errors (was silently swallowed with unwrap_or_default) - loadUsers: distinguish 401/403 (admin required) from other errors - Docs: fix users.id type (TEXT not UUID), remove "invitation flow" from V14 migration comment Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: i18n for Users tab, atomic user+token creation, transactional delete_user i18n: - Add 31 translation keys for all Users tab strings (en + zh-CN) - Wire data-i18n attributes on HTML elements (headings, buttons, inputs, table headers, empty state) - Replace all hard-coded strings in app.js with I18n.t() calls Atomic user+token creation: - Add create_user_with_token() to UserStore trait - PostgreSQL: wraps both INSERTs in conn.transaction() with auto-rollback - libSQL: wraps in explicit BEGIN/COMMIT with ROLLBACK on error - Handler uses single atomic call instead of two separate operations Transactional delete_user for libSQL: - Wrap multi-table DELETE cascade in BEGIN/COMMIT transaction - ROLLBACK on any error to prevent partial cleanup / inconsistent state - Matches the PostgreSQL implementation which already used transactions Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: revert V14 migration to match deployed checksum [skip-regression-check] Refinery checksums applied migrations — editing V14__users.sql after it was already applied causes deployment failures. Revert the cosmetic comment changes (added in df40b22f) to restore the original checksum. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: bootstrap onboarding flow for multi-tenant users The bootstrap greeting and workspace seeding only ran for the owner workspace at startup, so new users created via the admin API never received the welcome message or identity files (BOOTSTRAP.md, SOUL.md, AGENTS.md, USER.md, etc.). Three fixes: - tenant_ctx(): seed per-user workspace on first creation via seed_if_empty(), which writes identity files and sets bootstrap_pending when the workspace is truly fresh - handle_message(): check take_bootstrap_pending() on the tenant workspace (not the owner workspace) and persist the greeting to the user's own assistant conversation + broadcast via SSE - WorkspacePool: seed new per-user workspaces in the web gateway so memory tools also see identity files immediately The existing single-user bootstrap in Agent::run() is preserved for non-multi-tenant deployments. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address remaining PR review comments (round 3) - Docs: fix metadata description from "merge patch" to "full replacement" - Secrets: reject expires_in_days > 36500 with 400 (was silently clamped) - libSQL: CAST(SUM(cost) AS TEXT) in user_usage_stats and user_summary_stats to prevent SQLite numeric coercion from crashing get_text() — this was the root cause of the Copilot "SUM returns numeric type" comments - Add 3 regression tests: user_summary_stats (empty + with data) and user_usage_stats (multi-model aggregation) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * feat: add role change support for users (admin/member toggle) - Add update_user_role() to UserStore trait + both backends (PostgreSQL and libSQL) - Extend PATCH /api/admin/users/{id} to accept optional "role" field with validation (must be "admin" or "member") - Add "Make Admin" / "Make Member" toggle button in Users table actions - Add i18n keys for role change (en + zh-CN) - Update API docs to document the role field on PATCH - Fix test helpers to use fmt_ts() for timestamps (was using SQLite datetime('now') which produces incompatible format for string comparison) Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: show live LLM spend in Users table instead of only DB-recorded costs [skip-regression-check] Chat turns record LLM cost in CostGuard (in-memory) but don't create agent_jobs/llm_calls DB rows — those are only written for background jobs. The Users table was querying only from DB, so it showed $0.00 for users who only chatted. Now supplements DB stats with CostGuard.daily_spend_for_user() — the same source displayed in the status bar token counter. Shows whichever is larger (DB historical total vs live daily spend). Also falls back to last_login_at for "Last Active" when no DB job activity exists. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: persist chat LLM calls to DB and fix usage stats query Two root causes for zero usage stats: 1. ChatDelegate only recorded LLM costs to CostGuard (in-memory) — never to the llm_calls DB table. Added DB persistence via TenantScope.record_llm_call() after each chat LLM call, with job_id=NULL and conversation_id=thread_id. 2. user_summary_stats query only joined agent_jobs→llm_calls, missing chat calls (which have job_id=NULL). Redesigned query to start from llm_calls and resolve user_id via COALESCE(agent_jobs.user_id, conversations.user_id) — covers both job and chat LLM calls. Both PostgreSQL and libSQL queries updated. TenantScope gets record_llm_call() method. Tests updated for new query semantics. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address review comments — input validation, cost semantics, panic safety [skip-regression-check] - Validate display_name: trim whitespace, reject empty strings (create + update) - Validate metadata: must be a JSON object, return 400 if not (admin + profile) - secrets_list_handler: verify target user_id exists before listing - Cost display: use DB total directly (chat calls now persist to DB), remove confusing max(db,live) CostGuard fallback - CatchPanicLayer: truncate panic payload to 200 chars in log to limit potential sensitive data exposure Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: address Copilot round 5 — docs, secrets consistency, token name, provider field [skip-regression-check] - Docs: users.id note updated to "typically UUID v4 strings (bootstrap admin may use a custom ID)" - secrets_list_handler: return 503 when DB store is None (was falling through to list secrets without user validation) - tokens_create: trim + reject empty token name (matching display_name pattern) - LlmCallRecord.provider: use llm_backend ("nearai","openai") instead of model_name() which returns the model identifier - user_summary_stats zero-LLM users: acceptable — handler already falls back to 0 cost and last_login_at for missing entries Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: DB auth returns 503 on outage, scheduler counts only blocking jobs From serrrfirat review: - DB auth: return Err(()) on database errors so middleware returns 503 instead of silently returning Ok(None) → 401 (auth miss) - Scheduler: add parallel_blocking_count_for() that uses is_parallel_blocking() (Pending/InProgress/Stuck) instead of is_active() for per-user concurrency — Completed/Submitted jobs no longer count against MAX_JOBS_PER_USER From Copilot: - CLAUDE.md: fix secrets route paths from {id} to {user_id} - token_hash: use .as_slice() instead of .to_vec() to avoid heap allocation on every token auth/creation call Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: immediate auth cache invalidation on security-critical actions (zmanian review #6) Add DbAuthenticator::invalidate_user() that evicts all cached entries for a user. Called after: - Suspend user (immediate lockout, was 60s delay) - Activate user (immediate access restoration) - Role change (admin↔member takes effect immediately) - Token revocation (revoked token can't be reused from cache) The DbAuthenticator is shared (via Clone, which Arc-clones the cache) between the auth middleware and GatewayState, so handlers can evict entries from the same cache the middleware reads. Also from zmanian's review: - Items 1-5, 7-11 were already resolved in prior commits - Item 12 (String→enum for status/role) is deferred as a broader refactor Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: last-admin protection, usage stats for chat calls, UTF-8 safe panic truncation Last-admin protection: - Suspend, delete, and role-demotion of the last active admin now return 409 Conflict instead of succeeding and locking out the admin API - Helper is_last_admin() checks active admin count before destructive ops Usage stats: - user_usage_stats() now includes chat LLM calls (job_id=NULL) by joining via conversations.user_id, matching user_summary_stats() - Both PostgreSQL and libSQL queries updated Panic handler: - Use floor_char_boundary(200) instead of byte-index [..200] to prevent panic on multi-byte UTF-8 characters in panic messages Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: workspace seed race, bootstrap atomicity, email trim, secrets upsert response [skip-regression-check] - WorkspacePool: await seed_if_empty() synchronously after inserting into cache (drop lock first to avoid blocking), so callers see identity files immediately instead of racing a background task - Bootstrap admin: use create_user_with_token() for atomic user+token creation, matching the admin create endpoint - Email: trim whitespace, treat empty as None to prevent " " being stored and breaking uniqueness - Secrets PUT: report "updated" vs "created" based on prior existence - Last token_hash.to_vec() → .as_slice() in authenticate_token Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: disable unscoped webhook endpoint in multi-tenant mode [skip-regression-check] The original /api/webhooks/{path} endpoint looks up routines across all users. In multi-tenant mode, anyone who knows the webhook path + secret could trigger another user's routine. Now returns 410 Gone with a message pointing to the scoped endpoint /api/webhooks/u/{user_id}/{path}. Detection uses state.db_auth.is_some() — present only when DB-backed auth is enabled (multi-tenant). Single-user deployments are unaffected. From: standardtoaster review comment Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> * fix: webhook multi-tenant check, secrets error propagation, stale doc comment [skip-regression-check] - Webhook: use workspace_pool.is_some() instead of db_auth.is_some() for multi-tenant detection — db_auth is set for any DB deployment, workspace_pool is only set when has_any_users() was true at startup - Secrets: propagate exists() errors instead of unwrap_or(false) so backend outages surface as 500 rather than incorrect "created" status - Config: fix stale workspace_read_scopes comment referencing user_id Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]> --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
1569 lines
58 KiB
Rust
1569 lines
58 KiB
Rust
//! Generic adapter that bridges rig-core's `CompletionModel` trait to IronClaw's `LlmProvider`.
|
||
//!
|
||
//! This lets us use any rig-core provider (OpenAI, Anthropic, Ollama, etc.) as an
|
||
//! `Arc<dyn LlmProvider>` without changing any of the agent, reasoning, or tool code.
|
||
|
||
use crate::llm::config::CacheRetention;
|
||
use async_trait::async_trait;
|
||
use rig::OneOrMany;
|
||
use rig::completion::{
|
||
AssistantContent, CompletionModel, CompletionRequest as RigRequest,
|
||
ToolDefinition as RigToolDefinition, Usage as RigUsage,
|
||
};
|
||
use rig::message::{
|
||
DocumentSourceKind, Image, ImageMediaType, Message as RigMessage, MimeType,
|
||
ToolChoice as RigToolChoice, ToolFunction, ToolResult as RigToolResult, ToolResultContent,
|
||
UserContent,
|
||
};
|
||
use rust_decimal::Decimal;
|
||
use rust_decimal_macros::dec;
|
||
use serde::Serialize;
|
||
use serde::de::DeserializeOwned;
|
||
use serde_json::Value as JsonValue;
|
||
use sha2::{Digest, Sha256};
|
||
|
||
use std::collections::HashSet;
|
||
|
||
use crate::llm::costs;
|
||
use crate::llm::error::LlmError;
|
||
use crate::llm::provider::{
|
||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider,
|
||
ToolCall as IronToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||
ToolDefinition as IronToolDefinition, strip_unsupported_completion_params,
|
||
strip_unsupported_tool_params,
|
||
};
|
||
|
||
/// Adapter that wraps a rig-core `CompletionModel` and implements `LlmProvider`.
|
||
pub struct RigAdapter<M: CompletionModel> {
|
||
model: M,
|
||
model_name: String,
|
||
input_cost: Decimal,
|
||
output_cost: Decimal,
|
||
/// Prompt cache retention policy (Anthropic only).
|
||
/// When not `CacheRetention::None`, injects top-level `cache_control`
|
||
/// via `additional_params` for Anthropic automatic caching. Also controls
|
||
/// the cost multiplier for cache-creation tokens.
|
||
cache_retention: CacheRetention,
|
||
/// Parameter names that this provider does not support (e.g., `"temperature"`).
|
||
/// These are stripped from requests before sending to avoid 400 errors.
|
||
unsupported_params: HashSet<String>,
|
||
}
|
||
|
||
impl<M: CompletionModel> RigAdapter<M> {
|
||
/// Create a new adapter wrapping the given rig-core model.
|
||
pub fn new(model: M, model_name: impl Into<String>) -> Self {
|
||
let name = model_name.into();
|
||
let (input_cost, output_cost) =
|
||
costs::model_cost(&name).unwrap_or_else(costs::default_cost);
|
||
Self {
|
||
model,
|
||
model_name: name,
|
||
input_cost,
|
||
output_cost,
|
||
cache_retention: CacheRetention::None,
|
||
unsupported_params: HashSet::new(),
|
||
}
|
||
}
|
||
|
||
/// Set Anthropic prompt cache retention policy.
|
||
///
|
||
/// Controls both cache injection and cost tracking:
|
||
/// - `None` — no caching, no surcharge (1.0×).
|
||
/// - `Short` — 5-minute TTL via `{"type": "ephemeral"}`, 1.25× write surcharge.
|
||
/// - `Long` — 1-hour TTL via `{"type": "ephemeral", "ttl": "1h"}`, 2.0× write surcharge.
|
||
///
|
||
/// Cache injection uses Anthropic's **automatic caching** — a top-level
|
||
/// `cache_control` field in `additional_params` that gets `#[serde(flatten)]`'d
|
||
/// into the request body by rig-core.
|
||
///
|
||
/// If the configured model does not support caching (e.g. claude-2),
|
||
/// a warning is logged once at construction and caching is disabled.
|
||
pub fn with_cache_retention(mut self, retention: CacheRetention) -> Self {
|
||
if retention != CacheRetention::None && !supports_prompt_cache(&self.model_name) {
|
||
tracing::warn!(
|
||
model = %self.model_name,
|
||
"Prompt caching requested but model does not support it; disabling"
|
||
);
|
||
self.cache_retention = CacheRetention::None;
|
||
} else {
|
||
self.cache_retention = retention;
|
||
}
|
||
self
|
||
}
|
||
|
||
/// Set the list of unsupported parameter names for this provider.
|
||
///
|
||
/// Parameters in this set are stripped from requests before sending.
|
||
/// Supported parameter names: `"temperature"`, `"max_tokens"`, `"stop_sequences"`.
|
||
pub fn with_unsupported_params(mut self, params: Vec<String>) -> Self {
|
||
self.unsupported_params = params.into_iter().collect();
|
||
self
|
||
}
|
||
|
||
/// Strip unsupported fields from a `CompletionRequest` in place.
|
||
fn strip_unsupported_completion_params(&self, req: &mut CompletionRequest) {
|
||
strip_unsupported_completion_params(&self.unsupported_params, req);
|
||
}
|
||
|
||
/// Strip unsupported fields from a `ToolCompletionRequest` in place.
|
||
fn strip_unsupported_tool_params(&self, req: &mut ToolCompletionRequest) {
|
||
strip_unsupported_tool_params(&self.unsupported_params, req);
|
||
}
|
||
}
|
||
|
||
// -- Type conversion helpers --
|
||
|
||
/// Round an f32 to f64 without precision artifacts.
|
||
///
|
||
/// Direct `f32 as f64` preserves the binary representation, producing values
|
||
/// like `0.699999988079071` instead of `0.7`. Some providers (e.g. Zhipu/GLM)
|
||
/// reject these values with a 400 error. Rounding to 6 decimal places removes
|
||
/// the artifact while preserving all meaningful precision for temperature.
|
||
fn round_f32_to_f64(val: f32) -> f64 {
|
||
((val as f64) * 1_000_000.0).round() / 1_000_000.0
|
||
}
|
||
|
||
/// Normalize a JSON Schema for OpenAI strict mode compliance.
|
||
///
|
||
/// OpenAI strict function calling requires:
|
||
/// - Every object must have `"additionalProperties": false`
|
||
/// - `"required"` must list ALL property keys
|
||
/// - Optional fields use `"type": ["<original>", "null"]` instead of being omitted from `required`
|
||
/// - Nested objects and array items are recursively normalized
|
||
///
|
||
/// This is applied as a clone-and-transform at the provider boundary so the
|
||
/// original tool definitions remain unchanged for other providers.
|
||
pub(crate) fn normalize_schema_strict(schema: &JsonValue) -> JsonValue {
|
||
let mut schema = schema.clone();
|
||
normalize_schema_recursive(&mut schema);
|
||
schema
|
||
}
|
||
|
||
fn normalize_schema_recursive(schema: &mut JsonValue) {
|
||
let obj = match schema.as_object_mut() {
|
||
Some(o) => o,
|
||
None => return,
|
||
};
|
||
|
||
// Recurse into combinators: anyOf, oneOf, allOf
|
||
for key in &["anyOf", "oneOf", "allOf"] {
|
||
if let Some(JsonValue::Array(variants)) = obj.get_mut(*key) {
|
||
for variant in variants.iter_mut() {
|
||
normalize_schema_recursive(variant);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Recurse into array items
|
||
if let Some(items) = obj.get_mut("items") {
|
||
normalize_schema_recursive(items);
|
||
}
|
||
|
||
// Recurse into `not`, `if`, `then`, `else`
|
||
for key in &["not", "if", "then", "else"] {
|
||
if let Some(sub) = obj.get_mut(*key) {
|
||
normalize_schema_recursive(sub);
|
||
}
|
||
}
|
||
|
||
// Only apply object-level normalization if this schema has "properties"
|
||
// (explicit object schema) or type == "object"
|
||
let is_object = obj
|
||
.get("type")
|
||
.and_then(|t| t.as_str())
|
||
.map(|t| t == "object")
|
||
.unwrap_or(false);
|
||
let has_properties = obj.contains_key("properties");
|
||
|
||
if !is_object && !has_properties {
|
||
return;
|
||
}
|
||
|
||
// Ensure "type": "object" is present
|
||
if !obj.contains_key("type") && has_properties {
|
||
obj.insert("type".to_string(), JsonValue::String("object".to_string()));
|
||
}
|
||
|
||
// Force additionalProperties: false (overwrite any existing value)
|
||
obj.insert("additionalProperties".to_string(), JsonValue::Bool(false));
|
||
|
||
// Ensure "properties" exists
|
||
if !obj.contains_key("properties") {
|
||
obj.insert(
|
||
"properties".to_string(),
|
||
JsonValue::Object(serde_json::Map::new()),
|
||
);
|
||
}
|
||
|
||
// Collect current required set
|
||
let current_required: std::collections::HashSet<String> = obj
|
||
.get("required")
|
||
.and_then(|r| r.as_array())
|
||
.map(|arr| {
|
||
arr.iter()
|
||
.filter_map(|v| v.as_str().map(String::from))
|
||
.collect()
|
||
})
|
||
.unwrap_or_default();
|
||
|
||
// Get all property keys (sorted for deterministic output)
|
||
let all_keys: Vec<String> = obj
|
||
.get("properties")
|
||
.and_then(|p| p.as_object())
|
||
.map(|props| {
|
||
let mut keys: Vec<String> = props.keys().cloned().collect();
|
||
keys.sort();
|
||
keys
|
||
})
|
||
.unwrap_or_default();
|
||
|
||
// For properties NOT in the original required list, make them nullable
|
||
if let Some(JsonValue::Object(props)) = obj.get_mut("properties") {
|
||
for key in &all_keys {
|
||
// Recurse into each property's schema FIRST (before make_nullable,
|
||
// which may change the type to an array and prevent object detection)
|
||
if let Some(prop_schema) = props.get_mut(key) {
|
||
normalize_schema_recursive(prop_schema);
|
||
}
|
||
// Then make originally-optional properties nullable
|
||
if !current_required.contains(key)
|
||
&& let Some(prop_schema) = props.get_mut(key)
|
||
{
|
||
make_nullable(prop_schema);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Set required to ALL property keys
|
||
let required_value: Vec<JsonValue> = all_keys.into_iter().map(JsonValue::String).collect();
|
||
obj.insert("required".to_string(), JsonValue::Array(required_value));
|
||
}
|
||
|
||
/// Make a property schema nullable for OpenAI strict mode.
|
||
///
|
||
/// If it has a simple `"type": "<T>"`, converts to `"type": ["<T>", "null"]`.
|
||
/// If it already has an array type, adds "null" if not present.
|
||
/// Otherwise, wraps with `anyOf: [<existing>, {"type": "null"}]`.
|
||
fn make_nullable(schema: &mut JsonValue) {
|
||
let obj = match schema.as_object_mut() {
|
||
Some(o) => o,
|
||
None => return,
|
||
};
|
||
|
||
if let Some(type_val) = obj.get("type").cloned() {
|
||
match type_val {
|
||
// "type": "string" → "type": ["string", "null"]
|
||
JsonValue::String(ref t) if t != "null" => {
|
||
obj.insert("type".to_string(), serde_json::json!([t, "null"]));
|
||
}
|
||
// "type": ["string", "integer"] → add "null" if missing
|
||
JsonValue::Array(ref arr) => {
|
||
let has_null = arr.iter().any(|v| v.as_str() == Some("null"));
|
||
if !has_null {
|
||
let mut new_arr = arr.clone();
|
||
new_arr.push(JsonValue::String("null".to_string()));
|
||
obj.insert("type".to_string(), JsonValue::Array(new_arr));
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
} else {
|
||
// No "type" key — wrap with anyOf including null
|
||
// (handles enum-only, $ref, or combinator schemas)
|
||
let existing = JsonValue::Object(obj.clone());
|
||
obj.clear();
|
||
obj.insert(
|
||
"anyOf".to_string(),
|
||
serde_json::json!([existing, {"type": "null"}]),
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Convert IronClaw messages to rig-core format.
|
||
///
|
||
/// Returns `(preamble, chat_history)` where preamble is extracted from
|
||
/// any System message and chat_history contains the rest.
|
||
fn convert_messages(messages: &[ChatMessage]) -> (Option<String>, Vec<RigMessage>) {
|
||
let mut preamble: Option<String> = None;
|
||
let mut history = Vec::new();
|
||
|
||
for msg in messages {
|
||
match msg.role {
|
||
crate::llm::Role::System => {
|
||
// Concatenate system messages into preamble
|
||
match preamble {
|
||
Some(ref mut p) => {
|
||
p.push('\n');
|
||
p.push_str(&msg.content);
|
||
}
|
||
None => preamble = Some(msg.content.clone()),
|
||
}
|
||
}
|
||
crate::llm::Role::User => {
|
||
if msg.content_parts.is_empty() {
|
||
history.push(RigMessage::user(&msg.content));
|
||
} else {
|
||
// Build multimodal user message with text + image parts
|
||
let mut contents: Vec<UserContent> = vec![UserContent::text(&msg.content)];
|
||
for part in &msg.content_parts {
|
||
if let crate::llm::ContentPart::ImageUrl { image_url } = part {
|
||
// Parse data: URL for base64 images, or use raw URL
|
||
let image = if let Some(rest) = image_url.url.strip_prefix("data:") {
|
||
// Format: data:<mime>;base64,<data>
|
||
let (mime, b64) =
|
||
rest.split_once(";base64,").unwrap_or(("image/jpeg", rest));
|
||
Image {
|
||
data: DocumentSourceKind::base64(b64),
|
||
media_type: ImageMediaType::from_mime_type(mime),
|
||
detail: None,
|
||
additional_params: None,
|
||
}
|
||
} else {
|
||
Image {
|
||
data: DocumentSourceKind::url(&image_url.url),
|
||
media_type: None,
|
||
detail: None,
|
||
additional_params: None,
|
||
}
|
||
};
|
||
contents.push(UserContent::Image(image));
|
||
}
|
||
}
|
||
if let Ok(many) = OneOrMany::many(contents) {
|
||
history.push(RigMessage::User { content: many });
|
||
} else {
|
||
history.push(RigMessage::user(&msg.content));
|
||
}
|
||
}
|
||
}
|
||
crate::llm::Role::Assistant => {
|
||
if let Some(ref tool_calls) = msg.tool_calls {
|
||
// Assistant message with tool calls
|
||
let mut contents: Vec<AssistantContent> = Vec::new();
|
||
if !msg.content.is_empty() {
|
||
contents.push(AssistantContent::text(&msg.content));
|
||
}
|
||
for (idx, tc) in tool_calls.iter().enumerate() {
|
||
let tool_call_id =
|
||
normalized_tool_call_id(Some(tc.id.as_str()), history.len() + idx);
|
||
contents.push(AssistantContent::ToolCall(
|
||
rig::message::ToolCall::new(
|
||
tool_call_id.clone(),
|
||
ToolFunction::new(tc.name.clone(), tc.arguments.clone()),
|
||
)
|
||
.with_call_id(tool_call_id),
|
||
));
|
||
}
|
||
if let Ok(many) = OneOrMany::many(contents) {
|
||
history.push(RigMessage::Assistant {
|
||
id: None,
|
||
content: many,
|
||
});
|
||
} else {
|
||
// Shouldn't happen but fall back to text
|
||
history.push(RigMessage::assistant(&msg.content));
|
||
}
|
||
} else {
|
||
history.push(RigMessage::assistant(&msg.content));
|
||
}
|
||
}
|
||
crate::llm::Role::Tool => {
|
||
// Tool result message: wrap as User { ToolResult }.
|
||
// Merge consecutive tool results into a single User message
|
||
// so the API sees one multi-result message instead of
|
||
// multiple consecutive User messages (which Anthropic rejects).
|
||
let tool_id = normalized_tool_call_id(msg.tool_call_id.as_deref(), history.len());
|
||
let tool_result = UserContent::ToolResult(RigToolResult {
|
||
id: tool_id.clone(),
|
||
call_id: Some(tool_id),
|
||
content: OneOrMany::one(ToolResultContent::text(&msg.content)),
|
||
});
|
||
|
||
let should_merge = matches!(
|
||
history.last(),
|
||
Some(RigMessage::User { content }) if content.iter().all(|c| matches!(c, UserContent::ToolResult(_)))
|
||
);
|
||
|
||
if should_merge {
|
||
if let Some(RigMessage::User { content }) = history.last_mut() {
|
||
content.push(tool_result);
|
||
}
|
||
} else {
|
||
history.push(RigMessage::User {
|
||
content: OneOrMany::one(tool_result),
|
||
});
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
(preamble, history)
|
||
}
|
||
|
||
/// Responses-style providers require a non-empty tool call ID.
|
||
///
|
||
/// IDs must be compatible with providers like Mistral, which constrain IDs
|
||
/// to `[a-zA-Z0-9]{9}`. We therefore:
|
||
/// - pass through any non-empty raw ID that already matches this constraint;
|
||
/// - otherwise deterministically map the raw string into a provider-compliant ID;
|
||
/// - and when `raw` is empty/None, delegate to `generate_tool_call_id`.
|
||
fn normalized_tool_call_id(raw: Option<&str>, seed: usize) -> String {
|
||
// Trim and treat empty as None.
|
||
let trimmed = raw.and_then(|s| {
|
||
let t = s.trim();
|
||
if t.is_empty() { None } else { Some(t) }
|
||
});
|
||
|
||
if let Some(id) = trimmed {
|
||
// If the ID already satisfies `[a-zA-Z0-9]{9}`, pass it through unchanged.
|
||
if id.len() == 9 && id.chars().all(|c| c.is_ascii_alphanumeric()) {
|
||
return id.to_string();
|
||
}
|
||
|
||
// Otherwise, deterministically hash the raw ID and feed the hash-derived
|
||
// seed into the provider-level generator so that the encoding and any
|
||
// provider-specific constraints remain centralized in one place.
|
||
let digest = Sha256::digest(id.as_bytes());
|
||
// Derive a 64-bit value from the first 8 bytes of the digest, then
|
||
// split it into two usize seeds so we preserve all 64 bits of entropy
|
||
// even on 32-bit targets.
|
||
let hash64 = {
|
||
// SHA-256 always produces 32 bytes, so indexing the first 8 is safe.
|
||
let bytes: [u8; 8] = [
|
||
digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6],
|
||
digest[7],
|
||
];
|
||
u64::from_be_bytes(bytes)
|
||
};
|
||
let hi_seed: usize = (hash64 >> 32) as usize;
|
||
let lo_seed: usize = (hash64 & 0xFFFF_FFFF) as usize;
|
||
return super::provider::generate_tool_call_id(hi_seed, lo_seed);
|
||
}
|
||
|
||
// Fallback for missing/empty raw IDs: use the provider-level generator,
|
||
// which already produces compliant IDs.
|
||
super::provider::generate_tool_call_id(seed, 0)
|
||
}
|
||
|
||
/// Convert IronClaw tool definitions to rig-core format.
|
||
///
|
||
/// Applies OpenAI strict-mode schema normalization to ensure all tool
|
||
/// parameter schemas comply with OpenAI's function calling requirements.
|
||
fn convert_tools(tools: &[IronToolDefinition]) -> Vec<RigToolDefinition> {
|
||
tools
|
||
.iter()
|
||
.map(|t| RigToolDefinition {
|
||
name: t.name.clone(),
|
||
description: t.description.clone(),
|
||
parameters: normalize_schema_strict(&t.parameters),
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Convert IronClaw tool_choice string to rig-core ToolChoice.
|
||
fn convert_tool_choice(choice: Option<&str>) -> Option<RigToolChoice> {
|
||
match choice.map(|s| s.to_lowercase()).as_deref() {
|
||
Some("auto") => Some(RigToolChoice::Auto),
|
||
Some("required") => Some(RigToolChoice::Required),
|
||
Some("none") => Some(RigToolChoice::None),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
/// Extract text and tool calls from a rig-core completion response.
|
||
fn extract_response(
|
||
choice: &OneOrMany<AssistantContent>,
|
||
_usage: &RigUsage,
|
||
) -> (Option<String>, Vec<IronToolCall>, FinishReason) {
|
||
let mut text_parts: Vec<String> = Vec::new();
|
||
let mut tool_calls: Vec<IronToolCall> = Vec::new();
|
||
|
||
for content in choice.iter() {
|
||
match content {
|
||
AssistantContent::Text(t) => {
|
||
if !t.text.is_empty() {
|
||
text_parts.push(t.text.clone());
|
||
}
|
||
}
|
||
AssistantContent::ToolCall(tc) => {
|
||
tool_calls.push(IronToolCall {
|
||
id: tc.id.clone(),
|
||
name: tc.function.name.clone(),
|
||
arguments: tc.function.arguments.clone(),
|
||
reasoning: None,
|
||
});
|
||
}
|
||
// Reasoning and Image variants are not mapped to IronClaw types
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
let text = if text_parts.is_empty() {
|
||
None
|
||
} else {
|
||
Some(text_parts.join(""))
|
||
};
|
||
|
||
let finish = if !tool_calls.is_empty() {
|
||
FinishReason::ToolUse
|
||
} else {
|
||
FinishReason::Stop
|
||
};
|
||
|
||
(text, tool_calls, finish)
|
||
}
|
||
|
||
/// Saturate u64 to u32 for token counts.
|
||
fn saturate_u32(val: u64) -> u32 {
|
||
val.min(u32::MAX as u64) as u32
|
||
}
|
||
|
||
/// Returns `true` if the model supports Anthropic prompt caching.
|
||
///
|
||
/// Per Anthropic docs, only Claude 3+ models support prompt caching.
|
||
/// Unsupported: claude-2, claude-2.1, claude-instant-*.
|
||
fn supports_prompt_cache(name: &str) -> bool {
|
||
let lower = name.to_lowercase();
|
||
// Strip optional provider prefix (e.g. "anthropic/claude-...")
|
||
let model = lower.strip_prefix("anthropic/").unwrap_or(&lower);
|
||
// Only Claude 3+ families support prompt caching
|
||
model.starts_with("claude-3")
|
||
|| model.starts_with("claude-4")
|
||
|| model.starts_with("claude-sonnet")
|
||
|| model.starts_with("claude-opus")
|
||
|| model.starts_with("claude-haiku")
|
||
}
|
||
|
||
/// Extract `cache_creation_input_tokens` from the raw provider response.
|
||
///
|
||
/// Rig-core's unified `Usage` does not surface this field, but Anthropic's raw
|
||
/// response includes it at `usage.cache_creation_input_tokens`. We serialize the
|
||
/// raw response to JSON and attempt to read the value.
|
||
fn extract_cache_creation<T: Serialize>(raw: &T) -> u32 {
|
||
serde_json::to_value(raw)
|
||
.ok()
|
||
.and_then(|v| v.get("usage")?.get("cache_creation_input_tokens")?.as_u64())
|
||
.map(|n| n.min(u32::MAX as u64) as u32)
|
||
.unwrap_or(0)
|
||
}
|
||
|
||
/// Build a rig-core CompletionRequest from our internal types.
|
||
///
|
||
/// When `cache_retention` is not `None`, injects a top-level `cache_control`
|
||
/// field via `additional_params`. Rig-core's `AnthropicCompletionRequest`
|
||
/// uses `#[serde(flatten)]` on `additional_params`, so the field lands at
|
||
/// the request root — which is exactly what Anthropic's **automatic caching**
|
||
/// expects. The API auto-places the cache breakpoint at the last cacheable
|
||
/// block and moves it forward as conversations grow.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn build_rig_request(
|
||
preamble: Option<String>,
|
||
mut history: Vec<RigMessage>,
|
||
tools: Vec<RigToolDefinition>,
|
||
tool_choice: Option<RigToolChoice>,
|
||
temperature: Option<f32>,
|
||
max_tokens: Option<u32>,
|
||
cache_retention: CacheRetention,
|
||
) -> Result<RigRequest, LlmError> {
|
||
// rig-core requires at least one message in chat_history
|
||
if history.is_empty() {
|
||
history.push(RigMessage::user("Hello"));
|
||
}
|
||
|
||
let chat_history = OneOrMany::many(history).map_err(|e| LlmError::RequestFailed {
|
||
provider: "rig".to_string(),
|
||
reason: format!("Failed to build chat history: {}", e),
|
||
})?;
|
||
|
||
// Inject top-level cache_control for Anthropic automatic prompt caching.
|
||
let additional_params = match cache_retention {
|
||
CacheRetention::None => None,
|
||
CacheRetention::Short => Some(serde_json::json!({
|
||
"cache_control": {"type": "ephemeral"}
|
||
})),
|
||
CacheRetention::Long => Some(serde_json::json!({
|
||
"cache_control": {"type": "ephemeral", "ttl": "1h"}
|
||
})),
|
||
};
|
||
|
||
Ok(RigRequest {
|
||
preamble,
|
||
chat_history,
|
||
documents: Vec::new(),
|
||
tools,
|
||
temperature: temperature.map(round_f32_to_f64),
|
||
max_tokens: max_tokens.map(|t| t as u64),
|
||
tool_choice,
|
||
additional_params,
|
||
})
|
||
}
|
||
|
||
/// Inject a per-request model override into the rig request's `additional_params`.
|
||
///
|
||
/// Rig-core bakes the model name at construction time inside each provider's
|
||
/// `CompletionModel` implementation. This helper inserts a top-level `"model"`
|
||
/// key into `additional_params`, which rig-core flattens into the provider's
|
||
/// request payload via `#[serde(flatten)]`.
|
||
///
|
||
/// Whether the override takes effect depends on the downstream API server's
|
||
/// handling of duplicate JSON keys (most Python/Go servers use last-key-wins,
|
||
/// but this is not guaranteed by the JSON spec). The `effective_model_name()`
|
||
/// trait method should be consulted to determine the model actually used.
|
||
fn inject_model_override(rig_req: &mut RigRequest, model_override: Option<&str>) {
|
||
let Some(model) = model_override else {
|
||
return;
|
||
};
|
||
match rig_req.additional_params {
|
||
Some(ref mut params) => {
|
||
if let Some(obj) = params.as_object_mut() {
|
||
obj.insert("model".to_string(), serde_json::json!(model));
|
||
}
|
||
}
|
||
None => {
|
||
rig_req.additional_params = Some(serde_json::json!({ "model": model }));
|
||
}
|
||
}
|
||
}
|
||
|
||
#[async_trait]
|
||
impl<M> LlmProvider for RigAdapter<M>
|
||
where
|
||
M: CompletionModel + Send + Sync + 'static,
|
||
M::Response: Send + Sync + Serialize + DeserializeOwned,
|
||
{
|
||
fn model_name(&self) -> &str {
|
||
&self.model_name
|
||
}
|
||
|
||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||
(self.input_cost, self.output_cost)
|
||
}
|
||
|
||
fn cache_write_multiplier(&self) -> Decimal {
|
||
match self.cache_retention {
|
||
CacheRetention::None => Decimal::ONE,
|
||
CacheRetention::Short => Decimal::new(125, 2), // 1.25× (125% of input rate)
|
||
CacheRetention::Long => Decimal::TWO, // 2.0× (200% of input rate)
|
||
}
|
||
}
|
||
|
||
fn cache_read_discount(&self) -> Decimal {
|
||
if self.cache_retention != CacheRetention::None {
|
||
dec!(10) // Anthropic: 90% discount (cost = input_rate / 10)
|
||
} else {
|
||
Decimal::ONE
|
||
}
|
||
}
|
||
|
||
async fn complete(
|
||
&self,
|
||
mut request: CompletionRequest,
|
||
) -> Result<CompletionResponse, LlmError> {
|
||
let model_override = request.model.take();
|
||
|
||
self.strip_unsupported_completion_params(&mut request);
|
||
|
||
let mut messages = request.messages;
|
||
crate::llm::provider::sanitize_tool_messages(&mut messages);
|
||
let (preamble, history) = convert_messages(&messages);
|
||
|
||
let mut rig_req = build_rig_request(
|
||
preamble,
|
||
history,
|
||
Vec::new(),
|
||
None,
|
||
request.temperature,
|
||
request.max_tokens,
|
||
self.cache_retention,
|
||
)?;
|
||
|
||
inject_model_override(&mut rig_req, model_override.as_deref());
|
||
|
||
let response =
|
||
self.model
|
||
.completion(rig_req)
|
||
.await
|
||
.map_err(|e| LlmError::RequestFailed {
|
||
provider: self.model_name.clone(),
|
||
reason: e.to_string(),
|
||
})?;
|
||
|
||
let (text, _tool_calls, finish) = extract_response(&response.choice, &response.usage);
|
||
|
||
let resp = CompletionResponse {
|
||
content: text.unwrap_or_default(),
|
||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||
finish_reason: finish,
|
||
cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens),
|
||
cache_creation_input_tokens: extract_cache_creation(&response.raw_response),
|
||
};
|
||
|
||
if resp.cache_read_input_tokens > 0 {
|
||
tracing::debug!(
|
||
model = %self.model_name,
|
||
input = resp.input_tokens,
|
||
output = resp.output_tokens,
|
||
cache_read = resp.cache_read_input_tokens,
|
||
"prompt cache hit",
|
||
);
|
||
}
|
||
|
||
Ok(resp)
|
||
}
|
||
|
||
async fn complete_with_tools(
|
||
&self,
|
||
mut request: ToolCompletionRequest,
|
||
) -> Result<ToolCompletionResponse, LlmError> {
|
||
let model_override = request.model.take();
|
||
|
||
self.strip_unsupported_tool_params(&mut request);
|
||
|
||
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());
|
||
|
||
let mut rig_req = build_rig_request(
|
||
preamble,
|
||
history,
|
||
tools,
|
||
tool_choice,
|
||
request.temperature,
|
||
request.max_tokens,
|
||
self.cache_retention,
|
||
)?;
|
||
|
||
inject_model_override(&mut rig_req, model_override.as_deref());
|
||
|
||
let response =
|
||
self.model
|
||
.completion(rig_req)
|
||
.await
|
||
.map_err(|e| LlmError::RequestFailed {
|
||
provider: self.model_name.clone(),
|
||
reason: e.to_string(),
|
||
})?;
|
||
|
||
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;
|
||
}
|
||
}
|
||
|
||
let resp = ToolCompletionResponse {
|
||
content: text,
|
||
tool_calls,
|
||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||
finish_reason: finish,
|
||
cache_read_input_tokens: saturate_u32(response.usage.cached_input_tokens),
|
||
cache_creation_input_tokens: extract_cache_creation(&response.raw_response),
|
||
};
|
||
|
||
if resp.cache_read_input_tokens > 0 {
|
||
tracing::debug!(
|
||
model = %self.model_name,
|
||
input = resp.input_tokens,
|
||
output = resp.output_tokens,
|
||
cache_read = resp.cache_read_input_tokens,
|
||
"prompt cache hit",
|
||
);
|
||
}
|
||
|
||
Ok(resp)
|
||
}
|
||
|
||
fn active_model_name(&self) -> String {
|
||
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.
|
||
Err(LlmError::RequestFailed {
|
||
provider: self.model_name.clone(),
|
||
reason: "Runtime model switching not supported for rig-core providers. \
|
||
Restart with a different model configured."
|
||
.to_string(),
|
||
})
|
||
}
|
||
}
|
||
|
||
/// 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::*;
|
||
|
||
#[test]
|
||
fn test_round_f32_to_f64_no_precision_artifacts() {
|
||
// Direct f32->f64 cast produces 0.699999988079071 instead of 0.7
|
||
assert_eq!(round_f32_to_f64(0.7_f32), 0.7_f64);
|
||
assert_eq!(round_f32_to_f64(0.5_f32), 0.5_f64);
|
||
assert_eq!(round_f32_to_f64(1.0_f32), 1.0_f64);
|
||
assert_eq!(round_f32_to_f64(0.0_f32), 0.0_f64);
|
||
// Original cast produces artifacts — our fix should not
|
||
assert_ne!(0.7_f32 as f64, 0.7_f64);
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_messages_system_to_preamble() {
|
||
let messages = vec![
|
||
ChatMessage::system("You are a helpful assistant."),
|
||
ChatMessage::user("Hello"),
|
||
];
|
||
let (preamble, history) = convert_messages(&messages);
|
||
assert_eq!(preamble, Some("You are a helpful assistant.".to_string()));
|
||
assert_eq!(history.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_messages_multiple_systems_concatenated() {
|
||
let messages = vec![
|
||
ChatMessage::system("System 1"),
|
||
ChatMessage::system("System 2"),
|
||
ChatMessage::user("Hi"),
|
||
];
|
||
let (preamble, history) = convert_messages(&messages);
|
||
assert_eq!(preamble, Some("System 1\nSystem 2".to_string()));
|
||
assert_eq!(history.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_messages_tool_result() {
|
||
// Use a conforming 9-char alphanumeric ID so it passes through unchanged.
|
||
let messages = vec![ChatMessage::tool_result(
|
||
"abcDE1234",
|
||
"search",
|
||
"result text",
|
||
)];
|
||
let (preamble, history) = convert_messages(&messages);
|
||
assert!(preamble.is_none());
|
||
assert_eq!(history.len(), 1);
|
||
// Tool results become User messages in rig-core
|
||
match &history[0] {
|
||
RigMessage::User { content } => match content.first() {
|
||
UserContent::ToolResult(r) => {
|
||
assert_eq!(r.id, "abcDE1234");
|
||
assert_eq!(r.call_id.as_deref(), Some("abcDE1234"));
|
||
}
|
||
other => panic!("Expected tool result content, got: {:?}", other),
|
||
},
|
||
other => panic!("Expected User message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_messages_assistant_with_tool_calls() {
|
||
// Use a conforming 9-char alphanumeric ID so it passes through unchanged.
|
||
let tc = IronToolCall {
|
||
id: "Xt7mK9pQ2".to_string(),
|
||
name: "search".to_string(),
|
||
arguments: serde_json::json!({"query": "test"}),
|
||
reasoning: None,
|
||
};
|
||
let msg = ChatMessage::assistant_with_tool_calls(Some("thinking".to_string()), vec![tc]);
|
||
let messages = vec![msg];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
assert_eq!(history.len(), 1);
|
||
match &history[0] {
|
||
RigMessage::Assistant { content, .. } => {
|
||
// Should have both text and tool call
|
||
assert!(content.iter().count() >= 2);
|
||
for item in content.iter() {
|
||
if let AssistantContent::ToolCall(tc) = item {
|
||
assert_eq!(tc.call_id.as_deref(), Some("Xt7mK9pQ2"));
|
||
}
|
||
}
|
||
}
|
||
other => panic!("Expected Assistant message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_messages_tool_result_without_id_gets_fallback() {
|
||
let messages = vec![ChatMessage {
|
||
role: crate::llm::Role::Tool,
|
||
content: "result text".to_string(),
|
||
content_parts: Vec::new(),
|
||
tool_call_id: None,
|
||
name: Some("search".to_string()),
|
||
tool_calls: None,
|
||
}];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
match &history[0] {
|
||
RigMessage::User { content } => match content.first() {
|
||
UserContent::ToolResult(r) => {
|
||
// Missing ID → normalized_tool_call_id generates a 9-char alphanumeric ID.
|
||
assert_eq!(
|
||
r.id.len(),
|
||
9,
|
||
"fallback ID should be 9 chars, got: {}",
|
||
r.id
|
||
);
|
||
assert!(r.id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
assert_eq!(r.call_id.as_deref(), Some(r.id.as_str()));
|
||
}
|
||
other => panic!("Expected tool result content, got: {:?}", other),
|
||
},
|
||
other => panic!("Expected User message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_tools() {
|
||
let tools = vec![IronToolDefinition {
|
||
name: "search".to_string(),
|
||
description: "Search the web".to_string(),
|
||
parameters: serde_json::json!({
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {"type": "string"}
|
||
}
|
||
}),
|
||
}];
|
||
let rig_tools = convert_tools(&tools);
|
||
assert_eq!(rig_tools.len(), 1);
|
||
assert_eq!(rig_tools[0].name, "search");
|
||
assert_eq!(rig_tools[0].description, "Search the web");
|
||
}
|
||
|
||
#[test]
|
||
fn test_convert_tool_choice() {
|
||
assert!(matches!(
|
||
convert_tool_choice(Some("auto")),
|
||
Some(RigToolChoice::Auto)
|
||
));
|
||
assert!(matches!(
|
||
convert_tool_choice(Some("required")),
|
||
Some(RigToolChoice::Required)
|
||
));
|
||
assert!(matches!(
|
||
convert_tool_choice(Some("none")),
|
||
Some(RigToolChoice::None)
|
||
));
|
||
assert!(matches!(
|
||
convert_tool_choice(Some("AUTO")),
|
||
Some(RigToolChoice::Auto)
|
||
));
|
||
assert!(convert_tool_choice(None).is_none());
|
||
assert!(convert_tool_choice(Some("unknown")).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_extract_response_text_only() {
|
||
let content = OneOrMany::one(AssistantContent::text("Hello world"));
|
||
let usage = RigUsage::new();
|
||
let (text, calls, finish) = extract_response(&content, &usage);
|
||
assert_eq!(text, Some("Hello world".to_string()));
|
||
assert!(calls.is_empty());
|
||
assert_eq!(finish, FinishReason::Stop);
|
||
}
|
||
|
||
#[test]
|
||
fn test_extract_response_tool_call() {
|
||
let tc = AssistantContent::tool_call("call_1", "search", serde_json::json!({"q": "test"}));
|
||
let content = OneOrMany::one(tc);
|
||
let usage = RigUsage::new();
|
||
let (text, calls, finish) = extract_response(&content, &usage);
|
||
assert!(text.is_none());
|
||
assert_eq!(calls.len(), 1);
|
||
assert_eq!(calls[0].name, "search");
|
||
assert_eq!(finish, FinishReason::ToolUse);
|
||
}
|
||
|
||
#[test]
|
||
fn test_assistant_tool_call_empty_id_gets_generated() {
|
||
let tc = IronToolCall {
|
||
id: "".to_string(),
|
||
name: "search".to_string(),
|
||
arguments: serde_json::json!({"query": "test"}),
|
||
reasoning: None,
|
||
};
|
||
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
|
||
match &history[0] {
|
||
RigMessage::Assistant { content, .. } => {
|
||
let tool_call = content.iter().find_map(|c| match c {
|
||
AssistantContent::ToolCall(tc) => Some(tc),
|
||
_ => None,
|
||
});
|
||
let tc = tool_call.expect("should have a tool call");
|
||
// Empty ID → normalized_tool_call_id generates a 9-char alphanumeric ID.
|
||
assert_eq!(
|
||
tc.id.len(),
|
||
9,
|
||
"generated id should be 9 chars, got: {}",
|
||
tc.id
|
||
);
|
||
assert!(tc.id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
assert_eq!(tc.call_id.as_deref(), Some(tc.id.as_str()));
|
||
}
|
||
other => panic!("Expected Assistant message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_assistant_tool_call_whitespace_id_gets_generated() {
|
||
let tc = IronToolCall {
|
||
id: " ".to_string(),
|
||
name: "search".to_string(),
|
||
arguments: serde_json::json!({"query": "test"}),
|
||
reasoning: None,
|
||
};
|
||
let messages = vec![ChatMessage::assistant_with_tool_calls(None, vec![tc])];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
|
||
match &history[0] {
|
||
RigMessage::Assistant { content, .. } => {
|
||
let tool_call = content.iter().find_map(|c| match c {
|
||
AssistantContent::ToolCall(tc) => Some(tc),
|
||
_ => None,
|
||
});
|
||
let tc = tool_call.expect("should have a tool call");
|
||
// Whitespace-only ID → normalized_tool_call_id generates a 9-char alphanumeric ID.
|
||
assert_eq!(
|
||
tc.id.len(),
|
||
9,
|
||
"generated id should be 9 chars, got: {}",
|
||
tc.id
|
||
);
|
||
assert!(tc.id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
}
|
||
other => panic!("Expected Assistant message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_assistant_and_tool_result_missing_ids_share_generated_id() {
|
||
// Simulate: assistant emits a tool call with empty id, then tool
|
||
// result arrives without an id. Both should get deterministic
|
||
// generated ids that match (based on their position in history).
|
||
let tc = IronToolCall {
|
||
id: "".to_string(),
|
||
name: "search".to_string(),
|
||
arguments: serde_json::json!({"query": "test"}),
|
||
reasoning: None,
|
||
};
|
||
let assistant_msg = ChatMessage::assistant_with_tool_calls(None, vec![tc]);
|
||
let tool_result_msg = ChatMessage {
|
||
role: crate::llm::Role::Tool,
|
||
content: "search results here".to_string(),
|
||
content_parts: Vec::new(),
|
||
tool_call_id: None,
|
||
name: Some("search".to_string()),
|
||
tool_calls: None,
|
||
};
|
||
let messages = vec![assistant_msg, tool_result_msg];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
|
||
// Extract the generated call_id from the assistant tool call
|
||
let assistant_call_id = match &history[0] {
|
||
RigMessage::Assistant { content, .. } => {
|
||
let tc = content.iter().find_map(|c| match c {
|
||
AssistantContent::ToolCall(tc) => Some(tc),
|
||
_ => None,
|
||
});
|
||
tc.expect("should have tool call").id.clone()
|
||
}
|
||
other => panic!("Expected Assistant message, got: {:?}", other),
|
||
};
|
||
|
||
// Extract the generated call_id from the tool result
|
||
let tool_result_call_id = match &history[1] {
|
||
RigMessage::User { content } => match content.first() {
|
||
UserContent::ToolResult(r) => r
|
||
.call_id
|
||
.clone()
|
||
.expect("tool result call_id must be present"),
|
||
other => panic!("Expected ToolResult, got: {:?}", other),
|
||
},
|
||
other => panic!("Expected User message, got: {:?}", other),
|
||
};
|
||
|
||
assert!(
|
||
!assistant_call_id.is_empty(),
|
||
"assistant call_id must not be empty"
|
||
);
|
||
assert!(
|
||
!tool_result_call_id.is_empty(),
|
||
"tool result call_id must not be empty"
|
||
);
|
||
|
||
// NOTE: With the current seed-based generation, these IDs will differ
|
||
// because the assistant tool call uses seed=0 (history.len() at that
|
||
// point) and the tool result uses seed=1 (history.len() after the
|
||
// assistant message was pushed). This documents the current behavior.
|
||
// A future improvement could thread the assistant's generated ID into
|
||
// the tool result for exact matching.
|
||
assert_ne!(
|
||
assistant_call_id, tool_result_call_id,
|
||
"Current impl generates different IDs for assistant call and tool result \
|
||
because seeds differ; this documents the known limitation"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_saturate_u32() {
|
||
assert_eq!(saturate_u32(100), 100);
|
||
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");
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_rig_request_injects_cache_control_short() {
|
||
let req = build_rig_request(
|
||
Some("You are helpful.".to_string()),
|
||
vec![RigMessage::user("Hello")],
|
||
Vec::new(),
|
||
None,
|
||
None,
|
||
None,
|
||
CacheRetention::Short,
|
||
)
|
||
.unwrap();
|
||
|
||
let params = req
|
||
.additional_params
|
||
.expect("should have additional_params for Short retention");
|
||
assert_eq!(params["cache_control"]["type"], "ephemeral");
|
||
assert!(
|
||
params["cache_control"].get("ttl").is_none(),
|
||
"Short retention should not include ttl"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_rig_request_injects_cache_control_long() {
|
||
let req = build_rig_request(
|
||
Some("You are helpful.".to_string()),
|
||
vec![RigMessage::user("Hello")],
|
||
Vec::new(),
|
||
None,
|
||
None,
|
||
None,
|
||
CacheRetention::Long,
|
||
)
|
||
.unwrap();
|
||
|
||
let params = req
|
||
.additional_params
|
||
.expect("should have additional_params for Long retention");
|
||
assert_eq!(params["cache_control"]["type"], "ephemeral");
|
||
assert_eq!(params["cache_control"]["ttl"], "1h");
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_rig_request_no_cache_control_when_none() {
|
||
let req = build_rig_request(
|
||
Some("You are helpful.".to_string()),
|
||
vec![RigMessage::user("Hello")],
|
||
Vec::new(),
|
||
None,
|
||
None,
|
||
None,
|
||
CacheRetention::None,
|
||
)
|
||
.unwrap();
|
||
|
||
assert!(
|
||
req.additional_params.is_none(),
|
||
"additional_params should be None when cache is disabled"
|
||
);
|
||
}
|
||
|
||
/// Verify that the multiplier match arms in `RigAdapter::cache_write_multiplier`
|
||
/// produce the expected values. We use a standalone helper because constructing
|
||
/// a real `RigAdapter` requires a rig `Model` (which needs network/provider setup).
|
||
/// The helper mirrors the same match expression — if the impl drifts, the
|
||
/// `test_build_rig_request_*` tests will still catch regressions end-to-end.
|
||
#[test]
|
||
fn test_cache_write_multiplier_values() {
|
||
use rust_decimal::Decimal;
|
||
// None → 1.0× (no surcharge)
|
||
assert_eq!(
|
||
cache_write_multiplier_for(CacheRetention::None),
|
||
Decimal::ONE
|
||
);
|
||
// Short → 1.25× (25% surcharge)
|
||
assert_eq!(
|
||
cache_write_multiplier_for(CacheRetention::Short),
|
||
Decimal::new(125, 2)
|
||
);
|
||
// Long → 2.0× (100% surcharge)
|
||
assert_eq!(
|
||
cache_write_multiplier_for(CacheRetention::Long),
|
||
Decimal::TWO
|
||
);
|
||
}
|
||
|
||
fn cache_write_multiplier_for(retention: CacheRetention) -> rust_decimal::Decimal {
|
||
match retention {
|
||
CacheRetention::None => rust_decimal::Decimal::ONE,
|
||
CacheRetention::Short => rust_decimal::Decimal::new(125, 2),
|
||
CacheRetention::Long => rust_decimal::Decimal::TWO,
|
||
}
|
||
}
|
||
|
||
// -- supports_prompt_cache tests --
|
||
|
||
#[test]
|
||
fn test_supports_prompt_cache_supported_models() {
|
||
// All Claude 3+ models per Anthropic docs
|
||
assert!(supports_prompt_cache("claude-opus-4-6"));
|
||
assert!(supports_prompt_cache("claude-sonnet-4-6"));
|
||
assert!(supports_prompt_cache("claude-sonnet-4"));
|
||
assert!(supports_prompt_cache("claude-haiku-4-5"));
|
||
assert!(supports_prompt_cache("claude-3-5-sonnet-20241022"));
|
||
assert!(supports_prompt_cache("claude-haiku-3"));
|
||
assert!(supports_prompt_cache("Claude-Opus-4-5")); // case-insensitive
|
||
assert!(supports_prompt_cache("anthropic/claude-sonnet-4-6")); // provider prefix
|
||
}
|
||
|
||
#[test]
|
||
fn test_supports_prompt_cache_unsupported_models() {
|
||
// Legacy Claude models that predate caching
|
||
assert!(!supports_prompt_cache("claude-2"));
|
||
assert!(!supports_prompt_cache("claude-2.1"));
|
||
assert!(!supports_prompt_cache("claude-instant-1.2"));
|
||
// Non-Claude models
|
||
assert!(!supports_prompt_cache("gpt-4o"));
|
||
assert!(!supports_prompt_cache("llama3"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_with_unsupported_params_populates_set() {
|
||
use rig::client::CompletionClient;
|
||
use rig::providers::openai;
|
||
|
||
let client: openai::Client = openai::Client::builder()
|
||
.api_key("test-key")
|
||
.base_url("http://localhost:0")
|
||
.build()
|
||
.unwrap();
|
||
let client = client.completions_api();
|
||
let model = client.completion_model("test-model");
|
||
let adapter = RigAdapter::new(model, "test-model")
|
||
.with_unsupported_params(vec!["temperature".to_string()]);
|
||
|
||
assert!(adapter.unsupported_params.contains("temperature"));
|
||
assert!(!adapter.unsupported_params.contains("max_tokens"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_strip_unsupported_completion_params() {
|
||
use rig::client::CompletionClient;
|
||
use rig::providers::openai;
|
||
|
||
let client: openai::Client = openai::Client::builder()
|
||
.api_key("test-key")
|
||
.base_url("http://localhost:0")
|
||
.build()
|
||
.unwrap();
|
||
let client = client.completions_api();
|
||
let model = client.completion_model("test-model");
|
||
let adapter = RigAdapter::new(model, "test-model").with_unsupported_params(vec![
|
||
"temperature".to_string(),
|
||
"stop_sequences".to_string(),
|
||
]);
|
||
|
||
let mut req = CompletionRequest::new(vec![ChatMessage::user("hi")]);
|
||
req.temperature = Some(0.7);
|
||
req.max_tokens = Some(100);
|
||
req.stop_sequences = Some(vec!["STOP".to_string()]);
|
||
|
||
adapter.strip_unsupported_completion_params(&mut req);
|
||
|
||
assert!(req.temperature.is_none(), "temperature should be stripped");
|
||
assert_eq!(req.max_tokens, Some(100), "max_tokens should be preserved");
|
||
assert!(
|
||
req.stop_sequences.is_none(),
|
||
"stop_sequences should be stripped"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_strip_unsupported_tool_params() {
|
||
use rig::client::CompletionClient;
|
||
use rig::providers::openai;
|
||
|
||
let client: openai::Client = openai::Client::builder()
|
||
.api_key("test-key")
|
||
.base_url("http://localhost:0")
|
||
.build()
|
||
.unwrap();
|
||
let client = client.completions_api();
|
||
let model = client.completion_model("test-model");
|
||
let adapter = RigAdapter::new(model, "test-model")
|
||
.with_unsupported_params(vec!["temperature".to_string(), "max_tokens".to_string()]);
|
||
|
||
let mut req = ToolCompletionRequest::new(vec![ChatMessage::user("hi")], vec![]);
|
||
req.temperature = Some(0.5);
|
||
req.max_tokens = Some(200);
|
||
|
||
adapter.strip_unsupported_tool_params(&mut req);
|
||
|
||
assert!(req.temperature.is_none(), "temperature should be stripped");
|
||
assert!(req.max_tokens.is_none(), "max_tokens should be stripped");
|
||
}
|
||
|
||
#[test]
|
||
fn test_unsupported_params_empty_by_default() {
|
||
use rig::client::CompletionClient;
|
||
use rig::providers::openai;
|
||
|
||
let client: openai::Client = openai::Client::builder()
|
||
.api_key("test-key")
|
||
.base_url("http://localhost:0")
|
||
.build()
|
||
.unwrap();
|
||
let client = client.completions_api();
|
||
let model = client.completion_model("test-model");
|
||
let adapter = RigAdapter::new(model, "test-model");
|
||
|
||
assert!(adapter.unsupported_params.is_empty());
|
||
}
|
||
|
||
/// Regression test: consecutive tool_result messages from parallel tool
|
||
/// execution must be merged into a single User message with multiple
|
||
/// ToolResult content items. Without merging, APIs like Anthropic reject
|
||
/// the request due to consecutive User messages.
|
||
#[test]
|
||
fn test_consecutive_tool_results_merged_into_single_user_message() {
|
||
let tc1 = IronToolCall {
|
||
id: "call_a".to_string(),
|
||
name: "search".to_string(),
|
||
arguments: serde_json::json!({"q": "rust"}),
|
||
reasoning: None,
|
||
};
|
||
let tc2 = IronToolCall {
|
||
id: "call_b".to_string(),
|
||
name: "fetch".to_string(),
|
||
arguments: serde_json::json!({"url": "https://example.com"}),
|
||
reasoning: None,
|
||
};
|
||
let assistant = ChatMessage::assistant_with_tool_calls(None, vec![tc1, tc2]);
|
||
let result_a = ChatMessage::tool_result("call_a", "search", "search results");
|
||
let result_b = ChatMessage::tool_result("call_b", "fetch", "fetch results");
|
||
|
||
let messages = vec![assistant, result_a, result_b];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
|
||
// Should be: 1 assistant + 1 merged user (not 1 assistant + 2 users)
|
||
assert_eq!(
|
||
history.len(),
|
||
2,
|
||
"Expected 2 messages (assistant + merged user), got {}",
|
||
history.len()
|
||
);
|
||
|
||
// The second message should contain both tool results
|
||
match &history[1] {
|
||
RigMessage::User { content } => {
|
||
assert_eq!(
|
||
content.len(),
|
||
2,
|
||
"Expected 2 tool results in merged user message, got {}",
|
||
content.len()
|
||
);
|
||
for item in content.iter() {
|
||
assert!(
|
||
matches!(item, UserContent::ToolResult(_)),
|
||
"Expected ToolResult content"
|
||
);
|
||
}
|
||
}
|
||
other => panic!("Expected User message, got: {:?}", other),
|
||
}
|
||
}
|
||
|
||
/// Verify that a tool_result after a non-tool User message is NOT merged.
|
||
#[test]
|
||
fn test_tool_result_after_user_text_not_merged() {
|
||
let user_msg = ChatMessage::user("hello");
|
||
let tool_msg = ChatMessage::tool_result("call_1", "search", "results");
|
||
|
||
let messages = vec![user_msg, tool_msg];
|
||
let (_preamble, history) = convert_messages(&messages);
|
||
|
||
// Should be 2 separate User messages (text user + tool result user)
|
||
assert_eq!(history.len(), 2);
|
||
}
|
||
|
||
// -- normalized_tool_call_id tests --
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_conforming_passthrough() {
|
||
// A 9-char alphanumeric ID should pass through unchanged.
|
||
let id = normalized_tool_call_id(Some("abcDE1234"), 42);
|
||
assert_eq!(id, "abcDE1234");
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_non_conforming_hashed() {
|
||
// An ID that doesn't match [a-zA-Z0-9]{9} should be hashed into one.
|
||
let id = normalized_tool_call_id(Some("call_abc_long_id"), 0);
|
||
assert_eq!(id.len(), 9);
|
||
assert!(id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
// Should NOT be the raw input.
|
||
assert_ne!(id, "call_abc_l");
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_empty_input() {
|
||
let id = normalized_tool_call_id(Some(""), 5);
|
||
assert_eq!(id.len(), 9);
|
||
assert!(id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_whitespace_input() {
|
||
let id = normalized_tool_call_id(Some(" "), 5);
|
||
assert_eq!(id.len(), 9);
|
||
assert!(id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
// Empty and whitespace-only with the same seed should produce identical results.
|
||
let id_empty = normalized_tool_call_id(Some(""), 5);
|
||
assert_eq!(id, id_empty);
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_none_input() {
|
||
let id = normalized_tool_call_id(None, 7);
|
||
assert_eq!(id.len(), 9);
|
||
assert!(id.chars().all(|c| c.is_ascii_alphanumeric()));
|
||
// None and empty string with same seed should produce identical results.
|
||
let id_empty = normalized_tool_call_id(Some(""), 7);
|
||
assert_eq!(id, id_empty);
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_deterministic() {
|
||
let id1 = normalized_tool_call_id(Some("call_xyz_123"), 0);
|
||
let id2 = normalized_tool_call_id(Some("call_xyz_123"), 0);
|
||
assert_eq!(id1, id2, "same input must produce same output");
|
||
}
|
||
|
||
#[test]
|
||
fn test_normalized_tool_call_id_different_inputs_differ() {
|
||
let id_a = normalized_tool_call_id(Some("call_aaa"), 0);
|
||
let id_b = normalized_tool_call_id(Some("call_bbb"), 0);
|
||
assert_ne!(
|
||
id_a, id_b,
|
||
"different raw IDs should produce different hashed IDs"
|
||
);
|
||
}
|
||
|
||
fn make_rig_request(additional_params: Option<serde_json::Value>) -> RigRequest {
|
||
RigRequest {
|
||
preamble: None,
|
||
chat_history: OneOrMany::one(RigMessage::user("test")),
|
||
documents: Vec::new(),
|
||
tools: Vec::new(),
|
||
temperature: None,
|
||
max_tokens: None,
|
||
tool_choice: None,
|
||
additional_params,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn test_inject_model_override_creates_params_when_none() {
|
||
let mut req = make_rig_request(None);
|
||
inject_model_override(&mut req, Some("test-model"));
|
||
|
||
let params = req
|
||
.additional_params
|
||
.expect("additional_params should be Some");
|
||
assert_eq!(params, serde_json::json!({ "model": "test-model" }));
|
||
}
|
||
|
||
#[test]
|
||
fn test_inject_model_override_preserves_existing_params() {
|
||
let mut req = make_rig_request(Some(serde_json::json!({
|
||
"cache_control": { "type": "ephemeral" },
|
||
})));
|
||
inject_model_override(&mut req, Some("override-model"));
|
||
|
||
let params = req.additional_params.expect("should remain Some");
|
||
let obj = params.as_object().expect("should be object");
|
||
assert_eq!(
|
||
obj.get("cache_control"),
|
||
Some(&serde_json::json!({ "type": "ephemeral" }))
|
||
);
|
||
assert_eq!(obj.get("model"), Some(&serde_json::json!("override-model")));
|
||
}
|
||
|
||
#[test]
|
||
fn test_inject_model_override_noop_when_none() {
|
||
let mut req = make_rig_request(None);
|
||
inject_model_override(&mut req, None);
|
||
assert!(req.additional_params.is_none());
|
||
}
|
||
}
|