From f18fb5173b4f9b076ebbd6234c4782b4f409c51f Mon Sep 17 00:00:00 2001 From: Henry Park Date: Tue, 3 Mar 2026 06:23:30 -0800 Subject: [PATCH] feat(web): fix jobs UI parity for non-sandbox mode (#491) * feat(web): fix jobs UI parity for non-sandbox mode The web gateway Jobs UI was built primarily for sandbox (Docker) jobs. When running without sandbox (common for NEAR AI hosted envs), multiple features were broken. This change fixes all of them: - Agent jobs now broadcast live SSE events to the web UI (Activity tab) - Agent job restart via scheduler.dispatch_job (not chat message) - Follow-up prompts for agent jobs via WorkerMessage injection - Capability flags (can_restart, can_prompt, job_kind) in job detail API - Rate-limit retry with cap (10 consecutive) and Retry-After header parsing - Plan interruption on user message (breaks out of plan, re-evaluates) - Correct SSE status field in mark_completed/mark_failed/mark_stuck - SseManager preserved across rebuild_state calls Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix formatting in db/mod.rs and nearai_chat.rs Co-Authored-By: Claude Opus 4.6 (1M context) --------- Co-authored-by: Claude Opus 4.6 (1M context) --- src/agent/agent_loop.rs | 10 +- src/agent/dispatcher.rs | 3 + src/agent/scheduler.rs | 32 ++++ src/agent/worker.rs | 229 ++++++++++++++++++++-- src/channels/web/handlers/jobs.rs | 296 ++++++++++++++++++++--------- src/channels/web/mod.rs | 11 +- src/channels/web/server.rs | 2 + src/channels/web/sse.rs | 17 ++ src/channels/web/static/app.js | 32 ++-- src/channels/web/types.rs | 9 + src/channels/web/ws.rs | 1 + src/db/libsql/jobs.rs | 24 +++ src/db/mod.rs | 3 + src/db/postgres.rs | 7 + src/history/store.rs | 15 ++ src/llm/nearai_chat.rs | 25 ++- src/main.rs | 6 +- src/testing.rs | 1 + tests/openai_compat_integration.rs | 2 + tests/ws_gateway_integration.rs | 1 + 20 files changed, 593 insertions(+), 133 deletions(-) diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 38ae30d3..8d3f82bb 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -73,6 +73,8 @@ pub struct AgentDeps { pub hooks: Arc, /// Cost enforcement guardrails (daily budget, hourly rate limits). pub cost_guard: Arc, + /// SSE broadcast sender for live job event streaming to the web gateway. + pub sse_tx: Option>, } /// The main agent that coordinates all components. @@ -111,7 +113,7 @@ impl Agent { let session_manager = session_manager.unwrap_or_else(|| Arc::new(SessionManager::new())); - let scheduler = Arc::new(Scheduler::new( + let mut scheduler = Scheduler::new( config.clone(), context_manager.clone(), deps.llm.clone(), @@ -119,7 +121,11 @@ impl Agent { deps.tools.clone(), deps.store.clone(), deps.hooks.clone(), - )); + ); + if let Some(ref tx) = deps.sse_tx { + scheduler.set_sse_sender(tx.clone()); + } + let scheduler = Arc::new(scheduler); Self { config, diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index daa5da86..bdade0e0 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -982,6 +982,7 @@ mod tests { skills_config: SkillsConfig::default(), hooks: Arc::new(HookRegistry::new()), cost_guard: Arc::new(CostGuard::new(CostGuardConfig::default())), + sse_tx: None, }; Agent::new( @@ -1719,6 +1720,7 @@ mod tests { skills_config: SkillsConfig::default(), hooks: Arc::new(HookRegistry::new()), cost_guard: Arc::new(CostGuard::new(CostGuardConfig::default())), + sse_tx: None, }; Agent::new( @@ -1830,6 +1832,7 @@ mod tests { skills_config: SkillsConfig::default(), hooks: Arc::new(HookRegistry::new()), cost_guard: Arc::new(CostGuard::new(CostGuardConfig::default())), + sse_tx: None, }; Agent::new( diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 5950a8c7..17ffc644 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -10,6 +10,7 @@ use uuid::Uuid; use crate::agent::task::{Task, TaskContext, TaskOutput}; use crate::agent::worker::{Worker, WorkerDeps}; +use crate::channels::web::types::SseEvent; use crate::config::AgentConfig; use crate::context::{ContextManager, JobContext, JobState}; use crate::db::Database; @@ -28,6 +29,8 @@ pub enum WorkerMessage { Stop, /// Check health. Ping, + /// Inject a follow-up user message into the worker's reasoning context. + UserMessage(String), } /// Status of a scheduled job. @@ -51,6 +54,8 @@ pub struct Scheduler { tools: Arc, store: Option>, hooks: Arc, + /// SSE broadcast sender for live job event streaming. + sse_tx: Option>, /// Running jobs (main LLM-driven jobs). jobs: Arc>>, /// Running sub-tasks (tool executions, background tasks). @@ -76,11 +81,17 @@ impl Scheduler { tools, store, hooks, + sse_tx: None, jobs: Arc::new(RwLock::new(HashMap::new())), subtasks: Arc::new(RwLock::new(HashMap::new())), } } + /// Set the SSE broadcast sender for live job event streaming. + pub fn set_sse_sender(&mut self, tx: tokio::sync::broadcast::Sender) { + self.sse_tx = Some(tx); + } + /// Create, persist, and schedule a job in one shot. /// /// This is the preferred entry point for dispatching new jobs. It: @@ -169,6 +180,7 @@ impl Scheduler { hooks: self.hooks.clone(), timeout: self.config.job_timeout, use_planning: self.config.use_planning, + sse_tx: self.sse_tx.clone(), }; let worker = Worker::new(job_id, deps); @@ -500,6 +512,26 @@ impl Scheduler { Ok(()) } + /// Send a follow-up user message to a running job. + /// + /// Returns `Ok(())` if the message was queued, `Err` if the job is not running. + pub async fn send_message(&self, job_id: Uuid, content: String) -> Result<(), JobError> { + // Clone the sender while holding the lock, then release before the + // async send to avoid blocking scheduler writes during backpressure. + let tx = { + let jobs = self.jobs.read().await; + let scheduled = jobs.get(&job_id).ok_or(JobError::NotFound { id: job_id })?; + scheduled.tx.clone() + }; + tx.send(WorkerMessage::UserMessage(content)) + .await + .map_err(|_| JobError::Failed { + id: job_id, + reason: "Worker channel closed".to_string(), + })?; + Ok(()) + } + /// Check if a job is running. pub async fn is_running(&self, job_id: Uuid) -> bool { self.jobs.read().await.contains_key(&job_id) diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 67374b4d..70454e7e 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -9,6 +9,7 @@ use uuid::Uuid; use crate::agent::scheduler::WorkerMessage; use crate::agent::task::TaskOutput; +use crate::channels::web::types::SseEvent; use crate::context::{ContextManager, JobState}; use crate::db::Database; use crate::error::Error; @@ -34,6 +35,8 @@ pub struct WorkerDeps { pub hooks: Arc, pub timeout: Duration, pub use_planning: bool, + /// SSE broadcast sender for live job event streaming to the web gateway. + pub sse_tx: Option>, } /// Worker that executes a single job. @@ -98,18 +101,90 @@ impl Worker { } } - /// Fire-and-forget persistence of a job event. + /// Fire-and-forget persistence of a job event and SSE broadcast. fn log_event(&self, event_type: &str, data: serde_json::Value) { + let job_id = self.job_id; + + // Persist to DB if let Some(store) = self.store() { let store = store.clone(); - let job_id = self.job_id; - let event_type = event_type.to_string(); + let et = event_type.to_string(); + let d = data.clone(); tokio::spawn(async move { - if let Err(e) = store.save_job_event(job_id, &event_type, &data).await { + if let Err(e) = store.save_job_event(job_id, &et, &d).await { tracing::warn!("Failed to persist event for job {}: {}", job_id, e); } }); } + + // Broadcast SSE for live web UI updates + if let Some(ref tx) = self.deps.sse_tx { + let job_id_str = job_id.to_string(); + let event = match event_type { + "message" => Some(SseEvent::JobMessage { + job_id: job_id_str, + role: data + .get("role") + .and_then(|v| v.as_str()) + .unwrap_or("assistant") + .to_string(), + content: data + .get("content") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + }), + "tool_use" => Some(SseEvent::JobToolUse { + job_id: job_id_str, + tool_name: data + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string(), + input: data + .get("input") + .cloned() + .unwrap_or(serde_json::Value::Null), + }), + "tool_result" => Some(SseEvent::JobToolResult { + job_id: job_id_str, + tool_name: data + .get("tool_name") + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string(), + output: data + .get("output") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + }), + "status" => Some(SseEvent::JobStatus { + job_id: job_id_str, + message: data + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(), + }), + "result" => Some(SseEvent::JobResult { + job_id: job_id_str, + status: data + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("completed") + .to_string(), + session_id: data + .get("session_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()), + }), + _ => None, + }; + if let Some(event) = event { + let _ = tx.send(event); + } + } } /// Run the worker until the job is complete or stopped. @@ -123,7 +198,7 @@ impl Worker { tracing::debug!("Worker for job {} stopped before starting", self.job_id); return Ok(()); } - Some(WorkerMessage::Ping) => {} + Some(WorkerMessage::Ping) | Some(WorkerMessage::UserMessage(_)) => {} } // Get job context @@ -219,6 +294,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# .unwrap_or(50) as usize; let max_iterations = max_iterations.min(MAX_WORKER_ITERATIONS); let mut iteration = 0; + const MAX_CONSECUTIVE_RATE_LIMITS: usize = 10; + let mut consecutive_rate_limits = 0usize; // Initial tool definitions for planning (will be refreshed in loop) reason_ctx.available_tools = self.tools().tool_definitions().await; @@ -269,15 +346,27 @@ Report when the job is complete or if you encounter issues you cannot resolve."# None }; - // If we have a plan, execute it + // If we have a plan, execute it. Two exit paths: + // 1. Plan ran to completion → job is Completed or needs continuation + // (check state and only fall through if not terminal) + // 2. Plan was interrupted by UserMessage → fall through to direct loop if let Some(ref plan) = plan { - return self.execute_plan(rx, reasoning, reason_ctx, plan).await; + self.execute_plan(rx, reasoning, reason_ctx, plan).await?; + + // If the plan marked the job terminal, we're done. Only fall + // through to the direct selection loop if the plan was + // interrupted or explicitly left the job in-progress. + if let Ok(ctx) = self.context_manager().get_context(self.job_id).await + && (ctx.state.is_terminal() || ctx.state == JobState::Stuck) + { + return Ok(()); + } } - // Otherwise, use direct tool selection loop + // Direct tool selection loop (also used as fallback after plan interruption) loop { - // Check for stop signal - if let Ok(msg) = rx.try_recv() { + // Check for stop signal and injected user messages + while let Ok(msg) = rx.try_recv() { match msg { WorkerMessage::Stop => { tracing::debug!("Worker for job {} received stop signal", self.job_id); @@ -287,6 +376,20 @@ Report when the job is complete or if you encounter issues you cannot resolve."# tracing::trace!("Worker for job {} received ping", self.job_id); } WorkerMessage::Start => {} + WorkerMessage::UserMessage(content) => { + tracing::info!( + job_id = %self.job_id, + "Worker received follow-up user message" + ); + reason_ctx.messages.push(ChatMessage::user(&content)); + self.log_event( + "message", + serde_json::json!({ + "role": "user", + "content": content, + }), + ); + } } } @@ -307,12 +410,64 @@ Report when the job is complete or if you encounter issues you cannot resolve."# // Refresh tool definitions so newly built tools become visible reason_ctx.available_tools = self.tools().tool_definitions().await; - // Select next tool(s) to use - let selections = reasoning.select_tools(reason_ctx).await?; + // Select next tool(s) to use, with rate-limit retry. + let selections = match reasoning.select_tools(reason_ctx).await { + Ok(s) => s, + Err(crate::error::LlmError::RateLimited { retry_after, .. }) => { + consecutive_rate_limits += 1; + let wait = retry_after.unwrap_or(Duration::from_secs(5)); + tracing::warn!( + job_id = %self.job_id, + wait_secs = wait.as_secs(), + attempt = consecutive_rate_limits, + "LLM rate limited during tool selection, backing off" + ); + if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS { + self.mark_stuck("Persistent rate limiting").await?; + return Ok(()); + } + self.log_event( + "status", + serde_json::json!({ + "message": format!("Rate limited, retrying in {}s ({}/{})...", + wait.as_secs(), consecutive_rate_limits, MAX_CONSECUTIVE_RATE_LIMITS), + }), + ); + tokio::time::sleep(wait).await; + continue; + } + Err(e) => return Err(e.into()), + }; if selections.is_empty() { // No tools from select_tools, ask LLM directly (may still return tool calls) - let respond_output = reasoning.respond_with_tools(reason_ctx).await?; + let respond_output = match reasoning.respond_with_tools(reason_ctx).await { + Ok(o) => o, + Err(crate::error::LlmError::RateLimited { retry_after, .. }) => { + consecutive_rate_limits += 1; + let wait = retry_after.unwrap_or(Duration::from_secs(5)); + tracing::warn!( + job_id = %self.job_id, + wait_secs = wait.as_secs(), + attempt = consecutive_rate_limits, + "LLM rate limited during respond_with_tools, backing off" + ); + if consecutive_rate_limits >= MAX_CONSECUTIVE_RATE_LIMITS { + self.mark_stuck("Persistent rate limiting").await?; + return Ok(()); + } + self.log_event( + "status", + serde_json::json!({ + "message": format!("Rate limited, retrying in {}s ({}/{})...", + wait.as_secs(), consecutive_rate_limits, MAX_CONSECUTIVE_RATE_LIMITS), + }), + ); + tokio::time::sleep(wait).await; + continue; + } + Err(e) => return Err(e.into()), + }; match respond_output.result { RespondResult::Text(response) => { @@ -424,6 +579,11 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } + // Reset rate-limit counter after a successful iteration (all LLM + // calls succeeded). Placed here so alternating success/fail between + // select_tools and respond_with_tools cannot bypass the cap. + consecutive_rate_limits = 0; + // Small delay between iterations tokio::time::sleep(Duration::from_millis(100)).await; } @@ -836,8 +996,8 @@ Report when the job is complete or if you encounter issues you cannot resolve."# plan: &ActionPlan, ) -> Result<(), Error> { for (i, action) in plan.actions.iter().enumerate() { - // Check for stop signal - if let Ok(msg) = rx.try_recv() { + // Check for stop signal and injected user messages + while let Ok(msg) = rx.try_recv() { match msg { WorkerMessage::Stop => { tracing::debug!( @@ -850,6 +1010,29 @@ Report when the job is complete or if you encounter issues you cannot resolve."# tracing::trace!("Worker for job {} received ping", self.job_id); } WorkerMessage::Start => {} + WorkerMessage::UserMessage(content) => { + tracing::info!( + job_id = %self.job_id, + "User message received during plan execution, abandoning plan" + ); + reason_ctx.messages.push(ChatMessage::user(&content)); + self.log_event( + "message", + serde_json::json!({ + "role": "user", + "content": content, + }), + ); + self.log_event( + "status", + serde_json::json!({ + "message": "Plan interrupted by user message, re-evaluating...", + }), + ); + // Return Ok to break out of plan; caller falls through to + // the direct selection loop for LLM re-evaluation. + return Ok(()); + } } } @@ -902,14 +1085,18 @@ Report when the job is complete or if you encounter issues you cannot resolve."# if crate::util::llm_signals_completion(&response) { self.mark_completed().await?; } else { - // Job not complete, could re-plan or fall back to direct selection + // Job not complete — return Ok without marking terminal so the + // caller falls through to the direct selection loop for continuation. tracing::info!( "Job {} plan completed but work remains, falling back to direct selection", self.job_id ); - // Continue with standard execution loop by returning (will be picked up by main loop) - self.mark_stuck("Plan completed but job incomplete - needs re-planning") - .await?; + self.log_event( + "status", + serde_json::json!({ + "message": "Plan completed but job needs more work, continuing...", + }), + ); } Ok(()) @@ -940,6 +1127,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# self.log_event( "result", serde_json::json!({ + "status": "completed", "success": true, "message": "Job completed successfully", }), @@ -965,6 +1153,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# self.log_event( "result", serde_json::json!({ + "status": "failed", "success": false, "message": format!("Execution failed: {}", reason), }), @@ -985,6 +1174,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# self.log_event( "result", serde_json::json!({ + "status": "stuck", "success": false, "message": format!("Job stuck: {}", reason), }), @@ -1103,6 +1293,7 @@ mod tests { hooks: Arc::new(crate::hooks::HookRegistry::new()), timeout: Duration::from_secs(30), use_planning: false, + sse_tx: None, }; Worker::new(job_id, deps) diff --git a/src/channels/web/handlers/jobs.rs b/src/channels/web/handlers/jobs.rs index c32ccb85..8a127243 100644 --- a/src/channels/web/handlers/jobs.rs +++ b/src/channels/web/handlers/jobs.rs @@ -181,6 +181,9 @@ pub async fn jobs_detail_handler( }); } + let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten(); + let is_claude_code = mode.as_deref() == Some("claude_code"); + return Ok(Json(JobDetailResponse { id: job.id, title: job.task.clone(), @@ -193,11 +196,11 @@ pub async fn jobs_detail_handler( elapsed_secs, project_dir: Some(job.project_dir.clone()), browse_url: Some(format!("/projects/{}/", browse_id)), - job_mode: { - let mode = store.get_sandbox_job_mode(job.id).await.ok().flatten(); - mode.filter(|m| m != "worker") - }, + job_mode: mode.filter(|m| m != "worker"), transitions, + can_restart: state.job_manager.is_some(), + can_prompt: is_claude_code && state.prompt_queue.is_some(), + job_kind: Some("sandbox".to_string()), })); } @@ -208,6 +211,12 @@ pub async fn jobs_detail_handler( (end - start).num_seconds().max(0) as u64 }); + // Only show prompt bar for jobs that have a running worker (Pending/InProgress). + // Stuck jobs have no active worker loop, so messages would be silently dropped. + let is_promptable = matches!( + ctx.state, + crate::context::JobState::Pending | crate::context::JobState::InProgress + ); return Ok(Json(JobDetailResponse { id: ctx.job_id, title: ctx.title.clone(), @@ -222,6 +231,9 @@ pub async fn jobs_detail_handler( browse_url: None, job_mode: None, transitions: Vec::new(), + can_restart: state.scheduler.is_some(), + can_prompt: is_promptable && state.scheduler.is_some(), + job_kind: Some("agent".to_string()), })); } @@ -295,108 +307,164 @@ pub async fn jobs_restart_handler( StatusCode::SERVICE_UNAVAILABLE, "Database not available".to_string(), ))?; - let jm = state.job_manager.as_ref().ok_or(( - StatusCode::SERVICE_UNAVAILABLE, - "Sandbox not enabled".to_string(), - ))?; let old_job_id = Uuid::parse_str(&id) .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; - let old_job = store - .get_sandbox_job(old_job_id) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .ok_or((StatusCode::NOT_FOUND, "Job not found".to_string()))?; + // Try sandbox job restart first. + if let Ok(Some(old_job)) = store.get_sandbox_job(old_job_id).await { + if old_job.status != "interrupted" && old_job.status != "failed" { + return Err(( + StatusCode::CONFLICT, + format!("Cannot restart job in state '{}'", old_job.status), + )); + } - if old_job.status != "interrupted" && old_job.status != "failed" { - return Err(( - StatusCode::CONFLICT, - format!("Cannot restart job in state '{}'", old_job.status), - )); + let jm = state.job_manager.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Sandbox not enabled".to_string(), + ))?; + + // Enrich the task with failure context. + let task = if let Some(ref reason) = old_job.failure_reason { + format!( + "Previous attempt failed: {}. Retry: {}", + reason, old_job.task + ) + } else { + old_job.task.clone() + }; + + let new_job_id = Uuid::new_v4(); + let now = chrono::Utc::now(); + + let record = crate::history::SandboxJobRecord { + id: new_job_id, + task: task.clone(), + status: "creating".to_string(), + user_id: old_job.user_id.clone(), + project_dir: old_job.project_dir.clone(), + success: None, + failure_reason: None, + created_at: now, + started_at: None, + completed_at: None, + credential_grants_json: old_job.credential_grants_json.clone(), + }; + store + .save_sandbox_job(&record) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + let mode = match store.get_sandbox_job_mode(old_job_id).await { + Ok(Some(m)) if m == "claude_code" => { + crate::orchestrator::job_manager::JobMode::ClaudeCode + } + _ => crate::orchestrator::job_manager::JobMode::Worker, + }; + + let credential_grants: Vec = + serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| { + tracing::warn!( + job_id = %old_job.id, + "Failed to deserialize credential grants from stored job: {}. \ + Restarted job will have no credentials.", + e + ); + vec![] + }); + + let project_dir = std::path::PathBuf::from(&old_job.project_dir); + let _token = jm + .create_job( + new_job_id, + &task, + Some(project_dir), + mode, + credential_grants, + ) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to create container: {}", e), + ) + })?; + + store + .update_sandbox_job_status(new_job_id, "running", None, None, Some(now), None) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + return Ok(Json(serde_json::json!({ + "status": "restarted", + "old_job_id": old_job_id, + "new_job_id": new_job_id, + }))); } - // Create a new job with the same task and project_dir. - let new_job_id = Uuid::new_v4(); - let now = chrono::Utc::now(); + // Try agent job restart: dispatch a new job via the scheduler. + if let Ok(Some(old_job)) = store.get_job(old_job_id).await { + if old_job.state.is_active() { + return Err(( + StatusCode::CONFLICT, + format!("Cannot restart job in state '{}'", old_job.state), + )); + } - let record = crate::history::SandboxJobRecord { - id: new_job_id, - task: old_job.task.clone(), - status: "creating".to_string(), - user_id: old_job.user_id.clone(), - project_dir: old_job.project_dir.clone(), - success: None, - failure_reason: None, - created_at: now, - started_at: None, - completed_at: None, - credential_grants_json: old_job.credential_grants_json.clone(), - }; - store - .save_sandbox_job(&record) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + let slot = state.scheduler.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Scheduler not available".to_string(), + ))?; + let scheduler_guard = slot.read().await; + let scheduler = scheduler_guard.as_ref().ok_or(( + StatusCode::SERVICE_UNAVAILABLE, + "Agent not started yet".to_string(), + ))?; - // Look up the original job's mode so the restart uses the same mode. - let mode = match store.get_sandbox_job_mode(old_job_id).await { - Ok(Some(m)) if m == "claude_code" => crate::orchestrator::job_manager::JobMode::ClaudeCode, - _ => crate::orchestrator::job_manager::JobMode::Worker, - }; + // Look up failure reason (O(1) point lookup). + let failure_reason = store + .get_agent_job_failure_reason(old_job_id) + .await + .ok() + .flatten() + .unwrap_or_default(); - // Restore credential grants from the original job so the restarted container - // has access to the same secrets. - let credential_grants: Vec = - serde_json::from_str(&old_job.credential_grants_json).unwrap_or_else(|e| { - tracing::warn!( - job_id = %old_job.id, - "Failed to deserialize credential grants from stored job: {}. \ - Restarted job will have no credentials.", - e - ); - vec![] - }); - - let project_dir = std::path::PathBuf::from(&old_job.project_dir); - let _token = jm - .create_job( - new_job_id, - &old_job.task, - Some(project_dir), - mode, - credential_grants, - ) - .await - .map_err(|e| { - ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Failed to create container: {}", e), + let title = if !failure_reason.is_empty() { + format!( + "Previous attempt failed: {}. Retry: {}", + failure_reason, old_job.title ) - })?; + } else { + old_job.title.clone() + }; - store - .update_sandbox_job_status(new_job_id, "running", None, None, Some(now), None) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + let new_job_id = scheduler + .dispatch_job(&old_job.user_id, &title, &old_job.description, None) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; - Ok(Json(serde_json::json!({ - "status": "restarted", - "old_job_id": old_job_id, - "new_job_id": new_job_id, - }))) + return Ok(Json(serde_json::json!({ + "status": "restarted", + "old_job_id": old_job_id, + "new_job_id": new_job_id, + }))); + } + + Err((StatusCode::NOT_FOUND, "Job not found".to_string())) } -/// Submit a follow-up prompt to a running Claude Code sandbox job. +/// Submit a follow-up prompt to a running job. +/// +/// Routes to the appropriate backend: +/// - Claude Code sandbox jobs → prompt queue (polled by the bridge) +/// - Agent (non-sandbox) jobs → WorkerMessage injection via scheduler +/// - Worker-mode sandbox jobs → not supported (no mechanism to inject) pub async fn jobs_prompt_handler( State(state): State>, Path(id): Path, Json(body): Json, ) -> Result, (StatusCode, String)> { - let prompt_queue = state.prompt_queue.as_ref().ok_or(( - StatusCode::NOT_IMPLEMENTED, - "Claude Code not configured".to_string(), - ))?; - let job_id: uuid::Uuid = id .parse() .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid job ID".to_string()))?; @@ -412,17 +480,57 @@ pub async fn jobs_prompt_handler( let done = body.get("done").and_then(|v| v.as_bool()).unwrap_or(false); - let prompt = crate::orchestrator::api::PendingPrompt { content, done }; - + // Try sandbox job path: check if we have a sandbox record for this ID. + if let Some(ref s) = state.store + && let Ok(Some(_)) = s.get_sandbox_job(job_id).await { - let mut queue = prompt_queue.lock().await; - queue.entry(job_id).or_default().push_back(prompt); + // It's a sandbox job. Check if Claude Code mode. + let mode = s.get_sandbox_job_mode(job_id).await.ok().flatten(); + if mode.as_deref() == Some("claude_code") { + let prompt_queue = state.prompt_queue.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Claude Code not configured".to_string(), + ))?; + let prompt = crate::orchestrator::api::PendingPrompt { content, done }; + { + let mut queue = prompt_queue.lock().await; + queue.entry(job_id).or_default().push_back(prompt); + } + return Ok(Json(serde_json::json!({ + "status": "queued", + "job_id": job_id.to_string(), + }))); + } else { + return Err(( + StatusCode::NOT_IMPLEMENTED, + "Follow-up prompts are not supported for worker-mode sandbox jobs".to_string(), + )); + } } - Ok(Json(serde_json::json!({ - "status": "queued", - "job_id": job_id.to_string(), - }))) + // Try agent job path: send via scheduler. + let slot = state.scheduler.as_ref().ok_or(( + StatusCode::NOT_IMPLEMENTED, + "Agent job prompts require the scheduler to be configured".to_string(), + ))?; + let scheduler_guard = slot.read().await; + if let Some(ref scheduler) = *scheduler_guard + && scheduler.is_running(job_id).await + { + scheduler + .send_message(job_id, content) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + return Ok(Json(serde_json::json!({ + "status": "sent", + "job_id": job_id.to_string(), + }))); + } + + Err(( + StatusCode::NOT_FOUND, + "Job not found or not running".to_string(), + )) } /// Load persisted job events for a job (for history replay on page open). diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index 68165924..2fbb4dca 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -84,6 +84,7 @@ impl GatewayChannel { store: None, job_manager: None, prompt_queue: None, + scheduler: None, user_id: config.user_id.clone(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: Some(Arc::new(ws::WsConnectionTracker::new())), @@ -107,7 +108,8 @@ impl GatewayChannel { fn rebuild_state(&mut self, mutate: impl FnOnce(&mut GatewayState)) { let mut new_state = GatewayState { msg_tx: tokio::sync::RwLock::new(None), - sse: SseManager::new(), + // Preserve the existing broadcast channel so sender handles remain valid. + sse: SseManager::from_sender(self.state.sse.sender()), workspace: self.state.workspace.clone(), session_manager: self.state.session_manager.clone(), log_broadcaster: self.state.log_broadcaster.clone(), @@ -117,6 +119,7 @@ impl GatewayChannel { store: self.state.store.clone(), job_manager: self.state.job_manager.clone(), prompt_queue: self.state.prompt_queue.clone(), + scheduler: self.state.scheduler.clone(), user_id: self.state.user_id.clone(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: self.state.ws_tracker.clone(), @@ -196,6 +199,12 @@ impl GatewayChannel { self } + /// Inject the scheduler for sending follow-up messages to agent jobs. + pub fn with_scheduler(mut self, slot: crate::tools::builtin::SchedulerSlot) -> Self { + self.rebuild_state(|s| s.scheduler = Some(slot)); + self + } + /// Inject the skill registry for skill management API. pub fn with_skill_registry(mut self, sr: Arc>) -> Self { self.rebuild_state(|s| s.skill_registry = Some(sr)); diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index bbc15052..18b5f473 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -156,6 +156,8 @@ pub struct GatewayState { pub skill_registry: Option>>, /// Skill catalog for searching the ClawHub registry. pub skill_catalog: Option>, + /// Scheduler for sending follow-up messages to running agent jobs. + pub scheduler: Option, /// Rate limiter for chat endpoints (30 messages per 60 seconds). pub chat_rate_limiter: RateLimiter, /// Registry catalog entries for the available extensions API. diff --git a/src/channels/web/sse.rs b/src/channels/web/sse.rs index 0d5cf39a..e1e2b270 100644 --- a/src/channels/web/sse.rs +++ b/src/channels/web/sse.rs @@ -36,6 +36,23 @@ impl SseManager { } } + /// Create an SSE manager that reuses an existing broadcast sender. + /// + /// This preserves the broadcast channel across `rebuild_state` calls so + /// that sender handles captured by other components remain valid. + /// + /// **Important:** The connection counter is reset to zero. This method must + /// only be called before the server starts accepting connections (i.e., + /// during startup wiring). Calling it after connections are established + /// will break connection tracking and allow exceeding `MAX_CONNECTIONS`. + pub fn from_sender(tx: broadcast::Sender) -> Self { + Self { + tx, + connection_count: Arc::new(AtomicU64::new(0)), + max_connections: MAX_CONNECTIONS, + } + } + /// Broadcast an event to all connected clients. pub fn broadcast(&self, event: SseEvent) { // Ignore send errors (no receivers is fine) diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 4fddb20f..7e653408 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -2376,9 +2376,8 @@ function renderJobsList(jobs) { let actionBtns = ''; if (job.state === 'pending' || job.state === 'in_progress') { actionBtns = ''; - } else if (job.state === 'failed' || job.state === 'interrupted') { - actionBtns = ''; } + // Retry is only shown in the detail view where can_restart is available. return '' + '' + shortId + '' @@ -2445,8 +2444,8 @@ function renderJobDetail(job) { + '

' + escapeHtml(job.title) + '

' + '' + escapeHtml(job.state) + ''; - if (job.state === 'failed' || job.state === 'interrupted') { - headerHtml += ''; + if ((job.state === 'failed' || job.state === 'interrupted') && job.can_restart === true) { + headerHtml += ''; } if (job.browse_url) { headerHtml += 'Browse Files'; @@ -2693,7 +2692,7 @@ function renderJobActivity(container, job) { activityCurrentJobId = job ? job.id : null; activityRenderedLiveIndex = 0; - container.innerHTML = '
' + let html = '
' + '' + '' + '
' - + '
' - + '
' - + '' - + '' - + '' - + '
'; + + '
'; + + if (job && job.can_prompt === true) { + html += '
' + + '' + + '' + + '' + + '
'; + } + + container.innerHTML = html; document.getElementById('activity-type-filter').addEventListener('change', applyActivityFilter); @@ -2716,9 +2720,9 @@ function renderJobActivity(container, job) { const sendBtn = document.getElementById('activity-send-btn'); const doneBtn = document.getElementById('activity-done-btn'); - sendBtn.addEventListener('click', () => sendJobPrompt(job.id, false)); - doneBtn.addEventListener('click', () => sendJobPrompt(job.id, true)); - input.addEventListener('keydown', (e) => { + if (sendBtn) sendBtn.addEventListener('click', () => sendJobPrompt(job.id, false)); + if (doneBtn) doneBtn.addEventListener('click', () => sendJobPrompt(job.id, true)); + if (input) input.addEventListener('keydown', (e) => { if (e.key === 'Enter') sendJobPrompt(job.id, false); }); diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 564e9107..a01aed3a 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -332,6 +332,15 @@ pub struct JobDetailResponse { #[serde(skip_serializing_if = "Option::is_none")] pub job_mode: Option, pub transitions: Vec, + /// Whether this job can be restarted from the UI. + #[serde(default)] + pub can_restart: bool, + /// Whether follow-up prompts can be sent to this job. + #[serde(default)] + pub can_prompt: bool, + /// The kind of job: "sandbox" or "agent". + #[serde(skip_serializing_if = "Option::is_none")] + pub job_kind: Option, } // --- Project Files --- diff --git a/src/channels/web/ws.rs b/src/channels/web/ws.rs index 96c6f783..527daf4a 100644 --- a/src/channels/web/ws.rs +++ b/src/channels/web/ws.rs @@ -483,6 +483,7 @@ mod tests { store: None, job_manager: None, prompt_queue: None, + scheduler: None, user_id: "test".to_string(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: Some(Arc::new(WsConnectionTracker::new())), diff --git a/src/db/libsql/jobs.rs b/src/db/libsql/jobs.rs index 78a55b81..933d7f14 100644 --- a/src/db/libsql/jobs.rs +++ b/src/db/libsql/jobs.rs @@ -213,6 +213,30 @@ impl JobStore for LibSqlBackend { Ok(jobs) } + async fn get_agent_job_failure_reason( + &self, + id: Uuid, + ) -> Result, DatabaseError> { + let conn = self.connect().await?; + let mut rows = conn + .query( + "SELECT failure_reason FROM agent_jobs WHERE id = ?1", + [id.to_string()], + ) + .await + .map_err(|e| DatabaseError::Query(e.to_string()))?; + + if let Some(row) = rows + .next() + .await + .map_err(|e| DatabaseError::Query(e.to_string()))? + { + Ok(get_opt_text(&row, 0)) + } else { + Ok(None) + } + } + async fn agent_job_summary(&self) -> Result { let conn = self.connect().await?; let mut rows = conn diff --git a/src/db/mod.rs b/src/db/mod.rs index ee94f3ef..f065753a 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -177,6 +177,9 @@ pub trait JobStore: Send + Sync { async fn get_stuck_jobs(&self) -> Result, DatabaseError>; async fn list_agent_jobs(&self) -> Result, DatabaseError>; async fn agent_job_summary(&self) -> Result; + /// Get the failure reason for a single agent job (O(1) lookup). + async fn get_agent_job_failure_reason(&self, id: Uuid) + -> Result, DatabaseError>; async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError>; async fn get_job_actions(&self, job_id: Uuid) -> Result, DatabaseError>; async fn record_llm_call(&self, record: &LlmCallRecord<'_>) -> Result; diff --git a/src/db/postgres.rs b/src/db/postgres.rs index 27d5e70b..b73a81b4 100644 --- a/src/db/postgres.rs +++ b/src/db/postgres.rs @@ -223,6 +223,13 @@ impl JobStore for PgBackend { self.store.agent_job_summary().await } + async fn get_agent_job_failure_reason( + &self, + id: Uuid, + ) -> Result, DatabaseError> { + self.store.get_agent_job_failure_reason(id).await + } + async fn save_action(&self, job_id: Uuid, action: &ActionRecord) -> Result<(), DatabaseError> { self.store.save_action(job_id, action).await } diff --git a/src/history/store.rs b/src/history/store.rs index 2e5dc0c1..74f4aa9a 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -821,6 +821,21 @@ impl Store { .collect()) } + /// Get the failure reason for a single agent job. + pub async fn get_agent_job_failure_reason( + &self, + id: Uuid, + ) -> Result, DatabaseError> { + let conn = self.conn().await?; + let row = conn + .query_opt( + "SELECT failure_reason FROM agent_jobs WHERE id = $1", + &[&id], + ) + .await?; + Ok(row.and_then(|r| r.get::<_, Option>("failure_reason"))) + } + /// Summary counts for agent (non-sandbox) jobs. pub async fn agent_job_summary(&self) -> Result { let conn = self.conn().await?; diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index cac8d4a8..50895ecd 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -199,6 +199,29 @@ impl NearAiChatProvider { })?; let status = response.status(); + // Extract Retry-After header before consuming the response body. + // Supports both delay-seconds (RFC 7231 §7.1.3) and HTTP-date formats. + let retry_after_header = response + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|v| { + // Try delay-seconds first (most common from API providers) + if let Ok(secs) = v.trim().parse::() { + return Some(std::time::Duration::from_secs(secs)); + } + // Try HTTP-date (e.g. "Mon, 02 Mar 2026 18:00:00 GMT") + if let Ok(dt) = chrono::DateTime::parse_from_rfc2822(v.trim()) { + let now = chrono::Utc::now(); + let delta = dt.signed_duration_since(now); + // Use max(0) so past/present dates yield Duration::ZERO + // rather than None (which would cause an immediate retry). + return Some(std::time::Duration::from_secs( + delta.num_seconds().max(0) as u64 + )); + } + None + }); let response_text = response.text().await.map_err(|e| LlmError::RequestFailed { provider: "nearai_chat".to_string(), reason: format!("Failed to read response body: {}", e), @@ -230,7 +253,7 @@ impl NearAiChatProvider { if status_code == 429 { return Err(LlmError::RateLimited { provider: "nearai_chat".to_string(), - retry_after: None, + retry_after: retry_after_header, }); } diff --git a/src/main.rs b/src/main.rs index 0cc24a3b..a8cb0951 100644 --- a/src/main.rs +++ b/src/main.rs @@ -506,6 +506,7 @@ async fn async_main() -> anyhow::Result<()> { if let Some(ref jm) = container_job_manager { gw = gw.with_job_manager(Arc::clone(jm)); } + gw = gw.with_scheduler(scheduler_slot.clone()); if let Some(ref sr) = components.skill_registry { gw = gw.with_skill_registry(Arc::clone(sr)); } @@ -646,9 +647,9 @@ async fn async_main() -> anyhow::Result<()> { // Wire SSE sender into extension manager for broadcasting status events. if let Some(ref ext_mgr) = components.extension_manager - && let Some(sender) = sse_sender + && let Some(ref sender) = sse_sender { - ext_mgr.set_sse_sender(sender).await; + ext_mgr.set_sse_sender(sender.clone()).await; } let deps = AgentDeps { @@ -664,6 +665,7 @@ async fn async_main() -> anyhow::Result<()> { skills_config: config.skills.clone(), hooks: components.hooks, cost_guard: components.cost_guard, + sse_tx: sse_sender, }; let agent = Agent::new( diff --git a/src/testing.rs b/src/testing.rs index ededfbe4..dd9c8492 100644 --- a/src/testing.rs +++ b/src/testing.rs @@ -293,6 +293,7 @@ impl TestHarnessBuilder { skills_config: SkillsConfig::default(), hooks, cost_guard, + sse_tx: None, }; TestHarness { diff --git a/tests/openai_compat_integration.rs b/tests/openai_compat_integration.rs index f8b8631a..e70f895a 100644 --- a/tests/openai_compat_integration.rs +++ b/tests/openai_compat_integration.rs @@ -191,6 +191,7 @@ async fn start_test_server_with_provider( store: None, job_manager: None, prompt_queue: None, + scheduler: None, user_id: "test-user".to_string(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: Some(Arc::new(WsConnectionTracker::new())), @@ -679,6 +680,7 @@ async fn test_no_llm_provider_returns_503() { store: None, job_manager: None, prompt_queue: None, + scheduler: None, user_id: "test-user".to_string(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: Some(Arc::new(WsConnectionTracker::new())), diff --git a/tests/ws_gateway_integration.rs b/tests/ws_gateway_integration.rs index beb01859..307271d3 100644 --- a/tests/ws_gateway_integration.rs +++ b/tests/ws_gateway_integration.rs @@ -49,6 +49,7 @@ async fn start_test_server() -> ( store: None, job_manager: None, prompt_queue: None, + scheduler: None, user_id: "test-user".to_string(), shutdown_tx: tokio::sync::RwLock::new(None), ws_tracker: Some(Arc::new(WsConnectionTracker::new())),