1use std::{
2 collections::{HashMap, HashSet},
3 sync::{
4 Arc,
5 atomic::{AtomicBool, Ordering},
6 },
7 time::Duration,
8};
9
10use anyhow::Context as _;
11use chrono::{DateTime, Duration as ChronoDuration, Utc};
12use kcode_session_history::{SessionCommand, SessionRecord, SessionStopRequest};
13use serde_json::{Value, json};
14use tokio::sync::{Mutex, OnceCell, RwLock};
15use uuid::Uuid;
16
17use super::{
18 AgentMode, Api, ApiError, Config, Manuals, RuntimeModel, Session, SessionService, TurnDeadline,
19 TurnDeadlineKind,
20};
21use kcode_kennedy_sessions::SessionOptions;
22
23const POLL_INTERVAL: Duration = Duration::from_secs(1);
24const STARTUP_RETRY: Duration = Duration::from_secs(2);
25
26#[derive(Clone)]
27pub struct SessionRuntime {
28 manuals: Manuals,
29 pub model: RuntimeModel,
30 pub user_root_node_id: String,
31 pub kennedy_root_node_id: String,
32}
33
34pub enum TurnCompletion {
35 Finished,
36 Stopped,
37}
38
39pub struct Orchestrator {
40 config: Config,
41 api: Api,
42 sessions: SessionService,
43 runtime: OnceCell<SessionRuntime>,
44 initialization: Mutex<()>,
45 writer: Arc<Mutex<()>>,
46 writer_job_active: AtomicBool,
47 commands_in_flight: Mutex<HashSet<String>>,
48 active_operations: Mutex<HashMap<String, Uuid>>,
49 conversation_locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
50 last_poll_error: RwLock<Option<String>>,
51}
52
53impl Orchestrator {
54 pub fn new(config: Config, api: Api, sessions: SessionService) -> Self {
55 Self {
56 config,
57 api,
58 sessions,
59 runtime: OnceCell::new(),
60 initialization: Mutex::new(()),
61 writer: Arc::new(Mutex::new(())),
62 writer_job_active: AtomicBool::new(false),
63 commands_in_flight: Mutex::new(HashSet::new()),
64 active_operations: Mutex::new(HashMap::new()),
65 conversation_locks: Mutex::new(HashMap::new()),
66 last_poll_error: RwLock::new(None),
67 }
68 }
69
70 pub fn api(&self) -> &Api {
71 &self.api
72 }
73
74 pub fn writer(&self) -> &Arc<Mutex<()>> {
75 &self.writer
76 }
77
78 pub async fn run(self: Arc<Self>) -> anyhow::Result<()> {
79 self.initialize_until_ready().await;
80 loop {
81 match self.poll_once().await {
82 Ok(()) => *self.last_poll_error.write().await = None,
83 Err(error) => {
84 let message = error.to_string();
85 let mut previous = self.last_poll_error.write().await;
86 if previous.as_deref() != Some(message.as_str()) {
87 tracing::warn!(error=%error, "Backend orchestration poll will retry");
88 *previous = Some(message);
89 }
90 }
91 }
92 tokio::time::sleep(POLL_INTERVAL).await;
93 }
94 }
95
96 pub async fn initialize_until_ready(&self) {
97 if self.runtime.get().is_some() {
98 return;
99 }
100 let _initialization = self.initialization.lock().await;
101 if self.runtime.get().is_some() {
102 return;
103 }
104 let mut previous = None;
105 loop {
106 match self.initialize().await {
107 Ok(runtime) => {
108 let model = runtime.model.model.clone();
109 let _ = self.runtime.set(runtime);
110 tracing::info!(%model, "Native Rust orchestration worker ready");
111 return;
112 }
113 Err(error) => {
114 let message = error.to_string();
115 if previous.as_deref() != Some(message.as_str()) {
116 tracing::warn!(error=%error, "Waiting for Kennedy services before starting orchestration");
117 previous = Some(message);
118 }
119 tokio::time::sleep(STARTUP_RETRY).await;
120 }
121 }
122 }
123 }
124
125 async fn initialize(&self) -> anyhow::Result<SessionRuntime> {
126 self.api.kmap_node(self.api.user_root_node_id())?;
127 self.api.kmap_node(self.api.kennedy_root_node_id())?;
128 self.api.history_health()?;
129 let manuals = Manuals::open();
130 let runtime = SessionRuntime {
131 manuals,
132 model: self.config.runtime_model.clone(),
133 user_root_node_id: self.api.user_root_node_id().to_owned(),
134 kennedy_root_node_id: self.api.kennedy_root_node_id().to_owned(),
135 };
136 self.api.history_release_interrupted_ingress().await?;
137 Ok(runtime)
138 }
139
140 pub fn runtime(&self) -> anyhow::Result<&SessionRuntime> {
141 self.runtime
142 .get()
143 .context("orchestration runtime is not initialized")
144 }
145
146 pub async fn open_session(
147 &self,
148 runtime: SessionRuntime,
149 options: SessionOptions,
150 restored: Option<&Value>,
151 ) -> anyhow::Result<Session> {
152 let system_prompt = if matches!(options.mode, AgentMode::Ingress { .. }) {
153 runtime.manuals.compose_ingress(
154 &runtime.model,
155 options
156 .source_session_type
157 .as_deref()
158 .unwrap_or("conversation"),
159 )
160 } else {
161 let session_context = if options.session_type == "free-time" {
162 self_time_schedule(&options.free_time)
163 } else {
164 String::new()
165 };
166 runtime.manuals.compose_conversation(
167 &runtime.model,
168 &options.session_type,
169 &session_context,
170 )
171 };
172 Session::new(
173 self.sessions.clone(),
174 system_prompt,
175 runtime.manuals.subagent_codex_prompt().to_owned(),
176 runtime.model,
177 options,
178 restored,
179 )
180 .await
181 }
182
183 async fn poll_once(self: &Arc<Self>) -> anyhow::Result<()> {
184 let histories = self.list_history().await?;
185 self.signal_pending_stops().await?;
186 self.sync_conversation_commands().await?;
187 self.schedule_writer_job(&histories).await?;
188 self.api.synchronize_audio_ingress().await?;
189 Ok(())
190 }
191
192 async fn signal_pending_stops(&self) -> anyhow::Result<()> {
193 let command_conversations = self
194 .api
195 .history_command_heads()
196 .await?
197 .into_iter()
198 .map(|command| command.conversation_id)
199 .collect::<HashSet<_>>();
200 for request in self.pending_stops().await? {
201 if !self.operation_is_active(&request.session_id).await
202 && request.scope == "turn"
203 && !command_conversations.contains(&request.session_id)
204 {
205 self.finish_idle_turn_stop(&request.session_id).await?;
206 }
207 }
208 Ok(())
209 }
210
211 async fn pending_stops(&self) -> anyhow::Result<Vec<SessionStopRequest>> {
212 Ok(self.api.history_stop_heads().await?)
213 }
214
215 pub async fn pending_stop(
216 &self,
217 session_id: &str,
218 ) -> anyhow::Result<Option<SessionStopRequest>> {
219 Ok(self
220 .pending_stops()
221 .await?
222 .into_iter()
223 .find(|request| request.session_id == session_id))
224 }
225
226 pub async fn complete_pending_stop(
227 &self,
228 session_id: &str,
229 outcome: Value,
230 ) -> anyhow::Result<()> {
231 if let Some(request) = self.pending_stop(session_id).await? {
232 self.api.history_complete_stop(&request.id, outcome).await?;
233 }
234 Ok(())
235 }
236
237 pub async fn operation_is_active(&self, session_id: &str) -> bool {
238 self.active_operations.lock().await.contains_key(session_id)
239 }
240
241 async fn finish_idle_turn_stop(&self, session_id: &str) -> anyhow::Result<()> {
242 let lock = self.conversation_lock(session_id).await;
243 let _guard = lock.lock().await;
244 if self.operation_is_active(session_id).await {
245 return Ok(());
246 }
247 if self.pending_stop(session_id).await?.is_none() {
248 return Ok(());
249 }
250 let record = self.get_conversation(session_id).await?;
251 if record.phase != "active" {
252 return Ok(());
253 }
254 let record = Arc::new(Mutex::new(record));
255 let mut session = {
256 let locked = record.lock().await;
257 self.session_for_record(&locked).await?
258 };
259 let telegram_event = matches!(session.session_type.as_str(), "telegram" | "telegram-group")
260 .then(|| session.pending_external_event_id.clone())
261 .flatten();
262 session.interrupt_current_turn()?;
263 persist_record(&self.api, &record, session.snapshot()?, false).await?;
264 if let Some(event_id) = telegram_event {
265 self.api
266 .telegram_interrupt_event(&event_id, session_id)
267 .await?;
268 }
269 self.complete_pending_stop(session_id, json!({"status":"stopped","scope":"turn"}))
270 .await
271 }
272
273 pub async fn register_operation(&self, session_id: &str, operation_id: Uuid) {
274 self.active_operations
275 .lock()
276 .await
277 .insert(session_id.to_owned(), operation_id);
278 }
279
280 pub async fn remove_operation(&self, session_id: &str, operation_id: Uuid) {
281 let mut active = self.active_operations.lock().await;
282 if active
283 .get(session_id)
284 .is_some_and(|operation| *operation == operation_id)
285 {
286 active.remove(session_id);
287 }
288 }
289
290 pub async fn run_session_turn<C, F>(
291 &self,
292 session_id: &str,
293 session: &mut Session,
294 operation_id: Uuid,
295 turn_deadline: Option<TurnDeadline>,
296 checkpoint: C,
297 ) -> anyhow::Result<TurnCompletion>
298 where
299 C: FnMut(Value) -> F + Send,
300 F: std::future::Future<Output = anyhow::Result<()>> + Send,
301 {
302 self.register_operation(session_id, operation_id).await;
303 let stop = match self.api.history_listen_for_stop(session_id) {
304 Ok(stop) => stop,
305 Err(error) => {
306 self.remove_operation(session_id, operation_id).await;
307 return Err(error.into());
308 }
309 };
310 let result: anyhow::Result<TurnCompletion> = {
311 let turn = session.run_pending_turn(operation_id, turn_deadline, checkpoint);
312 tokio::pin!(turn);
313 tokio::select! {
314 biased;
315 _ = stop.requested() => Ok({
316 let _ = self.api.cancel_intelligence(operation_id);
317 TurnCompletion::Stopped
318 }),
319 result = &mut turn => result.map(|_| TurnCompletion::Finished),
320 }
321 };
322 self.remove_operation(session_id, operation_id).await;
323 result
324 }
325
326 pub async fn list_history(&self) -> anyhow::Result<Vec<SessionRecord>> {
327 Ok(self.api.history_list().await?)
328 }
329
330 pub async fn conversation_lock(&self, id: &str) -> Arc<Mutex<()>> {
331 self.conversation_locks
332 .lock()
333 .await
334 .entry(id.to_owned())
335 .or_insert_with(|| Arc::new(Mutex::new(())))
336 .clone()
337 }
338
339 async fn sync_conversation_commands(self: &Arc<Self>) -> anyhow::Result<()> {
340 let commands = self.api.history_command_heads().await?;
341 for command in commands {
342 let id = command.id.clone();
343 if command.cancel_requested && self.commands_in_flight.lock().await.contains(&id) {
344 continue;
345 }
346 let mut in_flight = self.commands_in_flight.lock().await;
347 if !in_flight.insert(id.clone()) {
348 continue;
349 }
350 drop(in_flight);
351 let worker = self.clone();
352 tokio::spawn(async move {
353 if let Err(error) = worker.process_conversation_command(command).await {
354 tracing::warn!(command_id=%id, error=%error, "Browser conversation command will retry");
355 }
356 worker.commands_in_flight.lock().await.remove(&id);
357 });
358 }
359 Ok(())
360 }
361
362 async fn process_conversation_command(&self, command: SessionCommand) -> anyhow::Result<()> {
363 let command_id = command.id.clone();
364 let conversation_id = command.conversation_id.clone();
365 let lock = self.conversation_lock(&conversation_id).await;
366 let _conversation_guard = lock.lock().await;
367 let command = if command.status == "pending" {
368 self.api.history_claim_command(&command_id).await?
369 } else {
370 command
371 };
372 let record = self.get_conversation(&conversation_id).await?;
373 if record.phase != "active" || !is_browser_conversation(&record) {
374 self.complete_command(&command_id, json!({"status":"conversation_closed"}))
375 .await?;
376 return Ok(());
377 }
378 let kind = command.kind.clone();
379 let payload = command.payload.clone();
380 let record = Arc::new(Mutex::new(record));
381 if kind == "end" {
382 let mut state = record.lock().await.state.clone();
383 let abandoned_pending_turn = state
384 .get("pendingTurn")
385 .and_then(Value::as_bool)
386 .unwrap_or(false);
387 state["orchestration"] = json!({
388 "owner":"backend",
389 "status":"ending",
390 "abandonedPendingTurn":abandoned_pending_turn,
391 });
392 if let Some(session_id) = state.get("rustLibSessionId").and_then(Value::as_str) {
393 self.api.release_managed_sources(session_id).await;
394 }
395 self.request_conversation_ingress(&record, Some(state))
396 .await?;
397 self.complete_command(&command_id, json!({"status":"closed"}))
398 .await?;
399 return Ok(());
400 }
401 let mut session = {
402 let locked = record.lock().await;
403 self.session_for_record(&locked).await?
404 };
405 if command.cancel_requested {
406 session.interrupt_current_turn()?;
407 persist_record(&self.api, &record, session.snapshot()?, false).await?;
408 self.complete_command(
409 &command_id,
410 json!({"status":"stopped","reason":"user_stopped"}),
411 )
412 .await?;
413 self.complete_pending_stop(
414 &conversation_id,
415 json!({"status":"stopped","scope":"turn"}),
416 )
417 .await?;
418 return Ok(());
419 }
420 if session.orchestration.get("owner").and_then(Value::as_str) != Some("backend") {
421 session.orchestration = json!({"owner":"backend","status":"idle"});
422 persist_record(&self.api, &record, session.snapshot()?, false).await?;
423 }
424 let external_event_id = format!("web:{command_id}");
425 let outcome = match kind.as_str() {
426 "message" => {
427 if session
428 .answer_for_external_event(&external_event_id)
429 .is_none()
430 {
431 if !session.pending_turn {
432 let mut metadata = payload
433 .get("metadata")
434 .cloned()
435 .unwrap_or_else(|| json!({}));
436 metadata["externalEventId"] = json!(external_event_id);
437 anyhow::ensure!(
438 session.begin_user_turn(
439 payload
440 .get("text")
441 .and_then(Value::as_str)
442 .unwrap_or_default(),
443 &metadata,
444 ),
445 "The queued message contained no usable input"
446 );
447 }
448 session.orchestration = json!({"owner":"backend","status":"working"});
449 persist_record(&self.api, &record, session.snapshot()?, true).await?;
450 let operation_id = Uuid::new_v4();
451 let api = self.api.clone();
452 let saved_record = record.clone();
453 let result = self
454 .run_session_turn(
455 &conversation_id,
456 &mut session,
457 operation_id,
458 None,
459 move |state| {
460 let api = api.clone();
461 let record = saved_record.clone();
462 async move {
463 persist_record(&api, &record, state, false).await?;
464 Ok(())
465 }
466 },
467 )
468 .await;
469 if matches!(&result, Ok(TurnCompletion::Stopped)) {
470 session.interrupt_current_turn()?;
471 persist_record(&self.api, &record, session.snapshot()?, false).await?;
472 self.complete_command(
473 &command_id,
474 json!({"status":"stopped","reason":"user_stopped"}),
475 )
476 .await?;
477 self.complete_pending_stop(
478 &conversation_id,
479 json!({"status":"stopped","scope":"turn"}),
480 )
481 .await?;
482 return Ok(());
483 }
484 if let Err(error) = result {
485 let round_limit = kcode_agent_runtime::is_session_round_limit(&error);
486 session.orchestration = if is_cancelled(&error) {
487 json!({"owner":"backend","status":"stopped"})
488 } else if round_limit {
489 json!({"owner":"backend","status":"stopped","lastError":bounded_error(&error)})
490 } else {
491 json!({"owner":"backend","status":"retrying","lastError":bounded_error(&error)})
492 };
493 let persisted =
494 persist_record(&self.api, &record, session.snapshot()?, false).await;
495 if round_limit {
496 persisted?;
497 tracing::warn!(command_id=%command_id, "Browser conversation stopped at the tool-loop round limit");
498 self.complete_command(
499 &command_id,
500 json!({"status":"stopped","reason":"tool_loop_round_limit"}),
501 )
502 .await?;
503 return Ok(());
504 }
505 persisted.ok();
506 if is_cancelled(&error) {
507 self.complete_command(&command_id, json!({"status":"stopped"}))
508 .await?;
509 return Ok(());
510 }
511 return Err(error);
512 }
513 if session.requires_history_ingress() {
514 session.orchestration =
515 json!({"owner":"backend","status":"ending","reason":"context-limit"});
516 persist_record(&self.api, &record, session.snapshot()?, false).await?;
517 self.request_conversation_ingress(&record, None).await?;
518 self.complete_command(
519 &command_id,
520 json!({"status":"closed","reason":"context_limit"}),
521 )
522 .await?;
523 return Ok(());
524 }
525 }
526 anyhow::ensure!(
527 session
528 .answer_for_external_event(&external_event_id)
529 .is_some(),
530 "Kennedy completed the web turn without a recoverable response"
531 );
532 session.orchestration = json!({"owner":"backend","status":"idle"});
533 persist_record(&self.api, &record, session.snapshot()?, false).await?;
534 json!({"status":"answered"})
535 }
536 "retry" => {
537 if session.pending_turn {
538 session.reset_exhausted_turn_rounds_for_retry();
539 session.orchestration = json!({"owner":"backend","status":"working"});
540 persist_record(&self.api, &record, session.snapshot()?, false).await?;
541 let operation_id = Uuid::new_v4();
542 let api = self.api.clone();
543 let saved_record = record.clone();
544 let result = self
545 .run_session_turn(
546 &conversation_id,
547 &mut session,
548 operation_id,
549 None,
550 move |state| {
551 let api = api.clone();
552 let record = saved_record.clone();
553 async move {
554 persist_record(&api, &record, state, false).await?;
555 Ok(())
556 }
557 },
558 )
559 .await;
560 if matches!(&result, Ok(TurnCompletion::Stopped)) {
561 session.interrupt_current_turn()?;
562 persist_record(&self.api, &record, session.snapshot()?, false).await?;
563 self.complete_command(
564 &command_id,
565 json!({"status":"stopped","reason":"user_stopped"}),
566 )
567 .await?;
568 self.complete_pending_stop(
569 &conversation_id,
570 json!({"status":"stopped","scope":"turn"}),
571 )
572 .await?;
573 return Ok(());
574 }
575 if let Err(error) = result {
576 let round_limit = kcode_agent_runtime::is_session_round_limit(&error);
577 session.orchestration = if is_cancelled(&error) {
578 json!({"owner":"backend","status":"stopped"})
579 } else if round_limit {
580 json!({"owner":"backend","status":"stopped","lastError":bounded_error(&error)})
581 } else {
582 json!({"owner":"backend","status":"retrying","lastError":bounded_error(&error)})
583 };
584 let persisted =
585 persist_record(&self.api, &record, session.snapshot()?, false).await;
586 if round_limit {
587 persisted?;
588 tracing::warn!(command_id=%command_id, "Browser conversation stopped at the tool-loop round limit");
589 self.complete_command(
590 &command_id,
591 json!({"status":"stopped","reason":"tool_loop_round_limit"}),
592 )
593 .await?;
594 return Ok(());
595 }
596 persisted.ok();
597 if is_cancelled(&error) {
598 self.complete_command(&command_id, json!({"status":"stopped"}))
599 .await?;
600 return Ok(());
601 }
602 return Err(error);
603 }
604 if session.requires_history_ingress() {
605 session.orchestration =
606 json!({"owner":"backend","status":"ending","reason":"context-limit"});
607 persist_record(&self.api, &record, session.snapshot()?, false).await?;
608 self.request_conversation_ingress(&record, None).await?;
609 self.complete_command(
610 &command_id,
611 json!({"status":"closed","reason":"context_limit"}),
612 )
613 .await?;
614 return Ok(());
615 }
616 }
617 session.orchestration = json!({"owner":"backend","status":"idle"});
618 persist_record(&self.api, &record, session.snapshot()?, false).await?;
619 json!({"status":"retried"})
620 }
621 "send-and-end" => {
622 anyhow::ensure!(
623 !session.pending_turn,
624 "The saved query must finish before this conversation can end"
625 );
626 if !session.transcript.iter().any(|item| {
627 item.get("externalEventId").and_then(Value::as_str) == Some(&external_event_id)
628 }) {
629 let mut metadata = payload
630 .get("metadata")
631 .cloned()
632 .unwrap_or_else(|| json!({}));
633 metadata["externalEventId"] = json!(external_event_id);
634 anyhow::ensure!(
635 session.append_final_user_message(
636 payload
637 .get("text")
638 .and_then(Value::as_str)
639 .unwrap_or_default(),
640 &metadata
641 ),
642 "The final conversation command contained no usable input"
643 );
644 }
645 persist_record(&self.api, &record, session.snapshot()?, true).await?;
646 self.close_conversation(&record, &session).await?;
647 json!({"status":"closed"})
648 }
649 _ => anyhow::bail!("Unsupported browser conversation command {kind}"),
650 };
651 self.complete_command(&command_id, outcome).await?;
652 self.complete_pending_stop(
653 &conversation_id,
654 json!({"status":"already-completed","scope":"turn"}),
655 )
656 .await?;
657 Ok(())
658 }
659
660 pub async fn session_for_record(&self, record: &SessionRecord) -> anyhow::Result<Session> {
661 let runtime = self.runtime()?.clone();
662 let mut state = record.state.clone();
663 let session_type = session_type(record);
664 if matches!(session_type.as_str(), "telegram" | "telegram-group") {
665 if !state.get("channel").is_some_and(Value::is_object) {
666 state["channel"] = json!({});
667 }
668 state["channel"]["maxObjectBytes"] = json!(self.config.telegram_max_media_bytes);
669 }
670 let roots = string_array(state.get("rootNodeIds"));
671 let roots = if roots.is_empty() {
672 vec![
673 runtime.user_root_node_id.clone(),
674 runtime.kennedy_root_node_id.clone(),
675 ]
676 } else {
677 roots
678 };
679 let mut options = SessionOptions::conversation(session_type.clone(), roots);
680 options.reference_root_node_ids = string_array(state.get("referenceRootNodeIds"));
681 options.channel = state.get("channel").cloned().unwrap_or(Value::Null);
682 options.free_time = state.get("freeTime").cloned().unwrap_or(Value::Null);
683 options.orchestration = state
684 .get("orchestration")
685 .cloned()
686 .unwrap_or_else(|| json!({"owner":"backend","status":"idle"}));
687 options.provenance_id = state
688 .get("provenanceId")
689 .and_then(Value::as_str)
690 .map(str::to_owned);
691 options.mode = match session_type.as_str() {
692 "free-time" => AgentMode::FreeTime,
693 "wakeup" => AgentMode::Wakeup,
694 _ => AgentMode::Conversation,
695 };
696 self.open_session(runtime, options, Some(&state)).await
697 }
698
699 pub async fn close_conversation(
700 &self,
701 record: &Arc<Mutex<SessionRecord>>,
702 session: &Session,
703 ) -> anyhow::Result<()> {
704 session.release_managed_sources().await;
705 self.request_conversation_ingress(record, None).await
706 }
707
708 pub async fn request_conversation_ingress(
709 &self,
710 record: &Arc<Mutex<SessionRecord>>,
711 state: Option<Value>,
712 ) -> anyhow::Result<()> {
713 let mut locked = record.lock().await;
714 let id = locked.id.clone();
715 let state = state.unwrap_or_else(|| locked.state.clone());
716 let response = self
717 .api
718 .history_request_ingress(
719 &id,
720 kcode_session_history::Checkpoint {
721 expected_version: locked.version,
722 state,
723 user_activity: false,
724 },
725 )
726 .await?;
727 *locked = response;
728 Ok(())
729 }
730
731 async fn complete_command(&self, id: &str, outcome: Value) -> anyhow::Result<()> {
732 self.api.history_complete_command(id, outcome).await?;
733 Ok(())
734 }
735
736 pub async fn get_conversation(&self, id: &str) -> anyhow::Result<SessionRecord> {
737 Ok(self.api.history_get_session(id).await?)
738 }
739
740 pub async fn get_listed_conversation(&self, id: &str) -> anyhow::Result<Option<SessionRecord>> {
741 match self.api.history_get_session(id).await {
742 Ok(record) => Ok(Some(record)),
743 Err(error) if listed_session_disappeared(&error) => Ok(None),
744 Err(error) => Err(error.into()),
745 }
746 }
747
748 async fn schedule_writer_job(
749 self: &Arc<Self>,
750 histories: &[SessionRecord],
751 ) -> anyhow::Result<()> {
752 if self.writer_job_active.load(Ordering::Acquire) {
753 return Ok(());
754 }
755 if let Some(record) = histories
756 .iter()
757 .find(|record| record.phase == "active" && session_type(record) == "free-time")
758 .cloned()
759 {
760 self.launch_writer_job("self time", move |worker| async move {
761 let id = record.id;
762 let Some(record) = worker.get_listed_conversation(&id).await? else {
763 return Ok(());
764 };
765 worker.process_self_time(record).await
766 })
767 .await;
768 return Ok(());
769 }
770 if let Some(record) = next_ingress(histories, Utc::now()).cloned() {
771 self.launch_writer_job("memory ingress", move |worker| async move {
772 let id = record.id;
773 let Some(record) = worker.get_listed_conversation(&id).await? else {
774 return Ok(());
775 };
776 worker.process_ingress(record).await
777 })
778 .await;
779 }
780 Ok(())
781 }
782
783 async fn launch_writer_job<F, Fut>(self: &Arc<Self>, label: &'static str, task: F)
784 where
785 F: FnOnce(Arc<Self>) -> Fut + Send + 'static,
786 Fut: std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
787 {
788 if self
789 .writer_job_active
790 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
791 .is_err()
792 {
793 return;
794 }
795 let worker = self.clone();
796 tokio::spawn(async move {
797 let _writer_guard = worker.writer.lock().await;
798 if let Err(error) = task(worker.clone()).await {
799 tracing::warn!(
800 %label,
801 error=%bounded_error(&error),
802 "Kmap writer job will retry"
803 );
804 }
805 worker.writer_job_active.store(false, Ordering::Release);
806 });
807 }
808
809 async fn process_ingress(&self, mut record: SessionRecord) -> anyhow::Result<()> {
810 let id = record.id.clone();
811 let rust_session_id = format!("kennedy:history-ingress:{id}");
812 let mut stage = "prepare";
813 let result = async {
814 if record.phase == "ingress_pending" {
815 record
816 .state
817 .get("sessionId")
818 .and_then(Value::as_str)
819 .context("The queued session has no Session History ID")?;
820 stage = "claim";
821 record = self
822 .api
823 .history_start_ingress(
824 &id,
825 kcode_session_history::StartIngress {
826 expected_version: record.version,
827 provenance_id: format!("session:{id}"),
828 },
829 )
830 .await?;
831 }
832 if record.phase != "ingress_in_progress" {
833 return Ok(());
834 }
835 stage = "model_loop";
836 let runtime = self.runtime()?.clone();
837 let state = record.state.clone();
838 let source_session_type = state
839 .get("sessionType")
840 .and_then(Value::as_str)
841 .unwrap_or("conversation")
842 .to_owned();
843 let roots = {
844 let roots = string_array(state.get("rootNodeIds"));
845 if roots.is_empty() {
846 vec![
847 runtime.user_root_node_id.clone(),
848 runtime.kennedy_root_node_id.clone(),
849 ]
850 } else {
851 roots
852 }
853 };
854 let options = SessionOptions {
855 session_type: "history-ingress".into(),
856 root_node_ids: roots,
857 reference_root_node_ids: string_array(state.get("referenceRootNodeIds")),
858 channel: state.get("channel").cloned().unwrap_or(Value::Null),
859 free_time: Value::Null,
860 orchestration: Value::Null,
861 provenance_id: None,
862 mode: AgentMode::Ingress {
863 record_id: Some(id.clone()),
864 },
865 source_session_type: Some(source_session_type),
866 group_context: state
867 .get("channel")
868 .and_then(|channel| channel.get("groupContext"))
869 .cloned()
870 .unwrap_or(Value::Null),
871 rust_lib_session_id: Some(rust_session_id.clone()),
872 };
873 let restored = ingress_restore_state(&state);
874 let mut session = self.open_session(runtime, options, Some(restored)).await?;
875 let record = Arc::new(Mutex::new(record));
876 persist_ingress_record(&self.api, &record, session.snapshot()?).await?;
877 if !session.completed {
878 session.pending_turn = true;
879 let api = self.api.clone();
880 let saved_record = record.clone();
881 let completion = self
882 .run_session_turn(
883 &id,
884 &mut session,
885 Uuid::new_v4(),
886 None,
887 move |session_state| {
888 let api = api.clone();
889 let record = saved_record.clone();
890 async move {
891 persist_ingress_record(&api, &record, session_state).await?;
892 Ok(())
893 }
894 },
895 )
896 .await?;
897 if matches!(completion, TurnCompletion::Stopped) {
898 session.interrupt_current_turn()?;
899 session.commit_current_write_session()?;
900 stage = "stop-completion";
901 }
902 }
903 persist_ingress_record(&self.api, &record, session.snapshot()?).await?;
904 stage = "completion";
905 let mut locked = record.lock().await;
906 let completed = self
907 .api
908 .history_complete_ingress(&id, locked.version)
909 .await?;
910 *locked = completed.clone();
911 Ok(())
912 }
913 .await;
914 if let Err(error) = result {
915 self.record_ingress_failure(&id, stage, &error).await.ok();
916 return Err(error);
917 }
918 self.api.release_managed_sources(&rust_session_id).await;
919 Ok(())
920 }
921
922 async fn record_ingress_failure(
923 &self,
924 id: &str,
925 stage: &str,
926 error: &anyhow::Error,
927 ) -> anyhow::Result<()> {
928 let latest = self.get_conversation(id).await?;
929 if !matches!(
930 latest.phase.as_str(),
931 "ingress_pending" | "ingress_in_progress"
932 ) {
933 return Ok(());
934 }
935 self.api
936 .history_fail_ingress(
937 id,
938 kcode_session_history::IngressFailure {
939 expected_version: latest.version,
940 stage: stage.to_owned(),
941 code: Some("ingress_error".into()),
942 message: bounded_error(error),
943 rounds_used: None,
944 context_tokens: None,
945 context_window_tokens: None,
946 },
947 )
948 .await?;
949 Ok(())
950 }
951
952 async fn process_self_time(&self, record: SessionRecord) -> anyhow::Result<()> {
953 let runtime = self.runtime()?.clone();
954 let id = record.id.clone();
955 let mut state = record.state.clone();
956 if state.get("freeTime").is_none() {
957 let intent = state
958 .get("selfTimeIntent")
959 .context("backend self-time record is missing its durable start intent")?;
960 let duration = intent
961 .get("durationMinutes")
962 .and_then(Value::as_f64)
963 .context("self-time duration is missing")?;
964 let requested = intent
965 .get("requestedAt")
966 .and_then(Value::as_str)
967 .or(Some(record.started_at.as_str()))
968 .context("self-time request time is missing")?;
969 let requested_at = DateTime::parse_from_rfc3339(requested)?.with_timezone(&Utc);
970 let deadline =
971 requested_at + ChronoDuration::milliseconds((duration * 60_000.0).round() as i64);
972 state["freeTime"] = json!({"runId":id,"runStartedAt":requested_at.to_rfc3339(),"deadlineAt":deadline.to_rfc3339(),"durationMinutes":duration,"customPrompt":intent.get("customPrompt").and_then(Value::as_str).unwrap_or(""),"sliceIndex":1});
973 state["orchestration"] = json!({"owner":"backend","status":"running"});
974 }
975 let mut options = SessionOptions::conversation(
976 "free-time",
977 vec![
978 runtime.user_root_node_id.clone(),
979 runtime.kennedy_root_node_id.clone(),
980 ],
981 );
982 options.mode = AgentMode::FreeTime;
983 options.free_time = state.get("freeTime").cloned().unwrap_or(Value::Null);
984 options.provenance_id = state
985 .get("provenanceId")
986 .and_then(Value::as_str)
987 .map(str::to_owned);
988 options.orchestration = json!({"owner":"backend","status":"running"});
989 let mut session = self
990 .open_session(runtime.clone(), options, Some(&state))
991 .await?;
992 session.stage_free_time_opening();
993 let record_arc = Arc::new(Mutex::new(record));
994 persist_record(&self.api, &record_arc, session.snapshot()?, true).await?;
995 let deadline = session
996 .free_time
997 .get("deadlineAt")
998 .and_then(Value::as_str)
999 .and_then(|value| DateTime::parse_from_rfc3339(value).ok())
1000 .map(|value| value.with_timezone(&Utc))
1001 .context("self-time deadline is invalid")?;
1002 let hard_stop_at = deadline + ChronoDuration::minutes(15);
1003 let timeout = (hard_stop_at - Utc::now())
1004 .to_std()
1005 .unwrap_or(Duration::ZERO);
1006 let operation_id = Uuid::new_v4();
1007 let api = self.api.clone();
1008 let saved = record_arc.clone();
1009 let result = tokio::time::timeout(
1010 timeout,
1011 self.run_session_turn(
1012 &id,
1013 &mut session,
1014 operation_id,
1015 Some(TurnDeadline {
1016 kind: TurnDeadlineKind::SelfTimeHardStop,
1017 at: hard_stop_at,
1018 }),
1019 move |state| {
1020 let api = api.clone();
1021 let record = saved.clone();
1022 async move {
1023 persist_record(&api, &record, state, false).await?;
1024 Ok(())
1025 }
1026 },
1027 ),
1028 )
1029 .await;
1030 let mut reason = match result {
1031 Ok(Ok(TurnCompletion::Stopped)) => "user-stop".into(),
1032 Ok(Ok(TurnCompletion::Finished)) => session
1033 .free_time
1034 .get("sliceEndedReason")
1035 .and_then(Value::as_str)
1036 .unwrap_or_else(|| {
1037 if Utc::now() >= deadline {
1038 "deadline"
1039 } else {
1040 "tool"
1041 }
1042 })
1043 .to_owned(),
1044 Ok(Err(error)) => return Err(error),
1045 Err(_) => {
1046 let _ = self.api.cancel_intelligence(operation_id);
1047 self.remove_operation(&id, operation_id).await;
1048 "hard-stop".into()
1049 }
1050 };
1051 if reason != "user-stop" && self.pending_stop(&id).await?.is_some() {
1052 reason = "user-stop".into();
1053 }
1054 if reason == "user-stop" {
1055 session.interrupt_current_turn()?;
1056 }
1057 session.finalize_free_time(&reason)?;
1058 session.commit_current_write_session()?;
1059 persist_record(&self.api, &record_arc, session.snapshot()?, false).await?;
1060 session.release_managed_sources().await;
1061 let mut locked = record_arc.lock().await;
1062 let completed = self
1063 .api
1064 .history_complete(
1065 &id,
1066 kcode_session_history::Checkpoint {
1067 expected_version: locked.version,
1068 state: locked.state.clone(),
1069 user_activity: false,
1070 },
1071 )
1072 .await?;
1073 *locked = completed;
1074 if reason != "user-stop" && deadline - Utc::now() >= ChronoDuration::minutes(5) {
1075 self.create_next_self_time_slice(
1076 &runtime,
1077 session.free_time.clone(),
1078 session.provenance_id.clone(),
1079 deadline,
1080 )
1081 .await?;
1082 }
1083 Ok(())
1084 }
1085
1086 async fn create_next_self_time_slice(
1087 &self,
1088 runtime: &SessionRuntime,
1089 mut free: Value,
1090 provenance_id: Option<String>,
1091 deadline: DateTime<Utc>,
1092 ) -> anyhow::Result<()> {
1093 free["sliceIndex"] = json!(
1094 free.get("sliceIndex")
1095 .and_then(Value::as_u64)
1096 .unwrap_or_default()
1097 + 1
1098 );
1099 if let Some(object) = free.as_object_mut() {
1100 object.remove("sliceEndedReason");
1101 object.remove("sliceEndedAt");
1102 object.remove("warningNoticeAt");
1103 object.remove("expiredNoticeAt");
1104 }
1105 if let Some(message) = free.get("nextSessionMessage").cloned() {
1106 free["handoffMessage"] = message;
1107 }
1108 if let Some(object) = free.as_object_mut() {
1109 object.remove("nextSessionMessage");
1110 }
1111 free["deadlineAt"] = json!(deadline.to_rfc3339());
1112 let mut options = SessionOptions::conversation(
1113 "free-time",
1114 vec![
1115 runtime.user_root_node_id.clone(),
1116 runtime.kennedy_root_node_id.clone(),
1117 ],
1118 );
1119 options.mode = AgentMode::FreeTime;
1120 options.free_time = free;
1121 options.provenance_id = provenance_id;
1122 options.orchestration = json!({"owner":"backend","status":"running"});
1123 let mut session = self.open_session(runtime.clone(), options, None).await?;
1124 session.stage_free_time_opening();
1125 let state = session.snapshot()?;
1126 self.api
1127 .history_register(kcode_session_history::RegisterSession {
1128 id: required_string(&state, "sessionId")?,
1129 started_at: session.started_at.clone(),
1130 state,
1131 })
1132 .await?;
1133 Ok(())
1134 }
1135}
1136
1137pub async fn persist_record(
1138 api: &Api,
1139 record: &Arc<Mutex<SessionRecord>>,
1140 state: Value,
1141 user_activity: bool,
1142) -> anyhow::Result<()> {
1143 let mut record = record.lock().await;
1144 let id = record.id.clone();
1145 let result = match api
1146 .history_checkpoint(
1147 &id,
1148 kcode_session_history::Checkpoint {
1149 expected_version: record.version,
1150 state: state.clone(),
1151 user_activity,
1152 },
1153 )
1154 .await
1155 {
1156 Ok(result) => result,
1157 Err(error) if error.code == "state_conflict" => {
1158 let latest = api.history_get_session(&id).await?;
1159 if latest.state == state {
1160 latest
1161 } else {
1162 return Err(error.into());
1163 }
1164 }
1165 Err(error) => return Err(error.into()),
1166 };
1167 *record = result;
1168 Ok(())
1169}
1170async fn persist_ingress_record(
1171 api: &Api,
1172 record: &Arc<Mutex<SessionRecord>>,
1173 archive: Value,
1174) -> anyhow::Result<()> {
1175 let mut record = record.lock().await;
1176 let id = record.id.clone();
1177 let mut state = record.state.clone();
1178 state["historyIngress"] = archive;
1179 let result = match api
1180 .history_checkpoint(
1181 &id,
1182 kcode_session_history::Checkpoint {
1183 expected_version: record.version,
1184 state: state.clone(),
1185 user_activity: false,
1186 },
1187 )
1188 .await
1189 {
1190 Ok(result) => result,
1191 Err(error) if error.code == "state_conflict" => {
1192 let latest = api.history_get_session(&id).await?;
1193 if latest.state == state {
1194 latest
1195 } else {
1196 return Err(error.into());
1197 }
1198 }
1199 Err(error) => return Err(error.into()),
1200 };
1201 *record = result;
1202 Ok(())
1203}
1204fn session_type(record: &SessionRecord) -> String {
1205 record
1206 .state
1207 .get("sessionType")
1208 .and_then(Value::as_str)
1209 .unwrap_or("conversation")
1210 .into()
1211}
1212
1213fn next_ingress(histories: &[SessionRecord], now: DateTime<Utc>) -> Option<&SessionRecord> {
1214 histories
1215 .iter()
1216 .filter(|record| match record.phase.as_str() {
1217 "ingress_in_progress" => true,
1218 "ingress_pending" => record
1219 .ingress_next_attempt_at
1220 .as_deref()
1221 .and_then(|value| DateTime::parse_from_rfc3339(value).ok())
1222 .is_none_or(|next| next.with_timezone(&Utc) <= now),
1223 _ => false,
1224 })
1225 .min_by(|left, right| ingress_record_order(left, right))
1226}
1227
1228fn ingress_record_order(left: &SessionRecord, right: &SessionRecord) -> std::cmp::Ordering {
1229 let rank = |record: &SessionRecord| {
1230 if record.phase == "ingress_in_progress" {
1231 0
1232 } else {
1233 1
1234 }
1235 };
1236 rank(left)
1237 .cmp(&rank(right))
1238 .then_with(|| ingress_record_time(left).cmp(&ingress_record_time(right)))
1239 .then_with(|| left.id.cmp(&right.id))
1240}
1241
1242fn ingress_record_time(record: &SessionRecord) -> DateTime<Utc> {
1243 [&record.updated_at, &record.started_at]
1244 .into_iter()
1245 .find_map(|value| {
1246 DateTime::parse_from_rfc3339(value)
1247 .ok()
1248 .map(|value| value.with_timezone(&Utc))
1249 })
1250 .unwrap_or(DateTime::<Utc>::MAX_UTC)
1251}
1252
1253fn is_browser_conversation(record: &SessionRecord) -> bool {
1254 session_type(record) == "conversation"
1255}
1256fn required_string(value: &Value, key: &str) -> anyhow::Result<String> {
1257 value
1258 .get(key)
1259 .and_then(Value::as_str)
1260 .filter(|value| !value.is_empty())
1261 .map(str::to_owned)
1262 .with_context(|| format!("backend response omitted {key}"))
1263}
1264fn string_array(value: Option<&Value>) -> Vec<String> {
1265 value
1266 .and_then(Value::as_array)
1267 .into_iter()
1268 .flatten()
1269 .filter_map(Value::as_str)
1270 .map(str::to_owned)
1271 .collect()
1272}
1273fn ingress_restore_state(state: &Value) -> &Value {
1274 state.get("historyIngress").unwrap_or(state)
1275}
1276fn self_time_schedule(value: &Value) -> String {
1277 value
1278 .get("deadlineAt")
1279 .and_then(Value::as_str)
1280 .and_then(|value| DateTime::parse_from_rfc3339(value).ok())
1281 .map(|deadline| {
1282 format!(
1283 "The self-time deadline is {}.",
1284 super::prompts::human_utc_datetime(deadline.with_timezone(&Utc))
1285 )
1286 })
1287 .unwrap_or_else(|| "The self-time deadline was not supplied.".into())
1288}
1289fn bounded_error(error: &anyhow::Error) -> String {
1290 format!("{error:#}").chars().take(1_000).collect()
1291}
1292fn listed_session_disappeared(error: &ApiError) -> bool {
1293 error.code == "not_found"
1294}
1295fn is_cancelled(error: &anyhow::Error) -> bool {
1296 error
1297 .downcast_ref::<super::ApiError>()
1298 .is_some_and(|error| error.code == "operation_cancelled")
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303 use super::*;
1304
1305 fn session_record(id: &str, phase: &str, updated_at: &str) -> SessionRecord {
1306 serde_json::from_value(json!({
1307 "id":id,
1308 "phase":phase,
1309 "started_at":updated_at,
1310 "updated_at":updated_at,
1311 "state":{},
1312 "provenance_id":null,
1313 "version":1,
1314 "last_user_message_at":null,
1315 "ended_at":null,
1316 "ingress_failure_count":0,
1317 "ingress_failures":[],
1318 "ingress_next_attempt_at":null
1319 }))
1320 .unwrap()
1321 }
1322
1323 #[test]
1324 fn ingress_resumes_claimed_work_before_pending_work() {
1325 let now = DateTime::parse_from_rfc3339("2026-07-25T03:00:00Z")
1326 .unwrap()
1327 .with_timezone(&Utc);
1328 let pending = session_record("pending", "ingress_pending", "2026-07-25T01:00:00Z");
1329 let claimed = session_record("claimed", "ingress_in_progress", "2026-07-25T02:00:00Z");
1330 assert_eq!(
1331 next_ingress(&[pending, claimed], now).map(|record| record.id.as_str()),
1332 Some("claimed")
1333 );
1334 }
1335
1336 #[test]
1337 fn ingress_restart_prefers_its_own_checkpoint() {
1338 let source = json!({
1339 "sessionType":"conversation",
1340 "historyIngress":{"sessionType":"history-ingress","completed":true}
1341 });
1342 assert_eq!(
1343 ingress_restore_state(&source)["sessionType"],
1344 "history-ingress"
1345 );
1346 assert_eq!(
1347 ingress_restore_state(&json!({"sessionType":"conversation"}))["sessionType"],
1348 "conversation"
1349 );
1350 }
1351
1352 #[test]
1353 fn bounded_errors_include_the_cause_chain() {
1354 let error = anyhow::anyhow!("inner cause").context("outer context");
1355 assert_eq!(bounded_error(&error), "outer context: inner cause");
1356 }
1357}