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
2467#[cfg(test)]
2468type SessionJournal = Session;
2469
2470#[cfg(test)]
2471mod tests {
2472 use std::path::PathBuf;
2473 use std::time::{SystemTime, UNIX_EPOCH};
2474
2475 use serde_json::json;
2476
2477 use super::*;
2478
2479 fn path(label: &str) -> PathBuf {
2480 std::env::temp_dir()
2481 .join(format!(
2482 "kennedy-chatend-{label}-{}-{}",
2483 std::process::id(),
2484 SystemTime::now()
2485 .duration_since(UNIX_EPOCH)
2486 .unwrap()
2487 .as_nanos()
2488 ))
2489 .join("session-1.session-log")
2490 }
2491
2492 fn metadata() -> SessionMetadata {
2493 SessionMetadata {
2494 session_id: "session-1".into(),
2495 kind: SessionKind::Conversation,
2496 created_at: "2026-07-23T00:00:00Z".into(),
2497 effective_context_tokens: 1_000,
2498 channel: json!({"kind":"test"}),
2499 }
2500 }
2501
2502 #[test]
2503 fn box_identity_continuations_staleness_and_replay_are_exact() {
2504 let path = path("boxes");
2505 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2506 let id = journal
2507 .create_box(
2508 "t1",
2509 "message",
2510 BoxOwner::User,
2511 BoxContent::text("original"),
2512 )
2513 .unwrap();
2514 assert_eq!(id, BoxId(1));
2515 journal.summarize_box("t2", id, "summary").unwrap();
2516 journal
2517 .update_box("t3", id, BoxContent::text("changed"))
2518 .unwrap();
2519 let state = journal.state().box_state(id).unwrap().clone();
2520 assert!(state.stale());
2521 assert_eq!(
2522 state.representation,
2523 Representation::Summarized {
2524 based_on: EventId(1),
2525 text: "summary".into()
2526 }
2527 );
2528 let projection = journal.state().projection();
2529 assert_eq!(projection.items[0].text, "[box updated]");
2530 assert_eq!(projection.items[1].text, "[box updated]");
2531 assert!(projection.items[2].text.contains("summary"));
2532 assert!(projection.items[2].stale);
2533 assert!(projection.footer.starts_with("[stale boxes: 1]\n"));
2534 assert!(projection.footer.contains("\n[current time: "));
2535 assert_eq!(
2536 projection.footer.lines().last().unwrap(),
2537 format!(
2538 "[current context size: {} | max context size: {}]",
2539 projection.estimated_tokens, projection.status.context_limit_tokens
2540 )
2541 );
2542 assert!(!projection.footer.contains("effective"));
2543 assert!(!projection.footer.contains("turn_limit"));
2544 assert!(projection.render().ends_with(&projection.footer));
2545 drop(journal);
2546 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2547 assert_eq!(reopened.state().box_state(id), Some(&state));
2548 std::fs::remove_file(path).unwrap();
2549 }
2550
2551 #[test]
2552 fn box_headers_hide_internal_ownership_and_timestamp_messages() {
2553 let path = path("box-headers");
2554 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2555 journal
2556 .create_box(
2557 "2026-07-23T00:00:01Z",
2558 "User message",
2559 BoxOwner::User,
2560 BoxContent::text("Hello"),
2561 )
2562 .unwrap();
2563 journal
2564 .create_box(
2565 "2026-07-23T00:00:02Z",
2566 "Kennedy message",
2567 BoxOwner::Kennedy,
2568 BoxContent::text("Hi"),
2569 )
2570 .unwrap();
2571 let mut content = BoxContent::text("Node ID: AAAAAAAB");
2572 content.use_concise_header();
2573 journal
2574 .create_box(
2575 "2026-07-23T00:00:03Z",
2576 "Kweb loaded node",
2577 BoxOwner::Tool {
2578 tool_instance: "kweb".into(),
2579 slot: "loaded".into(),
2580 },
2581 content,
2582 )
2583 .unwrap();
2584
2585 let projection = journal.state().projection();
2586 assert_eq!(
2587 projection.items[0].text,
2588 "[box 1 | User message | timestamp=2026-07-23T00:00:01Z | hydrated]\nHello"
2589 );
2590 assert_eq!(
2591 projection.items[1].text,
2592 "[box 2 | Kennedy message | timestamp=2026-07-23T00:00:02Z | hydrated]\nHi"
2593 );
2594 assert_eq!(
2595 projection.items[2].text,
2596 "[box 3 | Kweb loaded node | hydrated]\nNode ID: AAAAAAAB"
2597 );
2598 assert!(
2599 projection
2600 .items
2601 .iter()
2602 .all(|item| !item.text.contains("owner="))
2603 );
2604 let current_time = projection
2605 .footer
2606 .lines()
2607 .find_map(|line| line.strip_prefix("[current time: "))
2608 .and_then(|line| line.strip_suffix(']'))
2609 .unwrap();
2610 chrono::DateTime::parse_from_rfc3339(current_time).unwrap();
2611 assert!(current_time.starts_with("20"));
2612 std::fs::remove_file(path).unwrap();
2613 }
2614
2615 #[test]
2616 fn shared_pending_and_box_identity_space_never_overlaps() {
2617 let path = path("pending");
2618 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2619 let first = journal.allocate_pending_node("t1").unwrap();
2620 let box_id = journal
2621 .create_box("t2", "box", BoxOwner::Kennedy, BoxContent::text("hello"))
2622 .unwrap();
2623 let object = journal
2624 .stage_object(
2625 "t3",
2626 "application/octet-stream",
2627 None,
2628 Value::Null,
2629 b"\0binary\xff",
2630 )
2631 .unwrap();
2632 assert_eq!(first.to_string(), "pending:1");
2633 assert_eq!(box_id, BoxId(2));
2634 assert_eq!(object.to_string(), "pending:3");
2635 assert_eq!(journal.read_object(&object).unwrap(), b"\0binary\xff");
2636 let stored = journal.session_log();
2637 assert!(!stored.events[0].text.contains("pending_id"));
2638 assert!(!stored.events[1].text.contains("box_id"));
2639 assert!(!stored.events[2].text.contains("pending_id"));
2640 drop(journal);
2641 let mut reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2642 assert_eq!(reopened.read_object(&object).unwrap(), b"\0binary\xff");
2643 assert_eq!(reopened.state().next_id, 4);
2644 std::fs::remove_file(path).unwrap();
2645 }
2646
2647 #[test]
2648 fn kennedy_tool_completion_is_validated_before_storage_seals() {
2649 let path = path("seal-tool");
2650 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2651 journal
2652 .record(
2653 "t1",
2654 EventKind::ToolInvoked {
2655 tool_instance: "CreateNode:1".into(),
2656 tool_name: "CreateNode".into(),
2657 arguments: json!({}),
2658 invocation_id: None,
2659 },
2660 )
2661 .unwrap();
2662 assert!(journal.seal().is_err());
2663 journal
2664 .record(
2665 "t2",
2666 EventKind::ToolCompleted {
2667 tool_instance: "call_ktool".into(),
2668 tool_name: "call_ktool".into(),
2669 outcome: json!({"ok":true}),
2670 invocation_id: None,
2671 },
2672 )
2673 .unwrap();
2674 journal.seal().unwrap();
2675 assert!(journal.is_sealed());
2676 std::fs::remove_file(path).unwrap();
2677 }
2678
2679 #[test]
2680 fn interrupted_tools_are_recovered_by_identity_without_losing_legacy_journals() {
2681 let path = path("repair-tools");
2682 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2683 journal
2684 .record(
2685 "t1",
2686 EventKind::ToolInvoked {
2687 tool_instance: "WebSearch:legacy".into(),
2688 tool_name: "WebSearch".into(),
2689 arguments: json!({"question":"legacy"}),
2690 invocation_id: None,
2691 },
2692 )
2693 .unwrap();
2694 journal
2695 .record(
2696 "t2",
2697 EventKind::ToolInvoked {
2698 tool_instance: "WebSearch:new".into(),
2699 tool_name: "WebSearch".into(),
2700 arguments: json!({"question":"identified"}),
2701 invocation_id: Some("call-1".into()),
2702 },
2703 )
2704 .unwrap();
2705 assert!(journal.seal().is_err());
2706
2707 let repaired = journal.repair_unfinished_tools("recovery").unwrap();
2708 assert_eq!(repaired.len(), 2);
2709 assert!(
2710 journal
2711 .state()
2712 .events
2713 .iter()
2714 .rev()
2715 .take(2)
2716 .all(|event| matches!(
2717 &event.kind,
2718 EventKind::ToolCompleted { outcome, .. }
2719 if outcome.get("recovered").and_then(Value::as_bool) == Some(true)
2720 ))
2721 );
2722 journal.seal().unwrap();
2723 assert!(journal.is_sealed());
2724 std::fs::remove_file(path).unwrap();
2725 }
2726
2727 #[test]
2728 fn identified_tool_completions_can_arrive_out_of_order() {
2729 let path = path("tool-identity");
2730 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2731 for id in ["call-1", "call-2"] {
2732 journal
2733 .record(
2734 id,
2735 EventKind::ToolInvoked {
2736 tool_instance: format!("WebSearch:{id}"),
2737 tool_name: "WebSearch".into(),
2738 arguments: json!({"question":id}),
2739 invocation_id: Some(id.into()),
2740 },
2741 )
2742 .unwrap();
2743 }
2744 for id in ["call-1", "call-2"] {
2745 journal
2746 .record(
2747 format!("{id}-complete"),
2748 EventKind::ToolCompleted {
2749 tool_instance: format!("WebSearch:{id}"),
2750 tool_name: "WebSearch".into(),
2751 outcome: json!({"ok":true}),
2752 invocation_id: Some(id.into()),
2753 },
2754 )
2755 .unwrap();
2756 }
2757 journal.seal().unwrap();
2758 std::fs::remove_file(path).unwrap();
2759 }
2760
2761 #[test]
2762 fn invalid_box_operations_do_not_poison_the_append_only_journal() {
2763 let path = path("invalid-box-operation");
2764 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2765 let box_id = journal
2766 .create_box(
2767 "t1",
2768 "valid",
2769 BoxOwner::Controller,
2770 BoxContent::text("canonical"),
2771 )
2772 .unwrap();
2773 let valid_length = std::fs::metadata(&path).unwrap().len();
2774
2775 let result = journal.dehydrate_boxes("t2", &[box_id, BoxId(97)]);
2776 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2777 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2778 assert!(matches!(
2779 journal.state().box_state(box_id).unwrap().representation,
2780 Representation::Hydrated { .. }
2781 ));
2782
2783 for result in [
2784 journal.summarize_box("t3", BoxId(97), "summary"),
2785 journal.rehydrate_box("t4", BoxId(97)),
2786 journal.retire_box("t5", BoxId(97)),
2787 ] {
2788 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2789 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2790 }
2791
2792 drop(journal);
2793 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2794 assert_eq!(
2795 reopened
2796 .state()
2797 .box_state(box_id)
2798 .unwrap()
2799 .canonical
2800 .content,
2801 BoxContent::text("canonical")
2802 );
2803 std::fs::remove_file(path).unwrap();
2804 }
2805
2806 #[test]
2807 fn multiple_boxes_dehydrate_in_one_replayable_batch() {
2808 let path = path("dehydrate-boxes");
2809 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2810 let boxes = ["first", "second", "third"]
2811 .into_iter()
2812 .map(|text| {
2813 journal
2814 .create_box("t1", text, BoxOwner::Controller, BoxContent::text(text))
2815 .unwrap()
2816 })
2817 .collect::<Vec<_>>();
2818
2819 let event_ids = journal
2820 .dehydrate_boxes("t2", &[boxes[0], boxes[2]])
2821 .unwrap();
2822 assert_eq!(event_ids, [EventId(4), EventId(5)]);
2823 assert!(matches!(
2824 journal.state().box_state(boxes[0]).unwrap().representation,
2825 Representation::Dehydrated { .. }
2826 ));
2827 assert!(matches!(
2828 journal.state().box_state(boxes[1]).unwrap().representation,
2829 Representation::Hydrated { .. }
2830 ));
2831 assert!(matches!(
2832 journal.state().box_state(boxes[2]).unwrap().representation,
2833 Representation::Dehydrated { .. }
2834 ));
2835
2836 let valid_length = std::fs::metadata(&path).unwrap().len();
2837 assert_eq!(
2838 journal.dehydrate_boxes("t3", &[]).unwrap_err().to_string(),
2839 "at least one box must be selected for dehydration"
2840 );
2841 assert_eq!(
2842 journal
2843 .dehydrate_boxes("t3", &[boxes[1], boxes[1]])
2844 .unwrap_err()
2845 .to_string(),
2846 "box dehydration cannot contain duplicate box IDs"
2847 );
2848 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2849
2850 drop(journal);
2851 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2852 assert!(matches!(
2853 reopened.state().box_state(boxes[0]).unwrap().representation,
2854 Representation::Dehydrated { .. }
2855 ));
2856 assert!(matches!(
2857 reopened.state().box_state(boxes[2]).unwrap().representation,
2858 Representation::Dehydrated { .. }
2859 ));
2860 std::fs::remove_file(path).unwrap();
2861 }
2862
2863 #[test]
2864 fn status_estimates_the_context_with_every_active_box_hydrated() {
2865 let path = path("fully-hydrated-estimate");
2866 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2867 let box_id = journal
2868 .create_box(
2869 "t1",
2870 "large",
2871 BoxOwner::Controller,
2872 BoxContent::text("canonical material ".repeat(1_000)),
2873 )
2874 .unwrap();
2875 journal.dehydrate_boxes("t2", &[box_id]).unwrap();
2876
2877 let status = journal.state().projection().status;
2878 assert!(
2879 status.fully_hydrated_context_tokens > status.current_context_tokens,
2880 "the hydrated estimate should include the canonical body"
2881 );
2882 std::fs::remove_file(path).unwrap();
2883 }
2884
2885 #[test]
2886 fn provider_measurements_recalibrate_the_matching_manifest_then_track_deltas() {
2887 let path = path("calibration");
2888 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2889 journal
2890 .create_box(
2891 "t1",
2892 "system",
2893 BoxOwner::System,
2894 BoxContent::text("baseline provider content"),
2895 )
2896 .unwrap();
2897 let submitted = journal.state().projection();
2898 journal
2899 .record(
2900 "t2",
2901 EventKind::InferenceSubmitted {
2902 manifest_hash: "manifest-1".into(),
2903 estimated_input_tokens: submitted.estimated_tokens,
2904 raw_estimated_input_tokens: Some(submitted.raw_estimated_tokens),
2905 },
2906 )
2907 .unwrap();
2908 let tool_box = journal
2909 .create_box(
2910 "t3",
2911 "tool result",
2912 BoxOwner::Controller,
2913 BoxContent::text("x".repeat(3_000)),
2914 )
2915 .unwrap();
2916 let measured_context = journal.state().projection();
2917 journal
2918 .record(
2919 "t4",
2920 EventKind::ProviderReceipt {
2921 manifest_hash: "manifest-1".into(),
2922 input_tokens: Some(777),
2923 output_tokens: Some(3),
2924 context_bytes: Some(measured_context.context_bytes),
2925 raw_context_tokens: Some(measured_context.raw_estimated_tokens),
2926 provider_data: Value::Null,
2927 },
2928 )
2929 .unwrap();
2930 assert_eq!(journal.state().projection().estimated_tokens, 777);
2931 journal
2932 .create_box(
2933 "t5",
2934 "new material",
2935 BoxOwner::User,
2936 BoxContent::text("x".repeat(300)),
2937 )
2938 .unwrap();
2939 let expanded = journal.state().projection();
2940 assert_eq!(
2941 expanded.estimated_tokens,
2942 777 + expanded
2943 .context_bytes
2944 .abs_diff(measured_context.context_bytes)
2945 / ESTIMATED_BYTES_PER_TOKEN
2946 );
2947 assert!(expanded.raw_estimated_tokens > submitted.raw_estimated_tokens);
2948 journal.dehydrate_boxes("t6", &[tool_box]).unwrap();
2949 let shrunken = journal.state().projection();
2950 assert_eq!(
2951 shrunken.estimated_tokens,
2952 777_u64.saturating_sub(
2953 measured_context
2954 .context_bytes
2955 .abs_diff(shrunken.context_bytes)
2956 / ESTIMATED_BYTES_PER_TOKEN
2957 )
2958 );
2959 std::fs::remove_file(path).unwrap();
2960 }
2961
2962 #[test]
2963 fn session_status_keeps_provider_usage_categories_exact_and_exclusive() {
2964 let path = path("session-status");
2965 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2966 journal
2967 .create_box("t1", "system", BoxOwner::System, BoxContent::text("é"))
2968 .unwrap();
2969 let before_usage = journal.state().projection();
2970 assert_eq!(
2971 before_usage.raw_estimated_tokens,
2972 before_usage
2973 .context_bytes
2974 .div_ceil(ESTIMATED_BYTES_PER_TOKEN)
2975 );
2976 journal
2977 .record(
2978 "t2",
2979 EventKind::ProviderReceipt {
2980 manifest_hash: "manifest-1".into(),
2981 input_tokens: Some(100),
2982 output_tokens: Some(20),
2983 context_bytes: Some(before_usage.context_bytes),
2984 raw_context_tokens: Some(before_usage.raw_estimated_tokens),
2985 provider_data: json!({
2986 "usageIsDelta":true,
2987 "cachedInputTokens":40,
2988 "nonCachedInputTokens":60,
2989 "thinkingTokens":8,
2990 "outputTokens":12,
2991 "estimatedCostUsdNanos":12345
2992 }),
2993 },
2994 )
2995 .unwrap();
2996 let second_anchor = journal.state().projection();
2997 journal
2998 .record(
2999 "t3",
3000 EventKind::ProviderReceipt {
3001 manifest_hash: "manifest-1".into(),
3002 input_tokens: Some(120),
3003 output_tokens: Some(5),
3004 context_bytes: Some(second_anchor.context_bytes),
3005 raw_context_tokens: Some(second_anchor.raw_estimated_tokens),
3006 provider_data: json!({
3007 "usageIsDelta":true,
3008 "cachedInputTokens":10,
3009 "nonCachedInputTokens":10,
3010 "thinkingTokens":2,
3011 "outputTokens":3,
3012 "estimatedCostUsdNanos":6789
3013 }),
3014 },
3015 )
3016 .unwrap();
3017 let projection = journal.state().projection();
3018 assert_eq!(
3019 projection.status,
3020 SessionStatus {
3021 current_context_tokens: 120,
3022 fully_hydrated_context_tokens: 120,
3023 context_limit_tokens: 700,
3024 current_context_bytes: projection.context_bytes,
3025 cached_input_tokens: 50,
3026 non_cached_input_tokens: 70,
3027 thinking_tokens: 10,
3028 output_tokens: 15,
3029 estimated_cost_usd_nanos: 19_134,
3030 unpriced_provider_calls: 0,
3031 }
3032 );
3033 let archive: Value = serde_json::from_slice(&journal.archive_bytes().unwrap()).unwrap();
3034 assert_eq!(archive["context"]["status"]["cachedInputTokens"], 50);
3035 assert!(
3036 archive["chatendText"]
3037 .as_str()
3038 .unwrap()
3039 .ends_with(archive["context"]["footer"].as_str().unwrap())
3040 );
3041 std::fs::remove_file(path).unwrap();
3042 }
3043
3044 fn compatibility_cost(
3045 model: &str,
3046 metering: &ProviderMetering,
3047 ) -> Option<ProviderCostEstimate> {
3048 let ProviderMetering::Tokens(usage) = metering else {
3049 return None;
3050 };
3051 let base = match model {
3052 "gpt-5.6-sol" => 1_000,
3053 "gemini-3.1-pro-preview" => 2_000,
3054 _ => return None,
3055 };
3056 Some(ProviderCostEstimate {
3057 usd_nanos: base + usage.input_tokens,
3058 accuracy: json!("exact"),
3059 pricing_version: "test-prices".into(),
3060 })
3061 }
3062
3063 #[test]
3064 fn legacy_provider_costs_are_reconstructed_without_rewriting_history() {
3065 let path = path("legacy-provider-costs");
3066 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3067 journal
3068 .record(
3069 "t1",
3070 EventKind::ProviderReceipt {
3071 manifest_hash: "top".into(),
3072 input_tokens: Some(10),
3073 output_tokens: Some(1),
3074 context_bytes: None,
3075 raw_context_tokens: None,
3076 provider_data: json!({
3077 "usageIsDelta":true,
3078 "nonCachedInputTokens":10,
3079 "cachedInputTokens":0,
3080 "thinkingTokens":0,
3081 "outputTokens":1
3082 }),
3083 },
3084 )
3085 .unwrap();
3086 journal
3087 .record(
3088 "t2",
3089 EventKind::ToolInvoked {
3090 tool_instance: "RunSubagent:1".into(),
3091 tool_name: "RunSubagent".into(),
3092 arguments: json!({"model":"codex/gpt-5.6-sol"}),
3093 invocation_id: Some("subagent-1".into()),
3094 },
3095 )
3096 .unwrap();
3097 journal
3098 .record(
3099 "t3",
3100 EventKind::Note {
3101 label: "subagent_started".into(),
3102 value: json!({"providerModel":"gpt-5.6-sol"}),
3103 },
3104 )
3105 .unwrap();
3106 journal
3107 .record(
3108 "t4",
3109 EventKind::ProviderReceipt {
3110 manifest_hash: "subagent".into(),
3111 input_tokens: None,
3112 output_tokens: None,
3113 context_bytes: None,
3114 raw_context_tokens: None,
3115 provider_data: json!({
3116 "source":"subagent",
3117 "usageIsDelta":true,
3118 "nonCachedInputTokens":20,
3119 "cachedInputTokens":0,
3120 "thinkingTokens":0,
3121 "outputTokens":1
3122 }),
3123 },
3124 )
3125 .unwrap();
3126 journal
3127 .record(
3128 "t5",
3129 EventKind::ToolInvoked {
3130 tool_instance: "AnnotateMedia:1".into(),
3131 tool_name: "AnnotateMedia".into(),
3132 arguments: json!({"model":"gemini-3.1-pro-preview"}),
3133 invocation_id: Some("media-1".into()),
3134 },
3135 )
3136 .unwrap();
3137 journal
3138 .record(
3139 "t6",
3140 EventKind::ProviderReceipt {
3141 manifest_hash: "media".into(),
3142 input_tokens: None,
3143 output_tokens: None,
3144 context_bytes: None,
3145 raw_context_tokens: None,
3146 provider_data: json!({
3147 "source":"media_annotation",
3148 "usageIsDelta":true,
3149 "nonCachedInputTokens":30,
3150 "cachedInputTokens":0,
3151 "thinkingTokens":0,
3152 "outputTokens":1
3153 }),
3154 },
3155 )
3156 .unwrap();
3157 journal
3158 .record(
3159 "t7",
3160 EventKind::ToolCompleted {
3161 tool_instance: "AnnotateMedia:1".into(),
3162 tool_name: "AnnotateMedia".into(),
3163 outcome: json!({"ok":true}),
3164 invocation_id: Some("media-1".into()),
3165 },
3166 )
3167 .unwrap();
3168 journal
3169 .record(
3170 "t8",
3171 EventKind::ToolCompleted {
3172 tool_instance: "RunSubagent:1".into(),
3173 tool_name: "RunSubagent".into(),
3174 outcome: json!({"ok":true}),
3175 invocation_id: Some("subagent-1".into()),
3176 },
3177 )
3178 .unwrap();
3179 journal
3180 .record(
3181 "t9",
3182 EventKind::ProviderReceipt {
3183 manifest_hash: "unknown".into(),
3184 input_tokens: None,
3185 output_tokens: None,
3186 context_bytes: None,
3187 raw_context_tokens: None,
3188 provider_data: json!({
3189 "source":"web_search",
3190 "usageIsDelta":true,
3191 "nonCachedInputTokens":40,
3192 "cachedInputTokens":0,
3193 "thinkingTokens":0,
3194 "outputTokens":1
3195 }),
3196 },
3197 )
3198 .unwrap();
3199 drop(journal);
3200
3201 let compatible = SessionJournal::open_with_metadata_and_provider_costs(
3202 &path,
3203 metadata(),
3204 Some("gpt-5.6-sol"),
3205 Some(compatibility_cost),
3206 )
3207 .unwrap();
3208 let status = &compatible.state().projection().status;
3209 assert_eq!(status.estimated_cost_usd_nanos, 4_060);
3210 assert_eq!(status.unpriced_provider_calls, 1);
3211 let inferred_models = compatible
3212 .state()
3213 .events
3214 .iter()
3215 .filter_map(|event| match &event.kind {
3216 EventKind::ProviderReceipt { provider_data, .. } => provider_data
3217 .get("providerModel")
3218 .and_then(Value::as_str)
3219 .map(str::to_owned),
3220 _ => None,
3221 })
3222 .collect::<Vec<_>>();
3223 assert_eq!(
3224 inferred_models,
3225 vec!["gpt-5.6-sol", "gpt-5.6-sol", "gemini-3.1-pro-preview"]
3226 );
3227 drop(compatible);
3228
3229 let unchanged = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3230 assert_eq!(
3231 unchanged
3232 .state()
3233 .projection()
3234 .status
3235 .estimated_cost_usd_nanos,
3236 0
3237 );
3238 assert_eq!(
3239 unchanged
3240 .state()
3241 .projection()
3242 .status
3243 .unpriced_provider_calls,
3244 4
3245 );
3246 let archive: Value = serde_json::from_slice(&unchanged.archive_bytes().unwrap()).unwrap();
3247 let archive_summary = legacy_provider_cost_summary_for_archive(
3248 &archive,
3249 Some("gpt-5.6-sol"),
3250 compatibility_cost,
3251 )
3252 .unwrap();
3253 assert_eq!(archive_summary.estimated_cost_usd_nanos, 4_060);
3254 assert_eq!(archive_summary.unpriced_provider_calls, 1);
3255 assert_eq!(archive["context"]["status"]["estimatedCostUsdNanos"], 0);
3256 assert_eq!(archive["context"]["status"]["unpricedProviderCalls"], 4);
3257 drop(unchanged);
3258 std::fs::remove_file(path).unwrap();
3259 }
3260
3261 #[test]
3262 fn legacy_provider_receipts_without_a_raw_context_anchor_remain_readable() {
3263 let receipt: EventKind = serde_json::from_value(json!({
3264 "type":"provider_receipt",
3265 "manifest_hash":"legacy-manifest",
3266 "input_tokens":123,
3267 "output_tokens":4,
3268 "provider_data":null
3269 }))
3270 .unwrap();
3271 assert!(matches!(
3272 receipt,
3273 EventKind::ProviderReceipt {
3274 context_bytes: None,
3275 raw_context_tokens: None,
3276 ..
3277 }
3278 ));
3279 }
3280
3281 #[test]
3282 fn stateful_tool_slots_are_batched_append_only_and_do_not_see_summaries() {
3283 let path = path("slots");
3284 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3285 journal
3286 .apply_tool_slots(
3287 "t1",
3288 "rust-1",
3289 vec![
3290 ToolSlotInput {
3291 slot: "a.rs".into(),
3292 name: "a.rs".into(),
3293 content: BoxContent::text("a"),
3294 retired: false,
3295 },
3296 ToolSlotInput {
3297 slot: "b.rs".into(),
3298 name: "b.rs".into(),
3299 content: BoxContent::text("b"),
3300 retired: false,
3301 },
3302 ],
3303 )
3304 .unwrap();
3305 let a = journal.state().tools["rust-1"].slots[0].box_id;
3306 journal.summarize_box("t2", a, "Kennedy summary").unwrap();
3307 let before = journal.state().clone();
3308 assert!(
3309 journal
3310 .apply_tool_slots(
3311 "t3",
3312 "rust-1",
3313 vec![ToolSlotInput {
3314 slot: "b.rs".into(),
3315 name: "b.rs".into(),
3316 content: BoxContent::text("b"),
3317 retired: false,
3318 }]
3319 )
3320 .is_err()
3321 );
3322 assert_eq!(journal.state(), &before);
3323 journal
3324 .apply_tool_slots(
3325 "t4",
3326 "rust-1",
3327 vec![
3328 ToolSlotInput {
3329 slot: "a.rs".into(),
3330 name: "a.rs".into(),
3331 content: BoxContent::text("a2"),
3332 retired: false,
3333 },
3334 ToolSlotInput {
3335 slot: "b.rs".into(),
3336 name: "b.rs".into(),
3337 content: BoxContent::text("b"),
3338 retired: true,
3339 },
3340 ToolSlotInput {
3341 slot: "c.rs".into(),
3342 name: "c.rs".into(),
3343 content: BoxContent::text("c"),
3344 retired: false,
3345 },
3346 ],
3347 )
3348 .unwrap();
3349 let a_state = journal.state().box_state(a).unwrap();
3350 assert_eq!(a_state.canonical.content.text, "a2");
3351 assert!(matches!(
3352 a_state.representation,
3353 Representation::Summarized { .. }
3354 ));
3355 drop(journal);
3356 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3357 assert_eq!(reopened.state().tools["rust-1"].slots.len(), 3);
3358 assert!(reopened.state().tools["rust-1"].slots[1].retired);
3359 std::fs::remove_file(path).unwrap();
3360 }
3361
3362 #[test]
3363 fn tool_layout_orders_current_boxes_without_changing_their_identities() {
3364 let path = path("tool-layout");
3365 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3366 journal
3367 .apply_tool_slots_with_layout(
3368 "t1",
3369 "kweb",
3370 vec![
3371 ToolSlotInput {
3372 slot: "active".into(),
3373 name: "Active node".into(),
3374 content: BoxContent::text("ACTIVE NODE"),
3375 retired: false,
3376 },
3377 ToolSlotInput {
3378 slot: "direct".into(),
3379 name: "Direct node".into(),
3380 content: BoxContent::text("DIRECT NODE"),
3381 retired: false,
3382 },
3383 ],
3384 &["direct".into(), "active".into()],
3385 )
3386 .unwrap();
3387 let tool = &journal.state().tools["kweb"];
3388 let active_id = tool.slots[0].box_id;
3389 let direct_id = tool.slots[1].box_id;
3390 let projected_items = journal.state().projection().items;
3391 let rendered = projected_items
3392 .iter()
3393 .map(|item| item.text.as_str())
3394 .collect::<Vec<_>>()
3395 .join("\n\n");
3396 assert!(rendered.find("DIRECT NODE") < rendered.find("ACTIVE NODE"));
3397 assert_eq!(
3398 journal.state().tool_layouts["kweb"],
3399 vec![direct_id, active_id]
3400 );
3401 drop(journal);
3402 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3403 assert_eq!(reopened.state().projection().items, projected_items);
3404 std::fs::remove_file(path).unwrap();
3405 }
3406
3407 #[test]
3408 fn limits_use_exact_floor_percentages() {
3409 let state = Chatend::opened(SessionMetadata {
3410 effective_context_tokens: 101,
3411 ..metadata()
3412 });
3413 assert_eq!(state.live_context_limit(), 70);
3414 assert_eq!(state.forced_ingress_context_limit(), 75);
3415 assert_eq!(state.ingress_initial_context_limit(), 75);
3416 assert_eq!(state.ingress_context_limit(), 101);
3417 assert_eq!(state.active_context_limit(), 70);
3418 let ingress = Chatend::opened(SessionMetadata {
3419 kind: SessionKind::HistoryIngress,
3420 effective_context_tokens: 101,
3421 ..metadata()
3422 });
3423 assert_eq!(ingress.active_context_limit(), 101);
3424 assert_eq!(estimate_tokens("1234"), 1);
3425 }
3426
3427 #[test]
3428 fn new_box_projection_preview_matches_the_committed_projection() {
3429 let path = path("new-box-preview");
3430 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3431 let boxes = vec![
3432 (
3433 "User message".into(),
3434 BoxOwner::User,
3435 BoxContent::text("prospective user text"),
3436 ),
3437 (
3438 "attachment".into(),
3439 BoxOwner::User,
3440 BoxContent::text("prospective attachment text"),
3441 ),
3442 ];
3443 let preview = journal
3444 .state()
3445 .projection_with_new_boxes_at("t1", &boxes)
3446 .unwrap();
3447 for (name, owner, content) in boxes {
3448 journal.create_box("t1", name, owner, content).unwrap();
3449 }
3450 assert_eq!(journal.state().projection(), preview);
3451 std::fs::remove_file(path).unwrap();
3452 }
3453
3454 #[test]
3455 fn box_representation_preview_matches_the_batched_append() {
3456 let path = path("representation-plan");
3457 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3458 let hydrated = journal
3459 .create_box(
3460 "t1",
3461 "large result",
3462 BoxOwner::Controller,
3463 BoxContent::text("x".repeat(3_000)),
3464 )
3465 .unwrap();
3466 let summarized = journal
3467 .create_box(
3468 "t2",
3469 "Kennedy summary",
3470 BoxOwner::Kennedy,
3471 BoxContent::text("y".repeat(3_000)),
3472 )
3473 .unwrap();
3474 journal
3475 .summarize_box("t3", summarized, "important points")
3476 .unwrap();
3477 let desired = BTreeMap::from([
3478 (hydrated, BoxRepresentation::Dehydrated),
3479 (
3480 summarized,
3481 BoxRepresentation::Summarized("important points".into()),
3482 ),
3483 ]);
3484 let preview = journal
3485 .state()
3486 .projection_with_box_representations(&desired)
3487 .unwrap();
3488 let next_id = journal.state().next_id;
3489 let ids = journal.apply_box_representations("t4", &desired).unwrap();
3490 assert_eq!(ids, vec![EventId(next_id)]);
3491 assert_eq!(journal.state().projection(), preview);
3492 assert!(matches!(
3493 journal.state().box_state(hydrated).unwrap().representation,
3494 Representation::Dehydrated { .. }
3495 ));
3496 assert_eq!(
3497 journal
3498 .state()
3499 .box_state(summarized)
3500 .unwrap()
3501 .representation,
3502 Representation::Summarized {
3503 based_on: EventId(2),
3504 text: "important points".into(),
3505 }
3506 );
3507 std::fs::remove_file(path).unwrap();
3508 }
3509}