Skip to main content

kcode_kennedy_orchestration/
worker.rs

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