use std::{
collections::{HashMap, HashSet},
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
use anyhow::Context as _;
use chrono::{DateTime, Duration as ChronoDuration, Utc};
use kcode_session_history::{SessionCommand, SessionRecord, SessionStopRequest};
use serde_json::{Value, json};
use tokio::sync::{Mutex, OnceCell, RwLock};
use uuid::Uuid;
use super::{AgentMode, Api, ApiError, Config, Manuals, RuntimeModel, Session, SessionService};
use kcode_kennedy_sessions::SessionOptions;
const POLL_INTERVAL: Duration = Duration::from_secs(1);
const STARTUP_RETRY: Duration = Duration::from_secs(2);
#[derive(Clone)]
pub struct SessionRuntime {
manuals: Manuals,
pub model: RuntimeModel,
pub user_root_node_id: String,
pub kennedy_root_node_id: String,
}
pub enum TurnCompletion {
Finished,
Stopped,
}
pub struct Orchestrator {
config: Config,
api: Api,
sessions: SessionService,
runtime: OnceCell<SessionRuntime>,
initialization: Mutex<()>,
writer: Arc<Mutex<()>>,
writer_job_active: AtomicBool,
commands_in_flight: Mutex<HashSet<String>>,
active_operations: Mutex<HashMap<String, Uuid>>,
conversation_locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
last_poll_error: RwLock<Option<String>>,
}
impl Orchestrator {
pub fn new(config: Config, api: Api, sessions: SessionService) -> Self {
Self {
config,
api,
sessions,
runtime: OnceCell::new(),
initialization: Mutex::new(()),
writer: Arc::new(Mutex::new(())),
writer_job_active: AtomicBool::new(false),
commands_in_flight: Mutex::new(HashSet::new()),
active_operations: Mutex::new(HashMap::new()),
conversation_locks: Mutex::new(HashMap::new()),
last_poll_error: RwLock::new(None),
}
}
pub fn api(&self) -> &Api {
&self.api
}
pub fn writer(&self) -> &Arc<Mutex<()>> {
&self.writer
}
pub async fn run(self: Arc<Self>) -> anyhow::Result<()> {
self.initialize_until_ready().await;
loop {
match self.poll_once().await {
Ok(()) => *self.last_poll_error.write().await = None,
Err(error) => {
let message = error.to_string();
let mut previous = self.last_poll_error.write().await;
if previous.as_deref() != Some(message.as_str()) {
tracing::warn!(error=%error, "Backend orchestration poll will retry");
*previous = Some(message);
}
}
}
tokio::time::sleep(POLL_INTERVAL).await;
}
}
pub async fn initialize_until_ready(&self) {
if self.runtime.get().is_some() {
return;
}
let _initialization = self.initialization.lock().await;
if self.runtime.get().is_some() {
return;
}
let mut previous = None;
loop {
match self.initialize().await {
Ok(runtime) => {
let model = runtime.model.model.clone();
let _ = self.runtime.set(runtime);
tracing::info!(%model, "Native Rust orchestration worker ready");
return;
}
Err(error) => {
let message = error.to_string();
if previous.as_deref() != Some(message.as_str()) {
tracing::warn!(error=%error, "Waiting for Kennedy services before starting orchestration");
previous = Some(message);
}
tokio::time::sleep(STARTUP_RETRY).await;
}
}
}
}
async fn initialize(&self) -> anyhow::Result<SessionRuntime> {
self.api.kmap_node(self.api.user_root_node_id())?;
self.api.kmap_node(self.api.kennedy_root_node_id())?;
self.api.history_health()?;
let manuals = Manuals::load(&self.config.system_prompts_directory)?;
let runtime = SessionRuntime {
manuals,
model: self.config.runtime_model.clone(),
user_root_node_id: self.api.user_root_node_id().to_owned(),
kennedy_root_node_id: self.api.kennedy_root_node_id().to_owned(),
};
self.api.history_release_interrupted_ingress().await?;
Ok(runtime)
}
pub fn runtime(&self) -> anyhow::Result<&SessionRuntime> {
self.runtime
.get()
.context("orchestration runtime is not initialized")
}
pub async fn open_session(
&self,
runtime: SessionRuntime,
options: SessionOptions,
restored: Option<&Value>,
) -> anyhow::Result<Session> {
let system_prompt = if matches!(options.mode, AgentMode::Ingress { .. }) {
runtime.manuals.compose_ingress(
&runtime.model,
options
.source_session_type
.as_deref()
.unwrap_or("conversation"),
)?
} else {
let session_context = if options.session_type == "free-time" {
self_time_schedule(&options.free_time)
} else {
String::new()
};
runtime.manuals.compose_conversation(
&runtime.model,
&options.session_type,
&session_context,
)?
};
Session::new(
self.sessions.clone(),
system_prompt,
runtime.model,
options,
restored,
)
.await
}
async fn poll_once(self: &Arc<Self>) -> anyhow::Result<()> {
let histories = self.list_history().await?;
self.signal_pending_stops().await?;
self.sync_conversation_commands().await?;
self.schedule_writer_job(&histories).await?;
self.api.synchronize_audio_ingress().await?;
Ok(())
}
async fn signal_pending_stops(&self) -> anyhow::Result<()> {
let command_conversations = self
.api
.history_command_heads()
.await?
.into_iter()
.map(|command| command.conversation_id)
.collect::<HashSet<_>>();
for request in self.pending_stops().await? {
if !self.operation_is_active(&request.session_id).await
&& request.scope == "turn"
&& !command_conversations.contains(&request.session_id)
{
self.finish_idle_turn_stop(&request.session_id).await?;
}
}
Ok(())
}
async fn pending_stops(&self) -> anyhow::Result<Vec<SessionStopRequest>> {
Ok(self.api.history_stop_heads().await?)
}
pub async fn pending_stop(
&self,
session_id: &str,
) -> anyhow::Result<Option<SessionStopRequest>> {
Ok(self
.pending_stops()
.await?
.into_iter()
.find(|request| request.session_id == session_id))
}
pub async fn complete_pending_stop(
&self,
session_id: &str,
outcome: Value,
) -> anyhow::Result<()> {
if let Some(request) = self.pending_stop(session_id).await? {
self.api.history_complete_stop(&request.id, outcome).await?;
}
Ok(())
}
pub async fn operation_is_active(&self, session_id: &str) -> bool {
self.active_operations.lock().await.contains_key(session_id)
}
async fn finish_idle_turn_stop(&self, session_id: &str) -> anyhow::Result<()> {
let lock = self.conversation_lock(session_id).await;
let _guard = lock.lock().await;
if self.operation_is_active(session_id).await {
return Ok(());
}
if self.pending_stop(session_id).await?.is_none() {
return Ok(());
}
let record = self.get_conversation(session_id).await?;
if record.phase != "active" {
return Ok(());
}
let record = Arc::new(Mutex::new(record));
let mut session = {
let locked = record.lock().await;
self.session_for_record(&locked).await?
};
let telegram_event = matches!(session.session_type.as_str(), "telegram" | "telegram-group")
.then(|| session.pending_external_event_id.clone())
.flatten();
session.interrupt_current_turn()?;
persist_record(&self.api, &record, session.snapshot()?, false).await?;
if let Some(event_id) = telegram_event {
self.api
.telegram_interrupt_event(&event_id, session_id)
.await?;
}
self.complete_pending_stop(session_id, json!({"status":"stopped","scope":"turn"}))
.await
}
pub async fn register_operation(&self, session_id: &str, operation_id: Uuid) {
self.active_operations
.lock()
.await
.insert(session_id.to_owned(), operation_id);
}
pub async fn remove_operation(&self, session_id: &str, operation_id: Uuid) {
let mut active = self.active_operations.lock().await;
if active
.get(session_id)
.is_some_and(|operation| *operation == operation_id)
{
active.remove(session_id);
}
}
pub async fn run_session_turn<C, F>(
&self,
session_id: &str,
session: &mut Session,
operation_id: Uuid,
checkpoint: C,
) -> anyhow::Result<TurnCompletion>
where
C: FnMut(Value) -> F + Send,
F: std::future::Future<Output = anyhow::Result<()>> + Send,
{
self.register_operation(session_id, operation_id).await;
let stop = match self.api.history_listen_for_stop(session_id) {
Ok(stop) => stop,
Err(error) => {
self.remove_operation(session_id, operation_id).await;
return Err(error.into());
}
};
let result: anyhow::Result<TurnCompletion> = {
let turn = session.run_pending_turn(operation_id, checkpoint);
tokio::pin!(turn);
tokio::select! {
biased;
_ = stop.requested() => Ok({
let _ = self.api.cancel_intelligence(operation_id);
TurnCompletion::Stopped
}),
result = &mut turn => result.map(|_| TurnCompletion::Finished),
}
};
self.remove_operation(session_id, operation_id).await;
result
}
pub async fn list_history(&self) -> anyhow::Result<Vec<SessionRecord>> {
Ok(self.api.history_list().await?)
}
pub async fn conversation_lock(&self, id: &str) -> Arc<Mutex<()>> {
self.conversation_locks
.lock()
.await
.entry(id.to_owned())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
async fn sync_conversation_commands(self: &Arc<Self>) -> anyhow::Result<()> {
let commands = self.api.history_command_heads().await?;
for command in commands {
let id = command.id.clone();
if command.cancel_requested && self.commands_in_flight.lock().await.contains(&id) {
continue;
}
let mut in_flight = self.commands_in_flight.lock().await;
if !in_flight.insert(id.clone()) {
continue;
}
drop(in_flight);
let worker = self.clone();
tokio::spawn(async move {
if let Err(error) = worker.process_conversation_command(command).await {
tracing::warn!(command_id=%id, error=%error, "Browser conversation command will retry");
}
worker.commands_in_flight.lock().await.remove(&id);
});
}
Ok(())
}
async fn process_conversation_command(&self, command: SessionCommand) -> anyhow::Result<()> {
let command_id = command.id.clone();
let conversation_id = command.conversation_id.clone();
let lock = self.conversation_lock(&conversation_id).await;
let _conversation_guard = lock.lock().await;
let command = if command.status == "pending" {
self.api.history_claim_command(&command_id).await?
} else {
command
};
let record = self.get_conversation(&conversation_id).await?;
if record.phase != "active" || !is_browser_conversation(&record) {
self.complete_command(&command_id, json!({"status":"conversation_closed"}))
.await?;
return Ok(());
}
let kind = command.kind.clone();
let payload = command.payload.clone();
let record = Arc::new(Mutex::new(record));
if kind == "end" {
let mut state = record.lock().await.state.clone();
let abandoned_pending_turn = state
.get("pendingTurn")
.and_then(Value::as_bool)
.unwrap_or(false);
state["orchestration"] = json!({
"owner":"backend",
"status":"ending",
"abandonedPendingTurn":abandoned_pending_turn,
});
if let Some(session_id) = state.get("rustLibSessionId").and_then(Value::as_str) {
self.api.release_managed_sources(session_id).await;
}
self.request_conversation_ingress(&record, Some(state))
.await?;
self.complete_command(&command_id, json!({"status":"closed"}))
.await?;
return Ok(());
}
let mut session = {
let locked = record.lock().await;
self.session_for_record(&locked).await?
};
if command.cancel_requested {
session.interrupt_current_turn()?;
persist_record(&self.api, &record, session.snapshot()?, false).await?;
self.complete_command(
&command_id,
json!({"status":"stopped","reason":"user_stopped"}),
)
.await?;
self.complete_pending_stop(
&conversation_id,
json!({"status":"stopped","scope":"turn"}),
)
.await?;
return Ok(());
}
if session.orchestration.get("owner").and_then(Value::as_str) != Some("backend") {
session.orchestration = json!({"owner":"backend","status":"idle"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
}
let external_event_id = format!("web:{command_id}");
let outcome = match kind.as_str() {
"message" => {
if session
.answer_for_external_event(&external_event_id)
.is_none()
{
if !session.pending_turn {
let mut metadata = payload
.get("metadata")
.cloned()
.unwrap_or_else(|| json!({}));
metadata["externalEventId"] = json!(external_event_id);
anyhow::ensure!(
session.begin_user_turn(
payload
.get("text")
.and_then(Value::as_str)
.unwrap_or_default(),
&metadata,
),
"The queued message contained no usable input"
);
}
session.orchestration = json!({"owner":"backend","status":"working"});
persist_record(&self.api, &record, session.snapshot()?, true).await?;
let operation_id = Uuid::new_v4();
let api = self.api.clone();
let saved_record = record.clone();
let result = self
.run_session_turn(
&conversation_id,
&mut session,
operation_id,
move |state| {
let api = api.clone();
let record = saved_record.clone();
async move {
persist_record(&api, &record, state, false).await?;
Ok(())
}
},
)
.await;
if matches!(&result, Ok(TurnCompletion::Stopped)) {
session.interrupt_current_turn()?;
persist_record(&self.api, &record, session.snapshot()?, false).await?;
self.complete_command(
&command_id,
json!({"status":"stopped","reason":"user_stopped"}),
)
.await?;
self.complete_pending_stop(
&conversation_id,
json!({"status":"stopped","scope":"turn"}),
)
.await?;
return Ok(());
}
if let Err(error) = result {
let round_limit = kcode_agent_runtime::is_session_round_limit(&error);
session.orchestration = if is_cancelled(&error) {
json!({"owner":"backend","status":"stopped"})
} else if round_limit {
json!({"owner":"backend","status":"stopped","lastError":bounded_error(&error)})
} else {
json!({"owner":"backend","status":"retrying","lastError":bounded_error(&error)})
};
let persisted =
persist_record(&self.api, &record, session.snapshot()?, false).await;
if round_limit {
persisted?;
tracing::warn!(command_id=%command_id, "Browser conversation stopped at the tool-loop round limit");
self.complete_command(
&command_id,
json!({"status":"stopped","reason":"tool_loop_round_limit"}),
)
.await?;
return Ok(());
}
persisted.ok();
if is_cancelled(&error) {
self.complete_command(&command_id, json!({"status":"stopped"}))
.await?;
return Ok(());
}
return Err(error);
}
if session.requires_history_ingress() {
session.orchestration =
json!({"owner":"backend","status":"ending","reason":"context-limit"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
self.request_conversation_ingress(&record, None).await?;
self.complete_command(
&command_id,
json!({"status":"closed","reason":"context_limit"}),
)
.await?;
return Ok(());
}
}
anyhow::ensure!(
session
.answer_for_external_event(&external_event_id)
.is_some(),
"Kennedy completed the web turn without a recoverable response"
);
session.orchestration = json!({"owner":"backend","status":"idle"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
json!({"status":"answered"})
}
"retry" => {
if session.pending_turn {
session.reset_exhausted_turn_rounds_for_retry();
session.orchestration = json!({"owner":"backend","status":"working"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
let operation_id = Uuid::new_v4();
let api = self.api.clone();
let saved_record = record.clone();
let result = self
.run_session_turn(
&conversation_id,
&mut session,
operation_id,
move |state| {
let api = api.clone();
let record = saved_record.clone();
async move {
persist_record(&api, &record, state, false).await?;
Ok(())
}
},
)
.await;
if matches!(&result, Ok(TurnCompletion::Stopped)) {
session.interrupt_current_turn()?;
persist_record(&self.api, &record, session.snapshot()?, false).await?;
self.complete_command(
&command_id,
json!({"status":"stopped","reason":"user_stopped"}),
)
.await?;
self.complete_pending_stop(
&conversation_id,
json!({"status":"stopped","scope":"turn"}),
)
.await?;
return Ok(());
}
if let Err(error) = result {
let round_limit = kcode_agent_runtime::is_session_round_limit(&error);
session.orchestration = if is_cancelled(&error) {
json!({"owner":"backend","status":"stopped"})
} else if round_limit {
json!({"owner":"backend","status":"stopped","lastError":bounded_error(&error)})
} else {
json!({"owner":"backend","status":"retrying","lastError":bounded_error(&error)})
};
let persisted =
persist_record(&self.api, &record, session.snapshot()?, false).await;
if round_limit {
persisted?;
tracing::warn!(command_id=%command_id, "Browser conversation stopped at the tool-loop round limit");
self.complete_command(
&command_id,
json!({"status":"stopped","reason":"tool_loop_round_limit"}),
)
.await?;
return Ok(());
}
persisted.ok();
if is_cancelled(&error) {
self.complete_command(&command_id, json!({"status":"stopped"}))
.await?;
return Ok(());
}
return Err(error);
}
if session.requires_history_ingress() {
session.orchestration =
json!({"owner":"backend","status":"ending","reason":"context-limit"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
self.request_conversation_ingress(&record, None).await?;
self.complete_command(
&command_id,
json!({"status":"closed","reason":"context_limit"}),
)
.await?;
return Ok(());
}
}
session.orchestration = json!({"owner":"backend","status":"idle"});
persist_record(&self.api, &record, session.snapshot()?, false).await?;
json!({"status":"retried"})
}
"send-and-end" => {
anyhow::ensure!(
!session.pending_turn,
"The saved query must finish before this conversation can end"
);
if !session.transcript.iter().any(|item| {
item.get("externalEventId").and_then(Value::as_str) == Some(&external_event_id)
}) {
let mut metadata = payload
.get("metadata")
.cloned()
.unwrap_or_else(|| json!({}));
metadata["externalEventId"] = json!(external_event_id);
anyhow::ensure!(
session.append_final_user_message(
payload
.get("text")
.and_then(Value::as_str)
.unwrap_or_default(),
&metadata
),
"The final conversation command contained no usable input"
);
}
persist_record(&self.api, &record, session.snapshot()?, true).await?;
self.close_conversation(&record, &session).await?;
json!({"status":"closed"})
}
_ => anyhow::bail!("Unsupported browser conversation command {kind}"),
};
self.complete_command(&command_id, outcome).await?;
self.complete_pending_stop(
&conversation_id,
json!({"status":"already-completed","scope":"turn"}),
)
.await?;
Ok(())
}
pub async fn session_for_record(&self, record: &SessionRecord) -> anyhow::Result<Session> {
let runtime = self.runtime()?.clone();
let mut state = record.state.clone();
let session_type = session_type(record);
if matches!(session_type.as_str(), "telegram" | "telegram-group") {
if !state.get("channel").is_some_and(Value::is_object) {
state["channel"] = json!({});
}
state["channel"]["maxObjectBytes"] = json!(self.config.telegram_max_media_bytes);
}
let roots = string_array(state.get("rootNodeIds"));
let roots = if roots.is_empty() {
vec![
runtime.user_root_node_id.clone(),
runtime.kennedy_root_node_id.clone(),
]
} else {
roots
};
let mut options = SessionOptions::conversation(session_type.clone(), roots);
options.reference_root_node_ids = string_array(state.get("referenceRootNodeIds"));
options.channel = state.get("channel").cloned().unwrap_or(Value::Null);
options.free_time = state.get("freeTime").cloned().unwrap_or(Value::Null);
options.orchestration = state
.get("orchestration")
.cloned()
.unwrap_or_else(|| json!({"owner":"backend","status":"idle"}));
options.provenance_id = state
.get("provenanceId")
.and_then(Value::as_str)
.map(str::to_owned);
options.mode = match session_type.as_str() {
"free-time" => AgentMode::FreeTime,
"wakeup" => AgentMode::Wakeup,
_ => AgentMode::Conversation,
};
self.open_session(runtime, options, Some(&state)).await
}
pub async fn close_conversation(
&self,
record: &Arc<Mutex<SessionRecord>>,
session: &Session,
) -> anyhow::Result<()> {
session.release_managed_sources().await;
self.request_conversation_ingress(record, None).await
}
pub async fn request_conversation_ingress(
&self,
record: &Arc<Mutex<SessionRecord>>,
state: Option<Value>,
) -> anyhow::Result<()> {
let mut locked = record.lock().await;
let id = locked.id.clone();
let state = state.unwrap_or_else(|| locked.state.clone());
let response = self
.api
.history_request_ingress(
&id,
kcode_session_history::Checkpoint {
expected_version: locked.version,
state,
user_activity: false,
},
)
.await?;
*locked = response;
Ok(())
}
async fn complete_command(&self, id: &str, outcome: Value) -> anyhow::Result<()> {
self.api.history_complete_command(id, outcome).await?;
Ok(())
}
pub async fn get_conversation(&self, id: &str) -> anyhow::Result<SessionRecord> {
Ok(self.api.history_get_session(id).await?)
}
pub async fn get_listed_conversation(&self, id: &str) -> anyhow::Result<Option<SessionRecord>> {
match self.api.history_get_session(id).await {
Ok(record) => Ok(Some(record)),
Err(error) if listed_session_disappeared(&error) => Ok(None),
Err(error) => Err(error.into()),
}
}
async fn schedule_writer_job(
self: &Arc<Self>,
histories: &[SessionRecord],
) -> anyhow::Result<()> {
if self.writer_job_active.load(Ordering::Acquire) {
return Ok(());
}
if let Some(record) = histories
.iter()
.find(|record| record.phase == "active" && session_type(record) == "free-time")
.cloned()
{
self.launch_writer_job("self time", move |worker| async move {
let id = record.id;
let Some(record) = worker.get_listed_conversation(&id).await? else {
return Ok(());
};
worker.process_self_time(record).await
})
.await;
return Ok(());
}
if let Some(record) = next_ingress(histories, Utc::now()).cloned() {
self.launch_writer_job("memory ingress", move |worker| async move {
let id = record.id;
let Some(record) = worker.get_listed_conversation(&id).await? else {
return Ok(());
};
worker.process_ingress(record).await
})
.await;
}
Ok(())
}
async fn launch_writer_job<F, Fut>(self: &Arc<Self>, label: &'static str, task: F)
where
F: FnOnce(Arc<Self>) -> Fut + Send + 'static,
Fut: std::future::Future<Output = anyhow::Result<()>> + Send + 'static,
{
if self
.writer_job_active
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return;
}
let worker = self.clone();
tokio::spawn(async move {
let _writer_guard = worker.writer.lock().await;
if let Err(error) = task(worker.clone()).await {
tracing::warn!(
%label,
error=%bounded_error(&error),
"Kmap writer job will retry"
);
}
worker.writer_job_active.store(false, Ordering::Release);
});
}
async fn process_ingress(&self, mut record: SessionRecord) -> anyhow::Result<()> {
let id = record.id.clone();
let rust_session_id = format!("kennedy:history-ingress:{id}");
let mut stage = "prepare";
let result = async {
if record.phase == "ingress_pending" {
record
.state
.get("sessionId")
.and_then(Value::as_str)
.context("The queued session has no Session History ID")?;
stage = "claim";
record = self
.api
.history_start_ingress(
&id,
kcode_session_history::StartIngress {
expected_version: record.version,
provenance_id: format!("session:{id}"),
},
)
.await?;
}
if record.phase != "ingress_in_progress" {
return Ok(());
}
stage = "model_loop";
let runtime = self.runtime()?.clone();
let state = record.state.clone();
let source_session_type = state
.get("sessionType")
.and_then(Value::as_str)
.unwrap_or("conversation")
.to_owned();
let roots = {
let roots = string_array(state.get("rootNodeIds"));
if roots.is_empty() {
vec![
runtime.user_root_node_id.clone(),
runtime.kennedy_root_node_id.clone(),
]
} else {
roots
}
};
let options = SessionOptions {
session_type: "history-ingress".into(),
root_node_ids: roots,
reference_root_node_ids: string_array(state.get("referenceRootNodeIds")),
channel: state.get("channel").cloned().unwrap_or(Value::Null),
free_time: Value::Null,
orchestration: Value::Null,
provenance_id: None,
mode: AgentMode::Ingress {
record_id: Some(id.clone()),
},
source_session_type: Some(source_session_type),
group_context: state
.get("channel")
.and_then(|channel| channel.get("groupContext"))
.cloned()
.unwrap_or(Value::Null),
rust_lib_session_id: Some(rust_session_id.clone()),
};
let restored = ingress_restore_state(&state);
let mut session = self.open_session(runtime, options, Some(restored)).await?;
let record = Arc::new(Mutex::new(record));
persist_ingress_record(&self.api, &record, session.snapshot()?).await?;
if !session.completed {
session.pending_turn = true;
let api = self.api.clone();
let saved_record = record.clone();
let completion = self
.run_session_turn(&id, &mut session, Uuid::new_v4(), move |session_state| {
let api = api.clone();
let record = saved_record.clone();
async move {
persist_ingress_record(&api, &record, session_state).await?;
Ok(())
}
})
.await?;
if matches!(completion, TurnCompletion::Stopped) {
session.interrupt_current_turn()?;
session.commit_current_write_session()?;
stage = "stop-completion";
}
}
persist_ingress_record(&self.api, &record, session.snapshot()?).await?;
stage = "completion";
let mut locked = record.lock().await;
let completed = self
.api
.history_complete_ingress(&id, locked.version)
.await?;
*locked = completed.clone();
Ok(())
}
.await;
if let Err(error) = result {
self.record_ingress_failure(&id, stage, &error).await.ok();
return Err(error);
}
self.api.release_managed_sources(&rust_session_id).await;
Ok(())
}
async fn record_ingress_failure(
&self,
id: &str,
stage: &str,
error: &anyhow::Error,
) -> anyhow::Result<()> {
let latest = self.get_conversation(id).await?;
if !matches!(
latest.phase.as_str(),
"ingress_pending" | "ingress_in_progress"
) {
return Ok(());
}
self.api
.history_fail_ingress(
id,
kcode_session_history::IngressFailure {
expected_version: latest.version,
stage: stage.to_owned(),
code: Some("ingress_error".into()),
message: bounded_error(error),
rounds_used: None,
context_tokens: None,
context_window_tokens: None,
},
)
.await?;
Ok(())
}
async fn process_self_time(&self, record: SessionRecord) -> anyhow::Result<()> {
let runtime = self.runtime()?.clone();
let id = record.id.clone();
let mut state = record.state.clone();
if state.get("freeTime").is_none() {
let intent = state
.get("selfTimeIntent")
.context("backend self-time record is missing its durable start intent")?;
let duration = intent
.get("durationMinutes")
.and_then(Value::as_f64)
.context("self-time duration is missing")?;
let requested = intent
.get("requestedAt")
.and_then(Value::as_str)
.or(Some(record.started_at.as_str()))
.context("self-time request time is missing")?;
let requested_at = DateTime::parse_from_rfc3339(requested)?.with_timezone(&Utc);
let deadline =
requested_at + ChronoDuration::milliseconds((duration * 60_000.0).round() as i64);
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});
state["orchestration"] = json!({"owner":"backend","status":"running"});
}
let mut options = SessionOptions::conversation(
"free-time",
vec![
runtime.user_root_node_id.clone(),
runtime.kennedy_root_node_id.clone(),
],
);
options.mode = AgentMode::FreeTime;
options.free_time = state.get("freeTime").cloned().unwrap_or(Value::Null);
options.provenance_id = state
.get("provenanceId")
.and_then(Value::as_str)
.map(str::to_owned);
options.orchestration = json!({"owner":"backend","status":"running"});
let mut session = self
.open_session(runtime.clone(), options, Some(&state))
.await?;
session.stage_free_time_opening();
let record_arc = Arc::new(Mutex::new(record));
persist_record(&self.api, &record_arc, session.snapshot()?, true).await?;
let deadline = session
.free_time
.get("deadlineAt")
.and_then(Value::as_str)
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.map(|value| value.with_timezone(&Utc))
.context("self-time deadline is invalid")?;
let timeout = (deadline - Utc::now() + ChronoDuration::minutes(6))
.to_std()
.unwrap_or(Duration::ZERO);
let operation_id = Uuid::new_v4();
let api = self.api.clone();
let saved = record_arc.clone();
let result = tokio::time::timeout(
timeout,
self.run_session_turn(&id, &mut session, operation_id, move |state| {
let api = api.clone();
let record = saved.clone();
async move {
persist_record(&api, &record, state, false).await?;
Ok(())
}
}),
)
.await;
let mut reason = match result {
Ok(Ok(TurnCompletion::Stopped)) => "user-stop".into(),
Ok(Ok(TurnCompletion::Finished)) => session
.free_time
.get("sliceEndedReason")
.and_then(Value::as_str)
.unwrap_or_else(|| {
if Utc::now() >= deadline {
"deadline"
} else {
"tool"
}
})
.to_owned(),
Ok(Err(error)) => return Err(error),
Err(_) => {
let _ = self.api.cancel_intelligence(operation_id);
self.remove_operation(&id, operation_id).await;
"hard-stop".into()
}
};
if reason != "user-stop" && self.pending_stop(&id).await?.is_some() {
reason = "user-stop".into();
}
if reason == "user-stop" {
session.interrupt_current_turn()?;
}
session.finalize_free_time(&reason)?;
session.commit_current_write_session()?;
persist_record(&self.api, &record_arc, session.snapshot()?, false).await?;
session.release_managed_sources().await;
let mut locked = record_arc.lock().await;
let completed = self
.api
.history_complete(
&id,
kcode_session_history::Checkpoint {
expected_version: locked.version,
state: locked.state.clone(),
user_activity: false,
},
)
.await?;
*locked = completed;
if reason != "user-stop" && deadline - Utc::now() >= ChronoDuration::minutes(5) {
self.create_next_self_time_slice(
&runtime,
session.free_time.clone(),
session.provenance_id.clone(),
deadline,
)
.await?;
}
Ok(())
}
async fn create_next_self_time_slice(
&self,
runtime: &SessionRuntime,
mut free: Value,
provenance_id: Option<String>,
deadline: DateTime<Utc>,
) -> anyhow::Result<()> {
free["sliceIndex"] = json!(
free.get("sliceIndex")
.and_then(Value::as_u64)
.unwrap_or_default()
+ 1
);
if let Some(object) = free.as_object_mut() {
object.remove("sliceEndedReason");
object.remove("sliceEndedAt");
object.remove("warningNoticeAt");
object.remove("expiredNoticeAt");
}
if let Some(message) = free.get("nextSessionMessage").cloned() {
free["handoffMessage"] = message;
}
if let Some(object) = free.as_object_mut() {
object.remove("nextSessionMessage");
}
free["deadlineAt"] = json!(deadline.to_rfc3339());
let mut options = SessionOptions::conversation(
"free-time",
vec![
runtime.user_root_node_id.clone(),
runtime.kennedy_root_node_id.clone(),
],
);
options.mode = AgentMode::FreeTime;
options.free_time = free;
options.provenance_id = provenance_id;
options.orchestration = json!({"owner":"backend","status":"running"});
let mut session = self.open_session(runtime.clone(), options, None).await?;
session.stage_free_time_opening();
let state = session.snapshot()?;
self.api
.history_register(kcode_session_history::RegisterSession {
id: required_string(&state, "sessionId")?,
started_at: session.started_at.clone(),
state,
})
.await?;
Ok(())
}
}
pub async fn persist_record(
api: &Api,
record: &Arc<Mutex<SessionRecord>>,
state: Value,
user_activity: bool,
) -> anyhow::Result<()> {
let mut record = record.lock().await;
let id = record.id.clone();
let result = match api
.history_checkpoint(
&id,
kcode_session_history::Checkpoint {
expected_version: record.version,
state: state.clone(),
user_activity,
},
)
.await
{
Ok(result) => result,
Err(error) if error.code == "state_conflict" => {
let latest = api.history_get_session(&id).await?;
if latest.state == state {
latest
} else {
return Err(error.into());
}
}
Err(error) => return Err(error.into()),
};
*record = result;
Ok(())
}
async fn persist_ingress_record(
api: &Api,
record: &Arc<Mutex<SessionRecord>>,
archive: Value,
) -> anyhow::Result<()> {
let mut record = record.lock().await;
let id = record.id.clone();
let mut state = record.state.clone();
state["historyIngress"] = archive;
let result = match api
.history_checkpoint(
&id,
kcode_session_history::Checkpoint {
expected_version: record.version,
state: state.clone(),
user_activity: false,
},
)
.await
{
Ok(result) => result,
Err(error) if error.code == "state_conflict" => {
let latest = api.history_get_session(&id).await?;
if latest.state == state {
latest
} else {
return Err(error.into());
}
}
Err(error) => return Err(error.into()),
};
*record = result;
Ok(())
}
fn session_type(record: &SessionRecord) -> String {
record
.state
.get("sessionType")
.and_then(Value::as_str)
.unwrap_or("conversation")
.into()
}
fn next_ingress(histories: &[SessionRecord], now: DateTime<Utc>) -> Option<&SessionRecord> {
histories
.iter()
.filter(|record| match record.phase.as_str() {
"ingress_in_progress" => true,
"ingress_pending" => record
.ingress_next_attempt_at
.as_deref()
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.is_none_or(|next| next.with_timezone(&Utc) <= now),
_ => false,
})
.min_by(|left, right| ingress_record_order(left, right))
}
fn ingress_record_order(left: &SessionRecord, right: &SessionRecord) -> std::cmp::Ordering {
let rank = |record: &SessionRecord| {
if record.phase == "ingress_in_progress" {
0
} else {
1
}
};
rank(left)
.cmp(&rank(right))
.then_with(|| ingress_record_time(left).cmp(&ingress_record_time(right)))
.then_with(|| left.id.cmp(&right.id))
}
fn ingress_record_time(record: &SessionRecord) -> DateTime<Utc> {
[&record.updated_at, &record.started_at]
.into_iter()
.find_map(|value| {
DateTime::parse_from_rfc3339(value)
.ok()
.map(|value| value.with_timezone(&Utc))
})
.unwrap_or(DateTime::<Utc>::MAX_UTC)
}
fn is_browser_conversation(record: &SessionRecord) -> bool {
session_type(record) == "conversation"
}
fn required_string(value: &Value, key: &str) -> anyhow::Result<String> {
value
.get(key)
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.map(str::to_owned)
.with_context(|| format!("backend response omitted {key}"))
}
fn string_array(value: Option<&Value>) -> Vec<String> {
value
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_str)
.map(str::to_owned)
.collect()
}
fn ingress_restore_state(state: &Value) -> &Value {
state.get("historyIngress").unwrap_or(state)
}
fn self_time_schedule(value: &Value) -> String {
value
.get("deadlineAt")
.and_then(Value::as_str)
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.map(|deadline| {
format!(
"The self-time deadline is {}.",
super::prompts::human_utc_datetime(deadline.with_timezone(&Utc))
)
})
.unwrap_or_else(|| "The self-time deadline was not supplied.".into())
}
fn bounded_error(error: &anyhow::Error) -> String {
format!("{error:#}").chars().take(1_000).collect()
}
fn listed_session_disappeared(error: &ApiError) -> bool {
error.code == "not_found"
}
fn is_cancelled(error: &anyhow::Error) -> bool {
error
.downcast_ref::<super::ApiError>()
.is_some_and(|error| error.code == "operation_cancelled")
}
#[cfg(test)]
mod tests {
use super::*;
fn session_record(id: &str, phase: &str, updated_at: &str) -> SessionRecord {
serde_json::from_value(json!({
"id":id,
"phase":phase,
"started_at":updated_at,
"updated_at":updated_at,
"state":{},
"provenance_id":null,
"version":1,
"last_user_message_at":null,
"ended_at":null,
"ingress_failure_count":0,
"ingress_failures":[],
"ingress_next_attempt_at":null
}))
.unwrap()
}
#[test]
fn ingress_resumes_claimed_work_before_pending_work() {
let now = DateTime::parse_from_rfc3339("2026-07-25T03:00:00Z")
.unwrap()
.with_timezone(&Utc);
let pending = session_record("pending", "ingress_pending", "2026-07-25T01:00:00Z");
let claimed = session_record("claimed", "ingress_in_progress", "2026-07-25T02:00:00Z");
assert_eq!(
next_ingress(&[pending, claimed], now).map(|record| record.id.as_str()),
Some("claimed")
);
}
#[test]
fn ingress_restart_prefers_its_own_checkpoint() {
let source = json!({
"sessionType":"conversation",
"historyIngress":{"sessionType":"history-ingress","completed":true}
});
assert_eq!(
ingress_restore_state(&source)["sessionType"],
"history-ingress"
);
assert_eq!(
ingress_restore_state(&json!({"sessionType":"conversation"}))["sessionType"],
"conversation"
);
}
#[test]
fn bounded_errors_include_the_cause_chain() {
let error = anyhow::anyhow!("inner cause").context("outer context");
assert_eq!(bounded_error(&error), "outer context: inner cause");
}
}