Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

compact-improvements をマージ

Browse Source 
- 閾値の個別指定化 (compact_threshold / compact_request_threshold) と Option 化
- 占有量ソースを UsageRecord timeline に一本化 (last_input_tokens 撤去)
- retained_turns → retained_tokens
- compact worker をツール駆動に再設計 (mark_read_required / add_reference / write_summary / read_file)
- Auto-read budget と compact_worker_max_input_tokens の上限制御
- 新 history は system message のみで構成 [summary, auto-read..., references, retained...]
This commit is contained in:
Hare
2026-04-19 12:14:16 +09:00
parent ec3bf7324b 3c510860fa
commit 605e78468c
12 changed files with 1067 additions and 628 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/manifest/src/config.rs
+29 -7
View File
@@ -84,7 +84,13 @@ pub struct CompactionConfigPartial {
#[serde(default)]
pub compact_threshold: Option<u64>,
#[serde(default)]
pub compact_retained_turns: Option<usize>,
pub compact_request_threshold: Option<u64>,
#[serde(default)]
pub compact_retained_tokens: Option<u64>,
#[serde(default)]
pub compact_auto_read_budget: Option<u64>,
#[serde(default)]
pub compact_worker_max_input_tokens: Option<u64>,
#[serde(default)]
pub provider: Option<ProviderConfigPartial>,
}
@@ -236,9 +242,18 @@ impl CompactionConfigPartial {
prune_protected_turns: upper.prune_protected_turns.or(self.prune_protected_turns),
prune_min_savings: upper.prune_min_savings.or(self.prune_min_savings),
compact_threshold: upper.compact_threshold.or(self.compact_threshold),
compact_retained_turns: upper
.compact_retained_turns
.or(self.compact_retained_turns),
compact_request_threshold: upper
.compact_request_threshold
.or(self.compact_request_threshold),
compact_retained_tokens: upper
.compact_retained_tokens
.or(self.compact_retained_tokens),
compact_auto_read_budget: upper
.compact_auto_read_budget
.or(self.compact_auto_read_budget),
compact_worker_max_input_tokens: upper
.compact_worker_max_input_tokens
.or(self.compact_worker_max_input_tokens),
provider: merge_option(self.provider, upper.provider, ProviderConfigPartial::merge),
}
}
@@ -365,9 +380,16 @@ impl TryFrom<PodManifestConfig> for PodManifest {
.prune_min_savings
.unwrap_or(defaults::PRUNE_MIN_SAVINGS),
compact_threshold: c.compact_threshold,
compact_retained_turns: c
.compact_retained_turns
.unwrap_or(defaults::COMPACT_RETAINED_TURNS),
compact_request_threshold: c.compact_request_threshold,
compact_retained_tokens: c
.compact_retained_tokens
.unwrap_or(defaults::COMPACT_RETAINED_TOKENS),
compact_auto_read_budget: c
.compact_auto_read_budget
.unwrap_or(defaults::COMPACT_AUTO_READ_BUDGET),
compact_worker_max_input_tokens: c
.compact_worker_max_input_tokens
.unwrap_or(defaults::COMPACT_WORKER_MAX_INPUT_TOKENS),
provider: comp_provider,
})
})
crates/manifest/src/defaults.rs
+22 -3
View File
@@ -18,11 +18,30 @@ pub const PRUNE_PROTECTED_TURNS: usize = 3;
/// [`crate::CompactionConfig::prune_min_savings`].
pub const PRUNE_MIN_SAVINGS: u64 = 4096;
/// Number of most-recent turns retained after a compact. See
/// [`crate::CompactionConfig::compact_retained_turns`].
pub const COMPACT_RETAINED_TURNS: usize = 2;
/// Token budget retained (unchanged) at the tail of the history across
/// a compact. Items whose cumulative token count fits within this budget
/// starting from the end are kept verbatim; the rest are summarised.
/// See [`crate::CompactionConfig::compact_retained_tokens`].
pub const COMPACT_RETAINED_TOKENS: u64 = 8000;
/// Default instruction asset reference used when `worker.instruction`
/// is omitted. See the `PromptLoader` prefix addressing scheme for the
/// `$insomnia/` / `$user/` / `$workspace/` namespaces.
pub const DEFAULT_INSTRUCTION: &str = "$insomnia/default";
/// Token budget for auto-read file contents injected into the new
/// session after compaction. Limits how much raw file text the
/// compact worker can pull into the compacted context via
/// `mark_read_required`. See
/// [`crate::CompactionConfig::compact_auto_read_budget`].
pub const COMPACT_AUTO_READ_BUDGET: u64 = 8000;
/// Cumulative input-token cap for the compact worker's own LLM
/// calls. Exceeding this aborts the compact run (circuit-breaker
/// path). See
/// [`crate::CompactionConfig::compact_worker_max_input_tokens`].
pub const COMPACT_WORKER_MAX_INPUT_TOKENS: u64 = 50_000;
/// Number of recently-touched files fed to the compact worker as
/// default references.
pub const COMPACT_DEFAULT_REFERENCE_COUNT: usize = 5;
crates/manifest/src/lib.rs
+74 -8
View File
@@ -174,13 +174,42 @@ pub struct CompactionConfig {
#[serde(default = "default_prune_min_savings")]
pub prune_min_savings: u64,
/// When `input_tokens` exceeds this, run compact. `None` = compact disabled.
/// Proactive (between-turns) compaction threshold.
///
/// Checked by the Controller after each run. When current occupancy
/// exceeds this value, compact runs before the next turn. `None`
/// disables the between-turns check.
#[serde(default)]
pub compact_threshold: Option<u64>,
/// Number of recent turns retained after compaction.
#[serde(default = "default_compact_retained_turns")]
pub compact_retained_turns: usize,
/// Safety-net (between-requests) compaction threshold.
///
/// Checked by `PodInterceptor::pre_llm_request` inside a turn. When
/// current occupancy exceeds this value, the run yields so that the
/// Controller can compact before the next LLM request. `None`
/// disables the between-requests check.
///
/// Expected relation: `compact_threshold < compact_request_threshold`
/// (proactive triggers before safety net). A reversed configuration
/// is accepted but logged as a warning.
#[serde(default)]
pub compact_request_threshold: Option<u64>,
/// Token budget retained verbatim at the tail of the history after
/// compaction. Measured against the occupancy estimate from
/// `UsageRecord` history; turn boundaries are ignored.
#[serde(default = "default_compact_retained_tokens")]
pub compact_retained_tokens: u64,
/// Aggregate token budget for auto-read file contents injected into
/// the compacted session by the compact worker.
#[serde(default = "default_compact_auto_read_budget")]
pub compact_auto_read_budget: u64,
/// Cumulative input-token cap for the compact worker's own LLM
/// calls. Exceeding this aborts the compact run.
#[serde(default = "default_compact_worker_max_input_tokens")]
pub compact_worker_max_input_tokens: u64,
/// Optional provider for the compactor (summary) LLM.
/// If omitted, the main provider is cloned via `clone_boxed()`.
@@ -194,8 +223,14 @@ fn default_prune_protected_turns() -> usize {
fn default_prune_min_savings() -> u64 {
defaults::PRUNE_MIN_SAVINGS
}
fn default_compact_retained_turns() -> usize {
defaults::COMPACT_RETAINED_TURNS
fn default_compact_retained_tokens() -> u64 {
defaults::COMPACT_RETAINED_TOKENS
}
fn default_compact_auto_read_budget() -> u64 {
defaults::COMPACT_AUTO_READ_BUDGET
}
fn default_compact_worker_max_input_tokens() -> u64 {
defaults::COMPACT_WORKER_MAX_INPUT_TOKENS
}
impl Default for CompactionConfig {
@@ -204,7 +239,10 @@ impl Default for CompactionConfig {
prune_protected_turns: default_prune_protected_turns(),
prune_min_savings: default_prune_min_savings(),
compact_threshold: None,
compact_retained_turns: default_compact_retained_turns(),
compact_request_threshold: None,
compact_retained_tokens: default_compact_retained_tokens(),
compact_auto_read_budget: default_compact_auto_read_budget(),
compact_worker_max_input_tokens: default_compact_worker_max_input_tokens(),
provider: None,
}
}
@@ -338,7 +376,35 @@ model = "claude-sonnet-4-20250514"
assert_eq!(c.prune_protected_turns, 3);
assert_eq!(c.prune_min_savings, 4096);
assert_eq!(c.compact_threshold, Some(80000));
assert_eq!(c.compact_retained_turns, 2);
assert_eq!(c.compact_request_threshold, None);
assert_eq!(c.compact_retained_tokens, 8000);
}
#[test]
fn parse_compaction_both_thresholds() {
let toml = format!(
"{MINIMAL_REQUIRED}\n\
[compaction]\n\
compact_threshold = 80000\n\
compact_request_threshold = 90000\n"
);
let manifest = PodManifest::from_toml(&toml).unwrap();
let c = manifest.compaction.unwrap();
assert_eq!(c.compact_threshold, Some(80000));
assert_eq!(c.compact_request_threshold, Some(90000));
}
#[test]
fn parse_compaction_request_threshold_only() {
let toml = format!(
"{MINIMAL_REQUIRED}\n\
[compaction]\n\
compact_request_threshold = 90000\n"
);
let manifest = PodManifest::from_toml(&toml).unwrap();
let c = manifest.compaction.unwrap();
assert_eq!(c.compact_threshold, None);
assert_eq!(c.compact_request_threshold, Some(90000));
}
#[test]
crates/pod/src/compact_state.rs
+80 -63
View File
@@ -1,24 +1,32 @@
//! Shared state for compaction decisions.
//!
//! Holds atomic counters shared between:
//! - `on_usage` callback (writes `last_input_tokens`)
//! - `CompactInterceptor` (reads token count, checks thresholds)
//! - `Pod::run()`/`resume()` (circuit breaker, thrash detection)
//! Holds the two configured thresholds and circuit-breaker / thrash-detection
//! flags shared between:
//! - `PodInterceptor` (reads `request_threshold` — the *safety net* for
//! between-requests yielding)
//! - `Pod::try_post_run_compact` (reads `post_run_threshold` — the
//! *proactive* check between turns)
//! - `Pod::run()` / `resume()` (circuit breaker, thrash detection)
//!
//! Current occupancy (input-token count) is **not** stored here. The single
//! source of truth is `session_store::UsageRecord` (persisted per LLM call)
//! projected through `Pod::total_tokens()`. Callers pass the current
//! occupancy to `exceeds_*` at check time.
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
const MAX_COMPACT_FAILURES: usize = 3;
/// Shared mutable state for compaction decisions.
pub(crate) struct CompactState {
/// Last observed input_tokens from `on_usage` callback.
last_input_tokens: AtomicU64,
/// Proactive threshold — checked in `pre_llm_request` (between turns).
turn_threshold: u64,
/// Post-run threshold — checked by Controller after run completes.
post_run_threshold: u64,
/// Number of recent turns to retain after compaction.
retained_turns: usize,
/// Between-turns threshold (proactive). Checked by the Controller
/// after a run completes. `None` disables the post-run check.
post_run_threshold: Option<u64>,
/// Between-requests threshold (safety net). Checked inside a turn
/// before each LLM request. `None` disables the request check.
request_threshold: Option<u64>,
/// Token budget retained verbatim at the tail after compaction.
retained_tokens: u64,
/// Consecutive compact failures. At `MAX_COMPACT_FAILURES`, compaction is disabled.
consecutive_failures: AtomicUsize,
/// `true` immediately after a successful compact, cleared on next normal completion.
@@ -28,40 +36,29 @@ pub(crate) struct CompactState {
}
impl CompactState {
/// Create a new CompactState.
///
/// `turn_threshold` is the proactive (80%) threshold from the manifest.
/// `post_run_threshold` is derived as `turn_threshold * 9 / 8` (≈90%).
pub(crate) fn new(turn_threshold: u64, retained_turns: usize) -> Self {
pub(crate) fn new(
post_run_threshold: Option<u64>,
request_threshold: Option<u64>,
retained_tokens: u64,
) -> Self {
Self {
last_input_tokens: AtomicU64::new(0),
turn_threshold,
post_run_threshold: turn_threshold * 9 / 8,
retained_turns,
post_run_threshold,
request_threshold,
retained_tokens,
consecutive_failures: AtomicUsize::new(0),
just_compacted: AtomicBool::new(false),
disabled: AtomicBool::new(false),
}
}
/// Update the last observed input_tokens (called from `on_usage`).
pub(crate) fn update_input_tokens(&self, tokens: u64) {
self.last_input_tokens.store(tokens, Ordering::Relaxed);
/// Configured between-requests threshold (if any).
pub(crate) fn request_threshold(&self) -> Option<u64> {
self.request_threshold
}
/// Read the last observed input_tokens.
pub(crate) fn last_input_tokens(&self) -> u64 {
self.last_input_tokens.load(Ordering::Relaxed)
}
/// The between-turns threshold value.
pub(crate) fn turn_threshold(&self) -> u64 {
self.turn_threshold
}
/// Number of turns to retain after compaction.
pub(crate) fn retained_turns(&self) -> usize {
self.retained_turns
/// Token budget retained verbatim at the tail after compaction.
pub(crate) fn retained_tokens(&self) -> u64 {
self.retained_tokens
}
/// Whether compaction has been disabled by the circuit breaker.
@@ -69,14 +66,20 @@ impl CompactState {
self.disabled.load(Ordering::Relaxed)
}
/// Whether `last_input_tokens` exceeds the between-turns threshold.
pub(crate) fn exceeds_turn(&self) -> bool {
self.last_input_tokens() > self.turn_threshold
/// Whether `current_tokens` exceeds the between-requests threshold.
/// Returns `false` when `request_threshold` is unset.
pub(crate) fn exceeds_request(&self, current_tokens: u64) -> bool {
self.request_threshold
.map(|t| current_tokens > t)
.unwrap_or(false)
}
/// Whether `last_input_tokens` exceeds the post-run threshold.
pub(crate) fn exceeds_post_run(&self) -> bool {
self.last_input_tokens() > self.post_run_threshold
/// Whether `current_tokens` exceeds the post-run threshold.
/// Returns `false` when `post_run_threshold` is unset.
pub(crate) fn exceeds_post_run(&self, current_tokens: u64) -> bool {
self.post_run_threshold
.map(|t| current_tokens > t)
.unwrap_or(false)
}
/// Whether a compact just completed (for thrash detection).
@@ -109,31 +112,46 @@ mod tests {
use super::*;
#[test]
fn threshold_derivation() {
let state = CompactState::new(80_000, 2);
assert_eq!(state.turn_threshold, 80_000);
assert_eq!(state.post_run_threshold, 90_000);
assert_eq!(state.retained_turns(), 2);
fn both_thresholds_configured() {
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
assert_eq!(state.request_threshold(), Some(90_000));
assert_eq!(state.retained_tokens(), 8_000);
assert!(!state.exceeds_request(70_000));
assert!(!state.exceeds_post_run(70_000));
assert!(!state.exceeds_request(85_000));
assert!(state.exceeds_post_run(85_000));
assert!(state.exceeds_request(95_000));
assert!(state.exceeds_post_run(95_000));
}
#[test]
fn exceeds_checks() {
let state = CompactState::new(80_000, 2);
assert!(!state.exceeds_turn());
assert!(!state.exceeds_post_run());
fn post_run_only() {
let state = CompactState::new(Some(80_000), None, 8_000);
// request check always false when threshold is None.
assert!(!state.exceeds_request(1_000_000));
assert!(state.exceeds_post_run(85_000));
}
state.update_input_tokens(85_000);
assert!(state.exceeds_turn());
assert!(!state.exceeds_post_run());
#[test]
fn request_only() {
let state = CompactState::new(None, Some(90_000), 8_000);
assert!(!state.exceeds_post_run(1_000_000));
assert!(state.exceeds_request(95_000));
}
state.update_input_tokens(95_000);
assert!(state.exceeds_turn());
assert!(state.exceeds_post_run());
#[test]
fn both_none_disables_all_checks() {
let state = CompactState::new(None, None, 8_000);
assert!(!state.exceeds_request(1_000_000));
assert!(!state.exceeds_post_run(1_000_000));
}
#[test]
fn circuit_breaker_trips_after_max_failures() {
let state = CompactState::new(80_000, 2);
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
assert!(!state.is_disabled());
state.record_compact_failure();
@@ -146,7 +164,7 @@ mod tests {
#[test]
fn success_resets_failure_count() {
let state = CompactState::new(80_000, 2);
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
state.record_compact_failure();
state.record_compact_failure();
assert!(!state.is_disabled());
@@ -154,7 +172,6 @@ mod tests {
state.record_compact_success();
assert!(state.just_compacted());
// After success + 2 more failures, still not disabled (count was reset).
state.record_compact_failure();
state.record_compact_failure();
assert!(!state.is_disabled());
@@ -162,7 +179,7 @@ mod tests {
#[test]
fn just_compacted_lifecycle() {
let state = CompactState::new(80_000, 2);
let state = CompactState::new(Some(80_000), Some(90_000), 8_000);
assert!(!state.just_compacted());
state.record_compact_success();
crates/pod/src/compact_worker.rs
+384
View File
@@ -0,0 +1,384 @@
//! Compact worker state and the four tools that drive it.
//!
//! The compact worker is a disposable `Worker` instance spun up by
//! [`Pod::compact`]. It receives the history to summarise plus a list of
//! default reference files (from the session-lifetime `Tracker`) and runs
//! a tool-driven LLM loop. The tools here let it:
//!
//! - `read_file` — inspect referenced files (reuses `tools::read_tool`)
//! - `mark_read_required(path, offset?, limit?)` — nominate a file whose
//! contents should be injected into the compacted context as an
//! auto-read system message
//! - `add_reference(path)` — nominate a file the next session should
//! know about by name only (contents not included)
//! - `write_summary(text)` — deliver (or overwrite) the structured summary
//!
//! Everything the worker decides ends up in [`CompactWorkerContext`],
//! which `Pod::compact` drains after the loop and turns into the
//! compacted session's opening system messages.
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use llm_worker::Item;
use llm_worker::interceptor::{Interceptor, PreRequestAction};
use llm_worker::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use serde::Deserialize;
use tools::ScopedFs;
/// A file the compact worker has marked for auto-read in the new session.
#[derive(Debug, Clone)]
pub(crate) struct ReadRequirement {
pub path: PathBuf,
/// 0-based line offset. `None` means from the start of the file.
pub offset: Option<usize>,
/// Maximum number of lines. `None` means to the end of the file.
pub limit: Option<usize>,
}
/// Aggregated output of a compact worker run.
#[derive(Debug, Default, Clone)]
pub(crate) struct CompactWorkerContext {
pub read_required: Vec<ReadRequirement>,
pub references: Vec<PathBuf>,
pub summary: Option<String>,
/// Tokens already consumed by `mark_read_required` calls.
pub auto_read_consumed: u64,
/// Aggregate cap. `0` treats the budget as disabled.
pub auto_read_budget: u64,
}
impl CompactWorkerContext {
pub(crate) fn with_budget(auto_read_budget: u64) -> Self {
Self {
auto_read_budget,
..Self::default()
}
}
fn remaining_budget(&self) -> u64 {
self.auto_read_budget.saturating_sub(self.auto_read_consumed)
}
}
/// Input to `mark_read_required`.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct MarkParams {
/// Absolute path to the file.
pub file_path: PathBuf,
/// 0-based line offset.
#[serde(default)]
pub offset: Option<usize>,
/// Maximum number of lines to inject.
#[serde(default)]
pub limit: Option<usize>,
}
/// Input to `add_reference`.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct ReferenceParams {
/// Absolute path to the file.
pub file_path: PathBuf,
}
/// Input to `write_summary`.
#[derive(Debug, Deserialize, schemars::JsonSchema)]
struct SummaryParams {
/// Full structured summary text (overwrites any previous call).
pub text: String,
}
const MARK_DESCRIPTION: &str = "Inject a file's contents into the compacted context so the \
next session starts with it already read. Use this for files the next task needs in full. \
Optionally specify `offset` (0-based line) and `limit` (line count) to inject only a slice. \
Counts against `auto_read_budget`; overflow returns an error and the mark is not recorded. \
Paths must be absolute.";
const REFERENCE_DESCRIPTION: &str = "Record a file path as a named reference in the compacted \
context without injecting its contents. Use for files that are contextually relevant but \
whose current content the next session can fetch on demand.";
const SUMMARY_DESCRIPTION: &str = "Provide the final structured summary text. Subsequent calls \
replace the previous content; only the last call is used. Must be called before the compact run \
ends or compaction fails.";
struct MarkReadRequiredTool {
fs: ScopedFs,
ctx: Arc<Mutex<CompactWorkerContext>>,
}
#[async_trait]
impl Tool for MarkReadRequiredTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: MarkParams = serde_json::from_str(input_json).map_err(|e| {
ToolError::InvalidArgument(format!("invalid mark_read_required input: {e}"))
})?;
// Read the file through the shared ScopedFs so scope and I/O
// errors surface the same way the regular `read_file` tool does.
let bytes = self
.fs
.read_bytes(&params.file_path)
.map_err(|e| ToolError::ExecutionFailed(format!("read failed: {e}")))?;
let text = String::from_utf8_lossy(&bytes);
let slice = slice_lines(&text, params.offset.unwrap_or(0), params.limit);
let estimated_tokens = estimate_tokens(slice.len());
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
let budget = guard.auto_read_budget;
let would_consume = guard.auto_read_consumed.saturating_add(estimated_tokens);
if budget > 0 && would_consume > budget {
return Err(ToolError::ExecutionFailed(format!(
"auto-read budget exhausted ({budget} tokens). Remove an existing mark or use \
add_reference instead."
)));
}
guard.read_required.push(ReadRequirement {
path: params.file_path.clone(),
offset: params.offset,
limit: params.limit,
});
guard.auto_read_consumed = would_consume;
let remaining = guard.remaining_budget();
drop(guard);
let mut summary = format!(
"Marked {} for auto-read (≈{estimated_tokens} tokens). \
Budget: {remaining}/{budget} tokens remaining.",
params.file_path.display()
);
if budget > 0 && remaining * 2 <= budget {
summary.push_str(
"\nNote: auto-read budget is at least half consumed. \
Consider calling write_summary and finishing up soon.",
);
}
Ok(ToolOutput {
summary,
content: None,
})
}
}
struct AddReferenceTool {
ctx: Arc<Mutex<CompactWorkerContext>>,
}
#[async_trait]
impl Tool for AddReferenceTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: ReferenceParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid add_reference input: {e}")))?;
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
if !guard
.references
.iter()
.any(|p| p.as_path() == params.file_path.as_path())
{
guard.references.push(params.file_path.clone());
}
Ok(ToolOutput {
summary: format!("Added reference {}", params.file_path.display()),
content: None,
})
}
}
struct WriteSummaryTool {
ctx: Arc<Mutex<CompactWorkerContext>>,
}
#[async_trait]
impl Tool for WriteSummaryTool {
async fn execute(&self, input_json: &str) -> Result<ToolOutput, ToolError> {
let params: SummaryParams = serde_json::from_str(input_json)
.map_err(|e| ToolError::InvalidArgument(format!("invalid write_summary input: {e}")))?;
let mut guard = self.ctx.lock().expect("compact worker context poisoned");
let overwritten = guard.summary.is_some();
guard.summary = Some(params.text);
drop(guard);
let note = if overwritten {
"Summary replaced."
} else {
"Summary recorded."
};
Ok(ToolOutput {
summary: note.to_string(),
content: None,
})
}
}
pub(crate) fn mark_read_required_tool(
fs: ScopedFs,
ctx: Arc<Mutex<CompactWorkerContext>>,
) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(MarkParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("mark_read_required")
.description(MARK_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
fs: fs.clone(),
ctx: ctx.clone(),
});
(meta, tool)
})
}
pub(crate) fn add_reference_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(ReferenceParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("add_reference")
.description(REFERENCE_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(AddReferenceTool { ctx: ctx.clone() });
(meta, tool)
})
}
pub(crate) fn write_summary_tool(ctx: Arc<Mutex<CompactWorkerContext>>) -> ToolDefinition {
Arc::new(move || {
let schema = schemars::schema_for!(SummaryParams);
let schema_value = serde_json::to_value(schema).unwrap_or(serde_json::json!({}));
let meta = ToolMeta::new("write_summary")
.description(SUMMARY_DESCRIPTION)
.input_schema(schema_value);
let tool: Arc<dyn Tool> = Arc::new(WriteSummaryTool { ctx: ctx.clone() });
(meta, tool)
})
}
/// Interceptor that aborts the compact worker as soon as its cumulative
/// input-token count crosses `max_input_tokens`. Pairs with the
/// `on_usage` callback registered by `Pod::compact`, which is what
/// actually accumulates `input_so_far`.
pub(crate) struct CompactWorkerInterceptor {
pub input_so_far: Arc<AtomicU64>,
pub max_input_tokens: u64,
}
#[async_trait]
impl Interceptor for CompactWorkerInterceptor {
async fn pre_llm_request(&self, _context: &mut Vec<Item>) -> PreRequestAction {
if self.input_so_far.load(Ordering::Relaxed) > self.max_input_tokens {
return PreRequestAction::Cancel(format!(
"compact worker input exceeded {} tokens",
self.max_input_tokens
));
}
PreRequestAction::Continue
}
}
/// Crude bytes→tokens estimate; good enough for budget accounting.
fn estimate_tokens(bytes: usize) -> u64 {
(bytes as u64).div_ceil(4)
}
/// Return the slice of `text` covered by `offset` (line index) and
/// optional `limit` (line count), preserving the original newline
/// separation. Returns the whole file when both defaults apply.
pub(crate) fn slice_lines(text: &str, offset: usize, limit: Option<usize>) -> String {
let lines: Vec<&str> = text.lines().collect();
let start = offset.min(lines.len());
let end = limit
.map(|n| start.saturating_add(n).min(lines.len()))
.unwrap_or(lines.len());
lines[start..end].join("\n")
}
#[cfg(test)]
mod tests {
use super::*;
use manifest::Scope;
fn make_fs(tmp: &std::path::Path) -> ScopedFs {
let scope = Scope::writable(tmp.to_path_buf()).unwrap();
ScopedFs::new(scope, tmp.to_path_buf())
}
#[tokio::test]
async fn mark_read_required_records_and_deducts_budget() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("hello.txt");
std::fs::write(&path, "hello world\n").unwrap();
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(1_000)));
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
fs: make_fs(tmp.path()),
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
let out = tool.execute(&input).await.unwrap();
assert!(out.summary.starts_with("Marked"));
let guard = ctx.lock().unwrap();
assert_eq!(guard.read_required.len(), 1);
assert!(guard.auto_read_consumed > 0);
assert!(guard.auto_read_consumed <= 1_000);
}
#[tokio::test]
async fn mark_read_required_rejects_over_budget() {
let tmp = tempfile::TempDir::new().unwrap();
let path = tmp.path().join("big.txt");
std::fs::write(&path, "x".repeat(4_096)).unwrap(); // ≈1024 tokens
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(100)));
let tool: Arc<dyn Tool> = Arc::new(MarkReadRequiredTool {
fs: make_fs(tmp.path()),
ctx: ctx.clone(),
});
let input = serde_json::json!({ "file_path": path.to_str().unwrap() }).to_string();
let res = tool.execute(&input).await;
assert!(matches!(res, Err(ToolError::ExecutionFailed(_))));
let guard = ctx.lock().unwrap();
assert!(guard.read_required.is_empty());
assert_eq!(guard.auto_read_consumed, 0);
}
#[tokio::test]
async fn write_summary_overwrites_previous_call() {
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(0)));
let tool: Arc<dyn Tool> = Arc::new(WriteSummaryTool { ctx: ctx.clone() });
let first = serde_json::json!({ "text": "first" }).to_string();
let out1 = tool.execute(&first).await.unwrap();
assert!(out1.summary.contains("recorded"));
let second = serde_json::json!({ "text": "second" }).to_string();
let out2 = tool.execute(&second).await.unwrap();
assert!(out2.summary.contains("replaced"));
assert_eq!(ctx.lock().unwrap().summary.as_deref(), Some("second"));
}
#[tokio::test]
async fn add_reference_deduplicates() {
let ctx = Arc::new(Mutex::new(CompactWorkerContext::with_budget(0)));
let tool: Arc<dyn Tool> = Arc::new(AddReferenceTool { ctx: ctx.clone() });
let p = "/abs/path.rs";
let input = serde_json::json!({ "file_path": p }).to_string();
tool.execute(&input).await.unwrap();
tool.execute(&input).await.unwrap();
let guard = ctx.lock().unwrap();
assert_eq!(guard.references.len(), 1);
assert_eq!(guard.references[0], PathBuf::from(p));
}
#[test]
fn slice_lines_handles_offset_and_limit() {
let text = "a\nb\nc\nd";
assert_eq!(slice_lines(text, 0, None), "a\nb\nc\nd");
assert_eq!(slice_lines(text, 1, Some(2)), "b\nc");
assert_eq!(slice_lines(text, 10, None), "");
}
}
crates/pod/src/lib.rs
+1
View File
@@ -12,6 +12,7 @@ pub mod spawned_pod_registry;
mod agents_md;
mod compact_state;
mod compact_worker;
mod factory;
mod notification_buffer;
mod pod;
crates/pod/src/pod.rs
+350 -90
View File
@@ -47,18 +47,35 @@ impl Hook<PreLlmRequest> for UsageTrackingHook {
}
const SUMMARY_SYSTEM_PROMPT: &str = "\
You are a context compaction assistant. \
Summarise the conversation below into a structured summary. \
Preserve concrete details: file paths, function names, error messages, decisions made. \
Use the following format:\n\n\
## Original Task\n\
(the user's original request)\n\n\
## Completed Work\n\
- (what was done, with specifics)\n\n\
## Key Discoveries\n\
- (facts, constraints, errors found)\n\n\
## Current State\n\
- (files changed, remaining work)";
You are a context compaction assistant. Your job is to hand the next session a \
structured summary plus pointers to the files it actually needs.\n\n\
Tools you can call:\n\
- `read_file(file_path, offset?, limit?)` — inspect referenced files before deciding.\n\
- `mark_read_required(file_path, offset?, limit?)` — inject a file's contents into the \
next session as an auto-read system message. Counts against `auto_read_budget`.\n\
- `add_reference(file_path)` — record a file path the next session should know about \
without embedding its contents.\n\
- `write_summary(text)` — deliver the final structured summary. May be called multiple \
times; only the last call is kept.\n\n\
Always finish by calling `write_summary`. Produce the summary in this exact format:\n\n\
## Completed Tasks\n\
### (task name)\n\
- what was done (use concrete type / file / function names)\n\
- gotchas or facts that came up\n\n\
## Active Task\n\
### (task name)\n\
- goal\n\
- current state (what is done / not done)\n\
- next step\n\n\
## Key Decisions\n\
- (decision) — (reason)\n\n\
## User Directives\n\
- \"verbatim user line\" — only include directives whose wording the next session \
should not lose.\n\n\
## Current Work\n\
(2–3 lines on what was happening just before compaction).\n\n\
Keep code snippets and raw tool output OUT of the summary — that is what auto-read \
and references are for. Target 1000–2000 tokens.";
/// An independent agent execution unit.
///
@@ -405,9 +422,10 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Install the hook-based interceptor on the Worker if not already done.
///
/// When `compact_threshold` is configured in the manifest, wraps the
/// `HookInterceptor` in a [`CompactInterceptor`] and registers an
/// `on_usage` callback to track `input_tokens`.
/// When either compaction threshold (`compact_threshold` or
/// `compact_request_threshold`) is configured in the manifest, allocates
/// a shared [`CompactState`] and wires the interceptor to read current
/// occupancy through the `UsageRecord` timeline.
fn ensure_interceptor_installed(&mut self) {
if !self.interceptor_installed {
// Pre-LLM-request hook: record the item count at send time
@@ -420,41 +438,56 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
let builder = std::mem::take(&mut self.hook_builder);
let registry = Arc::new(builder.build());
let compact_threshold = self
let (post_run_threshold, request_threshold, retained) = self
.manifest
.compaction
.as_ref()
.and_then(|c| c.compact_threshold);
.map(|c| {
(
c.compact_threshold,
c.compact_request_threshold,
c.compact_retained_tokens,
)
})
.unwrap_or((None, None, manifest::defaults::COMPACT_RETAINED_TOKENS));
let tracker_for_usage = self.usage_tracker.clone();
self.worker_mut().on_usage(move |event| {
tracker_for_usage.record_usage(event);
});
let compact_state = if let Some(threshold) = compact_threshold {
let retained = self
.manifest
.compaction
.as_ref()
.map(|c| c.compact_retained_turns)
.unwrap_or(2);
let state = Arc::new(CompactState::new(threshold, retained));
let state_for_usage = state.clone();
self.worker_mut().on_usage(move |event| {
if let Some(tokens) = event.input_tokens {
state_for_usage.update_input_tokens(tokens);
let compact_state = if post_run_threshold.is_some() || request_threshold.is_some() {
if let (Some(post), Some(req)) = (post_run_threshold, request_threshold) {
if post > req {
warn!(
post_run_threshold = post,
request_threshold = req,
"compact_threshold > compact_request_threshold; \
proactive check will never fire before the safety net"
);
}
tracker_for_usage.record_usage(event);
});
}
let state = Arc::new(CompactState::new(
post_run_threshold,
request_threshold,
retained,
));
self.compact_state = Some(state.clone());
Some(state)
} else {
self.worker_mut().on_usage(move |event| {
tracker_for_usage.record_usage(event);
});
None
};
let interceptor =
PodInterceptor::new(registry, compact_state, self.pending_notifications.clone());
let usage_history_handle = compact_state
.as_ref()
.map(|_| self.usage_history.clone());
let interceptor = PodInterceptor::new(
registry,
compact_state,
usage_history_handle,
self.pending_notifications.clone(),
);
self.worker_mut().set_interceptor(interceptor);
self.interceptor_installed = true;
}
@@ -646,8 +679,8 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
let retained = self
.compact_state
.as_ref()
.map(|s| s.retained_turns())
.unwrap_or(2);
.map(|s| s.retained_tokens())
.unwrap_or(manifest::defaults::COMPACT_RETAINED_TOKENS);
match self.compact(retained).await {
Ok(new_session_id) => {
@@ -681,11 +714,15 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// Best-effort: failures are logged but do not propagate.
pub async fn try_post_run_compact(&mut self) -> Result<(), PodError> {
let state = match self.compact_state.as_ref() {
Some(s) if !s.is_disabled() && s.exceeds_post_run() && !s.just_compacted() => s.clone(),
Some(s) if !s.is_disabled() && !s.just_compacted() => s.clone(),
_ => return Ok(()),
};
let current_tokens = self.total_tokens().tokens;
if !state.exceeds_post_run(current_tokens) {
return Ok(());
}
let retained = state.retained_turns();
let retained = state.retained_tokens();
match self.compact(retained).await {
Ok(new_session_id) => {
info!(
@@ -785,60 +822,196 @@ impl<C: LlmClient, St: Store> Pod<C, St> {
/// - a clone of the main LlmClient via `clone_boxed()`.
///
/// Returns the new session ID.
pub async fn compact(&mut self, retained_turns: usize) -> Result<SessionId, PodError> {
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SessionId, PodError> {
use std::sync::atomic::{AtomicU64, Ordering};
use crate::compact_worker::{
CompactWorkerContext, CompactWorkerInterceptor, add_reference_tool,
mark_read_required_tool, slice_lines, write_summary_tool,
};
// Decide the cut point by projecting the UsageRecord timeline onto
// the current history: keep the tail whose estimated token count is
// within `retained_tokens`. Item-granular, turn boundaries ignored.
let cut = self.split_for_retained(retained_tokens);
let worker = self.worker.as_ref().expect("worker taken during run");
let history = worker.history();
// Identify turn boundaries (user message positions).
let turn_starts: Vec<usize> = history
.iter()
.enumerate()
.filter(|(_, item)| item.is_user_message())
.map(|(i, _)| i)
.collect();
// Items to retain: everything from `retained_turns` turns ago onward.
let retain_from = if turn_starts.len() > retained_turns {
turn_starts[turn_starts.len() - retained_turns]
} else {
0
};
let retain_from = cut.index.min(history.len());
let retained_items = history[retain_from..].to_vec();
let items_to_summarise = &history[..retain_from];
let items_to_summarise = history[..retain_from].to_vec();
// Build summary prompt.
let summary_prompt = build_summary_prompt(items_to_summarise);
// Compaction-related knobs. Fall through to manifest defaults when
// `[compaction]` is omitted entirely.
let (auto_read_budget, compact_worker_max_input_tokens) = self
.manifest
.compaction
.as_ref()
.map(|c| (c.compact_auto_read_budget, c.compact_worker_max_input_tokens))
.unwrap_or((
manifest::defaults::COMPACT_AUTO_READ_BUDGET,
manifest::defaults::COMPACT_WORKER_MAX_INPUT_TOKENS,
));
// Create a disposable summary Worker.
// Default references: the N most-recently-touched files in the
// session, surfaced so the compact worker can inspect them and
// decide which (if any) the next session needs.
let default_refs: Vec<PathBuf> = self
.tracker
.as_ref()
.map(|t| t.recent_files(manifest::defaults::COMPACT_DEFAULT_REFERENCE_COUNT))
.unwrap_or_default();
// Input text fed to the compact worker. Includes the default
// references and the (pruned) conversation text.
let summary_input = build_summary_input(&items_to_summarise, &default_refs);
// Worker-side state collected by the compact worker's tool calls.
let ctx = Arc::new(std::sync::Mutex::new(CompactWorkerContext::with_budget(
auto_read_budget,
)));
// Build an independent compact worker. Scope and pwd are shared
// with the main Pod (reads go through the same policy) but the
// Tracker is fresh — compact-time reads must not pollute the
// main session's recency list, which feeds `default_refs` above.
let scoped_fs = tools::ScopedFs::new(self.scope.clone(), self.pwd.clone());
let summary_tracker = tools::Tracker::new();
let summary_client: Box<dyn LlmClient> = self.build_compactor_client()?;
let mut summary_worker = Worker::new(summary_client)
.system_prompt(SUMMARY_SYSTEM_PROMPT)
.temperature(0.0);
summary_worker.set_max_tokens(2048);
summary_worker.set_max_tokens(4096);
// Cumulative input-token meter + interceptor. The meter is bumped
// from the on_usage callback and read on every pre_llm_request.
let input_so_far = Arc::new(AtomicU64::new(0));
{
let acc = input_so_far.clone();
summary_worker.on_usage(move |event| {
if let Some(tokens) = event.input_tokens {
acc.fetch_add(tokens, Ordering::Relaxed);
}
});
}
summary_worker.set_interceptor(CompactWorkerInterceptor {
input_so_far: input_so_far.clone(),
max_input_tokens: compact_worker_max_input_tokens,
});
// Tools: read_file (shared scope, fresh tracker) + the three
// compact-specific tools that populate `ctx`.
summary_worker.register_tool(tools::read_tool(scoped_fs.clone(), summary_tracker));
summary_worker
.register_tool(mark_read_required_tool(scoped_fs.clone(), ctx.clone()));
summary_worker.register_tool(add_reference_tool(ctx.clone()));
summary_worker.register_tool(write_summary_tool(ctx.clone()));
let out = summary_worker
.run(summary_prompt)
.run(summary_input)
.await
.map_err(PodError::Worker)?;
let summary_text = out
.worker
.history()
.iter()
.filter_map(|item| {
if item.is_assistant_message() {
item.as_text().map(String::from)
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n");
let mut locked_worker = out.worker;
// Build new history: [summary as user message, ...retained].
let mut new_history = Vec::with_capacity(retained_items.len() + 1);
// Guard: nudge the worker once more if the expected outputs
// (summary, and any auto-read nominations when default refs
// existed) were not produced on the first pass. `write_summary`
// is idempotent-by-overwrite so a second call is safe.
let nudge = {
let snapshot = ctx.lock().expect("compact ctx poisoned").clone();
if snapshot.summary.is_none() {
Some(
"You have not called `write_summary` yet. Deliver the structured \
summary now (Completed Tasks / Active Task / Key Decisions / \
User Directives / Current Work) and nominate any files the next \
session needs with `mark_read_required`."
.to_string(),
)
} else if snapshot.read_required.is_empty() && !default_refs.is_empty() {
Some(
"Summary received. If any of the referenced files are required \
for the next session to continue the task, call \
`mark_read_required` on them now. Otherwise reply briefly to \
close out."
.to_string(),
)
} else {
None
}
};
if let Some(prompt) = nudge {
let _ = locked_worker
.run(prompt)
.await
.map_err(PodError::Worker)?;
}
let final_ctx = ctx.lock().expect("compact ctx poisoned").clone();
let summary_text = final_ctx
.summary
.clone()
.ok_or(PodError::CompactSummaryMissing)?;
// Re-read each auto-read target through ScopedFs and render the
// requested slice. Errors are logged and skipped rather than
// aborting compaction — a missing / moved file should not fail
// the whole compact.
let mut auto_read_messages = Vec::new();
for req in &final_ctx.read_required {
match scoped_fs.read_bytes(&req.path) {
Ok(bytes) => {
let text = String::from_utf8_lossy(&bytes).into_owned();
let body = slice_lines(&text, req.offset.unwrap_or(0), req.limit);
let range = match (req.offset, req.limit) {
(None, None) => String::new(),
(Some(off), None) => format!(":{}-", off + 1),
(None, Some(lim)) => format!(":1-{lim}"),
(Some(off), Some(lim)) => {
format!(":{}-{}", off + 1, off.saturating_add(lim))
}
};
auto_read_messages.push(Item::system_message(format!(
"[Auto-read file: {}{range}]\n{body}",
req.path.display()
)));
}
Err(e) => {
warn!(
path = %req.path.display(),
error = %e,
"auto-read target could not be read; skipping",
);
}
}
}
// Reference list as a single system message; omitted when empty.
let reference_message = (!final_ctx.references.is_empty()).then(|| {
let list = final_ctx
.references
.iter()
.map(|p| format!("- {}", p.display()))
.collect::<Vec<_>>()
.join("\n");
Item::system_message(format!(
"[Referenced files — read before compaction, contents not included]\n\
{list}\n\
Use read_file to access current contents if needed."
))
});
// Build new history: [summary, ...auto-read, references, ...retained].
let mut new_history = Vec::with_capacity(
1 + auto_read_messages.len() + reference_message.is_some() as usize
+ retained_items.len(),
);
new_history.push(Item::system_message(format!(
"[Compacted context summary]\n\n{summary_text}"
)));
new_history.extend(auto_read_messages);
if let Some(msg) = reference_message {
new_history.push(msg);
}
new_history.extend(retained_items);
// Persist as a new compacted session.
@@ -1095,7 +1268,45 @@ impl From<WorkerResult> for PodRunResult {
}
}
/// Build the compact worker's input: default-reference instructions,
/// the list of recently-touched files, and the pruned conversation
/// produced by [`build_summary_prompt`].
fn build_summary_input(items: &[Item], default_refs: &[PathBuf]) -> String {
let mut out = String::new();
out.push_str(
"Summarise the conversation below into a structured summary and nominate \
files the next session needs.\n\n",
);
if !default_refs.is_empty() {
out.push_str(
"These files were touched recently in this session. Use `read_file` \
on them as needed, then call `mark_read_required` for any whose \
contents the next session must have, and `add_reference` for files \
it should know about by name only.\n\n## Referenced files\n",
);
for p in default_refs {
out.push_str("- ");
out.push_str(&p.display().to_string());
out.push('\n');
}
out.push('\n');
}
out.push_str("## Conversation\n");
out.push_str(&build_summary_prompt(items));
out.push_str(
"\n\nWhen you are done, call `write_summary` with the final 5-section text.",
);
out
}
/// Format conversation items into a text prompt for the summary Worker.
///
/// The summary should capture decisions and user intent, not recreate code.
/// File contents and tool IO belong in auto-read / references, not in the
/// summary input. So this strips:
/// - `ToolCall.arguments` (keep only the tool name)
/// - `ToolResult.content` (keep only the summary line)
/// - `Reasoning` entirely (intermediate thought, superseded by decisions)
fn build_summary_prompt(items: &[Item]) -> String {
let mut lines = Vec::new();
for item in items {
@@ -1113,20 +1324,13 @@ fn build_summary_prompt(items: &[Item]) -> String {
.join("");
lines.push(format!("[{role_label}] {text}"));
}
Item::ToolCall {
name, arguments, ..
} => {
lines.push(format!("[ToolCall] {name}({arguments})"));
Item::ToolCall { name, .. } => {
lines.push(format!("[ToolCall] {name}"));
}
Item::ToolResult {
summary, content, ..
} => match content {
Some(c) => lines.push(format!("[ToolResult] {summary}\n{c}")),
None => lines.push(format!("[ToolResult] {summary}")),
},
Item::Reasoning { text, .. } => {
lines.push(format!("[Reasoning] {text}"));
Item::ToolResult { summary, .. } => {
lines.push(format!("[ToolResult] {summary}"));
}
Item::Reasoning { .. } => {}
}
}
lines.join("\n\n")
@@ -1166,6 +1370,9 @@ pub enum PodError {
#[error("compaction thrash: context still exceeds threshold immediately after compact")]
CompactThrash,
#[error("compact worker did not produce a summary (write_summary was never called)")]
CompactSummaryMissing,
#[error("invalid system prompt template: {source}")]
InvalidSystemPromptTemplate {
#[source]
@@ -1196,3 +1403,56 @@ fn current_pwd() -> Result<PathBuf, PodError> {
source,
})
}
#[cfg(test)]
mod build_summary_prompt_tests {
use super::*;
#[test]
fn strips_tool_call_arguments() {
let items = vec![Item::tool_call_json(
"call-1",
"read_file",
serde_json::json!({ "path": "src/main.rs" }),
)];
let prompt = build_summary_prompt(&items);
assert_eq!(prompt, "[ToolCall] read_file");
assert!(!prompt.contains("src/main.rs"));
}
#[test]
fn strips_tool_result_content() {
let items = vec![Item::tool_result_with_content(
"call-1",
"read 3 lines",
"fn main() { println!(\"hello\"); }",
)];
let prompt = build_summary_prompt(&items);
assert_eq!(prompt, "[ToolResult] read 3 lines");
assert!(!prompt.contains("println"));
}
#[test]
fn drops_reasoning_entirely() {
let items = vec![
Item::user_message("hi"),
Item::reasoning("internal deliberation"),
Item::assistant_message("hello"),
];
let prompt = build_summary_prompt(&items);
assert!(prompt.contains("[User] hi"));
assert!(prompt.contains("[Assistant] hello"));
assert!(!prompt.contains("Reasoning"));
assert!(!prompt.contains("deliberation"));
}
#[test]
fn keeps_user_and_assistant_messages() {
let items = vec![
Item::user_message("fix the bug"),
Item::assistant_message("done"),
];
let prompt = build_summary_prompt(&items);
assert_eq!(prompt, "[User] fix the bug\n\n[Assistant] done");
}
}
crates/pod/src/pod_interceptor.rs
+101 -27
View File
@@ -7,8 +7,8 @@
//! read-only summary information and only return control-flow
//! decisions (continue / skip / abort / pause).
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use llm_worker::Item;
@@ -17,6 +17,7 @@ use llm_worker::interceptor::{
ToolResultInfo, TurnEndAction,
};
use llm_worker::tool::ToolOutput;
use session_store::UsageRecord;
use tracing::info;
use crate::compact_state::CompactState;
@@ -25,6 +26,7 @@ use crate::hook::{
TurnEndInfo,
};
use crate::notification_buffer::{NotificationBuffer, format_notification};
use crate::token_counter::total_tokens_impl;
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`.
const FINAL_TEXT_PREVIEW_LIMIT: usize = 512;
@@ -32,6 +34,10 @@ const FINAL_TEXT_PREVIEW_LIMIT: usize = 512;
pub(crate) struct PodInterceptor {
registry: Arc<HookRegistry>,
compact_state: Option<Arc<CompactState>>,
/// Shared view of the cumulative UsageRecord timeline. Used with the
/// per-request `context` to estimate current occupancy for threshold
/// checks. `None` when compaction is disabled (both thresholds unset).
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
/// Pending-notification buffer drained into the per-request
/// context at the head of `pre_llm_request`.
pending_notifications: NotificationBuffer,
@@ -45,11 +51,13 @@ impl PodInterceptor {
pub(crate) fn new(
registry: Arc<HookRegistry>,
compact_state: Option<Arc<CompactState>>,
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
pending_notifications: NotificationBuffer,
) -> Self {
Self {
registry,
compact_state,
usage_history,
pending_notifications,
next_turn_index: AtomicUsize::new(0),
tool_calls_this_turn: AtomicUsize::new(0),
@@ -61,6 +69,15 @@ impl PodInterceptor {
.load(Ordering::Relaxed)
.saturating_sub(1)
}
/// Estimate current input-token occupancy for `context`, projected
/// through the shared UsageRecord timeline. Returns `None` when
/// `usage_history` is not attached (compaction fully disabled).
fn estimated_tokens(&self, context: &[Item]) -> Option<u64> {
let handle = self.usage_history.as_ref()?;
let records = handle.lock().expect("usage_history poisoned").clone();
Some(total_tokens_impl(context, &records).tokens)
}
}
#[async_trait]
@@ -83,15 +100,20 @@ impl Interceptor for PodInterceptor {
}
async fn pre_llm_request(&self, context: &mut Vec<Item>) -> PreRequestAction {
// Internal mechanism: between-turns compaction trigger.
let current_tokens = self.estimated_tokens(context);
// Internal mechanism: between-requests compaction trigger (safety net).
if let Some(state) = self.compact_state.as_ref() {
if !state.is_disabled() && state.exceeds_turn() {
info!(
input_tokens = state.last_input_tokens(),
threshold = state.turn_threshold(),
"Between-turns compaction threshold exceeded, yielding"
);
return PreRequestAction::Yield;
if !state.is_disabled() {
let current = current_tokens.unwrap_or(0);
if state.exceeds_request(current) {
info!(
input_tokens = current,
threshold = state.request_threshold().unwrap_or(0),
"Between-requests compaction threshold exceeded, yielding"
);
return PreRequestAction::Yield;
}
}
}
@@ -105,7 +127,7 @@ impl Interceptor for PodInterceptor {
let info = PreRequestInfo {
item_count: context.len(),
estimated_tokens: self.compact_state.as_ref().map(|s| s.last_input_tokens()),
estimated_tokens: current_tokens,
turn_index: self.current_turn_index(),
tool_calls_this_turn: self.tool_calls_this_turn.load(Ordering::Relaxed),
};
@@ -232,16 +254,35 @@ mod tests {
Arc::new(builder.build())
}
/// Build a usage_history handle with a single record pinned at the
/// current `context_len` so that `total_tokens_impl` returns exactly
/// `tokens` (Measured, no interpolation or byte-based fallback).
fn usage_handle_with(context_len: usize, tokens: u64) -> Arc<Mutex<Vec<UsageRecord>>> {
Arc::new(Mutex::new(vec![UsageRecord {
history_len: context_len,
input_total_tokens: tokens,
cache_read_tokens: 0,
cache_write_tokens: 0,
output_tokens: 0,
}]))
}
#[tokio::test]
async fn pre_llm_request_yields_and_skips_hooks_when_compact_threshold_exceeded() {
async fn pre_llm_request_yields_and_skips_hooks_when_request_threshold_exceeded() {
let count = Arc::new(AtomicUsize::new(0));
let registry = registry_with_pre_llm_hook(count.clone());
let state = Arc::new(CompactState::new(100, 2));
state.update_input_tokens(200); // exceeds turn threshold
let state = Arc::new(CompactState::new(None, Some(100), 2));
let ctx_items = vec![Item::user_message("hi")];
let history = usage_handle_with(ctx_items.len(), 200);
let interceptor = PodInterceptor::new(registry, Some(state), NotificationBuffer::new());
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
let interceptor = PodInterceptor::new(
registry,
Some(state),
Some(history),
NotificationBuffer::new(),
);
let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await;
assert!(matches!(action, PreRequestAction::Yield));
@@ -254,11 +295,41 @@ mod tests {
let count = Arc::new(AtomicUsize::new(0));
let registry = registry_with_pre_llm_hook(count.clone());
let state = Arc::new(CompactState::new(100, 2));
// last_input_tokens stays at 0, well below threshold.
let state = Arc::new(CompactState::new(None, Some(100), 2));
let ctx_items = vec![Item::user_message("hi")];
let history = usage_handle_with(ctx_items.len(), 50);
let interceptor = PodInterceptor::new(registry, Some(state), NotificationBuffer::new());
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
let interceptor = PodInterceptor::new(
registry,
Some(state),
Some(history),
NotificationBuffer::new(),
);
let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await;
assert!(matches!(action, PreRequestAction::Continue));
assert_eq!(count.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn pre_llm_request_does_not_yield_when_only_post_run_threshold_set() {
// request_threshold = None → safety-net check is inert inside the turn
// even if current occupancy is huge. Post-run check runs elsewhere.
let count = Arc::new(AtomicUsize::new(0));
let registry = registry_with_pre_llm_hook(count.clone());
let state = Arc::new(CompactState::new(Some(100), None, 2));
let ctx_items = vec![Item::user_message("hi")];
let history = usage_handle_with(ctx_items.len(), 10_000);
let interceptor = PodInterceptor::new(
registry,
Some(state),
Some(history),
NotificationBuffer::new(),
);
let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await;
assert!(matches!(action, PreRequestAction::Continue));
@@ -270,7 +341,7 @@ mod tests {
let count = Arc::new(AtomicUsize::new(0));
let registry = registry_with_pre_llm_hook(count.clone());
let interceptor = PodInterceptor::new(registry, None, NotificationBuffer::new());
let interceptor = PodInterceptor::new(registry, None, None, NotificationBuffer::new());
let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await;
@@ -295,7 +366,7 @@ mod tests {
buffer.push("first".into());
buffer.push("second".into());
let interceptor = PodInterceptor::new(registry, None, buffer.clone());
let interceptor = PodInterceptor::new(registry, None, None, buffer.clone());
let mut ctx: Vec<Item> = vec![Item::user_message("hi")];
let action = interceptor.pre_llm_request(&mut ctx).await;
@@ -320,15 +391,18 @@ mod tests {
let buffer = NotificationBuffer::new();
buffer.push("msg".into());
let state = Arc::new(CompactState::new(100, 2));
state.update_input_tokens(200);
let state = Arc::new(CompactState::new(None, Some(100), 2));
let ctx_items = vec![Item::user_message("hi")];
let history = usage_handle_with(ctx_items.len(), 200);
let interceptor = PodInterceptor::new(registry, Some(state), buffer.clone());
let mut ctx: Vec<Item> = Vec::new();
let interceptor =
PodInterceptor::new(registry, Some(state), Some(history), buffer.clone());
let mut ctx = ctx_items;
let action = interceptor.pre_llm_request(&mut ctx).await;
assert!(matches!(action, PreRequestAction::Yield));
assert!(ctx.is_empty());
// Notifications were not drained (still held for post-compact resume).
assert_eq!(ctx.len(), 1);
assert_eq!(buffer.len(), 1);
}
@@ -341,7 +415,7 @@ mod tests {
builder.add_pre_llm_request(CountingHook(second_count.clone()));
let registry = Arc::new(builder.build());
let interceptor = PodInterceptor::new(registry, None, NotificationBuffer::new());
let interceptor = PodInterceptor::new(registry, None, None, NotificationBuffer::new());
let mut ctx: Vec<Item> = Vec::new();
let action = interceptor.pre_llm_request(&mut ctx).await;
crates/pod/src/token_counter.rs
+1 -1
View File
@@ -169,7 +169,7 @@ fn tokens_at(
}
}
fn total_tokens_impl(history: &[Item], records: &[UsageRecord]) -> TokenEstimate {
pub(crate) fn total_tokens_impl(history: &[Item], records: &[UsageRecord]) -> TokenEstimate {
let prefix = prefix_bytes(history);
tokens_at(history, records, history.len(), &prefix)
}
crates/pod/tests/system_prompt_template_test.rs
+25 -10
View File
@@ -61,6 +61,19 @@ fn single_text_events(text: &str) -> Vec<LlmEvent> {
]
}
/// Emit a single `write_summary(text=...)` tool call as one LLM response.
fn write_summary_tool_use_events(call_id: &str, text: &str) -> Vec<LlmEvent> {
let input = serde_json::json!({ "text": text }).to_string();
vec![
LlmEvent::tool_use_start(0, call_id, "write_summary"),
LlmEvent::tool_input_delta(0, input),
LlmEvent::tool_use_stop(0),
LlmEvent::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
const MINIMAL_MANIFEST_TOML: &str = r#"
[pod]
name = "test-pod"
@@ -233,10 +246,11 @@ async fn agents_md_absent_omits_trailing_section() {
#[tokio::test]
async fn agents_md_not_reread_after_compact() {
let client = MockClient::new(vec![
single_text_events("a"),
single_text_events("b"),
single_text_events("summary"),
single_text_events("c"),
single_text_events("a"), // pod.run("first")
single_text_events("b"), // pod.run("second")
write_summary_tool_use_events("call-1", "compacted summary"), // compact worker: tool_use
single_text_events("done"), // compact worker: close
single_text_events("c"), // pod.run("third")
]);
let (mut pod, pwd) = make_pod_with_body("BODY", client).await.unwrap();
let agents_path = pwd.join("AGENTS.md");
@@ -250,7 +264,7 @@ async fn agents_md_not_reread_after_compact() {
// Mutate the file after the first turn — must not affect the cached
// system prompt either on a subsequent turn or across compaction.
std::fs::write(&agents_path, "mutated").unwrap();
pod.compact(1).await.unwrap();
pod.compact(0).await.unwrap();
let after_compact = pod.worker().get_system_prompt().unwrap().to_string();
assert!(after_compact.contains("original"));
assert!(!after_compact.contains("mutated"));
@@ -264,10 +278,11 @@ async fn agents_md_not_reread_after_compact() {
#[tokio::test]
async fn compact_preserves_system_prompt() {
let client = MockClient::new(vec![
single_text_events("a"),
single_text_events("b"),
single_text_events("summary"),
single_text_events("c"),
single_text_events("a"), // pod.run("first")
single_text_events("b"), // pod.run("second")
write_summary_tool_use_events("call-1", "compacted summary"), // compact worker: tool_use
single_text_events("done"), // compact worker: close
single_text_events("c"), // pod.run("third")
]);
let (mut pod, _pwd) = make_pod_with_body("SP cwd={{ cwd }}", client)
.await
@@ -277,7 +292,7 @@ async fn compact_preserves_system_prompt() {
let before = pod.worker().get_system_prompt().unwrap().to_string();
pod.run("second").await.unwrap();
pod.compact(1).await.unwrap();
pod.compact(0).await.unwrap();
let after = pod.worker().get_system_prompt().unwrap().to_string();
assert_eq!(before, after);