1use std::{
8 collections::{BTreeMap, HashMap, HashSet},
9 path::Path,
10};
11
12use anyhow::{Context as _, ensure};
13use chrono::{SecondsFormat, Utc};
14use kcode_session_log::{
15 EventPosition, Role, Session as DurableSession, SessionStore as DurableSessionStore,
16};
17use serde::{Deserialize, Serialize};
18use serde_json::{Value, json};
19
20pub const FORMAT_VERSION: u32 = 1;
21pub const MAX_OBJECT_BYTES: u64 = 32 * 1024 * 1024 * 1024;
22pub const ESTIMATED_BYTES_PER_TOKEN: u64 = 4;
23
24#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
26pub struct ProviderTokenUsage {
27 pub input_tokens: u64,
28 pub cached_input_tokens: u64,
29 pub thinking_tokens: u64,
30 pub output_tokens: u64,
31}
32
33#[derive(Clone, Copy, Debug, PartialEq)]
35pub enum ProviderMetering {
36 Tokens(ProviderTokenUsage),
37 DurationSeconds { seconds: f64 },
38 Unavailable,
39}
40
41#[derive(Clone, Debug, Eq, PartialEq)]
43pub struct ProviderCostEstimate {
44 pub usd_nanos: u64,
45 pub accuracy: Value,
46 pub pricing_version: String,
47}
48
49pub type ProviderCostEstimator = fn(&str, &ProviderMetering) -> Option<ProviderCostEstimate>;
51
52#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
54pub struct ProviderCostSummary {
55 pub estimated_cost_usd_nanos: u64,
56 pub unpriced_provider_calls: u64,
57}
58
59pub(crate) fn legacy_provider_cost_summary_for_archive(
64 archive: &Value,
65 default_provider_model: Option<&str>,
66 estimator: ProviderCostEstimator,
67) -> anyhow::Result<ProviderCostSummary> {
68 let metadata = archive
69 .get("metadata")
70 .cloned()
71 .context("session archive is missing Chatend metadata")?;
72 let metadata = serde_json::from_value(metadata).context("decoding session archive metadata")?;
73 let log =
74 serde_json::from_value(archive.clone()).context("decoding immutable session archive")?;
75 let chatend = Chatend::replay_with_provider_costs(
76 metadata,
77 &log,
78 default_provider_model,
79 Some(estimator),
80 )?;
81 let status = chatend.projection().status;
82 Ok(ProviderCostSummary {
83 estimated_cost_usd_nanos: status.estimated_cost_usd_nanos,
84 unpriced_provider_calls: status.unpriced_provider_calls,
85 })
86}
87
88#[derive(
89 Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize,
90)]
91#[serde(transparent)]
92pub struct EventId(pub u64);
93
94impl std::fmt::Display for EventId {
95 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96 self.0.fmt(formatter)
97 }
98}
99
100#[derive(
101 Clone, Copy, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize,
102)]
103#[serde(transparent)]
104pub struct BoxId(pub u64);
105
106impl std::fmt::Display for BoxId {
107 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108 self.0.fmt(formatter)
109 }
110}
111
112#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize, Deserialize)]
113#[serde(transparent)]
114pub struct PendingId(String);
115
116impl PendingId {
117 pub fn from_event(id: EventId) -> Self {
118 Self(format!("pending:{}", id.0))
119 }
120
121 pub fn parse(value: impl Into<String>) -> anyhow::Result<Self> {
122 let value = value.into();
123 let number = value
124 .strip_prefix("pending:")
125 .context("pending identity must begin with `pending:`")?
126 .parse::<u64>()
127 .context("pending identity must end in an unsigned integer")?;
128 ensure!(number > 0, "pending identity zero is reserved");
129 Ok(Self(value))
130 }
131
132 pub fn number(&self) -> u64 {
133 self.0["pending:".len()..]
134 .parse()
135 .expect("validated PendingId")
136 }
137}
138
139impl std::fmt::Display for PendingId {
140 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
141 formatter.write_str(&self.0)
142 }
143}
144
145#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
146#[serde(rename_all = "snake_case")]
147pub enum SessionKind {
148 Conversation,
149 Telegram,
150 TelegramGroup,
151 SelfTime,
152 AudioIngress,
153 HistoryIngress,
154 Other(String),
155}
156
157#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
158#[serde(rename_all = "camelCase")]
159pub struct SessionMetadata {
160 pub session_id: String,
161 pub kind: SessionKind,
162 pub created_at: String,
163 pub effective_context_tokens: u64,
164 #[serde(default)]
165 pub channel: Value,
166}
167
168#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
169#[serde(tag = "kind", rename_all = "snake_case")]
170pub enum BoxOwner {
171 User,
172 Kennedy,
173 Controller,
174 System,
175 Tool { tool_instance: String, slot: String },
176}
177
178#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
179#[serde(rename_all = "camelCase")]
180pub struct BoxContent {
181 #[serde(default)]
182 pub text: String,
183 #[serde(default)]
184 pub objects: Vec<String>,
185 #[serde(default)]
186 pub metadata: Value,
187}
188
189impl BoxContent {
190 pub fn text(value: impl Into<String>) -> Self {
191 Self {
192 text: value.into(),
193 ..Self::default()
194 }
195 }
196
197 pub fn use_concise_header(&mut self) {
200 if !self.metadata.is_object() {
201 self.metadata = json!({});
202 }
203 self.metadata["chatendConciseHeader"] = json!(true);
204 }
205
206 fn render(&self) -> String {
207 let mut rendered = self.text.clone();
208 for object in &self.objects {
209 if !rendered.is_empty() && !rendered.ends_with('\n') {
210 rendered.push('\n');
211 }
212 rendered.push_str("Object provided: ");
213 rendered.push_str(object);
214 }
215 rendered
216 }
217}
218
219#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
220#[serde(tag = "kind", rename_all = "snake_case")]
221pub enum Representation {
222 Hydrated { canonical_event: EventId },
223 Dehydrated { based_on: EventId },
224 Summarized { based_on: EventId, text: String },
225}
226
227#[derive(Clone, Debug, Eq, PartialEq)]
228pub enum BoxRepresentation {
229 Hydrated,
230 Dehydrated,
231 Summarized(String),
232}
233
234impl Representation {
235 fn based_on(&self) -> EventId {
236 match self {
237 Self::Hydrated { canonical_event } => *canonical_event,
238 Self::Dehydrated { based_on } | Self::Summarized { based_on, .. } => *based_on,
239 }
240 }
241}
242
243#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
244#[serde(rename_all = "camelCase")]
245pub struct CanonicalRevision {
246 pub event_id: EventId,
247 pub content: BoxContent,
248}
249
250#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
251#[serde(rename_all = "camelCase")]
252pub struct BoxState {
253 pub id: BoxId,
254 pub name: String,
255 pub owner: BoxOwner,
256 pub created_at: EventId,
257 pub canonical: CanonicalRevision,
258 pub representation: Representation,
259 pub occurrence_events: Vec<EventId>,
260 pub active: bool,
261}
262
263impl BoxState {
264 pub fn stale(&self) -> bool {
265 !matches!(self.representation, Representation::Hydrated { .. })
266 && self.representation.based_on() != self.canonical.event_id
267 }
268}
269
270#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
271#[serde(rename_all = "snake_case")]
272pub enum PendingKind {
273 Node,
274 Object,
275}
276
277#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
278#[serde(tag = "type", rename_all = "snake_case")]
279pub enum EventKind {
280 SessionConfigured {
281 effective_context_tokens: u64,
282 kind: SessionKind,
283 },
284 BoxCreated {
285 box_id: BoxId,
286 name: String,
287 owner: BoxOwner,
288 content: BoxContent,
289 },
290 CanonicalUpdated {
291 box_id: BoxId,
292 content: BoxContent,
293 },
294 BoxRenamed {
295 box_id: BoxId,
296 name: String,
297 },
298 BoxDehydrated {
299 box_id: BoxId,
300 },
301 BoxSummarized {
302 box_id: BoxId,
303 text: String,
304 },
305 BoxRehydrated {
306 box_id: BoxId,
307 },
308 BoxRetired {
309 box_id: BoxId,
310 },
311 PendingAllocated {
312 pending_id: PendingId,
313 resource: PendingKind,
314 },
315 ToolInvoked {
316 tool_instance: String,
317 tool_name: String,
318 arguments: Value,
319 #[serde(default, skip_serializing_if = "Option::is_none")]
320 invocation_id: Option<String>,
321 },
322 ToolCompleted {
323 tool_instance: String,
324 tool_name: String,
325 outcome: Value,
326 #[serde(default, skip_serializing_if = "Option::is_none")]
327 invocation_id: Option<String>,
328 },
329 ToolLayoutChanged {
330 tool_instance: String,
331 box_ids: Vec<BoxId>,
332 },
333 InferenceSubmitted {
334 manifest_hash: String,
335 estimated_input_tokens: u64,
336 #[serde(default)]
337 raw_estimated_input_tokens: Option<u64>,
338 },
339 ProviderReceipt {
340 manifest_hash: String,
341 input_tokens: Option<u64>,
342 output_tokens: Option<u64>,
343 #[serde(default)]
346 context_bytes: Option<u64>,
347 #[serde(default)]
349 raw_context_tokens: Option<u64>,
350 provider_data: Value,
351 },
352 CapacityError {
353 attempted_operation: String,
354 projected_tokens: u64,
355 limit_tokens: u64,
356 },
357 SourceTerminated {
358 reason: String,
359 },
360 HistoryIngressStarted,
361 HistoryEventInspected {
362 source_event: EventId,
363 },
364 HistoryEventReleased {
365 source_event: EventId,
366 },
367 KwebPlanChanged {
368 operation: Value,
369 },
370 KwebCommitted {
371 transaction_id: String,
372 session_object_id: String,
373 mappings: Value,
374 },
375 SessionCompleted {
376 session_object_id: String,
377 },
378 Note {
379 label: String,
380 value: Value,
381 },
382}
383
384#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
385#[serde(rename_all = "camelCase")]
386pub struct Event {
387 pub id: EventId,
388 pub recorded_at: String,
389 pub kind: EventKind,
390}
391
392#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
393#[serde(rename_all = "camelCase")]
394pub struct Transition {
395 pub recorded_at: String,
396 pub events: Vec<Event>,
397}
398
399#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
400#[serde(rename_all = "camelCase")]
401struct PersistedContextEvent {
402 context_event_version: u32,
403 recorded_at: String,
404 kind: EventKind,
405}
406
407#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
408#[serde(rename_all = "camelCase")]
409struct PersistedContextEventWire {
410 context_event_version: u32,
411 recorded_at: String,
412 kind: Value,
413}
414
415#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
416#[serde(rename_all = "camelCase")]
417pub struct ObjectMetadata {
418 pub pending_id: PendingId,
419 pub event_id: EventId,
420 pub recorded_at: String,
421 pub media_type: String,
422 pub file_name: Option<String>,
423 #[serde(default)]
424 pub transport: Value,
425}
426
427#[derive(Clone, Debug, Eq, PartialEq)]
428pub struct ObjectLocation {
429 pub metadata: ObjectMetadata,
430 pub payload_offset: u64,
431 pub payload_len: u64,
432}
433
434#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
435#[serde(rename_all = "camelCase")]
436pub struct ToolSlot {
437 pub slot: String,
438 pub box_id: BoxId,
439 pub retired: bool,
440}
441
442#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
443#[serde(rename_all = "camelCase")]
444pub struct ToolState {
445 pub slots: Vec<ToolSlot>,
446}
447
448#[derive(Clone, Debug, Eq, PartialEq)]
449pub struct ToolSlotInput {
450 pub slot: String,
451 pub name: String,
452 pub content: BoxContent,
453 pub retired: bool,
454}
455
456#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
457#[serde(rename_all = "camelCase")]
458pub struct Chatend {
459 pub metadata: SessionMetadata,
460 pub next_id: u64,
461 pub events: Vec<Event>,
462 pub boxes: BTreeMap<BoxId, BoxState>,
463 pub pending: BTreeMap<PendingId, PendingKind>,
464 pub tools: BTreeMap<String, ToolState>,
465 #[serde(default)]
466 pub tool_layouts: BTreeMap<String, Vec<BoxId>>,
467 pub source_terminated: bool,
468 pub history_ingress_started: bool,
469 pub completed_session_object: Option<String>,
470}
471
472impl Chatend {
473 fn opened(metadata: SessionMetadata) -> Self {
474 Self {
475 metadata,
476 next_id: 1,
477 events: Vec::new(),
478 boxes: BTreeMap::new(),
479 pending: BTreeMap::new(),
480 tools: BTreeMap::new(),
481 tool_layouts: BTreeMap::new(),
482 source_terminated: false,
483 history_ingress_started: false,
484 completed_session_object: None,
485 }
486 }
487
488 pub(crate) fn replay(
489 metadata: SessionMetadata,
490 log: &kcode_session_log::SessionLog,
491 ) -> anyhow::Result<Self> {
492 Self::replay_with_provider_costs(metadata, log, None, None)
493 }
494
495 pub(crate) fn replay_with_provider_costs(
496 metadata: SessionMetadata,
497 log: &kcode_session_log::SessionLog,
498 default_provider_model: Option<&str>,
499 estimator: Option<ProviderCostEstimator>,
500 ) -> anyhow::Result<Self> {
501 ensure!(
502 metadata.session_id == log.header.session_id,
503 "session metadata and session-log identities differ"
504 );
505 ensure!(
506 metadata.created_at == log.header.created_at,
507 "session metadata and session-log creation times differ"
508 );
509 let mut chatend = Self::opened(metadata);
510 for (position, stored) in log.events.iter().enumerate() {
511 let persisted = decode_context_event(stored)?;
512 let id = EventId(position as u64 + 1);
513 let mut kind = persisted.kind;
514 normalize_derived_identity(&mut kind, id)?;
515 chatend.apply_transition(&Transition {
516 recorded_at: persisted.recorded_at.clone(),
517 events: vec![Event {
518 id,
519 recorded_at: persisted.recorded_at,
520 kind,
521 }],
522 })?;
523 }
524 if let Some(estimator) = estimator {
525 chatend.apply_legacy_provider_costs(default_provider_model, estimator);
526 }
527 Ok(chatend)
528 }
529
530 fn apply_legacy_provider_costs(
531 &mut self,
532 default_provider_model: Option<&str>,
533 estimator: ProviderCostEstimator,
534 ) {
535 let mut active_tool_models = Vec::<ActiveToolModel>::new();
536 let mut subagent_model = None::<String>;
537 for event in &mut self.events {
538 match &mut event.kind {
539 EventKind::ToolInvoked {
540 tool_instance,
541 tool_name,
542 arguments,
543 invocation_id,
544 } => {
545 let model = arguments
546 .get("model")
547 .and_then(Value::as_str)
548 .filter(|model| !model.trim().is_empty())
549 .map(normalize_requested_model);
550 active_tool_models.push(ActiveToolModel {
551 invocation_id: invocation_id.clone(),
552 tool_instance: tool_instance.clone(),
553 tool_name: tool_name.clone(),
554 model,
555 });
556 }
557 EventKind::ToolCompleted {
558 tool_instance,
559 tool_name,
560 invocation_id,
561 ..
562 } => {
563 if let Some(index) = active_tool_models.iter().rposition(|active| {
564 if let Some(invocation_id) = invocation_id {
565 active.invocation_id.as_ref() == Some(invocation_id)
566 } else {
567 active.tool_instance == *tool_instance && active.tool_name == *tool_name
568 }
569 }) {
570 active_tool_models.remove(index);
571 }
572 if tool_name == "RunSubagent" {
573 subagent_model = None;
574 }
575 }
576 EventKind::Note { label, value } if label == "subagent_started" => {
577 subagent_model = value
578 .get("providerModel")
579 .or_else(|| value.get("provider_model"))
580 .and_then(Value::as_str)
581 .filter(|model| !model.trim().is_empty())
582 .map(str::to_owned);
583 }
584 EventKind::ProviderReceipt { provider_data, .. }
585 if provider_data
586 .get("estimatedCostUsdNanos")
587 .and_then(Value::as_u64)
588 .is_none() =>
589 {
590 let source = provider_data.get("source").and_then(Value::as_str);
591 let model = provider_data
592 .get("providerModel")
593 .or_else(|| provider_data.get("provider_model"))
594 .and_then(Value::as_str)
595 .filter(|model| !model.trim().is_empty())
596 .map(str::to_owned)
597 .or_else(|| {
598 (source == Some("subagent"))
599 .then(|| subagent_model.clone())
600 .flatten()
601 })
602 .or_else(|| {
603 source
604 .is_some()
605 .then(|| {
606 active_tool_models
607 .iter()
608 .rev()
609 .find_map(|active| active.model.clone())
610 })
611 .flatten()
612 })
613 .or_else(|| {
614 source
615 .is_none()
616 .then(|| default_provider_model.map(str::to_owned))
617 .flatten()
618 });
619 let Some(model) = model else {
620 continue;
621 };
622 let metering = provider_metering(provider_data);
623 let Some(cost) = estimator(&model, &metering) else {
624 continue;
625 };
626 let Some(data) = provider_data.as_object_mut() else {
627 continue;
628 };
629 data.entry("providerModel")
630 .or_insert_with(|| Value::String(model));
631 data.insert("estimatedCostUsdNanos".into(), Value::from(cost.usd_nanos));
632 data.insert("costAccuracy".into(), cost.accuracy);
633 data.insert("pricingVersion".into(), Value::String(cost.pricing_version));
634 }
635 _ => {}
636 }
637 }
638 }
639
640 pub fn event(&self, id: EventId) -> Option<&Event> {
641 self.events.iter().find(|event| event.id == id)
642 }
643
644 pub fn box_state(&self, id: BoxId) -> Option<&BoxState> {
645 self.boxes.get(&id)
646 }
647
648 pub fn active_boxes(&self) -> impl Iterator<Item = &BoxState> {
649 self.boxes.values().filter(|state| state.active)
650 }
651
652 pub fn live_context_limit(&self) -> u64 {
653 self.metadata.effective_context_tokens.saturating_mul(70) / 100
654 }
655
656 pub fn forced_ingress_context_limit(&self) -> u64 {
657 self.metadata.effective_context_tokens.saturating_mul(75) / 100
658 }
659
660 pub fn ingress_initial_context_limit(&self) -> u64 {
661 self.metadata.effective_context_tokens.saturating_mul(75) / 100
662 }
663
664 pub fn ingress_context_limit(&self) -> u64 {
665 self.metadata.effective_context_tokens
666 }
667
668 pub fn active_context_limit(&self) -> u64 {
669 if matches!(
670 self.metadata.kind,
671 SessionKind::HistoryIngress | SessionKind::AudioIngress
672 ) {
673 self.ingress_context_limit()
674 } else {
675 self.live_context_limit()
676 }
677 }
678
679 pub fn projection_with_new_boxes(
680 &self,
681 boxes: &[(String, BoxOwner, BoxContent)],
682 ) -> anyhow::Result<ContextProjection> {
683 self.projection_with_new_boxes_at("preview", boxes)
684 }
685
686 pub fn projection_with_new_boxes_at(
687 &self,
688 recorded_at: &str,
689 boxes: &[(String, BoxOwner, BoxContent)],
690 ) -> anyhow::Result<ContextProjection> {
691 self.projection_with_new_boxes_and_updates_at(recorded_at, boxes, &BTreeMap::new())
692 }
693
694 pub fn projection_with_new_boxes_and_updates(
695 &self,
696 boxes: &[(String, BoxOwner, BoxContent)],
697 updates: &BTreeMap<BoxId, BoxContent>,
698 ) -> anyhow::Result<ContextProjection> {
699 self.projection_with_new_boxes_and_updates_at("preview", boxes, updates)
700 }
701
702 pub fn projection_with_new_boxes_and_updates_at(
703 &self,
704 recorded_at: &str,
705 boxes: &[(String, BoxOwner, BoxContent)],
706 updates: &BTreeMap<BoxId, BoxContent>,
707 ) -> anyhow::Result<ContextProjection> {
708 let mut preview = self.clone();
709 let mut next = preview.next_id;
710 let mut events = Vec::with_capacity(boxes.len() + updates.len());
711 for (name, owner, content) in boxes {
712 let id = EventId(next);
713 let box_id = BoxId(next);
714 next = next.checked_add(1).context("event identity overflow")?;
715 events.push(Event {
716 id,
717 recorded_at: recorded_at.into(),
718 kind: EventKind::BoxCreated {
719 box_id,
720 name: name.clone(),
721 owner: owner.clone(),
722 content: content.clone(),
723 },
724 });
725 }
726 for (box_id, content) in updates {
727 let state = preview
728 .box_state(*box_id)
729 .with_context(|| format!("box {box_id} does not exist"))?;
730 ensure!(state.active, "box {box_id} is retired");
731 if state.canonical.content == *content {
732 continue;
733 }
734 let id = EventId(next);
735 next = next.checked_add(1).context("event identity overflow")?;
736 events.push(Event {
737 id,
738 recorded_at: recorded_at.into(),
739 kind: EventKind::CanonicalUpdated {
740 box_id: *box_id,
741 content: content.clone(),
742 },
743 });
744 }
745 if !events.is_empty() {
746 preview.apply_transition(&Transition {
747 recorded_at: recorded_at.into(),
748 events,
749 })?;
750 }
751 Ok(preview.projection())
752 }
753
754 pub fn projection_with_box_representations(
755 &self,
756 desired: &BTreeMap<BoxId, BoxRepresentation>,
757 ) -> anyhow::Result<ContextProjection> {
758 let mut preview = self.clone();
759 let events = preview.box_representation_events("preview", desired)?;
760 if !events.is_empty() {
761 preview.apply_transition(&Transition {
762 recorded_at: "preview".into(),
763 events,
764 })?;
765 }
766 Ok(preview.projection())
767 }
768
769 fn box_representation_events(
770 &self,
771 recorded_at: &str,
772 desired: &BTreeMap<BoxId, BoxRepresentation>,
773 ) -> anyhow::Result<Vec<Event>> {
774 let mut next = self.next_id;
775 let mut events = Vec::new();
776 for (box_id, desired) in desired {
777 let state = self
778 .box_state(*box_id)
779 .with_context(|| format!("box {box_id} does not exist"))?;
780 ensure!(state.active, "box {box_id} is retired");
781 let kind = match (desired, &state.representation) {
782 (BoxRepresentation::Hydrated, Representation::Hydrated { .. })
783 | (BoxRepresentation::Dehydrated, Representation::Dehydrated { .. }) => None,
784 (
785 BoxRepresentation::Summarized(desired),
786 Representation::Summarized { text, .. },
787 ) if desired == text => None,
788 (BoxRepresentation::Hydrated, _) => {
789 Some(EventKind::BoxRehydrated { box_id: *box_id })
790 }
791 (BoxRepresentation::Dehydrated, _) => {
792 Some(EventKind::BoxDehydrated { box_id: *box_id })
793 }
794 (BoxRepresentation::Summarized(text), _) => Some(EventKind::BoxSummarized {
795 box_id: *box_id,
796 text: text.clone(),
797 }),
798 };
799 let Some(kind) = kind else {
800 continue;
801 };
802 events.push(Event {
803 id: EventId(next),
804 recorded_at: recorded_at.into(),
805 kind,
806 });
807 next = next.checked_add(1).context("event ID overflow")?;
808 }
809 Ok(events)
810 }
811
812 pub fn projection(&self) -> ContextProjection {
813 self.projection_at(¤t_time())
814 }
815
816 fn projection_at(&self, current_time: &str) -> ContextProjection {
817 let items = self.projection_items(false);
818 let stale_boxes = self
819 .active_boxes()
820 .filter(|state| state.stale())
821 .map(|state| state.id)
822 .collect::<Vec<_>>();
823 let stale_label = if stale_boxes.is_empty() {
824 "none".into()
825 } else {
826 stale_boxes
827 .iter()
828 .map(ToString::to_string)
829 .collect::<Vec<_>>()
830 .join(", ")
831 };
832 let mut estimated_tokens = estimate_bytes(projected_context_bytes(&items, ""));
833 let mut footer = String::new();
834 let mut context_bytes = 0;
835 let mut raw_estimated_tokens = 0;
836 for _ in 0..8 {
839 footer = format!(
840 "[stale boxes: {}]\n[current time: {}]\n[current context size: {} | max context size: {}]",
841 stale_label,
842 current_time,
843 estimated_tokens,
844 self.active_context_limit()
845 );
846 context_bytes = projected_context_bytes(&items, &footer);
847 raw_estimated_tokens = estimate_bytes(context_bytes);
848 let calibrated = self.calibrated_estimate(context_bytes, raw_estimated_tokens);
849 if calibrated == estimated_tokens {
850 break;
851 }
852 estimated_tokens = calibrated;
853 }
854 let fully_hydrated_context_tokens = self.fully_hydrated_context_tokens_at(current_time);
855 let usage = self.cumulative_token_usage();
856 let cost = self.cumulative_cost();
857 let status = SessionStatus {
858 current_context_tokens: estimated_tokens,
859 fully_hydrated_context_tokens,
860 context_limit_tokens: self.active_context_limit(),
861 current_context_bytes: context_bytes,
862 cached_input_tokens: usage.cached_input_tokens,
863 non_cached_input_tokens: usage.non_cached_input_tokens,
864 thinking_tokens: usage.thinking_tokens,
865 output_tokens: usage.output_tokens,
866 estimated_cost_usd_nanos: cost.estimated_cost_usd_nanos,
867 unpriced_provider_calls: cost.unpriced_provider_calls,
868 };
869 ContextProjection {
870 items,
871 stale_boxes,
872 footer,
873 estimated_tokens,
874 raw_estimated_tokens,
875 context_bytes,
876 status,
877 }
878 }
879
880 fn projection_items(&self, fully_hydrated: bool) -> Vec<ProjectionItem> {
881 let mut next_occurrence = HashMap::new();
882 for state in self.boxes.values() {
883 for pair in state.occurrence_events.windows(2) {
884 next_occurrence.insert(pair[0], pair[1]);
885 }
886 }
887 let mut items = Vec::new();
888 for event in &self.events {
889 let Some(box_id) = event_box_id(&event.kind) else {
890 continue;
891 };
892 let Some(state) = self.boxes.get(&box_id) else {
893 continue;
894 };
895 if next_occurrence.contains_key(&event.id) {
896 items.push(ProjectionItem::marker(
897 event.id,
898 box_id,
899 "[box updated]".into(),
900 ));
901 continue;
902 }
903 if !state.active || state.occurrence_events.last() != Some(&event.id) {
904 continue;
905 }
906 let (representation, body) = if fully_hydrated {
907 ("hydrated", state.canonical.content.render())
908 } else {
909 match &state.representation {
910 Representation::Hydrated { .. } => {
911 ("hydrated", state.canonical.content.render())
912 }
913 Representation::Dehydrated { .. } => (
914 "dehydrated",
915 format!(
916 "[contents dehydrated; hydrate box {} to inspect the latest canonical revision]",
917 box_id
918 ),
919 ),
920 Representation::Summarized { text, .. } => ("summarized", text.clone()),
921 }
922 };
923 let stale = !fully_hydrated && state.stale();
924 let mut header = vec![format!("box {box_id}"), state.name.clone()];
925 if matches!(
926 (&state.owner, state.name.as_str()),
927 (BoxOwner::User, "User message") | (BoxOwner::Kennedy, "Kennedy message")
928 ) {
929 header.push(format!("timestamp={}", event.recorded_at));
930 }
931 header.push(representation.into());
932 if stale {
933 header.push("stale".into());
934 }
935 let mut text = format!("[{}]\n{}", header.join(" | "), body);
936 if text.ends_with('\n') {
937 text.pop();
938 }
939 items.push(ProjectionItem {
940 event_id: event.id,
941 box_id,
942 marker: false,
943 stale,
944 approximate_tokens: estimate_tokens(&text),
945 text,
946 });
947 }
948 for box_ids in self.tool_layouts.values() {
949 arrange_tool_projection(&mut items, box_ids);
950 }
951 items
952 }
953
954 fn fully_hydrated_context_tokens_at(&self, current_time: &str) -> u64 {
955 let items = self.projection_items(true);
956 let mut estimated_tokens = estimate_bytes(projected_context_bytes(&items, ""));
957 for _ in 0..8 {
958 let footer = format!(
959 "[stale boxes: {}]\n[current time: {}]\n[current context size: {} | max context size: {}]",
960 "none",
961 current_time,
962 estimated_tokens,
963 self.active_context_limit()
964 );
965 let context_bytes = projected_context_bytes(&items, &footer);
966 let raw_estimated_tokens = estimate_bytes(context_bytes);
967 let calibrated = self.calibrated_estimate(context_bytes, raw_estimated_tokens);
968 if calibrated == estimated_tokens {
969 break;
970 }
971 estimated_tokens = calibrated;
972 }
973 estimated_tokens
974 }
975
976 fn calibrated_estimate(&self, current_bytes: u64, raw_current: u64) -> u64 {
977 let Some((manifest_hash, measured, bytes_at_receipt, raw_at_receipt)) =
978 self.events.iter().rev().find_map(|event| {
979 let EventKind::ProviderReceipt {
980 manifest_hash,
981 input_tokens: Some(input_tokens),
982 context_bytes,
983 raw_context_tokens,
984 ..
985 } = &event.kind
986 else {
987 return None;
988 };
989 Some((
990 manifest_hash,
991 *input_tokens,
992 *context_bytes,
993 *raw_context_tokens,
994 ))
995 })
996 else {
997 return raw_current;
998 };
999 if let Some(bytes_at_receipt) = bytes_at_receipt {
1000 let token_delta = current_bytes.abs_diff(bytes_at_receipt) / ESTIMATED_BYTES_PER_TOKEN;
1001 return if current_bytes >= bytes_at_receipt {
1002 measured.saturating_add(token_delta)
1003 } else {
1004 measured.saturating_sub(token_delta)
1005 };
1006 }
1007 let raw_at_measurement = match raw_at_receipt {
1008 Some(raw) => raw,
1009 None => {
1010 let Some(raw) = self.events.iter().rev().find_map(|event| {
1011 let EventKind::InferenceSubmitted {
1012 manifest_hash: submitted,
1013 estimated_input_tokens,
1014 raw_estimated_input_tokens,
1015 } = &event.kind
1016 else {
1017 return None;
1018 };
1019 (submitted == manifest_hash)
1020 .then_some(raw_estimated_input_tokens.unwrap_or(*estimated_input_tokens))
1021 }) else {
1022 return raw_current;
1023 };
1024 raw
1025 }
1026 };
1027 if raw_current >= raw_at_measurement {
1028 measured.saturating_add(raw_current - raw_at_measurement)
1029 } else {
1030 measured.saturating_sub(raw_at_measurement - raw_current)
1031 }
1032 }
1033
1034 fn cumulative_token_usage(&self) -> CumulativeTokenUsage {
1035 let mut total = CumulativeTokenUsage::default();
1036 for event in &self.events {
1037 let EventKind::ProviderReceipt { provider_data, .. } = &event.kind else {
1038 continue;
1039 };
1040 let cached = provider_u64(provider_data, &["cachedInputTokens", "cached_input_tokens"]);
1041 let thinking = provider_u64(
1042 provider_data,
1043 &[
1044 "thinkingTokens",
1045 "thinking_tokens",
1046 "reasoningOutputTokens",
1047 "reasoning_output_tokens",
1048 ],
1049 );
1050 let normalized_delta = provider_data
1051 .get("usageIsDelta")
1052 .and_then(Value::as_bool)
1053 .unwrap_or(false);
1054 let non_cached = if normalized_delta {
1055 provider_u64(
1056 provider_data,
1057 &["nonCachedInputTokens", "non_cached_input_tokens"],
1058 )
1059 } else {
1060 provider_u64(provider_data, &["inputTokens", "input_tokens"]).saturating_sub(cached)
1061 };
1062 let output = if normalized_delta {
1063 provider_u64(provider_data, &["outputTokens", "output_tokens"])
1064 } else {
1065 provider_u64(provider_data, &["outputTokens", "output_tokens"])
1066 .saturating_sub(thinking)
1067 };
1068 total.cached_input_tokens = total.cached_input_tokens.saturating_add(cached);
1069 total.non_cached_input_tokens =
1070 total.non_cached_input_tokens.saturating_add(non_cached);
1071 total.thinking_tokens = total.thinking_tokens.saturating_add(thinking);
1072 total.output_tokens = total.output_tokens.saturating_add(output);
1073 }
1074 total
1075 }
1076
1077 fn cumulative_cost(&self) -> CumulativeCost {
1078 let mut total = CumulativeCost::default();
1079 for event in &self.events {
1080 let EventKind::ProviderReceipt { provider_data, .. } = &event.kind else {
1081 continue;
1082 };
1083 if let Some(cost) = provider_data
1084 .get("estimatedCostUsdNanos")
1085 .and_then(Value::as_u64)
1086 {
1087 total.estimated_cost_usd_nanos =
1088 total.estimated_cost_usd_nanos.saturating_add(cost);
1089 } else {
1090 total.unpriced_provider_calls = total.unpriced_provider_calls.saturating_add(1);
1091 }
1092 }
1093 total
1094 }
1095
1096 pub fn render(&self) -> String {
1097 self.projection().render()
1098 }
1099
1100 fn apply_transition(&mut self, transition: &Transition) -> anyhow::Result<()> {
1101 ensure!(
1102 !transition.events.is_empty(),
1103 "a transition cannot be empty"
1104 );
1105 for event in &transition.events {
1106 self.apply_event(event)?;
1107 }
1108 Ok(())
1109 }
1110
1111 fn apply_event(&mut self, event: &Event) -> anyhow::Result<()> {
1112 ensure!(
1113 event.id.0 >= self.next_id,
1114 "event {} reuses an allocated identity (next is {})",
1115 event.id,
1116 self.next_id
1117 );
1118 self.next_id = event.id.0.checked_add(1).context("event ID overflow")?;
1119 match &event.kind {
1120 EventKind::SessionConfigured {
1121 effective_context_tokens,
1122 kind,
1123 } => {
1124 ensure!(
1125 *effective_context_tokens > 0,
1126 "effective context window must be positive"
1127 );
1128 self.metadata.effective_context_tokens = *effective_context_tokens;
1129 self.metadata.kind = kind.clone();
1130 }
1131 EventKind::BoxCreated {
1132 box_id,
1133 name,
1134 owner,
1135 content,
1136 } => {
1137 ensure!(
1138 box_id.0 == event.id.0,
1139 "BoxId must equal its creation EventId"
1140 );
1141 ensure!(
1142 !self.boxes.contains_key(box_id),
1143 "box {} already exists",
1144 box_id
1145 );
1146 self.boxes.insert(
1147 *box_id,
1148 BoxState {
1149 id: *box_id,
1150 name: name.clone(),
1151 owner: owner.clone(),
1152 created_at: event.id,
1153 canonical: CanonicalRevision {
1154 event_id: event.id,
1155 content: content.clone(),
1156 },
1157 representation: Representation::Hydrated {
1158 canonical_event: event.id,
1159 },
1160 occurrence_events: vec![event.id],
1161 active: true,
1162 },
1163 );
1164 if let BoxOwner::Tool {
1165 tool_instance,
1166 slot,
1167 } = owner
1168 {
1169 self.tools
1170 .entry(tool_instance.clone())
1171 .or_default()
1172 .slots
1173 .push(ToolSlot {
1174 slot: slot.clone(),
1175 box_id: *box_id,
1176 retired: false,
1177 });
1178 }
1179 }
1180 EventKind::CanonicalUpdated { box_id, content } => {
1181 let state = active_box_mut(&mut self.boxes, *box_id)?;
1182 state.canonical = CanonicalRevision {
1183 event_id: event.id,
1184 content: content.clone(),
1185 };
1186 if matches!(state.representation, Representation::Hydrated { .. }) {
1187 state.representation = Representation::Hydrated {
1188 canonical_event: event.id,
1189 };
1190 }
1191 state.occurrence_events.push(event.id);
1192 }
1193 EventKind::BoxRenamed { box_id, name } => {
1194 ensure!(!name.trim().is_empty(), "a box name cannot be empty");
1195 let state = active_box_mut(&mut self.boxes, *box_id)?;
1196 state.name = name.clone();
1197 state.occurrence_events.push(event.id);
1198 }
1199 EventKind::BoxDehydrated { box_id } => {
1200 let state = active_box_mut(&mut self.boxes, *box_id)?;
1201 state.representation = Representation::Dehydrated {
1202 based_on: state.canonical.event_id,
1203 };
1204 state.occurrence_events.push(event.id);
1205 }
1206 EventKind::BoxSummarized { box_id, text } => {
1207 ensure!(!text.trim().is_empty(), "a box summary cannot be empty");
1208 let state = active_box_mut(&mut self.boxes, *box_id)?;
1209 state.representation = Representation::Summarized {
1210 based_on: state.canonical.event_id,
1211 text: text.clone(),
1212 };
1213 state.occurrence_events.push(event.id);
1214 }
1215 EventKind::BoxRehydrated { box_id } => {
1216 let state = active_box_mut(&mut self.boxes, *box_id)?;
1217 state.representation = Representation::Hydrated {
1218 canonical_event: state.canonical.event_id,
1219 };
1220 state.occurrence_events.push(event.id);
1221 }
1222 EventKind::BoxRetired { box_id } => {
1223 let tool_instance = self.boxes.get(box_id).and_then(|state| {
1224 let BoxOwner::Tool { tool_instance, .. } = &state.owner else {
1225 return None;
1226 };
1227 Some(tool_instance.clone())
1228 });
1229 let state = active_box_mut(&mut self.boxes, *box_id)?;
1230 state.active = false;
1231 state.occurrence_events.push(event.id);
1232 if let Some(tool_instance) = tool_instance {
1233 let slot = self
1234 .tools
1235 .get_mut(&tool_instance)
1236 .and_then(|tool| tool.slots.iter_mut().find(|slot| slot.box_id == *box_id))
1237 .with_context(|| {
1238 format!("tool box {box_id} is missing from {tool_instance}")
1239 })?;
1240 slot.retired = true;
1241 }
1242 }
1243 EventKind::ToolLayoutChanged {
1244 tool_instance,
1245 box_ids,
1246 } => {
1247 let mut unique = std::collections::HashSet::new();
1248 for box_id in box_ids {
1249 ensure!(
1250 unique.insert(*box_id),
1251 "tool layout contains duplicate box {box_id}"
1252 );
1253 let state = self
1254 .boxes
1255 .get(box_id)
1256 .with_context(|| format!("tool layout references missing box {box_id}"))?;
1257 ensure!(state.active, "tool layout references retired box {box_id}");
1258 ensure!(
1259 matches!(
1260 &state.owner,
1261 BoxOwner::Tool {
1262 tool_instance: owner,
1263 ..
1264 } if owner == tool_instance
1265 ),
1266 "tool layout box {box_id} belongs to another tool"
1267 );
1268 }
1269 self.tool_layouts
1270 .insert(tool_instance.clone(), box_ids.clone());
1271 }
1272 EventKind::PendingAllocated {
1273 pending_id,
1274 resource,
1275 } => {
1276 ensure!(
1277 pending_id.number() == event.id.0,
1278 "pending identity must equal its allocation EventId"
1279 );
1280 ensure!(
1281 self.pending
1282 .insert(pending_id.clone(), resource.clone())
1283 .is_none(),
1284 "pending identity {} already exists",
1285 pending_id
1286 );
1287 }
1288 EventKind::SourceTerminated { .. } => self.source_terminated = true,
1289 EventKind::HistoryIngressStarted => {
1290 ensure!(
1291 self.source_terminated,
1292 "history ingress requires source termination"
1293 );
1294 self.history_ingress_started = true;
1295 }
1296 EventKind::SessionCompleted { session_object_id } => {
1297 self.completed_session_object = Some(session_object_id.clone());
1298 }
1299 EventKind::ToolInvoked { .. }
1300 | EventKind::ToolCompleted { .. }
1301 | EventKind::InferenceSubmitted { .. }
1302 | EventKind::ProviderReceipt { .. }
1303 | EventKind::CapacityError { .. }
1304 | EventKind::HistoryEventInspected { .. }
1305 | EventKind::HistoryEventReleased { .. }
1306 | EventKind::KwebPlanChanged { .. }
1307 | EventKind::KwebCommitted { .. }
1308 | EventKind::Note { .. } => {}
1309 }
1310 self.events.push(event.clone());
1311 Ok(())
1312 }
1313}
1314
1315fn active_box_mut(
1316 boxes: &mut BTreeMap<BoxId, BoxState>,
1317 box_id: BoxId,
1318) -> anyhow::Result<&mut BoxState> {
1319 let state = boxes
1320 .get_mut(&box_id)
1321 .with_context(|| format!("box {box_id} does not exist"))?;
1322 ensure!(state.active, "box {box_id} is retired");
1323 Ok(state)
1324}
1325
1326fn arrange_tool_projection(items: &mut Vec<ProjectionItem>, box_ids: &[BoxId]) {
1327 if box_ids.is_empty() {
1328 return;
1329 }
1330 let ranks = box_ids
1331 .iter()
1332 .enumerate()
1333 .map(|(rank, box_id)| (*box_id, rank))
1334 .collect::<HashMap<_, _>>();
1335 let insertion = items
1336 .iter()
1337 .position(|item| !item.marker && ranks.contains_key(&item.box_id));
1338 let Some(insertion) = insertion else {
1339 return;
1340 };
1341 let mut arranged = Vec::with_capacity(box_ids.len());
1342 let mut retained = Vec::with_capacity(items.len());
1343 for item in std::mem::take(items) {
1344 if !item.marker && ranks.contains_key(&item.box_id) {
1345 arranged.push(item);
1346 } else {
1347 retained.push(item);
1348 }
1349 }
1350 arranged.sort_by_key(|item| ranks[&item.box_id]);
1351 let insertion = insertion.min(retained.len());
1352 retained.splice(insertion..insertion, arranged);
1353 *items = retained;
1354}
1355
1356fn event_box_id(kind: &EventKind) -> Option<BoxId> {
1357 match kind {
1358 EventKind::BoxCreated { box_id, .. }
1359 | EventKind::CanonicalUpdated { box_id, .. }
1360 | EventKind::BoxRenamed { box_id, .. }
1361 | EventKind::BoxDehydrated { box_id }
1362 | EventKind::BoxSummarized { box_id, .. }
1363 | EventKind::BoxRehydrated { box_id }
1364 | EventKind::BoxRetired { box_id } => Some(*box_id),
1365 _ => None,
1366 }
1367}
1368
1369#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
1370#[serde(rename_all = "camelCase")]
1371pub struct ProjectionItem {
1372 pub event_id: EventId,
1373 pub box_id: BoxId,
1374 pub marker: bool,
1375 pub stale: bool,
1376 pub approximate_tokens: u64,
1377 pub text: String,
1378}
1379
1380impl ProjectionItem {
1381 fn marker(event_id: EventId, box_id: BoxId, text: String) -> Self {
1382 Self {
1383 event_id,
1384 box_id,
1385 marker: true,
1386 stale: false,
1387 approximate_tokens: estimate_tokens(&text),
1388 text,
1389 }
1390 }
1391}
1392
1393#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
1394#[serde(rename_all = "camelCase")]
1395pub struct ContextProjection {
1396 pub items: Vec<ProjectionItem>,
1397 pub stale_boxes: Vec<BoxId>,
1398 pub footer: String,
1399 pub estimated_tokens: u64,
1400 pub raw_estimated_tokens: u64,
1401 pub context_bytes: u64,
1402 pub status: SessionStatus,
1403}
1404
1405impl ContextProjection {
1406 pub fn render(&self) -> String {
1407 let mut blocks = self
1408 .items
1409 .iter()
1410 .map(|item| item.text.as_str())
1411 .collect::<Vec<_>>();
1412 blocks.push(&self.footer);
1413 blocks.join("\n\n")
1414 }
1415}
1416
1417#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
1418#[serde(rename_all = "camelCase")]
1419pub struct SessionStatus {
1420 pub current_context_tokens: u64,
1421 #[serde(default)]
1422 pub fully_hydrated_context_tokens: u64,
1423 pub context_limit_tokens: u64,
1424 pub current_context_bytes: u64,
1425 pub cached_input_tokens: u64,
1426 pub non_cached_input_tokens: u64,
1427 pub thinking_tokens: u64,
1428 pub output_tokens: u64,
1429 #[serde(default)]
1430 pub estimated_cost_usd_nanos: u64,
1431 #[serde(default)]
1432 pub unpriced_provider_calls: u64,
1433}
1434
1435#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1436struct CumulativeTokenUsage {
1437 cached_input_tokens: u64,
1438 non_cached_input_tokens: u64,
1439 thinking_tokens: u64,
1440 output_tokens: u64,
1441}
1442
1443#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1444struct CumulativeCost {
1445 estimated_cost_usd_nanos: u64,
1446 unpriced_provider_calls: u64,
1447}
1448
1449pub fn estimate_tokens(text: &str) -> u64 {
1450 estimate_bytes(text.len() as u64)
1451}
1452
1453fn estimate_bytes(bytes: u64) -> u64 {
1454 bytes.div_ceil(ESTIMATED_BYTES_PER_TOKEN)
1455}
1456
1457fn current_time() -> String {
1458 Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true)
1459}
1460
1461fn projected_context_bytes(items: &[ProjectionItem], footer: &str) -> u64 {
1462 items
1463 .iter()
1464 .fold(footer.len() as u64, |total, item| {
1465 total.saturating_add(item.text.len() as u64)
1466 })
1467 .saturating_add((items.len() as u64).saturating_mul(2))
1468}
1469
1470fn provider_u64(value: &Value, keys: &[&str]) -> u64 {
1471 keys.iter()
1472 .find_map(|key| value.get(*key).and_then(Value::as_u64))
1473 .unwrap_or_default()
1474}
1475
1476struct ActiveToolModel {
1477 invocation_id: Option<String>,
1478 tool_instance: String,
1479 tool_name: String,
1480 model: Option<String>,
1481}
1482
1483fn normalize_requested_model(model: &str) -> String {
1484 model.strip_prefix("codex/").unwrap_or(model).to_owned()
1485}
1486
1487fn provider_metering(provider_data: &Value) -> ProviderMetering {
1488 let metering = provider_data.get("metering");
1489 if metering
1490 .and_then(Value::as_str)
1491 .is_some_and(|kind| kind == "unavailable")
1492 {
1493 return ProviderMetering::Unavailable;
1494 }
1495 if let Some(metering_value) = metering
1496 && let Some(metering) = metering_value.as_object()
1497 {
1498 match metering.get("kind").and_then(Value::as_str) {
1499 Some("duration_seconds") => {
1500 return metering
1501 .get("seconds")
1502 .and_then(Value::as_f64)
1503 .map(|seconds| ProviderMetering::DurationSeconds { seconds })
1504 .unwrap_or(ProviderMetering::Unavailable);
1505 }
1506 Some("unavailable") => return ProviderMetering::Unavailable,
1507 Some("tokens") => {
1508 return ProviderMetering::Tokens(token_metering(metering_value));
1509 }
1510 _ => {}
1511 }
1512 }
1513 let has_tokens = metering.and_then(Value::as_str) == Some("tokens")
1514 || [
1515 "inputTokens",
1516 "input_tokens",
1517 "nonCachedInputTokens",
1518 "non_cached_input_tokens",
1519 "cachedInputTokens",
1520 "cached_input_tokens",
1521 "thinkingTokens",
1522 "thinking_tokens",
1523 "outputTokens",
1524 "output_tokens",
1525 ]
1526 .iter()
1527 .any(|key| provider_data.get(*key).and_then(Value::as_u64).is_some());
1528 if has_tokens {
1529 ProviderMetering::Tokens(token_metering(provider_data))
1530 } else {
1531 ProviderMetering::Unavailable
1532 }
1533}
1534
1535fn token_metering(value: &Value) -> ProviderTokenUsage {
1536 let cached_input_tokens = provider_u64(value, &["cachedInputTokens", "cached_input_tokens"]);
1537 let thinking_tokens = provider_u64(
1538 value,
1539 &[
1540 "thinkingTokens",
1541 "thinking_tokens",
1542 "reasoningOutputTokens",
1543 "reasoning_output_tokens",
1544 ],
1545 );
1546 let normalized_delta = value
1547 .get("usageIsDelta")
1548 .or_else(|| value.get("usage_is_delta"))
1549 .and_then(Value::as_bool)
1550 .unwrap_or(false);
1551 let input_tokens = if normalized_delta {
1552 provider_u64(value, &["nonCachedInputTokens", "non_cached_input_tokens"])
1553 } else {
1554 provider_u64(value, &["inputTokens", "input_tokens"]).saturating_sub(cached_input_tokens)
1555 };
1556 let output_tokens = if normalized_delta {
1557 provider_u64(value, &["outputTokens", "output_tokens"])
1558 } else {
1559 provider_u64(value, &["outputTokens", "output_tokens"]).saturating_sub(thinking_tokens)
1560 };
1561 ProviderTokenUsage {
1562 input_tokens,
1563 cached_input_tokens,
1564 thinking_tokens,
1565 output_tokens,
1566 }
1567}
1568
1569fn context_event_role(kind: &EventKind) -> Role {
1570 match kind {
1571 EventKind::BoxCreated { owner, content, .. } => match owner {
1572 BoxOwner::System | BoxOwner::Controller => {
1573 if content
1574 .metadata
1575 .get("capacityError")
1576 .and_then(Value::as_bool)
1577 .unwrap_or(false)
1578 {
1579 Role::SystemError
1580 } else {
1581 Role::SystemMessage
1582 }
1583 }
1584 BoxOwner::User => Role::UserMessage,
1585 BoxOwner::Kennedy => Role::KennedyMessage,
1586 BoxOwner::Tool { .. } => Role::ToolResult,
1587 },
1588 EventKind::ToolInvoked { .. } => Role::KennedyToolCall,
1589 EventKind::ToolCompleted { outcome, .. } => {
1590 if outcome.get("ok").and_then(Value::as_bool).unwrap_or(true) {
1591 Role::ToolResult
1592 } else {
1593 Role::ToolError
1594 }
1595 }
1596 EventKind::CapacityError { .. } => Role::SystemError,
1597 EventKind::PendingAllocated {
1598 resource: PendingKind::Object,
1599 ..
1600 } => Role::PendingObject,
1601 EventKind::BoxDehydrated { .. }
1602 | EventKind::BoxSummarized { .. }
1603 | EventKind::BoxRehydrated { .. }
1604 | EventKind::BoxRetired { .. } => Role::KennedyToolCall,
1605 _ => Role::SystemMessage,
1606 }
1607}
1608
1609fn encode_context_event(recorded_at: &str, kind: &EventKind) -> anyhow::Result<String> {
1610 let mut kind = serde_json::to_value(kind)?;
1611 let kind_object = kind
1612 .as_object_mut()
1613 .context("Kennedy context event kind must encode as an object")?;
1614 match kind_object.get("type").and_then(Value::as_str) {
1615 Some("box_created") => {
1616 kind_object.remove("box_id");
1617 }
1618 Some("pending_allocated") => {
1619 kind_object.remove("pending_id");
1620 }
1621 _ => {}
1622 }
1623 Ok(serde_json::to_string(&PersistedContextEventWire {
1624 context_event_version: FORMAT_VERSION,
1625 recorded_at: recorded_at.into(),
1626 kind,
1627 })?)
1628}
1629
1630fn decode_context_event(
1631 event: &kcode_session_log::SessionEvent,
1632) -> anyhow::Result<PersistedContextEvent> {
1633 let wire: PersistedContextEventWire = match serde_json::from_str(&event.text) {
1634 Ok(persisted) => persisted,
1635 Err(_) if event.role == Role::PendingObject => {
1636 return Ok(PersistedContextEvent {
1637 context_event_version: FORMAT_VERSION,
1638 recorded_at: String::new(),
1639 kind: EventKind::PendingAllocated {
1640 pending_id: PendingId::from_event(EventId(1)),
1641 resource: PendingKind::Object,
1642 },
1643 });
1644 }
1645 Err(error) => {
1646 return Err(error).context("decoding Kennedy context event from session log");
1647 }
1648 };
1649 ensure!(
1650 wire.context_event_version == FORMAT_VERSION,
1651 "unsupported Kennedy context event version {}",
1652 wire.context_event_version
1653 );
1654 let mut kind = wire.kind;
1655 let kind_object = kind
1656 .as_object_mut()
1657 .context("Kennedy context event kind must be an object")?;
1658 match kind_object.get("type").and_then(Value::as_str) {
1659 Some("box_created") if !kind_object.contains_key("box_id") => {
1660 kind_object.insert("box_id".into(), Value::from(0));
1661 }
1662 Some("pending_allocated") if !kind_object.contains_key("pending_id") => {
1663 kind_object.insert("pending_id".into(), Value::String("pending:1".into()));
1664 }
1665 _ => {}
1666 }
1667 Ok(PersistedContextEvent {
1668 context_event_version: wire.context_event_version,
1669 recorded_at: wire.recorded_at,
1670 kind: serde_json::from_value(kind).context("decoding Kennedy context event kind")?,
1671 })
1672}
1673
1674fn normalize_derived_identity(kind: &mut EventKind, id: EventId) -> anyhow::Result<()> {
1675 match kind {
1676 EventKind::BoxCreated { box_id, .. } => *box_id = BoxId(id.0),
1677 EventKind::PendingAllocated { pending_id, .. } => {
1678 *pending_id = PendingId::from_event(id);
1679 }
1680 _ => {}
1681 }
1682 Ok(())
1683}
1684
1685pub struct Session {
1686 durable: DurableSession,
1687 chatend: Chatend,
1688 objects: BTreeMap<PendingId, ObjectLocation>,
1689}
1690
1691struct UnfinishedToolInvocation {
1692 invocation_id: Option<String>,
1693 tool_instance: String,
1694 tool_name: String,
1695}
1696
1697impl Session {
1698 pub(crate) fn create(
1699 path: impl AsRef<Path>,
1700 metadata: SessionMetadata,
1701 ) -> anyhow::Result<Self> {
1702 ensure!(
1703 metadata.effective_context_tokens > 0,
1704 "effective context window must be positive"
1705 );
1706 let requested = path.as_ref();
1707 let directory = requested
1708 .parent()
1709 .filter(|path| !path.as_os_str().is_empty())
1710 .unwrap_or_else(|| Path::new("."));
1711 let durable = DurableSessionStore::new(directory)
1712 .create_session(&metadata.session_id, &metadata.created_at)?;
1713 Ok(Self {
1714 durable,
1715 chatend: Chatend::opened(metadata),
1716 objects: BTreeMap::new(),
1717 })
1718 }
1719
1720 pub(crate) fn open_with_metadata(
1721 path: impl AsRef<Path>,
1722 metadata: SessionMetadata,
1723 ) -> anyhow::Result<Self> {
1724 Self::open_with_metadata_and_provider_costs(path, metadata, None, None)
1725 }
1726
1727 pub(crate) fn open_with_metadata_and_provider_costs(
1728 path: impl AsRef<Path>,
1729 metadata: SessionMetadata,
1730 default_provider_model: Option<&str>,
1731 estimator: Option<ProviderCostEstimator>,
1732 ) -> anyhow::Result<Self> {
1733 let requested = path.as_ref();
1734 ensure!(
1735 requested.extension().and_then(|value| value.to_str()) == Some("session-log"),
1736 "{} is not a session-log path",
1737 requested.display()
1738 );
1739 let directory = requested
1740 .parent()
1741 .filter(|path| !path.as_os_str().is_empty())
1742 .unwrap_or_else(|| Path::new("."));
1743 let session_id = requested
1744 .file_stem()
1745 .and_then(|value| value.to_str())
1746 .context("session-log filename is not valid UTF-8")?;
1747 let durable = DurableSessionStore::new(directory).open_session(session_id)?;
1748 let log = durable.list();
1749 let chatend =
1750 Chatend::replay_with_provider_costs(metadata, &log, default_provider_model, estimator)?;
1751 let mut objects = BTreeMap::new();
1752 for (position, stored) in log.events.iter().enumerate() {
1753 let persisted = decode_context_event(stored)?;
1754 let id = EventId(position as u64 + 1);
1755 let mut kind = persisted.kind;
1756 normalize_derived_identity(&mut kind, id)?;
1757 if let EventKind::PendingAllocated {
1758 pending_id,
1759 resource: PendingKind::Object,
1760 } = kind
1761 {
1762 let object = durable.read_pending_object(EventPosition(position as u64))?;
1763 let metadata = ObjectMetadata {
1764 pending_id: pending_id.clone(),
1765 event_id: id,
1766 recorded_at: chatend
1767 .event(id)
1768 .map(|event| event.recorded_at.clone())
1769 .unwrap_or_default(),
1770 media_type: object.media_type,
1771 file_name: Some(object.file_name),
1772 transport: Value::Null,
1773 };
1774 objects.insert(
1775 pending_id,
1776 ObjectLocation {
1777 metadata,
1778 payload_offset: position as u64,
1779 payload_len: object.bytes.len() as u64,
1780 },
1781 );
1782 }
1783 }
1784 Ok(Self {
1785 durable,
1786 chatend,
1787 objects,
1788 })
1789 }
1790
1791 pub fn id(&self) -> &str {
1792 &self.chatend.metadata.session_id
1793 }
1794
1795 pub fn state(&self) -> &Chatend {
1796 &self.chatend
1797 }
1798
1799 pub fn objects(&self) -> &BTreeMap<PendingId, ObjectLocation> {
1800 &self.objects
1801 }
1802
1803 #[cfg(test)]
1804 fn session_log(&self) -> kcode_session_log::SessionLog {
1805 self.durable.list()
1806 }
1807
1808 pub fn archive_bytes(&self) -> anyhow::Result<Vec<u8>> {
1809 let mut archive =
1810 serde_json::to_value(self.durable.list()).context("serializing the session log")?;
1811 let object = archive
1812 .as_object_mut()
1813 .context("serialized session log is not an object")?;
1814 object.insert(
1815 "metadata".into(),
1816 serde_json::to_value(&self.chatend.metadata)?,
1817 );
1818 object.insert("boxes".into(), serde_json::to_value(&self.chatend.boxes)?);
1819 let projection = self.chatend.projection();
1820 object.insert("chatendText".into(), Value::String(projection.render()));
1821 object.insert("context".into(), serde_json::to_value(projection)?);
1822 serde_json::to_vec(&archive).context("serializing the session archive")
1823 }
1824
1825 pub fn is_sealed(&self) -> bool {
1826 self.durable.is_sealed()
1827 }
1828
1829 pub fn seal(&mut self) -> anyhow::Result<()> {
1830 let unfinished_tools = self.unfinished_tool_invocations()?;
1831 ensure!(
1832 unfinished_tools.is_empty(),
1833 "session ends with unfinished tools {}",
1834 unfinished_tools
1835 .iter()
1836 .map(|tool| tool.tool_name.as_str())
1837 .collect::<Vec<_>>()
1838 .join(", ")
1839 );
1840 self.durable.seal()?;
1841 Ok(())
1842 }
1843
1844 pub fn repair_unfinished_tools(
1845 &mut self,
1846 recorded_at: impl Into<String>,
1847 ) -> anyhow::Result<Vec<EventId>> {
1848 let unfinished = self.unfinished_tool_invocations()?;
1849 if unfinished.is_empty() {
1850 return Ok(Vec::new());
1851 }
1852 let recorded_at = recorded_at.into();
1853 let mut repaired = Vec::with_capacity(unfinished.len());
1854 for tool in unfinished.iter().rev() {
1855 let message = format!(
1856 "{} was interrupted before a durable completion was recorded; the abandoned invocation was closed during session recovery.",
1857 tool.tool_name
1858 );
1859 let kind = if let Some(invocation_id) = &tool.invocation_id {
1860 EventKind::ToolCompleted {
1861 tool_instance: tool.tool_instance.clone(),
1862 tool_name: tool.tool_name.clone(),
1863 outcome: serde_json::json!({"ok":false,"recovered":true,"result":message}),
1864 invocation_id: Some(invocation_id.clone()),
1865 }
1866 } else {
1867 EventKind::ToolCompleted {
1870 tool_instance: "call_ktool".into(),
1871 tool_name: "call_ktool".into(),
1872 outcome: serde_json::json!({"ok":false,"recovered":true,"result":message}),
1873 invocation_id: None,
1874 }
1875 };
1876 repaired.push(self.record(recorded_at.clone(), kind)?);
1877 }
1878 ensure!(
1879 self.unfinished_tool_invocations()?.is_empty(),
1880 "session tool recovery left unfinished invocations"
1881 );
1882 Ok(repaired)
1883 }
1884
1885 fn unfinished_tool_invocations(&self) -> anyhow::Result<Vec<UnfinishedToolInvocation>> {
1886 let mut identified = BTreeMap::<String, UnfinishedToolInvocation>::new();
1887 let mut legacy = Vec::<UnfinishedToolInvocation>::new();
1888 for event in &self.chatend.events {
1889 match &event.kind {
1890 EventKind::ToolInvoked {
1891 tool_instance,
1892 tool_name,
1893 invocation_id,
1894 ..
1895 } => {
1896 let pending = UnfinishedToolInvocation {
1897 invocation_id: invocation_id.clone(),
1898 tool_instance: tool_instance.clone(),
1899 tool_name: tool_name.clone(),
1900 };
1901 if let Some(invocation_id) = invocation_id {
1902 ensure!(
1903 identified.insert(invocation_id.clone(), pending).is_none(),
1904 "duplicate tool invocation identity {invocation_id}"
1905 );
1906 } else {
1907 legacy.push(pending);
1908 }
1909 }
1910 EventKind::ToolCompleted {
1911 tool_instance,
1912 tool_name,
1913 invocation_id,
1914 ..
1915 } => {
1916 if let Some(invocation_id) = invocation_id {
1917 let invoked = identified.remove(invocation_id).with_context(|| {
1918 format!(
1919 "tool {tool_name} completed without matching invocation {invocation_id}"
1920 )
1921 })?;
1922 ensure!(
1923 invoked.tool_instance.as_str() == tool_instance
1924 && invoked.tool_name.as_str() == tool_name,
1925 "tool completion {invocation_id} does not match its invocation"
1926 );
1927 } else if tool_name == "call_ktool" {
1928 legacy.pop();
1931 } else {
1932 let before = legacy.len();
1933 legacy.retain(|unfinished| unfinished.tool_name.as_str() != tool_name);
1934 ensure!(
1935 legacy.len() != before,
1936 "tool {tool_name} completed without a matching invocation"
1937 );
1938 }
1939 }
1940 _ => {}
1941 }
1942 }
1943 legacy.extend(identified.into_values());
1944 Ok(legacy)
1945 }
1946
1947 pub fn mark_completed(&mut self, session_object_id: String) {
1948 self.chatend.completed_session_object = Some(session_object_id);
1949 }
1950
1951 pub fn configure_context(&mut self, kind: SessionKind, effective_context_tokens: u64) {
1952 self.chatend.metadata.kind = kind;
1953 self.chatend.metadata.effective_context_tokens = effective_context_tokens;
1954 }
1955
1956 pub fn create_box(
1957 &mut self,
1958 recorded_at: impl Into<String>,
1959 name: impl Into<String>,
1960 owner: BoxOwner,
1961 content: BoxContent,
1962 ) -> anyhow::Result<BoxId> {
1963 let recorded_at = recorded_at.into();
1964 let id = EventId(self.chatend.next_id);
1965 let box_id = BoxId(id.0);
1966 self.commit_events(
1967 recorded_at.clone(),
1968 vec![Event {
1969 id,
1970 recorded_at,
1971 kind: EventKind::BoxCreated {
1972 box_id,
1973 name: name.into(),
1974 owner,
1975 content,
1976 },
1977 }],
1978 )?;
1979 Ok(box_id)
1980 }
1981
1982 pub fn update_box(
1983 &mut self,
1984 recorded_at: impl Into<String>,
1985 box_id: BoxId,
1986 content: BoxContent,
1987 ) -> anyhow::Result<Option<EventId>> {
1988 let state = self
1989 .chatend
1990 .boxes
1991 .get(&box_id)
1992 .with_context(|| format!("box {box_id} does not exist"))?;
1993 ensure!(state.active, "box {box_id} is retired");
1994 if state.canonical.content == content {
1995 return Ok(None);
1996 }
1997 let id = EventId(self.chatend.next_id);
1998 let recorded_at = recorded_at.into();
1999 self.commit_events(
2000 recorded_at.clone(),
2001 vec![Event {
2002 id,
2003 recorded_at,
2004 kind: EventKind::CanonicalUpdated { box_id, content },
2005 }],
2006 )?;
2007 Ok(Some(id))
2008 }
2009
2010 pub fn dehydrate_boxes(
2011 &mut self,
2012 recorded_at: impl Into<String>,
2013 box_ids: &[BoxId],
2014 ) -> anyhow::Result<Vec<EventId>> {
2015 ensure!(
2016 !box_ids.is_empty(),
2017 "at least one box must be selected for dehydration"
2018 );
2019 ensure!(
2020 box_ids.iter().copied().collect::<HashSet<_>>().len() == box_ids.len(),
2021 "box dehydration cannot contain duplicate box IDs"
2022 );
2023 let recorded_at = recorded_at.into();
2024 let events = box_ids
2025 .iter()
2026 .enumerate()
2027 .map(|(offset, box_id)| {
2028 let offset = u64::try_from(offset).context("box dehydration batch is too large")?;
2029 let id = self
2030 .chatend
2031 .next_id
2032 .checked_add(offset)
2033 .context("event ID overflow")?;
2034 Ok(Event {
2035 id: EventId(id),
2036 recorded_at: recorded_at.clone(),
2037 kind: EventKind::BoxDehydrated { box_id: *box_id },
2038 })
2039 })
2040 .collect::<anyhow::Result<Vec<_>>>()?;
2041 let ids = events.iter().map(|event| event.id).collect();
2042 self.commit_events(recorded_at, events)?;
2043 Ok(ids)
2044 }
2045
2046 pub fn summarize_box(
2047 &mut self,
2048 recorded_at: impl Into<String>,
2049 box_id: BoxId,
2050 text: impl Into<String>,
2051 ) -> anyhow::Result<EventId> {
2052 self.box_operation(
2053 recorded_at,
2054 EventKind::BoxSummarized {
2055 box_id,
2056 text: text.into(),
2057 },
2058 )
2059 }
2060
2061 pub fn rehydrate_box(
2062 &mut self,
2063 recorded_at: impl Into<String>,
2064 box_id: BoxId,
2065 ) -> anyhow::Result<EventId> {
2066 self.box_operation(recorded_at, EventKind::BoxRehydrated { box_id })
2067 }
2068
2069 pub fn retire_box(
2070 &mut self,
2071 recorded_at: impl Into<String>,
2072 box_id: BoxId,
2073 ) -> anyhow::Result<EventId> {
2074 self.box_operation(recorded_at, EventKind::BoxRetired { box_id })
2075 }
2076
2077 fn box_operation(
2078 &mut self,
2079 recorded_at: impl Into<String>,
2080 kind: EventKind,
2081 ) -> anyhow::Result<EventId> {
2082 let box_id = event_box_id(&kind).context("box operation has no box identity")?;
2083 let state = self
2084 .chatend
2085 .box_state(box_id)
2086 .with_context(|| format!("box {box_id} does not exist"))?;
2087 ensure!(state.active, "box {box_id} is retired");
2088 if let EventKind::BoxSummarized { text, .. } = &kind {
2089 ensure!(!text.trim().is_empty(), "a box summary cannot be empty");
2090 }
2091 if matches!(kind, EventKind::BoxRetired { .. })
2092 && let BoxOwner::Tool { tool_instance, .. } = &state.owner
2093 {
2094 ensure!(
2095 self.chatend
2096 .tools
2097 .get(tool_instance)
2098 .is_some_and(|tool| tool.slots.iter().any(|slot| slot.box_id == box_id)),
2099 "tool box {box_id} is missing from {tool_instance}"
2100 );
2101 }
2102 let id = EventId(self.chatend.next_id);
2103 let recorded_at = recorded_at.into();
2104 self.commit_events(
2105 recorded_at.clone(),
2106 vec![Event {
2107 id,
2108 recorded_at,
2109 kind,
2110 }],
2111 )?;
2112 Ok(id)
2113 }
2114
2115 pub fn allocate_pending_node(
2116 &mut self,
2117 recorded_at: impl Into<String>,
2118 ) -> anyhow::Result<PendingId> {
2119 let id = EventId(self.chatend.next_id);
2120 let pending_id = PendingId::from_event(id);
2121 let recorded_at = recorded_at.into();
2122 self.commit_events(
2123 recorded_at.clone(),
2124 vec![Event {
2125 id,
2126 recorded_at,
2127 kind: EventKind::PendingAllocated {
2128 pending_id: pending_id.clone(),
2129 resource: PendingKind::Node,
2130 },
2131 }],
2132 )?;
2133 Ok(pending_id)
2134 }
2135
2136 pub fn stage_object(
2137 &mut self,
2138 recorded_at: impl Into<String>,
2139 media_type: impl Into<String>,
2140 file_name: Option<String>,
2141 transport: Value,
2142 bytes: &[u8],
2143 ) -> anyhow::Result<PendingId> {
2144 ensure!(
2145 bytes.len() as u64 <= MAX_OBJECT_BYTES,
2146 "object exceeds the 32 GiB V1 limit"
2147 );
2148 let aggregate = self
2149 .objects
2150 .values()
2151 .try_fold(bytes.len() as u64, |total, object| {
2152 total.checked_add(object.payload_len)
2153 })
2154 .context("staged object aggregate length overflow")?;
2155 ensure!(
2156 aggregate <= MAX_OBJECT_BYTES,
2157 "session object payload total exceeds the 32 GiB V1 limit"
2158 );
2159 let event_id = EventId(self.chatend.next_id);
2160 let pending_id = PendingId::from_event(event_id);
2161 let metadata = ObjectMetadata {
2162 pending_id: pending_id.clone(),
2163 event_id,
2164 recorded_at: recorded_at.into(),
2165 media_type: media_type.into(),
2166 file_name,
2167 transport,
2168 };
2169 let kind = EventKind::PendingAllocated {
2170 pending_id: pending_id.clone(),
2171 resource: PendingKind::Object,
2172 };
2173 let text = encode_context_event(&metadata.recorded_at, &kind)?;
2174 let object_file_name = metadata
2175 .file_name
2176 .clone()
2177 .unwrap_or_else(|| format!("object-{}", event_id.0));
2178 let position = self.durable.add_pending_object(
2179 text,
2180 object_file_name,
2181 metadata.media_type.clone(),
2182 bytes,
2183 )?;
2184 ensure!(
2185 position.0 + 1 == event_id.0,
2186 "session-log event position diverged from Kennedy context identity"
2187 );
2188 let allocation = Event {
2189 id: event_id,
2190 recorded_at: metadata.recorded_at.clone(),
2191 kind,
2192 };
2193 self.chatend.apply_transition(&Transition {
2194 recorded_at: metadata.recorded_at.clone(),
2195 events: vec![allocation],
2196 })?;
2197 self.objects.insert(
2198 pending_id.clone(),
2199 ObjectLocation {
2200 metadata,
2201 payload_offset: position.0,
2202 payload_len: bytes.len() as u64,
2203 },
2204 );
2205 Ok(pending_id)
2206 }
2207
2208 pub fn read_object(&mut self, id: &PendingId) -> anyhow::Result<Vec<u8>> {
2209 let location = self
2210 .objects
2211 .get(id)
2212 .with_context(|| format!("staged object {id} does not exist"))?
2213 .clone();
2214 Ok(self
2215 .durable
2216 .read_pending_object(EventPosition(location.payload_offset))?
2217 .bytes)
2218 }
2219
2220 pub fn record(
2221 &mut self,
2222 recorded_at: impl Into<String>,
2223 kind: EventKind,
2224 ) -> anyhow::Result<EventId> {
2225 let id = EventId(self.chatend.next_id);
2226 let recorded_at = recorded_at.into();
2227 self.commit_events(
2228 recorded_at.clone(),
2229 vec![Event {
2230 id,
2231 recorded_at,
2232 kind,
2233 }],
2234 )?;
2235 Ok(id)
2236 }
2237
2238 pub fn commit_events(
2239 &mut self,
2240 recorded_at: impl Into<String>,
2241 events: Vec<Event>,
2242 ) -> anyhow::Result<()> {
2243 let transition = Transition {
2244 recorded_at: recorded_at.into(),
2245 events,
2246 };
2247 ensure!(
2248 !transition.events.is_empty(),
2249 "a transition cannot be empty"
2250 );
2251 ensure!(
2252 !transition.events.iter().any(|event| {
2253 matches!(
2254 event.kind,
2255 EventKind::PendingAllocated {
2256 resource: PendingKind::Object,
2257 ..
2258 }
2259 )
2260 }),
2261 "pending objects must be added through stage_object"
2262 );
2263 let mut preview = self.chatend.clone();
2264 preview.apply_transition(&transition)?;
2265 for event in &transition.events {
2266 let expected = self.durable.list().events.len() as u64 + 1;
2267 ensure!(
2268 event.id.0 == expected,
2269 "Kennedy context event {} does not match session-log position {}",
2270 event.id,
2271 expected - 1
2272 );
2273 self.durable.add_event(
2274 context_event_role(&event.kind),
2275 encode_context_event(&event.recorded_at, &event.kind)?,
2276 )?;
2277 }
2278 self.chatend = preview;
2279 Ok(())
2280 }
2281
2282 pub fn apply_tool_slots(
2283 &mut self,
2284 recorded_at: impl Into<String>,
2285 tool_instance: impl Into<String>,
2286 slots: Vec<ToolSlotInput>,
2287 ) -> anyhow::Result<Vec<EventId>> {
2288 self.apply_tool_slots_inner(recorded_at, tool_instance, slots, None)
2289 }
2290
2291 pub fn apply_tool_slots_with_layout(
2292 &mut self,
2293 recorded_at: impl Into<String>,
2294 tool_instance: impl Into<String>,
2295 slots: Vec<ToolSlotInput>,
2296 layout_slots: &[String],
2297 ) -> anyhow::Result<Vec<EventId>> {
2298 self.apply_tool_slots_inner(recorded_at, tool_instance, slots, Some(layout_slots))
2299 }
2300
2301 fn apply_tool_slots_inner(
2302 &mut self,
2303 recorded_at: impl Into<String>,
2304 tool_instance: impl Into<String>,
2305 slots: Vec<ToolSlotInput>,
2306 layout_slots: Option<&[String]>,
2307 ) -> anyhow::Result<Vec<EventId>> {
2308 let recorded_at = recorded_at.into();
2309 let tool_instance = tool_instance.into();
2310 let current = self
2311 .chatend
2312 .tools
2313 .get(&tool_instance)
2314 .cloned()
2315 .unwrap_or_default();
2316 ensure!(
2317 slots.len() >= current.slots.len(),
2318 "stateful tool slot sequence was truncated"
2319 );
2320 for (index, existing) in current.slots.iter().enumerate() {
2321 ensure!(
2322 slots[index].slot == existing.slot,
2323 "stateful tool slot sequence was reordered at index {index}"
2324 );
2325 ensure!(
2326 !existing.retired || slots[index].retired,
2327 "retired tool slot {} cannot be reactivated",
2328 existing.slot
2329 );
2330 }
2331 let mut events = Vec::new();
2332 let mut next = self.chatend.next_id;
2333 let mut next_state = current.clone();
2334 for (index, input) in slots.iter().enumerate() {
2335 if let Some(existing) = current.slots.get(index) {
2336 let state = self
2337 .chatend
2338 .boxes
2339 .get(&existing.box_id)
2340 .context("tool slot references a missing box")?;
2341 if input.retired && !existing.retired {
2342 let id = EventId(next);
2343 next += 1;
2344 events.push(Event {
2345 id,
2346 recorded_at: recorded_at.clone(),
2347 kind: EventKind::BoxRetired {
2348 box_id: existing.box_id,
2349 },
2350 });
2351 next_state.slots[index].retired = true;
2352 } else if !input.retired {
2353 if state.name != input.name {
2354 let id = EventId(next);
2355 next += 1;
2356 events.push(Event {
2357 id,
2358 recorded_at: recorded_at.clone(),
2359 kind: EventKind::BoxRenamed {
2360 box_id: existing.box_id,
2361 name: input.name.clone(),
2362 },
2363 });
2364 }
2365 if state.canonical.content != input.content {
2366 let id = EventId(next);
2367 next += 1;
2368 events.push(Event {
2369 id,
2370 recorded_at: recorded_at.clone(),
2371 kind: EventKind::CanonicalUpdated {
2372 box_id: existing.box_id,
2373 content: input.content.clone(),
2374 },
2375 });
2376 }
2377 }
2378 } else {
2379 ensure!(
2380 !input.retired,
2381 "a newly appended tool slot cannot start retired"
2382 );
2383 let id = EventId(next);
2384 next += 1;
2385 let box_id = BoxId(id.0);
2386 events.push(Event {
2387 id,
2388 recorded_at: recorded_at.clone(),
2389 kind: EventKind::BoxCreated {
2390 box_id,
2391 name: input.name.clone(),
2392 owner: BoxOwner::Tool {
2393 tool_instance: tool_instance.clone(),
2394 slot: input.slot.clone(),
2395 },
2396 content: input.content.clone(),
2397 },
2398 });
2399 next_state.slots.push(ToolSlot {
2400 slot: input.slot.clone(),
2401 box_id,
2402 retired: false,
2403 });
2404 }
2405 }
2406 if let Some(layout_slots) = layout_slots {
2407 let mut unique = std::collections::HashSet::new();
2408 let box_ids = layout_slots
2409 .iter()
2410 .map(|slot_name| {
2411 ensure!(
2412 unique.insert(slot_name),
2413 "tool layout contains duplicate slot {slot_name}"
2414 );
2415 let slot = next_state
2416 .slots
2417 .iter()
2418 .find(|slot| &slot.slot == slot_name)
2419 .with_context(|| {
2420 format!("tool layout references missing slot {slot_name}")
2421 })?;
2422 ensure!(
2423 !slot.retired,
2424 "tool layout references retired slot {slot_name}"
2425 );
2426 Ok(slot.box_id)
2427 })
2428 .collect::<anyhow::Result<Vec<_>>>()?;
2429 if self.chatend.tool_layouts.get(&tool_instance) != Some(&box_ids) {
2430 let id = EventId(next);
2431 events.push(Event {
2432 id,
2433 recorded_at: recorded_at.clone(),
2434 kind: EventKind::ToolLayoutChanged {
2435 tool_instance: tool_instance.clone(),
2436 box_ids,
2437 },
2438 });
2439 }
2440 }
2441 if events.is_empty() {
2442 return Ok(Vec::new());
2443 }
2444 let ids = events.iter().map(|event| event.id).collect::<Vec<_>>();
2445 self.commit_events(recorded_at, events)?;
2446 Ok(ids)
2447 }
2448
2449 pub fn apply_box_representations(
2450 &mut self,
2451 recorded_at: impl Into<String>,
2452 desired: &BTreeMap<BoxId, BoxRepresentation>,
2453 ) -> anyhow::Result<Vec<EventId>> {
2454 let recorded_at = recorded_at.into();
2455 let events = self
2456 .chatend
2457 .box_representation_events(&recorded_at, desired)?;
2458 if events.is_empty() {
2459 return Ok(Vec::new());
2460 }
2461 let ids = events.iter().map(|event| event.id).collect::<Vec<_>>();
2462 self.commit_events(recorded_at, events)?;
2463 Ok(ids)
2464 }
2465}
2466
2467pub struct SessionHistoryIntegration;
2474
2475impl SessionHistoryIntegration {
2476 pub fn create_session(
2478 path: impl AsRef<Path>,
2479 metadata: SessionMetadata,
2480 ) -> anyhow::Result<Session> {
2481 Session::create(path, metadata)
2482 }
2483
2484 pub fn open_session(
2490 path: impl AsRef<Path>,
2491 metadata: SessionMetadata,
2492 default_provider_model: Option<&str>,
2493 estimator: Option<ProviderCostEstimator>,
2494 ) -> anyhow::Result<Session> {
2495 match (default_provider_model, estimator) {
2496 (None, None) => Session::open_with_metadata(path, metadata),
2497 (default_provider_model, estimator) => Session::open_with_metadata_and_provider_costs(
2498 path,
2499 metadata,
2500 default_provider_model,
2501 estimator,
2502 ),
2503 }
2504 }
2505
2506 pub fn replay(
2511 metadata: SessionMetadata,
2512 log: &kcode_session_log::SessionLog,
2513 default_provider_model: Option<&str>,
2514 estimator: Option<ProviderCostEstimator>,
2515 ) -> anyhow::Result<Chatend> {
2516 match (default_provider_model, estimator) {
2517 (None, None) => Chatend::replay(metadata, log),
2518 (default_provider_model, estimator) => Chatend::replay_with_provider_costs(
2519 metadata,
2520 log,
2521 default_provider_model,
2522 estimator,
2523 ),
2524 }
2525 }
2526
2527 pub fn legacy_provider_cost_summary_for_archive(
2531 archive: &Value,
2532 default_provider_model: Option<&str>,
2533 estimator: ProviderCostEstimator,
2534 ) -> anyhow::Result<ProviderCostSummary> {
2535 legacy_provider_cost_summary_for_archive(archive, default_provider_model, estimator)
2536 }
2537}
2538
2539#[cfg(test)]
2540type SessionJournal = Session;
2541
2542#[cfg(test)]
2543mod tests {
2544 use std::path::PathBuf;
2545 use std::time::{SystemTime, UNIX_EPOCH};
2546
2547 use serde_json::json;
2548
2549 use super::*;
2550
2551 fn path(label: &str) -> PathBuf {
2552 std::env::temp_dir()
2553 .join(format!(
2554 "kennedy-chatend-{label}-{}-{}",
2555 std::process::id(),
2556 SystemTime::now()
2557 .duration_since(UNIX_EPOCH)
2558 .unwrap()
2559 .as_nanos()
2560 ))
2561 .join("session-1.session-log")
2562 }
2563
2564 fn metadata() -> SessionMetadata {
2565 SessionMetadata {
2566 session_id: "session-1".into(),
2567 kind: SessionKind::Conversation,
2568 created_at: "2026-07-23T00:00:00Z".into(),
2569 effective_context_tokens: 1_000,
2570 channel: json!({"kind":"test"}),
2571 }
2572 }
2573
2574 #[test]
2575 fn box_identity_continuations_staleness_and_replay_are_exact() {
2576 let path = path("boxes");
2577 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2578 let id = journal
2579 .create_box(
2580 "t1",
2581 "message",
2582 BoxOwner::User,
2583 BoxContent::text("original"),
2584 )
2585 .unwrap();
2586 assert_eq!(id, BoxId(1));
2587 journal.summarize_box("t2", id, "summary").unwrap();
2588 journal
2589 .update_box("t3", id, BoxContent::text("changed"))
2590 .unwrap();
2591 let state = journal.state().box_state(id).unwrap().clone();
2592 assert!(state.stale());
2593 assert_eq!(
2594 state.representation,
2595 Representation::Summarized {
2596 based_on: EventId(1),
2597 text: "summary".into()
2598 }
2599 );
2600 let projection = journal.state().projection();
2601 assert_eq!(projection.items[0].text, "[box updated]");
2602 assert_eq!(projection.items[1].text, "[box updated]");
2603 assert!(projection.items[2].text.contains("summary"));
2604 assert!(projection.items[2].stale);
2605 assert!(projection.footer.starts_with("[stale boxes: 1]\n"));
2606 assert!(projection.footer.contains("\n[current time: "));
2607 assert_eq!(
2608 projection.footer.lines().last().unwrap(),
2609 format!(
2610 "[current context size: {} | max context size: {}]",
2611 projection.estimated_tokens, projection.status.context_limit_tokens
2612 )
2613 );
2614 assert!(!projection.footer.contains("effective"));
2615 assert!(!projection.footer.contains("turn_limit"));
2616 assert!(projection.render().ends_with(&projection.footer));
2617 drop(journal);
2618 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2619 assert_eq!(reopened.state().box_state(id), Some(&state));
2620 std::fs::remove_file(path).unwrap();
2621 }
2622
2623 #[test]
2624 fn box_headers_hide_internal_ownership_and_timestamp_messages() {
2625 let path = path("box-headers");
2626 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2627 journal
2628 .create_box(
2629 "2026-07-23T00:00:01Z",
2630 "User message",
2631 BoxOwner::User,
2632 BoxContent::text("Hello"),
2633 )
2634 .unwrap();
2635 journal
2636 .create_box(
2637 "2026-07-23T00:00:02Z",
2638 "Kennedy message",
2639 BoxOwner::Kennedy,
2640 BoxContent::text("Hi"),
2641 )
2642 .unwrap();
2643 let mut content = BoxContent::text("Node ID: AAAAAAAB");
2644 content.use_concise_header();
2645 journal
2646 .create_box(
2647 "2026-07-23T00:00:03Z",
2648 "Kweb loaded node",
2649 BoxOwner::Tool {
2650 tool_instance: "kweb".into(),
2651 slot: "loaded".into(),
2652 },
2653 content,
2654 )
2655 .unwrap();
2656
2657 let projection = journal.state().projection();
2658 assert_eq!(
2659 projection.items[0].text,
2660 "[box 1 | User message | timestamp=2026-07-23T00:00:01Z | hydrated]\nHello"
2661 );
2662 assert_eq!(
2663 projection.items[1].text,
2664 "[box 2 | Kennedy message | timestamp=2026-07-23T00:00:02Z | hydrated]\nHi"
2665 );
2666 assert_eq!(
2667 projection.items[2].text,
2668 "[box 3 | Kweb loaded node | hydrated]\nNode ID: AAAAAAAB"
2669 );
2670 assert!(
2671 projection
2672 .items
2673 .iter()
2674 .all(|item| !item.text.contains("owner="))
2675 );
2676 let current_time = projection
2677 .footer
2678 .lines()
2679 .find_map(|line| line.strip_prefix("[current time: "))
2680 .and_then(|line| line.strip_suffix(']'))
2681 .unwrap();
2682 chrono::DateTime::parse_from_rfc3339(current_time).unwrap();
2683 assert!(current_time.starts_with("20"));
2684 std::fs::remove_file(path).unwrap();
2685 }
2686
2687 #[test]
2688 fn shared_pending_and_box_identity_space_never_overlaps() {
2689 let path = path("pending");
2690 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2691 let first = journal.allocate_pending_node("t1").unwrap();
2692 let box_id = journal
2693 .create_box("t2", "box", BoxOwner::Kennedy, BoxContent::text("hello"))
2694 .unwrap();
2695 let object = journal
2696 .stage_object(
2697 "t3",
2698 "application/octet-stream",
2699 None,
2700 Value::Null,
2701 b"\0binary\xff",
2702 )
2703 .unwrap();
2704 assert_eq!(first.to_string(), "pending:1");
2705 assert_eq!(box_id, BoxId(2));
2706 assert_eq!(object.to_string(), "pending:3");
2707 assert_eq!(journal.read_object(&object).unwrap(), b"\0binary\xff");
2708 let stored = journal.session_log();
2709 assert!(!stored.events[0].text.contains("pending_id"));
2710 assert!(!stored.events[1].text.contains("box_id"));
2711 assert!(!stored.events[2].text.contains("pending_id"));
2712 drop(journal);
2713 let mut reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2714 assert_eq!(reopened.read_object(&object).unwrap(), b"\0binary\xff");
2715 assert_eq!(reopened.state().next_id, 4);
2716 std::fs::remove_file(path).unwrap();
2717 }
2718
2719 #[test]
2720 fn kennedy_tool_completion_is_validated_before_storage_seals() {
2721 let path = path("seal-tool");
2722 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2723 journal
2724 .record(
2725 "t1",
2726 EventKind::ToolInvoked {
2727 tool_instance: "CreateNode:1".into(),
2728 tool_name: "CreateNode".into(),
2729 arguments: json!({}),
2730 invocation_id: None,
2731 },
2732 )
2733 .unwrap();
2734 assert!(journal.seal().is_err());
2735 journal
2736 .record(
2737 "t2",
2738 EventKind::ToolCompleted {
2739 tool_instance: "call_ktool".into(),
2740 tool_name: "call_ktool".into(),
2741 outcome: json!({"ok":true}),
2742 invocation_id: None,
2743 },
2744 )
2745 .unwrap();
2746 journal.seal().unwrap();
2747 assert!(journal.is_sealed());
2748 std::fs::remove_file(path).unwrap();
2749 }
2750
2751 #[test]
2752 fn interrupted_tools_are_recovered_by_identity_without_losing_legacy_journals() {
2753 let path = path("repair-tools");
2754 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2755 journal
2756 .record(
2757 "t1",
2758 EventKind::ToolInvoked {
2759 tool_instance: "WebSearch:legacy".into(),
2760 tool_name: "WebSearch".into(),
2761 arguments: json!({"question":"legacy"}),
2762 invocation_id: None,
2763 },
2764 )
2765 .unwrap();
2766 journal
2767 .record(
2768 "t2",
2769 EventKind::ToolInvoked {
2770 tool_instance: "WebSearch:new".into(),
2771 tool_name: "WebSearch".into(),
2772 arguments: json!({"question":"identified"}),
2773 invocation_id: Some("call-1".into()),
2774 },
2775 )
2776 .unwrap();
2777 assert!(journal.seal().is_err());
2778
2779 let repaired = journal.repair_unfinished_tools("recovery").unwrap();
2780 assert_eq!(repaired.len(), 2);
2781 assert!(
2782 journal
2783 .state()
2784 .events
2785 .iter()
2786 .rev()
2787 .take(2)
2788 .all(|event| matches!(
2789 &event.kind,
2790 EventKind::ToolCompleted { outcome, .. }
2791 if outcome.get("recovered").and_then(Value::as_bool) == Some(true)
2792 ))
2793 );
2794 journal.seal().unwrap();
2795 assert!(journal.is_sealed());
2796 std::fs::remove_file(path).unwrap();
2797 }
2798
2799 #[test]
2800 fn identified_tool_completions_can_arrive_out_of_order() {
2801 let path = path("tool-identity");
2802 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2803 for id in ["call-1", "call-2"] {
2804 journal
2805 .record(
2806 id,
2807 EventKind::ToolInvoked {
2808 tool_instance: format!("WebSearch:{id}"),
2809 tool_name: "WebSearch".into(),
2810 arguments: json!({"question":id}),
2811 invocation_id: Some(id.into()),
2812 },
2813 )
2814 .unwrap();
2815 }
2816 for id in ["call-1", "call-2"] {
2817 journal
2818 .record(
2819 format!("{id}-complete"),
2820 EventKind::ToolCompleted {
2821 tool_instance: format!("WebSearch:{id}"),
2822 tool_name: "WebSearch".into(),
2823 outcome: json!({"ok":true}),
2824 invocation_id: Some(id.into()),
2825 },
2826 )
2827 .unwrap();
2828 }
2829 journal.seal().unwrap();
2830 std::fs::remove_file(path).unwrap();
2831 }
2832
2833 #[test]
2834 fn invalid_box_operations_do_not_poison_the_append_only_journal() {
2835 let path = path("invalid-box-operation");
2836 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2837 let box_id = journal
2838 .create_box(
2839 "t1",
2840 "valid",
2841 BoxOwner::Controller,
2842 BoxContent::text("canonical"),
2843 )
2844 .unwrap();
2845 let valid_length = std::fs::metadata(&path).unwrap().len();
2846
2847 let result = journal.dehydrate_boxes("t2", &[box_id, BoxId(97)]);
2848 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2849 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2850 assert!(matches!(
2851 journal.state().box_state(box_id).unwrap().representation,
2852 Representation::Hydrated { .. }
2853 ));
2854
2855 for result in [
2856 journal.summarize_box("t3", BoxId(97), "summary"),
2857 journal.rehydrate_box("t4", BoxId(97)),
2858 journal.retire_box("t5", BoxId(97)),
2859 ] {
2860 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2861 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2862 }
2863
2864 drop(journal);
2865 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2866 assert_eq!(
2867 reopened
2868 .state()
2869 .box_state(box_id)
2870 .unwrap()
2871 .canonical
2872 .content,
2873 BoxContent::text("canonical")
2874 );
2875 std::fs::remove_file(path).unwrap();
2876 }
2877
2878 #[test]
2879 fn multiple_boxes_dehydrate_in_one_replayable_batch() {
2880 let path = path("dehydrate-boxes");
2881 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2882 let boxes = ["first", "second", "third"]
2883 .into_iter()
2884 .map(|text| {
2885 journal
2886 .create_box("t1", text, BoxOwner::Controller, BoxContent::text(text))
2887 .unwrap()
2888 })
2889 .collect::<Vec<_>>();
2890
2891 let event_ids = journal
2892 .dehydrate_boxes("t2", &[boxes[0], boxes[2]])
2893 .unwrap();
2894 assert_eq!(event_ids, [EventId(4), EventId(5)]);
2895 assert!(matches!(
2896 journal.state().box_state(boxes[0]).unwrap().representation,
2897 Representation::Dehydrated { .. }
2898 ));
2899 assert!(matches!(
2900 journal.state().box_state(boxes[1]).unwrap().representation,
2901 Representation::Hydrated { .. }
2902 ));
2903 assert!(matches!(
2904 journal.state().box_state(boxes[2]).unwrap().representation,
2905 Representation::Dehydrated { .. }
2906 ));
2907
2908 let valid_length = std::fs::metadata(&path).unwrap().len();
2909 assert_eq!(
2910 journal.dehydrate_boxes("t3", &[]).unwrap_err().to_string(),
2911 "at least one box must be selected for dehydration"
2912 );
2913 assert_eq!(
2914 journal
2915 .dehydrate_boxes("t3", &[boxes[1], boxes[1]])
2916 .unwrap_err()
2917 .to_string(),
2918 "box dehydration cannot contain duplicate box IDs"
2919 );
2920 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2921
2922 drop(journal);
2923 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2924 assert!(matches!(
2925 reopened.state().box_state(boxes[0]).unwrap().representation,
2926 Representation::Dehydrated { .. }
2927 ));
2928 assert!(matches!(
2929 reopened.state().box_state(boxes[2]).unwrap().representation,
2930 Representation::Dehydrated { .. }
2931 ));
2932 std::fs::remove_file(path).unwrap();
2933 }
2934
2935 #[test]
2936 fn status_estimates_the_context_with_every_active_box_hydrated() {
2937 let path = path("fully-hydrated-estimate");
2938 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2939 let box_id = journal
2940 .create_box(
2941 "t1",
2942 "large",
2943 BoxOwner::Controller,
2944 BoxContent::text("canonical material ".repeat(1_000)),
2945 )
2946 .unwrap();
2947 journal.dehydrate_boxes("t2", &[box_id]).unwrap();
2948
2949 let status = journal.state().projection().status;
2950 assert!(
2951 status.fully_hydrated_context_tokens > status.current_context_tokens,
2952 "the hydrated estimate should include the canonical body"
2953 );
2954 std::fs::remove_file(path).unwrap();
2955 }
2956
2957 #[test]
2958 fn provider_measurements_recalibrate_the_matching_manifest_then_track_deltas() {
2959 let path = path("calibration");
2960 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2961 journal
2962 .create_box(
2963 "t1",
2964 "system",
2965 BoxOwner::System,
2966 BoxContent::text("baseline provider content"),
2967 )
2968 .unwrap();
2969 let submitted = journal.state().projection();
2970 journal
2971 .record(
2972 "t2",
2973 EventKind::InferenceSubmitted {
2974 manifest_hash: "manifest-1".into(),
2975 estimated_input_tokens: submitted.estimated_tokens,
2976 raw_estimated_input_tokens: Some(submitted.raw_estimated_tokens),
2977 },
2978 )
2979 .unwrap();
2980 let tool_box = journal
2981 .create_box(
2982 "t3",
2983 "tool result",
2984 BoxOwner::Controller,
2985 BoxContent::text("x".repeat(3_000)),
2986 )
2987 .unwrap();
2988 let measured_context = journal.state().projection();
2989 journal
2990 .record(
2991 "t4",
2992 EventKind::ProviderReceipt {
2993 manifest_hash: "manifest-1".into(),
2994 input_tokens: Some(777),
2995 output_tokens: Some(3),
2996 context_bytes: Some(measured_context.context_bytes),
2997 raw_context_tokens: Some(measured_context.raw_estimated_tokens),
2998 provider_data: Value::Null,
2999 },
3000 )
3001 .unwrap();
3002 assert_eq!(journal.state().projection().estimated_tokens, 777);
3003 journal
3004 .create_box(
3005 "t5",
3006 "new material",
3007 BoxOwner::User,
3008 BoxContent::text("x".repeat(300)),
3009 )
3010 .unwrap();
3011 let expanded = journal.state().projection();
3012 assert_eq!(
3013 expanded.estimated_tokens,
3014 777 + expanded
3015 .context_bytes
3016 .abs_diff(measured_context.context_bytes)
3017 / ESTIMATED_BYTES_PER_TOKEN
3018 );
3019 assert!(expanded.raw_estimated_tokens > submitted.raw_estimated_tokens);
3020 journal.dehydrate_boxes("t6", &[tool_box]).unwrap();
3021 let shrunken = journal.state().projection();
3022 assert_eq!(
3023 shrunken.estimated_tokens,
3024 777_u64.saturating_sub(
3025 measured_context
3026 .context_bytes
3027 .abs_diff(shrunken.context_bytes)
3028 / ESTIMATED_BYTES_PER_TOKEN
3029 )
3030 );
3031 std::fs::remove_file(path).unwrap();
3032 }
3033
3034 #[test]
3035 fn session_status_keeps_provider_usage_categories_exact_and_exclusive() {
3036 let path = path("session-status");
3037 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3038 journal
3039 .create_box("t1", "system", BoxOwner::System, BoxContent::text("é"))
3040 .unwrap();
3041 let before_usage = journal.state().projection();
3042 assert_eq!(
3043 before_usage.raw_estimated_tokens,
3044 before_usage
3045 .context_bytes
3046 .div_ceil(ESTIMATED_BYTES_PER_TOKEN)
3047 );
3048 journal
3049 .record(
3050 "t2",
3051 EventKind::ProviderReceipt {
3052 manifest_hash: "manifest-1".into(),
3053 input_tokens: Some(100),
3054 output_tokens: Some(20),
3055 context_bytes: Some(before_usage.context_bytes),
3056 raw_context_tokens: Some(before_usage.raw_estimated_tokens),
3057 provider_data: json!({
3058 "usageIsDelta":true,
3059 "cachedInputTokens":40,
3060 "nonCachedInputTokens":60,
3061 "thinkingTokens":8,
3062 "outputTokens":12,
3063 "estimatedCostUsdNanos":12345
3064 }),
3065 },
3066 )
3067 .unwrap();
3068 let second_anchor = journal.state().projection();
3069 journal
3070 .record(
3071 "t3",
3072 EventKind::ProviderReceipt {
3073 manifest_hash: "manifest-1".into(),
3074 input_tokens: Some(120),
3075 output_tokens: Some(5),
3076 context_bytes: Some(second_anchor.context_bytes),
3077 raw_context_tokens: Some(second_anchor.raw_estimated_tokens),
3078 provider_data: json!({
3079 "usageIsDelta":true,
3080 "cachedInputTokens":10,
3081 "nonCachedInputTokens":10,
3082 "thinkingTokens":2,
3083 "outputTokens":3,
3084 "estimatedCostUsdNanos":6789
3085 }),
3086 },
3087 )
3088 .unwrap();
3089 let projection = journal.state().projection();
3090 assert_eq!(
3091 projection.status,
3092 SessionStatus {
3093 current_context_tokens: 120,
3094 fully_hydrated_context_tokens: 120,
3095 context_limit_tokens: 700,
3096 current_context_bytes: projection.context_bytes,
3097 cached_input_tokens: 50,
3098 non_cached_input_tokens: 70,
3099 thinking_tokens: 10,
3100 output_tokens: 15,
3101 estimated_cost_usd_nanos: 19_134,
3102 unpriced_provider_calls: 0,
3103 }
3104 );
3105 let archive: Value = serde_json::from_slice(&journal.archive_bytes().unwrap()).unwrap();
3106 assert_eq!(archive["context"]["status"]["cachedInputTokens"], 50);
3107 assert!(
3108 archive["chatendText"]
3109 .as_str()
3110 .unwrap()
3111 .ends_with(archive["context"]["footer"].as_str().unwrap())
3112 );
3113 std::fs::remove_file(path).unwrap();
3114 }
3115
3116 fn compatibility_cost(
3117 model: &str,
3118 metering: &ProviderMetering,
3119 ) -> Option<ProviderCostEstimate> {
3120 let ProviderMetering::Tokens(usage) = metering else {
3121 return None;
3122 };
3123 let base = match model {
3124 "gpt-5.6-sol" => 1_000,
3125 "gemini-3.1-pro-preview" => 2_000,
3126 _ => return None,
3127 };
3128 Some(ProviderCostEstimate {
3129 usd_nanos: base + usage.input_tokens,
3130 accuracy: json!("exact"),
3131 pricing_version: "test-prices".into(),
3132 })
3133 }
3134
3135 #[test]
3136 fn legacy_provider_costs_are_reconstructed_without_rewriting_history() {
3137 let path = path("legacy-provider-costs");
3138 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3139 journal
3140 .record(
3141 "t1",
3142 EventKind::ProviderReceipt {
3143 manifest_hash: "top".into(),
3144 input_tokens: Some(10),
3145 output_tokens: Some(1),
3146 context_bytes: None,
3147 raw_context_tokens: None,
3148 provider_data: json!({
3149 "usageIsDelta":true,
3150 "nonCachedInputTokens":10,
3151 "cachedInputTokens":0,
3152 "thinkingTokens":0,
3153 "outputTokens":1
3154 }),
3155 },
3156 )
3157 .unwrap();
3158 journal
3159 .record(
3160 "t2",
3161 EventKind::ToolInvoked {
3162 tool_instance: "RunSubagent:1".into(),
3163 tool_name: "RunSubagent".into(),
3164 arguments: json!({"model":"codex/gpt-5.6-sol"}),
3165 invocation_id: Some("subagent-1".into()),
3166 },
3167 )
3168 .unwrap();
3169 journal
3170 .record(
3171 "t3",
3172 EventKind::Note {
3173 label: "subagent_started".into(),
3174 value: json!({"providerModel":"gpt-5.6-sol"}),
3175 },
3176 )
3177 .unwrap();
3178 journal
3179 .record(
3180 "t4",
3181 EventKind::ProviderReceipt {
3182 manifest_hash: "subagent".into(),
3183 input_tokens: None,
3184 output_tokens: None,
3185 context_bytes: None,
3186 raw_context_tokens: None,
3187 provider_data: json!({
3188 "source":"subagent",
3189 "usageIsDelta":true,
3190 "nonCachedInputTokens":20,
3191 "cachedInputTokens":0,
3192 "thinkingTokens":0,
3193 "outputTokens":1
3194 }),
3195 },
3196 )
3197 .unwrap();
3198 journal
3199 .record(
3200 "t5",
3201 EventKind::ToolInvoked {
3202 tool_instance: "AnnotateMedia:1".into(),
3203 tool_name: "AnnotateMedia".into(),
3204 arguments: json!({"model":"gemini-3.1-pro-preview"}),
3205 invocation_id: Some("media-1".into()),
3206 },
3207 )
3208 .unwrap();
3209 journal
3210 .record(
3211 "t6",
3212 EventKind::ProviderReceipt {
3213 manifest_hash: "media".into(),
3214 input_tokens: None,
3215 output_tokens: None,
3216 context_bytes: None,
3217 raw_context_tokens: None,
3218 provider_data: json!({
3219 "source":"media_annotation",
3220 "usageIsDelta":true,
3221 "nonCachedInputTokens":30,
3222 "cachedInputTokens":0,
3223 "thinkingTokens":0,
3224 "outputTokens":1
3225 }),
3226 },
3227 )
3228 .unwrap();
3229 journal
3230 .record(
3231 "t7",
3232 EventKind::ToolCompleted {
3233 tool_instance: "AnnotateMedia:1".into(),
3234 tool_name: "AnnotateMedia".into(),
3235 outcome: json!({"ok":true}),
3236 invocation_id: Some("media-1".into()),
3237 },
3238 )
3239 .unwrap();
3240 journal
3241 .record(
3242 "t8",
3243 EventKind::ToolCompleted {
3244 tool_instance: "RunSubagent:1".into(),
3245 tool_name: "RunSubagent".into(),
3246 outcome: json!({"ok":true}),
3247 invocation_id: Some("subagent-1".into()),
3248 },
3249 )
3250 .unwrap();
3251 journal
3252 .record(
3253 "t9",
3254 EventKind::ProviderReceipt {
3255 manifest_hash: "unknown".into(),
3256 input_tokens: None,
3257 output_tokens: None,
3258 context_bytes: None,
3259 raw_context_tokens: None,
3260 provider_data: json!({
3261 "source":"web_search",
3262 "usageIsDelta":true,
3263 "nonCachedInputTokens":40,
3264 "cachedInputTokens":0,
3265 "thinkingTokens":0,
3266 "outputTokens":1
3267 }),
3268 },
3269 )
3270 .unwrap();
3271 drop(journal);
3272
3273 let compatible = SessionJournal::open_with_metadata_and_provider_costs(
3274 &path,
3275 metadata(),
3276 Some("gpt-5.6-sol"),
3277 Some(compatibility_cost),
3278 )
3279 .unwrap();
3280 let status = &compatible.state().projection().status;
3281 assert_eq!(status.estimated_cost_usd_nanos, 4_060);
3282 assert_eq!(status.unpriced_provider_calls, 1);
3283 let inferred_models = compatible
3284 .state()
3285 .events
3286 .iter()
3287 .filter_map(|event| match &event.kind {
3288 EventKind::ProviderReceipt { provider_data, .. } => provider_data
3289 .get("providerModel")
3290 .and_then(Value::as_str)
3291 .map(str::to_owned),
3292 _ => None,
3293 })
3294 .collect::<Vec<_>>();
3295 assert_eq!(
3296 inferred_models,
3297 vec!["gpt-5.6-sol", "gpt-5.6-sol", "gemini-3.1-pro-preview"]
3298 );
3299 drop(compatible);
3300
3301 let unchanged = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3302 assert_eq!(
3303 unchanged
3304 .state()
3305 .projection()
3306 .status
3307 .estimated_cost_usd_nanos,
3308 0
3309 );
3310 assert_eq!(
3311 unchanged
3312 .state()
3313 .projection()
3314 .status
3315 .unpriced_provider_calls,
3316 4
3317 );
3318 let archive: Value = serde_json::from_slice(&unchanged.archive_bytes().unwrap()).unwrap();
3319 let archive_summary = legacy_provider_cost_summary_for_archive(
3320 &archive,
3321 Some("gpt-5.6-sol"),
3322 compatibility_cost,
3323 )
3324 .unwrap();
3325 assert_eq!(archive_summary.estimated_cost_usd_nanos, 4_060);
3326 assert_eq!(archive_summary.unpriced_provider_calls, 1);
3327 assert_eq!(archive["context"]["status"]["estimatedCostUsdNanos"], 0);
3328 assert_eq!(archive["context"]["status"]["unpricedProviderCalls"], 4);
3329 drop(unchanged);
3330 std::fs::remove_file(path).unwrap();
3331 }
3332
3333 #[test]
3334 fn legacy_provider_receipts_without_a_raw_context_anchor_remain_readable() {
3335 let receipt: EventKind = serde_json::from_value(json!({
3336 "type":"provider_receipt",
3337 "manifest_hash":"legacy-manifest",
3338 "input_tokens":123,
3339 "output_tokens":4,
3340 "provider_data":null
3341 }))
3342 .unwrap();
3343 assert!(matches!(
3344 receipt,
3345 EventKind::ProviderReceipt {
3346 context_bytes: None,
3347 raw_context_tokens: None,
3348 ..
3349 }
3350 ));
3351 }
3352
3353 #[test]
3354 fn stateful_tool_slots_are_batched_append_only_and_do_not_see_summaries() {
3355 let path = path("slots");
3356 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3357 journal
3358 .apply_tool_slots(
3359 "t1",
3360 "rust-1",
3361 vec![
3362 ToolSlotInput {
3363 slot: "a.rs".into(),
3364 name: "a.rs".into(),
3365 content: BoxContent::text("a"),
3366 retired: false,
3367 },
3368 ToolSlotInput {
3369 slot: "b.rs".into(),
3370 name: "b.rs".into(),
3371 content: BoxContent::text("b"),
3372 retired: false,
3373 },
3374 ],
3375 )
3376 .unwrap();
3377 let a = journal.state().tools["rust-1"].slots[0].box_id;
3378 journal.summarize_box("t2", a, "Kennedy summary").unwrap();
3379 let before = journal.state().clone();
3380 assert!(
3381 journal
3382 .apply_tool_slots(
3383 "t3",
3384 "rust-1",
3385 vec![ToolSlotInput {
3386 slot: "b.rs".into(),
3387 name: "b.rs".into(),
3388 content: BoxContent::text("b"),
3389 retired: false,
3390 }]
3391 )
3392 .is_err()
3393 );
3394 assert_eq!(journal.state(), &before);
3395 journal
3396 .apply_tool_slots(
3397 "t4",
3398 "rust-1",
3399 vec![
3400 ToolSlotInput {
3401 slot: "a.rs".into(),
3402 name: "a.rs".into(),
3403 content: BoxContent::text("a2"),
3404 retired: false,
3405 },
3406 ToolSlotInput {
3407 slot: "b.rs".into(),
3408 name: "b.rs".into(),
3409 content: BoxContent::text("b"),
3410 retired: true,
3411 },
3412 ToolSlotInput {
3413 slot: "c.rs".into(),
3414 name: "c.rs".into(),
3415 content: BoxContent::text("c"),
3416 retired: false,
3417 },
3418 ],
3419 )
3420 .unwrap();
3421 let a_state = journal.state().box_state(a).unwrap();
3422 assert_eq!(a_state.canonical.content.text, "a2");
3423 assert!(matches!(
3424 a_state.representation,
3425 Representation::Summarized { .. }
3426 ));
3427 drop(journal);
3428 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3429 assert_eq!(reopened.state().tools["rust-1"].slots.len(), 3);
3430 assert!(reopened.state().tools["rust-1"].slots[1].retired);
3431 std::fs::remove_file(path).unwrap();
3432 }
3433
3434 #[test]
3435 fn tool_layout_orders_current_boxes_without_changing_their_identities() {
3436 let path = path("tool-layout");
3437 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3438 journal
3439 .apply_tool_slots_with_layout(
3440 "t1",
3441 "kweb",
3442 vec![
3443 ToolSlotInput {
3444 slot: "active".into(),
3445 name: "Active node".into(),
3446 content: BoxContent::text("ACTIVE NODE"),
3447 retired: false,
3448 },
3449 ToolSlotInput {
3450 slot: "direct".into(),
3451 name: "Direct node".into(),
3452 content: BoxContent::text("DIRECT NODE"),
3453 retired: false,
3454 },
3455 ],
3456 &["direct".into(), "active".into()],
3457 )
3458 .unwrap();
3459 let tool = &journal.state().tools["kweb"];
3460 let active_id = tool.slots[0].box_id;
3461 let direct_id = tool.slots[1].box_id;
3462 let projected_items = journal.state().projection().items;
3463 let rendered = projected_items
3464 .iter()
3465 .map(|item| item.text.as_str())
3466 .collect::<Vec<_>>()
3467 .join("\n\n");
3468 assert!(rendered.find("DIRECT NODE") < rendered.find("ACTIVE NODE"));
3469 assert_eq!(
3470 journal.state().tool_layouts["kweb"],
3471 vec![direct_id, active_id]
3472 );
3473 drop(journal);
3474 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3475 assert_eq!(reopened.state().projection().items, projected_items);
3476 std::fs::remove_file(path).unwrap();
3477 }
3478
3479 #[test]
3480 fn limits_use_exact_floor_percentages() {
3481 let state = Chatend::opened(SessionMetadata {
3482 effective_context_tokens: 101,
3483 ..metadata()
3484 });
3485 assert_eq!(state.live_context_limit(), 70);
3486 assert_eq!(state.forced_ingress_context_limit(), 75);
3487 assert_eq!(state.ingress_initial_context_limit(), 75);
3488 assert_eq!(state.ingress_context_limit(), 101);
3489 assert_eq!(state.active_context_limit(), 70);
3490 let ingress = Chatend::opened(SessionMetadata {
3491 kind: SessionKind::HistoryIngress,
3492 effective_context_tokens: 101,
3493 ..metadata()
3494 });
3495 assert_eq!(ingress.active_context_limit(), 101);
3496 assert_eq!(estimate_tokens("1234"), 1);
3497 }
3498
3499 #[test]
3500 fn new_box_projection_preview_matches_the_committed_projection() {
3501 let path = path("new-box-preview");
3502 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3503 let boxes = vec![
3504 (
3505 "User message".into(),
3506 BoxOwner::User,
3507 BoxContent::text("prospective user text"),
3508 ),
3509 (
3510 "attachment".into(),
3511 BoxOwner::User,
3512 BoxContent::text("prospective attachment text"),
3513 ),
3514 ];
3515 let preview = journal
3516 .state()
3517 .projection_with_new_boxes_at("t1", &boxes)
3518 .unwrap();
3519 for (name, owner, content) in boxes {
3520 journal.create_box("t1", name, owner, content).unwrap();
3521 }
3522 assert_eq!(journal.state().projection(), preview);
3523 std::fs::remove_file(path).unwrap();
3524 }
3525
3526 #[test]
3527 fn box_representation_preview_matches_the_batched_append() {
3528 let path = path("representation-plan");
3529 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3530 let hydrated = journal
3531 .create_box(
3532 "t1",
3533 "large result",
3534 BoxOwner::Controller,
3535 BoxContent::text("x".repeat(3_000)),
3536 )
3537 .unwrap();
3538 let summarized = journal
3539 .create_box(
3540 "t2",
3541 "Kennedy summary",
3542 BoxOwner::Kennedy,
3543 BoxContent::text("y".repeat(3_000)),
3544 )
3545 .unwrap();
3546 journal
3547 .summarize_box("t3", summarized, "important points")
3548 .unwrap();
3549 let desired = BTreeMap::from([
3550 (hydrated, BoxRepresentation::Dehydrated),
3551 (
3552 summarized,
3553 BoxRepresentation::Summarized("important points".into()),
3554 ),
3555 ]);
3556 let preview = journal
3557 .state()
3558 .projection_with_box_representations(&desired)
3559 .unwrap();
3560 let next_id = journal.state().next_id;
3561 let ids = journal.apply_box_representations("t4", &desired).unwrap();
3562 assert_eq!(ids, vec![EventId(next_id)]);
3563 assert_eq!(journal.state().projection(), preview);
3564 assert!(matches!(
3565 journal.state().box_state(hydrated).unwrap().representation,
3566 Representation::Dehydrated { .. }
3567 ));
3568 assert_eq!(
3569 journal
3570 .state()
3571 .box_state(summarized)
3572 .unwrap()
3573 .representation,
3574 Representation::Summarized {
3575 based_on: EventId(2),
3576 text: "important points".into(),
3577 }
3578 );
3579 std::fs::remove_file(path).unwrap();
3580 }
3581}