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_with_footer_lines(&[])
814 }
815
816 pub fn projection_with_footer_lines(&self, footer_lines: &[String]) -> ContextProjection {
818 self.projection_at(¤t_time(), footer_lines)
819 }
820
821 fn projection_at(&self, current_time: &str, footer_lines: &[String]) -> ContextProjection {
822 let items = self.projection_items(false);
823 let stale_boxes = self
824 .active_boxes()
825 .filter(|state| state.stale())
826 .map(|state| state.id)
827 .collect::<Vec<_>>();
828 let stale_label = if stale_boxes.is_empty() {
829 "none".into()
830 } else {
831 stale_boxes
832 .iter()
833 .map(ToString::to_string)
834 .collect::<Vec<_>>()
835 .join(", ")
836 };
837 let mut estimated_tokens = estimate_bytes(projected_context_bytes(&items, ""));
838 let mut footer = String::new();
839 let mut context_bytes = 0;
840 let mut raw_estimated_tokens = 0;
841 for _ in 0..8 {
844 footer = projection_footer(
845 &stale_label,
846 current_time,
847 footer_lines,
848 estimated_tokens,
849 self.active_context_limit(),
850 );
851 context_bytes = projected_context_bytes(&items, &footer);
852 raw_estimated_tokens = estimate_bytes(context_bytes);
853 let calibrated = self.calibrated_estimate(context_bytes, raw_estimated_tokens);
854 if calibrated == estimated_tokens {
855 break;
856 }
857 estimated_tokens = calibrated;
858 }
859 let fully_hydrated_context_tokens =
860 self.fully_hydrated_context_tokens_at(current_time, footer_lines);
861 let usage = self.cumulative_token_usage();
862 let cost = self.cumulative_cost();
863 let status = SessionStatus {
864 current_context_tokens: estimated_tokens,
865 fully_hydrated_context_tokens,
866 context_limit_tokens: self.active_context_limit(),
867 current_context_bytes: context_bytes,
868 cached_input_tokens: usage.cached_input_tokens,
869 non_cached_input_tokens: usage.non_cached_input_tokens,
870 thinking_tokens: usage.thinking_tokens,
871 output_tokens: usage.output_tokens,
872 estimated_cost_usd_nanos: cost.estimated_cost_usd_nanos,
873 unpriced_provider_calls: cost.unpriced_provider_calls,
874 };
875 ContextProjection {
876 items,
877 stale_boxes,
878 footer,
879 estimated_tokens,
880 raw_estimated_tokens,
881 context_bytes,
882 status,
883 }
884 }
885
886 fn projection_items(&self, fully_hydrated: bool) -> Vec<ProjectionItem> {
887 let mut next_occurrence = HashMap::new();
888 for state in self.boxes.values() {
889 for pair in state.occurrence_events.windows(2) {
890 next_occurrence.insert(pair[0], pair[1]);
891 }
892 }
893 let mut items = Vec::new();
894 for event in &self.events {
895 let Some(box_id) = event_box_id(&event.kind) else {
896 continue;
897 };
898 let Some(state) = self.boxes.get(&box_id) else {
899 continue;
900 };
901 if next_occurrence.contains_key(&event.id) {
902 items.push(ProjectionItem::marker(
903 event.id,
904 box_id,
905 "[box updated]".into(),
906 ));
907 continue;
908 }
909 if !state.active || state.occurrence_events.last() != Some(&event.id) {
910 continue;
911 }
912 let (representation, body) = if fully_hydrated {
913 ("hydrated", state.canonical.content.render())
914 } else {
915 match &state.representation {
916 Representation::Hydrated { .. } => {
917 ("hydrated", state.canonical.content.render())
918 }
919 Representation::Dehydrated { .. } => (
920 "dehydrated",
921 format!(
922 "[contents dehydrated; hydrate box {} to inspect the latest canonical revision]",
923 box_id
924 ),
925 ),
926 Representation::Summarized { text, .. } => ("summarized", text.clone()),
927 }
928 };
929 let stale = !fully_hydrated && state.stale();
930 let mut header = vec![format!("box {box_id}"), state.name.clone()];
931 if matches!(
932 (&state.owner, state.name.as_str()),
933 (BoxOwner::User, "User message") | (BoxOwner::Kennedy, "Kennedy message")
934 ) {
935 header.push(format!("timestamp={}", event.recorded_at));
936 }
937 header.push(representation.into());
938 if stale {
939 header.push("stale".into());
940 }
941 let mut text = format!("[{}]\n{}", header.join(" | "), body);
942 if text.ends_with('\n') {
943 text.pop();
944 }
945 items.push(ProjectionItem {
946 event_id: event.id,
947 box_id,
948 marker: false,
949 stale,
950 approximate_tokens: estimate_tokens(&text),
951 text,
952 });
953 }
954 for box_ids in self.tool_layouts.values() {
955 arrange_tool_projection(&mut items, box_ids);
956 }
957 items
958 }
959
960 fn fully_hydrated_context_tokens_at(&self, current_time: &str, footer_lines: &[String]) -> u64 {
961 let items = self.projection_items(true);
962 let mut estimated_tokens = estimate_bytes(projected_context_bytes(&items, ""));
963 for _ in 0..8 {
964 let footer = projection_footer(
965 "none",
966 current_time,
967 footer_lines,
968 estimated_tokens,
969 self.active_context_limit(),
970 );
971 let context_bytes = projected_context_bytes(&items, &footer);
972 let raw_estimated_tokens = estimate_bytes(context_bytes);
973 let calibrated = self.calibrated_estimate(context_bytes, raw_estimated_tokens);
974 if calibrated == estimated_tokens {
975 break;
976 }
977 estimated_tokens = calibrated;
978 }
979 estimated_tokens
980 }
981
982 fn calibrated_estimate(&self, current_bytes: u64, raw_current: u64) -> u64 {
983 let Some((manifest_hash, measured, bytes_at_receipt, raw_at_receipt)) =
984 self.events.iter().rev().find_map(|event| {
985 let EventKind::ProviderReceipt {
986 manifest_hash,
987 input_tokens: Some(input_tokens),
988 context_bytes,
989 raw_context_tokens,
990 ..
991 } = &event.kind
992 else {
993 return None;
994 };
995 Some((
996 manifest_hash,
997 *input_tokens,
998 *context_bytes,
999 *raw_context_tokens,
1000 ))
1001 })
1002 else {
1003 return raw_current;
1004 };
1005 if let Some(bytes_at_receipt) = bytes_at_receipt {
1006 let token_delta = current_bytes.abs_diff(bytes_at_receipt) / ESTIMATED_BYTES_PER_TOKEN;
1007 return if current_bytes >= bytes_at_receipt {
1008 measured.saturating_add(token_delta)
1009 } else {
1010 measured.saturating_sub(token_delta)
1011 };
1012 }
1013 let raw_at_measurement = match raw_at_receipt {
1014 Some(raw) => raw,
1015 None => {
1016 let Some(raw) = self.events.iter().rev().find_map(|event| {
1017 let EventKind::InferenceSubmitted {
1018 manifest_hash: submitted,
1019 estimated_input_tokens,
1020 raw_estimated_input_tokens,
1021 } = &event.kind
1022 else {
1023 return None;
1024 };
1025 (submitted == manifest_hash)
1026 .then_some(raw_estimated_input_tokens.unwrap_or(*estimated_input_tokens))
1027 }) else {
1028 return raw_current;
1029 };
1030 raw
1031 }
1032 };
1033 if raw_current >= raw_at_measurement {
1034 measured.saturating_add(raw_current - raw_at_measurement)
1035 } else {
1036 measured.saturating_sub(raw_at_measurement - raw_current)
1037 }
1038 }
1039
1040 fn cumulative_token_usage(&self) -> CumulativeTokenUsage {
1041 let mut total = CumulativeTokenUsage::default();
1042 for event in &self.events {
1043 let EventKind::ProviderReceipt { provider_data, .. } = &event.kind else {
1044 continue;
1045 };
1046 let cached = provider_u64(provider_data, &["cachedInputTokens", "cached_input_tokens"]);
1047 let thinking = provider_u64(
1048 provider_data,
1049 &[
1050 "thinkingTokens",
1051 "thinking_tokens",
1052 "reasoningOutputTokens",
1053 "reasoning_output_tokens",
1054 ],
1055 );
1056 let normalized_delta = provider_data
1057 .get("usageIsDelta")
1058 .and_then(Value::as_bool)
1059 .unwrap_or(false);
1060 let non_cached = if normalized_delta {
1061 provider_u64(
1062 provider_data,
1063 &["nonCachedInputTokens", "non_cached_input_tokens"],
1064 )
1065 } else {
1066 provider_u64(provider_data, &["inputTokens", "input_tokens"]).saturating_sub(cached)
1067 };
1068 let output = if normalized_delta {
1069 provider_u64(provider_data, &["outputTokens", "output_tokens"])
1070 } else {
1071 provider_u64(provider_data, &["outputTokens", "output_tokens"])
1072 .saturating_sub(thinking)
1073 };
1074 total.cached_input_tokens = total.cached_input_tokens.saturating_add(cached);
1075 total.non_cached_input_tokens =
1076 total.non_cached_input_tokens.saturating_add(non_cached);
1077 total.thinking_tokens = total.thinking_tokens.saturating_add(thinking);
1078 total.output_tokens = total.output_tokens.saturating_add(output);
1079 }
1080 total
1081 }
1082
1083 fn cumulative_cost(&self) -> CumulativeCost {
1084 let mut total = CumulativeCost::default();
1085 for event in &self.events {
1086 let EventKind::ProviderReceipt { provider_data, .. } = &event.kind else {
1087 continue;
1088 };
1089 if let Some(cost) = provider_data
1090 .get("estimatedCostUsdNanos")
1091 .and_then(Value::as_u64)
1092 {
1093 total.estimated_cost_usd_nanos =
1094 total.estimated_cost_usd_nanos.saturating_add(cost);
1095 } else {
1096 total.unpriced_provider_calls = total.unpriced_provider_calls.saturating_add(1);
1097 }
1098 }
1099 total
1100 }
1101
1102 pub fn render(&self) -> String {
1103 self.projection().render()
1104 }
1105
1106 fn apply_transition(&mut self, transition: &Transition) -> anyhow::Result<()> {
1107 ensure!(
1108 !transition.events.is_empty(),
1109 "a transition cannot be empty"
1110 );
1111 for event in &transition.events {
1112 self.apply_event(event)?;
1113 }
1114 Ok(())
1115 }
1116
1117 fn apply_event(&mut self, event: &Event) -> anyhow::Result<()> {
1118 ensure!(
1119 event.id.0 >= self.next_id,
1120 "event {} reuses an allocated identity (next is {})",
1121 event.id,
1122 self.next_id
1123 );
1124 self.next_id = event.id.0.checked_add(1).context("event ID overflow")?;
1125 match &event.kind {
1126 EventKind::SessionConfigured {
1127 effective_context_tokens,
1128 kind,
1129 } => {
1130 ensure!(
1131 *effective_context_tokens > 0,
1132 "effective context window must be positive"
1133 );
1134 self.metadata.effective_context_tokens = *effective_context_tokens;
1135 self.metadata.kind = kind.clone();
1136 }
1137 EventKind::BoxCreated {
1138 box_id,
1139 name,
1140 owner,
1141 content,
1142 } => {
1143 ensure!(
1144 box_id.0 == event.id.0,
1145 "BoxId must equal its creation EventId"
1146 );
1147 ensure!(
1148 !self.boxes.contains_key(box_id),
1149 "box {} already exists",
1150 box_id
1151 );
1152 self.boxes.insert(
1153 *box_id,
1154 BoxState {
1155 id: *box_id,
1156 name: name.clone(),
1157 owner: owner.clone(),
1158 created_at: event.id,
1159 canonical: CanonicalRevision {
1160 event_id: event.id,
1161 content: content.clone(),
1162 },
1163 representation: Representation::Hydrated {
1164 canonical_event: event.id,
1165 },
1166 occurrence_events: vec![event.id],
1167 active: true,
1168 },
1169 );
1170 if let BoxOwner::Tool {
1171 tool_instance,
1172 slot,
1173 } = owner
1174 {
1175 self.tools
1176 .entry(tool_instance.clone())
1177 .or_default()
1178 .slots
1179 .push(ToolSlot {
1180 slot: slot.clone(),
1181 box_id: *box_id,
1182 retired: false,
1183 });
1184 }
1185 }
1186 EventKind::CanonicalUpdated { box_id, content } => {
1187 let state = active_box_mut(&mut self.boxes, *box_id)?;
1188 state.canonical = CanonicalRevision {
1189 event_id: event.id,
1190 content: content.clone(),
1191 };
1192 if matches!(state.representation, Representation::Hydrated { .. }) {
1193 state.representation = Representation::Hydrated {
1194 canonical_event: event.id,
1195 };
1196 }
1197 state.occurrence_events.push(event.id);
1198 }
1199 EventKind::BoxRenamed { box_id, name } => {
1200 ensure!(!name.trim().is_empty(), "a box name cannot be empty");
1201 let state = active_box_mut(&mut self.boxes, *box_id)?;
1202 state.name = name.clone();
1203 state.occurrence_events.push(event.id);
1204 }
1205 EventKind::BoxDehydrated { box_id } => {
1206 let state = active_box_mut(&mut self.boxes, *box_id)?;
1207 state.representation = Representation::Dehydrated {
1208 based_on: state.canonical.event_id,
1209 };
1210 state.occurrence_events.push(event.id);
1211 }
1212 EventKind::BoxSummarized { box_id, text } => {
1213 ensure!(!text.trim().is_empty(), "a box summary cannot be empty");
1214 let state = active_box_mut(&mut self.boxes, *box_id)?;
1215 state.representation = Representation::Summarized {
1216 based_on: state.canonical.event_id,
1217 text: text.clone(),
1218 };
1219 state.occurrence_events.push(event.id);
1220 }
1221 EventKind::BoxRehydrated { box_id } => {
1222 let state = active_box_mut(&mut self.boxes, *box_id)?;
1223 state.representation = Representation::Hydrated {
1224 canonical_event: state.canonical.event_id,
1225 };
1226 state.occurrence_events.push(event.id);
1227 }
1228 EventKind::BoxRetired { box_id } => {
1229 let tool_instance = self.boxes.get(box_id).and_then(|state| {
1230 let BoxOwner::Tool { tool_instance, .. } = &state.owner else {
1231 return None;
1232 };
1233 Some(tool_instance.clone())
1234 });
1235 let state = active_box_mut(&mut self.boxes, *box_id)?;
1236 state.active = false;
1237 state.occurrence_events.push(event.id);
1238 if let Some(tool_instance) = tool_instance {
1239 let slot = self
1240 .tools
1241 .get_mut(&tool_instance)
1242 .and_then(|tool| tool.slots.iter_mut().find(|slot| slot.box_id == *box_id))
1243 .with_context(|| {
1244 format!("tool box {box_id} is missing from {tool_instance}")
1245 })?;
1246 slot.retired = true;
1247 }
1248 }
1249 EventKind::ToolLayoutChanged {
1250 tool_instance,
1251 box_ids,
1252 } => {
1253 let mut unique = std::collections::HashSet::new();
1254 for box_id in box_ids {
1255 ensure!(
1256 unique.insert(*box_id),
1257 "tool layout contains duplicate box {box_id}"
1258 );
1259 let state = self
1260 .boxes
1261 .get(box_id)
1262 .with_context(|| format!("tool layout references missing box {box_id}"))?;
1263 ensure!(state.active, "tool layout references retired box {box_id}");
1264 ensure!(
1265 matches!(
1266 &state.owner,
1267 BoxOwner::Tool {
1268 tool_instance: owner,
1269 ..
1270 } if owner == tool_instance
1271 ),
1272 "tool layout box {box_id} belongs to another tool"
1273 );
1274 }
1275 self.tool_layouts
1276 .insert(tool_instance.clone(), box_ids.clone());
1277 }
1278 EventKind::PendingAllocated {
1279 pending_id,
1280 resource,
1281 } => {
1282 ensure!(
1283 pending_id.number() == event.id.0,
1284 "pending identity must equal its allocation EventId"
1285 );
1286 ensure!(
1287 self.pending
1288 .insert(pending_id.clone(), resource.clone())
1289 .is_none(),
1290 "pending identity {} already exists",
1291 pending_id
1292 );
1293 }
1294 EventKind::SourceTerminated { .. } => self.source_terminated = true,
1295 EventKind::HistoryIngressStarted => {
1296 ensure!(
1297 self.source_terminated,
1298 "history ingress requires source termination"
1299 );
1300 self.history_ingress_started = true;
1301 }
1302 EventKind::SessionCompleted { session_object_id } => {
1303 self.completed_session_object = Some(session_object_id.clone());
1304 }
1305 EventKind::ToolInvoked { .. }
1306 | EventKind::ToolCompleted { .. }
1307 | EventKind::InferenceSubmitted { .. }
1308 | EventKind::ProviderReceipt { .. }
1309 | EventKind::CapacityError { .. }
1310 | EventKind::HistoryEventInspected { .. }
1311 | EventKind::HistoryEventReleased { .. }
1312 | EventKind::KwebPlanChanged { .. }
1313 | EventKind::KwebCommitted { .. }
1314 | EventKind::Note { .. } => {}
1315 }
1316 self.events.push(event.clone());
1317 Ok(())
1318 }
1319}
1320
1321fn active_box_mut(
1322 boxes: &mut BTreeMap<BoxId, BoxState>,
1323 box_id: BoxId,
1324) -> anyhow::Result<&mut BoxState> {
1325 let state = boxes
1326 .get_mut(&box_id)
1327 .with_context(|| format!("box {box_id} does not exist"))?;
1328 ensure!(state.active, "box {box_id} is retired");
1329 Ok(state)
1330}
1331
1332fn arrange_tool_projection(items: &mut Vec<ProjectionItem>, box_ids: &[BoxId]) {
1333 if box_ids.is_empty() {
1334 return;
1335 }
1336 let ranks = box_ids
1337 .iter()
1338 .enumerate()
1339 .map(|(rank, box_id)| (*box_id, rank))
1340 .collect::<HashMap<_, _>>();
1341 let insertion = items
1342 .iter()
1343 .position(|item| !item.marker && ranks.contains_key(&item.box_id));
1344 let Some(insertion) = insertion else {
1345 return;
1346 };
1347 let mut arranged = Vec::with_capacity(box_ids.len());
1348 let mut retained = Vec::with_capacity(items.len());
1349 for item in std::mem::take(items) {
1350 if !item.marker && ranks.contains_key(&item.box_id) {
1351 arranged.push(item);
1352 } else {
1353 retained.push(item);
1354 }
1355 }
1356 arranged.sort_by_key(|item| ranks[&item.box_id]);
1357 let insertion = insertion.min(retained.len());
1358 retained.splice(insertion..insertion, arranged);
1359 *items = retained;
1360}
1361
1362fn event_box_id(kind: &EventKind) -> Option<BoxId> {
1363 match kind {
1364 EventKind::BoxCreated { box_id, .. }
1365 | EventKind::CanonicalUpdated { box_id, .. }
1366 | EventKind::BoxRenamed { box_id, .. }
1367 | EventKind::BoxDehydrated { box_id }
1368 | EventKind::BoxSummarized { box_id, .. }
1369 | EventKind::BoxRehydrated { box_id }
1370 | EventKind::BoxRetired { box_id } => Some(*box_id),
1371 _ => None,
1372 }
1373}
1374
1375#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
1376#[serde(rename_all = "camelCase")]
1377pub struct ProjectionItem {
1378 pub event_id: EventId,
1379 pub box_id: BoxId,
1380 pub marker: bool,
1381 pub stale: bool,
1382 pub approximate_tokens: u64,
1383 pub text: String,
1384}
1385
1386impl ProjectionItem {
1387 fn marker(event_id: EventId, box_id: BoxId, text: String) -> Self {
1388 Self {
1389 event_id,
1390 box_id,
1391 marker: true,
1392 stale: false,
1393 approximate_tokens: estimate_tokens(&text),
1394 text,
1395 }
1396 }
1397}
1398
1399#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
1400#[serde(rename_all = "camelCase")]
1401pub struct ContextProjection {
1402 pub items: Vec<ProjectionItem>,
1403 pub stale_boxes: Vec<BoxId>,
1404 pub footer: String,
1405 pub estimated_tokens: u64,
1406 pub raw_estimated_tokens: u64,
1407 pub context_bytes: u64,
1408 pub status: SessionStatus,
1409}
1410
1411impl ContextProjection {
1412 pub fn render(&self) -> String {
1413 let mut blocks = self
1414 .items
1415 .iter()
1416 .map(|item| item.text.as_str())
1417 .collect::<Vec<_>>();
1418 blocks.push(&self.footer);
1419 blocks.join("\n\n")
1420 }
1421}
1422
1423#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
1424#[serde(rename_all = "camelCase")]
1425pub struct SessionStatus {
1426 pub current_context_tokens: u64,
1427 #[serde(default)]
1428 pub fully_hydrated_context_tokens: u64,
1429 pub context_limit_tokens: u64,
1430 pub current_context_bytes: u64,
1431 pub cached_input_tokens: u64,
1432 pub non_cached_input_tokens: u64,
1433 pub thinking_tokens: u64,
1434 pub output_tokens: u64,
1435 #[serde(default)]
1436 pub estimated_cost_usd_nanos: u64,
1437 #[serde(default)]
1438 pub unpriced_provider_calls: u64,
1439}
1440
1441#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1442struct CumulativeTokenUsage {
1443 cached_input_tokens: u64,
1444 non_cached_input_tokens: u64,
1445 thinking_tokens: u64,
1446 output_tokens: u64,
1447}
1448
1449#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1450struct CumulativeCost {
1451 estimated_cost_usd_nanos: u64,
1452 unpriced_provider_calls: u64,
1453}
1454
1455pub fn estimate_tokens(text: &str) -> u64 {
1456 estimate_bytes(text.len() as u64)
1457}
1458
1459fn estimate_bytes(bytes: u64) -> u64 {
1460 bytes.div_ceil(ESTIMATED_BYTES_PER_TOKEN)
1461}
1462
1463fn current_time() -> String {
1464 Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true)
1465}
1466
1467fn projection_footer(
1468 stale_label: &str,
1469 current_time: &str,
1470 transient_lines: &[String],
1471 estimated_tokens: u64,
1472 context_limit: u64,
1473) -> String {
1474 let mut lines = vec![
1475 format!("[stale boxes: {stale_label}]"),
1476 format!("[current time: {current_time}]"),
1477 ];
1478 lines.extend(transient_lines.iter().cloned());
1479 lines.push(format!(
1480 "[current context size: {estimated_tokens} | max context size: {context_limit}]"
1481 ));
1482 lines.join("\n")
1483}
1484
1485fn projected_context_bytes(items: &[ProjectionItem], footer: &str) -> u64 {
1486 items
1487 .iter()
1488 .fold(footer.len() as u64, |total, item| {
1489 total.saturating_add(item.text.len() as u64)
1490 })
1491 .saturating_add((items.len() as u64).saturating_mul(2))
1492}
1493
1494fn provider_u64(value: &Value, keys: &[&str]) -> u64 {
1495 keys.iter()
1496 .find_map(|key| value.get(*key).and_then(Value::as_u64))
1497 .unwrap_or_default()
1498}
1499
1500struct ActiveToolModel {
1501 invocation_id: Option<String>,
1502 tool_instance: String,
1503 tool_name: String,
1504 model: Option<String>,
1505}
1506
1507fn normalize_requested_model(model: &str) -> String {
1508 model.strip_prefix("codex/").unwrap_or(model).to_owned()
1509}
1510
1511fn provider_metering(provider_data: &Value) -> ProviderMetering {
1512 let metering = provider_data.get("metering");
1513 if metering
1514 .and_then(Value::as_str)
1515 .is_some_and(|kind| kind == "unavailable")
1516 {
1517 return ProviderMetering::Unavailable;
1518 }
1519 if let Some(metering_value) = metering
1520 && let Some(metering) = metering_value.as_object()
1521 {
1522 match metering.get("kind").and_then(Value::as_str) {
1523 Some("duration_seconds") => {
1524 return metering
1525 .get("seconds")
1526 .and_then(Value::as_f64)
1527 .map(|seconds| ProviderMetering::DurationSeconds { seconds })
1528 .unwrap_or(ProviderMetering::Unavailable);
1529 }
1530 Some("unavailable") => return ProviderMetering::Unavailable,
1531 Some("tokens") => {
1532 return ProviderMetering::Tokens(token_metering(metering_value));
1533 }
1534 _ => {}
1535 }
1536 }
1537 let has_tokens = metering.and_then(Value::as_str) == Some("tokens")
1538 || [
1539 "inputTokens",
1540 "input_tokens",
1541 "nonCachedInputTokens",
1542 "non_cached_input_tokens",
1543 "cachedInputTokens",
1544 "cached_input_tokens",
1545 "thinkingTokens",
1546 "thinking_tokens",
1547 "outputTokens",
1548 "output_tokens",
1549 ]
1550 .iter()
1551 .any(|key| provider_data.get(*key).and_then(Value::as_u64).is_some());
1552 if has_tokens {
1553 ProviderMetering::Tokens(token_metering(provider_data))
1554 } else {
1555 ProviderMetering::Unavailable
1556 }
1557}
1558
1559fn token_metering(value: &Value) -> ProviderTokenUsage {
1560 let cached_input_tokens = provider_u64(value, &["cachedInputTokens", "cached_input_tokens"]);
1561 let thinking_tokens = provider_u64(
1562 value,
1563 &[
1564 "thinkingTokens",
1565 "thinking_tokens",
1566 "reasoningOutputTokens",
1567 "reasoning_output_tokens",
1568 ],
1569 );
1570 let normalized_delta = value
1571 .get("usageIsDelta")
1572 .or_else(|| value.get("usage_is_delta"))
1573 .and_then(Value::as_bool)
1574 .unwrap_or(false);
1575 let input_tokens = if normalized_delta {
1576 provider_u64(value, &["nonCachedInputTokens", "non_cached_input_tokens"])
1577 } else {
1578 provider_u64(value, &["inputTokens", "input_tokens"]).saturating_sub(cached_input_tokens)
1579 };
1580 let output_tokens = if normalized_delta {
1581 provider_u64(value, &["outputTokens", "output_tokens"])
1582 } else {
1583 provider_u64(value, &["outputTokens", "output_tokens"]).saturating_sub(thinking_tokens)
1584 };
1585 ProviderTokenUsage {
1586 input_tokens,
1587 cached_input_tokens,
1588 thinking_tokens,
1589 output_tokens,
1590 }
1591}
1592
1593fn context_event_role(kind: &EventKind) -> Role {
1594 match kind {
1595 EventKind::BoxCreated { owner, content, .. } => match owner {
1596 BoxOwner::System | BoxOwner::Controller => {
1597 if content
1598 .metadata
1599 .get("capacityError")
1600 .and_then(Value::as_bool)
1601 .unwrap_or(false)
1602 {
1603 Role::SystemError
1604 } else {
1605 Role::SystemMessage
1606 }
1607 }
1608 BoxOwner::User => Role::UserMessage,
1609 BoxOwner::Kennedy => Role::KennedyMessage,
1610 BoxOwner::Tool { .. } => Role::ToolResult,
1611 },
1612 EventKind::ToolInvoked { .. } => Role::KennedyToolCall,
1613 EventKind::ToolCompleted { outcome, .. } => {
1614 if outcome.get("ok").and_then(Value::as_bool).unwrap_or(true) {
1615 Role::ToolResult
1616 } else {
1617 Role::ToolError
1618 }
1619 }
1620 EventKind::CapacityError { .. } => Role::SystemError,
1621 EventKind::PendingAllocated {
1622 resource: PendingKind::Object,
1623 ..
1624 } => Role::PendingObject,
1625 EventKind::BoxDehydrated { .. }
1626 | EventKind::BoxSummarized { .. }
1627 | EventKind::BoxRehydrated { .. }
1628 | EventKind::BoxRetired { .. } => Role::KennedyToolCall,
1629 _ => Role::SystemMessage,
1630 }
1631}
1632
1633fn encode_context_event(recorded_at: &str, kind: &EventKind) -> anyhow::Result<String> {
1634 let mut kind = serde_json::to_value(kind)?;
1635 let kind_object = kind
1636 .as_object_mut()
1637 .context("Kennedy context event kind must encode as an object")?;
1638 match kind_object.get("type").and_then(Value::as_str) {
1639 Some("box_created") => {
1640 kind_object.remove("box_id");
1641 }
1642 Some("pending_allocated") => {
1643 kind_object.remove("pending_id");
1644 }
1645 _ => {}
1646 }
1647 Ok(serde_json::to_string(&PersistedContextEventWire {
1648 context_event_version: FORMAT_VERSION,
1649 recorded_at: recorded_at.into(),
1650 kind,
1651 })?)
1652}
1653
1654fn decode_context_event(
1655 event: &kcode_session_log::SessionEvent,
1656) -> anyhow::Result<PersistedContextEvent> {
1657 let wire: PersistedContextEventWire = match serde_json::from_str(&event.text) {
1658 Ok(persisted) => persisted,
1659 Err(_) if event.role == Role::PendingObject => {
1660 return Ok(PersistedContextEvent {
1661 context_event_version: FORMAT_VERSION,
1662 recorded_at: String::new(),
1663 kind: EventKind::PendingAllocated {
1664 pending_id: PendingId::from_event(EventId(1)),
1665 resource: PendingKind::Object,
1666 },
1667 });
1668 }
1669 Err(error) => {
1670 return Err(error).context("decoding Kennedy context event from session log");
1671 }
1672 };
1673 ensure!(
1674 wire.context_event_version == FORMAT_VERSION,
1675 "unsupported Kennedy context event version {}",
1676 wire.context_event_version
1677 );
1678 let mut kind = wire.kind;
1679 let kind_object = kind
1680 .as_object_mut()
1681 .context("Kennedy context event kind must be an object")?;
1682 match kind_object.get("type").and_then(Value::as_str) {
1683 Some("box_created") if !kind_object.contains_key("box_id") => {
1684 kind_object.insert("box_id".into(), Value::from(0));
1685 }
1686 Some("pending_allocated") if !kind_object.contains_key("pending_id") => {
1687 kind_object.insert("pending_id".into(), Value::String("pending:1".into()));
1688 }
1689 _ => {}
1690 }
1691 Ok(PersistedContextEvent {
1692 context_event_version: wire.context_event_version,
1693 recorded_at: wire.recorded_at,
1694 kind: serde_json::from_value(kind).context("decoding Kennedy context event kind")?,
1695 })
1696}
1697
1698fn normalize_derived_identity(kind: &mut EventKind, id: EventId) -> anyhow::Result<()> {
1699 match kind {
1700 EventKind::BoxCreated { box_id, .. } => *box_id = BoxId(id.0),
1701 EventKind::PendingAllocated { pending_id, .. } => {
1702 *pending_id = PendingId::from_event(id);
1703 }
1704 _ => {}
1705 }
1706 Ok(())
1707}
1708
1709pub struct Session {
1710 durable: DurableSession,
1711 chatend: Chatend,
1712 objects: BTreeMap<PendingId, ObjectLocation>,
1713}
1714
1715struct UnfinishedToolInvocation {
1716 invocation_id: Option<String>,
1717 tool_instance: String,
1718 tool_name: String,
1719}
1720
1721impl Session {
1722 pub(crate) fn create(
1723 path: impl AsRef<Path>,
1724 metadata: SessionMetadata,
1725 ) -> anyhow::Result<Self> {
1726 ensure!(
1727 metadata.effective_context_tokens > 0,
1728 "effective context window must be positive"
1729 );
1730 let requested = path.as_ref();
1731 let directory = requested
1732 .parent()
1733 .filter(|path| !path.as_os_str().is_empty())
1734 .unwrap_or_else(|| Path::new("."));
1735 let durable = DurableSessionStore::new(directory)
1736 .create_session(&metadata.session_id, &metadata.created_at)?;
1737 Ok(Self {
1738 durable,
1739 chatend: Chatend::opened(metadata),
1740 objects: BTreeMap::new(),
1741 })
1742 }
1743
1744 pub(crate) fn open_with_metadata(
1745 path: impl AsRef<Path>,
1746 metadata: SessionMetadata,
1747 ) -> anyhow::Result<Self> {
1748 Self::open_with_metadata_and_provider_costs(path, metadata, None, None)
1749 }
1750
1751 pub(crate) fn open_with_metadata_and_provider_costs(
1752 path: impl AsRef<Path>,
1753 metadata: SessionMetadata,
1754 default_provider_model: Option<&str>,
1755 estimator: Option<ProviderCostEstimator>,
1756 ) -> anyhow::Result<Self> {
1757 let requested = path.as_ref();
1758 ensure!(
1759 requested.extension().and_then(|value| value.to_str()) == Some("session-log"),
1760 "{} is not a session-log path",
1761 requested.display()
1762 );
1763 let directory = requested
1764 .parent()
1765 .filter(|path| !path.as_os_str().is_empty())
1766 .unwrap_or_else(|| Path::new("."));
1767 let session_id = requested
1768 .file_stem()
1769 .and_then(|value| value.to_str())
1770 .context("session-log filename is not valid UTF-8")?;
1771 let durable = DurableSessionStore::new(directory).open_session(session_id)?;
1772 let log = durable.list();
1773 let chatend =
1774 Chatend::replay_with_provider_costs(metadata, &log, default_provider_model, estimator)?;
1775 let mut objects = BTreeMap::new();
1776 for (position, stored) in log.events.iter().enumerate() {
1777 let persisted = decode_context_event(stored)?;
1778 let id = EventId(position as u64 + 1);
1779 let mut kind = persisted.kind;
1780 normalize_derived_identity(&mut kind, id)?;
1781 if let EventKind::PendingAllocated {
1782 pending_id,
1783 resource: PendingKind::Object,
1784 } = kind
1785 {
1786 let object = durable.read_pending_object(EventPosition(position as u64))?;
1787 let metadata = ObjectMetadata {
1788 pending_id: pending_id.clone(),
1789 event_id: id,
1790 recorded_at: chatend
1791 .event(id)
1792 .map(|event| event.recorded_at.clone())
1793 .unwrap_or_default(),
1794 media_type: object.media_type,
1795 file_name: Some(object.file_name),
1796 transport: Value::Null,
1797 };
1798 objects.insert(
1799 pending_id,
1800 ObjectLocation {
1801 metadata,
1802 payload_offset: position as u64,
1803 payload_len: object.bytes.len() as u64,
1804 },
1805 );
1806 }
1807 }
1808 Ok(Self {
1809 durable,
1810 chatend,
1811 objects,
1812 })
1813 }
1814
1815 pub fn id(&self) -> &str {
1816 &self.chatend.metadata.session_id
1817 }
1818
1819 pub fn state(&self) -> &Chatend {
1820 &self.chatend
1821 }
1822
1823 pub fn objects(&self) -> &BTreeMap<PendingId, ObjectLocation> {
1824 &self.objects
1825 }
1826
1827 #[cfg(test)]
1828 fn session_log(&self) -> kcode_session_log::SessionLog {
1829 self.durable.list()
1830 }
1831
1832 pub fn archive_bytes(&self) -> anyhow::Result<Vec<u8>> {
1833 let mut archive =
1834 serde_json::to_value(self.durable.list()).context("serializing the session log")?;
1835 let object = archive
1836 .as_object_mut()
1837 .context("serialized session log is not an object")?;
1838 object.insert(
1839 "metadata".into(),
1840 serde_json::to_value(&self.chatend.metadata)?,
1841 );
1842 object.insert("boxes".into(), serde_json::to_value(&self.chatend.boxes)?);
1843 let projection = self.chatend.projection();
1844 object.insert("chatendText".into(), Value::String(projection.render()));
1845 object.insert("context".into(), serde_json::to_value(projection)?);
1846 serde_json::to_vec(&archive).context("serializing the session archive")
1847 }
1848
1849 pub fn is_sealed(&self) -> bool {
1850 self.durable.is_sealed()
1851 }
1852
1853 pub fn seal(&mut self) -> anyhow::Result<()> {
1854 let unfinished_tools = self.unfinished_tool_invocations()?;
1855 ensure!(
1856 unfinished_tools.is_empty(),
1857 "session ends with unfinished tools {}",
1858 unfinished_tools
1859 .iter()
1860 .map(|tool| tool.tool_name.as_str())
1861 .collect::<Vec<_>>()
1862 .join(", ")
1863 );
1864 self.durable.seal()?;
1865 Ok(())
1866 }
1867
1868 pub fn repair_unfinished_tools(
1869 &mut self,
1870 recorded_at: impl Into<String>,
1871 ) -> anyhow::Result<Vec<EventId>> {
1872 let unfinished = self.unfinished_tool_invocations()?;
1873 if unfinished.is_empty() {
1874 return Ok(Vec::new());
1875 }
1876 let recorded_at = recorded_at.into();
1877 let mut repaired = Vec::with_capacity(unfinished.len());
1878 for tool in unfinished.iter().rev() {
1879 let message = format!(
1880 "{} was interrupted before a durable completion was recorded; the abandoned invocation was closed during session recovery.",
1881 tool.tool_name
1882 );
1883 let kind = if let Some(invocation_id) = &tool.invocation_id {
1884 EventKind::ToolCompleted {
1885 tool_instance: tool.tool_instance.clone(),
1886 tool_name: tool.tool_name.clone(),
1887 outcome: serde_json::json!({"ok":false,"recovered":true,"result":message}),
1888 invocation_id: Some(invocation_id.clone()),
1889 }
1890 } else {
1891 EventKind::ToolCompleted {
1894 tool_instance: "call_ktool".into(),
1895 tool_name: "call_ktool".into(),
1896 outcome: serde_json::json!({"ok":false,"recovered":true,"result":message}),
1897 invocation_id: None,
1898 }
1899 };
1900 repaired.push(self.record(recorded_at.clone(), kind)?);
1901 }
1902 ensure!(
1903 self.unfinished_tool_invocations()?.is_empty(),
1904 "session tool recovery left unfinished invocations"
1905 );
1906 Ok(repaired)
1907 }
1908
1909 fn unfinished_tool_invocations(&self) -> anyhow::Result<Vec<UnfinishedToolInvocation>> {
1910 let mut identified = BTreeMap::<String, UnfinishedToolInvocation>::new();
1911 let mut legacy = Vec::<UnfinishedToolInvocation>::new();
1912 for event in &self.chatend.events {
1913 match &event.kind {
1914 EventKind::ToolInvoked {
1915 tool_instance,
1916 tool_name,
1917 invocation_id,
1918 ..
1919 } => {
1920 let pending = UnfinishedToolInvocation {
1921 invocation_id: invocation_id.clone(),
1922 tool_instance: tool_instance.clone(),
1923 tool_name: tool_name.clone(),
1924 };
1925 if let Some(invocation_id) = invocation_id {
1926 ensure!(
1927 identified.insert(invocation_id.clone(), pending).is_none(),
1928 "duplicate tool invocation identity {invocation_id}"
1929 );
1930 } else {
1931 legacy.push(pending);
1932 }
1933 }
1934 EventKind::ToolCompleted {
1935 tool_instance,
1936 tool_name,
1937 invocation_id,
1938 ..
1939 } => {
1940 if let Some(invocation_id) = invocation_id {
1941 let invoked = identified.remove(invocation_id).with_context(|| {
1942 format!(
1943 "tool {tool_name} completed without matching invocation {invocation_id}"
1944 )
1945 })?;
1946 ensure!(
1947 invoked.tool_instance.as_str() == tool_instance
1948 && invoked.tool_name.as_str() == tool_name,
1949 "tool completion {invocation_id} does not match its invocation"
1950 );
1951 } else if tool_name == "call_ktool" {
1952 legacy.pop();
1955 } else {
1956 let before = legacy.len();
1957 legacy.retain(|unfinished| unfinished.tool_name.as_str() != tool_name);
1958 ensure!(
1959 legacy.len() != before,
1960 "tool {tool_name} completed without a matching invocation"
1961 );
1962 }
1963 }
1964 _ => {}
1965 }
1966 }
1967 legacy.extend(identified.into_values());
1968 Ok(legacy)
1969 }
1970
1971 pub fn mark_completed(&mut self, session_object_id: String) {
1972 self.chatend.completed_session_object = Some(session_object_id);
1973 }
1974
1975 pub fn configure_context(&mut self, kind: SessionKind, effective_context_tokens: u64) {
1976 self.chatend.metadata.kind = kind;
1977 self.chatend.metadata.effective_context_tokens = effective_context_tokens;
1978 }
1979
1980 pub fn create_box(
1981 &mut self,
1982 recorded_at: impl Into<String>,
1983 name: impl Into<String>,
1984 owner: BoxOwner,
1985 content: BoxContent,
1986 ) -> anyhow::Result<BoxId> {
1987 let recorded_at = recorded_at.into();
1988 let id = EventId(self.chatend.next_id);
1989 let box_id = BoxId(id.0);
1990 self.commit_events(
1991 recorded_at.clone(),
1992 vec![Event {
1993 id,
1994 recorded_at,
1995 kind: EventKind::BoxCreated {
1996 box_id,
1997 name: name.into(),
1998 owner,
1999 content,
2000 },
2001 }],
2002 )?;
2003 Ok(box_id)
2004 }
2005
2006 pub fn update_box(
2007 &mut self,
2008 recorded_at: impl Into<String>,
2009 box_id: BoxId,
2010 content: BoxContent,
2011 ) -> anyhow::Result<Option<EventId>> {
2012 let state = self
2013 .chatend
2014 .boxes
2015 .get(&box_id)
2016 .with_context(|| format!("box {box_id} does not exist"))?;
2017 ensure!(state.active, "box {box_id} is retired");
2018 if state.canonical.content == content {
2019 return Ok(None);
2020 }
2021 let id = EventId(self.chatend.next_id);
2022 let recorded_at = recorded_at.into();
2023 self.commit_events(
2024 recorded_at.clone(),
2025 vec![Event {
2026 id,
2027 recorded_at,
2028 kind: EventKind::CanonicalUpdated { box_id, content },
2029 }],
2030 )?;
2031 Ok(Some(id))
2032 }
2033
2034 pub fn dehydrate_boxes(
2035 &mut self,
2036 recorded_at: impl Into<String>,
2037 box_ids: &[BoxId],
2038 ) -> anyhow::Result<Vec<EventId>> {
2039 ensure!(
2040 !box_ids.is_empty(),
2041 "at least one box must be selected for dehydration"
2042 );
2043 ensure!(
2044 box_ids.iter().copied().collect::<HashSet<_>>().len() == box_ids.len(),
2045 "box dehydration cannot contain duplicate box IDs"
2046 );
2047 let recorded_at = recorded_at.into();
2048 let events = box_ids
2049 .iter()
2050 .enumerate()
2051 .map(|(offset, box_id)| {
2052 let offset = u64::try_from(offset).context("box dehydration batch is too large")?;
2053 let id = self
2054 .chatend
2055 .next_id
2056 .checked_add(offset)
2057 .context("event ID overflow")?;
2058 Ok(Event {
2059 id: EventId(id),
2060 recorded_at: recorded_at.clone(),
2061 kind: EventKind::BoxDehydrated { box_id: *box_id },
2062 })
2063 })
2064 .collect::<anyhow::Result<Vec<_>>>()?;
2065 let ids = events.iter().map(|event| event.id).collect();
2066 self.commit_events(recorded_at, events)?;
2067 Ok(ids)
2068 }
2069
2070 pub fn summarize_box(
2071 &mut self,
2072 recorded_at: impl Into<String>,
2073 box_id: BoxId,
2074 text: impl Into<String>,
2075 ) -> anyhow::Result<EventId> {
2076 self.box_operation(
2077 recorded_at,
2078 EventKind::BoxSummarized {
2079 box_id,
2080 text: text.into(),
2081 },
2082 )
2083 }
2084
2085 pub fn rehydrate_box(
2086 &mut self,
2087 recorded_at: impl Into<String>,
2088 box_id: BoxId,
2089 ) -> anyhow::Result<EventId> {
2090 self.box_operation(recorded_at, EventKind::BoxRehydrated { box_id })
2091 }
2092
2093 pub fn retire_box(
2094 &mut self,
2095 recorded_at: impl Into<String>,
2096 box_id: BoxId,
2097 ) -> anyhow::Result<EventId> {
2098 self.box_operation(recorded_at, EventKind::BoxRetired { box_id })
2099 }
2100
2101 fn box_operation(
2102 &mut self,
2103 recorded_at: impl Into<String>,
2104 kind: EventKind,
2105 ) -> anyhow::Result<EventId> {
2106 let box_id = event_box_id(&kind).context("box operation has no box identity")?;
2107 let state = self
2108 .chatend
2109 .box_state(box_id)
2110 .with_context(|| format!("box {box_id} does not exist"))?;
2111 ensure!(state.active, "box {box_id} is retired");
2112 if let EventKind::BoxSummarized { text, .. } = &kind {
2113 ensure!(!text.trim().is_empty(), "a box summary cannot be empty");
2114 }
2115 if matches!(kind, EventKind::BoxRetired { .. })
2116 && let BoxOwner::Tool { tool_instance, .. } = &state.owner
2117 {
2118 ensure!(
2119 self.chatend
2120 .tools
2121 .get(tool_instance)
2122 .is_some_and(|tool| tool.slots.iter().any(|slot| slot.box_id == box_id)),
2123 "tool box {box_id} is missing from {tool_instance}"
2124 );
2125 }
2126 let id = EventId(self.chatend.next_id);
2127 let recorded_at = recorded_at.into();
2128 self.commit_events(
2129 recorded_at.clone(),
2130 vec![Event {
2131 id,
2132 recorded_at,
2133 kind,
2134 }],
2135 )?;
2136 Ok(id)
2137 }
2138
2139 pub fn allocate_pending_node(
2140 &mut self,
2141 recorded_at: impl Into<String>,
2142 ) -> anyhow::Result<PendingId> {
2143 let id = EventId(self.chatend.next_id);
2144 let pending_id = PendingId::from_event(id);
2145 let recorded_at = recorded_at.into();
2146 self.commit_events(
2147 recorded_at.clone(),
2148 vec![Event {
2149 id,
2150 recorded_at,
2151 kind: EventKind::PendingAllocated {
2152 pending_id: pending_id.clone(),
2153 resource: PendingKind::Node,
2154 },
2155 }],
2156 )?;
2157 Ok(pending_id)
2158 }
2159
2160 pub fn stage_object(
2161 &mut self,
2162 recorded_at: impl Into<String>,
2163 media_type: impl Into<String>,
2164 file_name: Option<String>,
2165 transport: Value,
2166 bytes: &[u8],
2167 ) -> anyhow::Result<PendingId> {
2168 ensure!(
2169 bytes.len() as u64 <= MAX_OBJECT_BYTES,
2170 "object exceeds the 32 GiB V1 limit"
2171 );
2172 let aggregate = self
2173 .objects
2174 .values()
2175 .try_fold(bytes.len() as u64, |total, object| {
2176 total.checked_add(object.payload_len)
2177 })
2178 .context("staged object aggregate length overflow")?;
2179 ensure!(
2180 aggregate <= MAX_OBJECT_BYTES,
2181 "session object payload total exceeds the 32 GiB V1 limit"
2182 );
2183 let event_id = EventId(self.chatend.next_id);
2184 let pending_id = PendingId::from_event(event_id);
2185 let metadata = ObjectMetadata {
2186 pending_id: pending_id.clone(),
2187 event_id,
2188 recorded_at: recorded_at.into(),
2189 media_type: media_type.into(),
2190 file_name,
2191 transport,
2192 };
2193 let kind = EventKind::PendingAllocated {
2194 pending_id: pending_id.clone(),
2195 resource: PendingKind::Object,
2196 };
2197 let text = encode_context_event(&metadata.recorded_at, &kind)?;
2198 let object_file_name = metadata
2199 .file_name
2200 .clone()
2201 .unwrap_or_else(|| format!("object-{}", event_id.0));
2202 let position = self.durable.add_pending_object(
2203 text,
2204 object_file_name,
2205 metadata.media_type.clone(),
2206 bytes,
2207 )?;
2208 ensure!(
2209 position.0 + 1 == event_id.0,
2210 "session-log event position diverged from Kennedy context identity"
2211 );
2212 let allocation = Event {
2213 id: event_id,
2214 recorded_at: metadata.recorded_at.clone(),
2215 kind,
2216 };
2217 self.chatend.apply_transition(&Transition {
2218 recorded_at: metadata.recorded_at.clone(),
2219 events: vec![allocation],
2220 })?;
2221 self.objects.insert(
2222 pending_id.clone(),
2223 ObjectLocation {
2224 metadata,
2225 payload_offset: position.0,
2226 payload_len: bytes.len() as u64,
2227 },
2228 );
2229 Ok(pending_id)
2230 }
2231
2232 pub fn read_object(&mut self, id: &PendingId) -> anyhow::Result<Vec<u8>> {
2233 let location = self
2234 .objects
2235 .get(id)
2236 .with_context(|| format!("staged object {id} does not exist"))?
2237 .clone();
2238 Ok(self
2239 .durable
2240 .read_pending_object(EventPosition(location.payload_offset))?
2241 .bytes)
2242 }
2243
2244 pub fn record(
2245 &mut self,
2246 recorded_at: impl Into<String>,
2247 kind: EventKind,
2248 ) -> anyhow::Result<EventId> {
2249 let id = EventId(self.chatend.next_id);
2250 let recorded_at = recorded_at.into();
2251 self.commit_events(
2252 recorded_at.clone(),
2253 vec![Event {
2254 id,
2255 recorded_at,
2256 kind,
2257 }],
2258 )?;
2259 Ok(id)
2260 }
2261
2262 pub fn commit_events(
2263 &mut self,
2264 recorded_at: impl Into<String>,
2265 events: Vec<Event>,
2266 ) -> anyhow::Result<()> {
2267 let transition = Transition {
2268 recorded_at: recorded_at.into(),
2269 events,
2270 };
2271 ensure!(
2272 !transition.events.is_empty(),
2273 "a transition cannot be empty"
2274 );
2275 ensure!(
2276 !transition.events.iter().any(|event| {
2277 matches!(
2278 event.kind,
2279 EventKind::PendingAllocated {
2280 resource: PendingKind::Object,
2281 ..
2282 }
2283 )
2284 }),
2285 "pending objects must be added through stage_object"
2286 );
2287 let mut preview = self.chatend.clone();
2288 preview.apply_transition(&transition)?;
2289 for event in &transition.events {
2290 let expected = self.durable.list().events.len() as u64 + 1;
2291 ensure!(
2292 event.id.0 == expected,
2293 "Kennedy context event {} does not match session-log position {}",
2294 event.id,
2295 expected - 1
2296 );
2297 self.durable.add_event(
2298 context_event_role(&event.kind),
2299 encode_context_event(&event.recorded_at, &event.kind)?,
2300 )?;
2301 }
2302 self.chatend = preview;
2303 Ok(())
2304 }
2305
2306 pub fn apply_tool_slots(
2307 &mut self,
2308 recorded_at: impl Into<String>,
2309 tool_instance: impl Into<String>,
2310 slots: Vec<ToolSlotInput>,
2311 ) -> anyhow::Result<Vec<EventId>> {
2312 self.apply_tool_slots_inner(recorded_at, tool_instance, slots, None)
2313 }
2314
2315 pub fn apply_tool_slots_with_layout(
2316 &mut self,
2317 recorded_at: impl Into<String>,
2318 tool_instance: impl Into<String>,
2319 slots: Vec<ToolSlotInput>,
2320 layout_slots: &[String],
2321 ) -> anyhow::Result<Vec<EventId>> {
2322 self.apply_tool_slots_inner(recorded_at, tool_instance, slots, Some(layout_slots))
2323 }
2324
2325 fn apply_tool_slots_inner(
2326 &mut self,
2327 recorded_at: impl Into<String>,
2328 tool_instance: impl Into<String>,
2329 slots: Vec<ToolSlotInput>,
2330 layout_slots: Option<&[String]>,
2331 ) -> anyhow::Result<Vec<EventId>> {
2332 let recorded_at = recorded_at.into();
2333 let tool_instance = tool_instance.into();
2334 let current = self
2335 .chatend
2336 .tools
2337 .get(&tool_instance)
2338 .cloned()
2339 .unwrap_or_default();
2340 ensure!(
2341 slots.len() >= current.slots.len(),
2342 "stateful tool slot sequence was truncated"
2343 );
2344 for (index, existing) in current.slots.iter().enumerate() {
2345 ensure!(
2346 slots[index].slot == existing.slot,
2347 "stateful tool slot sequence was reordered at index {index}"
2348 );
2349 ensure!(
2350 !existing.retired || slots[index].retired,
2351 "retired tool slot {} cannot be reactivated",
2352 existing.slot
2353 );
2354 }
2355 let mut events = Vec::new();
2356 let mut next = self.chatend.next_id;
2357 let mut next_state = current.clone();
2358 for (index, input) in slots.iter().enumerate() {
2359 if let Some(existing) = current.slots.get(index) {
2360 let state = self
2361 .chatend
2362 .boxes
2363 .get(&existing.box_id)
2364 .context("tool slot references a missing box")?;
2365 if input.retired && !existing.retired {
2366 let id = EventId(next);
2367 next += 1;
2368 events.push(Event {
2369 id,
2370 recorded_at: recorded_at.clone(),
2371 kind: EventKind::BoxRetired {
2372 box_id: existing.box_id,
2373 },
2374 });
2375 next_state.slots[index].retired = true;
2376 } else if !input.retired {
2377 if state.name != input.name {
2378 let id = EventId(next);
2379 next += 1;
2380 events.push(Event {
2381 id,
2382 recorded_at: recorded_at.clone(),
2383 kind: EventKind::BoxRenamed {
2384 box_id: existing.box_id,
2385 name: input.name.clone(),
2386 },
2387 });
2388 }
2389 if state.canonical.content != input.content {
2390 let id = EventId(next);
2391 next += 1;
2392 events.push(Event {
2393 id,
2394 recorded_at: recorded_at.clone(),
2395 kind: EventKind::CanonicalUpdated {
2396 box_id: existing.box_id,
2397 content: input.content.clone(),
2398 },
2399 });
2400 }
2401 }
2402 } else {
2403 ensure!(
2404 !input.retired,
2405 "a newly appended tool slot cannot start retired"
2406 );
2407 let id = EventId(next);
2408 next += 1;
2409 let box_id = BoxId(id.0);
2410 events.push(Event {
2411 id,
2412 recorded_at: recorded_at.clone(),
2413 kind: EventKind::BoxCreated {
2414 box_id,
2415 name: input.name.clone(),
2416 owner: BoxOwner::Tool {
2417 tool_instance: tool_instance.clone(),
2418 slot: input.slot.clone(),
2419 },
2420 content: input.content.clone(),
2421 },
2422 });
2423 next_state.slots.push(ToolSlot {
2424 slot: input.slot.clone(),
2425 box_id,
2426 retired: false,
2427 });
2428 }
2429 }
2430 if let Some(layout_slots) = layout_slots {
2431 let mut unique = std::collections::HashSet::new();
2432 let box_ids = layout_slots
2433 .iter()
2434 .map(|slot_name| {
2435 ensure!(
2436 unique.insert(slot_name),
2437 "tool layout contains duplicate slot {slot_name}"
2438 );
2439 let slot = next_state
2440 .slots
2441 .iter()
2442 .find(|slot| &slot.slot == slot_name)
2443 .with_context(|| {
2444 format!("tool layout references missing slot {slot_name}")
2445 })?;
2446 ensure!(
2447 !slot.retired,
2448 "tool layout references retired slot {slot_name}"
2449 );
2450 Ok(slot.box_id)
2451 })
2452 .collect::<anyhow::Result<Vec<_>>>()?;
2453 if self.chatend.tool_layouts.get(&tool_instance) != Some(&box_ids) {
2454 let id = EventId(next);
2455 events.push(Event {
2456 id,
2457 recorded_at: recorded_at.clone(),
2458 kind: EventKind::ToolLayoutChanged {
2459 tool_instance: tool_instance.clone(),
2460 box_ids,
2461 },
2462 });
2463 }
2464 }
2465 if events.is_empty() {
2466 return Ok(Vec::new());
2467 }
2468 let ids = events.iter().map(|event| event.id).collect::<Vec<_>>();
2469 self.commit_events(recorded_at, events)?;
2470 Ok(ids)
2471 }
2472
2473 pub fn apply_box_representations(
2474 &mut self,
2475 recorded_at: impl Into<String>,
2476 desired: &BTreeMap<BoxId, BoxRepresentation>,
2477 ) -> anyhow::Result<Vec<EventId>> {
2478 let recorded_at = recorded_at.into();
2479 let events = self
2480 .chatend
2481 .box_representation_events(&recorded_at, desired)?;
2482 if events.is_empty() {
2483 return Ok(Vec::new());
2484 }
2485 let ids = events.iter().map(|event| event.id).collect::<Vec<_>>();
2486 self.commit_events(recorded_at, events)?;
2487 Ok(ids)
2488 }
2489}
2490
2491pub struct SessionHistoryIntegration;
2498
2499impl SessionHistoryIntegration {
2500 pub fn create_session(
2502 path: impl AsRef<Path>,
2503 metadata: SessionMetadata,
2504 ) -> anyhow::Result<Session> {
2505 Session::create(path, metadata)
2506 }
2507
2508 pub fn open_session(
2514 path: impl AsRef<Path>,
2515 metadata: SessionMetadata,
2516 default_provider_model: Option<&str>,
2517 estimator: Option<ProviderCostEstimator>,
2518 ) -> anyhow::Result<Session> {
2519 match (default_provider_model, estimator) {
2520 (None, None) => Session::open_with_metadata(path, metadata),
2521 (default_provider_model, estimator) => Session::open_with_metadata_and_provider_costs(
2522 path,
2523 metadata,
2524 default_provider_model,
2525 estimator,
2526 ),
2527 }
2528 }
2529
2530 pub fn replay(
2535 metadata: SessionMetadata,
2536 log: &kcode_session_log::SessionLog,
2537 default_provider_model: Option<&str>,
2538 estimator: Option<ProviderCostEstimator>,
2539 ) -> anyhow::Result<Chatend> {
2540 match (default_provider_model, estimator) {
2541 (None, None) => Chatend::replay(metadata, log),
2542 (default_provider_model, estimator) => Chatend::replay_with_provider_costs(
2543 metadata,
2544 log,
2545 default_provider_model,
2546 estimator,
2547 ),
2548 }
2549 }
2550
2551 pub fn legacy_provider_cost_summary_for_archive(
2555 archive: &Value,
2556 default_provider_model: Option<&str>,
2557 estimator: ProviderCostEstimator,
2558 ) -> anyhow::Result<ProviderCostSummary> {
2559 legacy_provider_cost_summary_for_archive(archive, default_provider_model, estimator)
2560 }
2561}
2562
2563#[cfg(test)]
2564type SessionJournal = Session;
2565
2566#[cfg(test)]
2567mod tests {
2568 use std::path::PathBuf;
2569 use std::time::{SystemTime, UNIX_EPOCH};
2570
2571 use serde_json::json;
2572
2573 use super::*;
2574
2575 fn path(label: &str) -> PathBuf {
2576 std::env::temp_dir()
2577 .join(format!(
2578 "kennedy-chatend-{label}-{}-{}",
2579 std::process::id(),
2580 SystemTime::now()
2581 .duration_since(UNIX_EPOCH)
2582 .unwrap()
2583 .as_nanos()
2584 ))
2585 .join("session-1.session-log")
2586 }
2587
2588 fn metadata() -> SessionMetadata {
2589 SessionMetadata {
2590 session_id: "session-1".into(),
2591 kind: SessionKind::Conversation,
2592 created_at: "2026-07-23T00:00:00Z".into(),
2593 effective_context_tokens: 1_000,
2594 channel: json!({"kind":"test"}),
2595 }
2596 }
2597
2598 #[test]
2599 fn box_identity_continuations_staleness_and_replay_are_exact() {
2600 let path = path("boxes");
2601 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2602 let id = journal
2603 .create_box(
2604 "t1",
2605 "message",
2606 BoxOwner::User,
2607 BoxContent::text("original"),
2608 )
2609 .unwrap();
2610 assert_eq!(id, BoxId(1));
2611 journal.summarize_box("t2", id, "summary").unwrap();
2612 journal
2613 .update_box("t3", id, BoxContent::text("changed"))
2614 .unwrap();
2615 let state = journal.state().box_state(id).unwrap().clone();
2616 assert!(state.stale());
2617 assert_eq!(
2618 state.representation,
2619 Representation::Summarized {
2620 based_on: EventId(1),
2621 text: "summary".into()
2622 }
2623 );
2624 let projection = journal.state().projection();
2625 assert_eq!(projection.items[0].text, "[box updated]");
2626 assert_eq!(projection.items[1].text, "[box updated]");
2627 assert!(projection.items[2].text.contains("summary"));
2628 assert!(projection.items[2].stale);
2629 assert!(projection.footer.starts_with("[stale boxes: 1]\n"));
2630 assert!(projection.footer.contains("\n[current time: "));
2631 assert_eq!(
2632 projection.footer.lines().last().unwrap(),
2633 format!(
2634 "[current context size: {} | max context size: {}]",
2635 projection.estimated_tokens, projection.status.context_limit_tokens
2636 )
2637 );
2638 assert!(!projection.footer.contains("effective"));
2639 assert!(!projection.footer.contains("turn_limit"));
2640 assert!(projection.render().ends_with(&projection.footer));
2641
2642 let transient = vec![
2643 "[provider calls remaining after this call: 249 | max provider calls: 250]".into(),
2644 "[provider call time remaining: 3h 45m]".into(),
2645 ];
2646 let with_limits = journal.state().projection_with_footer_lines(&transient);
2647 assert_eq!(
2648 with_limits.footer.lines().nth(2),
2649 Some(transient[0].as_str())
2650 );
2651 assert_eq!(
2652 with_limits.footer.lines().nth(3),
2653 Some(transient[1].as_str())
2654 );
2655 assert_eq!(
2656 with_limits.footer.lines().last().unwrap(),
2657 format!(
2658 "[current context size: {} | max context size: {}]",
2659 with_limits.estimated_tokens, with_limits.status.context_limit_tokens
2660 )
2661 );
2662 assert!(with_limits.context_bytes > projection.context_bytes);
2663 drop(journal);
2664 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2665 assert_eq!(reopened.state().box_state(id), Some(&state));
2666 std::fs::remove_file(path).unwrap();
2667 }
2668
2669 #[test]
2670 fn box_headers_hide_internal_ownership_and_timestamp_messages() {
2671 let path = path("box-headers");
2672 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2673 journal
2674 .create_box(
2675 "2026-07-23T00:00:01Z",
2676 "User message",
2677 BoxOwner::User,
2678 BoxContent::text("Hello"),
2679 )
2680 .unwrap();
2681 journal
2682 .create_box(
2683 "2026-07-23T00:00:02Z",
2684 "Kennedy message",
2685 BoxOwner::Kennedy,
2686 BoxContent::text("Hi"),
2687 )
2688 .unwrap();
2689 let mut content = BoxContent::text("Node ID: AAAAAAAB");
2690 content.use_concise_header();
2691 journal
2692 .create_box(
2693 "2026-07-23T00:00:03Z",
2694 "Kweb loaded node",
2695 BoxOwner::Tool {
2696 tool_instance: "kweb".into(),
2697 slot: "loaded".into(),
2698 },
2699 content,
2700 )
2701 .unwrap();
2702
2703 let projection = journal.state().projection();
2704 assert_eq!(
2705 projection.items[0].text,
2706 "[box 1 | User message | timestamp=2026-07-23T00:00:01Z | hydrated]\nHello"
2707 );
2708 assert_eq!(
2709 projection.items[1].text,
2710 "[box 2 | Kennedy message | timestamp=2026-07-23T00:00:02Z | hydrated]\nHi"
2711 );
2712 assert_eq!(
2713 projection.items[2].text,
2714 "[box 3 | Kweb loaded node | hydrated]\nNode ID: AAAAAAAB"
2715 );
2716 assert!(
2717 projection
2718 .items
2719 .iter()
2720 .all(|item| !item.text.contains("owner="))
2721 );
2722 let current_time = projection
2723 .footer
2724 .lines()
2725 .find_map(|line| line.strip_prefix("[current time: "))
2726 .and_then(|line| line.strip_suffix(']'))
2727 .unwrap();
2728 chrono::DateTime::parse_from_rfc3339(current_time).unwrap();
2729 assert!(current_time.starts_with("20"));
2730 std::fs::remove_file(path).unwrap();
2731 }
2732
2733 #[test]
2734 fn shared_pending_and_box_identity_space_never_overlaps() {
2735 let path = path("pending");
2736 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2737 let first = journal.allocate_pending_node("t1").unwrap();
2738 let box_id = journal
2739 .create_box("t2", "box", BoxOwner::Kennedy, BoxContent::text("hello"))
2740 .unwrap();
2741 let object = journal
2742 .stage_object(
2743 "t3",
2744 "application/octet-stream",
2745 None,
2746 Value::Null,
2747 b"\0binary\xff",
2748 )
2749 .unwrap();
2750 assert_eq!(first.to_string(), "pending:1");
2751 assert_eq!(box_id, BoxId(2));
2752 assert_eq!(object.to_string(), "pending:3");
2753 assert_eq!(journal.read_object(&object).unwrap(), b"\0binary\xff");
2754 let stored = journal.session_log();
2755 assert!(!stored.events[0].text.contains("pending_id"));
2756 assert!(!stored.events[1].text.contains("box_id"));
2757 assert!(!stored.events[2].text.contains("pending_id"));
2758 drop(journal);
2759 let mut reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2760 assert_eq!(reopened.read_object(&object).unwrap(), b"\0binary\xff");
2761 assert_eq!(reopened.state().next_id, 4);
2762 std::fs::remove_file(path).unwrap();
2763 }
2764
2765 #[test]
2766 fn kennedy_tool_completion_is_validated_before_storage_seals() {
2767 let path = path("seal-tool");
2768 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2769 journal
2770 .record(
2771 "t1",
2772 EventKind::ToolInvoked {
2773 tool_instance: "CreateNode:1".into(),
2774 tool_name: "CreateNode".into(),
2775 arguments: json!({}),
2776 invocation_id: None,
2777 },
2778 )
2779 .unwrap();
2780 assert!(journal.seal().is_err());
2781 journal
2782 .record(
2783 "t2",
2784 EventKind::ToolCompleted {
2785 tool_instance: "call_ktool".into(),
2786 tool_name: "call_ktool".into(),
2787 outcome: json!({"ok":true}),
2788 invocation_id: None,
2789 },
2790 )
2791 .unwrap();
2792 journal.seal().unwrap();
2793 assert!(journal.is_sealed());
2794 std::fs::remove_file(path).unwrap();
2795 }
2796
2797 #[test]
2798 fn interrupted_tools_are_recovered_by_identity_without_losing_legacy_journals() {
2799 let path = path("repair-tools");
2800 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2801 journal
2802 .record(
2803 "t1",
2804 EventKind::ToolInvoked {
2805 tool_instance: "WebSearch:legacy".into(),
2806 tool_name: "WebSearch".into(),
2807 arguments: json!({"question":"legacy"}),
2808 invocation_id: None,
2809 },
2810 )
2811 .unwrap();
2812 journal
2813 .record(
2814 "t2",
2815 EventKind::ToolInvoked {
2816 tool_instance: "WebSearch:new".into(),
2817 tool_name: "WebSearch".into(),
2818 arguments: json!({"question":"identified"}),
2819 invocation_id: Some("call-1".into()),
2820 },
2821 )
2822 .unwrap();
2823 assert!(journal.seal().is_err());
2824
2825 let repaired = journal.repair_unfinished_tools("recovery").unwrap();
2826 assert_eq!(repaired.len(), 2);
2827 assert!(
2828 journal
2829 .state()
2830 .events
2831 .iter()
2832 .rev()
2833 .take(2)
2834 .all(|event| matches!(
2835 &event.kind,
2836 EventKind::ToolCompleted { outcome, .. }
2837 if outcome.get("recovered").and_then(Value::as_bool) == Some(true)
2838 ))
2839 );
2840 journal.seal().unwrap();
2841 assert!(journal.is_sealed());
2842 std::fs::remove_file(path).unwrap();
2843 }
2844
2845 #[test]
2846 fn identified_tool_completions_can_arrive_out_of_order() {
2847 let path = path("tool-identity");
2848 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2849 for id in ["call-1", "call-2"] {
2850 journal
2851 .record(
2852 id,
2853 EventKind::ToolInvoked {
2854 tool_instance: format!("WebSearch:{id}"),
2855 tool_name: "WebSearch".into(),
2856 arguments: json!({"question":id}),
2857 invocation_id: Some(id.into()),
2858 },
2859 )
2860 .unwrap();
2861 }
2862 for id in ["call-1", "call-2"] {
2863 journal
2864 .record(
2865 format!("{id}-complete"),
2866 EventKind::ToolCompleted {
2867 tool_instance: format!("WebSearch:{id}"),
2868 tool_name: "WebSearch".into(),
2869 outcome: json!({"ok":true}),
2870 invocation_id: Some(id.into()),
2871 },
2872 )
2873 .unwrap();
2874 }
2875 journal.seal().unwrap();
2876 std::fs::remove_file(path).unwrap();
2877 }
2878
2879 #[test]
2880 fn invalid_box_operations_do_not_poison_the_append_only_journal() {
2881 let path = path("invalid-box-operation");
2882 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2883 let box_id = journal
2884 .create_box(
2885 "t1",
2886 "valid",
2887 BoxOwner::Controller,
2888 BoxContent::text("canonical"),
2889 )
2890 .unwrap();
2891 let valid_length = std::fs::metadata(&path).unwrap().len();
2892
2893 let result = journal.dehydrate_boxes("t2", &[box_id, BoxId(97)]);
2894 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2895 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2896 assert!(matches!(
2897 journal.state().box_state(box_id).unwrap().representation,
2898 Representation::Hydrated { .. }
2899 ));
2900
2901 for result in [
2902 journal.summarize_box("t3", BoxId(97), "summary"),
2903 journal.rehydrate_box("t4", BoxId(97)),
2904 journal.retire_box("t5", BoxId(97)),
2905 ] {
2906 assert_eq!(result.unwrap_err().to_string(), "box 97 does not exist");
2907 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2908 }
2909
2910 drop(journal);
2911 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2912 assert_eq!(
2913 reopened
2914 .state()
2915 .box_state(box_id)
2916 .unwrap()
2917 .canonical
2918 .content,
2919 BoxContent::text("canonical")
2920 );
2921 std::fs::remove_file(path).unwrap();
2922 }
2923
2924 #[test]
2925 fn multiple_boxes_dehydrate_in_one_replayable_batch() {
2926 let path = path("dehydrate-boxes");
2927 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2928 let boxes = ["first", "second", "third"]
2929 .into_iter()
2930 .map(|text| {
2931 journal
2932 .create_box("t1", text, BoxOwner::Controller, BoxContent::text(text))
2933 .unwrap()
2934 })
2935 .collect::<Vec<_>>();
2936
2937 let event_ids = journal
2938 .dehydrate_boxes("t2", &[boxes[0], boxes[2]])
2939 .unwrap();
2940 assert_eq!(event_ids, [EventId(4), EventId(5)]);
2941 assert!(matches!(
2942 journal.state().box_state(boxes[0]).unwrap().representation,
2943 Representation::Dehydrated { .. }
2944 ));
2945 assert!(matches!(
2946 journal.state().box_state(boxes[1]).unwrap().representation,
2947 Representation::Hydrated { .. }
2948 ));
2949 assert!(matches!(
2950 journal.state().box_state(boxes[2]).unwrap().representation,
2951 Representation::Dehydrated { .. }
2952 ));
2953
2954 let valid_length = std::fs::metadata(&path).unwrap().len();
2955 assert_eq!(
2956 journal.dehydrate_boxes("t3", &[]).unwrap_err().to_string(),
2957 "at least one box must be selected for dehydration"
2958 );
2959 assert_eq!(
2960 journal
2961 .dehydrate_boxes("t3", &[boxes[1], boxes[1]])
2962 .unwrap_err()
2963 .to_string(),
2964 "box dehydration cannot contain duplicate box IDs"
2965 );
2966 assert_eq!(std::fs::metadata(&path).unwrap().len(), valid_length);
2967
2968 drop(journal);
2969 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
2970 assert!(matches!(
2971 reopened.state().box_state(boxes[0]).unwrap().representation,
2972 Representation::Dehydrated { .. }
2973 ));
2974 assert!(matches!(
2975 reopened.state().box_state(boxes[2]).unwrap().representation,
2976 Representation::Dehydrated { .. }
2977 ));
2978 std::fs::remove_file(path).unwrap();
2979 }
2980
2981 #[test]
2982 fn status_estimates_the_context_with_every_active_box_hydrated() {
2983 let path = path("fully-hydrated-estimate");
2984 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
2985 let box_id = journal
2986 .create_box(
2987 "t1",
2988 "large",
2989 BoxOwner::Controller,
2990 BoxContent::text("canonical material ".repeat(1_000)),
2991 )
2992 .unwrap();
2993 journal.dehydrate_boxes("t2", &[box_id]).unwrap();
2994
2995 let status = journal.state().projection().status;
2996 assert!(
2997 status.fully_hydrated_context_tokens > status.current_context_tokens,
2998 "the hydrated estimate should include the canonical body"
2999 );
3000 std::fs::remove_file(path).unwrap();
3001 }
3002
3003 #[test]
3004 fn provider_measurements_recalibrate_the_matching_manifest_then_track_deltas() {
3005 let path = path("calibration");
3006 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3007 journal
3008 .create_box(
3009 "t1",
3010 "system",
3011 BoxOwner::System,
3012 BoxContent::text("baseline provider content"),
3013 )
3014 .unwrap();
3015 let submitted = journal.state().projection();
3016 journal
3017 .record(
3018 "t2",
3019 EventKind::InferenceSubmitted {
3020 manifest_hash: "manifest-1".into(),
3021 estimated_input_tokens: submitted.estimated_tokens,
3022 raw_estimated_input_tokens: Some(submitted.raw_estimated_tokens),
3023 },
3024 )
3025 .unwrap();
3026 let tool_box = journal
3027 .create_box(
3028 "t3",
3029 "tool result",
3030 BoxOwner::Controller,
3031 BoxContent::text("x".repeat(3_000)),
3032 )
3033 .unwrap();
3034 let measured_context = journal.state().projection();
3035 journal
3036 .record(
3037 "t4",
3038 EventKind::ProviderReceipt {
3039 manifest_hash: "manifest-1".into(),
3040 input_tokens: Some(777),
3041 output_tokens: Some(3),
3042 context_bytes: Some(measured_context.context_bytes),
3043 raw_context_tokens: Some(measured_context.raw_estimated_tokens),
3044 provider_data: Value::Null,
3045 },
3046 )
3047 .unwrap();
3048 assert_eq!(journal.state().projection().estimated_tokens, 777);
3049 journal
3050 .create_box(
3051 "t5",
3052 "new material",
3053 BoxOwner::User,
3054 BoxContent::text("x".repeat(300)),
3055 )
3056 .unwrap();
3057 let expanded = journal.state().projection();
3058 assert_eq!(
3059 expanded.estimated_tokens,
3060 777 + expanded
3061 .context_bytes
3062 .abs_diff(measured_context.context_bytes)
3063 / ESTIMATED_BYTES_PER_TOKEN
3064 );
3065 assert!(expanded.raw_estimated_tokens > submitted.raw_estimated_tokens);
3066 journal.dehydrate_boxes("t6", &[tool_box]).unwrap();
3067 let shrunken = journal.state().projection();
3068 assert_eq!(
3069 shrunken.estimated_tokens,
3070 777_u64.saturating_sub(
3071 measured_context
3072 .context_bytes
3073 .abs_diff(shrunken.context_bytes)
3074 / ESTIMATED_BYTES_PER_TOKEN
3075 )
3076 );
3077 std::fs::remove_file(path).unwrap();
3078 }
3079
3080 #[test]
3081 fn session_status_keeps_provider_usage_categories_exact_and_exclusive() {
3082 let path = path("session-status");
3083 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3084 journal
3085 .create_box("t1", "system", BoxOwner::System, BoxContent::text("é"))
3086 .unwrap();
3087 let before_usage = journal.state().projection();
3088 assert_eq!(
3089 before_usage.raw_estimated_tokens,
3090 before_usage
3091 .context_bytes
3092 .div_ceil(ESTIMATED_BYTES_PER_TOKEN)
3093 );
3094 journal
3095 .record(
3096 "t2",
3097 EventKind::ProviderReceipt {
3098 manifest_hash: "manifest-1".into(),
3099 input_tokens: Some(100),
3100 output_tokens: Some(20),
3101 context_bytes: Some(before_usage.context_bytes),
3102 raw_context_tokens: Some(before_usage.raw_estimated_tokens),
3103 provider_data: json!({
3104 "usageIsDelta":true,
3105 "cachedInputTokens":40,
3106 "nonCachedInputTokens":60,
3107 "thinkingTokens":8,
3108 "outputTokens":12,
3109 "estimatedCostUsdNanos":12345
3110 }),
3111 },
3112 )
3113 .unwrap();
3114 let second_anchor = journal.state().projection();
3115 journal
3116 .record(
3117 "t3",
3118 EventKind::ProviderReceipt {
3119 manifest_hash: "manifest-1".into(),
3120 input_tokens: Some(120),
3121 output_tokens: Some(5),
3122 context_bytes: Some(second_anchor.context_bytes),
3123 raw_context_tokens: Some(second_anchor.raw_estimated_tokens),
3124 provider_data: json!({
3125 "usageIsDelta":true,
3126 "cachedInputTokens":10,
3127 "nonCachedInputTokens":10,
3128 "thinkingTokens":2,
3129 "outputTokens":3,
3130 "estimatedCostUsdNanos":6789
3131 }),
3132 },
3133 )
3134 .unwrap();
3135 let projection = journal.state().projection();
3136 assert_eq!(
3137 projection.status,
3138 SessionStatus {
3139 current_context_tokens: 120,
3140 fully_hydrated_context_tokens: 120,
3141 context_limit_tokens: 700,
3142 current_context_bytes: projection.context_bytes,
3143 cached_input_tokens: 50,
3144 non_cached_input_tokens: 70,
3145 thinking_tokens: 10,
3146 output_tokens: 15,
3147 estimated_cost_usd_nanos: 19_134,
3148 unpriced_provider_calls: 0,
3149 }
3150 );
3151 let archive: Value = serde_json::from_slice(&journal.archive_bytes().unwrap()).unwrap();
3152 assert_eq!(archive["context"]["status"]["cachedInputTokens"], 50);
3153 assert!(
3154 archive["chatendText"]
3155 .as_str()
3156 .unwrap()
3157 .ends_with(archive["context"]["footer"].as_str().unwrap())
3158 );
3159 std::fs::remove_file(path).unwrap();
3160 }
3161
3162 fn compatibility_cost(
3163 model: &str,
3164 metering: &ProviderMetering,
3165 ) -> Option<ProviderCostEstimate> {
3166 let ProviderMetering::Tokens(usage) = metering else {
3167 return None;
3168 };
3169 let base = match model {
3170 "gpt-5.6-sol" => 1_000,
3171 "gemini-3.1-pro-preview" => 2_000,
3172 _ => return None,
3173 };
3174 Some(ProviderCostEstimate {
3175 usd_nanos: base + usage.input_tokens,
3176 accuracy: json!("exact"),
3177 pricing_version: "test-prices".into(),
3178 })
3179 }
3180
3181 #[test]
3182 fn legacy_provider_costs_are_reconstructed_without_rewriting_history() {
3183 let path = path("legacy-provider-costs");
3184 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3185 journal
3186 .record(
3187 "t1",
3188 EventKind::ProviderReceipt {
3189 manifest_hash: "top".into(),
3190 input_tokens: Some(10),
3191 output_tokens: Some(1),
3192 context_bytes: None,
3193 raw_context_tokens: None,
3194 provider_data: json!({
3195 "usageIsDelta":true,
3196 "nonCachedInputTokens":10,
3197 "cachedInputTokens":0,
3198 "thinkingTokens":0,
3199 "outputTokens":1
3200 }),
3201 },
3202 )
3203 .unwrap();
3204 journal
3205 .record(
3206 "t2",
3207 EventKind::ToolInvoked {
3208 tool_instance: "RunSubagent:1".into(),
3209 tool_name: "RunSubagent".into(),
3210 arguments: json!({"model":"codex/gpt-5.6-sol"}),
3211 invocation_id: Some("subagent-1".into()),
3212 },
3213 )
3214 .unwrap();
3215 journal
3216 .record(
3217 "t3",
3218 EventKind::Note {
3219 label: "subagent_started".into(),
3220 value: json!({"providerModel":"gpt-5.6-sol"}),
3221 },
3222 )
3223 .unwrap();
3224 journal
3225 .record(
3226 "t4",
3227 EventKind::ProviderReceipt {
3228 manifest_hash: "subagent".into(),
3229 input_tokens: None,
3230 output_tokens: None,
3231 context_bytes: None,
3232 raw_context_tokens: None,
3233 provider_data: json!({
3234 "source":"subagent",
3235 "usageIsDelta":true,
3236 "nonCachedInputTokens":20,
3237 "cachedInputTokens":0,
3238 "thinkingTokens":0,
3239 "outputTokens":1
3240 }),
3241 },
3242 )
3243 .unwrap();
3244 journal
3245 .record(
3246 "t5",
3247 EventKind::ToolInvoked {
3248 tool_instance: "AnnotateMedia:1".into(),
3249 tool_name: "AnnotateMedia".into(),
3250 arguments: json!({"model":"gemini-3.1-pro-preview"}),
3251 invocation_id: Some("media-1".into()),
3252 },
3253 )
3254 .unwrap();
3255 journal
3256 .record(
3257 "t6",
3258 EventKind::ProviderReceipt {
3259 manifest_hash: "media".into(),
3260 input_tokens: None,
3261 output_tokens: None,
3262 context_bytes: None,
3263 raw_context_tokens: None,
3264 provider_data: json!({
3265 "source":"media_annotation",
3266 "usageIsDelta":true,
3267 "nonCachedInputTokens":30,
3268 "cachedInputTokens":0,
3269 "thinkingTokens":0,
3270 "outputTokens":1
3271 }),
3272 },
3273 )
3274 .unwrap();
3275 journal
3276 .record(
3277 "t7",
3278 EventKind::ToolCompleted {
3279 tool_instance: "AnnotateMedia:1".into(),
3280 tool_name: "AnnotateMedia".into(),
3281 outcome: json!({"ok":true}),
3282 invocation_id: Some("media-1".into()),
3283 },
3284 )
3285 .unwrap();
3286 journal
3287 .record(
3288 "t8",
3289 EventKind::ToolCompleted {
3290 tool_instance: "RunSubagent:1".into(),
3291 tool_name: "RunSubagent".into(),
3292 outcome: json!({"ok":true}),
3293 invocation_id: Some("subagent-1".into()),
3294 },
3295 )
3296 .unwrap();
3297 journal
3298 .record(
3299 "t9",
3300 EventKind::ProviderReceipt {
3301 manifest_hash: "unknown".into(),
3302 input_tokens: None,
3303 output_tokens: None,
3304 context_bytes: None,
3305 raw_context_tokens: None,
3306 provider_data: json!({
3307 "source":"web_search",
3308 "usageIsDelta":true,
3309 "nonCachedInputTokens":40,
3310 "cachedInputTokens":0,
3311 "thinkingTokens":0,
3312 "outputTokens":1
3313 }),
3314 },
3315 )
3316 .unwrap();
3317 drop(journal);
3318
3319 let compatible = SessionJournal::open_with_metadata_and_provider_costs(
3320 &path,
3321 metadata(),
3322 Some("gpt-5.6-sol"),
3323 Some(compatibility_cost),
3324 )
3325 .unwrap();
3326 let status = &compatible.state().projection().status;
3327 assert_eq!(status.estimated_cost_usd_nanos, 4_060);
3328 assert_eq!(status.unpriced_provider_calls, 1);
3329 let inferred_models = compatible
3330 .state()
3331 .events
3332 .iter()
3333 .filter_map(|event| match &event.kind {
3334 EventKind::ProviderReceipt { provider_data, .. } => provider_data
3335 .get("providerModel")
3336 .and_then(Value::as_str)
3337 .map(str::to_owned),
3338 _ => None,
3339 })
3340 .collect::<Vec<_>>();
3341 assert_eq!(
3342 inferred_models,
3343 vec!["gpt-5.6-sol", "gpt-5.6-sol", "gemini-3.1-pro-preview"]
3344 );
3345 drop(compatible);
3346
3347 let unchanged = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3348 assert_eq!(
3349 unchanged
3350 .state()
3351 .projection()
3352 .status
3353 .estimated_cost_usd_nanos,
3354 0
3355 );
3356 assert_eq!(
3357 unchanged
3358 .state()
3359 .projection()
3360 .status
3361 .unpriced_provider_calls,
3362 4
3363 );
3364 let archive: Value = serde_json::from_slice(&unchanged.archive_bytes().unwrap()).unwrap();
3365 let archive_summary = legacy_provider_cost_summary_for_archive(
3366 &archive,
3367 Some("gpt-5.6-sol"),
3368 compatibility_cost,
3369 )
3370 .unwrap();
3371 assert_eq!(archive_summary.estimated_cost_usd_nanos, 4_060);
3372 assert_eq!(archive_summary.unpriced_provider_calls, 1);
3373 assert_eq!(archive["context"]["status"]["estimatedCostUsdNanos"], 0);
3374 assert_eq!(archive["context"]["status"]["unpricedProviderCalls"], 4);
3375 drop(unchanged);
3376 std::fs::remove_file(path).unwrap();
3377 }
3378
3379 #[test]
3380 fn legacy_provider_receipts_without_a_raw_context_anchor_remain_readable() {
3381 let receipt: EventKind = serde_json::from_value(json!({
3382 "type":"provider_receipt",
3383 "manifest_hash":"legacy-manifest",
3384 "input_tokens":123,
3385 "output_tokens":4,
3386 "provider_data":null
3387 }))
3388 .unwrap();
3389 assert!(matches!(
3390 receipt,
3391 EventKind::ProviderReceipt {
3392 context_bytes: None,
3393 raw_context_tokens: None,
3394 ..
3395 }
3396 ));
3397 }
3398
3399 #[test]
3400 fn stateful_tool_slots_are_batched_append_only_and_do_not_see_summaries() {
3401 let path = path("slots");
3402 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3403 journal
3404 .apply_tool_slots(
3405 "t1",
3406 "rust-1",
3407 vec![
3408 ToolSlotInput {
3409 slot: "a.rs".into(),
3410 name: "a.rs".into(),
3411 content: BoxContent::text("a"),
3412 retired: false,
3413 },
3414 ToolSlotInput {
3415 slot: "b.rs".into(),
3416 name: "b.rs".into(),
3417 content: BoxContent::text("b"),
3418 retired: false,
3419 },
3420 ],
3421 )
3422 .unwrap();
3423 let a = journal.state().tools["rust-1"].slots[0].box_id;
3424 journal.summarize_box("t2", a, "Kennedy summary").unwrap();
3425 let before = journal.state().clone();
3426 assert!(
3427 journal
3428 .apply_tool_slots(
3429 "t3",
3430 "rust-1",
3431 vec![ToolSlotInput {
3432 slot: "b.rs".into(),
3433 name: "b.rs".into(),
3434 content: BoxContent::text("b"),
3435 retired: false,
3436 }]
3437 )
3438 .is_err()
3439 );
3440 assert_eq!(journal.state(), &before);
3441 journal
3442 .apply_tool_slots(
3443 "t4",
3444 "rust-1",
3445 vec![
3446 ToolSlotInput {
3447 slot: "a.rs".into(),
3448 name: "a.rs".into(),
3449 content: BoxContent::text("a2"),
3450 retired: false,
3451 },
3452 ToolSlotInput {
3453 slot: "b.rs".into(),
3454 name: "b.rs".into(),
3455 content: BoxContent::text("b"),
3456 retired: true,
3457 },
3458 ToolSlotInput {
3459 slot: "c.rs".into(),
3460 name: "c.rs".into(),
3461 content: BoxContent::text("c"),
3462 retired: false,
3463 },
3464 ],
3465 )
3466 .unwrap();
3467 let a_state = journal.state().box_state(a).unwrap();
3468 assert_eq!(a_state.canonical.content.text, "a2");
3469 assert!(matches!(
3470 a_state.representation,
3471 Representation::Summarized { .. }
3472 ));
3473 drop(journal);
3474 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3475 assert_eq!(reopened.state().tools["rust-1"].slots.len(), 3);
3476 assert!(reopened.state().tools["rust-1"].slots[1].retired);
3477 std::fs::remove_file(path).unwrap();
3478 }
3479
3480 #[test]
3481 fn tool_layout_orders_current_boxes_without_changing_their_identities() {
3482 let path = path("tool-layout");
3483 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3484 journal
3485 .apply_tool_slots_with_layout(
3486 "t1",
3487 "kweb",
3488 vec![
3489 ToolSlotInput {
3490 slot: "active".into(),
3491 name: "Active node".into(),
3492 content: BoxContent::text("ACTIVE NODE"),
3493 retired: false,
3494 },
3495 ToolSlotInput {
3496 slot: "direct".into(),
3497 name: "Direct node".into(),
3498 content: BoxContent::text("DIRECT NODE"),
3499 retired: false,
3500 },
3501 ],
3502 &["direct".into(), "active".into()],
3503 )
3504 .unwrap();
3505 let tool = &journal.state().tools["kweb"];
3506 let active_id = tool.slots[0].box_id;
3507 let direct_id = tool.slots[1].box_id;
3508 let projected_items = journal.state().projection().items;
3509 let rendered = projected_items
3510 .iter()
3511 .map(|item| item.text.as_str())
3512 .collect::<Vec<_>>()
3513 .join("\n\n");
3514 assert!(rendered.find("DIRECT NODE") < rendered.find("ACTIVE NODE"));
3515 assert_eq!(
3516 journal.state().tool_layouts["kweb"],
3517 vec![direct_id, active_id]
3518 );
3519 drop(journal);
3520 let reopened = SessionJournal::open_with_metadata(&path, metadata()).unwrap();
3521 assert_eq!(reopened.state().projection().items, projected_items);
3522 std::fs::remove_file(path).unwrap();
3523 }
3524
3525 #[test]
3526 fn limits_use_exact_floor_percentages() {
3527 let state = Chatend::opened(SessionMetadata {
3528 effective_context_tokens: 101,
3529 ..metadata()
3530 });
3531 assert_eq!(state.live_context_limit(), 70);
3532 assert_eq!(state.forced_ingress_context_limit(), 75);
3533 assert_eq!(state.ingress_initial_context_limit(), 75);
3534 assert_eq!(state.ingress_context_limit(), 101);
3535 assert_eq!(state.active_context_limit(), 70);
3536 let ingress = Chatend::opened(SessionMetadata {
3537 kind: SessionKind::HistoryIngress,
3538 effective_context_tokens: 101,
3539 ..metadata()
3540 });
3541 assert_eq!(ingress.active_context_limit(), 101);
3542 assert_eq!(estimate_tokens("1234"), 1);
3543 }
3544
3545 #[test]
3546 fn new_box_projection_preview_matches_the_committed_projection() {
3547 let path = path("new-box-preview");
3548 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3549 let boxes = vec![
3550 (
3551 "User message".into(),
3552 BoxOwner::User,
3553 BoxContent::text("prospective user text"),
3554 ),
3555 (
3556 "attachment".into(),
3557 BoxOwner::User,
3558 BoxContent::text("prospective attachment text"),
3559 ),
3560 ];
3561 let preview = journal
3562 .state()
3563 .projection_with_new_boxes_at("t1", &boxes)
3564 .unwrap();
3565 for (name, owner, content) in boxes {
3566 journal.create_box("t1", name, owner, content).unwrap();
3567 }
3568 assert_eq!(journal.state().projection(), preview);
3569 std::fs::remove_file(path).unwrap();
3570 }
3571
3572 #[test]
3573 fn box_representation_preview_matches_the_batched_append() {
3574 let path = path("representation-plan");
3575 let mut journal = SessionJournal::create(&path, metadata()).unwrap();
3576 let hydrated = journal
3577 .create_box(
3578 "t1",
3579 "large result",
3580 BoxOwner::Controller,
3581 BoxContent::text("x".repeat(3_000)),
3582 )
3583 .unwrap();
3584 let summarized = journal
3585 .create_box(
3586 "t2",
3587 "Kennedy summary",
3588 BoxOwner::Kennedy,
3589 BoxContent::text("y".repeat(3_000)),
3590 )
3591 .unwrap();
3592 journal
3593 .summarize_box("t3", summarized, "important points")
3594 .unwrap();
3595 let desired = BTreeMap::from([
3596 (hydrated, BoxRepresentation::Dehydrated),
3597 (
3598 summarized,
3599 BoxRepresentation::Summarized("important points".into()),
3600 ),
3601 ]);
3602 let preview = journal
3603 .state()
3604 .projection_with_box_representations(&desired)
3605 .unwrap();
3606 let next_id = journal.state().next_id;
3607 let ids = journal.apply_box_representations("t4", &desired).unwrap();
3608 assert_eq!(ids, vec![EventId(next_id)]);
3609 assert_eq!(journal.state().projection(), preview);
3610 assert!(matches!(
3611 journal.state().box_state(hydrated).unwrap().representation,
3612 Representation::Dehydrated { .. }
3613 ));
3614 assert_eq!(
3615 journal
3616 .state()
3617 .box_state(summarized)
3618 .unwrap()
3619 .representation,
3620 Representation::Summarized {
3621 based_on: EventId(2),
3622 text: "important points".into(),
3623 }
3624 );
3625 std::fs::remove_file(path).unwrap();
3626 }
3627}