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