Resolve review2 regressions across wait deadlines, lane projections, and turn context
128 unmodified linesExpand all129130131132133134135136128 unmodified linesExpand all{"id":"int-7522a87e70f6fbbc3a67cb49fdbb30fa","kind":"field_change","created_at":"2026-09-06T14:38:58.564197Z","actor":"dawn","issue_id":"klbr-h37","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-76681c810ef3778cea80ba29b5c876e6","kind":"field_change","created_at":"2026-09-06T14:38:59.945928Z","actor":"dawn","issue_id":"klbr-06i","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-af125e68fb1857ce6a635781b5b963b5","kind":"field_change","created_at":"2026-09-06T14:45:19.598535Z","actor":"dawn","issue_id":"klbr-2d0","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-234b77c3332e105337eb5a361c9564f3","kind":"field_change","created_at":"2026-09-06T14:55:35.56433Z","actor":"dawn","issue_id":"klbr-99b","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-a892d32f7e821faaf5fd18d11d02df6d","kind":"field_change","created_at":"2026-09-06T16:14:48.750121Z","actor":"dawn","issue_id":"klbr-e0c","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-11f2a2429c81ffbace8247bb2eac5abc","kind":"field_change","created_at":"2026-09-06T16:15:12.603959Z","actor":"dawn","issue_id":"klbr-xdj","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}{"id":"int-6d3313ee6f1f5c50739fc79ee141f1ce","kind":"field_change","created_at":"2026-09-06T16:20:38.406803Z","actor":"dawn","issue_id":"klbr-9ap","extra":{"field":"status","new_value":"closed","old_value":"in_progress"}}{"id":"int-e8d141517617a0312ed515db6e6927b3","kind":"field_change","created_at":"2026-09-06T16:22:33.833817Z","actor":"dawn","issue_id":"klbr-uyt","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Closed"}}757 unmodified linesExpand all75875976076176276376476576690 unmodified lines85785885986086186286386486586686786886953 unmodified lines92392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031292 unmodified linesExpand all757 unmodified linesExpand all
let mut buf = String::new(); let mut partial_calls: HashMap<usize, PartialCall> = HashMap::new(); let mut tool_call_started_sent = false; let mut terminal_finish_reason: Option<String> = None;
while let Some(chunk) = res.next().await { let chunk = match chunk { Ok(c) => c,90 unmodified lines }; if choices.is_empty() { continue; } if let Some(fr) = choices[0].get("finish_reason").and_then(|v| v.as_str()) { if !fr.is_empty() && fr != "null" { terminal_finish_reason = Some(fr.to_string()); } } let delta = &choices[0]["delta"];
// accumulate tool call deltas if let Some(tc_arr) = delta["tool_calls"].as_array() {53 unmodified lines } } }
if !buf.trim().is_empty() { let line = buf.trim().to_string(); if let Some(data) = sse_data_payload(&line) { if data == "[DONE]" { if !partial_calls.is_empty() { let mut calls: Vec<(usize, ToolCall)> = partial_calls .into_iter() .map(|(idx, p)| { ( idx, ToolCall { id: p.id, kind: "function".into(), function: ToolCallFunction { name: p.name, arguments: p.arguments, }, }, ) }) .collect(); calls.sort_by_key(|(i, _)| *i); let calls: Vec<ToolCall> = calls.into_iter().map(|(_, c)| c).collect(); let _ = tok_tx.send(LlmEvent::ToolCalls(calls)).await; }
let elapsed = start_time.elapsed(); let mut u = final_usage.clone().unwrap_or_default(); if let Some(ttfc) = first_response_time { u.prompt_processing_ms = ttfc.as_millis() as u64; u.generation_ms = elapsed.saturating_sub(ttfc).as_millis() as u64; } else { u.prompt_processing_ms = elapsed.as_millis() as u64; u.generation_ms = 0; } u.bridge_processing_ms = bridge_setup_ms; let _ = tok_tx.send(LlmEvent::Usage(u)).await; let _ = tok_tx.send(LlmEvent::Completed).await;
return Ok(()); } } }
if terminal_finish_reason.is_none() { let err_msg = "stream terminated prematurely without [DONE]".to_string(); tracing::error!("{}", err_msg); let _ = tok_tx.send(LlmEvent::Error(err_msg.clone())).await; return Err(anyhow::anyhow!(err_msg)); }
if !partial_calls.is_empty() { let mut calls: Vec<(usize, ToolCall)> = partial_calls .into_iter() .map(|(idx, p)| { ( idx, ToolCall { id: p.id, kind: "function".into(), function: ToolCallFunction { name: p.name, arguments: p.arguments, }, }, ) }) .collect(); calls.sort_by_key(|(i, _)| *i); let calls: Vec<ToolCall> = calls.into_iter().map(|(_, c)| c).collect(); let _ = tok_tx.send(LlmEvent::ToolCalls(calls)).await; }
let elapsed = start_time.elapsed(); let mut u = final_usage.unwrap_or_default(); if let Some(ttfc) = first_response_time { u.prompt_processing_ms = ttfc.as_millis() as u64;1292 unmodified linesExpand all1234456782 unmodified lines111213141515161718192021222323242526272829303132293334353633373839404116 unmodified lines585960615859606162636465666263646566676869707172737475767778798081828384858687888989909192909192939495969798991001011021031041051061071081099495969798991001101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464471141151161171181191201211221231241251261274484491294504514524534544554564574584594604614624634644654664674684694701434714724734744751481491501511521531541551561571581591601611621634764774784794804814824834844854864874884891654904914924934943 unmodified lines49849950050150250350450550634 unmodified lines54154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258322923023123258458558658758858959059159259323423523623723823959459559659759859960060160260360460560660760860961061161261361461561661761861962062162262356 unmodified lines680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725118 unmodified linesExpand all84484584684784884943285085185285385485585685734 unmodified lines89289389489589689789889990090190290390490590690747690890991091191291391491591691791833 unmodified lines95295395495595695795895996096196252996396496596696725 unmodified lines9939949959965639979989991000100110021003100410051006100753 unmodified lines106110621063106463010651066106710681069107010711072107310741075305 unmodified linesExpand alluse std::path::Path;use std::sync::Arc;use std::time::Duration;use anyhow::{bail, Result};use anyhow::Result;use futures::{SinkExt, StreamExt};use tokio::{ net::{TcpListener, TcpStream}, sync::{broadcast, mpsc, watch},2 unmodified lines
use klbr_core::{ interrupt::{ExternalEvent, Interrupt}, memory::MemoryStore, models::{LlmClient, Message as CoreMessage}, models::{LlmClient, Message as CoreMessage, Role, ToolCall, ToolCallFunction}, runtime_turn::{ integrator::run_turn_loop, state::{CellOutputStream, ScratchpadEvent, TurnEvent}, }, AgentEvent, MetricsSnapshot,};use klbr_core::runtime_turn::state::SpeechStatus;use klbr_ipc::{ v2::{InputPayload, Lane, LanePayload, LifecyclePayload, V2ClientMsg, V2EventRecord, V2ServerMsg}, v2::{ CellPayload, CellState, CellStreamKind, CellStreamOutput, InputPayload, Lane, LanePayload, LifecyclePayload, ScratchpadPayload, SpeechPayload, V2ClientMsg, V2EventRecord, V2ServerMsg, }, ClientMsg, ServerMsg,};use klbr_runtime::{ behavior::Release, protocol::{OperationId, SessionId}, store::{Store, StoreAdmit}, store::{Event, Store, StoreAdmit},};
use crate::ingress::{admit_input, SourceRef};use crate::session_registry::{SessionRegistry, SessionState};use crate::session_registry::SessionRegistry;
const HISTORY_PAGE: usize = 50;const WS_BIND_ADDR: &str = "127.0.0.1:8765";
16 unmodified lines sanitized.parse::<SessionId>().unwrap_or_else(|_| SessionId::fresh()) })}
pub async fn query_events_after( store: &Store, session_id: &SessionId, after_seq: i64,) -> Result<Vec<V2EventRecord>> { let events = store.history(session_id, after_seq, 1024).await?; let mut records = Vec::new(); for ev in events { let payload = if ev.kind == "input" {pub fn project_event(session_id: &SessionId, ev: Event) -> V2EventRecord { let payload = match ev.kind.as_str() { "input" => { let text = ev .payload .get("content") .or_else(|| ev.payload.get("text")) .and_then(|c| c.as_str()) .unwrap_or("") .to_string(); let source = ev .payload .get("source_kind") .or_else(|| ev.payload.get("source")) .and_then(|s| s.as_str()) .unwrap_or("unknown") .to_string(); let author = ev .payload .get("author_name") .or_else(|| ev.payload.get("author")) .and_then(|a| a.as_str()) .map(|s| s.to_string()); LanePayload::Input(InputPayload { source, text, author, }) } else { LanePayload::Lifecycle(LifecyclePayload { kind: ev.kind, data: ev.payload, } "local.sent" => { let effect_id = ev .payload .get("effect_id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let text = ev .payload .get("text") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); LanePayload::Speech(SpeechPayload::Sent { effect_id, text, event_id: ev.id, status: "confirmed".to_string(), }) }; records.push(V2EventRecord::new( ev.id, session_id.to_string(), payload, ev.created_ms, )); } "delivery.held" => { let effect_id = ev .payload .get("effect_id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let text = ev .payload .get("text") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let reason = ev .payload .get("reason") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); LanePayload::Speech(SpeechPayload::Held { effect_id, text, status: "held".to_string(), reason, }) } "scratchpad" => { let text = ev .payload .get("text") .or_else(|| ev.payload.get("content")) .or_else(|| ev.payload.get("thoughts")) .and_then(|t| t.as_str()) .map(|s| s.to_string()) .unwrap_or_else(|| { if let Some(status) = ev.payload.get("status").and_then(|s| s.as_str()) { format!("tool_result: {}", status) } else { serde_json::to_string(&ev.payload).unwrap_or_default() } }); let execution_id = ev .payload .get("execution_id") .and_then(|e| e.as_str()) .map(|s| s.to_string()); LanePayload::Scratchpad(ScratchpadPayload { text, execution_id, delta: None, }) } "execution.started" => { let execution_id = ev .payload .get("execution_id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let generation = ev .payload .get("generation") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let revision = ev .payload .get("revision") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let code = ev .payload .get("code") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); LanePayload::Cell(CellPayload::Started { execution_id, generation, revision, code, }) } "execution.completed" => { let execution_id = ev .payload .get("execution_id") .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); let state = ev .payload .get("state") .and_then(|s| s.as_str()) .map(|s| match s { "error" => CellState::Error, "yielded" => CellState::Yielded, "interrupted" => CellState::Interrupted, _ => CellState::Ok, }) .unwrap_or(CellState::Ok);
let mut outputs = Vec::new(); let mut truncated_bytes = 0; let mut error = None;
if let Some(outcome) = ev.payload.get("outcome") { if let Some(outs) = outcome.get("output").and_then(|o| o.as_array()) { for out in outs { let stream_str = out.get("stream").and_then(|s| s.as_str()).unwrap_or("stdout"); let stream = match stream_str { "stderr" => CellStreamKind::Stderr, "result" => CellStreamKind::Result, _ => CellStreamKind::Stdout, }; let text = out.get("text").and_then(|t| t.as_str()).unwrap_or("").to_string(); outputs.push(CellStreamOutput { stream, text }); } } truncated_bytes = outcome .get("truncated_bytes") .and_then(|b| b.as_u64()) .unwrap_or(0) as usize; error = outcome.get("error").and_then(|e| e.as_str()).map(|s| s.to_string()); }
LanePayload::Cell(CellPayload::Completed { execution_id, state, output: outputs, truncated_bytes, error, }) } _ => LanePayload::Lifecycle(LifecyclePayload { kind: ev.kind, data: ev.payload, }), }; V2EventRecord::new(ev.id, session_id.to_string(), payload, ev.created_ms)}
pub async fn query_events_after( store: &Store, session_id: &SessionId, after_seq: i64,) -> Result<Vec<V2EventRecord>> { let mut current_after = after_seq; let mut records = Vec::new(); loop { let events = store.history(session_id, current_after, 1024).await?; if events.is_empty() { break; } let batch_len = events.len(); for ev in events { current_after = ev.id; records.push(project_event(session_id, ev)); } if batch_len < 1024 { break; } } Ok(records)}
async fn assemble_conversation_messages( store: &Store, session_id: &SessionId, pending_ids: &[i64], soul: &str,) -> Result<Vec<CoreMessage>> { let mut messages = Vec::new(); if !soul.is_empty() { messages.push(CoreMessage::system(soul)); }
let mut all_events = Vec::new(); let mut after = 0; loop { let page = store.history(session_id, after, 1024).await?; if page.is_empty() { break; } if let Some(last) = page.last() { after = last.id; } let page_len = page.len(); all_events.extend(page); if page_len < 1024 { break; } }
for ev in all_events { if pending_ids.contains(&ev.id) { continue; }
if ev.kind == "input" { let content = ev .payload .get("content") .and_then(|c| c.as_str()) .unwrap_or(""); if !content.is_empty() { messages.push(CoreMessage::user(content)); } } else if ev.kind == "scratchpad" { let event_type = ev .payload .get("type") .and_then(|t| t.as_str()) .unwrap_or("");
if event_type == "model_response" { let prose = ev .payload .get("content") .and_then(|c| c.as_str()) .map(String::from) .or_else(|| { ev.payload .get("text") .and_then(|t| t.as_str()) .map(String::from) }); let thoughts = ev .payload .get("thoughts") .and_then(|t| t.as_str()) .map(String::from);
let tool_calls: Option<Vec<ToolCall>> = ev .payload .get("tool_calls") .and_then(|tc| tc.as_array()) .map(|arr| { arr.iter() .filter_map(|item| { let id = item.get("id")?.as_str()?.to_string(); let name = item.get("name")?.as_str()?.to_string(); let arguments = item.get("arguments")?.as_str()?.to_string(); Some(ToolCall { id, kind: "function".to_string(), function: ToolCallFunction { name, arguments }, }) }) .collect() }) .filter(|calls: &Vec<ToolCall>| !calls.is_empty());
if let Some(calls) = tool_calls { messages.push(CoreMessage::with_tool_calls(calls, prose, thoughts)); } else { messages.push(CoreMessage { role: Role::Assistant, content: prose, reasoning_content: thoughts, tool_calls: None, tool_call_id: None, }); } } else if event_type == "tool_result" { let call_id = ev .payload .get("call_id") .and_then(|c| c.as_str()) .unwrap_or(""); let content = ev .payload .get("content") .and_then(|c| c.as_str()) .map(String::from) .or_else(|| { if let Some(err) = ev.payload.get("error").and_then(|e| e.as_str()) { Some(format!("Error: {err}")) } else if let Some(reason) = ev.payload.get("reason").and_then(|r| r.as_str()) { Some(format!("Skipped: {reason}")) } else { ev.payload .get("status") .and_then(|s| s.as_str()) .map(|s| format!("Tool execution {s}")) } }) .unwrap_or_default();
messages.push(CoreMessage::tool_result(call_id, content)); } } }
for id in pending_ids { if let Ok(Some(ev)) = store.get_event(session_id, *id).await { let content = ev .payload .get("content") .and_then(|c| c.as_str()) .unwrap_or(""); messages.push(CoreMessage::user(content)); } }
Ok(messages)}
pub async fn trigger_session_turn( registry: Arc<SessionRegistry>, llm: LlmClient, session_id: SessionId, initial_release: Release, v2_event_tx: broadcast::Sender<V2ServerMsg>,) -> Result<()> { trigger_session_turn_with_soul( registry, llm, session_id, initial_release, v2_event_tx, None, ) .await}
pub async fn trigger_session_turn_with_soul( registry: Arc<SessionRegistry>, llm: LlmClient, session_id: SessionId, initial_release: Release, v2_event_tx: broadcast::Sender<V2ServerMsg>, soul: Option<String>,) -> Result<()> { let entry = registry.get_or_open(&session_id, &initial_release).await?;
let state = entry.state().await; match state { SessionState::Draining | SessionState::Stopped => { bail!("session {session_id} is in state {state:?}, cannot run turn"); } SessionState::Executing { .. } => { return Ok(()); } SessionState::Waiting { .. } => { entry.transition_to_ready_with_store(registry.store()).await?; } SessionState::Ready => {} }
let op_id = OperationId::fresh(); let (abort_handle, abort_reg) = futures::future::AbortHandle::new_pair(); entry.register_runner(abort_handle).await;
if !entry.try_start_runner(op_id.clone(), abort_handle).await? { return Ok(()); } let _ = registry.store().set_session_state(&session_id, "ready").await;
let entry_clone = entry.clone(); let reg_clone = registry.clone(); let s_id = session_id.clone(); let tx_clone = v2_event_tx.clone(); let soul_str = soul .filter(|s| !s.is_empty()) .unwrap_or_else(|| klbr_core::config::DEFAULT_SOUL.to_string());
tokio::spawn(async move { let runner_future = async { loop { if !entry_clone.can_admit().await { break; }
let pending_ids = match reg_clone.store().pending(&s_id, 10).await { let pending_ids = match reg_clone.store().claim_pending(&s_id, 10).await { Ok(ids) if !ids.is_empty() => ids, _ => break, };
let op_id = OperationId::fresh(); let epoch = entry_clone.runtime.hook_epoch().await; if let Err(e) = entry_clone.transition_to_executing(epoch, op_id).await { tracing::warn!(session = %s_id, err = %e, "cannot transition to executing"); break; }
let mut messages = Vec::new(); for id in &pending_ids { if let Ok(Some(ev)) = reg_clone.store().get_event(&s_id, *id).await { let content = ev .payload .get("content") .and_then(|c| c.as_str()) .unwrap_or(""); messages.push(CoreMessage::user(content)); let messages = match assemble_conversation_messages( reg_clone.store(), &s_id, &pending_ids, &soul_str, ) .await { Ok(msgs) => msgs, Err(e) => { tracing::error!(session = %s_id, err = %e, "failed to assemble conversation context"); reg_clone.store().unclaim_inputs(&s_id, &pending_ids); break; } } };
let turn_stream = run_turn_loop( &llm, &entry_clone.runtime,3 unmodified lines ); tokio::pin!(turn_stream);
let mut yielded = None; let mut failed_error = None; while let Some(turn_event) = turn_stream.next().await { match &turn_event { TurnEvent::Scratchpad(sp) => { let text = match sp {34 unmodified lines execution_id: None, stream: None, delta: sp.text.clone(), }); if let Some(event_id) = sp.event_id { let payload = match sp.status { SpeechStatus::Confirmed => SpeechPayload::Sent { effect_id: sp.effect_id.clone(), text: sp.text.clone(), event_id, status: "confirmed".to_string(), }, _ => SpeechPayload::Held { effect_id: sp.effect_id.clone(), text: sp.text.clone(), status: "held".to_string(), reason: sp.reason.clone().unwrap_or_default(), }, }; let record = V2EventRecord::speech( event_id, s_id.to_string(), payload, chrono::Utc::now().timestamp_millis(), ); let _ = tx_clone.send(V2ServerMsg::Event { record }); } } TurnEvent::Yielded(y) => { yielded = Some(y.clone()); } TurnEvent::Failed(err) => { failed_error = Some(err.clone()); } TurnEvent::Interrupted | TurnEvent::Cancelled | TurnEvent::Timeout => { failed_error = Some(format!("{turn_event:?}")); } _ => {} } }
let _ = reg_clone.store().acknowledge_inputs(&s_id, pending_ids).await;
if let Some(y) = yielded { let wake_at = y .seconds .map(|s| chrono::Utc::now().timestamp_millis() + (s * 1000.0) as i64); if let Some(err) = failed_error { let _ = reg_clone.store().append(&s_id, "scratchpad", serde_json::json!({ "type": "turn_failed", "error": err, })).await; let _ = tx_clone.send(V2ServerMsg::Error { code: "turn_failed".into(), message: err, }); let _ = entry_clone .transition_to_waiting_with_store( reg_clone.store(), None, wake_at, Some(y.reason), ) .transition_to_ready_with_store(reg_clone.store()) .await; break; }
if let Some(y) = yielded { let has_new_pending = !reg_clone.store().pending(&s_id, 1).await.unwrap_or_default().is_empty(); if has_new_pending { let _ = entry_clone .transition_to_ready_with_store(reg_clone.store()) .await; continue; } else { let wake_at = y .seconds .map(|s| chrono::Utc::now().timestamp_millis() + (s * 1000.0) as i64); let _ = entry_clone .transition_to_waiting_with_store( reg_clone.store(), None, wake_at, Some(y.reason), ) .await; break; } } else { let _ = entry_clone .transition_to_ready_with_store(reg_clone.store()) .await;56 unmodified lines initial_release: Option<Release>, v2_mode: bool, mut shutdown: watch::Receiver<bool>,) -> Result<()> { let driver_reg = registry.clone(); let driver_llm = llm.clone(); let driver_release = initial_release.clone(); let driver_v2_tx = v2_event_tx.clone(); let mut driver_shutdown = shutdown.clone();
tokio::spawn(async move { let mut interval = tokio::time::interval(Duration::from_millis(50)); loop { tokio::select! { _ = driver_shutdown.changed() => break, _ = interval.tick() => { let now = match klbr_runtime::store::now_ms() { Ok(n) => n, Err(_) => chrono::Utc::now().timestamp_millis(), }; if let Ok(woken_sessions) = driver_reg.store().wake_due(now).await { for sid_str in woken_sessions { let s_id = parse_session_id(&sid_str); if let Some(entry) = driver_reg.get(&s_id).await { let _ = entry.transition_to_ready().await; } if let (Some(l), Some(rel)) = (&driver_llm, &driver_release) { let _ = trigger_session_turn( driver_reg.clone(), l.clone(), s_id, rel.clone(), driver_v2_tx.clone(), ).await; } } } } } } });
loop { let accepted = tokio::select! { _ = shutdown.changed() => break, accepted = listener.accept() => accepted,118 unmodified linesExpand all let _ = send_v2_msg(&mut ws_tx, &msg).await; } } } else { match &msg { V2ServerMsg::Pong { .. } | V2ServerMsg::Error { .. } => { let _ = send_v2_msg(&mut ws_tx, &msg).await; let _ = send_v2_msg(&mut ws_tx, &msg).await; } _ => {} } } } } ws_msg = ws_rx.next() => {34 unmodified lines seq, client_request_id: client_request_id.clone(), duplicate: Some(false), }).await; let record = V2EventRecord::input( seq, session_id.clone(), InputPayload { source: source.clone(), text: text.clone(), author: None, }, chrono::Utc::now().timestamp_millis(), ); let _ = v2_bcast.send(V2ServerMsg::Event { record }); if let (Some(l), Some(rel)) = (&llm, &initial_release) { let _ = trigger_session_turn( let _ = trigger_session_turn_with_soul( registry.clone(), l.clone(), s_id, rel.clone(), v2_bcast.clone(), Some(default_soul.clone()), ).await; } } Ok(StoreAdmit::Duplicate(seq)) => {33 unmodified lines V2ClientMsg::Interrupt { session_id, .. } => { let s_id = parse_session_id(&session_id); if let Some(entry) = registry.get(&s_id).await { entry.abort_runner().await; let _ = entry.transition_to_ready_with_store(registry.store()).await; } } V2ClientMsg::ResetKernel { session_id } => { let s_id = parse_session_id(&session_id); if let Some(entry) = registry.get(&s_id).await { entry.abort_runner().await; entry.runtime.shutdown().await; let _ = entry.transition_to_ready_with_store(registry.store()).await; } } } continue;25 unmodified lines if !v2_mode { let int = Interrupt::message(source, content); let _ = interrupt_tx.send(int).await; } else if let (Some(l), Some(rel)) = (&llm, &initial_release) { let _ = trigger_session_turn( let _ = trigger_session_turn_with_soul( registry.clone(), l.clone(), session_id, rel.clone(), v2_bcast.clone(), Some(default_soul.clone()), ).await; } } else if let StoreAdmit::Conflict(msg) = outcome { let _ = send_msg(&mut ws_tx, &ServerMsg::Error {53 unmodified lines metadata, })) .await; } else if let (Some(l), Some(rel)) = (&llm, &initial_release) { let _ = trigger_session_turn( let _ = trigger_session_turn_with_soul( registry.clone(), l.clone(), session_id, rel.clone(), v2_bcast.clone(), Some(default_soul.clone()), ).await; } } else if let StoreAdmit::Conflict(msg) = outcome { let _ = send_msg(&mut ws_tx, &ServerMsg::Error {305 unmodified linesExpand all144 unmodified linesExpand all14514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817915515618018118218318418518618718818919019119 unmodified lines21121221321418621521621721821918 unmodified lines23823924024121324224324424524682 unmodified lines32933033133230433333430633533633733833934 unmodified lines3743753763773493783793803813825 unmodified lines388389390391363392393365394395367396369397398399400401402300 unmodified linesExpand all144 unmodified linesExpand all let mut guard = self.active_runner_abort.lock().await; *guard = Some(handle); }
pub async fn try_start_runner( &self, op_id: OperationId, handle: futures::future::AbortHandle, ) -> Result<bool> { let mut state_guard = self.state.write().await; match &*state_guard { SessionState::Draining | SessionState::Stopped => { bail!("session {} is in state {:?}, cannot run turn", self.id, *state_guard); } SessionState::Executing { .. } => { return Ok(false); } SessionState::Waiting { .. } => { state_guard.transition_to_ready()?; } SessionState::Ready => {} } let epoch = self.runtime.hook_epoch().await; state_guard.transition_to_executing(epoch, op_id)?; let mut runner_guard = self.active_runner_abort.lock().await; *runner_guard = Some(handle); Ok(true) }
pub async fn clear_runner(&self) { let mut guard = self.active_runner_abort.lock().await; *guard = None; }
pub async fn abort_runner(&self) { let mut guard = self.active_runner_abort.lock().await; if let Some(handle) = guard.take() { let handle = { let mut guard = self.active_runner_abort.lock().await; guard.take() }; if let Some(handle) = handle { handle.abort(); } self.runtime.interrupt().await; }
pub async fn state(&self) -> SessionState { self.state.read().await.clone()19 unmodified lines }
pub async fn transition_to_ready_with_store(&self, store: &Store) -> Result<()> { self.state.write().await.transition_to_ready()?; let _ = store.set_session_state(&self.id, "ready").await; store.set_session_state(&self.id, "ready").await?; Ok(()) }
pub async fn transition_to_waiting(18 unmodified lines self.state .write() .await .transition_to_waiting(receipt, wake_at_ms, reason)?; let _ = store.set_session_state(&self.id, "waiting").await; store.wait_session(&self.id, wake_at_ms).await?; Ok(()) }
pub async fn transition_to_draining(&self) -> Result<()> {82 unmodified lines let now_ms = chrono::Utc::now().timestamp_millis(); let initial_state = match self.store.session_state(id).await? { Some((ref st, _)) if st == "stopped" => SessionState::Stopped, Some((ref st, _)) if st == "draining" => SessionState::Draining, Some((ref st, _)) if st == "waiting" => SessionState::Waiting { Some((ref st, wake_at)) if st == "waiting" => SessionState::Waiting { receipt: None, wake_at_ms: None, wake_at_ms: wake_at, reason: None, }, _ => SessionState::Ready, };34 unmodified lines .with_context(|| format!("session not found: {id}"))?;
entry.abort_runner().await; entry.transition_to_stopped().await?; let _ = self.store.set_session_state(id, "stopped").await; self.store.set_session_state(id, "stopped").await?; entry.runtime.shutdown().await; Ok(()) }
5 unmodified lines
// If not already draining, transition to draining if entry.state().await != SessionState::Draining { entry.transition_to_draining().await?; let _ = self.store.set_session_state(id, "draining").await; self.store.set_session_state(id, "draining").await?; } entry.runtime.wait_for_idle().await; let _ = tokio::time::timeout(std::time::Duration::from_secs(5), entry.runtime.wait_for_idle()).await; entry.abort_runner().await; entry.runtime.shutdown().await; entry.transition_to_stopped().await?; let _ = self.store.set_session_state(id, "stopped").await; self.store.set_session_state(id, "stopped").await?; entry.runtime.shutdown().await; Ok(()) }
pub async fn activate(300 unmodified linesExpand all8 unmodified lines910111213131415161760 unmodified lines8 unmodified lines) STRICT;
CREATE TABLE sessions ( id TEXT PRIMARY KEY, state TEXT NOT NULL DEFAULT 'ready' CHECK(state IN ('ready','waiting','draining','stopped')), state TEXT NOT NULL DEFAULT 'ready' CHECK(state IN ('ready','waiting')), wake_at_ms INTEGER, hook_revision TEXT REFERENCES releases(id), hook_epoch INTEGER NOT NULL DEFAULT 0 CHECK(hook_epoch >= 0), CHECK(state = 'waiting' OR wake_at_ms IS NULL)60 unmodified lines1234567891011121314151617181920-- Migration 002: expand sessions table CHECK constraint to include 'draining' and 'stopped'.BEGIN IMMEDIATE;
CREATE TABLE sessions_v2 ( id TEXT PRIMARY KEY, state TEXT NOT NULL DEFAULT 'ready' CHECK(state IN ('ready','waiting','draining','stopped')), wake_at_ms INTEGER, hook_revision TEXT REFERENCES releases(id), hook_epoch INTEGER NOT NULL DEFAULT 0 CHECK(hook_epoch >= 0), CHECK(state = 'waiting' OR wake_at_ms IS NULL)) STRICT;
INSERT INTO sessions_v2(id, state, wake_at_ms, hook_revision, hook_epoch)SELECT id, state, wake_at_ms, hook_revision, hook_epoch FROM sessions;
DROP TABLE sessions;ALTER TABLE sessions_v2 RENAME TO sessions;
PRAGMA user_version = 2;COMMIT;