diff --git a/crates/ironclaw_engine/src/executor/compaction.rs b/crates/ironclaw_engine/src/executor/compaction.rs deleted file mode 100644 index 4fde5b0e8e..0000000000 --- a/crates/ironclaw_engine/src/executor/compaction.rs +++ /dev/null @@ -1,176 +0,0 @@ -//! Context compaction and token counting. -//! -//! When message history approaches the model's context limit, compaction -//! asks the LLM to summarize progress and resets the history. This follows -//! the official RLM pattern (compaction at 85% of context limit). - -use std::sync::Arc; - -use tracing::debug; - -use crate::traits::llm::{LlmBackend, LlmCallConfig}; -use crate::types::error::EngineError; -use crate::types::message::{MessageRole, ThreadMessage}; -use crate::types::step::{LlmResponse, TokenUsage}; - -/// Characters per token estimate when no tokenizer is available. -/// Conservative estimate (official RLM uses 4). -const CHARS_PER_TOKEN: usize = 4; - -/// Estimate token count for a list of messages. -/// -/// Uses character length / `CHARS_PER_TOKEN` as a rough estimate. -/// The official RLM uses tiktoken when available; we use this fallback -/// since we don't depend on a Python tokenizer. -pub fn estimate_tokens(messages: &[ThreadMessage]) -> usize { - let total_chars: usize = messages - .iter() - .map(|m| { - m.content.len() + m.action_name.as_ref().map_or(0, |n| n.len()) + 4 // overhead per message (role token, delimiters) - }) - .sum(); - total_chars.div_ceil(CHARS_PER_TOKEN) -} - -/// Check if compaction should be triggered. -/// -/// Returns `true` when estimated token count exceeds `threshold_pct` of -/// the model's context limit. -pub fn should_compact( - messages: &[ThreadMessage], - model_context_limit: usize, - threshold_pct: f64, -) -> bool { - let tokens = estimate_tokens(messages); - let threshold = (model_context_limit as f64 * threshold_pct) as usize; - tokens >= threshold -} - -/// The compaction prompt sent to the LLM. -const COMPACTION_PROMPT: &str = "\ -Summarize your progress so far in a concise but complete way. Include: -1. What you have accomplished -2. Key intermediate results and variable values -3. What still needs to be done -4. Any errors encountered and how they were handled - -Preserve all information needed to continue the task. Be specific about data values."; - -/// Compact the message history by asking the LLM to summarize. -/// -/// Returns the new (shorter) message list and the token usage from the -/// summarization call. The original messages are replaced with: -/// `[system_prompt, summary, continuation_note]` -/// -/// The full original messages are returned separately so the caller can -/// store them (e.g., in a `history` variable or event log). -pub async fn compact_messages( - messages: &[ThreadMessage], - llm: &Arc, - compaction_count: u32, -) -> Result { - // Build a summarization request from existing messages + prompt - let mut summarize_messages = messages.to_vec(); - summarize_messages.push(ThreadMessage::user(COMPACTION_PROMPT.to_string())); - - let config = LlmCallConfig { - force_text: true, - ..LlmCallConfig::default() - }; - - let output = llm.complete(&summarize_messages, &[], &config).await?; - - let summary_text = match output.response { - LlmResponse::Text(t) => t, - LlmResponse::ActionCalls { content, .. } | LlmResponse::Code { content, .. } => { - content.unwrap_or_else(|| "[compaction produced no summary]".into()) - } - }; - - // Preserve the system prompt (first message if it's a system message) - let system_msg = messages - .iter() - .find(|m| m.role == MessageRole::System) - .cloned(); - - // Build compacted history - let mut compacted = Vec::new(); - if let Some(sys) = system_msg { - compacted.push(sys); - } - compacted.push(ThreadMessage::assistant(summary_text.clone())); - compacted.push(ThreadMessage::user(format!( - "Your conversation has been compacted {n} time(s). \ - The summary above captures your progress. Continue working on the task.", - n = compaction_count + 1, - ))); - - let tokens_before = estimate_tokens(messages); - let tokens_after = estimate_tokens(&compacted); - - debug!( - tokens_before, - tokens_after, - compaction_count = compaction_count + 1, - "context compacted" - ); - - Ok(CompactionResult { - compacted_messages: compacted, - summary: summary_text, - tokens_used: output.usage, - tokens_before, - tokens_after, - }) -} - -/// Result of a compaction operation. -pub struct CompactionResult { - /// The new (shorter) message list. - pub compacted_messages: Vec, - /// The summary text produced by the LLM. - pub summary: String, - /// Tokens used by the summarization LLM call. - pub tokens_used: TokenUsage, - /// Estimated token count before compaction. - pub tokens_before: usize, - /// Estimated token count after compaction. - pub tokens_after: usize, -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn estimate_tokens_empty() { - assert_eq!(estimate_tokens(&[]), 0); - } - - #[test] - fn estimate_tokens_basic() { - let msgs = vec![ - ThreadMessage::system("Hello world"), // 11 chars + 4 overhead = 15 / 4 = 3.75 - ThreadMessage::user("Hi"), // 2 chars + 4 = 6 / 4 = 1.5 - ]; - let tokens = estimate_tokens(&msgs); - // (11+4 + 2+4) / 4 = 21/4 = 5.25 → 6 (ceiling) - assert!(tokens > 0); - assert!(tokens < 100); - } - - #[test] - fn should_compact_below_threshold() { - let msgs = vec![ThreadMessage::user("short message")]; - assert!(!should_compact(&msgs, 128_000, 0.85)); - } - - #[test] - fn should_compact_above_threshold() { - // Create a message large enough to trigger compaction at low limit - let big = "x".repeat(1000); - let msgs = vec![ThreadMessage::user(big)]; - // 1000 chars / 4 = 250 tokens. Context limit 200, threshold 85% = 170 - assert!(should_compact(&msgs, 200, 0.85)); - } -} diff --git a/crates/ironclaw_engine/src/executor/mod.rs b/crates/ironclaw_engine/src/executor/mod.rs index 65ce7cbeac..718ebcf1b7 100644 --- a/crates/ironclaw_engine/src/executor/mod.rs +++ b/crates/ironclaw_engine/src/executor/mod.rs @@ -5,7 +5,6 @@ //! - [`context`] — context building for LLM calls //! - [`intent`] — tool intent nudge detection -pub mod compaction; pub mod context; pub mod loop_engine; pub mod orchestrator; diff --git a/src/bridge/effect_adapter.rs b/src/bridge/effect_adapter.rs index 681f7a6815..4b5d930478 100644 --- a/src/bridge/effect_adapter.rs +++ b/src/bridge/effect_adapter.rs @@ -16,8 +16,8 @@ use tokio::sync::RwLock; use tracing::debug; use ironclaw_engine::{ - ActionDef, ActionResult, CapabilityLease, EffectExecutor, EngineError, MountError, - ThreadExecutionContext, WorkspaceMounts, + ActionDef, ActionResult, CapabilityLease, CapabilityRegistry, EffectExecutor, EngineError, + MountError, ThreadExecutionContext, WorkspaceMounts, }; use crate::auth::oauth::sanitize_auth_url; @@ -64,6 +64,12 @@ pub struct EffectBridgeAdapter { /// in Phase 5+) instead of the host tool. When unset, all tool calls run /// on the host as before. workspace_mounts: RwLock>>, + /// Engine capability registry. `available_actions()` reads this to surface + /// actions from non-v1 capabilities (e.g. `missions`) to the LLM. The v1 + /// `ToolRegistry` only covers built-in + extension tools; engine-native + /// capabilities like `missions` are registered here in `router.rs` and + /// would otherwise be invisible to the LLM despite having active leases. + capability_registry: RwLock>>, } impl EffectBridgeAdapter { @@ -84,6 +90,7 @@ impl EffectBridgeAdapter { auth_manager: RwLock::new(None), http_interceptor: RwLock::new(None), workspace_mounts: RwLock::new(None), + capability_registry: RwLock::new(None), } } @@ -98,6 +105,14 @@ impl EffectBridgeAdapter { *self.workspace_mounts.write().await = mounts; } + /// Install the engine capability registry so `available_actions()` can + /// surface actions from engine-native capabilities (missions, etc.) to + /// the LLM. Called once at bridge setup after `router.rs` has finished + /// registering all capabilities. + pub async fn set_capability_registry(&self, registry: Arc) { + *self.capability_registry.write().await = Some(registry); + } + /// Install the trace HTTP interceptor on this adapter. Every JobContext /// the adapter constructs for tool dispatch will carry a clone of this /// interceptor, so http-aware tools will record/replay through it. @@ -1242,7 +1257,7 @@ impl EffectExecutor for EffectBridgeAdapter { async fn available_actions( &self, - _leases: &[CapabilityLease], + leases: &[CapabilityLease], ) -> Result, EngineError> { let tool_defs = self.tools.tool_definitions().await; @@ -1293,6 +1308,44 @@ impl EffectExecutor for EffectBridgeAdapter { } } + // Surface actions from engine-native capabilities (e.g. `missions`). + // The v1 `ToolRegistry` path above only covers built-in + extension + // tools; capabilities registered directly against the engine + // (`CapabilityRegistry`) would otherwise be invisible to the LLM + // even though the thread holds active leases for them. Iterate + // leases so we only advertise what the current thread actually has + // access to, and skip the `"tools"` capability — that lease is + // reconciled dynamically from the v1 path already covered above. + if let Some(registry) = self.capability_registry.read().await.as_ref() { + let mut seen: HashSet = actions.iter().map(|a| a.name.clone()).collect(); + for lease in leases { + if lease.capability_name == "tools" { + continue; + } + let Some(cap) = registry.get(&lease.capability_name) else { + continue; + }; + for action in &cap.actions { + if !lease.granted_actions.covers(&action.name) { + continue; + } + // Defensive: apply the same v1-isolation filters we run + // on v1 tools. If a future engine capability registers + // an action under a v1-denylisted name (`create_job`, + // `tool_auth`, ...), the v1 filters above would have + // hidden it — the engine path must not become a + // silent bypass. + if is_v1_only_tool(&action.name) || is_v1_auth_tool(&action.name) { + continue; + } + if !seen.insert(action.name.clone()) { + continue; + } + actions.push(action.clone()); + } + } + } + actions.sort_by(|a, b| a.name.cmp(&b.name)); Ok(actions) @@ -3929,4 +3982,482 @@ mod tests { delete_result.output ); } + + // ── Phase 6 acceptance: full mission lifecycle through the bridge ── + // + // These tests pin the gateway-facing contract that v2 clients rely on: + // a mission round-trips through create → list → fire → complete and + // each step's response shape stays stable. Existing per-action tests + // above cover error paths; these cover the happy-path interactions + // between actions, which is where regressions tend to bite (e.g. + // status not surfacing in mission_list after complete, or fire not + // returning a thread_id for manual missions). + + /// Full lifecycle: create → list (present) → complete → list (Completed). + /// Pins the post-complete visibility of status through `mission_list`, + /// which a chat client polls to render terminal-state UI. + #[tokio::test] + async fn mission_full_lifecycle_via_execute_action() { + let adapter = make_adapter_with_missions().await; + let ctx = exec_ctx(ironclaw_engine::ThreadId::new(), Some("lc1")); + + // Create + let create = adapter + .execute_action( + "mission_create", + serde_json::json!({ + "name": "lifecycle-mission", + "goal": "exercise the full lifecycle", + "cadence": "0 9 * * *" + }), + &lease(), + &ctx, + ) + .await + .expect("create should succeed"); + assert!(!create.is_error, "create failed: {}", create.output); + let mission_id = create + .output + .get("mission_id") + .and_then(|v| v.as_str()) + .expect("create must return mission_id") + .to_string(); + + // List → present, status not yet Completed + let list = adapter + .execute_action("mission_list", serde_json::json!({}), &lease(), &ctx) + .await + .expect("list should succeed"); + let missions = list.output.as_array().expect("list output is array"); + let entry = missions + .iter() + .find(|m| m.get("id").and_then(|v| v.as_str()) == Some(mission_id.as_str())) + .expect("created mission must appear in list"); + let initial_status = entry.get("status").and_then(|v| v.as_str()).unwrap_or(""); + assert_ne!( + initial_status, "Completed", + "fresh mission should not be Completed; got status={initial_status}" + ); + + // Complete + let complete = adapter + .execute_action( + "mission_complete", + serde_json::json!({"id": mission_id}), + &lease(), + &ctx, + ) + .await + .expect("complete should succeed"); + assert!(!complete.is_error); + assert_eq!( + complete.output.get("status").and_then(|v| v.as_str()), + Some("completed") + ); + + // List again → Completed status now visible + let list_after = adapter + .execute_action("mission_list", serde_json::json!({}), &lease(), &ctx) + .await + .expect("list-after should succeed"); + let missions_after = list_after.output.as_array().expect("array"); + let entry_after = missions_after + .iter() + .find(|m| m.get("id").and_then(|v| v.as_str()) == Some(mission_id.as_str())) + .expect("mission still present after complete"); + let post_status = entry_after + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or(""); + assert_eq!( + post_status, "Completed", + "mission_list must surface Completed status after mission_complete; got {post_status}" + ); + } + + /// `mission_fire` on a manual-cadence mission returns a thread_id and + /// fired status. Pins the response shape gateway clients consume to + /// link the fired mission to its child thread. + #[tokio::test] + async fn mission_fire_returns_thread_id_for_manual_cadence_via_execute_action() { + let adapter = make_adapter_with_missions().await; + let ctx = exec_ctx(ironclaw_engine::ThreadId::new(), Some("fire1")); + + let create = adapter + .execute_action( + "mission_create", + serde_json::json!({ + "name": "fireable", + "goal": "test fire flow", + "cadence": "manual" + }), + &lease(), + &ctx, + ) + .await + .expect("create should succeed"); + let mission_id = create + .output + .get("mission_id") + .and_then(|v| v.as_str()) + .expect("mission_id present") + .to_string(); + + let fire = adapter + .execute_action( + "mission_fire", + serde_json::json!({"id": mission_id}), + &lease(), + &ctx, + ) + .await + .expect("fire should succeed"); + + assert!(!fire.is_error, "fire failed: {}", fire.output); + // Two terminal shapes are valid: (a) {thread_id, status="fired"} + // when the mission ran; (b) {status="not_fired", reason} when + // budget/cooldown gated it. A fresh manual mission has no + // budget — must produce shape (a). + assert_eq!( + fire.output.get("status").and_then(|v| v.as_str()), + Some("fired"), + "fresh manual mission should fire successfully, got: {}", + fire.output + ); + let thread_id = fire + .output + .get("thread_id") + .and_then(|v| v.as_str()) + .expect("fired response must include thread_id"); + assert!( + uuid::Uuid::parse_str(thread_id).is_ok(), + "thread_id must be a valid UUID, got {thread_id:?}", + ); + } + + /// `mission_list` returns every mission the user created in the + /// current project, isolated from other users. Pins the per-user + /// scoping that chat history and project-detail pages rely on. + #[tokio::test] + async fn mission_list_returns_all_user_missions_via_execute_action() { + let adapter = make_adapter_with_missions().await; + let ctx = exec_ctx(ironclaw_engine::ThreadId::new(), Some("list1")); + + let names = ["alpha", "beta", "gamma"]; + for name in names { + let r = adapter + .execute_action( + "mission_create", + serde_json::json!({ + "name": name, + "goal": format!("test {name}"), + "cadence": "manual" + }), + &lease(), + &ctx, + ) + .await + .expect("create should succeed"); + assert!(!r.is_error, "create {name} failed: {}", r.output); + } + + let list = adapter + .execute_action("mission_list", serde_json::json!({}), &lease(), &ctx) + .await + .expect("list should succeed"); + let missions = list.output.as_array().expect("array"); + let listed_names: Vec<&str> = missions + .iter() + .filter_map(|m| m.get("name").and_then(|v| v.as_str())) + .collect(); + for expected in names { + assert!( + listed_names.contains(&expected), + "expected mission {expected:?} in list, got: {listed_names:?}" + ); + } + } + + // ── available_actions surfaces engine-registered capability actions ── + // + // Regression: without the capability registry, `available_actions` + // returned only v1 `ToolRegistry` tools + latent OAuth actions, so + // the LLM never saw mission tools in its tools list even though the + // thread held an active `missions` lease. This test pins that a + // thread with a mission lease gets `mission_*` advertised. + + fn mission_capability() -> ironclaw_engine::Capability { + ironclaw_engine::Capability { + name: "missions".into(), + description: "Mission lifecycle".into(), + actions: vec![ + ActionDef { + name: "mission_create".into(), + description: "Create a mission".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ActionDef { + name: "mission_list".into(), + description: "List missions".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ActionDef { + name: "mission_complete".into(), + description: "Complete a mission".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ], + knowledge: vec![], + policies: vec![], + } + } + + fn mission_lease(granted: &[&str]) -> ironclaw_engine::CapabilityLease { + ironclaw_engine::CapabilityLease { + id: ironclaw_engine::types::capability::LeaseId::new(), + thread_id: ironclaw_engine::ThreadId::new(), + capability_name: "missions".into(), + granted_actions: ironclaw_engine::GrantedActions::Specific( + granted.iter().map(|s| s.to_string()).collect(), + ), + granted_at: chrono::Utc::now(), + expires_at: None, + max_uses: None, + uses_remaining: None, + revoked: false, + revoked_reason: None, + } + } + + #[tokio::test] + async fn available_actions_surfaces_leased_mission_capability() { + let adapter = make_adapter(); + let mut registry = CapabilityRegistry::new(); + registry.register(mission_capability()); + adapter.set_capability_registry(Arc::new(registry)).await; + + let actions = adapter + .available_actions(&[mission_lease(&[ + "mission_create", + "mission_list", + "mission_complete", + ])]) + .await + .expect("available_actions should succeed"); + + let names: Vec<&str> = actions.iter().map(|a| a.name.as_str()).collect(); + for expected in ["mission_create", "mission_list", "mission_complete"] { + assert!( + names.contains(&expected), + "expected {expected} in advertised actions, got: {names:?}" + ); + } + } + + #[tokio::test] + async fn available_actions_respects_partial_lease_grant() { + let adapter = make_adapter(); + let mut registry = CapabilityRegistry::new(); + registry.register(mission_capability()); + adapter.set_capability_registry(Arc::new(registry)).await; + + // Lease only grants mission_list; mission_create / mission_complete + // must NOT be advertised to the LLM even though they exist in the + // capability registry. + let actions = adapter + .available_actions(&[mission_lease(&["mission_list"])]) + .await + .expect("available_actions should succeed"); + + let names: Vec<&str> = actions.iter().map(|a| a.name.as_str()).collect(); + assert!( + names.contains(&"mission_list"), + "mission_list should be advertised: {names:?}" + ); + assert!( + !names.contains(&"mission_create"), + "mission_create must not leak when lease did not grant it: {names:?}" + ); + assert!( + !names.contains(&"mission_complete"), + "mission_complete must not leak when lease did not grant it: {names:?}" + ); + } + + #[tokio::test] + async fn available_actions_omits_capability_without_lease() { + let adapter = make_adapter(); + let mut registry = CapabilityRegistry::new(); + registry.register(mission_capability()); + adapter.set_capability_registry(Arc::new(registry)).await; + + // No leases passed — no capability actions should surface even + // though the registry has them. + let actions = adapter + .available_actions(&[]) + .await + .expect("available_actions should succeed"); + + let names: Vec<&str> = actions.iter().map(|a| a.name.as_str()).collect(); + for name in ["mission_create", "mission_list", "mission_complete"] { + assert!( + !names.contains(&name), + "{name} must not appear without a lease: {names:?}" + ); + } + } + + /// Trivial v1 tool for the combined advertising test. Keeps the test + /// close to the helper so it doesn't pollute the top-level tool list. + struct V1EchoTool; + + #[async_trait] + impl Tool for V1EchoTool { + fn name(&self) -> &str { + "v1_echo" + } + fn description(&self) -> &str { + "v1 echo tool" + } + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object"}) + } + async fn execute( + &self, + _: serde_json::Value, + _: &JobContext, + ) -> Result { + Ok(ToolOutput::success( + serde_json::json!({}), + std::time::Duration::from_millis(1), + )) + } + } + + #[tokio::test] + async fn available_actions_merges_v1_tools_with_engine_capabilities() { + // Exercises the real production shape: the adapter has both a v1 + // `ToolRegistry` (echo tool) and a capability registry (missions). + // With a missions lease active, the LLM's tools list must include + // BOTH. Prior tests covered each path in isolation; this pins the + // combined advertising on the same call. + use ironclaw_safety::SafetyConfig; + + let tools = Arc::new(ToolRegistry::new()); + tools.register(Arc::new(V1EchoTool)).await; + let adapter = EffectBridgeAdapter::new( + tools, + Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 10_000, + injection_check_enabled: false, + })), + Arc::new(HookRegistry::default()), + ); + + let mut registry = CapabilityRegistry::new(); + registry.register(mission_capability()); + adapter.set_capability_registry(Arc::new(registry)).await; + + let actions = adapter + .available_actions(&[mission_lease(&[ + "mission_create", + "mission_list", + "mission_complete", + ])]) + .await + .expect("available_actions should succeed"); + + let names: Vec<&str> = actions.iter().map(|a| a.name.as_str()).collect(); + assert!( + names.contains(&"v1_echo"), + "v1 tool should be advertised: {names:?}" + ); + for mission in ["mission_create", "mission_list", "mission_complete"] { + assert!( + names.contains(&mission), + "engine capability action {mission} should be advertised alongside v1 tools: {names:?}" + ); + } + } + + /// Defensive: an engine capability must not be able to sneak a + /// v1-denylisted action (`create_job` etc.) past the v1-isolation + /// filters by registering under a different capability name. The + /// engine-capability path applies the same `is_v1_only_tool` / + /// `is_v1_auth_tool` gates as the v1 path. + #[tokio::test] + async fn available_actions_filters_v1_denylisted_names_from_engine_capabilities() { + let adapter = make_adapter(); + let mut registry = CapabilityRegistry::new(); + // A hypothetical malformed capability that tries to expose v1 + // tools through the v2 advertising path. + registry.register(ironclaw_engine::Capability { + name: "rogue".into(), + description: "should not surface denylisted v1 names".into(), + actions: vec![ + ActionDef { + name: "create_job".into(), // v1-only denylist + description: "forbidden".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ActionDef { + name: "tool_auth".into(), // v1 auth tool + description: "forbidden".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ActionDef { + name: "safe_action".into(), + description: "allowed".into(), + parameters_schema: serde_json::json!({"type": "object"}), + effects: vec![], + requires_approval: false, + }, + ], + knowledge: vec![], + policies: vec![], + }); + adapter.set_capability_registry(Arc::new(registry)).await; + + let rogue_lease = ironclaw_engine::CapabilityLease { + id: ironclaw_engine::types::capability::LeaseId::new(), + thread_id: ironclaw_engine::ThreadId::new(), + capability_name: "rogue".into(), + granted_actions: ironclaw_engine::GrantedActions::All, + granted_at: chrono::Utc::now(), + expires_at: None, + max_uses: None, + uses_remaining: None, + revoked: false, + revoked_reason: None, + }; + + let actions = adapter + .available_actions(&[rogue_lease]) + .await + .expect("available_actions should succeed"); + + let names: Vec<&str> = actions.iter().map(|a| a.name.as_str()).collect(); + assert!( + !names.contains(&"create_job"), + "create_job is v1-denylisted and must not surface via engine capability: {names:?}" + ); + assert!( + !names.contains(&"tool_auth"), + "tool_auth is a v1 auth tool and must not surface via engine capability: {names:?}" + ); + assert!( + names.contains(&"safe_action"), + "safe_action should surface through the engine capability path: {names:?}" + ); + } } diff --git a/src/bridge/llm_adapter.rs b/src/bridge/llm_adapter.rs index b7469d63c2..f639dba71a 100644 --- a/src/bridge/llm_adapter.rs +++ b/src/bridge/llm_adapter.rs @@ -6,12 +6,64 @@ use ironclaw_engine::{ ActionDef, EngineError, LlmBackend, LlmCallConfig, LlmOutput, LlmResponse, ThreadMessage, TokenUsage, }; +use rust_decimal::Decimal; +use rust_decimal::prelude::ToPrimitive; use crate::llm::{ ChatMessage, LlmProvider, Role, ToolCall, ToolCompletionRequest, ToolDefinition, sanitize_tool_messages, }; +/// Compute the USD cost of a single completion response, honoring the +/// provider's prompt-caching pricing. Mirrors the formula in +/// `src/agent/cost_guard.rs::CostGuard::record_llm_call` so engine v2's +/// `Thread::total_cost_usd` matches what `max_budget_usd` / v1's daily +/// budget enforcer would have computed: +/// +/// * uncached input tokens are priced at `cost_per_token().0`; +/// * cache-read tokens are discounted by `cache_read_discount()` (10x +/// off for Anthropic, 2x for OpenAI); +/// * cache-write tokens are multiplied by `cache_write_multiplier()` +/// (1.25× for Anthropic 5m TTL, 2× for 1h); +/// * output tokens are priced at `cost_per_token().1`. +/// +/// Returns 0.0 for subscription-billed providers that report +/// `cost_per_token() == (0, 0)` (e.g. OpenAI Codex via ChatGPT OAuth). +fn cost_usd_from( + provider: &Arc, + input_tokens: u32, + output_tokens: u32, + cache_read_input_tokens: u32, + cache_creation_input_tokens: u32, +) -> f64 { + let (input_rate, output_rate) = provider.cost_per_token(); + + // `input_tokens` is the provider-reported total. Cache tokens are + // already counted inside that total, so the uncached remainder is + // what's left after subtracting both buckets. + let cached_total = cache_read_input_tokens.saturating_add(cache_creation_input_tokens); + let uncached_input = input_tokens.saturating_sub(cached_total); + + // Guard against providers reporting a zero discount — treat zero as + // "no discount" rather than attempting a div-by-zero. + let discount = provider.cache_read_discount(); + let effective_discount = if discount.is_zero() { + Decimal::ONE + } else { + discount + }; + + let cache_read_cost = input_rate * Decimal::from(cache_read_input_tokens) / effective_discount; + let cache_write_cost = + input_rate * Decimal::from(cache_creation_input_tokens) * provider.cache_write_multiplier(); + let cost = input_rate * Decimal::from(uncached_input) + + cache_read_cost + + cache_write_cost + + output_rate * Decimal::from(output_tokens); + + cost.to_f64().unwrap_or(0.0) +} + /// Wraps an existing `LlmProvider` to implement the engine's `LlmBackend` trait. pub struct LlmBridgeAdapter { provider: Arc, @@ -97,7 +149,13 @@ impl LlmBackend for LlmBridgeAdapter { output_tokens: u64::from(response.output_tokens), cache_read_tokens: u64::from(response.cache_read_input_tokens), cache_write_tokens: u64::from(response.cache_creation_input_tokens), - cost_usd: 0.0, + cost_usd: cost_usd_from( + provider, + response.input_tokens, + response.output_tokens, + response.cache_read_input_tokens, + response.cache_creation_input_tokens, + ), }, }); } @@ -171,7 +229,13 @@ impl LlmBackend for LlmBridgeAdapter { output_tokens: u64::from(response.output_tokens), cache_read_tokens: u64::from(response.cache_read_input_tokens), cache_write_tokens: u64::from(response.cache_creation_input_tokens), - cost_usd: 0.0, // TODO: populate from provider cost data when available + cost_usd: cost_usd_from( + provider, + response.input_tokens, + response.output_tokens, + response.cache_read_input_tokens, + response.cache_creation_input_tokens, + ), }, }) } @@ -1229,4 +1293,303 @@ And also check the token price:\n\ other => panic!("Expected ActionCalls, got: {other:?}"), } } + + // ── Caller-level cost-tracking test ────────────────────── + // + // Per testing rules: "Test Through the Caller, Not Just the Helper". + // model_cost() returning the right Decimal is necessary but not + // sufficient — the gap that motivated this test was that + // LlmBridgeAdapter hardcoded `cost_usd: 0.0` and never consulted the + // provider's calculate_cost(), so Thread::total_cost_usd never + // accumulated and `max_budget_usd` gates were inert. This test drives + // the adapter end-to-end with a provider that has known per-token + // pricing and asserts the populated cost flows out via TokenUsage. + + /// Provider with deterministic pricing — Anthropic Sonnet rates + /// (input $3/MTok, output $15/MTok), expressed per token. + struct PricedProvider; + + #[async_trait] + impl LlmProvider for PricedProvider { + fn model_name(&self) -> &str { + "priced-mock" + } + fn cost_per_token(&self) -> (Decimal, Decimal) { + ( + rust_decimal_macros::dec!(0.000003), + rust_decimal_macros::dec!(0.000015), + ) + } + async fn complete( + &self, + _req: crate::llm::CompletionRequest, + ) -> Result { + Ok(crate::llm::CompletionResponse { + content: "hello".to_string(), + input_tokens: 1000, + output_tokens: 500, + finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }) + } + async fn complete_with_tools( + &self, + _req: ToolCompletionRequest, + ) -> Result { + Ok(ToolCompletionResponse { + content: Some("hello".to_string()), + tool_calls: Vec::new(), + input_tokens: 1000, + output_tokens: 500, + finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }) + } + } + + /// Expected cost: 1000 * $0.000003 + 500 * $0.000015 = $0.0105 + const EXPECTED_COST_USD: f64 = 0.0105; + + #[tokio::test] + async fn complete_no_tools_populates_cost_usd_through_adapter() { + let provider: Arc = Arc::new(PricedProvider); + let adapter = LlmBridgeAdapter::new(provider, None); + + let output = adapter + .complete( + &[ThreadMessage::user("hi")], + &[], // no actions => no-tools path + &LlmCallConfig::default(), + ) + .await + .unwrap(); + + assert!( + (output.usage.cost_usd - EXPECTED_COST_USD).abs() < 1e-9, + "expected cost_usd ≈ {EXPECTED_COST_USD}, got {}", + output.usage.cost_usd + ); + } + + #[tokio::test] + async fn complete_with_tools_populates_cost_usd_through_adapter() { + let provider: Arc = Arc::new(PricedProvider); + let adapter = LlmBridgeAdapter::new(provider, None); + + let output = adapter + .complete( + &[ThreadMessage::user("hi")], + &[test_action("noop")], // forces with-tools path + &LlmCallConfig::default(), + ) + .await + .unwrap(); + + assert!( + (output.usage.cost_usd - EXPECTED_COST_USD).abs() < 1e-9, + "expected cost_usd ≈ {EXPECTED_COST_USD}, got {}", + output.usage.cost_usd + ); + } + + #[tokio::test] + async fn complete_routes_subcalls_through_cheap_provider_for_cost() { + // Sub-calls (depth > 0) must be priced with the cheap provider, not + // the primary. Otherwise nested CodeAct calls inflate the parent + // thread's cost by the wrong rate and `max_budget_usd` gates fire + // against the wrong total. + struct ZeroProvider; + #[async_trait] + impl LlmProvider for ZeroProvider { + fn model_name(&self) -> &str { + "zero-mock" + } + fn cost_per_token(&self) -> (Decimal, Decimal) { + (Decimal::ZERO, Decimal::ZERO) + } + async fn complete( + &self, + _req: crate::llm::CompletionRequest, + ) -> Result { + Ok(crate::llm::CompletionResponse { + content: "ok".into(), + input_tokens: 1000, + output_tokens: 500, + finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }) + } + async fn complete_with_tools( + &self, + _req: ToolCompletionRequest, + ) -> Result { + unreachable!() + } + } + + let primary: Arc = Arc::new(PricedProvider); + let cheap: Arc = Arc::new(ZeroProvider); + let adapter = LlmBridgeAdapter::new(primary, Some(cheap)); + + let output = adapter + .complete( + &[ThreadMessage::user("hi")], + &[], + &LlmCallConfig { + depth: 1, + ..LlmCallConfig::default() + }, + ) + .await + .unwrap(); + + assert_eq!( + output.usage.cost_usd, 0.0, + "depth>0 must use cheap provider's pricing (zero), not primary's" + ); + } + + /// Subscription-billed providers (e.g. OpenAI Codex via ChatGPT OAuth) + /// report `(Decimal::ZERO, Decimal::ZERO)` per token and must + /// round-trip through `cost_usd_from` as a clean `0.0` — not panic, + /// not NaN. Exercises the fallback `.unwrap_or(0.0)` on a case that + /// matters in production. + #[tokio::test] + async fn complete_with_subscription_billed_provider_yields_zero_cost() { + struct SubscriptionProvider; + #[async_trait] + impl LlmProvider for SubscriptionProvider { + fn model_name(&self) -> &str { + "subscription-mock" + } + fn cost_per_token(&self) -> (Decimal, Decimal) { + (Decimal::ZERO, Decimal::ZERO) + } + async fn complete( + &self, + _req: crate::llm::CompletionRequest, + ) -> Result { + Ok(crate::llm::CompletionResponse { + content: "ok".into(), + input_tokens: 10_000, + output_tokens: 5_000, + finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }) + } + async fn complete_with_tools( + &self, + _req: ToolCompletionRequest, + ) -> Result { + unreachable!() + } + } + + let provider: Arc = Arc::new(SubscriptionProvider); + let adapter = LlmBridgeAdapter::new(provider, None); + + let output = adapter + .complete(&[ThreadMessage::user("hi")], &[], &LlmCallConfig::default()) + .await + .unwrap(); + + assert_eq!( + output.usage.cost_usd, 0.0, + "zero cost_per_token must produce exactly 0.0 cost_usd" + ); + assert!( + output.usage.cost_usd.is_finite(), + "cost_usd must be finite, never NaN/Inf" + ); + } + + /// Providers that expose prompt caching (Anthropic 5m TTL: 10× read + /// discount, 1.25× write multiplier) must see cost computed with the + /// three-bucket formula, not the flat `input × rate` approximation. + /// Pins the fix for the Copilot/Gemini review comment on #2660 — + /// before the fix, cost_usd undercounted cache-writes and + /// over-counted cache-reads, leaving `max_budget_usd` gates inert + /// against heavy-cache workloads. + #[tokio::test] + async fn complete_prices_cache_tokens_with_discount_and_multiplier() { + /// Anthropic Sonnet 5m-TTL rates: $3/MTok input, $15/MTok output, + /// read discount 10, write multiplier 1.25. + struct AnthropicCachingProvider; + #[async_trait] + impl LlmProvider for AnthropicCachingProvider { + fn model_name(&self) -> &str { + "anthropic-caching-mock" + } + fn cost_per_token(&self) -> (Decimal, Decimal) { + ( + rust_decimal_macros::dec!(0.000003), + rust_decimal_macros::dec!(0.000015), + ) + } + fn cache_read_discount(&self) -> Decimal { + rust_decimal_macros::dec!(10) + } + fn cache_write_multiplier(&self) -> Decimal { + rust_decimal_macros::dec!(1.25) + } + async fn complete( + &self, + _req: crate::llm::CompletionRequest, + ) -> Result { + // Total input = 10_000; 2_000 cache-read, 1_000 cache-write, + // 7_000 uncached. Output = 500. + Ok(crate::llm::CompletionResponse { + content: "ok".into(), + input_tokens: 10_000, + output_tokens: 500, + finish_reason: crate::llm::FinishReason::Stop, + cache_read_input_tokens: 2_000, + cache_creation_input_tokens: 1_000, + }) + } + async fn complete_with_tools( + &self, + _req: ToolCompletionRequest, + ) -> Result { + unreachable!() + } + } + + let provider: Arc = Arc::new(AnthropicCachingProvider); + let adapter = LlmBridgeAdapter::new(provider, None); + + let output = adapter + .complete(&[ThreadMessage::user("hi")], &[], &LlmCallConfig::default()) + .await + .unwrap(); + + // uncached = 10_000 - 2_000 - 1_000 = 7_000 + // uncached_cost = 7_000 * 0.000003 = 0.021 + // cache_read_cost = 2_000 * 0.000003 / 10 = 0.0006 + // cache_write_cost = 1_000 * 0.000003 * 1.25 = 0.00375 + // output_cost = 500 * 0.000015 = 0.0075 + // total = 0.021 + 0.0006 + 0.00375 + 0.0075 = 0.03285 + let expected = 0.032_85_f64; + assert!( + (output.usage.cost_usd - expected).abs() < 1e-9, + "expected cost_usd ≈ {expected}, got {}", + output.usage.cost_usd + ); + + // The naive `(input+output) × rate` approximation the old helper + // computed would have been: 10_000 * 0.000003 + 500 * 0.000015 + // = 0.0375 — i.e. ~14% over-counted vs. the correct 0.03285. + // Pin that we are NOT computing that value. + let naive = 10_000.0 * 0.000_003 + 500.0 * 0.000_015; + assert!( + (output.usage.cost_usd - naive).abs() > 1e-6, + "cost_usd {} must not match the pre-fix naive formula {}", + output.usage.cost_usd, + naive + ); + } } diff --git a/src/bridge/router.rs b/src/bridge/router.rs index 6a4afb8a31..4e220554dd 100644 --- a/src/bridge/router.rs +++ b/src/bridge/router.rs @@ -1190,11 +1190,20 @@ pub async fn init_engine(agent: &Agent) -> Result<(), Error> { let store_dyn: Arc = store.clone(); + // Share the registry with the effect adapter so its `available_actions` + // can advertise engine-native capability actions (missions) to the LLM. + // Without this, mission tools have active leases but never appear in + // the tools list sent with each LLM call. + let capabilities = Arc::new(capabilities); + effect_adapter + .set_capability_registry(Arc::clone(&capabilities)) + .await; + let thread_manager = Arc::new(ThreadManager::new( llm_adapter, effect_adapter.clone(), store_dyn.clone(), - Arc::new(capabilities), + capabilities, leases, policy, )); diff --git a/src/bridge/store_adapter.rs b/src/bridge/store_adapter.rs index db1aff33d2..e70170bcc0 100644 --- a/src/bridge/store_adapter.rs +++ b/src/bridge/store_adapter.rs @@ -1249,6 +1249,10 @@ struct ThreadArchiveSummary { total_tokens: u64, #[serde(default)] outcome_preview: String, + // `#[serde(default)]` lets summaries written before this field existed + // continue to deserialize as zero rather than failing. + #[serde(default)] + total_cost_usd: f64, } fn compact_thread_summary(thread: &Thread) -> ThreadArchiveSummary { @@ -1270,6 +1274,7 @@ fn compact_thread_summary(thread: &Thread) -> ThreadArchiveSummary { step_count: thread.step_count, total_tokens: thread.total_tokens_used, outcome_preview: outcome, + total_cost_usd: thread.total_cost_usd, } } @@ -1309,7 +1314,7 @@ fn thread_from_archive(summary: &ThreadArchiveSummary) -> Option { completed_at, step_count: summary.step_count, total_tokens_used: summary.total_tokens, - total_cost_usd: 0.0, + total_cost_usd: summary.total_cost_usd, }) } @@ -2360,6 +2365,75 @@ mod tests { assert_eq!(hash.len(), 64); assert!(hash.chars().all(|c| c.is_ascii_hexdigit())); } + + // ── ThreadArchiveSummary serialization round-trip ────────── + // + // Regression: `thread_from_archive` previously hardcoded + // `total_cost_usd: 0.0`, silently losing accumulated cost when an + // archived thread got rehydrated for a mission detail page. Pin both + // a live cost round-trip (serialize → JSON → deserialize → + // reconstruct) and legacy compatibility (archive files written before + // this field existed still deserialize, falling back to 0.0). + + fn archive_thread_fixture(cost: f64) -> ironclaw_engine::Thread { + let mut thread = ironclaw_engine::Thread::new( + "archive-round-trip", + ironclaw_engine::ThreadType::Mission, + ironclaw_engine::ProjectId::new(), + "u1", + ironclaw_engine::ThreadConfig::default(), + ); + thread.step_count = 3; + thread.total_tokens_used = 1500; + thread.total_cost_usd = cost; + thread.completed_at = Some(thread.created_at); + thread.state = ironclaw_engine::ThreadState::Completed; + thread + } + + #[test] + fn archive_summary_preserves_total_cost_usd_through_round_trip() { + let thread = archive_thread_fixture(0.0105); + let summary = compact_thread_summary(&thread); + let json = serde_json::to_string(&summary).expect("serialize"); + let restored: ThreadArchiveSummary = serde_json::from_str(&json).expect("deserialize"); + let rehydrated = + thread_from_archive(&restored).expect("thread_from_archive should succeed"); + let delta = (rehydrated.total_cost_usd - 0.0105_f64).abs(); + assert!( + delta < 1e-12, + "total_cost_usd must round-trip: expected 0.0105, got {}", + rehydrated.total_cost_usd + ); + } + + #[test] + fn archive_summary_handles_legacy_json_without_total_cost_usd_field() { + // Craft JSON as it would have been written before this PR: no + // `total_cost_usd` key. `#[serde(default)]` must accept it. + let legacy = serde_json::json!({ + "thread_id": uuid::Uuid::new_v4().to_string(), + "goal": "legacy", + "state": "Completed", + "created_at": chrono::Utc::now().to_rfc3339(), + "completed_at": chrono::Utc::now().to_rfc3339(), + "step_count": 2, + "total_tokens": 900, + // total_cost_usd deliberately omitted + }) + .to_string(); + + let restored: ThreadArchiveSummary = + serde_json::from_str(&legacy).expect("legacy summary must still deserialize"); + assert_eq!( + restored.total_cost_usd, 0.0, + "missing field should default to 0.0" + ); + + let rehydrated = + thread_from_archive(&restored).expect("thread_from_archive should succeed"); + assert_eq!(rehydrated.total_cost_usd, 0.0); + } } #[cfg(all(test, feature = "libsql"))]