diff --git a/.env.example b/.env.example index 1200400d..5c21e995 100644 --- a/.env.example +++ b/.env.example @@ -115,6 +115,8 @@ AGENT_NAME=ironclaw AGENT_MAX_PARALLEL_JOBS=5 AGENT_JOB_TIMEOUT_SECS=3600 AGENT_STUCK_THRESHOLD_SECS=300 +# Maximum tokens per job (0 = unlimited, also settable via settings.json agent.max_tokens_per_job) +# AGENT_MAX_TOKENS_PER_JOB=0 # Enable planning phase before tool execution (default: true) AGENT_USE_PLANNING=true diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 6118ec7d..99feed9d 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -1205,6 +1205,7 @@ mod tests { max_tool_iterations: 50, auto_approve_tools: false, default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, }, deps, Arc::new(ChannelManager::new()), @@ -2043,6 +2044,7 @@ mod tests { max_tool_iterations, auto_approve_tools: true, default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, }, deps, Arc::new(ChannelManager::new()), @@ -2159,6 +2161,7 @@ mod tests { max_tool_iterations: max_iter, auto_approve_tools: true, default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, }, deps, Arc::new(ChannelManager::new()), diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 99386d7f..85f3f6eb 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -160,6 +160,13 @@ 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 + .as_ref() + .and_then(|m| m.get("max_tokens")) + .and_then(|v| v.as_u64()) + .unwrap_or(self.config.max_tokens_per_job); + // Apply metadata if provided if let Some(meta) = metadata { self.context_manager @@ -169,6 +176,15 @@ impl Scheduler { .await?; } + // Set token budget (separate update to avoid overwriting metadata) + if max_tokens > 0 { + self.context_manager + .update_context(job_id, |ctx| { + ctx.max_tokens = max_tokens; + }) + .await?; + } + // Persist to DB before scheduling so the worker's FK references are valid if let Some(ref store) = self.store { let ctx = self.context_manager.get_context(job_id).await?; diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 3604cea9..19bfc8e5 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -417,7 +417,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# iteration += 1; if iteration > max_iterations { - self.mark_stuck("Maximum iterations exceeded").await?; + self.mark_failed("Maximum iterations exceeded: job hit the iteration cap") + .await?; return Ok(()); } @@ -437,7 +438,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# "LLM rate limited during tool selection, backing off" ); if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS { - self.mark_stuck("Persistent rate limiting").await?; + self.mark_failed("Persistent rate limiting: exceeded retry limit") + .await?; return Ok(()); } self.log_event( @@ -467,7 +469,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# "LLM rate limited during respond_with_tools, backing off" ); if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS { - self.mark_stuck("Persistent rate limiting").await?; + self.mark_failed("Persistent rate limiting: exceeded retry limit") + .await?; return Ok(()); } self.log_event( @@ -483,6 +486,20 @@ Report when the job is complete or if you encounter issues you cannot resolve."# Err(e) => return Err(e.into()), }; + // Track token usage from LLM call against the job budget. + // NOTE: select_tools() also makes LLM calls but doesn't expose + // TokenUsage; only respond_with_tools() usage is tracked here. + let total_tokens = respond_output.usage.total() as u64; + if total_tokens > 0 + && let Err(msg) = self + .context_manager() + .update_context(self.job_id, |ctx| ctx.add_tokens(total_tokens)) + .await? + { + self.mark_failed(&msg).await?; + return Ok(()); + } + match respond_output.result { RespondResult::Text(response) => { // Check for explicit completion phrases. Use word-boundary @@ -1762,4 +1779,84 @@ mod tests { "Always tool should be allowed with permission" ); } + + #[tokio::test] + async fn test_token_budget_exceeded_fails_job() { + let worker = make_worker(vec![]).await; + + // Transition to InProgress (required for mark_failed) + worker + .context_manager() + .update_context(worker.job_id, |ctx| { + ctx.transition_to(JobState::InProgress, None) + }) + .await + .unwrap() + .unwrap(); + + // Set a token budget + worker + .context_manager() + .update_context(worker.job_id, |ctx| { + ctx.max_tokens = 100; + }) + .await + .unwrap(); + + // Simulate adding tokens that exceed the budget + let budget_result = worker + .context_manager() + .update_context(worker.job_id, |ctx| ctx.add_tokens(200)) + .await + .unwrap(); + + assert!( + budget_result.is_err(), + "Should return error when token budget exceeded" + ); + + // Verify that mark_failed transitions job to Failed + worker + .mark_failed(&budget_result.unwrap_err()) + .await + .unwrap(); + let ctx = worker + .context_manager() + .get_context(worker.job_id) + .await + .unwrap(); + assert_eq!(ctx.state, JobState::Failed); + } + + #[tokio::test] + async fn test_iteration_cap_marks_failed_not_stuck() { + let worker = make_worker(vec![]).await; + + // Transition to InProgress (required for mark_failed) + worker + .context_manager() + .update_context(worker.job_id, |ctx| { + ctx.transition_to(JobState::InProgress, None) + }) + .await + .unwrap() + .unwrap(); + + // Simulate what the execution loop does when max_iterations is exceeded + worker + .mark_failed("Maximum iterations exceeded: job hit the iteration cap") + .await + .unwrap(); + + let ctx = worker + .context_manager() + .get_context(worker.job_id) + .await + .unwrap(); + assert_eq!( + ctx.state, + JobState::Failed, + "Iteration cap should transition to Failed, not Stuck" + ); + } } diff --git a/src/channels/web/handlers/jobs.rs b/src/channels/web/handlers/jobs.rs index 8a127243..5a94e055 100644 --- a/src/channels/web/handlers/jobs.rs +++ b/src/channels/web/handlers/jobs.rs @@ -276,11 +276,25 @@ pub async fn jobs_cancel_handler( }))); } - // Fall back to agent job cancellation via DB status update. + // Fall back to agent job cancellation: stop the worker via the scheduler + // (which updates the in-memory ContextManager AND aborts the task handle), + // then persist the status to the DB as a fallback. if let Some(ref store) = state.store && let Ok(Some(job)) = store.get_job(job_id).await { if job.state.is_active() { + // Try to stop via scheduler (aborts the worker task + updates + // in-memory ContextManager). This is best-effort — the job may + // not be in the scheduler map if it already finished. + if let Some(ref slot) = state.scheduler + && let Some(ref scheduler) = *slot.read().await + { + let _ = scheduler.stop(job_id).await; + } + + // Always persist cancellation to the DB so the state is + // consistent even if the scheduler wasn't available or the + // job wasn't in its in-memory map. store .update_job_status( job_id, diff --git a/src/channels/web/handlers/routines.rs b/src/channels/web/handlers/routines.rs index d7c4f764..8fbcc97b 100644 --- a/src/channels/web/handlers/routines.rs +++ b/src/channels/web/handlers/routines.rs @@ -108,6 +108,7 @@ pub async fn routines_detail_handler( status: format!("{:?}", run.status), result_summary: run.result_summary.clone(), tokens_used: run.tokens_used, + job_id: run.job_id, }) .collect(); @@ -252,6 +253,7 @@ pub async fn routines_runs_handler( status: format!("{:?}", run.status), result_summary: run.result_summary.clone(), tokens_used: run.tokens_used, + job_id: run.job_id, }) .collect(); diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index d6605eee..f454363d 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -2017,6 +2017,7 @@ async fn routines_detail_handler( status: format!("{:?}", run.status), result_summary: run.result_summary.clone(), tokens_used: run.tokens_used, + job_id: run.job_id, }) .collect(); @@ -2169,6 +2170,7 @@ async fn routines_runs_handler( status: format!("{:?}", run.status), result_summary: run.result_summary.clone(), tokens_used: run.tokens_used, + job_id: run.job_id, }) .collect(); diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 4d85c671..b6d0d05a 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -776,6 +776,7 @@ pub struct RoutineRunInfo { pub status: String, pub result_summary: Option, pub tokens_used: Option, + pub job_id: Option, } // --- Settings --- diff --git a/src/config/agent.rs b/src/config/agent.rs index 096c141f..cb09707d 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -29,6 +29,8 @@ pub struct AgentConfig { pub auto_approve_tools: bool, /// Default timezone for new sessions (IANA name, e.g. "America/New_York"). pub default_timezone: String, + /// Maximum tokens per job (0 = unlimited). + pub max_tokens_per_job: u64, } impl AgentConfig { @@ -50,6 +52,7 @@ impl AgentConfig { max_tool_iterations: 10, auto_approve_tools: true, default_timezone: "UTC".to_string(), + max_tokens_per_job: 0, } } @@ -105,6 +108,10 @@ impl AgentConfig { } tz }, + max_tokens_per_job: parse_optional_env( + "AGENT_MAX_TOKENS_PER_JOB", + settings.agent.max_tokens_per_job, + )?, }) } } diff --git a/src/db/libsql/jobs.rs b/src/db/libsql/jobs.rs index d5172360..0750873d 100644 --- a/src/db/libsql/jobs.rs +++ b/src/db/libsql/jobs.rs @@ -28,14 +28,16 @@ impl JobStore for LibSqlBackend { r#" INSERT INTO agent_jobs ( 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) + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18) ON CONFLICT (id) DO UPDATE SET title = excluded.title, description = excluded.description, category = excluded.category, status = excluded.status, + user_id = excluded.user_id, estimated_cost = excluded.estimated_cost, estimated_time_secs = excluded.estimated_time_secs, actual_cost = excluded.actual_cost, @@ -51,6 +53,7 @@ impl JobStore for LibSqlBackend { opt_text(ctx.category.as_deref()), status, "direct", + ctx.user_id.as_str(), opt_text_owned(ctx.budget.map(|d| d.to_string())), opt_text(ctx.budget_token.as_deref()), opt_text_owned(ctx.bid_amount.map(|d| d.to_string())), diff --git a/src/db/libsql/mod.rs b/src/db/libsql/mod.rs index 2845c757..404441e6 100644 --- a/src/db/libsql/mod.rs +++ b/src/db/libsql/mod.rs @@ -482,6 +482,24 @@ mod tests { assert_eq!(timeout, 5000); } + /// Regression test: save_job must persist user_id and get_job must return it. + #[tokio::test] + async fn test_save_job_persists_user_id() { + use crate::context::JobContext; + use crate::db::JobStore; + + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("test_user_id.db"); + let backend = LibSqlBackend::new_local(&db_path).await.unwrap(); + backend.run_migrations().await.unwrap(); + + let ctx = JobContext::with_user("test-user-42", "Test Job", "A test job"); + backend.save_job(&ctx).await.unwrap(); + + let loaded = backend.get_job(ctx.job_id).await.unwrap().unwrap(); + assert_eq!(loaded.user_id, "test-user-42"); + } + #[tokio::test] async fn test_concurrent_writes_succeed() { // Use a temp file so connections share state (in-memory DBs are connection-local) diff --git a/src/history/store.rs b/src/history/store.rs index f0b0b144..1153f3e4 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -149,14 +149,16 @@ impl Store { r#" INSERT INTO agent_jobs ( 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) + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18) ON CONFLICT (id) DO UPDATE SET title = EXCLUDED.title, description = EXCLUDED.description, category = EXCLUDED.category, status = EXCLUDED.status, + user_id = EXCLUDED.user_id, estimated_cost = EXCLUDED.estimated_cost, estimated_time_secs = EXCLUDED.estimated_time_secs, actual_cost = EXCLUDED.actual_cost, @@ -172,6 +174,7 @@ impl Store { &ctx.category, &status, &"direct", // source + &ctx.user_id, &ctx.budget, &ctx.budget_token, &ctx.bid_amount, @@ -2133,4 +2136,36 @@ mod tests { assert_eq!(summary.channel, ch); } } + + /// Regression test: save_job must persist user_id and get_job must return it. + /// Requires a running PostgreSQL instance (integration tier). + #[cfg(feature = "postgres")] + #[tokio::test] + #[ignore] + async fn test_save_job_persists_user_id() { + use crate::config::Config; + use crate::context::JobContext; + + let _ = dotenvy::dotenv(); + let config = Config::from_env().await.expect("Failed to load config"); + let store = Store::new(&config.database) + .await + .expect("Failed to connect to database"); + store + .run_migrations() + .await + .expect("Failed to run migrations"); + + let ctx = JobContext::with_user("test-user-42", "PG user_id test", "regression test"); + store.save_job(&ctx).await.unwrap(); + + let loaded = store.get_job(ctx.job_id).await.unwrap().unwrap(); + assert_eq!(loaded.user_id, "test-user-42"); + + // Clean up + let conn = store.conn().await.unwrap(); + conn.execute("DELETE FROM agent_jobs WHERE id = $1", &[&ctx.job_id]) + .await + .unwrap(); + } } diff --git a/src/settings.rs b/src/settings.rs index 836d1d2c..63535aef 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -386,6 +386,10 @@ pub struct AgentSettings { /// Default timezone for new sessions (IANA name, e.g. "America/New_York"). #[serde(default = "default_timezone")] pub default_timezone: String, + + /// Maximum tokens per job (0 = unlimited). + #[serde(default)] + pub max_tokens_per_job: u64, } fn default_agent_name() -> String { @@ -442,6 +446,7 @@ impl Default for AgentSettings { max_tool_iterations: default_max_tool_iterations(), auto_approve_tools: false, default_timezone: default_timezone(), + max_tokens_per_job: 0, } } }