mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +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: 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]> --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
1519 lines
57 KiB
Rust
1519 lines
57 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. The actual HTTP request body includes a
|
||
/// `model` field set by the provider. Rig-core's `#[serde(flatten)]` on
|
||
/// `additional_params` emits these fields AFTER the provider's own fields.
|
||
/// Most API servers (Python, Go) use last-key-wins when deserializing
|
||
/// duplicate JSON keys, so the injected `model` value takes effect.
|
||
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"
|
||
);
|
||
}
|
||
}
|