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