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.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}