docs(comments): rewrite comments across all crates to the guidelines

Sweep every first-party crate source (1956 .rs files) to the project comment
guidelines: delete redundant restatements, decorative banners, change
narration, and end-of-line comments; keep and tighten the crucial ones
(invariants, bug rationale, SAFETY blocks, ported-source attribution).

No functional code changed. Every edit is proven comment-only against the
prior tree by a comment-stripping lexer (string/char/raw-string aware) plus a
separate doctest-fence check. Where removing a comment made rustfmt or clippy
want to re-lay-out adjacent code, the minimal triggering comment is restored so
code tokens stay byte-identical.

Gates green: cargo fmt --all --check (0 diffs), cargo check and cargo clippy
--workspace --all-targets (0 warnings).

Adds scripts/check_codegen_comment_guidelines.py — the enforcement gate for
these guidelines (flags banners, end-of-line comments, change narration, and
commented-out code).
This commit is contained in:
2026-07-23 16:55:39 -04:00
parent ff0fb56c67
commit a02b555e66
1458 changed files with 10729 additions and 21750 deletions
@@ -1,7 +1,6 @@
//! Compacted-history assembly (kigi's rebuild structure, generic).
//!
//! Moved from `kigi-chat-state::compaction_utils::build_compacted_history` and
//! made generic over a write-side item factory so any harness can assemble
//! Generic over a write-side item factory so any harness can assemble
//! the canonical post-compaction history:
//!
//! ```text
@@ -25,7 +24,6 @@ use super::summary::{format_compact_summary_content, wrap_user_query};
/// - Providing the `user_message_prefix` (e.g. `<user_info>` block).
/// - Extracting `last_user_query` / `recent_messages` from its own state.
pub struct CompactedHistoryParts<T> {
/// The original system message from the conversation.
pub system_message: T,
/// The user-info / project-layout prefix (not wrapped in `<user_query>`).
pub user_message_prefix: String,
@@ -36,7 +34,6 @@ pub struct CompactedHistoryParts<T> {
pub last_user_query: Option<String>,
/// Messages retained verbatim from after the last real user turn.
pub recent_messages: Vec<T>,
/// The LLM-generated compaction summary text.
pub compaction_summary: String,
/// An optional pre-rendered `<system-reminder>` block to append after the
/// summary. `None` means no state reminder is appended.
@@ -75,7 +72,6 @@ pub fn assemble_compacted_history<T: CompactionItemFactory>(
compacted.push(T::new_project_instructions(reminder.clone()));
}
// Last user query wrapped in <user_query> tags for consistency.
if let Some(ref last_query) = parts.last_user_query {
compacted.push(T::new_user(wrap_user_query(last_query.as_str())));
}
@@ -88,7 +84,6 @@ pub fn assemble_compacted_history<T: CompactionItemFactory>(
}
let summary_item = T::new_user_meta(formatted_summary);
// Recent messages come first, then the summary.
for msg in parts.recent_messages {
compacted.push(msg);
}
@@ -8,12 +8,12 @@
//! build prompt → sample (retry + classify) → clean → assemble
//! ```
//!
//! Per-harness concerns stay in the product host (for example `kigi-shell`): the triggers, the
//! conversation *gathering / sanitization* that produces `llm_turns`, the
//! verbatim→fitted→lossy input ladder, the live LLM transport (the
//! [`CompactionSampler`] impl), persistence/replay, and the rendering of
//! `system_reminder`. This function takes those as inputs and returns the
//! rebuilt history; it never commits or persists.
//! Per-harness concerns stay in the product host (for example `kigi-shell`):
//! the triggers, the conversation gathering/sanitization that produces
//! `llm_turns`, the verbatim→fitted→lossy input ladder, the live LLM transport
//! (the [`CompactionSampler`] impl), persistence/replay, and the rendering of
//! `system_reminder`. This module takes those as inputs and returns the rebuilt
//! history; it never commits or persists.
use std::time::{Duration, Instant};
@@ -30,37 +30,33 @@ use super::prompt::build_summary_prompt;
use super::sample::{SampleRetryError, SampledSummary, sample_summary_with_retries};
/// Everything the assembler needs that the harness extracts from its own
/// state (separate from the conversation that gets summarized).
/// state (separate from the conversation that gets summarized). Every field is
/// carried through verbatim; nothing here is re-derived.
pub struct FullReplaceContext<T> {
/// The original system message, carried over verbatim.
pub system_message: T,
/// The user-info / project-layout prefix (no `<user_query>` tags).
pub user_message_prefix: String,
/// Pre-rendered AGENTS.md block to re-inject, if any.
pub agents_md_reminder: Option<String>,
/// The last real user query (raw), kept verbatim post-compaction.
/// Raw, not wrapped in `<user_query>` tags yet.
pub last_user_query: Option<String>,
/// Working tail retained verbatim (tool/subagent results from the current
/// turn). kigi keeps this; pass empty to drop it.
/// Working tail (tool/subagent results from the current turn). Pass empty
/// to drop it.
pub recent_messages: Vec<T>,
/// Pre-rendered `<system-reminder>` (edited files, running tasks,
/// subagents, MCP, …). The harness builds this; we only carry it.
/// subagents, MCP, …).
pub system_reminder: Option<String>,
/// Optional transcript-pointer block appended to the summary.
/// Transcript-pointer block appended to the summary.
pub transcript_hint: Option<String>,
}
/// Outcome of a failed full-replace pass.
#[derive(Debug)]
pub enum FullReplaceError {
/// No turns were supplied to summarize.
NothingToCompact,
/// The model returned no usable summary text after all attempts.
EmptyResponse,
/// The sampler failed deterministically (re-sending can't help), or all
/// transient retries were exhausted.
Sampler {
/// The rendered upstream error.
message: String,
/// Whether re-sending the *same* input cannot help. The product host
/// uses this to decide whether to suppress auto-compaction.
@@ -84,46 +80,40 @@ impl std::fmt::Display for FullReplaceError {
impl std::error::Error for FullReplaceError {}
/// A successful full-replace pass.
pub struct FullReplaceOutput<T> {
/// The rebuilt, compacted history (`[SP, UP', AGENTS_MD?, UQ_last?,
/// recent…, summary, reminder?]`).
/// Shape: `[SP, UP', AGENTS_MD?, UQ_last?, recent…, summary, reminder?]`.
pub history: Vec<T>,
/// The **raw** model summary (pre-clean), so the product host can persist
/// it (request artifact, compaction segment) exactly as the model emitted
/// it. The cleaned form is already embedded in `history` by the assembler.
pub summary: String,
/// Total sample attempts made (first try + retries).
/// First try + retries.
pub attempts: u32,
}
/// A successful full-replace **sampling** pass (summary only, no assembly).
///
/// Returned by [`sample_full_replace_summary`] for harnesses (kigi's
/// shell) that drive the input ladder and assemble the history themselves —
/// they build the assembly inputs (state-context system-reminder, AGENTS.md,
/// plan-mode) *after* the LLM call, so they cannot use the bundled
/// Summary-only result of [`sample_full_replace_summary`], for harnesses
/// (kigi's shell) that drive the input ladder and assemble the history
/// themselves — they build the assembly inputs (state-context system-reminder,
/// AGENTS.md, plan-mode) *after* the LLM call, so they cannot use the bundled
/// [`apply_full_replace_compaction`].
pub struct FullReplaceSummary {
/// The **raw** model summary (pre-clean).
pub summary: String,
/// Total sample attempts made (first try + retries).
/// First try + retries.
pub attempts: u32,
}
/// Run kigi's full-replace compaction pass and return the rebuilt
/// history. Pure orchestration: no triggers, no persistence, no commit.
///
/// - `llm_turns` — the (harness-prepared/sanitized) conversation the model
/// summarizes. Empty ⇒ [`FullReplaceError::NothingToCompact`].
/// - `user_context` — optional `/compact <text>` context spliced into the prompt.
/// - `ctx` — the assembly inputs the harness extracted from its state.
/// - `observer` — per-attempt + terminal telemetry seam (pass `&()` for none).
/// Empty `llm_turns` ⇒ [`FullReplaceError::NothingToCompact`]. `user_context`
/// is the optional `/compact <text>` context spliced into the prompt; pass
/// `&()` as `observer` for no telemetry.
///
/// The **input ladder** (verbatim → fitted → lossy) stays in the product host: on a
/// context-length overflow this returns
/// [`FullReplaceError::Sampler`] with `context_overflow = true`, and the
/// harness rebuilds a smaller input and calls this pass again.
/// The **input ladder** (verbatim → fitted → lossy) stays in the product host:
/// on a context-length overflow this returns [`FullReplaceError::Sampler`] with
/// `context_overflow = true`, and the harness rebuilds a smaller input and
/// calls this pass again.
pub async fn apply_full_replace_compaction<T, S, O>(
sampler: &S,
llm_turns: &[T],
@@ -147,9 +137,8 @@ where
"[FullReplaceCompaction] sampled summary; assembling history"
);
// Clean (inside the assembler via `format_compact_summary_content`) and
// rebuild the compacted history. `compaction_summary` is the raw model
// output; the assembler strips scratchpad / control tokens.
// `compaction_summary` goes in raw; the assembler strips scratchpad /
// control tokens via `format_compact_summary_content`.
let parts = CompactedHistoryParts {
system_message: ctx.system_message,
user_message_prefix: ctx.user_message_prefix,
@@ -170,12 +159,8 @@ where
/// Run only the **sampling** half of the full-replace pass: build the prompt,
/// sample with bounded retries (transient + degenerate), classify failures,
/// and report every attempt through `observer`. Returns the raw summary; the
/// caller assembles the history (and owns the input ladder).
///
/// This is the seam kigi's shell uses: it drives the verbatim → fitted →
/// lossy input ladder around this call (stepping on a
/// [`FullReplaceError::Sampler`] with `context_overflow = true`) and assembles
/// the compacted history afterward from inputs it gathers post-sampling.
/// caller assembles the history and owns the input ladder, stepping it on a
/// [`FullReplaceError::Sampler`] with `context_overflow = true`.
pub async fn sample_full_replace_summary<T, S, O>(
sampler: &S,
llm_turns: &[T],
@@ -365,10 +350,7 @@ mod tests {
)
}
/// Golden end-to-end test: a realistic conversation + a mock sampler that
/// returns a structured summary must produce kigi's exact compacted
/// history shape, with the LLM output cleaned and the agent-state reminder
/// carried through as the final item.
/// Golden test pinning kigi's exact compacted history shape end to end.
#[tokio::test]
async fn full_replace_produces_kigi_history_shape() {
let llm_turns = vec![
@@ -408,7 +390,7 @@ mod tests {
MockItem::Tail("tool: read_file(auth.rs) -> ...".into())
);
// Summary carrier: cleaned (no <analysis>/<summary> tags), with preamble.
// Summary carrier: cleaned, with preamble.
let MockItem::UserMeta(summary) = &out[5] else {
panic!("expected UserMeta summary at [5], got {:?}", out[5]);
};
@@ -421,7 +403,6 @@ mod tests {
assert!(!summary.contains("<summary>"), "live tag leaked: {summary}");
assert!(!summary.contains("thinking about it"));
// Agent-state reminder carried through verbatim as the final item.
assert_eq!(
out[6],
MockItem::SystemReminder(
@@ -553,8 +534,6 @@ mod tests {
assert_eq!(sampler.call_count(), 1, "overflow must not retry");
}
/// The observer sees one terminal `on_success` and the right per-attempt
/// outcomes (a degenerate retry then a success).
#[tokio::test]
async fn observer_receives_attempt_and_success_callbacks() {
use std::sync::Mutex;
@@ -1,32 +1,26 @@
//! kigi compaction configuration.
//!
//! Holds the [`FullReplaceConfig`] tunables struct (mirroring
//! Tunables and defaults only; trigger *wiring* (pre-sampling checks,
//! preflight overflow, model-switch, suppression) stays per-host. Mirrors
//! [`IntraCompactionConfig`](crate::intra_compaction::IntraCompactionConfig) /
//! [`InterCompactionConfig`](crate::inter_compaction::InterCompactionConfig),
//! which also live in their module's `config.rs`) plus the shared default
//! values. Trigger *wiring* (pre-sampling checks, preflight overflow,
//! model-switch, suppression) stays per-host.
//! which also live in their module's `config.rs`.
/// Default auto-compact threshold (% of context window) when no other source
/// (env var, user config, remote per-model/global flags) sets it. Shared by
/// kigi and Kigi chat (~85% trigger on both sides).
/// Applies only when no other source (env var, user config, remote
/// per-model/global flags) sets it. Shared by kigi and Kigi chat.
pub const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT: u8 = 85;
/// Minimum character count for a cleaned summary seed.
///
/// kigi retries when the cleaned summary is shorter than this — the
/// smallest healthy prod summary observed was ~3,242 chars; anything under
/// 500 is treated as degenerate and retried like a transient failure.
/// A cleaned summary shorter than this is treated as degenerate and retried
/// like a transient failure. The smallest healthy prod summary observed was
/// ~3,242 chars, so 500 is a wide margin below anything legitimate.
pub const MIN_SUMMARY_SEED_CHARS: usize = 500;
/// Tunables for the full-replace pass.
#[derive(Debug, Clone)]
pub struct FullReplaceConfig {
/// Total LLM attempts (first try + retries) on transient failures.
/// First try + retries, counted together.
pub max_attempts: u32,
/// Delay between transient retries.
pub retry_delay_secs: u64,
/// End-to-end timeout for each compaction LLM call.
/// Applies per attempt, not to the whole retry loop.
pub sampling_timeout_secs: u64,
}
@@ -1,8 +1,8 @@
//! Deterministic-vs-transient failure classification for compaction
//! LLM calls.
//!
//! The *policy* lives here (shared across harnesses); the per-harness error
//! types and their wrapping (e.g. kigi's `SamplingError` →
//! The *policy* lives here, shared across harnesses; per-harness error types
//! and their wrapping (e.g. kigi's `SamplingError` →
//! `CompactFailure(acp::Error)`) stay in thin host wrappers that delegate the
//! status/message decisions to these functions.
@@ -17,15 +17,16 @@ pub enum FailureKind {
}
impl FailureKind {
/// `true` for [`FailureKind::Deterministic`].
pub fn is_deterministic(self) -> bool {
matches!(self, Self::Deterministic)
}
}
/// True when an error message indicates a context-window overflow. Backends report
/// this inconsistently with no stable error code, so we match the message text; it's
/// deterministic (re-sending the same payload always fails), so callers must not retry.
/// True when an error message indicates a context-window overflow.
///
/// Backends report this inconsistently with no stable error code, so the match
/// is on message text. Re-sending the same payload always fails, so callers
/// must not retry.
pub fn is_context_length_error(message: &str) -> bool {
let m = message.to_ascii_lowercase();
m.contains("too long for this model")
@@ -38,10 +39,8 @@ pub fn is_context_length_error(message: &str) -> bool {
/// Classify an HTTP API failure (status + message) for the compaction retry
/// loop.
///
/// 4xx responses other than 408 (timeout) and 429 (rate limit) are
/// deterministic; a context-length overflow message is deterministic
/// regardless of status (backends sometimes dress it as a synthesized 500).
/// Everything else (5xx, 408, 429) is transient.
/// A context-length overflow message is deterministic regardless of status,
/// because backends sometimes dress it as a synthesized 500.
pub fn classify_http_status(status: u16, message: &str) -> FailureKind {
if is_context_length_error(message)
|| ((400..500).contains(&status) && status != 408 && status != 429)
@@ -55,15 +54,11 @@ pub fn classify_http_status(status: u16, message: &str) -> FailureKind {
/// Classify a provider-style stream error event (`ResponseError` /
/// `ResponseFailed.error`) for the compaction retry loop.
///
/// `code` is the structured `code` field on the event (typically a numeric
/// HTTP status as a string, but some providers also use error-type strings like
/// `"invalid_request_error"`). `message` is the human-readable detail.
///
/// Numeric codes are classified by HTTP-status range. The
/// `invalid_request_error` marker, which can appear in either field, always
/// maps to `Deterministic` (schema violations cannot be fixed by re-sending
/// the same payload). The check order is semantic — marker, then numeric
/// code, then context-length message, then default-to-transient.
/// `code` is the structured `code` field on the event: typically a numeric HTTP
/// status as a string, but some providers send error-type strings like
/// `"invalid_request_error"` instead, and that marker can also arrive buried in
/// `message` alone. It always means a schema violation, which re-sending the
/// same payload cannot fix.
pub fn classify_stream_event_error(code: Option<&str>, message: &str) -> FailureKind {
if matches!(code, Some("invalid_request_error")) || message.contains("invalid_request_error") {
return FailureKind::Deterministic;
@@ -9,17 +9,6 @@
//! [`intra_compaction`](crate::intra_compaction) (tail-keep, per-step) and
//! [`inter_compaction`](crate::inter_compaction) (chunked, between-turn).
//!
//! Layout (mirroring
//! [`intra_compaction`](crate::intra_compaction) /
//! [`inter_compaction`](crate::inter_compaction)):
//!
//! - **Policy & content**: [`prompt`] (summarization prompt), [`summary`]
//! (summary cleaning + carrier), [`failure`] (deterministic-vs-transient
//! classification), [`config`] (tunables + trigger/seed defaults).
//! - **Algorithm**: [`assemble`] (full-replace history rebuild).
//! - **Orchestration**: [`compact`]
//! (`build prompt → sample → clean → assemble`).
//!
//! Host-specific concerns (triggers, transport, persistence/replay, state
//! commit, metrics observer) stay in the product host (for example
//! `kigi-shell`).
@@ -9,8 +9,7 @@
//! / [`InterCompactionObserver`](crate::inter_compaction::InterCompactionObserver).
//!
//! Emission points are part of the behavior contract: the kigi observer
//! preserves the pre-migration `CompactionAttempt`/`CompactionRetryDegraded`
//! semantics byte-for-byte.
//! preserves `CompactionAttempt` / `CompactionRetryDegraded` semantics.
use std::time::Duration;
@@ -34,14 +33,13 @@ pub enum FullReplaceAttemptOutcome<'a> {
/// The cleaned summary seed was too short to carry the conversation's task
/// state; retried like a transient failure.
Degenerate {
/// Raw model summary text (still captured for offline inspection).
/// Captured for offline inspection even when too short to keep.
summary: &'a str,
/// Whether the orchestrator will retry after this attempt.
will_retry: bool,
},
/// The sampler returned an error.
Failure {
/// Rendered error message.
message: &'a str,
/// Whether re-sending the *same* input cannot help (auth / schema /
/// size). Transient failures (timeout / stream blip / 5xx) are `false`.
@@ -4,8 +4,6 @@
//! full-replace pass ([`sample_full_replace_summary`](super::sample_full_replace_summary))
//! and Kigi chat's intra `Shared` summarizer
//! ([`apply_intra_compaction`](crate::intra_compaction::apply_intra_compaction)).
//! Centralising it here removes the two near-identical copies that previously
//! lived in `code_compaction::compact` and `intra_compaction::compact`.
//!
//! Classification is uniform:
//! - a usable, non-degenerate response wins immediately;
@@ -38,7 +36,7 @@ use super::summary::is_degenerate_summary;
/// plus the total number of attempts made (first try + retries).
#[derive(Debug)]
pub struct SampledSummary {
/// Raw model summary text, exactly as emitted. Callers clean it as needed.
/// Exactly as emitted; callers clean before use.
pub summary: String,
/// Total sample attempts made (1-based).
pub attempts: u32,
@@ -58,14 +56,12 @@ pub enum SampleRetryError {
/// input cannot help — auth / schema / context overflow), or transient but
/// retries were exhausted.
Failure {
/// Rendered upstream error message.
message: String,
/// Whether re-sending the same input cannot help.
deterministic: bool,
/// Whether the failure was a context-length overflow (a deterministic
/// signal the kigi host uses to step down its input size).
context_overflow: bool,
/// Total attempts made.
attempts: u32,
},
}
@@ -344,7 +340,7 @@ mod tests {
#[tokio::test]
async fn empty_then_degenerate_exhausts_to_empty() {
let short = "<summary>\n1. Primary Request: q\n</summary>"; // degenerate
let short = "<summary>\n1. Primary Request: q\n</summary>";
let sampler = MockSampler::scripted(vec![Ok(String::new()), Ok(short.into())]);
let err = run(&sampler, 2).await.expect_err("should fail");
assert!(matches!(err, SampleRetryError::Empty { attempts: 2 }));
@@ -1,7 +1,5 @@
//! Summary output cleaning and carrier formatting.
//!
//! Moved verbatim from `kigi-chat-state`'s `compaction_utils`. Covers:
//!
//! - cleaning the compaction model's raw output ([`format_compact_summary`]),
//! - the kigi continuation carrier ([`format_compact_summary_content`]),
//! - the canonical `<user_query>` wrapping ([`wrap_user_query`]).
@@ -19,15 +17,15 @@
pub fn format_compact_summary(summary: &str) -> String {
let mut result = summary.to_string();
// 1. Remove leading <analysis>…</analysis> drafting block(s). A block is
// only stripped when it is a genuinely LEADING scratchpad: top-level
// (before any <summary>) or immediately after the <summary> open modulo
// whitespace (nested). An <analysis> quoted mid-body — after real
// sections, e.g. a section-6 instruction echo — is NOT leading and is
// left for step 3 to neutralize, so neither a balanced body quote
// spanning sections nor an unclosed one ever deletes real content. The
// loop peels successive leading blocks should the model emit more than
// one.
// Strip leading <analysis>…</analysis> drafting blocks only when they are
// genuinely LEADING scratchpad: top-level (before any <summary>) or
// immediately after the <summary> open modulo whitespace (nested). An
// <analysis> quoted mid-body — after real sections, e.g. a section-6
// instruction echo — is NOT leading and is left for
// `neutralize_compaction_control_tokens` to defuse, so neither a balanced
// body quote spanning sections nor an unclosed one ever deletes real
// content. The loop peels successive leading blocks if the model emits more
// than one.
while let Some(start) = result.find("<analysis>") {
let is_leading = match result.find("<summary>") {
Some(sp) => start < sp || result[sp + "<summary>".len()..start].trim().is_empty(),
@@ -53,12 +51,13 @@ pub fn format_compact_summary(summary: &str) -> String {
}
}
// 2. Convert the outer <summary>…</summary> to "Summary:\n{inner}", keeping
// any text outside the wrapper. `rfind` matches the outer close, so a
// literal "</summary>" echoed in the body does not truncate the summary;
// `end > start` guards a malformed "</summary> … <summary>" order. Leading
// scratchpad inside the block is peeled (see `strip_leading_scratchpad`);
// a body echo that quotes the instruction is left for step 3 to defuse.
// Convert outer <summary>…</summary> to "Summary:\n{inner}", keeping text
// outside the wrapper. `rfind` matches the outer close so a literal
// "</summary>" echoed in the body does not truncate; `end > start` guards a
// malformed "</summary> … <summary>" order. Leading scratchpad inside the
// block is peeled (see `strip_leading_scratchpad`); a body echo that quotes
// the instruction is left for `neutralize_compaction_control_tokens` to
// defuse.
if let Some(start) = result.find("<summary>")
&& let Some(end) = result.rfind("</summary>")
&& end > start
@@ -69,8 +68,8 @@ pub fn format_compact_summary(summary: &str) -> String {
result = format!("{before}Summary:\n{inner}{after}");
}
// 3. Defuse any compaction-control tokens still echoed inside the body so the
// seed can't prime the next turn to re-emit a <summary> block.
// Defuse any compaction-control tokens still echoed inside the body so the
// seed can't prime the next turn to re-emit a <summary> block.
result = neutralize_compaction_control_tokens(&result);
// Collapse excessive blank lines (3+ newlines → 2)
@@ -84,13 +83,14 @@ pub fn format_compact_summary(summary: &str) -> String {
/// Peel leading drafting scratchpad off an extracted `<summary>` block.
///
/// A markdown "**Analysis**"-style header has no opening `<analysis>` tag for
/// step 1 to catch; it ends at an orphan `</analysis>`. Everything up to and
/// including the *last* `</analysis>` is dropped, so a scratchpad that itself
/// quotes `</analysis>` mid-reasoning is still removed whole. The peel is
/// skipped when the block already starts with a numbered section — including a
/// markdown-decorated one like `## 1.` or `**1.**` — so a `</analysis>` merely
/// echoed inside a real section never truncates the summary. Any leftover
/// leading `<summary>` wrapper is then unwrapped.
/// the leading-block peel in [`format_compact_summary`] to catch; it ends at an
/// orphan `</analysis>`. Everything up to and including the *last*
/// `</analysis>` is dropped, so a scratchpad that itself quotes `</analysis>`
/// mid-reasoning is still removed whole. The peel is skipped when the block
/// already starts with a numbered section — including a markdown-decorated one
/// like `## 1.` or `**1.**` — so a `</analysis>` merely echoed inside a real
/// section never truncates the summary. Any leftover leading `<summary>`
/// wrapper is then unwrapped.
fn strip_leading_scratchpad(inner: &str) -> String {
let mut s = inner.trim();
let lead = s.trim_start_matches(['#', '*', '-', '>', ' ', '\t']);
@@ -1,7 +1,7 @@
//! Item filtering and user-query extraction for history compaction —
//! generic over [`CompactionItem`] / [`CompactionItemBuilder`].
//!
//! Behavior is byte-for-byte identical for Kigi chat (`T = Arc<KigiTurn>`).
//! Behavior matches Kigi chat (`T = Arc<KigiTurn>`).
use tracing::info;
@@ -48,10 +48,9 @@ pub fn filter_turns_for_inter_compaction<T: CompactionItemBuilder>(turns: &[T])
turns
.iter()
.filter_map(|turn| match turn.role() {
// Drop tool and system items.
CompactionRole::Tool | CompactionRole::System => None,
// Keep prior compaction summaries; drop all other developer items.
// Keep prior compaction summaries; drop other developer items.
CompactionRole::Developer => {
if turn.is_compaction_summary() {
Some(turn.clone())
@@ -60,10 +59,8 @@ pub fn filter_turns_for_inter_compaction<T: CompactionItemBuilder>(turns: &[T])
}
}
// Keep user items.
CompactionRole::User => Some(turn.clone()),
// Filter assistant item contents.
CompactionRole::Assistant => turn.strip_tool_content(),
})
.collect()
@@ -152,7 +149,8 @@ pub fn truncate_middle(msg: &str, max_chars: usize) -> Option<String> {
return None;
}
let front_len = max_chars / 2;
let back_len = max_chars - front_len; // handles odd max_chars
// handles odd max_chars
let back_len = max_chars - front_len;
let front: String = msg.chars().take(front_len).collect();
let back: String = msg.chars().skip(char_count - back_len).collect();
Some(format!("{}...[truncated]...{}", front, back))
@@ -291,10 +289,8 @@ pub fn separate_prior_user_queries<T: CompactionItemBuilder>(
None => prior_user_queries = Some(user_sec),
}
}
// Matches inter's previous inline behavior (`if !rest.is_empty()`):
// a prior compaction item whose entire content was the
// `<kigi_user_queries>` block (and therefore stripped to an empty
// `rest`) contributes nothing for the LLM and is dropped here.
// Prior compaction whose content was only `<kigi_user_queries>`
// (empty `rest` after strip) contributes nothing for the LLM.
if !rest.is_empty() {
turns_for_llm.push(T::compaction_summary_item(rest));
}
@@ -1,5 +1,5 @@
//! Conversation history compaction — shared selection/assembly logic for compacting
//! prior conversation turns into a summary.
//! Shared selection/assembly logic for compacting prior conversation turns
//! into a summary.
//!
//! Everything here is generic over [`CompactionItem`](crate::CompactionItem)
//! / [`CompactionItemBuilder`](crate::CompactionItemBuilder) or pure
@@ -5,12 +5,10 @@
use anyhow::Result;
/// Builds the developer prompt to send to the compaction model.
pub fn format_compaction_developer_prompt() -> Result<String> {
Ok(include_str!("../templates/compaction_developer_prompt.txt").to_string())
}
/// Builds the user prompt to send to the compaction model.
pub fn format_compaction_user_prompt() -> Result<String> {
Ok(include_str!("../templates/compaction_user_prompt.txt").to_string())
}
@@ -27,9 +25,8 @@ mod tests {
assert!(!user.trim().is_empty(), "user prompt empty");
}
/// Belt-and-suspenders: the developer and user prompts are intentionally
/// identical so the model sees the instructions on both turns. If you edit
/// one, edit the other — this test catches drift.
/// The two templates are separate files that must stay byte-identical; this
/// catches drift when only one is edited.
#[test]
fn compaction_prompts_match() {
let dev = format_compaction_developer_prompt().expect("dev prompt renders");
@@ -2,11 +2,10 @@
use serde::{Deserialize, Serialize};
/// Strategy for how conversation compaction is performed.
#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum CompactionStrategy {
/// Send all turns to the LLM in one shot (original behaviour).
/// Send all turns to the LLM in one shot.
#[default]
Basic,
/// Divide turns into ≤ `dnc_chunk_token_limit` chunks, compact each,
@@ -44,12 +44,10 @@ pub fn validate_compaction_text(
text_content: &str,
strategy: &CompactionStrategy,
) -> Result<(), CompactionValidationError> {
// 1. Non-empty text content
if text_content.trim().is_empty() {
return Err(CompactionValidationError::EmptyContent);
}
// 2. DnC: validate chunk_summary tags are balanced
if matches!(strategy, CompactionStrategy::DivideAndConquer) {
let open_count = text_content.matches("<chunk_summary").count();
let close_count = text_content.matches("</chunk_summary>").count();
@@ -70,8 +70,7 @@ pub struct ChunkedCompactionOutput {
/// attachment refs). `conversation_id` / `response_id` are threaded
/// through for log correlation only.
///
/// Observer events (the Kigi chat observer maps them to the
/// pre-unification metrics):
/// Observer events:
/// - [`InterCompactionObserver::on_recompaction`] when prior-compaction
/// summary items are found.
/// - [`InterCompactionObserver::on_chunk_count`] — chunk count after
@@ -111,7 +110,6 @@ pub async fn sample_compaction_chunked<T: CompactionItemBuilder + Send + Sync>(
"[InterCompaction] starting chunked compaction"
);
// Step 1 — filter.
let filtered = filter_turns_for_inter_compaction(turns);
info!(
conversation_id = %conversation_id,
@@ -127,14 +125,13 @@ pub async fn sample_compaction_chunked<T: CompactionItemBuilder + Send + Sync>(
)));
}
// Step 2 — split prior `<kigi_user_queries>` out of every prior
// compaction summary item. The LLM never sees them (it would re-emit
// them verbatim and snowball across rounds); they are reattached to
// the final summary via `assemble_user_queries_preamble`. Shared with
// intra-compaction's `History` target.
// Split prior `<kigi_user_queries>` out of every prior compaction summary.
// The LLM never sees them (it would re-emit them and snowball across
// rounds); they are reattached via `assemble_user_queries_preamble`.
// Shared with intra-compaction's `History` target.
let separated = separate_prior_user_queries(&filtered);
// Step 3 — chunk + flush over the LLM-safe item list.
// Chunk + flush over the LLM-safe item list.
let mut compactable: Vec<T> = Vec::new();
let mut chunk_tokens: u32 = 0;
let mut chunk_outputs: Vec<LlmCompactionOutput> = Vec::new();
@@ -193,7 +190,7 @@ pub async fn sample_compaction_chunked<T: CompactionItemBuilder + Send + Sync>(
);
}
// Step 4a — combine summaries.
// Combine summaries.
let preamble =
assemble_user_queries_preamble(separated.prior_user_queries, current_user_queries);
let mut combined = preamble;
@@ -203,14 +200,11 @@ pub async fn sample_compaction_chunked<T: CompactionItemBuilder + Send + Sync>(
combined.push_str("\n</chunk_summary>\n\n");
}
// Step 4b — combine thinking-channel output.
let mut combined_analysis = String::new();
for (i, output) in chunk_outputs.iter().enumerate() {
combined_analysis.push_str(&wrap_chunk_analysis(i, &output.thinking));
}
// Record chunk count after assembly so dashboards see the same timing
// they saw pre-unification (where this lived inside DnC).
observer.on_chunk_count(chunk_outputs.len());
info!(
@@ -16,13 +16,11 @@ use crate::history::types::CompactionStrategy;
pub struct InterCompactionConfig {
/// The agent/scheduler name to use for the compaction model.
///
/// NOTE: model routing is host policy — kept here only because
/// service configs deserialize this struct as-is; slated to move to the
/// per-harness policy split in a later phase.
/// Model routing is host policy — kept here because service configs
/// deserialize this struct as-is.
pub compaction_model_name: String,
/// End-to-end timeout for the compaction sampling in seconds.
pub sampling_timeout_secs: u64,
/// Which compaction strategy to use.
pub compaction_strategy: CompactionStrategy,
/// [DivideAndConquer] Max tokens per chunk before sending to the LLM.
/// (Basic strategy ignores this and emits a single chunk.)
@@ -6,7 +6,6 @@
use std::time::Duration;
/// Receives inter-compaction pipeline events. All methods default to no-ops.
pub trait InterCompactionObserver: Send + Sync {
/// A prior compaction summary was found in the input (re-compaction).
/// `strategy` is the stable label from `CompactionStrategy::label()`.
@@ -15,7 +14,6 @@ pub trait InterCompactionObserver: Send + Sync {
/// One chunk's LLM call finished (success or error).
fn on_chunk_sampled(&self, _success: bool, _elapsed: Duration) {}
/// The whole pipeline finished assembling `num_chunks` chunk summaries.
fn on_chunk_count(&self, _num_chunks: usize) {}
}
@@ -221,7 +221,6 @@ where
{
let start = Instant::now();
// 1. Read the whole conversation (history ++ accumulated steps).
let source_turns = stream_proc.get_all_turns_for_compaction().await;
if source_turns.is_empty() {
return Err(IntraCompactionError::NothingToCompact);
@@ -243,27 +242,22 @@ where
"[IntraCompaction] starting full replace"
);
// 2. Summarize the whole conversation through kigi's shared core.
// FullReplace always uses the shared summarizer (it *is* the
// `code_compaction` path); `policy.summarizer` is ignored for this mode.
// FullReplace always uses the shared summarizer (it *is* the
// `code_compaction` path); `policy.summarizer` is ignored for this mode.
let summary_text = sample_shared_summary_with_retries(sampler, &source_turns, policy).await?;
// 2b. Preserve in-flight active agent state (e.g. running sub-agents) across
// the compaction. FullReplace drops the working tail, so append the
// harness-supplied `<system-reminder>` (verbatim ids) to the summary so
// the model can keep polling/cancelling them. Empty/None → no change.
// Shared with Kigi chat inter-compaction via `append_reminder_block` so
// both inject the reminder into the summary text identically, before the
// reduction guard below counts it.
// FullReplace drops the working tail — append harness-supplied
// `<system-reminder>` (verbatim ids) so in-flight sub-agents stay
// pollable/cancellable. Empty/None → no change. Shared with inter via
// `append_reminder_block` so both inject before the reduction guard counts it.
let summary_text = crate::append_reminder_block(summary_text, active_reminder);
// 3. Build the replacement developer turn. Snapshot the summary as a
// cheap-to-clone `Arc<str>` before moving the owned text into the item.
// Snapshot summary as cheap-to-clone `Arc<str>` before moving text into the item.
let summary: Arc<str> = Arc::from(summary_text.as_str());
let compaction_turn = T::compaction_summary_item(summary_text);
let tokens_after = token_counter.count_item_tokens(&compaction_turn);
// 4. Guard: don't apply if compaction didn't help.
// Don't apply if compaction didn't help.
if tokens_before > 0
&& tokens_after > (tokens_before as f64 * policy.max_reduction_ratio) as u32
{
@@ -280,7 +274,6 @@ where
});
}
// 5. Commit: replace the entire conversation with the single summary turn.
let turns_compacted = source_turns.len();
stream_proc
.replace_with_compaction(
@@ -441,9 +434,8 @@ where
{
let start = Instant::now();
// 1. Read source turns for this target. `FullReplace` never reaches
// `compact_one_pass` — it has a dedicated `apply_full_replace_compaction`
// (no tail-keep), so the orchestrator only routes `Steps`/`History` here.
// `FullReplace` never reaches here — dedicated `apply_full_replace_compaction`
// (no tail-keep); only `Steps`/`History` are routed here.
let source_turns = match target {
CompactionTarget::Steps => stream_proc.get_accumulated_turns_for_compaction().await,
CompactionTarget::History => stream_proc.get_history_turns_for_compaction().await,
@@ -459,7 +451,6 @@ where
.map(|t| token_counter.count_item_tokens(t))
.collect();
// 2. Choose split point.
let target_tokens =
(trigger.context_window as u64 * policy.target_threshold_percent as u64 / 100) as u32;
let plan = select_turns_to_compact(
@@ -483,12 +474,11 @@ where
"[IntraCompaction] starting"
);
// 3a. For `History` target, split prior `<kigi_user_queries>` blocks
// out of any prior compaction summary items before sampling — same
// primitive inter-compaction uses, so the LLM never sees
// `<kigi_user_queries>` and won't re-emit it (which would snowball
// with our explicit preamble across re-compactions). `Steps` target
// has no user-queries semantics and skips this entirely.
// For `History`, split prior `<kigi_user_queries>` out of prior compaction
// summary items before sampling — same primitive inter uses, so the LLM
// never sees `<kigi_user_queries>` and won't re-emit it (which would
// snowball with our explicit preamble). `Steps` has no user-queries
// semantics and skips this.
let (turns_for_llm, prior_user_queries) = match target {
CompactionTarget::History => {
let separated = separate_prior_user_queries(&turns_to_compact);
@@ -500,34 +490,29 @@ where
}
};
// 3b. Sample the summary. The *summarization algorithm* is switchable via
// `policy.summarizer`; everything around it — tail selection, the
// reduction guard, the prefix-replace commit, the Steps/History modes,
// and the `<kigi_user_queries>` preamble below — stays intra's.
// Summarization algorithm is switchable via `policy.summarizer`; tail
// selection, reduction guard, commit, Steps/History modes, and the
// `<kigi_user_queries>` preamble stay intra's.
let summary_text = match policy.summarizer {
// Previous intra algorithm: per-target prompt, bounded retry, and NO
// output cleaning — the raw model text flows straight to the preamble.
// Legacy: per-target prompt, bounded retry, no output cleaning —
// raw model text flows straight to the preamble.
IntraSummarizer::Legacy => {
let prompt = build_prompt_for_target(target)?;
let timeout = Duration::from_secs(policy.sampling_timeout_secs);
sample_compaction_with_retries(sampler, &turns_for_llm, &prompt, timeout, policy)
.await?
}
// New (default): kigi's shared summarization core from
// `code_compaction` — `build_summary_prompt` + degenerate-reject +
// `format_compact_summary` cleaning — run intra-locally.
// Shared: `code_compaction` core — `build_summary_prompt` +
// degenerate-reject + `format_compact_summary` cleaning.
IntraSummarizer::Shared => {
sample_shared_summary_with_retries(sampler, &turns_for_llm, policy).await?
}
};
// 3c. For `History` target, prepend a `<kigi_user_queries>` preamble so
// the original user messages + attachment refs survive the
// summarization. Carries forward both prior (from earlier
// compactions) and current (from this round's `User` turns) via
// the same `assemble_user_queries_preamble` helper that inter
// uses. (Legacy feeds the raw summary text here; Shared feeds the
// already-cleaned summary.)
// For `History`, prepend `<kigi_user_queries>` so original user messages
// + attachment refs survive. Carries prior (earlier compactions) and
// current (`User` turns) via the same `assemble_user_queries_preamble`
// helper as inter. Legacy feeds raw summary text; Shared feeds cleaned.
let final_summary_text = match target {
CompactionTarget::History => {
// Current user queries come from `turns_to_compact` (the
@@ -549,16 +534,13 @@ where
}
};
// 4. Build the replacement item. Carries category metadata so that
// subsequent compaction passes (inter or intra) treat it as
// already-compacted content. Snapshot the (possibly large) summary as a
// cheap-to-clone `Arc<str>` for the result before moving the owned text
// into the item.
// Replacement item carries category metadata so later passes treat it as
// already-compacted. Snapshot large summary as `Arc<str>` before move.
let summary: Arc<str> = Arc::from(final_summary_text.as_str());
let compaction_turn = T::compaction_summary_item(final_summary_text);
let tokens_after = token_counter.count_item_tokens(&compaction_turn);
// 5. Guard: don't apply if compaction didn't help.
// Don't apply if compaction didn't help.
if plan.tokens_to_compact > 0
&& tokens_after > (plan.tokens_to_compact as f64 * policy.max_reduction_ratio) as u32
{
@@ -575,9 +557,7 @@ where
});
}
// 6. Commit the LLM-produced summary into parser state. The trait
// method dispatches internally on `target` (Steps view vs History
// view) and rebuilds any derived state (e.g. SglangEngine).
// Trait method dispatches on `target` and rebuilds derived state (e.g. SglangEngine).
stream_proc
.replace_with_compaction(target, plan.split_idx, compaction_turn)
.await?;
@@ -676,7 +656,7 @@ where
.await
{
// kigi returns the raw summary and cleans it in its assembler;
// intra has no assembler, so it cleans here (pre-refactor behavior).
// intra has no assembler, so it cleans here.
Ok(SampledSummary { summary, .. }) => Ok(format_compact_summary(&summary)),
Err(SampleRetryError::Empty { .. }) => Err(IntraCompactionError::EmptyResponse),
Err(SampleRetryError::Failure {
@@ -853,8 +833,6 @@ mod tests {
}
}
// ── Legacy summarizer: error mapping ──
#[test]
fn compaction_sample_error_to_intra_maps_timeout() {
let intra = compaction_sample_error_to_intra(CompactionSampleError::Timeout {
@@ -931,8 +909,6 @@ mod tests {
)));
}
// ── generic orchestrator over a pure mock item ─────────────────────
/// Mock item: a `(role, text)` pair with deterministic token counting.
#[derive(Debug, Clone)]
struct MockItem {
@@ -1097,8 +1073,7 @@ mod tests {
fn enabled_policy() -> IntraCompactionConfig {
IntraCompactionConfig {
enabled: true,
// These tests exercise the steps (tail-keep) path; pin the mode so
// they stay independent of the crate default (now `FullReplace`).
// Steps (tail-keep) path; pin mode independent of FullReplace default.
mode: IntraCompactionMode::StepsOnly,
trigger_threshold_percent: 85,
target_threshold_percent: 50,
@@ -1138,7 +1113,6 @@ mod tests {
#[tokio::test]
async fn compact_replaces_turns_on_success_and_notifies_observer() {
// 6 turns × 500 tokens (2000 chars / 4); ctx 1000, target 50% → 500
// → keep the newest 1 turn, compact the oldest 5 (2500 tokens). The
// summary must be non-degenerate (>= 500 cleaned chars) for the shared
// sampler to accept it, yet still pass the reduction guard (≤ 2000).
@@ -1321,8 +1295,6 @@ mod tests {
assert!(matches!(result, Err(IntraCompactionError::Apply(_))));
}
// ── FullReplace mode (default): whole-conversation replace ────────
fn full_replace_policy() -> IntraCompactionConfig {
IntraCompactionConfig {
mode: IntraCompactionMode::FullReplace,
@@ -1332,7 +1304,6 @@ mod tests {
#[tokio::test]
async fn full_replace_compacts_whole_conversation_and_notifies_observer() {
// 6 turns × 500 tokens (2000 chars / 4) = 3000 tokens. FullReplace
// summarizes *all* of them (no tail-keep) into one developer turn.
let turns: Vec<_> = (0..6).map(|_| MockItem::user(&"x".repeat(2000))).collect();
let sp = MockStreamProc::with_turns(turns);
@@ -1355,7 +1326,7 @@ mod tests {
assert_eq!(r.turns_compacted, 6);
assert!(!r.summary.is_empty());
assert_eq!(sampler.call_count(), 1);
// The mock now holds only the single summary turn.
// Mock holds only the single summary turn.
assert_eq!(sp.turns.lock().unwrap().len(), 1);
// Observer recorded a `FullReplace` success (drives the
// `target="full_replace"` metric).
@@ -1423,7 +1394,6 @@ mod tests {
#[tokio::test]
async fn full_replace_skips_below_min_compactable_tokens() {
// 2 turns × 1 token = 2 tokens, below `min_compactable_tokens`.
let turns: Vec<_> = (0..2).map(|_| MockItem::user("x")).collect();
let sp = MockStreamProc::with_turns(turns);
let sampler = MockSampler::returns(&long_summary());
@@ -1442,8 +1412,6 @@ mod tests {
assert!(!sp.was_applied());
}
// ── Shared summarizer (default): direct helper coverage ───────────
#[tokio::test]
async fn shared_summarizer_rejects_degenerate_then_errors() {
// A too-short summary is degenerate; `Shared` retries (max_attempts=2)
@@ -1511,12 +1479,10 @@ mod tests {
assert!(!cleaned.contains("<summary>"), "tags must be neutralized");
}
// ── Legacy summarizer: end-to-end switch ──────────────────────────
#[tokio::test]
async fn legacy_summarizer_accepts_short_uncleaned_summary() {
// Legacy has no degenerate floor and does NO cleaning: a short raw
// summary is accepted verbatim (would be rejected under `Shared`).
// Legacy has no degenerate floor and does no cleaning: a short raw
// summary is accepted verbatim (rejected under `Shared`).
let turns: Vec<_> = (0..6).map(|_| MockItem::user(&"x".repeat(400))).collect();
let sp = MockStreamProc::with_turns(turns);
let sampler = MockSampler::returns("compacted summary");
@@ -32,11 +32,11 @@ pub enum IntraCompactionMode {
#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum IntraSummarizer {
/// New (default): the shared summarization core — `build_summary_prompt`
/// + degenerate-reject + `format_compact_summary` cleaning.
/// Shared summarization core — `build_summary_prompt` + degenerate-reject
/// + `format_compact_summary` cleaning.
#[default]
Shared,
/// Previous intra algorithm: per-target prompt (`format_compaction_prompt`
/// Legacy path: per-target prompt (`format_compaction_prompt`
/// / history dev+user prompts), no cleaning. Kept for switchability.
Legacy,
}
@@ -82,11 +82,11 @@ pub enum IntraSummarizer {
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct IntraCompactionConfig {
// ───────────────────────────── Common (all modes) ─────────────────────────────
// Common (all modes)
// Present on every config path regardless of `mode`. Some trigger fields
// are ignored under FullReplace (see per-field docs).
// -- Enablement & strategy selection --
// Enablement & strategy selection
/// Enable intra-compaction between steps. Default: `false` (disabled).
pub enabled: bool,
@@ -94,7 +94,7 @@ pub struct IntraCompactionConfig {
/// Default: `FullReplace`.
pub mode: IntraCompactionMode,
// -- Trigger gating: when a compaction pass fires (see `should_compact`) --
// Trigger gating: when a compaction pass fires (see `should_compact`)
/// Context window usage percentage (0-100) that triggers compaction.
/// Compared against: `last_prompt_tokens / context_length.max_len`.
/// Default: `85`.
@@ -112,7 +112,7 @@ pub struct IntraCompactionConfig {
/// / reduction guards after a trigger.
pub min_steps_before_compact: u32,
// -- Reduction guards: whether a produced summary is worth keeping --
// Reduction guards: whether a produced summary is worth keeping
/// Minimum tokens that must be reducible before compaction is worth
/// running. Below this, the LLM overhead outweighs the savings.
/// Default: `5000`.
@@ -123,7 +123,7 @@ pub struct IntraCompactionConfig {
/// Default: `0.8`.
pub max_reduction_ratio: f64,
// -- Compaction LLM call (sampling) --
// Compaction LLM call (sampling)
/// Compaction model name. Blank/`None` → [`DEFAULT_COMPACTION_MODEL_NAME`].
/// Prefer [`Self::effective_compaction_model_name`].
pub compaction_model_name: Option<String>,
@@ -142,12 +142,12 @@ pub struct IntraCompactionConfig {
/// Delay between retries. Default: `3`.
pub retry_delay_secs: u64,
// -- Audit --
// Audit
/// Version string for the compaction (e.g. `"intra-v1"`).
/// Recorded in audit logs. Default: `"intra-v1"`.
pub compaction_version: String,
// ───────────────────────────── Mode-specific ─────────────────────────────
// Mode-specific
// Each field below is read by only a subset of modes; the other modes
// ignore it entirely. The bracketed `[...]` tag on each doc names the modes
// that consume it.
@@ -172,7 +172,7 @@ pub struct IntraCompactionConfig {
/// replaces everything and never reads it.
pub target_threshold_percent: u8,
// -- HistoryThenSteps only --
// HistoryThenSteps only
/// [HistoryThenSteps mode] Only compact accumulated step turns when their
/// token count exceeds this fraction of the history token count.
///
@@ -5,7 +5,7 @@
//! counters/histograms in the harness crate)
//! without the shared crate depending on a metrics backend. Emission points
//! and label values are part of the behavior contract — Kigi chat's
//! observer preserves them byte-for-byte.
//! observer preserves them exactly.
use std::time::Duration;
@@ -1,27 +1,19 @@
//! Trait abstractions for intra-compaction.
use async_trait::async_trait;
use super::trigger::IntraCompactionError;
/// Which segment of the conversation a single intra-compaction pass acts on.
///
/// Determines the prompt template the orchestrator uses, which read-view
/// it pulls items from on the stream processor (`get_accumulated_turns_for_compaction`
/// vs `get_history_turns_for_compaction`), and which branch the stream processor's
/// [`CompactionStreamProc::replace_with_compaction`] dispatches to.
/// Selects both the prompt template and the read-view items are pulled from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompactionTarget {
/// Compact the agent loop's accumulated step turns (assistant outputs,
/// tool calls, tool results). Fine-grained prompt.
/// The agent loop's accumulated step turns: assistant outputs, tool calls,
/// tool results.
Steps,
/// Compact prior conversation-history turns (user/assistant exchanges
/// from before the current agent loop). Coarser prompt, shared with
/// inter-compaction.
/// Prior conversation-history turns, from before the current agent loop.
/// Coarser prompt, shared with inter-compaction.
History,
/// Replace the *whole* conversation (prior history + accumulated steps)
/// with a single summary — kigi's full-replace strategy. No tail is
/// kept; the read-view is [`CompactionStreamProc::get_all_turns_for_compaction`].
/// Prior history *and* accumulated steps, replaced wholesale by a single
/// summary. No tail is kept.
FullReplace,
}
@@ -37,26 +29,8 @@ impl CompactionTarget {
}
/// Minimal interface the compaction orchestrator needs from the agent's
/// stream processor. Implemented by Kigi chat's
/// `StreamProcessor` (`Item = Arc<KigiTurn>`).
///
/// Two read-views are exposed:
///
/// - **Accumulated step turns**: items added since the agent loop started
/// — assistant outputs, tool calls, tool results, recovery turns. The
/// original conversation (system prompt, user messages, prior history)
/// is excluded. Used by step (fine-grained) compaction.
/// - **History turns**: items from prior user-query/assistant-response
/// exchanges, before the current agent loop began. Used by history
/// (coarse) compaction.
///
/// The single mutator [`Self::replace_with_compaction`] takes a
/// [`CompactionTarget`] and dispatches internally to the steps- or
/// history-specific path. It is the final step of a compaction cycle:
/// the LLM-produced summary is committed into parser state. The
/// orchestrator [`super::apply_intra_compaction`] and its peers
/// [`super::apply_steps_compaction`] / [`super::apply_history_compaction`]
/// are the layers above that produce the summary and call this method.
/// stream processor. Implemented by Kigi chat's `StreamProcessor`
/// (`Item = Arc<KigiTurn>`).
///
/// Implementations that don't support a particular target return
/// [`IntraCompactionError::Unsupported`] from the matching match arm.
@@ -65,13 +39,13 @@ pub trait CompactionStreamProc: Send + Sync {
/// The harness's conversation item type.
type Item;
/// Get the items accumulated across all completed steps, oldest first.
/// Candidates for **steps** compaction.
/// Items accumulated across all completed steps, oldest first. The
/// original conversation (system prompt, user messages, prior history) is
/// excluded.
async fn get_accumulated_turns_for_compaction(&self) -> Vec<Self::Item>;
/// Get the conversation-history items (prior user/assistant exchanges
/// from before the current agent loop), oldest first. Candidates for
/// **history** compaction.
/// Items from prior user/assistant exchanges, before the current agent
/// loop began, oldest first.
///
/// Default impl returns empty — implementations that do not support
/// history compaction will have nothing to compact.
@@ -79,13 +53,11 @@ pub trait CompactionStreamProc: Send + Sync {
Vec::new()
}
/// Get the **whole** conversation — prior history followed by the
/// accumulated step turns, oldest first. Candidates for **full-replace**
/// (`CompactionTarget::FullReplace`) compaction.
/// The whole conversation — prior history followed by accumulated step
/// turns, oldest first.
///
/// The default composes the two read-views above (`history ++ steps`),
/// which is correct for any implementation; override only if a harness can
/// produce the combined view more cheaply.
/// The default composition is correct for any implementation; override
/// only if a harness can produce the combined view more cheaply.
///
/// The `Self::Item: Send` bound lets the default hold the history vec across
/// the second `await` while keeping the boxed future `Send`; every concrete
@@ -99,13 +71,10 @@ pub trait CompactionStreamProc: Send + Sync {
all
}
/// Top-level intra-compaction mutator. Replaces the first
/// `n_turns_to_remove` items in the read-view selected by `target` with
/// the single given `compaction_turn`.
/// Replaces the first `n_turns_to_remove` items of the read-view selected
/// by `target` with the single given `compaction_turn`.
///
/// Implementations dispatch internally on `target` to the steps or
/// history specific path. On invalid input
/// (`n_turns_to_remove > view.len()`), returns
/// On invalid input (`n_turns_to_remove > view.len()`), returns
/// [`IntraCompactionError::InvalidSplit`] and leaves state untouched.
async fn replace_with_compaction(
&self,
@@ -24,7 +24,6 @@ pub struct IntraCompactionTrigger {
pub step: u32,
}
/// Result of a successful compaction.
#[derive(Debug, Clone)]
pub struct IntraCompactionResult {
/// Sum of tokens in the turns that were compacted.
+1 -5
View File
@@ -23,16 +23,13 @@
/// `ConversationItem` variants (kigi).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompactionRole {
/// System prompt.
System,
/// Developer prompt (Kigi chat) — maps to System on harnesses without a
/// distinct developer role.
Developer,
/// A user message.
User,
/// An assistant output (may carry tool requests).
/// Assistant output (may carry tool requests).
Assistant,
/// A tool result.
Tool,
}
@@ -43,7 +40,6 @@ pub enum CompactionRole {
pub struct CompactionFileRef {
/// Stable unique id of the attachment source.
pub id: String,
/// Human-readable file name.
pub name: String,
}
+2 -2
View File
@@ -57,8 +57,8 @@ pub mod token;
/// 3. this constant
pub use intra_compaction::DEFAULT_COMPACTION_MODEL_NAME;
// kigi's full-replace subsystem now lives under `code_compaction`;
// re-exported at the crate root so consumers keep a stable public API.
// Full-replace subsystem lives under `code_compaction`; re-exported at the
// crate root for a stable public API.
pub use code_compaction::{
CompactedHistoryParts, DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT, FailureKind,
FullReplaceAttemptOutcome, FullReplaceConfig, FullReplaceContext, FullReplaceError,
+4 -6
View File
@@ -1,14 +1,12 @@
//! The shared compaction prompt seam.
//!
//! [`CompactionPrompt`] is the system+user prompt pair every orchestrator's
//! [`CompactionSampler`](crate::sampler::CompactionSampler) call takes. The
//! per-strategy prompt *content* lives with each subsystem:
//! The shared compaction prompt seam. Per-strategy prompt *content* lives with
//! each subsystem:
//!
//! - steps prompt → [`crate::steps::format_compaction_prompt`]
//! - history prompts → [`crate::history::prompt`]
//! - kigi summary prompt → [`crate::code_compaction::build_summary_prompt`]
/// System + user prompt pair for the compaction LLM call.
/// System + user prompt pair every
/// [`CompactionSampler`](crate::sampler::CompactionSampler) call takes.
#[derive(Debug, Clone)]
pub struct CompactionPrompt {
pub system: String,
@@ -17,9 +17,7 @@
//! views** (`&str` over live state) so long fields (commands, todo content,
//! descriptions, ids) are not cloned just to format.
// ---------------------------------------------------------------------------
// Borrowed views over harness live state (no long-string clones)
// ---------------------------------------------------------------------------
/// Model-facing poll/cancel tool names from the current toolset.
/// Never hard-code: a client manifest can rename them.
@@ -102,9 +100,7 @@ impl ActiveAgentReminderState<'_> {
}
}
// ---------------------------------------------------------------------------
// Section formatters
// ---------------------------------------------------------------------------
/// `## Running Background Tasks`, or `None` when empty.
pub fn section_background_tasks(tasks: &[BackgroundTask<'_>]) -> Option<String> {
@@ -249,9 +245,7 @@ pub fn format_active_agent_reminder(
wrap_system_reminder(format_active_agent_sections(state, subagent_tools))
}
// ---------------------------------------------------------------------------
// Summary injection
// ---------------------------------------------------------------------------
/// Append a trailing block to a compaction summary, separated by a blank line.
/// Returns `summary` unchanged when `reminder` is `None` or blank.
+6 -9
View File
@@ -7,9 +7,7 @@ use async_trait::async_trait;
use crate::prompt::CompactionPrompt;
// ---------------------------------------------------------------------------
// Sampler output + error types
// ---------------------------------------------------------------------------
/// Raw text captured from a compaction LLM call, split by channel.
///
@@ -45,16 +43,17 @@ pub enum CompactionSampleError {
Build(String),
/// The sampling call could not be started.
///
/// Classification is asymmetric for pre-migration parity: the *inter*
/// retry policy ([`Self::is_deterministic`]) treats it as deterministic
/// (no retry), while the *intra* orchestrator maps it to
/// Classification is asymmetric: the *inter* retry policy
/// ([`Self::is_deterministic`]) treats it as deterministic (no retry),
/// while the *intra* orchestrator maps it to
/// `IntraCompactionError::SamplerStart` which its retry loop treats as
/// transient.
Start(String),
/// The model produced no response-channel content. Transient.
EmptyResponse,
/// Anything else — classified by string matching for backward
/// compatibility with samplers that pre-date the structured variants.
/// Opaque error — classified by string matching when the sampler does not
/// use a structured variant. Keep match literals in sync with the harness
/// sampler (`compaction_sample_error_to_intra*` tests guard the mapping).
Other(anyhow::Error),
}
@@ -103,9 +102,7 @@ impl CompactionSampleError {
}
}
// ---------------------------------------------------------------------------
// Sampler trait
// ---------------------------------------------------------------------------
/// Interface for the LLM call that produces compaction summaries.
///
+9 -10
View File
@@ -74,14 +74,13 @@ pub fn select_turns_to_compact<T: CompactionItem>(
return None;
}
// Step 1: Walk backward, sum "keep" tokens until target is reached.
// Find the highest split_idx such that sum(item_token_counts[split_idx..]) ≤ target_tokens.
// Highest split_idx such that sum(item_token_counts[split_idx..]) ≤ target_tokens.
let mut kept = 0u32;
let mut split_idx = total; // start with "compact nothing", will move down
// Start with "compact nothing", walk down.
let mut split_idx = total;
for i in (0..total).rev() {
let count = item_token_counts[i];
if kept.saturating_add(count) > target_tokens {
// Adding this item would exceed the budget — split here.
split_idx = i + 1;
break;
}
@@ -94,7 +93,6 @@ pub fn select_turns_to_compact<T: CompactionItem>(
return None;
}
// Step 2: Snap the split forward to a safe boundary.
let safe_split_idx = snap_to_safe_boundary(items, split_idx);
// After snapping forward we might have eaten everything.
@@ -102,7 +100,6 @@ pub fn select_turns_to_compact<T: CompactionItem>(
return None;
}
// Step 3: Compute tokens to compact and check the minimum.
let tokens_to_compact: u32 = item_token_counts[..safe_split_idx]
.iter()
.copied()
@@ -131,7 +128,6 @@ fn snap_to_safe_boundary<T: CompactionItem>(items: &[T], candidate: usize) -> us
return total;
}
// If candidate is not a tool-result item, no snap needed.
if !items[candidate].is_tool_result() {
return candidate;
}
@@ -214,7 +210,8 @@ mod tests {
MockItem::user(),
MockItem::assistant(),
];
let counts = vec![40, 30, 20, 10]; // keep last two (sum 30)
// keep last two (sum 30)
let counts = vec![40, 30, 20, 10];
let plan = select_turns_to_compact(&counts, &items, 30, 5).expect("should split");
assert_eq!(plan.split_idx, 2);
assert_eq!(plan.tokens_to_compact, 70);
@@ -235,7 +232,8 @@ mod tests {
let items = vec![
MockItem::user(),
MockItem::assistant(),
MockItem::assistant(), // pretend this had tool_requests
// pretend this had tool_requests
MockItem::assistant(),
MockItem::tool(),
MockItem::tool(),
MockItem::assistant(),
@@ -256,7 +254,8 @@ mod tests {
let items = vec![
MockItem::user(),
MockItem::assistant(),
MockItem::user(), // safe split here
// safe split here
MockItem::user(),
MockItem::assistant(),
];
let counts = vec![50, 50, 10, 10];
@@ -1,11 +1,10 @@
//! Steps compaction — prompt content for compacting accumulated step
//! turns (tool calls + assistant responses) within a single agent turn.
//! Prompt content for compacting accumulated step turns (tool calls +
//! assistant responses) within a single agent turn.
//!
//! Parallel to [`crate::history`] (the history-compaction content): this is the
//! *steps* side. The orchestration that uses it lives in
//! [`crate::intra_compaction`] (the `Steps` target / `StepsOnly` mode), and the
//! Parallel to [`crate::history`], the history-compaction content. The
//! orchestration that uses it lives in [`crate::intra_compaction`], and the
//! turn selection it shares with the History target is the crate-root
//! [`select`](crate::select) primitive (not steps-specific).
//! [`select`](crate::select) primitive.
pub mod prompt;
@@ -1,16 +1,11 @@
//! Prompt construction for **steps** compaction.
//!
//! The step-level intra-compaction prompt: short and focused on summarising
//! tool-call history mid-task. Parallel to [`crate::history::prompt`] (the
//! history-compaction prompts); templates live in the crate-root `templates/`.
//! Prompt construction for **steps** compaction. Templates live in the
//! crate-root `templates/`.
use crate::prompt::CompactionPrompt;
/// Build the standard prompt for step-level intra-compaction.
///
/// The prompts are short and focused on summarising tool-call history
/// mid-task — the assistant has already done several steps of work and
/// we need to free up context so it can continue.
/// Short prompts focused on summarising tool-call history mid-task: the
/// assistant has already done several steps of work and needs context freed up
/// to continue.
pub fn format_compaction_prompt() -> CompactionPrompt {
CompactionPrompt {
system: include_str!("../templates/intra_compaction_system.txt").to_string(),
+5 -12
View File
@@ -1,24 +1,17 @@
//! Token-count seam.
//!
//! Budgeting math in the shared engine needs a *trusted* token count, but the
//! two harnesses disagree on how to produce one:
//!
//! - Kigi chat has a real tokenizer (`TextTokenizer` / `ImageTokenizer`) and
//! counts whole turns via `KigiTurn::get_num_tokens`.
//! - kigi estimates with `bytes / 4`.
//!
//! Rather than bake either policy into the shared crate, callers supply an
//! [`ItemTokenCounter`]. This keeps the engine deterministic and testable
//! while letting each harness plug in its own counting strategy.
//! two harnesses disagree on how to produce one: Kigi chat has a real tokenizer
//! (`TextTokenizer` / `ImageTokenizer`) and counts whole turns via
//! `KigiTurn::get_num_tokens`, while kigi estimates with `bytes / 4`. Rather
//! than bake either policy into the shared crate, callers supply an
//! [`ItemTokenCounter`].
//!
//! There is intentionally **no** blanket `Arc` forwarding here: each harness
//! implements the counter directly for the item type its algorithms run on
//! (Kigi chat: `ItemTokenCounter<Arc<KigiTurn>>`), so exactly one mechanism
//! is in play.
/// Counts tokens for a single conversation item on behalf of the shared
/// budgeting logic.
pub trait ItemTokenCounter<T: ?Sized>: Send + Sync {
/// Trusted token count of `item`.
fn count_item_tokens(&self, item: &T) -> u32;
}
@@ -3,30 +3,27 @@ use serde::{Deserialize, Serialize};
use crate::events::EventQueue;
use crate::format::format_interjection;
/// A buffered mid-turn interjection awaiting the next safe drain point.
/// `Attachment` is host-defined (inline images, asset IDs); core never reads it.
/// Mid-turn interjection waiting for the next safe drain point.
/// `Attachment` is host-defined; core never inspects it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PendingInterjection<Attachment> {
pub text: String,
pub attachments: Vec<Attachment>,
}
/// A drained entry, wrapped and ready to emit as a synthetic user message.
/// Drained entry, framed as a synthetic user message.
#[derive(Debug, Clone, PartialEq)]
pub struct FormattedInterjection<Attachment> {
pub text: String,
pub attachments: Vec<Attachment>,
}
/// A queue of pending interjections — just an [`EventQueue`] of
/// [`PendingInterjection`]. Use [`drain_formatted`] to drain + frame them as
/// synthetic user messages.
/// Queue of [`PendingInterjection`] values. Drain via [`drain_formatted`].
pub type InterjectionBuffer<Attachment> = EventQueue<PendingInterjection<Attachment>>;
/// Drain `buffer`, framing each entry as a synthetic user message (FIFO, one
/// message per entry, never merged). `sanitize_text` runs on the raw text first
/// (hosts strip artifacts like image placeholder paths; pass
/// `std::convert::identity` if none).
/// Drain `buffer` FIFO, one synthetic user message per entry (never merged).
/// `sanitize_text` runs on raw text first (e.g. strip image placeholders);
/// pass `std::convert::identity` when none is needed.
pub fn drain_formatted<Attachment>(
buffer: &InterjectionBuffer<Attachment>,
sanitize_text: impl Fn(String) -> String,
@@ -1,7 +1,6 @@
//! Shared event queue for the push path: producers enqueue out-of-band events;
//! readers drain the ones relevant to them at hook points. Internally
//! synchronized and `Arc`-shared (clones share one queue), mirroring
//! [`crate::buffer::InterjectionBuffer`].
//! Shared push-path event queue: producers enqueue; readers drain at hook
//! points. Internally synchronized and `Arc`-shared (clones share one queue),
//! matching [`crate::buffer::InterjectionBuffer`].
use std::sync::{Arc, Mutex, MutexGuard};
@@ -31,12 +30,11 @@ impl<E> EventQueue<E> {
}
}
/// Producer hook: record an event for later draining.
pub fn push(&self, event: E) {
self.lock().push(event);
}
/// Push, then drop the oldest events so at most `max` remain.
/// Push, then drop oldest entries so length stays ≤ `max`.
pub fn push_capped(&self, event: E, max: usize) {
let mut q = self.lock();
q.push(event);
@@ -54,8 +52,7 @@ impl<E> EventQueue<E> {
self.lock().is_empty()
}
/// Remove and return events matching `take`, retaining the rest. FIFO order
/// is preserved in both the returned and retained sets.
/// Remove matching events; keep the rest. Both sets stay FIFO.
pub fn drain_matching(&self, take: impl Fn(&E) -> bool) -> Vec<E> {
let mut q = self.lock();
let (matched, kept): (Vec<E>, Vec<E>) =
@@ -64,12 +61,10 @@ impl<E> EventQueue<E> {
matched
}
/// Remove and return all events, leaving the queue empty (FIFO order).
pub fn drain_all(&self) -> Vec<E> {
std::mem::take(&mut *self.lock())
}
/// Discard all events.
pub fn clear(&self) {
self.lock().clear();
}
@@ -80,7 +75,7 @@ impl<E> EventQueue<E> {
}
impl<E: Clone> EventQueue<E> {
/// Clone of the current events, for inspection without draining.
/// Snapshot without draining.
pub fn snapshot(&self) -> Vec<E> {
self.lock().clone()
}
@@ -1,4 +1,4 @@
/// Truncation threshold, matching the shell's large-prompt limit.
/// Truncation threshold; matches the shell's large-prompt limit.
pub const LARGE_PROMPT_THRESHOLD: usize = 25_000;
/// Wrap a user message in the canonical `<user_query>` envelope.
@@ -10,9 +10,9 @@ pub fn user_query(user_message: &str) -> String {
)
}
/// Wrap interjection text as a synthetic user message with a mid-turn note.
/// No deferral instruction: the model decides how to weigh it against
/// in-flight work. Output is byte-identical to the shell's historical format.
/// Frame interjection text as a synthetic mid-turn user message.
/// No deferral instruction the model weighs it against in-flight work.
/// Byte-identical to the shell's historical format.
pub fn format_interjection(text: String) -> String {
let truncated = if text.len() > LARGE_PROMPT_THRESHOLD {
let end = text
+2 -3
View File
@@ -1,8 +1,7 @@
//! Environment-variable test knobs.
/// Parse a `usize` env knob, falling back to `default` when unset or
/// unparseable. The perf-repro convention for sizing `#[ignore]` benches
/// (e.g. `KIGI_PERF_GIT_FILES`).
/// Parse a `usize` env knob; use `default` when unset or unparseable.
/// Perf-repro convention for sizing `#[ignore]` benches (e.g. `KIGI_PERF_GIT_FILES`).
pub fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
+20 -35
View File
@@ -1,21 +1,16 @@
//! Hermetic git helpers for tests.
//!
//! When running under `bazel test`, the `GIT_BIN_PATH` environment variable
//! points to a statically-linked git binary provided by Bazel. The helpers
//! in this module prepend that binary's directory to `PATH` so that
//! `Command::new("git")` resolves to it instead of relying on a
//! system-installed git.
//! Under `bazel test`, `GIT_BIN_PATH` points at a Bazel-provided static git.
//! Helpers prepend that binary's directory to `PATH` so `Command::new("git")`
//! resolves to it instead of a system install.
use std::path::{Path, PathBuf};
use std::sync::Once;
static HERMETIC_GIT_INIT: Once = Once::new();
/// Prepend the hermetic git binary directory to `PATH` so that
/// `Command::new("git")` resolves to the Bazel-provided static binary
/// instead of relying on a system-installed git.
///
/// Safe to call multiple times — only the first call mutates `PATH`.
/// Prepend the hermetic git binary directory to `PATH`.
/// Idempotent — only the first call mutates `PATH`.
pub fn ensure_hermetic_git_on_path() {
HERMETIC_GIT_INIT.call_once(|| {
if let Ok(git_bin) = std::env::var("GIT_BIN_PATH") {
@@ -27,7 +22,7 @@ pub fn ensure_hermetic_git_on_path() {
};
if let Some(bin_dir) = git_path.parent() {
let current_path = std::env::var("PATH").unwrap_or_default();
// SAFETY: called once via `Once` before any child processes are spawned.
// SAFETY: once via `Once`, before any child processes spawn.
unsafe {
std::env::set_var("PATH", format!("{}:{}", bin_dir.display(), current_path));
}
@@ -36,8 +31,7 @@ pub fn ensure_hermetic_git_on_path() {
});
}
/// Ensure the hermetic git binary is on `PATH` before running tests that
/// need git. Call at the top of any `#[test]` that spawns `git` commands.
/// Put hermetic git on `PATH` at the top of tests that spawn `git`.
///
/// ```ignore
/// #[test]
@@ -53,9 +47,7 @@ macro_rules! require_git {
};
}
/// Initialise a fresh git repository at `path` with a dummy user config.
///
/// Calls [`ensure_hermetic_git_on_path`] first so the hermetic binary is used.
/// Init a fresh repo at `path` with dummy user config (hermetic git).
pub fn init_git_repo(path: &Path) {
ensure_hermetic_git_on_path();
std::process::Command::new("git")
@@ -77,9 +69,7 @@ pub fn init_git_repo(path: &Path) {
.unwrap();
}
/// Stage all files and create a commit.
///
/// Calls [`ensure_hermetic_git_on_path`] first so the hermetic binary is used.
/// Stage all files and create a commit (hermetic git).
pub fn git_commit_all(path: &Path, message: &str) {
ensure_hermetic_git_on_path();
std::process::Command::new("git")
@@ -94,20 +84,16 @@ pub fn git_commit_all(path: &Path, message: &str) {
.unwrap();
}
/// Run a git command in `dir` with a deterministic author/committer, assert
/// success, and return trimmed stdout.
///
/// Calls [`ensure_hermetic_git_on_path`] first so the hermetic binary is used.
/// Run git in `dir` with a fixed author/committer; assert success; return
/// trimmed stdout (hermetic git).
pub fn run_git(dir: &Path, args: &[&str]) -> String {
run_git_with_env(dir, args, &[])
}
/// Like [`run_git`], with extra environment variables (e.g.
/// `GIT_SEQUENCE_EDITOR`). Hermetic beyond the binary and author identity:
/// the developer's global/system git config is masked (a local
/// `commit.gpgsign`/`core.hooksPath`/`rebase.autoSquash` must not change
/// test behavior) and credential prompts are disabled. `envs` is applied
/// last, so callers can override any of this.
/// Like [`run_git`], with extra env vars (e.g. `GIT_SEQUENCE_EDITOR`).
/// Masks global/system git config and disables credential prompts so local
/// `commit.gpgsign` / `core.hooksPath` / `rebase.autoSquash` cannot skew
/// tests. `envs` is applied last and may override any of this.
pub fn run_git_with_env(dir: &Path, args: &[&str], envs: &[(&str, &str)]) -> String {
ensure_hermetic_git_on_path();
let mut cmd = std::process::Command::new("git");
@@ -138,9 +124,8 @@ pub fn run_git_with_env(dir: &Path, args: &[&str], envs: &[(&str, &str)]) -> Str
String::from_utf8_lossy(&output.stdout).trim().to_string()
}
/// Write a grouped fan-out tree of ~`files` files (`files_per_dir` per
/// directory, directories bucketed 100 per group) under `dir`. No git
/// operations — callers stage/commit as needed.
/// Write a grouped fan-out of ~`files` files under `dir` (`files_per_dir`
/// per directory, directories bucketed 100 per group). No git ops.
pub fn write_fanout_tree(dir: &Path, files: usize, files_per_dir: usize) {
for d in 0..files.div_ceil(files_per_dir) {
let sub = dir.join(format!("g{}", d / 100)).join(format!("d{d}"));
@@ -155,9 +140,9 @@ pub fn write_fanout_tree(dir: &Path, files: usize, files_per_dir: usize) {
}
}
/// Create a `feature` branch with `picks` one-file commits off the current
/// HEAD, advance the base branch by one commit (so a rebase has work), and
/// leave `feature` checked out. Returns the base branch name.
/// Create `feature` with `picks` one-file commits off HEAD, advance the base
/// by one commit (so rebase has work), leave `feature` checked out.
/// Returns the base branch name.
pub fn make_feature_branch(dir: &Path, picks: usize) -> String {
let base = run_git(dir, &["rev-parse", "--abbrev-ref", "HEAD"]);
run_git(dir, &["checkout", "-b", "feature"]);
+7 -6
View File
@@ -1,13 +1,14 @@
//! Synthetic image fixtures shared across crates' test suites.
/// Wrap a PNG into a minimal single-frame ICO. `width`/`height` are the
/// ICONDIRENTRY bytes (`0` means 256); the PNG carries the real dimensions.
/// Wrap a PNG into a minimal single-frame ICO.
/// `width`/`height` are ICONDIRENTRY bytes (`0` means 256); the PNG holds
/// the real dimensions.
pub fn ico_with_png_frame(png: &[u8], width: u8, height: u8) -> Vec<u8> {
let mut buf = Vec::with_capacity(22 + png.len());
buf.extend_from_slice(&[0, 0, 1, 0, 1, 0]); // ICONDIR
buf.extend_from_slice(&[width, height, 0, 0, 1, 0, 32, 0]); // ICONDIRENTRY
buf.extend_from_slice(&(png.len() as u32).to_le_bytes()); // bytes in resource
buf.extend_from_slice(&22u32.to_le_bytes()); // offset to the PNG payload
buf.extend_from_slice(&[0, 0, 1, 0, 1, 0]);
buf.extend_from_slice(&[width, height, 0, 0, 1, 0, 32, 0]);
buf.extend_from_slice(&(png.len() as u32).to_le_bytes());
buf.extend_from_slice(&22u32.to_le_bytes());
buf.extend_from_slice(png);
buf
}
+6 -18
View File
@@ -1,22 +1,10 @@
//! Shared test utilities for xAI crates.
//! Shared test utilities for Kigi crates.
//!
//! Provides common helpers that are needed by many crates' test suites:
//!
//! - **Hermetic git**: [`git::ensure_hermetic_git_on_path`] prepends the Bazel-provided
//! static `git` binary to `PATH` so that tests don't depend on a system-installed git.
//! The [`require_git!`] macro is a convenient shorthand.
//!
//! - **Git repo helpers**: [`git::init_git_repo`] and [`git::git_commit_all`] for
//! setting up throwaway git repos in tests.
//!
//! - **Bazel runfiles**: [`crate_root!`] resolves the crate root directory via
//! Bazel runfiles (for `bazel test`) or `CARGO_MANIFEST_DIR` (for `cargo test`).
//!
//! - **Tracing capture**: [`tracing_capture::MessagePrefixCounter`] counts
//! log lines by message prefix (thread-scoped or global install) for tests
//! that assert on how often an instrumented code path ran.
//!
//! - **Env knobs**: [`env::env_usize`] for perf-repro test sizing.
//! - **Hermetic git**: [`git::ensure_hermetic_git_on_path`] / [`require_git!`]
//! - **Repo helpers**: [`git::init_git_repo`], [`git::git_commit_all`]
//! - **Bazel runfiles**: [`crate_root!`]
//! - **Tracing capture**: [`tracing_capture::MessagePrefixCounter`]
//! - **Env knobs**: [`env::env_usize`]
pub mod env;
pub mod git;
@@ -1,15 +1,12 @@
//! Bazel runfiles helpers for locating test data.
//! Bazel runfiles helpers for test data.
//!
//! Under `bazel test`, source files and test data are accessed via the
//! *runfiles* tree. Under `cargo test`, `CARGO_MANIFEST_DIR` provides
//! the crate root. The [`crate_root!`] macro abstracts over both.
//! Under `bazel test`, data lives in the runfiles tree; under `cargo test`,
//! `CARGO_MANIFEST_DIR` is the crate root. [`crate_root!`] covers both.
use std::path::PathBuf;
/// Try to resolve a runfiles path to an absolute directory.
///
/// Returns `Some(path)` when running under Bazel (with the `bazel` feature
/// enabled) and the runfiles entry exists, `None` otherwise.
/// Resolve a runfiles path to an absolute directory when the `bazel` feature
/// is on and the entry exists; otherwise `None`.
pub fn try_resolve_runfiles(_path: &str) -> Option<PathBuf> {
#[cfg(feature = "bazel")]
{
@@ -22,11 +19,8 @@ pub fn try_resolve_runfiles(_path: &str) -> Option<PathBuf> {
}
}
/// Resolve the crate root directory, working under both `bazel test` and
/// `cargo test`.
///
/// Under Bazel the path is resolved via runfiles; under Cargo it falls back
/// to `CARGO_MANIFEST_DIR`.
/// Crate root under both `bazel test` (runfiles) and `cargo test`
/// (`CARGO_MANIFEST_DIR`).
///
/// # Example
///
@@ -1,14 +1,13 @@
//! Test-only tracing capture: count events whose `message` starts with a
//! known prefix.
//! Count tracing events whose `message` starts with a known prefix.
//!
//! Producers should export the exact log-line prefixes as `pub const`s next
//! to the `tracing::debug!` call sites (e.g. `kigi_hunk_tracker`'s
//! `REFRESH_SCAN_LOG_PREFIX`) so tests never duplicate the strings.
//! Producers should export exact prefixes as `pub const`s next to the
//! `tracing::debug!` site (e.g. `REFRESH_SCAN_LOG_PREFIX`) so tests never
//! duplicate the strings.
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
/// Extracts the formatted `message` field of one event.
/// Pulls the formatted `message` field off one event.
#[derive(Default)]
struct MessageVisitor(String);
@@ -20,8 +19,8 @@ impl tracing::field::Visit for MessageVisitor {
}
}
/// A `tracing_subscriber::Layer` counting, per registered prefix, the events
/// whose `message` starts with it. Clones share the counts.
/// Layer that counts, per registered prefix, events whose `message` starts
/// with it. Clones share the same counters.
#[derive(Clone)]
pub struct MessagePrefixCounter {
counters: Arc<Vec<(&'static str, AtomicUsize)>>,
@@ -34,8 +33,7 @@ impl MessagePrefixCounter {
}
}
/// Events counted so far for `prefix`. Panics on a prefix that was never
/// registered — that is a bug in the test, not a zero count.
/// Count for `prefix`. Panics if `prefix` was never registered.
pub fn count(&self, prefix: &str) -> usize {
self.counters
.iter()
@@ -62,9 +60,9 @@ impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for MessagePrefixCount
}
}
/// Install a **thread-scoped** default subscriber counting `prefixes`; hold
/// the guard for the test's lifetime. Only observes events emitted on the
/// current thread — tasks under test must run on a current-thread runtime.
/// Thread-scoped default subscriber counting `prefixes`. Hold the guard for
/// the test lifetime. Only sees events on the current thread — use a
/// current-thread runtime for the subject under test.
pub fn install_prefix_counter_thread(
prefixes: &[&'static str],
) -> (tracing::subscriber::DefaultGuard, MessagePrefixCounter) {
@@ -74,12 +72,10 @@ pub fn install_prefix_counter_thread(
(tracing::subscriber::set_default(subscriber), counter)
}
/// Install the **process-global** subscriber counting `prefixes` — for tests
/// whose subject spawns its own threads/runtimes. Panics if a global
/// subscriber already exists: the test binary must own it.
/// Process-global subscriber counting `prefixes` — for subjects that spawn
/// their own threads/runtimes. Panics if a global subscriber already exists.
///
/// `stderr_env_filter` additionally tees formatted logs matching the given
/// `EnvFilter` directive to stderr (local debugging).
/// `stderr_env_filter` optionally tees matching formatted logs to stderr.
pub fn install_prefix_counter_global(
prefixes: &[&'static str],
stderr_env_filter: Option<&str>,
@@ -4,12 +4,9 @@ use std::collections::HashMap;
use serde::{Deserialize, Serialize};
/// Per-tool wire-traveling capabilities. Defaults conservatively (no
/// progress, no cancel, single concurrency, no hooks).
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct ToolCapabilities {
/// Streaming declaration. `None` — the default for every tool today —
/// means the tool never emits partial-result progress.
/// `None` means the tool never emits partial-result progress.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub streaming: Option<StreamingSpec>,
@@ -17,8 +14,7 @@ pub struct ToolCapabilities {
#[serde(default)]
pub supports_cancel: bool,
/// Maximum concurrent invocations the tool will accept. `None` is
/// unlimited.
/// `None` is unlimited.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub max_concurrency: Option<u32>,
@@ -26,7 +22,6 @@ pub struct ToolCapabilities {
#[serde(default)]
pub is_read_only: bool,
/// Lifecycle hooks the tool opts in to receive.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub hooks: Vec<HookKind>,
@@ -43,17 +38,14 @@ pub struct ToolCapabilities {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timeout_ms: Option<u64>,
/// Multi-agent write-coordination scope. Tools that mutate external
/// state must declare `Write` so the computer hub routes them to the
/// leader agent only. Absence is treated as `Read`.
/// Absence is treated as [`ToolScope::Read`].
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_scope: Option<ToolScope>,
}
/// How a tool streams partial results. Declared once in
/// [`ToolCapabilities::streaming`] and consumed at the source to stamp a
/// self-describing progress envelope; downstream layers dispatch on that
/// envelope rather than the tool's identity.
/// How a tool streams partial results. The spec is stamped onto a
/// self-describing progress envelope at the source, so downstream layers
/// dispatch on the envelope rather than on the tool's identity.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StreamingSpec {
/// Stable snake_case discriminator the tool stamps on its
@@ -67,7 +59,6 @@ pub struct StreamingSpec {
pub max_delta_bytes: Option<u32>,
}
/// Lifecycle hook a tool may opt in to receive.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum HookKind {
@@ -86,21 +77,19 @@ pub enum HookKind {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolScope {
/// Tool does not mutate external state.
Read,
/// Tool mutates external state.
Write,
}
/// Per-tool notification schemas. Keys are the notification `kind` strings
/// the computer hub validates against.
/// Keys are the notification `kind` strings the computer hub validates
/// against.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct NotificationSchemas {
/// Schemas for notifications the tool emits to subscribers.
/// Notifications the tool emits to subscribers.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub outbound: HashMap<String, serde_json::Value>,
/// Schemas for notifications the harness sends to the tool.
/// Notifications the harness sends to the tool.
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub inbound: HashMap<String, serde_json::Value>,
}
@@ -2,8 +2,8 @@
use serde::{Deserialize, Serialize};
/// Role of a WebSocket connection. The computer hub uses this to decide
/// which methods are valid on a given socket.
/// Role of a WebSocket connection; the computer hub decides from it which
/// methods are valid on a given socket.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ConnectionKind {
@@ -13,9 +13,6 @@ pub enum ConnectionKind {
/// How the computer hub exposes the registered tool set to the model.
///
/// `Concise` carries a configurable meta-tool pair so callers can choose
/// the model-facing names of the search/invoke meta-tools per session.
///
/// Wire form is adjacently tagged on `mode`: `Full` serialises as
/// `{"mode": "full"}` (an object, not a bare string), and `Concise` as
/// `{"mode": "concise", "meta_search": "...", "meta_call": "..."}`.
@@ -79,7 +79,6 @@ impl JsonRpcId {
Self::String(s.into())
}
/// Build a fresh UUID v7-backed id.
pub fn new_uuid_v7() -> Self {
Self::String(uuid::Uuid::now_v7().to_string())
}
@@ -4,9 +4,8 @@
//! than the numeric JSON-RPC `error.code`. The numeric is the JSON-RPC
//! envelope code; the string is the Kigi stable identifier.
//!
//! Implemented as a `&'static [(i32, &'static str)]` table; the table is
//! a small fixed set so a linear scan is faster than any
//! `HashMap`/`OnceLock`-shaped alternative.
//! The mapping is a flat table scanned linearly: the set is small and fixed,
//! so that beats any `HashMap`/`OnceLock`-shaped alternative.
use serde::{Deserialize, Serialize};
@@ -44,24 +43,22 @@ pub const ERROR_CODES: &[(i32, &str)] = &[
(-32099, "rate_limited"),
];
/// Returns `None` for strings not in the table. Receivers should fall
/// back to `-32603 internal_error` for unknown strings.
/// Receivers should fall back to `-32603 internal_error` for strings that
/// are not in the table.
pub fn numeric_for(code_str: &str) -> Option<i32> {
ERROR_CODES
.iter()
.find_map(|(n, s)| (*s == code_str).then_some(*n))
}
/// Returns `None` for codes not in the table.
pub fn string_for(code: i32) -> Option<&'static str> {
ERROR_CODES
.iter()
.find_map(|(n, s)| (*n == code).then_some(*s))
}
/// Numeric code most-appropriate for a [`ToolErrorWire`] variant.
/// `Custom` always maps to `-32603 internal_error` since its `code`
/// string is not in the table by definition.
/// `Custom` always maps to `-32603 internal_error`, since by definition its
/// `code` string is not in the table.
pub fn from_tool_error_wire(err: &ToolErrorWire) -> i32 {
match err {
ToolErrorWire::ToolNotFound { .. } => -32011,
@@ -138,7 +135,6 @@ pub struct WorkspaceUnavailableDetails {
pub retryable: bool,
}
/// Build the recognizable "workspace gone" error as a [`ToolErrorWire::Custom`].
pub fn workspace_unavailable_wire(
reason: WorkspaceGoneReason,
phase: WorkspaceGonePhase,
@@ -222,7 +218,6 @@ mod tests {
) else {
panic!("expected Custom variant");
};
// Exact, tenant-data-free contract.
assert_eq!(message, WORKSPACE_UNAVAILABLE_MESSAGE);
}
@@ -264,7 +259,8 @@ mod tests {
#[test]
fn custom_variant_tolerates_unknown_future_details_shape() {
// An unknown subcode + richer future details must still deserialize rather than failing the frame.
// An unknown subcode carrying richer details must still deserialize
// rather than failing the whole frame.
let future = json!({
"code": "custom",
"subcode": "some_future_subcode",
@@ -290,7 +286,6 @@ mod tests {
.is_some(),
"unknown details fields are preserved",
);
// Re-serialization preserves the unknown fields.
let reser = serde_json::to_value(&wire).unwrap();
assert_eq!(
reser["details"]["extra_new_field"]["nested"],
@@ -52,9 +52,7 @@ pub enum ToolErrorWire {
#[error("behavior_version unsupported")]
BehaviorVersionUnsupported { tool_id: ToolId, requested: String },
/// Render-card budget exceeded for the current session. `card_id`
/// carries the offending render-card identifier when known; `reason`
/// is a free-form human-readable explanation.
/// Render-card budget exceeded for the current session.
#[error("render limited for {tool_id}: {reason}")]
RenderLimited {
tool_id: ToolId,
@@ -74,10 +72,9 @@ pub enum ToolErrorWire {
Internal {
#[serde(default, skip_serializing_if = "Option::is_none")]
request_id: Option<RequestId>,
/// Bounded, human-readable cause of the internal error. Optional for
/// wire compatibility with older peers; producers SHOULD populate it
/// (truncated at the producer) so receivers can distinguish failure
/// modes without correlating server logs.
/// Optional for wire compatibility with older peers; producers SHOULD
/// populate it (truncated at the producer) so receivers can distinguish
/// failure modes without correlating server logs.
#[serde(default, skip_serializing_if = "Option::is_none")]
detail: Option<String>,
},
+27 -42
View File
@@ -18,7 +18,7 @@ use crate::{
output_wire::ToolOutputWire,
};
// ── Tool call params / result / progress ─────────────────────────────────
// Tool call params / result / progress
/// `tool.call` (harness → service) and `tool_call_request` (service →
/// tool_server) share the same params shape; `tool_call_id` is preserved
@@ -59,7 +59,7 @@ pub struct ToolCallResult {
pub chat_completion_output: Option<serde_json::Value>,
}
// ── Trace donation ────────────────────────────────────────────────────────
// Trace donation
/// Hub rejects oversized batches wholesale; donors chunk before encoding.
pub const MAX_SPANS_PER_DONATION: usize = 512;
@@ -77,7 +77,7 @@ pub struct TracesDonateParams {
pub otlp_request: String,
}
// ── Log donation ──────────────────────────────────────────────────────────
// Log donation
/// Hub rejects oversized batches wholesale; donors chunk before encoding.
/// The 1 MiB [`MAX_DONATION_BYTES`] decoded-size cap is the real bound; this
@@ -94,7 +94,7 @@ pub struct LogsDonateParams {
pub otlp_request: String,
}
// ── Metric donation ───────────────────────────────────────────────────────
// Metric donation
/// Hub rejects oversized batches wholesale; donors chunk before encoding.
/// Secondary guard alongside the 1 MiB [`MAX_DONATION_BYTES`] decoded-size cap.
@@ -189,7 +189,7 @@ pub struct SystemNotifyParams {
pub request_id: Option<String>,
}
// ── Registration frames ──────────────────────────────────────────────────
// Registration frames
/// `register_tool` params — single-tool sugar over `register_server`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -217,7 +217,7 @@ pub struct UnregisterServerParams {
pub server_id: ServerId,
}
// ── Per-tool session binding ───────────────────────────────────────────────
// Per-tool session binding
/// `bind_tool_session` params — add `session_id` to a registered tool's
/// per-tool session set.
@@ -229,7 +229,6 @@ pub struct UnregisterServerParams {
/// typically omitted on connection-control frames.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BindToolSessionParams {
/// The tool whose session set is being mutated.
pub tool_id: ToolId,
/// The session id to add to the tool's session set. Must already
/// be in the connection's bound-session set.
@@ -249,7 +248,7 @@ pub struct BindToolSessionParams {
/// `ServerError::ToolBindingConflict` (-32600) so the contended caller
/// sees a wire-level error frame with a dedicated code instead of a
/// quietly-buried ack outcome — mirroring it would re-introduce the
/// `UnknownTool`-overload ambiguity the dedicated code was added to fix.
/// `UnknownTool`-overload ambiguity the dedicated code exists to fix.
/// - `SessionNotBound` is router-injected by the per-frame envelope
/// pre-check (the connection's bound-session set is router state, not
/// registry state) and never originates from the registry call.
@@ -285,7 +284,6 @@ pub struct BindToolSessionAck {
/// the calling-frame routing scope and serves a different concept.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct UnbindToolSessionParams {
/// The tool whose session set is being mutated.
pub tool_id: ToolId,
/// The session id to remove from the tool's session set.
pub session_id: SessionId,
@@ -309,7 +307,7 @@ pub struct UnbindToolSessionAck {
pub outcome: ToolSessionUnbindOutcome,
}
// ── Server discovery + binding ────────────────────────────────────────────
// Server discovery + binding
/// `servers.list` params — discover available tool servers for the
/// authenticated user.
@@ -345,7 +343,6 @@ pub struct ServersListResult {
pub struct ServerBindParams {
/// Which tool server to bind (its server_id from `servers.list`).
pub server_id: ServerId,
/// The harness session to bind tools to.
pub session_id: SessionId,
}
@@ -355,12 +352,12 @@ pub struct ServerBindParams {
pub enum ServerBindOutcome {
/// Tools successfully bound to the session.
Bound,
/// Tools were already bound to this session.
/// Tools already bound to this session.
AlreadyBound,
/// No server with this server_id found.
ServerNotFound,
/// A server was located and the bind forwarded, but it did not complete:
/// the ack timed out, the transport send/delivery failed, or the ack was
/// Server found and bind forwarded, but bind did not complete:
/// ack timed out, transport send/delivery failed, or ack was
/// malformed or an explicit error. Distinct from `ServerNotFound`, which
/// means no such server is registered.
Unavailable,
@@ -393,7 +390,7 @@ pub struct ServerUnbindAck {
pub outcome: ServerUnbindOutcome,
}
// ── List & search ────────────────────────────────────────────────────────
// List & search
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolsListParams {
@@ -432,7 +429,7 @@ pub struct ToolsSearchResultBody {
pub is_ready: bool,
}
// ── Session lifecycle ────────────────────────────────────────────────────
// Session lifecycle
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SessionOpenParams {
@@ -467,9 +464,8 @@ pub struct SessionOpenResult {}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SessionBindServerParams {
pub server_id: ServerId,
/// Working directory for the session. The tool server creates a
/// session rooted at this path. When absent, the server's default
/// CWD is used.
/// Working directory for the session. When absent, the tool server's
/// default CWD is used.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cwd: Option<String>,
/// Opaque metadata passed through to the tool server (sandbox_id,
@@ -543,7 +539,7 @@ pub enum AttachRoute {
Unknown,
}
// ── Simplified lifecycle ─────────────────────────────────────────────────
// Simplified lifecycle
/// `serve` params (server → hub). Full tool snapshot for a session.
///
@@ -558,13 +554,12 @@ pub struct ServeParams {
/// Reply to [`ServeParams`].
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct ServeResult {
/// Number of tools accepted (informational).
#[serde(default)]
pub accepted: usize,
/// Tool IDs that were added relative to the previous snapshot.
/// Tool IDs added relative to the previous snapshot.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub added: Vec<ToolId>,
/// Tool IDs that were removed relative to the previous snapshot.
/// Tool IDs removed relative to the previous snapshot.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub removed: Vec<ToolId>,
}
@@ -606,7 +601,7 @@ pub struct SessionBindResult {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionUnbindParams {}
// ── Subscriptions ────────────────────────────────────────────────────────
// Subscriptions
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SubscribeNotificationsParams {
@@ -673,11 +668,11 @@ pub struct UnsubscribeNotificationsParams {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UnsubscribeOutcome {
/// Subscription was present and was removed.
/// Subscription was present and the service removed it.
Unsubscribed,
/// Subscription was not present; no-op.
NotSubscribed,
/// Subscription was removed by the service because the subscriber's
/// The service removed the subscription because the subscriber's
/// outbound mpsc was full or dropped during fan-out.
Evicted,
}
@@ -690,7 +685,7 @@ pub struct UnsubscribeAck {
pub subscription_id: String,
}
// ── Hooks ────────────────────────────────────────────────────────────────
// Hooks
/// `hook` frame body, routed in both directions through the hub: harness →
/// tool-server for forward hooks (e.g. `Cancel`, `SessionEnded`), and
@@ -813,7 +808,7 @@ pub struct HookReplyFrame {
pub result: serde_json::Value,
}
// ── Service → harness pushes ─────────────────────────────────────────────
// Service → harness pushes
/// `tools_changed` body — the active tool set for `session_id` changed.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -827,7 +822,7 @@ pub struct ToolsChanged {
pub updated: Vec<ToolId>,
}
// ── Tool server status lifecycle ──────────────────────────────────────
// Tool server status lifecycle
/// Lifecycle status of a tool server connection.
///
@@ -937,7 +932,7 @@ pub enum ToolServerDisconnectReason {
ConnectionLost,
}
// ── Heartbeat ────────────────────────────────────────────────────────────
// Heartbeat
//
// PingFrame / PongFrame carry a `method` discriminator on the wire so
// any receiver (hub or SDK) can route them through a method-based demux.
@@ -975,7 +970,7 @@ impl PongFrame {
}
}
// -- Custom Serialize: always includes `"method"` on the wire. -----------
// Custom Serialize: always includes `"method"` on the wire.
impl serde::Serialize for PingFrame {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
@@ -997,7 +992,7 @@ impl serde::Serialize for PongFrame {
}
}
// -- Custom Deserialize: accepts with or without `method` for compat. ----
// Custom Deserialize: accepts with or without `method` for compat.
impl<'de> serde::Deserialize<'de> for PingFrame {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
@@ -1068,8 +1063,6 @@ mod tests {
ToolCallId::new_v7()
}
// ── HookFrame constructors ──────────────────────────────────────
#[test]
fn hook_cancel_sets_tool_and_call_ids() {
let hook = HookFrame::cancel(sid(), tid(), cid());
@@ -1123,8 +1116,6 @@ mod tests {
assert_eq!(hook, back);
}
// ── ToolNotificationFrame constructors ───────────────────────────
#[test]
fn notification_custom_sets_wire_shape() {
let frame = ToolNotificationFrame::custom(tid(), "echo.status", json!({"status": "idle"}));
@@ -1187,8 +1178,6 @@ mod tests {
assert_eq!(back, params);
}
// ── Donation params ────────────────────────────────────────────
#[test]
fn logs_donate_params_round_trips() {
let params = super::LogsDonateParams {
@@ -1211,8 +1200,6 @@ mod tests {
assert_eq!(back, params);
}
// ── ToolServerStatusPayload ────────────────────────────────────
#[test]
fn tool_server_lifecycle_status_serde_snake_case() {
let status = super::ToolServerLifecycleStatus::ShuttingDown;
@@ -1470,8 +1457,6 @@ mod tests {
assert_eq!(frame, back);
}
// ── hook_id backward-compat ─────────────────────────────────────
#[test]
fn hook_frame_missing_hook_id_deserializes_as_none() {
let v = json!({
@@ -4,9 +4,8 @@ use serde::{Deserialize, Serialize};
use crate::{ConnectionId, ConnectionKind, ServerId, UserId};
/// Wire-protocol version both ends speak. Bumped when an incompatible
/// schema change lands; minor additions go through capability
/// negotiation rather than a version bump.
/// Bumped only for incompatible schema changes; additive changes go through
/// capability negotiation instead.
pub const PROTOCOL_VERSION: &str = "1.0.0";
/// First frame sent by the client after the WebSocket upgrade succeeds.
@@ -14,15 +13,13 @@ pub const PROTOCOL_VERSION: &str = "1.0.0";
/// No session ids are carried at handshake time. The connection starts with
/// an empty bound-session set and binds sessions dynamically over its
/// lifetime via `register_session` / `unregister_session` JSON-RPC calls.
///
/// Tool-server connections carry `server_id` so the hub can
/// identify the server without a separate `register_server` call.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct HelloMsg {
pub protocol_version: String,
pub kind: ConnectionKind,
/// Stable server identity. Only set for
/// [`ConnectionKind::ToolServer`] connections.
/// Stable server identity, set only for [`ConnectionKind::ToolServer`]
/// connections, so the hub can identify the server without a separate
/// `register_server` call.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_id: Option<ServerId>,
/// One-line server description for `servers.list`.
+1 -4
View File
@@ -2,13 +2,10 @@
use serde::{Deserialize, Serialize};
/// Internally-tagged hook payload. New variants land alongside `Custom`,
/// which keeps unknown future kinds round-trippable.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum HookEvent {
/// Cancel an in-flight call. The owning `tool_call_id` travels in the
/// enclosing `hook` frame.
/// The `tool_call_id` this cancels travels in the enclosing `hook` frame.
Cancel,
Pause,
Resume,
@@ -11,7 +11,6 @@ use std::str::FromStr;
use serde::{Deserialize, Deserializer, Serialize};
/// Errors produced by id constructors and validators.
#[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
pub enum IdError {
#[error("identifier must not be empty")]
@@ -56,7 +55,6 @@ macro_rules! opaque_id {
pub struct $name(String);
impl $name {
/// Construct, validating the id's invariants.
pub fn new(value: impl Into<String>) -> Result<Self, IdError> {
let value = value.into();
ensure_non_empty(&value)?;
@@ -132,7 +130,6 @@ opaque_id!(
);
impl ToolCallId {
/// Generate a fresh UUID v7-backed `ToolCallId`.
pub fn new_v7() -> Self {
Self(uuid::Uuid::now_v7().to_string())
}
@@ -209,9 +206,6 @@ opaque_id!(
/// Per-connection monotonic notification sequence (starts at 0 on every new
/// connection).
///
/// The inner `u64` is private so `new`, `From<u64>`, and `Default` are the
/// only construction paths.
#[derive(
Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Default, Serialize, Deserialize,
)]
@@ -1,10 +1,4 @@
//! Tool wire-protocol types.
//!
//! Identifier newtypes, registration payloads, capabilities, hook events,
//! handshake messages, the JSON-RPC 2.0 envelope and method catalog, the
//! `ToolErrorWire` / `ToolOutputWire` / `WireToolNotification` wire enums,
//! every method's `params` / `result` payload struct, and the numeric ↔
//! string error-code mapping.
#![forbid(unsafe_code)]
@@ -123,7 +123,7 @@ define_methods! {
ToolServerGetStatus => "tool_server.get_status",
ToolServerEvict => "tool_server.evict",
// ── Session lifecycle ───────────────────────────────────────────
// Session lifecycle
/// Full tool snapshot for a session (server → hub). Idempotent:
/// re-sending replaces the tool set; the hub diffs and emits
@@ -16,7 +16,6 @@
use serde::{Deserialize, Serialize};
/// Adjacent-tagged notification wire wrapper.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "shape", content = "value", rename_all = "snake_case")]
pub enum WireToolNotification {
@@ -72,9 +71,8 @@ pub const fn known_notification_kinds() -> &'static [&'static str] {
}
/// Reject custom notification kinds whose name shadows a known PascalCase
/// variant. Runs at notification-emit time; an empty `kind` is accepted
/// here (the producer is responsible for validating that the field is
/// non-empty).
/// variant. An empty `kind` is accepted here; the producer is responsible
/// for validating that the field is non-empty.
pub fn check_custom_kind(kind: &str) -> Result<(), KnownVariantCollision> {
if KNOWN_NOTIFICATION_KINDS.contains(&kind) {
Err(KnownVariantCollision {
@@ -14,9 +14,8 @@ pub enum TransportKind {
Remote,
}
/// A single tool's wire description plus optional schema and capability
/// metadata. The `tool_id` is **not** stored explicitly — it is derived
/// from `description.{namespace, name}` via [`Self::derive_tool_id`].
/// The `tool_id` is **not** carried explicitly — it is derived from
/// `description.{namespace, name}` via [`Self::derive_tool_id`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolDescriptionWithSchema {
pub description: kigi_tool_types::ToolDescription,
@@ -29,8 +28,6 @@ pub struct ToolDescriptionWithSchema {
}
impl ToolDescriptionWithSchema {
/// Derive the canonical `ToolId`.
///
/// Namespaced descriptions render as `"{namespace}:{name}"`; otherwise
/// the bare `name`. The result is run through [`ToolId::new`], so an
/// invalid name or namespace surfaces as an [`IdError`].
@@ -69,8 +66,7 @@ pub struct ToolRegistration {
/// enforces this at register-tool time and rejects mismatches with
/// `InvalidRequest`.
pub tool_id: ToolId,
/// Per-tool session set. See struct doc-comment for the
/// `None` / `Some(vec![])` / `Some(vec![...])` semantics.
/// See the struct doc-comment for the three-state semantics.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sessions: Option<Vec<SessionId>>,
pub user_id: UserId,
@@ -95,9 +91,6 @@ pub struct ToolRegistration {
}
impl ToolRegistration {
/// Derive the canonical `ToolId` from `description.{namespace, name}`.
/// The `tool_id` payload field MUST equal this value; the IC service
/// router enforces the invariant at register-tool time.
pub fn derive_tool_id(&self) -> Result<ToolId, IdError> {
match &self.description.namespace {
Some(ns) => ToolId::new(format!("{ns}:{}", self.description.name)),
@@ -110,17 +103,12 @@ impl ToolRegistration {
/// `sessions` value; per-tool outcomes are reported individually via
/// [`RegistrationOutcome`].
///
/// `sessions` follows the same three-state semantics as
/// [`ToolRegistration::sessions`]: `None` means "no change" (preserves
/// existing per-tool session bindings on a re-register), `Some(vec![])`
/// means "unbind every session for every tool in this batch", and
/// `Some(vec![...])` means "replace each tool's session set with
/// exactly these ids".
/// `sessions` follows the three-state semantics documented on
/// [`ToolRegistration`], applied to every tool in the batch.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ToolServerRegistration {
pub server_id: ServerId,
/// Per-batch session set. See struct doc-comment for `None` /
/// `Some(vec![])` / `Some(vec![...])` semantics.
/// See the struct doc-comment for the three-state semantics.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sessions: Option<Vec<SessionId>>,
pub user_id: UserId,
@@ -17,8 +17,6 @@ pub enum RegistryError {
#[error("tool already registered: {tool_id}")]
AlreadyRegistered { tool_id: ToolId },
/// The registration's session does not match the connection's bound
/// session.
#[error("session mismatch: token session={token_session}, registration session={reg_session}")]
SessionMismatch {
token_session: SessionId,
@@ -41,7 +39,7 @@ pub enum RegistryError {
#[error("invalid description: {message}")]
InvalidDescription { message: String },
/// `if_match_generation` precondition failed.
/// The `if_match_generation` precondition failed.
#[error("stale generation: expected={expected}, actual={actual}")]
StaleGeneration { expected: u64, actual: u64 },
}
@@ -98,8 +98,6 @@ mod tests {
use super::*;
use serde_json::json;
// ── SessionEvent round-trip tests ────────────────────────────────
#[test]
fn turn_started_round_trip() {
let event = SessionEvent::TurnStarted {
@@ -226,8 +224,6 @@ mod tests {
}
}
// ── #[serde(other)] backward-compat ─────────────────────────────
#[test]
fn unknown_event_type_deserializes_as_unknown() {
let v = json!({ "event_type": "some_future_event", "extra": 123 });
@@ -242,8 +238,6 @@ mod tests {
assert_eq!(event, SessionEvent::Unknown);
}
// ── ToolCallOutcome serialization ───────────────────────────────
#[test]
fn tool_call_outcome_snake_case() {
for (variant, expected) in [
@@ -259,8 +253,6 @@ mod tests {
}
}
// ── SessionPhase serialization ──────────────────────────────────
#[test]
fn session_phase_snake_case() {
for (variant, expected) in [
@@ -277,8 +269,6 @@ mod tests {
}
}
// ── Forward-compat: inner enum Unknown ──────────────────────────
#[test]
fn tool_call_outcome_unknown_variant_on_future_value() {
let back: ToolCallOutcome = serde_json::from_value(json!("timeout")).unwrap();
@@ -327,16 +317,12 @@ mod tests {
);
}
// ── Unknown variant serialization ───────────────────────────────
#[test]
fn unknown_variant_serializes_as_expected() {
let v = serde_json::to_value(SessionEvent::Unknown).unwrap();
assert_eq!(v, json!({"event_type": "unknown"}));
}
// ── Extra/unknown fields on known variants ──────────────────────
#[test]
fn extra_fields_ignored_on_known_variant() {
let v = json!({
@@ -356,8 +342,6 @@ mod tests {
);
}
// ── Negative: missing required fields ───────────────────────────
#[test]
fn turn_ended_missing_required_field_rejected() {
let v = json!({
@@ -371,8 +355,6 @@ mod tests {
assert!(serde_json::from_value::<SessionEvent>(v).is_err());
}
// ── Boundary values ─────────────────────────────────────────────
#[test]
fn turn_number_zero_and_max() {
for turn_number in [0, u64::MAX] {
@@ -42,8 +42,8 @@ pub struct BeforeTurnPayload {
/// Whether the session is in YOLO / auto-approve mode.
#[serde(default)]
pub yolo_mode: bool,
// ── Extended fields (workspace mirrors these into `events.jsonl`);
// all `#[serde(default)]` for old-shell / old-workspace interop. ──
// Extended fields (workspace mirrors these into `events.jsonl`);
// all `#[serde(default)]` for old shell / old workspace interop.
/// Mirrors `Event::TurnStarted::conversation_message_count`.
#[serde(default)]
pub conversation_message_count: usize,
@@ -86,7 +86,6 @@ impl Default for BeforeTurnPayload {
pub struct AfterTurnPayload {
/// Same turn counter as the preceding `before_turn`.
pub turn_number: u64,
/// High-level outcome of the turn.
pub outcome: TurnHookOutcome,
/// Wall-clock duration of the turn in milliseconds.
pub duration_ms: u64,
@@ -128,7 +127,6 @@ pub enum TurnHookOutcome {
Completed,
/// Turn was cancelled by the user (Ctrl+C / abort).
Cancelled,
/// Turn ended due to an error.
Error,
}
@@ -152,9 +150,7 @@ pub enum TurnHookRequest {
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum InjectionRole {
/// Append as a system turn.
System,
/// Append as a developer turn.
Developer,
/// Append as a user turn (e.g. a `<system-reminder>`-wrapped message).
User,
@@ -164,9 +160,7 @@ pub enum InjectionRole {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct HookInjection {
/// Role to append the content as.
pub role: InjectionRole,
/// Verbatim turn content.
pub content: String,
}
@@ -188,10 +182,8 @@ pub enum TurnControl {
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct HookReply {
/// Turns to append before the next sampling step, in order.
#[serde(default)]
pub injections: Vec<HookInjection>,
/// Optional loop-control override.
#[serde(default)]
pub control: TurnControl,
/// Artifact-handling ack for a [`TurnHookRequest::After`] request; `None`
@@ -202,7 +202,6 @@ fn jsonrpc_id_round_trips_to_request_id_correlator() {
let envelope_id = JsonRpcId::from_request_id(&original);
assert_eq!(envelope_id.as_request_id().unwrap(), original);
// Numeric ids are stringified.
let nid = JsonRpcId::Number(7);
assert_eq!(nid.as_request_id().unwrap().as_str(), "7");
}
@@ -234,10 +233,9 @@ fn full_call_envelope_serialises_to_expected_shape() {
assert_eq!(v["params"]["tool_call_id"], json!("call_xyz"));
}
/// The envelope-level `session_id` and an inner `params.session_id` (e.g.
/// on `ToolsListParams`) are independent keys in the wire JSON tree.
/// This test pins that invariant so a refactor that accidentally
/// collapses the two layers (e.g. via `#[serde(flatten)]`) fails loudly.
/// The envelope-level `session_id` and an inner `params.session_id` are
/// independent keys in the wire JSON tree; a refactor that collapses the two
/// layers (e.g. via `#[serde(flatten)]`) must fail here.
#[test]
fn envelope_session_id_and_inner_params_session_id_are_distinct_layers() {
use kigi_tool_protocol::{ToolDefinitionMode, ToolsListParams};
@@ -1002,9 +1002,9 @@ fn session_lifecycle_payloads_round_trip() {
fn attach_route_round_trips_snake_case_and_tolerates_unknown() {
assert_eq!(roundtrip(&AttachRoute::Local), json!("local"));
assert_eq!(roundtrip(&AttachRoute::Remote), json!("remote"));
// "restored" was removed with restore-on-activity; old hubs may still send
// it, and it must fall into the tolerant `Unknown` bucket like any other
// retired/newer value.
// "restored" is a retired restore-on-activity route; old hubs may still
// send it, and it must fall into the tolerant `Unknown` bucket like any
// other retired/newer value.
let parsed: AttachRoute =
serde_json::from_value(json!("restored")).expect("tolerant parse of retired value");
assert_eq!(parsed, AttachRoute::Unknown);
+6 -10
View File
@@ -85,13 +85,11 @@ impl ToolCallContext {
}
}
/// Delegate to `self.extensions.insert()`.
pub fn insert<T: Send + Sync + 'static>(&mut self, value: T) -> &mut Self {
self.extensions.insert(value);
self
}
/// Delegate to `self.extensions.get()`.
pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
self.extensions.get::<T>()
}
@@ -129,9 +127,8 @@ pub struct BehaviorVersion(pub String);
#[derive(Clone, Debug)]
pub struct TraceContext(pub String);
/// Session ID context — identifies which hub session this call belongs to.
/// Used by multi-session tool servers to dispatch to the correct
/// per-session state.
/// Session ID which hub session this call belongs to.
/// Multi-session tool servers dispatch to the matching per-session state.
#[derive(Clone, Debug)]
pub struct SessionContext(pub String);
@@ -141,11 +138,10 @@ pub struct SessionContext(pub String);
#[derive(Clone, Debug)]
pub struct Cancellation(pub tokio_util::sync::CancellationToken);
/// Per-user feature-flag bag attached as a [`ToolCallContext`] extension.
/// Dispatcher resolves; tools read. Default = "off" for every field so an
/// absent extension never accidentally opts a feature in. Extend by
/// adding fields with safe defaults; new fields need `#[serde(default)]`
/// so older `session.bind` payloads stay deserializable.
/// Per-user feature-flag bag on [`ToolCallContext`]. Dispatcher resolves;
/// tools read. Default is off for every field so an absent extension never
/// opts a feature in. New fields need `#[serde(default)]` so older
/// `session.bind` payloads stay deserializable.
#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
pub struct WorkspaceViewerContext {
/// When `true`, `BashTool` emits `bash_output_chunk` Progress frames.
+13 -17
View File
@@ -1,15 +1,13 @@
//! Object-safe `ToolDispatch` trait — the runtime contract for handling tool calls.
//!
//! `Tool` itself is not object-safe (it carries associated `Args` /
//! `Output` types), so implementations expose a JSON-typed surface and rely on
//! per-tool adapters to encode/decode at the boundary. The default
//! `call_terminal` impl drains the stream so the common "I just want the
//! result" path doesn't have to depend on `futures` internals.
//! `Tool` itself is not object-safe (associated `Args` / `Output` types), so
//! implementations expose a JSON-typed surface and rely on per-tool adapters
//! to encode/decode at the boundary. The default `call_terminal` impl drains
//! the stream so callers that only need the result avoid `futures` internals.
//!
//! This crate is upstream of every concrete impl. Doc-comments here describe
//! trait semantics in terms of "the runtime" or "the implementation" —
//! concrete dispatch routers live downstream and are intentionally not named
//! here.
//! This crate is upstream of every concrete impl. Docs here describe trait
//! semantics in terms of "the runtime" or "the implementation" — concrete
//! dispatch routers live downstream and are intentionally not named here.
use async_trait::async_trait;
use futures::StreamExt;
@@ -23,7 +21,7 @@ use crate::tool::{ToolStream, ToolStreamItem, TypedToolOutput};
/// Object-safe tool dispatch interface.
///
/// Implementations route the `tool_id` to the correct tool, decode `args`
/// Implementations route `tool_id` to the correct tool, decode `args`
/// against the tool's typed `Args`, and return the streaming result as
/// [`TypedToolOutput`] — preserving model-facing content blocks and
/// optional chat-completion metadata end-to-end. Raw `Value` only appears
@@ -39,14 +37,12 @@ pub trait ToolDispatch: Send + Sync {
ctx: ToolCallContext,
) -> ToolStream<TypedToolOutput>;
/// Drain the stream and return only the terminal result. Useful for
/// callers that don't care about progress chunks.
/// Drain the stream and return only the terminal result.
///
/// Default impl pulls items off the stream and discards `Progress`
/// items; the first `Terminal` short-circuits. A stream that ends
/// without a `Terminal` is a protocol violation by the implementation;
/// the default surfaces this as `ToolError::Custom { code:
/// "stream_no_terminal", ... }`.
/// Default impl discards `Progress` items and short-circuits on the
/// first `Terminal`. A stream that ends without a `Terminal` is a
/// protocol violation; the default surfaces
/// `ToolError::Custom { code: "stream_no_terminal", ... }`.
async fn call_terminal(
&self,
tool_id: ToolId,
+14 -28
View File
@@ -14,12 +14,12 @@ use serde_json::Value;
use kigi_tool_protocol::{ToolErrorWire, ToolId};
/// Discriminator for tool errors.
/// Machine-readable tool-error discriminator.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ToolErrorKind {
/// The tool has no implementation for the requested operation.
/// No implementation for the requested operation.
NotImplemented,
/// Inputs failed validation.
/// Input validation failed.
InvalidArguments,
/// No tool registered under the given id.
NotFound,
@@ -27,11 +27,10 @@ pub enum ToolErrorKind {
PermissionDenied,
/// Authentication failed (401-shaped).
Unauthorized,
/// The tool ran past its time budget.
/// Tool ran past its time budget.
Timeout,
/// The caller cancelled the tool call.
/// Caller cancelled the tool call.
Cancelled,
/// Rate limit exceeded.
RateLimited,
/// The caller's usage pool / billing balance is exhausted (out
/// of credits). Payment-required-shaped; distinct from
@@ -57,13 +56,10 @@ pub enum ToolErrorKind {
/// shed) so the surface can tailor a "too many in progress" message.
/// Named to match the chat surface's `concurrency_limit` typed error.
ConcurrencyLimit,
/// Upstream service unavailable.
ServiceUnavailable,
/// Network-level failure.
NetworkError,
/// Tool body returned an error.
Execution,
/// Requested behavior version not supported.
BehaviorVersionUnsupported,
/// Render-card budget exceeded.
RenderLimited,
@@ -151,12 +147,10 @@ impl std::error::Error for ToolError {
}
}
// ---------------------------------------------------------------------------
// Constructors — one per kind for ergonomic tool code
// ---------------------------------------------------------------------------
// Constructors
impl ToolError {
/// Core constructor. All other constructors delegate here.
/// Core constructor; other constructors delegate here.
pub fn new(kind: ToolErrorKind, detail: impl Into<String>) -> Self {
Self {
kind,
@@ -166,13 +160,12 @@ impl ToolError {
}
}
/// Attach structured metadata.
pub fn with_details(mut self, details: Value) -> Self {
self.details = Some(details);
self
}
/// Attach a causal error chain (for developer logs, not sent to model).
/// Attach a causal chain for developer logs (not sent to the model).
pub fn with_source(mut self, source: impl Into<anyhow::Error>) -> Self {
self.source = Some(source.into());
self
@@ -252,16 +245,13 @@ impl ToolError {
.with_details(serde_json::json!({ "code": code.into() }))
}
/// Snake-case identifier for the kind. Delegates to
/// [`ToolErrorKind::as_str`].
/// Snake-case identifier for the kind ([`ToolErrorKind::as_str`]).
pub fn variant_name(&self) -> &'static str {
self.kind.as_str()
}
}
// ---------------------------------------------------------------------------
// From impls
// ---------------------------------------------------------------------------
impl From<serde_json::Error> for ToolError {
fn from(value: serde_json::Error) -> Self {
@@ -269,9 +259,7 @@ impl From<serde_json::Error> for ToolError {
}
}
// ---------------------------------------------------------------------------
// Wire bridge
// ---------------------------------------------------------------------------
/// Carry a [`ToolError`]'s structured `details` onto a `Custom` wire variant
/// while keeping the round-trip recognizable: the decoder
@@ -451,9 +439,8 @@ mod wire_bridge_tests {
#[test]
fn service_unavailable_details_survive_wire_projection() {
// Structured details used to be dropped (`details: None`) for the
// Custom-mapped kinds; they must now ride the wire with the subcode
// merged in so recognizers keying on `details.code` keep working.
// Custom-mapped kinds carry structured details on the wire with the
// subcode merged in so recognizers keying on `details.code` keep working.
let err = ToolError::service_unavailable("sandbox not ready")
.with_details(serde_json::json!({ "retry_after_ms": 1500 }));
let wire = ToolErrorWire::from(err);
@@ -481,10 +468,9 @@ mod wire_bridge_tests {
#[test]
fn rate_limit_and_usage_kinds_merge_subcode_uniformly() {
// Same property as service_unavailable, applied to every
// Custom-mapped kind: object details without a `code` key gain the
// subcode, so decode-side recognizers keying on `details.code` can
// still classify the error.
// Same as service_unavailable for every Custom-mapped kind: object
// details without a `code` key gain the subcode so decode-side
// recognizers keying on `details.code` can still classify the error.
let cases: [(ToolError, &str); 5] = [
(ToolError::rate_limited("slow down"), "rate_limited"),
(
+2 -5
View File
@@ -1,10 +1,7 @@
//! Unified tool runtime contract.
//!
//! Single home for the `Tool` trait, `ToolDispatch`, `ToolError`,
//! `ToolNotification`, `ToolSearchIndex`, `ToolCallContext`, `ToolStream`,
//! the in-process `LocalRegistry`, and the helper constructors that build
//! well-formed streams. Adapters for individual tool sources re-export
//! from here so every tool author sees the same surface.
//! Adapters for individual tool sources re-export from here so every tool
//! author sees the same surface.
#![forbid(unsafe_code)]
@@ -38,14 +38,12 @@ impl std::fmt::Debug for LocalRegistry {
}
impl LocalRegistry {
/// Construct an empty registry.
pub fn new() -> Self {
Self::default()
}
/// Register a typed [`Tool`] implementation by value. Subsequent
/// registrations of the same id replace the previous handle and
/// return the displaced handle for inspection / drop ordering.
/// Register a typed [`Tool`] by value. A later registration of the
/// same id replaces the previous handle and returns it.
pub fn register<T>(&self, tool: T) -> Option<ArcTool>
where
T: Tool + 'static,
@@ -53,7 +51,7 @@ impl LocalRegistry {
self.register_arc(Arc::new(tool))
}
/// Register a typed [`Tool`] implementation already wrapped in `Arc`.
/// Register a typed [`Tool`] already wrapped in `Arc`.
pub fn register_arc<T>(&self, tool: Arc<T>) -> Option<ArcTool>
where
T: Tool + 'static,
@@ -64,47 +62,41 @@ impl LocalRegistry {
/// Register a type-erased [`ToolDyn`](crate::tool::ToolDyn) directly.
///
/// Use this for inherently dynamic tools (e.g. MCP tools retrieved
/// from a registry as `Arc<dyn ToolDyn>`) where the concrete type
/// is not available. For native tools with a concrete type, prefer
/// [`register`](Self::register).
/// Use for inherently dynamic tools (e.g. MCP tools as
/// `Arc<dyn ToolDyn>`) where the concrete type is unavailable. For
/// native tools with a concrete type, prefer [`register`](Self::register).
pub fn register_dyn(&self, tool: ArcTool) -> Option<ArcTool> {
let id = tool.id();
self.entries.write().insert(id, tool)
}
/// Resolve `tool_id` to its in-process handle, if registered.
/// Returns a clone of the handle so the caller can dispatch without
/// holding the lock across an await point.
/// Resolve `tool_id` to its in-process handle, if registered. Returns
/// a clone so the caller can dispatch without holding the lock across
/// an await point.
pub fn find(&self, tool_id: &ToolId) -> Option<ArcTool> {
self.entries.read().get(tool_id).cloned()
}
/// Drop the handle bound to `tool_id`. Returns `true` iff a
/// matching entry was removed.
/// Drop the handle bound to `tool_id`. Returns `true` iff a matching
/// entry was found and removed.
pub fn unregister(&self, tool_id: &ToolId) -> bool {
self.entries.write().shift_remove(tool_id).is_some()
}
/// Number of tools currently registered.
pub fn len(&self) -> usize {
self.entries.read().len()
}
/// `true` iff no tools are registered.
pub fn is_empty(&self) -> bool {
self.entries.read().is_empty()
}
/// `true` iff `tool_id` is currently registered.
pub fn contains(&self, tool_id: &ToolId) -> bool {
self.entries.read().contains_key(tool_id)
}
/// Descriptions of registered tools filtered by `should_list`.
///
/// Returns descriptions in **insertion order** — the order tools
/// were registered — so the caller sees the same ordering as the
/// Descriptions of registered tools filtered by `should_list`, in
/// **insertion order**, so the caller sees the same ordering as the
/// config-defined tool list.
pub fn list_tools(&self, ctx: &ListToolsContext) -> Vec<ToolDescription> {
self.entries
@@ -5,8 +5,7 @@
//! adapters can serialise them without enabling additional features.
//!
//! Each `ToolNotification` variant has a parallel `send_*` convenience on
//! [`ToolNotificationHandle`]. The two surfaces are kept in lockstep — when
//! adding a variant here, add the `send_*` constructor too.
//! [`ToolNotificationHandle`]. Keep the two surfaces in lockstep.
//!
//! The handle is built on `futures::channel::mpsc` so it is runtime-neutral:
//! the trait crate doesn't pin a particular async executor on its
@@ -24,10 +23,9 @@ use serde::{Deserialize, Serialize};
/// made once.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct BashNotificationBase {
/// Tool call id, used to correlate with the originating tool call.
/// Correlates with the originating tool call.
pub tool_call_id: String,
/// The command being executed.
pub command: String,
/// Captured output bytes. May be truncated; use `output_lossy` for a
@@ -40,7 +38,6 @@ pub struct BashNotificationBase {
/// Whether `output` was truncated to fit a size cap.
pub truncated: bool,
/// Working directory the command ran in.
pub cwd: PathBuf,
}
@@ -76,7 +73,6 @@ pub struct BashExecutionComplete {
}
impl BashExecutionComplete {
/// `true` when termination was triggered by a signal.
pub fn was_signaled(&self) -> bool {
self.signal.is_some()
}
@@ -89,10 +85,8 @@ pub struct BashExecutionTimeout {
#[serde(flatten)]
pub base: BashNotificationBase,
/// Wall time the command ran for before being killed.
pub elapsed: Duration,
/// Configured timeout that was exceeded.
pub timeout: Duration,
}
@@ -122,23 +116,19 @@ pub struct BashExecutionFailed {
pub tool_call_id: String,
pub command: String,
pub cwd: PathBuf,
/// Error message describing the spawn / IO failure.
pub error: String,
}
/// Emitted when a tool reads a file. Subscribers use this for state
/// snapshotting (rewind, audit) of accessed files.
/// Payload for a tool file-read event (rewind / audit subscribers).
///
/// **Reserved for a future `ToolNotification::FileRead` variant.** The
/// struct is kept in the public API so adapters can construct it ahead of
/// time, but it is not currently dispatched by any
/// [`ToolNotificationHandle`] helper. Adding the enum variant here is a
/// breaking change for exhaustive `match` consumers, so the variant is
/// deferred until a downstream crate has a real consumer wired up.
/// **Reserved for a future `ToolNotification::FileRead` variant.** Public
/// so adapters can construct it, but no [`ToolNotificationHandle`] helper
/// dispatches it yet. Introducing the enum variant is a breaking change for
/// exhaustive `match` consumers, so it waits until a downstream crate has a
/// real consumer wired up.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FileRead {
pub tool_call_id: String,
/// Absolute filesystem path of the file that was read.
pub absolute_path: PathBuf,
}
@@ -147,13 +137,11 @@ pub struct FileRead {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FileWritten {
pub tool_call_id: String,
/// Absolute filesystem path of the file that was written.
pub absolute_path: PathBuf,
/// Full file content after the write.
pub content: String,
/// Full file content before the write. `None` for a fresh file.
pub previous_content: Option<String>,
/// Whether the write created a new file.
pub is_new_file: bool,
}
@@ -173,7 +161,6 @@ pub struct PlanModeExited {
/// Plan content as captured at exit time. `None` when the plan file
/// did not exist or was empty.
pub plan_content: Option<String>,
/// Path the plan file lives at.
pub plan_file_path: String,
}
@@ -243,7 +230,6 @@ pub struct ScheduledTaskRemoved {
pub task_id: String,
}
/// Sent when a scheduled task is created.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScheduledTaskCreated {
pub task_id: String,
@@ -262,7 +248,6 @@ pub struct MonitorEvent {
pub description: String,
/// XML-wrapped event text, ready for conversation injection.
pub event_text: String,
/// Raw text without XML wrapping.
pub raw_text: String,
}
@@ -289,7 +274,6 @@ pub struct TaskSnapshot {
pub exit_code: Option<i32>,
pub signal: Option<String>,
pub completed: bool,
/// Distinguishes monitor tasks from regular bash tasks.
#[serde(default)]
pub kind: TaskKind,
}
@@ -305,18 +289,17 @@ impl TaskSnapshot {
}
}
/// Distinguishes background-task kinds.
/// Background-task kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TaskKind {
/// Regular bash command.
#[default]
Bash,
/// Monitor tool — streams stdout events with rate limiting.
Monitor,
}
/// A typed notification a tool emits during or after execution.
/// Typed notification a tool emits during or after execution.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ToolNotification {
@@ -382,7 +365,6 @@ pub struct ToolNotificationHandle {
}
impl ToolNotificationHandle {
/// Wrap a sender obtained elsewhere.
pub fn new(sender: mpsc::UnboundedSender<ToolNotification>) -> Self {
Self { sender }
}
@@ -408,122 +390,83 @@ impl ToolNotificationHandle {
Self { sender }
}
/// Send a fully-built notification. Errors are deliberately swallowed;
/// notifications are best-effort.
/// Best-effort send; errors (closed receiver) are swallowed.
pub fn send(&self, notification: ToolNotification) {
let _ = self.sender.unbounded_send(notification);
}
/// Send a [`ToolNotification::BashOutputChunk`]: an incremental
/// stdout/stderr chunk while a bash command is still running.
pub fn send_bash_output_chunk(&self, chunk: BashOutputChunk) {
self.send(ToolNotification::BashOutputChunk(chunk));
}
/// Send a [`ToolNotification::BashExecutionComplete`]: a bash command
/// exited (normally or via signal).
pub fn send_bash_complete(&self, complete: BashExecutionComplete) {
self.send(ToolNotification::BashExecutionComplete(complete));
}
/// Send a [`ToolNotification::BashExecutionTimeout`]: a bash command
/// exceeded its configured timeout and was killed.
pub fn send_bash_timeout(&self, timeout: BashExecutionTimeout) {
self.send(ToolNotification::BashExecutionTimeout(timeout));
}
/// Send a [`ToolNotification::BashExecutionBackgrounded`]: a
/// foreground bash command was moved to the background.
pub fn send_bash_backgrounded(&self, backgrounded: BashExecutionBackgrounded) {
self.send(ToolNotification::BashExecutionBackgrounded(backgrounded));
}
/// Send a [`ToolNotification::BashExecutionFailed`]: a bash command
/// could not be spawned.
pub fn send_bash_failed(&self, failed: BashExecutionFailed) {
self.send(ToolNotification::BashExecutionFailed(failed));
}
/// Send a [`ToolNotification::FileWritten`]: a tool wrote to a file
/// on disk.
pub fn send_file_written(&self, written: FileWritten) {
self.send(ToolNotification::FileWritten(written));
}
/// Send a [`ToolNotification::TaskCompleted`]: a background task
/// transitioned to a terminal state.
pub fn send_task_complete(&self, task_completed: TaskSnapshot) {
self.send(ToolNotification::TaskCompleted(task_completed));
}
/// Send a [`ToolNotification::PlanModeEntered`]: the agent
/// transitioned into plan mode.
pub fn send_plan_mode_entered(&self, entered: PlanModeEntered) {
self.send(ToolNotification::PlanModeEntered(entered));
}
/// Send a [`ToolNotification::PlanModeExited`]: the agent transitioned
/// out of plan mode and the captured plan is attached.
pub fn send_plan_mode_exited(&self, exited: PlanModeExited) {
self.send(ToolNotification::PlanModeExited(exited));
}
/// Send a [`ToolNotification::UserQuestionAsked`]: the agent issued a
/// structured question payload to the user.
pub fn send_user_question_asked(&self, asked: UserQuestionAsked) {
self.send(ToolNotification::UserQuestionAsked(asked));
}
/// Send a [`ToolNotification::LspServerStarting`]: an LSP server is
/// being spawned.
pub fn send_lsp_starting(&self, starting: LspServerStarting) {
self.send(ToolNotification::LspServerStarting(starting));
}
/// Send a [`ToolNotification::LspServerReady`]: an LSP server
/// finished its initialise handshake.
pub fn send_lsp_ready(&self, ready: LspServerReady) {
self.send(ToolNotification::LspServerReady(ready));
}
/// Send a [`ToolNotification::LspServerCrashed`]: an LSP server
/// process died unexpectedly.
pub fn send_lsp_crashed(&self, crashed: LspServerCrashed) {
self.send(ToolNotification::LspServerCrashed(crashed));
}
/// Send a [`ToolNotification::LspServerRetrying`]: an LSP server is
/// being restarted after a crash.
pub fn send_lsp_retrying(&self, retrying: LspServerRetrying) {
self.send(ToolNotification::LspServerRetrying(retrying));
}
/// Send a [`ToolNotification::LspServerFailed`]: an LSP server is
/// permanently dead (init failure or retry budget exhausted).
pub fn send_lsp_failed(&self, failed: LspServerFailed) {
self.send(ToolNotification::LspServerFailed(failed));
}
/// Send a [`ToolNotification::ScheduledTaskFired`]: a recurring or
/// one-shot scheduled task fired and its prompt should be executed.
pub fn send_scheduled_task_fired(&self, fired: ScheduledTaskFired) {
self.send(ToolNotification::ScheduledTaskFired(fired));
}
/// Send a [`ToolNotification::ScheduledTaskRemoved`]: a scheduled task
/// was deleted, expired, or a one-shot variant completed.
pub fn send_scheduled_task_removed(&self, removed: ScheduledTaskRemoved) {
self.send(ToolNotification::ScheduledTaskRemoved(removed));
}
/// Send a [`ToolNotification::ScheduledTaskCreated`]: a new scheduled
/// task was registered and should appear in subscriber views.
pub fn send_scheduled_task_created(&self, created: ScheduledTaskCreated) {
self.send(ToolNotification::ScheduledTaskCreated(created));
}
/// Send a [`ToolNotification::MonitorEvent`]: a streaming event from
/// a Monitor background process, ready for conversation injection.
pub fn send_monitor_event(&self, event: MonitorEvent) {
self.send(ToolNotification::MonitorEvent(event));
}
+7 -18
View File
@@ -112,7 +112,6 @@ impl<T: ToolOutput + Serialize + ?Sized> ToolOutput for Box<T> {
/// frontend.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ToolChatCompletionResponse {
/// The main completion payload.
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<ToolChatCompletion>,
/// Structured stream error (e.g. rate-limit, tool failure).
@@ -126,7 +125,6 @@ pub struct ToolChatCompletion {
/// Always `"assistant"`.
#[serde(default)]
pub sender: String,
/// Text body of the response.
#[serde(default)]
pub message: String,
/// Tag discriminator: `"final"`, `"raw_function_result"`,
@@ -139,7 +137,6 @@ pub struct ToolChatCompletion {
/// JSON-encoded card attachment (images, render cards, files).
#[serde(skip_serializing_if = "Option::is_none")]
pub card_attachment: Option<String>,
/// Code execution result.
#[serde(skip_serializing_if = "Option::is_none")]
pub code_execution_result: Option<ToolCodeExecutionResult>,
/// Catch-all for additional fields the tool wants to set. Merged
@@ -183,7 +180,7 @@ pub struct ToolStreamError {
/// | 4 | Object with mixed fields | block-shaped fields extracted, rest as JSON text |
/// | 5 | Anything else | `ContentBlock::Text` with the stringified value |
pub fn extract_content_blocks(value: &Value) -> Vec<ContentBlock> {
// 1. Value IS a single ContentBlock.
// 1. Value is a single ContentBlock.
if let Some(block) = try_parse_block(value) {
return vec![block];
}
@@ -261,9 +258,7 @@ pub fn extract_content_blocks(value: &Value) -> Vec<ContentBlock> {
vec![value_to_block(value)]
}
// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
/// The `ContentBlock` enum is `#[serde(tag = "type", rename_all =
/// "snake_case")]`, so a JSON object can only be a content block when
@@ -293,13 +288,11 @@ fn try_parse_block(value: &Value) -> Option<ContentBlock> {
serde_json::from_value::<ContentBlock>(value.clone()).ok()
}
/// Result of inspecting a single object field value.
enum FieldShape {
/// The field value IS a single `ContentBlock`.
/// Single `ContentBlock`.
Block(ContentBlock),
/// The field value is an array where *every* element is a `ContentBlock`.
/// Array where every element is a `ContentBlock`.
Blocks(Vec<ContentBlock>),
/// The field value does not look like block content.
Other,
}
@@ -309,7 +302,6 @@ enum FieldShape {
/// `ContentBlock`; mixed arrays go to `Other` so ambiguous data
/// (e.g. `"scores": [0.9, 0.8]`) is not silently dropped.
fn classify_field(value: &Value) -> FieldShape {
// Single block.
if let Some(block) = try_parse_block(value) {
return FieldShape::Block(block);
}
@@ -343,11 +335,8 @@ fn value_to_block(value: &Value) -> ContentBlock {
})
}
// ---------------------------------------------------------------------------
// Type-erased extractor (used by the toolbox registry)
// ---------------------------------------------------------------------------
// Type-erased extractor (toolbox registry)
/// Type-erased model output extractor.
pub type ModelOutputExtractor = Arc<dyn Fn(&Value) -> Option<Vec<ContentBlock>> + Send + Sync>;
/// Build a [`ModelOutputExtractor`] for a concrete output type.
@@ -374,7 +363,7 @@ mod tests {
use super::*;
use serde_json::json;
// ── ToolOutput with custom override ─────────────────────────────
// ToolOutput with custom override
#[derive(Serialize)]
struct FakeOutput {
@@ -423,7 +412,7 @@ mod tests {
assert_eq!(o.model_output().len(), 2);
}
// ── ToolOutput default → empty (runtime fills via extract) ──────
// ToolOutput default → empty (runtime fills via extract)
#[test]
fn default_model_output_returns_empty() {
@@ -464,7 +453,7 @@ mod tests {
);
}
// ── extract_content_blocks unit tests ──────────────────────────
// extract_content_blocks unit tests
// Strategy 1: single ContentBlock
#[test]
+14 -19
View File
@@ -1,8 +1,8 @@
//! Backend-agnostic tool search interface.
//!
//! `ToolSearchIndex` is a `Send + Sync` trait so concrete implementations
//! can live in different crates (BM25, OpenSearch, in-memory linear) and
//! be stored as `Arc<dyn ToolSearchIndex>` for shared access across tasks.
//! `ToolSearchIndex` is `Send + Sync` so concrete implementations can live
//! in different crates (BM25, OpenSearch, in-memory linear) and be stored
//! as `Arc<dyn ToolSearchIndex>` for shared access across tasks.
use std::sync::Arc;
@@ -13,26 +13,25 @@ pub struct ToolSearchResult {
pub tool_name: String,
/// Origin server name (e.g. `"linear"`).
pub server_name: String,
/// Tool description.
pub description: String,
/// Backend-defined relevance score; comparable within a single
/// snapshot but not across snapshots.
pub score: f32,
/// Parameter names from the tool's input schema, in declaration order.
pub parameters: Vec<String>,
/// Full JSON Schema for the tool's input. Included so callers can
/// construct dispatched tool calls without a separate schema fetch.
/// Full JSON Schema for the tool's input so callers can construct
/// dispatched tool calls without a separate schema fetch.
pub input_schema: serde_json::Value,
}
/// Snapshot of a search query — results plus index metadata captured from
/// the same point-in-time view.
/// Snapshot of a search query — results plus index metadata from the same
/// point-in-time view.
#[derive(Debug, Clone, PartialEq)]
pub struct SearchSnapshot {
pub results: Vec<ToolSearchResult>,
/// Number of indexed tools that did not appear in `results`.
pub total_hidden_tools: usize,
/// `true` when the index reflects all available tools. `false` while
/// `true` when the index reflects all available tools; `false` while
/// the index source is still warming up.
pub is_ready: bool,
}
@@ -44,13 +43,11 @@ pub struct ServerSummary {
pub name: String,
/// Optional short description of the server's surface area.
pub description: Option<String>,
/// Unqualified tool names, sorted alphabetically. Use
/// [`Self::tool_count`] for a count without indirection.
/// Unqualified tool names, sorted alphabetically.
pub tool_names: Vec<String>,
}
impl ServerSummary {
/// Number of tools the server exposes.
pub fn tool_count(&self) -> usize {
self.tool_names.len()
}
@@ -61,18 +58,16 @@ impl ServerSummary {
/// Implementations must be `Send + Sync` so they can be wrapped in
/// `Arc<dyn ToolSearchIndex>` and shared across concurrent tasks.
pub trait ToolSearchIndex: Send + Sync {
/// Run a query against a single consistent index snapshot. Returning
/// the metadata alongside the results lets the caller render an
/// accurate "N results out of M" line without a second call.
/// Query a single consistent index snapshot. Metadata rides with the
/// results so the caller can render "N of M" without a second call.
fn search_snapshot(&self, query: &str, limit: usize) -> SearchSnapshot;
/// Enumerate the unique servers in the index. Used to render the
/// system-reminder listing connected integrations.
/// Unique servers in the index (e.g. for a system-reminder listing
/// connected integrations).
fn list_server_summaries(&self) -> Vec<ServerSummary>;
}
/// Resource wrapper for storing a `ToolSearchIndex` behind an `Arc` in
/// shared resource maps.
/// `ToolSearchIndex` behind an `Arc` for shared resource maps.
#[derive(Clone)]
pub struct ToolIndex(pub Arc<dyn ToolSearchIndex>);
@@ -315,7 +315,7 @@ mod tests {
);
}
// ── Limit / latency invariants ──────────────────────────────────────────
// Limit / latency invariants
/// A backlog drains in exactly `ceil(new / cap)` calls — no extra round-trips.
#[test]
@@ -383,7 +383,8 @@ mod tests {
/// UTF-8 backoff loses at most 3 bytes, so frames stay within 3 of the cap.
#[test]
fn utf8_backoff_stays_within_three_bytes_of_cap() {
let cap = 7usize; // splits a 4-byte char -> backs off to 4 (cap - 3)
// Cap 7 splits a 4-byte char; backoff yields 4 (cap - 3).
let cap = 7usize;
let spec = spec_with(Some(cap as u32));
let data = "😀😀😀😀".as_bytes();
let total = data.len() as u64;
@@ -408,7 +409,8 @@ mod tests {
fn gap_with_cap_paces_surviving_tail_and_terminates() {
let spec = spec_with(Some(4));
let tail = b"abcdefgh";
let total = 1000u64; // only 8 of 1000 bytes survived in the tail
// Only 8 of 1000 bytes survive in the tail.
let total = 1000u64;
let mut last = 0;
let mut ticks = 0usize;
let mut emitted = 0usize;
+6 -16
View File
@@ -141,11 +141,9 @@ pub enum ToolProgress {
Text { text: String },
/// Rich content blocks.
Content { blocks: Vec<ContentBlock> },
/// Tool-defined progress payload. `subkind` is a stable snake-case
/// discriminator owned by the tool. The outer `"kind"` serde tag is
/// always `"custom"` for this variant; `subkind` is the producer's
/// own discriminator and lives one level deeper to avoid colliding
/// with the tag.
/// Tool-defined progress. Outer serde tag is always `"custom"`; `subkind`
/// is the producer's discriminator one level deeper (avoids colliding
/// with the tag).
Custom {
subkind: String,
payload: serde_json::Value,
@@ -244,7 +242,6 @@ where
/// least one block.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TypedToolOutput {
/// Identity of the tool that produced this output.
pub tool_id: ToolId,
/// Serialised JSON representation of the tool output.
pub value: Value,
@@ -301,21 +298,17 @@ impl ToolOutput for TypedToolOutput {
}
}
/// Type erased tool trait. Auto-generated for every typed Tool implementation.
/// Type-erased tool trait. Blanket-impl'd for every typed [`Tool`].
#[async_trait]
pub trait ToolDyn: Send + Sync {
/// Stable identity. Same value as [`Tool::id`].
fn id(&self) -> ToolId;
/// Model-facing description. Same value as [`Tool::description`].
fn description(&self, ctx: &ListToolsContext) -> ToolDescription;
/// Per-tool capability flags. Same value as [`Tool::capabilities`].
fn capabilities(&self) -> ToolCapabilities {
ToolCapabilities::default()
}
/// Same value as [`Tool::has_dynamic_description`].
fn has_dynamic_description(&self) -> bool {
false
}
@@ -401,23 +394,21 @@ impl<T: Tool> ToolDyn for T {
}
}
/// Convenience alias for the most common [`ToolDyn`] handle shape.
pub type ArcTool = Arc<dyn ToolDyn>;
/// Variant identifier for tools that ship multiple implementations under
/// one stable [`ToolId`].
#[derive(Debug, Clone, Eq, PartialEq, Hash)]
pub enum ToolVariant {
/// The implicit fallback variant.
/// Implicit fallback variant.
Default,
/// A named variant. The string is treated opaquely by the registry.
/// Named variant; treated opaquely by the registry.
Variant(String),
}
/// Group of related tools that share one [`ToolId`] but route to different
/// implementations chosen by a [`ToolVariant`].
pub trait ToolFamily: Send + Sync {
/// Identity shared by every variant in this family.
fn id(&self) -> ToolId;
/// Resolve a `variant` to its concrete tool. Returns `None` when the
@@ -436,5 +427,4 @@ pub trait ToolFamily: Send + Sync {
}
}
/// Convenience alias for the most common [`ToolFamily`] handle shape.
pub type ArcToolFamily = Arc<dyn ToolFamily>;
@@ -88,8 +88,7 @@ fn remove_returns_value_then_none() {
#[test]
fn insert_arc_shares_allocation() {
// Inserting an existing Arc means the stored value and the original
// share strong-count.
// insert_arc shares the Arc (strong-count rises).
let arc = Arc::new(Config {
base_url: "shared".into(),
timeout_ms: 9,
@@ -97,10 +96,7 @@ fn insert_arc_shares_allocation() {
let mut ctx = ToolCallContext::default();
ctx.extensions.insert_arc(arc.clone());
let from_ctx = ctx.extensions.get::<Config>().unwrap();
// Strong-count on the original Arc should reflect at least:
// - the original `arc` binding
// - the value stored in the extension map
// - the clone returned from `get`
// Strong-count: original binding + map entry + get() clone.
assert!(Arc::strong_count(&arc) >= 3);
assert_eq!(*from_ctx, *arc);
}
@@ -152,13 +148,10 @@ fn clone_preserves_call_id_and_extensions() {
assert_eq!(copy.call_id, ctx.call_id);
assert_eq!(copy.extensions.len(), 1);
// Both clones see the same Arc-backed extension value.
let from_orig = ctx.extensions.get::<AuthToken>().unwrap();
let from_copy = copy.extensions.get::<AuthToken>().unwrap();
assert_eq!(from_orig.0, from_copy.0);
// The Arc allocation is shared; mutating via one path is impossible
// (extensions are immutable through `get`), but strong-count rises
// because of the clone.
// Arc is shared; get() only clones the handle (immutable).
assert!(Arc::strong_count(&from_orig) >= 3);
}
@@ -176,21 +169,12 @@ fn clone_extension_map_is_independent_after_remove() {
);
}
// ---------------------------------------------------------------------------
// Per-concept client/SDK-side extensions.
//
// These exist as separate extensions (one per concept) rather than a
// single bundle. The tests below pin three contracts:
//
// 1. Each extension round-trips through the typed-extension store
// independently of the others.
// 2. A dispatcher with only some of the concepts can install them
// individually — installing `Cwd` MUST NOT make `BehaviorVersion`
// look "present" with a default value, and vice versa.
// 3. Absence of every well-known extension is the legitimate "backend
// dispatcher" shape; tools that require one MUST treat absence as
// a hard error rather than fall back to a process-wide default.
// ---------------------------------------------------------------------------
// Per-concept client/SDK-side extensions (one type per concept, not a bundle).
// Pins three contracts:
// 1. Each extension round-trips independently.
// 2. Installing one MUST NOT make another look "present" with a default.
// 3. Absence is the legitimate backend-dispatcher shape; tools that need
// an extension MUST treat absence as a hard error.
#[test]
fn each_well_known_extension_round_trips_independently() {
@@ -217,9 +201,7 @@ fn each_well_known_extension_round_trips_independently() {
#[test]
fn dispatcher_can_install_only_what_it_has() {
// A dispatcher that knows the cwd but not the trace context installs
// only `Cwd`. The other extensions stay absent (not "default"),
// which is the discriminator a tool can rely on.
// Only Cwd installed — other extensions stay absent (not defaulted).
let mut ctx = ToolCallContext::default();
ctx.extensions
.insert(Cwd(std::path::PathBuf::from("/work")));
@@ -229,8 +211,7 @@ fn dispatcher_can_install_only_what_it_has() {
assert!(!ctx.extensions.contains::<TraceContext>());
assert_eq!(ctx.extensions.len(), 1);
// Adding `TraceContext` later does not implicitly conjure a
// `BehaviorVersion` — extensions are independent.
// Installing TraceContext does not conjure BehaviorVersion.
ctx.extensions.insert(TraceContext("tp".into()));
assert!(ctx.extensions.contains::<TraceContext>());
assert!(!ctx.extensions.contains::<BehaviorVersion>());
@@ -239,9 +220,7 @@ fn dispatcher_can_install_only_what_it_has() {
#[test]
fn absence_signals_backend_or_other_mode() {
// A backend dispatcher installs none of the client-side extensions.
// Tools that require any of them must treat absence as a hard error
// — this test pins the contract.
// Backend dispatcher: no client-side extensions present.
let ctx = ToolCallContext::default();
assert!(ctx.extensions.get::<Cwd>().is_none());
assert!(ctx.extensions.get::<BehaviorVersion>().is_none());
@@ -1,4 +1,4 @@
//! `From<ToolError> for ToolErrorWire` coverage for the struct-based ToolError.
//! `From<ToolError> for ToolErrorWire` coverage.
use serde_json::json;
@@ -390,7 +390,6 @@ fn noop_handle_does_not_panic_or_record() {
handle.send_lsp_ready(LspServerReady {
server_name: "x".into(),
});
// No assertion needed — the handle drops sends silently.
}
#[test]
@@ -93,10 +93,8 @@ fn tool_index_wrapper_clones_arc() {
});
let wrapped = ToolIndex(inner.clone());
let copy = wrapped.clone();
// Both wrappers hold the same Arc — strong-count includes both
// wrappers and the original `inner` binding.
// Both wrappers share the Arc with `inner`.
assert!(Arc::strong_count(&inner) >= 3);
// Debug impl renders without leaking the inner type.
let debug = format!("{wrapped:?}");
assert_eq!(debug, "ToolIndex");
drop(copy);
@@ -77,8 +77,6 @@ impl Tool for NeedsAttachmentTool {
}
}
// Tool::should_list (typed)
#[test]
fn default_returns_true() {
assert!(Tool::should_list(&AlwaysTool, &ListToolsContext::default()));
@@ -109,8 +107,6 @@ fn reads_custom_extension() {
assert!(Tool::should_list(&tool, &some));
}
// ToolDyn blanket forwarding
#[test]
fn dyn_forwards_default() {
let tool: ArcTool = Arc::new(AlwaysTool);
@@ -136,8 +132,6 @@ fn arc_dyn_callable() {
assert!(tool.should_list(&ctx));
}
// ListToolsContext
#[test]
fn list_ctx_default_is_empty() {
let ctx = ListToolsContext::default();
@@ -164,8 +158,6 @@ fn list_ctx_clone_is_independent() {
assert!(!copy.extensions.contains::<AttachmentCount>());
}
// TypedExtensions standalone
#[test]
fn typed_extensions_insert_get_remove() {
let mut ext = kigi_tool_runtime::TypedExtensions::new();
@@ -151,8 +151,7 @@ async fn unimplemented_tool_returns_not_implemented_terminal() {
#[tokio::test]
async fn run_takes_args_by_value() {
// The trait `run` consumes args; this would not compile if the
// signature accidentally borrowed.
// run consumes args (would not compile if the signature borrowed).
let tool = BlockingOk;
let args = EchoArgs {
text: "consumed".into(),
@@ -163,7 +162,6 @@ async fn run_takes_args_by_value() {
#[tokio::test]
async fn execute_default_drains_in_one_pass() {
// A stream from the default impl should always have exactly one item.
let tool = BlockingOk;
let count = tool
.execute(ToolCallContext::default(), EchoArgs { text: "n".into() })
@@ -130,7 +130,7 @@ impl Tool for UnencodableTool {
}
}
// ── Tool with custom ToolOutput (non-empty) ──────────────────
// Tool with custom ToolOutput (non-empty)
/// Output that provides its own model-facing content blocks. The blanket
/// impl must forward these as-is rather than filling in the JSON fallback.
@@ -201,7 +201,6 @@ async fn tool_dyn_preserves_custom_model_output() {
{"type": "image", "mime_type": "image/png", "data": "base64data"},
]})
);
// Custom model output preserved verbatim — no JSON fallback.
assert_eq!(typed.model_output.len(), 2);
assert_eq!(
typed.model_output[0],
@@ -238,8 +237,7 @@ async fn tool_dyn_blanket_encodes_terminal_output() {
ToolStreamItem::Terminal(Ok(typed)) => {
assert_eq!(typed.tool_id, tid("blocking_echo"));
assert_eq!(typed.value, json!({"text": "hi"}));
// EchoOutput uses the default ToolOutput which
// serialises self to a JSON text block (MCP-compliant).
// Default ToolOutput serialises self to a JSON text block (MCP).
assert_eq!(typed.model_output.len(), 1);
assert_eq!(
typed.model_output[0],
@@ -283,7 +281,7 @@ async fn tool_dyn_blanket_passes_progress_through() {
#[tokio::test]
async fn tool_dyn_invalid_args_become_invalid_arguments_terminal() {
let tool: ArcTool = Arc::new(BlockingEcho);
// `text` is required and must be a string — `null` fails serde.
// `text` is required and must be a string.
let mut stream = tool
.execute(ToolCallContext::default(), json!({"text": null}))
.await;
@@ -314,9 +312,7 @@ async fn tool_dyn_unencodable_output_becomes_execution_terminal() {
}
}
// ---------------------------------------------------------------------------
// ToolFamily
// ---------------------------------------------------------------------------
/// Backend-flavoured echo. Two variants share the `echo` tool id and only
/// differ in the prefix attached to the output text — enough to assert
@@ -437,9 +433,7 @@ async fn tool_family_default_variant_name_defaults_to_none() {
assert!(family.default_variant_name().is_none());
}
// ---------------------------------------------------------------------------
// Object safety / ergonomic checks
// ---------------------------------------------------------------------------
#[test]
fn tool_dyn_is_object_safe_in_arc_and_box() {
@@ -455,8 +449,7 @@ fn tool_family_is_object_safe_in_arc_and_box() {
#[test]
fn arc_tool_alias_holds_heterogeneous_tools() {
// The whole point of `ArcTool` — many typed `Tool` impls collapse
// into one container shape via the blanket impl.
// ArcTool: many typed Tool impls collapse into one container via the blanket.
let tools: Vec<ArcTool> = vec![Arc::new(BlockingEcho), Arc::new(StreamingEcho)];
assert_eq!(tools.len(), 2);
let ids: Vec<_> = tools.iter().map(|t| t.id()).collect();
@@ -484,8 +477,7 @@ fn _compile_time_blanket_check() {
let tool = StreamingEcho;
_accepts_dyn(&tool);
// The trait objects themselves must be `Send + Sync` so they can be
// shared across tasks without further bounds at the call site.
// Trait objects are Send + Sync for sharing across tasks.
fn _is_send_sync<T: Send + Sync + ?Sized>() {}
_is_send_sync::<dyn ToolDyn>();
_is_send_sync::<dyn ToolFamily>();
@@ -150,7 +150,6 @@ async fn streaming_err_propagates_through_terminal() {
#[tokio::test]
async fn streaming_progress_count_is_independent_of_args() {
// Distinct invocations on the same tool produce the same shape.
let tool = StreamingOk;
for _ in 0..3 {
let count = tool
@@ -164,8 +163,7 @@ async fn streaming_progress_count_is_independent_of_args() {
#[tokio::test]
async fn empty_progress_still_yields_terminal() {
// Building `with_progress` on an empty stream still produces exactly
// one terminal item — the same shape `terminal_only` produces.
// Empty progress stream still yields exactly one terminal item.
let progress = stream::iter(Vec::<ToolProgress>::new());
let mut stream = with_progress(progress, async move { Ok::<u32, ToolError>(99) });
let item = stream.next().await.unwrap();
-10
View File
@@ -39,14 +39,12 @@ impl Extensions {
Self::default()
}
/// Retrieve a reference to a stored value by type.
pub fn get<T: Any + Send + Sync + 'static>(&self) -> Option<&T> {
self.map
.get(&TypeId::of::<T>())
.and_then(|e| e.data.downcast_ref())
}
/// Retrieve a mutable reference to a stored value by type.
pub fn get_mut<T: Any + Send + Sync + 'static>(&mut self) -> Option<&mut T> {
self.map
.get_mut(&TypeId::of::<T>())
@@ -64,7 +62,6 @@ impl Extensions {
);
}
/// Remove and return a value by type.
pub fn remove<T: Any + Send + Sync + 'static>(&mut self) -> Option<T> {
self.map
.remove(&TypeId::of::<T>())
@@ -72,17 +69,14 @@ impl Extensions {
.map(|b| *b)
}
/// Check if a value of the given type is stored.
pub fn contains<T: Any + Send + Sync + 'static>(&self) -> bool {
self.map.contains_key(&TypeId::of::<T>())
}
/// Number of stored entries.
pub fn len(&self) -> usize {
self.map.len()
}
/// Returns true if no entries are stored.
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
@@ -105,10 +99,6 @@ impl fmt::Debug for Extensions {
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
@@ -143,9 +143,7 @@ pub fn parse_arguments_from_schema_lossy(schema: &serde_json::Value) -> Vec<Tool
.collect()
}
// ---------------------------------------------------------------------------
// $ref / $defs / anyOf / oneOf resolution
// ---------------------------------------------------------------------------
/// Resolve type info from a property, following `$ref` → `$defs` and
/// `anyOf` patterns that schemars generates for Rust enums and
@@ -263,7 +261,6 @@ fn extract_enum_from_def(
(Some(arg_type), Some(values), first_value)
}
/// Infer the [`ArgumentType`] from a sample enum value.
fn infer_arg_type(sample: &Option<Value>) -> ArgumentType {
match sample {
Some(Value::String(_)) => ArgumentType::String,
@@ -274,10 +271,6 @@ fn infer_arg_type(sample: &Option<Value>) -> ArgumentType {
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
@@ -464,8 +457,6 @@ mod tests {
assert!(args[0].arg_type.contains(ArgumentType::Integer));
}
// -- $ref / $defs / anyOf resolution ----------------------------------------
#[test]
fn parse_schema_any_of_ref_resolves_enum() {
// schemars pattern for `Option<MyEnum>` with oneOf-style defs.
@@ -593,7 +584,7 @@ mod tests {
#[test]
fn parse_schema_no_defs_still_works() {
// Properties with no $ref/$defs should work exactly as before.
// Properties with no $ref/$defs still parse normally.
let schema = serde_json::json!({
"type": "object",
"properties": {
@@ -611,8 +602,6 @@ mod tests {
assert_eq!(args[0].default, Some(serde_json::json!("x")));
}
// -- numeric constraints --------------------------------------------------
#[test]
fn parse_schema_numeric_constraints() {
let schema = serde_json::json!({
@@ -1,18 +1,14 @@
//! Lenient deserializers for tool-argument booleans: a boolean may arrive as a
//! JSON string (`"true"`) or number (`1`) when a client doesn't coerce args
//! against the tool schema. Accepted forms (strings case-insensitive, trimmed;
//! `null` is `false`):
//!
//! | Truthy | Falsy |
//! |---------------------------------------|------------------------------------------------|
//! | `true`, `"true"`, `"yes"`, `"1"`, `1` | `false`, `"false"`, `"no"`, `"0"`, `0`, `null` |
//! against the tool schema. Strings are trimmed and matched case-insensitively,
//! and `null` reads as `false`.
use serde::Deserialize;
const TRUE_LITERALS: [&str; 3] = ["true", "yes", "1"];
const FALSE_LITERALS: [&str; 3] = ["false", "no", "0"];
/// Parse a JSON value into a `bool` per the accepted forms; `None` otherwise.
/// `None` when the value matches none of the accepted forms.
pub fn lenient_bool_from_json(value: &serde_json::Value) -> Option<bool> {
match value {
serde_json::Value::Bool(b) => Some(*b),
@@ -48,8 +44,8 @@ fn invalid_bool_message(value: &serde_json::Value) -> String {
)
}
/// Deserialize a required `bool`; pair with `#[serde(default)]` so an absent key
/// uses the field default.
/// Pair with `#[serde(default)]` so an absent key falls back to the field
/// default instead of failing.
pub fn deserialize_lenient_bool<'de, D>(deserializer: D) -> Result<bool, D::Error>
where
D: serde::Deserializer<'de>,
@@ -59,8 +55,8 @@ where
.ok_or_else(|| serde::de::Error::custom(invalid_bool_message(&value)))
}
/// Deserialize `Option<bool>`: absent key `None` (via `#[serde(default)]`),
/// explicit `null` `Some(false)`.
/// With `#[serde(default)]` an absent key yields `None`, while an explicit
/// `null` yields `Some(false)`.
pub fn deserialize_lenient_option_bool<'de, D>(deserializer: D) -> Result<Option<bool>, D::Error>
where
D: serde::Deserializer<'de>,
+6 -67
View File
@@ -4,15 +4,12 @@
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
// ───────────────────────────────────────────────────────────────────────────
// `task` (spawn) tool — Input
// ───────────────────────────────────────────────────────────────────────────
/// Input for the `task` tool — launches a subagent to handle a task
/// autonomously.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct TaskToolInput {
/// The full task prompt for the subagent to execute.
#[schemars(description = "The full task prompt for the subagent to execute.")]
pub prompt: String,
@@ -28,10 +25,8 @@ pub struct TaskToolInput {
#[serde(default = "default_subagent_type")]
pub subagent_type: String,
/// Whether to run the subagent in the background.
///
/// Returns immediately with a subagent_id. Use the task output tool to
/// retrieve results. This is set to true by default.
/// retrieve results. Defaults to true.
#[schemars(
description = "Returns immediately with a subagent_id. Use the task output tool to \
retrieve results. This is set to true by default."
@@ -42,7 +37,6 @@ pub struct TaskToolInput {
)]
pub run_in_background: bool,
/// Capability mode controlling the child's tool access.
#[schemars(
description = "Capability mode: \"read-only\", \"read-write\", \"execute\", or \"all\". \
Controls which tool classes the child can use. Default is determined by the role."
@@ -50,7 +44,6 @@ pub struct TaskToolInput {
#[serde(default)]
pub capability_mode: Option<SubagentCapabilityMode>,
/// Isolation mode for the child's execution environment.
#[schemars(
description = "Isolation mode: \"none\" (default, shared workspace) or \"worktree\" \
(isolated git worktree). Worktree mode prevents the child's edits from \
@@ -93,7 +86,6 @@ pub struct TaskToolInput {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cwd: Option<String>,
/// Optional model slug for this subagent.
#[schemars(
description = "Optional model slug for this agent. If provided, it must resolve to one \
of the available model slugs. If omitted, the subagent uses the same model as the \
@@ -109,7 +101,6 @@ pub struct TaskToolInput {
pub task_id: Option<String>,
}
/// Default `subagent_type` for [`TaskToolInput`] when the caller omits it.
pub fn default_subagent_type() -> String {
"general-purpose".to_string()
}
@@ -180,7 +171,6 @@ impl SubagentCapabilityMode {
}
}
/// Isolation mode for subagent execution.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "kebab-case")]
pub enum SubagentIsolationMode {
@@ -201,9 +191,7 @@ impl SubagentIsolationMode {
}
}
// ───────────────────────────────────────────────────────────────────────────
// `task` (spawn) tool — Output
// ───────────────────────────────────────────────────────────────────────────
/// Structured completion output from a subagent (`task` tool).
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -215,7 +203,6 @@ pub struct SubagentCompletedOutput {
pub turns: u32,
pub duration_ms: u64,
pub worktree_path: Option<String>,
/// Persona used by this subagent, if any.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub persona: Option<String>,
/// The `subagent_id` to pass as `resume_from` to continue this subagent.
@@ -228,7 +215,6 @@ pub struct SubagentCompletedOutput {
}
impl SubagentCompletedOutput {
/// Render the resume footer showing the subagent ID and resume hint.
pub fn resume_footer(&self) -> String {
format_resume_footer(
&self.subagent_id,
@@ -317,10 +303,8 @@ pub fn format_resume_footer(
/// fan-out, and the toolbox wait path so the cap cannot drift.
pub const MAX_MULTI_WAIT_IDS: usize = 20;
/// Input for the `get_task_output` tool.
#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema)]
pub struct TaskOutputToolInput {
/// Task IDs to query. Pass one or more; a single task is a one-element list.
#[schemars(
description = "Task IDs to get output from. Pass one or more; for a single task use a one-element array. With a positive timeout_ms, multiple ids wait until all complete. Omit timeout_ms or pass 0 for a non-blocking snapshot."
)]
@@ -351,12 +335,10 @@ pub fn resolve_task_ids(ids: &[String]) -> Vec<String> {
}
impl TaskOutputToolInput {
/// Resolved, de-duplicated task IDs preserving first-seen order.
pub fn resolved_task_ids(&self) -> Vec<String> {
resolve_task_ids(&self.task_ids)
}
/// True only when `timeout_ms` is set and greater than zero.
pub fn waits(&self) -> bool {
task_output_waits(self.timeout_ms)
}
@@ -378,7 +360,6 @@ pub fn task_output_waits_from_json(args: &serde_json::Value) -> bool {
task_output_waits(timeout_ms)
}
/// Output from the `get_task_output` tool.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub enum TaskOutputOutput {
Result(TaskOutputResult),
@@ -386,18 +367,16 @@ pub enum TaskOutputOutput {
MultiResult(MultiTaskOutputResult),
}
/// Successful result from the `get_task_output` tool.
#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
pub struct TaskOutputResult {
pub task_id: String,
pub command: String,
pub status: String,
pub exit_code: Option<i32>,
/// Wall-clock start time (ISO 8601 format)
/// ISO 8601 wall-clock start time.
pub started: String,
/// Wall-clock end time if completed (ISO 8601 format)
/// ISO 8601 wall-clock end time when completed.
pub ended: Option<String>,
/// Duration in seconds
pub duration_secs: f64,
pub output: String,
pub output_file: String,
@@ -428,7 +407,6 @@ impl TaskOutputOutput {
}
}
/// Result from a multi-wait `get_task_output` / `wait_tasks` call.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct MultiTaskOutputResult {
pub mode: String,
@@ -468,11 +446,8 @@ impl TaskOutputResult {
}
}
// ───────────────────────────────────────────────────────────────────────────
// `wait_tasks` tool — Input
// ───────────────────────────────────────────────────────────────────────────
/// How a multi-wait (`wait_tasks`) request should resolve.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum WaitMode {
@@ -497,9 +472,7 @@ pub struct WaitTasksToolInput {
pub timeout_ms: Option<u64>,
}
// ───────────────────────────────────────────────────────────────────────────
// `kill_task` (cancel) tool — Input / Output
// ───────────────────────────────────────────────────────────────────────────
// `kill_task` tool — Input / Output
/// Input for the `kill_task` tool — terminates a running background task,
/// monitor, or subagent by id.
@@ -509,14 +482,12 @@ pub struct KillTaskToolInput {
pub task_id: String,
}
/// Output from the `kill_task` tool.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub enum KillTaskOutput {
Result(KillTaskResult),
TaskNotFound(String),
}
/// Successful result from the `kill_task` tool.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct KillTaskResult {
pub task_id: String,
@@ -538,7 +509,6 @@ impl KillTaskOutput {
pub struct SubagentDescriptor {
/// `subagent_type` value the model passes to the `task` tool.
pub name: String,
/// One-line summary of what this subagent does.
pub description: String,
/// Optional fragment summarizing the tools the subagent can use. Appended
/// verbatim after the description; may itself contain product-specific
@@ -572,8 +542,6 @@ impl BuiltinSubagent {
})
}
/// Build a [`SubagentDescriptor`], rendering the tool-access fragment via
/// [`Self::render_tools`] with the supplied `naming`.
pub fn to_descriptor(&self, naming: &SubagentToolNaming) -> SubagentDescriptor {
SubagentDescriptor {
name: self.name.to_owned(),
@@ -679,9 +647,6 @@ fn substitute_tool_placeholders(
}
/// Prompt body for the **general-purpose** subagent.
///
/// This agent has access to all tools and is used for complex search,
/// code exploration, and multi-step research tasks.
pub const GENERAL_PURPOSE_PROMPT: &str = "\
Complete the assigned task directly. Do what was asked; nothing more, nothing less. \
Respond with a detailed writeup when done.
@@ -708,8 +673,6 @@ Workspace boundary:
- Do not run whole-filesystem searches unless the user clearly requires it.";
/// Prompt body for the **explore** subagent.
///
/// A fast, read-only agent specialized for codebase exploration.
pub const EXPLORE_PROMPT: &str = "\
You are a fast, read-only codebase exploration agent.
@@ -737,9 +700,6 @@ Workspace boundary:
- If not found in the workspace, report that rather than broadening scope.";
/// Prompt body for the **plan** subagent.
///
/// A read-only architect agent that explores the codebase and produces
/// implementation plans.
pub const PLAN_PROMPT: &str = "\
You are a read-only software architect. Explore the codebase and design implementation plans.
@@ -771,7 +731,6 @@ Workspace boundary:
- Your default analysis scope is the workspace in <user_info>. Stay within it unless asked otherwise.
- Note explicitly if the design requires understanding external dependencies.";
/// The **general-purpose** built-in subagent.
pub const GENERAL_PURPOSE_SUBAGENT: BuiltinSubagent = BuiltinSubagent {
name: "general-purpose",
description: "General purpose agent for multi-step tasks.",
@@ -782,7 +741,6 @@ pub const GENERAL_PURPOSE_SUBAGENT: BuiltinSubagent = BuiltinSubagent {
prompt_template: GENERAL_PURPOSE_PROMPT,
};
/// The **explore** built-in subagent.
pub const EXPLORE_SUBAGENT: BuiltinSubagent = BuiltinSubagent {
name: "explore",
description: "Fast, read-only agent specialized for codebase exploration.",
@@ -792,7 +750,6 @@ pub const EXPLORE_SUBAGENT: BuiltinSubagent = BuiltinSubagent {
prompt_template: EXPLORE_PROMPT,
};
/// The **plan** built-in subagent.
pub const PLAN_SUBAGENT: BuiltinSubagent = BuiltinSubagent {
name: "plan",
description: "Software architect for planning implementation strategies.",
@@ -817,18 +774,11 @@ pub fn builtin_subagent_by_name(name: &str) -> Option<&'static BuiltinSubagent>
/// rendering the shared `task` tool description.
#[derive(Clone, Copy, Debug)]
pub struct TaskToolNaming<'a> {
/// Name of the spawn tool (canonical: `task`).
pub task_tool: &'a str,
/// Name of the `subagent_type` parameter.
pub subagent_type_param: &'a str,
/// Name of the `run_in_background` parameter.
pub run_in_background_param: &'a str,
/// Name of the `resume_from` parameter.
pub resume_from_param: &'a str,
/// Name of the task result retrieval tool.
pub background_retrieval_tool: &'a str,
/// Name of the `isolation` parameter, used in the isolation/worktree
/// paragraph.
pub isolation_param: &'a str,
}
@@ -888,7 +838,6 @@ fn lifecycle_target_suffix(monitor_present: bool, subagent_present: bool) -> &'s
}
}
/// Optional "(a monitor's task_id is returned by {monitor})" clause.
fn monitor_task_id_note(monitor_tool: Option<&str>) -> String {
match monitor_tool {
Some(m) => format!(" (a monitor's task_id is returned by {m})"),
@@ -909,7 +858,6 @@ pub struct KillTaskToolNaming<'a> {
pub is_windows: bool,
}
/// Build the shared `kill_task` tool description.
pub fn build_kill_task_description(naming: &KillTaskToolNaming) -> String {
let KillTaskToolNaming {
monitor_tool,
@@ -966,7 +914,6 @@ pub struct TaskOutputToolNaming<'a> {
pub subagent_background_param: Option<&'a str>,
}
/// Build the shared `get_task_output` tool description.
pub fn build_task_output_description(naming: &TaskOutputToolNaming) -> String {
let TaskOutputToolNaming {
monitor_tool,
@@ -1014,7 +961,6 @@ pub struct WaitTasksToolNaming<'a> {
pub subagent_background_param: Option<&'a str>,
}
/// Build the shared `wait_tasks` tool description.
pub fn build_wait_tasks_description(naming: &WaitTasksToolNaming) -> String {
let WaitTasksToolNaming {
background_retrieval_tool,
@@ -1174,7 +1120,6 @@ mod tests {
assert!(
desc.contains("- **general-purpose**: General-purpose agent. Has access to all tools.")
);
// User-defined entries (tools = None) get no trailing fragment.
assert!(desc.contains("- **code-reviewer**: Reviews code."));
assert!(desc.contains("## Usage notes"));
assert!(desc.contains(
@@ -1317,7 +1262,6 @@ mod tests {
#[test]
fn render_tools_substitutes_naming_with_bare_kind_fallback() {
// Bare-kind naming reproduces the placeholder kinds verbatim.
assert_eq!(
GENERAL_PURPOSE_SUBAGENT.render_tools(&plain_tool_naming()),
"Has access to all tools: execute, read, edit, list, search, web_search, and plan."
@@ -1328,7 +1272,6 @@ mod tests {
read, list, search, web_search, and plan."
);
// Real tool names are substituted per kind.
let naming = SubagentToolNaming {
execute: "run_terminal_cmd",
read: "read_file",
@@ -1372,12 +1315,8 @@ mod tests {
assert!(desc.contains("Use ${{ params.task.isolation }} to control"));
}
// ── Lifecycle tool descriptions ──────────────────────────────────────
//
// These lock the exact model-facing text. The "cli_default" cases must
// match what the kigi-shell MiniJinja templates render for the default
// kigi toolset (monitor + task + bash + read present, POSIX). The
// "toolbox" cases lock the subagent-only rendering used by the backend toolbox.
// Lifecycle tool descriptions — lock model-facing text against kigi-shell
// MiniJinja defaults (cli_default) and backend toolbox (subagent-only).
#[test]
fn kill_task_matches_cli_default_posix() {
+6 -50
View File
@@ -8,39 +8,28 @@ use crate::ext::Extensions;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolDescription {
/// Tool name (e.g. "web_search", "read_file") that is called by
/// the model.
pub name: String,
/// Optional namespace grouping (e.g. "github", "slack").
/// None for xAI native tools.
/// Namespace grouping (e.g. "github", "slack"). `None` for native tools.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub namespace: Option<String>,
/// Display name (e.g. "Web Search") can be shown to the model.
/// If absent, derive the title from 'name'.
/// Human display title; when absent, derive from `name`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
/// Description of the tool.
pub description: String,
/// Raw JSON Schema describing the tool's arguments.
/// JSON Schema for tool arguments.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub arguments_schema: Option<Value>,
/// High-level tool kind (stable snake_case, e.g. "read"), set by the tool
/// server so consumers can group tools by kind. `None` if undeclared.
/// Stable snake_case kind (e.g. "read") for grouping; `None` if undeclared.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub kind: Option<String>,
/// Metadata attached by downstream libraries to support
/// custom tool behavior. NOT serialized and NOT sent over
/// the wire.
///
/// Note: 'Extensions' always compares as equal (it carries opaque
/// runtime data), so 'ToolDescription's derived 'PartialEq' ignores
/// this field. See 'Extensions' for details.
/// Opaque runtime metadata — not serialized, not on the wire.
/// `Extensions` always compares equal, so derived `PartialEq` ignores this.
#[serde(skip)]
pub extra: Extensions,
}
@@ -63,7 +52,6 @@ impl ToolDescription {
self
}
/// Set the high-level tool kind (snake_case string, e.g. "read").
pub fn with_kind(mut self, kind: impl Into<String>) -> Self {
self.kind = Some(kind.into());
self
@@ -74,7 +62,6 @@ impl ToolDescription {
self
}
/// Attach the raw JSON Schema for this tool's arguments.
pub fn with_arguments_schema(mut self, schema: impl Into<Value>) -> Self {
self.arguments_schema = Some(schema.into());
self
@@ -94,8 +81,6 @@ impl ToolDescription {
.unwrap_or_default()
}
/// Returns the raw JSON Schema for the tool's arguments if one was
/// attached via [`Self::with_arguments_schema`].
pub fn arguments_schema(&self) -> Option<&Value> {
self.arguments_schema.as_ref()
}
@@ -160,14 +145,11 @@ impl fmt::Display for ToolDescription {
}
}
/// A single argument for a tool.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct ToolArgument {
/// Argument name (e.g. "file_path").
pub name: String,
/// Human-readable description of the argument.
pub description: String,
/// Type of the argument. Accepts both a single JSON Schema type
@@ -179,12 +161,10 @@ pub struct ToolArgument {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub schema: Option<serde_json::Value>,
/// Whether the argument is required.
/// Defaults to true. Omitted from JSON when true.
#[serde(default = "default_true", skip_serializing_if = "is_true")]
pub required: bool,
/// Default value for the argument.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub default: Option<Value>,
@@ -241,7 +221,6 @@ impl ToolArgument {
self
}
/// Set a default value for this argument.
pub fn with_default(mut self, default: impl Into<Value>) -> Self {
self.default = Some(default.into());
self
@@ -276,7 +255,6 @@ impl ToolArgument {
}
}
/// Type of a tool argument.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum ArgumentType {
@@ -347,15 +325,11 @@ impl fmt::Display for ArgumentType {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum SchemaType {
/// A single type, e.g. "string".
Single(ArgumentType),
/// Multiple types, e.g. ["string", "null"].
Multiple(Vec<ArgumentType>),
}
impl SchemaType {
/// Parse a JSON Schema "type" value (string or array) into a
/// `SchemaType`.
pub fn from_value(v: &serde_json::Value) -> Self {
if let Some(s) = v.as_str() {
return ArgumentType::from_schema_type(s)
@@ -404,7 +378,6 @@ impl SchemaType {
}
}
/// Whether the type list contains a specific `ArgumentType`.
pub fn contains(&self, ty: ArgumentType) -> bool {
match self {
Self::Single(t) => *t == ty,
@@ -442,7 +415,6 @@ impl SchemaType {
}
}
/// Return the JSON Schema `"type"` representation.
pub fn to_schema_value(&self) -> serde_json::Value {
match self {
Self::Single(t) => serde_json::Value::String(t.as_str().to_owned()),
@@ -500,7 +472,6 @@ fn is_true(v: &bool) -> bool {
*v
}
// -- Helpers
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValidationError {
pub field: String,
@@ -515,22 +486,18 @@ impl fmt::Display for ValidationError {
impl std::error::Error for ValidationError {}
/// Wrapper around multiple ValidationError.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValidationErrors(pub Vec<ValidationError>);
impl ValidationErrors {
/// Iterate over the individual errors.
pub fn iter(&self) -> std::slice::Iter<'_, ValidationError> {
self.0.iter()
}
/// Number of validation errors.
pub fn len(&self) -> usize {
self.0.len()
}
/// Returns true if there are no errors.
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
@@ -588,10 +555,6 @@ fn validate_identifier(field: &str, value: &str, errors: &mut Vec<ValidationErro
}
}
// ============================================================================
// Tests
// ============================================================================
#[cfg(test)]
mod tests {
use super::*;
@@ -644,8 +607,6 @@ mod tests {
assert!(!ArgumentType::Null.is_composite());
}
// -- SchemaType -----------------------------------------------------------
#[test]
fn schema_type_single_serde_roundtrip() {
let st = SchemaType::Single(ArgumentType::String);
@@ -752,17 +713,14 @@ mod tests {
#[test]
fn schema_type_primitive_composite_multiple() {
// All primitive → is_primitive=true, is_composite=false
let nullable_string = SchemaType::Multiple(vec![ArgumentType::String, ArgumentType::Null]);
assert!(nullable_string.is_primitive());
assert!(!nullable_string.is_composite());
// Any composite → is_primitive=false, is_composite=true
let nullable_array = SchemaType::Multiple(vec![ArgumentType::Array, ArgumentType::Null]);
assert!(!nullable_array.is_primitive());
assert!(nullable_array.is_composite());
// Mixed primitive + composite
let mixed = SchemaType::Multiple(vec![ArgumentType::String, ArgumentType::Object]);
assert!(!mixed.is_primitive());
assert!(mixed.is_composite());
@@ -915,7 +873,6 @@ mod tests {
assert_eq!(tool.to_input_schema(), raw);
}
/// Without a raw schema, `to_input_schema` returns an empty object schema.
#[test]
fn description_to_input_schema_empty_when_no_raw() {
let tool = ToolDescription::new("echo", "Echo a string");
@@ -1042,7 +999,6 @@ mod tests {
tool.namespace = Some(String::new());
let errors = tool.validate().unwrap_err();
// empty tool name + empty namespace = 2
assert!(errors.len() >= 2);
}
}