diff --git a/migrations/V12__job_token_budget.sql b/migrations/V12__job_token_budget.sql new file mode 100644 index 00000000..fbda73e3 --- /dev/null +++ b/migrations/V12__job_token_budget.sql @@ -0,0 +1,7 @@ +-- Add token budget tracking columns to agent_jobs. +-- +-- Tracks max_tokens (configured limit per job) and total_tokens_used (running total) +-- to enforce job-level token budgets and prevent budget bypass via user-supplied metadata. + +ALTER TABLE agent_jobs ADD COLUMN max_tokens BIGINT NOT NULL DEFAULT 0; +ALTER TABLE agent_jobs ADD COLUMN total_tokens_used BIGINT NOT NULL DEFAULT 0; diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 971842e5..5e4bf01a 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -160,24 +160,36 @@ impl Scheduler { .create_job_for_user(user_id, title, description) .await?; - // Apply token budget from config, allowing per-job metadata override. - let max_tokens = metadata + // Apply metadata and token budget in a single atomic update. + // This prevents concurrent workers from observing partial state. + // Cap user-supplied max_tokens at the configured limit (Issue #815). + let user_max_tokens = metadata .as_ref() .and_then(|m| m.get("max_tokens")) - .and_then(|v| v.as_u64()) + .and_then(|v| v.as_u64()); + + let max_tokens = user_max_tokens + .map(|user_val| { + if self.config.max_tokens_per_job == 0 { + // Config is "unlimited": use the user-supplied value directly. + user_val + } else { + std::cmp::min(user_val, self.config.max_tokens_per_job) + } + }) .unwrap_or(self.config.max_tokens_per_job); - // Apply metadata if provided + // Apply both metadata and token budget in one closure (Issue #813: atomic update) if let Some(meta) = metadata { self.context_manager .update_context(job_id, |ctx| { ctx.metadata = meta; + if max_tokens > 0 { + ctx.max_tokens = max_tokens; + } }) .await?; - } - - // Set token budget (separate update to avoid overwriting metadata) - if max_tokens > 0 { + } else if max_tokens > 0 { self.context_manager .update_context(job_id, |ctx| { ctx.max_tokens = max_tokens; @@ -685,8 +697,140 @@ impl Scheduler { mod tests { use super::*; use crate::config::SafetyConfig; + use crate::llm::{ + CompletionRequest, CompletionResponse, LlmError, LlmProvider, ToolCompletionRequest, + ToolCompletionResponse, + }; use crate::safety::SafetyLayer; use crate::tools::{ApprovalRequirement, Tool, ToolError, ToolOutput}; + use rust_decimal_macros::dec; + + /// Minimal LLM provider stub for scheduler tests that don't exercise LLM calls. + struct StubLlm; + + #[async_trait::async_trait] + impl LlmProvider for StubLlm { + fn model_name(&self) -> &str { + "stub" + } + fn cost_per_token(&self) -> (rust_decimal::Decimal, rust_decimal::Decimal) { + (dec!(0), dec!(0)) + } + async fn complete(&self, _req: CompletionRequest) -> Result { + Err(LlmError::RequestFailed { + provider: "stub".into(), + reason: "not implemented".into(), + }) + } + async fn complete_with_tools( + &self, + _req: ToolCompletionRequest, + ) -> Result { + Err(LlmError::RequestFailed { + provider: "stub".into(), + reason: "not implemented".into(), + }) + } + } + + /// Create a Scheduler for token-budget tests. The LLM stub will fail if a + /// worker actually tries to call it, but `dispatch_job` sets the token + /// budget *before* spawning the worker so we can inspect the context + /// immediately after dispatch. + fn make_test_scheduler(max_tokens_per_job: u64) -> Scheduler { + let config = AgentConfig { + name: "test".to_string(), + max_parallel_jobs: 5, + job_timeout: std::time::Duration::from_secs(30), + stuck_threshold: std::time::Duration::from_secs(300), + repair_check_interval: std::time::Duration::from_secs(3600), + max_repair_attempts: 0, + use_planning: false, + session_idle_timeout: std::time::Duration::from_secs(3600), + allow_local_tools: true, + max_cost_per_day_cents: None, + max_actions_per_hour: None, + max_tool_iterations: 10, + auto_approve_tools: true, + default_timezone: "UTC".to_string(), + max_tokens_per_job, + }; + let cm = Arc::new(ContextManager::new(5)); + let llm: Arc = Arc::new(StubLlm); + let safety = Arc::new(SafetyLayer::new(&SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: false, + })); + let tools = Arc::new(ToolRegistry::new()); + let hooks = Arc::new(HookRegistry::default()); + + Scheduler::new(config, cm, llm, safety, tools, None, hooks) + } + + #[tokio::test] + async fn test_dispatch_job_caps_user_max_tokens() { + let sched = make_test_scheduler(1000); + let meta = serde_json::json!({ "max_tokens": 5000 }); + let job_id = sched + .dispatch_job("user1", "test", "desc", Some(meta)) + .await + .unwrap(); + + let ctx = sched.context_manager.get_context(job_id).await.unwrap(); + assert_eq!(ctx.max_tokens, 1000, "should cap at configured limit"); + } + + #[tokio::test] + async fn test_dispatch_job_unlimited_config_preserves_user_tokens() { + let sched = make_test_scheduler(0); // 0 = unlimited + let meta = serde_json::json!({ "max_tokens": 5000 }); + let job_id = sched + .dispatch_job("user1", "test", "desc", Some(meta)) + .await + .unwrap(); + + let ctx = sched.context_manager.get_context(job_id).await.unwrap(); + assert_eq!( + ctx.max_tokens, 5000, + "unlimited config should preserve user value" + ); + } + + #[tokio::test] + async fn test_dispatch_job_no_user_tokens_uses_config() { + let sched = make_test_scheduler(2000); + let job_id = sched + .dispatch_job("user1", "test", "desc", None) + .await + .unwrap(); + + let ctx = sched.context_manager.get_context(job_id).await.unwrap(); + assert_eq!( + ctx.max_tokens, 2000, + "should use config default when no user value" + ); + } + + #[tokio::test] + async fn test_dispatch_job_atomic_metadata_and_tokens() { + let sched = make_test_scheduler(10_000); + let meta = serde_json::json!({ + "max_tokens": 3000, + "custom_key": "custom_value" + }); + let job_id = sched + .dispatch_job("user1", "test", "desc", Some(meta)) + .await + .unwrap(); + + let ctx = sched.context_manager.get_context(job_id).await.unwrap(); + assert_eq!(ctx.max_tokens, 3000, "should use user value within limit"); + assert_eq!( + ctx.metadata.get("custom_key").and_then(|v| v.as_str()), + Some("custom_value"), + "metadata should be set atomically with token budget" + ); + } #[test] fn test_scheduler_creation() { diff --git a/src/channels/http.rs b/src/channels/http.rs index af0fafcf..e40e251b 100644 --- a/src/channels/http.rs +++ b/src/channels/http.rs @@ -395,9 +395,14 @@ async fn process_message( None }; - // Send message to the channel - let tx_guard = state.tx.read().await; - if let Some(tx) = tx_guard.as_ref() { + // Clone sender while holding read lock, then release lock before async send. + // This prevents blocking other webhook handlers during the async I/O. + let tx = { + let guard = state.tx.read().await; + guard.as_ref().cloned() + }; + + if let Some(tx) = tx { if tx.send(msg).await.is_err() { return ( StatusCode::INTERNAL_SERVER_ERROR, @@ -418,7 +423,6 @@ async fn process_message( }), ); } - drop(tx_guard); // Wait for response if requested let response = if let Some(rx) = response_rx { diff --git a/src/db/libsql/jobs.rs b/src/db/libsql/jobs.rs index 0750873d..3db3ab30 100644 --- a/src/db/libsql/jobs.rs +++ b/src/db/libsql/jobs.rs @@ -30,8 +30,9 @@ impl JobStore for LibSqlBackend { id, conversation_id, title, description, category, status, source, user_id, budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18) + actual_cost, repair_attempts, max_tokens, total_tokens_used, + created_at, started_at, completed_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20) ON CONFLICT (id) DO UPDATE SET title = excluded.title, description = excluded.description, @@ -42,6 +43,8 @@ impl JobStore for LibSqlBackend { estimated_time_secs = excluded.estimated_time_secs, actual_cost = excluded.actual_cost, repair_attempts = excluded.repair_attempts, + max_tokens = excluded.max_tokens, + total_tokens_used = excluded.total_tokens_used, started_at = excluded.started_at, completed_at = excluded.completed_at "#, @@ -61,6 +64,8 @@ impl JobStore for LibSqlBackend { estimated_time_secs, ctx.actual_cost.to_string(), ctx.repair_attempts as i64, + ctx.max_tokens as i64, + ctx.total_tokens_used as i64, fmt_ts(&ctx.created_at), fmt_opt_ts(&ctx.started_at), fmt_opt_ts(&ctx.completed_at), @@ -78,7 +83,8 @@ impl JobStore for LibSqlBackend { r#" SELECT id, conversation_id, title, description, category, status, user_id, budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at + actual_cost, repair_attempts, max_tokens, total_tokens_used, + created_at, started_at, completed_at FROM agent_jobs WHERE id = ?1 "#, params![id.to_string()], @@ -111,12 +117,12 @@ impl JobStore for LibSqlBackend { estimated_duration: estimated_time_secs .map(|s| std::time::Duration::from_secs(s as u64)), actual_cost: get_decimal(&row, 12), - total_tokens_used: 0, - max_tokens: 0, + max_tokens: get_i64(&row, 14) as u64, + total_tokens_used: get_i64(&row, 15) as u64, repair_attempts: get_i64(&row, 13) as u32, - created_at: get_ts(&row, 14), - started_at: get_opt_ts(&row, 15), - completed_at: get_opt_ts(&row, 16), + created_at: get_ts(&row, 16), + started_at: get_opt_ts(&row, 17), + completed_at: get_opt_ts(&row, 18), transitions: Vec::new(), metadata: serde_json::Value::Null, extra_env: std::sync::Arc::new(std::collections::HashMap::new()), diff --git a/src/db/libsql_migrations.rs b/src/db/libsql_migrations.rs index 02c4c9b2..fc445b7c 100644 --- a/src/db/libsql_migrations.rs +++ b/src/db/libsql_migrations.rs @@ -583,20 +583,21 @@ INSERT OR IGNORE INTO leak_detection_patterns (id, name, pattern, severity, acti /// /// Each entry is `(version, name, sql)`. Migrations are idempotent: the /// `_migrations` table tracks which versions have been applied. -pub const INCREMENTAL_MIGRATIONS: &[(i64, &str, &str)] = &[( - 9, - "flexible_embedding_dimension", - // Rebuild memory_chunks to remove the fixed F32_BLOB(1536) type - // constraint so any embedding dimension works. Existing embeddings - // are preserved; users only need to re-embed if they change models. - // - // The vector index (libsql_vector_idx) requires a fixed-dimension - // F32_BLOB(N), so we drop it entirely. Vector search falls back to - // brute-force cosine distance which is fast enough for personal - // assistant workspaces. This matches PostgreSQL after its V9 migration. - // - // SQLite cannot ALTER COLUMN types, so we recreate the table. - r#" +pub const INCREMENTAL_MIGRATIONS: &[(i64, &str, &str)] = &[ + ( + 9, + "flexible_embedding_dimension", + // Rebuild memory_chunks to remove the fixed F32_BLOB(1536) type + // constraint so any embedding dimension works. Existing embeddings + // are preserved; users only need to re-embed if they change models. + // + // The vector index (libsql_vector_idx) requires a fixed-dimension + // F32_BLOB(N), so we drop it entirely. Vector search falls back to + // brute-force cosine distance which is fast enough for personal + // assistant workspaces. This matches PostgreSQL after its V9 migration. + // + // SQLite cannot ALTER COLUMN types, so we recreate the table. + r#" -- Drop vector index (requires fixed F32_BLOB(N), incompatible with flexible dimensions) DROP INDEX IF EXISTS idx_memory_chunks_embedding; @@ -644,7 +645,18 @@ CREATE TRIGGER IF NOT EXISTS memory_chunks_fts_update AFTER UPDATE ON memory_chu INSERT INTO memory_chunks_fts(rowid, content) VALUES (new._rowid, new.content); END; "#, -)]; + ), + ( + 12, + "job_token_budget", + // Add token budget tracking columns to agent_jobs. + // SQLite supports ALTER TABLE ADD COLUMN, so no table rebuild needed. + r#" +ALTER TABLE agent_jobs ADD COLUMN max_tokens INTEGER NOT NULL DEFAULT 0; +ALTER TABLE agent_jobs ADD COLUMN total_tokens_used INTEGER NOT NULL DEFAULT 0; +"#, + ), +]; /// Run incremental migrations that haven't been applied yet. /// diff --git a/src/history/store.rs b/src/history/store.rs index f35f31a6..e877cbbf 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -151,8 +151,9 @@ impl Store { id, conversation_id, title, description, category, status, source, user_id, budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18) + actual_cost, repair_attempts, max_tokens, total_tokens_used, + created_at, started_at, completed_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20) ON CONFLICT (id) DO UPDATE SET title = EXCLUDED.title, description = EXCLUDED.description, @@ -163,6 +164,8 @@ impl Store { estimated_time_secs = EXCLUDED.estimated_time_secs, actual_cost = EXCLUDED.actual_cost, repair_attempts = EXCLUDED.repair_attempts, + max_tokens = EXCLUDED.max_tokens, + total_tokens_used = EXCLUDED.total_tokens_used, started_at = EXCLUDED.started_at, completed_at = EXCLUDED.completed_at "#, @@ -182,6 +185,8 @@ impl Store { &estimated_time_secs, &ctx.actual_cost, &(ctx.repair_attempts as i32), + &(ctx.max_tokens as i64), + &(ctx.total_tokens_used as i64), &ctx.created_at, &ctx.started_at, &ctx.completed_at, @@ -201,7 +206,8 @@ impl Store { r#" SELECT id, conversation_id, title, description, category, status, user_id, budget_amount, budget_token, bid_amount, estimated_cost, estimated_time_secs, - actual_cost, repair_attempts, created_at, started_at, completed_at + actual_cost, repair_attempts, max_tokens, total_tokens_used, + created_at, started_at, completed_at FROM agent_jobs WHERE id = $1 "#, &[&id], @@ -237,8 +243,9 @@ impl Store { completed_at: row.get("completed_at"), transitions: Vec::new(), // Not loaded from DB for now metadata: serde_json::Value::Null, - total_tokens_used: 0, - max_tokens: 0, + max_tokens: row.get::<_, Option>("max_tokens").unwrap_or(0) as u64, + total_tokens_used: row.get::<_, Option>("total_tokens_used").unwrap_or(0) + as u64, extra_env: std::sync::Arc::new(std::collections::HashMap::new()), http_interceptor: None, tool_output_stash: std::sync::Arc::new(tokio::sync::RwLock::new( diff --git a/src/main.rs b/src/main.rs index 58f7769e..581e3500 100644 --- a/src/main.rs +++ b/src/main.rs @@ -768,10 +768,10 @@ async fn async_main() -> anyhow::Result<()> { } }; - // Restart listener if addr changed + // Restart listener if addr changed. + // Minimize lock scope: acquire, read old addr, release, then restart. let mut restart_failed = false; if let Some(ref ws_arc) = sighup_webhook_server { - // Read old address while holding lock, then drop immediately let old_addr = { let ws = ws_arc.lock().await; ws.current_addr() @@ -783,8 +783,10 @@ async fn async_main() -> anyhow::Result<()> { old_addr, new_addr ); - // Wait for restart to complete before proceeding with secret update. - // This ensures atomicity: if restart fails, secret is not updated (partial state corruption). + // NOTE: Lock is held across restart_with_addr().await. This is + // acceptable because SIGHUP is infrequent and restart is fast. A full + // fix would require refactoring restart_with_addr to separate state + // mutation from async I/O. let mut ws = ws_arc.lock().await; match ws.restart_with_addr(new_addr).await { Ok(()) => {