返回 CodeWhale
session.rs
根目录 / crates / tui / src / core / session.rs
1 //! Session state management for the core engine.
2 //!
3 //! Tracks conversation history, token usage, and session metadata.
4
5 use crate::project_context::{ProjectContext, load_project_context_with_parents};
6 use crate::prompt_zones::{AppendLog, FrozenPrefix};
7 use crate::working_set::WorkingSet;
8 use codewhale_core::prefix_cache::PrefixStabilityManager;
9 use codewhale_execpolicy::ApprovalMode;
10 use codewhale_models::{Message, SystemPrompt, Usage};
11 use std::collections::{HashSet, VecDeque};
12 use std::path::PathBuf;
13
14 /// Maximum number of deferred schemas a conversation may keep in its active
15 /// toolbox. The permanent `read`/`write`/`edit`/`bash`/`agent`/`tool_search`
16 /// router surface is not counted here.
17 pub(crate) const TOOL_ACTIVATION_CACHE_MAX_NAMES: usize = 8;
18 /// Maximum serialized bytes added to requests by cached deferred schemas.
19 pub(crate) const TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES: usize = 16 * 1024;
20
21 /// Bounded, process-local conversation cache for tools activated by
22 /// `tool_search`.
23 ///
24 /// Only names are retained. Every turn revalidates them against the currently
25 /// filtered catalog, so a disconnected MCP server, changed allow/deny rule, or
26 /// mode switch cannot resurrect a tool from an older authority posture.
27 #[derive(Debug, Clone, Default)]
28 pub(crate) struct ToolActivationCache {
29 /// Least-recently used at the front, most-recently used at the back.
30 names: VecDeque<String>,
31 }
32
33 #[derive(Debug, Clone, Default, Eq, PartialEq)]
34 pub(crate) struct ToolActivationDelta {
35 pub(crate) admitted: Vec<String>,
36 pub(crate) evicted: Vec<String>,
37 pub(crate) rejected: Vec<String>,
38 }
39
40 impl ToolActivationCache {
41 /// Forget every deferred tool activated by the current conversation.
42 ///
43 /// `Op::SyncSession` calls this before installing another conversation's
44 /// identity, history, and workspace. The cache is intentionally
45 /// process-local, but it must still be conversation-local.
46 pub(crate) fn clear(&mut self) {
47 self.names.clear();
48 }
49
50 fn catalog_tool<'a>(
51 catalog: &'a [codewhale_models::Tool],
52 name: &str,
53 ) -> Option<&'a codewhale_models::Tool> {
54 catalog
55 .iter()
56 .find(|tool| tool.name == name && tool.defer_loading.unwrap_or(false))
57 }
58
59 fn serialized_bytes(tool: &codewhale_models::Tool) -> usize {
60 serde_json::to_vec(tool).map_or(usize::MAX, |bytes| bytes.len())
61 }
62
63 fn total_serialized_bytes(&self, catalog: &[codewhale_models::Tool]) -> usize {
64 self.names
65 .iter()
66 .filter_map(|name| Self::catalog_tool(catalog, name))
67 .map(Self::serialized_bytes)
68 .fold(0usize, usize::saturating_add)
69 }
70
71 /// Drop entries that are no longer deferred members of this turn's
72 /// filtered catalog and enforce both cache bounds.
73 pub(crate) fn revalidate(&mut self, catalog: &[codewhale_models::Tool]) -> Vec<String> {
74 let mut evicted = Vec::new();
75 self.names.retain(|name| {
76 let keep = Self::catalog_tool(catalog, name).is_some_and(|tool| {
77 Self::serialized_bytes(tool) <= TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES
78 });
79 if !keep {
80 evicted.push(name.clone());
81 }
82 keep
83 });
84 while self.names.len() > TOOL_ACTIVATION_CACHE_MAX_NAMES
85 || self.total_serialized_bytes(catalog) > TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES
86 {
87 if let Some(name) = self.names.pop_front() {
88 evicted.push(name);
89 } else {
90 break;
91 }
92 }
93 evicted
94 }
95
96 /// Touch requested deferred tools in search-result order. An oversized
97 /// schema is rejected; otherwise least-recently-used entries are evicted
98 /// until both bounds hold.
99 pub(crate) fn activate(
100 &mut self,
101 catalog: &[codewhale_models::Tool],
102 requested: &[String],
103 ) -> ToolActivationDelta {
104 let mut delta = ToolActivationDelta {
105 evicted: self.revalidate(catalog),
106 ..ToolActivationDelta::default()
107 };
108 let mut seen = HashSet::new();
109 for name in requested {
110 if !seen.insert(name.clone()) {
111 continue;
112 }
113 let Some(tool) = Self::catalog_tool(catalog, name) else {
114 delta.rejected.push(name.clone());
115 continue;
116 };
117 if Self::serialized_bytes(tool) > TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES {
118 delta.rejected.push(name.clone());
119 continue;
120 }
121 if let Some(index) = self.names.iter().position(|cached| cached == name) {
122 self.names.remove(index);
123 }
124 self.names.push_back(name.clone());
125 while self.names.len() > TOOL_ACTIVATION_CACHE_MAX_NAMES
126 || self.total_serialized_bytes(catalog) > TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES
127 {
128 if let Some(evicted) = self.names.pop_front() {
129 delta.evicted.push(evicted);
130 }
131 }
132 }
133
134 let retained = self.names.iter().collect::<HashSet<_>>();
135 delta.admitted = requested
136 .iter()
137 .filter(|name| retained.contains(name))
138 .cloned()
139 .collect();
140 delta.evicted.sort();
141 delta.evicted.dedup();
142 delta.rejected.sort();
143 delta.rejected.dedup();
144 delta
145 }
146
147 pub(crate) fn names(&self) -> impl Iterator<Item = &str> {
148 self.names.iter().map(String::as_str)
149 }
150 }
151
152 /// Session state for the engine.
153 #[derive(Debug, Clone)]
154 pub struct Session {
155 /// Model being used
156 pub model: String,
157
158 /// Reasoning-effort tier for DeepSeek thinking mode:
159 /// `"off" | "low" | "medium" | "high" | "max"`. `None` lets the provider
160 /// apply its own defaults.
161 pub reasoning_effort: Option<String>,
162 /// Whether the user selected automatic reasoning effort.
163 pub reasoning_effort_auto: bool,
164
165 /// Whether the user selected automatic model routing.
166 pub auto_model: bool,
167
168 /// Workspace directory
169 pub workspace: PathBuf,
170
171 /// System prompt (optional)
172 pub system_prompt: Option<SystemPrompt>,
173 /// True when `system_prompt` is a persisted/runtime-supplied prefix that
174 /// should not be replaced by mode/context refreshes.
175 pub system_prompt_override: bool,
176 /// Hash of the last assembled stable system prompt. Used to avoid
177 /// replacing `system_prompt` when unchanged.
178 pub last_system_prompt_hash: Option<u64>,
179 /// Reason the pinned prefix will move on the next model request, set by
180 /// an explicit header-change op (`/model`, mode change, goal edit, session
181 /// sync) when it actually alters the system-prompt bytes. Consumed by the
182 /// turn loop's prefix check so a declared change re-pins under a logged
183 /// reason while an undeclared mid-loop change is reported as drift and the
184 /// pin holds. `None` means "no declared change since the last request".
185 pub pending_prefix_change_reason: Option<String>,
186 /// The explicit prompt inputs (model, mode, goal, route, translation,
187 /// verbosity) the pinned system prompt was composed from. At each new user
188 /// turn the engine recomposes: if these inputs are unchanged but the
189 /// composed bytes differ, that is workspace/instruction/skills/memory
190 /// drift and is delivered to the model as a `<context_update>` message,
191 /// never by moving the pinned header.
192 pub(crate) pinned_prompt_context: Option<crate::core::engine::NextTurnPromptContext>,
193 /// Flat text of the session context the model has last been shown, either
194 /// as the pinned header or through `<context_update>` messages. New-turn
195 /// deltas are computed against this so an update is delivered exactly once.
196 pub(crate) context_update_baseline: Option<String>,
197 /// Host-persistence copy of the history checkpoint generated by context
198 /// compaction. This is never part of the standing system prompt.
199 pub compaction_summary_prompt: Option<SystemPrompt>,
200
201 /// Conversation history (API format), backed by AppendLog (#2264).
202 pub messages: AppendLog,
203
204 /// Total tokens used in this session
205 pub total_usage: SessionUsage,
206
207 /// Whether shell execution is allowed
208 pub allow_shell: bool,
209
210 /// Whether to trust paths outside workspace
211 pub trust_mode: bool,
212
213 /// Whether the current session should auto-approve tool safety checks.
214 pub auto_approve: bool,
215
216 /// Live UI approval policy used to steer the system prompt.
217 pub approval_mode: ApprovalMode,
218
219 /// Notes file path
220 pub notes_path: PathBuf,
221
222 /// MCP config path
223 pub mcp_config_path: PathBuf,
224
225 /// Session ID (for tracking)
226 pub id: String,
227
228 /// Project context loaded from AGENTS.md, etc.
229 pub project_context: Option<ProjectContext>,
230
231 /// Repo-aware working set for context management.
232 pub working_set: WorkingSet,
233
234 /// Prefix-cache stability monitor (inspired by Reasonix's Pillar 1).
235 /// Tracks the immutable prefix fingerprint and detects drift across turns.
236 /// Set during engine construction; None until the first system prompt assembly.
237 pub prefix_stability: Option<PrefixStabilityManager>,
238
239 /// Three-zone immutable prefix baseline (#2264). Frozen on the first
240 /// request of the session; verified against the current system+tool
241 /// state before every subsequent request. None until the first turn.
242 pub frozen_prefix: Option<FrozenPrefix>,
243
244 /// Deferred tools explicitly discovered during this conversation. Names
245 /// are revalidated against the live catalog before each request.
246 pub(super) tool_activation_cache: ToolActivationCache,
247
248 /// Monotonic counter bumped on every direct mutation of `messages`.
249 /// Consumed by the engine token-estimate cache
250 /// to memoize the per-turn token estimate without re-walking the message
251 /// list. Defaults to 0; bumped in [`Session::add_message`],
252 /// [`Session::replace_messages`], and at other mutation sites in
253 /// `core/engine.rs`.
254 pub messages_revision: u64,
255
256 /// Provider-billed prompt tokens from the most recent parent-route
257 /// request that still describes the live message list. Survives turn
258 /// boundaries so the next send can compact on last-turn pressure
259 /// (#5577). Cleared when history is rewritten (compaction, restore).
260 pub(crate) latest_parent_input_tokens: Option<u32>,
261 }
262
263 /// Cumulative usage statistics for a session.
264 #[derive(Debug, Clone, Default)]
265 #[allow(clippy::struct_field_names)]
266 pub struct SessionUsage {
267 pub input_tokens: u64,
268 pub output_tokens: u64,
269 /// Cache creation (write) tokens. `None` when never observed by the API —
270 /// do NOT display as 0, which would be indistinguishable from "no writes".
271 pub cache_creation_input_tokens: Option<u64>,
272 /// Cache read (hit) tokens. `None` when never observed by the API —
273 /// do NOT display as 0, which would be indistinguishable from "no hits".
274 pub cache_read_input_tokens: Option<u64>,
275 }
276
277 impl SessionUsage {
278 /// Add usage from a turn
279 pub fn add(&mut self, usage: &Usage) {
280 self.input_tokens += u64::from(usage.input_tokens);
281 self.output_tokens += u64::from(usage.output_tokens);
282 if let Some(tokens) = usage.prompt_cache_write_tokens {
283 self.cache_creation_input_tokens =
284 Some(self.cache_creation_input_tokens.unwrap_or(0) + u64::from(tokens));
285 }
286 if let Some(tokens) = usage.prompt_cache_hit_tokens {
287 self.cache_read_input_tokens =
288 Some(self.cache_read_input_tokens.unwrap_or(0) + u64::from(tokens));
289 }
290 }
291 }
292
293 impl Session {
294 /// Create a new session
295 pub fn new(
296 model: String,
297 workspace: PathBuf,
298 allow_shell: bool,
299 trust_mode: bool,
300 notes_path: PathBuf,
301 mcp_config_path: PathBuf,
302 ) -> Self {
303 // Load project context from AGENTS.md, CLAUDE.md, etc.
304 let project_context = load_project_context_with_parents(&workspace);
305 let has_context = project_context.has_instructions();
306
307 Self {
308 model,
309 reasoning_effort: None,
310 reasoning_effort_auto: false,
311 auto_model: false,
312 workspace,
313 system_prompt: None,
314 system_prompt_override: false,
315 compaction_summary_prompt: None,
316 messages: AppendLog::new(),
317 total_usage: SessionUsage::default(),
318 allow_shell,
319 trust_mode,
320 auto_approve: false,
321 approval_mode: ApprovalMode::Suggest,
322 notes_path,
323 mcp_config_path,
324 id: uuid::Uuid::new_v4().to_string(),
325 project_context: if has_context {
326 Some(project_context)
327 } else {
328 None
329 },
330 last_system_prompt_hash: None,
331 pending_prefix_change_reason: None,
332 pinned_prompt_context: None,
333 context_update_baseline: None,
334 working_set: WorkingSet::default(),
335 prefix_stability: None,
336 frozen_prefix: None,
337 tool_activation_cache: ToolActivationCache::default(),
338 messages_revision: 0,
339 latest_parent_input_tokens: None,
340 }
341 }
342
343 /// Add a message to the conversation
344 pub fn add_message(&mut self, message: Message) {
345 self.messages.push(message);
346 self.messages_revision = self.messages_revision.saturating_add(1);
347 }
348
349 /// Replace the entire message history. Used by session resume and
350 /// compaction. Bumps `messages_revision` exactly once even when the new
351 /// history has a different length, so downstream caches invalidate
352 /// atomically.
353 pub fn replace_messages(&mut self, messages: Vec<Message>) {
354 self.messages = messages.into();
355 self.messages_revision = self.messages_revision.saturating_add(1);
356 self.latest_parent_input_tokens = None;
357 }
358
359 /// Bump `messages_revision` without otherwise mutating the message list.
360 /// Reserved for sites that mutate the message list in place (e.g. an
361 /// in-place rewrite of a content block). Most call sites do not need
362 /// this — prefer [`add_message`](Self::add_message) and
363 /// [`replace_messages`](Self::replace_messages).
364 pub fn bump_messages_revision(&mut self) {
365 self.messages_revision = self.messages_revision.saturating_add(1);
366 }
367
368 /// Rebuild the working set from current messages (best effort).
369 pub fn rebuild_working_set(&mut self) {
370 self.working_set
371 .rebuild_from_messages(&self.messages, &self.workspace);
372 }
373 }
374
375 #[cfg(test)]
376 mod tests {
377 use super::*;
378 use serde_json::json;
379
380 fn deferred_tool(name: &str, description_bytes: usize) -> codewhale_models::Tool {
381 codewhale_models::Tool {
382 tool_type: None,
383 name: name.to_string(),
384 description: "x".repeat(description_bytes),
385 input_schema: json!({"type": "object", "properties": {}}),
386 allowed_callers: None,
387 defer_loading: Some(true),
388 input_examples: None,
389 strict: None,
390 cache_control: None,
391 }
392 }
393
394 #[test]
395 fn session_usage_cache_starts_none() {
396 let usage = SessionUsage::default();
397 assert!(usage.cache_creation_input_tokens.is_none());
398 assert!(usage.cache_read_input_tokens.is_none());
399 }
400
401 #[test]
402 fn session_usage_cache_remains_none_when_api_omits_cache() {
403 let mut usage = SessionUsage::default();
404 let api_usage = Usage {
405 input_tokens: 100,
406 output_tokens: 50,
407 prompt_cache_hit_tokens: None,
408 prompt_cache_miss_tokens: None,
409 prompt_cache_write_tokens: None,
410 reasoning_tokens: None,
411 reasoning_replay_tokens: None,
412 server_tool_use: None,
413 };
414 usage.add(&api_usage);
415 assert!(usage.cache_creation_input_tokens.is_none());
416 assert!(usage.cache_read_input_tokens.is_none());
417 }
418
419 #[test]
420 fn session_usage_cache_accumulates_when_reported() {
421 let mut usage = SessionUsage::default();
422 let api_usage = Usage {
423 input_tokens: 100,
424 output_tokens: 50,
425 prompt_cache_hit_tokens: Some(30),
426 prompt_cache_miss_tokens: Some(50),
427 prompt_cache_write_tokens: Some(20),
428 reasoning_tokens: None,
429 reasoning_replay_tokens: None,
430 server_tool_use: None,
431 };
432 usage.add(&api_usage);
433 assert_eq!(usage.cache_read_input_tokens, Some(30));
434 assert_eq!(usage.cache_creation_input_tokens, Some(20));
435 usage.add(&api_usage);
436 assert_eq!(usage.cache_read_input_tokens, Some(60));
437 assert_eq!(usage.cache_creation_input_tokens, Some(40));
438 }
439
440 #[test]
441 fn session_usage_cache_preserves_explicit_zero() {
442 let mut usage = SessionUsage::default();
443 let api_usage = Usage {
444 input_tokens: 100,
445 output_tokens: 50,
446 prompt_cache_hit_tokens: Some(0), // explicit zero from provider
447 prompt_cache_miss_tokens: Some(50),
448 prompt_cache_write_tokens: Some(1234),
449 reasoning_tokens: None,
450 reasoning_replay_tokens: None,
451 server_tool_use: None,
452 };
453 usage.add(&api_usage);
454 // 0 is a valid observed value, must NOT be converted to None
455 assert_eq!(usage.cache_read_input_tokens, Some(0));
456 assert_eq!(usage.cache_creation_input_tokens, Some(1234));
457 }
458
459 #[test]
460 fn tool_activation_cache_is_lru_bounded_to_eight_names() {
461 let catalog = (0..10)
462 .map(|index| deferred_tool(&format!("tool_{index}"), 8))
463 .collect::<Vec<_>>();
464 let requested = catalog
465 .iter()
466 .map(|tool| tool.name.clone())
467 .collect::<Vec<_>>();
468 let mut cache = ToolActivationCache::default();
469 let delta = cache.activate(&catalog, &requested);
470
471 assert_eq!(cache.names().count(), TOOL_ACTIVATION_CACHE_MAX_NAMES);
472 assert_eq!(
473 cache.names().collect::<Vec<_>>(),
474 vec![
475 "tool_2", "tool_3", "tool_4", "tool_5", "tool_6", "tool_7", "tool_8", "tool_9"
476 ]
477 );
478 assert_eq!(delta.admitted.len(), TOOL_ACTIVATION_CACHE_MAX_NAMES);
479 assert!(delta.evicted.contains(&"tool_0".to_string()));
480 assert!(delta.evicted.contains(&"tool_1".to_string()));
481 }
482
483 #[test]
484 fn touching_a_cached_tool_makes_it_most_recent() {
485 let catalog = (0..9)
486 .map(|index| deferred_tool(&format!("tool_{index}"), 8))
487 .collect::<Vec<_>>();
488 let first_eight = catalog[..8]
489 .iter()
490 .map(|tool| tool.name.clone())
491 .collect::<Vec<_>>();
492 let mut cache = ToolActivationCache::default();
493 cache.activate(&catalog, &first_eight);
494 cache.activate(&catalog, &["tool_0".to_string()]);
495 cache.activate(&catalog, &["tool_8".to_string()]);
496
497 let names = cache.names().collect::<Vec<_>>();
498 assert!(names.contains(&"tool_0"));
499 assert!(!names.contains(&"tool_1"));
500 assert_eq!(names.last().copied(), Some("tool_8"));
501 }
502
503 #[test]
504 fn oversized_schema_is_never_admitted() {
505 let catalog = vec![deferred_tool(
506 "huge",
507 TOOL_ACTIVATION_CACHE_MAX_SCHEMA_BYTES + 1,
508 )];
509 let mut cache = ToolActivationCache::default();
510 let delta = cache.activate(&catalog, &["huge".to_string()]);
511
512 assert_eq!(cache.names().count(), 0);
513 assert_eq!(delta.rejected, vec!["huge"]);
514 }
515
516 #[test]
517 fn revalidate_drops_removed_denied_or_eager_tools() {
518 let catalog = vec![deferred_tool("kept", 8), deferred_tool("gone", 8)];
519 let mut cache = ToolActivationCache::default();
520 cache.activate(&catalog, &["kept".to_string(), "gone".to_string()]);
521 let mut next_catalog = vec![deferred_tool("kept", 8), deferred_tool("gone", 8)];
522 next_catalog[1].defer_loading = Some(false);
523
524 let evicted = cache.revalidate(&next_catalog);
525 assert_eq!(cache.names().collect::<Vec<_>>(), vec!["kept"]);
526 assert_eq!(evicted, vec!["gone"]);
527 }
528
529 #[test]
530 fn clearing_for_session_sync_forgets_all_activated_tools() {
531 let catalog = vec![deferred_tool("one", 8), deferred_tool("two", 8)];
532 let mut cache = ToolActivationCache::default();
533 cache.activate(&catalog, &["one".to_string(), "two".to_string()]);
534
535 cache.clear();
536
537 assert_eq!(cache.names().count(), 0);
538 }
539 }
540
540 lines RUST