mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 07:20:19 +00:00
fix: staging CI review issues (batch 1) (#883)
* fix: address staging-ci-review issues (batch 1) - #811: Fix unreachable error handling in worker — restructure .await? to explicit match on nested Result so token budget errors are properly logged and marked as failed - #813: Combine metadata + token budget into single update_context() call to prevent concurrent worker observing partial state - #814: Persist max_tokens and total_tokens_used to both PostgreSQL and libSQL backends — add V12 migration, update save_job/get_job - #815: Cap user-supplied max_tokens at configured max_tokens_per_job to prevent budget bypass via metadata injection - #869: Release locks before async I/O in webhook handler (http.rs) and SIGHUP handler (main.rs) to prevent blocking concurrent requests Fixes: #811, #813, #814, #815, #869 Co-Authored-By: Claude Opus 4.6 <[email protected]> * fix: address PR #883 review feedback - Fix min(user_val, 0) bug: guard for unlimited config (max_tokens_per_job == 0) - Remove duplicate columns from libSQL base SCHEMA (v12 migration is sole source) - Use get_i64() helper for consistency in libsql/jobs.rs - Add regression tests for scheduler token budget capping Co-Authored-By: Claude Opus 4.6 <[email protected]> --------- Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Henry Park
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
1f5b582c5f
commit
873322f2fb
@@ -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;
|
||||
+152
-8
@@ -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<CompletionResponse, LlmError> {
|
||||
Err(LlmError::RequestFailed {
|
||||
provider: "stub".into(),
|
||||
reason: "not implemented".into(),
|
||||
})
|
||||
}
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_req: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
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<dyn LlmProvider> = 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() {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+14
-8
@@ -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()),
|
||||
|
||||
+27
-15
@@ -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.
|
||||
///
|
||||
|
||||
+12
-5
@@ -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<i64>>("max_tokens").unwrap_or(0) as u64,
|
||||
total_tokens_used: row.get::<_, Option<i64>>("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(
|
||||
|
||||
+6
-4
@@ -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(()) => {
|
||||
|
||||
Reference in New Issue
Block a user