1pub mod engine;
2pub mod fragments;
3pub mod ids;
4pub mod journal;
5pub mod request;
6pub mod session;
7
8use std::collections::{HashMap, HashSet};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12use std::time::Duration;
13
14use anyhow::Result;
15use codewhale_agent::ModelRegistry;
16use codewhale_config::{CliRuntimeOverrides, ConfigToml, ProviderKind};
17use codewhale_execpolicy::{
18 AskForApproval, ExecApprovalRequirement, ExecPolicyContext, ExecPolicyDecision,
19 ExecPolicyEngine,
20};
21use codewhale_hooks::{HookDispatcher, HookEvent};
22use codewhale_mcp::{
23 McpManager, McpStartupCompleteEvent, McpStartupStatus as McpManagerStartupStatus,
24};
25use codewhale_protocol::{
26 AppResponse, EventFrame, ExecApprovalRequestEvent, PromptRequest, PromptResponse,
27 ResponseChannel, ReviewDecision, Status, Thread, ThreadForkParams, ThreadGoal,
28 ThreadGoalClearParams, ThreadGoalGetParams, ThreadGoalProgressParams, ThreadGoalSetParams,
29 ThreadGoalStatus, ThreadListParams, ThreadReadParams, ThreadRequest, ThreadResponse,
30 ThreadResumeParams, ThreadSetNameParams, ThreadStatus, ToolPayload, UserInputRequestEvent,
31};
32use codewhale_state::{
33 JobStateRecord, JobStateStatus, SessionSource, StateStore, ThreadGoalRecord,
34 ThreadGoalStatus as PersistedThreadGoalStatus, ThreadListFilters, ThreadMetadata,
35 ThreadStatus as PersistedThreadStatus,
36};
37use codewhale_tools::{ToolCall, ToolRegistry};
38use serde_json::{Value, json};
39use tokio::time;
40use uuid::Uuid;
41
42fn tool_dispatch_timeout() -> Duration {
45 if cfg!(test) {
46 Duration::from_millis(50)
47 } else {
48 Duration::from_secs(300)
49 }
50}
51
52#[derive(Debug, Clone)]
54pub enum InitialHistory {
55 New,
57 Forked(Vec<Value>),
59 Resumed {
61 conversation_id: String,
62 history: Vec<Value>,
63 rollout_path: PathBuf,
64 },
65}
66
67#[derive(Debug, Clone)]
69pub struct NewThread {
70 pub thread: Thread,
72 pub model: String,
74 pub model_provider: String,
76 pub cwd: PathBuf,
78 pub approval_policy: Option<String>,
80 pub sandbox: Option<String>,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum JobStatus {
87 Queued,
89 Running,
91 Paused,
93 Completed,
95 Failed,
97 Cancelled,
99}
100
101impl Status for JobStatus {
102 fn is_terminal(&self) -> bool {
103 matches!(self, Self::Completed | Self::Failed | Self::Cancelled)
104 }
105 fn is_active(&self) -> bool {
106 matches!(self, Self::Queued | Self::Running)
107 }
108 fn is_paused(&self) -> bool {
109 matches!(self, Self::Paused)
110 }
111}
112
113const JOB_DETAIL_SCHEMA_VERSION: u8 = 1;
114const DEFAULT_JOB_MAX_ATTEMPTS: u32 = 3;
115const DEFAULT_JOB_BACKOFF_BASE_MS: u64 = 500;
116const MAX_JOB_HISTORY_ENTRIES: usize = 64;
117
118#[derive(Debug, Clone)]
120pub struct JobRetryMetadata {
121 pub attempt: u32,
123 pub max_attempts: u32,
125 pub backoff_base_ms: u64,
127 pub next_backoff_ms: u64,
129 pub next_retry_at: Option<i64>,
131}
132
133impl Default for JobRetryMetadata {
134 fn default() -> Self {
135 Self {
136 attempt: 0,
137 max_attempts: DEFAULT_JOB_MAX_ATTEMPTS,
138 backoff_base_ms: DEFAULT_JOB_BACKOFF_BASE_MS,
139 next_backoff_ms: 0,
140 next_retry_at: None,
141 }
142 }
143}
144
145#[derive(Debug, Clone)]
147pub struct JobHistoryEntry {
148 pub at: i64,
150 pub phase: String,
152 pub status: JobStatus,
154 pub progress: Option<u8>,
156 pub detail: Option<String>,
158 pub retry: JobRetryMetadata,
160}
161
162#[derive(Debug, Clone)]
163struct PersistedJobDetail {
164 pub status: JobStatus,
165 pub detail: Option<String>,
166 pub retry: JobRetryMetadata,
167 pub history: Vec<JobHistoryEntry>,
168}
169
170#[derive(Debug, Clone)]
172pub struct JobRecord {
173 pub id: String,
175 pub name: String,
177 pub status: JobStatus,
179 pub progress: Option<u8>,
181 pub detail: Option<String>,
183 pub retry: JobRetryMetadata,
185 pub history: Vec<JobHistoryEntry>,
187 pub created_at: i64,
189 pub updated_at: i64,
191}
192
193#[must_use]
201pub fn job_record_to_agent_run(
202 record: &JobRecord,
203) -> codewhale_protocol::agent_run::AgentRunSnapshot {
204 use codewhale_protocol::agent_run::{
205 AgentRunSnapshot, BudgetSummary, RunSource, RunState, TerminalOutcome, TerminalSummary,
206 };
207
208 let (state, terminal) = match record.status {
209 JobStatus::Queued => (RunState::Queued, None),
210 JobStatus::Running => (RunState::Running, None),
211 JobStatus::Paused => (RunState::Paused, None),
212 JobStatus::Completed | JobStatus::Failed | JobStatus::Cancelled => {
213 let outcome = match record.status {
214 JobStatus::Completed => TerminalOutcome::Completed,
215 JobStatus::Failed => TerminalOutcome::Failed,
216 _ => TerminalOutcome::Cancelled,
217 };
218 (
219 RunState::Terminal,
220 Some(TerminalSummary {
221 outcome,
222 ended_at_ms: record.updated_at.checked_mul(1000),
223 detail: None,
224 }),
225 )
226 }
227 };
228
229 AgentRunSnapshot {
230 run_id: record.id.clone(),
231 parent: None,
232 source: RunSource::CoreJob,
233 state,
234 budget: BudgetSummary::default(),
235 terminal,
236 refs: Vec::new(),
237 }
238}
239
240#[derive(Debug, Default)]
242pub struct JobManager {
243 jobs: HashMap<String, JobRecord>,
244}
245
246impl JobManager {
247 fn now_ts() -> i64 {
248 chrono::Utc::now().timestamp()
249 }
250
251 fn deterministic_backoff_ms(retry: &JobRetryMetadata) -> u64 {
252 if retry.attempt == 0 {
253 return 0;
254 }
255 let exponent = retry.attempt.saturating_sub(1).min(20);
256 let multiplier = 1u64.checked_shl(exponent).unwrap_or(u64::MAX);
257 retry.backoff_base_ms.saturating_mul(multiplier)
258 }
259
260 fn clear_retry_schedule(retry: &mut JobRetryMetadata) {
261 retry.next_backoff_ms = 0;
262 retry.next_retry_at = None;
263 }
264
265 fn push_history(job: &mut JobRecord, phase: &str) {
266 job.history.push(JobHistoryEntry {
267 at: job.updated_at,
268 phase: phase.to_string(),
269 status: job.status,
270 progress: job.progress,
271 detail: job.detail.clone(),
272 retry: job.retry.clone(),
273 });
274 if job.history.len() > MAX_JOB_HISTORY_ENTRIES {
275 let to_drain = job.history.len() - MAX_JOB_HISTORY_ENTRIES;
276 job.history.drain(0..to_drain);
277 }
278 }
279
280 fn parse_persisted_detail(raw: Option<&str>) -> Option<PersistedJobDetail> {
281 let raw = raw?;
282 let parsed: Value = serde_json::from_str(raw).ok()?;
283 let status = parsed
284 .get("status")
285 .and_then(Value::as_str)
286 .and_then(job_status_from_str)?;
287 let detail = parsed.get("detail").and_then(json_optional_string);
288 let retry = parse_retry_metadata(parsed.get("retry"));
289 let history = parsed
290 .get("history")
291 .and_then(Value::as_array)
292 .map(|items| {
293 items
294 .iter()
295 .filter_map(parse_history_entry)
296 .collect::<Vec<_>>()
297 })
298 .unwrap_or_default();
299 Some(PersistedJobDetail {
300 status,
301 detail,
302 retry,
303 history,
304 })
305 }
306
307 fn encode_persisted_detail(job: &JobRecord) -> Result<Option<String>> {
308 let encoded = json!({
309 "schema_version": JOB_DETAIL_SCHEMA_VERSION,
310 "status": job_status_to_str(job.status),
311 "detail": job.detail.clone(),
312 "retry": job_retry_to_value(&job.retry),
313 "history": job.history.iter().map(job_history_to_value).collect::<Vec<_>>()
314 })
315 .to_string();
316 Ok(Some(encoded))
317 }
318
319 pub fn enqueue(&mut self, name: impl Into<String>) -> JobRecord {
321 let now = Self::now_ts();
322 let id = format!("job-{}", Uuid::new_v4());
323 let mut job = JobRecord {
324 id: id.clone(),
325 name: name.into(),
326 status: JobStatus::Queued,
327 progress: Some(0),
328 detail: None,
329 retry: JobRetryMetadata::default(),
330 history: Vec::new(),
331 created_at: now,
332 updated_at: now,
333 };
334 Self::push_history(&mut job, "created");
335 self.jobs.insert(id, job.clone());
336 job
337 }
338
339 pub fn set_running(&mut self, id: &str) {
341 if let Some(job) = self.jobs.get_mut(id) {
342 job.status = JobStatus::Running;
343 Self::clear_retry_schedule(&mut job.retry);
344 job.updated_at = Self::now_ts();
345 Self::push_history(job, "running");
346 }
347 }
348
349 pub fn update_progress(&mut self, id: &str, progress: u8, detail: Option<String>) {
351 if let Some(job) = self.jobs.get_mut(id) {
352 job.progress = Some(progress.min(100));
353 job.detail = detail;
354 job.updated_at = Self::now_ts();
355 Self::push_history(job, "progress_updated");
356 }
357 }
358
359 pub fn complete(&mut self, id: &str) {
361 if let Some(job) = self.jobs.get_mut(id) {
362 job.status = JobStatus::Completed;
363 job.progress = Some(100);
364 Self::clear_retry_schedule(&mut job.retry);
365 job.updated_at = Self::now_ts();
366 Self::push_history(job, "completed");
367 }
368 }
369
370 pub fn fail(&mut self, id: &str, detail: impl Into<String>) {
372 if let Some(job) = self.jobs.get_mut(id) {
373 let now = Self::now_ts();
374 job.status = JobStatus::Failed;
375 job.detail = Some(detail.into());
376 if job.retry.attempt < job.retry.max_attempts {
377 job.retry.attempt += 1;
378 job.retry.next_backoff_ms = Self::deterministic_backoff_ms(&job.retry);
379 let delay_secs = ((job.retry.next_backoff_ms.saturating_add(999)) / 1000)
380 .min(i64::MAX as u64) as i64;
381 job.retry.next_retry_at = Some(now.saturating_add(delay_secs));
382 } else {
383 Self::clear_retry_schedule(&mut job.retry);
384 }
385 job.updated_at = now;
386 Self::push_history(job, "failed");
387 }
388 }
389
390 pub fn cancel(&mut self, id: &str) {
392 if let Some(job) = self.jobs.get_mut(id) {
393 job.status = JobStatus::Cancelled;
394 Self::clear_retry_schedule(&mut job.retry);
395 job.updated_at = Self::now_ts();
396 Self::push_history(job, "cancelled");
397 }
398 }
399
400 pub fn pause(&mut self, id: &str, detail: Option<String>) {
402 if let Some(job) = self.jobs.get_mut(id) {
403 job.status = JobStatus::Paused;
404 if detail.is_some() {
405 job.detail = detail;
406 }
407 job.updated_at = Self::now_ts();
408 Self::push_history(job, "paused");
409 }
410 }
411
412 pub fn resume(&mut self, id: &str, detail: Option<String>) {
414 if let Some(job) = self.jobs.get_mut(id) {
415 job.status = JobStatus::Running;
416 if detail.is_some() {
417 job.detail = detail;
418 }
419 Self::clear_retry_schedule(&mut job.retry);
420 job.updated_at = Self::now_ts();
421 Self::push_history(job, "resumed");
422 }
423 }
424
425 pub fn list(&self) -> Vec<JobRecord> {
427 let mut out = self.jobs.values().cloned().collect::<Vec<_>>();
428 out.sort_by_key(|job| std::cmp::Reverse(job.updated_at));
429 out
430 }
431
432 pub fn history(&self, id: &str) -> Vec<JobHistoryEntry> {
434 self.jobs
435 .get(id)
436 .map(|job| job.history.clone())
437 .unwrap_or_default()
438 }
439
440 pub fn resume_pending(&mut self) -> Vec<JobRecord> {
442 let mut resumed = Vec::new();
443 for job in self.jobs.values_mut() {
444 if matches!(job.status, JobStatus::Queued | JobStatus::Running) {
445 job.status = JobStatus::Queued;
446 job.updated_at = Self::now_ts();
447 Self::push_history(job, "queued_after_resume");
448 resumed.push(job.clone());
449 }
450 }
451 resumed
452 }
453
454 pub fn load_from_store(&mut self, store: &StateStore) -> Result<()> {
456 let persisted = store.list_jobs(Some(500))?;
457 for job in persisted {
458 let fallback_status = job_state_status_to_runtime(job.status);
459 let parsed = Self::parse_persisted_detail(job.detail.as_deref());
460 let (status, detail, retry, history) = if let Some(detail_state) = parsed {
461 (
462 detail_state.status,
463 detail_state.detail,
464 detail_state.retry,
465 detail_state.history,
466 )
467 } else {
468 (
469 fallback_status,
470 job.detail,
471 JobRetryMetadata::default(),
472 Vec::new(),
473 )
474 };
475 self.jobs.insert(
476 job.id.clone(),
477 JobRecord {
478 id: job.id,
479 name: job.name,
480 status,
481 progress: job.progress,
482 detail,
483 retry,
484 history,
485 created_at: job.created_at,
486 updated_at: job.updated_at,
487 },
488 );
489 }
490 Ok(())
491 }
492
493 pub fn persist_job(&self, store: &StateStore, id: &str) -> Result<()> {
495 let Some(job) = self.jobs.get(id) else {
496 return Ok(());
497 };
498 let encoded_detail = Self::encode_persisted_detail(job)?;
499 store.upsert_job(&JobStateRecord {
500 id: job.id.clone(),
501 name: job.name.clone(),
502 status: runtime_status_to_job_state(job.status),
503 progress: job.progress,
504 detail: encoded_detail,
505 created_at: job.created_at,
506 updated_at: job.updated_at,
507 })
508 }
509
510 pub fn persist_all(&self, store: &StateStore) -> Result<()> {
512 for id in self.jobs.keys() {
513 self.persist_job(store, id)?;
514 }
515 Ok(())
516 }
517}
518
519pub struct ThreadManager {
521 store: StateStore,
522 running_threads: HashMap<String, Thread>,
523 cli_version: String,
524}
525
526impl ThreadManager {
527 pub fn new(store: StateStore) -> Self {
529 Self {
530 store,
531 running_threads: HashMap::new(),
532 cli_version: env!("CARGO_PKG_VERSION").to_string(),
533 }
534 }
535
536 pub fn state_store(&self) -> &StateStore {
538 &self.store
539 }
540
541 pub fn spawn_thread_with_history(
543 &mut self,
544 model_provider: String,
545 cwd: PathBuf,
546 initial_history: InitialHistory,
547 persist_extended_history: bool,
548 ) -> Result<NewThread> {
549 let id = format!("thread-{}", Uuid::new_v4());
550 let now = chrono::Utc::now().timestamp();
551 let preview = preview_from_initial_history(&initial_history);
552 let source = match initial_history {
553 InitialHistory::New => SessionSource::Interactive,
554 InitialHistory::Forked(_) => SessionSource::Fork,
555 InitialHistory::Resumed { .. } => SessionSource::Resume,
556 };
557 let thread = Thread {
558 id: id.clone(),
559 preview,
560 ephemeral: !persist_extended_history,
561 model_provider: model_provider.clone(),
562 created_at: now,
563 updated_at: now,
564 status: ThreadStatus::Running,
565 path: None,
566 cwd: cwd.clone(),
567 cli_version: self.cli_version.clone(),
568 source: match source {
569 SessionSource::Interactive => codewhale_protocol::SessionSource::Interactive,
570 SessionSource::Resume => codewhale_protocol::SessionSource::Resume,
571 SessionSource::Fork => codewhale_protocol::SessionSource::Fork,
572 SessionSource::Api => codewhale_protocol::SessionSource::Api,
573 SessionSource::Unknown => codewhale_protocol::SessionSource::Unknown,
574 },
575 name: None,
576 };
577 self.persist_thread(&thread, None)?;
578 match &initial_history {
579 InitialHistory::Forked(items) => {
580 for item in items {
581 self.store.append_message(
582 &thread.id,
583 "history",
584 &item.to_string(),
585 Some(item.clone()),
586 )?;
587 }
588 }
589 InitialHistory::Resumed { history, .. } => {
590 for item in history {
591 self.store.append_message(
592 &thread.id,
593 "history",
594 &item.to_string(),
595 Some(item.clone()),
596 )?;
597 }
598 }
599 InitialHistory::New => {}
600 }
601 self.running_threads
602 .insert(thread.id.clone(), thread.clone());
603 Ok(NewThread {
604 thread,
605 model: "auto".to_string(),
606 model_provider,
607 cwd,
608 approval_policy: None,
609 sandbox: None,
610 })
611 }
612
613 pub fn resume_thread_with_history(
615 &mut self,
616 params: &ThreadResumeParams,
617 fallback_cwd: &Path,
618 model_provider: String,
619 ) -> Result<Option<NewThread>> {
620 if params.history.is_none()
621 && let Some(thread) = self.running_threads.get(¶ms.thread_id).cloned()
622 {
623 return Ok(Some(NewThread {
624 model: params.model.clone().unwrap_or_else(|| "auto".to_string()),
625 model_provider: params.model_provider.clone().unwrap_or(model_provider),
626 cwd: params.cwd.clone().unwrap_or_else(|| thread.cwd.clone()),
627 approval_policy: params.approval_policy.clone(),
628 sandbox: params.sandbox.clone(),
629 thread,
630 }));
631 }
632
633 let persisted = self.store.get_thread(¶ms.thread_id)?;
634 let Some(metadata) = persisted else {
635 return Ok(None);
636 };
637 let mut thread = to_protocol_thread(metadata);
638 thread.status = ThreadStatus::Running;
639 thread.updated_at = chrono::Utc::now().timestamp();
640 thread.cwd = params
641 .cwd
642 .clone()
643 .unwrap_or_else(|| fallback_cwd.to_path_buf());
644 self.persist_thread(&thread, None)?;
645 self.running_threads
646 .insert(thread.id.clone(), thread.clone());
647 if let Some(history) = params.history.as_ref() {
648 let mut seen: HashSet<String> = self
655 .store
656 .list_messages(&thread.id, None)?
657 .into_iter()
658 .map(|message| {
659 message
660 .item
661 .as_ref()
662 .map_or(message.content.clone(), |item| item.to_string())
663 })
664 .collect();
665 for item in history {
666 let fingerprint = item.to_string();
667 if !seen.insert(fingerprint.clone()) {
668 continue;
669 }
670 self.store.append_message(
671 &thread.id,
672 "history",
673 &fingerprint,
674 Some(item.clone()),
675 )?;
676 }
677 }
678
679 Ok(Some(NewThread {
680 model: params.model.clone().unwrap_or_else(|| "auto".to_string()),
681 model_provider: params.model_provider.clone().unwrap_or(model_provider),
682 cwd: thread.cwd.clone(),
683 approval_policy: params.approval_policy.clone(),
684 sandbox: params.sandbox.clone(),
685 thread,
686 }))
687 }
688
689 pub fn fork_thread(
691 &mut self,
692 params: &ThreadForkParams,
693 fallback_cwd: &Path,
694 ) -> Result<Option<NewThread>> {
695 let parent = self.store.get_thread(¶ms.thread_id)?;
696 let Some(parent) = parent else {
697 return Ok(None);
698 };
699 let parent_thread = to_protocol_thread(parent);
700 let new = self.spawn_thread_with_history(
701 params
702 .model_provider
703 .clone()
704 .unwrap_or_else(|| parent_thread.model_provider.clone()),
705 params
706 .cwd
707 .clone()
708 .unwrap_or_else(|| fallback_cwd.to_path_buf()),
709 InitialHistory::Forked(vec![json!({
710 "type": "fork",
711 "from_thread_id": parent_thread.id
712 })]),
713 params.persist_extended_history,
714 )?;
715 Ok(Some(new))
716 }
717
718 pub fn list_threads(&self, params: &ThreadListParams) -> Result<Vec<Thread>> {
720 let list = self.store.list_threads(ThreadListFilters {
721 include_archived: params.include_archived,
722 limit: params.limit,
723 })?;
724 Ok(list.into_iter().map(to_protocol_thread).collect())
725 }
726
727 pub fn read_thread(&self, params: &ThreadReadParams) -> Result<Option<Thread>> {
729 Ok(self
730 .store
731 .get_thread(¶ms.thread_id)?
732 .map(to_protocol_thread))
733 }
734
735 pub fn set_thread_name(&mut self, params: &ThreadSetNameParams) -> Result<Option<Thread>> {
737 let Some(mut metadata) = self.store.get_thread(¶ms.thread_id)? else {
738 return Ok(None);
739 };
740 metadata.name = Some(params.name.clone());
741 metadata.updated_at = chrono::Utc::now().timestamp();
742 self.store.upsert_thread(&metadata)?;
743 let updated = to_protocol_thread(metadata);
744 self.running_threads
745 .insert(updated.id.clone(), updated.clone());
746 Ok(Some(updated))
747 }
748
749 pub fn set_thread_goal(&mut self, params: &ThreadGoalSetParams) -> Result<Option<ThreadGoal>> {
751 if self.store.get_thread(¶ms.thread_id)?.is_none() {
752 return Ok(None);
753 }
754 let now = chrono::Utc::now().timestamp();
755 let goal = ThreadGoalRecord {
756 thread_id: params.thread_id.clone(),
757 goal_id: format!("goal-{}", Uuid::new_v4()),
758 objective: params.objective.clone(),
759 status: PersistedThreadGoalStatus::Active,
760 token_budget: params.token_budget,
761 tokens_used: 0,
762 time_used_seconds: 0,
763 continuation_count: 0,
764 created_at: now,
765 updated_at: now,
766 };
767 self.store.upsert_thread_goal(&goal)?;
768 Ok(Some(to_protocol_goal(goal)))
769 }
770
771 pub fn get_thread_goal(&self, params: &ThreadGoalGetParams) -> Result<Option<ThreadGoal>> {
773 Ok(self
774 .store
775 .get_thread_goal(¶ms.thread_id)?
776 .map(to_protocol_goal))
777 }
778
779 pub fn record_thread_goal_progress(
781 &mut self,
782 params: &ThreadGoalProgressParams,
783 ) -> Result<Option<ThreadGoal>> {
784 if self.store.get_thread(¶ms.thread_id)?.is_none() {
785 return Ok(None);
786 }
787
788 let now = chrono::Utc::now().timestamp();
789 let mut goal = if params.token_delta != 0 || params.time_delta_seconds != 0 {
790 self.store.record_thread_goal_usage(
791 ¶ms.thread_id,
792 params.token_delta,
793 params.time_delta_seconds,
794 now,
795 )?
796 } else {
797 self.store.get_thread_goal(¶ms.thread_id)?
798 };
799
800 if params.record_continuation {
801 goal = self
802 .store
803 .record_thread_goal_continuation(¶ms.thread_id, now)?;
804 }
805
806 Ok(goal.map(to_protocol_goal))
807 }
808
809 pub fn clear_thread_goal(&mut self, params: &ThreadGoalClearParams) -> Result<bool> {
811 self.store.delete_thread_goal(¶ms.thread_id)
812 }
813
814 pub fn archive_thread(&mut self, thread_id: &str) -> Result<()> {
816 self.store.mark_archived(thread_id)?;
817 if let Some(thread) = self.running_threads.get_mut(thread_id) {
818 thread.status = ThreadStatus::Archived;
819 }
820 Ok(())
821 }
822
823 pub fn unarchive_thread(&mut self, thread_id: &str) -> Result<()> {
825 self.store.mark_unarchived(thread_id)?;
826 if let Some(metadata) = self.store.get_thread(thread_id)? {
827 let thread = to_protocol_thread(metadata);
828 if let Some(cached) = self.running_threads.get_mut(thread_id) {
829 *cached = thread;
830 }
831 }
832 Ok(())
833 }
834
835 pub fn touch_message(&mut self, thread_id: &str, input: &str) -> Result<()> {
837 let Some(mut metadata) = self.store.get_thread(thread_id)? else {
838 return Ok(());
839 };
840 metadata.updated_at = chrono::Utc::now().timestamp();
841 metadata.preview = truncate_preview(input);
842 metadata.status = PersistedThreadStatus::Running;
843 self.store.upsert_thread(&metadata)?;
844 if let Some(thread) = self.running_threads.get_mut(thread_id) {
845 thread.updated_at = metadata.updated_at;
846 thread.preview = metadata.preview;
847 thread.status = ThreadStatus::Running;
848 }
849 let message_id = self.store.append_message(thread_id, "user", input, None)?;
850 self.store.save_checkpoint(
851 thread_id,
852 "latest",
853 &json!({
854 "reason": "thread_message",
855 "message_id": message_id,
856 "role": "user",
857 "preview": truncate_preview(input),
858 "updated_at": metadata.updated_at
859 }),
860 )?;
861 Ok(())
862 }
863
864 fn persist_thread(&self, thread: &Thread, rollout_path: Option<PathBuf>) -> Result<()> {
865 let existing = self.store.get_thread(&thread.id)?;
869 self.store.upsert_thread(&ThreadMetadata {
870 id: thread.id.clone(),
871 rollout_path,
872 preview: thread.preview.clone(),
873 ephemeral: thread.ephemeral,
874 model_provider: thread.model_provider.clone(),
875 created_at: thread.created_at,
876 updated_at: thread.updated_at,
877 status: to_persisted_status(&thread.status),
878 path: thread.path.clone(),
879 cwd: thread.cwd.clone(),
880 cli_version: thread.cli_version.clone(),
881 source: to_persisted_source(&thread.source),
882 name: thread.name.clone(),
883 sandbox_policy: existing
884 .as_ref()
885 .and_then(|metadata| metadata.sandbox_policy.clone()),
886 approval_mode: existing
887 .as_ref()
888 .and_then(|metadata| metadata.approval_mode.clone()),
889 archived: matches!(thread.status, ThreadStatus::Archived),
890 archived_at: None,
891 git_sha: None,
892 git_branch: None,
893 git_origin_url: None,
894 memory_mode: None,
895 current_leaf_id: None,
896 })
897 }
898}
899
900pub struct Runtime {
902 pub config: ConfigToml,
904 pub model_registry: ModelRegistry,
906 pub thread_manager: ThreadManager,
908 pub tool_registry: Arc<ToolRegistry>,
910 pub mcp_manager: Arc<McpManager>,
912 pub exec_policy: ExecPolicyEngine,
914 pub hooks: HookDispatcher,
916 pub jobs: JobManager,
918}
919
920impl Runtime {
921 pub fn new(
923 config: ConfigToml,
924 model_registry: ModelRegistry,
925 state: StateStore,
926 tool_registry: Arc<ToolRegistry>,
927 mcp_manager: Arc<McpManager>,
928 exec_policy: ExecPolicyEngine,
929 hooks: HookDispatcher,
930 ) -> Self {
931 let mut jobs = JobManager::default();
932 if let Err(e) = jobs.load_from_store(&state) {
933 tracing::warn!("Failed to load job store, starting with empty job list: {e}");
934 }
935 Self {
936 config,
937 model_registry,
938 thread_manager: ThreadManager::new(state),
939 tool_registry,
940 mcp_manager,
941 exec_policy,
942 hooks,
943 jobs,
944 }
945 }
946
947 pub fn update_config(&mut self, config: ConfigToml) {
959 self.config = config;
960 }
961
962 pub fn reload_config_and_policy(&mut self, config: ConfigToml, exec_policy: ExecPolicyEngine) {
983 self.config = config;
984 self.exec_policy = exec_policy;
985 }
986
987 fn persisted_thread_data(&self, thread_id: &str) -> Result<Value> {
988 let history = self
989 .thread_manager
990 .state_store()
991 .list_messages(thread_id, Some(500))?
992 .into_iter()
993 .map(|message| {
994 json!({
995 "id": message.id,
996 "role": message.role,
997 "content": message.content,
998 "item": message.item,
999 "created_at": message.created_at
1000 })
1001 })
1002 .collect::<Vec<_>>();
1003
1004 let checkpoint = self
1005 .thread_manager
1006 .state_store()
1007 .load_checkpoint(thread_id, None)?
1008 .map(|record| {
1009 json!({
1010 "checkpoint_id": record.checkpoint_id,
1011 "state": record.state,
1012 "created_at": record.created_at
1013 })
1014 });
1015
1016 let goal = self
1017 .thread_manager
1018 .state_store()
1019 .get_thread_goal(thread_id)?
1020 .map(to_protocol_goal);
1021
1022 Ok(json!({
1023 "history": history,
1024 "checkpoint": checkpoint,
1025 "goal": goal
1026 }))
1027 }
1028
1029 fn persist_latest_checkpoint(&self, thread_id: &str, reason: &str, state: Value) -> Result<()> {
1030 self.thread_manager.state_store().save_checkpoint(
1031 thread_id,
1032 "latest",
1033 &json!({
1034 "reason": reason,
1035 "saved_at": chrono::Utc::now().timestamp(),
1036 "state": state
1037 }),
1038 )
1039 }
1040
1041 pub async fn handle_thread(&mut self, req: ThreadRequest) -> Result<ThreadResponse> {
1043 match req {
1044 ThreadRequest::Create { .. } => {
1045 let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
1046 let new = self.thread_manager.spawn_thread_with_history(
1047 "deepseek".to_string(),
1048 cwd,
1049 InitialHistory::New,
1050 false,
1051 )?;
1052 let mut response = thread_response_from_new("created", new);
1053 response.data = self.persisted_thread_data(&response.thread_id)?;
1054 Ok(response)
1055 }
1056 ThreadRequest::Start(params) => {
1057 let cwd = params.cwd.clone().unwrap_or_else(|| {
1058 std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."))
1059 });
1060 let new = self.thread_manager.spawn_thread_with_history(
1061 params
1062 .model_provider
1063 .clone()
1064 .unwrap_or_else(|| "deepseek".to_string()),
1065 cwd,
1066 InitialHistory::New,
1067 params.persist_extended_history,
1068 )?;
1069 let mut response = thread_response_from_new("started", new);
1070 response.data = self.persisted_thread_data(&response.thread_id)?;
1071 Ok(response)
1072 }
1073 ThreadRequest::Resume(params) => {
1074 let fallback_cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
1075 if let Some(new) = self.thread_manager.resume_thread_with_history(
1076 ¶ms,
1077 &fallback_cwd,
1078 "deepseek".to_string(),
1079 )? {
1080 let mut response = thread_response_from_new("resumed", new);
1081 response.data = self.persisted_thread_data(&response.thread_id)?;
1082 Ok(response)
1083 } else {
1084 Ok(ThreadResponse {
1085 thread_id: params.thread_id,
1086 status: "missing".to_string(),
1087 thread: None,
1088 threads: Vec::new(),
1089 goal: None,
1090 model: None,
1091 model_provider: None,
1092 cwd: None,
1093 approval_policy: params.approval_policy,
1094 sandbox: params.sandbox,
1095 events: Vec::new(),
1096 data: json!({"error":"thread not found"}),
1097 })
1098 }
1099 }
1100 ThreadRequest::Fork(params) => {
1101 let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from("."));
1102 if let Some(new) = self.thread_manager.fork_thread(¶ms, &cwd)? {
1103 let mut response = thread_response_from_new("forked", new);
1104 response.data = self.persisted_thread_data(&response.thread_id)?;
1105 Ok(response)
1106 } else {
1107 Ok(ThreadResponse {
1108 thread_id: params.thread_id,
1109 status: "missing".to_string(),
1110 thread: None,
1111 threads: Vec::new(),
1112 goal: None,
1113 model: None,
1114 model_provider: None,
1115 cwd: None,
1116 approval_policy: params.approval_policy,
1117 sandbox: params.sandbox,
1118 events: Vec::new(),
1119 data: json!({"error":"thread not found"}),
1120 })
1121 }
1122 }
1123 ThreadRequest::List(params) => Ok(ThreadResponse {
1124 thread_id: "list".to_string(),
1125 status: "ok".to_string(),
1126 thread: None,
1127 threads: self.thread_manager.list_threads(¶ms)?,
1128 goal: None,
1129 model: None,
1130 model_provider: None,
1131 cwd: None,
1132 approval_policy: None,
1133 sandbox: None,
1134 events: Vec::new(),
1135 data: json!({}),
1136 }),
1137 ThreadRequest::Read(params) => {
1138 let id = params.thread_id.clone();
1139 let data = self.persisted_thread_data(&id)?;
1140 Ok(ThreadResponse {
1141 thread_id: id,
1142 status: "ok".to_string(),
1143 thread: self.thread_manager.read_thread(¶ms)?,
1144 threads: Vec::new(),
1145 goal: self.thread_manager.get_thread_goal(&ThreadGoalGetParams {
1146 thread_id: params.thread_id,
1147 })?,
1148 model: None,
1149 model_provider: None,
1150 cwd: None,
1151 approval_policy: None,
1152 sandbox: None,
1153 events: Vec::new(),
1154 data,
1155 })
1156 }
1157 ThreadRequest::SetName(params) => Ok(ThreadResponse {
1158 thread_id: params.thread_id.clone(),
1159 status: "ok".to_string(),
1160 thread: self.thread_manager.set_thread_name(¶ms)?,
1161 threads: Vec::new(),
1162 goal: None,
1163 model: None,
1164 model_provider: None,
1165 cwd: None,
1166 approval_policy: None,
1167 sandbox: None,
1168 events: Vec::new(),
1169 data: json!({}),
1170 }),
1171 ThreadRequest::GoalSet(params) => {
1172 let thread_id = params.thread_id.clone();
1173 if let Some(goal) = self.thread_manager.set_thread_goal(¶ms)? {
1174 Ok(ThreadResponse {
1175 thread_id,
1176 status: "ok".to_string(),
1177 thread: None,
1178 threads: Vec::new(),
1179 goal: Some(goal.clone()),
1180 model: None,
1181 model_provider: None,
1182 cwd: None,
1183 approval_policy: None,
1184 sandbox: None,
1185 events: vec![EventFrame::ThreadGoalUpdated { goal: goal.clone() }],
1186 data: json!({ "goal": goal }),
1187 })
1188 } else {
1189 Ok(ThreadResponse {
1190 thread_id,
1191 status: "missing".to_string(),
1192 thread: None,
1193 threads: Vec::new(),
1194 goal: None,
1195 model: None,
1196 model_provider: None,
1197 cwd: None,
1198 approval_policy: None,
1199 sandbox: None,
1200 events: Vec::new(),
1201 data: json!({"error":"thread not found"}),
1202 })
1203 }
1204 }
1205 ThreadRequest::GoalGet(params) => {
1206 let goal = self.thread_manager.get_thread_goal(¶ms)?;
1207 Ok(ThreadResponse {
1208 thread_id: params.thread_id,
1209 status: "ok".to_string(),
1210 thread: None,
1211 threads: Vec::new(),
1212 goal: goal.clone(),
1213 model: None,
1214 model_provider: None,
1215 cwd: None,
1216 approval_policy: None,
1217 sandbox: None,
1218 events: Vec::new(),
1219 data: json!({ "goal": goal }),
1220 })
1221 }
1222 ThreadRequest::GoalClear(params) => {
1223 let thread_id = params.thread_id.clone();
1224 let cleared = self.thread_manager.clear_thread_goal(¶ms)?;
1225 Ok(ThreadResponse {
1226 thread_id: thread_id.clone(),
1227 status: if cleared { "cleared" } else { "empty" }.to_string(),
1228 thread: None,
1229 threads: Vec::new(),
1230 goal: None,
1231 model: None,
1232 model_provider: None,
1233 cwd: None,
1234 approval_policy: None,
1235 sandbox: None,
1236 events: if cleared {
1237 vec![EventFrame::ThreadGoalCleared { thread_id }]
1238 } else {
1239 Vec::new()
1240 },
1241 data: json!({ "cleared": cleared }),
1242 })
1243 }
1244 ThreadRequest::GoalRecordProgress(params) => {
1245 let thread_id = params.thread_id.clone();
1246 if let Some(goal) = self.thread_manager.record_thread_goal_progress(¶ms)? {
1247 Ok(ThreadResponse {
1248 thread_id,
1249 status: "ok".to_string(),
1250 thread: None,
1251 threads: Vec::new(),
1252 goal: Some(goal.clone()),
1253 model: None,
1254 model_provider: None,
1255 cwd: None,
1256 approval_policy: None,
1257 sandbox: None,
1258 events: vec![EventFrame::ThreadGoalUpdated { goal: goal.clone() }],
1259 data: json!({ "goal": goal }),
1260 })
1261 } else {
1262 Ok(ThreadResponse {
1263 thread_id,
1264 status: "missing".to_string(),
1265 thread: None,
1266 threads: Vec::new(),
1267 goal: None,
1268 model: None,
1269 model_provider: None,
1270 cwd: None,
1271 approval_policy: None,
1272 sandbox: None,
1273 events: Vec::new(),
1274 data: json!({"error":"thread or goal not found"}),
1275 })
1276 }
1277 }
1278 ThreadRequest::Archive { thread_id } => {
1279 self.thread_manager.archive_thread(&thread_id)?;
1280 Ok(ThreadResponse {
1281 thread_id,
1282 status: "archived".to_string(),
1283 thread: None,
1284 threads: Vec::new(),
1285 goal: None,
1286 model: None,
1287 model_provider: None,
1288 cwd: None,
1289 approval_policy: None,
1290 sandbox: None,
1291 events: Vec::new(),
1292 data: json!({}),
1293 })
1294 }
1295 ThreadRequest::Unarchive { thread_id } => {
1296 self.thread_manager.unarchive_thread(&thread_id)?;
1297 Ok(ThreadResponse {
1298 thread_id,
1299 status: "unarchived".to_string(),
1300 thread: None,
1301 threads: Vec::new(),
1302 goal: None,
1303 model: None,
1304 model_provider: None,
1305 cwd: None,
1306 approval_policy: None,
1307 sandbox: None,
1308 events: Vec::new(),
1309 data: json!({}),
1310 })
1311 }
1312 ThreadRequest::Message { thread_id, input } => {
1313 self.thread_manager.touch_message(&thread_id, &input)?;
1314 let response_id = format!("resp-{}", Uuid::new_v4());
1318 self.hooks
1319 .emit(HookEvent::ResponseStart {
1320 response_id: response_id.clone(),
1321 })
1322 .await;
1323 self.hooks
1324 .emit(HookEvent::ResponseEnd {
1325 response_id: response_id.clone(),
1326 })
1327 .await;
1328
1329 Ok(ThreadResponse {
1330 thread_id,
1331 status: "accepted".to_string(),
1332 thread: None,
1333 threads: Vec::new(),
1334 goal: None,
1335 model: None,
1336 model_provider: None,
1337 cwd: None,
1338 approval_policy: None,
1339 sandbox: None,
1340 events: vec![
1341 EventFrame::ResponseStart {
1342 response_id: response_id.clone(),
1343 },
1344 EventFrame::ResponseDelta {
1345 response_id: response_id.clone(),
1346 delta: "queued".to_string(),
1347 channel: ResponseChannel::Text,
1348 },
1349 EventFrame::ResponseEnd { response_id },
1350 ],
1351 data: json!({}),
1352 })
1353 }
1354 }
1355 }
1356
1357 pub async fn handle_prompt(
1359 &mut self,
1360 req: PromptRequest,
1361 cli_overrides: &CliRuntimeOverrides,
1362 ) -> Result<PromptResponse> {
1363 let resolved = self.config.resolve_runtime_options(cli_overrides);
1364 let requested_model = req.model.clone().unwrap_or_else(|| resolved.model.clone());
1365 let selection = self
1366 .model_registry
1367 .resolve(Some(&requested_model), Some(resolved.provider));
1368 let resolved_model = selection.resolved.id.clone();
1369 let response_id = format!("resp-{}", Uuid::new_v4());
1370
1371 self.hooks
1372 .emit(HookEvent::ResponseStart {
1373 response_id: response_id.clone(),
1374 })
1375 .await;
1376 self.hooks
1377 .emit(HookEvent::ResponseDelta {
1378 response_id: response_id.clone(),
1379 delta: "model-selected".to_string(),
1380 })
1381 .await;
1382 self.hooks
1383 .emit(HookEvent::ResponseEnd {
1384 response_id: response_id.clone(),
1385 })
1386 .await;
1387
1388 let payload = json!({
1389 "provider": resolved.provider.as_str(),
1390 "model": resolved_model.clone(),
1391 "prompt": req.prompt,
1392 "telemetry": resolved.telemetry,
1393 "base_url": resolved.base_url,
1394 "has_api_key": resolved.api_key.as_ref().is_some_and(|k| !k.trim().is_empty()),
1395 "approval_policy": resolved.approval_policy,
1396 "sandbox_mode": resolved.sandbox_mode
1397 });
1398 if let Some(thread_id) = req.thread_id.as_ref() {
1399 self.thread_manager.touch_message(thread_id, &req.prompt)?;
1400 let assistant_message_id = self.thread_manager.store.append_message(
1401 thread_id,
1402 "assistant",
1403 &payload.to_string(),
1404 Some(payload.clone()),
1405 )?;
1406 self.persist_latest_checkpoint(
1407 thread_id,
1408 "prompt_response",
1409 json!({
1410 "response_id": response_id.clone(),
1411 "model": resolved_model.clone(),
1412 "provider": resolved.provider.as_str(),
1413 "assistant_message_id": assistant_message_id
1414 }),
1415 )?;
1416 }
1417
1418 Ok(PromptResponse {
1419 output: payload.to_string(),
1420 model: resolved_model,
1421 events: vec![
1422 EventFrame::ResponseStart {
1423 response_id: response_id.clone(),
1424 },
1425 EventFrame::ResponseDelta {
1426 response_id: response_id.clone(),
1427 delta: "model-selected".to_string(),
1428 channel: ResponseChannel::Text,
1429 },
1430 EventFrame::ResponseEnd { response_id },
1431 ],
1432 })
1433 }
1434
1435 pub async fn invoke_tool(
1437 &self,
1438 call: ToolCall,
1439 approval_mode: AskForApproval,
1440 cwd: &Path,
1441 ) -> Result<Value> {
1442 let fallback_cwd = cwd.display().to_string();
1443 let (command, policy_cwd, execution_kind) = call.execution_subject(&fallback_cwd);
1444 let policy_tool = match &call.payload {
1445 ToolPayload::LocalShell { .. } => "exec_shell",
1446 _ => call.name.as_str(),
1447 };
1448 let policy_path = permission_path_for_call(&call);
1449 let decision = self.exec_policy.check(ExecPolicyContext {
1450 command: &command,
1451 cwd: &policy_cwd,
1452 tool: Some(policy_tool),
1453 path: policy_path.as_deref(),
1454 ask_for_approval: approval_mode,
1455 sandbox_mode: None,
1456 })?;
1457 let precheck = policy_precheck_payload(&decision, &command, &policy_cwd, execution_kind);
1458 let response_id = format!("tool-{}", Uuid::new_v4());
1459 let call_id = call
1460 .raw_tool_call_id
1461 .clone()
1462 .unwrap_or_else(|| format!("tool-call-{}", Uuid::new_v4()));
1463 self.hooks
1464 .emit(HookEvent::ToolLifecycle {
1465 response_id: response_id.clone(),
1466 tool_name: call.name.clone(),
1467 phase: "precheck".to_string(),
1468 payload: precheck.clone(),
1469 })
1470 .await;
1471
1472 if !decision.allow {
1473 let reason = decision.reason().to_string();
1474 let approval_id = format!("approval-{}", Uuid::new_v4());
1475 let error_frame = EventFrame::Error {
1476 response_id: response_id.clone(),
1477 message: reason.clone(),
1478 };
1479 self.hooks
1480 .emit(HookEvent::ApprovalLifecycle {
1481 approval_id,
1482 phase: "denied".to_string(),
1483 reason: Some(reason.clone()),
1484 })
1485 .await;
1486 self.hooks
1487 .emit(HookEvent::GenericEventFrame {
1488 frame: Box::new(error_frame.clone()),
1489 })
1490 .await;
1491 return Ok(json!({
1492 "ok": false,
1493 "status": "denied",
1494 "execution_kind": execution_kind,
1495 "response_id": response_id,
1496 "precheck": precheck,
1497 "error": reason,
1498 "events": [event_frame_payload(&error_frame)],
1499 }));
1500 }
1501
1502 if decision.requires_approval {
1503 let approval_id = format!("approval-{}", Uuid::new_v4());
1504 let reason = decision.reason().to_string();
1505 let maybe_approval_frame = approval_request_frame(
1506 &decision.requirement,
1507 decision.matched_rule.as_deref(),
1508 call_id,
1509 approval_id.clone(),
1510 response_id.clone(),
1511 command.clone(),
1512 policy_cwd.clone(),
1513 );
1514 self.hooks
1515 .emit(HookEvent::ApprovalLifecycle {
1516 approval_id: approval_id.clone(),
1517 phase: "requested".to_string(),
1518 reason: Some(reason.clone()),
1519 })
1520 .await;
1521 let mut events = Vec::new();
1522 if let Some(frame) = maybe_approval_frame {
1523 self.hooks
1524 .emit(HookEvent::GenericEventFrame {
1525 frame: Box::new(frame.clone()),
1526 })
1527 .await;
1528 events.push(event_frame_payload(&frame));
1529 }
1530 return Ok(json!({
1531 "ok": false,
1532 "status": "approval_required",
1533 "execution_kind": execution_kind,
1534 "response_id": response_id,
1535 "approval_id": approval_id,
1536 "precheck": precheck,
1537 "error": reason,
1538 "events": events,
1539 }));
1540 }
1541
1542 if call.name == REQUEST_USER_INPUT_TOOL_NAME {
1551 let request_id = format!("user-input-{}", Uuid::new_v4());
1552 let arguments = match &call.payload {
1553 ToolPayload::Function { arguments } => arguments.as_str(),
1554 _ => "",
1557 };
1558 let maybe_frame = user_input_request_frame(
1559 call_id.clone(),
1560 response_id.clone(),
1561 request_id.clone(),
1562 arguments,
1563 );
1564 let mut events = Vec::new();
1565 if let Some(frame) = maybe_frame {
1566 self.hooks
1567 .emit(HookEvent::GenericEventFrame {
1568 frame: Box::new(frame.clone()),
1569 })
1570 .await;
1571 events.push(event_frame_payload(&frame));
1572 }
1573 return Ok(json!({
1574 "ok": false,
1575 "status": "user_input_required",
1576 "execution_kind": execution_kind,
1577 "response_id": response_id,
1578 "request_id": request_id,
1579 "precheck": precheck,
1580 "events": events,
1581 }));
1582 }
1583
1584 let start_frame = EventFrame::ToolCallStart {
1585 response_id: response_id.clone(),
1586 tool_name: call.name.clone(),
1587 arguments: tool_payload_value(&call.payload),
1588 };
1589 self.hooks
1590 .emit(HookEvent::GenericEventFrame {
1591 frame: Box::new(start_frame.clone()),
1592 })
1593 .await;
1594 self.hooks
1595 .emit(HookEvent::ToolLifecycle {
1596 response_id: response_id.clone(),
1597 tool_name: call.name.clone(),
1598 phase: "dispatching".to_string(),
1599 payload: json!({
1600 "call_id": call_id,
1601 "execution_kind": execution_kind
1602 }),
1603 })
1604 .await;
1605
1606 match time::timeout(
1607 tool_dispatch_timeout(),
1608 self.tool_registry.dispatch(call.clone(), true),
1609 )
1610 .await
1611 {
1612 Ok(Ok(tool_output)) => {
1613 let success = tool_output.success();
1614 let status = if success { "completed" } else { "failed" };
1615 let result_frame = EventFrame::ToolCallResult {
1616 response_id: response_id.clone(),
1617 tool_name: call.name.clone(),
1618 output: tool_output_value(&tool_output),
1619 };
1620 self.hooks
1621 .emit(HookEvent::GenericEventFrame {
1622 frame: Box::new(result_frame.clone()),
1623 })
1624 .await;
1625 self.hooks
1626 .emit(HookEvent::ToolLifecycle {
1627 response_id: response_id.clone(),
1628 tool_name: call.name,
1629 phase: status.to_string(),
1630 payload: json!({ "ok": success }),
1631 })
1632 .await;
1633 Ok(json!({
1634 "ok": success,
1635 "status": status,
1636 "execution_kind": execution_kind,
1637 "response_id": response_id,
1638 "precheck": precheck,
1639 "output": tool_output,
1640 "events": [
1641 event_frame_payload(&start_frame),
1642 event_frame_payload(&result_frame)
1643 ]
1644 }))
1645 }
1646 Ok(Err(err)) => {
1647 let message = format!("{err:?}");
1648 let error_frame = EventFrame::Error {
1649 response_id: response_id.clone(),
1650 message: message.clone(),
1651 };
1652 self.hooks
1653 .emit(HookEvent::GenericEventFrame {
1654 frame: Box::new(error_frame.clone()),
1655 })
1656 .await;
1657 self.hooks
1658 .emit(HookEvent::ToolLifecycle {
1659 response_id: response_id.clone(),
1660 tool_name: call.name,
1661 phase: "failed".to_string(),
1662 payload: json!({ "error": message.clone() }),
1663 })
1664 .await;
1665 Ok(json!({
1666 "ok": false,
1667 "status": "failed",
1668 "execution_kind": execution_kind,
1669 "response_id": response_id,
1670 "precheck": precheck,
1671 "error": message,
1672 "events": [
1673 event_frame_payload(&start_frame),
1674 event_frame_payload(&error_frame)
1675 ]
1676 }))
1677 }
1678 Err(_elapsed) => {
1679 let seconds = tool_dispatch_timeout().as_secs().max(1);
1680 let message = format!("Tool '{}' timed out after {seconds}s", call.name);
1681 let error_frame = EventFrame::Error {
1682 response_id: response_id.clone(),
1683 message: message.clone(),
1684 };
1685 self.hooks
1686 .emit(HookEvent::GenericEventFrame {
1687 frame: Box::new(error_frame.clone()),
1688 })
1689 .await;
1690 self.hooks
1691 .emit(HookEvent::ToolLifecycle {
1692 response_id: response_id.clone(),
1693 tool_name: call.name,
1694 phase: "failed".to_string(),
1695 payload: json!({ "error": message.clone(), "timeout": true }),
1696 })
1697 .await;
1698 Ok(json!({
1699 "ok": false,
1700 "status": "timeout",
1701 "execution_kind": execution_kind,
1702 "response_id": response_id,
1703 "precheck": precheck,
1704 "error": message,
1705 "events": [
1706 event_frame_payload(&start_frame),
1707 event_frame_payload(&error_frame)
1708 ]
1709 }))
1710 }
1711 }
1712 }
1713
1714 pub async fn mcp_startup(&self) -> McpStartupCompleteEvent {
1716 let mut updates = Vec::new();
1717 let summary = self.mcp_manager.start_all(|update| {
1718 updates.push(update);
1719 });
1720 for update in updates {
1721 let status = match update.status {
1722 McpManagerStartupStatus::Starting => codewhale_protocol::McpStartupStatus::Starting,
1723 McpManagerStartupStatus::Ready => codewhale_protocol::McpStartupStatus::Ready,
1724 McpManagerStartupStatus::Failed { error } => {
1725 codewhale_protocol::McpStartupStatus::Failed { error }
1726 }
1727 McpManagerStartupStatus::Cancelled => {
1728 codewhale_protocol::McpStartupStatus::Cancelled
1729 }
1730 };
1731 self.hooks
1732 .emit(HookEvent::GenericEventFrame {
1733 frame: Box::new(EventFrame::McpStartupUpdate {
1734 update: codewhale_protocol::McpStartupUpdateEvent {
1735 server_name: update.server_name,
1736 status,
1737 },
1738 }),
1739 })
1740 .await;
1741 }
1742 self.hooks
1743 .emit(HookEvent::GenericEventFrame {
1744 frame: Box::new(EventFrame::McpStartupComplete {
1745 summary: codewhale_protocol::McpStartupCompleteEvent {
1746 ready: summary.ready.clone(),
1747 failed: summary
1748 .failed
1749 .iter()
1750 .map(|f| codewhale_protocol::McpStartupFailure {
1751 server_name: f.server_name.clone(),
1752 error: f.error.clone(),
1753 })
1754 .collect(),
1755 cancelled: summary.cancelled.clone(),
1756 },
1757 }),
1758 })
1759 .await;
1760 summary
1761 }
1762
1763 pub fn app_status(&self) -> AppResponse {
1765 let jobs = self.jobs.list();
1766 let events = jobs
1767 .iter()
1768 .flat_map(|job| {
1769 job.history.iter().map(|entry| EventFrame::ResponseDelta {
1770 response_id: job.id.clone(),
1771 delta: json!({
1772 "kind": "job_transition",
1773 "job_id": job.id.clone(),
1774 "phase": entry.phase.clone(),
1775 "status": job_status_to_str(entry.status),
1776 "progress": entry.progress,
1777 "detail": entry.detail.clone(),
1778 "retry": job_retry_to_value(&entry.retry),
1779 "at": entry.at
1780 })
1781 .to_string(),
1782 channel: ResponseChannel::Text,
1783 })
1784 })
1785 .collect::<Vec<_>>();
1786 AppResponse {
1787 ok: true,
1788 data: json!({
1789 "jobs": jobs.into_iter().map(|job| {
1790 json!({
1791 "id": job.id,
1792 "name": job.name,
1793 "status": job_status_to_str(job.status),
1794 "progress": job.progress,
1795 "detail": job.detail,
1796 "retry": job_retry_to_value(&job.retry),
1797 "history": job.history.iter().map(job_history_to_value).collect::<Vec<_>>()
1798 })
1799 }).collect::<Vec<_>>()
1800 }),
1801 events,
1802 }
1803 }
1804
1805 pub fn provider_default(&self) -> ProviderKind {
1807 self.config.provider
1808 }
1809
1810 pub fn save_thread_checkpoint(
1812 &self,
1813 thread_id: &str,
1814 checkpoint_id: &str,
1815 state: &Value,
1816 ) -> Result<()> {
1817 self.thread_manager
1818 .state_store()
1819 .save_checkpoint(thread_id, checkpoint_id, state)
1820 }
1821
1822 pub fn load_thread_checkpoint(
1824 &self,
1825 thread_id: &str,
1826 checkpoint_id: Option<&str>,
1827 ) -> Result<Option<Value>> {
1828 Ok(self
1829 .thread_manager
1830 .state_store()
1831 .load_checkpoint(thread_id, checkpoint_id)?
1832 .map(|checkpoint| checkpoint.state))
1833 }
1834
1835 pub fn enqueue_job(&mut self, name: impl Into<String>) -> Result<JobRecord> {
1837 let job = self.jobs.enqueue(name);
1838 self.jobs
1839 .persist_job(self.thread_manager.state_store(), &job.id)?;
1840 Ok(job)
1841 }
1842
1843 pub fn set_job_running(&mut self, job_id: &str) -> Result<()> {
1845 self.jobs.set_running(job_id);
1846 self.jobs
1847 .persist_job(self.thread_manager.state_store(), job_id)
1848 }
1849
1850 pub fn update_job_progress(
1852 &mut self,
1853 job_id: &str,
1854 progress: u8,
1855 detail: Option<String>,
1856 ) -> Result<()> {
1857 self.jobs.update_progress(job_id, progress, detail);
1858 self.jobs
1859 .persist_job(self.thread_manager.state_store(), job_id)
1860 }
1861
1862 pub fn complete_job(&mut self, job_id: &str) -> Result<()> {
1864 self.jobs.complete(job_id);
1865 self.jobs
1866 .persist_job(self.thread_manager.state_store(), job_id)
1867 }
1868
1869 pub fn fail_job(&mut self, job_id: &str, detail: impl Into<String>) -> Result<()> {
1871 self.jobs.fail(job_id, detail);
1872 self.jobs
1873 .persist_job(self.thread_manager.state_store(), job_id)
1874 }
1875
1876 pub fn cancel_job(&mut self, job_id: &str) -> Result<()> {
1878 self.jobs.cancel(job_id);
1879 self.jobs
1880 .persist_job(self.thread_manager.state_store(), job_id)
1881 }
1882
1883 pub fn pause_job(&mut self, job_id: &str, detail: Option<String>) -> Result<()> {
1885 self.jobs.pause(job_id, detail);
1886 self.jobs
1887 .persist_job(self.thread_manager.state_store(), job_id)
1888 }
1889
1890 pub fn resume_job(&mut self, job_id: &str, detail: Option<String>) -> Result<()> {
1892 self.jobs.resume(job_id, detail);
1893 self.jobs
1894 .persist_job(self.thread_manager.state_store(), job_id)
1895 }
1896
1897 pub fn job_history(&self, job_id: &str) -> Vec<JobHistoryEntry> {
1899 self.jobs.history(job_id)
1900 }
1901}
1902
1903fn thread_response_from_new(status: &str, new: NewThread) -> ThreadResponse {
1904 ThreadResponse {
1905 thread_id: new.thread.id.clone(),
1906 status: status.to_string(),
1907 thread: Some(new.thread),
1908 threads: Vec::new(),
1909 goal: None,
1910 model: Some(new.model),
1911 model_provider: Some(new.model_provider),
1912 cwd: Some(new.cwd),
1913 approval_policy: new.approval_policy,
1914 sandbox: new.sandbox,
1915 events: Vec::new(),
1916 data: json!({}),
1917 }
1918}
1919
1920fn preview_from_initial_history(initial_history: &InitialHistory) -> String {
1921 match initial_history {
1922 InitialHistory::New => "New conversation".to_string(),
1923 InitialHistory::Forked(items) => truncate_preview(
1924 &items
1925 .first()
1926 .map(Value::to_string)
1927 .unwrap_or_else(|| "Forked conversation".to_string()),
1928 ),
1929 InitialHistory::Resumed { history, .. } => truncate_preview(
1930 &history
1931 .first()
1932 .map(Value::to_string)
1933 .unwrap_or_else(|| "Resumed conversation".to_string()),
1934 ),
1935 }
1936}
1937
1938fn permission_path_for_call(call: &ToolCall) -> Option<String> {
1939 match &call.payload {
1940 ToolPayload::Function { arguments } => serde_json::from_str::<Value>(arguments)
1941 .ok()
1942 .and_then(|value| {
1943 value
1944 .get("path")
1945 .and_then(Value::as_str)
1946 .map(str::to_string)
1947 }),
1948 ToolPayload::Mcp { raw_arguments, .. } => raw_arguments
1949 .get("path")
1950 .and_then(Value::as_str)
1951 .map(str::to_string),
1952 ToolPayload::Custom { .. } | ToolPayload::LocalShell { .. } => None,
1953 }
1954}
1955
1956fn truncate_preview(value: &str) -> String {
1957 value.chars().take(120).collect()
1958}
1959
1960fn to_protocol_thread(thread: ThreadMetadata) -> Thread {
1961 Thread {
1962 id: thread.id,
1963 preview: thread.preview,
1964 ephemeral: thread.ephemeral,
1965 model_provider: thread.model_provider,
1966 created_at: thread.created_at,
1967 updated_at: thread.updated_at,
1968 status: match thread.status {
1969 PersistedThreadStatus::Running => ThreadStatus::Running,
1970 PersistedThreadStatus::Idle => ThreadStatus::Idle,
1971 PersistedThreadStatus::Completed => ThreadStatus::Completed,
1972 PersistedThreadStatus::Failed => ThreadStatus::Failed,
1973 PersistedThreadStatus::Paused => ThreadStatus::Paused,
1974 PersistedThreadStatus::Archived => ThreadStatus::Archived,
1975 },
1976 path: thread.path,
1977 cwd: thread.cwd,
1978 cli_version: thread.cli_version,
1979 source: match thread.source {
1980 SessionSource::Interactive => codewhale_protocol::SessionSource::Interactive,
1981 SessionSource::Resume => codewhale_protocol::SessionSource::Resume,
1982 SessionSource::Fork => codewhale_protocol::SessionSource::Fork,
1983 SessionSource::Api => codewhale_protocol::SessionSource::Api,
1984 SessionSource::Unknown => codewhale_protocol::SessionSource::Unknown,
1985 },
1986 name: thread.name,
1987 }
1988}
1989
1990fn to_protocol_goal(goal: ThreadGoalRecord) -> ThreadGoal {
1991 ThreadGoal {
1992 thread_id: goal.thread_id,
1993 goal_id: goal.goal_id,
1994 objective: goal.objective,
1995 status: to_protocol_goal_status(goal.status),
1996 token_budget: goal.token_budget,
1997 tokens_used: goal.tokens_used,
1998 time_used_seconds: goal.time_used_seconds,
1999 continuation_count: goal.continuation_count,
2000 created_at: goal.created_at,
2001 updated_at: goal.updated_at,
2002 }
2003}
2004
2005fn to_protocol_goal_status(status: PersistedThreadGoalStatus) -> ThreadGoalStatus {
2006 match status {
2007 PersistedThreadGoalStatus::Active => ThreadGoalStatus::Active,
2008 PersistedThreadGoalStatus::Paused => ThreadGoalStatus::Paused,
2009 PersistedThreadGoalStatus::Blocked => ThreadGoalStatus::Blocked,
2010 PersistedThreadGoalStatus::UsageLimited => ThreadGoalStatus::UsageLimited,
2011 PersistedThreadGoalStatus::BudgetLimited => ThreadGoalStatus::BudgetLimited,
2012 PersistedThreadGoalStatus::Complete => ThreadGoalStatus::Complete,
2013 }
2014}
2015
2016fn to_persisted_status(status: &ThreadStatus) -> PersistedThreadStatus {
2017 match status {
2018 ThreadStatus::Running => PersistedThreadStatus::Running,
2019 ThreadStatus::Idle => PersistedThreadStatus::Idle,
2020 ThreadStatus::Completed => PersistedThreadStatus::Completed,
2021 ThreadStatus::Failed => PersistedThreadStatus::Failed,
2022 ThreadStatus::Paused => PersistedThreadStatus::Paused,
2023 ThreadStatus::Archived => PersistedThreadStatus::Archived,
2024 }
2025}
2026
2027fn to_persisted_source(source: &codewhale_protocol::SessionSource) -> SessionSource {
2028 match source {
2029 codewhale_protocol::SessionSource::Interactive => SessionSource::Interactive,
2030 codewhale_protocol::SessionSource::Resume => SessionSource::Resume,
2031 codewhale_protocol::SessionSource::Fork => SessionSource::Fork,
2032 codewhale_protocol::SessionSource::Api => SessionSource::Api,
2033 codewhale_protocol::SessionSource::Unknown => SessionSource::Unknown,
2034 }
2035}
2036
2037fn approval_request_frame(
2038 requirement: &ExecApprovalRequirement,
2039 matched_rule: Option<&str>,
2040 call_id: String,
2041 approval_id: String,
2042 turn_id: String,
2043 command: String,
2044 cwd: String,
2045) -> Option<EventFrame> {
2046 let ExecApprovalRequirement::NeedsApproval {
2047 reason,
2048 proposed_execpolicy_amendment,
2049 proposed_network_policy_amendments,
2050 } = requirement
2051 else {
2052 return None;
2053 };
2054
2055 let mut available_decisions = vec![
2056 ReviewDecision::Approved,
2057 ReviewDecision::ApprovedForSession,
2058 ReviewDecision::Denied,
2059 ReviewDecision::Abort,
2060 ];
2061 if proposed_execpolicy_amendment
2062 .as_ref()
2063 .is_some_and(|amendment| !amendment.prefixes.is_empty())
2064 {
2065 available_decisions.push(ReviewDecision::ApprovedExecpolicyAmendment);
2066 }
2067 available_decisions.extend(proposed_network_policy_amendments.iter().cloned().map(
2068 |amendment| ReviewDecision::NetworkPolicyAmendment {
2069 host: amendment.host,
2070 action: amendment.action,
2071 },
2072 ));
2073
2074 Some(EventFrame::ExecApprovalRequest {
2075 request: ExecApprovalRequestEvent {
2076 call_id,
2077 approval_id,
2078 turn_id,
2079 command,
2080 cwd,
2081 reason: reason.clone(),
2082 matched_rule: matched_rule.map(|rule| rule.to_string().into_boxed_str()),
2083 network_approval_context: None,
2084 proposed_execpolicy_amendment: proposed_execpolicy_amendment
2085 .as_ref()
2086 .map(|amendment| amendment.prefixes.clone())
2087 .unwrap_or_default(),
2088 proposed_network_policy_amendments: proposed_network_policy_amendments.clone(),
2089 additional_permissions: Vec::new(),
2090 available_decisions,
2091 },
2092 })
2093}
2094
2095fn user_input_request_frame(
2103 call_id: String,
2104 turn_id: String,
2105 request_id: String,
2106 arguments: &str,
2107) -> Option<EventFrame> {
2108 let parsed: Value = serde_json::from_str(arguments).ok()?;
2109 let questions = parsed.get("questions").cloned().filter(Value::is_array)?;
2113 let request = UserInputRequestEvent {
2114 call_id,
2115 turn_id,
2116 request_id,
2117 questions: serde_json::from_value(questions).ok()?,
2118 };
2119 Some(EventFrame::UserInputRequest { request })
2120}
2121
2122fn approval_requirement_payload(requirement: &ExecApprovalRequirement) -> Value {
2123 match requirement {
2124 ExecApprovalRequirement::Skip {
2125 bypass_sandbox,
2126 proposed_execpolicy_amendment,
2127 } => json!({
2128 "type": "skip",
2129 "bypass_sandbox": bypass_sandbox,
2130 "reason": requirement.reason(),
2131 "proposed_execpolicy_amendment": proposed_execpolicy_amendment
2132 .as_ref()
2133 .map(|amendment| amendment.prefixes.clone())
2134 .unwrap_or_default()
2135 }),
2136 ExecApprovalRequirement::NeedsApproval {
2137 reason,
2138 proposed_execpolicy_amendment,
2139 proposed_network_policy_amendments,
2140 } => json!({
2141 "type": "needs_approval",
2142 "reason": reason,
2143 "proposed_execpolicy_amendment": proposed_execpolicy_amendment
2144 .as_ref()
2145 .map(|amendment| amendment.prefixes.clone())
2146 .unwrap_or_default(),
2147 "proposed_network_policy_amendments": proposed_network_policy_amendments
2148 }),
2149 ExecApprovalRequirement::Forbidden { reason } => json!({
2150 "type": "forbidden",
2151 "reason": reason
2152 }),
2153 }
2154}
2155
2156fn policy_precheck_payload(
2157 decision: &ExecPolicyDecision,
2158 command: &str,
2159 cwd: &str,
2160 execution_kind: &str,
2161) -> Value {
2162 json!({
2163 "execution_kind": execution_kind,
2164 "command": command,
2165 "cwd": cwd,
2166 "allow": decision.allow,
2167 "requires_approval": decision.requires_approval,
2168 "matched_rule": decision.matched_rule.clone(),
2169 "phase": decision.requirement.phase(),
2170 "reason": decision.reason(),
2171 "requirement": approval_requirement_payload(&decision.requirement)
2172 })
2173}
2174
2175fn tool_payload_value(payload: &ToolPayload) -> Value {
2176 serde_json::to_value(payload).unwrap_or_else(
2177 |_| json!({"type":"serialization_error","message":"tool payload unavailable"}),
2178 )
2179}
2180
2181fn tool_output_value(output: &codewhale_protocol::ToolOutput) -> Value {
2182 serde_json::to_value(output).unwrap_or_else(
2183 |_| json!({"type":"serialization_error","message":"tool output unavailable"}),
2184 )
2185}
2186
2187fn event_frame_payload(frame: &EventFrame) -> Value {
2188 serde_json::to_value(frame)
2189 .unwrap_or_else(|_| json!({"event":"error","message":"failed to encode event frame"}))
2190}
2191
2192const REQUEST_USER_INPUT_TOOL_NAME: &str = "request_user_input";
2198
2199fn json_optional_string(value: &Value) -> Option<String> {
2200 if value.is_null() {
2201 None
2202 } else {
2203 value.as_str().map(ToString::to_string)
2204 }
2205}
2206
2207fn parse_retry_metadata(value: Option<&Value>) -> JobRetryMetadata {
2208 let Some(value) = value else {
2209 return JobRetryMetadata::default();
2210 };
2211 JobRetryMetadata {
2212 attempt: value
2213 .get("attempt")
2214 .and_then(Value::as_u64)
2215 .unwrap_or(0)
2216 .min(u32::MAX as u64) as u32,
2217 max_attempts: value
2218 .get("max_attempts")
2219 .and_then(Value::as_u64)
2220 .unwrap_or(DEFAULT_JOB_MAX_ATTEMPTS as u64)
2221 .min(u32::MAX as u64) as u32,
2222 backoff_base_ms: value
2223 .get("backoff_base_ms")
2224 .and_then(Value::as_u64)
2225 .unwrap_or(DEFAULT_JOB_BACKOFF_BASE_MS),
2226 next_backoff_ms: value
2227 .get("next_backoff_ms")
2228 .and_then(Value::as_u64)
2229 .unwrap_or(0),
2230 next_retry_at: value.get("next_retry_at").and_then(Value::as_i64),
2231 }
2232}
2233
2234fn parse_history_entry(value: &Value) -> Option<JobHistoryEntry> {
2235 let status = value
2236 .get("status")
2237 .and_then(Value::as_str)
2238 .and_then(job_status_from_str)?;
2239 Some(JobHistoryEntry {
2240 at: value.get("at").and_then(Value::as_i64).unwrap_or(0),
2241 phase: value
2242 .get("phase")
2243 .and_then(Value::as_str)
2244 .unwrap_or("unknown")
2245 .to_string(),
2246 status,
2247 progress: value
2248 .get("progress")
2249 .and_then(Value::as_u64)
2250 .map(|v| v.min(u8::MAX as u64) as u8),
2251 detail: value.get("detail").and_then(json_optional_string),
2252 retry: parse_retry_metadata(value.get("retry")),
2253 })
2254}
2255
2256fn job_status_to_str(status: JobStatus) -> &'static str {
2257 match status {
2258 JobStatus::Queued => "queued",
2259 JobStatus::Running => "running",
2260 JobStatus::Paused => "paused",
2261 JobStatus::Completed => "completed",
2262 JobStatus::Failed => "failed",
2263 JobStatus::Cancelled => "cancelled",
2264 }
2265}
2266
2267fn job_status_from_str(value: &str) -> Option<JobStatus> {
2268 match value {
2269 "queued" => Some(JobStatus::Queued),
2270 "running" => Some(JobStatus::Running),
2271 "paused" => Some(JobStatus::Paused),
2272 "completed" => Some(JobStatus::Completed),
2273 "failed" => Some(JobStatus::Failed),
2274 "cancelled" => Some(JobStatus::Cancelled),
2275 _ => None,
2276 }
2277}
2278
2279fn job_retry_to_value(retry: &JobRetryMetadata) -> Value {
2280 json!({
2281 "attempt": retry.attempt,
2282 "max_attempts": retry.max_attempts,
2283 "backoff_base_ms": retry.backoff_base_ms,
2284 "next_backoff_ms": retry.next_backoff_ms,
2285 "next_retry_at": retry.next_retry_at
2286 })
2287}
2288
2289fn job_history_to_value(entry: &JobHistoryEntry) -> Value {
2290 json!({
2291 "at": entry.at,
2292 "phase": entry.phase.clone(),
2293 "status": job_status_to_str(entry.status),
2294 "progress": entry.progress,
2295 "detail": entry.detail.clone(),
2296 "retry": job_retry_to_value(&entry.retry)
2297 })
2298}
2299
2300fn runtime_status_to_job_state(status: JobStatus) -> JobStateStatus {
2301 match status {
2302 JobStatus::Queued => JobStateStatus::Queued,
2303 JobStatus::Running => JobStateStatus::Running,
2304 JobStatus::Paused => JobStateStatus::Paused,
2305 JobStatus::Completed => JobStateStatus::Completed,
2306 JobStatus::Failed => JobStateStatus::Failed,
2307 JobStatus::Cancelled => JobStateStatus::Cancelled,
2308 }
2309}
2310
2311fn job_state_status_to_runtime(status: JobStateStatus) -> JobStatus {
2312 match status {
2313 JobStateStatus::Queued => JobStatus::Queued,
2314 JobStateStatus::Running => JobStatus::Running,
2315 JobStateStatus::Paused => JobStatus::Paused,
2316 JobStateStatus::Completed => JobStatus::Completed,
2317 JobStateStatus::Failed => JobStatus::Failed,
2318 JobStateStatus::Cancelled => JobStatus::Cancelled,
2319 }
2320}
2321
2322#[cfg(test)]
2323mod tests {
2324 use super::*;
2325 use codewhale_protocol::ThreadResumeParams;
2326 use codewhale_tools::ToolCallSource;
2327
2328 fn temp_core_state(name: &str) -> StateStore {
2329 let dir =
2330 std::env::temp_dir().join(format!("codewhale-core-{name}-{}", Uuid::new_v4().simple()));
2331 std::fs::create_dir_all(&dir).expect("create temp state dir");
2332 StateStore::open(Some(dir.join("state.db"))).expect("open state store")
2333 }
2334
2335 fn test_thread_metadata(id: &str) -> ThreadMetadata {
2336 ThreadMetadata {
2337 id: id.to_string(),
2338 rollout_path: None,
2339 preview: "test thread".to_string(),
2340 ephemeral: false,
2341 model_provider: "deepseek".to_string(),
2342 created_at: 10,
2343 updated_at: 10,
2344 status: PersistedThreadStatus::Running,
2345 path: None,
2346 cwd: PathBuf::from("/tmp/codewhale"),
2347 cli_version: "0.0.0-test".to_string(),
2348 source: SessionSource::Interactive,
2349 name: None,
2350 sandbox_policy: None,
2351 approval_mode: None,
2352 archived: false,
2353 archived_at: None,
2354 git_sha: None,
2355 git_branch: None,
2356 git_origin_url: None,
2357 memory_mode: None,
2358 current_leaf_id: None,
2359 }
2360 }
2361
2362 #[test]
2365 fn permission_path_for_call_extracts_function_path_argument() {
2366 let call = ToolCall {
2367 name: "read_file".to_string(),
2368 payload: ToolPayload::Function {
2369 arguments: json!({ "path": "README.md" }).to_string(),
2370 },
2371 source: ToolCallSource::Direct,
2372 raw_tool_call_id: None,
2373 };
2374
2375 assert_eq!(
2376 permission_path_for_call(&call).as_deref(),
2377 Some("README.md")
2378 );
2379 }
2380
2381 #[test]
2382 fn permission_path_for_call_extracts_mcp_path_argument() {
2383 let call = ToolCall {
2384 name: "mcp_fs_read".to_string(),
2385 payload: ToolPayload::Mcp {
2386 server: "fs".to_string(),
2387 tool: "read".to_string(),
2388 raw_arguments: json!({ "path": "secrets/token.txt" }),
2389 raw_tool_call_id: None,
2390 },
2391 source: ToolCallSource::Direct,
2392 raw_tool_call_id: None,
2393 };
2394
2395 assert_eq!(
2396 permission_path_for_call(&call).as_deref(),
2397 Some("secrets/token.txt")
2398 );
2399 }
2400
2401 #[test]
2402 fn permission_path_for_call_ignores_shell_payload() {
2403 let call = ToolCall {
2404 name: "exec_shell".to_string(),
2405 payload: ToolPayload::LocalShell {
2406 params: codewhale_protocol::LocalShellParams {
2407 command: "cargo test".to_string(),
2408 cwd: None,
2409 timeout_ms: None,
2410 },
2411 },
2412 source: ToolCallSource::Direct,
2413 raw_tool_call_id: None,
2414 };
2415
2416 assert_eq!(permission_path_for_call(&call), None);
2417 }
2418
2419 #[test]
2420 fn thread_goal_progress_accumulates_durable_accounting() {
2421 let store = temp_core_state("thread-goal-progress");
2422 store
2423 .upsert_thread(&test_thread_metadata("thread-1"))
2424 .expect("upsert thread");
2425 let mut manager = ThreadManager::new(store);
2426 manager
2427 .set_thread_goal(&ThreadGoalSetParams {
2428 thread_id: "thread-1".to_string(),
2429 objective: "Carry the goal across turns".to_string(),
2430 token_budget: Some(2_000),
2431 })
2432 .expect("set goal")
2433 .expect("goal exists");
2434
2435 let updated = manager
2436 .record_thread_goal_progress(&ThreadGoalProgressParams {
2437 thread_id: "thread-1".to_string(),
2438 token_delta: 750,
2439 time_delta_seconds: 12,
2440 record_continuation: true,
2441 })
2442 .expect("record progress")
2443 .expect("goal exists");
2444
2445 assert_eq!(updated.tokens_used, 750);
2446 assert_eq!(updated.time_used_seconds, 12);
2447 assert_eq!(updated.continuation_count, 1);
2448
2449 let persisted = manager
2450 .get_thread_goal(&ThreadGoalGetParams {
2451 thread_id: "thread-1".to_string(),
2452 })
2453 .expect("read goal")
2454 .expect("goal exists");
2455 assert_eq!(persisted.tokens_used, 750);
2456 assert_eq!(persisted.time_used_seconds, 12);
2457 assert_eq!(persisted.continuation_count, 1);
2458 }
2459
2460 #[test]
2461 fn approval_request_frame_includes_matched_rule() {
2462 let requirement = ExecApprovalRequirement::NeedsApproval {
2463 reason: "Typed ask rule 'tool=exec_shell command=cargo test' requires approval."
2464 .to_string(),
2465 proposed_execpolicy_amendment: None,
2466 proposed_network_policy_amendments: Vec::new(),
2467 };
2468
2469 let frame = approval_request_frame(
2470 &requirement,
2471 Some("tool=exec_shell command=cargo test"),
2472 "call-1".to_string(),
2473 "approval-1".to_string(),
2474 "turn-1".to_string(),
2475 "cargo test --workspace".to_string(),
2476 "/repo".to_string(),
2477 )
2478 .expect("approval frame");
2479
2480 let EventFrame::ExecApprovalRequest { request } = frame else {
2481 panic!("expected exec approval request frame");
2482 };
2483 assert_eq!(
2484 request.matched_rule.as_deref(),
2485 Some("tool=exec_shell command=cargo test")
2486 );
2487 assert_eq!(request.reason, requirement.reason());
2488 }
2489
2490 #[test]
2491 fn user_input_request_frame_lifts_questions_from_arguments() {
2492 let arguments = r#"{"questions":[{"header":"Scope","id":"scope","question":"Which?","options":[{"label":"A","description":"a"},{"label":"B","description":"b"}],"allow_free_text":true}]}"#;
2496 let frame = user_input_request_frame(
2497 "call-1".to_string(),
2498 "turn-1".to_string(),
2499 "ui-1".to_string(),
2500 arguments,
2501 )
2502 .expect("user input frame");
2503
2504 let EventFrame::UserInputRequest { request } = frame else {
2505 panic!("expected user_input_request frame");
2506 };
2507 assert_eq!(request.call_id, "call-1");
2508 assert_eq!(request.turn_id, "turn-1");
2509 assert_eq!(request.request_id, "ui-1");
2510 assert_eq!(request.questions.len(), 1);
2511 assert_eq!(request.questions[0].id, "scope");
2512 assert!(request.questions[0].allow_free_text);
2513 assert!(!request.questions[0].multi_select);
2515 assert_eq!(request.questions[0].options.len(), 2);
2516 }
2517
2518 #[test]
2519 fn user_input_request_frame_returns_none_on_invalid_arguments() {
2520 let frame = user_input_request_frame(
2523 "call-1".to_string(),
2524 "turn-1".to_string(),
2525 "ui-1".to_string(),
2526 "not json",
2527 );
2528 assert!(frame.is_none());
2529
2530 let frame = user_input_request_frame(
2532 "call-1".to_string(),
2533 "turn-1".to_string(),
2534 "ui-1".to_string(),
2535 r#"{"foo":"bar"}"#,
2536 );
2537 assert!(frame.is_none());
2538 }
2539
2540 #[test]
2541 fn enqueue_creates_queued_job_with_zero_progress() {
2542 let mut jm = JobManager::default();
2543 let job = jm.enqueue("build");
2544 assert_eq!(job.name, "build");
2545 assert_eq!(job.status, JobStatus::Queued);
2546 assert_eq!(job.progress, Some(0));
2547 assert!(job.detail.is_none());
2548 assert_eq!(job.history.len(), 1);
2549 assert_eq!(job.history[0].phase, "created");
2550 }
2551
2552 #[test]
2553 fn set_running_transitions_from_queued() {
2554 let mut jm = JobManager::default();
2555 let job = jm.enqueue("deploy");
2556 let id = job.id.clone();
2557 jm.set_running(&id);
2558 let jobs = jm.list();
2559 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2560 assert_eq!(updated.status, JobStatus::Running);
2561 assert_eq!(updated.history.last().unwrap().phase, "running");
2562 }
2563
2564 #[test]
2565 fn update_progress_clamps_to_100() {
2566 let mut jm = JobManager::default();
2567 let job = jm.enqueue("task");
2568 let id = job.id.clone();
2569 jm.update_progress(&id, 150, Some("over".to_string()));
2570 let jobs = jm.list();
2571 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2572 assert_eq!(updated.progress, Some(100));
2573 }
2574
2575 #[test]
2576 fn complete_sets_progress_to_100() {
2577 let mut jm = JobManager::default();
2578 let job = jm.enqueue("task");
2579 let id = job.id.clone();
2580 jm.set_running(&id);
2581 jm.complete(&id);
2582 let jobs = jm.list();
2583 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2584 assert_eq!(updated.status, JobStatus::Completed);
2585 assert_eq!(updated.progress, Some(100));
2586 }
2587
2588 #[test]
2589 fn fail_increments_attempt_and_sets_backoff() {
2590 let mut jm = JobManager::default();
2591 let job = jm.enqueue("fragile");
2592 let id = job.id.clone();
2593 jm.set_running(&id);
2594 jm.fail(&id, "crashed");
2595 let jobs = jm.list();
2596 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2597 assert_eq!(updated.status, JobStatus::Failed);
2598 assert_eq!(updated.retry.attempt, 1);
2599 assert!(updated.retry.next_backoff_ms > 0);
2600 assert!(updated.retry.next_retry_at.is_some());
2601 assert_eq!(updated.detail.as_deref(), Some("crashed"));
2602 }
2603
2604 #[test]
2605 fn fail_clears_retry_after_max_attempts() {
2606 let mut jm = JobManager::default();
2607 let job = jm.enqueue("fragile");
2608 let id = job.id.clone();
2609 for _ in 0..=DEFAULT_JOB_MAX_ATTEMPTS {
2610 jm.set_running(&id);
2611 jm.fail(&id, "boom");
2612 }
2613 let jobs = jm.list();
2614 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2615 assert_eq!(updated.retry.attempt, DEFAULT_JOB_MAX_ATTEMPTS);
2616 assert_eq!(updated.retry.next_backoff_ms, 0);
2617 assert!(updated.retry.next_retry_at.is_none());
2618 }
2619
2620 #[test]
2621 fn cancel_sets_status_and_clears_retry() {
2622 let mut jm = JobManager::default();
2623 let job = jm.enqueue("task");
2624 let id = job.id.clone();
2625 jm.cancel(&id);
2626 let jobs = jm.list();
2627 let updated = jobs.iter().find(|j| j.id == id).unwrap();
2628 assert_eq!(updated.status, JobStatus::Cancelled);
2629 assert_eq!(updated.retry.next_backoff_ms, 0);
2630 }
2631
2632 #[test]
2633 fn pause_and_resume_round_trip() {
2634 let mut jm = JobManager::default();
2635 let job = jm.enqueue("task");
2636 let id = job.id.clone();
2637 jm.set_running(&id);
2638 jm.pause(&id, Some("waiting".to_string()));
2639 let jobs = jm.list();
2640 let paused = jobs.iter().find(|j| j.id == id).unwrap();
2641 assert_eq!(paused.status, JobStatus::Paused);
2642 assert_eq!(paused.detail.as_deref(), Some("waiting"));
2643
2644 jm.resume(&id, None);
2645 let jobs = jm.list();
2646 let resumed = jobs.iter().find(|j| j.id == id).unwrap();
2647 assert_eq!(resumed.status, JobStatus::Running);
2648 assert_eq!(resumed.history.last().unwrap().phase, "resumed");
2649 }
2650
2651 #[test]
2652 fn list_returns_jobs_sorted_by_updated_at_desc() {
2653 let mut jm = JobManager::default();
2654 jm.enqueue("first");
2655 jm.enqueue("second");
2656 jm.enqueue("third");
2657 let jobs = jm.list();
2658 assert_eq!(jobs.len(), 3);
2659 for window in jobs.windows(2) {
2660 assert!(window[0].updated_at >= window[1].updated_at);
2661 }
2662 }
2663
2664 #[test]
2665 fn history_returns_entries_for_existing_job() {
2666 let mut jm = JobManager::default();
2667 let job = jm.enqueue("task");
2668 let id = job.id.clone();
2669 jm.set_running(&id);
2670 jm.complete(&id);
2671 let history = jm.history(&id);
2672 assert_eq!(history.len(), 3); assert_eq!(history[0].phase, "created");
2674 assert_eq!(history[1].phase, "running");
2675 assert_eq!(history[2].phase, "completed");
2676 }
2677
2678 #[test]
2679 fn history_returns_empty_for_unknown_job() {
2680 let jm = JobManager::default();
2681 assert!(jm.history("nonexistent").is_empty());
2682 }
2683
2684 #[test]
2685 fn resume_pending_requeues_running_and_queued() {
2686 let mut jm = JobManager::default();
2687 let _j1 = jm.enqueue("queued_task");
2688 let j2 = jm.enqueue("running_task");
2689 let j3 = jm.enqueue("completed_task");
2690 let id2 = j2.id.clone();
2691 let id3 = j3.id.clone();
2692 jm.set_running(&id2);
2693 jm.set_running(&id3);
2694 jm.complete(&id3);
2695
2696 let resumed = jm.resume_pending();
2697 assert_eq!(resumed.len(), 2);
2698 for job in &resumed {
2699 assert_eq!(job.status, JobStatus::Queued);
2700 }
2701 }
2702
2703 #[test]
2706 fn deterministic_backoff_zero_on_first_attempt() {
2707 let retry = JobRetryMetadata {
2708 attempt: 0,
2709 ..Default::default()
2710 };
2711 assert_eq!(JobManager::deterministic_backoff_ms(&retry), 0);
2712 }
2713
2714 #[test]
2715 fn deterministic_backoff_exponential_growth() {
2716 let base = DEFAULT_JOB_BACKOFF_BASE_MS;
2717 for attempt in 1..=5 {
2718 let retry = JobRetryMetadata {
2719 attempt,
2720 backoff_base_ms: base,
2721 ..Default::default()
2722 };
2723 let expected = base * 2u64.pow(attempt.saturating_sub(1).min(20));
2724 assert_eq!(
2725 JobManager::deterministic_backoff_ms(&retry),
2726 expected,
2727 "attempt {attempt}"
2728 );
2729 }
2730 }
2731
2732 #[test]
2733 fn deterministic_backoff_saturates_at_high_exponent() {
2734 let retry = JobRetryMetadata {
2735 attempt: 63,
2736 backoff_base_ms: 1000,
2737 ..Default::default()
2738 };
2739 let _ = JobManager::deterministic_backoff_ms(&retry);
2741 }
2742
2743 #[test]
2746 fn push_history_truncates_beyond_max() {
2747 let mut jm = JobManager::default();
2748 let job = jm.enqueue("task");
2749 let id = job.id.clone();
2750 for i in 0..(MAX_JOB_HISTORY_ENTRIES + 20) {
2752 jm.update_progress(&id, (i % 100) as u8, Some(format!("step {i}")));
2753 }
2754 let history = jm.history(&id);
2755 assert_eq!(history.len(), MAX_JOB_HISTORY_ENTRIES);
2756 }
2757
2758 #[test]
2761 fn encode_and_parse_persisted_detail_round_trip() {
2762 let mut jm = JobManager::default();
2763 let job = jm.enqueue("task");
2764 let id = job.id.clone();
2765 jm.set_running(&id);
2766 jm.fail(&id, "oops");
2767 let job = jm.list().into_iter().find(|j| j.id == id).unwrap();
2768
2769 let encoded = JobManager::encode_persisted_detail(&job).unwrap().unwrap();
2770 let parsed = JobManager::parse_persisted_detail(Some(&encoded)).unwrap();
2771
2772 assert_eq!(parsed.status, job.status);
2773 assert_eq!(parsed.detail, job.detail);
2774 assert_eq!(parsed.retry.attempt, job.retry.attempt);
2775 assert_eq!(parsed.history.len(), job.history.len());
2776 }
2777
2778 #[test]
2779 fn parse_persisted_detail_returns_none_for_none_input() {
2780 assert!(JobManager::parse_persisted_detail(None).is_none());
2781 }
2782
2783 #[test]
2784 fn parse_persisted_detail_returns_none_for_invalid_json() {
2785 assert!(JobManager::parse_persisted_detail(Some("not json")).is_none());
2786 }
2787
2788 #[test]
2791 fn job_status_round_trip_str() {
2792 let statuses = [
2793 JobStatus::Queued,
2794 JobStatus::Running,
2795 JobStatus::Paused,
2796 JobStatus::Completed,
2797 JobStatus::Failed,
2798 JobStatus::Cancelled,
2799 ];
2800 for status in &statuses {
2801 let s = job_status_to_str(*status);
2802 let parsed = job_status_from_str(s);
2803 assert_eq!(parsed, Some(*status), "round-trip failed for {s:?}");
2804 }
2805 }
2806
2807 #[test]
2808 fn job_status_from_str_returns_none_for_unknown() {
2809 assert_eq!(job_status_from_str("unknown"), None);
2810 assert_eq!(job_status_from_str(""), None);
2811 }
2812
2813 #[test]
2814 fn truncate_preview_limits_to_120_chars() {
2815 let long = "a".repeat(200);
2816 let truncated = truncate_preview(&long);
2817 assert_eq!(truncated.len(), 120);
2818 }
2819
2820 #[test]
2821 fn truncate_preview_preserves_short_strings() {
2822 let short = "hello";
2823 assert_eq!(truncate_preview(short), "hello");
2824 }
2825
2826 #[test]
2827 fn runtime_status_to_job_state_maps_correctly() {
2828 assert_eq!(
2829 runtime_status_to_job_state(JobStatus::Queued),
2830 JobStateStatus::Queued
2831 );
2832 assert_eq!(
2833 runtime_status_to_job_state(JobStatus::Running),
2834 JobStateStatus::Running
2835 );
2836 assert_eq!(
2837 runtime_status_to_job_state(JobStatus::Paused),
2838 JobStateStatus::Paused
2839 );
2840 assert_eq!(
2841 runtime_status_to_job_state(JobStatus::Completed),
2842 JobStateStatus::Completed
2843 );
2844 assert_eq!(
2845 runtime_status_to_job_state(JobStatus::Failed),
2846 JobStateStatus::Failed
2847 );
2848 assert_eq!(
2849 runtime_status_to_job_state(JobStatus::Cancelled),
2850 JobStateStatus::Cancelled
2851 );
2852 }
2853
2854 #[test]
2855 fn job_state_status_to_runtime_maps_correctly() {
2856 assert_eq!(
2857 job_state_status_to_runtime(JobStateStatus::Queued),
2858 JobStatus::Queued
2859 );
2860 assert_eq!(
2861 job_state_status_to_runtime(JobStateStatus::Running),
2862 JobStatus::Running
2863 );
2864 assert_eq!(
2865 job_state_status_to_runtime(JobStateStatus::Paused),
2866 JobStatus::Paused
2867 );
2868 assert_eq!(
2869 job_state_status_to_runtime(JobStateStatus::Completed),
2870 JobStatus::Completed
2871 );
2872 assert_eq!(
2873 job_state_status_to_runtime(JobStateStatus::Failed),
2874 JobStatus::Failed
2875 );
2876 assert_eq!(
2877 job_state_status_to_runtime(JobStateStatus::Cancelled),
2878 JobStatus::Cancelled
2879 );
2880 }
2881
2882 #[test]
2883 fn preview_from_initial_history_new() {
2884 let preview = preview_from_initial_history(&InitialHistory::New);
2885 assert_eq!(preview, "New conversation");
2886 }
2887
2888 #[test]
2889 fn preview_from_initial_history_forked() {
2890 let preview = preview_from_initial_history(&InitialHistory::Forked(vec![json!("hello")]));
2891 assert!(preview.contains("hello"));
2892 }
2893
2894 #[test]
2895 fn preview_from_initial_history_resumed() {
2896 let preview = preview_from_initial_history(&InitialHistory::Resumed {
2897 conversation_id: "test".to_string(),
2898 history: vec![json!("world")],
2899 rollout_path: PathBuf::from("/tmp/test"),
2900 });
2901 assert!(preview.contains("world"));
2902 }
2903
2904 #[test]
2905 fn json_optional_string_handles_null() {
2906 assert!(json_optional_string(&Value::Null).is_none());
2907 }
2908
2909 #[test]
2910 fn json_optional_string_handles_string() {
2911 assert_eq!(
2912 json_optional_string(&Value::String("hello".to_string())),
2913 Some("hello".to_string())
2914 );
2915 }
2916
2917 #[test]
2918 fn json_optional_string_handles_non_string() {
2919 assert!(json_optional_string(&json!(42)).is_none());
2920 }
2921
2922 #[test]
2923 fn parse_retry_metadata_returns_default_for_none() {
2924 let retry = parse_retry_metadata(None);
2925 assert_eq!(retry.attempt, 0);
2926 assert_eq!(retry.max_attempts, DEFAULT_JOB_MAX_ATTEMPTS);
2927 assert_eq!(retry.backoff_base_ms, DEFAULT_JOB_BACKOFF_BASE_MS);
2928 }
2929
2930 #[test]
2931 fn parse_retry_metadata_parses_fields() {
2932 let value = json!({
2933 "attempt": 2,
2934 "max_attempts": 5,
2935 "backoff_base_ms": 1000,
2936 "next_backoff_ms": 2000,
2937 "next_retry_at": 1234567890i64
2938 });
2939 let retry = parse_retry_metadata(Some(&value));
2940 assert_eq!(retry.attempt, 2);
2941 assert_eq!(retry.max_attempts, 5);
2942 assert_eq!(retry.backoff_base_ms, 1000);
2943 assert_eq!(retry.next_backoff_ms, 2000);
2944 assert_eq!(retry.next_retry_at, Some(1234567890));
2945 }
2946
2947 #[test]
2948 fn parse_history_entry_returns_none_without_status() {
2949 let value = json!({"at": 1, "phase": "test"});
2950 assert!(parse_history_entry(&value).is_none());
2951 }
2952
2953 #[test]
2954 fn parse_history_entry_parses_valid_entry() {
2955 let value = json!({
2956 "at": 100,
2957 "phase": "running",
2958 "status": "running",
2959 "progress": 50,
2960 "detail": "working",
2961 "retry": {"attempt": 0, "max_attempts": 3, "backoff_base_ms": 500}
2962 });
2963 let entry = parse_history_entry(&value).unwrap();
2964 assert_eq!(entry.at, 100);
2965 assert_eq!(entry.phase, "running");
2966 assert_eq!(entry.status, JobStatus::Running);
2967 assert_eq!(entry.progress, Some(50));
2968 assert_eq!(entry.detail.as_deref(), Some("working"));
2969 }
2970
2971 #[test]
2972 fn paused_job_persists_as_paused_not_running() {
2973 let store = temp_core_state("paused-persist");
2974 let mut jm = JobManager::default();
2975 let job = jm.enqueue("task");
2976 let id = job.id.clone();
2977 jm.set_running(&id);
2978 jm.pause(&id, Some("waiting".to_string()));
2979 jm.persist_job(&store, &id).expect("persist paused job");
2980
2981 let persisted = store.list_jobs(Some(10)).expect("list jobs");
2982 let record = persisted.iter().find(|job| job.id == id).unwrap();
2983 assert_eq!(record.status, JobStateStatus::Paused);
2984
2985 let mut reloaded = JobManager::default();
2986 reloaded.load_from_store(&store).expect("reload jobs");
2987 let jobs = reloaded.list();
2988 let reloaded_job = jobs.iter().find(|job| job.id == id).unwrap();
2989 assert_eq!(reloaded_job.status, JobStatus::Paused);
2990 }
2991
2992 fn sample_job_record(status: JobStatus, detail: Option<&str>) -> JobRecord {
2995 JobRecord {
2996 id: "job-o1-1".to_string(),
2997 name: "sample".to_string(),
2998 status,
2999 progress: None,
3000 detail: detail.map(str::to_string),
3001 retry: JobRetryMetadata {
3002 attempt: 0,
3003 max_attempts: DEFAULT_JOB_MAX_ATTEMPTS,
3004 backoff_base_ms: DEFAULT_JOB_BACKOFF_BASE_MS,
3005 next_backoff_ms: 0,
3006 next_retry_at: None,
3007 },
3008 history: Vec::new(),
3009 created_at: 1_700_000_000,
3010 updated_at: 1_700_000_042,
3011 }
3012 }
3013
3014 #[test]
3015 fn job_record_to_agent_run_maps_non_terminal_states() {
3016 use codewhale_protocol::agent_run::RunState;
3017
3018 for (status, expected) in [
3019 (JobStatus::Queued, RunState::Queued),
3020 (JobStatus::Running, RunState::Running),
3021 (JobStatus::Paused, RunState::Paused),
3022 ] {
3023 let snapshot = job_record_to_agent_run(&sample_job_record(status, None));
3024 assert!(snapshot.is_coherent());
3025 assert_eq!(snapshot.run_id, "job-o1-1");
3026 assert_eq!(snapshot.parent, None);
3027 assert_eq!(
3028 snapshot.source,
3029 codewhale_protocol::agent_run::RunSource::CoreJob
3030 );
3031 assert_eq!(snapshot.state, expected);
3032 assert!(snapshot.terminal.is_none());
3033 assert!(snapshot.refs.is_empty());
3034 assert_eq!(
3035 snapshot.budget,
3036 codewhale_protocol::agent_run::BudgetSummary::default()
3037 );
3038 }
3039 }
3040
3041 #[test]
3042 fn job_record_to_agent_run_maps_terminal_states_without_fabricating_fields() {
3043 use codewhale_protocol::agent_run::{RunState, TerminalOutcome};
3044
3045 let cases = [
3046 (
3047 JobStatus::Completed,
3048 TerminalOutcome::Completed,
3049 Some("done"),
3050 ),
3051 (JobStatus::Failed, TerminalOutcome::Failed, Some("boom")),
3052 (JobStatus::Cancelled, TerminalOutcome::Cancelled, None),
3053 ];
3054
3055 for (status, outcome, detail) in cases {
3056 let snapshot = job_record_to_agent_run(&sample_job_record(status, detail));
3057 assert!(snapshot.is_coherent());
3058 assert_eq!(snapshot.state, RunState::Terminal);
3059 let terminal = snapshot.terminal.expect("terminal summary");
3060 assert_eq!(terminal.outcome, outcome);
3061 assert_eq!(terminal.ended_at_ms, Some(1_700_000_042_000));
3062 assert_eq!(terminal.detail, None);
3063 assert_eq!(
3064 snapshot.budget,
3065 codewhale_protocol::agent_run::BudgetSummary::default()
3066 );
3067 assert!(snapshot.refs.is_empty());
3068 assert_eq!(snapshot.parent, None);
3069 }
3070 }
3071
3072 #[test]
3073 fn job_record_to_agent_run_does_not_export_unclassified_detail() {
3074 let record = sample_job_record(JobStatus::Failed, Some("owner-private diagnostic"));
3075 let snapshot = job_record_to_agent_run(&record);
3076 let terminal = snapshot.terminal.as_ref().expect("terminal summary");
3077 assert_eq!(terminal.detail, None);
3078 let serialized = serde_json::to_string(&snapshot).expect("serialize snapshot");
3079 assert!(!serialized.contains("owner-private diagnostic"));
3080 }
3081
3082 #[test]
3083 fn job_record_to_agent_run_omits_ended_at_on_updated_at_overflow() {
3084 let mut record = sample_job_record(JobStatus::Completed, Some("ok"));
3085 record.updated_at = i64::MAX;
3086 let snapshot = job_record_to_agent_run(&record);
3087 assert!(snapshot.is_coherent());
3088 let terminal = snapshot.terminal.expect("terminal summary");
3089 assert_eq!(terminal.ended_at_ms, None);
3090 }
3091
3092 #[test]
3093 fn unarchive_thread_updates_running_threads_cache() {
3094 let store = temp_core_state("unarchive-cache");
3095 let mut manager = ThreadManager::new(store);
3096 let spawned = manager
3097 .spawn_thread_with_history(
3098 "deepseek".to_string(),
3099 PathBuf::from("/tmp/codewhale"),
3100 InitialHistory::New,
3101 true,
3102 )
3103 .expect("spawn thread");
3104 let thread_id = spawned.thread.id.clone();
3105 let resume_params = ThreadResumeParams {
3106 thread_id: thread_id.clone(),
3107 history: None,
3108 path: None,
3109 model: None,
3110 model_provider: None,
3111 cwd: None,
3112 approval_policy: None,
3113 sandbox: None,
3114 config: None,
3115 base_instructions: None,
3116 developer_instructions: None,
3117 personality: None,
3118 persist_extended_history: false,
3119 };
3120
3121 manager.archive_thread(&thread_id).expect("archive thread");
3122 let archived = manager
3123 .resume_thread_with_history(
3124 &resume_params,
3125 Path::new("/tmp/codewhale"),
3126 "deepseek".to_string(),
3127 )
3128 .expect("resume archived thread")
3129 .expect("thread in cache");
3130 assert_eq!(archived.thread.status, ThreadStatus::Archived);
3131
3132 manager
3133 .unarchive_thread(&thread_id)
3134 .expect("unarchive thread");
3135 let restored = manager
3136 .resume_thread_with_history(
3137 &resume_params,
3138 Path::new("/tmp/codewhale"),
3139 "deepseek".to_string(),
3140 )
3141 .expect("resume unarchived thread")
3142 .expect("thread in cache");
3143 assert_eq!(restored.thread.status, ThreadStatus::Idle);
3144 }
3145
3146 #[test]
3147 fn resume_with_history_does_not_reappend_persisted_messages() {
3148 let store = temp_core_state("resume-history-dedup");
3152 let mut manager = ThreadManager::new(store);
3153 let history = vec![
3154 json!({"type": "user_message", "message": "hello"}),
3155 json!({"type": "assistant_message", "message": "hi there"}),
3156 ];
3157 let spawned = manager
3158 .spawn_thread_with_history(
3159 "deepseek".to_string(),
3160 PathBuf::from("/tmp/codewhale"),
3161 InitialHistory::Forked(history.clone()),
3162 true,
3163 )
3164 .expect("spawn thread");
3165 let thread_id = spawned.thread.id.clone();
3166 let message_count = |manager: &ThreadManager| {
3167 manager
3168 .state_store()
3169 .list_messages(&thread_id, None)
3170 .expect("list messages")
3171 .len()
3172 };
3173 assert_eq!(message_count(&manager), 2);
3174
3175 let resume_params = ThreadResumeParams {
3176 thread_id: thread_id.clone(),
3177 history: Some(history.clone()),
3178 path: None,
3179 model: None,
3180 model_provider: None,
3181 cwd: None,
3182 approval_policy: None,
3183 sandbox: None,
3184 config: None,
3185 base_instructions: None,
3186 developer_instructions: None,
3187 personality: None,
3188 persist_extended_history: false,
3189 };
3190
3191 for _ in 0..2 {
3193 manager
3194 .resume_thread_with_history(
3195 &resume_params,
3196 Path::new("/tmp/codewhale"),
3197 "deepseek".to_string(),
3198 )
3199 .expect("resume thread")
3200 .expect("thread found");
3201 }
3202 assert_eq!(
3203 message_count(&manager),
3204 2,
3205 "resume re-appended messages already on the persisted chain"
3206 );
3207
3208 let mut extended = history.clone();
3210 extended.push(json!({"type": "user_message", "message": "something new"}));
3211 let resume_params = ThreadResumeParams {
3212 history: Some(extended),
3213 ..resume_params
3214 };
3215 manager
3216 .resume_thread_with_history(
3217 &resume_params,
3218 Path::new("/tmp/codewhale"),
3219 "deepseek".to_string(),
3220 )
3221 .expect("resume thread")
3222 .expect("thread found");
3223 assert_eq!(message_count(&manager), 3);
3224 }
3225
3226 #[test]
3227 fn persist_thread_preserves_stored_policy() {
3228 let store = temp_core_state("persist-policy");
3232 let mut metadata = test_thread_metadata("thread-policy");
3233 metadata.sandbox_policy = Some("workspace-write".to_string());
3234 metadata.approval_mode = Some("on-request".to_string());
3235 store.upsert_thread(&metadata).expect("seed thread");
3236
3237 let mut manager = ThreadManager::new(store);
3240 let resume_params = ThreadResumeParams {
3241 thread_id: "thread-policy".to_string(),
3242 history: None,
3243 path: None,
3244 model: None,
3245 model_provider: None,
3246 cwd: None,
3247 approval_policy: None,
3248 sandbox: None,
3249 config: None,
3250 base_instructions: None,
3251 developer_instructions: None,
3252 personality: None,
3253 persist_extended_history: false,
3254 };
3255 manager
3256 .resume_thread_with_history(
3257 &resume_params,
3258 Path::new("/tmp/codewhale"),
3259 "deepseek".to_string(),
3260 )
3261 .expect("resume thread")
3262 .expect("thread found");
3263
3264 let persisted = manager
3265 .state_store()
3266 .get_thread("thread-policy")
3267 .expect("read thread")
3268 .expect("thread persisted");
3269 assert_eq!(persisted.sandbox_policy.as_deref(), Some("workspace-write"));
3270 assert_eq!(persisted.approval_mode.as_deref(), Some("on-request"));
3271 }
3272
3273 #[tokio::test]
3274 async fn invoke_tool_returns_timeout_status_for_slow_tools() {
3275 use async_trait::async_trait;
3276 use codewhale_agent::ModelRegistry;
3277 use codewhale_config::ConfigToml;
3278 use codewhale_execpolicy::{AskForApproval, ExecPolicyEngine};
3279 use codewhale_hooks::HookDispatcher;
3280 use codewhale_mcp::McpManager;
3281 use codewhale_protocol::{ToolKind, ToolOutput, ToolPayload};
3282 use codewhale_tools::{FunctionCallError, ToolDescriptor, ToolHandler, ToolInvocation};
3283
3284 struct SlowTool;
3285 #[async_trait]
3286 impl ToolHandler for SlowTool {
3287 fn kind(&self) -> ToolKind {
3288 ToolKind::Function
3289 }
3290
3291 async fn handle(
3292 &self,
3293 _invocation: ToolInvocation,
3294 ) -> std::result::Result<ToolOutput, FunctionCallError> {
3295 time::sleep(Duration::from_millis(200)).await;
3296 Ok(ToolOutput::Function {
3297 body: Some(json!("late")),
3298 success: true,
3299 })
3300 }
3301 }
3302
3303 let mut registry = ToolRegistry::default();
3304 registry
3305 .register(
3306 ToolDescriptor {
3307 name: "slow_tool".to_string(),
3308 input_schema: json!({"type":"object"}),
3309 output_schema: json!({"type":"object"}),
3310 supports_parallel_tool_calls: true,
3311 timeout_ms: None,
3312 },
3313 Arc::new(SlowTool),
3314 )
3315 .expect("register slow tool");
3316
3317 let runtime = Runtime::new(
3318 ConfigToml::default(),
3319 ModelRegistry::default(),
3320 temp_core_state("invoke-tool-timeout"),
3321 Arc::new(registry),
3322 Arc::new(McpManager::default()),
3323 ExecPolicyEngine::new(vec![], vec![]),
3324 HookDispatcher::default(),
3325 );
3326
3327 let result = runtime
3328 .invoke_tool(
3329 ToolCall {
3330 name: "slow_tool".to_string(),
3331 payload: ToolPayload::Function {
3332 arguments: "{}".to_string(),
3333 },
3334 source: ToolCallSource::Direct,
3335 raw_tool_call_id: None,
3336 },
3337 AskForApproval::Never,
3338 Path::new("/tmp/codewhale"),
3339 )
3340 .await
3341 .expect("invoke tool");
3342
3343 assert_eq!(result["status"], "timeout");
3344 assert_eq!(result["ok"], false);
3345 }
3346
3347 #[tokio::test]
3348 async fn thread_message_response_ids_are_unique_for_equal_length_inputs() {
3349 let mut runtime = Runtime::new(
3353 ConfigToml::default(),
3354 ModelRegistry::default(),
3355 temp_core_state("response-id-unique"),
3356 Arc::new(ToolRegistry::default()),
3357 Arc::new(McpManager::default()),
3358 ExecPolicyEngine::new(vec![], vec![]),
3359 HookDispatcher::default(),
3360 );
3361 let spawned = runtime
3362 .thread_manager
3363 .spawn_thread_with_history(
3364 "deepseek".to_string(),
3365 PathBuf::from("/tmp/codewhale"),
3366 InitialHistory::New,
3367 true,
3368 )
3369 .expect("spawn thread");
3370 let thread_id = spawned.thread.id.clone();
3371
3372 let mut response_ids = Vec::new();
3373 for input in ["aaaa", "bbbb"] {
3374 let response = runtime
3375 .handle_thread(ThreadRequest::Message {
3376 thread_id: thread_id.clone(),
3377 input: input.to_string(),
3378 })
3379 .await
3380 .expect("handle message");
3381 let response_id = response
3382 .events
3383 .iter()
3384 .find_map(|frame| match frame {
3385 EventFrame::ResponseStart { response_id } => Some(response_id.clone()),
3386 _ => None,
3387 })
3388 .expect("response start event");
3389 response_ids.push(response_id);
3390 }
3391
3392 assert_ne!(
3393 response_ids[0], response_ids[1],
3394 "equal-length inputs must not share a response_id"
3395 );
3396 assert!(
3397 response_ids.iter().all(|id| id.starts_with("resp-")),
3398 "response ids should use the resp-<uuid> shape: {response_ids:?}"
3399 );
3400 }
3401}