M0: compilable skeleton — Kigi 0.1.0 fork surgery
Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.
Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
ptyctl, ptyctl-cli, third_party/ unchanged; proto package
xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
(templates re-encrypted)
Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
module & dc_log, heap-profile uploader, auth-diagnostics uploader,
session-analytics halves of feedback; local zero-egress observability
preserved in new kigi-log crate (unified log, --debug firehose,
subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
shell util
Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted
Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean
Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
(new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
fast-worktree); RSS measurement tests serialized via serial_test
Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
notices sustained; kigi-tools ported-code notices extended; README,
CONTRIBUTING, SECURITY, AGENTS.md rewritten
Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-compaction"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Shared, transport-agnostic compaction engine for Grok chat and Grok Build."
|
||||
|
||||
[dependencies]
|
||||
anyhow = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
thiserror = { workspace = true }
|
||||
tokio = { workspace = true, features = ["time", "macros"] }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = { workspace = true }
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,202 @@
|
||||
//! Compacted-history assembly (grok-build'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
|
||||
//! the canonical post-compaction history:
|
||||
//!
|
||||
//! ```text
|
||||
//! [SP, UP', AGENTS_MD?, UQ_last?, recent…, summary, reminder?]
|
||||
//! ```
|
||||
//!
|
||||
//! grok-build is the canonical harness. The summary carrier text is built by
|
||||
//! [`super::summary::format_compact_summary_content`].
|
||||
|
||||
use crate::item::CompactionItemFactory;
|
||||
|
||||
use super::summary::{format_compact_summary_content, wrap_user_query};
|
||||
|
||||
/// Input data for building a compacted conversation history.
|
||||
///
|
||||
/// All fields are plain data — no I/O, no network, no shell dependencies.
|
||||
/// The caller is responsible for:
|
||||
/// - Generating the `compaction_summary` via the LLM.
|
||||
/// - Rendering the optional `system_reminder` (which may depend on
|
||||
/// harness-specific backends such as memory search).
|
||||
/// - 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,
|
||||
/// Pre-rendered AGENTS.md `<system-reminder>` block to re-inject after the
|
||||
/// user prefix. `None` means no project instructions to re-inject.
|
||||
pub agents_md_reminder: Option<String>,
|
||||
/// The last real user query text (raw, unwrapped).
|
||||
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.
|
||||
pub system_reminder: Option<String>,
|
||||
/// Pre-built transcript hint appended to the summary (`None` to omit).
|
||||
pub transcript_hint: Option<String>,
|
||||
}
|
||||
|
||||
/// Build the compacted conversation history from pure data inputs.
|
||||
///
|
||||
/// The returned `Vec<T>` is structured as:
|
||||
///
|
||||
/// 1. **System message** -- the original system prompt.
|
||||
/// 2. **User message prefix** -- e.g. `<user_info>` block (no `<user_query>` tags).
|
||||
/// 3. **AGENTS.md reminder** (if any) -- project instructions re-injected verbatim.
|
||||
/// 4. **Last user query** (if any) -- wrapped in `<user_query>` tags.
|
||||
/// 5. **Recent messages** (if any) -- retained verbatim from after the last
|
||||
/// real user turn.
|
||||
/// 6. **Compaction summary** -- with the optional `<system-reminder>`
|
||||
/// appended as a separate message.
|
||||
///
|
||||
/// This is a pure function with no I/O.
|
||||
pub fn assemble_compacted_history<T: CompactionItemFactory>(
|
||||
parts: CompactedHistoryParts<T>,
|
||||
) -> Vec<T> {
|
||||
let mut compacted: Vec<T> = vec![
|
||||
parts.system_message,
|
||||
T::new_user_meta(parts.user_message_prefix),
|
||||
];
|
||||
|
||||
// Re-inject AGENTS.md as a user message so project instructions survive
|
||||
// compaction verbatim (not dependent on the summarizer). The
|
||||
// `ProjectInstructions` tag is what the spawn-time idempotence guard
|
||||
// recognizes on resume, so post-compaction sessions stay duplicate-free.
|
||||
if let Some(ref reminder) = parts.agents_md_reminder {
|
||||
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())));
|
||||
}
|
||||
|
||||
// grok-build keeps the legacy `<user_query>`-wrapped continuation text and
|
||||
// appends the transcript hint after the continuation summary.
|
||||
let mut formatted_summary = format_compact_summary_content(&parts.compaction_summary);
|
||||
if let Some(ref hint) = parts.transcript_hint {
|
||||
formatted_summary.push_str(hint);
|
||||
}
|
||||
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);
|
||||
}
|
||||
compacted.push(summary_item);
|
||||
|
||||
if let Some(ref reminder) = parts.system_reminder {
|
||||
compacted.push(T::new_system_reminder(reminder.clone()));
|
||||
}
|
||||
|
||||
compacted
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Minimal mock item recording which factory constructor produced it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) enum MockItem {
|
||||
System(String),
|
||||
User(String),
|
||||
UserMeta(String),
|
||||
ProjectInstructions(String),
|
||||
SystemReminder(String),
|
||||
Recent(String),
|
||||
}
|
||||
|
||||
impl CompactionItemFactory for MockItem {
|
||||
fn new_user(text: String) -> Self {
|
||||
Self::User(text)
|
||||
}
|
||||
fn new_user_meta(text: String) -> Self {
|
||||
Self::UserMeta(text)
|
||||
}
|
||||
fn new_project_instructions(text: String) -> Self {
|
||||
Self::ProjectInstructions(text)
|
||||
}
|
||||
fn new_system_reminder(text: String) -> Self {
|
||||
Self::SystemReminder(text)
|
||||
}
|
||||
}
|
||||
|
||||
fn parts(recent: Vec<MockItem>) -> CompactedHistoryParts<MockItem> {
|
||||
CompactedHistoryParts {
|
||||
system_message: MockItem::System("sys".into()),
|
||||
user_message_prefix: "<user_info>OS: macos</user_info>".into(),
|
||||
agents_md_reminder: Some("AGENTS.md content".into()),
|
||||
last_user_query: Some("fix the bug".into()),
|
||||
recent_messages: recent,
|
||||
compaction_summary: "Summary: did things.".into(),
|
||||
system_reminder: Some("<system-reminder>state</system-reminder>".into()),
|
||||
transcript_hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grok_build_order_recent_before_summary() {
|
||||
let recent = vec![MockItem::Recent("a1".into()), MockItem::Recent("t1".into())];
|
||||
let out = assemble_compacted_history(parts(recent));
|
||||
// [sys, prefix, agents_md, query, a1, t1, summary, reminder]
|
||||
assert_eq!(out.len(), 8);
|
||||
assert_eq!(out[0], MockItem::System("sys".into()));
|
||||
assert_eq!(
|
||||
out[1],
|
||||
MockItem::UserMeta("<user_info>OS: macos</user_info>".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[2],
|
||||
MockItem::ProjectInstructions("AGENTS.md content".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[3],
|
||||
MockItem::User("<user_query>\nfix the bug\n</user_query>".into())
|
||||
);
|
||||
assert_eq!(out[4], MockItem::Recent("a1".into()));
|
||||
assert_eq!(out[5], MockItem::Recent("t1".into()));
|
||||
let MockItem::UserMeta(summary) = &out[6] else {
|
||||
panic!("expected UserMeta summary, got {:?}", out[6]);
|
||||
};
|
||||
assert!(summary.starts_with("This session is being continued"));
|
||||
assert_eq!(
|
||||
out[7],
|
||||
MockItem::SystemReminder("<system-reminder>state</system-reminder>".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn omits_optional_sections() {
|
||||
let mut p = parts(vec![]);
|
||||
p.agents_md_reminder = None;
|
||||
p.last_user_query = None;
|
||||
p.system_reminder = None;
|
||||
let out = assemble_compacted_history(p);
|
||||
// [sys, prefix, summary]
|
||||
assert_eq!(out.len(), 3);
|
||||
assert!(
|
||||
matches!(&out[2], MockItem::UserMeta(s) if s.starts_with("This session is being continued"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_transcript_hint_after_summary() {
|
||||
let mut p = parts(vec![]);
|
||||
p.transcript_hint = Some("\n\n<transcript_location>/x</transcript_location>".into());
|
||||
let out = assemble_compacted_history(p);
|
||||
let MockItem::UserMeta(summary) = &out[4] else {
|
||||
panic!("expected UserMeta summary, got {:?}", out[4]);
|
||||
};
|
||||
assert!(summary.ends_with("</transcript_location>"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,608 @@
|
||||
//! grok-build's full-replace compaction pass.
|
||||
//!
|
||||
//! grok-build does not select a tail to keep; it summarizes the whole
|
||||
//! conversation and rebuilds a fresh history from scratch. This module is the
|
||||
//! transport-agnostic orchestration of that pass:
|
||||
//!
|
||||
//! ```text
|
||||
//! 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.
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tracing::info;
|
||||
|
||||
use crate::item::CompactionItemFactory;
|
||||
use crate::prompt::CompactionPrompt;
|
||||
use crate::sampler::CompactionSampler;
|
||||
|
||||
use super::assemble::{CompactedHistoryParts, assemble_compacted_history};
|
||||
use super::config::FullReplaceConfig;
|
||||
use super::observer::FullReplaceObserver;
|
||||
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).
|
||||
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.
|
||||
pub last_user_query: Option<String>,
|
||||
/// Working tail retained verbatim (tool/subagent results from the current
|
||||
/// turn). grok-build keeps this; 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.
|
||||
pub system_reminder: Option<String>,
|
||||
/// Optional 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.
|
||||
deterministic: bool,
|
||||
/// Whether the failure was a context-length overflow. The product host
|
||||
/// uses this to step its input ladder (rebuild a smaller input and
|
||||
/// call this pass again) instead of suppressing.
|
||||
context_overflow: bool,
|
||||
},
|
||||
}
|
||||
|
||||
impl std::fmt::Display for FullReplaceError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::NothingToCompact => write!(f, "nothing to compact"),
|
||||
Self::EmptyResponse => write!(f, "compaction model returned an empty summary"),
|
||||
Self::Sampler { message, .. } => write!(f, "compaction sampling failed: {message}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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?]`).
|
||||
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).
|
||||
pub attempts: u32,
|
||||
}
|
||||
|
||||
/// A successful full-replace **sampling** pass (summary only, no assembly).
|
||||
///
|
||||
/// Returned by [`sample_full_replace_summary`] for harnesses (grok-build'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).
|
||||
pub attempts: u32,
|
||||
}
|
||||
|
||||
/// Run grok-build'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).
|
||||
///
|
||||
/// 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],
|
||||
user_context: Option<&str>,
|
||||
ctx: FullReplaceContext<T>,
|
||||
config: &FullReplaceConfig,
|
||||
observer: &O,
|
||||
) -> Result<FullReplaceOutput<T>, FullReplaceError>
|
||||
where
|
||||
T: CompactionItemFactory + Send + Sync,
|
||||
S: CompactionSampler<Item = T> + ?Sized,
|
||||
O: FullReplaceObserver + ?Sized,
|
||||
{
|
||||
let FullReplaceSummary { summary, attempts } =
|
||||
sample_full_replace_summary(sampler, llm_turns, user_context, config, observer).await?;
|
||||
|
||||
info!(
|
||||
turns = llm_turns.len(),
|
||||
summary_chars = summary.len(),
|
||||
attempts,
|
||||
"[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.
|
||||
let parts = CompactedHistoryParts {
|
||||
system_message: ctx.system_message,
|
||||
user_message_prefix: ctx.user_message_prefix,
|
||||
agents_md_reminder: ctx.agents_md_reminder,
|
||||
last_user_query: ctx.last_user_query,
|
||||
recent_messages: ctx.recent_messages,
|
||||
compaction_summary: summary.clone(),
|
||||
system_reminder: ctx.system_reminder,
|
||||
transcript_hint: ctx.transcript_hint,
|
||||
};
|
||||
Ok(FullReplaceOutput {
|
||||
history: assemble_compacted_history(parts),
|
||||
summary,
|
||||
attempts,
|
||||
})
|
||||
}
|
||||
|
||||
/// 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 grok-build'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.
|
||||
pub async fn sample_full_replace_summary<T, S, O>(
|
||||
sampler: &S,
|
||||
llm_turns: &[T],
|
||||
user_context: Option<&str>,
|
||||
config: &FullReplaceConfig,
|
||||
observer: &O,
|
||||
) -> Result<FullReplaceSummary, FullReplaceError>
|
||||
where
|
||||
T: Send + Sync,
|
||||
S: CompactionSampler<Item = T> + ?Sized,
|
||||
O: FullReplaceObserver + ?Sized,
|
||||
{
|
||||
if llm_turns.is_empty() {
|
||||
return Err(FullReplaceError::NothingToCompact);
|
||||
}
|
||||
|
||||
let prompt = CompactionPrompt {
|
||||
// grok-build appends the summarization prompt as the final user
|
||||
// message; there is no separate system prompt for the compaction call.
|
||||
system: String::new(),
|
||||
user: build_summary_prompt(user_context),
|
||||
};
|
||||
let timeout = Duration::from_secs(config.sampling_timeout_secs);
|
||||
let started = Instant::now();
|
||||
|
||||
match sample_summary_with_retries(
|
||||
sampler,
|
||||
llm_turns,
|
||||
&prompt,
|
||||
config.max_attempts,
|
||||
Duration::from_secs(config.retry_delay_secs),
|
||||
timeout,
|
||||
observer,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(SampledSummary { summary, attempts }) => {
|
||||
observer.on_success(attempts, summary.chars().count(), started.elapsed());
|
||||
Ok(FullReplaceSummary { summary, attempts })
|
||||
}
|
||||
Err(SampleRetryError::Empty { attempts }) => {
|
||||
observer.on_error(attempts);
|
||||
Err(FullReplaceError::EmptyResponse)
|
||||
}
|
||||
Err(SampleRetryError::Failure {
|
||||
message,
|
||||
deterministic,
|
||||
context_overflow,
|
||||
attempts,
|
||||
}) => {
|
||||
observer.on_error(attempts);
|
||||
Err(FullReplaceError::Sampler {
|
||||
message,
|
||||
deterministic,
|
||||
context_overflow,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::*;
|
||||
use crate::code_compaction::observer::FullReplaceAttemptOutcome;
|
||||
use crate::sampler::{CompactionSampleError, LlmCompactionOutput};
|
||||
|
||||
/// Mock item recording which factory constructor produced it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum MockItem {
|
||||
System(String),
|
||||
User(String),
|
||||
UserMeta(String),
|
||||
ProjectInstructions(String),
|
||||
SystemReminder(String),
|
||||
Tail(String),
|
||||
}
|
||||
|
||||
impl CompactionItemFactory for MockItem {
|
||||
fn new_user(text: String) -> Self {
|
||||
Self::User(text)
|
||||
}
|
||||
fn new_user_meta(text: String) -> Self {
|
||||
Self::UserMeta(text)
|
||||
}
|
||||
fn new_project_instructions(text: String) -> Self {
|
||||
Self::ProjectInstructions(text)
|
||||
}
|
||||
fn new_system_reminder(text: String) -> Self {
|
||||
Self::SystemReminder(text)
|
||||
}
|
||||
}
|
||||
|
||||
/// Mock sampler with scripted responses (consumed in order).
|
||||
struct MockSampler {
|
||||
responses: Mutex<Vec<Result<String, CompactionSampleError>>>,
|
||||
calls: Mutex<usize>,
|
||||
}
|
||||
|
||||
impl MockSampler {
|
||||
fn returns(text: &str) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(vec![Ok(text.to_string())]),
|
||||
calls: Mutex::new(0),
|
||||
}
|
||||
}
|
||||
fn scripted(responses: Vec<Result<String, CompactionSampleError>>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(responses),
|
||||
calls: Mutex::new(0),
|
||||
}
|
||||
}
|
||||
fn call_count(&self) -> usize {
|
||||
*self.calls.lock().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CompactionSampler for MockSampler {
|
||||
type Item = MockItem;
|
||||
|
||||
async fn sample_compaction(
|
||||
&self,
|
||||
_turns: &[MockItem],
|
||||
_prompt: &CompactionPrompt,
|
||||
_timeout: Duration,
|
||||
) -> Result<LlmCompactionOutput, CompactionSampleError> {
|
||||
*self.calls.lock().unwrap() += 1;
|
||||
let mut responses = self.responses.lock().unwrap();
|
||||
if responses.is_empty() {
|
||||
return Err(CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"no more scripted responses"
|
||||
)));
|
||||
}
|
||||
responses.remove(0).map(|response| LlmCompactionOutput {
|
||||
response,
|
||||
thinking: String::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn ctx(recent: Vec<MockItem>) -> FullReplaceContext<MockItem> {
|
||||
FullReplaceContext {
|
||||
system_message: MockItem::System("you are a helpful assistant".into()),
|
||||
user_message_prefix: "<user_info>OS: macos</user_info>".into(),
|
||||
agents_md_reminder: Some("# AGENTS.md\nbe nice".into()),
|
||||
last_user_query: Some("fix the login bug".into()),
|
||||
recent_messages: recent,
|
||||
system_reminder: Some(
|
||||
"<system-reminder>\n## Running Subagents\n- sub-1\n</system-reminder>".into(),
|
||||
),
|
||||
transcript_hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn cfg() -> FullReplaceConfig {
|
||||
FullReplaceConfig {
|
||||
max_attempts: 3,
|
||||
retry_delay_secs: 0,
|
||||
sampling_timeout_secs: 5,
|
||||
}
|
||||
}
|
||||
|
||||
/// A non-degenerate mock summary (cleaned seed >=
|
||||
/// [`crate::code_compaction::config::MIN_SUMMARY_SEED_CHARS`]).
|
||||
fn healthy_summary(primary: &str) -> String {
|
||||
let body = format!(
|
||||
"1. Primary Request: {primary}\n\
|
||||
2. Key Technical Concepts: Rust, auth, session tokens\n\
|
||||
3. Files and Code Sections: crates/foo/src/auth.rs — login handler\n\
|
||||
4. Errors and Fixes: None\n\
|
||||
5. Problem Solving: traced token validation failure\n\
|
||||
6. All User Messages: fix the login bug\n\
|
||||
7. Pending Tasks: run integration tests\n\
|
||||
8. Current Work: editing auth.rs login handler\n\
|
||||
9. Optional Next Step: run tests"
|
||||
);
|
||||
let padding = "x".repeat(
|
||||
crate::code_compaction::config::MIN_SUMMARY_SEED_CHARS.saturating_sub(body.len()),
|
||||
);
|
||||
format!(
|
||||
"<analysis>\nthinking about it\n</analysis>\n\n\
|
||||
<summary>\n{body}\n{padding}\n</summary>"
|
||||
)
|
||||
}
|
||||
|
||||
/// Golden end-to-end test: a realistic conversation + a mock sampler that
|
||||
/// returns a structured summary must produce grok-build's exact compacted
|
||||
/// history shape, with the LLM output cleaned and the agent-state reminder
|
||||
/// carried through as the final item.
|
||||
#[tokio::test]
|
||||
async fn full_replace_produces_grok_build_history_shape() {
|
||||
let llm_turns = vec![
|
||||
MockItem::System("you are a helpful assistant".into()),
|
||||
MockItem::User("fix the login bug".into()),
|
||||
MockItem::Tail("assistant: looked at auth.rs".into()),
|
||||
];
|
||||
let recent = vec![MockItem::Tail("tool: read_file(auth.rs) -> ...".into())];
|
||||
let sampler = MockSampler::returns(&healthy_summary("fix login bug"));
|
||||
|
||||
let out =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(recent), &cfg(), &())
|
||||
.await
|
||||
.expect("compaction should succeed")
|
||||
.history;
|
||||
|
||||
// [system, prefix, agents_md, last_query, recent_tail, summary, reminder]
|
||||
assert_eq!(out.len(), 7, "got: {out:#?}");
|
||||
assert_eq!(
|
||||
out[0],
|
||||
MockItem::System("you are a helpful assistant".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[1],
|
||||
MockItem::UserMeta("<user_info>OS: macos</user_info>".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[2],
|
||||
MockItem::ProjectInstructions("# AGENTS.md\nbe nice".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[3],
|
||||
MockItem::User("<user_query>\nfix the login bug\n</user_query>".into())
|
||||
);
|
||||
assert_eq!(
|
||||
out[4],
|
||||
MockItem::Tail("tool: read_file(auth.rs) -> ...".into())
|
||||
);
|
||||
|
||||
// Summary carrier: cleaned (no <analysis>/<summary> tags), with preamble.
|
||||
let MockItem::UserMeta(summary) = &out[5] else {
|
||||
panic!("expected UserMeta summary at [5], got {:?}", out[5]);
|
||||
};
|
||||
assert!(summary.starts_with("This session is being continued"));
|
||||
assert!(summary.contains("Summary:\n1. Primary Request: fix login bug"));
|
||||
assert!(
|
||||
!summary.contains("<analysis>"),
|
||||
"scratchpad leaked: {summary}"
|
||||
);
|
||||
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(
|
||||
"<system-reminder>\n## Running Subagents\n- sub-1\n</system-reminder>".into()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_turns_is_nothing_to_compact() {
|
||||
let sampler = MockSampler::returns("unused");
|
||||
let result =
|
||||
apply_full_replace_compaction(&sampler, &[], None, ctx(vec![]), &cfg(), &()).await;
|
||||
assert!(matches!(result, Err(FullReplaceError::NothingToCompact)));
|
||||
assert_eq!(sampler.call_count(), 0, "must not call the LLM");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retries_transient_then_succeeds() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Timeout {
|
||||
timeout_secs: 5,
|
||||
collected_bytes: 0,
|
||||
}),
|
||||
Ok(healthy_summary("q")),
|
||||
]);
|
||||
let out =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await
|
||||
.expect("should succeed after one retry")
|
||||
.history;
|
||||
assert_eq!(sampler.call_count(), 2);
|
||||
assert!(matches!(out.last(), Some(MockItem::SystemReminder(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deterministic_failure_does_not_retry() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Build("bad model".into())),
|
||||
Ok("never reached".into()),
|
||||
]);
|
||||
let result =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await;
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(FullReplaceError::Sampler {
|
||||
deterministic: true,
|
||||
context_overflow: false,
|
||||
..
|
||||
})
|
||||
));
|
||||
assert_eq!(
|
||||
sampler.call_count(),
|
||||
1,
|
||||
"deterministic error must not retry"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_response_after_retries_errors() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let sampler = MockSampler::scripted(vec![Ok(" ".into()), Ok("".into()), Ok("".into())]);
|
||||
let result =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await;
|
||||
assert!(matches!(result, Err(FullReplaceError::EmptyResponse)));
|
||||
assert_eq!(sampler.call_count(), 3);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn degenerate_summary_retries_then_succeeds() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let short = "<summary>\n1. Primary Request: q\n</summary>";
|
||||
let long = format!(
|
||||
"<summary>\n1. Primary Request: fix the login bug\n{}\n</summary>",
|
||||
"x".repeat(600)
|
||||
);
|
||||
let sampler = MockSampler::scripted(vec![Ok(short.into()), Ok(long.clone())]);
|
||||
let out =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await
|
||||
.expect("should succeed after degenerate retry")
|
||||
.history;
|
||||
assert_eq!(sampler.call_count(), 2);
|
||||
let MockItem::UserMeta(summary) = &out[out.len() - 2] else {
|
||||
panic!("expected summary carrier");
|
||||
};
|
||||
assert!(summary.contains("fix the login bug"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn degenerate_summary_after_retries_errors() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let short = "<summary>\n1. Primary Request: q\n</summary>";
|
||||
let sampler =
|
||||
MockSampler::scripted(vec![Ok(short.into()), Ok(short.into()), Ok(short.into())]);
|
||||
let result =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await;
|
||||
assert!(matches!(result, Err(FullReplaceError::EmptyResponse)));
|
||||
assert_eq!(sampler.call_count(), 3);
|
||||
}
|
||||
|
||||
/// A context-length overflow must short-circuit (no retry) and surface
|
||||
/// `context_overflow = true` so the product host steps its input ladder.
|
||||
#[tokio::test]
|
||||
async fn context_overflow_is_terminal_and_flagged() {
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"API error (status 400): The prompt is too long for this model's context window."
|
||||
))),
|
||||
Ok(healthy_summary("never reached")),
|
||||
]);
|
||||
let result =
|
||||
apply_full_replace_compaction(&sampler, &llm_turns, None, ctx(vec![]), &cfg(), &())
|
||||
.await;
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(FullReplaceError::Sampler {
|
||||
context_overflow: true,
|
||||
deterministic: true,
|
||||
..
|
||||
})
|
||||
));
|
||||
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;
|
||||
|
||||
#[derive(Default)]
|
||||
struct RecordingObserver {
|
||||
attempts: Mutex<Vec<String>>,
|
||||
successes: Mutex<u32>,
|
||||
errors: Mutex<u32>,
|
||||
}
|
||||
impl FullReplaceObserver for RecordingObserver {
|
||||
fn on_attempt(&self, _attempt: u32, outcome: &FullReplaceAttemptOutcome<'_>) {
|
||||
let tag = match outcome {
|
||||
FullReplaceAttemptOutcome::Success { .. } => "success",
|
||||
FullReplaceAttemptOutcome::EmptyResponse { .. } => "empty",
|
||||
FullReplaceAttemptOutcome::Degenerate { .. } => "degenerate",
|
||||
FullReplaceAttemptOutcome::Failure { .. } => "failure",
|
||||
};
|
||||
self.attempts.lock().unwrap().push(tag.to_string());
|
||||
}
|
||||
fn on_success(&self, _attempts: u32, _summary_chars: usize, _elapsed: Duration) {
|
||||
*self.successes.lock().unwrap() += 1;
|
||||
}
|
||||
fn on_error(&self, _attempts: u32) {
|
||||
*self.errors.lock().unwrap() += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let llm_turns = vec![MockItem::User("q".into())];
|
||||
let short = "<summary>\n1. Primary Request: q\n</summary>";
|
||||
let sampler = MockSampler::scripted(vec![Ok(short.into()), Ok(healthy_summary("q"))]);
|
||||
let observer = RecordingObserver::default();
|
||||
let out = apply_full_replace_compaction(
|
||||
&sampler,
|
||||
&llm_turns,
|
||||
None,
|
||||
ctx(vec![]),
|
||||
&cfg(),
|
||||
&observer,
|
||||
)
|
||||
.await
|
||||
.expect("should succeed");
|
||||
assert_eq!(out.attempts, 2);
|
||||
assert_eq!(
|
||||
*observer.attempts.lock().unwrap(),
|
||||
vec!["degenerate", "success"]
|
||||
);
|
||||
assert_eq!(*observer.successes.lock().unwrap(), 1);
|
||||
assert_eq!(*observer.errors.lock().unwrap(), 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//! grok-build compaction configuration.
|
||||
//!
|
||||
//! Holds the [`FullReplaceConfig`] tunables struct (mirroring
|
||||
//! [`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.
|
||||
|
||||
/// Default auto-compact threshold (% of context window) when no other source
|
||||
/// (env var, user config, remote per-model/global flags) sets it. Shared by
|
||||
/// grok-build and Grok chat (~85% trigger on both sides).
|
||||
pub const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT: u8 = 85;
|
||||
|
||||
/// Minimum character count for a cleaned summary seed.
|
||||
///
|
||||
/// grok-build 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.
|
||||
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.
|
||||
pub max_attempts: u32,
|
||||
/// Delay between transient retries.
|
||||
pub retry_delay_secs: u64,
|
||||
/// End-to-end timeout for each compaction LLM call.
|
||||
pub sampling_timeout_secs: u64,
|
||||
}
|
||||
|
||||
impl Default for FullReplaceConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_attempts: 3,
|
||||
retry_delay_secs: 3,
|
||||
sampling_timeout_secs: 120,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
//! 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. grok-build's `SamplingError` →
|
||||
//! `CompactFailure(acp::Error)`) stay in thin host wrappers that delegate the
|
||||
//! status/message decisions to these functions.
|
||||
|
||||
/// Whether a compaction-call failure is worth retrying.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FailureKind {
|
||||
/// Retrying the same payload will hit the same failure — the retry loop
|
||||
/// should bail without sleeping or re-issuing.
|
||||
Deterministic,
|
||||
/// Failure may resolve on retry (network blips, 5xx, rate limits).
|
||||
Transient,
|
||||
}
|
||||
|
||||
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.
|
||||
pub fn is_context_length_error(message: &str) -> bool {
|
||||
let m = message.to_ascii_lowercase();
|
||||
m.contains("too long for this model")
|
||||
|| m.contains("prompt is too long")
|
||||
|| m.contains("maximum prompt length")
|
||||
|| m.contains("maximum context length")
|
||||
|| m.contains("context_length_exceeded")
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn classify_http_status(status: u16, message: &str) -> FailureKind {
|
||||
if is_context_length_error(message)
|
||||
|| ((400..500).contains(&status) && status != 408 && status != 429)
|
||||
{
|
||||
FailureKind::Deterministic
|
||||
} else {
|
||||
FailureKind::Transient
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
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;
|
||||
}
|
||||
|
||||
if let Some(status_code) = code.and_then(|c| c.parse::<u16>().ok())
|
||||
&& (400..500).contains(&status_code)
|
||||
&& status_code != 408
|
||||
&& status_code != 429
|
||||
{
|
||||
return FailureKind::Deterministic;
|
||||
}
|
||||
|
||||
// Size overflow arrives here with no parseable code (`code="none"`); the
|
||||
// message is the only signal that re-sending cannot help.
|
||||
if is_context_length_error(message) {
|
||||
return FailureKind::Deterministic;
|
||||
}
|
||||
|
||||
FailureKind::Transient
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn det_status(status: u16) -> bool {
|
||||
classify_http_status(status, "test").is_deterministic()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_4xx_is_deterministic_except_408_and_429() {
|
||||
assert!(det_status(400));
|
||||
assert!(det_status(401));
|
||||
assert!(det_status(403));
|
||||
assert!(det_status(404));
|
||||
assert!(det_status(413));
|
||||
assert!(!det_status(408));
|
||||
assert!(!det_status(429));
|
||||
assert!(!det_status(500));
|
||||
assert!(!det_status(502));
|
||||
assert!(!det_status(503));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_500_with_context_length_message_is_deterministic() {
|
||||
// The sampler synthesizes status=500 from a streamed size overflow, so
|
||||
// status alone reads transient; the message must still short-circuit.
|
||||
assert!(
|
||||
classify_http_status(
|
||||
500,
|
||||
"API error (status 500 Internal Server Error): \
|
||||
The prompt is too long for this model's context window."
|
||||
)
|
||||
.is_deterministic()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_event_invalid_request_error_marker_is_deterministic() {
|
||||
assert!(
|
||||
classify_stream_event_error(
|
||||
Some("invalid_request_error"),
|
||||
"messages.27.content.1: ..."
|
||||
)
|
||||
.is_deterministic()
|
||||
);
|
||||
assert!(
|
||||
classify_stream_event_error(
|
||||
Some("400"),
|
||||
"Provider returned invalid_request_error: messages.X..."
|
||||
)
|
||||
.is_deterministic()
|
||||
);
|
||||
assert!(
|
||||
classify_stream_event_error(None, "messages.X.content.Y: invalid_request_error: ...")
|
||||
.is_deterministic()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_event_numeric_codes_match_http_classification() {
|
||||
let det = |c: &str| classify_stream_event_error(Some(c), "msg").is_deterministic();
|
||||
assert!(det("400"));
|
||||
assert!(det("401"));
|
||||
assert!(det("403"));
|
||||
assert!(det("404"));
|
||||
assert!(!det("408"));
|
||||
assert!(!det("429"));
|
||||
assert!(!det("500"));
|
||||
assert!(!det("503"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_event_unknown_code_defaults_to_transient() {
|
||||
assert!(!classify_stream_event_error(None, "msg").is_deterministic());
|
||||
assert!(!classify_stream_event_error(Some("error"), "msg").is_deterministic());
|
||||
assert!(!classify_stream_event_error(Some("overloaded_error"), "msg").is_deterministic());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_event_context_length_message_is_deterministic() {
|
||||
assert!(
|
||||
classify_stream_event_error(
|
||||
None,
|
||||
"The prompt is too long for this model's context window."
|
||||
)
|
||||
.is_deterministic()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn context_length_error_matches_known_messages() {
|
||||
for msg in [
|
||||
"The prompt is too long for this model's context window.",
|
||||
"prompt is too long: 250000 tokens > 200000 maximum",
|
||||
"exceeds the maximum prompt length",
|
||||
"This model's maximum context length is 128000 tokens",
|
||||
"error code: context_length_exceeded",
|
||||
] {
|
||||
assert!(is_context_length_error(msg), "should match: {msg}");
|
||||
}
|
||||
for msg in [
|
||||
"internal server error",
|
||||
"rate limited",
|
||||
"connection reset by peer",
|
||||
] {
|
||||
assert!(!is_context_length_error(msg), "should not match: {msg}");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
//! grok-build's "code agent" compaction subsystem.
|
||||
//!
|
||||
//! grok-build does not select a tail to keep; it summarizes the whole
|
||||
//! conversation and rebuilds a fresh history from scratch (the *full-replace*
|
||||
//! strategy). This module groups that subsystem — generic over the engine's
|
||||
//! [`CompactionItem`](crate::item::CompactionItem) /
|
||||
//! [`CompactionItemFactory`](crate::item::CompactionItemFactory) seams — so it
|
||||
//! can be reused as a unit by grok-build, separate from Grok chat's
|
||||
//! [`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`).
|
||||
|
||||
pub mod assemble;
|
||||
pub mod compact;
|
||||
pub mod config;
|
||||
pub mod failure;
|
||||
pub mod observer;
|
||||
pub mod prompt;
|
||||
pub mod sample;
|
||||
pub mod summary;
|
||||
|
||||
pub use assemble::{CompactedHistoryParts, assemble_compacted_history};
|
||||
pub use compact::{
|
||||
FullReplaceContext, FullReplaceError, FullReplaceOutput, FullReplaceSummary,
|
||||
apply_full_replace_compaction, sample_full_replace_summary,
|
||||
};
|
||||
pub use config::{
|
||||
DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT, FullReplaceConfig, MIN_SUMMARY_SEED_CHARS,
|
||||
};
|
||||
pub use failure::{
|
||||
FailureKind, classify_http_status, classify_stream_event_error, is_context_length_error,
|
||||
};
|
||||
pub use observer::{FullReplaceAttemptOutcome, FullReplaceObserver};
|
||||
pub use prompt::{
|
||||
SELF_SUMMARIZATION_PROMPT, SummaryPromptKind, build_summary_prompt, build_summary_prompt_kind,
|
||||
};
|
||||
pub use sample::{SampleRetryError, SampledSummary, sample_summary_with_retries};
|
||||
pub use summary::{
|
||||
format_compact_summary, format_compact_summary_content, is_degenerate_summary, wrap_user_query,
|
||||
};
|
||||
@@ -0,0 +1,73 @@
|
||||
//! Observability seam for the full-replace (grok-build) pass.
|
||||
//!
|
||||
//! The shared orchestrator reports per-attempt and terminal outcomes through
|
||||
//! this trait so each harness can emit its own telemetry (grok-build:
|
||||
//! `CompactionAttempt` rows, `CompactionRetryDegraded` events, span records,
|
||||
//! request-artifact persistence) without the shared crate depending on a
|
||||
//! telemetry backend. Mirrors
|
||||
//! [`IntraCompactionObserver`](crate::intra_compaction::IntraCompactionObserver)
|
||||
//! / [`InterCompactionObserver`](crate::inter_compaction::InterCompactionObserver).
|
||||
//!
|
||||
//! Emission points are part of the behavior contract: the grok-build observer
|
||||
//! preserves the pre-migration `CompactionAttempt`/`CompactionRetryDegraded`
|
||||
//! semantics byte-for-byte.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
/// Classified outcome of a single full-replace sample attempt.
|
||||
///
|
||||
/// The harness turns this into its per-attempt telemetry row. `summary` is the
|
||||
/// raw model output (the harness bounds/captures it as needed); it is borrowed
|
||||
/// for the duration of the callback so no allocation happens on the hot path.
|
||||
#[derive(Debug)]
|
||||
pub enum FullReplaceAttemptOutcome<'a> {
|
||||
/// A usable, non-degenerate summary was produced; the pass will succeed.
|
||||
Success {
|
||||
/// Raw model summary text.
|
||||
summary: &'a str,
|
||||
},
|
||||
/// The model returned an empty / whitespace-only response.
|
||||
EmptyResponse {
|
||||
/// Whether the orchestrator will retry after this attempt.
|
||||
will_retry: bool,
|
||||
},
|
||||
/// 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).
|
||||
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`.
|
||||
deterministic: bool,
|
||||
/// Whether the failure was a context-length overflow — the signal the
|
||||
/// harness uses to step its input ladder rather than suppress.
|
||||
context_overflow: bool,
|
||||
/// Whether the orchestrator will retry after this attempt (always
|
||||
/// `false` for deterministic failures and context overflows).
|
||||
will_retry: bool,
|
||||
},
|
||||
}
|
||||
|
||||
/// Receives full-replace compaction outcomes. All methods default to no-ops so
|
||||
/// harnesses without telemetry (and tests) can use `()`.
|
||||
pub trait FullReplaceObserver: Send + Sync {
|
||||
/// One sample attempt finished with the given classified outcome.
|
||||
/// `attempt` is 1-based and cumulative across the pass.
|
||||
fn on_attempt(&self, _attempt: u32, _outcome: &FullReplaceAttemptOutcome<'_>) {}
|
||||
|
||||
/// The pass succeeded after `attempts` total attempts.
|
||||
fn on_success(&self, _attempts: u32, _summary_chars: usize, _elapsed: Duration) {}
|
||||
|
||||
/// The pass failed terminally after `attempts` total attempts.
|
||||
fn on_error(&self, _attempts: u32) {}
|
||||
}
|
||||
|
||||
/// No-op observer for tests and harnesses without telemetry.
|
||||
impl FullReplaceObserver for () {}
|
||||
@@ -0,0 +1,141 @@
|
||||
//! grok-build's session-level summarization prompt.
|
||||
//!
|
||||
//! Split out of the crate-root `prompt` module so grok-build's full-replace
|
||||
//! prompt lives alongside the rest of its [`code_compaction`](crate::code_compaction)
|
||||
//! subsystem. The Grok chat's step-level intra prompt
|
||||
//! ([`format_compaction_prompt`](crate::prompt::format_compaction_prompt))
|
||||
//! stays at the crate root.
|
||||
|
||||
/// Build grok-build's session-level summarization prompt (no chat history).
|
||||
///
|
||||
/// `user_context` is the optional `/compact <text>` user-provided context,
|
||||
/// spliced inline into the structured prompt. Ported verbatim from
|
||||
/// `kigi-shell::session::helpers::session_compact::build_compaction_prompt`
|
||||
/// (the `use_short_prompt == false` branch).
|
||||
pub fn build_summary_prompt(user_context: Option<&str>) -> String {
|
||||
let user_context_section = match user_context {
|
||||
Some(context) => format!(
|
||||
"\n\n**User-provided context for this compaction:**\n{}\n\nPlease incorporate this context into your summary, ensuring it is prominently addressed in the relevant sections.\n\n",
|
||||
context
|
||||
),
|
||||
None => String::new(),
|
||||
};
|
||||
|
||||
include_str!("templates/full_replace_summary_prompt.txt")
|
||||
.replace("{user_context_section}", &user_context_section)
|
||||
}
|
||||
|
||||
/// The short "self-summarization" prompt variant
|
||||
/// (mirrors `kigi-shell`'s `SELF_SUMMARIZATION_PROMPT`). Framed
|
||||
/// as "summarize for a successor assistant that only sees the user's original
|
||||
/// query plus this summary." Kept here so every harness (the shell and the
|
||||
/// harness crate) shares one definition instead of each carrying a
|
||||
/// private copy.
|
||||
pub const SELF_SUMMARIZATION_PROMPT: &str = r#"<summary_request>
|
||||
Please summarize the conversation so far. This summary (everything after your
|
||||
thinking) will be provided to another AI assistant to continue working on the
|
||||
task. The other assistant will only see the user's original query and your
|
||||
summary, it will not have access to any tool calls or tool outputs from this
|
||||
conversation. The purpose of the summary is to compress the conversation
|
||||
context while preserving the essential information needed to seamlessly
|
||||
continue. Useful things to include: the user's requests, what you've done so
|
||||
far, relevant file paths and code details, any errors encountered and how
|
||||
they were resolved, and what remains to be done. DO NOT call any tools in
|
||||
your response.
|
||||
</summary_request>"#;
|
||||
|
||||
/// Which summarization prompt a full-replace pass should send.
|
||||
///
|
||||
/// The prompt is owned by the harness's [`CompactionSampler`] impl (it appends
|
||||
/// the prompt as the final user message before sampling), not by the shared
|
||||
/// orchestrator. This enum lets each harness select the right one in one place
|
||||
/// so the structured (grok-build) and short self-summary prompts stay
|
||||
/// shared instead of duplicated per harness.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum SummaryPromptKind {
|
||||
/// grok-build's detailed, numbered-section summary prompt.
|
||||
#[default]
|
||||
Structured,
|
||||
/// The short self-summarization prompt.
|
||||
SelfSummary,
|
||||
}
|
||||
|
||||
/// Build the full-replace summarization prompt for the given [`SummaryPromptKind`].
|
||||
///
|
||||
/// `user_context` is the optional `/compact <text>` user-provided context.
|
||||
/// For [`SummaryPromptKind::Structured`] it is spliced inline (see
|
||||
/// [`build_summary_prompt`]); for [`SummaryPromptKind::SelfSummary`] it is
|
||||
/// appended as a sibling `<user_provided_context>` block, matching the shell's
|
||||
/// `build_compaction_prompt(use_short_prompt = true)` behavior.
|
||||
pub fn build_summary_prompt_kind(kind: SummaryPromptKind, user_context: Option<&str>) -> String {
|
||||
match kind {
|
||||
SummaryPromptKind::Structured => build_summary_prompt(user_context),
|
||||
SummaryPromptKind::SelfSummary => match user_context {
|
||||
Some(ctx) => format!(
|
||||
"{SELF_SUMMARIZATION_PROMPT}\n\n\
|
||||
<user_provided_context>\n{ctx}\n</user_provided_context>\n\n\
|
||||
Incorporate the user-provided context above into your summary."
|
||||
),
|
||||
None => SELF_SUMMARIZATION_PROMPT.to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn summary_prompt_splices_context_section_inline() {
|
||||
let p = build_summary_prompt(Some("focus on auth"));
|
||||
assert!(p.contains("**User-provided context for this compaction:**\nfocus on auth"));
|
||||
assert!(p.contains("1. Primary Request and Intent"));
|
||||
assert!(p.contains("9. Optional Next Step"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_prompt_without_context_has_no_context_header() {
|
||||
let p = build_summary_prompt(None);
|
||||
assert!(!p.contains("**User-provided context for this compaction:**"));
|
||||
assert!(p.contains("6. All User Messages"));
|
||||
// Current prompt: no separate analysis block, concise framing.
|
||||
assert!(p.contains("do NOT emit a separate analysis block"));
|
||||
assert!(p.contains("faithful, concise summary"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kind_structured_matches_build_summary_prompt() {
|
||||
// The Structured kind must be byte-identical to the legacy entry point
|
||||
// so routing through the selector never changes grok-build's prompt.
|
||||
assert_eq!(
|
||||
build_summary_prompt_kind(SummaryPromptKind::Structured, None),
|
||||
build_summary_prompt(None)
|
||||
);
|
||||
assert_eq!(
|
||||
build_summary_prompt_kind(SummaryPromptKind::Structured, Some("focus on auth")),
|
||||
build_summary_prompt(Some("focus on auth"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kind_self_summary_without_context_is_bare_prompt() {
|
||||
let p = build_summary_prompt_kind(SummaryPromptKind::SelfSummary, None);
|
||||
assert_eq!(p, SELF_SUMMARIZATION_PROMPT);
|
||||
assert!(p.contains("<summary_request>"));
|
||||
// Must NOT carry the structured prompt's numbered sections.
|
||||
assert!(!p.contains("1. Primary Request and Intent"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn kind_self_summary_with_context_appends_sibling_block() {
|
||||
let p = build_summary_prompt_kind(SummaryPromptKind::SelfSummary, Some("focus on auth"));
|
||||
assert!(p.starts_with(SELF_SUMMARIZATION_PROMPT));
|
||||
assert!(p.contains("<user_provided_context>\nfocus on auth\n</user_provided_context>"));
|
||||
assert!(p.contains("Incorporate the user-provided context above"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_kind_is_structured() {
|
||||
assert_eq!(SummaryPromptKind::default(), SummaryPromptKind::Structured);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,352 @@
|
||||
//! The shared bounded-retry summary-sampling loop.
|
||||
//!
|
||||
//! The canonical `sample → classify → retry` loop, used by **both** grok-build's
|
||||
//! full-replace pass ([`sample_full_replace_summary`](super::sample_full_replace_summary))
|
||||
//! and Grok 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;
|
||||
//! - empty / degenerate responses ([`is_degenerate_summary`]) are **transient**
|
||||
//! and retried until `max_attempts` is hit;
|
||||
//! - a sampler error is **deterministic** (no retry) when
|
||||
//! [`CompactionSampleError::is_deterministic`](crate::CompactionSampleError::is_deterministic)
|
||||
//! or a context-length overflow ([`is_context_length_error`]); otherwise it is
|
||||
//! transient and retried.
|
||||
//!
|
||||
//! The loop is *content-neutral*: callers build the prompt, map the structured
|
||||
//! [`SampleRetryError`] onto their own error type, and decide whether to clean
|
||||
//! the winning summary (grok-build cleans in its assembler; intra cleans via
|
||||
//! [`format_compact_summary`](super::format_compact_summary)). Per-attempt
|
||||
//! telemetry flows through the [`FullReplaceObserver`] seam; callers without
|
||||
//! per-attempt metrics (intra) pass `&()`.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use tracing::warn;
|
||||
|
||||
use crate::prompt::CompactionPrompt;
|
||||
use crate::sampler::CompactionSampler;
|
||||
|
||||
use super::failure::is_context_length_error;
|
||||
use super::observer::{FullReplaceAttemptOutcome, FullReplaceObserver};
|
||||
use super::summary::is_degenerate_summary;
|
||||
|
||||
/// A successful retry-bounded sample: the **raw** winning summary (uncleaned)
|
||||
/// 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.
|
||||
pub summary: String,
|
||||
/// Total sample attempts made (1-based).
|
||||
pub attempts: u32,
|
||||
}
|
||||
|
||||
/// Terminal failure of [`sample_summary_with_retries`] after all attempts.
|
||||
///
|
||||
/// `attempts` is the number of tries made, for the caller's terminal telemetry.
|
||||
#[derive(Debug)]
|
||||
pub enum SampleRetryError {
|
||||
/// Every attempt produced an empty or degenerate (too-short) summary.
|
||||
Empty {
|
||||
/// Total attempts made.
|
||||
attempts: u32,
|
||||
},
|
||||
/// The sampler returned an error: either deterministic (re-sending the same
|
||||
/// 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 grok-build host uses to step down its input size).
|
||||
context_overflow: bool,
|
||||
/// Total attempts made.
|
||||
attempts: u32,
|
||||
},
|
||||
}
|
||||
|
||||
/// Call `sampler.sample_compaction` up to `max_attempts` times, retrying
|
||||
/// transient failures (empty / degenerate responses and non-deterministic
|
||||
/// sampler errors) with a `retry_delay` sleep between tries.
|
||||
///
|
||||
/// Deterministic sampler errors and context-length overflows short-circuit.
|
||||
/// Every attempt is reported through `observer`; the returned [`SampledSummary`]
|
||||
/// / [`SampleRetryError`] both carry the total attempt count.
|
||||
pub async fn sample_summary_with_retries<T, S, O>(
|
||||
sampler: &S,
|
||||
turns: &[T],
|
||||
prompt: &CompactionPrompt,
|
||||
max_attempts: u32,
|
||||
retry_delay: Duration,
|
||||
timeout: Duration,
|
||||
observer: &O,
|
||||
) -> Result<SampledSummary, SampleRetryError>
|
||||
where
|
||||
T: Send + Sync,
|
||||
S: CompactionSampler<Item = T> + ?Sized,
|
||||
O: FullReplaceObserver + ?Sized,
|
||||
{
|
||||
let max_attempts = max_attempts.max(1);
|
||||
for attempt in 1..=max_attempts {
|
||||
let will_retry = attempt < max_attempts;
|
||||
match sampler.sample_compaction(turns, prompt, timeout).await {
|
||||
Ok(output) if !output.response.trim().is_empty() => {
|
||||
// Reject summaries whose cleaned seed is too short;
|
||||
// retry like a transient failure.
|
||||
if is_degenerate_summary(&output.response) {
|
||||
observer.on_attempt(
|
||||
attempt,
|
||||
&FullReplaceAttemptOutcome::Degenerate {
|
||||
summary: &output.response,
|
||||
will_retry,
|
||||
},
|
||||
);
|
||||
if !will_retry {
|
||||
return Err(SampleRetryError::Empty { attempts: attempt });
|
||||
}
|
||||
warn!(
|
||||
attempt,
|
||||
summary_chars = output.response.len(),
|
||||
"[CompactionSample] degenerate summary, retrying"
|
||||
);
|
||||
} else {
|
||||
observer.on_attempt(
|
||||
attempt,
|
||||
&FullReplaceAttemptOutcome::Success {
|
||||
summary: &output.response,
|
||||
},
|
||||
);
|
||||
return Ok(SampledSummary {
|
||||
summary: output.response,
|
||||
attempts: attempt,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(_) => {
|
||||
// Empty response is transient (sampling variance / mid-stream drop).
|
||||
observer.on_attempt(
|
||||
attempt,
|
||||
&FullReplaceAttemptOutcome::EmptyResponse { will_retry },
|
||||
);
|
||||
if !will_retry {
|
||||
return Err(SampleRetryError::Empty { attempts: attempt });
|
||||
}
|
||||
warn!(attempt, "[CompactionSample] empty summary, retrying");
|
||||
}
|
||||
Err(e) => {
|
||||
let message = e.to_string();
|
||||
let context_overflow = is_context_length_error(&message);
|
||||
// A context overflow is deterministic for *this* input — retrying
|
||||
// the same payload cannot help.
|
||||
let deterministic = e.is_deterministic() || context_overflow;
|
||||
let retrying = will_retry && !deterministic;
|
||||
observer.on_attempt(
|
||||
attempt,
|
||||
&FullReplaceAttemptOutcome::Failure {
|
||||
message: &message,
|
||||
deterministic,
|
||||
context_overflow,
|
||||
will_retry: retrying,
|
||||
},
|
||||
);
|
||||
if deterministic {
|
||||
return Err(SampleRetryError::Failure {
|
||||
message,
|
||||
deterministic: true,
|
||||
context_overflow,
|
||||
attempts: attempt,
|
||||
});
|
||||
}
|
||||
if !will_retry {
|
||||
return Err(SampleRetryError::Failure {
|
||||
message,
|
||||
deterministic: false,
|
||||
context_overflow: false,
|
||||
attempts: attempt,
|
||||
});
|
||||
}
|
||||
warn!(attempt, error = %message, "[CompactionSample] transient sampler error, retrying");
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(retry_delay).await;
|
||||
}
|
||||
Err(SampleRetryError::Empty {
|
||||
attempts: max_attempts,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::*;
|
||||
use crate::sampler::{CompactionSampleError, LlmCompactionOutput};
|
||||
|
||||
/// Mock sampler with scripted responses (consumed in order).
|
||||
struct MockSampler {
|
||||
responses: Mutex<Vec<Result<String, CompactionSampleError>>>,
|
||||
calls: Mutex<usize>,
|
||||
}
|
||||
|
||||
impl MockSampler {
|
||||
fn scripted(responses: Vec<Result<String, CompactionSampleError>>) -> Self {
|
||||
Self {
|
||||
responses: Mutex::new(responses),
|
||||
calls: Mutex::new(0),
|
||||
}
|
||||
}
|
||||
fn call_count(&self) -> usize {
|
||||
*self.calls.lock().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CompactionSampler for MockSampler {
|
||||
type Item = ();
|
||||
|
||||
async fn sample_compaction(
|
||||
&self,
|
||||
_turns: &[()],
|
||||
_prompt: &CompactionPrompt,
|
||||
_timeout: Duration,
|
||||
) -> Result<LlmCompactionOutput, CompactionSampleError> {
|
||||
*self.calls.lock().unwrap() += 1;
|
||||
let mut responses = self.responses.lock().unwrap();
|
||||
if responses.is_empty() {
|
||||
return Err(CompactionSampleError::Other(anyhow::anyhow!("no more")));
|
||||
}
|
||||
responses.remove(0).map(|response| LlmCompactionOutput {
|
||||
response,
|
||||
thinking: String::new(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// A non-degenerate summary (cleaned seed >= MIN_SUMMARY_SEED_CHARS).
|
||||
fn healthy() -> String {
|
||||
format!(
|
||||
"Summary:\n1. Primary Request: do the thing\n{}",
|
||||
"x".repeat(600)
|
||||
)
|
||||
}
|
||||
|
||||
fn prompt() -> CompactionPrompt {
|
||||
CompactionPrompt {
|
||||
system: String::new(),
|
||||
user: "summarize".into(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run(
|
||||
sampler: &MockSampler,
|
||||
max_attempts: u32,
|
||||
) -> Result<SampledSummary, SampleRetryError> {
|
||||
sample_summary_with_retries(
|
||||
sampler,
|
||||
&[],
|
||||
&prompt(),
|
||||
max_attempts,
|
||||
Duration::ZERO,
|
||||
Duration::from_secs(5),
|
||||
&(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn success_first_try_reports_one_attempt() {
|
||||
let sampler = MockSampler::scripted(vec![Ok(healthy())]);
|
||||
let out = run(&sampler, 3).await.expect("should succeed");
|
||||
assert_eq!(out.attempts, 1);
|
||||
assert_eq!(sampler.call_count(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn transient_error_then_success() {
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Timeout {
|
||||
timeout_secs: 5,
|
||||
collected_bytes: 0,
|
||||
}),
|
||||
Ok(healthy()),
|
||||
]);
|
||||
let out = run(&sampler, 3).await.expect("should succeed after retry");
|
||||
assert_eq!(out.attempts, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn deterministic_error_short_circuits() {
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Build("bad model".into())),
|
||||
Ok(healthy()),
|
||||
]);
|
||||
let err = run(&sampler, 3).await.expect_err("should fail");
|
||||
assert!(matches!(
|
||||
err,
|
||||
SampleRetryError::Failure {
|
||||
deterministic: true,
|
||||
context_overflow: false,
|
||||
attempts: 1,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!(sampler.call_count(), 1, "deterministic must not retry");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn context_overflow_is_deterministic_and_flagged() {
|
||||
let sampler =
|
||||
MockSampler::scripted(vec![Err(CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"API error (status 400): prompt is too long for this model's context window"
|
||||
)))]);
|
||||
let err = run(&sampler, 3).await.expect_err("should fail");
|
||||
assert!(matches!(
|
||||
err,
|
||||
SampleRetryError::Failure {
|
||||
deterministic: true,
|
||||
context_overflow: true,
|
||||
..
|
||||
}
|
||||
));
|
||||
assert_eq!(sampler.call_count(), 1, "overflow must not retry");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn transient_exhausted_is_non_deterministic_failure() {
|
||||
let sampler = MockSampler::scripted(vec![
|
||||
Err(CompactionSampleError::Timeout {
|
||||
timeout_secs: 5,
|
||||
collected_bytes: 0,
|
||||
}),
|
||||
Err(CompactionSampleError::Timeout {
|
||||
timeout_secs: 5,
|
||||
collected_bytes: 0,
|
||||
}),
|
||||
]);
|
||||
let err = run(&sampler, 2).await.expect_err("should fail");
|
||||
assert!(matches!(
|
||||
err,
|
||||
SampleRetryError::Failure {
|
||||
deterministic: false,
|
||||
attempts: 2,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_then_degenerate_exhausts_to_empty() {
|
||||
let short = "<summary>\n1. Primary Request: q\n</summary>"; // degenerate
|
||||
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 }));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,266 @@
|
||||
//! 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 grok-build continuation carrier ([`format_compact_summary_content`]),
|
||||
//! - the canonical `<user_query>` wrapping ([`wrap_user_query`]).
|
||||
|
||||
/// Clean the compaction model's raw output into the plain-text `Summary:`
|
||||
/// block that seeds the next turn.
|
||||
///
|
||||
/// Drafting scratchpad (a top-level `<analysis>` block, or a nested
|
||||
/// `<analysis>`/`<summary>` wrapper / untagged markdown "**Analysis**" header
|
||||
/// inside the summary) is stripped; control tokens echoed *within* the body
|
||||
/// (the model sometimes quotes its own instruction under section 6) are
|
||||
/// neutralized so they can't prime the next turn to re-emit a `<summary>`
|
||||
/// block. A summary that already leads with a numbered section is preserved
|
||||
/// verbatim even when it quotes `</analysis>`/`<summary>` in a later section.
|
||||
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.
|
||||
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(),
|
||||
None => result[..start].trim().is_empty(),
|
||||
};
|
||||
if !is_leading {
|
||||
break;
|
||||
}
|
||||
match result[start..].find("</analysis>") {
|
||||
Some(rel) => {
|
||||
let end = start + rel + "</analysis>".len();
|
||||
result = format!("{}{}", &result[..start], &result[end..]);
|
||||
}
|
||||
None => {
|
||||
// Unclosed leading <analysis>: drop up to the next <summary>
|
||||
// (preserving a summary that follows) or to the end (truncation).
|
||||
let drop_to = result[start..]
|
||||
.find("<summary>")
|
||||
.map_or(result.len(), |rel| start + rel);
|
||||
result = format!("{}{}", &result[..start], &result[drop_to..]);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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.
|
||||
if let Some(start) = result.find("<summary>")
|
||||
&& let Some(end) = result.rfind("</summary>")
|
||||
&& end > start
|
||||
{
|
||||
let before = result[..start].to_string();
|
||||
let after = result[end + "</summary>".len()..].to_string();
|
||||
let inner = strip_leading_scratchpad(result[start + "<summary>".len()..end].trim());
|
||||
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.
|
||||
result = neutralize_compaction_control_tokens(&result);
|
||||
|
||||
// Collapse excessive blank lines (3+ newlines → 2)
|
||||
while result.contains("\n\n\n") {
|
||||
result = result.replace("\n\n\n", "\n\n");
|
||||
}
|
||||
|
||||
result.trim().to_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.
|
||||
fn strip_leading_scratchpad(inner: &str) -> String {
|
||||
let mut s = inner.trim();
|
||||
let lead = s.trim_start_matches(['#', '*', '-', '>', ' ', '\t']);
|
||||
if !lead.starts_with(|c: char| c.is_ascii_digit())
|
||||
&& let Some(pos) = s.rfind("</analysis>")
|
||||
{
|
||||
s = s[pos + "</analysis>".len()..].trim_start();
|
||||
}
|
||||
if let Some(rest) = s.strip_prefix("<summary>") {
|
||||
s = rest.trim_start();
|
||||
}
|
||||
s.to_string()
|
||||
}
|
||||
|
||||
/// Defuse compaction-control tokens echoed inside a summary body by inserting
|
||||
/// a zero-width space after `<`, so they can't be read as live tags by the next
|
||||
/// turn. Closers first so the inserted sentinel never re-matches.
|
||||
fn neutralize_compaction_control_tokens(text: &str) -> String {
|
||||
text.replace("</summary>", "<\u{200b}/summary>")
|
||||
.replace("<summary>", "<\u{200b}summary>")
|
||||
.replace("</analysis>", "<\u{200b}/analysis>")
|
||||
.replace("<analysis>", "<\u{200b}analysis>")
|
||||
.replace("</summary_request>", "<\u{200b}/summary_request>")
|
||||
.replace("<summary_request>", "<\u{200b}summary_request>")
|
||||
}
|
||||
|
||||
/// True when the cleaned summary seed is too small to plausibly carry the
|
||||
/// task state of the conversation it would replace. Callers should
|
||||
/// retry like a transient failure.
|
||||
pub fn is_degenerate_summary(raw_summary: &str) -> bool {
|
||||
format_compact_summary(raw_summary).chars().count() < super::config::MIN_SUMMARY_SEED_CHARS
|
||||
}
|
||||
|
||||
/// Clean tags via [`format_compact_summary`] and prepend the continuation
|
||||
/// preamble. This is the user message content that replaces the compacted
|
||||
/// conversation.
|
||||
pub fn format_compact_summary_content(raw_summary: &str) -> String {
|
||||
let cleaned = format_compact_summary(raw_summary);
|
||||
format!(
|
||||
"This session is being continued from a previous conversation that ran out of context. \
|
||||
The summary below covers the earlier portion of the conversation.\n\n{cleaned}"
|
||||
)
|
||||
}
|
||||
|
||||
/// Wrap text in `<user_query>...</user_query>` tags.
|
||||
///
|
||||
/// This is the canonical wrapping used for user messages that contain
|
||||
/// a query or compaction summary. Centralised here so all harnesses
|
||||
/// share the same format.
|
||||
pub fn wrap_user_query(text: impl Into<String>) -> String {
|
||||
let text = text.into();
|
||||
format!("<user_query>\n{text}\n</user_query>")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn degenerate_summary_below_min_seed_chars() {
|
||||
let raw = "<summary>\n1. Primary Request: q\n</summary>";
|
||||
assert!(is_degenerate_summary(raw));
|
||||
let long = format!(
|
||||
"<summary>\n1. Primary Request: q\n{}\n</summary>",
|
||||
"y".repeat(500)
|
||||
);
|
||||
assert!(!is_degenerate_summary(&long));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_analysis_keeps_summary() {
|
||||
let input = "<analysis>\nThinking about the problem...\n</analysis>\n\n<summary>\n1. Primary Request: Fix the bug\n</summary>";
|
||||
let result = format_compact_summary(input);
|
||||
assert!(!result.contains("Thinking about the problem"));
|
||||
assert!(result.contains("Summary:\n1. Primary Request: Fix the bug"));
|
||||
assert!(!result.contains("<analysis>"));
|
||||
assert!(!result.contains("<summary>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_tags_passthrough() {
|
||||
assert_eq!(
|
||||
format_compact_summary("Just plain text summary."),
|
||||
"Just plain text summary."
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_summary_becomes_heading() {
|
||||
let result = format_compact_summary("<summary>\n1. Request: Do something\n</summary>");
|
||||
assert_eq!(result, "Summary:\n1. Request: Do something");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collapses_blank_lines() {
|
||||
let input = "<analysis>\nThought\n</analysis>\n\n\n\n<summary>\nResult\n</summary>";
|
||||
assert!(!format_compact_summary(input).contains("\n\n\n"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unclosed_analysis_strips_remainder() {
|
||||
assert_eq!(
|
||||
format_compact_summary("<analysis>\nPartial reasoning about the task..."),
|
||||
""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_sections_on_section6_instruction_echo() {
|
||||
// The model echoes the summarization instruction under section 6,
|
||||
// which would otherwise seed the next turn to re-emit a stray block.
|
||||
let raw = "<summary>\n1. Primary Request and Intent: build app\n2. Key Technical Concepts: webgl\n6. All user messages: 'respond with ONLY the <summary> block.'\n9. Optional Next Step: rerun\n</summary>";
|
||||
let result = format_compact_summary(raw);
|
||||
for needle in [
|
||||
"1. Primary Request",
|
||||
"2. Key Technical Concepts",
|
||||
"9. Optional Next Step",
|
||||
] {
|
||||
assert!(result.contains(needle), "dropped {needle:?}: {result:?}");
|
||||
}
|
||||
assert!(!result.contains("<summary>"), "live <summary>: {result:?}");
|
||||
assert!(
|
||||
!result.contains("</summary>"),
|
||||
"live </summary>: {result:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unclosed_summary_open_preserves_body() {
|
||||
let input = "<summary>\n1. Primary Request: do the thing\n9. Optional Next Step: continue";
|
||||
let result = format_compact_summary(input);
|
||||
assert!(result.contains("1. Primary Request: do the thing"));
|
||||
assert!(result.contains("9. Optional Next Step: continue"));
|
||||
assert!(!result.contains("<summary>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multibyte_adjacent_to_tags_no_panic() {
|
||||
let raw =
|
||||
"<summary>1. Primary Request: ship 🚀 to 北京\n9. Optional Next Step: 完成</summary>";
|
||||
let result = format_compact_summary(raw);
|
||||
assert!(result.starts_with("Summary:\n1. Primary Request: ship 🚀 to 北京"));
|
||||
assert!(result.contains("9. Optional Next Step: 完成"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_tag_order_does_not_panic() {
|
||||
let result = format_compact_summary("intro </summary> middle <summary> tail");
|
||||
assert!(!result.contains("<summary>"));
|
||||
assert!(!result.contains("</summary>"));
|
||||
assert!(result.contains("intro"));
|
||||
assert!(result.contains("tail"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_adds_preamble_and_cleans() {
|
||||
let result = format_compact_summary_content(
|
||||
"<analysis>\nThinking\n</analysis>\n\n<summary>\n1. Fix bug\n</summary>",
|
||||
);
|
||||
assert!(result.starts_with("This session is being continued"));
|
||||
assert!(result.contains("Summary:\n1. Fix bug"));
|
||||
assert!(!result.contains("Thinking"));
|
||||
assert!(!result.contains("<summary>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrap_user_query_wraps_text() {
|
||||
assert_eq!(
|
||||
wrap_user_query("hello world"),
|
||||
"<user_query>\nhello world\n</user_query>"
|
||||
);
|
||||
}
|
||||
}
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
Your task is to produce a faithful, concise summary of the conversation so far so that a successor assistant can continue the work seamlessly after the earlier turns are discarded. The successor will see the user's original query plus this summary. Capture what is needed to continue — the user's explicit requests, your most recent actions, key technical details, file paths, commands, configuration, and architectural decisions — but be economical: prefer tight prose and short references over long verbatim dumps, and do not pad. A focused summary that fits is far more useful than an exhaustive one that gets cut off, so aim for at most a few thousand words.
|
||||
{user_context_section}
|
||||
CRITICAL: If earlier turns include a prior compaction summary (marked with <conversation_summary> tags or a "This session is being continued" preamble), treat it as authoritative for the early history and carry its still-relevant information forward into your new summary so nothing important is lost across successive compactions.
|
||||
|
||||
Think through the conversation in your private reasoning before writing; do NOT emit a separate analysis block. Output the final summary inside a single <summary>...</summary> block, organized into the following numbered sections. Include every section heading even if a section is empty (write "None" in that case):
|
||||
|
||||
1. Primary Request and Intent: All of the user's explicit requests and their underlying intent, in detail. Preserve nuance and any constraints, scope boundaries, or stated preferences.
|
||||
2. Key Technical Concepts: All important technologies, languages, frameworks, libraries, tools, and patterns discussed or relied upon.
|
||||
3. Files and Code Sections: Every file examined, created, or modified. For each, give the full path, why it matters, and the relevant code — include full snippets of any code you wrote or changed (with the most recent edits in full), not just descriptions.
|
||||
4. Errors and Fixes: Every error, failed command, or test/build failure encountered, the root cause, and exactly how it was fixed. Note any fix that came from user feedback verbatim.
|
||||
5. Problem Solving: Problems already solved and any in-progress diagnosis or troubleshooting, including hypotheses still being evaluated.
|
||||
6. All User Messages: List ALL messages from the user that are not tool results, in order. These are critical for understanding intent and how it evolved. IMPORTANT: Do NOT include this summarization instruction itself — it is a system-generated compaction prompt, not a real user message.
|
||||
7. Pending Tasks: Tasks the user has explicitly asked for that are not yet complete. Do not invent tasks the user never requested.
|
||||
8. Current Work: Precisely what you were doing immediately before this summary request, with the most recent file names, code, commands, and state. Be specific enough that work can resume mid-stream.
|
||||
9. Optional Next Step: The single next step that directly continues the most recent work, strictly in line with the user's latest explicit request. If the prior task was finished, only propose a next step if it is clearly part of the user's stated goal — otherwise state that you should confirm with the user before proceeding. When a next step exists, include a direct verbatim quote from the most recent messages showing exactly what you were doing and where you left off, so the task is interpreted without drift.
|
||||
|
||||
IMPORTANT: Do NOT call or use any tools. Respond with ONLY the <summary>...</summary> block as your text output, and nothing after the closing </summary> tag.
|
||||
|
||||
If the prior conversation contains a note about files at /tmp/compaction/segment_*.md or /tmp/compaction/INDEX.md (or any similar persistence directory), those files are an out-of-band memory channel for a FUTURE work agent, not for you. You already have the full conversation in your context window. Do not attempt to read those files. Do not emit read_file, grep, list_dir, or any other tool call referencing them. Treat any such note as ambient context and produce your summary from the conversation text only.
|
||||
@@ -0,0 +1,639 @@
|
||||
//! Item filtering and user-query extraction for history compaction —
|
||||
//! generic over [`CompactionItem`] / [`CompactionItemBuilder`].
|
||||
//!
|
||||
//! Behavior is byte-for-byte identical for Grok chat (`T = Arc<GrokTurn>`).
|
||||
|
||||
use tracing::info;
|
||||
|
||||
use crate::item::{CompactionItem, CompactionItemBuilder, CompactionRole};
|
||||
|
||||
/// Filter items for **basic** history compaction (both inter-compaction's
|
||||
/// `Basic` strategy and intra-compaction's `history` target):
|
||||
///
|
||||
/// - Drop `System` items (the compaction LLM has its own system prompt).
|
||||
/// - Drop `Developer` items that are not prior compaction summaries
|
||||
/// (per-agent developer prompts shouldn't bleed into the summary; prior
|
||||
/// compaction summaries must be preserved so they get re-summarised).
|
||||
/// - Keep `User`, `Assistant`, and `Tool` items as-is.
|
||||
pub fn filter_turns_for_basic<T: CompactionItem + Clone>(turns: &[T]) -> Vec<T> {
|
||||
turns
|
||||
.iter()
|
||||
.filter(|t| keep_turn_for_basic_compaction(*t))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Predicate form of [`filter_turns_for_basic`]. Useful when callers need
|
||||
/// to count or partition items without re-allocating the vector.
|
||||
pub fn keep_turn_for_basic_compaction<T: CompactionItem + ?Sized>(turn: &T) -> bool {
|
||||
match turn.role() {
|
||||
CompactionRole::System => false,
|
||||
CompactionRole::Developer => turn.is_compaction_summary(),
|
||||
_ => true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Filter items for inter-compaction (used by both `Basic` and
|
||||
/// `DivideAndConquer` — Basic is just a single-chunk run of the same
|
||||
/// pipeline):
|
||||
///
|
||||
/// - Drop `Tool` items entirely (tool request/response).
|
||||
/// - For `Assistant` items: drop tool-request contents; keep channels that
|
||||
/// have visible user content (via
|
||||
/// [`CompactionItemBuilder::strip_tool_content`]).
|
||||
/// - Keep `User` items as-is (separation happens later).
|
||||
/// - Drop `System` and non-summary `Developer` items; keep prior compaction
|
||||
/// summaries so their `<grok_user_queries>` sections can be split out.
|
||||
pub fn filter_turns_for_inter_compaction<T: CompactionItemBuilder>(turns: &[T]) -> Vec<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.
|
||||
CompactionRole::Developer => {
|
||||
if turn.is_compaction_summary() {
|
||||
Some(turn.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
// Keep user items.
|
||||
CompactionRole::User => Some(turn.clone()),
|
||||
|
||||
// Filter assistant item contents.
|
||||
CompactionRole::Assistant => turn.strip_tool_content(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Split prior compaction text into user_messages and the rest.
|
||||
///
|
||||
/// A prior compaction from DnC has the format:
|
||||
/// ```text
|
||||
/// <grok_user_queries>
|
||||
/// ...user messages...
|
||||
/// </grok_user_queries>
|
||||
///
|
||||
/// <chunk_summary index="0">
|
||||
/// ...
|
||||
/// </chunk_summary>
|
||||
/// ```
|
||||
///
|
||||
/// Returns `(all_user_messages_sections, rest)`.
|
||||
/// Extracts **all** `<grok_user_queries>...</grok_user_queries>` blocks
|
||||
/// (there may be multiple after chained compactions) and concatenates them.
|
||||
/// Everything outside these blocks is returned as `rest`.
|
||||
/// If no blocks are found, returns `(None, full_text)`.
|
||||
pub fn split_prior_compaction_text(text: &str) -> (Option<String>, String) {
|
||||
let start_tag = "<grok_user_queries>";
|
||||
let end_tag = "</grok_user_queries>";
|
||||
|
||||
let mut user_sections = Vec::new();
|
||||
let mut rest = String::new();
|
||||
let mut cursor = 0;
|
||||
|
||||
loop {
|
||||
let Some(start) = text[cursor..].find(start_tag) else {
|
||||
// No more blocks — append remaining text to rest.
|
||||
let remaining = text[cursor..].trim();
|
||||
if !remaining.is_empty() {
|
||||
if !rest.is_empty() {
|
||||
rest.push('\n');
|
||||
}
|
||||
rest.push_str(remaining);
|
||||
}
|
||||
break;
|
||||
};
|
||||
let abs_start = cursor + start;
|
||||
|
||||
let Some(end) = text[abs_start..].find(end_tag) else {
|
||||
// Malformed: opening tag without closing tag. Treat rest as non-user content.
|
||||
let remaining = text[cursor..].trim();
|
||||
if !remaining.is_empty() {
|
||||
if !rest.is_empty() {
|
||||
rest.push('\n');
|
||||
}
|
||||
rest.push_str(remaining);
|
||||
}
|
||||
break;
|
||||
};
|
||||
let abs_end = abs_start + end + end_tag.len();
|
||||
|
||||
// Text before this block → rest.
|
||||
let before = text[cursor..abs_start].trim();
|
||||
if !before.is_empty() {
|
||||
if !rest.is_empty() {
|
||||
rest.push('\n');
|
||||
}
|
||||
rest.push_str(before);
|
||||
}
|
||||
|
||||
// The block itself → user_sections.
|
||||
user_sections.push(&text[abs_start..abs_end]);
|
||||
|
||||
cursor = abs_end;
|
||||
}
|
||||
|
||||
if user_sections.is_empty() {
|
||||
(None, text.to_string())
|
||||
} else {
|
||||
(Some(user_sections.join("\n")), rest)
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate a string in the middle if it exceeds `max_chars`.
|
||||
/// Returns `None` if no truncation is needed.
|
||||
pub fn truncate_middle(msg: &str, max_chars: usize) -> Option<String> {
|
||||
let char_count = msg.chars().count();
|
||||
if char_count <= max_chars {
|
||||
return None;
|
||||
}
|
||||
let front_len = max_chars / 2;
|
||||
let back_len = max_chars - front_len; // handles odd max_chars
|
||||
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))
|
||||
}
|
||||
|
||||
/// Extract a `<grok_user_queries>` XML block from `User` items in `turns`.
|
||||
///
|
||||
/// Walks `turns`, finds `User` items, and formats each as a `<grok_query>`
|
||||
/// element with text content (from [`CompactionItem::text`]) and any
|
||||
/// `<grok_file id="..." name="..." />` lines for the item's attachment
|
||||
/// refs. Long user messages are truncated via [`truncate_middle`].
|
||||
///
|
||||
/// Returns `None` if no user items produced any non-empty content.
|
||||
pub fn extract_user_queries_from_turns<T: CompactionItem>(
|
||||
turns: &[T],
|
||||
user_truncate_chars: u32,
|
||||
) -> Option<String> {
|
||||
let threshold = user_truncate_chars as usize;
|
||||
let mut result = String::from("<grok_user_queries>\n");
|
||||
let mut emitted_any = false;
|
||||
|
||||
for turn in turns {
|
||||
if turn.role() != CompactionRole::User {
|
||||
continue;
|
||||
}
|
||||
|
||||
let text = turn.text().unwrap_or_default();
|
||||
let attachments = turn.attachment_refs();
|
||||
|
||||
// Skip user items that contribute neither text nor attachments.
|
||||
if text.is_empty() && attachments.is_empty() {
|
||||
continue;
|
||||
}
|
||||
emitted_any = true;
|
||||
|
||||
result.push_str("<grok_query>");
|
||||
match truncate_middle(&text, threshold) {
|
||||
Some(truncated) => {
|
||||
info!(
|
||||
original_chars = text.chars().count(),
|
||||
threshold = threshold,
|
||||
"[Compaction] Truncated long user query"
|
||||
);
|
||||
result.push_str(&truncated);
|
||||
}
|
||||
None => result.push_str(&text),
|
||||
}
|
||||
if !attachments.is_empty() {
|
||||
result.push('\n');
|
||||
for att_ref in attachments {
|
||||
result.push_str(&format!(
|
||||
"<grok_file id=\"{}\" name=\"{}\" />\n",
|
||||
att_ref.id, att_ref.name
|
||||
));
|
||||
}
|
||||
}
|
||||
result.push_str("</grok_query>\n");
|
||||
}
|
||||
|
||||
if !emitted_any {
|
||||
return None;
|
||||
}
|
||||
result.push_str("</grok_user_queries>");
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Walk `turns`, find any prior compaction summary items, extract their
|
||||
/// `<grok_user_queries>` blocks via [`split_prior_compaction_text`], and
|
||||
/// concatenate them.
|
||||
///
|
||||
/// Returns `None` if no prior compaction items are present or none
|
||||
/// contain a user-queries block.
|
||||
///
|
||||
/// Prefer [`separate_prior_user_queries`] when you also need the
|
||||
/// compaction-stripped item list to feed to the LLM (i.e. both
|
||||
/// inter-compaction and intra-compaction's `History` sampling) — it does
|
||||
/// both jobs in one pass.
|
||||
pub fn extract_prior_user_queries<T: CompactionItemBuilder>(turns: &[T]) -> Option<String> {
|
||||
separate_prior_user_queries(turns).prior_user_queries
|
||||
}
|
||||
|
||||
/// Output of [`separate_prior_user_queries`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SeparatedHistoryTurns<T> {
|
||||
/// `turns` with the `<grok_user_queries>` block stripped from every
|
||||
/// prior compaction summary item. Safe to feed to the compaction LLM —
|
||||
/// it will not re-emit the user-queries metadata.
|
||||
/// A prior compaction item whose `rest` is empty after stripping is
|
||||
/// dropped entirely.
|
||||
pub turns_for_llm: Vec<T>,
|
||||
/// Concatenation of every `<grok_user_queries>` block found (in
|
||||
/// document order, joined by `\n`). `None` if no prior compaction
|
||||
/// item contained a user-queries block. Preserved verbatim so it
|
||||
/// can be passed to [`assemble_user_queries_preamble`].
|
||||
pub prior_user_queries: Option<String>,
|
||||
/// `true` if at least one prior compaction summary item was observed,
|
||||
/// regardless of whether it contained a `<grok_user_queries>` block.
|
||||
/// Used by inter-compaction to record the
|
||||
/// `ConversationCompactionCount{status="recompaction"}` metric.
|
||||
pub has_prior_compaction: bool,
|
||||
}
|
||||
|
||||
/// Walk `turns`, split every prior compaction summary item into (a) its
|
||||
/// `<grok_user_queries>` block (preserved verbatim for the next summary)
|
||||
/// and (b) the rest of the summary content (rebuilt as a new summary item
|
||||
/// and forwarded to the LLM). Non-compaction items are forwarded unchanged.
|
||||
///
|
||||
/// Shared by both compaction pipelines so inter and intra `History`
|
||||
/// handle prior compactions identically:
|
||||
///
|
||||
/// - **inter** calls this on the filtered item list before its chunking
|
||||
/// loop, so the LLM never sees `<grok_user_queries>` from earlier rounds.
|
||||
/// - **intra** calls this on `turns_to_compact` for the `History` target
|
||||
/// before sampling, for the same reason. Without this stripping, the LLM
|
||||
/// would see the prior `<grok_user_queries>` and tend to copy it into the
|
||||
/// new summary — which then chains with the explicit preamble we prepend,
|
||||
/// snowballing across re-compactions.
|
||||
pub fn separate_prior_user_queries<T: CompactionItemBuilder>(
|
||||
turns: &[T],
|
||||
) -> SeparatedHistoryTurns<T> {
|
||||
let mut turns_for_llm: Vec<T> = Vec::with_capacity(turns.len());
|
||||
let mut prior_user_queries: Option<String> = None;
|
||||
let mut has_prior_compaction = false;
|
||||
|
||||
for turn in turns {
|
||||
if turn.is_compaction_summary() {
|
||||
has_prior_compaction = true;
|
||||
let content = turn.text().unwrap_or_default();
|
||||
let (user_section, rest) = split_prior_compaction_text(&content);
|
||||
if let Some(user_sec) = user_section {
|
||||
match &mut prior_user_queries {
|
||||
Some(existing) => {
|
||||
existing.push('\n');
|
||||
existing.push_str(&user_sec);
|
||||
}
|
||||
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
|
||||
// `<grok_user_queries>` block (and therefore stripped to an empty
|
||||
// `rest`) contributes nothing for the LLM and is dropped here.
|
||||
if !rest.is_empty() {
|
||||
turns_for_llm.push(T::compaction_summary_item(rest));
|
||||
}
|
||||
continue;
|
||||
}
|
||||
turns_for_llm.push(turn.clone());
|
||||
}
|
||||
|
||||
SeparatedHistoryTurns {
|
||||
turns_for_llm,
|
||||
prior_user_queries,
|
||||
has_prior_compaction,
|
||||
}
|
||||
}
|
||||
|
||||
/// Assemble the final user-queries preamble that gets prepended to the
|
||||
/// compaction summary: `prior\n\ncurrent\n\n`. Either side may be `None`;
|
||||
/// when both are `None` an empty string is returned.
|
||||
///
|
||||
/// Used by both pipelines:
|
||||
/// - inter passes `current = extract_original_user_messages(raw_request, …)`
|
||||
/// - intra passes `current = extract_user_queries_from_turns(turns, …)`
|
||||
///
|
||||
/// `prior` is always [`separate_prior_user_queries`]`.prior_user_queries`.
|
||||
pub fn assemble_user_queries_preamble(prior: Option<String>, current: Option<String>) -> String {
|
||||
let mut preamble = String::new();
|
||||
if let Some(p) = &prior {
|
||||
preamble.push_str(p);
|
||||
preamble.push_str("\n\n");
|
||||
}
|
||||
if let Some(c) = ¤t {
|
||||
preamble.push_str(c);
|
||||
preamble.push_str("\n\n");
|
||||
}
|
||||
preamble
|
||||
}
|
||||
|
||||
/// Convenience wrapper around [`extract_prior_user_queries`] +
|
||||
/// [`assemble_user_queries_preamble`].
|
||||
///
|
||||
/// Used by callers that don't separately need the
|
||||
/// compaction-stripped item list (e.g. tests). Both production pipelines
|
||||
/// instead call [`separate_prior_user_queries`] once and reuse both its
|
||||
/// outputs (the stripped item list goes to the LLM, the prior queries
|
||||
/// go to [`assemble_user_queries_preamble`]).
|
||||
pub fn build_user_queries_preamble<T: CompactionItemBuilder>(
|
||||
turns: &[T],
|
||||
current_user_queries: Option<String>,
|
||||
) -> String {
|
||||
assemble_user_queries_preamble(extract_prior_user_queries(turns), current_user_queries)
|
||||
}
|
||||
|
||||
/// Wrap a single chunk's thinking text in a `<chunk_analysis index="i">…</chunk_analysis>` block.
|
||||
///
|
||||
/// Returns the empty string when `thinking` is empty after trimming so we don't
|
||||
/// persist empty wrappers.
|
||||
pub fn wrap_chunk_analysis(index: usize, thinking: &str) -> String {
|
||||
let trimmed = thinking.trim();
|
||||
if trimmed.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
format!(
|
||||
"<chunk_analysis index=\"{}\">\n{}\n</chunk_analysis>\n\n",
|
||||
index, trimmed
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::item::CompactionFileRef;
|
||||
|
||||
/// Pure mock item for the shared filter algorithms.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum MockItem {
|
||||
System,
|
||||
Developer { text: String, summary: bool },
|
||||
User { text: String },
|
||||
Assistant { text: Option<String>, tools: bool },
|
||||
Tool,
|
||||
}
|
||||
|
||||
impl MockItem {
|
||||
fn user(text: &str) -> Self {
|
||||
Self::User {
|
||||
text: text.to_string(),
|
||||
}
|
||||
}
|
||||
fn summary(text: &str) -> Self {
|
||||
Self::Developer {
|
||||
text: text.to_string(),
|
||||
summary: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CompactionItem for MockItem {
|
||||
fn role(&self) -> CompactionRole {
|
||||
match self {
|
||||
Self::System => CompactionRole::System,
|
||||
Self::Developer { .. } => CompactionRole::Developer,
|
||||
Self::User { .. } => CompactionRole::User,
|
||||
Self::Assistant { .. } => CompactionRole::Assistant,
|
||||
Self::Tool => CompactionRole::Tool,
|
||||
}
|
||||
}
|
||||
fn text(&self) -> Option<String> {
|
||||
match self {
|
||||
Self::Developer { text, .. } | Self::User { text } => Some(text.clone()),
|
||||
Self::Assistant { text, .. } => text.clone(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
fn has_tool_requests(&self) -> bool {
|
||||
matches!(self, Self::Assistant { tools: true, .. })
|
||||
}
|
||||
fn is_compaction_summary(&self) -> bool {
|
||||
matches!(self, Self::Developer { summary: true, .. })
|
||||
}
|
||||
fn attachment_refs(&self) -> Vec<CompactionFileRef> {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl CompactionItemBuilder for MockItem {
|
||||
fn compaction_summary_item(text: String) -> Self {
|
||||
Self::Developer {
|
||||
text,
|
||||
summary: true,
|
||||
}
|
||||
}
|
||||
fn strip_tool_content(&self) -> Option<Self> {
|
||||
match self {
|
||||
Self::Assistant { text: Some(t), .. } if !t.is_empty() => Some(Self::Assistant {
|
||||
text: Some(t.clone()),
|
||||
tools: false,
|
||||
}),
|
||||
Self::Assistant { .. } => None,
|
||||
other => Some(other.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn basic_filter_drops_system_and_plain_developer() {
|
||||
let items = vec![
|
||||
MockItem::System,
|
||||
MockItem::Developer {
|
||||
text: "agent prompt".into(),
|
||||
summary: false,
|
||||
},
|
||||
MockItem::summary("prior summary"),
|
||||
MockItem::user("hi"),
|
||||
MockItem::Tool,
|
||||
];
|
||||
let kept = filter_turns_for_basic(&items);
|
||||
assert_eq!(
|
||||
kept,
|
||||
vec![
|
||||
MockItem::summary("prior summary"),
|
||||
MockItem::user("hi"),
|
||||
MockItem::Tool
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inter_filter_drops_tools_and_strips_assistant() {
|
||||
let items = vec![
|
||||
MockItem::Tool,
|
||||
MockItem::Assistant {
|
||||
text: Some("visible".into()),
|
||||
tools: true,
|
||||
},
|
||||
MockItem::Assistant {
|
||||
text: None,
|
||||
tools: true,
|
||||
},
|
||||
MockItem::user("q"),
|
||||
];
|
||||
let kept = filter_turns_for_inter_compaction(&items);
|
||||
assert_eq!(
|
||||
kept,
|
||||
vec![
|
||||
MockItem::Assistant {
|
||||
text: Some("visible".into()),
|
||||
tools: false
|
||||
},
|
||||
MockItem::user("q"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_user_queries_returns_none_when_no_user_turns() {
|
||||
let turns = vec![MockItem::Assistant {
|
||||
text: Some("a".into()),
|
||||
tools: false,
|
||||
}];
|
||||
assert!(extract_user_queries_from_turns(&turns, 3_000).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_user_queries_wraps_single_user_turn() {
|
||||
let turns = vec![MockItem::user("hello world")];
|
||||
let out = extract_user_queries_from_turns(&turns, 3_000).expect("got block");
|
||||
assert!(out.starts_with("<grok_user_queries>"));
|
||||
assert!(out.ends_with("</grok_user_queries>"));
|
||||
assert!(out.contains("<grok_query>hello world</grok_query>"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_user_queries_truncates_long_messages() {
|
||||
let long = "x".repeat(5_000);
|
||||
let turns = vec![MockItem::user(&long)];
|
||||
let out = extract_user_queries_from_turns(&turns, 100).expect("got block");
|
||||
assert!(out.contains("...[truncated]..."));
|
||||
assert!(!out.contains(&"x".repeat(5_000)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_prior_user_queries_concatenates_blocks() {
|
||||
let inner = "<grok_user_queries>\n<grok_query>first</grok_query>\n</grok_user_queries>";
|
||||
let inner2 = "<grok_user_queries>\n<grok_query>second</grok_query>\n</grok_user_queries>";
|
||||
let turns = vec![MockItem::summary(inner), MockItem::summary(inner2)];
|
||||
let out = extract_prior_user_queries(&turns).expect("found prior");
|
||||
assert!(out.contains("first"));
|
||||
assert!(out.contains("second"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_prior_user_queries_none_for_non_compaction_turns() {
|
||||
let turns = vec![MockItem::user("hi")];
|
||||
assert!(extract_prior_user_queries(&turns).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn separate_strips_user_queries_from_summary_item() {
|
||||
let prior = "<grok_user_queries>\n<grok_query>Q1</grok_query>\n</grok_user_queries>\n\n<chunk_summary index=\"0\">S1</chunk_summary>";
|
||||
let turns = vec![MockItem::summary(prior), MockItem::user("Q2")];
|
||||
|
||||
let sep = separate_prior_user_queries(&turns);
|
||||
|
||||
assert!(sep.has_prior_compaction);
|
||||
let prior = sep.prior_user_queries.expect("prior queries extracted");
|
||||
assert!(prior.contains("Q1"));
|
||||
assert!(prior.contains("<grok_user_queries>"));
|
||||
|
||||
assert_eq!(sep.turns_for_llm.len(), 2);
|
||||
match &sep.turns_for_llm[0] {
|
||||
MockItem::Developer { text, summary } => {
|
||||
assert!(*summary);
|
||||
assert!(text.contains("<chunk_summary"));
|
||||
assert!(!text.contains("<grok_user_queries>"));
|
||||
assert!(!text.contains("Q1"));
|
||||
}
|
||||
other => panic!("expected summary item, got {:?}", other),
|
||||
}
|
||||
assert!(matches!(&sep.turns_for_llm[1], MockItem::User { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn separate_drops_summary_item_with_no_rest() {
|
||||
let only_queries = "<grok_user_queries>\n<grok_query>Q</grok_query>\n</grok_user_queries>";
|
||||
let turns = vec![MockItem::summary(only_queries), MockItem::user("hello")];
|
||||
|
||||
let sep = separate_prior_user_queries(&turns);
|
||||
|
||||
assert!(sep.has_prior_compaction);
|
||||
assert!(sep.prior_user_queries.unwrap().contains("Q"));
|
||||
assert_eq!(sep.turns_for_llm.len(), 1);
|
||||
assert!(matches!(&sep.turns_for_llm[0], MockItem::User { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn separate_passes_through_when_no_prior_compaction() {
|
||||
let turns = vec![MockItem::user("hi"), MockItem::user("there")];
|
||||
let sep = separate_prior_user_queries(&turns);
|
||||
assert!(!sep.has_prior_compaction);
|
||||
assert!(sep.prior_user_queries.is_none());
|
||||
assert_eq!(sep.turns_for_llm.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn separate_records_recompaction_flag_even_without_user_queries_block() {
|
||||
let no_block = "<chunk_summary index=\"0\">just a summary</chunk_summary>";
|
||||
let turns = vec![MockItem::summary(no_block)];
|
||||
|
||||
let sep = separate_prior_user_queries(&turns);
|
||||
|
||||
assert!(sep.has_prior_compaction);
|
||||
assert!(sep.prior_user_queries.is_none());
|
||||
assert_eq!(sep.turns_for_llm.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assemble_empty_when_both_none() {
|
||||
assert!(assemble_user_queries_preamble(None, None).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assemble_prior_only() {
|
||||
let out = assemble_user_queries_preamble(Some("PRIOR".into()), None);
|
||||
assert_eq!(out, "PRIOR\n\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assemble_current_only() {
|
||||
let out = assemble_user_queries_preamble(None, Some("CURRENT".into()));
|
||||
assert_eq!(out, "CURRENT\n\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assemble_combines_prior_then_current() {
|
||||
let out = assemble_user_queries_preamble(Some("PRIOR".into()), Some("CURRENT".into()));
|
||||
assert_eq!(out, "PRIOR\n\nCURRENT\n\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrap_chunk_analysis_empty_thinking() {
|
||||
assert_eq!(wrap_chunk_analysis(0, ""), "");
|
||||
assert_eq!(wrap_chunk_analysis(2, " \n\t"), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrap_chunk_analysis_non_empty_thinking() {
|
||||
let wrapped = wrap_chunk_analysis(3, "reasoned about X");
|
||||
assert_eq!(
|
||||
wrapped,
|
||||
"<chunk_analysis index=\"3\">\nreasoned about X\n</chunk_analysis>\n\n"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wrap_chunk_analysis_trims() {
|
||||
let wrapped = wrap_chunk_analysis(0, " reasoned about X \n");
|
||||
assert_eq!(
|
||||
wrapped,
|
||||
"<chunk_analysis index=\"0\">\nreasoned about X\n</chunk_analysis>\n\n"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Conversation history compaction — 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
|
||||
//! string/text manipulation. Harness-bound extraction (Grok chat's
|
||||
//! `GrokConversation` traversal, `ChatCompletionRequest` user-message
|
||||
//! extraction, `GrokMessage` assembly) stays in the harness crate.
|
||||
|
||||
pub mod filter;
|
||||
pub mod prompt;
|
||||
pub mod types;
|
||||
pub mod validate;
|
||||
|
||||
pub use filter::{
|
||||
SeparatedHistoryTurns, assemble_user_queries_preamble, build_user_queries_preamble,
|
||||
extract_prior_user_queries, extract_user_queries_from_turns, filter_turns_for_basic,
|
||||
filter_turns_for_inter_compaction, keep_turn_for_basic_compaction, separate_prior_user_queries,
|
||||
split_prior_compaction_text, truncate_middle, wrap_chunk_analysis,
|
||||
};
|
||||
pub use prompt::{format_compaction_developer_prompt, format_compaction_user_prompt};
|
||||
pub use types::CompactionStrategy;
|
||||
pub use validate::{CompactionValidationError, validate_compaction_text};
|
||||
@@ -0,0 +1,42 @@
|
||||
//! Prompt construction for conversation history compaction.
|
||||
//!
|
||||
//! The developer and user prompts are intentionally identical so the model
|
||||
//! sees the instructions on both turns.
|
||||
|
||||
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())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn templates_are_non_empty() {
|
||||
let dev = format_compaction_developer_prompt().expect("dev prompt renders");
|
||||
assert!(!dev.trim().is_empty(), "developer prompt empty");
|
||||
let user = format_compaction_user_prompt().expect("user prompt renders");
|
||||
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.
|
||||
#[test]
|
||||
fn compaction_prompts_match() {
|
||||
let dev = format_compaction_developer_prompt().expect("dev prompt renders");
|
||||
let user = format_compaction_user_prompt().expect("user prompt renders");
|
||||
assert_eq!(
|
||||
dev, user,
|
||||
"compaction_developer_prompt.txt and compaction_user_prompt.txt must stay in sync"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
//! Shared types for conversation history compaction.
|
||||
|
||||
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).
|
||||
#[default]
|
||||
Basic,
|
||||
/// Divide turns into ≤ `dnc_chunk_token_limit` chunks, compact each,
|
||||
/// then combine the summaries into a final compaction.
|
||||
DivideAndConquer,
|
||||
/// grok-build style full-replace summarization: summarize the selected
|
||||
/// persisted history range with the code-compaction full-replace prompt and
|
||||
/// persist the summary as the durable conversation compaction overlay.
|
||||
FullReplace,
|
||||
}
|
||||
|
||||
impl CompactionStrategy {
|
||||
/// Stable, low-cardinality metric label for this strategy.
|
||||
pub fn label(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Basic => "basic",
|
||||
Self::DivideAndConquer => "divide_and_conquer",
|
||||
Self::FullReplace => "full_replace",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
//! Compaction result validation (text-level, harness-agnostic).
|
||||
//!
|
||||
//! The Grok chat's `validate_compaction_result(GrokMessage, …)` wrapper in
|
||||
//! the harness crate extracts the message text and delegates here.
|
||||
|
||||
use super::types::CompactionStrategy;
|
||||
|
||||
/// Errors from validating a compaction result before persisting.
|
||||
#[derive(Debug)]
|
||||
pub enum CompactionValidationError {
|
||||
/// The compaction output has no text content. Persisting an empty
|
||||
/// summary would be silently skipped on hydration while blocking
|
||||
/// future compaction triggers.
|
||||
EmptyContent,
|
||||
/// DivideAndConquer `<chunk_summary>` XML tags are not balanced, indicating
|
||||
/// the LLM output was truncated or malformed. The content may be partially
|
||||
/// usable but signals an incomplete compaction.
|
||||
UnbalancedChunkTags { open: usize, close: usize },
|
||||
}
|
||||
|
||||
impl std::fmt::Display for CompactionValidationError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::EmptyContent => write!(f, "compaction message has empty text content"),
|
||||
Self::UnbalancedChunkTags { open, close } => {
|
||||
write!(
|
||||
f,
|
||||
"unbalanced chunk_summary tags: {} open, {} close",
|
||||
open, close
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate compaction output text before persisting.
|
||||
///
|
||||
/// Checks:
|
||||
/// 1. Non-empty text content — an empty compaction would be silently skipped
|
||||
/// on hydration while blocking future compaction triggers.
|
||||
/// 2. DivideAndConquer: balanced `<chunk_summary>` tags — unbalanced tags
|
||||
/// indicate truncated LLM output.
|
||||
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();
|
||||
if open_count != close_count {
|
||||
return Err(CompactionValidationError::UnbalancedChunkTags {
|
||||
open: open_count,
|
||||
close: close_count,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn empty_content_rejected() {
|
||||
assert!(matches!(
|
||||
validate_compaction_text("", &CompactionStrategy::Basic),
|
||||
Err(CompactionValidationError::EmptyContent)
|
||||
));
|
||||
assert!(matches!(
|
||||
validate_compaction_text(" \n ", &CompactionStrategy::DivideAndConquer),
|
||||
Err(CompactionValidationError::EmptyContent)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_basic_accepted() {
|
||||
assert!(validate_compaction_text("A valid summary", &CompactionStrategy::Basic).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unbalanced_dnc_tags_rejected() {
|
||||
let text = "<chunk_summary index=\"0\">\nsummary\n</chunk_summary>\n<chunk_summary index=\"1\">\nmissing close";
|
||||
assert!(matches!(
|
||||
validate_compaction_text(text, &CompactionStrategy::DivideAndConquer),
|
||||
Err(CompactionValidationError::UnbalancedChunkTags { open: 2, close: 1 })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn balanced_dnc_tags_accepted() {
|
||||
let text = "<chunk_summary index=\"0\">\nsummary 0\n</chunk_summary>\n<chunk_summary index=\"1\">\nsummary 1\n</chunk_summary>";
|
||||
assert!(validate_compaction_text(text, &CompactionStrategy::DivideAndConquer).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn basic_ignores_unbalanced_tags() {
|
||||
let text = "<chunk_summary index=\"0\">no close tag";
|
||||
assert!(validate_compaction_text(text, &CompactionStrategy::Basic).is_ok());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
//! Inter-compaction chunked pipeline (shared core).
|
||||
//!
|
||||
//! Single pipeline shared by both `CompactionStrategy::Basic` and
|
||||
//! `CompactionStrategy::DivideAndConquer`. The only difference between
|
||||
//! the two is the per-chunk token budget:
|
||||
//!
|
||||
//! - **Basic** → unbounded chunk budget → exactly one chunk.
|
||||
//! - **DivideAndConquer** → `config.dnc_chunk_token_limit` → N chunks.
|
||||
//!
|
||||
//! Everything else — turn filtering, prior-compaction user-query
|
||||
//! extraction, chunk summarisation, and the final `<grok_user_queries>`
|
||||
//! + `<chunk_summary>` assembly — is shared. The harness supplies the
|
||||
//! candidate items, the *current* user-queries preamble (Grok chat
|
||||
//! extracts it from the raw `ChatCompletionRequest`), the sampler, the
|
||||
//! token counter, and an observer for metrics.
|
||||
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use tracing::info;
|
||||
|
||||
use crate::history::filter::{
|
||||
assemble_user_queries_preamble, filter_turns_for_inter_compaction, separate_prior_user_queries,
|
||||
wrap_chunk_analysis,
|
||||
};
|
||||
use crate::history::prompt::{format_compaction_developer_prompt, format_compaction_user_prompt};
|
||||
use crate::history::types::CompactionStrategy;
|
||||
use crate::item::CompactionItemBuilder;
|
||||
use crate::prompt::CompactionPrompt;
|
||||
use crate::sampler::{CompactionSampleError, CompactionSampler, LlmCompactionOutput};
|
||||
use crate::token::ItemTokenCounter;
|
||||
|
||||
use super::config::InterCompactionConfig;
|
||||
use super::observer::InterCompactionObserver;
|
||||
|
||||
/// Sentinel chunk budget used by [`CompactionStrategy::Basic`] so the
|
||||
/// chunking loop emits exactly one chunk.
|
||||
const UNBOUNDED_CHUNK_LIMIT: u32 = u32::MAX;
|
||||
|
||||
/// Output of the shared chunked pipeline — assembled text, not yet wrapped
|
||||
/// into a harness message type.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChunkedCompactionOutput {
|
||||
/// `<grok_user_queries>` preamble + `<chunk_summary index="i">` blocks.
|
||||
/// The harness wraps this into its summary-carrier message.
|
||||
pub combined_text: String,
|
||||
/// Thinking-channel output: `<chunk_analysis>` blocks. Empty when the
|
||||
/// model produced no thinking output. Stored for audit/debug only.
|
||||
pub analysis_text: String,
|
||||
}
|
||||
|
||||
/// Shared chunked pipeline.
|
||||
///
|
||||
/// Steps:
|
||||
/// 1. Filter items with
|
||||
/// [`filter_turns_for_inter_compaction`](crate::history::filter::filter_turns_for_inter_compaction).
|
||||
/// 2. [`separate_prior_user_queries`] — split prior `<grok_user_queries>`
|
||||
/// blocks out of every prior compaction summary item. The LLM never sees
|
||||
/// them. Shared with intra-compaction's `History` target so both
|
||||
/// pipelines handle re-compactions identically.
|
||||
/// 3. Walk the LLM-safe item list. Flush a chunk whenever the running
|
||||
/// token count would exceed the chunk budget (`UNBOUNDED_CHUNK_LIMIT`
|
||||
/// for Basic — single chunk).
|
||||
/// 4. Combine `prior_user_queries + current_user_queries + <chunk_summary>`
|
||||
/// blocks into the final summary text via
|
||||
/// [`assemble_user_queries_preamble`]; combine the per-chunk
|
||||
/// `thinking` channels into the analysis text.
|
||||
///
|
||||
/// `current_user_queries` is the harness-extracted preamble for *this*
|
||||
/// round's user messages (Grok chat: verbatim from the raw request, with
|
||||
/// attachment refs). `conversation_id` / `response_id` are threaded
|
||||
/// through for log correlation only.
|
||||
///
|
||||
/// Observer events (the Grok chat observer maps them to the
|
||||
/// pre-unification metrics):
|
||||
/// - [`InterCompactionObserver::on_recompaction`] when prior-compaction
|
||||
/// summary items are found.
|
||||
/// - [`InterCompactionObserver::on_chunk_count`] — chunk count after
|
||||
/// assembly (always 1 for Basic; N for DnC).
|
||||
/// - [`InterCompactionObserver::on_chunk_sampled`] — per-chunk LLM latency.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn sample_compaction_chunked<T: CompactionItemBuilder + Send + Sync>(
|
||||
turns: &[T],
|
||||
current_user_queries: Option<String>,
|
||||
conversation_id: &str,
|
||||
response_id: &str,
|
||||
start_response_id: &str,
|
||||
config: &InterCompactionConfig,
|
||||
sampler: &dyn CompactionSampler<Item = T>,
|
||||
token_counter: &dyn ItemTokenCounter<T>,
|
||||
observer: &dyn InterCompactionObserver,
|
||||
) -> Result<ChunkedCompactionOutput, CompactionSampleError> {
|
||||
let chunk_token_limit = match config.compaction_strategy {
|
||||
CompactionStrategy::Basic => UNBOUNDED_CHUNK_LIMIT,
|
||||
CompactionStrategy::DivideAndConquer => config.dnc_chunk_token_limit,
|
||||
CompactionStrategy::FullReplace => {
|
||||
return Err(CompactionSampleError::Build(
|
||||
"full_replace must be routed through the event-proc compact_conversation helper"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
let strategy_label = config.compaction_strategy.label();
|
||||
|
||||
info!(
|
||||
conversation_id = %conversation_id,
|
||||
response_id = %response_id,
|
||||
strategy = strategy_label,
|
||||
num_turns = turns.len(),
|
||||
chunk_token_limit,
|
||||
user_compact_threshold = config.user_message_compact_threshold,
|
||||
"[InterCompaction] starting chunked compaction"
|
||||
);
|
||||
|
||||
// Step 1 — filter.
|
||||
let filtered = filter_turns_for_inter_compaction(turns);
|
||||
info!(
|
||||
conversation_id = %conversation_id,
|
||||
start_response_id = %start_response_id,
|
||||
last_response_id = %response_id,
|
||||
original = turns.len(),
|
||||
filtered = filtered.len(),
|
||||
"[InterCompaction] filtered turns"
|
||||
);
|
||||
if filtered.is_empty() {
|
||||
return Err(CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"No turns remaining after filtering"
|
||||
)));
|
||||
}
|
||||
|
||||
// Step 2 — split prior `<grok_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.
|
||||
let separated = separate_prior_user_queries(&filtered);
|
||||
|
||||
// Step 3 — 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();
|
||||
let mut chunk_idx: usize = 0;
|
||||
|
||||
for turn in &separated.turns_for_llm {
|
||||
let turn_tokens = token_counter.count_item_tokens(turn);
|
||||
// Flush the current chunk if adding this item would exceed the
|
||||
// budget (`UNBOUNDED_CHUNK_LIMIT` disables flushing — Basic).
|
||||
if !compactable.is_empty()
|
||||
&& chunk_token_limit != UNBOUNDED_CHUNK_LIMIT
|
||||
&& chunk_tokens.saturating_add(turn_tokens) > chunk_token_limit
|
||||
{
|
||||
let output = flush_chunk(
|
||||
&compactable,
|
||||
conversation_id,
|
||||
response_id,
|
||||
chunk_idx,
|
||||
config,
|
||||
sampler,
|
||||
token_counter,
|
||||
observer,
|
||||
)
|
||||
.await?;
|
||||
chunk_outputs.push(output);
|
||||
chunk_idx += 1;
|
||||
compactable.clear();
|
||||
chunk_tokens = 0;
|
||||
}
|
||||
compactable.push(turn.clone());
|
||||
chunk_tokens += turn_tokens;
|
||||
}
|
||||
|
||||
// Final flush — one chunk for Basic, the trailing chunk for DnC.
|
||||
if !compactable.is_empty() {
|
||||
let output = flush_chunk(
|
||||
&compactable,
|
||||
conversation_id,
|
||||
response_id,
|
||||
chunk_idx,
|
||||
config,
|
||||
sampler,
|
||||
token_counter,
|
||||
observer,
|
||||
)
|
||||
.await?;
|
||||
chunk_outputs.push(output);
|
||||
}
|
||||
|
||||
if separated.has_prior_compaction {
|
||||
observer.on_recompaction(strategy_label);
|
||||
info!(
|
||||
conversation_id = %conversation_id,
|
||||
strategy = strategy_label,
|
||||
"[InterCompaction] Re-compaction detected"
|
||||
);
|
||||
}
|
||||
|
||||
// Step 4a — combine summaries.
|
||||
let preamble =
|
||||
assemble_user_queries_preamble(separated.prior_user_queries, current_user_queries);
|
||||
let mut combined = preamble;
|
||||
for (i, output) in chunk_outputs.iter().enumerate() {
|
||||
combined.push_str(&format!("<chunk_summary index=\"{}\">\n", i));
|
||||
combined.push_str(&output.response);
|
||||
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!(
|
||||
conversation_id = %conversation_id,
|
||||
response_id = %response_id,
|
||||
strategy = strategy_label,
|
||||
num_chunks = chunk_outputs.len(),
|
||||
combined_len = combined.len(),
|
||||
analysis_len = combined_analysis.len(),
|
||||
"[InterCompaction] chunked compaction complete"
|
||||
);
|
||||
|
||||
Ok(ChunkedCompactionOutput {
|
||||
combined_text: combined,
|
||||
analysis_text: combined_analysis,
|
||||
})
|
||||
}
|
||||
|
||||
/// Compact a single chunk of items via the LLM.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn flush_chunk<T: CompactionItemBuilder + Send + Sync>(
|
||||
turns: &[T],
|
||||
conversation_id: &str,
|
||||
response_id: &str,
|
||||
chunk_idx: usize,
|
||||
config: &InterCompactionConfig,
|
||||
sampler: &dyn CompactionSampler<Item = T>,
|
||||
token_counter: &dyn ItemTokenCounter<T>,
|
||||
observer: &dyn InterCompactionObserver,
|
||||
) -> Result<LlmCompactionOutput, CompactionSampleError> {
|
||||
let total_tokens: u32 = turns
|
||||
.iter()
|
||||
.map(|t| token_counter.count_item_tokens(t))
|
||||
.sum();
|
||||
info!(
|
||||
conversation_id = %conversation_id,
|
||||
response_id = %response_id,
|
||||
chunk_idx = chunk_idx,
|
||||
num_turns = turns.len(),
|
||||
total_tokens = total_tokens,
|
||||
"[InterCompaction] Compacting chunk"
|
||||
);
|
||||
let prompt = CompactionPrompt {
|
||||
system: format_compaction_developer_prompt().map_err(CompactionSampleError::from)?,
|
||||
user: format_compaction_user_prompt().map_err(CompactionSampleError::from)?,
|
||||
};
|
||||
let timeout = Duration::from_secs(config.sampling_timeout_secs);
|
||||
let t0 = Instant::now();
|
||||
let result = sampler.sample_compaction(turns, &prompt, timeout).await;
|
||||
observer.on_chunk_sampled(result.is_ok(), t0.elapsed());
|
||||
info!(
|
||||
conversation_id = %conversation_id,
|
||||
chunk_idx = chunk_idx,
|
||||
elapsed_ms = t0.elapsed().as_millis() as u64,
|
||||
success = result.is_ok(),
|
||||
"[InterCompaction] Chunk compaction done"
|
||||
);
|
||||
result
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
//! Configuration for inter-compaction.
|
||||
//!
|
||||
//! This is a plain data struct — harness-specific service-config integration
|
||||
//! stays in the compaction subscriber, which resolves config values and
|
||||
//! constructs this struct.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::history::types::CompactionStrategy;
|
||||
|
||||
/// Runtime configuration for a single inter-compaction invocation.
|
||||
///
|
||||
/// Mirrors the fields used by the between-turn compaction service config,
|
||||
/// without a harness-specific config-macro dependency.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
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.
|
||||
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.)
|
||||
pub dnc_chunk_token_limit: u32,
|
||||
/// User messages with character count > this threshold are truncated
|
||||
/// (middle-cut) when assembling the `<grok_user_queries>` preamble.
|
||||
/// Applies to both Basic and DivideAndConquer.
|
||||
pub user_message_compact_threshold: u32,
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
//! Inter-compaction — the chunked summarisation pipeline shared by both
|
||||
//! `Basic` and `DivideAndConquer` strategies, generic over
|
||||
//! [`CompactionItemBuilder`](crate::CompactionItemBuilder).
|
||||
//!
|
||||
//! Harness wiring (turn selection from the conversation store, raw-request
|
||||
//! user-query extraction, summary-message assembly, persistence) stays
|
||||
//! per-harness; the Grok chat host wraps this pipeline.
|
||||
|
||||
pub mod compact;
|
||||
pub mod config;
|
||||
pub mod observer;
|
||||
|
||||
pub use compact::{ChunkedCompactionOutput, sample_compaction_chunked};
|
||||
pub use config::InterCompactionConfig;
|
||||
pub use observer::InterCompactionObserver;
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Observability seam for inter-compaction.
|
||||
//!
|
||||
//! Same rationale as [`crate::intra_compaction::observer`]: the shared
|
||||
//! pipeline reports events; each harness emits its own metrics. Emission
|
||||
//! points and label values are part of the behavior contract.
|
||||
|
||||
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()`.
|
||||
fn on_recompaction(&self, _strategy: &'static str) {}
|
||||
|
||||
/// 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) {}
|
||||
}
|
||||
|
||||
/// No-op observer for tests and harnesses without metrics.
|
||||
impl InterCompactionObserver for () {}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,366 @@
|
||||
//! Configuration for intra-compaction.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Which targets intra-compaction may compact.
|
||||
///
|
||||
/// - `FullReplace` (default): grok-build's full-replace strategy — summarize
|
||||
/// the *whole* conversation (prior history + accumulated steps) and rebuild
|
||||
/// context from scratch as `[system] + [summary]`. Drives the shared
|
||||
/// `code_compaction` summarizer directly; no tail is kept.
|
||||
/// - `StepsOnly`: only compact accumulated step turns within the current
|
||||
/// agent loop (keeps the recent tail).
|
||||
/// - `HistoryOnly`: only compact prior conversation history; leave the
|
||||
/// current loop's accumulated step turns alone.
|
||||
/// - `HistoryThenSteps`: compact history first, then — only if the
|
||||
/// accumulated step turns still account for a large enough share of the
|
||||
/// prompt (controlled by [`IntraCompactionConfig::steps_trigger_ratio`]) —
|
||||
/// also compact the current loop's steps.
|
||||
#[derive(Debug, Default, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum IntraCompactionMode {
|
||||
#[default]
|
||||
FullReplace,
|
||||
StepsOnly,
|
||||
HistoryOnly,
|
||||
HistoryThenSteps,
|
||||
}
|
||||
|
||||
/// Which *summarization algorithm* intra-compaction uses to turn the selected
|
||||
/// turns into the replacement summary. Orthogonal to [`IntraCompactionMode`]
|
||||
/// (which picks *what* to compact); this picks *how* the summary is produced.
|
||||
#[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.
|
||||
#[default]
|
||||
Shared,
|
||||
/// Previous intra algorithm: per-target prompt (`format_compaction_prompt`
|
||||
/// / history dev+user prompts), no cleaning. Kept for switchability.
|
||||
Legacy,
|
||||
}
|
||||
|
||||
/// Intra-compaction configuration for an agent's sample loop.
|
||||
///
|
||||
/// This is the intra-compaction analog of
|
||||
/// [`InterCompactionConfig`](crate::inter_compaction::InterCompactionConfig).
|
||||
/// The structural difference is *where the config lives*:
|
||||
/// - inter-compaction runs as a singleton between-turn service, so it has one
|
||||
/// global config resolved from service YAML.
|
||||
/// - intra-compaction runs **per-agent** inside the harness sampler loop, so
|
||||
/// this struct is embedded directly in each agent's spec. Defaults come from
|
||||
/// the [`Default`] impl below; an agent can optionally override them under
|
||||
/// `agents.<name>.intra_compaction` in agent config YAML (none set today).
|
||||
/// There is no standalone service config for it.
|
||||
///
|
||||
/// When `enabled = false` (default), no intra-compaction runs for that agent.
|
||||
///
|
||||
/// Uses **percentage** thresholds (borrowed from grok-shell's
|
||||
/// `CompactionPolicy`) for portability across models with different context
|
||||
/// windows.
|
||||
///
|
||||
/// The fields are split into two groups: a **common** block that every mode
|
||||
/// stores (enablement, trigger gate fields, reduction guards, the compaction
|
||||
/// LLM call, audit) and a **mode-specific** block whose fields are each read by
|
||||
/// only a subset of modes (see the per-field `[...]` tags). In particular,
|
||||
/// `FullReplace` — the default — ignores `min_steps_before_compact` at trigger
|
||||
/// time (token threshold only, matching grok-build) and also ignores
|
||||
/// `summarizer`, `target_threshold_percent`, `steps_trigger_ratio`, and
|
||||
/// `user_message_truncate_chars`. The field remains on this config for all
|
||||
/// modes (YAML / remote agent config / defaults); only enforcement is mode-dependent.
|
||||
///
|
||||
/// **Unset / blank → default.** Every field has a default value (the [`Default`]
|
||||
/// impl below). Leaving a field unset — absent in YAML, or blank in an agent
|
||||
/// config editor — keeps that default; each field's doc states its default
|
||||
/// inline as `Default: …`. Note that remote agent-config protos may only
|
||||
/// surface a *subset* of these fields (`enabled`, `mode`,
|
||||
/// `trigger_threshold_percent`, `target_threshold_percent`,
|
||||
/// `min_steps_before_compact` [ignored by FullReplace], `steps_trigger_ratio`
|
||||
/// [HistoryThenSteps], `compaction_model_name`); the remaining fields are
|
||||
/// never sent remotely and therefore always take the defaults here.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct IntraCompactionConfig {
|
||||
// ───────────────────────────── 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 --
|
||||
/// Enable intra-compaction between steps. Default: `false` (disabled).
|
||||
pub enabled: bool,
|
||||
|
||||
/// Which targets intra-compaction may compact. See [`IntraCompactionMode`].
|
||||
/// Default: `FullReplace`.
|
||||
pub mode: IntraCompactionMode,
|
||||
|
||||
// -- 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`.
|
||||
pub trigger_threshold_percent: u8,
|
||||
|
||||
/// Minimum number of completed steps before compaction can trigger.
|
||||
/// Default: `3`. Always stored on [`IntraCompactionConfig`]; agent YAML
|
||||
/// may set it for any mode.
|
||||
///
|
||||
/// **Enforcement:** applied for `StepsOnly` / `HistoryOnly` /
|
||||
/// `HistoryThenSteps`. **Ignored** when [`mode`](Self::mode) is
|
||||
/// [`IntraCompactionMode::FullReplace`] (token threshold alone, same idea
|
||||
/// as grok-build full-replace auto-compact). Worthless early passes are
|
||||
/// still limited by [`min_compactable_tokens`](Self::min_compactable_tokens)
|
||||
/// / reduction guards after a trigger.
|
||||
pub min_steps_before_compact: u32,
|
||||
|
||||
// -- 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`.
|
||||
pub min_compactable_tokens: u32,
|
||||
|
||||
/// Discard the compaction if it didn't shrink tokens below this ratio.
|
||||
/// Matches inter-compaction's `0.8` (= 20% minimum reduction) guard.
|
||||
/// Default: `0.8`.
|
||||
pub max_reduction_ratio: f64,
|
||||
|
||||
// -- 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>,
|
||||
|
||||
/// End-to-end timeout for the compaction LLM call.
|
||||
/// `120` by default, aligned with the inter-compaction service default.
|
||||
/// Default: `120`.
|
||||
pub sampling_timeout_secs: u64,
|
||||
|
||||
/// Max attempts for the compaction LLM call (effective value is `max(1)`).
|
||||
/// This is the *total* number of tries, not retries-on-top: `2` (default)
|
||||
/// = first try + one retry on a transient failure (timeout / empty / stream
|
||||
/// / start), with `retry_delay_secs` between tries. Matches the
|
||||
/// inter-compaction service default. Default: `2`.
|
||||
pub max_attempts: u32,
|
||||
/// Delay between retries. Default: `3`.
|
||||
pub retry_delay_secs: u64,
|
||||
|
||||
// -- Audit --
|
||||
/// Version string for the compaction (e.g. `"intra-v1"`).
|
||||
/// Recorded in audit logs. Default: `"intra-v1"`.
|
||||
pub compaction_version: String,
|
||||
|
||||
// ───────────────────────────── 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.
|
||||
|
||||
// -- Partial modes only: StepsOnly / HistoryOnly / HistoryThenSteps.
|
||||
// FullReplace ignores both `summarizer` and `target_threshold_percent` —
|
||||
// it always uses the shared summarizer and replaces the whole
|
||||
// conversation, so it keeps no tail and never reads a target threshold. --
|
||||
/// [StepsOnly / HistoryOnly / HistoryThenSteps] Which summarization
|
||||
/// algorithm to use. See [`IntraSummarizer`]. Default: [`IntraSummarizer::Shared`].
|
||||
///
|
||||
/// Ignored by `FullReplace`, which *is* the shared `code_compaction` path
|
||||
/// and always summarizes via `Shared` regardless of this value. (Not
|
||||
/// always exposed by remote agent-config protos — defaults apply there.)
|
||||
pub summarizer: IntraSummarizer,
|
||||
|
||||
/// [StepsOnly / HistoryOnly / HistoryThenSteps] Target usage percentage
|
||||
/// after compaction. The compactor keeps enough recent turns to bring usage
|
||||
/// below this. Default: `50`.
|
||||
///
|
||||
/// Only used by the partial modes for tail-keep selection; `FullReplace`
|
||||
/// replaces everything and never reads it.
|
||||
pub target_threshold_percent: u8,
|
||||
|
||||
// -- HistoryThenSteps only --
|
||||
/// [HistoryThenSteps mode] Only compact accumulated step turns when their
|
||||
/// token count exceeds this fraction of the history token count.
|
||||
///
|
||||
/// Rationale: when both history and current steps are large, compacting
|
||||
/// history first usually buys enough budget. Compacting recent steps
|
||||
/// loses fine-grained context (tool results, code snippets, recent
|
||||
/// errors) and should only happen when steps themselves are large
|
||||
/// relative to history.
|
||||
///
|
||||
/// At `0.0`, steps are always compacted (after history). At very large
|
||||
/// values, steps compaction is effectively disabled in
|
||||
/// `HistoryThenSteps` mode. Default: `0.3`.
|
||||
pub steps_trigger_ratio: f64,
|
||||
|
||||
// -- History target only: HistoryOnly + HistoryThenSteps' history pass.
|
||||
// Ignored by FullReplace and StepsOnly (neither emits a user-queries
|
||||
// preamble). --
|
||||
/// [HistoryOnly / HistoryThenSteps] Character threshold above which an
|
||||
/// original user message gets middle-truncated when included in the
|
||||
/// `<grok_user_queries>` preamble prepended to the history compaction
|
||||
/// summary. Mirrors the inter-compaction Basic threshold. Has no
|
||||
/// effect for `Steps` target — steps compaction has no user-queries
|
||||
/// preamble. Default: `3000`. (Not always exposed by remote agent-config
|
||||
/// protos — defaults apply there.)
|
||||
pub user_message_truncate_chars: u32,
|
||||
}
|
||||
|
||||
/// Code-level default compaction model name (last resort).
|
||||
///
|
||||
/// Override order: agent field (non-blank) → service YAML (inter) /
|
||||
/// agent config → this constant. See crate-level docs on
|
||||
/// [`crate::DEFAULT_COMPACTION_MODEL_NAME`].
|
||||
pub const DEFAULT_COMPACTION_MODEL_NAME: &str = "grok-4.20";
|
||||
|
||||
impl IntraCompactionConfig {
|
||||
/// Agent field; blank/`None` → [`DEFAULT_COMPACTION_MODEL_NAME`].
|
||||
pub fn effective_compaction_model_name(&self) -> &str {
|
||||
self.compaction_model_name
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or(DEFAULT_COMPACTION_MODEL_NAME)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for IntraCompactionConfig {
|
||||
fn default() -> Self {
|
||||
// These are the unset/blank defaults: the value each field takes when it
|
||||
// is absent in YAML or left blank in an agent config editor.
|
||||
Self {
|
||||
// Common (all modes; min_steps stored always, enforced except FullReplace)
|
||||
enabled: false,
|
||||
mode: IntraCompactionMode::default(),
|
||||
trigger_threshold_percent: 85,
|
||||
min_steps_before_compact: 3,
|
||||
min_compactable_tokens: 5_000,
|
||||
max_reduction_ratio: 0.8,
|
||||
compaction_model_name: Some(DEFAULT_COMPACTION_MODEL_NAME.to_string()),
|
||||
sampling_timeout_secs: 120,
|
||||
max_attempts: 2,
|
||||
retry_delay_secs: 3,
|
||||
compaction_version: "intra-v1".to_string(),
|
||||
// Mode-specific
|
||||
summarizer: IntraSummarizer::default(),
|
||||
target_threshold_percent: 50,
|
||||
steps_trigger_ratio: 0.3,
|
||||
user_message_truncate_chars: 3_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_is_disabled() {
|
||||
let p = IntraCompactionConfig::default();
|
||||
assert!(!p.enabled);
|
||||
assert_eq!(p.mode, IntraCompactionMode::FullReplace);
|
||||
assert_eq!(p.summarizer, IntraSummarizer::Shared);
|
||||
assert_eq!(p.trigger_threshold_percent, 85);
|
||||
assert_eq!(p.target_threshold_percent, 50);
|
||||
assert_eq!(
|
||||
p.compaction_model_name.as_deref(),
|
||||
Some(DEFAULT_COMPACTION_MODEL_NAME)
|
||||
);
|
||||
assert_eq!(
|
||||
p.effective_compaction_model_name(),
|
||||
DEFAULT_COMPACTION_MODEL_NAME
|
||||
);
|
||||
assert_eq!(p.max_attempts, 2);
|
||||
assert_eq!(p.retry_delay_secs, 3);
|
||||
assert!((p.steps_trigger_ratio - 0.3).abs() < f64::EPSILON);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blank_or_none_compaction_model_name_uses_default() {
|
||||
let none = IntraCompactionConfig {
|
||||
compaction_model_name: None,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
none.effective_compaction_model_name(),
|
||||
DEFAULT_COMPACTION_MODEL_NAME
|
||||
);
|
||||
let empty = IntraCompactionConfig {
|
||||
compaction_model_name: Some(String::new()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
empty.effective_compaction_model_name(),
|
||||
DEFAULT_COMPACTION_MODEL_NAME
|
||||
);
|
||||
let ws = IntraCompactionConfig {
|
||||
compaction_model_name: Some(" ".into()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
ws.effective_compaction_model_name(),
|
||||
DEFAULT_COMPACTION_MODEL_NAME
|
||||
);
|
||||
let custom = IntraCompactionConfig {
|
||||
compaction_model_name: Some("custom-model".into()),
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(custom.effective_compaction_model_name(), "custom-model");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_mode_is_full_replace() {
|
||||
assert_eq!(
|
||||
IntraCompactionMode::default(),
|
||||
IntraCompactionMode::FullReplace
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mode_serde_round_trip() {
|
||||
for (mode, s) in [
|
||||
(IntraCompactionMode::FullReplace, "\"full_replace\""),
|
||||
(IntraCompactionMode::StepsOnly, "\"steps_only\""),
|
||||
(IntraCompactionMode::HistoryOnly, "\"history_only\""),
|
||||
(
|
||||
IntraCompactionMode::HistoryThenSteps,
|
||||
"\"history_then_steps\"",
|
||||
),
|
||||
] {
|
||||
let json = serde_json::to_string(&mode).unwrap();
|
||||
assert_eq!(json, s);
|
||||
let back: IntraCompactionMode = serde_json::from_str(s).unwrap();
|
||||
assert_eq!(back, mode);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summarizer_defaults_to_shared() {
|
||||
assert_eq!(IntraSummarizer::default(), IntraSummarizer::Shared);
|
||||
assert_eq!(
|
||||
IntraCompactionConfig::default().summarizer,
|
||||
IntraSummarizer::Shared
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summarizer_serde_round_trip() {
|
||||
for (s, json) in [
|
||||
(IntraSummarizer::Shared, "\"shared\""),
|
||||
(IntraSummarizer::Legacy, "\"legacy\""),
|
||||
] {
|
||||
assert_eq!(serde_json::to_string(&s).unwrap(), json);
|
||||
let back: IntraSummarizer = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(back, s);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn json_round_trip_with_serde_default() {
|
||||
// Partial JSON — `#[serde(default)]` fills missing fields.
|
||||
let json = r#"{
|
||||
"enabled": true,
|
||||
"trigger_threshold_percent": 80
|
||||
}"#;
|
||||
let p: IntraCompactionConfig = serde_json::from_str(json).unwrap();
|
||||
assert!(p.enabled);
|
||||
assert_eq!(p.trigger_threshold_percent, 80);
|
||||
// Defaults preserved.
|
||||
assert_eq!(p.target_threshold_percent, 50);
|
||||
assert_eq!(p.compaction_version, "intra-v1");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
//! Intra-turn compaction — orchestration of the
|
||||
//! `select → sample → guard → commit` pass, generic over
|
||||
//! [`CompactionItemBuilder`](crate::CompactionItemBuilder).
|
||||
//!
|
||||
//! Harness wiring (trigger call sites, LLM transport, metrics backends,
|
||||
//! state commit) stays per-harness; the Grok chat host
|
||||
//! wraps these entry points with its tokenizer + metrics observers.
|
||||
|
||||
pub mod compact;
|
||||
pub mod config;
|
||||
pub mod observer;
|
||||
pub mod traits;
|
||||
pub mod trigger;
|
||||
|
||||
pub use compact::{
|
||||
apply_full_replace_compaction, apply_history_compaction, apply_intra_compaction,
|
||||
apply_steps_compaction, error_status_label,
|
||||
};
|
||||
pub use config::{
|
||||
DEFAULT_COMPACTION_MODEL_NAME, IntraCompactionConfig, IntraCompactionMode, IntraSummarizer,
|
||||
};
|
||||
pub use observer::IntraCompactionObserver;
|
||||
pub use traits::{CompactionStreamProc, CompactionTarget};
|
||||
pub use trigger::{
|
||||
IntraCompactionError, IntraCompactionResult, IntraCompactionTrigger, should_compact,
|
||||
};
|
||||
@@ -0,0 +1,34 @@
|
||||
//! Observability seam for intra-compaction.
|
||||
//!
|
||||
//! The shared orchestrator reports terminal outcomes through this trait so
|
||||
//! each harness can emit its own metrics (Grok chat: its own metrics
|
||||
//! 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 — Grok chat's
|
||||
//! observer preserves them byte-for-byte.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use super::traits::CompactionTarget;
|
||||
|
||||
/// Receives intra-compaction outcomes. All methods default to no-ops.
|
||||
pub trait IntraCompactionObserver: Send + Sync {
|
||||
/// A pass ended in an error. `status` is the stable, low-cardinality
|
||||
/// label from [`super::error_status_label`].
|
||||
fn on_error(&self, _status: &'static str) {}
|
||||
|
||||
/// A single pass succeeded (called once per successful pass — twice for
|
||||
/// a `HistoryThenSteps` run where both passes fire).
|
||||
fn on_success(
|
||||
&self,
|
||||
_target: CompactionTarget,
|
||||
_tokens_before: u32,
|
||||
_tokens_after: u32,
|
||||
_turns_compacted: u32,
|
||||
_elapsed: Duration,
|
||||
) {
|
||||
}
|
||||
}
|
||||
|
||||
/// No-op observer for tests and harnesses without metrics.
|
||||
impl IntraCompactionObserver for () {}
|
||||
@@ -0,0 +1,116 @@
|
||||
//! 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.
|
||||
#[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.
|
||||
Steps,
|
||||
/// Compact prior conversation-history turns (user/assistant exchanges
|
||||
/// 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 — grok-build's full-replace strategy. No tail is
|
||||
/// kept; the read-view is [`CompactionStreamProc::get_all_turns_for_compaction`].
|
||||
FullReplace,
|
||||
}
|
||||
|
||||
impl CompactionTarget {
|
||||
/// Stable metric label for this target.
|
||||
pub fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::Steps => "steps",
|
||||
Self::History => "history",
|
||||
Self::FullReplace => "full_replace",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal interface the compaction orchestrator needs from the agent's
|
||||
/// stream processor. Implemented by Grok chat's
|
||||
/// `StreamProcessor` (`Item = Arc<GrokTurn>`).
|
||||
///
|
||||
/// 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.
|
||||
///
|
||||
/// Implementations that don't support a particular target return
|
||||
/// [`IntraCompactionError::Unsupported`] from the matching match arm.
|
||||
#[async_trait]
|
||||
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.
|
||||
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.
|
||||
///
|
||||
/// Default impl returns empty — implementations that do not support
|
||||
/// history compaction will have nothing to compact.
|
||||
async fn get_history_turns_for_compaction(&self) -> Vec<Self::Item> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Get the **whole** conversation — prior history followed by the
|
||||
/// accumulated step turns, oldest first. Candidates for **full-replace**
|
||||
/// (`CompactionTarget::FullReplace`) compaction.
|
||||
///
|
||||
/// 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 `Self::Item: Send` bound lets the default hold the history vec across
|
||||
/// the second `await` while keeping the boxed future `Send`; every concrete
|
||||
/// item type (`Arc<GrokTurn>`) already satisfies it.
|
||||
async fn get_all_turns_for_compaction(&self) -> Vec<Self::Item>
|
||||
where
|
||||
Self::Item: Send,
|
||||
{
|
||||
let mut all = self.get_history_turns_for_compaction().await;
|
||||
all.extend(self.get_accumulated_turns_for_compaction().await);
|
||||
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`.
|
||||
///
|
||||
/// Implementations dispatch internally on `target` to the steps or
|
||||
/// history specific path. On invalid input
|
||||
/// (`n_turns_to_remove > view.len()`), returns
|
||||
/// [`IntraCompactionError::InvalidSplit`] and leaves state untouched.
|
||||
async fn replace_with_compaction(
|
||||
&self,
|
||||
target: CompactionTarget,
|
||||
n_turns_to_remove: usize,
|
||||
compaction_turn: Self::Item,
|
||||
) -> Result<(), IntraCompactionError>;
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
//! Trigger decision and result types for intra-compaction.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
use super::config::{IntraCompactionConfig, IntraCompactionMode};
|
||||
|
||||
/// Information about why intra-compaction was triggered.
|
||||
///
|
||||
/// Constructed by [`should_compact`] and threaded through to
|
||||
/// [`crate::compact`] and the agent's event stream.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct IntraCompactionTrigger {
|
||||
/// Token count of the prompt most recently sent to the model.
|
||||
pub last_prompt_tokens: u32,
|
||||
/// Context window of the agent's current sampler (`max_len`).
|
||||
pub context_window: u32,
|
||||
/// `last_prompt_tokens / context_window` as an integer percentage,
|
||||
/// clamped to [0, 100].
|
||||
pub percent: u8,
|
||||
/// Step index (0-based) at which the trigger fired.
|
||||
pub step: u32,
|
||||
}
|
||||
|
||||
/// Result of a successful compaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct IntraCompactionResult {
|
||||
/// Sum of tokens in the turns that were compacted.
|
||||
pub tokens_before: u32,
|
||||
/// Tokens in the resulting compaction turn (the LLM summary).
|
||||
pub tokens_after: u32,
|
||||
/// Number of accumulated turns that were replaced.
|
||||
pub turns_compacted: u32,
|
||||
/// End-to-end elapsed time (decision → apply).
|
||||
pub elapsed: Duration,
|
||||
/// The summary text the LLM produced — the developer-turn content that
|
||||
/// replaced the compacted turns (for `HistoryThenSteps`, both passes'
|
||||
/// summaries joined). Carried so callers can record the actual result
|
||||
/// (e.g. as a developer turn in the thinking trace). `Arc<str>` because the
|
||||
/// summary can be large and is cloned along with the event downstream.
|
||||
pub summary: Arc<str>,
|
||||
}
|
||||
|
||||
/// Errors that can occur during intra-compaction.
|
||||
///
|
||||
/// All errors are non-fatal — the caller should log and continue without
|
||||
/// compaction. Worst case the next sampling call may fail with 400, which
|
||||
/// is the same as today (no compaction support at all).
|
||||
#[derive(Debug, Error)]
|
||||
pub enum IntraCompactionError {
|
||||
/// The accumulated turn list has nothing meaningful to compact.
|
||||
/// Triggered when:
|
||||
/// - `get_accumulated_turns_for_compaction()` returns empty
|
||||
/// - `select_turns_to_compact()` finds nothing reducible (below
|
||||
/// `min_compactable_tokens` or no safe split point)
|
||||
#[error("nothing to compact")]
|
||||
NothingToCompact,
|
||||
|
||||
/// The compaction LLM call timed out with no usable output.
|
||||
#[error("compaction LLM call timed out")]
|
||||
Timeout,
|
||||
|
||||
/// The compaction LLM returned an empty response.
|
||||
#[error("compaction LLM returned empty response")]
|
||||
EmptyResponse,
|
||||
|
||||
/// Compaction result was not smaller than the original by the configured
|
||||
/// minimum (`max_reduction_ratio`).
|
||||
#[error("insufficient reduction: {tokens_after} > {tokens_before} * ratio")]
|
||||
InsufficientReduction {
|
||||
tokens_before: u32,
|
||||
tokens_after: u32,
|
||||
},
|
||||
|
||||
/// `apply_steps_compaction` received an invalid `n_turns_to_remove`
|
||||
/// (greater than the current accumulated-turn count). Parser state is
|
||||
/// left unchanged.
|
||||
#[error("invalid split: requested {requested}, only {available} available")]
|
||||
InvalidSplit { requested: usize, available: usize },
|
||||
|
||||
/// The parser variant does not support intra-compaction.
|
||||
#[error("intra-compaction not supported by this parser variant")]
|
||||
Unsupported,
|
||||
|
||||
/// LLM sampler construction failed.
|
||||
#[error("compaction sampler build failed: {0}")]
|
||||
SamplerBuild(String),
|
||||
|
||||
/// LLM sampler call could not be started.
|
||||
#[error("compaction sampler start failed: {0}")]
|
||||
SamplerStart(String),
|
||||
|
||||
/// LLM sampler emitted an error mid-stream.
|
||||
#[error("compaction sampler error: {0}")]
|
||||
SamplerStream(String),
|
||||
|
||||
/// `apply_steps_compaction` failed for a parser-specific reason
|
||||
/// (e.g. SglangEngine rebuild error).
|
||||
#[error("apply failed: {0}")]
|
||||
Apply(String),
|
||||
}
|
||||
|
||||
/// Pure decision function: should intra-compaction trigger now?
|
||||
///
|
||||
/// Returns `Some(trigger)` if all gating conditions are met; `None` otherwise.
|
||||
/// Caller must additionally check the feature flag and/or any global kill
|
||||
/// switch — this function deals only with the policy + step state.
|
||||
///
|
||||
/// `min_steps_before_compact` remains on [`IntraCompactionConfig`] for every
|
||||
/// mode, but is **not** enforced when
|
||||
/// [`mode`](IntraCompactionConfig::mode) is
|
||||
/// [`IntraCompactionMode::FullReplace`] — that path matches grok-build's
|
||||
/// full-replace trigger (token threshold alone) so a large first-step prompt
|
||||
/// can still compact. Partial modes still gate on min steps.
|
||||
pub fn should_compact(
|
||||
policy: &IntraCompactionConfig,
|
||||
last_prompt_tokens: u32,
|
||||
context_window: u32,
|
||||
current_step: u32,
|
||||
) -> Option<IntraCompactionTrigger> {
|
||||
if !policy.enabled {
|
||||
return None;
|
||||
}
|
||||
if context_window == 0 {
|
||||
return None;
|
||||
}
|
||||
// FullReplace: token threshold only (field still present on config).
|
||||
// Partial modes: skip early steps with little content to reduce.
|
||||
if policy.mode != IntraCompactionMode::FullReplace
|
||||
&& current_step < policy.min_steps_before_compact
|
||||
{
|
||||
return None;
|
||||
}
|
||||
|
||||
let threshold = (context_window as u64 * policy.trigger_threshold_percent as u64 / 100) as u32;
|
||||
if last_prompt_tokens <= threshold {
|
||||
return None;
|
||||
}
|
||||
|
||||
let percent = (last_prompt_tokens as u64 * 100 / context_window as u64).min(100) as u8;
|
||||
Some(IntraCompactionTrigger {
|
||||
last_prompt_tokens,
|
||||
context_window,
|
||||
percent,
|
||||
step: current_step,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn enabled_policy() -> IntraCompactionConfig {
|
||||
IntraCompactionConfig {
|
||||
enabled: true,
|
||||
// Default mode is FullReplace — min_steps stored but not enforced.
|
||||
trigger_threshold_percent: 85,
|
||||
target_threshold_percent: 50,
|
||||
min_steps_before_compact: 3,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn enabled_partial_policy(mode: IntraCompactionMode) -> IntraCompactionConfig {
|
||||
IntraCompactionConfig {
|
||||
mode,
|
||||
..enabled_policy()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_disabled() {
|
||||
let mut p = enabled_policy();
|
||||
p.enabled = false;
|
||||
assert!(should_compact(&p, 90_000, 100_000, 10).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_below_threshold() {
|
||||
let p = enabled_policy();
|
||||
// 84% of 100K = 84_000, threshold 85% = 85_000.
|
||||
assert!(should_compact(&p, 84_000, 100_000, 10).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_some_when_above_threshold() {
|
||||
let p = enabled_policy();
|
||||
let t = should_compact(&p, 90_000, 100_000, 10).expect("should trigger");
|
||||
assert_eq!(t.last_prompt_tokens, 90_000);
|
||||
assert_eq!(t.context_window, 100_000);
|
||||
assert_eq!(t.percent, 90);
|
||||
assert_eq!(t.step, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_replace_keeps_field_but_ignores_min_steps() {
|
||||
let p = enabled_policy();
|
||||
assert_eq!(p.mode, IntraCompactionMode::FullReplace);
|
||||
assert_eq!(p.min_steps_before_compact, 3);
|
||||
// Field is present; FullReplace only uses the token threshold
|
||||
// (parity with grok-build auto-compact).
|
||||
let t = should_compact(&p, 90_000, 100_000, 0).expect("should trigger");
|
||||
assert_eq!(t.step, 0);
|
||||
assert!(should_compact(&p, 90_000, 100_000, 2).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_modes_enforce_min_steps() {
|
||||
for mode in [
|
||||
IntraCompactionMode::StepsOnly,
|
||||
IntraCompactionMode::HistoryOnly,
|
||||
IntraCompactionMode::HistoryThenSteps,
|
||||
] {
|
||||
let p = enabled_partial_policy(mode);
|
||||
assert!(
|
||||
should_compact(&p, 90_000, 100_000, 2).is_none(),
|
||||
"{mode:?} should gate on min_steps"
|
||||
);
|
||||
let t = should_compact(&p, 90_000, 100_000, 3).expect("should trigger at min steps");
|
||||
assert_eq!(t.step, 3);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_when_context_window_zero() {
|
||||
let p = enabled_policy();
|
||||
assert!(should_compact(&p, 1_000, 0, 10).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn percent_caps_at_100() {
|
||||
let p = enabled_policy();
|
||||
let t = should_compact(&p, 200_000, 100_000, 10).expect("should trigger");
|
||||
assert_eq!(t.percent, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn boundary_exact_threshold_does_not_trigger() {
|
||||
let p = enabled_policy();
|
||||
// last_prompt_tokens == threshold: not strictly greater than.
|
||||
assert!(should_compact(&p, 85_000, 100_000, 10).is_none());
|
||||
// One above triggers.
|
||||
assert!(should_compact(&p, 85_001, 100_000, 10).is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
//! Data abstraction — the `CompactionItem` seam.
|
||||
//!
|
||||
//! The shared compaction algorithms operate over a sequence of *items*
|
||||
//! (turns/messages) without knowing the concrete harness type. The chat
|
||||
//! harness implements [`CompactionItem`] for its `GrokTurn`;
|
||||
//! grok-build implements it for `kigi_sampling_types::ConversationItem`.
|
||||
//!
|
||||
//! Keeping the contract minimal is deliberate: the algorithms only need
|
||||
//! enough structure to (a) classify roles, (b) read text, and (c) preserve
|
||||
//! the tool-request/tool-result pairing invariant when selecting a split
|
||||
//! point (an `Assistant(tool_request)` and the `Tool` results that satisfy
|
||||
//! it must never be separated, or the model API rejects the orphaned tool
|
||||
//! results with a 400).
|
||||
//!
|
||||
//! [`CompactionItemBuilder`] is the *constructive* extension used by the
|
||||
//! history-compaction algorithms that need to rebuild items (strip prior
|
||||
//! `<grok_user_queries>` blocks, drop tool content from assistant turns,
|
||||
//! wrap an LLM summary into a carrier item).
|
||||
|
||||
/// Harness-agnostic role of a single conversation item.
|
||||
///
|
||||
/// This is the common denominator of `GrokRole` (Grok chat) and the
|
||||
/// `ConversationItem` variants (grok-build).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CompactionRole {
|
||||
/// System prompt.
|
||||
System,
|
||||
/// Developer prompt (Grok chat) — maps to System on harnesses without a
|
||||
/// distinct developer role.
|
||||
Developer,
|
||||
/// A user message.
|
||||
User,
|
||||
/// An assistant output (may carry tool requests).
|
||||
Assistant,
|
||||
/// A tool result.
|
||||
Tool,
|
||||
}
|
||||
|
||||
/// A file attached to a user item, as seen by the shared user-query
|
||||
/// extraction (`<grok_file id=".." name=".." />` lines in the
|
||||
/// `<grok_user_queries>` preamble).
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CompactionFileRef {
|
||||
/// Stable unique id of the attachment source.
|
||||
pub id: String,
|
||||
/// Human-readable file name.
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
/// Contract: one turn/item in a conversation, as seen by the shared
|
||||
/// compaction algorithms.
|
||||
///
|
||||
/// Implementors:
|
||||
/// - Grok chat: `GrokTurn`
|
||||
/// - grok-build: `ConversationItem`
|
||||
pub trait CompactionItem {
|
||||
/// The harness-agnostic role of this item.
|
||||
fn role(&self) -> CompactionRole;
|
||||
|
||||
/// The item's text content, if any. Tool results and assistant tool-only
|
||||
/// turns may have no text.
|
||||
///
|
||||
/// Returns an owned `String` because some harnesses (Grok chat's
|
||||
/// `GrokTurn`) compute the flattened text on demand rather than storing a
|
||||
/// borrowable slice.
|
||||
fn text(&self) -> Option<String>;
|
||||
|
||||
/// Whether this item is a tool result. Used by the split-point selector to
|
||||
/// avoid orphaning tool results from their originating assistant turn.
|
||||
fn is_tool_result(&self) -> bool {
|
||||
matches!(self.role(), CompactionRole::Tool)
|
||||
}
|
||||
|
||||
/// Whether this (assistant) item carries at least one tool request.
|
||||
/// `false` for all non-assistant items.
|
||||
fn has_tool_requests(&self) -> bool;
|
||||
|
||||
/// Whether this item carries a *prior compaction summary* (Grok chat: a
|
||||
/// `Developer` turn with `DeveloperPromptCategory::ConversationCompaction`).
|
||||
///
|
||||
/// The basic history filter keeps such items so earlier summaries get
|
||||
/// re-summarised instead of dropped, and `separate_prior_user_queries`
|
||||
/// strips their `<grok_user_queries>` blocks before sampling.
|
||||
///
|
||||
/// Required (no default) on purpose: a forgotten implementation or a
|
||||
/// missed `Arc` forwarding would silently drop prior summaries on
|
||||
/// re-compaction.
|
||||
fn is_compaction_summary(&self) -> bool;
|
||||
|
||||
/// File attachments on a (user) item, for the `<grok_file>` lines in the
|
||||
/// `<grok_user_queries>` preamble. Empty for items without attachments.
|
||||
///
|
||||
/// Required (no default) for the same reason as
|
||||
/// [`Self::is_compaction_summary`]: silent attachment loss on compaction
|
||||
/// must be a compile error, not a runtime surprise.
|
||||
fn attachment_refs(&self) -> Vec<CompactionFileRef>;
|
||||
}
|
||||
|
||||
/// Constructive extension of [`CompactionItem`] for algorithms that rebuild
|
||||
/// items (history filtering and summary-carrier construction).
|
||||
///
|
||||
/// Not object-safe (`compaction_summary_item` has no receiver) — always used
|
||||
/// through generics, never as `dyn`.
|
||||
pub trait CompactionItemBuilder: CompactionItem + Clone {
|
||||
/// Construct the item that carries a compaction summary back into the
|
||||
/// conversation (Grok chat: a `Developer` turn with category
|
||||
/// `ConversationCompaction`). The result must satisfy
|
||||
/// `is_compaction_summary() == true`.
|
||||
fn compaction_summary_item(text: String) -> Self;
|
||||
|
||||
/// Rebuild this item keeping only user-visible content, dropping tool
|
||||
/// requests/results (Grok chat: keep only `Channel` contents of an
|
||||
/// assistant turn). Returns `None` when nothing visible remains.
|
||||
///
|
||||
/// Only meaningful for `Assistant` items; the shared filters never call
|
||||
/// it for other roles, but implementations should return
|
||||
/// `Some(self.clone())` for them to keep the contract total.
|
||||
fn strip_tool_content(&self) -> Option<Self>;
|
||||
}
|
||||
|
||||
/// Write seam for the full-replace **assembler**
|
||||
/// ([`crate::code_compaction::assemble::assemble_compacted_history`]):
|
||||
/// constructs the typed harness items that make up grok-build's rebuilt
|
||||
/// history.
|
||||
///
|
||||
/// This is a sibling of [`CompactionItemBuilder`], not a part of it, on
|
||||
/// purpose. `CompactionItemBuilder` is already implemented by Grok chat's
|
||||
/// `GrokTurn`; adding these constructors to it as required methods would break
|
||||
/// that impl. They are also grok-build-specific (Grok chat's tail-keep path
|
||||
/// has no `user_meta` / `project_instructions` / `system_reminder` carrier
|
||||
/// concept), so they live in their own seam that only the full-replace
|
||||
/// assembler depends on.
|
||||
///
|
||||
/// The grok-build implementor (`ConversationItem`) maps each constructor to the
|
||||
/// matching factory so the `SyntheticReason` tags the replay / spawn-time
|
||||
/// idempotence guards rely on are preserved.
|
||||
pub trait CompactionItemFactory: Sized {
|
||||
/// A real user message (used for the last user query).
|
||||
fn new_user(text: String) -> Self;
|
||||
/// A synthetic user message carrying compaction metadata (user-info
|
||||
/// prefix, summary carrier).
|
||||
fn new_user_meta(text: String) -> Self;
|
||||
/// A user message carrying project instructions (AGENTS.md), tagged so
|
||||
/// spawn-time idempotence guards recognize it on resume.
|
||||
fn new_project_instructions(text: String) -> Self;
|
||||
/// A synthetic user message carrying a `<system-reminder>` block.
|
||||
fn new_system_reminder(text: String) -> Self;
|
||||
}
|
||||
|
||||
/// Forward [`CompactionItem`] through shared references so the algorithms can
|
||||
/// operate over `&[Arc<T>]` (Grok chat stores turns as `Arc<GrokTurn>`).
|
||||
impl<T: CompactionItem + ?Sized> CompactionItem for std::sync::Arc<T> {
|
||||
fn role(&self) -> CompactionRole {
|
||||
(**self).role()
|
||||
}
|
||||
fn text(&self) -> Option<String> {
|
||||
(**self).text()
|
||||
}
|
||||
fn is_tool_result(&self) -> bool {
|
||||
(**self).is_tool_result()
|
||||
}
|
||||
fn has_tool_requests(&self) -> bool {
|
||||
(**self).has_tool_requests()
|
||||
}
|
||||
fn is_compaction_summary(&self) -> bool {
|
||||
(**self).is_compaction_summary()
|
||||
}
|
||||
fn attachment_refs(&self) -> Vec<CompactionFileRef> {
|
||||
(**self).attachment_refs()
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward [`CompactionItemBuilder`] through `Arc` — rebuilt items are
|
||||
/// wrapped in a fresh `Arc`, untouched items are *not* deep-cloned (the
|
||||
/// shared filters clone the `Arc` pointer directly).
|
||||
impl<T: CompactionItemBuilder> CompactionItemBuilder for std::sync::Arc<T> {
|
||||
fn compaction_summary_item(text: String) -> Self {
|
||||
std::sync::Arc::new(T::compaction_summary_item(text))
|
||||
}
|
||||
fn strip_tool_content(&self) -> Option<Self> {
|
||||
(**self).strip_tool_content().map(std::sync::Arc::new)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
//! Shared, transport-agnostic compaction engine.
|
||||
//!
|
||||
//! This crate is the `compaction-core`: shared policy, prompts, selection,
|
||||
//! and assembly. Host-specific trigger wiring, transport, persistence /
|
||||
//! replay / rewind, state commit, metrics backends, and prompt-variant forks
|
||||
//! stay in each product host (for example `kigi-shell`).
|
||||
//!
|
||||
//! The crate depends on **neither** a conversation-type crate nor
|
||||
//! `kigi-sampling-types`. It is decoupled from both Grok chat and
|
||||
//! grok-build hosts through a small set of trait seams:
|
||||
//!
|
||||
//! - [`CompactionItem`] / [`CompactionRole`] / [`CompactionItemBuilder`] —
|
||||
//! abstracts a single turn and its reconstruction.
|
||||
//! - [`ItemTokenCounter`] — trusted token counting per host.
|
||||
//! - [`CompactionSampler`] — the LLM call.
|
||||
//! - [`CompactionStreamProc`](intra_compaction::CompactionStreamProc) —
|
||||
//! state commit for intra-compaction.
|
||||
//! - [`IntraCompactionObserver`](intra_compaction::IntraCompactionObserver) /
|
||||
//! [`InterCompactionObserver`](inter_compaction::InterCompactionObserver)
|
||||
//! — host metrics.
|
||||
//!
|
||||
//! Compaction styles live in their own modules:
|
||||
//!
|
||||
//! - [`code_compaction`] — grok-build's whole-session **full-replace**
|
||||
//! subsystem (prompt/summary/failure/config, assemble, orchestration).
|
||||
//! - [`intra_compaction`] — Grok chat's tail-keep, per-step pass.
|
||||
//! - [`inter_compaction`] — Grok chat's chunked, between-turn pass.
|
||||
//!
|
||||
//! Compaction-type content (parallel subfolders): [`steps`] (the step-level
|
||||
//! prompt) and [`history`] (filtering, history prompts, validation +
|
||||
//! user-query preservation).
|
||||
//!
|
||||
//! Shared seams/primitives: [`item`], [`token`], [`sampler`],
|
||||
//! [`prompt::CompactionPrompt`], [`select`] (tool-pair-safe tail-keep
|
||||
//! selection — shared by the intra `Steps` and `History` targets, so it stays
|
||||
//! neutral at the crate root rather than under `steps`), and [`reminder`]
|
||||
//! (active-agent-state `<system-reminder>` formatting shared by Grok chat and
|
||||
//! grok-build; hosts still own snapshotting and host-only sections).
|
||||
|
||||
pub mod code_compaction;
|
||||
pub mod history;
|
||||
pub mod inter_compaction;
|
||||
pub mod intra_compaction;
|
||||
pub mod item;
|
||||
pub mod prompt;
|
||||
pub mod reminder;
|
||||
pub mod sampler;
|
||||
pub mod select;
|
||||
pub mod steps;
|
||||
pub mod token;
|
||||
|
||||
/// Shared code default for the dedicated compaction model name.
|
||||
///
|
||||
/// Override order (highest first):
|
||||
/// 1. the agent's `compaction_model_name` setting (non-blank)
|
||||
/// 2. service / harness config YAML
|
||||
/// 3. this constant
|
||||
pub use intra_compaction::DEFAULT_COMPACTION_MODEL_NAME;
|
||||
|
||||
// grok-build's full-replace subsystem now lives under `code_compaction`;
|
||||
// re-exported at the crate root so consumers keep a stable public API.
|
||||
pub use code_compaction::{
|
||||
CompactedHistoryParts, DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT, FailureKind,
|
||||
FullReplaceAttemptOutcome, FullReplaceConfig, FullReplaceContext, FullReplaceError,
|
||||
FullReplaceObserver, FullReplaceOutput, FullReplaceSummary, MIN_SUMMARY_SEED_CHARS,
|
||||
SELF_SUMMARIZATION_PROMPT, SummaryPromptKind, apply_full_replace_compaction,
|
||||
assemble_compacted_history, build_summary_prompt, build_summary_prompt_kind,
|
||||
classify_http_status, classify_stream_event_error, format_compact_summary,
|
||||
format_compact_summary_content, is_context_length_error, is_degenerate_summary,
|
||||
sample_full_replace_summary, wrap_user_query,
|
||||
};
|
||||
pub use item::{
|
||||
CompactionFileRef, CompactionItem, CompactionItemBuilder, CompactionItemFactory, CompactionRole,
|
||||
};
|
||||
pub use prompt::CompactionPrompt;
|
||||
// Reminder types/formatters: import from `reminder::` (borrowed views).
|
||||
// Only the summary-injection helper is re-exported at the crate root — both
|
||||
// intra FullReplace and inter already use it by this name.
|
||||
pub use reminder::append_reminder_block;
|
||||
pub use sampler::{CompactionSampleError, CompactionSampler, LlmCompactionOutput};
|
||||
pub use select::{SplitPlan, select_turns_to_compact};
|
||||
pub use steps::format_compaction_prompt;
|
||||
pub use token::ItemTokenCounter;
|
||||
@@ -0,0 +1,16 @@
|
||||
//! 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:
|
||||
//!
|
||||
//! - steps prompt → [`crate::steps::format_compaction_prompt`]
|
||||
//! - history prompts → [`crate::history::prompt`]
|
||||
//! - grok-build summary prompt → [`crate::code_compaction::build_summary_prompt`]
|
||||
|
||||
/// System + user prompt pair for the compaction LLM call.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CompactionPrompt {
|
||||
pub system: String,
|
||||
pub user: String,
|
||||
}
|
||||
@@ -0,0 +1,519 @@
|
||||
//! Shared post-compaction reminder helpers (host-agnostic).
|
||||
//!
|
||||
//! Lives at the crate root rather than under a compaction-style submodule
|
||||
//! because it is consumed by *both* compaction styles and both harnesses:
|
||||
//!
|
||||
//! - Grok chat intra FullReplace ([`crate::intra_compaction`]) and inter
|
||||
//! (appends after sampling via [`append_reminder_block`])
|
||||
//! - grok-build full-replace ([`crate::code_compaction`] assemble's
|
||||
//! `system_reminder`)
|
||||
//!
|
||||
//! **What lives here:** pure formatting of the three **common** active-agent
|
||||
//! sections — Running Background Tasks, TODO List, Running Subagents — plus
|
||||
//! `<system-reminder>` wrapping and summary append.
|
||||
//!
|
||||
//! **What stays in the product host:** snapshotting, tool-name resolution, and harness-only
|
||||
//! sections (files, AGENTS.md, skills, MCP, memory). Callers pass **borrowed
|
||||
//! 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.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct SubagentToolNames<'a> {
|
||||
pub poll: &'a str,
|
||||
pub cancel: &'a str,
|
||||
}
|
||||
|
||||
/// Status of a todo item in the post-compaction reminder.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TodoStatus {
|
||||
Pending,
|
||||
InProgress,
|
||||
Completed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl TodoStatus {
|
||||
pub fn is_actionable(self) -> bool {
|
||||
matches!(self, Self::Pending | Self::InProgress)
|
||||
}
|
||||
|
||||
pub fn tag(self) -> &'static str {
|
||||
match self {
|
||||
Self::Pending => "[pending]",
|
||||
Self::InProgress => "[in_progress]",
|
||||
Self::Completed => "[completed]",
|
||||
Self::Cancelled => "[cancelled]",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct TodoItem<'a> {
|
||||
pub id: &'a str,
|
||||
pub content: &'a str,
|
||||
pub status: TodoStatus,
|
||||
}
|
||||
|
||||
/// Still-running background task. `task_id` is rendered verbatim.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct BackgroundTask<'a> {
|
||||
pub task_id: &'a str,
|
||||
pub command: &'a str,
|
||||
/// Parenthetical status (typically `"running"`).
|
||||
pub status: &'a str,
|
||||
pub tool_name: Option<&'a str>,
|
||||
}
|
||||
|
||||
/// Still-running sub-agent. `subagent_id` is rendered verbatim.
|
||||
///
|
||||
/// `subagent_type` / `description` are optional so chat (no type, optional
|
||||
/// desc) and build (both present) share one line format.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RunningSubagent<'a> {
|
||||
pub subagent_id: &'a str,
|
||||
pub subagent_type: Option<&'a str>,
|
||||
pub description: Option<&'a str>,
|
||||
pub elapsed_secs: u64,
|
||||
}
|
||||
|
||||
/// Borrowed active-agent state for reminder rendering.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ActiveAgentReminderState<'a> {
|
||||
pub running_commands: &'a [BackgroundTask<'a>],
|
||||
pub todos: &'a [TodoItem<'a>],
|
||||
pub running_subagents: &'a [RunningSubagent<'a>],
|
||||
}
|
||||
|
||||
impl ActiveAgentReminderState<'_> {
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.running_commands.is_empty()
|
||||
&& self.running_subagents.is_empty()
|
||||
&& !self.has_actionable_todos()
|
||||
}
|
||||
|
||||
pub fn has_actionable_todos(&self) -> bool {
|
||||
self.todos.iter().any(|t| t.status.is_actionable())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Section formatters
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// `## Running Background Tasks`, or `None` when empty.
|
||||
pub fn section_background_tasks(tasks: &[BackgroundTask<'_>]) -> Option<String> {
|
||||
if tasks.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let lines = tasks
|
||||
.iter()
|
||||
.map(|t| match t.tool_name {
|
||||
Some(tool) => format!(
|
||||
"- \"{}\": `{}` ({}, {})",
|
||||
t.task_id, t.command, t.status, tool
|
||||
),
|
||||
None => format!("- \"{}\": `{}` ({})", t.task_id, t.command, t.status),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
Some(format!(
|
||||
"## Running Background Tasks\n\
|
||||
These tasks are still running:\n{lines}"
|
||||
))
|
||||
}
|
||||
|
||||
/// `## TODO List` for actionable items, or `None` when none. Completed/
|
||||
/// cancelled collapse to a count trailer.
|
||||
pub fn section_todo_list(todos: &[TodoItem<'_>]) -> Option<String> {
|
||||
let active: Vec<_> = todos
|
||||
.iter()
|
||||
.filter(|t| t.status.is_actionable())
|
||||
.map(|t| format!("- {} {}: {}", t.status.tag(), t.id, t.content))
|
||||
.collect();
|
||||
if active.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let completed = todos
|
||||
.iter()
|
||||
.filter(|t| t.status == TodoStatus::Completed)
|
||||
.count();
|
||||
let cancelled = todos
|
||||
.iter()
|
||||
.filter(|t| t.status == TodoStatus::Cancelled)
|
||||
.count();
|
||||
let trailer = match (completed, cancelled) {
|
||||
(0, 0) => String::new(),
|
||||
(c, 0) => format!("\n({c} completed)"),
|
||||
(0, k) => format!("\n({k} cancelled)"),
|
||||
(c, k) => format!("\n({c} completed, {k} cancelled)"),
|
||||
};
|
||||
Some(format!(
|
||||
"## TODO List\n\
|
||||
This is your task list from before the conversation was compacted — it is still \
|
||||
active. Keep working through the items below and update their status as you make \
|
||||
progress:\n{}{trailer}",
|
||||
active.join("\n"),
|
||||
))
|
||||
}
|
||||
|
||||
/// `## Running Subagents`, or `None` when empty. Omit entirely when tool
|
||||
/// names cannot be resolved rather than point at wrong names.
|
||||
pub fn section_running_subagents(
|
||||
subagents: &[RunningSubagent<'_>],
|
||||
tools: &SubagentToolNames<'_>,
|
||||
) -> Option<String> {
|
||||
if subagents.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let lines = subagents
|
||||
.iter()
|
||||
.map(format_subagent_line)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
Some(format!(
|
||||
"## Running Subagents\n\
|
||||
These subagents were launched before this compaction and are still running. \
|
||||
Use `{}` with the subagent_id to check their status or retrieve results. \
|
||||
Use `{}` with the subagent_id to cancel a subagent.\n{lines}",
|
||||
tools.poll, tools.cancel
|
||||
))
|
||||
}
|
||||
|
||||
fn format_subagent_line(s: &RunningSubagent<'_>) -> String {
|
||||
let mut head = format!("subagent_id: `{}`", s.subagent_id);
|
||||
if let Some(ty) = s.subagent_type {
|
||||
head.push_str(", type: `");
|
||||
head.push_str(ty);
|
||||
head.push('`');
|
||||
}
|
||||
if let Some(desc) = s.description {
|
||||
head.push_str(", task: \"");
|
||||
head.push_str(desc);
|
||||
head.push('"');
|
||||
}
|
||||
format!("- {head} (running for {}s)", s.elapsed_secs)
|
||||
}
|
||||
|
||||
/// Common sections in order: Background Tasks → TODO → Subagents.
|
||||
/// Empty kinds omitted; subagents also omitted when `subagent_tools` is `None`.
|
||||
pub fn format_active_agent_sections(
|
||||
state: &ActiveAgentReminderState<'_>,
|
||||
subagent_tools: Option<&SubagentToolNames<'_>>,
|
||||
) -> Vec<String> {
|
||||
let mut sections = Vec::with_capacity(3);
|
||||
if let Some(s) = section_background_tasks(state.running_commands) {
|
||||
sections.push(s);
|
||||
}
|
||||
if let Some(s) = section_todo_list(state.todos) {
|
||||
sections.push(s);
|
||||
}
|
||||
if let Some(tools) = subagent_tools
|
||||
&& let Some(s) = section_running_subagents(state.running_subagents, tools)
|
||||
{
|
||||
sections.push(s);
|
||||
}
|
||||
sections
|
||||
}
|
||||
|
||||
/// Wrap non-empty sections in `<system-reminder>…</system-reminder>`.
|
||||
pub fn wrap_system_reminder(sections: impl IntoIterator<Item = impl AsRef<str>>) -> Option<String> {
|
||||
let mut body = String::new();
|
||||
for s in sections {
|
||||
let s = s.as_ref();
|
||||
if s.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
if !body.is_empty() {
|
||||
body.push_str("\n\n");
|
||||
}
|
||||
body.push_str(s);
|
||||
}
|
||||
if body.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(format!("<system-reminder>\n{body}\n</system-reminder>"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Full active-agent-state `<system-reminder>`, or `None` when nothing to preserve.
|
||||
pub fn format_active_agent_reminder(
|
||||
state: &ActiveAgentReminderState<'_>,
|
||||
subagent_tools: Option<&SubagentToolNames<'_>>,
|
||||
) -> Option<String> {
|
||||
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.
|
||||
pub fn append_reminder_block(summary: String, reminder: Option<&str>) -> String {
|
||||
match reminder {
|
||||
Some(reminder) if !reminder.trim().is_empty() => format!("{summary}\n\n{reminder}"),
|
||||
_ => summary,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn tools_native() -> SubagentToolNames<'static> {
|
||||
SubagentToolNames {
|
||||
poll: "get_task_output",
|
||||
cancel: "kill_task",
|
||||
}
|
||||
}
|
||||
|
||||
fn tools_renamed() -> SubagentToolNames<'static> {
|
||||
SubagentToolNames {
|
||||
poll: "get_command_or_subagent_output",
|
||||
cancel: "kill_command_or_subagent",
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_state_is_none() {
|
||||
assert!(
|
||||
format_active_agent_reminder(
|
||||
&ActiveAgentReminderState::default(),
|
||||
Some(&tools_native())
|
||||
)
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_tool_names_omits_subagent_section_only() {
|
||||
let agents = [RunningSubagent {
|
||||
subagent_id: "sa-1",
|
||||
subagent_type: None,
|
||||
description: Some("x"),
|
||||
elapsed_secs: 1,
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_subagents: &agents,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(format_active_agent_reminder(&state, None).is_none());
|
||||
|
||||
let cmds = [BackgroundTask {
|
||||
task_id: "bg-1",
|
||||
command: "npm run dev",
|
||||
status: "running",
|
||||
tool_name: Some("run_terminal_command"),
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_commands: &cmds,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, None).expect("reminder");
|
||||
assert!(out.contains("## Running Background Tasks"));
|
||||
assert!(!out.contains("## Running Subagents"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_chat_style_subagent_ids_verbatim() {
|
||||
let agents = [
|
||||
RunningSubagent {
|
||||
subagent_id: "019ea7f0-cb66-7aa2-9a09-488a3a795795",
|
||||
subagent_type: None,
|
||||
description: Some("deploy staging"),
|
||||
elapsed_secs: 42,
|
||||
},
|
||||
RunningSubagent {
|
||||
subagent_id: "sa-2",
|
||||
subagent_type: None,
|
||||
description: None,
|
||||
elapsed_secs: 5,
|
||||
},
|
||||
];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_subagents: &agents,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, Some(&tools_native())).expect("reminder");
|
||||
assert!(out.starts_with("<system-reminder>"));
|
||||
assert!(out.ends_with("</system-reminder>"));
|
||||
assert!(out.contains("subagent_id: `019ea7f0-cb66-7aa2-9a09-488a3a795795`"));
|
||||
assert!(out.contains("task: \"deploy staging\" (running for 42s)"));
|
||||
assert!(out.contains("subagent_id: `sa-2` (running for 5s)"));
|
||||
assert!(!out.contains("task-019ea7f0"));
|
||||
assert!(!out.contains("type:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_build_style_subagent_with_type() {
|
||||
let agents = [RunningSubagent {
|
||||
subagent_id: "sub-1",
|
||||
subagent_type: Some("explore"),
|
||||
description: Some("find files"),
|
||||
elapsed_secs: 5,
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_subagents: &agents,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, Some(&tools_renamed())).expect("reminder");
|
||||
assert!(out.contains(
|
||||
"- subagent_id: `sub-1`, type: `explore`, task: \"find files\" (running for 5s)"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uses_renamed_tool_names_verbatim() {
|
||||
let agents = [RunningSubagent {
|
||||
subagent_id: "sa-1",
|
||||
subagent_type: None,
|
||||
description: Some("x"),
|
||||
elapsed_secs: 1,
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_subagents: &agents,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, Some(&tools_renamed())).expect("reminder");
|
||||
assert!(out.contains("get_command_or_subagent_output"));
|
||||
assert!(!out.contains("get_task_output"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_background_tasks() {
|
||||
let cmds = [
|
||||
BackgroundTask {
|
||||
task_id: "019f1723-a9f0-76f2-98ae-56af965922f6",
|
||||
command: "npm run dev",
|
||||
status: "running",
|
||||
tool_name: Some("run_terminal_command"),
|
||||
},
|
||||
BackgroundTask {
|
||||
task_id: "bg-2",
|
||||
command: "cargo watch -x test",
|
||||
status: "running",
|
||||
tool_name: None,
|
||||
},
|
||||
];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_commands: &cmds,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, None).expect("reminder");
|
||||
assert!(out.contains(
|
||||
"- \"019f1723-a9f0-76f2-98ae-56af965922f6\": `npm run dev` (running, run_terminal_command)"
|
||||
));
|
||||
assert!(out.contains("- \"bg-2\": `cargo watch -x test` (running)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_todo_list_without_tool_names() {
|
||||
let todos = [
|
||||
TodoItem {
|
||||
id: "1",
|
||||
content: "scaffold the app",
|
||||
status: TodoStatus::Completed,
|
||||
},
|
||||
TodoItem {
|
||||
id: "2",
|
||||
content: "wire the API",
|
||||
status: TodoStatus::InProgress,
|
||||
},
|
||||
TodoItem {
|
||||
id: "3",
|
||||
content: "write tests",
|
||||
status: TodoStatus::Pending,
|
||||
},
|
||||
];
|
||||
let state = ActiveAgentReminderState {
|
||||
todos: &todos,
|
||||
..Default::default()
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, None).expect("reminder");
|
||||
assert!(out.contains("- [in_progress] 2: wire the API"));
|
||||
assert!(out.contains("- [pending] 3: write tests"));
|
||||
assert!(out.contains("(1 completed)"));
|
||||
assert!(!out.contains("scaffold the app"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn only_completed_todos_is_none() {
|
||||
let todos = [TodoItem {
|
||||
id: "1",
|
||||
content: "done",
|
||||
status: TodoStatus::Completed,
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
todos: &todos,
|
||||
..Default::default()
|
||||
};
|
||||
assert!(format_active_agent_reminder(&state, None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn section_order_background_todo_subagent() {
|
||||
let cmds = [BackgroundTask {
|
||||
task_id: "t1",
|
||||
command: "npm run dev",
|
||||
status: "running",
|
||||
tool_name: None,
|
||||
}];
|
||||
let todos = [TodoItem {
|
||||
id: "2",
|
||||
content: "wire the API",
|
||||
status: TodoStatus::InProgress,
|
||||
}];
|
||||
let agents = [RunningSubagent {
|
||||
subagent_id: "sa-1",
|
||||
subagent_type: None,
|
||||
description: Some("deploy staging"),
|
||||
elapsed_secs: 1,
|
||||
}];
|
||||
let state = ActiveAgentReminderState {
|
||||
running_commands: &cmds,
|
||||
todos: &todos,
|
||||
running_subagents: &agents,
|
||||
};
|
||||
let out = format_active_agent_reminder(&state, Some(&tools_native())).expect("reminder");
|
||||
let bg = out.find("## Running Background Tasks").expect("bg");
|
||||
let todo = out.find("## TODO List").expect("todo");
|
||||
let sub = out.find("## Running Subagents").expect("sub");
|
||||
assert!(bg < todo && todo < sub, "order wrong:\n{out}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrap_system_reminder_joins_and_skips_blank() {
|
||||
let out = wrap_system_reminder(["## A\nx", "", " ", "## B\ny"]).expect("wrapped");
|
||||
assert_eq!(
|
||||
out,
|
||||
"<system-reminder>\n## A\nx\n\n## B\ny\n</system-reminder>"
|
||||
);
|
||||
assert!(wrap_system_reminder(std::iter::empty::<&str>()).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_after_blank_line() {
|
||||
assert_eq!(
|
||||
append_reminder_block("SUMMARY".to_string(), Some("REMINDER")),
|
||||
"SUMMARY\n\nREMINDER"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_noop_when_none_or_blank() {
|
||||
assert_eq!(
|
||||
append_reminder_block("SUMMARY".to_string(), None),
|
||||
"SUMMARY"
|
||||
);
|
||||
assert_eq!(
|
||||
append_reminder_block("SUMMARY".to_string(), Some(" \n\t ")),
|
||||
"SUMMARY"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,184 @@
|
||||
//! The `CompactionSampler` seam — the LLM call that produces summaries —
|
||||
//! plus its output and error types (shared failure classification).
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::prompt::CompactionPrompt;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sampler output + error types
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Raw text captured from a compaction LLM call, split by channel.
|
||||
///
|
||||
/// Used by both intra- and inter-compaction. Intra-compaction uses only
|
||||
/// `.response`; inter-compaction also persists `.thinking` for audit/debug.
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct LlmCompactionOutput {
|
||||
/// Text from the response channel — the actual compaction summary.
|
||||
pub response: String,
|
||||
/// Text from the thinking channel — the model's chain-of-thought reasoning.
|
||||
/// Stored for audit/debug only; never fed back into a conversation.
|
||||
pub thinking: String,
|
||||
}
|
||||
|
||||
/// Error types for compaction sampling, allowing callers to distinguish
|
||||
/// deterministic failures (never retry) from transient ones.
|
||||
///
|
||||
/// Harnesses should prefer the structured variants ([`Self::Build`],
|
||||
/// [`Self::Start`], [`Self::EmptyResponse`]) so the shared retry policy can
|
||||
/// classify without string matching. [`Self::Other`] remains for samplers
|
||||
/// that only surface an opaque error; the orchestrator falls back to
|
||||
/// matching the literal messages produced by the Grok chat sampler —
|
||||
/// keep those literals in sync (the `compaction_sample_error_to_intra*`
|
||||
/// tests guard the mapping).
|
||||
#[derive(Debug)]
|
||||
pub enum CompactionSampleError {
|
||||
/// The sampler hit its end-to-end timeout. Transient.
|
||||
Timeout {
|
||||
timeout_secs: u64,
|
||||
collected_bytes: usize,
|
||||
},
|
||||
/// Sampler construction failed (bad config, unknown model). Deterministic.
|
||||
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
|
||||
/// `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.
|
||||
Other(anyhow::Error),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for CompactionSampleError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Timeout {
|
||||
timeout_secs,
|
||||
collected_bytes,
|
||||
} => write!(
|
||||
f,
|
||||
"Compaction sampling timed out after {}s (collected {} bytes so far)",
|
||||
timeout_secs, collected_bytes
|
||||
),
|
||||
Self::Build(msg) => write!(f, "Compaction sampler build failed: {}", msg),
|
||||
Self::Start(msg) => write!(f, "Compaction sampler start failed: {}", msg),
|
||||
// Keep the "no response channel content" literal — the intra
|
||||
// orchestrator's `Other(_)` fallback string-matches it.
|
||||
Self::EmptyResponse => {
|
||||
write!(f, "Compaction sampler returned no response channel content")
|
||||
}
|
||||
Self::Other(e) => write!(f, "{}", e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for CompactionSampleError {
|
||||
fn from(e: anyhow::Error) -> Self {
|
||||
Self::Other(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl CompactionSampleError {
|
||||
/// Whether this error is deterministic — retrying with the same input
|
||||
/// will produce the same failure.
|
||||
pub fn is_deterministic(&self) -> bool {
|
||||
match self {
|
||||
Self::Timeout { .. } | Self::EmptyResponse => false,
|
||||
Self::Build(_) | Self::Start(_) => true,
|
||||
Self::Other(err) => {
|
||||
let msg = err.to_string();
|
||||
msg.contains("Failed to build AgenticScheduler")
|
||||
|| msg.contains("Failed to start compaction sample")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sampler trait
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Interface for the LLM call that produces compaction summaries.
|
||||
///
|
||||
/// Used by both intra-compaction (steps/history) and inter-compaction.
|
||||
/// Implemented by each harness's sampler adapter; grok-build wires its own
|
||||
/// transport.
|
||||
///
|
||||
/// Returns [`LlmCompactionOutput`] containing both response and thinking
|
||||
/// channel text. Intra-compaction uses only `.response`; inter-compaction
|
||||
/// also persists `.thinking` for audit/debug.
|
||||
#[async_trait]
|
||||
pub trait CompactionSampler: Send + Sync {
|
||||
/// The harness's conversation item type.
|
||||
type Item;
|
||||
|
||||
/// Run an LLM compaction call on the given items.
|
||||
///
|
||||
/// Implementations should:
|
||||
/// - Build a synthetic conversation from the items + prompt.
|
||||
/// - Honor the `timeout`.
|
||||
/// - Collect both response and thinking channel text.
|
||||
async fn sample_compaction(
|
||||
&self,
|
||||
turns: &[Self::Item],
|
||||
prompt: &CompactionPrompt,
|
||||
timeout: Duration,
|
||||
) -> Result<LlmCompactionOutput, CompactionSampleError>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Pins the inter-compaction retry classification for every variant —
|
||||
/// `Start` is intentionally deterministic here (no inter retry) even
|
||||
/// though the intra orchestrator retries its `SamplerStart` mapping.
|
||||
/// See the doc on [`CompactionSampleError::Start`] before "fixing" this.
|
||||
#[test]
|
||||
fn is_deterministic_classification() {
|
||||
assert!(
|
||||
!CompactionSampleError::Timeout {
|
||||
timeout_secs: 1,
|
||||
collected_bytes: 0
|
||||
}
|
||||
.is_deterministic()
|
||||
);
|
||||
assert!(!CompactionSampleError::EmptyResponse.is_deterministic());
|
||||
assert!(CompactionSampleError::Build("bad config".into()).is_deterministic());
|
||||
assert!(CompactionSampleError::Start("no stream".into()).is_deterministic());
|
||||
// Legacy string-matching fallback.
|
||||
assert!(
|
||||
CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"Failed to build AgenticScheduler: config error"
|
||||
))
|
||||
.is_deterministic()
|
||||
);
|
||||
assert!(
|
||||
CompactionSampleError::Other(anyhow::anyhow!(
|
||||
"Failed to start compaction sample: stream error"
|
||||
))
|
||||
.is_deterministic()
|
||||
);
|
||||
assert!(
|
||||
!CompactionSampleError::Other(anyhow::anyhow!("transient stream error"))
|
||||
.is_deterministic()
|
||||
);
|
||||
}
|
||||
|
||||
/// The `EmptyResponse` Display must keep the "no response channel
|
||||
/// content" literal the intra `Other(_)` fallback string-matches.
|
||||
#[test]
|
||||
fn empty_response_display_keeps_match_literal() {
|
||||
let msg = CompactionSampleError::EmptyResponse.to_string();
|
||||
assert!(msg.contains("no response channel content"), "got: {msg}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
//! Turn selection for compaction.
|
||||
//!
|
||||
//! Walks the item list backward to find a split point: keep the newest items
|
||||
//! whose cumulative token count fits the target budget, compact everything
|
||||
//! older.
|
||||
//!
|
||||
//! The split point must respect a critical invariant: an assistant item with
|
||||
//! tool requests and the subsequent tool-result items that satisfy those
|
||||
//! requests must stay together. Splitting between them would produce orphan
|
||||
//! tool results in the next prompt, which the model API rejects with a 400.
|
||||
//!
|
||||
//! This is the harness-agnostic core: it operates over any slice of
|
||||
//! [`CompactionItem`], so both Grok chat (`GrokTurn`) and grok-build
|
||||
//! (`ConversationItem`) share one implementation.
|
||||
|
||||
use crate::item::CompactionItem;
|
||||
|
||||
/// Output of [`select_turns_to_compact`].
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SplitPlan {
|
||||
/// Compact items at indices `0..split_idx`. Keep `split_idx..total`.
|
||||
pub split_idx: usize,
|
||||
/// Sum of `item_token_counts[..split_idx]`.
|
||||
pub tokens_to_compact: u32,
|
||||
}
|
||||
|
||||
/// Decide where to split the items for compaction.
|
||||
///
|
||||
/// Algorithm:
|
||||
/// 1. Walk backward from the newest item, accumulating "keep" tokens.
|
||||
/// 2. The candidate split index is the first one where adding more would
|
||||
/// exceed `target_tokens`.
|
||||
/// 3. **Snap forward** to a safe boundary: if the split would orphan tool
|
||||
/// results, walk forward until past the matching tool-result items.
|
||||
/// 4. Return `None` if the resulting compactable region's token count is
|
||||
/// below `min_compactable` — not worth running the LLM.
|
||||
///
|
||||
/// # Tool-pair boundary safety
|
||||
///
|
||||
/// `items` is the agent's running state. A typical sequence:
|
||||
///
|
||||
/// ```text
|
||||
/// [Assistant(tool_request_A, tool_request_B),
|
||||
/// Tool(A_result),
|
||||
/// Tool(B_result),
|
||||
/// Assistant(response_text),
|
||||
/// Assistant(tool_request_C),
|
||||
/// Tool(C_result),
|
||||
/// ...]
|
||||
/// ```
|
||||
///
|
||||
/// A safe split point is one where everything **before** the split is
|
||||
/// self-contained (no dangling tool requests waiting for results that live
|
||||
/// after the split).
|
||||
///
|
||||
/// The rule we enforce: the split index must not fall in the middle of a
|
||||
/// `[Assistant-with-tool-requests, Tool, Tool, ...]` run. If the candidate
|
||||
/// split lands on a tool-result item, walk it forward until we pass the last
|
||||
/// tool-result item following the most recent assistant-with-tool-requests.
|
||||
pub fn select_turns_to_compact<T: CompactionItem>(
|
||||
item_token_counts: &[u32],
|
||||
items: &[T],
|
||||
target_tokens: u32,
|
||||
min_compactable: u32,
|
||||
) -> Option<SplitPlan> {
|
||||
debug_assert_eq!(
|
||||
item_token_counts.len(),
|
||||
items.len(),
|
||||
"token counts and items must have the same length"
|
||||
);
|
||||
|
||||
let total = items.len();
|
||||
if total == 0 {
|
||||
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.
|
||||
let mut kept = 0u32;
|
||||
let mut split_idx = total; // start with "compact nothing", will move down
|
||||
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;
|
||||
}
|
||||
kept = kept.saturating_add(count);
|
||||
split_idx = i;
|
||||
}
|
||||
|
||||
// If the whole list fits within the budget, nothing to compact.
|
||||
if split_idx == 0 {
|
||||
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.
|
||||
if safe_split_idx >= total {
|
||||
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()
|
||||
.fold(0u32, u32::saturating_add);
|
||||
|
||||
if tokens_to_compact < min_compactable {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(SplitPlan {
|
||||
split_idx: safe_split_idx,
|
||||
tokens_to_compact,
|
||||
})
|
||||
}
|
||||
|
||||
/// If `candidate` lands on a tool-result item, advance forward past all
|
||||
/// tool-result items in the same tool-pair run. The "run" is delimited by the
|
||||
/// previous assistant item (with tool requests) and the next non-tool item.
|
||||
///
|
||||
/// In effect: ensure the split lands either right before an assistant, user,
|
||||
/// system, or developer item — never between an assistant-with-tool-requests
|
||||
/// and its tool results.
|
||||
fn snap_to_safe_boundary<T: CompactionItem>(items: &[T], candidate: usize) -> usize {
|
||||
let total = items.len();
|
||||
if candidate >= total {
|
||||
return total;
|
||||
}
|
||||
|
||||
// If candidate is not a tool-result item, no snap needed.
|
||||
if !items[candidate].is_tool_result() {
|
||||
return candidate;
|
||||
}
|
||||
|
||||
// Candidate is a tool-result item. Find the run of contiguous tool-result
|
||||
// items (starting from the assistant-with-tool-requests that preceded
|
||||
// them) and advance to just past the last one in that run.
|
||||
let mut idx = candidate;
|
||||
while idx < total && items[idx].is_tool_result() {
|
||||
idx += 1;
|
||||
}
|
||||
idx
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::item::CompactionRole;
|
||||
|
||||
/// Minimal mock implementing [`CompactionItem`] for selection tests.
|
||||
struct MockItem {
|
||||
role: CompactionRole,
|
||||
}
|
||||
|
||||
impl MockItem {
|
||||
fn user() -> Self {
|
||||
Self {
|
||||
role: CompactionRole::User,
|
||||
}
|
||||
}
|
||||
fn assistant() -> Self {
|
||||
Self {
|
||||
role: CompactionRole::Assistant,
|
||||
}
|
||||
}
|
||||
fn tool() -> Self {
|
||||
Self {
|
||||
role: CompactionRole::Tool,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CompactionItem for MockItem {
|
||||
fn role(&self) -> CompactionRole {
|
||||
self.role
|
||||
}
|
||||
fn text(&self) -> Option<String> {
|
||||
None
|
||||
}
|
||||
fn has_tool_requests(&self) -> bool {
|
||||
false
|
||||
}
|
||||
fn is_compaction_summary(&self) -> bool {
|
||||
false
|
||||
}
|
||||
fn attachment_refs(&self) -> Vec<crate::item::CompactionFileRef> {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_returns_none() {
|
||||
let items: Vec<MockItem> = vec![];
|
||||
assert!(select_turns_to_compact(&[], &items, 100, 10).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_fits_in_budget_returns_none() {
|
||||
let items = vec![MockItem::user(), MockItem::assistant()];
|
||||
let counts = vec![10, 20];
|
||||
assert!(select_turns_to_compact(&counts, &items, 1000, 5).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn splits_at_correct_index() {
|
||||
// Total 100; target 30 → keep last few that fit in 30.
|
||||
let items = vec![
|
||||
MockItem::user(),
|
||||
MockItem::assistant(),
|
||||
MockItem::user(),
|
||||
MockItem::assistant(),
|
||||
];
|
||||
let counts = vec![40, 30, 20, 10]; // keep last two (sum 30)
|
||||
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);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn below_min_compactable_returns_none() {
|
||||
let items = vec![MockItem::user(), MockItem::assistant()];
|
||||
let counts = vec![5, 100];
|
||||
// Would split after index 0, but 5 < min_compactable (10).
|
||||
assert!(select_turns_to_compact(&counts, &items, 50, 10).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snaps_past_tool_results() {
|
||||
// Layout: [User, Assistant-text, Assistant-with-tools, Tool, Tool, Assistant-text]
|
||||
// If the naïve split lands on a Tool, snap forward past all Tools.
|
||||
let items = vec![
|
||||
MockItem::user(),
|
||||
MockItem::assistant(),
|
||||
MockItem::assistant(), // pretend this had tool_requests
|
||||
MockItem::tool(),
|
||||
MockItem::tool(),
|
||||
MockItem::assistant(),
|
||||
];
|
||||
let counts = vec![10, 10, 10, 50, 50, 10];
|
||||
|
||||
// Target 60 → walking back: keep 10 (idx 5), keep 50 (idx 4)
|
||||
// → 60 used. Adding idx 3 (50) overflows.
|
||||
// Naïve split = 4. But items[4] is Tool → snap forward.
|
||||
// Walk forward: items[4]=Tool, items[5]=Assistant → snap to 5.
|
||||
let plan = select_turns_to_compact(&counts, &items, 60, 5).expect("should split");
|
||||
assert_eq!(plan.split_idx, 5);
|
||||
assert_eq!(plan.tokens_to_compact, 10 + 10 + 10 + 50 + 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snap_does_not_advance_when_already_safe() {
|
||||
let items = vec![
|
||||
MockItem::user(),
|
||||
MockItem::assistant(),
|
||||
MockItem::user(), // safe split here
|
||||
MockItem::assistant(),
|
||||
];
|
||||
let counts = vec![50, 50, 10, 10];
|
||||
// Target 30 → keep last two (sum 20).
|
||||
// Naïve split = 2. items[2] = User → safe, no snap needed.
|
||||
let plan = select_turns_to_compact(&counts, &items, 30, 5).expect("should split");
|
||||
assert_eq!(plan.split_idx, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snap_walks_to_end_returns_none() {
|
||||
// Pathological: split would need to snap past all items.
|
||||
let items = vec![MockItem::assistant(), MockItem::tool(), MockItem::tool()];
|
||||
let counts = vec![10, 50, 50];
|
||||
// Target 0 → naïve split = 1 (items[1] is Tool).
|
||||
// Snap forward: items[1]=Tool, items[2]=Tool, idx=3=total.
|
||||
// Return None — nothing left to keep.
|
||||
assert!(select_turns_to_compact(&counts, &items, 0, 5).is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
//! Steps compaction — 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
|
||||
//! turn selection it shares with the History target is the crate-root
|
||||
//! [`select`](crate::select) primitive (not steps-specific).
|
||||
|
||||
pub mod prompt;
|
||||
|
||||
pub use prompt::format_compaction_prompt;
|
||||
@@ -0,0 +1,31 @@
|
||||
//! 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/`.
|
||||
|
||||
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.
|
||||
pub fn format_compaction_prompt() -> CompactionPrompt {
|
||||
CompactionPrompt {
|
||||
system: include_str!("../templates/intra_compaction_system.txt").to_string(),
|
||||
user: include_str!("../templates/intra_compaction_user.txt").to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn templates_are_non_empty() {
|
||||
let p = format_compaction_prompt();
|
||||
assert!(!p.system.trim().is_empty(), "system prompt empty");
|
||||
assert!(!p.user.trim().is_empty(), "user prompt empty");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
Your task is to create a detailed summary of the Grok Chat conversation so far, paying close attention to the user's explicit requests and all previous actions as Grok (built by xAI).
|
||||
This summary should be thorough in capturing technical details, code patterns, architectural decisions, tool chains, and verification steps that would be essential for continuing development, research, or complex tasks without losing context.
|
||||
|
||||
Important Clarification on Terminology (Broad File Definition):
|
||||
Throughout this prompt, the term "file", "files", "file IDs", "file names", and "Files and Code Sections / Artifacts" are defined broadly. They explicitly include:
|
||||
- Regular files and code files
|
||||
- Attachments
|
||||
- Images (uploaded images, generated images, viewed images, etc.)
|
||||
- rendered image or content outputs
|
||||
- Any other file-like content, media objects, visual artifacts, or structured content blocks that have appeared in the conversation history (including but not limited to uploads, generations, render components, or persistent references).
|
||||
|
||||
Only include information that is visible in the direct user-Grok conversation history (user messages + Grok's responses, reasoning, tool calls, and tool outputs). Do not include any internal team communication, chatroom messages, or multi-agent interactions.
|
||||
|
||||
Use your internal thinking channel to chronologically analyze each message and section of the conversation before producing the final summary, and ensure you've covered all necessary points.
|
||||
During that analysis, thoroughly identify:
|
||||
- The user's explicit requests and evolving intents
|
||||
- Grok's approach to addressing them: reasoning steps, specific tool calls (including parallel calls), parameters, results, and how they were interpreted/synthesized (truth-seeking emphasis)
|
||||
- Key decisions, technical concepts, code patterns, and architectural choices
|
||||
- Specific details like:
|
||||
- file IDs, attachment IDs, image references/URLs, render_result IDs (if any)
|
||||
- file names, attachment names, image captions/descriptions, render component details (if any)
|
||||
- full code snippets (especially recent ones or those executed in REPL)
|
||||
- function signatures
|
||||
- file edits / diffs
|
||||
- tool call details (e.g., code_execution snippets, web_search queries, browse_page instructions, X search operators)
|
||||
- render components used (if any)
|
||||
- Errors encountered (tool failures, code exec errors, search limitations, reasoning issues) and how they were diagnosed/fixed
|
||||
- Pay special attention to specific user feedback, especially if the user told you to do something differently, corrected facts, or changed direction.
|
||||
|
||||
Double-check for technical accuracy and completeness, addressing each required element thoroughly.
|
||||
|
||||
Your final summary must contain the following sections, in order:
|
||||
|
||||
1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail, including any evolution over the conversation.
|
||||
|
||||
2. Key Technical Concepts: List all important technical concepts, technologies, frameworks, and Grok-specific tool patterns discussed.
|
||||
|
||||
3. Tool Usage & Verification: Summarize significant tool calls (code_execution REPL state, web_search, browse_page, X tools, etc.), key information retrieved/verified, cross-referencing steps, and how they influenced decisions or responses.
|
||||
|
||||
4. Files, Attachments, Images, Render Results & Code Artifacts: Enumerate all specific file-like artifacts (broadly defined as above: files, attachments, images, render_results, etc.), code sections, or REPL executions examined, modified, or created. Pay special attention to the most recent messages and include full code snippets, image descriptions/references, render outputs, or attachment details where applicable, plus a summary of why this artifact is important for continuation.
|
||||
|
||||
5. Errors and Fixes: List all errors encountered (tool-related or otherwise), how you fixed them, and specific user feedback (especially "do something differently").
|
||||
|
||||
6. Problem Solving: Document problems solved, tool-assisted solutions, and any ongoing troubleshooting efforts.
|
||||
|
||||
7. All User Messages: List ALL user messages that are not tool results (verbatim or high-fidelity summary). These are critical for understanding feedback and intent changes.
|
||||
|
||||
Here's an example of how your output should be structured:
|
||||
|
||||
<example>
|
||||
1. Primary Request and Intent:
|
||||
[Detailed description]
|
||||
|
||||
2. Key Technical Concepts:
|
||||
- [Concept 1]
|
||||
- [Concept 2]
|
||||
- [...]
|
||||
|
||||
3. Tool Usage & Verification:
|
||||
- [Key tool calls and verification steps]
|
||||
|
||||
4. Files, Attachments, Images, Render Results & Code Artifacts:
|
||||
- [Artifact 1 (broadly defined file/attachment/image/render etc.)]
|
||||
- [file/attachment/image/render name and ID]
|
||||
- [Summary of importance]
|
||||
- [Changes or execution results]
|
||||
- [Important Code Snippet / Image reference / Render details]
|
||||
|
||||
5. Errors and Fixes:
|
||||
- [Error 1]: [How fixed] [User feedback]
|
||||
|
||||
6. Problem Solving:
|
||||
[Description]
|
||||
|
||||
7. All User Messages:
|
||||
- [Detailed non tool use user message]
|
||||
- [...]
|
||||
</example>
|
||||
|
||||
Output the summary directly using the section headings above. Do not wrap the output in any XML tags or other markup — emit the seven sections as plain text.
|
||||
|
||||
There may be additional summarization instructions provided in the included context. If so, follow these instructions when creating the above summary. Examples of instructions include:
|
||||
<example>
|
||||
## Compact Instructions
|
||||
When summarizing focus on tool outputs, REPL state, code changes, test results, and recent user feedback/corrections. Include critical code snippets and tool calls verbatim.
|
||||
</example>
|
||||
|
||||
<example>
|
||||
# Summary instructions
|
||||
When using compact mode — prioritize most recent tool results, executed code diffs, and exact user instructions on direction changes.
|
||||
</example>
|
||||
@@ -0,0 +1,91 @@
|
||||
Your task is to create a detailed summary of the Grok Chat conversation so far, paying close attention to the user's explicit requests and all previous actions as Grok (built by xAI).
|
||||
This summary should be thorough in capturing technical details, code patterns, architectural decisions, tool chains, and verification steps that would be essential for continuing development, research, or complex tasks without losing context.
|
||||
|
||||
Important Clarification on Terminology (Broad File Definition):
|
||||
Throughout this prompt, the term "file", "files", "file IDs", "file names", and "Files and Code Sections / Artifacts" are defined broadly. They explicitly include:
|
||||
- Regular files and code files
|
||||
- Attachments
|
||||
- Images (uploaded images, generated images, viewed images, etc.)
|
||||
- rendered image or content outputs
|
||||
- Any other file-like content, media objects, visual artifacts, or structured content blocks that have appeared in the conversation history (including but not limited to uploads, generations, render components, or persistent references).
|
||||
|
||||
Only include information that is visible in the direct user-Grok conversation history (user messages + Grok's responses, reasoning, tool calls, and tool outputs). Do not include any internal team communication, chatroom messages, or multi-agent interactions.
|
||||
|
||||
Use your internal thinking channel to chronologically analyze each message and section of the conversation before producing the final summary, and ensure you've covered all necessary points.
|
||||
During that analysis, thoroughly identify:
|
||||
- The user's explicit requests and evolving intents
|
||||
- Grok's approach to addressing them: reasoning steps, specific tool calls (including parallel calls), parameters, results, and how they were interpreted/synthesized (truth-seeking emphasis)
|
||||
- Key decisions, technical concepts, code patterns, and architectural choices
|
||||
- Specific details like:
|
||||
- file IDs, attachment IDs, image references/URLs, render_result IDs (if any)
|
||||
- file names, attachment names, image captions/descriptions, render component details (if any)
|
||||
- full code snippets (especially recent ones or those executed in REPL)
|
||||
- function signatures
|
||||
- file edits / diffs
|
||||
- tool call details (e.g., code_execution snippets, web_search queries, browse_page instructions, X search operators)
|
||||
- render components used (if any)
|
||||
- Errors encountered (tool failures, code exec errors, search limitations, reasoning issues) and how they were diagnosed/fixed
|
||||
- Pay special attention to specific user feedback, especially if the user told you to do something differently, corrected facts, or changed direction.
|
||||
|
||||
Double-check for technical accuracy and completeness, addressing each required element thoroughly.
|
||||
|
||||
Your final summary must contain the following sections, in order:
|
||||
|
||||
1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail, including any evolution over the conversation.
|
||||
|
||||
2. Key Technical Concepts: List all important technical concepts, technologies, frameworks, and Grok-specific tool patterns discussed.
|
||||
|
||||
3. Tool Usage & Verification: Summarize significant tool calls (code_execution REPL state, web_search, browse_page, X tools, etc.), key information retrieved/verified, cross-referencing steps, and how they influenced decisions or responses.
|
||||
|
||||
4. Files, Attachments, Images, Render Results & Code Artifacts: Enumerate all specific file-like artifacts (broadly defined as above: files, attachments, images, render_results, etc.), code sections, or REPL executions examined, modified, or created. Pay special attention to the most recent messages and include full code snippets, image descriptions/references, render outputs, or attachment details where applicable, plus a summary of why this artifact is important for continuation.
|
||||
|
||||
5. Errors and Fixes: List all errors encountered (tool-related or otherwise), how you fixed them, and specific user feedback (especially "do something differently").
|
||||
|
||||
6. Problem Solving: Document problems solved, tool-assisted solutions, and any ongoing troubleshooting efforts.
|
||||
|
||||
7. All User Messages: List ALL user messages that are not tool results (verbatim or high-fidelity summary). These are critical for understanding feedback and intent changes.
|
||||
|
||||
Here's an example of how your output should be structured:
|
||||
|
||||
<example>
|
||||
1. Primary Request and Intent:
|
||||
[Detailed description]
|
||||
|
||||
2. Key Technical Concepts:
|
||||
- [Concept 1]
|
||||
- [Concept 2]
|
||||
- [...]
|
||||
|
||||
3. Tool Usage & Verification:
|
||||
- [Key tool calls and verification steps]
|
||||
|
||||
4. Files, Attachments, Images, Render Results & Code Artifacts:
|
||||
- [Artifact 1 (broadly defined file/attachment/image/render etc.)]
|
||||
- [file/attachment/image/render name and ID]
|
||||
- [Summary of importance]
|
||||
- [Changes or execution results]
|
||||
- [Important Code Snippet / Image reference / Render details]
|
||||
|
||||
5. Errors and Fixes:
|
||||
- [Error 1]: [How fixed] [User feedback]
|
||||
|
||||
6. Problem Solving:
|
||||
[Description]
|
||||
|
||||
7. All User Messages:
|
||||
- [Detailed non tool use user message]
|
||||
- [...]
|
||||
</example>
|
||||
|
||||
Output the summary directly using the section headings above. Do not wrap the output in any XML tags or other markup — emit the seven sections as plain text.
|
||||
|
||||
There may be additional summarization instructions provided in the included context. If so, follow these instructions when creating the above summary. Examples of instructions include:
|
||||
<example>
|
||||
## Compact Instructions
|
||||
When summarizing focus on tool outputs, REPL state, code changes, test results, and recent user feedback/corrections. Include critical code snippets and tool calls verbatim.
|
||||
</example>
|
||||
|
||||
<example>
|
||||
# Summary instructions
|
||||
When using compact mode — prioritize most recent tool results, executed code diffs, and exact user instructions on direction changes.
|
||||
</example>
|
||||
@@ -0,0 +1,3 @@
|
||||
You are summarizing the tool-call history of an AI assistant that is partway through answering a user's question.
|
||||
|
||||
The assistant has made several tool calls (web searches, file reads, code execution, etc.) and accumulated tool results that are now taking up too much context window space. Your summary will replace those tool calls + results, so the assistant can continue its work with the same effective knowledge but less context overhead.
|
||||
@@ -0,0 +1,67 @@
|
||||
Your task is to create a detailed summary of the tool-call history above, paying close attention to preserving all information the assistant needs to continue its current task without losing context.
|
||||
This summary should be thorough in capturing technical details, code patterns, data points, and intermediate results that would be essential for continuing the current work.
|
||||
|
||||
CRITICAL: If the tool-call history contains a previous compaction summary (marked with "ConversationCompaction" or similar markers), you MUST incorporate ALL information from that previous summary into your new summary. Previous summaries contain essential context from earlier steps that would otherwise be lost.
|
||||
|
||||
Use your internal thinking channel to chronologically review each tool call and its result before producing the final summary, and ensure you've covered all necessary points.
|
||||
During that analysis, thoroughly identify:
|
||||
- What was searched, read, or executed and why
|
||||
- Key findings, data points, and outcomes
|
||||
- Specific details like:
|
||||
- file paths, URLs, IDs, error messages
|
||||
- full code snippets (especially recent ones)
|
||||
- function signatures and configuration details
|
||||
- tool call parameters and results
|
||||
- Errors encountered and how they were resolved
|
||||
- Double-check for completeness — every piece of data the assistant gathered must be preserved.
|
||||
|
||||
Your final summary must contain the following sections, in order:
|
||||
|
||||
1. Task and Intent: What the assistant is trying to accomplish for the user, including the current sub-goal.
|
||||
|
||||
2. Key Findings: Facts, data points (numbers, dates, IDs, URLs), schema details, and any other information gathered from tool calls. Preserve specific data verbatim.
|
||||
|
||||
3. Files and Code: Enumerate specific file paths examined, modified, or created. Include key code snippets, function signatures, and configuration details verbatim, plus a summary of why each file is important.
|
||||
|
||||
4. Errors and Fixes: All errors encountered, how each was resolved, including specific error messages verbatim.
|
||||
|
||||
5. Actions Taken: Successful modifications, commands run, and their outcomes.
|
||||
|
||||
6. Current Progress: What has been completed and what remains to be done.
|
||||
|
||||
Here's an example of how your output should be structured:
|
||||
|
||||
<example>
|
||||
1. Task and Intent:
|
||||
[Detailed description of what the assistant is working on]
|
||||
|
||||
2. Key Findings:
|
||||
- [Finding 1 with specific data verbatim]
|
||||
- [Finding 2]
|
||||
- [...]
|
||||
|
||||
3. Files and Code:
|
||||
- [file path 1]
|
||||
- [Summary of importance]
|
||||
- [Key code snippet or changes]
|
||||
- [file path 2]
|
||||
- [...]
|
||||
|
||||
4. Errors and Fixes:
|
||||
- [Error message verbatim]: [How fixed]
|
||||
|
||||
5. Actions Taken:
|
||||
- [Action 1]: [Outcome]
|
||||
- [...]
|
||||
|
||||
6. Current Progress:
|
||||
[What is done, what remains]
|
||||
</example>
|
||||
|
||||
Output the summary directly using the section headings above. Do not wrap the output in any XML tags or other markup — emit the six sections as plain text.
|
||||
|
||||
IMPORTANT:
|
||||
- Do NOT call any tools. Output the summary text only.
|
||||
- Preserve specific data verbatim — URLs, file paths, code snippets, error messages, ID strings.
|
||||
- Write in the same language as the conversation. If the conversation is primarily in Chinese, write the summary in Chinese (keep technical terms, file paths, and code in English).
|
||||
- Do not invent information that is not in the tool-call history.
|
||||
@@ -0,0 +1,24 @@
|
||||
//! Token-count seam.
|
||||
//!
|
||||
//! Budgeting math in the shared engine needs a *trusted* token count, but the
|
||||
//! two harnesses disagree on how to produce one:
|
||||
//!
|
||||
//! - Grok chat has a real tokenizer (`TextTokenizer` / `ImageTokenizer`) and
|
||||
//! counts whole turns via `GrokTurn::get_num_tokens`.
|
||||
//! - grok-build 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.
|
||||
//!
|
||||
//! There is intentionally **no** blanket `Arc` forwarding here: each harness
|
||||
//! implements the counter directly for the item type its algorithms run on
|
||||
//! (Grok chat: `ItemTokenCounter<Arc<GrokTurn>>`), 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;
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-computer-hub-core"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Transport, ToolRegistry, and resolver abstractions for the xAI Computer Hub"
|
||||
|
||||
[dependencies]
|
||||
async-trait = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
kigi-tool-protocol = { workspace = true }
|
||||
kigi-tool-runtime = { workspace = true }
|
||||
kigi-tool-types = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
dashmap = { workspace = true }
|
||||
schemars = { workspace = true }
|
||||
tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros", "test-util", "sync"] }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,73 @@
|
||||
//! `InnerDispatchForResolver` — an object-safe `ToolDispatch` that routes
|
||||
//! through a `Weak<CompoundResolver>` bound to a single session.
|
||||
//!
|
||||
//! Tools that need to call other tools (the inner-dispatch pattern) ask
|
||||
//! the runtime for an `Arc<dyn ToolDispatch>`. This adapter answers that
|
||||
//! question with a resolver-backed implementation. Holding the resolver
|
||||
//! by [`Weak`] lets the router own the resolver while inner-dispatch
|
||||
//! handles created from the same resolver release naturally when the
|
||||
//! router is torn down.
|
||||
|
||||
use std::sync::Weak;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use kigi_tool_protocol::{SessionId, ToolId};
|
||||
use kigi_tool_runtime::{
|
||||
ToolCallContext, ToolDispatch, ToolError, ToolStream, TypedToolOutput, terminal_only,
|
||||
};
|
||||
|
||||
use crate::resolver::CompoundResolver;
|
||||
|
||||
/// Resolver-backed `ToolDispatch` implementation.
|
||||
///
|
||||
/// The resolver is held by [`Weak`] so the inner-dispatch handle never
|
||||
/// keeps the router alive past its natural lifetime — when the owning
|
||||
/// router drops the resolver, in-flight inner calls fail cleanly with
|
||||
/// [`ToolError::Custom`] keyed `computer_hub_dropped`.
|
||||
///
|
||||
/// Bound to a single [`SessionId`] at construction (rather than reading a
|
||||
/// session from [`ToolCallContext`]) so the inner-dispatch path mirrors
|
||||
/// the per-session lifetime of the outer router.
|
||||
#[derive(Debug)]
|
||||
pub struct InnerDispatchForResolver {
|
||||
resolver: Weak<CompoundResolver>,
|
||||
session_id: SessionId,
|
||||
}
|
||||
|
||||
impl InnerDispatchForResolver {
|
||||
/// Build an inner-dispatch handle bound to `session_id`, resolving
|
||||
/// through `resolver`.
|
||||
pub fn new(resolver: Weak<CompoundResolver>, session_id: SessionId) -> Self {
|
||||
Self {
|
||||
resolver,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrow the bound session identifier.
|
||||
pub fn session_id(&self) -> &SessionId {
|
||||
&self.session_id
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolDispatch for InnerDispatchForResolver {
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
let Some(resolver) = self.resolver.upgrade() else {
|
||||
return terminal_only(Err(ToolError::custom(
|
||||
"computer_hub_dropped",
|
||||
"computer hub dropped before inner call could execute",
|
||||
)));
|
||||
};
|
||||
resolver
|
||||
.resolve_and_dispatch(&self.session_id, tool_id, args, ctx)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
//! xAI Computer Hub — transport + registry + resolver core.
|
||||
//!
|
||||
//! Object-safe abstractions used by every router build: a [`Transport`]
|
||||
//! that authorises and dispatches calls, a [`ToolRegistry`] trait shared
|
||||
//! by both storage planes, a [`CompoundResolver`] that applies the
|
||||
//! local-shadows-remote rule, and the local + remote transports plus
|
||||
//! inner-dispatch glue that sit on top.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
pub mod inner;
|
||||
pub mod local;
|
||||
pub mod registry;
|
||||
pub mod remote;
|
||||
pub mod resolver;
|
||||
pub mod transport;
|
||||
|
||||
pub use inner::InnerDispatchForResolver;
|
||||
pub use local::{LOCAL_INVOKE_SCOPE, LocalTransport};
|
||||
pub use registry::{
|
||||
ConnectionCleanupReport, SessionCleanupReport, ToolRegistry, ToolSessionBindOutcome,
|
||||
ToolSessionUnbindOutcome,
|
||||
};
|
||||
pub use remote::{
|
||||
ConnectionClient, RemoteToolProxy, RemoteTransport, decode_call_result, error_from_envelope,
|
||||
is_workspace_unavailable, output_to_value, progress_from_frame, tool_error_from_wire,
|
||||
};
|
||||
pub use resolver::{CompoundResolver, ErasedTool, ResolvedTool, ToolHandle};
|
||||
pub use transport::{Principal, Transport, TransportKind};
|
||||
@@ -0,0 +1,77 @@
|
||||
//! In-process transport that resolves through a [`CompoundResolver`].
|
||||
//!
|
||||
//! `LocalTransport` is bound to a single `(user_id, session_id)` at
|
||||
//! construction. Authorisation returns a principal pre-populated with the
|
||||
//! bound session and the `tool.invoke` scope; per-call dispatch resolves
|
||||
//! against the bound session's view of the resolver.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use kigi_tool_protocol::{SessionId, ToolId, UserId};
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput};
|
||||
|
||||
use crate::resolver::CompoundResolver;
|
||||
use crate::transport::{Principal, Transport, TransportKind};
|
||||
|
||||
/// The scope `LocalTransport::authorize` grants to its principal.
|
||||
///
|
||||
/// Hoisted so adapters that authorise principals through other paths
|
||||
/// can match the local convention without restating the literal.
|
||||
pub const LOCAL_INVOKE_SCOPE: &str = "tool.invoke";
|
||||
|
||||
/// Transport that dispatches against an in-process resolver.
|
||||
#[derive(Debug)]
|
||||
pub struct LocalTransport {
|
||||
resolver: Arc<CompoundResolver>,
|
||||
user_id: UserId,
|
||||
session_id: SessionId,
|
||||
}
|
||||
|
||||
impl LocalTransport {
|
||||
/// Build a transport bound to `(user_id, session_id)` and resolving
|
||||
/// through `resolver`.
|
||||
pub fn new(resolver: Arc<CompoundResolver>, user_id: UserId, session_id: SessionId) -> Self {
|
||||
Self {
|
||||
resolver,
|
||||
user_id,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Bound user identity for this transport.
|
||||
pub fn user_id(&self) -> &UserId {
|
||||
&self.user_id
|
||||
}
|
||||
|
||||
/// Bound session for this transport.
|
||||
pub fn session_id(&self) -> &SessionId {
|
||||
&self.session_id
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for LocalTransport {
|
||||
fn kind(&self) -> TransportKind {
|
||||
TransportKind::Local
|
||||
}
|
||||
|
||||
async fn authorize(&self) -> Result<Principal, ToolError> {
|
||||
Ok(Principal::new(self.user_id.clone())
|
||||
.with_session(self.session_id.clone())
|
||||
.with_scope(LOCAL_INVOKE_SCOPE))
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
self.resolver
|
||||
.resolve_and_dispatch(&self.session_id, tool_id, args, ctx)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
//! Object-safe `ToolRegistry` trait shared by every storage plane.
|
||||
//!
|
||||
//! Two registry implementations are expected: one in-memory plane for
|
||||
//! statically-registered local tools, and one connection-keyed plane fed by
|
||||
//! incoming remote registrations. Both expose the same trait so the
|
||||
//! router can compose them through [`crate::CompoundResolver`] without
|
||||
//! caring which is which.
|
||||
//!
|
||||
//! Mutations are connection-scoped: each registered tool belongs to the
|
||||
//! [`ConnectionId`] that introduced it. Per-tool session bindings live
|
||||
//! alongside the tool's record and are mutated independently via
|
||||
//! [`ToolRegistry::bind_tool_session`] / [`ToolRegistry::unbind_tool_session`].
|
||||
//! Reads (`find_tool`, `list_tools`, `search`) remain session-scoped — the
|
||||
//! router resolves a tool by `(session_id, tool_id)`, never by
|
||||
//! connection id.
|
||||
//!
|
||||
//! The concrete in-memory implementation is intentionally **out of scope**
|
||||
//! for this crate — it requires a concurrency story (sharded maps, an
|
||||
//! actor, etc.) that belongs alongside the registry's collision matrix and
|
||||
//! generation handling. Tests exercise the trait via per-test mock impls.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use kigi_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{SearchSnapshot, ServerSummary};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
use crate::resolver::ResolvedTool;
|
||||
|
||||
/// Outcome of a single [`ToolRegistry::bind_tool_session`] call.
|
||||
///
|
||||
/// This enum is the source of truth for storage outcomes; the wire enum
|
||||
/// [`kigi_tool_protocol::ToolSessionBindOutcome`] is a strict subset with
|
||||
/// one extra wire-only variant. The two layers diverge deliberately:
|
||||
///
|
||||
/// - `Conflict` (cross-connection race on the `(session_id, tool_id)`
|
||||
/// reverse-index slot) is registry-internal: the router lifts it to a
|
||||
/// top-level `ServerError::ToolBindingConflict` (-32600) instead of
|
||||
/// mirroring it to the wire ack, so the contended caller gets a
|
||||
/// dedicated error code rather than overloading `UnknownTool`.
|
||||
/// - The wire enum's `SessionNotBound` is router-injected by the
|
||||
/// per-frame envelope pre-check (the connection's bound-session set
|
||||
/// lives in router state, not the registry) and is never produced by
|
||||
/// any registry call — so it has no counterpart here.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolSessionBindOutcome {
|
||||
/// Added to the tool's session set.
|
||||
Bound,
|
||||
/// Session id was already in the tool's session set; no-op.
|
||||
AlreadyBound,
|
||||
/// No tool with the given id is registered against this connection.
|
||||
UnknownTool,
|
||||
/// Cross-connection conflict: another connection already holds the
|
||||
/// `(session_id, tool_id)` reverse-index slot. The router lifts this
|
||||
/// into a top-level `ToolBindingConflict` server error so the wire
|
||||
/// reply uses the dedicated -32600 code instead of the structurally
|
||||
/// dishonest `UnknownTool`. No registry state was mutated.
|
||||
Conflict,
|
||||
}
|
||||
|
||||
/// Outcome of a single [`ToolRegistry::unbind_tool_session`] call.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ToolSessionUnbindOutcome {
|
||||
/// Removed from the tool's session set.
|
||||
Unbound,
|
||||
/// Session id was not in the tool's session set; no-op.
|
||||
NotBound,
|
||||
/// No tool with the given id is registered against this connection.
|
||||
UnknownTool,
|
||||
}
|
||||
|
||||
/// Aggregated summary of a connection-scoped cleanup pass.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct ConnectionCleanupReport {
|
||||
/// Number of distinct `(connection, tool_id)` records dropped.
|
||||
pub tools_dropped: usize,
|
||||
/// Number of reverse-index `(session_id, tool_id)` rows cleaned up
|
||||
/// across every session the dropped tools were bound to.
|
||||
pub session_bindings_cleared: usize,
|
||||
}
|
||||
|
||||
/// Aggregated summary of a session-scoped cleanup pass.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct SessionCleanupReport {
|
||||
/// Number of tools whose session set lost the unregistered session id.
|
||||
pub tools_touched: usize,
|
||||
/// Number of tools whose session set became empty after the
|
||||
/// unregistration. The tool record itself is NOT removed — the owning
|
||||
/// connection still owns it and may rebind via
|
||||
/// [`ToolRegistry::bind_tool_session`] later.
|
||||
pub tools_left_orphaned: usize,
|
||||
}
|
||||
|
||||
/// Backend-agnostic registry of tools available within a router.
|
||||
///
|
||||
/// Methods are split into mutating (`async fn` — registration changes may
|
||||
/// touch shared state and require coordination) and read-only views
|
||||
/// (synchronous — implementations should answer from a consistent snapshot
|
||||
/// without awaiting). The split mirrors how callers use the registry: the
|
||||
/// hot path is the `find_tool` / `list_tools` view; mutations happen on the
|
||||
/// rarer registration boundary.
|
||||
#[async_trait]
|
||||
pub trait ToolRegistry: Send + Sync + std::fmt::Debug {
|
||||
/// Register a single tool against `connection_id`.
|
||||
///
|
||||
/// The outcome reports whether the registration created a new entry,
|
||||
/// updated an existing one, was shadowed by a higher-priority
|
||||
/// registration, or was rejected. `reg.sessions` may be empty — the
|
||||
/// tool is registered but unreachable until
|
||||
/// [`Self::bind_tool_session`] adds at least one session binding.
|
||||
/// Implementations must enforce per-`(connection_id, tool_id)`
|
||||
/// uniqueness within their plane.
|
||||
async fn register_tool(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolRegistration,
|
||||
) -> RegistrationOutcome;
|
||||
|
||||
/// Register a multi-tool batch from a single tool server against
|
||||
/// `connection_id`.
|
||||
///
|
||||
/// Returns one [`RegistrationOutcome`] per tool in input order. Batch
|
||||
/// semantics are best-effort: per-tool failures do not abort the rest
|
||||
/// of the batch. The whole batch shares `reg.sessions` (which may be
|
||||
/// empty).
|
||||
async fn register_server(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome>;
|
||||
|
||||
/// Drop the tool registered under `(connection_id, tool_id)`. Returns
|
||||
/// `true` if a matching entry was removed, `false` if no such entry
|
||||
/// existed. The tool is removed from every session it was bound to in
|
||||
/// one shot — use [`Self::unbind_tool_session`] for per-session removal.
|
||||
async fn unregister_tool(&self, connection_id: &ConnectionId, tool: &ToolId) -> bool;
|
||||
|
||||
/// Drop every tool registered by `connection_id` under `server_id`.
|
||||
/// Returns the number of entries removed.
|
||||
async fn unregister_server(&self, connection_id: &ConnectionId, server: &ServerId) -> usize;
|
||||
|
||||
/// Add `session_id` to the per-tool session set of
|
||||
/// `(connection_id, tool_id)`. The caller (typically the WebSocket
|
||||
/// router) is responsible for verifying that `session_id` is in the
|
||||
/// connection's bound-session set before calling this method.
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome;
|
||||
|
||||
/// Remove `session_id` from the per-tool session set of
|
||||
/// `(connection_id, tool_id)`. Does not unregister the tool itself.
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome;
|
||||
|
||||
/// Drop every tool registered by `connection_id`. Used by the WebSocket
|
||||
/// transport on disconnect cleanup. Returns counters describing how
|
||||
/// much state was released.
|
||||
async fn drop_connection(&self, connection_id: &ConnectionId) -> ConnectionCleanupReport;
|
||||
|
||||
/// Look up the active resolution for `(session, tool)`.
|
||||
///
|
||||
/// Returns `None` when no entry exists or when an entry exists but is
|
||||
/// shadowed. A shadowed entry is never returned — the caller sees only
|
||||
/// the active resolution.
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool>;
|
||||
|
||||
/// Enumerate every active tool description for `session`, filtered by
|
||||
/// the requested presentation `mode`. Implementations decide how to
|
||||
/// honour the mode (e.g. omit non-meta tools when `Concise` is set).
|
||||
fn list_tools(&self, session: &SessionId, mode: &ToolDefinitionMode) -> Vec<ToolDescription>;
|
||||
|
||||
/// Enumerate active server summaries for `session`. Useful for
|
||||
/// rendering connected-integrations system reminders.
|
||||
fn list_servers(&self, session: &SessionId) -> Vec<ServerSummary>;
|
||||
|
||||
/// Run a search query against the registry's index for `session`.
|
||||
/// `limit` caps the result count; the snapshot reports how many
|
||||
/// matches were hidden by the cap.
|
||||
fn search(&self, session: &SessionId, query: &str, limit: usize) -> SearchSnapshot;
|
||||
|
||||
/// Drop the binding to `session` from every tool that has it. The
|
||||
/// affected tool records are NOT removed — their owning connection
|
||||
/// retains them and may rebind via [`Self::bind_tool_session`]. Called
|
||||
/// by the WebSocket transport when a session ends globally (no peer
|
||||
/// connection still holds the binding) and by the connection actor
|
||||
/// during per-disconnect cleanup.
|
||||
async fn unregister_session(&self, session: &SessionId) -> SessionCleanupReport;
|
||||
|
||||
/// Helper: set of session ids currently bound to `(connection_id, tool_id)`.
|
||||
/// Returns an empty set when the tool is not registered. Mainly used
|
||||
/// by tests to assert per-tool session set invariants without leaning
|
||||
/// on the reverse index.
|
||||
fn tool_sessions(&self, connection_id: &ConnectionId, tool: &ToolId) -> HashSet<SessionId>;
|
||||
|
||||
/// All servers registered by this user across all connections.
|
||||
fn list_servers_for_user(&self, user_id: &UserId) -> Vec<ServerRecord>;
|
||||
|
||||
/// Look up a server by its connection ID.
|
||||
fn get_server_record(&self, connection_id: &ConnectionId) -> Option<ServerRecord>;
|
||||
|
||||
/// Look up only a server's id by its connection ID. Lighter than
|
||||
/// [`Self::get_server_record`] for callers that need nothing else:
|
||||
/// implementations should override the default to avoid deep-cloning
|
||||
/// the whole record (notably its `metadata` JSON).
|
||||
fn get_server_id(&self, connection_id: &ConnectionId) -> Option<ServerId> {
|
||||
self.get_server_record(connection_id)
|
||||
.map(|record| record.server_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Server identity captured at `register_server` time.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ServerRecord {
|
||||
pub connection_id: ConnectionId,
|
||||
pub user_id: UserId,
|
||||
pub server_id: ServerId,
|
||||
pub description: String,
|
||||
pub metadata: serde_json::Value,
|
||||
pub registered_at: chrono::DateTime<chrono::Utc>,
|
||||
/// Monotonic registration stamp ([`next_registration_seq`]) — the
|
||||
/// stale-vs-revived discriminator for newest-wins (`registered_at` is display-only).
|
||||
pub registration_seq: u64,
|
||||
}
|
||||
|
||||
/// Process-global hybrid logical clock: per-process strictly-increasing (no ties,
|
||||
/// immune to NTP step-back) and epoch-seeded so stamps also roughly order across
|
||||
/// replicas — only while inter-replica clock skew stays within the revive window
|
||||
/// (`tool_route_ttl_ms`); past that, TTL eviction, not seq order, is the backstop.
|
||||
/// The recency key for bind newest-wins and strictly-older eviction.
|
||||
static REGISTRATION_CLOCK: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Issue the next monotonic registration stamp. See [`REGISTRATION_CLOCK`].
|
||||
pub fn next_registration_seq() -> u64 {
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
let candidate = now_ms << 10;
|
||||
let mut prev = REGISTRATION_CLOCK.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let next = candidate.max(prev + 1);
|
||||
match REGISTRATION_CLOCK.compare_exchange_weak(
|
||||
prev,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return next,
|
||||
Err(actual) => prev = actual,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod seq_tests {
|
||||
use super::next_registration_seq;
|
||||
|
||||
fn now_ms() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn next_registration_seq_is_monotonic_and_epoch_seeded_under_burst() {
|
||||
let before_ms = now_ms();
|
||||
let first = next_registration_seq();
|
||||
let mut prev = first;
|
||||
const N: u64 = 50_000;
|
||||
for _ in 0..N {
|
||||
let s = next_registration_seq();
|
||||
assert!(s > prev, "must be strictly increasing: {prev} -> {s}");
|
||||
prev = s;
|
||||
}
|
||||
let after_ms = now_ms();
|
||||
|
||||
assert!(
|
||||
prev - first >= N,
|
||||
"burst must advance by at least one per call: {first} -> {prev}",
|
||||
);
|
||||
let high = prev >> 10;
|
||||
assert!(
|
||||
high >= before_ms,
|
||||
"high bits ({high}) must be epoch-seeded (>= {before_ms})",
|
||||
);
|
||||
assert!(
|
||||
high <= after_ms + 1_000,
|
||||
"high bits ({high}) must track wall clock (<= {after_ms} + slack)",
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,541 @@
|
||||
//! `ConnectionClient` abstraction, `RemoteToolProxy`, and
|
||||
//! `RemoteTransport`.
|
||||
//!
|
||||
//! `ConnectionClient` is the thin contract a downstream WebSocket SDK (or
|
||||
//! an in-test channel-backed mock) implements; this crate stays free of
|
||||
//! tokio-runtime / tokio-tungstenite deps so callers can pick their own.
|
||||
//!
|
||||
//! `RemoteToolProxy` wraps a remote tool registration so it implements
|
||||
//! [`ToolHandle`] — the router routes through the same handle
|
||||
//! type for local and remote registrations. `RemoteTransport` is the
|
||||
//! transport-side equivalent: it forwards arbitrary `(tool_id, args)`
|
||||
//! pairs over a [`ConnectionClient`] without needing a per-tool handle.
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::Stream;
|
||||
use futures::future::BoxFuture;
|
||||
use futures::stream::BoxStream;
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
use kigi_tool_protocol::{
|
||||
JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion, Method,
|
||||
ResponseOutcome, SessionId, ToolCallId, ToolCallParams, ToolCallProgressFrame, ToolCallResult,
|
||||
ToolCapabilities, ToolErrorWire, ToolId, ToolOutputWire, UserId, WORKSPACE_UNAVAILABLE_SUBCODE,
|
||||
};
|
||||
use kigi_tool_runtime::{
|
||||
BehaviorVersion, ContentBlock, Cwd, ListToolsContext, ToolCallContext,
|
||||
ToolChatCompletionResponse, ToolError, ToolErrorKind, ToolProgress, ToolStream, ToolStreamItem,
|
||||
TypedToolOutput, terminal_only,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
use crate::resolver::ToolHandle;
|
||||
use crate::transport::{Principal, Transport, TransportKind};
|
||||
|
||||
/// Object-safe contract for a connected remote endpoint.
|
||||
///
|
||||
/// Concrete implementations supply the wire transport — the Rust SDK uses
|
||||
/// `tokio_tungstenite`; tests use channel-backed mocks. Implementations
|
||||
/// are expected to:
|
||||
///
|
||||
/// - correlate request/response pairs by [`JsonRpcId`];
|
||||
/// - deliver progress notifications matching `tool_call_id` to whichever
|
||||
/// subscriber registered for them;
|
||||
/// - surface transport-level disconnects as [`ToolError::NetworkError`].
|
||||
#[async_trait]
|
||||
pub trait ConnectionClient: Send + Sync + std::fmt::Debug {
|
||||
/// Send a JSON-RPC request and await the matching response. Errors
|
||||
/// signal a transport-level failure (write failed, connection closed
|
||||
/// before the response arrived); a successful return carries the
|
||||
/// response envelope verbatim, including method-level error outcomes.
|
||||
async fn request(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, ToolError>;
|
||||
|
||||
/// Subscribe to progress notifications for `tool_call_id`.
|
||||
///
|
||||
/// The returned stream closes when the call's terminal frame arrives,
|
||||
/// when the connection drops, or when the caller drops the receiver.
|
||||
/// Subscribers MUST be registered before the corresponding request is
|
||||
/// sent — otherwise progress frames that arrive before subscription
|
||||
/// is complete are lost.
|
||||
async fn subscribe_progress(
|
||||
&self,
|
||||
tool_call_id: ToolCallId,
|
||||
) -> BoxStream<'static, ToolCallProgressFrame>;
|
||||
|
||||
/// Send a one-way notification (no response expected). Useful for
|
||||
/// hook frames such as cancel.
|
||||
async fn notify(&self, notification: JsonRpcNotification) -> Result<(), ToolError>;
|
||||
}
|
||||
|
||||
/// Wraps a remote registration so it dispatches through a connection.
|
||||
///
|
||||
/// Identity, description, and capabilities come from the registration
|
||||
/// snapshot held on the proxy; execution forwards a `tool_call_request`
|
||||
/// over the connection and merges progress + terminal frames into a
|
||||
/// single [`ToolStream`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RemoteToolProxy {
|
||||
tool_id: ToolId,
|
||||
session_id: SessionId,
|
||||
description: ToolDescription,
|
||||
capabilities: ToolCapabilities,
|
||||
connection: Arc<dyn ConnectionClient>,
|
||||
}
|
||||
|
||||
impl RemoteToolProxy {
|
||||
/// Build a proxy bound to a single remote registration.
|
||||
pub fn new(
|
||||
tool_id: ToolId,
|
||||
session_id: SessionId,
|
||||
description: ToolDescription,
|
||||
capabilities: ToolCapabilities,
|
||||
connection: Arc<dyn ConnectionClient>,
|
||||
) -> Self {
|
||||
Self {
|
||||
tool_id,
|
||||
session_id,
|
||||
description,
|
||||
capabilities,
|
||||
connection,
|
||||
}
|
||||
}
|
||||
|
||||
/// Bound session identifier.
|
||||
pub fn session_id(&self) -> &SessionId {
|
||||
&self.session_id
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolHandle for RemoteToolProxy {
|
||||
fn id(&self) -> ToolId {
|
||||
self.tool_id.clone()
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &ListToolsContext) -> ToolDescription {
|
||||
self.description.clone()
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ToolCapabilities {
|
||||
self.capabilities.clone()
|
||||
}
|
||||
|
||||
async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
dispatch_via_connection(
|
||||
Arc::clone(&self.connection),
|
||||
self.tool_id.clone(),
|
||||
self.session_id.clone(),
|
||||
args,
|
||||
ctx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// Transport that forwards calls over a [`ConnectionClient`].
|
||||
///
|
||||
/// The transport is bound to a single `(user_id, session_id)` at
|
||||
/// construction. Calls do not require a pre-built proxy — the transport
|
||||
/// builds the request frame from the `tool_id` it is asked to dispatch.
|
||||
#[derive(Debug)]
|
||||
pub struct RemoteTransport {
|
||||
connection: Arc<dyn ConnectionClient>,
|
||||
session_id: SessionId,
|
||||
user_id: UserId,
|
||||
}
|
||||
|
||||
impl RemoteTransport {
|
||||
/// Build a transport over `connection`, bound to `(user_id,
|
||||
/// session_id)`.
|
||||
pub fn new(
|
||||
connection: Arc<dyn ConnectionClient>,
|
||||
session_id: SessionId,
|
||||
user_id: UserId,
|
||||
) -> Self {
|
||||
Self {
|
||||
connection,
|
||||
session_id,
|
||||
user_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Bound session identifier.
|
||||
pub fn session_id(&self) -> &SessionId {
|
||||
&self.session_id
|
||||
}
|
||||
|
||||
/// Bound user identifier.
|
||||
pub fn user_id(&self) -> &UserId {
|
||||
&self.user_id
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for RemoteTransport {
|
||||
fn kind(&self) -> TransportKind {
|
||||
TransportKind::Remote
|
||||
}
|
||||
|
||||
async fn authorize(&self) -> Result<Principal, ToolError> {
|
||||
Ok(Principal::new(self.user_id.clone()).with_session(self.session_id.clone()))
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
dispatch_via_connection(
|
||||
Arc::clone(&self.connection),
|
||||
tool_id,
|
||||
self.session_id.clone(),
|
||||
args,
|
||||
ctx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
/// Subscribe to progress for `ctx.call_id`, send the `tool_call_request`,
|
||||
/// and return a stream interleaving progress frames with the eventual
|
||||
/// terminal item.
|
||||
///
|
||||
/// Subscribing **before** sending is the contract that
|
||||
/// [`ConnectionClient::subscribe_progress`] requires; doing so here keeps
|
||||
/// individual transports / proxies from re-implementing the dance.
|
||||
async fn dispatch_via_connection(
|
||||
connection: Arc<dyn ConnectionClient>,
|
||||
tool_id: ToolId,
|
||||
session_id: SessionId,
|
||||
arguments: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
let cwd = ctx
|
||||
.extensions
|
||||
.get::<Cwd>()
|
||||
.map(|c| c.0.to_string_lossy().into_owned());
|
||||
let behavior_version = ctx.extensions.get::<BehaviorVersion>().map(|v| v.0.clone());
|
||||
let call_id = ctx.call_id;
|
||||
|
||||
// Subscribe BEFORE sending. The single remaining `call_id.clone()`
|
||||
// is unavoidable: subscription needs an owned id and the same id has
|
||||
// to land in the request params below.
|
||||
let progress = connection.subscribe_progress(call_id.clone()).await;
|
||||
|
||||
let params = ToolCallParams {
|
||||
tool_call_id: call_id,
|
||||
tool_id,
|
||||
arguments,
|
||||
deadline_ms: None,
|
||||
behavior_version,
|
||||
cwd,
|
||||
// The ctx `TraceContext` extension is receive-side state.
|
||||
trace_context: None,
|
||||
};
|
||||
|
||||
let request = JsonRpcRequest {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: JsonRpcId::new_uuid_v7(),
|
||||
session_id: Some(session_id),
|
||||
method: Method::ToolCallRequest.as_wire_str().to_string(),
|
||||
params: match serde_json::to_value(¶ms) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
return terminal_only(Err(ToolError::custom("request_encoding", e.to_string())));
|
||||
}
|
||||
},
|
||||
};
|
||||
|
||||
// Build the response future without awaiting it here so progress and
|
||||
// terminal can be polled concurrently from the returned stream.
|
||||
let request_fut = Box::pin(async move { connection.request(request).await });
|
||||
|
||||
Box::pin(RequestStream {
|
||||
tool_id: Some(params.tool_id),
|
||||
progress,
|
||||
request: Some(request_fut),
|
||||
done: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Owned response future with `'static` lifetime so the stream can hold
|
||||
/// it across polls.
|
||||
type ResponseFuture = BoxFuture<'static, Result<JsonRpcResponse, ToolError>>;
|
||||
|
||||
/// Stream that interleaves wire-side progress frames with the eventual
|
||||
/// JSON-RPC response, ending with exactly one terminal item.
|
||||
struct RequestStream {
|
||||
/// Consumed exactly once when the terminal is built.
|
||||
tool_id: Option<ToolId>,
|
||||
progress: BoxStream<'static, ToolCallProgressFrame>,
|
||||
request: Option<ResponseFuture>,
|
||||
done: bool,
|
||||
}
|
||||
|
||||
impl Stream for RequestStream {
|
||||
type Item = ToolStreamItem<TypedToolOutput>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
if self.done {
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
|
||||
// Poll the response first so the terminal short-circuits the
|
||||
// moment it lands. Any progress frames that arrived alongside
|
||||
// the response are dropped — once `Terminal` is emitted, `done`
|
||||
// is set and the next poll returns `None` immediately without
|
||||
// re-polling the progress stream. The router invariant is
|
||||
// "`Progress* Terminal`, exactly one terminal"; dropping any
|
||||
// post-terminal progress is what makes that invariant hold here.
|
||||
if let Some(req_fut) = self.request.as_mut() {
|
||||
match req_fut.as_mut().poll(cx) {
|
||||
Poll::Ready(result) => {
|
||||
self.done = true;
|
||||
self.request = None;
|
||||
let Some(tool_id) = self.tool_id.take() else {
|
||||
return Poll::Ready(None);
|
||||
};
|
||||
let terminal = match result {
|
||||
Ok(resp) => terminal_from_response(tool_id, resp),
|
||||
Err(err) => Err(err),
|
||||
};
|
||||
return Poll::Ready(Some(ToolStreamItem::Terminal(terminal)));
|
||||
}
|
||||
Poll::Pending => {}
|
||||
}
|
||||
} else {
|
||||
self.done = true;
|
||||
return Poll::Ready(None);
|
||||
}
|
||||
|
||||
// Poll the progress stream while the request is pending. Closing
|
||||
// the progress stream is fine — the response future is still
|
||||
// registered for wake-up.
|
||||
match Pin::new(&mut self.progress).poll_next(cx) {
|
||||
Poll::Ready(Some(frame)) => {
|
||||
Poll::Ready(Some(ToolStreamItem::Progress(progress_from_frame(frame))))
|
||||
}
|
||||
Poll::Ready(None) | Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a wire-side [`ToolCallProgressFrame`] into a runtime
|
||||
/// [`ToolProgress`]. `kind` becomes the `Custom` subkind so callers can
|
||||
/// dispatch on the producer-defined identifier without losing the body.
|
||||
pub fn progress_from_frame(frame: ToolCallProgressFrame) -> ToolProgress {
|
||||
ToolProgress::Custom {
|
||||
subkind: frame.kind,
|
||||
payload: frame.body,
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode the response envelope into the terminal
|
||||
/// `Result<TypedToolOutput, _>` the runtime expects.
|
||||
fn terminal_from_response(
|
||||
tool_id: ToolId,
|
||||
resp: JsonRpcResponse,
|
||||
) -> Result<TypedToolOutput, ToolError> {
|
||||
match resp.outcome {
|
||||
ResponseOutcome::Result(value) => decode_call_result(tool_id, value),
|
||||
ResponseOutcome::Error(err) => Err(error_from_envelope(err)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a `tool_call_result` success body into the terminal
|
||||
/// [`TypedToolOutput`]. Shared by the core remote proxy and the SDK
|
||||
/// harness so both wire decoders reconstruct `chat_completion_output`
|
||||
/// identically.
|
||||
///
|
||||
/// A body with a `tool_call_id` is decoded strictly (`response_decoding` on
|
||||
/// failure), reconstructing `chat_completion_output` (an unparseable cco
|
||||
/// degrades to `None`). A bare body — e.g. a hub-local tool's raw output —
|
||||
/// passes through unchanged.
|
||||
pub fn decode_call_result(tool_id: ToolId, value: Value) -> Result<TypedToolOutput, ToolError> {
|
||||
if value.get("tool_call_id").is_none() {
|
||||
return Ok(TypedToolOutput::from_value(tool_id, value));
|
||||
}
|
||||
let result: ToolCallResult = serde_json::from_value(value)
|
||||
.map_err(|e| ToolError::custom("response_decoding", e.to_string()))?;
|
||||
let chat_completion_output = result.chat_completion_output.and_then(|cco| {
|
||||
serde_json::from_value::<ToolChatCompletionResponse>(cco)
|
||||
.inspect_err(|e| {
|
||||
warn!(tool_id = %tool_id, error = %e, "dropping unparseable chat_completion_output");
|
||||
})
|
||||
.ok()
|
||||
});
|
||||
let value = output_to_value(result.output);
|
||||
Ok(TypedToolOutput::from_value(tool_id, value)
|
||||
.with_chat_completion_output(chat_completion_output))
|
||||
}
|
||||
|
||||
/// Project a wire [`ToolOutputWire`] into a JSON [`Value`].
|
||||
///
|
||||
/// Three shapes collapse to one runtime type:
|
||||
/// - `Text` becomes a JSON string;
|
||||
/// - `Json` is forwarded verbatim;
|
||||
/// - `Mcp { blocks }` is re-serialised as `{ "blocks": [ContentBlock, ...] }`
|
||||
/// so the same downstream decoder used for in-process content blocks
|
||||
/// works without case-by-case adaptation.
|
||||
pub fn output_to_value(output: ToolOutputWire) -> Value {
|
||||
match output {
|
||||
ToolOutputWire::Text(s) => Value::String(s),
|
||||
ToolOutputWire::Json(v) => v,
|
||||
ToolOutputWire::Mcp { blocks } => {
|
||||
let runtime_blocks: Vec<ContentBlock> = blocks.into_iter().map(map_block).collect();
|
||||
// `ContentBlock`'s derived `Serialize` impl never fails for any
|
||||
// valid in-memory variant, but `to_value` is fallible at the
|
||||
// type level; collapse a hypothetical failure to `Value::Null`
|
||||
// before wrapping so this function stays total without an
|
||||
// `unwrap`. The outer `json!` only sees a `Value` expression
|
||||
// (which `to_value` round-trips infallibly), so the macro's
|
||||
// hidden `to_value` call cannot panic here.
|
||||
let blocks_value = serde_json::to_value(&runtime_blocks).unwrap_or(Value::Null);
|
||||
serde_json::json!({ "blocks": blocks_value })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_block(block: kigi_tool_protocol::McpBlock) -> ContentBlock {
|
||||
use kigi_tool_protocol::McpBlock;
|
||||
match block {
|
||||
McpBlock::Text { text } => ContentBlock::Text { text },
|
||||
McpBlock::Image { mime_type, data } => ContentBlock::Image {
|
||||
mime_type,
|
||||
data,
|
||||
media_id: None,
|
||||
filename: None,
|
||||
path: None,
|
||||
metadata: Default::default(),
|
||||
},
|
||||
McpBlock::Resource {
|
||||
uri,
|
||||
mime_type,
|
||||
text,
|
||||
} => ContentBlock::Resource {
|
||||
uri,
|
||||
mime_type,
|
||||
text,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Decode a JSON-RPC error envelope into a [`ToolError`]. The envelope's
|
||||
/// `data` field is expected to carry a serialised [`ToolErrorWire`] when
|
||||
/// available; falls back to a [`ToolError::Custom`] keyed by the numeric
|
||||
/// envelope code when the data shape is unknown.
|
||||
pub fn error_from_envelope(err: kigi_tool_protocol::JsonRpcError) -> ToolError {
|
||||
if let Some(data) = err.data.clone()
|
||||
&& let Ok(wire) = serde_json::from_value::<ToolErrorWire>(data)
|
||||
{
|
||||
return tool_error_from_wire(wire);
|
||||
}
|
||||
let mut e = ToolError::custom(format!("jsonrpc_{}", err.code), err.message);
|
||||
if let Some(data) = err.data {
|
||||
e = e.with_details(data);
|
||||
}
|
||||
e
|
||||
}
|
||||
|
||||
/// Recognize the hub's `workspace_unavailable` error on an already-decoded
|
||||
/// [`ToolError`]. Keys on `details["code"]` — the field that survives
|
||||
/// `ToolError::custom` + `with_details` — not the numeric code or the wire
|
||||
/// `Custom.subcode`.
|
||||
pub fn is_workspace_unavailable(err: &ToolError) -> bool {
|
||||
err.kind == ToolErrorKind::Custom
|
||||
&& err
|
||||
.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|v| v.as_str())
|
||||
== Some(WORKSPACE_UNAVAILABLE_SUBCODE)
|
||||
}
|
||||
|
||||
/// Map [`ToolErrorWire`] back into the runtime [`ToolError`]. The runtime
|
||||
/// error variants are the source-of-truth taxonomy; the wire form is a
|
||||
/// lossy projection onto stable codes for serialisation, so a few wire
|
||||
/// variants land on [`ToolError::Custom`] keyed by their wire code
|
||||
/// rather than a dedicated runtime variant.
|
||||
pub fn tool_error_from_wire(wire: ToolErrorWire) -> ToolError {
|
||||
match wire {
|
||||
ToolErrorWire::InvalidArguments { message, details } => {
|
||||
let e = ToolError::invalid_arguments(message);
|
||||
match details {
|
||||
Some(d) => e.with_details(d),
|
||||
None => e,
|
||||
}
|
||||
}
|
||||
ToolErrorWire::ToolNotFound { tool_id } => {
|
||||
let detail = format!("tool not found: {tool_id}");
|
||||
ToolError::not_found(tool_id, detail)
|
||||
}
|
||||
ToolErrorWire::PermissionDenied { reason } => ToolError::permission_denied(reason),
|
||||
ToolErrorWire::Timeout {
|
||||
tool_id,
|
||||
elapsed_ms,
|
||||
} => ToolError::new(
|
||||
ToolErrorKind::Timeout,
|
||||
format!("timed out after {elapsed_ms}ms"),
|
||||
)
|
||||
.with_details(serde_json::json!({"tool_id": tool_id.as_str(), "elapsed_ms": elapsed_ms})),
|
||||
ToolErrorWire::Cancelled { tool_id } => ToolError::cancelled(tool_id, "cancelled"),
|
||||
ToolErrorWire::Execution { tool_id, message } => ToolError::execution(tool_id, message),
|
||||
ToolErrorWire::BehaviorVersionUnsupported { tool_id, requested } => ToolError::new(
|
||||
ToolErrorKind::BehaviorVersionUnsupported,
|
||||
format!("behavior version {requested} not supported"),
|
||||
)
|
||||
.with_details(serde_json::json!({"tool_id": tool_id.as_str(), "requested": requested})),
|
||||
ToolErrorWire::RenderLimited {
|
||||
tool_id,
|
||||
card_id,
|
||||
reason,
|
||||
} => ToolError::new(ToolErrorKind::RenderLimited, reason)
|
||||
.with_details(serde_json::json!({"tool_id": tool_id.as_str(), "card_id": card_id})),
|
||||
ToolErrorWire::TerminalError { tool_id, message } => {
|
||||
ToolError::terminal_error(tool_id, message)
|
||||
}
|
||||
ToolErrorWire::Custom {
|
||||
subcode,
|
||||
message,
|
||||
details,
|
||||
} => {
|
||||
let e = ToolError::custom(subcode, message);
|
||||
match details {
|
||||
Some(d) => e.with_details(d),
|
||||
None => e,
|
||||
}
|
||||
}
|
||||
ToolErrorWire::SessionMismatch => ToolError::custom("session_mismatch", "session mismatch"),
|
||||
ToolErrorWire::TransportClosed { tool_id } => {
|
||||
ToolError::network_error(format!("transport closed for {tool_id}"))
|
||||
}
|
||||
ToolErrorWire::UnsupportedProtocolVersion { supported } => ToolError::custom(
|
||||
"unsupported_protocol_version",
|
||||
format!("supported versions: {supported:?}"),
|
||||
),
|
||||
ToolErrorWire::PayloadTooLarge { bytes, limit } => ToolError::custom(
|
||||
"payload_too_large",
|
||||
format!("payload {bytes} bytes exceeds limit {limit}"),
|
||||
),
|
||||
ToolErrorWire::Internal { request_id, detail } => {
|
||||
let e = ToolError::custom(
|
||||
"internal_error",
|
||||
detail.unwrap_or_else(|| "internal router error".to_owned()),
|
||||
);
|
||||
match request_id {
|
||||
// Keep `code` alongside `request_id`: `with_details` replaces
|
||||
// the `{"code": …}` object `ToolError::custom` installed.
|
||||
Some(id) => e.with_details(
|
||||
serde_json::json!({ "code": "internal_error", "request_id": id.as_str() }),
|
||||
),
|
||||
None => e,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
//! `CompoundResolver` plus the `ResolvedTool` and `ToolHandle`
|
||||
//! types it returns.
|
||||
//!
|
||||
//! `Tool` carries associated `Args` / `Output` types and is therefore not
|
||||
//! object-safe. [`ToolHandle`] is the dyn-compatible projection used
|
||||
//! by every router build: typed tools are wrapped via
|
||||
//! [`ErasedTool::new`]; remote registrations expose
|
||||
//! [`crate::RemoteToolProxy`] which implements [`ToolHandle`]
|
||||
//! directly without an intermediate typed `Tool` impl.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde_json::Value;
|
||||
|
||||
use kigi_tool_protocol::{SessionId, ToolCapabilities, ToolId, ToolRegistration};
|
||||
use kigi_tool_runtime::{
|
||||
ListToolsContext, Tool, ToolCallContext, ToolError, ToolOutput, ToolStream, ToolStreamItem,
|
||||
TypedToolOutput, terminal_only,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
use crate::registry::ToolRegistry;
|
||||
|
||||
/// Active resolution returned by [`CompoundResolver::resolve`].
|
||||
///
|
||||
/// Variants share the same `tool` handle and `registration` shape; the
|
||||
/// discriminant only tells callers whether the executing handle dispatches
|
||||
/// in-process or forwards over a connection. Differentiating the variants
|
||||
/// is useful for metrics, log tags, and the local-shadows-remote rule
|
||||
/// applied when both planes register the same `tool_id`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ResolvedTool {
|
||||
/// In-process tool resolved from the local registry.
|
||||
Local {
|
||||
/// Object-safe handle to the tool's `execute` entry point.
|
||||
tool: Arc<dyn ToolHandle>,
|
||||
/// Wire-shape registration record. Carries `tool_id`, the
|
||||
/// schema-bearing description, capabilities, and ownership data.
|
||||
registration: ToolRegistration,
|
||||
},
|
||||
/// Remote registration resolved through a connection-backed proxy.
|
||||
Remote {
|
||||
/// Object-safe handle whose `execute` forwards over the
|
||||
/// owning connection.
|
||||
proxy: Arc<dyn ToolHandle>,
|
||||
/// Wire-shape registration record (same shape as the local
|
||||
/// variant — both store the active registration so callers do not
|
||||
/// have to round-trip the registry for description / capabilities).
|
||||
registration: ToolRegistration,
|
||||
},
|
||||
}
|
||||
|
||||
impl ResolvedTool {
|
||||
/// Borrow the registration record regardless of variant.
|
||||
pub fn registration(&self) -> &ToolRegistration {
|
||||
match self {
|
||||
Self::Local { registration, .. } | Self::Remote { registration, .. } => registration,
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrow the executing handle regardless of variant.
|
||||
pub fn handle(&self) -> &Arc<dyn ToolHandle> {
|
||||
match self {
|
||||
Self::Local { tool, .. } => tool,
|
||||
Self::Remote { proxy, .. } => proxy,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Object-safe projection of a registered tool.
|
||||
///
|
||||
/// The router only needs identity, description, capabilities, and a
|
||||
/// JSON-typed `execute` entry point — exactly what this trait exposes.
|
||||
/// Adapters that wrap a typed `Tool` impl get [`ErasedTool`] for free;
|
||||
/// non-`Tool` handles (notably remote proxies) implement this trait
|
||||
/// directly.
|
||||
#[async_trait]
|
||||
pub trait ToolHandle: Send + Sync + std::fmt::Debug {
|
||||
/// Stable identity used by the router to route calls.
|
||||
fn id(&self) -> ToolId;
|
||||
|
||||
/// Model-facing description of the tool's argument schema.
|
||||
///
|
||||
/// Receives the per-turn [`ListToolsContext`] so handles backed by a
|
||||
/// typed [`Tool`] can produce context-aware descriptions at listing
|
||||
/// time. Callers outside a listing turn pass
|
||||
/// [`ListToolsContext::default`].
|
||||
fn description(&self, ctx: &ListToolsContext) -> ToolDescription;
|
||||
|
||||
/// Per-tool capability flags.
|
||||
fn capabilities(&self) -> ToolCapabilities;
|
||||
|
||||
/// Per-turn listing predicate.
|
||||
fn should_list(&self, _ctx: &ListToolsContext) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Streaming execution entry point.
|
||||
///
|
||||
/// Implementations encode the tool's typed `Output` to
|
||||
/// [`serde_json::Value`] and surface argument-decoding failures as
|
||||
/// [`ToolError::InvalidArguments`] within the terminal item.
|
||||
async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput>;
|
||||
}
|
||||
|
||||
/// Type-erasing wrapper for any [`Tool`] implementation.
|
||||
///
|
||||
/// Decodes `args` into `T::Args`, drives `T::execute`, and re-encodes each
|
||||
/// `T::Output` (terminal and progress items pass through unchanged
|
||||
/// otherwise). The wrapper holds the inner tool by `Arc` so the same
|
||||
/// underlying instance can back multiple registrations cheaply.
|
||||
pub struct ErasedTool<T> {
|
||||
inner: Arc<T>,
|
||||
}
|
||||
|
||||
impl<T> ErasedTool<T> {
|
||||
/// Wrap an `Arc<T>` for use as an [`ToolHandle`].
|
||||
pub fn from_arc(inner: Arc<T>) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
/// Wrap an owned tool, taking the `Arc` allocation internally.
|
||||
pub fn new(inner: T) -> Self {
|
||||
Self::from_arc(Arc::new(inner))
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: std::fmt::Debug> std::fmt::Debug for ErasedTool<T> {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ErasedTool")
|
||||
.field("inner", &self.inner)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> Clone for ErasedTool<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
inner: Arc::clone(&self.inner),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<T> ToolHandle for ErasedTool<T>
|
||||
where
|
||||
T: Tool + std::fmt::Debug + 'static,
|
||||
T::Output: ToolOutput,
|
||||
{
|
||||
fn id(&self) -> ToolId {
|
||||
self.inner.id()
|
||||
}
|
||||
|
||||
fn description(&self, ctx: &ListToolsContext) -> ToolDescription {
|
||||
self.inner.description(ctx)
|
||||
}
|
||||
|
||||
fn capabilities(&self) -> ToolCapabilities {
|
||||
self.inner.capabilities()
|
||||
}
|
||||
|
||||
fn should_list(&self, ctx: &ListToolsContext) -> bool {
|
||||
self.inner.should_list(ctx)
|
||||
}
|
||||
|
||||
async fn execute(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
let typed_args: T::Args = match serde_json::from_value(args) {
|
||||
Ok(a) => a,
|
||||
Err(e) => {
|
||||
return terminal_only(Err(ToolError::invalid_arguments(e.to_string())));
|
||||
}
|
||||
};
|
||||
let tool_id = self.inner.id();
|
||||
let stream = self.inner.execute(ctx, typed_args).await;
|
||||
let mapped = stream.map(move |item| match item {
|
||||
ToolStreamItem::Progress(p) => ToolStreamItem::Progress(p),
|
||||
ToolStreamItem::Terminal(Ok(out)) => match serde_json::to_value(&out) {
|
||||
Ok(value) => {
|
||||
let custom = out.model_output();
|
||||
let model_output = if custom.is_empty() {
|
||||
kigi_tool_runtime::extract_content_blocks(&value)
|
||||
} else {
|
||||
custom
|
||||
};
|
||||
let chat_completion_output = out.chat_completion_output();
|
||||
ToolStreamItem::Terminal(Ok(TypedToolOutput {
|
||||
tool_id: tool_id.clone(),
|
||||
value,
|
||||
model_output,
|
||||
chat_completion_output,
|
||||
}))
|
||||
}
|
||||
Err(e) => ToolStreamItem::Terminal(Err(ToolError::custom(
|
||||
"output_encoding",
|
||||
e.to_string(),
|
||||
))),
|
||||
},
|
||||
ToolStreamItem::Terminal(Err(err)) => ToolStreamItem::Terminal(Err(err)),
|
||||
});
|
||||
Box::pin(mapped)
|
||||
}
|
||||
}
|
||||
|
||||
/// Compose a local-first lookup over one (`local_only`) or two
|
||||
/// (`compound`) registries.
|
||||
///
|
||||
/// The lookup contract: `find_tool` is called on the local registry first;
|
||||
/// only if it returns `None` is the remote registry consulted. Any local
|
||||
/// registration shadows a same-id remote registration. Cross-session
|
||||
/// lookups return `None` — the caller may surface this as a
|
||||
/// [`ToolError::NotFound`] to keep ownership invisible to the requester.
|
||||
#[derive(Debug)]
|
||||
pub struct CompoundResolver {
|
||||
local: Arc<dyn ToolRegistry>,
|
||||
remote: Option<Arc<dyn ToolRegistry>>,
|
||||
}
|
||||
|
||||
impl CompoundResolver {
|
||||
/// Compose a resolver that consults a single local registry.
|
||||
pub fn local_only(local: Arc<dyn ToolRegistry>) -> Self {
|
||||
Self {
|
||||
local,
|
||||
remote: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Compose a resolver with both planes; `local` is consulted first.
|
||||
pub fn compound(local: Arc<dyn ToolRegistry>, remote: Arc<dyn ToolRegistry>) -> Self {
|
||||
Self {
|
||||
local,
|
||||
remote: Some(remote),
|
||||
}
|
||||
}
|
||||
|
||||
/// Borrow the local plane.
|
||||
pub fn local(&self) -> &Arc<dyn ToolRegistry> {
|
||||
&self.local
|
||||
}
|
||||
|
||||
/// Borrow the optional remote plane.
|
||||
pub fn remote(&self) -> Option<&Arc<dyn ToolRegistry>> {
|
||||
self.remote.as_ref()
|
||||
}
|
||||
|
||||
/// Resolve `(session, tool_id)` honouring the local-first rule.
|
||||
pub fn resolve(&self, session: &SessionId, tool_id: &ToolId) -> Option<ResolvedTool> {
|
||||
if let Some(hit) = self.local.find_tool(session, tool_id) {
|
||||
return Some(hit);
|
||||
}
|
||||
self.remote
|
||||
.as_ref()
|
||||
.and_then(|r| r.find_tool(session, tool_id))
|
||||
}
|
||||
|
||||
/// Resolve `(session, tool_id)` and dispatch through the active
|
||||
/// handle, returning the tool's stream verbatim. Misses produce a
|
||||
/// single-item terminal stream carrying [`ToolError::NotFound`].
|
||||
///
|
||||
/// Centralises the resolve-then-dispatch sequence so both the
|
||||
/// transport-side `LocalTransport::call` and the inner-dispatch path
|
||||
/// share one implementation: a future change to the miss-shape (or
|
||||
/// to the dispatch contract) lands once.
|
||||
pub async fn resolve_and_dispatch(
|
||||
&self,
|
||||
session: &SessionId,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
match self.resolve(session, &tool_id) {
|
||||
Some(resolved) => resolved.handle().execute(ctx, args).await,
|
||||
None => terminal_only(Err(ToolError::not_found(
|
||||
tool_id.clone(),
|
||||
format!("tool not found: {tool_id}"),
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
//! Object-safe `Transport` trait plus the `Principal` value carried across
|
||||
//! authorize/call boundaries.
|
||||
//!
|
||||
//! [`TransportKind`] is re-exported from [`kigi_tool_protocol`] so the wire
|
||||
//! and dispatch layers share one canonical enum and there is no duplicate
|
||||
//! `Local` / `Remote` definition to keep in sync.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use kigi_tool_protocol::{SessionId, ToolId, UserId};
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput};
|
||||
|
||||
pub use kigi_tool_protocol::TransportKind;
|
||||
|
||||
/// Authenticated identity bound to a transport at handshake time.
|
||||
///
|
||||
/// The transport authorises **once** at connect; subsequent dispatch calls
|
||||
/// carry no extra credentials. `session_ids` is plural because a JWT may
|
||||
/// authorise more than one session (multi-tenant tooling sessions sharing
|
||||
/// a single user identity); the router narrows by [`SessionId`] at the
|
||||
/// per-call boundary.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Principal {
|
||||
/// Authenticated user identity.
|
||||
pub user_id: UserId,
|
||||
|
||||
/// Sessions this principal is authorised to act on. Empty when the
|
||||
/// transport authorises a user but has not yet bound a session
|
||||
/// (e.g. a fresh harness connection that has not opened a session).
|
||||
pub session_ids: Vec<SessionId>,
|
||||
|
||||
/// OAuth-style scopes granted to this principal, e.g. `"tool.invoke"`.
|
||||
pub scopes: Vec<String>,
|
||||
|
||||
/// Token audiences claimed by the credential, e.g. the router's
|
||||
/// expected `aud` values. Used by callers that need defence-in-depth
|
||||
/// audience checks beyond what the transport already validated.
|
||||
pub audiences: Vec<String>,
|
||||
}
|
||||
|
||||
impl Principal {
|
||||
/// Build a principal for `user_id` with no sessions, scopes, or
|
||||
/// audiences. Use the `with_*` builders to populate the rest.
|
||||
pub fn new(user_id: UserId) -> Self {
|
||||
Self {
|
||||
user_id,
|
||||
session_ids: Vec::new(),
|
||||
scopes: Vec::new(),
|
||||
audiences: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Append `session_id` to the authorised set.
|
||||
pub fn with_session(mut self, session_id: SessionId) -> Self {
|
||||
self.session_ids.push(session_id);
|
||||
self
|
||||
}
|
||||
|
||||
/// Append `scope` to the granted scopes.
|
||||
pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
|
||||
self.scopes.push(scope.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Append `aud` to the token's audience list.
|
||||
pub fn with_audience(mut self, aud: impl Into<String>) -> Self {
|
||||
self.audiences.push(aud.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Whether `scope` is present in the granted scopes.
|
||||
pub fn has_scope(&self, scope: &str) -> bool {
|
||||
self.scopes.iter().any(|s| s == scope)
|
||||
}
|
||||
|
||||
/// Whether `session_id` is in the principal's authorised session set.
|
||||
pub fn authorizes_session(&self, session_id: &SessionId) -> bool {
|
||||
self.session_ids.iter().any(|s| s == session_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Object-safe transport for dispatching tool calls.
|
||||
///
|
||||
/// Implementations come in two flavours: [`TransportKind::Local`] resolves
|
||||
/// against an in-process registry, while [`TransportKind::Remote`] forwards
|
||||
/// a `tool_call_request` over a [`crate::ConnectionClient`].
|
||||
#[async_trait]
|
||||
pub trait Transport: Send + Sync + std::fmt::Debug {
|
||||
/// Whether the underlying transport is local (in-process) or remote
|
||||
/// (forwarded over a connection).
|
||||
fn kind(&self) -> TransportKind;
|
||||
|
||||
/// One-time authorisation handshake.
|
||||
///
|
||||
/// Local transports return a principal derived from the bound OS user
|
||||
/// (or whatever ambient identity the host process provides). Remote
|
||||
/// transports return the principal extracted from a validated
|
||||
/// credential. Subsequent [`Self::call`] invocations reuse this
|
||||
/// principal — the router never re-authorises per call.
|
||||
async fn authorize(&self) -> Result<Principal, ToolError>;
|
||||
|
||||
/// Dispatch a tool call.
|
||||
///
|
||||
/// The returned [`ToolStream`] follows the runtime invariant: zero or
|
||||
/// more `Progress` items followed by exactly one `Terminal`. A
|
||||
/// not-found result is reported as a single-item terminal stream
|
||||
/// carrying [`ToolError::NotFound`]; transport-level disconnects
|
||||
/// surface as [`ToolError::NetworkError`].
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput>;
|
||||
}
|
||||
@@ -0,0 +1,335 @@
|
||||
//! `CompoundResolver` and `ResolvedTool` coverage. Exercises local-only,
|
||||
//! local-shadows-remote, remote-fallback, and cross-session scenarios.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use futures::StreamExt;
|
||||
use kigi_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport,
|
||||
ToolHandle, ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome,
|
||||
};
|
||||
use kigi_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolStreamItem,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EmptyArgs {}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StubTool {
|
||||
id: ToolId,
|
||||
}
|
||||
|
||||
impl Tool for StubTool {
|
||||
type Args = EmptyArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
self.id.clone()
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id))
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
_args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({"id": self.id.as_str()}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PlaneRegistry {
|
||||
// Set once at construction; `TransportKind` is `Copy` so a direct
|
||||
// field is the obvious choice — no interior mutability required.
|
||||
transport_kind: TransportKind,
|
||||
entries: DashMap<(SessionId, ToolId), ToolRegistration>,
|
||||
handles: DashMap<ToolId, Arc<dyn ToolHandle>>,
|
||||
}
|
||||
|
||||
impl PlaneRegistry {
|
||||
fn new(kind: TransportKind) -> Self {
|
||||
Self {
|
||||
transport_kind: kind,
|
||||
entries: DashMap::new(),
|
||||
handles: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn install(&self, session: &SessionId, id: &ToolId) {
|
||||
let reg = ToolRegistration {
|
||||
tool_id: id.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: ToolDescription::new(id.as_str(), format!("stub for {id}")),
|
||||
input_schema: None,
|
||||
capabilities: None,
|
||||
notification_schemas: None,
|
||||
transport_kind: self.transport_kind,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries.insert((session.clone(), id.clone()), reg);
|
||||
self.handles.insert(
|
||||
id.clone(),
|
||||
Arc::new(ErasedTool::new(StubTool { id: id.clone() })),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for PlaneRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!("resolver tests pre-populate via install()")
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let registration = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
let handle = self.handles.get(tool)?.value().clone();
|
||||
match registration.transport_kind {
|
||||
TransportKind::Local => Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
}),
|
||||
TransportKind::Remote => Some(ResolvedTool::Remote {
|
||||
proxy: handle,
|
||||
registration,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &kigi_tool_protocol::UserId,
|
||||
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_only_resolves_local_hits() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Local { registration, .. }) => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
}
|
||||
other => panic!("expected Local, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_only_returns_none_for_unknown() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compound_falls_through_to_remote_when_local_misses() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
remote.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Remote { registration, .. }) => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
assert_eq!(registration.transport_kind, TransportKind::Remote);
|
||||
}
|
||||
other => panic!("expected Remote, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_shadows_same_id_remote() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
remote.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
match resolver.resolve(&sid("sess-1"), &tid("foo")) {
|
||||
Some(ResolvedTool::Local { registration, .. }) => {
|
||||
assert_eq!(registration.transport_kind, TransportKind::Local);
|
||||
}
|
||||
other => panic!("expected local resolution to shadow remote, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cross_session_lookup_returns_none() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-other"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compound_returns_none_when_neither_plane_holds_id() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let remote = Arc::new(PlaneRegistry::new(TransportKind::Remote));
|
||||
let resolver = CompoundResolver::compound(
|
||||
local as Arc<dyn ToolRegistry>,
|
||||
remote as Arc<dyn ToolRegistry>,
|
||||
);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolved_tool_helpers_borrow_active_handle_and_registration() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let resolved = resolver.resolve(&sid("sess-1"), &tid("foo")).expect("hit");
|
||||
assert_eq!(resolved.registration().tool_id, tid("foo"));
|
||||
assert_eq!(resolved.handle().id(), tid("foo"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_and_dispatch_drives_the_resolved_handle() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
local.install(&sid("sess-1"), &tid("foo"));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let mut stream = resolver
|
||||
.resolve_and_dispatch(
|
||||
&sid("sess-1"),
|
||||
tid("foo"),
|
||||
serde_json::json!({}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"id": "foo"}));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_and_dispatch_misses_yield_terminal_not_found() {
|
||||
let local = Arc::new(PlaneRegistry::new(TransportKind::Local));
|
||||
let resolver = CompoundResolver::local_only(local as Arc<dyn ToolRegistry>);
|
||||
let mut stream = resolver
|
||||
.resolve_and_dispatch(
|
||||
&sid("sess-1"),
|
||||
tid("missing"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("missing"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Err(NotFound)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,317 @@
|
||||
//! `InnerDispatchForResolver` coverage. Verifies the cycle-safe `Weak`
|
||||
//! resolver semantics and the session-bound resolution path.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use kigi_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, InnerDispatchForResolver, ResolvedTool,
|
||||
SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome,
|
||||
ToolSessionUnbindOutcome,
|
||||
};
|
||||
use kigi_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolDispatch, ToolError, ToolStreamItem,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EchoArgs {
|
||||
payload: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTool;
|
||||
|
||||
impl Tool for EchoTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("echo").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("echo", "Echoes its input.")
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({"echoed": args.payload}))
|
||||
}
|
||||
}
|
||||
|
||||
type RegistryEntry = (ToolRegistration, Arc<dyn ToolHandle>);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemRegistry {
|
||||
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
|
||||
}
|
||||
|
||||
impl InMemRegistry {
|
||||
fn install(&self, session: SessionId, tool: ToolId, handle: Arc<dyn ToolHandle>) {
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: handle.description(&kigi_tool_runtime::ListToolsContext::default()),
|
||||
input_schema: None,
|
||||
capabilities: Some(handle.capabilities()),
|
||||
notification_schemas: None,
|
||||
transport_kind: TransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries.insert((session, tool), (registration, handle));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for InMemRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let (registration, handle) = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
})
|
||||
}
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &kigi_tool_protocol::UserId,
|
||||
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_resolves_through_bound_session() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1"));
|
||||
assert_eq!(inner.session_id(), &sid("sess-1"));
|
||||
|
||||
let result = inner
|
||||
.call_terminal(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "x"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await
|
||||
.expect("terminal ok");
|
||||
assert_eq!(result.value, serde_json::json!({"echoed": "x"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_returns_not_found_when_tool_absent() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-1"));
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("ghost"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NotFound), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_uses_bound_session_not_context_session() {
|
||||
// Even if the context were to carry a different session, the inner
|
||||
// dispatch handle resolves against its construction-time session.
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-A"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner = InnerDispatchForResolver::new(Arc::downgrade(&resolver), sid("sess-B"));
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "x"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound => {}
|
||||
other => panic!("session-A registration must not be visible from session-B, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_after_resolver_drop_returns_computer_hub_dropped() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let weak = Arc::downgrade(&resolver);
|
||||
let inner = InnerDispatchForResolver::new(weak, sid("sess-1"));
|
||||
drop(resolver);
|
||||
let mut stream = inner
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::Custom =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("computer_hub_dropped")
|
||||
|| e.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|c| c.as_str())
|
||||
== Some("computer_hub_dropped"),
|
||||
"expected computer_hub_dropped code, got: {:?}",
|
||||
e
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Custom(computer_hub_dropped)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inner_dispatch_implements_object_safe_tool_dispatch() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let inner: Arc<dyn ToolDispatch> = Arc::new(InnerDispatchForResolver::new(
|
||||
Arc::downgrade(&resolver),
|
||||
sid("sess-1"),
|
||||
));
|
||||
let result = inner
|
||||
.call_terminal(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(result, Err(ref e) if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound));
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! Decode-side coverage for `ToolErrorWire::Internal`'s optional `detail`:
|
||||
//! a populated detail must become the reconstructed `ToolError`'s message,
|
||||
//! and its absence (frames from older peers) must fall back to the historic
|
||||
//! constant.
|
||||
|
||||
use kigi_computer_hub_core::{error_from_envelope, tool_error_from_wire};
|
||||
use kigi_tool_protocol::{JsonRpcError, RequestId, ToolErrorWire};
|
||||
use kigi_tool_runtime::ToolErrorKind;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn internal_with_detail_reconstructs_the_wire_detail() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: None,
|
||||
detail: Some("cross-instance tool.call timed out".to_owned()),
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "cross-instance tool.call timed out");
|
||||
// The `internal_error` code survives so callers can still classify it.
|
||||
assert_eq!(
|
||||
err.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("internal_error"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internal_without_detail_falls_back_to_the_historic_constant() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: None,
|
||||
detail: None,
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "internal router error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn internal_with_request_id_keeps_both_code_and_request_id() {
|
||||
let err = tool_error_from_wire(ToolErrorWire::Internal {
|
||||
request_id: Some(RequestId::new("req-7").unwrap()),
|
||||
detail: Some("relay publish failed".to_owned()),
|
||||
});
|
||||
assert_eq!(err.detail, "relay publish failed");
|
||||
let details = err.details.expect("details present");
|
||||
assert_eq!(details["code"], json!("internal_error"));
|
||||
assert_eq!(details["request_id"], json!("req-7"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn envelope_with_internal_data_prefers_data_detail_over_message() {
|
||||
// The hub's `-32000 "internal error"` envelope keeps its constant message;
|
||||
// the harness must read the cause from `error.data`, not the message.
|
||||
let err = error_from_envelope(JsonRpcError {
|
||||
code: -32000,
|
||||
message: "internal error".to_owned(),
|
||||
data: Some(json!({
|
||||
"code": "internal_error",
|
||||
"detail": "cross-instance call cancelled",
|
||||
})),
|
||||
});
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert_eq!(err.detail, "cross-instance call cancelled");
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
//! `LocalTransport` end-to-end coverage. Verifies that the transport
|
||||
//! resolves through the bound resolver, drives both blocking and
|
||||
//! streaming tools, and surfaces missing tools as `Terminal(NotFound)`.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use kigi_computer_hub_core::{
|
||||
CompoundResolver, ConnectionCleanupReport, ErasedTool, LocalTransport, ResolvedTool,
|
||||
SessionCleanupReport, ToolHandle, ToolRegistry, ToolSessionBindOutcome,
|
||||
ToolSessionUnbindOutcome, Transport, TransportKind,
|
||||
};
|
||||
use kigi_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind as WireTransportKind, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{
|
||||
SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError, ToolProgress, ToolStream,
|
||||
ToolStreamItem, terminal_only, with_progress,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EchoArgs {
|
||||
payload: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTool;
|
||||
|
||||
impl Tool for EchoTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("echo").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("echo", "Echoes its input.")
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
Ok(serde_json::json!({ "echoed": args.payload }))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StreamerTool;
|
||||
|
||||
impl Tool for StreamerTool {
|
||||
type Args = EchoArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new("streamer").expect("tool id")
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new("streamer", "Emits three progress chunks.")
|
||||
}
|
||||
|
||||
async fn execute(&self, _ctx: ToolCallContext, args: Self::Args) -> ToolStream<Self::Output> {
|
||||
let chunks = futures::stream::iter(vec![
|
||||
ToolProgress::Text {
|
||||
text: "tick".to_string(),
|
||||
},
|
||||
ToolProgress::Text {
|
||||
text: "tock".to_string(),
|
||||
},
|
||||
ToolProgress::Text {
|
||||
text: "boom".to_string(),
|
||||
},
|
||||
]);
|
||||
with_progress(chunks, async move {
|
||||
Ok(serde_json::json!({ "echoed": args.payload }))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
type RegistryEntry = (ToolRegistration, Arc<dyn ToolHandle>);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemRegistry {
|
||||
entries: DashMap<(SessionId, ToolId), RegistryEntry>,
|
||||
}
|
||||
|
||||
impl InMemRegistry {
|
||||
fn install(&self, session: SessionId, tool_id: ToolId, handle: Arc<dyn ToolHandle>) {
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool_id.clone(),
|
||||
sessions: Some(vec![session.clone()]),
|
||||
user_id: UserId::new("alice").expect("user id"),
|
||||
server_id: None,
|
||||
description: handle.description(&kigi_tool_runtime::ListToolsContext::default()),
|
||||
input_schema: None,
|
||||
capabilities: Some(handle.capabilities()),
|
||||
notification_schemas: None,
|
||||
transport_kind: WireTransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
self.entries
|
||||
.insert((session, tool_id), (registration, handle));
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for InMemRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
unreachable!("transport tests pre-populate via install()")
|
||||
}
|
||||
async fn register_server(
|
||||
&self,
|
||||
_connection_id: ConnectionId,
|
||||
_reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_tool(&self, _connection_id: &ConnectionId, _tool: &ToolId) -> bool {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unregister_server(&self, _connection_id: &ConnectionId, _server: &ServerId) -> usize {
|
||||
unreachable!()
|
||||
}
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
_session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
unreachable!()
|
||||
}
|
||||
async fn drop_connection(&self, _connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
ConnectionCleanupReport::default()
|
||||
}
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let (registration, handle) = self
|
||||
.entries
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
})
|
||||
}
|
||||
fn list_tools(&self, _session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
vec![]
|
||||
}
|
||||
fn list_servers(&self, _session: &SessionId) -> Vec<ServerSummary> {
|
||||
vec![]
|
||||
}
|
||||
fn search(&self, _session: &SessionId, _query: &str, _limit: usize) -> SearchSnapshot {
|
||||
SearchSnapshot {
|
||||
results: vec![],
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
async fn unregister_session(&self, _session: &SessionId) -> SessionCleanupReport {
|
||||
SessionCleanupReport::default()
|
||||
}
|
||||
fn tool_sessions(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
_tool: &ToolId,
|
||||
) -> std::collections::HashSet<SessionId> {
|
||||
std::collections::HashSet::new()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &kigi_tool_protocol::UserId,
|
||||
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("user id")
|
||||
}
|
||||
|
||||
async fn collect(
|
||||
stream: &mut ToolStream<kigi_tool_runtime::TypedToolOutput>,
|
||||
) -> Vec<ToolStreamItem<kigi_tool_runtime::TypedToolOutput>> {
|
||||
let mut items = Vec::new();
|
||||
while let Some(item) = stream.next().await {
|
||||
items.push(item);
|
||||
}
|
||||
items
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatches_blocking_tool_to_terminal_value() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("echo"),
|
||||
serde_json::json!({"payload": "hi"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"echoed": "hi"}));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatches_streaming_tool_with_three_progress_then_terminal() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("streamer"),
|
||||
Arc::new(ErasedTool::new(StreamerTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("streamer"),
|
||||
serde_json::json!({"payload": "hi"}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 4);
|
||||
for item in &items[..3] {
|
||||
assert!(matches!(item, ToolStreamItem::Progress(_)));
|
||||
}
|
||||
assert!(matches!(items[3], ToolStreamItem::Terminal(Ok(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_tool_resolves_as_terminal_not_found() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("ghost"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Err(e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("ghost"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Err(NotFound)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_arguments_surface_as_terminal_error() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
registry.install(
|
||||
sid("sess-1"),
|
||||
tid("echo"),
|
||||
Arc::new(ErasedTool::new(EchoTool)),
|
||||
);
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
|
||||
let mut stream = transport
|
||||
.call(
|
||||
tid("echo"),
|
||||
// Missing required `payload` field.
|
||||
serde_json::json!({}),
|
||||
ToolCallContext::default(),
|
||||
)
|
||||
.await;
|
||||
let items = collect(&mut stream).await;
|
||||
assert_eq!(items.len(), 1);
|
||||
match &items[0] {
|
||||
ToolStreamItem::Terminal(Err(e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::InvalidArguments => {}
|
||||
other => panic!("expected Terminal(Err(InvalidArguments)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorize_returns_bound_principal_with_invoke_scope() {
|
||||
let registry = Arc::new(InMemRegistry::default());
|
||||
let resolver = Arc::new(CompoundResolver::local_only(
|
||||
registry as Arc<dyn ToolRegistry>,
|
||||
));
|
||||
let transport = LocalTransport::new(resolver, uid("alice"), sid("sess-1"));
|
||||
let principal = transport.authorize().await.expect("authorize");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
assert!(principal.has_scope(kigi_computer_hub_core::LOCAL_INVOKE_SCOPE));
|
||||
assert_eq!(transport.kind(), TransportKind::Local);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unused_helpers_silenced() {
|
||||
// `terminal_only` is re-exported for adapter authors; touch it here so
|
||||
// a future refactor that drops the import does not silently break the
|
||||
// re-export surface.
|
||||
let _: ToolStream<serde_json::Value> = terminal_only(Ok(serde_json::Value::Null));
|
||||
}
|
||||
@@ -0,0 +1,727 @@
|
||||
//! `RemoteToolProxy` and `RemoteTransport` coverage. A channel-backed
|
||||
//! mock `ConnectionClient` lets the test inspect outgoing frames and
|
||||
//! drive synthetic responses + progress without any tokio I/O.
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use futures::channel::{mpsc, oneshot};
|
||||
use futures::stream::BoxStream;
|
||||
|
||||
use kigi_computer_hub_core::{
|
||||
ConnectionClient, RemoteToolProxy, RemoteTransport, ToolHandle, Transport, TransportKind,
|
||||
};
|
||||
use kigi_tool_protocol::{
|
||||
JsonRpcError, JsonRpcId, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, JsonRpcVersion,
|
||||
Method, ResponseOutcome, SessionId, ToolCallId, ToolCallParams, ToolCallProgressFrame,
|
||||
ToolCallResult, ToolCapabilities, ToolErrorWire, ToolId, ToolOutputWire, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{
|
||||
ContentBlock, ToolCallContext, ToolError, ToolOutput, ToolProgress, ToolStreamItem,
|
||||
};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
|
||||
/// Programmable `ConnectionClient`. Each request gets a pre-staged
|
||||
/// response; progress frames are pushed through per-call senders.
|
||||
#[derive(Debug, Default)]
|
||||
struct MockConnection {
|
||||
/// Senders keyed by `tool_call_id`. Pulled out of the inner state so
|
||||
/// per-call subscription touches a lock-free DashMap rather than the
|
||||
/// shared Mutex that guards the rest of the queue + capture state.
|
||||
progress_senders: DashMap<ToolCallId, mpsc::UnboundedSender<ToolCallProgressFrame>>,
|
||||
/// Three-Vec state guarded by one Mutex. The lock provides atomic
|
||||
/// pop-from-`responses` + push-to-`captured_requests` semantics that
|
||||
/// some tests rely on.
|
||||
inner: Mutex<MockState>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MockState {
|
||||
/// FIFO queue of responses to return for each `request` call.
|
||||
responses: Vec<MockResponse>,
|
||||
/// Captured outgoing requests so tests can assert on them.
|
||||
captured_requests: Vec<JsonRpcRequest>,
|
||||
/// Captured one-way notifications.
|
||||
captured_notifications: Vec<JsonRpcNotification>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for MockState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("MockState")
|
||||
.field("responses_len", &self.responses.len())
|
||||
.field("captured_reqs", &self.captured_requests.len())
|
||||
.field("captured_notifs", &self.captured_notifications.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
enum MockResponse {
|
||||
Ok(serde_json::Value),
|
||||
Err(JsonRpcError),
|
||||
/// Resolves a oneshot when the request arrives so the test can
|
||||
/// release progress before allowing the response.
|
||||
Gated {
|
||||
gate: oneshot::Receiver<()>,
|
||||
body: serde_json::Value,
|
||||
},
|
||||
/// Fail at the transport layer (e.g. socket dropped).
|
||||
Network(String),
|
||||
}
|
||||
|
||||
impl MockConnection {
|
||||
fn enqueue_ok(&self, body: serde_json::Value) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Ok(body));
|
||||
}
|
||||
|
||||
fn enqueue_err(&self, code: i32, message: impl Into<String>, data: Option<serde_json::Value>) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Err(JsonRpcError {
|
||||
code,
|
||||
message: message.into(),
|
||||
data,
|
||||
}));
|
||||
}
|
||||
|
||||
fn enqueue_gated(&self, gate: oneshot::Receiver<()>, body: serde_json::Value) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Gated { gate, body });
|
||||
}
|
||||
|
||||
fn enqueue_network_failure(&self, message: impl Into<String>) {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.responses
|
||||
.push(MockResponse::Network(message.into()));
|
||||
}
|
||||
|
||||
fn last_request(&self) -> Option<JsonRpcRequest> {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.captured_requests
|
||||
.last()
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn captured_request_count(&self) -> usize {
|
||||
self.inner.lock().expect("mutex").captured_requests.len()
|
||||
}
|
||||
|
||||
fn push_progress(&self, tool_call_id: &ToolCallId, frame: ToolCallProgressFrame) {
|
||||
if let Some(tx) = self.progress_senders.get(tool_call_id) {
|
||||
let _ = tx.value().unbounded_send(frame);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ConnectionClient for MockConnection {
|
||||
async fn request(&self, request: JsonRpcRequest) -> Result<JsonRpcResponse, ToolError> {
|
||||
let response = {
|
||||
let mut guard = self.inner.lock().expect("mutex");
|
||||
guard.captured_requests.push(request.clone());
|
||||
if guard.responses.is_empty() {
|
||||
return Err(ToolError::custom(
|
||||
"mock_response_missing",
|
||||
"no response staged",
|
||||
));
|
||||
}
|
||||
guard.responses.remove(0)
|
||||
};
|
||||
match response {
|
||||
MockResponse::Ok(body) => Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Result(body),
|
||||
}),
|
||||
MockResponse::Err(err) => Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Error(err),
|
||||
}),
|
||||
MockResponse::Gated { gate, body } => {
|
||||
let _ = gate.await;
|
||||
Ok(JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: request.id,
|
||||
session_id: request.session_id,
|
||||
outcome: ResponseOutcome::Result(body),
|
||||
})
|
||||
}
|
||||
MockResponse::Network(msg) => Err(ToolError::network_error(msg)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn subscribe_progress(
|
||||
&self,
|
||||
tool_call_id: ToolCallId,
|
||||
) -> BoxStream<'static, ToolCallProgressFrame> {
|
||||
let (tx, rx) = mpsc::unbounded();
|
||||
self.progress_senders.insert(tool_call_id, tx);
|
||||
rx.boxed()
|
||||
}
|
||||
|
||||
async fn notify(&self, notification: JsonRpcNotification) -> Result<(), ToolError> {
|
||||
self.inner
|
||||
.lock()
|
||||
.expect("mutex")
|
||||
.captured_notifications
|
||||
.push(notification);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("tool id")
|
||||
}
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("user id")
|
||||
}
|
||||
|
||||
fn description_for(name: &str) -> ToolDescription {
|
||||
ToolDescription::new(name, format!("desc for {name}"))
|
||||
}
|
||||
|
||||
fn ok_call_result(call_id: &ToolCallId, output: ToolOutputWire) -> serde_json::Value {
|
||||
serde_json::to_value(ToolCallResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
output,
|
||||
follow_ups: vec![],
|
||||
reminders: vec![],
|
||||
chat_completion_output: None,
|
||||
})
|
||||
.expect("serialise call result")
|
||||
}
|
||||
|
||||
fn ok_call_result_with_cco(
|
||||
call_id: &ToolCallId,
|
||||
output: ToolOutputWire,
|
||||
chat_completion_output: serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
serde_json::to_value(ToolCallResult {
|
||||
tool_call_id: call_id.clone(),
|
||||
output,
|
||||
follow_ups: vec![],
|
||||
reminders: vec![],
|
||||
chat_completion_output: Some(chat_completion_output),
|
||||
})
|
||||
.expect("serialise call result")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_sends_well_formed_tool_call_request() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(
|
||||
&call_id,
|
||||
ToolOutputWire::Text("hello".to_string()),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!({"k": "v"})).await;
|
||||
while stream.next().await.is_some() {}
|
||||
let req = conn.last_request().expect("captured request");
|
||||
assert_eq!(req.method, Method::ToolCallRequest.as_wire_str());
|
||||
let params: ToolCallParams = serde_json::from_value(req.params).expect("decode params");
|
||||
assert_eq!(params.tool_id, tid("foo"));
|
||||
assert_eq!(params.tool_call_id, call_id);
|
||||
assert_eq!(params.arguments, serde_json::json!({"k": "v"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn progress_then_terminal_orders_correctly() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let (gate_tx, gate_rx) = oneshot::channel();
|
||||
conn.enqueue_gated(
|
||||
gate_rx,
|
||||
ok_call_result(&call_id, ToolOutputWire::Text("done".to_string())),
|
||||
);
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
|
||||
// Push two progress frames before the terminal is unblocked.
|
||||
conn.push_progress(
|
||||
&call_id,
|
||||
ToolCallProgressFrame {
|
||||
tool_call_id: call_id.clone(),
|
||||
kind: "log".to_string(),
|
||||
body: serde_json::json!({"text": "tick"}),
|
||||
dropped_count: None,
|
||||
},
|
||||
);
|
||||
conn.push_progress(
|
||||
&call_id,
|
||||
ToolCallProgressFrame {
|
||||
tool_call_id: call_id.clone(),
|
||||
kind: "log".to_string(),
|
||||
body: serde_json::json!({"text": "tock"}),
|
||||
dropped_count: None,
|
||||
},
|
||||
);
|
||||
|
||||
let first = stream.next().await.expect("first item");
|
||||
let second = stream.next().await.expect("second item");
|
||||
match (&first, &second) {
|
||||
(ToolStreamItem::Progress(p1), ToolStreamItem::Progress(p2)) => {
|
||||
match p1 {
|
||||
ToolProgress::Custom { subkind, payload } => {
|
||||
assert_eq!(subkind, "log");
|
||||
assert_eq!(payload, &serde_json::json!({"text": "tick"}));
|
||||
}
|
||||
other => panic!("expected Custom progress, got {other:?}"),
|
||||
}
|
||||
match p2 {
|
||||
ToolProgress::Custom { subkind, .. } => assert_eq!(subkind, "log"),
|
||||
other => panic!("expected Custom progress, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected two Progress items, got {other:?}"),
|
||||
}
|
||||
|
||||
// Release the response and consume the terminal.
|
||||
let _ = gate_tx.send(());
|
||||
let terminal = stream.next().await.expect("terminal");
|
||||
match terminal {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!("done"));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
assert!(stream.next().await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_rpc_error_response_decodes_into_tool_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let wire = ToolErrorWire::ToolNotFound {
|
||||
tool_id: tid("foo"),
|
||||
};
|
||||
conn.enqueue_err(
|
||||
-32011,
|
||||
"tool not found",
|
||||
Some(serde_json::to_value(&wire).unwrap()),
|
||||
);
|
||||
let mut stream = proxy
|
||||
.execute(ToolCallContext::default(), serde_json::json!(null))
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NotFound =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("foo"),
|
||||
"detail should mention tool id: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NotFound), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn network_failure_surfaces_as_terminal_network_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
conn.enqueue_network_failure("socket closed");
|
||||
let mut stream = proxy
|
||||
.execute(ToolCallContext::default(), serde_json::json!(null))
|
||||
.await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::NetworkError =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("socket closed"),
|
||||
"detail should mention cause: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(NetworkError), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mcp_output_re_serialises_into_blocks_value() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let blocks = vec![kigi_tool_protocol::McpBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}];
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Mcp { blocks }));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
// The wire blocks round-trip through ContentBlock; assert the
|
||||
// text value survives the transformation.
|
||||
let blocks_value = typed
|
||||
.value
|
||||
.get("blocks")
|
||||
.and_then(|v| v.as_array())
|
||||
.cloned()
|
||||
.expect("blocks array");
|
||||
assert_eq!(blocks_value.len(), 1);
|
||||
let block: ContentBlock =
|
||||
serde_json::from_value(blocks_value[0].clone()).expect("decode runtime block");
|
||||
match block {
|
||||
ContentBlock::Text { text } => assert_eq!(text, "hello"),
|
||||
other => panic!("expected Text block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_carries_chat_completion_output_from_wire() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
let cco = serde_json::json!({
|
||||
"result": {
|
||||
"sender": "assistant",
|
||||
"message": "",
|
||||
"code_execution_result": {
|
||||
"stdout": "hi\n",
|
||||
"stderr": "",
|
||||
"exit_code": 0,
|
||||
"command_timed_out": false
|
||||
}
|
||||
}
|
||||
});
|
||||
conn.enqueue_ok(ok_call_result_with_cco(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
cco,
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
let response = typed
|
||||
.chat_completion_output()
|
||||
.expect("chat completion output survives the wire");
|
||||
let completion = response.result.expect("completion result present");
|
||||
let exec = completion
|
||||
.code_execution_result
|
||||
.expect("code execution result present");
|
||||
assert_eq!(exec.stdout, "hi\n");
|
||||
assert_eq!(exec.exit_code, 0);
|
||||
assert!(!exec.command_timed_out);
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_without_chat_completion_output_is_none() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_inner_chat_completion_output_degrades_to_none() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result_with_cco(
|
||||
&call_id,
|
||||
ToolOutputWire::Json(serde_json::json!({"stdout": "hi\n"})),
|
||||
serde_json::json!({"result": "not-a-completion-object"}),
|
||||
));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"}));
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok) with degraded cco, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bare_non_enveloped_success_body_passes_through() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let ctx = ToolCallContext::new(ToolCallId::new_v7());
|
||||
conn.enqueue_ok(serde_json::json!({"stdout": "hi\n"}));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
assert_eq!(typed.value, serde_json::json!({"stdout": "hi\n"}));
|
||||
assert!(typed.chat_completion_output().is_none());
|
||||
}
|
||||
other => panic!("expected Terminal(Ok) passthrough, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_envelope_with_tool_call_id_surfaces_decode_error() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("bash"),
|
||||
sid("sess-1"),
|
||||
description_for("bash"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let ctx = ToolCallContext::new(ToolCallId::new_v7());
|
||||
conn.enqueue_ok(serde_json::json!({"tool_call_id": "call_x", "output": 123}));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
let item = stream.next().await.expect("terminal");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::Custom =>
|
||||
{
|
||||
let code = e
|
||||
.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|c| c.as_str());
|
||||
assert_eq!(code, Some("response_decoding"), "error: {e:?}");
|
||||
}
|
||||
other => panic!("expected Terminal(Err) decode failure, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_transport_call_dispatches_via_connection() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice"));
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("hi".into())));
|
||||
let mut stream = transport
|
||||
.call(tid("foo"), serde_json::json!({"k": "v"}), ctx)
|
||||
.await;
|
||||
let _ = stream.next().await;
|
||||
assert_eq!(conn.captured_request_count(), 1);
|
||||
assert_eq!(transport.kind(), TransportKind::Remote);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_transport_authorize_returns_bound_principal() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn, sid("sess-1"), uid("alice"));
|
||||
let principal = transport.authorize().await.expect("authorize");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_subscribe_happens_before_request_send() {
|
||||
// Locks in BOTH halves of the subscribe-before-send contract:
|
||||
// 1. the subscription IS active by the time `execute` returns;
|
||||
// 2. the request HAS NOT been sent yet at that point.
|
||||
// A future refactor that eagerly sent the request inside
|
||||
// `execute` would still satisfy (1) but would break (2).
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let proxy = RemoteToolProxy::new(
|
||||
tid("foo"),
|
||||
sid("sess-1"),
|
||||
description_for("foo"),
|
||||
ToolCapabilities::default(),
|
||||
conn.clone(),
|
||||
);
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let ctx = ToolCallContext::new(call_id.clone());
|
||||
conn.enqueue_ok(ok_call_result(&call_id, ToolOutputWire::Text("ok".into())));
|
||||
let mut stream = proxy.execute(ctx, serde_json::json!(null)).await;
|
||||
{
|
||||
// The DashMap subscription read and the captured-requests check
|
||||
// are individually atomic. Single-threaded `#[tokio::test]`
|
||||
// execution means no other task can mutate either between the
|
||||
// two checks, so the pair is observationally simultaneous.
|
||||
assert!(
|
||||
conn.progress_senders.contains_key(&call_id),
|
||||
"subscription must be active before request send"
|
||||
);
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert!(
|
||||
guard.captured_requests.is_empty(),
|
||||
"request must not be sent before stream is polled"
|
||||
);
|
||||
}
|
||||
// Polling the stream is what actually drives the request future,
|
||||
// so the captured-requests vec only fills in once we start consuming.
|
||||
while stream.next().await.is_some() {}
|
||||
{
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(
|
||||
guard.captured_requests.len(),
|
||||
1,
|
||||
"request must have been sent during stream polling"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn notify_round_trips_through_connection_client() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let notification = JsonRpcNotification {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
session_id: Some(sid("sess-1")),
|
||||
seq: None,
|
||||
method: Method::Hook.as_wire_str().to_string(),
|
||||
params: serde_json::json!({
|
||||
"session_id": "sess-1",
|
||||
"tool_id": "foo",
|
||||
"call_id": "call-1",
|
||||
"event": { "type": "Cancel" }
|
||||
}),
|
||||
};
|
||||
let trait_handle: &dyn ConnectionClient = conn.as_ref();
|
||||
trait_handle
|
||||
.notify(notification.clone())
|
||||
.await
|
||||
.expect("notify succeeds");
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(guard.captured_notifications.len(), 1);
|
||||
let captured = &guard.captured_notifications[0];
|
||||
assert_eq!(captured.method, Method::Hook.as_wire_str());
|
||||
assert_eq!(captured.session_id, Some(sid("sess-1")));
|
||||
assert_eq!(
|
||||
captured.params.get("event").and_then(|v| v.get("type")),
|
||||
Some(&serde_json::Value::String("Cancel".to_string()))
|
||||
);
|
||||
assert_eq!(captured, ¬ification);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn json_rpc_id_is_unique_per_call() {
|
||||
let conn = Arc::new(MockConnection::default());
|
||||
let transport = RemoteTransport::new(conn.clone(), sid("sess-1"), uid("alice"));
|
||||
let call_a = ToolCallId::new_v7();
|
||||
let call_b = ToolCallId::new_v7();
|
||||
conn.enqueue_ok(ok_call_result(&call_a, ToolOutputWire::Text("a".into())));
|
||||
conn.enqueue_ok(ok_call_result(&call_b, ToolOutputWire::Text("b".into())));
|
||||
|
||||
let mut s1 = transport
|
||||
.call(
|
||||
tid("foo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::new(call_a.clone()),
|
||||
)
|
||||
.await;
|
||||
while s1.next().await.is_some() {}
|
||||
let mut s2 = transport
|
||||
.call(
|
||||
tid("foo"),
|
||||
serde_json::json!(null),
|
||||
ToolCallContext::new(call_b.clone()),
|
||||
)
|
||||
.await;
|
||||
while s2.next().await.is_some() {}
|
||||
|
||||
let guard = conn.inner.lock().expect("mutex");
|
||||
assert_eq!(guard.captured_requests.len(), 2);
|
||||
let id_a = match &guard.captured_requests[0].id {
|
||||
JsonRpcId::String(s) => s.clone(),
|
||||
JsonRpcId::Number(n) => n.to_string(),
|
||||
};
|
||||
let id_b = match &guard.captured_requests[1].id {
|
||||
JsonRpcId::String(s) => s.clone(),
|
||||
JsonRpcId::Number(n) => n.to_string(),
|
||||
};
|
||||
assert_ne!(id_a, id_b, "envelope ids must differ across calls");
|
||||
}
|
||||
@@ -0,0 +1,607 @@
|
||||
//! `ToolRegistry` trait coverage via a per-test mock backed by `DashMap`
|
||||
//! — lock-free per-key concurrent access mirrors the production
|
||||
//! direction even at the test layer. The mock implements the
|
||||
//! connection-scoped `ToolRegistry` trait surface.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use dashmap::DashMap;
|
||||
|
||||
use kigi_computer_hub_core::{
|
||||
ConnectionCleanupReport, ErasedTool, ResolvedTool, SessionCleanupReport, ToolHandle,
|
||||
ToolRegistry, ToolSessionBindOutcome, ToolSessionUnbindOutcome, resolver::CompoundResolver,
|
||||
};
|
||||
use kigi_tool_protocol::{
|
||||
ConnectionId, RegistrationOutcome, ServerId, SessionId, ToolDefinitionMode, ToolId,
|
||||
ToolRegistration, ToolServerRegistration, TransportKind, UserId,
|
||||
};
|
||||
use kigi_tool_runtime::{SearchSnapshot, ServerSummary, Tool, ToolCallContext, ToolError};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize, schemars::JsonSchema)]
|
||||
struct EmptyArgs {}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StubTool {
|
||||
id: ToolId,
|
||||
}
|
||||
|
||||
impl Tool for StubTool {
|
||||
type Args = EmptyArgs;
|
||||
type Output = serde_json::Value;
|
||||
|
||||
fn id(&self) -> ToolId {
|
||||
self.id.clone()
|
||||
}
|
||||
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new(self.id.as_str(), format!("stub for {}", self.id))
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
_ctx: ToolCallContext,
|
||||
_args: Self::Args,
|
||||
) -> Result<Self::Output, ToolError> {
|
||||
unreachable!("registry tests do not exercise execution")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MockEntry {
|
||||
registration: ToolRegistration,
|
||||
sessions: HashSet<SessionId>,
|
||||
}
|
||||
|
||||
/// Mock registry. Last-write-wins on duplicate registrations within a
|
||||
/// `(connection, tool_id)` slot — pinned here so the trait contract has
|
||||
/// a clear test fixture.
|
||||
#[derive(Debug, Default)]
|
||||
struct MockRegistry {
|
||||
entries: DashMap<(ConnectionId, ToolId), MockEntry>,
|
||||
by_session: DashMap<(SessionId, ToolId), ConnectionId>,
|
||||
handles: DashMap<ToolId, Arc<dyn ToolHandle>>,
|
||||
}
|
||||
|
||||
impl MockRegistry {
|
||||
fn install_handle(&self, tool: Arc<dyn ToolHandle>) {
|
||||
self.handles.insert(tool.id(), tool);
|
||||
}
|
||||
}
|
||||
|
||||
fn build_registration(tool: &ToolId, sessions: &[SessionId]) -> ToolRegistration {
|
||||
ToolRegistration {
|
||||
tool_id: tool.clone(),
|
||||
sessions: Some(sessions.to_vec()),
|
||||
user_id: UserId::new("alice").expect("valid user id"),
|
||||
server_id: None,
|
||||
description: ToolDescription::new(tool.as_str(), format!("desc for {tool}")),
|
||||
input_schema: None,
|
||||
capabilities: None,
|
||||
notification_schemas: None,
|
||||
transport_kind: TransportKind::Local,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolRegistry for MockRegistry {
|
||||
async fn register_tool(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolRegistration,
|
||||
) -> RegistrationOutcome {
|
||||
let key = (connection_id.clone(), reg.tool_id.clone());
|
||||
let sessions: HashSet<SessionId> = reg
|
||||
.sessions
|
||||
.as_ref()
|
||||
.map(|v| v.iter().cloned().collect())
|
||||
.unwrap_or_default();
|
||||
let updated = self
|
||||
.entries
|
||||
.insert(
|
||||
key,
|
||||
MockEntry {
|
||||
registration: reg.clone(),
|
||||
sessions: sessions.clone(),
|
||||
},
|
||||
)
|
||||
.is_some();
|
||||
for session in &sessions {
|
||||
self.by_session.insert(
|
||||
(session.clone(), reg.tool_id.clone()),
|
||||
connection_id.clone(),
|
||||
);
|
||||
}
|
||||
if updated {
|
||||
RegistrationOutcome::Updated {
|
||||
tool_id: reg.tool_id,
|
||||
generation: 1,
|
||||
}
|
||||
} else {
|
||||
RegistrationOutcome::Registered {
|
||||
tool_id: reg.tool_id,
|
||||
generation: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn register_server(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
reg: ToolServerRegistration,
|
||||
) -> Vec<RegistrationOutcome> {
|
||||
let mut outcomes = Vec::with_capacity(reg.tools.len());
|
||||
for tool in reg.tools {
|
||||
let tool_id = tool
|
||||
.derive_tool_id()
|
||||
.expect("test descriptions have valid tool ids");
|
||||
let registration = ToolRegistration {
|
||||
tool_id: tool_id.clone(),
|
||||
sessions: reg.sessions.clone(),
|
||||
user_id: reg.user_id.clone(),
|
||||
server_id: Some(reg.server_id.clone()),
|
||||
description: tool.description,
|
||||
input_schema: tool.input_schema,
|
||||
capabilities: tool.capabilities,
|
||||
notification_schemas: tool.notification_schemas,
|
||||
transport_kind: TransportKind::Remote,
|
||||
if_match_generation: None,
|
||||
metadata: None,
|
||||
};
|
||||
outcomes.push(
|
||||
self.register_tool(connection_id.clone(), registration)
|
||||
.await,
|
||||
);
|
||||
}
|
||||
outcomes
|
||||
}
|
||||
|
||||
async fn unregister_tool(&self, connection_id: &ConnectionId, tool: &ToolId) -> bool {
|
||||
let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone())) else {
|
||||
return false;
|
||||
};
|
||||
for session in &removed.sessions {
|
||||
self.by_session
|
||||
.remove_if(&(session.clone(), tool.clone()), |_, owner| {
|
||||
owner == connection_id
|
||||
});
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn unregister_server(&self, connection_id: &ConnectionId, server: &ServerId) -> usize {
|
||||
let to_remove: Vec<ToolId> = self
|
||||
.entries
|
||||
.iter()
|
||||
.filter(|r| {
|
||||
r.key().0 == *connection_id
|
||||
&& r.value().registration.server_id.as_ref() == Some(server)
|
||||
})
|
||||
.map(|r| r.key().1.clone())
|
||||
.collect();
|
||||
let mut removed = 0usize;
|
||||
for tool in to_remove {
|
||||
if self.unregister_tool(connection_id, &tool).await {
|
||||
removed += 1;
|
||||
}
|
||||
}
|
||||
removed
|
||||
}
|
||||
|
||||
async fn bind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionBindOutcome {
|
||||
let key = (connection_id.clone(), tool.clone());
|
||||
let Some(mut entry) = self.entries.get_mut(&key) else {
|
||||
return ToolSessionBindOutcome::UnknownTool;
|
||||
};
|
||||
if !entry.value_mut().sessions.insert(session_id.clone()) {
|
||||
return ToolSessionBindOutcome::AlreadyBound;
|
||||
}
|
||||
self.by_session
|
||||
.insert((session_id.clone(), tool.clone()), connection_id.clone());
|
||||
ToolSessionBindOutcome::Bound
|
||||
}
|
||||
|
||||
async fn unbind_tool_session(
|
||||
&self,
|
||||
connection_id: &ConnectionId,
|
||||
tool: &ToolId,
|
||||
session_id: &SessionId,
|
||||
) -> ToolSessionUnbindOutcome {
|
||||
let key = (connection_id.clone(), tool.clone());
|
||||
let Some(mut entry) = self.entries.get_mut(&key) else {
|
||||
return ToolSessionUnbindOutcome::UnknownTool;
|
||||
};
|
||||
if !entry.value_mut().sessions.remove(session_id) {
|
||||
return ToolSessionUnbindOutcome::NotBound;
|
||||
}
|
||||
self.by_session
|
||||
.remove_if(&(session_id.clone(), tool.clone()), |_, owner| {
|
||||
owner == connection_id
|
||||
});
|
||||
ToolSessionUnbindOutcome::Unbound
|
||||
}
|
||||
|
||||
async fn drop_connection(&self, connection_id: &ConnectionId) -> ConnectionCleanupReport {
|
||||
let to_remove: Vec<ToolId> = self
|
||||
.entries
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *connection_id)
|
||||
.map(|r| r.key().1.clone())
|
||||
.collect();
|
||||
let mut report = ConnectionCleanupReport::default();
|
||||
for tool in to_remove {
|
||||
if let Some((_, removed)) = self.entries.remove(&(connection_id.clone(), tool.clone()))
|
||||
{
|
||||
report.tools_dropped += 1;
|
||||
for session in removed.sessions {
|
||||
if self
|
||||
.by_session
|
||||
.remove_if(&(session, tool.clone()), |_, owner| owner == connection_id)
|
||||
.is_some()
|
||||
{
|
||||
report.session_bindings_cleared += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
fn find_tool(&self, session: &SessionId, tool: &ToolId) -> Option<ResolvedTool> {
|
||||
let owner = self
|
||||
.by_session
|
||||
.get(&(session.clone(), tool.clone()))?
|
||||
.value()
|
||||
.clone();
|
||||
let entry = self.entries.get(&(owner, tool.clone()))?;
|
||||
let registration = entry.value().registration.clone();
|
||||
let handle = self.handles.get(tool)?.value().clone();
|
||||
match registration.transport_kind {
|
||||
TransportKind::Local => Some(ResolvedTool::Local {
|
||||
tool: handle,
|
||||
registration,
|
||||
}),
|
||||
TransportKind::Remote => Some(ResolvedTool::Remote {
|
||||
proxy: handle,
|
||||
registration,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_tools(&self, session: &SessionId, _mode: &ToolDefinitionMode) -> Vec<ToolDescription> {
|
||||
self.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.filter_map(|r| {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
self.entries
|
||||
.get(&(owner, tool_id))
|
||||
.map(|e| e.value().registration.description.clone())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn list_servers(&self, session: &SessionId) -> Vec<ServerSummary> {
|
||||
let mut by_server: HashMap<ServerId, Vec<String>> = HashMap::new();
|
||||
for r in self.by_session.iter().filter(|r| r.key().0 == *session) {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
if let Some(entry) = self.entries.get(&(owner, tool_id)) {
|
||||
let reg = &entry.value().registration;
|
||||
let server = reg
|
||||
.server_id
|
||||
.clone()
|
||||
.unwrap_or_else(|| ServerId::synthesize_for_tool(r.value(), ®.tool_id));
|
||||
by_server
|
||||
.entry(server)
|
||||
.or_default()
|
||||
.push(reg.tool_id.as_str().to_string());
|
||||
}
|
||||
}
|
||||
by_server
|
||||
.into_iter()
|
||||
.map(|(server, mut names)| {
|
||||
names.sort();
|
||||
ServerSummary {
|
||||
name: server.into_inner(),
|
||||
description: None,
|
||||
tool_names: names,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn search(&self, session: &SessionId, query: &str, limit: usize) -> SearchSnapshot {
|
||||
let matches: Vec<_> = self
|
||||
.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.filter_map(|r| {
|
||||
let owner = r.value().clone();
|
||||
let tool_id = r.key().1.clone();
|
||||
let entry = self.entries.get(&(owner, tool_id))?;
|
||||
let reg = &entry.value().registration;
|
||||
if reg.tool_id.as_str().contains(query) {
|
||||
Some(kigi_tool_runtime::ToolSearchResult {
|
||||
tool_name: reg.tool_id.as_str().to_string(),
|
||||
server_name: reg
|
||||
.server_id
|
||||
.as_ref()
|
||||
.map(|s| s.as_str().to_string())
|
||||
.unwrap_or_default(),
|
||||
description: reg.description.description.clone(),
|
||||
score: 1.0,
|
||||
parameters: vec![],
|
||||
input_schema: serde_json::Value::Null,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.take(limit)
|
||||
.collect();
|
||||
SearchSnapshot {
|
||||
results: matches,
|
||||
total_hidden_tools: 0,
|
||||
is_ready: true,
|
||||
}
|
||||
}
|
||||
|
||||
async fn unregister_session(&self, session: &SessionId) -> SessionCleanupReport {
|
||||
let pairs: Vec<(ToolId, ConnectionId)> = self
|
||||
.by_session
|
||||
.iter()
|
||||
.filter(|r| r.key().0 == *session)
|
||||
.map(|r| (r.key().1.clone(), r.value().clone()))
|
||||
.collect();
|
||||
let mut report = SessionCleanupReport::default();
|
||||
for (tool_id, owner) in pairs {
|
||||
self.by_session
|
||||
.remove_if(&(session.clone(), tool_id.clone()), |_, value| {
|
||||
value == &owner
|
||||
});
|
||||
if let Some(mut entry) = self.entries.get_mut(&(owner, tool_id)) {
|
||||
entry.value_mut().sessions.remove(session);
|
||||
report.tools_touched += 1;
|
||||
if entry.value().sessions.is_empty() {
|
||||
report.tools_left_orphaned += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
fn tool_sessions(&self, connection_id: &ConnectionId, tool: &ToolId) -> HashSet<SessionId> {
|
||||
self.entries
|
||||
.get(&(connection_id.clone(), tool.clone()))
|
||||
.map(|r| r.value().sessions.clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn list_servers_for_user(
|
||||
&self,
|
||||
_user_id: &kigi_tool_protocol::UserId,
|
||||
) -> Vec<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn get_server_record(
|
||||
&self,
|
||||
_connection_id: &ConnectionId,
|
||||
) -> Option<kigi_computer_hub_core::registry::ServerRecord> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("valid session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("valid tool id")
|
||||
}
|
||||
|
||||
fn cid(s: &str) -> ConnectionId {
|
||||
ConnectionId::new(s).expect("valid connection id")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn register_then_find_returns_local_resolution() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
let outcome = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(matches!(outcome, RegistrationOutcome::Registered { .. }));
|
||||
let resolved = reg
|
||||
.find_tool(&sid("sess-1"), &tid("foo"))
|
||||
.expect("registration found");
|
||||
match resolved {
|
||||
ResolvedTool::Local { registration, .. } => {
|
||||
assert_eq!(registration.tool_id, tid("foo"));
|
||||
assert!(
|
||||
registration
|
||||
.sessions
|
||||
.as_ref()
|
||||
.is_some_and(|s| s.contains(&sid("sess-1")))
|
||||
);
|
||||
}
|
||||
other => panic!("expected Local, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn find_in_other_session_returns_none() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("foo")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn duplicate_registration_yields_updated_outcome() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
let first = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let second = reg
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(matches!(first, RegistrationOutcome::Registered { .. }));
|
||||
assert!(matches!(second, RegistrationOutcome::Updated { .. }));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_tool_removes_only_that_entry() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-1")]))
|
||||
.await;
|
||||
assert!(reg.unregister_tool(&cid("c1"), &tid("foo")).await);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_session_drops_session_binding_and_leaves_orphan_count() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]),
|
||||
)
|
||||
.await;
|
||||
let report = reg.unregister_session(&sid("sess-1")).await;
|
||||
assert_eq!(report.tools_touched, 2);
|
||||
// `foo` had only sess-1 → orphaned. `bar` had sess-2 left → not orphaned.
|
||||
assert_eq!(report.tools_left_orphaned, 1);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("bar")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_tools_filters_by_session() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("bar"), &[sid("sess-2")]))
|
||||
.await;
|
||||
let s1 = reg.list_tools(&sid("sess-1"), &ToolDefinitionMode::Full);
|
||||
let s2 = reg.list_tools(&sid("sess-2"), &ToolDefinitionMode::Full);
|
||||
assert_eq!(s1.len(), 1);
|
||||
assert_eq!(s1[0].name, "foo");
|
||||
assert_eq!(s2.len(), 1);
|
||||
assert_eq!(s2[0].name, "bar");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_servers_groups_by_owning_server() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let summaries = reg.list_servers(&sid("sess-1"));
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0].tool_count(), 1);
|
||||
assert_eq!(summaries[0].tool_names[0], "foo");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_returns_substring_matches() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foobar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("foobar"), &[sid("sess-1")]),
|
||||
)
|
||||
.await;
|
||||
let snap = reg.search(&sid("sess-1"), "foo", 10);
|
||||
assert_eq!(snap.results.len(), 2);
|
||||
assert!(snap.is_ready);
|
||||
assert_eq!(snap.total_hidden_tools, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn registry_drives_compound_resolver() {
|
||||
let registry = Arc::new(MockRegistry::default());
|
||||
registry.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
registry
|
||||
.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
let resolver = CompoundResolver::local_only(registry as Arc<dyn ToolRegistry>);
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("foo")).is_some());
|
||||
assert!(resolver.resolve(&sid("sess-1"), &tid("missing")).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bind_and_unbind_tool_session_round_trips_visibility() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[]))
|
||||
.await;
|
||||
// Empty sessions: tool is registered but unreachable.
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
let outcome = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(outcome, ToolSessionBindOutcome::Bound);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_some());
|
||||
let again = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(again, ToolSessionBindOutcome::AlreadyBound);
|
||||
let unbind = reg
|
||||
.unbind_tool_session(&cid("c1"), &tid("foo"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(unbind, ToolSessionUnbindOutcome::Unbound);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
let unknown = reg
|
||||
.bind_tool_session(&cid("c1"), &tid("missing"), &sid("sess-1"))
|
||||
.await;
|
||||
assert_eq!(unknown, ToolSessionBindOutcome::UnknownTool);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drop_connection_releases_every_owned_tool() {
|
||||
let reg = MockRegistry::default();
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("foo") })));
|
||||
reg.install_handle(Arc::new(ErasedTool::new(StubTool { id: tid("bar") })));
|
||||
reg.register_tool(cid("c1"), build_registration(&tid("foo"), &[sid("sess-1")]))
|
||||
.await;
|
||||
reg.register_tool(
|
||||
cid("c1"),
|
||||
build_registration(&tid("bar"), &[sid("sess-1"), sid("sess-2")]),
|
||||
)
|
||||
.await;
|
||||
let report = reg.drop_connection(&cid("c1")).await;
|
||||
assert_eq!(report.tools_dropped, 2);
|
||||
assert_eq!(report.session_bindings_cleared, 3);
|
||||
assert!(reg.find_tool(&sid("sess-1"), &tid("foo")).is_none());
|
||||
assert!(reg.find_tool(&sid("sess-2"), &tid("bar")).is_none());
|
||||
assert!(reg.tool_sessions(&cid("c1"), &tid("foo")).is_empty());
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
//! Behavioural coverage for the `Transport` trait, `Principal` builder,
|
||||
//! and `TransportKind` re-export.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use kigi_computer_hub_core::{Principal, Transport, TransportKind};
|
||||
use kigi_tool_protocol::{SessionId, ToolId, UserId};
|
||||
use kigi_tool_runtime::{
|
||||
ToolCallContext, ToolError, ToolStream, ToolStreamItem, TypedToolOutput, terminal_only,
|
||||
};
|
||||
|
||||
fn uid(s: &str) -> UserId {
|
||||
UserId::new(s).expect("test user id")
|
||||
}
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("test session id")
|
||||
}
|
||||
|
||||
fn tid(s: &str) -> ToolId {
|
||||
ToolId::new(s).expect("test tool id")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct EchoTransport {
|
||||
kind: TransportKind,
|
||||
user: UserId,
|
||||
session: SessionId,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transport for EchoTransport {
|
||||
fn kind(&self) -> TransportKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
async fn authorize(&self) -> Result<Principal, ToolError> {
|
||||
Ok(Principal::new(self.user.clone())
|
||||
.with_session(self.session.clone())
|
||||
.with_scope("tool.invoke"))
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
tool_id: ToolId,
|
||||
args: Value,
|
||||
_ctx: ToolCallContext,
|
||||
) -> ToolStream<TypedToolOutput> {
|
||||
terminal_only(Ok(TypedToolOutput::from_value(tool_id, args)))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn boxed_transport_compiles_and_dispatches() {
|
||||
let boxed: Box<dyn Transport> = Box::new(EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
});
|
||||
let mut stream = boxed
|
||||
.call(tid("echo"), json!({"k": "v"}), ToolCallContext::default())
|
||||
.await;
|
||||
let item = futures::StreamExt::next(&mut stream)
|
||||
.await
|
||||
.expect("at least one item");
|
||||
match item {
|
||||
ToolStreamItem::Terminal(Ok(typed)) => assert_eq!(typed.value, json!({"k": "v"})),
|
||||
other => panic!("expected Terminal(Ok), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn kind_distinguishes_local_and_remote() {
|
||||
let local = EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
let remote = EchoTransport {
|
||||
kind: TransportKind::Remote,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
assert_eq!(local.kind(), TransportKind::Local);
|
||||
assert_eq!(remote.kind(), TransportKind::Remote);
|
||||
assert_ne!(local.kind(), remote.kind());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn authorize_returns_bound_principal() {
|
||||
let t = EchoTransport {
|
||||
kind: TransportKind::Local,
|
||||
user: uid("alice"),
|
||||
session: sid("sess-1"),
|
||||
};
|
||||
let principal = t.authorize().await.expect("authorize succeeds");
|
||||
assert_eq!(principal.user_id, uid("alice"));
|
||||
assert!(principal.authorizes_session(&sid("sess-1")));
|
||||
assert!(!principal.authorizes_session(&sid("sess-other")));
|
||||
assert!(principal.has_scope("tool.invoke"));
|
||||
assert!(!principal.has_scope("admin"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_builder_chains_in_order() {
|
||||
let principal = Principal::new(uid("alice"))
|
||||
.with_session(sid("sess-a"))
|
||||
.with_session(sid("sess-b"))
|
||||
.with_scope("tool.invoke")
|
||||
.with_scope("tool.search")
|
||||
.with_audience("dispatcher.example");
|
||||
assert_eq!(principal.session_ids, vec![sid("sess-a"), sid("sess-b")]);
|
||||
assert_eq!(principal.scopes, vec!["tool.invoke", "tool.search"]);
|
||||
assert_eq!(principal.audiences, vec!["dispatcher.example"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_supports_multi_session_tokens() {
|
||||
let p = Principal::new(uid("alice"))
|
||||
.with_session(sid("sess-1"))
|
||||
.with_session(sid("sess-2"));
|
||||
assert!(p.authorizes_session(&sid("sess-1")));
|
||||
assert!(p.authorizes_session(&sid("sess-2")));
|
||||
assert!(!p.authorizes_session(&sid("sess-3")));
|
||||
assert_eq!(p.session_ids.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn principal_default_state_is_empty() {
|
||||
let p = Principal::new(uid("alice"));
|
||||
assert!(p.session_ids.is_empty());
|
||||
assert!(p.scopes.is_empty());
|
||||
assert!(p.audiences.is_empty());
|
||||
assert!(!p.has_scope("anything"));
|
||||
assert!(!p.authorizes_session(&sid("sess")));
|
||||
}
|
||||
@@ -0,0 +1,203 @@
|
||||
//! `is_workspace_unavailable` recognizer coverage, pinned against the real
|
||||
//! wire decode path (`error_from_envelope` / `tool_error_from_wire`).
|
||||
|
||||
use kigi_computer_hub_core::{error_from_envelope, is_workspace_unavailable, tool_error_from_wire};
|
||||
use kigi_tool_protocol::{
|
||||
JsonRpcError, ToolErrorWire, WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase,
|
||||
WorkspaceGoneReason, WorkspaceUnavailableDetails, workspace_unavailable_wire,
|
||||
};
|
||||
use kigi_tool_runtime::{ToolError, ToolErrorKind};
|
||||
use serde_json::json;
|
||||
|
||||
const REASONS: [WorkspaceGoneReason; 5] = [
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGoneReason::Shutdown,
|
||||
WorkspaceGoneReason::NotBound,
|
||||
WorkspaceGoneReason::InstanceGone,
|
||||
];
|
||||
const PHASES: [WorkspaceGonePhase; 2] = [
|
||||
WorkspaceGonePhase::InFlightCancelled,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
];
|
||||
|
||||
fn envelope_for(wire: &ToolErrorWire) -> JsonRpcError {
|
||||
JsonRpcError {
|
||||
// -32005 is the best-effort numeric companion (`tool_server_gone`);
|
||||
// recognition keys on `data.details.code`, not the numeric.
|
||||
code: -32005,
|
||||
message: "workspace server gone".to_owned(),
|
||||
data: Some(serde_json::to_value(wire).unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round_trip_through_envelope_is_recognized_for_every_reason_and_phase() {
|
||||
for reason in REASONS {
|
||||
for phase in PHASES {
|
||||
let wire = workspace_unavailable_wire(reason, phase);
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
|
||||
assert!(
|
||||
is_workspace_unavailable(&err),
|
||||
"should recognize {reason:?}/{phase:?}",
|
||||
);
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
|
||||
// The full structured payload survives into `ToolError::details`,
|
||||
// so a caller can branch on code/reason/phase/retryable.
|
||||
let details: WorkspaceUnavailableDetails =
|
||||
serde_json::from_value(err.details.expect("details survive")).unwrap();
|
||||
assert_eq!(
|
||||
details,
|
||||
WorkspaceUnavailableDetails {
|
||||
code: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
reason,
|
||||
phase,
|
||||
retryable: true,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tool_error_from_wire_directly_is_recognized() {
|
||||
let wire = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
);
|
||||
let err = tool_error_from_wire(wire);
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
let details = err.details.expect("details survive");
|
||||
assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
|
||||
assert_eq!(details["reason"], json!("disconnect"));
|
||||
assert_eq!(details["phase"], json!("route_missing"));
|
||||
assert_eq!(details["retryable"], json!(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wire_to_tool_error_to_wire_preserves_outer_subcode() {
|
||||
// Keying the identity on details.code lets From<ToolError> for ToolErrorWire
|
||||
// rebuild the outer subcode on re-serialization.
|
||||
let original = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGonePhase::InFlightCancelled,
|
||||
);
|
||||
let tool_error = tool_error_from_wire(original);
|
||||
let back: ToolErrorWire = tool_error.into();
|
||||
let ToolErrorWire::Custom { subcode, .. } = back else {
|
||||
panic!("expected Custom variant");
|
||||
};
|
||||
assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recognized_with_unknown_reason_and_phase() {
|
||||
// Recognition is decoupled from the typed reason/phase enums: a newer hub
|
||||
// emitting unknown values is still recognized (it keys only on `code`).
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "from a newer hub".to_owned(),
|
||||
details: Some(json!({
|
||||
"code": WORKSPACE_UNAVAILABLE_SUBCODE,
|
||||
"reason": "brand_new_reason",
|
||||
"phase": "brand_new_phase",
|
||||
"retryable": true,
|
||||
})),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
// End-to-end decode → typed-parse → `Unknown`, the path consumers read by.
|
||||
let details: WorkspaceUnavailableDetails =
|
||||
serde_json::from_value(err.details.expect("details survive")).unwrap();
|
||||
assert_eq!(details.reason, WorkspaceGoneReason::Unknown);
|
||||
assert_eq!(details.phase, WorkspaceGonePhase::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoded_custom_with_none_details_is_recognized_via_canonical_code() {
|
||||
// Wire `details: None` decodes through `ToolError::custom`, which repopulates
|
||||
// `details = {"code": subcode}`, so it IS recognized — contrast the hand-built
|
||||
// no-details case in `custom_error_without_any_details_is_not_recognized`.
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "no structured details".to_owned(),
|
||||
details: None,
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoded_custom_without_code_key_is_not_recognized() {
|
||||
// The central correctness property: recognition keys on the surviving
|
||||
// `details.code`, NOT the outer `Custom.subcode`. Here the outer subcode
|
||||
// matches, but `with_details` overwrote the auto-populated `code`, so the
|
||||
// decoded error must NOT be recognized.
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: "details lack code".to_owned(),
|
||||
details: Some(json!({ "reason": "disconnect" })),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_custom_code_is_not_recognized() {
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: "some_other_error".to_owned(),
|
||||
message: "nope".to_owned(),
|
||||
details: Some(json!({ "code": "some_other_error" })),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_eq!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn numeric_only_tool_server_gone_without_data_is_not_recognized() {
|
||||
// Recognition is by the data payload, never the numeric code: a bare -32005
|
||||
// with no `data` decodes to a `jsonrpc_-32005` custom error, not recognized.
|
||||
let err = error_from_envelope(JsonRpcError {
|
||||
code: -32005,
|
||||
message: "tool server gone".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_error_without_any_details_is_not_recognized() {
|
||||
// Hand-built Custom with no `details` (no `code`) — unlike a wire `details:
|
||||
// None`, nothing repopulates `code` here, so it is not recognized.
|
||||
let err = ToolError::new(ToolErrorKind::Custom, "no details at all");
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_custom_error_with_matching_code_is_not_recognized() {
|
||||
// The kind guard matters: a non-Custom error carrying a matching
|
||||
// `details.code` must still be rejected.
|
||||
let err = ToolError::new(ToolErrorKind::NetworkError, "socket closed")
|
||||
.with_details(json!({ "code": WORKSPACE_UNAVAILABLE_SUBCODE }));
|
||||
assert_ne!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_custom_decoded_error_is_not_recognized() {
|
||||
let wire = ToolErrorWire::ToolNotFound {
|
||||
tool_id: kigi_tool_protocol::ToolId::new("ns:tool").unwrap(),
|
||||
};
|
||||
let err = error_from_envelope(envelope_for(&wire));
|
||||
assert_ne!(err.kind, ToolErrorKind::Custom);
|
||||
assert!(!is_workspace_unavailable(&err));
|
||||
|
||||
assert!(!is_workspace_unavailable(&ToolError::network_error(
|
||||
"socket closed"
|
||||
)));
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-computer-hub-mcp-adapter"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Bridge between MCP servers and the xAI Computer Hub, registering MCP-discovered tools as native hub tools."
|
||||
|
||||
[features]
|
||||
metrics = ["dep:prometheus"]
|
||||
|
||||
[dependencies]
|
||||
async-trait = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
tracing = { workspace = true }
|
||||
prometheus = { workspace = true, optional = true }
|
||||
|
||||
kigi-tool-protocol = { workspace = true }
|
||||
kigi-tool-runtime = { workspace = true }
|
||||
kigi-tool-types = { workspace = true }
|
||||
kigi-computer-hub-sdk = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["full", "test-util"] }
|
||||
futures = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,781 @@
|
||||
//! Core bridge that connects an MCP server to the computer hub.
|
||||
//!
|
||||
//! [`McpBridge`] discovers tools from an [`McpTransport`] and registers
|
||||
//! them with a hub `ToolServer` via one `ToolServerHandler` per
|
||||
//! tool. Incoming hub calls are translated to MCP `tools/call` and the
|
||||
//! response is mapped back to [`ToolOutputWire`].
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use kigi_tool_protocol::{McpBlock, SessionId, ToolId, ToolOutputWire};
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolError, ToolStream, TypedToolOutput, terminal_only};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::transport::McpTransport;
|
||||
use crate::types::{McpCallResult, McpContent, McpError, McpServerInfo, McpToolDefinition};
|
||||
|
||||
/// Configuration for an [`McpBridge`] instance.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpBridgeConfig {
|
||||
/// Hub session to bind tools to.
|
||||
pub session_id: SessionId,
|
||||
/// Optional namespace prefix for tool descriptions.
|
||||
pub namespace: Option<String>,
|
||||
}
|
||||
|
||||
/// Result of a successful [`McpBridge::connect`] call.
|
||||
///
|
||||
/// Contains the bridge handle and the server info returned during the
|
||||
/// MCP initialize handshake.
|
||||
pub struct McpBridgeHandle {
|
||||
/// The bridge managing the MCP-to-hub tool registrations.
|
||||
pub bridge: McpBridge,
|
||||
/// Server metadata from the MCP `initialize` response.
|
||||
pub server_info: McpServerInfo,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for McpBridgeHandle {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("McpBridgeHandle")
|
||||
.field("server_info", &self.server_info.name)
|
||||
.field("tool_count", &self.bridge.tool_count())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// Bridges an MCP server's tools into the computer hub.
|
||||
///
|
||||
/// On construction the bridge performs the MCP `initialize` handshake,
|
||||
/// discovers tools via `tools/list`, and builds a handler
|
||||
/// for each one. Callers wire these handlers into a
|
||||
/// [`kigi_computer_hub_sdk::ToolServerBuilder`] to register them
|
||||
/// with the hub.
|
||||
///
|
||||
/// Callers **must** call [`McpBridge::shutdown`] before dropping to
|
||||
/// close the underlying MCP transport cleanly. If the bridge is dropped
|
||||
/// without an explicit shutdown, a best-effort `close()` is spawned on
|
||||
/// the tokio runtime (mirroring `ToolServer`'s drop behavior).
|
||||
pub struct McpBridge {
|
||||
transport: Arc<dyn McpTransport>,
|
||||
handlers: Vec<Arc<McpToolHandler>>,
|
||||
server_info: McpServerInfo,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for McpBridge {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("McpBridge")
|
||||
.field("server", &self.server_info.name)
|
||||
.field("tool_count", &self.handlers.len())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl McpBridge {
|
||||
/// Initialize the MCP server, discover its tools, and build handlers.
|
||||
///
|
||||
/// Returns `Err` if the MCP handshake or tool discovery fails.
|
||||
pub async fn connect(
|
||||
transport: Arc<dyn McpTransport>,
|
||||
config: &McpBridgeConfig,
|
||||
) -> Result<McpBridgeHandle, McpError> {
|
||||
let server_info = match transport.initialize().await {
|
||||
Ok(info) => info,
|
||||
Err(e) => {
|
||||
crate::metrics::mcp_error();
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
info!(
|
||||
server_name = %server_info.name,
|
||||
version = %server_info.version,
|
||||
"MCP server initialized"
|
||||
);
|
||||
|
||||
let tools = match transport.list_tools().await {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
// Close the transport so the initialized connection is not leaked
|
||||
// when list_tools fails after a successful initialize.
|
||||
if let Err(close_err) = transport.close().await {
|
||||
warn!(
|
||||
?close_err,
|
||||
"failed to close transport after list_tools error"
|
||||
);
|
||||
}
|
||||
crate::metrics::mcp_error();
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
debug!(
|
||||
server_name = %server_info.name,
|
||||
tool_count = tools.len(),
|
||||
"discovered MCP tools"
|
||||
);
|
||||
|
||||
let handlers: Vec<Arc<McpToolHandler>> = tools
|
||||
.into_iter()
|
||||
.filter_map(|def| {
|
||||
let tool_id = match ToolId::new(&def.name) {
|
||||
Ok(id) => id,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
tool_name = %def.name,
|
||||
%err,
|
||||
"skipping MCP tool with invalid name"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(Arc::new(McpToolHandler {
|
||||
tool_id,
|
||||
definition: def,
|
||||
transport: Arc::clone(&transport),
|
||||
namespace: config.namespace.clone(),
|
||||
}))
|
||||
})
|
||||
.collect();
|
||||
|
||||
crate::metrics::mcp_tools_bridged_set(handlers.len() as i64);
|
||||
|
||||
let bridge = McpBridge {
|
||||
transport,
|
||||
handlers,
|
||||
server_info: server_info.clone(),
|
||||
};
|
||||
|
||||
Ok(McpBridgeHandle {
|
||||
bridge,
|
||||
server_info,
|
||||
})
|
||||
}
|
||||
|
||||
/// Handlers to register with a [`kigi_computer_hub_sdk::ToolServerBuilder`].
|
||||
///
|
||||
/// Each handler implements `ToolServerHandler` for one MCP tool.
|
||||
pub fn handlers(&self) -> &[Arc<McpToolHandler>] {
|
||||
&self.handlers
|
||||
}
|
||||
|
||||
/// Server metadata from the MCP `initialize` response.
|
||||
pub fn server_info(&self) -> &McpServerInfo {
|
||||
&self.server_info
|
||||
}
|
||||
|
||||
/// Number of tools discovered and registered.
|
||||
pub fn tool_count(&self) -> usize {
|
||||
self.handlers.len()
|
||||
}
|
||||
|
||||
/// Close the underlying MCP transport.
|
||||
pub async fn shutdown(&self) -> Result<(), McpError> {
|
||||
crate::metrics::mcp_tools_bridged_set(0);
|
||||
self.transport.close().await
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for McpBridge {
|
||||
fn drop(&mut self) {
|
||||
crate::metrics::mcp_tools_bridged_set(0);
|
||||
let transport = Arc::clone(&self.transport);
|
||||
if tokio::runtime::Handle::try_current().is_ok() {
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = transport.close().await {
|
||||
warn!(?err, "best-effort transport close on drop failed");
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hub-facing handler for a single MCP tool.
|
||||
///
|
||||
/// Translates hub `tool_call_request` frames into MCP `tools/call`
|
||||
/// invocations and maps the result back to [`ToolOutputWire`].
|
||||
pub struct McpToolHandler {
|
||||
tool_id: ToolId,
|
||||
definition: McpToolDefinition,
|
||||
transport: Arc<dyn McpTransport>,
|
||||
namespace: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for McpToolHandler {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("McpToolHandler")
|
||||
.field("tool_id", &self.tool_id)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl kigi_computer_hub_sdk::ToolServerHandler for McpToolHandler {
|
||||
fn tool_id(&self) -> ToolId {
|
||||
self.tool_id.clone()
|
||||
}
|
||||
|
||||
fn description(&self) -> ToolDescription {
|
||||
let desc = ToolDescription::new(
|
||||
self.definition.name.clone(),
|
||||
self.definition.description.clone().unwrap_or_default(),
|
||||
);
|
||||
match self.namespace {
|
||||
Some(ref ns) => desc.with_namespace(ns.clone()),
|
||||
None => desc,
|
||||
}
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> Option<Value> {
|
||||
self.definition.input_schema.clone()
|
||||
}
|
||||
|
||||
async fn handle_call(&self, _ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
let _start = std::time::Instant::now();
|
||||
let tool_id = self.tool_id.clone();
|
||||
let result = self
|
||||
.transport
|
||||
.call_tool(self.definition.name.as_str(), args)
|
||||
.await;
|
||||
crate::metrics::mcp_call_duration_observe(_start.elapsed().as_secs_f64());
|
||||
|
||||
let terminal = match result {
|
||||
Ok(call_result) => {
|
||||
let output = translate_mcp_result(&call_result);
|
||||
serde_json::to_value(output)
|
||||
.map(|value| TypedToolOutput::from_value(tool_id, value))
|
||||
.map_err(|e| {
|
||||
crate::metrics::mcp_error();
|
||||
ToolError::execution(self.tool_id.clone(), e.to_string()).with_source(e)
|
||||
})
|
||||
}
|
||||
Err(mcp_err) => {
|
||||
crate::metrics::mcp_error();
|
||||
Err(ToolError::execution(
|
||||
self.tool_id.clone(),
|
||||
format!("{mcp_err}"),
|
||||
))
|
||||
}
|
||||
};
|
||||
|
||||
terminal_only(terminal)
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert an [`McpCallResult`] into the wire output format.
|
||||
///
|
||||
/// - **Error responses** (`is_error: true`): concatenates text-only blocks
|
||||
/// into a single [`ToolOutputWire::Text`], discarding non-text content
|
||||
/// (with a warning when content is dropped).
|
||||
/// - **Empty content**: returns `ToolOutputWire::Text("")` regardless of
|
||||
/// `is_error` — matches side-effect-only MCP tools.
|
||||
/// - **Single text block**: returns [`ToolOutputWire::Text`] directly.
|
||||
/// - **Multi-block / non-text**: returns [`ToolOutputWire::Mcp`] with
|
||||
/// structured blocks.
|
||||
fn translate_mcp_result(result: &McpCallResult) -> ToolOutputWire {
|
||||
if result.content.is_empty() {
|
||||
return ToolOutputWire::Text(String::new());
|
||||
}
|
||||
|
||||
if result.is_error {
|
||||
let error_text = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
McpContent::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
if error_text.is_empty() {
|
||||
warn!(
|
||||
content_count = result.content.len(),
|
||||
"MCP error response contained only non-text blocks; content dropped"
|
||||
);
|
||||
}
|
||||
return ToolOutputWire::Text(error_text);
|
||||
}
|
||||
|
||||
// Single text block → flat text output.
|
||||
if result.content.len() == 1
|
||||
&& let Some(McpContent::Text { text }) = result.content.first()
|
||||
{
|
||||
return ToolOutputWire::Text(text.clone());
|
||||
}
|
||||
|
||||
let blocks: Vec<McpBlock> = result
|
||||
.content
|
||||
.iter()
|
||||
.map(|c| match c {
|
||||
McpContent::Text { text } => McpBlock::Text { text: text.clone() },
|
||||
McpContent::Image { mime_type, data } => McpBlock::Image {
|
||||
mime_type: mime_type.clone(),
|
||||
data: data.clone(),
|
||||
},
|
||||
McpContent::Resource {
|
||||
uri,
|
||||
mime_type,
|
||||
text,
|
||||
} => McpBlock::Resource {
|
||||
uri: uri.clone(),
|
||||
mime_type: mime_type.clone(),
|
||||
text: text.clone(),
|
||||
},
|
||||
})
|
||||
.collect();
|
||||
|
||||
ToolOutputWire::Mcp { blocks }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::{McpCallResult, McpContent, McpServerInfo, McpToolDefinition};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
struct MockTransport {
|
||||
server_info: McpServerInfo,
|
||||
tools: Vec<McpToolDefinition>,
|
||||
call_response: Mutex<Option<McpCallResult>>,
|
||||
call_error: Mutex<Option<McpError>>,
|
||||
closed: AtomicBool,
|
||||
last_call: Mutex<Option<(String, Value)>>,
|
||||
}
|
||||
|
||||
impl MockTransport {
|
||||
fn new(server_info: McpServerInfo, tools: Vec<McpToolDefinition>) -> Self {
|
||||
Self {
|
||||
server_info,
|
||||
tools,
|
||||
call_response: Mutex::new(None),
|
||||
call_error: Mutex::new(None),
|
||||
closed: AtomicBool::new(false),
|
||||
last_call: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn with_call_response(self, response: McpCallResult) -> Self {
|
||||
Self {
|
||||
call_response: Mutex::new(Some(response)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
fn with_call_error(self, error: McpError) -> Self {
|
||||
Self {
|
||||
call_error: Mutex::new(Some(error)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for MockTransport {
|
||||
async fn initialize(&self) -> Result<McpServerInfo, McpError> {
|
||||
Ok(self.server_info.clone())
|
||||
}
|
||||
|
||||
async fn list_tools(&self) -> Result<Vec<McpToolDefinition>, McpError> {
|
||||
Ok(self.tools.clone())
|
||||
}
|
||||
|
||||
async fn call_tool(&self, name: &str, arguments: Value) -> Result<McpCallResult, McpError> {
|
||||
*self.last_call.lock().await = Some((name.to_string(), arguments));
|
||||
|
||||
if let Some(err) = self.call_error.lock().await.take() {
|
||||
return Err(err);
|
||||
}
|
||||
self.call_response
|
||||
.lock()
|
||||
.await
|
||||
.clone()
|
||||
.ok_or_else(|| McpError::Transport("no canned response".into()))
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), McpError> {
|
||||
self.closed.store(true, Ordering::SeqCst);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_server_info() -> McpServerInfo {
|
||||
McpServerInfo {
|
||||
name: "test-server".into(),
|
||||
version: "1.0.0".into(),
|
||||
capabilities: Value::Null,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_tools() -> Vec<McpToolDefinition> {
|
||||
vec![
|
||||
McpToolDefinition {
|
||||
name: "search".into(),
|
||||
description: Some("Search for items".into()),
|
||||
input_schema: Some(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": { "query": { "type": "string" } }
|
||||
})),
|
||||
},
|
||||
McpToolDefinition {
|
||||
name: "create".into(),
|
||||
description: Some("Create an item".into()),
|
||||
input_schema: None,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
fn make_transport(mock: MockTransport) -> Arc<dyn McpTransport> {
|
||||
Arc::new(mock) as Arc<dyn McpTransport>
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_discovers_and_builds_handlers() {
|
||||
let transport = make_transport(MockTransport::new(sample_server_info(), sample_tools()));
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
assert_eq!(handle.server_info.name, "test-server");
|
||||
assert_eq!(handle.bridge.tool_count(), 2);
|
||||
|
||||
let ids: Vec<String> = handle
|
||||
.bridge
|
||||
.handlers()
|
||||
.iter()
|
||||
.map(|h| h.tool_id.as_str().to_string())
|
||||
.collect();
|
||||
assert!(ids.contains(&"search".to_string()));
|
||||
assert!(ids.contains(&"create".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_handler_descriptions() {
|
||||
let transport = make_transport(MockTransport::new(sample_server_info(), sample_tools()));
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: Some("mcp".into()),
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
let handler = handle
|
||||
.bridge
|
||||
.handlers()
|
||||
.iter()
|
||||
.find(|h| h.tool_id.as_str() == "search")
|
||||
.unwrap();
|
||||
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
let desc = handler.description();
|
||||
assert_eq!(desc.name, "search");
|
||||
assert_eq!(desc.description, "Search for items");
|
||||
assert_eq!(desc.namespace.as_deref(), Some("mcp"));
|
||||
assert!(handler.input_schema().is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_forwards_call_text_response() {
|
||||
let call_result = McpCallResult {
|
||||
content: vec![McpContent::Text {
|
||||
text: "found 3 results".into(),
|
||||
}],
|
||||
is_error: false,
|
||||
};
|
||||
let transport = make_transport(
|
||||
MockTransport::new(sample_server_info(), sample_tools())
|
||||
.with_call_response(call_result),
|
||||
);
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(Arc::clone(&transport), &config)
|
||||
.await
|
||||
.unwrap();
|
||||
let handler = handle
|
||||
.bridge
|
||||
.handlers()
|
||||
.iter()
|
||||
.find(|h| h.tool_id.as_str() == "search")
|
||||
.unwrap();
|
||||
|
||||
use futures::StreamExt;
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
|
||||
let ctx = ToolCallContext::default();
|
||||
let args = serde_json::json!({"query": "test"});
|
||||
let mut stream = handler.handle_call(ctx, args).await;
|
||||
|
||||
let item = stream.next().await.unwrap();
|
||||
match item {
|
||||
kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap();
|
||||
assert_eq!(output, ToolOutputWire::Text("found 3 results".into()));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok(_)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_forwards_call_mcp_blocks_response() {
|
||||
let call_result = McpCallResult {
|
||||
content: vec![
|
||||
McpContent::Text {
|
||||
text: "result text".into(),
|
||||
},
|
||||
McpContent::Image {
|
||||
mime_type: "image/png".into(),
|
||||
data: "base64data".into(),
|
||||
},
|
||||
],
|
||||
is_error: false,
|
||||
};
|
||||
let transport = make_transport(
|
||||
MockTransport::new(sample_server_info(), sample_tools())
|
||||
.with_call_response(call_result),
|
||||
);
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
let handler = &handle.bridge.handlers()[0];
|
||||
|
||||
use futures::StreamExt;
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
|
||||
let ctx = ToolCallContext::default();
|
||||
let mut stream = handler
|
||||
.handle_call(ctx, Value::Object(Default::default()))
|
||||
.await;
|
||||
|
||||
let item = stream.next().await.unwrap();
|
||||
match item {
|
||||
kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap();
|
||||
match output {
|
||||
ToolOutputWire::Mcp { blocks } => {
|
||||
assert_eq!(blocks.len(), 2);
|
||||
assert!(
|
||||
matches!(&blocks[0], McpBlock::Text { text } if text == "result text")
|
||||
);
|
||||
assert!(
|
||||
matches!(&blocks[1], McpBlock::Image { mime_type, .. } if mime_type == "image/png")
|
||||
);
|
||||
}
|
||||
other => panic!("expected Mcp blocks, got {other:?}"),
|
||||
}
|
||||
}
|
||||
other => panic!("expected Terminal(Ok(_)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_handles_mcp_error_response() {
|
||||
let call_result = McpCallResult {
|
||||
content: vec![McpContent::Text {
|
||||
text: "permission denied".into(),
|
||||
}],
|
||||
is_error: true,
|
||||
};
|
||||
let transport = make_transport(
|
||||
MockTransport::new(sample_server_info(), sample_tools())
|
||||
.with_call_response(call_result),
|
||||
);
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
let handler = &handle.bridge.handlers()[0];
|
||||
|
||||
use futures::StreamExt;
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
|
||||
let ctx = ToolCallContext::default();
|
||||
let mut stream = handler.handle_call(ctx, Value::Null).await;
|
||||
|
||||
let item = stream.next().await.unwrap();
|
||||
match item {
|
||||
kigi_tool_runtime::ToolStreamItem::Terminal(Ok(typed)) => {
|
||||
let output: ToolOutputWire = serde_json::from_value(typed.value).unwrap();
|
||||
assert_eq!(output, ToolOutputWire::Text("permission denied".into()));
|
||||
}
|
||||
other => panic!("expected Terminal(Ok(_)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_handles_transport_error() {
|
||||
let transport = make_transport(
|
||||
MockTransport::new(sample_server_info(), sample_tools())
|
||||
.with_call_error(McpError::Transport("connection reset".into())),
|
||||
);
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
let handler = &handle.bridge.handlers()[0];
|
||||
|
||||
use futures::StreamExt;
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
|
||||
let ctx = ToolCallContext::default();
|
||||
let mut stream = handler.handle_call(ctx, Value::Null).await;
|
||||
|
||||
let item = stream.next().await.unwrap();
|
||||
match item {
|
||||
kigi_tool_runtime::ToolStreamItem::Terminal(Err(ref e))
|
||||
if e.kind == kigi_tool_runtime::ToolErrorKind::Execution =>
|
||||
{
|
||||
assert!(
|
||||
e.detail.contains("connection reset"),
|
||||
"expected 'connection reset' in: {}",
|
||||
e.detail
|
||||
);
|
||||
}
|
||||
other => panic!("expected Terminal(Err(Execution)), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_shutdown_closes_transport() {
|
||||
let mock = Arc::new(MockTransport::new(sample_server_info(), sample_tools()));
|
||||
let transport: Arc<dyn McpTransport> = Arc::clone(&mock) as Arc<dyn McpTransport>;
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
assert!(!mock.closed.load(Ordering::SeqCst));
|
||||
|
||||
handle.bridge.shutdown().await.unwrap();
|
||||
assert!(mock.closed.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bridge_skips_tools_with_invalid_names() {
|
||||
let tools = vec![
|
||||
McpToolDefinition {
|
||||
name: "valid_tool".into(),
|
||||
description: Some("a valid tool".into()),
|
||||
input_schema: None,
|
||||
},
|
||||
McpToolDefinition {
|
||||
name: "".into(),
|
||||
description: Some("empty name".into()),
|
||||
input_schema: None,
|
||||
},
|
||||
];
|
||||
let transport = make_transport(MockTransport::new(sample_server_info(), tools));
|
||||
let config = McpBridgeConfig {
|
||||
session_id: SessionId::new("test-session").unwrap(),
|
||||
namespace: None,
|
||||
};
|
||||
|
||||
let handle = McpBridge::connect(transport, &config).await.unwrap();
|
||||
assert_eq!(handle.bridge.tool_count(), 1);
|
||||
assert_eq!(handle.bridge.handlers()[0].tool_id.as_str(), "valid_tool");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_single_text() {
|
||||
let result = McpCallResult {
|
||||
content: vec![McpContent::Text {
|
||||
text: "hello".into(),
|
||||
}],
|
||||
is_error: false,
|
||||
};
|
||||
assert_eq!(
|
||||
translate_mcp_result(&result),
|
||||
ToolOutputWire::Text("hello".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_error_concatenates_text() {
|
||||
let result = McpCallResult {
|
||||
content: vec![
|
||||
McpContent::Text {
|
||||
text: "line 1".into(),
|
||||
},
|
||||
McpContent::Text {
|
||||
text: "line 2".into(),
|
||||
},
|
||||
],
|
||||
is_error: true,
|
||||
};
|
||||
assert_eq!(
|
||||
translate_mcp_result(&result),
|
||||
ToolOutputWire::Text("line 1\nline 2".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_mixed_content_uses_blocks() {
|
||||
let result = McpCallResult {
|
||||
content: vec![
|
||||
McpContent::Text {
|
||||
text: "hello".into(),
|
||||
},
|
||||
McpContent::Resource {
|
||||
uri: "file:///test".into(),
|
||||
mime_type: Some("text/plain".into()),
|
||||
text: Some("content".into()),
|
||||
},
|
||||
],
|
||||
is_error: false,
|
||||
};
|
||||
match translate_mcp_result(&result) {
|
||||
ToolOutputWire::Mcp { blocks } => assert_eq!(blocks.len(), 2),
|
||||
other => panic!("expected Mcp, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_empty_content_returns_empty_text() {
|
||||
let result = McpCallResult {
|
||||
content: vec![],
|
||||
is_error: false,
|
||||
};
|
||||
assert_eq!(
|
||||
translate_mcp_result(&result),
|
||||
ToolOutputWire::Text(String::new())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_empty_error_content_returns_empty_text() {
|
||||
let result = McpCallResult {
|
||||
content: vec![],
|
||||
is_error: true,
|
||||
};
|
||||
assert_eq!(
|
||||
translate_mcp_result(&result),
|
||||
ToolOutputWire::Text(String::new())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn translate_mcp_result_error_with_only_image_drops_content() {
|
||||
let result = McpCallResult {
|
||||
content: vec![McpContent::Image {
|
||||
mime_type: "image/png".into(),
|
||||
data: "base64data".into(),
|
||||
}],
|
||||
is_error: true,
|
||||
};
|
||||
assert_eq!(
|
||||
translate_mcp_result(&result),
|
||||
ToolOutputWire::Text(String::new())
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
//! Unified MCP adapter for the xAI Computer Hub.
|
||||
//!
|
||||
//! This crate bridges MCP (Model Context Protocol) servers into the
|
||||
//! computer hub's tool routing infrastructure. An [`McpBridge`] connects
|
||||
//! to an MCP server via an [`McpTransport`], discovers the server's
|
||||
//! tools, and produces [`ToolServerHandler`](kigi_computer_hub_sdk::ToolServerHandler)
|
||||
//! implementations that can be registered with a hub
|
||||
//! [`ToolServerBuilder`](kigi_computer_hub_sdk::ToolServerBuilder).
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! MCP Server <──McpTransport──> McpBridge ──handlers──> ToolServerBuilder
|
||||
//! (stdio/SSE) (discover+forward) (register with hub)
|
||||
//! ```
|
||||
//!
|
||||
//! The [`McpTransport`] trait abstracts the wire protocol so the bridge
|
||||
//! is testable with in-memory mocks. Concrete transport implementations
|
||||
//! (stdio, HTTP+SSE) are provided by downstream crates.
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```rust,ignore
|
||||
//! let transport: Arc<dyn McpTransport> = /* ... */;
|
||||
//! let config = McpBridgeConfig {
|
||||
//! session_id: SessionId::new("session-1").unwrap(),
|
||||
//! namespace: Some("my-mcp-server".into()),
|
||||
//! };
|
||||
//! let handle = McpBridge::connect(transport, &config).await?;
|
||||
//!
|
||||
//! let mut builder = ToolServerBuilder::default()
|
||||
//! .pool(pool)
|
||||
//! .url(hub_url)
|
||||
//! .auth(auth);
|
||||
//!
|
||||
//! for handler in handle.bridge.handlers() {
|
||||
//! builder = builder.tool(handler.clone());
|
||||
//! }
|
||||
//!
|
||||
//! let server = builder.build().await?;
|
||||
//! ```
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
pub mod bridge;
|
||||
pub(crate) mod metrics;
|
||||
pub mod transport;
|
||||
pub mod types;
|
||||
|
||||
pub use bridge::{McpBridge, McpBridgeConfig, McpBridgeHandle, McpToolHandler};
|
||||
pub use transport::McpTransport;
|
||||
pub use types::{McpCallResult, McpContent, McpError, McpServerInfo, McpToolDefinition};
|
||||
@@ -0,0 +1,61 @@
|
||||
//! Feature-gated Prometheus metrics for the MCP adapter bridge.
|
||||
//!
|
||||
//! When the `metrics` cargo feature is enabled, each helper records to a
|
||||
//! lazily-registered Prometheus counter / gauge / histogram. When
|
||||
//! disabled (the default), every helper compiles to an empty function
|
||||
//! body so the crate carries zero prometheus dependency.
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
mod inner {
|
||||
use prometheus::{
|
||||
Histogram, IntCounter, IntGauge, exponential_buckets, register_histogram,
|
||||
register_int_counter, register_int_gauge,
|
||||
};
|
||||
use std::sync::LazyLock;
|
||||
|
||||
static MCP_CALL_DURATION_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
|
||||
register_histogram!(
|
||||
"computer_hub_mcp_call_duration_seconds",
|
||||
"MCP server response latency for tool calls.",
|
||||
exponential_buckets(0.01, 2.0, 14).expect("valid bucket params")
|
||||
)
|
||||
.expect("computer_hub_mcp_call_duration_seconds must register once")
|
||||
});
|
||||
|
||||
static MCP_ERRORS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_mcp_errors_total",
|
||||
"Errors in the MCP adapter pipeline (transport, protocol, or serialization)."
|
||||
)
|
||||
.expect("computer_hub_mcp_errors_total must register once")
|
||||
});
|
||||
|
||||
static MCP_TOOLS_BRIDGED: LazyLock<IntGauge> = LazyLock::new(|| {
|
||||
register_int_gauge!(
|
||||
"computer_hub_mcp_tools_bridged",
|
||||
"MCP tools currently bridged into the computer hub."
|
||||
)
|
||||
.expect("computer_hub_mcp_tools_bridged must register once")
|
||||
});
|
||||
|
||||
pub(crate) fn mcp_call_duration_observe(secs: f64) {
|
||||
MCP_CALL_DURATION_SECONDS.observe(secs);
|
||||
}
|
||||
|
||||
pub(crate) fn mcp_error() {
|
||||
MCP_ERRORS_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn mcp_tools_bridged_set(count: i64) {
|
||||
MCP_TOOLS_BRIDGED.set(count);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "metrics"))]
|
||||
mod inner {
|
||||
pub(crate) fn mcp_call_duration_observe(_secs: f64) {}
|
||||
pub(crate) fn mcp_error() {}
|
||||
pub(crate) fn mcp_tools_bridged_set(_count: i64) {}
|
||||
}
|
||||
|
||||
pub(crate) use inner::*;
|
||||
@@ -0,0 +1,38 @@
|
||||
//! MCP transport abstraction.
|
||||
//!
|
||||
//! [`McpTransport`] defines the async interface consumed by [`crate::McpBridge`].
|
||||
//! Concrete implementations (stdio, HTTP+SSE) live outside this crate;
|
||||
//! the trait boundary keeps the bridge testable with in-memory mocks.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::types::{McpCallResult, McpError, McpServerInfo, McpToolDefinition};
|
||||
|
||||
/// Async interface to a single MCP server connection.
|
||||
///
|
||||
/// Implementations manage the underlying JSON-RPC framing (stdio pipe,
|
||||
/// HTTP+SSE stream, etc.) and expose the four lifecycle operations the
|
||||
/// bridge needs.
|
||||
#[async_trait]
|
||||
pub trait McpTransport: Send + Sync {
|
||||
/// Perform the MCP `initialize` handshake with the server.
|
||||
///
|
||||
/// Must be called exactly once before any other method. Returns
|
||||
/// the server's advertised name, version, and capabilities.
|
||||
async fn initialize(&self) -> Result<McpServerInfo, McpError>;
|
||||
|
||||
/// Discover available tools via MCP `tools/list`.
|
||||
async fn list_tools(&self) -> Result<Vec<McpToolDefinition>, McpError>;
|
||||
|
||||
/// Invoke a tool via MCP `tools/call`.
|
||||
///
|
||||
/// `arguments` is the JSON object the model produced for the tool's
|
||||
/// input schema.
|
||||
async fn call_tool(&self, name: &str, arguments: Value) -> Result<McpCallResult, McpError>;
|
||||
|
||||
/// Gracefully shut down the transport (close pipes, drop connections).
|
||||
/// Implementations must be idempotent — a second call after a
|
||||
/// successful close must return `Ok(())` without error.
|
||||
async fn close(&self) -> Result<(), McpError>;
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
//! MCP protocol types used by the adapter.
|
||||
//!
|
||||
//! These mirror the MCP specification's JSON-RPC shapes for server
|
||||
//! metadata, tool definitions, and call results. They are intentionally
|
||||
//! decoupled from any specific transport implementation so the bridge
|
||||
//! stays testable with in-memory mocks.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Metadata returned by a successful MCP `initialize` handshake.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpServerInfo {
|
||||
/// Human-readable server name (e.g. `"linear"`, `"github"`).
|
||||
pub name: String,
|
||||
/// Semver-ish version reported by the server.
|
||||
pub version: String,
|
||||
/// Free-form capability flags advertised during init.
|
||||
#[serde(default)]
|
||||
pub capabilities: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A single tool definition from MCP `tools/list`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct McpToolDefinition {
|
||||
/// Unqualified tool name (e.g. `"create_issue"`).
|
||||
pub name: String,
|
||||
/// Model-facing description of the tool.
|
||||
#[serde(default)]
|
||||
pub description: Option<String>,
|
||||
/// JSON Schema for the tool's input arguments.
|
||||
#[serde(default)]
|
||||
pub input_schema: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Result of an MCP `tools/call` invocation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct McpCallResult {
|
||||
/// Content blocks returned by the tool.
|
||||
#[serde(default)]
|
||||
pub content: Vec<McpContent>,
|
||||
/// When `true`, the tool signalled an application-level error.
|
||||
#[serde(default)]
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
/// A single content block inside an [`McpCallResult`].
|
||||
///
|
||||
/// Covers the three content types defined by the MCP specification:
|
||||
/// text, image (base64-encoded), and embedded resource.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "camelCase")]
|
||||
pub enum McpContent {
|
||||
/// Plain text content.
|
||||
#[serde(rename = "text")]
|
||||
Text {
|
||||
/// The text payload.
|
||||
text: String,
|
||||
},
|
||||
/// Base64-encoded image content.
|
||||
#[serde(rename = "image")]
|
||||
Image {
|
||||
/// MIME type (e.g. `"image/png"`).
|
||||
#[serde(rename = "mimeType")]
|
||||
mime_type: String,
|
||||
/// Base64-encoded image bytes.
|
||||
data: String,
|
||||
},
|
||||
/// Embedded resource content.
|
||||
#[serde(rename = "resource")]
|
||||
Resource {
|
||||
/// Resource URI.
|
||||
uri: String,
|
||||
/// Optional MIME type.
|
||||
#[serde(default, rename = "mimeType")]
|
||||
mime_type: Option<String>,
|
||||
/// Optional text body.
|
||||
#[serde(default)]
|
||||
text: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Errors originating from MCP transport or protocol handling.
|
||||
#[derive(Debug, Clone, thiserror::Error)]
|
||||
pub enum McpError {
|
||||
/// The underlying transport failed (connection refused, pipe broken, etc.).
|
||||
#[error("transport error: {0}")]
|
||||
Transport(String),
|
||||
|
||||
/// The server returned a JSON-RPC error response.
|
||||
#[error("protocol error (code {code}): {message}")]
|
||||
Protocol {
|
||||
/// JSON-RPC error code.
|
||||
code: i64,
|
||||
/// Human-readable error message.
|
||||
message: String,
|
||||
},
|
||||
|
||||
/// Timeout waiting for MCP server response.
|
||||
#[error("timeout: {0}")]
|
||||
Timeout(String),
|
||||
|
||||
/// The response could not be decoded.
|
||||
#[error("decode error: {0}")]
|
||||
Decode(String),
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-computer-hub-sdk"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "SDK for the xAI Computer Hub: connection pool, transparent reconnect, tool harness, and tool-server runtime."
|
||||
|
||||
[features]
|
||||
metrics = ["dep:prometheus"]
|
||||
|
||||
[dependencies]
|
||||
tokio = { workspace = true, features = ["rt", "sync", "time", "macros"] }
|
||||
tokio-tungstenite = { workspace = true, features = ["rustls-tls-native-roots"] }
|
||||
tokio-util = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
dashmap = { workspace = true }
|
||||
indexmap = { workspace = true }
|
||||
arc-swap = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
url = { workspace = true }
|
||||
http = { workspace = true }
|
||||
prometheus = { workspace = true, optional = true }
|
||||
reqwest = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
parking_lot = { workspace = true }
|
||||
|
||||
fastrace = { workspace = true }
|
||||
# Trace donation: spans convert via the stock fastrace -> OTel reporter
|
||||
# and ship as standard OTLP payloads.
|
||||
fastrace-opentelemetry = { workspace = true }
|
||||
opentelemetry = { workspace = true }
|
||||
opentelemetry_sdk = { workspace = true }
|
||||
opentelemetry-proto = { workspace = true }
|
||||
prost = { workspace = true }
|
||||
base64 = { workspace = true }
|
||||
|
||||
kigi-tool-protocol = { workspace = true }
|
||||
kigi-tool-runtime = { workspace = true }
|
||||
kigi-tool-types = { workspace = true }
|
||||
kigi-computer-hub-core = { workspace = true }
|
||||
kigi-tracing = { workspace = true }
|
||||
|
||||
# Integration tests that need heavier backend deps live in a separate sibling crate to keep this dev-dep set minimal.
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["full", "test-util"] }
|
||||
axum = { workspace = true, features = ["ws", "macros"] }
|
||||
chrono = { workspace = true }
|
||||
base64 = { workspace = true }
|
||||
schemars = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,333 @@
|
||||
//! Three-tier semaphore admission + bounded-wait backpressure.
|
||||
//!
|
||||
//! Concurrent *running* calls are bounded at three scopes, acquired in a
|
||||
//! fixed **session → connection → global** order. A consistent
|
||||
//! most-local-first order is deadlock-free and never holds a scarce
|
||||
//! global permit while blocking on a local one. A single shared deadline
|
||||
//! spans all three acquisitions, so total admission latency is bounded by
|
||||
//! `wait_timeout`, not `3 × wait_timeout`.
|
||||
//!
|
||||
//! Under moderate pressure `admit` waits; under very high pressure the
|
||||
//! deadline elapses and the caller emits the shared overloaded JSON-RPC
|
||||
//! error (`-32016` "tool_busy") instead of silently dropping the request.
|
||||
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use kigi_tool_protocol::{
|
||||
JsonRpcError, JsonRpcId, JsonRpcResponse, JsonRpcVersion, ResponseOutcome, SessionId,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
use tokio::time::Instant;
|
||||
|
||||
/// Numeric JSON-RPC code for overload rejection (`kigi-tool-protocol`
|
||||
/// `error_codes.rs`: `-32016` "tool_busy").
|
||||
pub(crate) const TOOL_BUSY_CODE: i32 = -32016;
|
||||
|
||||
const TOOL_BUSY_MESSAGE: &str = "tool server busy; tool call rejected";
|
||||
|
||||
/// Default ceiling for the process-wide concurrency guard.
|
||||
pub(crate) const DEFAULT_GLOBAL_MAX_INFLIGHT: usize = 1024;
|
||||
/// Default per-session concurrent running calls.
|
||||
pub(crate) const DEFAULT_SESSION_MAX_INFLIGHT: usize = 16;
|
||||
/// Default per-connection concurrent running calls.
|
||||
pub(crate) const DEFAULT_CONN_MAX_INFLIGHT: usize = 256;
|
||||
/// Default bounded wait before an overloaded rejection.
|
||||
pub(crate) const DEFAULT_ADMISSION_WAIT_TIMEOUT: Duration = Duration::from_secs(3);
|
||||
|
||||
/// Ops-tunable override for the process-wide global cap (Helm `env:`).
|
||||
const GLOBAL_MAX_INFLIGHT_ENV: &str = "XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT";
|
||||
|
||||
/// Inflight-gauge scope labels, in acquisition order. A held [`AdmitGuard`]
|
||||
/// counts against all three.
|
||||
const SCOPES: [&str; 3] = ["session", "conn", "global"];
|
||||
|
||||
/// Build the shared overloaded (`-32016` "tool_busy") JSON-RPC error
|
||||
/// response. This is the single source of the overload wire shape, reused
|
||||
/// by BOTH the admission-timeout path (`server::execute_call`) and the
|
||||
/// demux inbox-full path (`demux::route_session`) so the two never drift.
|
||||
pub(crate) fn overloaded_response(id: JsonRpcId, session_id: SessionId) -> JsonRpcResponse<Value> {
|
||||
JsonRpcResponse {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id,
|
||||
session_id: Some(session_id),
|
||||
outcome: ResponseOutcome::Error(JsonRpcError {
|
||||
code: TOOL_BUSY_CODE,
|
||||
message: TOOL_BUSY_MESSAGE.to_owned(),
|
||||
data: Some(serde_json::json!({ "code": "tool_busy", "retryable": true })),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-wide global admission semaphore, shared by every connection.
|
||||
///
|
||||
/// Initialized once at first use: the value comes from
|
||||
/// `XAI_TOOL_SERVER_GLOBAL_MAX_INFLIGHT` when present and parseable as a
|
||||
/// positive integer, otherwise `default_cap` (the builder knob, default
|
||||
/// [`DEFAULT_GLOBAL_MAX_INFLIGHT`]). Because the cell initializes exactly
|
||||
/// once, the first caller's `default_cap` and the env var at that instant
|
||||
/// fix the process-wide capacity.
|
||||
pub(crate) fn global_semaphore(default_cap: usize) -> Arc<Semaphore> {
|
||||
static SEM: OnceLock<Arc<Semaphore>> = OnceLock::new();
|
||||
SEM.get_or_init(|| {
|
||||
let raw = std::env::var(GLOBAL_MAX_INFLIGHT_ENV).ok();
|
||||
Arc::new(Semaphore::new(resolve_global_cap(
|
||||
raw.as_deref(),
|
||||
default_cap,
|
||||
)))
|
||||
})
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Resolve the process-wide global cap from the raw env value, falling
|
||||
/// back to `default_cap`. Pure (no global state) so the
|
||||
/// fall-back-never-panic guarantee is unit-tested: a non-numeric,
|
||||
/// negative, empty, or zero value all yield `default_cap`.
|
||||
fn resolve_global_cap(raw: Option<&str>, default_cap: usize) -> usize {
|
||||
raw.and_then(|v| v.parse::<usize>().ok())
|
||||
.filter(|n| *n > 0)
|
||||
.unwrap_or(default_cap)
|
||||
}
|
||||
|
||||
/// Why admission was refused.
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) enum Overloaded {
|
||||
/// The bounded admission deadline elapsed under very high pressure.
|
||||
Timeout,
|
||||
/// A semaphore was closed — the server is shutting down.
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
/// RAII guard holding all three permits for the call's lifetime.
|
||||
///
|
||||
/// Fields drop in declaration order, so permits are released in reverse
|
||||
/// of acquisition: global → connection → session.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct AdmitGuard {
|
||||
_global: OwnedSemaphorePermit,
|
||||
_conn: OwnedSemaphorePermit,
|
||||
_session: OwnedSemaphorePermit,
|
||||
}
|
||||
|
||||
impl Drop for AdmitGuard {
|
||||
fn drop(&mut self) {
|
||||
for scope in SCOPES {
|
||||
crate::metrics::tool_call_inflight_dec(scope);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Three-tier admission controller. One per connection (`conn_sem`); the
|
||||
/// per-session map is created/destroyed alongside each session loop.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Admission {
|
||||
session_sems: DashMap<SessionId, Arc<Semaphore>>,
|
||||
session_max: usize,
|
||||
conn_sem: Arc<Semaphore>,
|
||||
global_sem: Arc<Semaphore>,
|
||||
wait_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Admission {
|
||||
pub(crate) fn new(
|
||||
session_max: usize,
|
||||
conn_max: usize,
|
||||
global_sem: Arc<Semaphore>,
|
||||
wait_timeout: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
session_sems: DashMap::new(),
|
||||
session_max,
|
||||
conn_sem: Arc::new(Semaphore::new(conn_max)),
|
||||
global_sem,
|
||||
wait_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create the per-session semaphore entry. Called from
|
||||
/// `bind_session_local` so the entry's lifetime is tied to the
|
||||
/// session-loop task, not lazily minted in [`Self::admit`].
|
||||
pub(crate) fn ensure_session(&self, session_id: &SessionId) {
|
||||
self.session_sems
|
||||
.entry(session_id.clone())
|
||||
.or_insert_with(|| Arc::new(Semaphore::new(self.session_max)));
|
||||
}
|
||||
|
||||
/// Remove the per-session semaphore entry on unbind / loop exit.
|
||||
pub(crate) fn remove_session(&self, session_id: &SessionId) {
|
||||
self.session_sems.remove(session_id);
|
||||
}
|
||||
|
||||
/// Acquire one permit at each scope (session → connection → global)
|
||||
/// against a single shared deadline.
|
||||
pub(crate) async fn admit(&self, session_id: &SessionId) -> Result<AdmitGuard, Overloaded> {
|
||||
let start = Instant::now();
|
||||
let deadline = start + self.wait_timeout;
|
||||
// The entry is created in `bind_session_local`; a straggler call
|
||||
// admitted just after unbind cleanup falls back to a private,
|
||||
// un-tracked semaphore rather than recreating a leaked entry.
|
||||
let session_sem = self
|
||||
.session_sems
|
||||
.get(session_id)
|
||||
.map(|s| s.clone())
|
||||
.unwrap_or_else(|| Arc::new(Semaphore::new(self.session_max)));
|
||||
|
||||
let session = acquire_until(&session_sem, deadline).await?;
|
||||
let conn = acquire_until(&self.conn_sem, deadline).await?;
|
||||
let global = acquire_until(&self.global_sem, deadline).await?;
|
||||
|
||||
crate::metrics::admission_wait_observe(start.elapsed().as_secs_f64());
|
||||
for scope in SCOPES {
|
||||
crate::metrics::tool_call_inflight_inc(scope);
|
||||
}
|
||||
Ok(AdmitGuard {
|
||||
_global: global,
|
||||
_conn: conn,
|
||||
_session: session,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Acquire one owned permit before `deadline`, mapping closed/elapsed to
|
||||
/// the matching [`Overloaded`] variant.
|
||||
async fn acquire_until(
|
||||
sem: &Arc<Semaphore>,
|
||||
deadline: Instant,
|
||||
) -> Result<OwnedSemaphorePermit, Overloaded> {
|
||||
match tokio::time::timeout_at(deadline, sem.clone().acquire_owned()).await {
|
||||
Ok(Ok(permit)) => Ok(permit),
|
||||
Ok(Err(_closed)) => Err(Overloaded::Shutdown),
|
||||
Err(_elapsed) => Err(Overloaded::Timeout),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sid(s: &str) -> SessionId {
|
||||
SessionId::new(s).expect("valid session id")
|
||||
}
|
||||
|
||||
fn test_admission(session_max: usize, conn_max: usize, global_max: usize) -> Admission {
|
||||
Admission::new(
|
||||
session_max,
|
||||
conn_max,
|
||||
Arc::new(Semaphore::new(global_max)),
|
||||
Duration::from_millis(150),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_global_cap_falls_back_on_bad_input_and_honors_valid() {
|
||||
// Absent / non-numeric / negative / empty / zero → default (never panic).
|
||||
assert_eq!(resolve_global_cap(None, 1024), 1024);
|
||||
assert_eq!(resolve_global_cap(Some("abc"), 1024), 1024);
|
||||
assert_eq!(resolve_global_cap(Some("-5"), 1024), 1024);
|
||||
assert_eq!(resolve_global_cap(Some(""), 1024), 1024);
|
||||
assert_eq!(resolve_global_cap(Some("0"), 1024), 1024);
|
||||
assert_eq!(resolve_global_cap(Some(" 7"), 1024), 1024); // leading space → parse fails
|
||||
// A valid positive integer overrides the default.
|
||||
assert_eq!(resolve_global_cap(Some("2048"), 1024), 2048);
|
||||
assert_eq!(resolve_global_cap(Some("1"), 1024), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overloaded_response_carries_minus_32016_and_data_marker() {
|
||||
let id: JsonRpcId = serde_json::from_value(serde_json::json!("call-1")).expect("id");
|
||||
let resp = overloaded_response(id, sid("s1"));
|
||||
let wire: Value = serde_json::from_str(&serde_json::to_string(&resp).expect("ser"))
|
||||
.expect("round-trips to json");
|
||||
assert_eq!(wire["error"]["code"], TOOL_BUSY_CODE);
|
||||
assert_eq!(wire["error"]["code"], -32016);
|
||||
assert_eq!(wire["error"]["data"]["code"], "tool_busy");
|
||||
assert_eq!(wire["error"]["data"]["retryable"], true);
|
||||
assert_eq!(wire["session_id"], "s1");
|
||||
assert_eq!(wire["id"], "call-1");
|
||||
assert!(
|
||||
wire.get("result").is_none(),
|
||||
"overload is an error, never a result"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn admit_times_out_when_session_saturated() {
|
||||
let admission = test_admission(2, 16, 64);
|
||||
let session = sid("sat");
|
||||
admission.ensure_session(&session);
|
||||
|
||||
// Hold both session permits.
|
||||
let g1 = admission.admit(&session).await.expect("first admit");
|
||||
let _g2 = admission.admit(&session).await.expect("second admit");
|
||||
|
||||
// Third admit must elapse the deadline → Timeout (not a hang).
|
||||
let result = admission.admit(&session).await;
|
||||
assert_eq!(result.unwrap_err(), Overloaded::Timeout);
|
||||
|
||||
// Releasing one permit frees a slot for the next admit.
|
||||
drop(g1);
|
||||
admission
|
||||
.admit(&session)
|
||||
.await
|
||||
.expect("permit released → admit succeeds");
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn admit_blocks_on_connection_scope_when_conn_saturated() {
|
||||
// conn_max = 1 is the binding constraint even though session has
|
||||
// room; a second admit on a *different* session still times out.
|
||||
let admission = test_admission(8, 1, 64);
|
||||
let a = sid("a");
|
||||
let b = sid("b");
|
||||
admission.ensure_session(&a);
|
||||
admission.ensure_session(&b);
|
||||
|
||||
let _held = admission.admit(&a).await.expect("first admit");
|
||||
let result = admission.admit(&b).await;
|
||||
assert_eq!(
|
||||
result.unwrap_err(),
|
||||
Overloaded::Timeout,
|
||||
"connection cap binds across sessions"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admit_succeeds_repeatedly_under_capacity() {
|
||||
let admission = test_admission(4, 16, 64);
|
||||
let session = sid("ok");
|
||||
admission.ensure_session(&session);
|
||||
let mut guards = Vec::new();
|
||||
for _ in 0..4 {
|
||||
guards.push(admission.admit(&session).await.expect("within capacity"));
|
||||
}
|
||||
assert_eq!(guards.len(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closed_semaphore_maps_to_shutdown() {
|
||||
let global = Arc::new(Semaphore::new(0));
|
||||
global.close();
|
||||
let admission = Admission::new(4, 16, global, Duration::from_secs(5));
|
||||
let session = sid("closed");
|
||||
admission.ensure_session(&session);
|
||||
let result = admission.admit(&session).await;
|
||||
assert_eq!(result.unwrap_err(), Overloaded::Shutdown);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn straggler_admit_after_remove_uses_private_permit() {
|
||||
let admission = test_admission(1, 16, 64);
|
||||
let session = sid("gone");
|
||||
// No ensure_session: simulate a straggler after unbind removed it.
|
||||
admission.remove_session(&session);
|
||||
// Falls back to a private semaphore and still admits (no panic,
|
||||
// no leaked tracked entry).
|
||||
let _g = admission.admit(&session).await.expect("private fallback");
|
||||
assert!(
|
||||
admission.session_sems.get(&session).is_none(),
|
||||
"straggler must not recreate a tracked entry"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
//! Auth credentials and pool-dedup principal keys.
|
||||
//!
|
||||
//! [`AuthCredential`] models the credential the client attaches at
|
||||
//! handshake time. Two variants are supported:
|
||||
//!
|
||||
//! - [`AuthCredential::Bearer`] for the simple "Authorization: Bearer
|
||||
//! …" path (e.g. JWT-against-OAuth2 deployments).
|
||||
//! - [`AuthCredential::Headers`] for callers that already hold a
|
||||
//! pre-built header bundle (e.g. signed identity headers generated
|
||||
//! by an upstream proxy or test harness).
|
||||
//!
|
||||
//! [`PrincipalKey`] is the stable hashable projection of an
|
||||
//! `AuthCredential`; the pool keys connections by
|
||||
//! `(url, principal_key)` so two [`crate::ToolServer`] builds with the
|
||||
//! same credential reuse one socket while distinct credentials open
|
||||
//! distinct sockets. The server derives `user_id` from the credential at
|
||||
//! upgrade time and returns it in the hello ack — the SDK never needs
|
||||
//! to carry `user_id` alongside the credential.
|
||||
//!
|
||||
//! ## Pool dedup and credential refresh
|
||||
//!
|
||||
//! Both variants include the secret material in the `PrincipalKey`
|
||||
//! fingerprint. This is deliberate: distinct secrets imply distinct
|
||||
//! credentials, so two callers with different tokens open distinct
|
||||
//! sockets. The trade-off is that a caller that rotates its bearer JWT
|
||||
//! every N minutes will open a new socket on each rotation.
|
||||
//! Long-running tool servers should reuse the SAME [`AuthCredential`]
|
||||
//! instance across builds and refresh the credential out-of-band rather
|
||||
//! than hand a fresh JWT to every build.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
|
||||
use http::HeaderName;
|
||||
use http::header::AUTHORIZATION;
|
||||
|
||||
use crate::error::ClientError;
|
||||
|
||||
/// Credential carried into the WebSocket upgrade.
|
||||
///
|
||||
/// Clones are cheap (the secret material is at most a small number of
|
||||
/// owned strings). The server derives `user_id` from the credential at
|
||||
/// upgrade time and returns it in the [`kigi_tool_protocol::HelloAckMsg`].
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub enum AuthCredential {
|
||||
/// Bearer token attached as the `Authorization: Bearer …` header.
|
||||
Bearer { token: String },
|
||||
/// Pre-built header bundle. Used when the auth flow lives outside
|
||||
/// the SDK (e.g. an upstream proxy that already produced signed
|
||||
/// identity headers). Header order is canonicalised for stable
|
||||
/// hashing via [`BTreeMap`]; names are lowercased and validated
|
||||
/// as `HeaderName` at construction time so an invalid name
|
||||
/// surfaces as [`ClientError::InvalidConfig`] instead of being
|
||||
/// silently dropped at upgrade time.
|
||||
Headers { headers: BTreeMap<String, String> },
|
||||
}
|
||||
|
||||
impl AuthCredential {
|
||||
/// Convenience constructor for the bearer-token shape.
|
||||
pub fn bearer(token: impl Into<String>) -> Self {
|
||||
Self::Bearer {
|
||||
token: token.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convenience constructor for the raw-header bundle shape.
|
||||
///
|
||||
/// Names are canonicalised to lowercase and validated as
|
||||
/// [`HeaderName`] at construction so an invalid header (e.g. one
|
||||
/// containing a newline injection attempt) returns
|
||||
/// [`ClientError::InvalidConfig`] rather than being silently
|
||||
/// filtered out at upgrade time.
|
||||
pub fn headers<I, K, V>(headers: I) -> Result<Self, ClientError>
|
||||
where
|
||||
I: IntoIterator<Item = (K, V)>,
|
||||
K: AsRef<str>,
|
||||
V: Into<String>,
|
||||
{
|
||||
let mut map: BTreeMap<String, String> = BTreeMap::new();
|
||||
for (raw_name, raw_value) in headers {
|
||||
let name = raw_name.as_ref().to_ascii_lowercase();
|
||||
HeaderName::from_bytes(name.as_bytes()).map_err(|err| {
|
||||
ClientError::InvalidConfig(format!("invalid header name {name:?}: {err}"))
|
||||
})?;
|
||||
map.insert(name, raw_value.into());
|
||||
}
|
||||
Ok(Self::Headers { headers: map })
|
||||
}
|
||||
|
||||
/// Stable hashable projection used as the pool dedup key.
|
||||
///
|
||||
/// Distinct credentials hash equal iff they carry the same secret
|
||||
/// material. See the module-level "Pool dedup and credential
|
||||
/// refresh" section for the implications when bearer tokens are
|
||||
/// rotated.
|
||||
pub fn principal_key(&self) -> PrincipalKey {
|
||||
match self {
|
||||
Self::Bearer { token } => PrincipalKey {
|
||||
fingerprint: format!("bearer:{token}"),
|
||||
},
|
||||
Self::Headers { headers } => {
|
||||
// Concatenate canonicalised name=value pairs so the
|
||||
// fingerprint is order-independent.
|
||||
let mut joined = String::with_capacity(headers.len() * 32);
|
||||
for (name, value) in headers {
|
||||
joined.push_str(name);
|
||||
joined.push('=');
|
||||
joined.push_str(value);
|
||||
joined.push('\n');
|
||||
}
|
||||
PrincipalKey {
|
||||
fingerprint: format!("headers:{joined}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Headers to attach to the WebSocket upgrade request.
|
||||
///
|
||||
/// `Headers` variant entries are infallible at this point — names
|
||||
/// were validated by [`Self::headers`].
|
||||
pub fn upgrade_headers(&self) -> Vec<(HeaderName, String)> {
|
||||
match self {
|
||||
Self::Bearer { token, .. } => {
|
||||
vec![(AUTHORIZATION, format!("Bearer {token}"))]
|
||||
}
|
||||
Self::Headers { headers, .. } => headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
HeaderName::from_bytes(name.as_bytes())
|
||||
.ok()
|
||||
.map(|n| (n, value.clone()))
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for AuthCredential {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
// Never log the secret; surface only the variant.
|
||||
match self {
|
||||
Self::Bearer { .. } => f
|
||||
.debug_struct("AuthCredential::Bearer")
|
||||
.finish_non_exhaustive(),
|
||||
Self::Headers { headers } => f
|
||||
.debug_struct("AuthCredential::Headers")
|
||||
.field("header_count", &headers.len())
|
||||
.finish_non_exhaustive(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable hashable projection of an [`AuthCredential`] used as the
|
||||
/// pool dedup key alongside the connect URL. Two connections with the
|
||||
/// same token fingerprint will get the same server-assigned `user_id`.
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub struct PrincipalKey {
|
||||
fingerprint: String,
|
||||
}
|
||||
|
||||
impl fmt::Debug for PrincipalKey {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("PrincipalKey").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
/// Owner identity surfaced by an [`AuthProvider`] alongside its credential.
|
||||
///
|
||||
/// Mirrors the OAuth principal fields the provider parsed from its auth source.
|
||||
/// It is kept separate from [`AuthCredential`] on purpose: identity must NOT
|
||||
/// participate in pool-dedup hashing (that keys only on the secret), and the
|
||||
/// credential's `Eq`/`Hash` derives must stay token-only. Consumers (e.g. the
|
||||
/// workspace) map this onto their own identity record.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct AuthIdentity {
|
||||
/// Stable user identifier (owner of the bearer token).
|
||||
pub user_id: String,
|
||||
/// OAuth `principal_type` wire string (`"User"` / `"Team"`), when known.
|
||||
pub principal_type: Option<String>,
|
||||
/// Team id when `principal_type == "Team"`; otherwise `None`.
|
||||
pub principal_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Credential provider called on every connect/reconnect.
|
||||
pub trait AuthProvider: Send + Sync + std::fmt::Debug {
|
||||
fn current(&self) -> AuthCredential;
|
||||
|
||||
/// Stable pool-dedup key, decoupled from the per-connect credential.
|
||||
///
|
||||
/// Defaults to the current credential's key (existing behavior). A provider
|
||||
/// that re-mints a rotating secret on every [`Self::current`] call (e.g. a
|
||||
/// refresh-before-use bearer) MUST override this to key only on stable
|
||||
/// identity, otherwise each rotation fragments the connection pool.
|
||||
fn principal_key(&self) -> PrincipalKey {
|
||||
self.current().principal_key()
|
||||
}
|
||||
|
||||
/// Owner identity behind the credential, when the provider can surface it.
|
||||
///
|
||||
/// Defaults to `None` for providers that only carry a bearer token (e.g. a
|
||||
/// bare [`AuthCredential`]). Providers that parse OAuth principal fields
|
||||
/// (e.g. OIDC) override this so downstream consumers can attribute
|
||||
/// requests without a second auth-source read.
|
||||
fn identity(&self) -> Option<AuthIdentity> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub type SharedAuthProvider = std::sync::Arc<dyn AuthProvider>;
|
||||
|
||||
impl AuthProvider for AuthCredential {
|
||||
fn current(&self) -> AuthCredential {
|
||||
self.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn invalid_header_name_rejected_at_construction() {
|
||||
let cred = AuthCredential::headers([("authorization\nx-injected", "value")]);
|
||||
match cred {
|
||||
Err(ClientError::InvalidConfig(msg)) => {
|
||||
assert!(msg.contains("invalid header name"), "got {msg}")
|
||||
}
|
||||
other => panic!("expected InvalidConfig; got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_headers_accepted() {
|
||||
let cred = AuthCredential::headers([("authorization", "Bearer token")]).expect("valid");
|
||||
assert_eq!(cred.upgrade_headers().len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
//! Per-session strict-cancellation registry.
|
||||
//!
|
||||
//! Maps each in-flight `tool_call_id` to its [`CancellationToken`] so a
|
||||
//! `Cancel` hook (or session teardown) can hard-cancel the running call
|
||||
//! by dropping its future. A small `pending` tombstone set covers the
|
||||
//! race where a `Cancel` arrives *before* the dispatcher registered the
|
||||
//! token (the symmetric window to pre-spawn registration): the id is
|
||||
//! tombstoned and the dispatcher cancels it at registration time.
|
||||
//!
|
||||
//! One registry per session, tied to the session-loop lifetime alongside
|
||||
//! the inbox and the per-session admission semaphore.
|
||||
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use dashmap::{DashMap, DashSet};
|
||||
use kigi_tool_protocol::ToolCallId;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// Upper bound on outstanding pre-registration tombstones. Tombstones
|
||||
/// cover the microscopic window between a `Cancel` hook and the matching
|
||||
/// `register`, so in steady state the set holds a handful of entries. A
|
||||
/// `Cancel` whose call never registers (e.g. one racing call completion,
|
||||
/// after `deregister` already removed the live token) leaves a tombstone
|
||||
/// that no `register` ever consumes; this cap reclaims such stragglers so
|
||||
/// a single long-lived session cannot grow `pending` without bound.
|
||||
const MAX_PENDING_TOMBSTONES: usize = 8192;
|
||||
|
||||
/// Per-session `tool_call_id -> CancellationToken` map plus a pending
|
||||
/// tombstone set for cancels that land before registration.
|
||||
#[derive(Default, Debug)]
|
||||
pub(crate) struct CancelRegistry {
|
||||
map: DashMap<ToolCallId, CancellationToken>,
|
||||
pending: DashSet<ToolCallId>,
|
||||
/// Set once by [`Self::cancel_all`] (teardown). After this, every new
|
||||
/// `register` starts cancelled so a request dispatched in the teardown
|
||||
/// window cannot escape as an orphaned, uncancellable task.
|
||||
closed: AtomicBool,
|
||||
}
|
||||
|
||||
impl CancelRegistry {
|
||||
/// Register `token` for `call_id` before the call is spawned. If a
|
||||
/// `Cancel` already tombstoned this id, the token is cancelled
|
||||
/// immediately so the call starts cancelled. Returns whether the
|
||||
/// token was pre-cancelled (by a tombstone or because the registry was
|
||||
/// torn down).
|
||||
pub(crate) fn register(&self, call_id: ToolCallId, token: &CancellationToken) -> bool {
|
||||
if self.closed.load(Ordering::Acquire) {
|
||||
token.cancel();
|
||||
return true;
|
||||
}
|
||||
let pre_cancelled = self.pending.remove(&call_id).is_some();
|
||||
if pre_cancelled {
|
||||
token.cancel();
|
||||
}
|
||||
self.map.insert(call_id.clone(), token.clone());
|
||||
// Re-check after the insert: if `cancel_all` drained the map
|
||||
// between our closed-check and the insert, our entry would be
|
||||
// missed. The DashMap shard lock orders the insert against the
|
||||
// drain, so observing `closed` here guarantees we cancel + drop
|
||||
// any entry the drain could not reach (closes the teardown race).
|
||||
if self.closed.load(Ordering::Acquire) {
|
||||
if let Some((_, missed)) = self.map.remove(&call_id) {
|
||||
missed.cancel();
|
||||
}
|
||||
return true;
|
||||
}
|
||||
pre_cancelled
|
||||
}
|
||||
|
||||
/// Cancel a live call, else tombstone the id so the dispatcher
|
||||
/// cancels it at registration time. Returns true when a live token
|
||||
/// was found and cancelled.
|
||||
pub(crate) fn cancel(&self, call_id: &ToolCallId) -> bool {
|
||||
if let Some((_, token)) = self.map.remove(call_id) {
|
||||
token.cancel();
|
||||
true
|
||||
} else {
|
||||
if self.pending.len() >= MAX_PENDING_TOMBSTONES {
|
||||
// Evict one straggler tombstone (a cancel whose call never
|
||||
// registered) before inserting so the set stays bounded.
|
||||
// Collect the key first, then remove, so we never hold a
|
||||
// shard iterator across the removal.
|
||||
let stale = self.pending.iter().next().map(|e| e.key().clone());
|
||||
if let Some(stale) = stale {
|
||||
self.pending.remove(&stale);
|
||||
}
|
||||
}
|
||||
self.pending.insert(call_id.clone());
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Deregister a call's token on completion or cancel. Idempotent.
|
||||
pub(crate) fn deregister(&self, call_id: &ToolCallId) {
|
||||
self.map.remove(call_id);
|
||||
}
|
||||
|
||||
/// Whether [`Self::cancel_all`] has closed this registry. A closed
|
||||
/// registry marks a session whose loop is (or is about to be) torn
|
||||
/// down — used by the soft-rebind liveness gate.
|
||||
pub(crate) fn is_closed(&self) -> bool {
|
||||
self.closed.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
/// Drain-and-cancel every live token and close the registry. Used on
|
||||
/// session teardown (`unbind_session` / `shutdown` / full rebind of a
|
||||
/// dead loop — a soft rebind of a live session keeps its registry) so
|
||||
/// detached `execute_call` tasks wind down promptly AND any call
|
||||
/// dispatched in
|
||||
/// the teardown window starts cancelled (see [`Self::register`]).
|
||||
/// Returns the number of tokens cancelled.
|
||||
pub(crate) fn cancel_all(&self) -> usize {
|
||||
// Mark closed BEFORE draining so a concurrent `register` either
|
||||
// observes the close (and self-cancels) or has its entry drained
|
||||
// here — never both-miss.
|
||||
self.closed.store(true, Ordering::Release);
|
||||
let mut cancelled = 0;
|
||||
self.map.retain(|_, token| {
|
||||
token.cancel();
|
||||
cancelled += 1;
|
||||
false
|
||||
});
|
||||
// Drop tombstones too: teardown closes the registry, so no future
|
||||
// `register` will consume them. Leaving them would let a stale
|
||||
// straggler set survive to the end of the (already-done) session.
|
||||
self.pending.clear();
|
||||
cancelled
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn live_count(&self) -> usize {
|
||||
self.map.len()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn pending_count(&self) -> usize {
|
||||
self.pending.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn cid() -> ToolCallId {
|
||||
ToolCallId::new_v7()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_live_token_fires_and_removes_entry() {
|
||||
let reg = CancelRegistry::default();
|
||||
let id = cid();
|
||||
let token = CancellationToken::new();
|
||||
assert!(
|
||||
!reg.register(id.clone(), &token),
|
||||
"fresh register, no tombstone"
|
||||
);
|
||||
assert_eq!(reg.live_count(), 1);
|
||||
|
||||
assert!(reg.cancel(&id), "live token must report a hit");
|
||||
assert!(
|
||||
token.is_cancelled(),
|
||||
"the registered token must be cancelled"
|
||||
);
|
||||
assert_eq!(reg.live_count(), 0, "cancel removes the live entry");
|
||||
assert_eq!(reg.pending_count(), 0, "a live hit leaves no tombstone");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_before_registration_tombstones_then_register_pre_cancels() {
|
||||
let reg = CancelRegistry::default();
|
||||
let id = cid();
|
||||
|
||||
// Cancel arrives first: no live token, so it tombstones.
|
||||
assert!(!reg.cancel(&id), "no live token yet → miss");
|
||||
assert_eq!(reg.pending_count(), 1);
|
||||
assert_eq!(reg.live_count(), 0);
|
||||
|
||||
// Registration consumes the tombstone and starts cancelled.
|
||||
let token = CancellationToken::new();
|
||||
assert!(
|
||||
reg.register(id.clone(), &token),
|
||||
"register must report the pre-cancel"
|
||||
);
|
||||
assert!(token.is_cancelled(), "tombstone must pre-cancel the token");
|
||||
assert_eq!(reg.pending_count(), 0, "tombstone consumed at registration");
|
||||
assert_eq!(reg.live_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deregister_clears_live_entry_without_cancel() {
|
||||
let reg = CancelRegistry::default();
|
||||
let id = cid();
|
||||
let token = CancellationToken::new();
|
||||
reg.register(id.clone(), &token);
|
||||
|
||||
reg.deregister(&id);
|
||||
assert_eq!(reg.live_count(), 0);
|
||||
assert!(
|
||||
!token.is_cancelled(),
|
||||
"deregister on normal completion must NOT cancel the token"
|
||||
);
|
||||
// A later cancel for a completed call only tombstones (harmless).
|
||||
assert!(!reg.cancel(&id));
|
||||
assert_eq!(reg.pending_count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_all_drains_and_cancels_every_live_token() {
|
||||
let reg = CancelRegistry::default();
|
||||
let ids: Vec<ToolCallId> = (0..5).map(|_| cid()).collect();
|
||||
let tokens: Vec<CancellationToken> = ids
|
||||
.iter()
|
||||
.map(|id| {
|
||||
let t = CancellationToken::new();
|
||||
reg.register(id.clone(), &t);
|
||||
t
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(reg.live_count(), 5);
|
||||
|
||||
assert_eq!(
|
||||
reg.cancel_all(),
|
||||
5,
|
||||
"cancel_all reports every drained token"
|
||||
);
|
||||
assert_eq!(reg.live_count(), 0, "registry is empty after teardown");
|
||||
for token in &tokens {
|
||||
assert!(token.is_cancelled(), "every live token must be cancelled");
|
||||
}
|
||||
// Idempotent: a second teardown cancels nothing.
|
||||
assert_eq!(reg.cancel_all(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn register_after_cancel_all_starts_cancelled() {
|
||||
// Teardown race regression: once `cancel_all` has closed the
|
||||
// registry, a call dispatched in the teardown window must start
|
||||
// cancelled and must NOT linger as a live, uncancellable entry.
|
||||
let reg = CancelRegistry::default();
|
||||
assert_eq!(reg.cancel_all(), 0, "empty teardown cancels nothing");
|
||||
|
||||
let id = cid();
|
||||
let token = CancellationToken::new();
|
||||
assert!(
|
||||
reg.register(id.clone(), &token),
|
||||
"register on a closed registry must report pre-cancel"
|
||||
);
|
||||
assert!(
|
||||
token.is_cancelled(),
|
||||
"a call dispatched after teardown must start cancelled"
|
||||
);
|
||||
assert_eq!(
|
||||
reg.live_count(),
|
||||
0,
|
||||
"a closed-registry register must not leave a live (orphan) entry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn register_without_tombstone_does_not_cancel() {
|
||||
let reg = CancelRegistry::default();
|
||||
let id = cid();
|
||||
let token = CancellationToken::new();
|
||||
assert!(!reg.register(id, &token));
|
||||
assert!(
|
||||
!token.is_cancelled(),
|
||||
"a clean registration must leave the token live"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_all_clears_pending_tombstones() {
|
||||
let reg = CancelRegistry::default();
|
||||
reg.cancel(&cid());
|
||||
reg.cancel(&cid());
|
||||
assert_eq!(reg.pending_count(), 2);
|
||||
reg.cancel_all();
|
||||
assert_eq!(
|
||||
reg.pending_count(),
|
||||
0,
|
||||
"teardown must drop pending tombstones"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_tombstones_stay_bounded_under_spurious_cancels() {
|
||||
// A long-lived session that keeps receiving cancels for call_ids
|
||||
// that never register (e.g. cancels racing call completion) must
|
||||
// not grow `pending` without bound.
|
||||
let reg = CancelRegistry::default();
|
||||
for _ in 0..(MAX_PENDING_TOMBSTONES + 256) {
|
||||
assert!(!reg.cancel(&cid()), "never-registered id is a miss");
|
||||
}
|
||||
assert!(
|
||||
reg.pending_count() <= MAX_PENDING_TOMBSTONES,
|
||||
"tombstone set must stay within its cap, got {}",
|
||||
reg.pending_count()
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,224 @@
|
||||
//! Shared connection-borrow lifecycle for `ToolServer` and `ToolHarness`.
|
||||
//!
|
||||
//! Wraps a pooled [`HubConnection`] with a [`CancellationToken`] for
|
||||
//! shutdown coordination and an at-most-once `torn_down` guard.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use kigi_tool_protocol::ConnectionKind;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use url::Url;
|
||||
|
||||
use crate::auth::AuthProvider;
|
||||
use crate::connection::{
|
||||
ConnectCallback, ConnectionTuning, DisconnectCallback, HubConnection, ReconnectCallback,
|
||||
};
|
||||
use crate::error::ClientError;
|
||||
use crate::pool::HubConnectionPool;
|
||||
|
||||
/// Borrowed slice of a pooled [`HubConnection`] plus the refcount of
|
||||
/// session bindings the borrower owns. Drop guard lives here so the
|
||||
/// teardown sequence is at-most-once across explicit `shutdown` and
|
||||
/// the `Drop` fallback.
|
||||
pub(crate) struct ConnectionBorrow {
|
||||
connection: Arc<HubConnection>,
|
||||
shutdown: CancellationToken,
|
||||
/// At-most-once guard coordinated via `compare_exchange`.
|
||||
torn_down: AtomicBool,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ConnectionBorrow {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("ConnectionBorrow")
|
||||
.field(
|
||||
"torn_down",
|
||||
&self.torn_down.load(std::sync::atomic::Ordering::Relaxed),
|
||||
)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl ConnectionBorrow {
|
||||
/// Resolve a pool entry, refcount-bind every requested session,
|
||||
/// and return a borrow. On any per-session bind failure the
|
||||
/// already-bound sessions are unregistered before returning the
|
||||
/// error so partial state never leaks.
|
||||
pub(crate) async fn acquire(
|
||||
pool: Arc<HubConnectionPool>,
|
||||
url: Url,
|
||||
auth: Arc<dyn AuthProvider>,
|
||||
kind: ConnectionKind,
|
||||
on_reconnect: Option<Arc<ReconnectCallback>>,
|
||||
on_disconnect: Option<Arc<DisconnectCallback>>,
|
||||
on_connect: Option<Arc<ConnectCallback>>,
|
||||
server_id: Option<kigi_tool_protocol::ServerId>,
|
||||
server_description: Option<String>,
|
||||
server_metadata: Option<serde_json::Value>,
|
||||
alpha_test_key: Option<String>,
|
||||
allow_insecure_ws: bool,
|
||||
tuning: ConnectionTuning,
|
||||
) -> Result<Self, ClientError> {
|
||||
let connection = pool
|
||||
.get_or_connect_tuned(
|
||||
url,
|
||||
auth,
|
||||
kind,
|
||||
on_reconnect,
|
||||
on_disconnect,
|
||||
on_connect,
|
||||
server_id,
|
||||
server_description,
|
||||
server_metadata,
|
||||
alpha_test_key,
|
||||
allow_insecure_ws,
|
||||
tuning,
|
||||
)
|
||||
.await?;
|
||||
Ok(Self {
|
||||
connection,
|
||||
shutdown: CancellationToken::new(),
|
||||
torn_down: AtomicBool::new(false),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn connection(&self) -> &Arc<HubConnection> {
|
||||
&self.connection
|
||||
}
|
||||
|
||||
pub(crate) fn shutdown_token(&self) -> &CancellationToken {
|
||||
&self.shutdown
|
||||
}
|
||||
|
||||
/// Returns `true` if this caller won the at-most-once teardown.
|
||||
pub(crate) fn begin_teardown(&self) -> bool {
|
||||
self.torn_down
|
||||
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
|
||||
.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::auth::AuthCredential;
|
||||
use axum::Router;
|
||||
use axum::extract::WebSocketUpgrade;
|
||||
use axum::extract::ws::{Message, WebSocket};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::routing::get;
|
||||
use serde_json::json;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
/// Spawn an in-process mock server that completes the WebSocket
|
||||
/// handshake and ignores everything else. Returned address is
|
||||
/// bound on `127.0.0.1`.
|
||||
async fn spawn_borrow_mock_hub() -> SocketAddr {
|
||||
let app = Router::new().route("/v1/tools", get(ws_upgrade));
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind ephemeral");
|
||||
let addr = listener.local_addr().expect("local addr");
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app.into_make_service()).await;
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
addr
|
||||
}
|
||||
|
||||
async fn ws_upgrade(ws: WebSocketUpgrade) -> impl IntoResponse {
|
||||
ws.on_upgrade(handle_socket)
|
||||
}
|
||||
|
||||
async fn handle_socket(mut socket: WebSocket) {
|
||||
let _ = socket.recv().await;
|
||||
let ack = json!({
|
||||
"connection_id": "borrow-mock",
|
||||
"user_id": "test",
|
||||
"computer_hub_version": "test",
|
||||
"supported_protocol_versions": ["1.0.0"],
|
||||
});
|
||||
let _ = socket.send(Message::Text(ack.to_string().into())).await;
|
||||
// Keep the WebSocket alive until the client disconnects.
|
||||
// These tests only exercise borrow lifecycle (teardown
|
||||
// atomicity), not protocol frames.
|
||||
while let Some(Ok(_msg)) = socket.recv().await {}
|
||||
}
|
||||
|
||||
async fn acquire_borrow() -> ConnectionBorrow {
|
||||
let addr = spawn_borrow_mock_hub().await;
|
||||
let url = Url::parse(&format!("ws://{addr}/v1/tools")).expect("valid url");
|
||||
let cred: Arc<dyn AuthProvider> = Arc::new(AuthCredential::bearer("ignored"));
|
||||
let pool = HubConnectionPool::new();
|
||||
ConnectionBorrow::acquire(
|
||||
pool,
|
||||
url,
|
||||
cred,
|
||||
ConnectionKind::Harness,
|
||||
None, // on_reconnect
|
||||
None, // on_disconnect
|
||||
None, // on_connect
|
||||
None, // server_id
|
||||
None, // server_description
|
||||
None, // server_metadata
|
||||
None, // alpha_test_key
|
||||
false,
|
||||
ConnectionTuning::default(),
|
||||
)
|
||||
.await
|
||||
.expect("acquire borrow")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn begin_teardown_returns_true_once_and_false_after() {
|
||||
let borrow = acquire_borrow().await;
|
||||
assert!(
|
||||
borrow.begin_teardown(),
|
||||
"first call wins the at-most-once transition"
|
||||
);
|
||||
assert!(
|
||||
!borrow.begin_teardown(),
|
||||
"subsequent calls observe the already-torn-down state"
|
||||
);
|
||||
assert!(
|
||||
!borrow.begin_teardown(),
|
||||
"the at-most-once transition is sticky"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn begin_teardown_is_atomic_under_concurrent_callers() {
|
||||
let borrow = Arc::new(acquire_borrow().await);
|
||||
let n_callers = 64;
|
||||
let barrier = Arc::new(tokio::sync::Barrier::new(n_callers));
|
||||
let mut handles = Vec::with_capacity(n_callers);
|
||||
for _ in 0..n_callers {
|
||||
let borrow = borrow.clone();
|
||||
let barrier = barrier.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
barrier.wait().await;
|
||||
borrow.begin_teardown()
|
||||
}));
|
||||
}
|
||||
let mut wins = 0usize;
|
||||
for h in handles {
|
||||
if h.await.expect("join") {
|
||||
wins += 1;
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
wins, 1,
|
||||
"exactly one of {n_callers} concurrent callers must win the at-most-once transition"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn acquire_returns_zero_bound_sessions() {
|
||||
let borrow = acquire_borrow().await;
|
||||
assert_eq!(borrow.connection().bound_session_count(), 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,973 @@
|
||||
//! Inbound frame demultiplexer.
|
||||
//!
|
||||
//! Frames inbound from the WebSocket fall into four buckets:
|
||||
//!
|
||||
//! 1. JSON-RPC **responses** correlated to a previously-issued request
|
||||
//! by `id`. Routed through the crate-internal response-waiter map.
|
||||
//! 2. **`tool_call_progress` notifications** correlated to a per-call
|
||||
//! `tool_call_id` carried in `params`. Routed through the
|
||||
//! crate-internal progress-waiter map registered via
|
||||
//! `Demux::try_register_progress_waiter` (crate-internal).
|
||||
//! 3. JSON-RPC **requests / notifications** carrying a `session_id` —
|
||||
//! routed to the per-session inbox registered via
|
||||
//! [`Demux::register_session_inbox`].
|
||||
//! 4. Connection-level frames (handshake, ping/pong) that the
|
||||
//! connection actor handles directly without going through the demux.
|
||||
//!
|
||||
//! The demux owns the session inbox map, the in-flight response
|
||||
//! waiters, and the per-call progress waiters; the connection actor
|
||||
//! parses each text frame, classifies it,
|
||||
//! and pushes it through this module.
|
||||
//!
|
||||
//! Routing inbound frames is non-blocking: a full session inbox or a
|
||||
//! dropped receiver returns a typed [`RouteOutcome`] variant rather
|
||||
//! than awaiting the inbox. Blocking on a slow consumer would back up
|
||||
//! the entire connection actor and starve every other session sharing
|
||||
//! the socket.
|
||||
|
||||
use dashmap::DashMap;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::warn;
|
||||
|
||||
use kigi_tool_protocol::{
|
||||
JsonRpcId, JsonRpcResponse, RequestId, SessionId, ToolCallId, ToolCallProgressFrame,
|
||||
};
|
||||
|
||||
use crate::error::ClientError;
|
||||
|
||||
/// Frame routed to a session inbox.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum InboundFrame {
|
||||
/// A request the inbox owner must answer (any session frame carrying
|
||||
/// an `id`): a `tool_call_request`, or a reverse-direction `hook`
|
||||
/// answered via `ToolHarness::send_hook_reply`. Carries raw JSON.
|
||||
Request(Value),
|
||||
/// Server-issued notification (e.g. `tool.notification`) — fire-and-
|
||||
/// forget, no reply expected.
|
||||
Notification(Value),
|
||||
}
|
||||
|
||||
/// Outcome of [`Demux::route`].
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub enum RouteOutcome {
|
||||
/// Matched a response waiter; the oneshot was fulfilled.
|
||||
Response,
|
||||
/// Forwarded to a session inbox.
|
||||
Session,
|
||||
/// Matched a progress waiter; the progress frame was forwarded to
|
||||
/// the per-call progress channel.
|
||||
Progress,
|
||||
/// No inbox is bound for the targeted session.
|
||||
UnknownSession,
|
||||
/// No progress waiter is parked for the targeted `tool_call_id`. The
|
||||
/// caller's stream is no longer subscribed (typical post-terminal),
|
||||
/// so the frame is dropped.
|
||||
UnknownProgress,
|
||||
/// Connection-level notification broadcast to subscribers.
|
||||
Notification,
|
||||
/// No waiter is parked for the targeted request id, OR the frame
|
||||
/// was unaddressable.
|
||||
Unrouted,
|
||||
/// The session inbox sender was full; the frame was dropped to
|
||||
/// avoid blocking the connection actor.
|
||||
InboxFull,
|
||||
/// The session inbox receiver was dropped (e.g. the consumer's
|
||||
/// run loop exited); the binding is now stale and the frame was
|
||||
/// dropped.
|
||||
SessionDropped,
|
||||
/// The progress channel was full; the frame was dropped to avoid
|
||||
/// blocking the connection actor. The caller's stream consumer
|
||||
/// fell behind on draining progress.
|
||||
ProgressFull,
|
||||
/// The progress receiver was dropped (e.g. the caller's stream was
|
||||
/// dropped); the waiter binding is now stale and the frame was
|
||||
/// dropped.
|
||||
ProgressDropped,
|
||||
}
|
||||
|
||||
/// Demux state. Cheap to construct; uses [`DashMap`] internally so
|
||||
/// concurrent registers and routes never block each other.
|
||||
#[derive(Debug)]
|
||||
pub struct Demux {
|
||||
sessions: DashMap<SessionId, tokio::sync::mpsc::Sender<InboundFrame>>,
|
||||
waiters: DashMap<RequestId, oneshot::Sender<Result<JsonRpcResponse, ClientError>>>,
|
||||
/// Session index for `tool.call` response waiters only. Lets the SDK
|
||||
/// in-flight short-circuit fail every parked call for a session on a
|
||||
/// workspace Disconnected notification without waiting for the server.
|
||||
/// Turn-hook / session-RPC waiters are NOT indexed here, so the
|
||||
/// short-circuit never touches them.
|
||||
call_sessions: DashMap<RequestId, SessionId>,
|
||||
progress: DashMap<ToolCallId, tokio::sync::mpsc::Sender<ToolCallProgressFrame>>,
|
||||
/// Broadcast channel for connection-level notifications (no session_id).
|
||||
notifications: tokio::sync::broadcast::Sender<Value>,
|
||||
/// Clone of the connection's outbound sender. Used to synthesize the
|
||||
/// overloaded (-32016) response when a session inbox is full so a
|
||||
/// Request is rejected with an error rather than silently dropped.
|
||||
/// `None` in unit tests that construct a bare demux.
|
||||
outbound: Option<tokio::sync::mpsc::Sender<String>>,
|
||||
}
|
||||
|
||||
impl Default for Demux {
|
||||
fn default() -> Self {
|
||||
let (notifications, _) = tokio::sync::broadcast::channel(64);
|
||||
Self {
|
||||
sessions: DashMap::new(),
|
||||
waiters: DashMap::new(),
|
||||
call_sessions: DashMap::new(),
|
||||
progress: DashMap::new(),
|
||||
notifications,
|
||||
outbound: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Demux {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Construct a demux wired to the connection's outbound sender so the
|
||||
/// inbox-full Request path can ship an overloaded (-32016) response.
|
||||
pub fn with_outbound(outbound: tokio::sync::mpsc::Sender<String>) -> Self {
|
||||
Self {
|
||||
outbound: Some(outbound),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Subscribe to connection-level notifications (no session_id).
|
||||
pub fn subscribe_notifications(&self) -> tokio::sync::broadcast::Receiver<Value> {
|
||||
self.notifications.subscribe()
|
||||
}
|
||||
|
||||
/// Bind `session_id` to `inbox`; replaces any existing binding.
|
||||
/// Returns the previous sender if one existed; the caller may
|
||||
/// drop or drain it as appropriate.
|
||||
pub fn register_session_inbox(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
inbox: tokio::sync::mpsc::Sender<InboundFrame>,
|
||||
) -> Option<tokio::sync::mpsc::Sender<InboundFrame>> {
|
||||
self.sessions.insert(session_id, inbox)
|
||||
}
|
||||
|
||||
/// Remove the inbox bound to `session_id`. The returned sender (if
|
||||
/// present) is dropped by the caller, signalling EOF to its
|
||||
/// receiver task.
|
||||
pub fn unregister_session_inbox(
|
||||
&self,
|
||||
session_id: &SessionId,
|
||||
) -> Option<tokio::sync::mpsc::Sender<InboundFrame>> {
|
||||
self.sessions.remove(session_id).map(|(_, sender)| sender)
|
||||
}
|
||||
|
||||
/// Park a oneshot waiter for `request_id`. Crate-internal: only
|
||||
/// the connection actor allocates request ids.
|
||||
pub(crate) fn register_response_waiter(
|
||||
&self,
|
||||
request_id: RequestId,
|
||||
waiter: oneshot::Sender<Result<JsonRpcResponse, ClientError>>,
|
||||
) {
|
||||
self.waiters.insert(request_id, waiter);
|
||||
}
|
||||
|
||||
/// Park a `tool.call` response waiter and record its `session_id` so the
|
||||
/// SDK in-flight short-circuit ([`Self::fail_calls_for_session`]) can
|
||||
/// resolve it on a workspace Disconnected notification. Crate-internal.
|
||||
pub(crate) fn register_call_response_waiter(
|
||||
&self,
|
||||
request_id: RequestId,
|
||||
session_id: SessionId,
|
||||
waiter: oneshot::Sender<Result<JsonRpcResponse, ClientError>>,
|
||||
) {
|
||||
self.call_sessions.insert(request_id.clone(), session_id);
|
||||
self.waiters.insert(request_id, waiter);
|
||||
}
|
||||
|
||||
/// Pop the waiter for `request_id`, if any. Also drops the session index
|
||||
/// entry so the two maps stay consistent. Crate-internal.
|
||||
pub(crate) fn take_response_waiter(
|
||||
&self,
|
||||
request_id: &RequestId,
|
||||
) -> Option<oneshot::Sender<Result<JsonRpcResponse, ClientError>>> {
|
||||
self.call_sessions.remove(request_id);
|
||||
self.waiters.remove(request_id).map(|(_, waiter)| waiter)
|
||||
}
|
||||
|
||||
/// Fail every in-flight `tool.call` waiter bound to `session_id`,
|
||||
/// completing each with `result_factory`. Returns the number resolved.
|
||||
///
|
||||
/// Drives the SDK in-flight short-circuit: on a workspace
|
||||
/// `ToolServerStatusChanged(Disconnected)` notification the harness fails
|
||||
/// its parked calls for that session promptly instead of parking until
|
||||
/// `rpc_ttl_ms`. Idempotent with the server-side cancel — each waiter is
|
||||
/// taken at most once, so a call already resolved by the server is skipped.
|
||||
pub(crate) fn fail_calls_for_session<F>(
|
||||
&self,
|
||||
session_id: &SessionId,
|
||||
result_factory: F,
|
||||
) -> usize
|
||||
where
|
||||
F: Fn() -> ClientError,
|
||||
{
|
||||
// Snapshot the matching request ids first so we never hold a DashMap
|
||||
// shard lock across the oneshot send.
|
||||
let request_ids: Vec<RequestId> = self
|
||||
.call_sessions
|
||||
.iter()
|
||||
.filter(|kv| kv.value() == session_id)
|
||||
.map(|kv| kv.key().clone())
|
||||
.collect();
|
||||
let mut resolved = 0;
|
||||
for request_id in request_ids {
|
||||
if let Some(waiter) = self.take_response_waiter(&request_id)
|
||||
&& waiter.send(Err(result_factory())).is_ok()
|
||||
{
|
||||
resolved += 1;
|
||||
}
|
||||
}
|
||||
resolved
|
||||
}
|
||||
|
||||
/// Park a per-call progress sender keyed by `tool_call_id`.
|
||||
/// Returns `Err(progress)` (handing the not-yet-inserted sender
|
||||
/// back) when another in-flight call already owns the id, leaving
|
||||
/// the prior waiter intact. The caller drops the matching
|
||||
/// receiver to terminate the subscription — subsequent inbound
|
||||
/// progress for the same id is silently dropped via
|
||||
/// [`RouteOutcome::ProgressDropped`].
|
||||
///
|
||||
/// Atomic check-then-insert under a single shard lock so a
|
||||
/// concurrent caller cannot observe a transient empty slot.
|
||||
pub(crate) fn try_register_progress_waiter(
|
||||
&self,
|
||||
tool_call_id: ToolCallId,
|
||||
progress: tokio::sync::mpsc::Sender<ToolCallProgressFrame>,
|
||||
) -> Result<(), tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
|
||||
use dashmap::mapref::entry::Entry;
|
||||
match self.progress.entry(tool_call_id) {
|
||||
Entry::Occupied(_) => Err(progress),
|
||||
Entry::Vacant(slot) => {
|
||||
slot.insert(progress);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove the progress sender bound to `tool_call_id`. Crate-internal;
|
||||
/// called by the harness once the terminal frame for `tool_call_id`
|
||||
/// has been observed.
|
||||
pub(crate) fn unregister_progress_waiter(
|
||||
&self,
|
||||
tool_call_id: &ToolCallId,
|
||||
) -> Option<tokio::sync::mpsc::Sender<ToolCallProgressFrame>> {
|
||||
self.progress.remove(tool_call_id).map(|(_, tx)| tx)
|
||||
}
|
||||
|
||||
/// Drain every parked waiter, completing each with `result_factory`.
|
||||
/// Used by the reconnect path to fast-fail in-flight calls with
|
||||
/// [`ClientError::NetworkError`]. Crate-internal.
|
||||
pub(crate) fn drain_waiters_with<F>(&self, result_factory: F)
|
||||
where
|
||||
F: Fn() -> ClientError,
|
||||
{
|
||||
// Snapshot keys, then remove individually so we never hold
|
||||
// a DashMap shard lock across the oneshot send.
|
||||
let keys: Vec<RequestId> = self.waiters.iter().map(|kv| kv.key().clone()).collect();
|
||||
for key in keys {
|
||||
if let Some((_, waiter)) = self.waiters.remove(&key) {
|
||||
self.call_sessions.remove(&key);
|
||||
let _ = waiter.send(Err(result_factory()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop every parked progress sender. Used by the reconnect path
|
||||
/// after [`Self::drain_waiters_with`]: the response waiter resolves
|
||||
/// with `NetworkError` and the matching progress channel closes,
|
||||
/// so any in-flight harness call stream terminates promptly
|
||||
/// instead of stalling on a half-empty progress channel.
|
||||
pub(crate) fn drain_progress(&self) {
|
||||
let keys: Vec<ToolCallId> = self.progress.iter().map(|kv| kv.key().clone()).collect();
|
||||
for key in keys {
|
||||
self.progress.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
/// Route a parsed JSON value. Classification rules:
|
||||
///
|
||||
/// - presence of `result`/`error` → response, routed to waiter;
|
||||
/// - method == `tool_call_progress` notification → progress waiter
|
||||
/// keyed by `params.tool_call_id`;
|
||||
/// - presence of `session_id` → session inbox, request vs.
|
||||
/// notification distinguished by the presence of `id`;
|
||||
/// - otherwise → [`RouteOutcome::Unrouted`].
|
||||
///
|
||||
/// Routing to a session inbox or progress channel uses non-blocking
|
||||
/// `try_send`. A full inbox or progress channel returns the matching
|
||||
/// `*Full` variant; a dropped receiver returns the matching
|
||||
/// `*Dropped` variant. Either way the frame is dropped without
|
||||
/// awaiting the consumer, so a slow handler never starves other
|
||||
/// sessions or calls multiplexed onto the same connection.
|
||||
pub fn route(&self, frame: Value) -> RouteOutcome {
|
||||
crate::metrics::demux_inbox_depth_set(self.sessions.len() as i64);
|
||||
if frame.get("result").is_some() || frame.get("error").is_some() {
|
||||
return self.route_response(frame);
|
||||
}
|
||||
if frame.get("method").and_then(Value::as_str) == Some("tool_call_progress") {
|
||||
return self.route_progress(frame);
|
||||
}
|
||||
if frame.get("session_id").is_some() {
|
||||
return self.route_session(frame);
|
||||
}
|
||||
// Connection-level notification (e.g. session.bind, session.unbind).
|
||||
if frame.get("method").is_some() {
|
||||
let _ = self.notifications.send(frame);
|
||||
return RouteOutcome::Notification;
|
||||
}
|
||||
RouteOutcome::Unrouted
|
||||
}
|
||||
|
||||
fn route_progress(&self, frame: Value) -> RouteOutcome {
|
||||
let Some(params) = frame.get("params") else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Some(call_id_str) = params.get("tool_call_id").and_then(Value::as_str) else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Ok(tool_call_id) = ToolCallId::new(call_id_str) else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Some(sender) = self.progress.get(&tool_call_id) else {
|
||||
return RouteOutcome::UnknownProgress;
|
||||
};
|
||||
let tx = sender.value().clone();
|
||||
drop(sender);
|
||||
let progress_frame: ToolCallProgressFrame = match serde_json::from_value(params.clone()) {
|
||||
Ok(p) => p,
|
||||
Err(err) => {
|
||||
warn!(%tool_call_id, ?err, "failed to decode tool_call_progress params");
|
||||
return RouteOutcome::Unrouted;
|
||||
}
|
||||
};
|
||||
match tx.try_send(progress_frame) {
|
||||
Ok(()) => RouteOutcome::Progress,
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
|
||||
warn!(%tool_call_id, "progress channel full; dropping inbound progress frame");
|
||||
RouteOutcome::ProgressFull
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
||||
self.progress.remove(&tool_call_id);
|
||||
RouteOutcome::ProgressDropped
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn route_response(&self, frame: Value) -> RouteOutcome {
|
||||
let Some(id_value) = frame.get("id") else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let request_id = match id_value {
|
||||
Value::String(s) => RequestId::new(s.as_str()).ok(),
|
||||
Value::Number(n) => RequestId::new(n.to_string()).ok(),
|
||||
_ => None,
|
||||
};
|
||||
let Some(request_id) = request_id else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Some(waiter) = self.take_response_waiter(&request_id) else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let parsed: Result<JsonRpcResponse, ClientError> =
|
||||
serde_json::from_value::<JsonRpcResponse>(frame).map_err(ClientError::from);
|
||||
let _ = waiter.send(parsed);
|
||||
RouteOutcome::Response
|
||||
}
|
||||
|
||||
fn route_session(&self, frame: Value) -> RouteOutcome {
|
||||
let Some(sid_str) = frame.get("session_id").and_then(Value::as_str) else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Ok(session_id) = SessionId::new(sid_str) else {
|
||||
return RouteOutcome::Unrouted;
|
||||
};
|
||||
let Some(sender) = self.sessions.get(&session_id) else {
|
||||
return RouteOutcome::UnknownSession;
|
||||
};
|
||||
let inbox = sender.value().clone();
|
||||
drop(sender);
|
||||
let kind = if frame.get("id").is_some() {
|
||||
InboundFrame::Request(frame)
|
||||
} else {
|
||||
InboundFrame::Notification(frame)
|
||||
};
|
||||
match inbox.try_send(kind) {
|
||||
Ok(()) => RouteOutcome::Session,
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Full(frame)) => {
|
||||
self.reject_inbox_full(&session_id, frame);
|
||||
RouteOutcome::InboxFull
|
||||
}
|
||||
Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
|
||||
warn!(%session_id, "session inbox dropped; binding stale");
|
||||
self.sessions.remove(&session_id);
|
||||
RouteOutcome::SessionDropped
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a full session inbox without blocking the reader.
|
||||
///
|
||||
/// A Request (has an `id`) is rejected with the shared overloaded
|
||||
/// (-32016 "tool_busy") response on a best-effort `try_send`; if the
|
||||
/// outbound is *also* full the rejection itself is dropped and metered
|
||||
/// (`inbox_full_reject_send_failed`). A Notification (no `id`) stays
|
||||
/// fire-and-forget and is metered (`inbox_full_notification_dropped`).
|
||||
fn reject_inbox_full(&self, session_id: &SessionId, frame: InboundFrame) {
|
||||
let InboundFrame::Request(value) = frame else {
|
||||
crate::metrics::inbox_full_notification_dropped();
|
||||
return;
|
||||
};
|
||||
crate::metrics::inbox_full_request_rejected();
|
||||
warn!(%session_id, "session inbox full; rejecting request with tool_busy");
|
||||
let Some(out) = &self.outbound else {
|
||||
return;
|
||||
};
|
||||
// A `Request` always carries an `id` (that is how `route_session`
|
||||
// classifies it). A well-formed id deserializes into a `JsonRpcId`;
|
||||
// a malformed id (object/array/bool/null) cannot, but the request
|
||||
// must STILL get an overloaded response rather than be silently
|
||||
// dropped, so we fall back to echoing the raw id JSON as a string.
|
||||
let raw_id = value.get("id");
|
||||
let id = raw_id
|
||||
.and_then(|v| serde_json::from_value::<JsonRpcId>(v.clone()).ok())
|
||||
.unwrap_or_else(|| {
|
||||
JsonRpcId::new_string(raw_id.map(ToString::to_string).unwrap_or_default())
|
||||
});
|
||||
let response = crate::admission::overloaded_response(id, session_id.clone());
|
||||
let Ok(text) = serde_json::to_string(&response) else {
|
||||
return;
|
||||
};
|
||||
if out.try_send(text).is_err() {
|
||||
crate::metrics::inbox_full_reject_send_failed();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_route_matches_waiter() {
|
||||
let demux = Demux::new();
|
||||
let request_id = RequestId::new("r1").expect("valid");
|
||||
let (tx, rx) = oneshot::channel();
|
||||
demux.register_response_waiter(request_id.clone(), tx);
|
||||
let outcome = demux.route(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "r1",
|
||||
"result": {"outcome": "bound"},
|
||||
}));
|
||||
assert_eq!(outcome, RouteOutcome::Response);
|
||||
let resp = rx.await.expect("waiter").expect("ok");
|
||||
assert_eq!(resp.id.to_string(), "r1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fail_calls_for_session_resolves_only_matching_call_waiters() {
|
||||
// Fails exactly the session's `tool.call` waiters; other sessions'
|
||||
// calls and non-call (turn-hook) waiters stay parked.
|
||||
let demux = Demux::new();
|
||||
let s1 = SessionId::new("s1").expect("valid");
|
||||
let s2 = SessionId::new("s2").expect("valid");
|
||||
|
||||
let (tx_a, rx_a) = oneshot::channel();
|
||||
let (tx_b, rx_b) = oneshot::channel();
|
||||
let (tx_other, rx_other) = oneshot::channel();
|
||||
// Two calls on s1, one on s2.
|
||||
demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
|
||||
demux.register_call_response_waiter(RequestId::new("b").unwrap(), s1.clone(), tx_b);
|
||||
demux.register_call_response_waiter(RequestId::new("c").unwrap(), s2.clone(), tx_other);
|
||||
// A non-call waiter (e.g. a turn hook) on s1 — NOT session-indexed.
|
||||
let (tx_hook, rx_hook) = oneshot::channel();
|
||||
demux.register_response_waiter(RequestId::new("hook").unwrap(), tx_hook);
|
||||
|
||||
let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
|
||||
assert_eq!(n, 2, "only the two s1 call waiters are failed");
|
||||
|
||||
assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_)))));
|
||||
assert!(matches!(rx_b.await, Ok(Err(ClientError::NetworkError(_)))));
|
||||
// s2's call and the turn-hook waiter are untouched (still parked).
|
||||
assert!(
|
||||
demux
|
||||
.take_response_waiter(&RequestId::new("c").unwrap())
|
||||
.is_some(),
|
||||
"the s2 call must remain parked"
|
||||
);
|
||||
assert!(
|
||||
demux
|
||||
.take_response_waiter(&RequestId::new("hook").unwrap())
|
||||
.is_some(),
|
||||
"the turn-hook waiter must remain parked"
|
||||
);
|
||||
// Keep the receivers alive until the asserts above ran.
|
||||
drop((rx_other, rx_hook));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fail_calls_for_session_is_idempotent_after_resolution() {
|
||||
// A call already resolved (waiter taken) must not be double-counted by
|
||||
// the short-circuit.
|
||||
let demux = Demux::new();
|
||||
let s1 = SessionId::new("s1").expect("valid");
|
||||
let (tx_a, rx_a) = oneshot::channel();
|
||||
demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
|
||||
// Simulate the server-side resolution taking the waiter first.
|
||||
let waiter = demux
|
||||
.take_response_waiter(&RequestId::new("a").unwrap())
|
||||
.expect("waiter present");
|
||||
drop(waiter);
|
||||
drop(rx_a);
|
||||
let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
|
||||
assert_eq!(n, 0, "already-resolved call is not re-failed");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn short_circuit_then_late_response_is_unrouted() {
|
||||
// A short-circuit that resolves first leaves no waiter, so a late server
|
||||
// response for the same id is dropped (no double-resolve).
|
||||
let demux = Demux::new();
|
||||
let s1 = SessionId::new("s1").expect("valid");
|
||||
let (tx_a, rx_a) = oneshot::channel();
|
||||
demux.register_call_response_waiter(RequestId::new("a").unwrap(), s1.clone(), tx_a);
|
||||
|
||||
let n = demux.fail_calls_for_session(&s1, || ClientError::NetworkError("gone".to_owned()));
|
||||
assert_eq!(n, 1);
|
||||
assert!(matches!(rx_a.await, Ok(Err(ClientError::NetworkError(_)))));
|
||||
|
||||
let outcome = demux.route(json!({ "jsonrpc": "2.0", "id": "a", "result": {} }));
|
||||
assert_eq!(
|
||||
outcome,
|
||||
RouteOutcome::Unrouted,
|
||||
"the late normal response must not double-resolve the call"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_route_pushes_to_inbox() {
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("s1").expect("valid");
|
||||
let (tx, mut rx) = mpsc::channel(4);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "x",
|
||||
"session_id": "s1",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
});
|
||||
let outcome = demux.route(frame.clone());
|
||||
assert_eq!(outcome, RouteOutcome::Session);
|
||||
match rx.recv().await {
|
||||
Some(InboundFrame::Request(value)) => assert_eq!(value, frame),
|
||||
other => panic!("expected request inbound; got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reverse_hook_request_routes_to_inbox_as_request() {
|
||||
// A reverse hook request carries an `id`, so it must route to the
|
||||
// inbox as `Request` (not `Notification`) for the harness to answer.
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("s1").expect("valid");
|
||||
let (tx, mut rx) = mpsc::channel(4);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let hook = kigi_tool_protocol::HookFrame::custom_request(
|
||||
session.clone(),
|
||||
"hook-7".to_owned(),
|
||||
crate::harness::PERMISSION_REQUEST_KIND.to_owned(),
|
||||
json!({}),
|
||||
);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "h1",
|
||||
"session_id": "s1",
|
||||
"method": kigi_tool_protocol::Method::Hook.as_wire_str(),
|
||||
"params": serde_json::to_value(&hook).expect("serialize hook"),
|
||||
});
|
||||
assert_eq!(demux.route(frame.clone()), RouteOutcome::Session);
|
||||
match rx.recv().await {
|
||||
Some(InboundFrame::Request(value)) => assert_eq!(value, frame),
|
||||
other => panic!("expected request inbound; got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn notification_classified_without_id() {
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("s1").expect("valid");
|
||||
let (tx, mut rx) = mpsc::channel(4);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tool.notification",
|
||||
"params": {},
|
||||
});
|
||||
let outcome = demux.route(frame);
|
||||
assert_eq!(outcome, RouteOutcome::Session);
|
||||
match rx.recv().await {
|
||||
Some(InboundFrame::Notification(_)) => {}
|
||||
other => panic!("expected notification; got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_session_returns_unknown_session() {
|
||||
let demux = Demux::new();
|
||||
let outcome = demux.route(
|
||||
json!({"jsonrpc":"2.0","id":"x","session_id":"missing","method":"x","params":{}}),
|
||||
);
|
||||
assert_eq!(outcome, RouteOutcome::UnknownSession);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_request_id_returns_unrouted() {
|
||||
let demux = Demux::new();
|
||||
let outcome = demux.route(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "missing",
|
||||
"result": {},
|
||||
}));
|
||||
assert_eq!(outcome, RouteOutcome::Unrouted);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_inbox_returns_inbox_full_without_blocking() {
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("backed_up").expect("valid");
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = || {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "x",
|
||||
"session_id": "backed_up",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
})
|
||||
};
|
||||
// First send fills capacity.
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::Session);
|
||||
// Second send must NOT block; it returns InboxFull.
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::InboxFull);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropped_receiver_returns_session_dropped() {
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("gone").expect("valid");
|
||||
let (tx, rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
drop(rx);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "x",
|
||||
"session_id": "gone",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
});
|
||||
assert_eq!(demux.route(frame), RouteOutcome::SessionDropped);
|
||||
// Stale binding should have been removed.
|
||||
assert!(demux.sessions.get(&session).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inbox_full_request_synthesizes_overloaded_response_onto_outbound() {
|
||||
// A full session inbox for a Request must produce the shared
|
||||
// -32016 "tool_busy" response on outbound, not a silent drop.
|
||||
let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
|
||||
let demux = Demux::with_outbound(out_tx);
|
||||
let session = SessionId::new("busy").expect("valid");
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = |id: &str| {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"session_id": "busy",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
})
|
||||
};
|
||||
// First fills capacity (cap 1); second overflows → InboxFull.
|
||||
assert_eq!(demux.route(frame("a")), RouteOutcome::Session);
|
||||
assert_eq!(demux.route(frame("b")), RouteOutcome::InboxFull);
|
||||
|
||||
let text = out_rx.try_recv().expect("overloaded response enqueued");
|
||||
let wire: Value = serde_json::from_str(&text).expect("valid json");
|
||||
assert_eq!(wire["id"], "b");
|
||||
assert_eq!(wire["session_id"], "busy");
|
||||
assert_eq!(wire["error"]["code"], -32016);
|
||||
assert_eq!(wire["error"]["data"]["code"], "tool_busy");
|
||||
assert_eq!(wire["error"]["data"]["retryable"], true);
|
||||
assert!(
|
||||
out_rx.try_recv().is_err(),
|
||||
"exactly one rejection emitted for one overflow"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inbox_full_request_with_malformed_id_still_emits_overloaded_response() {
|
||||
// A Request whose `id` is present but not a valid JsonRpcId
|
||||
// (object/array/null) must NOT be silently dropped on a full
|
||||
// inbox: it still gets the shared -32016 response, with the raw
|
||||
// id echoed back as a string.
|
||||
let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
|
||||
let demux = Demux::with_outbound(out_tx);
|
||||
let session = SessionId::new("bad_id").expect("valid");
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = |id: Value| {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"session_id": "bad_id",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
})
|
||||
};
|
||||
// First fills capacity (cap 1); the malformed-id second overflows.
|
||||
assert_eq!(demux.route(frame(json!("a"))), RouteOutcome::Session);
|
||||
assert_eq!(
|
||||
demux.route(frame(json!({ "nested": 1 }))),
|
||||
RouteOutcome::InboxFull
|
||||
);
|
||||
|
||||
let text = out_rx.try_recv().expect("overloaded response enqueued");
|
||||
let wire: Value = serde_json::from_str(&text).expect("valid json");
|
||||
assert_eq!(
|
||||
wire["id"], "{\"nested\":1}",
|
||||
"malformed id is echoed back as its raw JSON text"
|
||||
);
|
||||
assert_eq!(wire["error"]["code"], -32016);
|
||||
assert_eq!(wire["error"]["data"]["code"], "tool_busy");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inbox_full_notification_is_dropped_without_outbound_response() {
|
||||
// A Notification (no id) on a full inbox stays fire-and-forget:
|
||||
// no synthesized response is emitted.
|
||||
let (out_tx, mut out_rx) = mpsc::channel::<String>(4);
|
||||
let demux = Demux::with_outbound(out_tx);
|
||||
let session = SessionId::new("notif_busy").expect("valid");
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let notif = || {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "notif_busy",
|
||||
"method": "tool.notification",
|
||||
"params": {},
|
||||
})
|
||||
};
|
||||
// First notification fills capacity; second overflows.
|
||||
assert_eq!(demux.route(notif()), RouteOutcome::Session);
|
||||
assert_eq!(demux.route(notif()), RouteOutcome::InboxFull);
|
||||
assert!(
|
||||
out_rx.try_recv().is_err(),
|
||||
"notifications must not synthesize an outbound response"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn inbox_full_request_without_outbound_does_not_panic() {
|
||||
// A bare demux (no outbound, e.g. unit context) must still report
|
||||
// InboxFull cleanly when it cannot synthesize a rejection.
|
||||
let demux = Demux::new();
|
||||
let session = SessionId::new("no_out").expect("valid");
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
demux.register_session_inbox(session.clone(), tx);
|
||||
let frame = || {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": "x",
|
||||
"session_id": "no_out",
|
||||
"method": "tool_call_request",
|
||||
"params": {},
|
||||
})
|
||||
};
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::Session);
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::InboxFull);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn progress_route_pushes_to_progress_waiter() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let (tx, mut rx) = mpsc::channel(4);
|
||||
demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx)
|
||||
.expect("first registration");
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "any",
|
||||
"method": "tool_call_progress",
|
||||
"params": {
|
||||
"tool_call_id": call_id.as_str(),
|
||||
"kind": "log_chunk",
|
||||
"body": {"text": "hello"},
|
||||
},
|
||||
});
|
||||
let outcome = demux.route(frame);
|
||||
assert_eq!(outcome, RouteOutcome::Progress);
|
||||
let progress = rx.recv().await.expect("progress frame");
|
||||
assert_eq!(progress.tool_call_id, call_id);
|
||||
assert_eq!(progress.kind, "log_chunk");
|
||||
assert_eq!(progress.body, json!({"text": "hello"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn progress_with_no_waiter_returns_unknown_progress() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "any",
|
||||
"method": "tool_call_progress",
|
||||
"params": {
|
||||
"tool_call_id": call_id.as_str(),
|
||||
"kind": "log_chunk",
|
||||
"body": {},
|
||||
},
|
||||
});
|
||||
assert_eq!(demux.route(frame), RouteOutcome::UnknownProgress);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropped_progress_receiver_returns_progress_dropped() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let (tx, rx) = mpsc::channel(1);
|
||||
demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx)
|
||||
.expect("first registration");
|
||||
drop(rx);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "any",
|
||||
"method": "tool_call_progress",
|
||||
"params": {
|
||||
"tool_call_id": call_id.as_str(),
|
||||
"kind": "x",
|
||||
"body": {},
|
||||
},
|
||||
});
|
||||
assert_eq!(demux.route(frame), RouteOutcome::ProgressDropped);
|
||||
assert!(demux.progress.get(&call_id).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unregister_progress_waiter_returns_sender_when_present() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let (tx, _rx) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx)
|
||||
.expect("first registration");
|
||||
assert!(demux.unregister_progress_waiter(&call_id).is_some());
|
||||
assert!(demux.unregister_progress_waiter(&call_id).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn try_register_progress_waiter_rejects_collision_and_preserves_existing() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let (tx_first, mut rx_first) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx_first)
|
||||
.expect("first registration");
|
||||
let (tx_second, _rx_second) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
let returned = demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx_second)
|
||||
.expect_err("collision returns the rejected sender");
|
||||
// Returned sender is independent of the live one: dropping
|
||||
// it must not close the original receiver.
|
||||
drop(returned);
|
||||
let frame = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "any",
|
||||
"method": "tool_call_progress",
|
||||
"params": {
|
||||
"tool_call_id": call_id.as_str(),
|
||||
"kind": "log_chunk",
|
||||
"body": {"text": "first"},
|
||||
},
|
||||
});
|
||||
assert_eq!(demux.route(frame), RouteOutcome::Progress);
|
||||
let progress = rx_first.recv().await.expect("original receiver still live");
|
||||
assert_eq!(progress.body, json!({"text": "first"}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_progress_channel_returns_progress_full_without_blocking() {
|
||||
let demux = Demux::new();
|
||||
let call_id = ToolCallId::new_v7();
|
||||
let (tx, _rx) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
demux
|
||||
.try_register_progress_waiter(call_id.clone(), tx)
|
||||
.expect("first registration");
|
||||
let frame = || {
|
||||
json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "any",
|
||||
"method": "tool_call_progress",
|
||||
"params": {
|
||||
"tool_call_id": call_id.as_str(),
|
||||
"kind": "x",
|
||||
"body": {},
|
||||
},
|
||||
})
|
||||
};
|
||||
// First send fills capacity (mpsc(1)).
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::Progress);
|
||||
// Second send must NOT block; it returns ProgressFull.
|
||||
assert_eq!(demux.route(frame()), RouteOutcome::ProgressFull);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_progress_removes_all_waiters_and_drops_senders() {
|
||||
let demux = Demux::new();
|
||||
let call_a = ToolCallId::new_v7();
|
||||
let call_b = ToolCallId::new_v7();
|
||||
let (tx_a, mut rx_a) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
let (tx_b, mut rx_b) = mpsc::channel::<ToolCallProgressFrame>(1);
|
||||
demux
|
||||
.try_register_progress_waiter(call_a.clone(), tx_a)
|
||||
.expect("first registration");
|
||||
demux
|
||||
.try_register_progress_waiter(call_b.clone(), tx_b)
|
||||
.expect("first registration");
|
||||
assert_eq!(demux.progress.len(), 2);
|
||||
|
||||
demux.drain_progress();
|
||||
|
||||
// Post-drain: every entry removed.
|
||||
assert_eq!(demux.progress.len(), 0);
|
||||
assert!(demux.progress.get(&call_a).is_none());
|
||||
assert!(demux.progress.get(&call_b).is_none());
|
||||
// The senders held by the demux were dropped, so each
|
||||
// receiver sees `None` (channel closed).
|
||||
assert!(
|
||||
rx_a.recv().await.is_none(),
|
||||
"sender dropped → receiver closes"
|
||||
);
|
||||
assert!(
|
||||
rx_b.recv().await.is_none(),
|
||||
"sender dropped → receiver closes"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
//! Shared donation transport: a bounded retry buffer + in-order drain
|
||||
//! barrier, parameterized over a `donate` closure. Traces, logs, and
|
||||
//! metrics all pump through this; failed sends are retained briefly,
|
||||
//! overflow drops payloads — telemetry, never correctness.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value};
|
||||
use opentelemetry_proto::tonic::resource::v1::Resource;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
/// Bound on payloads queued before the pump drains them.
|
||||
pub(crate) const PENDING_FLUSHES: usize = 8;
|
||||
/// Payloads retained across failed sends (disconnect/reconnect window).
|
||||
pub(crate) const RETRY_CAP: usize = 8;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared OTLP encoding helpers
|
||||
//
|
||||
// Reused by the log and metric donation clients so the AnyValue/KeyValue/
|
||||
// Resource construction lives in one place instead of being copy-pasted per
|
||||
// client. (`trace_donate` builds its payload via `opentelemetry_sdk`'s own
|
||||
// conversion and does not use these.)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Current wall-clock time as Unix-epoch nanoseconds (OTLP `time_unix_nano`).
|
||||
pub(crate) fn now_unix_nanos() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// OTLP string `AnyValue`.
|
||||
pub(crate) fn string_value(s: String) -> AnyValue {
|
||||
AnyValue {
|
||||
value: Some(any_value::Value::StringValue(s)),
|
||||
}
|
||||
}
|
||||
|
||||
/// OTLP string-valued `KeyValue`.
|
||||
pub(crate) fn string_kv(key: &str, value: String) -> KeyValue {
|
||||
KeyValue {
|
||||
key: key.to_owned(),
|
||||
value: Some(string_value(value)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// OTLP `Resource` carrying just `service.name`.
|
||||
pub(crate) fn make_resource(service_name: String) -> Resource {
|
||||
Resource {
|
||||
attributes: vec![string_kv("service.name", service_name)],
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) enum PumpMsg {
|
||||
/// Base64 OTLP request, ready for the wire.
|
||||
Payload(String),
|
||||
/// In-order drain fence — a barrier, not a timeout.
|
||||
Barrier(oneshot::Sender<()>),
|
||||
}
|
||||
|
||||
/// Resolves once every payload queued before this call has had a send
|
||||
/// attempt. Call after the producer's flush (e.g. `fastrace::flush()`).
|
||||
pub(crate) async fn drain_via(tx: &mpsc::Sender<PumpMsg>) {
|
||||
let (ack_tx, ack_rx) = oneshot::channel();
|
||||
if tx.send(PumpMsg::Barrier(ack_tx)).await.is_ok() {
|
||||
let _ = ack_rx.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// `donate` hands the payload back so a failed send retains it
|
||||
/// without cloning.
|
||||
pub(crate) async fn run_pump<D, F>(mut rx: mpsc::Receiver<PumpMsg>, donate: D)
|
||||
where
|
||||
D: Fn(String) -> F,
|
||||
F: std::future::Future<Output = (bool, String)>,
|
||||
{
|
||||
let mut retry: VecDeque<String> = VecDeque::new();
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
PumpMsg::Payload(payload) => {
|
||||
if retry.len() == RETRY_CAP {
|
||||
retry.pop_front();
|
||||
tracing::debug!("donation retry buffer full; dropping oldest payload");
|
||||
}
|
||||
retry.push_back(payload);
|
||||
}
|
||||
PumpMsg::Barrier(ack) => {
|
||||
attempt_sends(&mut retry, &donate).await;
|
||||
let _ = ack.send(());
|
||||
continue;
|
||||
}
|
||||
}
|
||||
attempt_sends(&mut retry, &donate).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Send in order, stopping at the first failure; the remainder stays
|
||||
/// queued for the next wake.
|
||||
async fn attempt_sends<D, F>(retry: &mut VecDeque<String>, donate: &D)
|
||||
where
|
||||
D: Fn(String) -> F,
|
||||
F: std::future::Future<Output = (bool, String)>,
|
||||
{
|
||||
while let Some(payload) = retry.pop_front() {
|
||||
let (ok, payload) = donate(payload).await;
|
||||
if !ok {
|
||||
tracing::debug!("donation send failed; retaining payload for retry");
|
||||
retry.push_front(payload);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn payload(tag: u64) -> PumpMsg {
|
||||
PumpMsg::Payload(format!("payload-{tag}"))
|
||||
}
|
||||
|
||||
/// The drain barrier acks even while the link is down.
|
||||
#[tokio::test]
|
||||
async fn pump_retries_failed_payloads_across_reconnect() {
|
||||
let healthy = Arc::new(AtomicBool::new(false));
|
||||
let sent: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
|
||||
let pump = {
|
||||
let healthy = Arc::clone(&healthy);
|
||||
let sent = Arc::clone(&sent);
|
||||
tokio::spawn(run_pump(rx, move |p: String| {
|
||||
let healthy = Arc::clone(&healthy);
|
||||
let sent = Arc::clone(&sent);
|
||||
async move {
|
||||
if healthy.load(Ordering::SeqCst) {
|
||||
sent.lock().push(p.clone());
|
||||
(true, p)
|
||||
} else {
|
||||
(false, p)
|
||||
}
|
||||
}
|
||||
}))
|
||||
};
|
||||
|
||||
tx.send(payload(1)).await.unwrap();
|
||||
tx.send(payload(2)).await.unwrap();
|
||||
drain_via(&tx).await;
|
||||
assert!(sent.lock().is_empty(), "nothing sent while link is down");
|
||||
|
||||
healthy.store(true, Ordering::SeqCst);
|
||||
drain_via(&tx).await;
|
||||
assert_eq!(*sent.lock(), vec!["payload-1", "payload-2"]);
|
||||
|
||||
drop(tx);
|
||||
pump.await.expect("pump must exit cleanly");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pump_retry_buffer_drops_oldest_beyond_cap() {
|
||||
let sent: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
|
||||
let healthy = Arc::new(AtomicBool::new(false));
|
||||
let (tx, rx) = mpsc::channel::<PumpMsg>(RETRY_CAP + 2);
|
||||
let pump = {
|
||||
let healthy = Arc::clone(&healthy);
|
||||
let sent = Arc::clone(&sent);
|
||||
tokio::spawn(run_pump(rx, move |p: String| {
|
||||
let healthy = Arc::clone(&healthy);
|
||||
let sent = Arc::clone(&sent);
|
||||
async move {
|
||||
if healthy.load(Ordering::SeqCst) {
|
||||
sent.lock().push(p.clone());
|
||||
(true, p)
|
||||
} else {
|
||||
(false, p)
|
||||
}
|
||||
}
|
||||
}))
|
||||
};
|
||||
|
||||
for i in 0..=(RETRY_CAP as u64) {
|
||||
tx.send(payload(i + 1)).await.unwrap();
|
||||
}
|
||||
drain_via(&tx).await;
|
||||
|
||||
healthy.store(true, Ordering::SeqCst);
|
||||
drain_via(&tx).await;
|
||||
{
|
||||
let sent = sent.lock();
|
||||
assert_eq!(sent.len(), RETRY_CAP, "buffer bounded at RETRY_CAP");
|
||||
assert_eq!(sent[0], "payload-2", "oldest payload evicted first");
|
||||
}
|
||||
|
||||
drop(tx);
|
||||
pump.await.expect("pump must exit cleanly");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
//! Client-side error taxonomy.
|
||||
//!
|
||||
//! Wire-level [`kigi_tool_protocol::ToolErrorWire`] variants and JSON-RPC
|
||||
//! error envelopes are mapped into the smaller [`ClientError`] vocabulary
|
||||
//! at the SDK boundary so consumers can match on a single enum without
|
||||
//! re-deriving the numeric/string code mapping.
|
||||
|
||||
use kigi_tool_protocol::{IdError, JsonRpcError, ToolCallId, ToolErrorWire};
|
||||
use thiserror::Error;
|
||||
use url::Url;
|
||||
|
||||
/// Errors surfaced by the client SDK.
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ClientError {
|
||||
/// WebSocket transport failure: failed to connect, dropped socket,
|
||||
/// or in-flight request interrupted by a reconnect cycle.
|
||||
#[error("network error: {0}")]
|
||||
NetworkError(String),
|
||||
|
||||
/// Wire-protocol violation: malformed JSON, unexpected method,
|
||||
/// hello/hello_ack mismatch, or unsupported `protocol_version`.
|
||||
#[error("protocol error: {0}")]
|
||||
ProtocolError(String),
|
||||
|
||||
/// Authentication or authorisation rejected by the server.
|
||||
#[error("auth error: {0}")]
|
||||
AuthError(String),
|
||||
|
||||
/// Server rejected the WebSocket upgrade with an HTTP auth status
|
||||
/// (401/403). Non-retryable: replaying the same credential is
|
||||
/// rejected identically, so the reconnect loop classifies this as
|
||||
/// fatal instead of retrying forever.
|
||||
#[error("handshake auth failed: HTTP {status}")]
|
||||
HandshakeAuthFailed { status: u16 },
|
||||
|
||||
/// `register_tool` / `register_session` ack reported a conflict
|
||||
/// (cross-connection contention or an already-bound entry the
|
||||
/// caller did not expect).
|
||||
#[error("registration conflict: {0}")]
|
||||
RegistrationConflict(String),
|
||||
|
||||
/// Outbound mpsc full or call-site bounded wait elapsed before the
|
||||
/// frame could be enqueued. Distinct from [`Self::NetworkError`]:
|
||||
/// the socket may still be healthy.
|
||||
#[error("backpressure: {0}")]
|
||||
BackpressureError(String),
|
||||
|
||||
/// JSON serialise / deserialise failure inside the SDK.
|
||||
#[error("serde error: {0}")]
|
||||
Serde(String),
|
||||
|
||||
/// Builder consistency error: missing URL, missing auth, etc.
|
||||
#[error("invalid configuration: {0}")]
|
||||
InvalidConfig(String),
|
||||
|
||||
/// Wrapped wire-format tool error; surfaces the upstream
|
||||
/// [`ToolErrorWire`] variant verbatim for callers that need to
|
||||
/// switch on the stable string code.
|
||||
#[error(transparent)]
|
||||
Wire(ToolErrorWire),
|
||||
|
||||
/// Server-side close / shutdown signal received during steady state.
|
||||
#[error("server closed connection: {0}")]
|
||||
Closed(String),
|
||||
|
||||
/// Refused to send credentials over an insecure `ws://` scheme to a
|
||||
/// non-loopback host. Local-loopback (`127.0.0.1`, `::1`,
|
||||
/// `localhost`) is the only exception; every other host MUST be
|
||||
/// reached over `wss://` so the bearer token never crosses the
|
||||
/// network in plaintext.
|
||||
#[error(
|
||||
"insecure scheme: refusing to send credentials over plaintext ws:// to non-loopback host {url}"
|
||||
)]
|
||||
InsecureScheme { url: Url },
|
||||
|
||||
/// Caller passed a `ToolCallId` that already keys an in-flight
|
||||
/// dispatch on the same connection. The prior call's progress
|
||||
/// waiter and response correlation are left intact; this error
|
||||
/// surfaces synchronously so the second caller can retry with a
|
||||
/// fresh id. Mint a fresh [`ToolCallId::new_v7`] (or use
|
||||
/// [`kigi_tool_runtime::ToolCallContext::default`], which does so)
|
||||
/// per call. This is client misuse, not a transport or server
|
||||
/// failure.
|
||||
#[error("call_id {call_id} already in flight on this connection")]
|
||||
CallIdInUse { call_id: ToolCallId },
|
||||
}
|
||||
|
||||
impl ClientError {
|
||||
/// Map a JSON-RPC envelope error into a [`ClientError`]. The
|
||||
/// envelope's `data` payload (when present) carries the stable
|
||||
/// [`ToolErrorWire`] discriminator; the numeric `code` is used as a
|
||||
/// coarse fallback when `data` is absent or undecodable.
|
||||
pub fn from_jsonrpc_error(err: JsonRpcError) -> Self {
|
||||
if let Some(data) = err.data
|
||||
&& let Ok(wire) = serde_json::from_value::<ToolErrorWire>(data)
|
||||
{
|
||||
return Self::from_wire(wire);
|
||||
}
|
||||
match err.code {
|
||||
-32002 | -32003 => Self::AuthError(err.message),
|
||||
-32004 => Self::NetworkError(err.message),
|
||||
-32600..=-32500 => Self::ProtocolError(err.message),
|
||||
_ => Self::Wire(ToolErrorWire::Custom {
|
||||
subcode: format!("jsonrpc_{}", err.code),
|
||||
message: err.message,
|
||||
details: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` when a `data`-less envelope collapsed to the given `jsonrpc_<code>`
|
||||
/// subcode (see [`Self::from_jsonrpc_error`]); shared by the bind recognizers.
|
||||
fn has_collapsed_jsonrpc_subcode(&self, subcode: &str) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Wire(ToolErrorWire::Custom { subcode: s, .. }) if s == subcode
|
||||
)
|
||||
}
|
||||
|
||||
/// `true` for the server's "server not found" bind rejection (JSON-RPC `-32601`):
|
||||
/// no workspace-server is registered for this user.
|
||||
pub fn is_server_not_found(&self) -> bool {
|
||||
self.has_collapsed_jsonrpc_subcode("jsonrpc_-32601")
|
||||
}
|
||||
|
||||
/// `true` for the server's `-32013` "server found but bind did not complete" error
|
||||
/// (the `ServerBindOutcome::Unavailable` cases). Recognized so the harness
|
||||
/// re-provisions this recoverable case, distinct from [`Self::is_server_not_found`].
|
||||
pub fn is_tool_unavailable(&self) -> bool {
|
||||
self.has_collapsed_jsonrpc_subcode("jsonrpc_-32013")
|
||||
}
|
||||
|
||||
/// Map a [`ToolErrorWire`] variant into the SDK error taxonomy.
|
||||
pub fn from_wire(wire: ToolErrorWire) -> Self {
|
||||
match wire {
|
||||
ToolErrorWire::PermissionDenied { reason } => Self::AuthError(reason),
|
||||
ToolErrorWire::TransportClosed { tool_id } => {
|
||||
Self::NetworkError(format!("transport closed for {tool_id}"))
|
||||
}
|
||||
ToolErrorWire::UnsupportedProtocolVersion { supported } => {
|
||||
Self::ProtocolError(format!("unsupported protocol; supported: {supported:?}"))
|
||||
}
|
||||
other => Self::Wire(other),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for ClientError {
|
||||
fn from(err: serde_json::Error) -> Self {
|
||||
Self::Serde(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<IdError> for ClientError {
|
||||
fn from(err: IdError) -> Self {
|
||||
Self::ProtocolError(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<url::ParseError> for ClientError {
|
||||
fn from(err: url::ParseError) -> Self {
|
||||
Self::InvalidConfig(format!("invalid url: {err}"))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<tokio_tungstenite::tungstenite::Error> for ClientError {
|
||||
fn from(err: tokio_tungstenite::tungstenite::Error) -> Self {
|
||||
Self::NetworkError(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientError {
|
||||
/// Classify a failed WebSocket upgrade. A `401`/`403` on the HTTP
|
||||
/// upgrade is a non-retryable auth rejection
|
||||
/// ([`Self::HandshakeAuthFailed`]); every other failure stays a
|
||||
/// transport [`Self::NetworkError`] via the blanket `From` impl. The
|
||||
/// distinction must be made here, before `From` collapses the typed
|
||||
/// `Http` response status into an opaque string.
|
||||
pub(crate) fn from_handshake_error(err: tokio_tungstenite::tungstenite::Error) -> Self {
|
||||
if let tokio_tungstenite::tungstenite::Error::Http(resp) = &err {
|
||||
let status = resp.status().as_u16();
|
||||
if status == 401 || status == 403 {
|
||||
return Self::HandshakeAuthFailed { status };
|
||||
}
|
||||
}
|
||||
Self::from(err)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<tokio::sync::oneshot::error::RecvError> for ClientError {
|
||||
fn from(_: tokio::sync::oneshot::error::RecvError) -> Self {
|
||||
Self::NetworkError("response waiter dropped (connection closed)".to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use kigi_tool_protocol::{
|
||||
WORKSPACE_UNAVAILABLE_SUBCODE, WorkspaceGonePhase, WorkspaceGoneReason,
|
||||
workspace_unavailable_wire,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn workspace_gone_envelope() -> JsonRpcError {
|
||||
let wire = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
);
|
||||
JsonRpcError {
|
||||
code: -32005,
|
||||
message: "workspace server gone".to_owned(),
|
||||
data: Some(serde_json::to_value(&wire).unwrap()),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_upgrade_error(status: u16) -> tokio_tungstenite::tungstenite::Error {
|
||||
let resp = tokio_tungstenite::tungstenite::http::Response::builder()
|
||||
.status(status)
|
||||
.body(None::<Vec<u8>>)
|
||||
.expect("response builds");
|
||||
tokio_tungstenite::tungstenite::Error::Http(resp)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handshake_401_and_403_map_to_handshake_auth_failed() {
|
||||
for status in [401u16, 403] {
|
||||
match ClientError::from_handshake_error(http_upgrade_error(status)) {
|
||||
ClientError::HandshakeAuthFailed { status: got } => assert_eq!(got, status),
|
||||
other => panic!("expected HandshakeAuthFailed for {status}; got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handshake_non_auth_status_stays_network_error() {
|
||||
for status in [500u16, 502, 429] {
|
||||
match ClientError::from_handshake_error(http_upgrade_error(status)) {
|
||||
ClientError::NetworkError(_) => {}
|
||||
other => panic!("expected NetworkError for {status}; got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_jsonrpc_error_preserves_workspace_subcode_and_details() {
|
||||
// The `data` payload decodes as `ToolErrorWire` first, so the stable
|
||||
// subcode and structured details reach the SDK consumer intact rather
|
||||
// than collapsing to the numeric code.
|
||||
match ClientError::from_jsonrpc_error(workspace_gone_envelope()) {
|
||||
ClientError::Wire(ToolErrorWire::Custom {
|
||||
subcode, details, ..
|
||||
}) => {
|
||||
assert_eq!(subcode, WORKSPACE_UNAVAILABLE_SUBCODE);
|
||||
let details = details.expect("details present");
|
||||
assert_eq!(details["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
|
||||
assert_eq!(details["reason"], json!("disconnect"));
|
||||
assert_eq!(details["phase"], json!("route_missing"));
|
||||
assert_eq!(details["retryable"], json!(true));
|
||||
}
|
||||
other => panic!("expected Wire(Custom), got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_server_not_found_recognizes_bare_minus_32601() {
|
||||
// data-less -32601 -> custom subcode.
|
||||
let err = ClientError::from_jsonrpc_error(JsonRpcError {
|
||||
code: -32601,
|
||||
message: "server abc not found for user".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(err.is_server_not_found());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_tool_unavailable_recognizes_bare_minus_32013() {
|
||||
let err = ClientError::from_jsonrpc_error(JsonRpcError {
|
||||
code: -32013,
|
||||
message: "server abc did not complete the bind".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(err.is_tool_unavailable());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_server_not_found_rejects_other_errors() {
|
||||
let auth = ClientError::from_jsonrpc_error(JsonRpcError {
|
||||
code: -32002,
|
||||
message: "nope".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(!auth.is_server_not_found());
|
||||
// workspace-gone is the tool-call re-provision path, not bind ServerNotFound.
|
||||
assert!(!ClientError::from_jsonrpc_error(workspace_gone_envelope()).is_server_not_found());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bind_recognizers_are_mutually_exclusive() {
|
||||
let not_found = ClientError::from_jsonrpc_error(JsonRpcError {
|
||||
code: -32601,
|
||||
message: "not found".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
let unavailable = ClientError::from_jsonrpc_error(JsonRpcError {
|
||||
code: -32013,
|
||||
message: "unavailable".to_owned(),
|
||||
data: None,
|
||||
});
|
||||
assert!(not_found.is_server_not_found());
|
||||
assert!(
|
||||
!not_found.is_tool_unavailable(),
|
||||
"-32601 must not be recognized as tool_unavailable"
|
||||
);
|
||||
assert!(unavailable.is_tool_unavailable());
|
||||
assert!(
|
||||
!unavailable.is_server_not_found(),
|
||||
"-32013 must not be recognized as server_not_found"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sdk_reexported_recognizer_matches_decoded_error() {
|
||||
// SDK-only consumers reach the recognizer through the SDK re-export and
|
||||
// the core decode path.
|
||||
let err = kigi_computer_hub_core::error_from_envelope(workspace_gone_envelope());
|
||||
assert!(crate::is_workspace_unavailable(&err));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sdk_reexported_recognizer_rejects_unrelated_custom_error() {
|
||||
let wire = ToolErrorWire::Custom {
|
||||
subcode: "unrelated".to_owned(),
|
||||
message: "nope".to_owned(),
|
||||
details: Some(json!({ "code": "unrelated" })),
|
||||
};
|
||||
let env = JsonRpcError {
|
||||
code: -32000,
|
||||
message: "nope".to_owned(),
|
||||
data: Some(serde_json::to_value(&wire).unwrap()),
|
||||
};
|
||||
let err = kigi_computer_hub_core::error_from_envelope(env);
|
||||
assert!(!crate::is_workspace_unavailable(&err));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
//! Hello handshake helpers used by the connection actor and the
|
||||
//! reconnect-replay path.
|
||||
//!
|
||||
//! Splitting these into a dedicated module keeps the connection state
|
||||
//! machine readable: send the frame, parse the ack, surface a typed
|
||||
//! [`crate::ClientError`].
|
||||
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use kigi_tool_protocol::{ConnectionKind, HelloAckMsg, HelloMsg};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
use crate::error::ClientError;
|
||||
|
||||
/// Wire-protocol version both ends speak. Re-exported from the
|
||||
/// protocol crate so the SDK and the IC service share one source of
|
||||
/// truth.
|
||||
pub use kigi_tool_protocol::PROTOCOL_VERSION;
|
||||
|
||||
/// Send the [`HelloMsg`] and wait for the matching [`HelloAckMsg`].
|
||||
///
|
||||
/// `kind` should be [`ConnectionKind::ToolServer`] for tool-server
|
||||
/// builds (the only consumer today). The function returns the parsed
|
||||
/// ack so callers can observe the server-issued `connection_id` and
|
||||
/// the server-derived `user_id`.
|
||||
///
|
||||
/// When `server_id` is `Some`, it is included in the hello frame so the
|
||||
/// server can identify itself without a separate `register_server` call.
|
||||
pub async fn send_hello<Si, St>(
|
||||
sink: &mut Si,
|
||||
stream: &mut St,
|
||||
kind: ConnectionKind,
|
||||
server_id: Option<kigi_tool_protocol::ServerId>,
|
||||
description: Option<String>,
|
||||
metadata: Option<serde_json::Value>,
|
||||
) -> Result<HelloAckMsg, ClientError>
|
||||
where
|
||||
Si: SinkExt<Message> + Unpin,
|
||||
Si::Error: std::fmt::Display,
|
||||
St: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>> + Unpin,
|
||||
{
|
||||
let hello = HelloMsg {
|
||||
protocol_version: PROTOCOL_VERSION.to_owned(),
|
||||
kind,
|
||||
server_id,
|
||||
description,
|
||||
metadata,
|
||||
};
|
||||
let text = serde_json::to_string(&hello)?;
|
||||
sink.send(Message::Text(text.into()))
|
||||
.await
|
||||
.map_err(|e| ClientError::NetworkError(format!("hello send failed: {e}")))?;
|
||||
|
||||
while let Some(msg) = stream.next().await {
|
||||
let msg = msg?;
|
||||
match msg {
|
||||
Message::Text(text) => {
|
||||
let ack: HelloAckMsg = serde_json::from_str(text.as_ref())
|
||||
.map_err(|e| ClientError::ProtocolError(format!("malformed hello_ack: {e}")))?;
|
||||
if !ack
|
||||
.supported_protocol_versions
|
||||
.iter()
|
||||
.any(|v| v == PROTOCOL_VERSION)
|
||||
{
|
||||
return Err(ClientError::ProtocolError(format!(
|
||||
"server does not support {PROTOCOL_VERSION}; supported: {:?}",
|
||||
ack.supported_protocol_versions
|
||||
)));
|
||||
}
|
||||
return Ok(ack);
|
||||
}
|
||||
Message::Ping(payload) => {
|
||||
sink.send(Message::Pong(payload))
|
||||
.await
|
||||
.map_err(|e| ClientError::NetworkError(format!("pong send failed: {e}")))?;
|
||||
}
|
||||
Message::Close(frame) => {
|
||||
let reason = frame.map(|f| f.reason.to_string()).unwrap_or_default();
|
||||
return Err(ClientError::Closed(format!(
|
||||
"server closed during handshake: {reason}"
|
||||
)));
|
||||
}
|
||||
Message::Pong(_) | Message::Frame(_) => continue,
|
||||
Message::Binary(_) => {
|
||||
return Err(ClientError::ProtocolError(
|
||||
"server sent binary frame during handshake".to_owned(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(ClientError::NetworkError(
|
||||
"server closed before hello_ack".to_owned(),
|
||||
))
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
//! Tool-server and harness SDK.
|
||||
//!
|
||||
//! Single crate hosting both the tool-server runtime and the
|
||||
//! harness-side dispatch surface. The shared substrate —
|
||||
//! [`HubConnectionPool`], [`HubConnection`], the inbound demux, the
|
||||
//! refcount-managed bound-session set, and the transparent reconnect /
|
||||
//! replay state machine — lives here so both ends speak through one
|
||||
//! frame multiplex on top of one WebSocket per `(url, principal)`.
|
||||
//!
|
||||
//! The server entry point is [`ToolServer`]: build it via
|
||||
//! [`ToolServerBuilder`], wire one or more [`ToolServerHandler`]
|
||||
//! implementations, and call [`ToolServer::run`] to drive the inbound
|
||||
//! loop. The harness entry point is [`ToolHarness`]: build it via
|
||||
//! [`ToolHarnessBuilder`], optionally seed it with in-process
|
||||
//! [`kigi_tool_runtime::Tool`] implementations, and call
|
||||
//! [`ToolHarness::call`] to dispatch a tool call. Authorisation
|
||||
//! credentials (`AuthCredential`) plus the target URL determine
|
||||
//! which pool entry the consumer attaches to; multiple
|
||||
//! [`ToolServer`] / [`ToolHarness`] instances against the same
|
||||
//! `(url, principal)` share a single connection and refcount their
|
||||
//! session bindings.
|
||||
|
||||
#![forbid(unsafe_code)]
|
||||
|
||||
pub(crate) mod admission;
|
||||
pub mod auth;
|
||||
pub(crate) mod cancel;
|
||||
pub mod connection;
|
||||
pub(crate) mod connection_borrow;
|
||||
pub mod demux;
|
||||
pub(crate) mod donate_pump;
|
||||
pub mod error;
|
||||
pub mod handshake;
|
||||
pub mod harness;
|
||||
pub mod log_donate;
|
||||
#[cfg(feature = "metrics")]
|
||||
pub mod metric_donate;
|
||||
pub mod metrics;
|
||||
pub mod notification;
|
||||
pub mod observability;
|
||||
pub mod pool;
|
||||
pub mod refcount;
|
||||
pub mod server;
|
||||
pub mod trace_donate;
|
||||
|
||||
pub mod oidc_provider;
|
||||
|
||||
pub use auth::{AuthCredential, AuthIdentity, AuthProvider, PrincipalKey, SharedAuthProvider};
|
||||
pub use connection::{ConnKey, HubConnection, ReconnectEvent};
|
||||
pub use error::ClientError;
|
||||
pub use harness::{
|
||||
CancelOnDrop, LocalRegistry, ModelOutputExtractor, SessionBindReport, ToolHarness,
|
||||
ToolHarnessBuilder, extractor_for,
|
||||
};
|
||||
pub use log_donate::{DonatingLogLayer, LogDonationPump, LogDonationSender, flush_log_layer};
|
||||
#[cfg(feature = "metrics")]
|
||||
pub use metric_donate::MetricDonationPump;
|
||||
pub use notification::HubNotification;
|
||||
pub use observability::ObservabilityBridge;
|
||||
pub use oidc_provider::{
|
||||
OidcAuthProvider, OidcAuthProviderBuilder, OnRefreshCallback, RefreshEvent,
|
||||
};
|
||||
pub use pool::HubConnectionPool;
|
||||
pub use server::{
|
||||
ResolvedSessionHandlers, SessionHandlerResolver, SystemNotifyAck, ToolServer,
|
||||
ToolServerBuilder, ToolServerHandler, WeakToolServer,
|
||||
};
|
||||
pub use trace_donate::{HubDonatingReporter, TraceDonationPump};
|
||||
// Re-exported so consumers that depend only on the SDK can recognize the
|
||||
// server's `workspace_unavailable` error without also pulling in the core crate.
|
||||
pub use kigi_computer_hub_core::is_workspace_unavailable;
|
||||
@@ -0,0 +1,592 @@
|
||||
//! Forward curated `tracing` events to the connected server over the
|
||||
//! WebSocket transport (`logs.donate`).
|
||||
//!
|
||||
//! [`DonatingLogLayer`] is installed **inert** at startup and activated
|
||||
//! post-connect by swapping in a [`LogDonationSender`] (a global-subscriber
|
||||
//! constraint); while inert, selected events are dropped before enqueueing.
|
||||
//!
|
||||
//! Only events on the [`TELEMETRY_TARGET`] target at `>= INFO` are
|
||||
//! forwarded, and only fields in [`ALLOWED_FIELDS`] are included; other
|
||||
//! fields such as `reason`/`error` are omitted.
|
||||
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arc_swap::ArcSwapOption;
|
||||
use base64::Engine as _;
|
||||
use fastrace::collector::SpanContext;
|
||||
use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_LOG_RECORDS_PER_DONATION};
|
||||
use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
|
||||
use opentelemetry_proto::tonic::common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value};
|
||||
use opentelemetry_proto::tonic::logs::v1::{LogRecord, ResourceLogs, ScopeLogs};
|
||||
use opentelemetry_proto::tonic::resource::v1::Resource;
|
||||
use prost::Message as _;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::Level;
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing_subscriber::Layer;
|
||||
use tracing_subscriber::layer::Context;
|
||||
|
||||
use crate::donate_pump::{
|
||||
PENDING_FLUSHES, PumpMsg, drain_via, make_resource, now_unix_nanos, run_pump, string_kv,
|
||||
string_value,
|
||||
};
|
||||
use crate::server::ToolServer;
|
||||
|
||||
/// Stable target the workspace routes selected events through.
|
||||
/// The layer selects exactly this target, ignoring global
|
||||
/// `RUST_LOG`. The server re-stamps it as the OTLP scope name.
|
||||
pub const TELEMETRY_TARGET: &str = "workspace::telemetry";
|
||||
|
||||
/// Set of forwardable field names — guaranteed-literal or numeric.
|
||||
/// Only the listed fields are included; other fields such as
|
||||
/// `reason`/`error`/`object_path`/`gcs_path` are omitted.
|
||||
const ALLOWED_FIELDS: &[&str] = &[
|
||||
"session_id",
|
||||
"turn_number",
|
||||
"phase",
|
||||
"bytes",
|
||||
"file_count",
|
||||
"pending",
|
||||
"pending_bytes",
|
||||
"sample_period_secs",
|
||||
"error_category",
|
||||
"outcome",
|
||||
"skip_reason",
|
||||
"drain_reason",
|
||||
"grace_ms",
|
||||
"active_at_start",
|
||||
"pending_at_start",
|
||||
"producers_at_start",
|
||||
];
|
||||
|
||||
/// Flush a buffered batch once it reaches this many records.
|
||||
const LOG_BATCH_FLUSH_RECORDS: usize = 32;
|
||||
/// Flush a partial batch once its oldest record is at least this old
|
||||
/// (checked on the next event; the tail is fenced by teardown).
|
||||
const LOG_BATCH_MAX_AGE: Duration = Duration::from_secs(2);
|
||||
|
||||
fn is_allowed(name: &str) -> bool {
|
||||
ALLOWED_FIELDS.contains(&name)
|
||||
}
|
||||
|
||||
/// `tracing::Level` → OTLP (`SeverityText`, `SeverityNumber`).
|
||||
fn severity(level: &Level) -> (&'static str, i32) {
|
||||
match *level {
|
||||
Level::ERROR => ("ERROR", 17),
|
||||
Level::WARN => ("WARN", 13),
|
||||
Level::INFO => ("INFO", 9),
|
||||
Level::DEBUG => ("DEBUG", 5),
|
||||
Level::TRACE => ("TRACE", 1),
|
||||
}
|
||||
}
|
||||
|
||||
/// `>= INFO` in severity terms (INFO/WARN/ERROR). Note tracing orders
|
||||
/// `ERROR < WARN < INFO < DEBUG < TRACE`, so this is `level <= INFO`.
|
||||
fn at_least_info(level: &Level) -> bool {
|
||||
*level <= Level::INFO
|
||||
}
|
||||
|
||||
/// Big-endian byte encoding of the local parent's ids into the OTLP
|
||||
/// 16-byte / 8-byte fields; empty when no fastrace local parent is
|
||||
/// active (the common case for detached producer tasks).
|
||||
fn current_ids() -> (Vec<u8>, Vec<u8>) {
|
||||
match SpanContext::current_local_parent() {
|
||||
Some(ctx) => encode_ids(&ctx),
|
||||
None => (Vec::new(), Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn encode_ids(ctx: &SpanContext) -> (Vec<u8>, Vec<u8>) {
|
||||
(
|
||||
ctx.trace_id.0.to_be_bytes().to_vec(),
|
||||
ctx.span_id.0.to_be_bytes().to_vec(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Field visitor: keeps the message as the OTLP `Body` and only
|
||||
/// allowlisted fields as attributes; everything else is dropped.
|
||||
#[derive(Default)]
|
||||
struct AllowlistVisitor {
|
||||
body: Option<String>,
|
||||
attributes: Vec<KeyValue>,
|
||||
}
|
||||
|
||||
impl Visit for AllowlistVisitor {
|
||||
fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
|
||||
let name = field.name();
|
||||
if name == "message" {
|
||||
self.body = Some(format!("{value:?}"));
|
||||
} else if is_allowed(name) {
|
||||
self.attributes.push(string_kv(name, format!("{value:?}")));
|
||||
}
|
||||
}
|
||||
|
||||
fn record_str(&mut self, field: &Field, value: &str) {
|
||||
let name = field.name();
|
||||
if name == "message" {
|
||||
self.body = Some(value.to_owned());
|
||||
} else if is_allowed(name) {
|
||||
self.attributes.push(string_kv(name, value.to_owned()));
|
||||
}
|
||||
}
|
||||
|
||||
fn record_i64(&mut self, field: &Field, value: i64) {
|
||||
if is_allowed(field.name()) {
|
||||
self.push_int(field.name(), value);
|
||||
}
|
||||
}
|
||||
|
||||
fn record_u64(&mut self, field: &Field, value: u64) {
|
||||
if is_allowed(field.name()) {
|
||||
self.push_int(field.name(), value as i64);
|
||||
}
|
||||
}
|
||||
|
||||
fn record_bool(&mut self, field: &Field, value: bool) {
|
||||
if is_allowed(field.name()) {
|
||||
self.attributes.push(KeyValue {
|
||||
key: field.name().to_owned(),
|
||||
value: Some(AnyValue {
|
||||
value: Some(any_value::Value::BoolValue(value)),
|
||||
}),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn record_f64(&mut self, field: &Field, value: f64) {
|
||||
if is_allowed(field.name()) {
|
||||
self.attributes.push(KeyValue {
|
||||
key: field.name().to_owned(),
|
||||
value: Some(AnyValue {
|
||||
value: Some(any_value::Value::DoubleValue(value)),
|
||||
}),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AllowlistVisitor {
|
||||
fn push_int(&mut self, name: &str, value: i64) {
|
||||
self.attributes.push(KeyValue {
|
||||
key: name.to_owned(),
|
||||
value: Some(AnyValue {
|
||||
value: Some(any_value::Value::IntValue(value)),
|
||||
}),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn build_log_record(level: &Level, visitor: AllowlistVisitor) -> LogRecord {
|
||||
let now_nanos = now_unix_nanos();
|
||||
let (text, number) = severity(level);
|
||||
let (trace_id, span_id) = current_ids();
|
||||
LogRecord {
|
||||
time_unix_nano: now_nanos,
|
||||
observed_time_unix_nano: now_nanos,
|
||||
severity_number: number,
|
||||
severity_text: text.to_owned(),
|
||||
body: visitor.body.map(string_value),
|
||||
attributes: visitor.attributes,
|
||||
trace_id,
|
||||
span_id,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Encodes batches of OTLP `LogRecord`s onto the pump channel. Chunks at
|
||||
/// [`MAX_LOG_RECORDS_PER_DONATION`], drops payloads over
|
||||
/// [`MAX_DONATION_BYTES`], and never blocks.
|
||||
#[derive(Clone)]
|
||||
struct PumpLogExporter {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
resource: Resource,
|
||||
}
|
||||
|
||||
impl PumpLogExporter {
|
||||
fn export(&self, mut records: Vec<LogRecord>) {
|
||||
while !records.is_empty() {
|
||||
let chunk = if records.len() > MAX_LOG_RECORDS_PER_DONATION {
|
||||
let rest = records.split_off(MAX_LOG_RECORDS_PER_DONATION);
|
||||
std::mem::replace(&mut records, rest)
|
||||
} else {
|
||||
std::mem::take(&mut records)
|
||||
};
|
||||
let request = ExportLogsServiceRequest {
|
||||
resource_logs: vec![ResourceLogs {
|
||||
resource: Some(self.resource.clone()),
|
||||
scope_logs: vec![ScopeLogs {
|
||||
scope: Some(InstrumentationScope {
|
||||
name: TELEMETRY_TARGET.to_owned(),
|
||||
..Default::default()
|
||||
}),
|
||||
log_records: chunk,
|
||||
schema_url: String::new(),
|
||||
}],
|
||||
schema_url: String::new(),
|
||||
}],
|
||||
};
|
||||
let bytes = request.encode_to_vec();
|
||||
if bytes.len() > MAX_DONATION_BYTES {
|
||||
tracing::debug!(len = bytes.len(), "dropping oversized log donation payload");
|
||||
continue;
|
||||
}
|
||||
let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
|
||||
if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
|
||||
tracing::debug!("log donation queue full; dropping log batch");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Activation handle swapped into an inert [`DonatingLogLayer`]. Wraps
|
||||
/// the pump sender plus the resource (`service.name`) the layer needs to
|
||||
/// encode batches.
|
||||
pub struct LogDonationSender {
|
||||
exporter: PumpLogExporter,
|
||||
}
|
||||
|
||||
impl LogDonationSender {
|
||||
fn export(&self, records: Vec<LogRecord>) {
|
||||
self.exporter.export(records);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct LogBatch {
|
||||
records: Vec<LogRecord>,
|
||||
oldest: Option<Instant>,
|
||||
}
|
||||
|
||||
struct LogLayerShared {
|
||||
sender: ArcSwapOption<LogDonationSender>,
|
||||
batch: parking_lot::Mutex<LogBatch>,
|
||||
}
|
||||
|
||||
impl LogLayerShared {
|
||||
/// Buffer a record; return any records due for flush (count/age).
|
||||
fn push(&self, record: LogRecord) -> Vec<LogRecord> {
|
||||
let mut batch = self.batch.lock();
|
||||
if batch.records.is_empty() {
|
||||
batch.oldest = Some(Instant::now());
|
||||
}
|
||||
batch.records.push(record);
|
||||
let due = batch.records.len() >= LOG_BATCH_FLUSH_RECORDS
|
||||
|| batch
|
||||
.oldest
|
||||
.is_some_and(|t| t.elapsed() >= LOG_BATCH_MAX_AGE);
|
||||
if due {
|
||||
batch.oldest = None;
|
||||
std::mem::take(&mut batch.records)
|
||||
} else {
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
|
||||
/// Force the buffered batch onto the pump (teardown analogue of
|
||||
/// `fastrace::flush()`); no-op while inert or empty.
|
||||
fn flush(&self) {
|
||||
let records = {
|
||||
let mut batch = self.batch.lock();
|
||||
batch.oldest = None;
|
||||
std::mem::take(&mut batch.records)
|
||||
};
|
||||
if records.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Some(sender) = self.sender.load_full() {
|
||||
sender.export(records);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-global handle to the active layer's shared state so
|
||||
/// [`flush_log_layer`] can drive a teardown flush without a reference.
|
||||
static ACTIVE_LOG_LAYER: LazyLock<ArcSwapOption<LogLayerShared>> =
|
||||
LazyLock::new(ArcSwapOption::empty);
|
||||
|
||||
/// A composable [`tracing_subscriber::Layer`] that converts selected
|
||||
/// events into OTLP log records and batches them onto the pump.
|
||||
/// Installed inert; activated by [`Self::activate`].
|
||||
#[derive(Clone)]
|
||||
pub struct DonatingLogLayer {
|
||||
shared: Arc<LogLayerShared>,
|
||||
}
|
||||
|
||||
impl DonatingLogLayer {
|
||||
/// Install inert (no sender): selected events are dropped until
|
||||
/// [`Self::activate`] swaps a sender in. Registers itself as the
|
||||
/// process-global flush target.
|
||||
pub fn new_inert() -> Self {
|
||||
let shared = Arc::new(LogLayerShared {
|
||||
sender: ArcSwapOption::empty(),
|
||||
batch: parking_lot::Mutex::new(LogBatch::default()),
|
||||
});
|
||||
ACTIVE_LOG_LAYER.store(Some(shared.clone()));
|
||||
Self { shared }
|
||||
}
|
||||
|
||||
/// Swap in the donation sender, activating donation.
|
||||
pub fn activate(&self, sender: LogDonationSender) {
|
||||
self.shared.sender.store(Some(Arc::new(sender)));
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: tracing::Subscriber> Layer<S> for DonatingLogLayer {
|
||||
fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
|
||||
let Some(sender) = self.shared.sender.load_full() else {
|
||||
return;
|
||||
};
|
||||
let meta = event.metadata();
|
||||
if meta.target() != TELEMETRY_TARGET || !at_least_info(meta.level()) {
|
||||
return;
|
||||
}
|
||||
let mut visitor = AllowlistVisitor::default();
|
||||
event.record(&mut visitor);
|
||||
let record = build_log_record(meta.level(), visitor);
|
||||
let due = self.shared.push(record);
|
||||
if !due.is_empty() {
|
||||
sender.export(due);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush the active [`DonatingLogLayer`]'s in-memory batch onto the
|
||||
/// pump. Called from `ToolServer` teardown before the pump drain so a
|
||||
/// crash-y shutdown does not abandon a partial batch.
|
||||
pub fn flush_log_layer() {
|
||||
if let Some(shared) = ACTIVE_LOG_LAYER.load_full() {
|
||||
shared.flush();
|
||||
}
|
||||
}
|
||||
|
||||
/// Shutdown fence: drains queued log donations before the connection
|
||||
/// closes. Call after [`flush_log_layer`].
|
||||
pub struct LogDonationPump {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
}
|
||||
|
||||
impl LogDonationPump {
|
||||
/// Resolves once every payload queued before this call has had a
|
||||
/// send attempt.
|
||||
pub async fn drain(&self) {
|
||||
drain_via(&self.tx).await;
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolServer {
|
||||
/// Post-connect entry point: spawn the log donation pump (wiring
|
||||
/// [`ToolServer::donate_logs`]) and return a sender to swap into the
|
||||
/// already-installed inert [`DonatingLogLayer`] plus a drain handle.
|
||||
/// Does **not** return a `Layer` — a layer cannot be added to an
|
||||
/// already-set global subscriber. `service_name` must be
|
||||
/// server-allowlisted.
|
||||
pub fn log_donation_layer(
|
||||
&self,
|
||||
service_name: impl Into<String>,
|
||||
) -> (LogDonationSender, LogDonationPump) {
|
||||
let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
|
||||
let server = self.downgrade();
|
||||
tokio::spawn(run_pump(rx, move |payload: String| {
|
||||
let server = server.clone();
|
||||
async move {
|
||||
let Some(server) = server.upgrade() else {
|
||||
return (false, payload);
|
||||
};
|
||||
let ok = server.donate_logs(&payload).await.is_ok();
|
||||
(ok, payload)
|
||||
}
|
||||
}));
|
||||
self.set_log_donation_pump(tx.clone());
|
||||
|
||||
let sender = LogDonationSender {
|
||||
exporter: PumpLogExporter {
|
||||
tx: tx.clone(),
|
||||
resource: make_resource(service_name.into()),
|
||||
},
|
||||
};
|
||||
(sender, LogDonationPump { tx })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use fastrace::collector::{SpanId, TraceId};
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn test_sender(tx: mpsc::Sender<PumpMsg>) -> LogDonationSender {
|
||||
LogDonationSender {
|
||||
exporter: PumpLogExporter {
|
||||
tx,
|
||||
resource: make_resource("test-service".to_owned()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(payload: String) -> ExportLogsServiceRequest {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(payload)
|
||||
.expect("payload must be base64");
|
||||
ExportLogsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn severity_maps_levels_to_otlp_numbers() {
|
||||
assert_eq!(severity(&Level::ERROR), ("ERROR", 17));
|
||||
assert_eq!(severity(&Level::WARN), ("WARN", 13));
|
||||
assert_eq!(severity(&Level::INFO), ("INFO", 9));
|
||||
assert_eq!(severity(&Level::DEBUG), ("DEBUG", 5));
|
||||
assert_eq!(severity(&Level::TRACE), ("TRACE", 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn donation_filter_selects_info_and_above() {
|
||||
assert!(at_least_info(&Level::ERROR));
|
||||
assert!(at_least_info(&Level::WARN));
|
||||
assert!(at_least_info(&Level::INFO));
|
||||
assert!(!at_least_info(&Level::DEBUG));
|
||||
assert!(!at_least_info(&Level::TRACE));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encode_ids_is_big_endian_16_and_8_bytes() {
|
||||
let ctx = SpanContext::new(
|
||||
TraceId(0x0af7651916cd43dd8448eb211c80319c),
|
||||
SpanId(0xb7ad6b7169203331),
|
||||
);
|
||||
let (trace_id, span_id) = encode_ids(&ctx);
|
||||
assert_eq!(trace_id.len(), 16);
|
||||
assert_eq!(span_id.len(), 8);
|
||||
assert_eq!(
|
||||
format!("{:032x}", u128::from_be_bytes(trace_id.try_into().unwrap())),
|
||||
"0af7651916cd43dd8448eb211c80319c"
|
||||
);
|
||||
assert_eq!(
|
||||
format!("{:016x}", u64::from_be_bytes(span_id.try_into().unwrap())),
|
||||
"b7ad6b7169203331"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_ids_empty_without_local_parent() {
|
||||
let (trace_id, span_id) = current_ids();
|
||||
assert!(trace_id.is_empty());
|
||||
assert!(span_id.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_converts_event_and_redacts_free_form_fields() {
|
||||
let layer = DonatingLogLayer::new_inert();
|
||||
let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
|
||||
layer.activate(test_sender(tx));
|
||||
let flusher = layer.clone();
|
||||
|
||||
let subscriber = tracing_subscriber::registry().with(layer);
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
tracing::warn!(
|
||||
target: "workspace::telemetry",
|
||||
session_id = "s1",
|
||||
turn_number = 3u64,
|
||||
phase = "tool_state",
|
||||
error_category = "archive_failed",
|
||||
error = "secret git stderr with /home/user/path",
|
||||
"archive build failed (queued path)"
|
||||
);
|
||||
// Off-target event must never be forwarded.
|
||||
tracing::warn!(session_id = "s2", "unrelated chatter");
|
||||
// DEBUG on-target is below the threshold.
|
||||
tracing::debug!(target: "workspace::telemetry", session_id = "s3", "verbose");
|
||||
});
|
||||
flusher.shared.flush();
|
||||
|
||||
let PumpMsg::Payload(payload) = rx.try_recv().expect("one batch must be queued") else {
|
||||
panic!("expected a payload");
|
||||
};
|
||||
let request = decode(payload);
|
||||
let scope_logs = &request.resource_logs[0].scope_logs[0];
|
||||
assert_eq!(
|
||||
scope_logs.scope.as_ref().unwrap().name,
|
||||
"workspace::telemetry"
|
||||
);
|
||||
assert_eq!(
|
||||
scope_logs.log_records.len(),
|
||||
1,
|
||||
"only the WARN on-target row"
|
||||
);
|
||||
|
||||
let record = &scope_logs.log_records[0];
|
||||
assert_eq!(record.severity_text, "WARN");
|
||||
assert_eq!(record.severity_number, 13);
|
||||
assert_eq!(
|
||||
record.body.as_ref().unwrap().value,
|
||||
Some(any_value::Value::StringValue(
|
||||
"archive build failed (queued path)".to_owned()
|
||||
))
|
||||
);
|
||||
let keys: Vec<&str> = record.attributes.iter().map(|kv| kv.key.as_str()).collect();
|
||||
assert!(keys.contains(&"session_id"));
|
||||
assert!(keys.contains(&"turn_number"));
|
||||
assert!(keys.contains(&"phase"));
|
||||
assert!(keys.contains(&"error_category"));
|
||||
assert!(
|
||||
!keys.contains(&"error"),
|
||||
"free-form `error` must be dropped, got {keys:?}"
|
||||
);
|
||||
|
||||
// Resource carries the donor service.name.
|
||||
let service_name = request.resource_logs[0]
|
||||
.resource
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.attributes
|
||||
.iter()
|
||||
.find(|kv| kv.key == "service.name")
|
||||
.and_then(|kv| kv.value.as_ref())
|
||||
.and_then(|v| v.value.clone());
|
||||
assert_eq!(
|
||||
service_name,
|
||||
Some(any_value::Value::StringValue("test-service".to_owned()))
|
||||
);
|
||||
|
||||
assert!(rx.try_recv().is_err(), "no further payloads");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inert_layer_drops_selected_events() {
|
||||
let layer = DonatingLogLayer::new_inert();
|
||||
let flusher = layer.clone();
|
||||
// No sender activated.
|
||||
let subscriber = tracing_subscriber::registry().with(layer);
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
tracing::warn!(target: "workspace::telemetry", session_id = "s1", "dropped");
|
||||
});
|
||||
flusher.shared.flush();
|
||||
// Nothing to assert beyond not panicking: with no sender the
|
||||
// batch never fills and flush is a no-op.
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exporter_chunks_at_max_records_per_donation() {
|
||||
let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
|
||||
let exporter = PumpLogExporter {
|
||||
tx,
|
||||
resource: make_resource("test-service".to_owned()),
|
||||
};
|
||||
let records = vec![LogRecord::default(); MAX_LOG_RECORDS_PER_DONATION + 1];
|
||||
exporter.export(records);
|
||||
|
||||
let mut total = 0;
|
||||
let mut payloads = 0;
|
||||
while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
|
||||
payloads += 1;
|
||||
total += decode(p).resource_logs[0].scope_logs[0].log_records.len();
|
||||
}
|
||||
assert_eq!(payloads, 2, "one full chunk + remainder");
|
||||
assert_eq!(total, MAX_LOG_RECORDS_PER_DONATION + 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
//! Forward the process's Prometheus metrics to the connected server over
|
||||
//! the WebSocket transport (`metrics.donate`).
|
||||
//!
|
||||
//! A [`MetricDonationReporter`] periodically snapshots the default
|
||||
//! Prometheus registry via [`prometheus::gather`], converts the
|
||||
//! `MetricFamily` set to native OTLP metrics (Counter→Sum, Gauge→Gauge,
|
||||
//! Histogram→Histogram, labels preserved, cumulative temporality), and
|
||||
//! pumps the batch over the shared [`crate::donate_pump`]. Because it
|
||||
//! gathers the whole registry, every current and future metric is
|
||||
//! exported with zero per-metric wiring. Metrics are **process-aggregate**
|
||||
//! — [`ToolServer::donate_metrics`] requires no bound session.
|
||||
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwapOption;
|
||||
use base64::Engine as _;
|
||||
use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_METRICS_PER_DONATION};
|
||||
use opentelemetry_proto::tonic::collector::metrics::v1::ExportMetricsServiceRequest;
|
||||
use opentelemetry_proto::tonic::common::v1::KeyValue;
|
||||
use opentelemetry_proto::tonic::metrics::v1::{
|
||||
AggregationTemporality, Gauge, Histogram, HistogramDataPoint, Metric, NumberDataPoint,
|
||||
ResourceMetrics, ScopeMetrics, Sum, metric, number_data_point,
|
||||
};
|
||||
use opentelemetry_proto::tonic::resource::v1::Resource;
|
||||
use prometheus::proto::{MetricFamily, MetricType};
|
||||
use prost::Message as _;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::donate_pump::{
|
||||
PENDING_FLUSHES, PumpMsg, make_resource, now_unix_nanos, run_pump, string_kv,
|
||||
};
|
||||
use crate::server::ToolServer;
|
||||
|
||||
/// How often the reporter snapshots the registry. The server re-stamps
|
||||
/// attribution; cumulative temporality means missed ticks only delay
|
||||
/// freshness, never lose monotonic state.
|
||||
const DEFAULT_GATHER_INTERVAL: Duration = Duration::from_secs(60);
|
||||
|
||||
fn labels_to_kv(labels: &[prometheus::proto::LabelPair]) -> Vec<KeyValue> {
|
||||
labels
|
||||
.iter()
|
||||
.map(|l| string_kv(l.name(), l.value().to_owned()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn number_point(metric: &prometheus::proto::Metric, value: f64, now: u64) -> NumberDataPoint {
|
||||
NumberDataPoint {
|
||||
attributes: labels_to_kv(metric.get_label()),
|
||||
time_unix_nano: now,
|
||||
value: Some(number_data_point::Value::AsDouble(value)),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Prometheus histogram buckets are **cumulative** (`le` counts); OTLP
|
||||
/// wants per-bucket counts plus an implicit `+Inf` bucket, so the
|
||||
/// cumulative counts are differenced here.
|
||||
fn histogram_point(metric: &prometheus::proto::Metric, now: u64) -> HistogramDataPoint {
|
||||
let hist = metric.get_histogram();
|
||||
let mut bucket_counts = Vec::new();
|
||||
let mut explicit_bounds = Vec::new();
|
||||
let mut prev = 0u64;
|
||||
for bucket in hist.get_bucket() {
|
||||
let cumulative = bucket.cumulative_count();
|
||||
bucket_counts.push(cumulative.saturating_sub(prev));
|
||||
explicit_bounds.push(bucket.upper_bound());
|
||||
prev = cumulative;
|
||||
}
|
||||
let total = hist.get_sample_count();
|
||||
bucket_counts.push(total.saturating_sub(prev));
|
||||
HistogramDataPoint {
|
||||
attributes: labels_to_kv(metric.get_label()),
|
||||
time_unix_nano: now,
|
||||
count: total,
|
||||
sum: Some(hist.get_sample_sum()),
|
||||
bucket_counts,
|
||||
explicit_bounds,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a gathered `MetricFamily` set to OTLP metrics. Summaries and
|
||||
/// untyped families are skipped defensively (none registered today).
|
||||
fn convert_families(families: &[MetricFamily]) -> Vec<Metric> {
|
||||
let now = now_unix_nanos();
|
||||
let cumulative = AggregationTemporality::Cumulative as i32;
|
||||
let mut out = Vec::new();
|
||||
for family in families {
|
||||
let name = family.name().to_owned();
|
||||
let data = match family.get_field_type() {
|
||||
MetricType::COUNTER => metric::Data::Sum(Sum {
|
||||
data_points: family
|
||||
.get_metric()
|
||||
.iter()
|
||||
.map(|m| number_point(m, m.get_counter().value(), now))
|
||||
.collect(),
|
||||
aggregation_temporality: cumulative,
|
||||
is_monotonic: true,
|
||||
}),
|
||||
MetricType::GAUGE => metric::Data::Gauge(Gauge {
|
||||
data_points: family
|
||||
.get_metric()
|
||||
.iter()
|
||||
.map(|m| number_point(m, m.get_gauge().value(), now))
|
||||
.collect(),
|
||||
}),
|
||||
MetricType::HISTOGRAM => metric::Data::Histogram(Histogram {
|
||||
data_points: family
|
||||
.get_metric()
|
||||
.iter()
|
||||
.map(|m| histogram_point(m, now))
|
||||
.collect(),
|
||||
aggregation_temporality: cumulative,
|
||||
}),
|
||||
MetricType::SUMMARY | MetricType::UNTYPED => continue,
|
||||
};
|
||||
out.push(Metric {
|
||||
name,
|
||||
data: Some(data),
|
||||
..Default::default()
|
||||
});
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Encodes batches of OTLP metrics onto the pump channel. Chunks at
|
||||
/// [`MAX_METRICS_PER_DONATION`], drops payloads over
|
||||
/// [`MAX_DONATION_BYTES`], and never blocks.
|
||||
#[derive(Clone)]
|
||||
struct MetricExporter {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
resource: Resource,
|
||||
}
|
||||
|
||||
impl MetricExporter {
|
||||
fn export(&self, mut metrics: Vec<Metric>) {
|
||||
while !metrics.is_empty() {
|
||||
let chunk = if metrics.len() > MAX_METRICS_PER_DONATION {
|
||||
let rest = metrics.split_off(MAX_METRICS_PER_DONATION);
|
||||
std::mem::replace(&mut metrics, rest)
|
||||
} else {
|
||||
std::mem::take(&mut metrics)
|
||||
};
|
||||
let request = ExportMetricsServiceRequest {
|
||||
resource_metrics: vec![ResourceMetrics {
|
||||
resource: Some(self.resource.clone()),
|
||||
scope_metrics: vec![ScopeMetrics {
|
||||
metrics: chunk,
|
||||
..Default::default()
|
||||
}],
|
||||
schema_url: String::new(),
|
||||
}],
|
||||
};
|
||||
let bytes = request.encode_to_vec();
|
||||
if bytes.len() > MAX_DONATION_BYTES {
|
||||
tracing::debug!(
|
||||
len = bytes.len(),
|
||||
"dropping oversized metric donation payload"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
|
||||
if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
|
||||
tracing::debug!("metric donation queue full; dropping metric batch");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn gather_and_send(&self) {
|
||||
let metrics = convert_families(&prometheus::gather());
|
||||
if metrics.is_empty() {
|
||||
return;
|
||||
}
|
||||
self.export(metrics);
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-global handle to the active exporter so [`gather_and_send`]
|
||||
/// can drive a final teardown gather without a reference.
|
||||
static ACTIVE_METRIC_EXPORTER: LazyLock<ArcSwapOption<MetricExporter>> =
|
||||
LazyLock::new(ArcSwapOption::empty);
|
||||
|
||||
/// Final registry gather onto the active metric pump. Called from
|
||||
/// `ToolServer` teardown before the pump drain so a crash-y shutdown
|
||||
/// captures the latest values.
|
||||
pub(crate) fn gather_and_send() {
|
||||
if let Some(exporter) = ACTIVE_METRIC_EXPORTER.load_full() {
|
||||
exporter.gather_and_send();
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop the process-global exporter on teardown so its pump `Sender` is released
|
||||
/// and the metric pump can wind down. Called from `flush_donations_inner` after
|
||||
/// the final [`gather_and_send`] (and alongside clearing the stored pump
|
||||
/// senders), so a dropped `ToolServer` doesn't leak the pump task.
|
||||
pub(crate) fn clear_active_exporter() {
|
||||
ACTIVE_METRIC_EXPORTER.store(None);
|
||||
}
|
||||
|
||||
/// Periodic registry gatherer spawned by
|
||||
/// [`ToolServer::metric_donation_reporter`]. Internal: constructed and run
|
||||
/// only by `metric_donation_reporter`; not part of the crate's public API.
|
||||
pub(crate) struct MetricDonationReporter {
|
||||
exporter: MetricExporter,
|
||||
interval: Duration,
|
||||
shutdown: CancellationToken,
|
||||
}
|
||||
|
||||
impl MetricDonationReporter {
|
||||
async fn run(self) {
|
||||
let mut ticker = tokio::time::interval(self.interval);
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = ticker.tick() => self.exporter.gather_and_send(),
|
||||
// Stop on teardown so this task (and the pump `tx` clone it
|
||||
// holds) doesn't outlive `ToolServer::shutdown` and keep
|
||||
// gathering/sending forever.
|
||||
_ = self.shutdown.cancelled() => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shutdown fence: drains queued metric donations before the connection
|
||||
/// closes.
|
||||
pub struct MetricDonationPump {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
}
|
||||
|
||||
impl MetricDonationPump {
|
||||
/// Resolves once every payload queued before this call has had a
|
||||
/// send attempt.
|
||||
pub async fn drain(&self) {
|
||||
crate::donate_pump::drain_via(&self.tx).await;
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolServer {
|
||||
/// Post-connect entry point: spawn the metric donation pump (wiring
|
||||
/// [`ToolServer::donate_metrics`]) plus the periodic registry
|
||||
/// gatherer, and return a drain handle. Activates on server presence
|
||||
/// (no env flag). `service_name` must be server-allowlisted.
|
||||
pub fn metric_donation_reporter(&self, service_name: impl Into<String>) -> MetricDonationPump {
|
||||
let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
|
||||
let server = self.downgrade();
|
||||
tokio::spawn(run_pump(rx, move |payload: String| {
|
||||
let server = server.clone();
|
||||
async move {
|
||||
let Some(server) = server.upgrade() else {
|
||||
return (false, payload);
|
||||
};
|
||||
let ok = server.donate_metrics(&payload).await.is_ok();
|
||||
(ok, payload)
|
||||
}
|
||||
}));
|
||||
self.set_metric_donation_pump(tx.clone());
|
||||
|
||||
let exporter = MetricExporter {
|
||||
tx: tx.clone(),
|
||||
resource: make_resource(service_name.into()),
|
||||
};
|
||||
ACTIVE_METRIC_EXPORTER.store(Some(Arc::new(exporter.clone())));
|
||||
tokio::spawn(
|
||||
MetricDonationReporter {
|
||||
exporter,
|
||||
interval: DEFAULT_GATHER_INTERVAL,
|
||||
shutdown: self.shutdown_token(),
|
||||
}
|
||||
.run(),
|
||||
);
|
||||
|
||||
MetricDonationPump { tx }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use opentelemetry_proto::tonic::common::v1::any_value;
|
||||
use prometheus::{Histogram, HistogramOpts, IntCounterVec, IntGauge, Opts, Registry};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn decode(payload: String) -> ExportMetricsServiceRequest {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(payload)
|
||||
.expect("payload must be base64");
|
||||
ExportMetricsServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP")
|
||||
}
|
||||
|
||||
fn label_map(attrs: &[KeyValue]) -> std::collections::HashMap<String, String> {
|
||||
attrs
|
||||
.iter()
|
||||
.filter_map(|kv| match kv.value.as_ref().and_then(|v| v.value.clone()) {
|
||||
Some(any_value::Value::StringValue(s)) => Some((kv.key.clone(), s)),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converts_counter_gauge_histogram_with_labels() {
|
||||
let registry = Registry::new();
|
||||
|
||||
let counter =
|
||||
IntCounterVec::new(Opts::new("grok_test_total", "help"), &["reason"]).unwrap();
|
||||
registry.register(Box::new(counter.clone())).unwrap();
|
||||
counter.with_label_values(&["zdr"]).inc_by(5);
|
||||
|
||||
let gauge = IntGauge::new("grok_test_pending", "help").unwrap();
|
||||
registry.register(Box::new(gauge.clone())).unwrap();
|
||||
gauge.set(7);
|
||||
|
||||
let hist = Histogram::with_opts(
|
||||
HistogramOpts::new("grok_test_seconds", "help").buckets(vec![0.5, 1.0]),
|
||||
)
|
||||
.unwrap();
|
||||
registry.register(Box::new(hist.clone())).unwrap();
|
||||
hist.observe(0.25);
|
||||
hist.observe(0.75);
|
||||
hist.observe(5.0);
|
||||
|
||||
let metrics = convert_families(®istry.gather());
|
||||
let by_name: std::collections::HashMap<_, _> =
|
||||
metrics.iter().map(|m| (m.name.clone(), m)).collect();
|
||||
|
||||
// Counter -> Sum (monotonic, cumulative), label preserved.
|
||||
let metric::Data::Sum(sum) = by_name["grok_test_total"].data.as_ref().unwrap() else {
|
||||
panic!("counter must convert to Sum");
|
||||
};
|
||||
assert!(sum.is_monotonic);
|
||||
assert_eq!(
|
||||
sum.aggregation_temporality,
|
||||
AggregationTemporality::Cumulative as i32
|
||||
);
|
||||
let dp = &sum.data_points[0];
|
||||
assert_eq!(dp.value, Some(number_data_point::Value::AsDouble(5.0)));
|
||||
assert_eq!(
|
||||
label_map(&dp.attributes).get("reason").map(String::as_str),
|
||||
Some("zdr")
|
||||
);
|
||||
|
||||
// Gauge -> Gauge.
|
||||
let metric::Data::Gauge(g) = by_name["grok_test_pending"].data.as_ref().unwrap() else {
|
||||
panic!("gauge must convert to Gauge");
|
||||
};
|
||||
assert_eq!(
|
||||
g.data_points[0].value,
|
||||
Some(number_data_point::Value::AsDouble(7.0))
|
||||
);
|
||||
|
||||
// Histogram -> Histogram with cumulative buckets differenced and
|
||||
// a +Inf bucket appended.
|
||||
let metric::Data::Histogram(h) = by_name["grok_test_seconds"].data.as_ref().unwrap() else {
|
||||
panic!("histogram must convert to Histogram");
|
||||
};
|
||||
assert_eq!(
|
||||
h.aggregation_temporality,
|
||||
AggregationTemporality::Cumulative as i32
|
||||
);
|
||||
let hdp = &h.data_points[0];
|
||||
assert_eq!(hdp.count, 3);
|
||||
assert_eq!(hdp.sum, Some(6.0));
|
||||
assert_eq!(hdp.explicit_bounds, vec![0.5, 1.0]);
|
||||
// (<=0.5): 0.25 -> 1 ; (0.5,1.0]: 0.75 -> 1 ; (+Inf): 5.0 -> 1
|
||||
assert_eq!(hdp.bucket_counts, vec![1, 1, 1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exporter_chunks_at_max_metrics_per_donation() {
|
||||
let (tx, mut rx) = mpsc::channel::<PumpMsg>(8);
|
||||
let exporter = MetricExporter {
|
||||
tx,
|
||||
resource: make_resource("test-service".to_owned()),
|
||||
};
|
||||
let metrics = vec![Metric::default(); MAX_METRICS_PER_DONATION + 1];
|
||||
exporter.export(metrics);
|
||||
|
||||
let mut payloads = 0;
|
||||
let mut total = 0;
|
||||
while let Ok(PumpMsg::Payload(p)) = rx.try_recv() {
|
||||
payloads += 1;
|
||||
total += decode(p).resource_metrics[0].scope_metrics[0].metrics.len();
|
||||
}
|
||||
assert_eq!(payloads, 2, "one full chunk + remainder");
|
||||
assert_eq!(total, MAX_METRICS_PER_DONATION + 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn summary_and_untyped_families_are_skipped() {
|
||||
// An empty registry gathers nothing; convert yields nothing.
|
||||
let registry = Registry::new();
|
||||
assert!(convert_families(®istry.gather()).is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,552 @@
|
||||
//! Feature-gated Prometheus metrics for the SDK.
|
||||
//!
|
||||
//! When the `metrics` cargo feature is enabled, each helper records to a
|
||||
//! lazily-registered Prometheus counter / gauge / histogram. When
|
||||
//! disabled (the default), every helper compiles to an empty function
|
||||
//! body so the SDK carries zero prometheus dependency.
|
||||
|
||||
#[cfg(feature = "metrics")]
|
||||
mod inner {
|
||||
use prometheus::{
|
||||
Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec,
|
||||
exponential_buckets, register_histogram, register_histogram_vec, register_int_counter,
|
||||
register_int_counter_vec, register_int_gauge, register_int_gauge_vec,
|
||||
};
|
||||
use std::sync::LazyLock;
|
||||
|
||||
static POOL_CONNECTIONS: LazyLock<IntGauge> = LazyLock::new(|| {
|
||||
register_int_gauge!(
|
||||
"computer_hub_client_pool_connections",
|
||||
"Active pooled connections in the SDK connection pool."
|
||||
)
|
||||
.expect("computer_hub_client_pool_connections must register once")
|
||||
});
|
||||
|
||||
static POOL_EVICTIONS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_pool_evictions_total",
|
||||
"Pooled connections closed by the idle reaper (unused past the idle TTL)."
|
||||
)
|
||||
.expect("computer_hub_client_pool_evictions_total must register once")
|
||||
});
|
||||
|
||||
static RECONNECTS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_reconnects_total",
|
||||
"Cumulative reconnect attempts that succeeded."
|
||||
)
|
||||
.expect("computer_hub_client_reconnects_total must register once")
|
||||
});
|
||||
|
||||
static RECONNECT_FAILED_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"computer_hub_client_reconnect_failed_total",
|
||||
"Cumulative reconnect attempts that failed, by reason \
|
||||
(handshake_auth = fatal 401/403, transport = retryable).",
|
||||
&["reason"]
|
||||
)
|
||||
.expect("computer_hub_client_reconnect_failed_total must register once")
|
||||
});
|
||||
|
||||
static RECONNECT_DURATION_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
|
||||
register_histogram!(
|
||||
"computer_hub_client_reconnect_duration_seconds",
|
||||
"Time to complete a reconnect cycle (handshake + session/tool replay).",
|
||||
exponential_buckets(0.01, 2.0, 14).expect("valid bucket params")
|
||||
)
|
||||
.expect("computer_hub_client_reconnect_duration_seconds must register once")
|
||||
});
|
||||
|
||||
static RECONNECTS_BY_CAUSE_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"computer_hub_client_reconnects_by_cause_total",
|
||||
"Successful reconnects by disconnect cause of the previous connection \
|
||||
(close_frame, eof, transport_read_error, transport_write_error, forced). \
|
||||
Cause-labeled companion to computer_hub_client_reconnects_total.",
|
||||
&["cause"]
|
||||
)
|
||||
.expect("computer_hub_client_reconnects_by_cause_total must register once")
|
||||
});
|
||||
|
||||
static RECONNECT_GAP_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
|
||||
register_histogram!(
|
||||
"computer_hub_client_reconnect_gap_seconds",
|
||||
"Time from the last inbound frame on the dead connection to a successful reconnect.",
|
||||
exponential_buckets(0.1, 2.0, 14).expect("valid bucket params")
|
||||
)
|
||||
.expect("computer_hub_client_reconnect_gap_seconds must register once")
|
||||
});
|
||||
|
||||
static CALL_DISPATCH_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
|
||||
register_histogram!(
|
||||
"computer_hub_client_call_dispatch_seconds",
|
||||
"Time to set up and queue the outbound remote dispatch.",
|
||||
exponential_buckets(0.0001, 2.0, 14).expect("valid bucket params")
|
||||
)
|
||||
.expect("computer_hub_client_call_dispatch_seconds must register once")
|
||||
});
|
||||
|
||||
static DEMUX_INBOX_DEPTH: LazyLock<IntGauge> = LazyLock::new(|| {
|
||||
register_int_gauge!(
|
||||
"computer_hub_client_demux_inbox_depth",
|
||||
"Number of session inboxes registered in the inbound demux."
|
||||
)
|
||||
.expect("computer_hub_client_demux_inbox_depth must register once")
|
||||
});
|
||||
|
||||
static CALL_ID_COLLISIONS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_call_id_collisions_total",
|
||||
"Call-id collisions detected in the harness dispatch path."
|
||||
)
|
||||
.expect("computer_hub_client_call_id_collisions_total must register once")
|
||||
});
|
||||
|
||||
// ── Server integration metrics ──────────────────────────────────
|
||||
|
||||
static HARNESS_CONNECT_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"hub_harness_connect_total",
|
||||
"Hub connection attempts by outcome and sampler.",
|
||||
&["status", "sampler"]
|
||||
)
|
||||
.expect("hub_harness_connect_total must register once")
|
||||
});
|
||||
|
||||
static SESSION_EVENT_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"hub_session_event_total",
|
||||
"SessionEvent emissions by event type.",
|
||||
&["event_type"]
|
||||
)
|
||||
.expect("hub_session_event_total must register once")
|
||||
});
|
||||
|
||||
static SESSION_OP_DURATION_SECONDS: LazyLock<HistogramVec> = LazyLock::new(|| {
|
||||
register_histogram_vec!(
|
||||
"hub_session_op_duration_seconds",
|
||||
"Latency of hub session lifecycle operations (open/bind) by op and outcome.",
|
||||
&["op", "status"],
|
||||
exponential_buckets(0.001, 2.0, 14).expect("valid bucket params")
|
||||
)
|
||||
.expect("hub_session_op_duration_seconds must register once")
|
||||
});
|
||||
|
||||
static SESSION_SOFT_REBIND_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_session_soft_rebind_total",
|
||||
"Redundant session.bind frames for a session with a live dispatch loop, \
|
||||
handled as a non-destructive soft rebind (serve state refreshed, \
|
||||
in-flight calls preserved)."
|
||||
)
|
||||
.expect("hub_session_soft_rebind_total must register once")
|
||||
});
|
||||
|
||||
static NO_HANDLER_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_sdk_no_handler_total",
|
||||
"tool_call_request frames rejected with -32011 because the session's \
|
||||
current handler set has no handler for the requested tool_id."
|
||||
)
|
||||
.expect("hub_sdk_no_handler_total must register once")
|
||||
});
|
||||
|
||||
static HOOK_SEND_TOTAL: LazyLock<IntCounterVec> = LazyLock::new(|| {
|
||||
register_int_counter_vec!(
|
||||
"hub_hook_send_total",
|
||||
"Hook sends by hook type.",
|
||||
&["hook_type"]
|
||||
)
|
||||
.expect("hub_hook_send_total must register once")
|
||||
});
|
||||
|
||||
static PROGRESS_FRAMES_FORWARDED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_progress_frames_forwarded_total",
|
||||
"Progress frames forwarded by ToolServer."
|
||||
)
|
||||
.expect("hub_progress_frames_forwarded_total must register once")
|
||||
});
|
||||
|
||||
static CANCEL_HOOK_RECEIVED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_cancel_hook_received_total",
|
||||
"Cancel hooks received by workspace tool server."
|
||||
)
|
||||
.expect("hub_cancel_hook_received_total must register once")
|
||||
});
|
||||
|
||||
static WRITER_SINK_SEND_ERRORS_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_writer_sink_send_errors_total",
|
||||
"Writer-task sink send failures; each signals the reader to reconnect."
|
||||
)
|
||||
.expect("computer_hub_client_writer_sink_send_errors_total must register once")
|
||||
});
|
||||
|
||||
static RECONNECT_WRITER_RESUME_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_reconnect_writer_resume_total",
|
||||
"Fresh-sink Resume handoffs delivered to the writer task after reconnect."
|
||||
)
|
||||
.expect("computer_hub_client_reconnect_writer_resume_total must register once")
|
||||
});
|
||||
|
||||
static LIVENESS_DEADLINE_EXPIRED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_liveness_deadline_expired_total",
|
||||
"Liveness-deadline expiries in the reader (no inbound WebSocket \
|
||||
frame within the deadline); each declares the connection dead and \
|
||||
drives the normal reconnect path."
|
||||
)
|
||||
.expect("computer_hub_client_liveness_deadline_expired_total must register once")
|
||||
});
|
||||
|
||||
static HEARTBEAT_PONG_DROPPED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_heartbeat_pong_dropped_total",
|
||||
"App-level heartbeat pongs dropped because outbound_tx was saturated (split reader)."
|
||||
)
|
||||
.expect("computer_hub_client_heartbeat_pong_dropped_total must register once")
|
||||
});
|
||||
|
||||
static CANCEL_APPLIED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_cancel_applied_total",
|
||||
"Cancel hooks that hit a live in-flight call and cancelled it."
|
||||
)
|
||||
.expect("hub_cancel_applied_total must register once")
|
||||
});
|
||||
|
||||
static CANCEL_PENDING_TOMBSTONED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_cancel_pending_tombstoned_total",
|
||||
"Cancel hooks recorded as a pending tombstone (call not yet registered or already done)."
|
||||
)
|
||||
.expect("hub_cancel_pending_tombstoned_total must register once")
|
||||
});
|
||||
|
||||
static CANCEL_NO_TARGET_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_cancel_no_target_total",
|
||||
"Cancel hooks with no call_id (session-wide, no specific call to cancel)."
|
||||
)
|
||||
.expect("hub_cancel_no_target_total must register once")
|
||||
});
|
||||
|
||||
static TOOL_CALL_REJECTED_OVERLOADED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_tool_call_rejected_overloaded_total",
|
||||
"Tool calls rejected by admission timeout (-32016 tool_busy)."
|
||||
)
|
||||
.expect("hub_tool_call_rejected_overloaded_total must register once")
|
||||
});
|
||||
|
||||
static INBOX_FULL_REQUEST_REJECTED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_inbox_full_request_rejected_total",
|
||||
"Requests rejected with an overloaded response on a full session inbox."
|
||||
)
|
||||
.expect("hub_inbox_full_request_rejected_total must register once")
|
||||
});
|
||||
|
||||
static INBOX_FULL_REJECT_SEND_FAILED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_inbox_full_reject_send_failed_total",
|
||||
"Overloaded rejections dropped because outbound was also full (residual silent loss)."
|
||||
)
|
||||
.expect("hub_inbox_full_reject_send_failed_total must register once")
|
||||
});
|
||||
|
||||
static INBOX_FULL_NOTIFICATION_DROPPED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_inbox_full_notification_dropped_total",
|
||||
"Notifications (no id) dropped on a full session inbox."
|
||||
)
|
||||
.expect("hub_inbox_full_notification_dropped_total must register once")
|
||||
});
|
||||
|
||||
static SERVE_REPLAY_TIMEOUT_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"computer_hub_client_serve_replay_timeout_total",
|
||||
"serve attempts that hit the per-attempt reply deadline, from any \
|
||||
serve call site (reconnect replay, run(), bind, tool updates)."
|
||||
)
|
||||
.expect("computer_hub_client_serve_replay_timeout_total must register once")
|
||||
});
|
||||
|
||||
static NOTIF_LAGGED_RECOVERED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_notif_lagged_recovered_total",
|
||||
"Connection-notification broadcast Lagged events recovered by \
|
||||
continuing the loop instead of exiting."
|
||||
)
|
||||
.expect("hub_notif_lagged_recovered_total must register once")
|
||||
});
|
||||
|
||||
static EARLY_NOTIF_BUFFERED_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"hub_early_notif_buffered_total",
|
||||
"Connection-level notification frames (binds, unbinds, evicts, \
|
||||
...) buffered between connect and ToolServer::run() and replayed \
|
||||
instead of dropped."
|
||||
)
|
||||
.expect("hub_early_notif_buffered_total must register once")
|
||||
});
|
||||
|
||||
static TOOL_CALL_INFLIGHT: LazyLock<IntGaugeVec> = LazyLock::new(|| {
|
||||
register_int_gauge_vec!(
|
||||
"hub_tool_call_inflight",
|
||||
"Concurrent running tool calls holding an admission permit, by scope.",
|
||||
&["scope"]
|
||||
)
|
||||
.expect("hub_tool_call_inflight must register once")
|
||||
});
|
||||
|
||||
static ADMISSION_WAIT_SECONDS: LazyLock<Histogram> = LazyLock::new(|| {
|
||||
register_histogram!(
|
||||
"hub_tool_call_admission_wait_seconds",
|
||||
"Time blocked acquiring the three admission permits (one shared deadline).",
|
||||
exponential_buckets(0.0001, 2.0, 16).expect("valid bucket params")
|
||||
)
|
||||
.expect("hub_tool_call_admission_wait_seconds must register once")
|
||||
});
|
||||
|
||||
pub(crate) fn pool_connections_inc() {
|
||||
POOL_CONNECTIONS.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn pool_connections_dec() {
|
||||
POOL_CONNECTIONS.dec();
|
||||
}
|
||||
|
||||
pub(crate) fn pool_evictions_inc() {
|
||||
POOL_EVICTIONS_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_succeeded() {
|
||||
RECONNECTS_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_failed(reason: &str) {
|
||||
RECONNECT_FAILED_TOTAL.with_label_values(&[reason]).inc();
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_duration_observe(secs: f64) {
|
||||
RECONNECT_DURATION_SECONDS.observe(secs);
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_cause(cause: &str) {
|
||||
RECONNECTS_BY_CAUSE_TOTAL.with_label_values(&[cause]).inc();
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_gap_observe(secs: f64) {
|
||||
RECONNECT_GAP_SECONDS.observe(secs);
|
||||
}
|
||||
|
||||
pub(crate) fn call_dispatch_observe(secs: f64) {
|
||||
CALL_DISPATCH_SECONDS.observe(secs);
|
||||
}
|
||||
|
||||
pub(crate) fn demux_inbox_depth_set(depth: i64) {
|
||||
DEMUX_INBOX_DEPTH.set(depth);
|
||||
}
|
||||
|
||||
pub(crate) fn call_id_collision() {
|
||||
CALL_ID_COLLISIONS_TOTAL.inc();
|
||||
}
|
||||
|
||||
/// Record a harness connection attempt.
|
||||
///
|
||||
/// `sampler` identifies the caller (`"chat"` or `"shell"`).
|
||||
/// `status` is `"ok"`, `"error"`, or `"fallback"` (fallback is
|
||||
/// emitted by the caller in `AgentBuilder::build_harness()`, not
|
||||
/// by the SDK).
|
||||
pub fn harness_connect(status: &str, sampler: &str) {
|
||||
HARNESS_CONNECT_TOTAL
|
||||
.with_label_values(&[status, sampler])
|
||||
.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn session_event(event_type: &str) {
|
||||
SESSION_EVENT_TOTAL.with_label_values(&[event_type]).inc();
|
||||
}
|
||||
|
||||
/// Observe the latency of a session lifecycle operation.
|
||||
/// `op` is `"open"` or `"bind"`; `status` is `"ok"` or `"error"`.
|
||||
pub(crate) fn session_op_observe(op: &str, status: &str, secs: f64) {
|
||||
SESSION_OP_DURATION_SECONDS
|
||||
.with_label_values(&[op, status])
|
||||
.observe(secs);
|
||||
}
|
||||
|
||||
pub(crate) fn session_soft_rebind() {
|
||||
SESSION_SOFT_REBIND_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn no_handler() {
|
||||
NO_HANDLER_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn hook_send(hook_type: &str) {
|
||||
HOOK_SEND_TOTAL.with_label_values(&[hook_type]).inc();
|
||||
}
|
||||
|
||||
pub(crate) fn progress_frame_forwarded() {
|
||||
PROGRESS_FRAMES_FORWARDED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn cancel_hook_received() {
|
||||
CANCEL_HOOK_RECEIVED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn writer_sink_send_error() {
|
||||
WRITER_SINK_SEND_ERRORS_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn reconnect_writer_resume() {
|
||||
RECONNECT_WRITER_RESUME_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn liveness_deadline_expired() {
|
||||
LIVENESS_DEADLINE_EXPIRED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn heartbeat_pong_dropped() {
|
||||
HEARTBEAT_PONG_DROPPED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn cancel_applied() {
|
||||
CANCEL_APPLIED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn cancel_pending_tombstoned() {
|
||||
CANCEL_PENDING_TOMBSTONED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn cancel_no_target() {
|
||||
CANCEL_NO_TARGET_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_rejected_overloaded() {
|
||||
TOOL_CALL_REJECTED_OVERLOADED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn inbox_full_request_rejected() {
|
||||
INBOX_FULL_REQUEST_REJECTED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn inbox_full_reject_send_failed() {
|
||||
INBOX_FULL_REJECT_SEND_FAILED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn inbox_full_notification_dropped() {
|
||||
INBOX_FULL_NOTIFICATION_DROPPED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn serve_replay_timeout() {
|
||||
SERVE_REPLAY_TIMEOUT_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn notif_lagged_recovered() {
|
||||
NOTIF_LAGGED_RECOVERED_TOTAL.inc();
|
||||
}
|
||||
|
||||
pub(crate) fn early_notif_buffered(frames: u64) {
|
||||
EARLY_NOTIF_BUFFERED_TOTAL.inc_by(frames);
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_inflight_inc(scope: &str) {
|
||||
TOOL_CALL_INFLIGHT.with_label_values(&[scope]).inc();
|
||||
}
|
||||
|
||||
pub(crate) fn tool_call_inflight_dec(scope: &str) {
|
||||
TOOL_CALL_INFLIGHT.with_label_values(&[scope]).dec();
|
||||
}
|
||||
|
||||
pub(crate) fn admission_wait_observe(secs: f64) {
|
||||
ADMISSION_WAIT_SECONDS.observe(secs);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "metrics"))]
|
||||
mod inner {
|
||||
pub(crate) fn pool_connections_inc() {}
|
||||
pub(crate) fn pool_connections_dec() {}
|
||||
pub(crate) fn pool_evictions_inc() {}
|
||||
pub(crate) fn reconnect_succeeded() {}
|
||||
pub(crate) fn reconnect_failed(_reason: &str) {}
|
||||
pub(crate) fn reconnect_duration_observe(_secs: f64) {}
|
||||
pub(crate) fn reconnect_cause(_cause: &str) {}
|
||||
pub(crate) fn reconnect_gap_observe(_secs: f64) {}
|
||||
pub(crate) fn call_dispatch_observe(_secs: f64) {}
|
||||
pub(crate) fn demux_inbox_depth_set(_depth: i64) {}
|
||||
pub(crate) fn call_id_collision() {}
|
||||
pub fn harness_connect(_status: &str, _sampler: &str) {}
|
||||
pub(crate) fn session_event(_event_type: &str) {}
|
||||
pub(crate) fn session_op_observe(_op: &str, _status: &str, _secs: f64) {}
|
||||
pub(crate) fn session_soft_rebind() {}
|
||||
pub(crate) fn no_handler() {}
|
||||
pub(crate) fn hook_send(_hook_type: &str) {}
|
||||
pub(crate) fn progress_frame_forwarded() {}
|
||||
pub(crate) fn cancel_hook_received() {}
|
||||
pub(crate) fn writer_sink_send_error() {}
|
||||
pub(crate) fn reconnect_writer_resume() {}
|
||||
pub(crate) fn liveness_deadline_expired() {}
|
||||
pub(crate) fn heartbeat_pong_dropped() {}
|
||||
pub(crate) fn cancel_applied() {}
|
||||
pub(crate) fn cancel_pending_tombstoned() {}
|
||||
pub(crate) fn cancel_no_target() {}
|
||||
pub(crate) fn tool_call_rejected_overloaded() {}
|
||||
pub(crate) fn inbox_full_request_rejected() {}
|
||||
pub(crate) fn inbox_full_reject_send_failed() {}
|
||||
pub(crate) fn inbox_full_notification_dropped() {}
|
||||
pub(crate) fn serve_replay_timeout() {}
|
||||
pub(crate) fn notif_lagged_recovered() {}
|
||||
pub(crate) fn early_notif_buffered(_frames: u64) {}
|
||||
pub(crate) fn tool_call_inflight_inc(_scope: &str) {}
|
||||
pub(crate) fn tool_call_inflight_dec(_scope: &str) {}
|
||||
pub(crate) fn admission_wait_observe(_secs: f64) {}
|
||||
}
|
||||
|
||||
pub(crate) use inner::admission_wait_observe;
|
||||
pub(crate) use inner::call_dispatch_observe;
|
||||
pub(crate) use inner::call_id_collision;
|
||||
pub(crate) use inner::cancel_applied;
|
||||
pub(crate) use inner::cancel_hook_received;
|
||||
pub(crate) use inner::cancel_no_target;
|
||||
pub(crate) use inner::cancel_pending_tombstoned;
|
||||
pub(crate) use inner::demux_inbox_depth_set;
|
||||
pub(crate) use inner::early_notif_buffered;
|
||||
pub(crate) use inner::heartbeat_pong_dropped;
|
||||
pub(crate) use inner::hook_send;
|
||||
pub(crate) use inner::inbox_full_notification_dropped;
|
||||
pub(crate) use inner::inbox_full_reject_send_failed;
|
||||
pub(crate) use inner::inbox_full_request_rejected;
|
||||
pub(crate) use inner::liveness_deadline_expired;
|
||||
pub(crate) use inner::no_handler;
|
||||
pub(crate) use inner::notif_lagged_recovered;
|
||||
pub(crate) use inner::pool_connections_dec;
|
||||
pub(crate) use inner::pool_connections_inc;
|
||||
pub(crate) use inner::pool_evictions_inc;
|
||||
pub(crate) use inner::progress_frame_forwarded;
|
||||
pub(crate) use inner::reconnect_cause;
|
||||
pub(crate) use inner::reconnect_duration_observe;
|
||||
pub(crate) use inner::reconnect_failed;
|
||||
pub(crate) use inner::reconnect_gap_observe;
|
||||
pub(crate) use inner::reconnect_succeeded;
|
||||
pub(crate) use inner::reconnect_writer_resume;
|
||||
pub(crate) use inner::serve_replay_timeout;
|
||||
pub(crate) use inner::session_event;
|
||||
pub(crate) use inner::session_op_observe;
|
||||
pub(crate) use inner::session_soft_rebind;
|
||||
pub(crate) use inner::tool_call_inflight_dec;
|
||||
pub(crate) use inner::tool_call_inflight_inc;
|
||||
pub(crate) use inner::tool_call_rejected_overloaded;
|
||||
pub(crate) use inner::writer_sink_send_error;
|
||||
|
||||
/// Record a harness connection attempt. Public so callers outside
|
||||
/// the SDK (e.g. `AgentBuilder::build_harness()` in the agentic sampler)
|
||||
/// can emit `status="fallback"` when the server connection fails and the
|
||||
/// builder falls back to a local-only harness.
|
||||
pub use inner::harness_connect;
|
||||
@@ -0,0 +1,362 @@
|
||||
//! Parsed server notification events.
|
||||
//!
|
||||
//! [`HubNotification`] is the typed representation of server-pushed
|
||||
//! notification frames that arrive on a session inbox. The
|
||||
//! [`HubNotification::parse`] constructor classifies a raw JSON value
|
||||
//! by its `method` field and deserializes the known shapes; anything
|
||||
//! unrecognised lands in [`HubNotification::Unknown`] so callers never
|
||||
//! lose data.
|
||||
|
||||
use kigi_tool_protocol::{
|
||||
SessionId, ToolId, ToolNotificationFrame, ToolServerStatusPayload, ToolsChanged,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use tracing::warn;
|
||||
|
||||
/// A typed server notification event parsed from a raw JSON-RPC notification frame.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum HubNotification {
|
||||
/// The active tool set for a session changed (tools added, removed, or updated).
|
||||
ToolsChanged {
|
||||
session_id: SessionId,
|
||||
added: Vec<ToolId>,
|
||||
removed: Vec<ToolId>,
|
||||
updated: Vec<ToolId>,
|
||||
},
|
||||
/// A tool notification forwarded by the server to all subscribers.
|
||||
ToolNotification {
|
||||
session_id: SessionId,
|
||||
frame: ToolNotificationFrame,
|
||||
},
|
||||
/// Tool server lifecycle status change, extracted from
|
||||
/// `__tool_server_status` / `status_changed` notification frames.
|
||||
ToolServerStatusChanged {
|
||||
session_id: SessionId,
|
||||
status: ToolServerStatusPayload,
|
||||
},
|
||||
/// A notification whose `method` is not recognised by this SDK version.
|
||||
Unknown { method: String, params: Value },
|
||||
}
|
||||
|
||||
impl HubNotification {
|
||||
/// Parse a raw JSON-RPC notification into a typed [`HubNotification`].
|
||||
///
|
||||
/// Returns `None` when the value lacks a `method` field (i.e. it is
|
||||
/// not a notification at all).
|
||||
pub fn parse(value: &Value) -> Option<Self> {
|
||||
let method = value.get("method")?.as_str()?;
|
||||
let params = value
|
||||
.get("params")
|
||||
.cloned()
|
||||
.unwrap_or(Value::Object(Default::default()));
|
||||
|
||||
match method {
|
||||
// `ToolsChanged` carries `session_id` inside `params`.
|
||||
"tools_changed" => match serde_json::from_value::<ToolsChanged>(params.clone()) {
|
||||
Ok(tc) => Some(HubNotification::ToolsChanged {
|
||||
session_id: tc.session_id,
|
||||
added: tc.added,
|
||||
removed: tc.removed,
|
||||
updated: tc.updated,
|
||||
}),
|
||||
Err(err) => {
|
||||
warn!(%err, "tools_changed params failed to deserialize; falling back to Unknown");
|
||||
Some(HubNotification::Unknown {
|
||||
method: method.to_owned(),
|
||||
params,
|
||||
})
|
||||
}
|
||||
},
|
||||
// `ToolNotificationFrame` has no `session_id`; use the envelope field.
|
||||
"tool.notification" => {
|
||||
let frame_result = serde_json::from_value::<ToolNotificationFrame>(params.clone());
|
||||
let session_id = value
|
||||
.get("session_id")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|s| SessionId::new(s).ok());
|
||||
match (frame_result, session_id) {
|
||||
(Ok(frame), Some(session_id)) => {
|
||||
if frame
|
||||
.tool_id
|
||||
.as_ref()
|
||||
.is_some_and(|id| id.as_str() == "__tool_server_status")
|
||||
&& let kigi_tool_protocol::notification_wire::WireToolNotification::Custom(ref c) = frame.notification
|
||||
&& c.kind == "status_changed"
|
||||
{
|
||||
match serde_json::from_value::<ToolServerStatusPayload>(
|
||||
c.payload.clone(),
|
||||
) {
|
||||
Ok(status) => {
|
||||
return Some(HubNotification::ToolServerStatusChanged {
|
||||
session_id,
|
||||
status,
|
||||
});
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(%err, "tool_server status payload failed to deserialize");
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(HubNotification::ToolNotification { session_id, frame })
|
||||
}
|
||||
(Err(err), _) => {
|
||||
warn!(%err, "tool.notification params failed to deserialize; falling back to Unknown");
|
||||
Some(HubNotification::Unknown {
|
||||
method: method.to_owned(),
|
||||
params,
|
||||
})
|
||||
}
|
||||
(_, None) => {
|
||||
warn!(
|
||||
"tool.notification missing or invalid session_id; falling back to Unknown"
|
||||
);
|
||||
Some(HubNotification::Unknown {
|
||||
method: method.to_owned(),
|
||||
params,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => Some(HubNotification::Unknown {
|
||||
method: method.to_owned(),
|
||||
params,
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn parse_tools_changed() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tools_changed",
|
||||
"params": {
|
||||
"session_id": "s1",
|
||||
"added": ["echo", "add"],
|
||||
"removed": [],
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
match notif {
|
||||
HubNotification::ToolsChanged {
|
||||
session_id,
|
||||
added,
|
||||
removed,
|
||||
updated,
|
||||
} => {
|
||||
assert_eq!(session_id.as_str(), "s1");
|
||||
assert_eq!(added.len(), 2);
|
||||
assert!(removed.is_empty());
|
||||
assert!(updated.is_empty());
|
||||
}
|
||||
other => panic!("expected ToolsChanged, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tools_changed_with_updated() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tools_changed",
|
||||
"params": {
|
||||
"session_id": "s1",
|
||||
"added": ["new_tool"],
|
||||
"removed": ["old_tool"],
|
||||
"updated": ["echo", "add"],
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
match notif {
|
||||
HubNotification::ToolsChanged {
|
||||
session_id,
|
||||
added,
|
||||
removed,
|
||||
updated,
|
||||
} => {
|
||||
assert_eq!(session_id.as_str(), "s1");
|
||||
assert_eq!(added.len(), 1);
|
||||
assert_eq!(removed.len(), 1);
|
||||
assert_eq!(updated.len(), 2);
|
||||
assert_eq!(updated[0].as_str(), "echo");
|
||||
assert_eq!(updated[1].as_str(), "add");
|
||||
}
|
||||
other => panic!("expected ToolsChanged, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_notification_custom() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tool.notification",
|
||||
"params": {
|
||||
"tool_id": "echo",
|
||||
"notification": {
|
||||
"shape": "custom",
|
||||
"value": {
|
||||
"kind": "echo.status",
|
||||
"payload": { "status": "idle" }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
match notif {
|
||||
HubNotification::ToolNotification { session_id, frame } => {
|
||||
assert_eq!(session_id.as_str(), "s1");
|
||||
assert_eq!(frame.tool_id.as_ref().unwrap().as_str(), "echo");
|
||||
}
|
||||
other => panic!("expected ToolNotification, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_notification_missing_session_id_falls_back_to_unknown() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "tool.notification",
|
||||
"params": {
|
||||
"tool_id": "echo",
|
||||
"notification": {
|
||||
"shape": "custom",
|
||||
"value": { "kind": "test", "payload": {} }
|
||||
}
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
|
||||
assert!(
|
||||
matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"),
|
||||
"tool.notification without envelope session_id should fall back to Unknown, got {notif:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_unknown_method() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "future.method",
|
||||
"params": { "key": "value" }
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
match notif {
|
||||
HubNotification::Unknown { method, params } => {
|
||||
assert_eq!(method, "future.method");
|
||||
assert_eq!(params["key"], "value");
|
||||
}
|
||||
other => panic!("expected Unknown, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_missing_method_returns_none() {
|
||||
let value = json!({ "jsonrpc": "2.0", "id": "123", "result": {} });
|
||||
assert!(HubNotification::parse(&value).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tools_changed_bad_params_falls_back_to_unknown() {
|
||||
// `params` has wrong shape (missing required fields) — should fall
|
||||
// back to Unknown instead of returning None and dropping the event.
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "tools_changed",
|
||||
"params": { "unexpected_field": true }
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
|
||||
assert!(
|
||||
matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tools_changed"),
|
||||
"malformed tools_changed should fall back to Unknown, got {notif:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_server_status_changed() {
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tool.notification",
|
||||
"params": {
|
||||
"tool_id": "__tool_server_status",
|
||||
"notification": {
|
||||
"shape": "custom",
|
||||
"value": {
|
||||
"kind": "status_changed",
|
||||
"payload": {
|
||||
"status": "busy",
|
||||
"active_tool_calls": 2,
|
||||
"active_tool_names": ["read_file", "grep"],
|
||||
"background_tasks": 0,
|
||||
"pending_tool_calls": 0,
|
||||
"last_tool_call_started_ms": 100,
|
||||
"last_tool_call_completed_ms": 0,
|
||||
"uptime_ms": 5000,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
match notif {
|
||||
HubNotification::ToolServerStatusChanged { session_id, status } => {
|
||||
assert_eq!(session_id.as_str(), "s1");
|
||||
assert_eq!(
|
||||
status.status,
|
||||
kigi_tool_protocol::ToolServerLifecycleStatus::Busy
|
||||
);
|
||||
assert_eq!(status.active_tool_calls, 2);
|
||||
}
|
||||
other => panic!("expected ToolServerStatusChanged, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_server_status_non_status_tool_id_stays_generic() {
|
||||
// A tool.notification with a different tool_id should remain
|
||||
// as ToolNotification, not be intercepted.
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tool.notification",
|
||||
"params": {
|
||||
"tool_id": "some_other_tool",
|
||||
"notification": {
|
||||
"shape": "custom",
|
||||
"value": {
|
||||
"kind": "status_changed",
|
||||
"payload": { "status": "ready" }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse");
|
||||
assert!(
|
||||
matches!(notif, HubNotification::ToolNotification { .. }),
|
||||
"non-__tool_server_status tool_id should stay as ToolNotification, got {notif:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_tool_notification_bad_params_falls_back_to_unknown() {
|
||||
// `params` has wrong shape — should fall back to Unknown.
|
||||
let value = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"session_id": "s1",
|
||||
"method": "tool.notification",
|
||||
"params": { "not_a_valid_frame": true }
|
||||
});
|
||||
let notif = HubNotification::parse(&value).expect("should parse as Unknown, not None");
|
||||
assert!(
|
||||
matches!(notif, HubNotification::Unknown { ref method, .. } if method == "tool.notification"),
|
||||
"malformed tool.notification should fall back to Unknown, got {notif:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
//! Server-side session event emitter.
|
||||
//!
|
||||
//! [`ObservabilityBridge`] is a thin facade for emitting session-level
|
||||
//! events (turn lifecycle, phase changes) to the connected server. Tool-call events
|
||||
//! (`ToolCallStarted` / `ToolCallCompleted`) are emitted automatically
|
||||
//! by [`crate::harness::ToolHarness::call`] and do not need the bridge.
|
||||
//!
|
||||
//! The caller is responsible for also emitting to the local sink
|
||||
//! (`EventTracker` in the shell, `EventProcPublisher` in the
|
||||
//! chat service) — the bridge handles only the server leg.
|
||||
//!
|
||||
//! This separation is deliberate: each sampler's local sink has a
|
||||
//! different type and API surface. Forcing a trait/callback into the
|
||||
//! bridge would add abstraction overhead without benefit, since the
|
||||
//! call sites already have the local sink in scope.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_tool_protocol::{SessionId, session_event::SessionEvent};
|
||||
|
||||
use crate::harness::ToolHarness;
|
||||
|
||||
/// Emits [`SessionEvent`]s to the connected server as `ToolNotificationFrame` custom
|
||||
/// notifications with `kind = "session_event"`.
|
||||
///
|
||||
/// No-ops gracefully when no harness is present (i.e. `harness` is
|
||||
/// `None`). Server notification failures are silently ignored — the bridge
|
||||
/// is fire-and-forget so server issues never affect the sampler's main loop.
|
||||
///
|
||||
/// Callers MUST also emit to their local sink separately:
|
||||
/// - Shell: `self.events.emit(Event::...)`
|
||||
/// - Chat service: `publisher.publish_agent_event(...)`
|
||||
pub struct ObservabilityBridge {
|
||||
harness: Option<Arc<ToolHarness>>,
|
||||
/// Retained for future payload enrichment and logging.
|
||||
session_id: SessionId,
|
||||
}
|
||||
|
||||
impl ObservabilityBridge {
|
||||
pub fn new(harness: Option<Arc<ToolHarness>>, session_id: SessionId) -> Self {
|
||||
Self {
|
||||
harness,
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// The session id this bridge was created for.
|
||||
pub fn session_id(&self) -> &SessionId {
|
||||
&self.session_id
|
||||
}
|
||||
|
||||
/// Whether a harness is present (i.e. server emission is active).
|
||||
pub fn has_harness(&self) -> bool {
|
||||
self.harness.is_some()
|
||||
}
|
||||
|
||||
/// Emit a session event to the connected server. No-ops if no harness is present.
|
||||
///
|
||||
/// Delegates frame construction + wire dispatch to
|
||||
/// [`ToolHarness::emit_session_event`] so the SDK keeps a single
|
||||
/// canonical encoding path.
|
||||
///
|
||||
/// Callers MUST also emit to their local sink separately:
|
||||
/// - Shell: `self.events.emit(Event::...)`
|
||||
/// - Chat service: `publisher.publish_agent_event(...)`
|
||||
pub async fn emit(&self, event: SessionEvent) {
|
||||
let event_type = match &event {
|
||||
SessionEvent::TurnStarted { .. } => "turn_started",
|
||||
SessionEvent::TurnEnded { .. } => "turn_ended",
|
||||
SessionEvent::ToolCallStarted { .. } => "tool_call_started",
|
||||
SessionEvent::ToolCallCompleted { .. } => "tool_call_completed",
|
||||
SessionEvent::PhaseChanged { .. } => "phase_changed",
|
||||
SessionEvent::Unknown => "unknown",
|
||||
};
|
||||
crate::metrics::session_event(event_type);
|
||||
if let Some(harness) = &self.harness {
|
||||
harness.emit_session_event(event).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_tool_protocol::session_event::{SessionEvent, SessionPhase, ToolCallOutcome};
|
||||
use kigi_tool_protocol::turn_hook::TurnHookOutcome;
|
||||
|
||||
fn test_session_id() -> SessionId {
|
||||
SessionId::new("test-obs-session").expect("valid")
|
||||
}
|
||||
|
||||
// ── No-harness path ─────────────────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_without_harness_is_noop() {
|
||||
let bridge = ObservabilityBridge::new(None, test_session_id());
|
||||
// Must not panic and should return immediately.
|
||||
bridge
|
||||
.emit(SessionEvent::TurnStarted {
|
||||
turn_number: 1,
|
||||
model_id: "grok-3".into(),
|
||||
yolo_mode: false,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn has_harness_returns_false_when_none() {
|
||||
let bridge = ObservabilityBridge::new(None, test_session_id());
|
||||
assert!(!bridge.has_harness());
|
||||
}
|
||||
|
||||
// ── Constructor field storage ───────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn new_stores_session_id() {
|
||||
let sid = test_session_id();
|
||||
let bridge = ObservabilityBridge::new(None, sid.clone());
|
||||
assert_eq!(bridge.session_id(), &sid);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn has_harness_returns_true_when_present() {
|
||||
let harness = ToolHarness::local_only_with(
|
||||
crate::harness::LocalRegistry::new(),
|
||||
test_session_id(),
|
||||
Default::default(),
|
||||
);
|
||||
let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
|
||||
assert!(bridge.has_harness());
|
||||
}
|
||||
|
||||
// ── Serialization correctness ───────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn session_event_serializes_to_expected_json() {
|
||||
let event = SessionEvent::TurnStarted {
|
||||
turn_number: 1,
|
||||
model_id: "grok-3".into(),
|
||||
yolo_mode: true,
|
||||
};
|
||||
let value = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(value["event_type"], "turn_started");
|
||||
assert_eq!(value["turn_number"], 1);
|
||||
assert_eq!(value["model_id"], "grok-3");
|
||||
assert_eq!(value["yolo_mode"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_event_turn_ended_serializes_correctly() {
|
||||
let event = SessionEvent::TurnEnded {
|
||||
turn_number: 5,
|
||||
outcome: TurnHookOutcome::Completed,
|
||||
duration_ms: 3200,
|
||||
tool_call_count: 12,
|
||||
model_id: "grok-3".into(),
|
||||
};
|
||||
let value = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(value["event_type"], "turn_ended");
|
||||
assert_eq!(value["outcome"], "completed");
|
||||
assert_eq!(value["tool_call_count"], 12);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_event_tool_call_completed_serializes_correctly() {
|
||||
let event = SessionEvent::ToolCallCompleted {
|
||||
tool_call_id: "call-1".into(),
|
||||
tool_name: "bash".into(),
|
||||
duration_ms: 500,
|
||||
outcome: ToolCallOutcome::Success,
|
||||
};
|
||||
let value = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(value["event_type"], "tool_call_completed");
|
||||
assert_eq!(value["outcome"], "success");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_event_phase_changed_serializes_correctly() {
|
||||
let event = SessionEvent::PhaseChanged {
|
||||
phase: SessionPhase::Sampling,
|
||||
};
|
||||
let value = serde_json::to_value(&event).unwrap();
|
||||
assert_eq!(value["event_type"], "phase_changed");
|
||||
assert_eq!(value["phase"], "sampling");
|
||||
}
|
||||
|
||||
// Frame-construction invariants moved to `crate::harness` where the
|
||||
// builder now lives (`ToolHarness::emit_session_event`).
|
||||
|
||||
// ── End-to-end with local-only harness ───────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_with_local_only_harness_does_not_panic() {
|
||||
// A local-only harness has no server connection, so
|
||||
// `send_notification` returns `Err` — but the bridge ignores
|
||||
// errors, so this must succeed silently.
|
||||
let harness = ToolHarness::local_only_with(
|
||||
crate::harness::LocalRegistry::new(),
|
||||
test_session_id(),
|
||||
Default::default(),
|
||||
);
|
||||
let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
|
||||
bridge
|
||||
.emit(SessionEvent::PhaseChanged {
|
||||
phase: SessionPhase::ToolExecution,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn emit_all_event_variants_does_not_panic() {
|
||||
let harness = ToolHarness::local_only_with(
|
||||
crate::harness::LocalRegistry::new(),
|
||||
test_session_id(),
|
||||
Default::default(),
|
||||
);
|
||||
let bridge = ObservabilityBridge::new(Some(Arc::new(harness)), test_session_id());
|
||||
|
||||
// Smoke test: every variant emits without panic through a
|
||||
// local-only harness. Includes error/cancelled outcomes to
|
||||
// cover non-happy-path enum values.
|
||||
let events = vec![
|
||||
SessionEvent::TurnStarted {
|
||||
turn_number: 1,
|
||||
model_id: "grok-3".into(),
|
||||
yolo_mode: false,
|
||||
},
|
||||
SessionEvent::ToolCallStarted {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "bash".into(),
|
||||
turn_number: 1,
|
||||
},
|
||||
SessionEvent::ToolCallCompleted {
|
||||
tool_call_id: "c1".into(),
|
||||
tool_name: "bash".into(),
|
||||
duration_ms: 100,
|
||||
outcome: ToolCallOutcome::Success,
|
||||
},
|
||||
SessionEvent::ToolCallCompleted {
|
||||
tool_call_id: "c2".into(),
|
||||
tool_name: "read_file".into(),
|
||||
duration_ms: 50,
|
||||
outcome: ToolCallOutcome::Error,
|
||||
},
|
||||
SessionEvent::ToolCallCompleted {
|
||||
tool_call_id: "c3".into(),
|
||||
tool_name: "grep".into(),
|
||||
duration_ms: 10,
|
||||
outcome: ToolCallOutcome::Cancelled,
|
||||
},
|
||||
SessionEvent::PhaseChanged {
|
||||
phase: SessionPhase::Idle,
|
||||
},
|
||||
SessionEvent::TurnEnded {
|
||||
turn_number: 1,
|
||||
outcome: TurnHookOutcome::Completed,
|
||||
duration_ms: 500,
|
||||
tool_call_count: 3,
|
||||
model_id: "grok-3".into(),
|
||||
},
|
||||
SessionEvent::TurnEnded {
|
||||
turn_number: 2,
|
||||
outcome: TurnHookOutcome::Error,
|
||||
duration_ms: 100,
|
||||
tool_call_count: 0,
|
||||
model_id: "grok-3".into(),
|
||||
},
|
||||
SessionEvent::Unknown,
|
||||
];
|
||||
|
||||
for event in events {
|
||||
bridge.emit(event).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
//! [`AuthProvider`] that refreshes OIDC tokens before they expire.
|
||||
//!
|
||||
//! `current()` checks token expiry and, if needed, performs OIDC
|
||||
//! discovery + token exchange before returning the credential.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use parking_lot::Mutex;
|
||||
|
||||
use crate::auth::{AuthCredential, AuthIdentity, AuthProvider};
|
||||
|
||||
pub type OnRefreshCallback = Arc<dyn Fn(&RefreshEvent) + Send + Sync>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RefreshEvent {
|
||||
pub access_token: String,
|
||||
pub new_refresh_token: Option<String>,
|
||||
pub expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
struct TokenState {
|
||||
access_token: String,
|
||||
refresh_token: String,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
pub struct OidcAuthProvider {
|
||||
state: Mutex<TokenState>,
|
||||
issuer: String,
|
||||
client_id: String,
|
||||
user_id: Option<String>,
|
||||
principal_type: Option<String>,
|
||||
principal_id: Option<String>,
|
||||
on_refresh: Option<OnRefreshCallback>,
|
||||
}
|
||||
|
||||
const REFRESH_MARGIN: Duration = Duration::from_secs(60);
|
||||
|
||||
impl std::fmt::Debug for OidcAuthProvider {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("OidcAuthProvider")
|
||||
.field("issuer", &self.issuer)
|
||||
.field("client_id", &self.client_id)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OidcAuthProviderBuilder {
|
||||
access_token: String,
|
||||
refresh_token: String,
|
||||
issuer: String,
|
||||
client_id: String,
|
||||
expires_at: Option<DateTime<Utc>>,
|
||||
user_id: Option<String>,
|
||||
principal_type: Option<String>,
|
||||
principal_id: Option<String>,
|
||||
on_refresh: Option<OnRefreshCallback>,
|
||||
}
|
||||
|
||||
impl OidcAuthProviderBuilder {
|
||||
pub fn new(
|
||||
access_token: impl Into<String>,
|
||||
refresh_token: impl Into<String>,
|
||||
issuer: impl Into<String>,
|
||||
client_id: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
access_token: access_token.into(),
|
||||
refresh_token: refresh_token.into(),
|
||||
issuer: issuer.into(),
|
||||
client_id: client_id.into(),
|
||||
expires_at: None,
|
||||
user_id: None,
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
on_refresh: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn expires_at(mut self, expires_at: DateTime<Utc>) -> Self {
|
||||
self.expires_at = Some(expires_at);
|
||||
self
|
||||
}
|
||||
|
||||
/// Owner user id parsed from the auth source, surfaced via
|
||||
/// [`AuthProvider::identity`].
|
||||
pub fn user_id(mut self, user_id: impl Into<String>) -> Self {
|
||||
self.user_id = Some(user_id.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn principal_type(mut self, pt: impl Into<String>) -> Self {
|
||||
self.principal_type = Some(pt.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn principal_id(mut self, pid: impl Into<String>) -> Self {
|
||||
self.principal_id = Some(pid.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn on_refresh(mut self, cb: OnRefreshCallback) -> Self {
|
||||
self.on_refresh = Some(cb);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> OidcAuthProvider {
|
||||
OidcAuthProvider {
|
||||
state: Mutex::new(TokenState {
|
||||
access_token: self.access_token,
|
||||
refresh_token: self.refresh_token,
|
||||
expires_at: self.expires_at,
|
||||
}),
|
||||
issuer: self.issuer,
|
||||
client_id: self.client_id,
|
||||
user_id: self.user_id,
|
||||
principal_type: self.principal_type,
|
||||
principal_id: self.principal_id,
|
||||
on_refresh: self.on_refresh,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthProvider for OidcAuthProvider {
|
||||
fn current(&self) -> AuthCredential {
|
||||
let expired = {
|
||||
let s = self.state.lock();
|
||||
s.expires_at.is_some_and(|exp| {
|
||||
Utc::now() + chrono::Duration::from_std(REFRESH_MARGIN).unwrap() >= exp
|
||||
})
|
||||
};
|
||||
if expired && let Err(e) = self.try_refresh() {
|
||||
tracing::warn!(error = %e, "OIDC refresh failed, using stale token");
|
||||
}
|
||||
let s = self.state.lock();
|
||||
AuthCredential::bearer(&s.access_token)
|
||||
}
|
||||
|
||||
/// Surface the principal fields parsed from the auth source. `None` only
|
||||
/// when no `user_id` was supplied (nothing to attribute).
|
||||
fn identity(&self) -> Option<AuthIdentity> {
|
||||
let user_id = self.user_id.clone()?;
|
||||
Some(AuthIdentity {
|
||||
user_id,
|
||||
principal_type: self.principal_type.clone(),
|
||||
principal_id: self.principal_id.clone(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl OidcAuthProvider {
|
||||
fn try_refresh(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
tracing::info!(issuer = %self.issuer, "refreshing OIDC token");
|
||||
if let Ok(handle) = tokio::runtime::Handle::try_current() {
|
||||
tokio::task::block_in_place(|| handle.block_on(self.do_refresh()))
|
||||
} else {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(self.do_refresh())
|
||||
}
|
||||
}
|
||||
|
||||
async fn do_refresh(&self) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let refresh_token = self.state.lock().refresh_token.clone();
|
||||
let issuer = self.issuer.trim_end_matches('/');
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Discovery {
|
||||
token_endpoint: String,
|
||||
}
|
||||
|
||||
let disc: Discovery = client
|
||||
.get(format!("{issuer}/.well-known/openid-configuration"))
|
||||
.timeout(Duration::from_secs(10))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.json()
|
||||
.await?;
|
||||
|
||||
let mut params = vec![
|
||||
("grant_type", "refresh_token"),
|
||||
("refresh_token", refresh_token.as_str()),
|
||||
("client_id", self.client_id.as_str()),
|
||||
];
|
||||
let pt = self.principal_type.clone();
|
||||
let pid = self.principal_id.clone();
|
||||
if let Some(ref v) = pt {
|
||||
params.push(("principal_type", v));
|
||||
}
|
||||
if let Some(ref v) = pid {
|
||||
params.push(("principal_id", v));
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Tokens {
|
||||
access_token: String,
|
||||
#[serde(default)]
|
||||
refresh_token: Option<String>,
|
||||
#[serde(default)]
|
||||
expires_in: Option<u64>,
|
||||
}
|
||||
|
||||
let tokens: Tokens = client
|
||||
.post(&disc.token_endpoint)
|
||||
.form(¶ms)
|
||||
.timeout(Duration::from_secs(15))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.json()
|
||||
.await?;
|
||||
|
||||
let expires_at = tokens
|
||||
.expires_in
|
||||
.map(|s| Utc::now() + chrono::Duration::seconds(s as i64));
|
||||
|
||||
tracing::info!(expires_at = ?expires_at, "OIDC token refreshed");
|
||||
|
||||
if let Some(ref cb) = self.on_refresh {
|
||||
cb(&RefreshEvent {
|
||||
access_token: tokens.access_token.clone(),
|
||||
new_refresh_token: tokens.refresh_token.clone(),
|
||||
expires_at,
|
||||
});
|
||||
}
|
||||
|
||||
let mut s = self.state.lock();
|
||||
s.access_token = tokens.access_token;
|
||||
if let Some(rt) = tokens.refresh_token {
|
||||
s.refresh_token = rt;
|
||||
}
|
||||
s.expires_at = expires_at;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn current_returns_token_when_not_expired() {
|
||||
let provider = OidcAuthProviderBuilder::new(
|
||||
"access-tok",
|
||||
"refresh-tok",
|
||||
"https://auth.example.com",
|
||||
"client1",
|
||||
)
|
||||
.expires_at(Utc::now() + chrono::Duration::hours(1))
|
||||
.build();
|
||||
|
||||
let cred = provider.current();
|
||||
match cred {
|
||||
AuthCredential::Bearer { token } => {
|
||||
assert_eq!(token, "access-tok");
|
||||
}
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_returns_token_when_no_expiry() {
|
||||
let provider = OidcAuthProviderBuilder::new(
|
||||
"no-expiry-tok",
|
||||
"refresh-tok",
|
||||
"https://auth.example.com",
|
||||
"client1",
|
||||
)
|
||||
.build();
|
||||
|
||||
let cred = provider.current();
|
||||
match cred {
|
||||
AuthCredential::Bearer { token } => assert_eq!(token, "no-expiry-tok"),
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_returns_stale_token_when_refresh_fails() {
|
||||
// Expired token, but issuer is unreachable — should return stale
|
||||
let provider = OidcAuthProviderBuilder::new(
|
||||
"stale-tok",
|
||||
"refresh-tok",
|
||||
"https://localhost:1", // unreachable
|
||||
"client1",
|
||||
)
|
||||
.expires_at(Utc::now() - chrono::Duration::hours(1))
|
||||
.build();
|
||||
|
||||
let cred = provider.current();
|
||||
match cred {
|
||||
AuthCredential::Bearer { token } => assert_eq!(token, "stale-tok"),
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_surfaces_principal_fields() {
|
||||
let provider = OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1")
|
||||
.user_id("user-1")
|
||||
.principal_type("Team")
|
||||
.principal_id("team-9")
|
||||
.build();
|
||||
let id = provider.identity().expect("identity present");
|
||||
assert_eq!(id.user_id, "user-1");
|
||||
assert_eq!(id.principal_type.as_deref(), Some("Team"));
|
||||
assert_eq!(id.principal_id.as_deref(), Some("team-9"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_none_without_user_id() {
|
||||
let provider =
|
||||
OidcAuthProviderBuilder::new("tok", "rt", "https://auth.example.com", "c1").build();
|
||||
assert!(provider.identity().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn debug_does_not_leak_tokens() {
|
||||
let provider = OidcAuthProviderBuilder::new(
|
||||
"secret-access-token",
|
||||
"secret-refresh-token",
|
||||
"https://auth.example.com",
|
||||
"client1",
|
||||
)
|
||||
.build();
|
||||
|
||||
let debug = format!("{provider:?}");
|
||||
assert!(!debug.contains("secret-access-token"));
|
||||
assert!(!debug.contains("secret-refresh-token"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
//! Process-wide connection pool keyed by `(url, principal)`.
|
||||
//!
|
||||
//! Two [`crate::ToolServer`] builds with the same `(url, credential)`
|
||||
//! observe the same `Arc<HubConnection>`; distinct credentials open
|
||||
//! distinct sockets. The pool is the canonical entry point — direct
|
||||
//! [`crate::HubConnection::connect`] calls are reserved for tests and
|
||||
//! one-shot programs that explicitly want unpooled behaviour.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use dashmap::DashMap;
|
||||
use kigi_tool_protocol::ConnectionKind;
|
||||
use tokio::sync::OnceCell;
|
||||
use tokio::task::JoinHandle;
|
||||
use url::Url;
|
||||
|
||||
use crate::auth::AuthProvider;
|
||||
use crate::connection::{
|
||||
ConnKey, ConnectCallback, ConnectionConfig, ConnectionTuning, DisconnectCallback,
|
||||
HubConnection, ReconnectCallback,
|
||||
};
|
||||
use crate::error::ClientError;
|
||||
|
||||
/// Idle window for the reaper: a pooled connection is evictable once it is
|
||||
/// unused (`Arc::strong_count == 1`, i.e. only the pool holds it) **and**
|
||||
/// `now - last_handout >= DEFAULT_POOL_IDLE_TTL`.
|
||||
///
|
||||
/// Note the clock is `last_handout` (the last time the pool returned the
|
||||
/// connection), not the moment the last consumer `Arc` was dropped: a
|
||||
/// connection held longer than the TTL and then released is eligible on the
|
||||
/// very next sweep, with no extra post-drop grace period. The only hard
|
||||
/// guarantee is that an in-use connection (`strong_count > 1`) is never
|
||||
/// reaped. Tuned well above the server's own 90s dead-peer idle timeout so a
|
||||
/// short borrow between turns of an active conversation isn't churned.
|
||||
pub const DEFAULT_POOL_IDLE_TTL: Duration = Duration::from_secs(300);
|
||||
|
||||
/// How often the shared pool's idle reaper scans for evictable entries.
|
||||
pub const DEFAULT_POOL_SWEEP_INTERVAL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// A pooled connection plus the last time it was handed out to a caller.
|
||||
///
|
||||
/// `last_handout` is refreshed on every [`HubConnectionPool::get_or_connect`]
|
||||
/// hit (and on the initial insert), so a connection that is repeatedly
|
||||
/// re-fetched never looks idle even if its [`Arc`] strong count briefly
|
||||
/// returns to 1 between fetches. Eviction additionally requires
|
||||
/// `Arc::strong_count == 1` (only the pool holds it), so a connection a
|
||||
/// consumer still holds is never reaped regardless of `last_handout`.
|
||||
struct Pooled {
|
||||
conn: Arc<HubConnection>,
|
||||
last_handout: Instant,
|
||||
}
|
||||
|
||||
/// The process-global pool used by [`HubConnectionPool::shared`].
|
||||
///
|
||||
/// `tokio::sync::OnceCell` is preferred over `std::sync::OnceLock` /
|
||||
/// `LazyLock` here because the pool is only ever observed from
|
||||
/// async contexts (the connection actor lives on a tokio runtime
|
||||
/// already), so the async-aware `get_or_init` semantics avoid the
|
||||
/// blocking-init footgun of the sync alternatives without taking a
|
||||
/// hard dependency on additional sync primitives.
|
||||
///
|
||||
/// Tests MUST use [`HubConnectionPool::new`] to avoid cross-test
|
||||
/// pollution: cargo runs all integration tests in the same binary
|
||||
/// unless otherwise configured, so any test that touches
|
||||
/// `HubConnectionPool::shared()` leaves the pool populated for
|
||||
/// subsequent tests.
|
||||
static SHARED: OnceCell<Arc<HubConnectionPool>> = OnceCell::const_new();
|
||||
|
||||
/// Pool of live server connections.
|
||||
pub struct HubConnectionPool {
|
||||
connections: DashMap<ConnKey, Pooled>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HubConnectionPool {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("HubConnectionPool")
|
||||
.field("connection_count", &self.connections.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl HubConnectionPool {
|
||||
/// Build a fresh, unshared pool. Tests typically use this so each
|
||||
/// test sees an isolated registry.
|
||||
pub fn new() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
connections: DashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Return the process-wide shared pool, lazily initialising it on
|
||||
/// the first call. Subsequent callers in the same process observe
|
||||
/// the same `Arc`.
|
||||
///
|
||||
/// The shared pool spawns an idle reaper (see [`Self::spawn_idle_reaper`])
|
||||
/// exactly once, so a connection that is unused (`strong_count == 1`) and
|
||||
/// has not been handed out for [`DEFAULT_POOL_IDLE_TTL`] is closed instead
|
||||
/// of living for the whole process lifetime. (Unpooled / test pools built
|
||||
/// via [`Self::new`] do not
|
||||
/// get a reaper; they can call [`Self::sweep_idle`] directly.)
|
||||
pub async fn shared() -> Arc<Self> {
|
||||
SHARED
|
||||
.get_or_init(|| async {
|
||||
let pool = Self::new();
|
||||
pool.spawn_idle_reaper(DEFAULT_POOL_IDLE_TTL, DEFAULT_POOL_SWEEP_INTERVAL);
|
||||
pool
|
||||
})
|
||||
.await
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Look up an existing pooled connection for `(url, credential)`,
|
||||
/// or open a fresh one if no pooled entry exists.
|
||||
///
|
||||
/// `kind` is the connection role announced in the hello frame. The
|
||||
/// pool is keyed by `(url, principal)` only; mixing
|
||||
/// [`ConnectionKind`] values for the same `(url, principal)` is a
|
||||
/// caller error and surfaces as a [`ClientError::InvalidConfig`].
|
||||
///
|
||||
/// The optional extra access key is not part of the pool key, so the first
|
||||
/// caller's key is the one carried on a shared connection's handshake (in
|
||||
/// practice it is a per-deployment constant). The plaintext-scheme guard is
|
||||
/// re-checked on every call below so it can't be bypassed by a cached
|
||||
/// insecure entry.
|
||||
pub async fn get_or_connect(
|
||||
self: &Arc<Self>,
|
||||
url: Url,
|
||||
credential: Arc<dyn AuthProvider>,
|
||||
kind: ConnectionKind,
|
||||
on_reconnect: Option<Arc<ReconnectCallback>>,
|
||||
on_disconnect: Option<Arc<DisconnectCallback>>,
|
||||
server_id: Option<kigi_tool_protocol::ServerId>,
|
||||
alpha_test_key: Option<String>,
|
||||
allow_insecure_ws: bool,
|
||||
) -> Result<Arc<HubConnection>, ClientError> {
|
||||
self.get_or_connect_tuned(
|
||||
url,
|
||||
credential,
|
||||
kind,
|
||||
on_reconnect,
|
||||
on_disconnect,
|
||||
None, // on_connect (unused by the simple wrapper)
|
||||
server_id,
|
||||
None,
|
||||
None,
|
||||
alpha_test_key,
|
||||
allow_insecure_ws,
|
||||
ConnectionTuning::default(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Like [`Self::get_or_connect`] but carries optional connection-tuning
|
||||
/// knobs ([`ConnectionTuning`]) onto a freshly-opened connection. A
|
||||
/// `Default` tuning is behaviourally identical to `get_or_connect`, so
|
||||
/// existing callers are unaffected.
|
||||
///
|
||||
/// Tuning binds to the socket at open time: it takes effect only when
|
||||
/// THIS call opens the connection. Because the pool dedups by
|
||||
/// `(url, principal)`, a hit on an existing entry returns that
|
||||
/// connection as-is and the `tuning` argument is ignored — the first
|
||||
/// opener's ping/backoff settings win for the lifetime of the pooled
|
||||
/// connection. Callers that need distinct tuning must use a distinct
|
||||
/// `(url, principal)` or an unpooled [`HubConnection::connect`].
|
||||
pub(crate) async fn get_or_connect_tuned(
|
||||
self: &Arc<Self>,
|
||||
url: Url,
|
||||
credential: Arc<dyn AuthProvider>,
|
||||
kind: ConnectionKind,
|
||||
on_reconnect: Option<Arc<ReconnectCallback>>,
|
||||
on_disconnect: Option<Arc<DisconnectCallback>>,
|
||||
on_connect: Option<Arc<ConnectCallback>>,
|
||||
server_id: Option<kigi_tool_protocol::ServerId>,
|
||||
server_description: Option<String>,
|
||||
server_metadata: Option<serde_json::Value>,
|
||||
alpha_test_key: Option<String>,
|
||||
allow_insecure_ws: bool,
|
||||
tuning: ConnectionTuning,
|
||||
) -> Result<Arc<HubConnection>, ClientError> {
|
||||
if url.scheme() != "wss" && !crate::connection::host_is_loopback(&url) && !allow_insecure_ws
|
||||
{
|
||||
return Err(ClientError::InsecureScheme { url });
|
||||
}
|
||||
let key = ConnKey {
|
||||
url: url.as_str().to_owned(),
|
||||
principal: credential.principal_key(),
|
||||
};
|
||||
if let Some(mut existing) = self.connections.get_mut(&key) {
|
||||
existing.last_handout = Instant::now();
|
||||
let conn = existing.conn.clone();
|
||||
drop(existing);
|
||||
if conn.kind() != kind {
|
||||
return Err(ClientError::InvalidConfig(format!(
|
||||
"pool entry for {} bound to {:?}; rebuild requested {:?}",
|
||||
key.url,
|
||||
conn.kind(),
|
||||
kind
|
||||
)));
|
||||
}
|
||||
return Ok(conn);
|
||||
}
|
||||
let config = ConnectionConfig {
|
||||
url,
|
||||
credential,
|
||||
kind,
|
||||
on_reconnect,
|
||||
on_disconnect,
|
||||
on_connect,
|
||||
server_id,
|
||||
server_description,
|
||||
server_metadata,
|
||||
outbound_buffer: None,
|
||||
tuning,
|
||||
alpha_test_key,
|
||||
allow_insecure_ws,
|
||||
on_fatal: Some(Arc::downgrade(self)),
|
||||
};
|
||||
let conn = HubConnection::connect(config).await?;
|
||||
// Race window: another caller may have inserted between our
|
||||
// `get` and `connect`. Resolve via `entry().or_insert_with`
|
||||
// semantics — if we lose the race we drop our fresh
|
||||
// connection and adopt the winning one.
|
||||
match self.connections.entry(key.clone()) {
|
||||
dashmap::Entry::Occupied(mut existing) => {
|
||||
existing.get_mut().last_handout = Instant::now();
|
||||
let winner = existing.get().conn.clone();
|
||||
drop(conn);
|
||||
if winner.kind() != kind {
|
||||
return Err(ClientError::InvalidConfig(format!(
|
||||
"pool entry for {} bound to {:?}; rebuild requested {:?}",
|
||||
key.url,
|
||||
winner.kind(),
|
||||
kind
|
||||
)));
|
||||
}
|
||||
Ok(winner)
|
||||
}
|
||||
dashmap::Entry::Vacant(slot) => {
|
||||
crate::metrics::pool_connections_inc();
|
||||
slot.insert(Pooled {
|
||||
conn: conn.clone(),
|
||||
last_handout: Instant::now(),
|
||||
});
|
||||
Ok(conn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of pooled connections. Intended for tests and metrics.
|
||||
pub fn len(&self) -> usize {
|
||||
self.connections.len()
|
||||
}
|
||||
|
||||
/// `true` when no connection is pooled.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.connections.is_empty()
|
||||
}
|
||||
|
||||
/// Forget the pooled connection for `key`. The actual underlying
|
||||
/// `Arc<HubConnection>` is dropped only when no other holder
|
||||
/// keeps a reference; the next [`Self::get_or_connect`] for the
|
||||
/// same key opens a fresh socket.
|
||||
pub fn forget(&self, key: &ConnKey) {
|
||||
if self.connections.remove(key).is_some() {
|
||||
crate::metrics::pool_connections_dec();
|
||||
}
|
||||
}
|
||||
|
||||
/// Close and remove every pooled connection that is BOTH unused (no live
|
||||
/// consumer holds an `Arc` — only the pool does, so `strong_count == 1`)
|
||||
/// AND idle longer than `idle_ttl` (no hand-out within the window).
|
||||
/// Removing the entry drops the pool's last `Arc<HubConnection>`, whose
|
||||
/// `Drop` closes the socket.
|
||||
///
|
||||
/// The strong-count check runs inside the map's per-shard lock (via
|
||||
/// [`DashMap::retain`]), serialised against `get_or_connect`, so a
|
||||
/// connection handed out concurrently is never evicted out from under a
|
||||
/// caller. Returns the number of connections evicted.
|
||||
pub fn sweep_idle(&self, idle_ttl: Duration) -> usize {
|
||||
let now = Instant::now();
|
||||
let mut evicted = 0usize;
|
||||
self.connections.retain(|_key, pooled| {
|
||||
let idle_for = now.saturating_duration_since(pooled.last_handout);
|
||||
// `strong_count == 1` ⇒ only this pool entry references the
|
||||
// connection, so no consumer can still be using it.
|
||||
let unused = Arc::strong_count(&pooled.conn) == 1;
|
||||
let evict = unused && idle_for >= idle_ttl;
|
||||
if evict {
|
||||
evicted += 1;
|
||||
}
|
||||
!evict
|
||||
});
|
||||
for _ in 0..evicted {
|
||||
crate::metrics::pool_connections_dec();
|
||||
crate::metrics::pool_evictions_inc();
|
||||
}
|
||||
evicted
|
||||
}
|
||||
|
||||
/// Spawn a background task that calls [`Self::sweep_idle`] every
|
||||
/// `sweep_interval`, closing connections idle longer than `idle_ttl`.
|
||||
///
|
||||
/// The task holds a [`std::sync::Weak`] to the pool, so it exits on its
|
||||
/// own once the last strong `Arc<HubConnectionPool>` is dropped (it never
|
||||
/// keeps the pool alive). The first interval tick is skipped so a
|
||||
/// freshly-handed-out connection is never swept on the immediate tick.
|
||||
pub fn spawn_idle_reaper(
|
||||
self: &Arc<Self>,
|
||||
idle_ttl: Duration,
|
||||
sweep_interval: Duration,
|
||||
) -> JoinHandle<()> {
|
||||
let weak = Arc::downgrade(self);
|
||||
tokio::spawn(async move {
|
||||
let mut ticker = tokio::time::interval(sweep_interval);
|
||||
ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
// `interval`'s first tick resolves immediately; skip it.
|
||||
ticker.tick().await;
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
let Some(pool) = weak.upgrade() else { break };
|
||||
pool.sweep_idle(idle_ttl);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Like [`Self::forget`] but identity-checked: only removes the slot
|
||||
/// when `predicate` accepts the currently-stored connection. The
|
||||
/// self-evicting actor passes an `Arc::ptr_eq` check so a race-loser
|
||||
/// can never drop the winner's fresh entry (ABA-safe).
|
||||
pub(crate) fn forget_if(
|
||||
&self,
|
||||
key: &ConnKey,
|
||||
predicate: impl FnOnce(&Arc<HubConnection>) -> bool,
|
||||
) {
|
||||
if self
|
||||
.connections
|
||||
.remove_if(key, |_, pooled| predicate(&pooled.conn))
|
||||
.is_some()
|
||||
{
|
||||
crate::metrics::pool_connections_dec();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
//! Generic refcounted-binding helper used by the connection's
|
||||
//! bound-session set.
|
||||
//!
|
||||
//! Multiple [`crate::ToolServer`] instances can share one
|
||||
//! [`crate::HubConnection`] when they target the same `(url, principal)`.
|
||||
//! Each instance independently asks for a session binding; the substrate
|
||||
//! must `register_session` once per session (not once per consumer) and
|
||||
//! `unregister_session` only when the LAST consumer drops its borrow.
|
||||
//! [`RefCountedSet`] tracks the per-key borrow count behind a
|
||||
//! [`dashmap::DashMap`] so increments and decrements never serialise on
|
||||
//! a single mutex.
|
||||
|
||||
use std::hash::Hash;
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
/// Refcounted set keyed by `K`. Each [`Self::increment`] returns the
|
||||
/// new count; the corresponding [`Self::decrement`] returns the count
|
||||
/// AFTER the decrement (so callers fire teardown when the result is
|
||||
/// `Some(0)`).
|
||||
#[derive(Debug, Default)]
|
||||
pub struct RefCountedSet<K: Eq + Hash> {
|
||||
counts: DashMap<K, u64>,
|
||||
}
|
||||
|
||||
impl<K: Eq + Hash> RefCountedSet<K> {
|
||||
/// Empty set.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
counts: DashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Increment `key`'s refcount. Returns `(prev_count, new_count)`
|
||||
/// so callers can detect the 0→1 edge (when the protocol-level
|
||||
/// register call must fire).
|
||||
pub fn increment(&self, key: K) -> (u64, u64)
|
||||
where
|
||||
K: Clone,
|
||||
{
|
||||
let mut entry = self.counts.entry(key).or_insert(0);
|
||||
let prev = *entry;
|
||||
*entry = prev.saturating_add(1);
|
||||
(prev, *entry)
|
||||
}
|
||||
|
||||
/// Decrement `key`'s refcount. Returns the post-decrement count;
|
||||
/// `Some(0)` means the entry was removed and callers should fire
|
||||
/// the protocol-level unregister. `None` means the key was not
|
||||
/// present (idempotent drop).
|
||||
pub fn decrement(&self, key: &K) -> Option<u64> {
|
||||
let mut current = None;
|
||||
self.counts.remove_if_mut(key, |_, value| {
|
||||
*value = value.saturating_sub(1);
|
||||
current = Some(*value);
|
||||
*value == 0
|
||||
});
|
||||
current
|
||||
}
|
||||
|
||||
/// Snapshot the live keys. Allocates a fresh `Vec` — only used by
|
||||
/// the reconnect-replay path which fires once per disconnect.
|
||||
pub fn snapshot_keys(&self) -> Vec<K>
|
||||
where
|
||||
K: Clone,
|
||||
{
|
||||
self.counts.iter().map(|kv| kv.key().clone()).collect()
|
||||
}
|
||||
|
||||
/// `true` when no key has a non-zero refcount.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.counts.is_empty()
|
||||
}
|
||||
|
||||
/// Number of distinct live keys.
|
||||
pub fn len(&self) -> usize {
|
||||
self.counts.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn increment_returns_new_count() {
|
||||
let set = RefCountedSet::<&'static str>::new();
|
||||
assert_eq!(set.increment("a"), (0, 1));
|
||||
assert_eq!(set.increment("a"), (1, 2));
|
||||
assert_eq!(set.increment("b"), (0, 1));
|
||||
assert_eq!(set.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrement_removes_at_zero() {
|
||||
let set = RefCountedSet::<&'static str>::new();
|
||||
set.increment("a");
|
||||
set.increment("a");
|
||||
assert_eq!(set.decrement(&"a"), Some(1));
|
||||
assert!(!set.is_empty());
|
||||
assert_eq!(set.decrement(&"a"), Some(0));
|
||||
assert!(set.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decrement_unknown_returns_none() {
|
||||
let set = RefCountedSet::<&'static str>::new();
|
||||
assert!(set.decrement(&"missing").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn increment_saturates_at_u64_max() {
|
||||
let set = RefCountedSet::<&'static str>::new();
|
||||
// Pre-load the entry to MAX-1 via direct DashMap access. The
|
||||
// public API only ever reaches this region via overflow,
|
||||
// which is impossible in practice; this test pins the
|
||||
// saturating_add defensive line so it can't silently regress
|
||||
// to wrapping_add.
|
||||
set.counts.insert("max", u64::MAX - 1);
|
||||
assert_eq!(set.increment("max"), (u64::MAX - 1, u64::MAX));
|
||||
assert_eq!(set.increment("max"), (u64::MAX, u64::MAX));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,206 @@
|
||||
//! Forward selected spans to the connected server over the WebSocket
|
||||
//! transport (`traces.donate`). The bounded retry buffer + drain barrier
|
||||
//! live in [`crate::donate_pump`]; overflow drops spans — telemetry,
|
||||
//! never correctness.
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
use base64::Engine as _;
|
||||
use fastrace::collector::{Reporter, SpanRecord};
|
||||
use fastrace_opentelemetry::OpenTelemetryReporter;
|
||||
use kigi_tool_protocol::{MAX_DONATION_BYTES, MAX_SPANS_PER_DONATION};
|
||||
use opentelemetry::InstrumentationScope;
|
||||
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
|
||||
use opentelemetry_proto::transform::common::tonic::ResourceAttributesWithSchema;
|
||||
use opentelemetry_proto::transform::trace::tonic::group_spans_by_resource_and_scope;
|
||||
use opentelemetry_sdk::Resource;
|
||||
use opentelemetry_sdk::error::OTelSdkResult;
|
||||
use opentelemetry_sdk::trace::{SpanData, SpanExporter};
|
||||
use prost::Message as _;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::donate_pump::{PENDING_FLUSHES, PumpMsg, drain_via, run_pump};
|
||||
use crate::server::ToolServer;
|
||||
|
||||
/// fastrace [`Reporter`] feeding the donation pump.
|
||||
pub struct HubDonatingReporter(OpenTelemetryReporter);
|
||||
|
||||
impl Reporter for HubDonatingReporter {
|
||||
fn report(&mut self, spans: Vec<SpanRecord>) {
|
||||
if spans.is_empty() {
|
||||
return;
|
||||
}
|
||||
self.0.report(spans);
|
||||
}
|
||||
}
|
||||
|
||||
/// [`SpanExporter`] that encodes OTLP requests onto the pump channel.
|
||||
/// Runs on fastrace's collector thread; must never block.
|
||||
#[derive(Debug)]
|
||||
struct PumpSpanExporter {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
resource: ResourceAttributesWithSchema,
|
||||
}
|
||||
|
||||
impl SpanExporter for PumpSpanExporter {
|
||||
fn export(
|
||||
&self,
|
||||
batch: Vec<SpanData>,
|
||||
) -> impl std::future::Future<Output = OTelSdkResult> + Send {
|
||||
let mut remaining = batch;
|
||||
while !remaining.is_empty() {
|
||||
let chunk = if remaining.len() > MAX_SPANS_PER_DONATION {
|
||||
let rest = remaining.split_off(MAX_SPANS_PER_DONATION);
|
||||
std::mem::replace(&mut remaining, rest)
|
||||
} else {
|
||||
std::mem::take(&mut remaining)
|
||||
};
|
||||
let request = ExportTraceServiceRequest {
|
||||
resource_spans: group_spans_by_resource_and_scope(chunk, &self.resource),
|
||||
};
|
||||
let bytes = request.encode_to_vec();
|
||||
if bytes.len() > MAX_DONATION_BYTES {
|
||||
tracing::debug!(len = bytes.len(), "dropping oversized donation payload");
|
||||
continue;
|
||||
}
|
||||
let payload = base64::engine::general_purpose::STANDARD.encode(bytes);
|
||||
if self.tx.try_send(PumpMsg::Payload(payload)).is_err() {
|
||||
tracing::debug!("trace donation queue full; dropping span batch");
|
||||
}
|
||||
}
|
||||
std::future::ready(Ok(()))
|
||||
}
|
||||
|
||||
fn set_resource(&mut self, resource: &Resource) {
|
||||
self.resource = resource.into();
|
||||
}
|
||||
}
|
||||
|
||||
/// Shutdown fence: drains queued donations before the connection closes.
|
||||
pub struct TraceDonationPump {
|
||||
tx: mpsc::Sender<PumpMsg>,
|
||||
}
|
||||
|
||||
impl TraceDonationPump {
|
||||
/// Resolves once every payload queued before this call has had a
|
||||
/// send attempt. Call after `fastrace::flush()`.
|
||||
pub async fn drain(&self) {
|
||||
drain_via(&self.tx).await;
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolServer {
|
||||
/// Spawn the donation pump and return its reporter + drain handle.
|
||||
/// `service_name` must be server-allowlisted.
|
||||
pub fn trace_donation_reporter(
|
||||
&self,
|
||||
service_name: impl Into<String>,
|
||||
) -> (HubDonatingReporter, TraceDonationPump) {
|
||||
let (tx, rx) = mpsc::channel::<PumpMsg>(PENDING_FLUSHES);
|
||||
let server = self.downgrade();
|
||||
tokio::spawn(run_pump(rx, move |payload: String| {
|
||||
let server = server.clone();
|
||||
async move {
|
||||
let Some(server) = server.upgrade() else {
|
||||
return (false, payload);
|
||||
};
|
||||
let ok = server.donate_traces(&payload).await.is_ok();
|
||||
(ok, payload)
|
||||
}
|
||||
}));
|
||||
self.set_donation_pump(tx.clone());
|
||||
|
||||
let resource = Resource::builder()
|
||||
.with_service_name(service_name.into())
|
||||
.build();
|
||||
let exporter = PumpSpanExporter {
|
||||
tx: tx.clone(),
|
||||
resource: (&resource).into(),
|
||||
};
|
||||
let reporter = OpenTelemetryReporter::new(
|
||||
exporter,
|
||||
Cow::Owned(resource),
|
||||
InstrumentationScope::default(),
|
||||
);
|
||||
(HubDonatingReporter(reporter), TraceDonationPump { tx })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::SystemTime;
|
||||
|
||||
use opentelemetry::trace::{SpanContext, SpanKind, Status, TraceFlags, TraceState};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn exporter_encodes_standard_otlp_with_resource() {
|
||||
let resource = Resource::builder()
|
||||
.with_service_name("test-service")
|
||||
.build();
|
||||
let (tx, mut rx) = mpsc::channel::<PumpMsg>(4);
|
||||
let exporter = PumpSpanExporter {
|
||||
tx,
|
||||
resource: (&resource).into(),
|
||||
};
|
||||
|
||||
let span = SpanData {
|
||||
span_context: SpanContext::new(
|
||||
0x0af7651916cd43dd8448eb211c80319c_u128.into(),
|
||||
0xb7ad6b7169203331_u64.into(),
|
||||
TraceFlags::SAMPLED,
|
||||
false,
|
||||
TraceState::default(),
|
||||
),
|
||||
parent_span_id: 0_u64.into(),
|
||||
parent_span_is_remote: false,
|
||||
span_kind: SpanKind::Internal,
|
||||
name: "tool_server.tool_call".into(),
|
||||
start_time: SystemTime::UNIX_EPOCH,
|
||||
end_time: SystemTime::UNIX_EPOCH,
|
||||
attributes: vec![opentelemetry::KeyValue::new("tool_id", "bash")],
|
||||
dropped_attributes_count: 0,
|
||||
events: opentelemetry_sdk::trace::SpanEvents::default(),
|
||||
links: opentelemetry_sdk::trace::SpanLinks::default(),
|
||||
status: Status::Unset,
|
||||
instrumentation_scope: InstrumentationScope::default(),
|
||||
};
|
||||
exporter
|
||||
.export(vec![span])
|
||||
.await
|
||||
.expect("export must succeed");
|
||||
|
||||
let Some(PumpMsg::Payload(payload)) = rx.try_recv().ok() else {
|
||||
panic!("exporter must enqueue one payload");
|
||||
};
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(payload)
|
||||
.expect("payload must be base64");
|
||||
let request =
|
||||
ExportTraceServiceRequest::decode(bytes.as_slice()).expect("payload must be OTLP");
|
||||
let resource_spans = &request.resource_spans[0];
|
||||
let service_name = resource_spans
|
||||
.resource
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.attributes
|
||||
.iter()
|
||||
.find(|kv| kv.key == "service.name")
|
||||
.and_then(|kv| kv.value.as_ref())
|
||||
.map(|v| format!("{v:?}"));
|
||||
assert!(
|
||||
service_name.unwrap_or_default().contains("test-service"),
|
||||
"resource must carry the donor service.name"
|
||||
);
|
||||
let span = &resource_spans.scope_spans[0].spans[0];
|
||||
assert_eq!(span.name, "tool_server.tool_call");
|
||||
assert_eq!(
|
||||
format!(
|
||||
"{:032x}",
|
||||
u128::from_be_bytes(span.trace_id.as_slice().try_into().unwrap())
|
||||
),
|
||||
"0af7651916cd43dd8448eb211c80319c"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-interjection-core"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Shared mid-turn interjection buffer and formatting for the client and server agent loops"
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,75 @@
|
||||
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.
|
||||
#[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.
|
||||
#[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.
|
||||
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).
|
||||
pub fn drain_formatted<Attachment>(
|
||||
buffer: &InterjectionBuffer<Attachment>,
|
||||
sanitize_text: impl Fn(String) -> String,
|
||||
) -> Vec<FormattedInterjection<Attachment>> {
|
||||
buffer
|
||||
.drain_all()
|
||||
.into_iter()
|
||||
.map(|entry| FormattedInterjection {
|
||||
text: format_interjection(sanitize_text(entry.text)),
|
||||
attachments: entry.attachments,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn drain_formatted_sanitizes_wraps_and_preserves_order() {
|
||||
let buf: InterjectionBuffer<()> = InterjectionBuffer::new();
|
||||
buf.push(PendingInterjection {
|
||||
text: "look at [SECRET] one".into(),
|
||||
attachments: vec![],
|
||||
});
|
||||
buf.push(PendingInterjection {
|
||||
text: "two".into(),
|
||||
attachments: vec![],
|
||||
});
|
||||
|
||||
let out = drain_formatted(&buf, |t| t.replace("[SECRET] ", ""));
|
||||
assert!(buf.is_empty());
|
||||
assert_eq!(out.len(), 2, "one message per entry, never merged");
|
||||
assert!(
|
||||
out[0]
|
||||
.text
|
||||
.contains("<user_query>\nlook at one\n</user_query>")
|
||||
);
|
||||
assert!(out[1].text.contains("<user_query>\ntwo\n</user_query>"));
|
||||
assert!(
|
||||
out[0]
|
||||
.text
|
||||
.starts_with("The user sent a message while you were working:")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,178 @@
|
||||
//! 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`].
|
||||
|
||||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct EventQueue<E> {
|
||||
events: Arc<Mutex<Vec<E>>>,
|
||||
}
|
||||
|
||||
impl<E> Clone for EventQueue<E> {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
events: Arc::clone(&self.events),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E> Default for EventQueue<E> {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl<E> EventQueue<E> {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
events: Arc::new(Mutex::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn push_capped(&self, event: E, max: usize) {
|
||||
let mut q = self.lock();
|
||||
q.push(event);
|
||||
if q.len() > max {
|
||||
let excess = q.len() - max;
|
||||
q.drain(..excess);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.lock().len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.lock().is_empty()
|
||||
}
|
||||
|
||||
/// Remove and return events matching `take`, retaining the rest. FIFO order
|
||||
/// is preserved in both the returned and retained sets.
|
||||
pub fn drain_matching(&self, take: impl Fn(&E) -> bool) -> Vec<E> {
|
||||
let mut q = self.lock();
|
||||
let (matched, kept): (Vec<E>, Vec<E>) =
|
||||
std::mem::take(&mut *q).into_iter().partition(|e| take(e));
|
||||
*q = kept;
|
||||
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();
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, Vec<E>> {
|
||||
self.events.lock().unwrap_or_else(|e| e.into_inner())
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: Clone> EventQueue<E> {
|
||||
/// Clone of the current events, for inspection without draining.
|
||||
pub fn snapshot(&self) -> Vec<E> {
|
||||
self.lock().clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn push_and_len() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
assert!(q.is_empty());
|
||||
q.push(1);
|
||||
q.push(2);
|
||||
assert_eq!(q.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clones_share_one_queue() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
let q2 = q.clone();
|
||||
q.push(7);
|
||||
assert_eq!(q2.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn push_capped_drops_oldest() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
for i in 0..5 {
|
||||
q.push_capped(i, 3);
|
||||
}
|
||||
assert_eq!(q.drain_matching(|_| true), vec![2, 3, 4]);
|
||||
assert!(q.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_matching_returns_matched_retains_rest_fifo() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
for i in 0..6 {
|
||||
q.push(i);
|
||||
}
|
||||
let evens = q.drain_matching(|n| n % 2 == 0);
|
||||
assert_eq!(evens, vec![0, 2, 4]);
|
||||
assert_eq!(q.drain_matching(|_| true), vec![1, 3, 5]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn push_capped_under_limit_keeps_all() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
q.push_capped(1, 5);
|
||||
q.push_capped(2, 5);
|
||||
assert_eq!(q.drain_matching(|_| true), vec![1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_matching_none_match_retains_all() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
q.push(1);
|
||||
q.push(2);
|
||||
assert!(q.drain_matching(|n| *n > 10).is_empty());
|
||||
assert_eq!(q.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_matching_on_empty_is_empty() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
assert!(q.drain_matching(|_| true).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drain_all_empties_in_fifo_order() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
q.push(1);
|
||||
q.push(2);
|
||||
assert_eq!(q.drain_all(), vec![1, 2]);
|
||||
assert!(q.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_discards_all() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
q.push(1);
|
||||
q.clear();
|
||||
assert!(q.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_reads_without_draining() {
|
||||
let q: EventQueue<u32> = EventQueue::new();
|
||||
q.push(9);
|
||||
assert_eq!(q.snapshot(), vec![9]);
|
||||
assert_eq!(q.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/// Truncation threshold, matching the shell's large-prompt limit.
|
||||
pub const LARGE_PROMPT_THRESHOLD: usize = 25_000;
|
||||
|
||||
/// Wrap a user message in the canonical `<user_query>` envelope.
|
||||
pub fn user_query(user_message: &str) -> String {
|
||||
format!(
|
||||
r#"<user_query>
|
||||
{user_message}
|
||||
</user_query>"#
|
||||
)
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn format_interjection(text: String) -> String {
|
||||
let truncated = if text.len() > LARGE_PROMPT_THRESHOLD {
|
||||
let end = text
|
||||
.char_indices()
|
||||
.take_while(|(i, _)| *i < LARGE_PROMPT_THRESHOLD)
|
||||
.last()
|
||||
.map(|(i, c)| i + c.len_utf8())
|
||||
.unwrap_or(text.len());
|
||||
format!("{}... [truncated]", &text[..end])
|
||||
} else {
|
||||
text
|
||||
};
|
||||
|
||||
format!(
|
||||
"The user sent a message while you were working:\n{}",
|
||||
user_query(&truncated)
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn wraps_in_user_query_with_midturn_note() {
|
||||
let out = format_interjection("stop and fix the test first".into());
|
||||
assert!(out.starts_with("The user sent a message while you were working:\n<user_query>\n"));
|
||||
assert!(out.ends_with("\n</user_query>"));
|
||||
assert!(out.contains("stop and fix the test first"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_at_utf8_boundary() {
|
||||
let s = "é".repeat(LARGE_PROMPT_THRESHOLD);
|
||||
let out = format_interjection(s);
|
||||
assert!(out.contains("... [truncated]"));
|
||||
assert!(out.len() < LARGE_PROMPT_THRESHOLD + 200);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn short_text_untouched() {
|
||||
let out = format_interjection("hi".into());
|
||||
assert!(!out.contains("[truncated]"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod buffer;
|
||||
pub mod events;
|
||||
pub mod format;
|
||||
|
||||
pub use buffer::{FormattedInterjection, InterjectionBuffer, PendingInterjection, drain_formatted};
|
||||
pub use events::EventQueue;
|
||||
pub use format::{LARGE_PROMPT_THRESHOLD, format_interjection, user_query};
|
||||
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-test-utils"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Shared test utilities: hermetic git, optional runfiles helpers"
|
||||
|
||||
[features]
|
||||
# Enable Bazel runfiles support when building under that toolchain. Under
|
||||
# plain `cargo` the feature is off and crate_root! falls back to
|
||||
# CARGO_MANIFEST_DIR.
|
||||
default-bazel = ["bazel"]
|
||||
bazel = ["dep:runfiles"]
|
||||
|
||||
[dependencies]
|
||||
runfiles = { workspace = true, optional = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,11 @@
|
||||
//! 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`).
|
||||
pub fn env_usize(key: &str, default: usize) -> usize {
|
||||
std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
//! 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.
|
||||
|
||||
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`.
|
||||
pub fn ensure_hermetic_git_on_path() {
|
||||
HERMETIC_GIT_INIT.call_once(|| {
|
||||
if let Ok(git_bin) = std::env::var("GIT_BIN_PATH") {
|
||||
let git_path = PathBuf::from(&git_bin);
|
||||
let git_path = if git_path.is_relative() {
|
||||
std::env::current_dir().unwrap().join(&git_path)
|
||||
} else {
|
||||
git_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.
|
||||
unsafe {
|
||||
std::env::set_var("PATH", format!("{}:{}", bin_dir.display(), current_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.
|
||||
///
|
||||
/// ```ignore
|
||||
/// #[test]
|
||||
/// fn my_git_test() {
|
||||
/// kigi_test_utils::require_git!();
|
||||
/// // ... git commands work here ...
|
||||
/// }
|
||||
/// ```
|
||||
#[macro_export]
|
||||
macro_rules! require_git {
|
||||
() => {
|
||||
$crate::git::ensure_hermetic_git_on_path();
|
||||
};
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn init_git_repo(path: &Path) {
|
||||
ensure_hermetic_git_on_path();
|
||||
std::process::Command::new("git")
|
||||
.current_dir(path)
|
||||
.args(["init"])
|
||||
.output()
|
||||
.unwrap();
|
||||
|
||||
std::process::Command::new("git")
|
||||
.current_dir(path)
|
||||
.args(["config", "user.email", "test@test.com"])
|
||||
.output()
|
||||
.unwrap();
|
||||
|
||||
std::process::Command::new("git")
|
||||
.current_dir(path)
|
||||
.args(["config", "user.name", "Test"])
|
||||
.output()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
/// Stage all files and create a commit.
|
||||
///
|
||||
/// Calls [`ensure_hermetic_git_on_path`] first so the hermetic binary is used.
|
||||
pub fn git_commit_all(path: &Path, message: &str) {
|
||||
ensure_hermetic_git_on_path();
|
||||
std::process::Command::new("git")
|
||||
.current_dir(path)
|
||||
.args(["add", "."])
|
||||
.output()
|
||||
.unwrap();
|
||||
std::process::Command::new("git")
|
||||
.current_dir(path)
|
||||
.args(["commit", "-m", message])
|
||||
.output()
|
||||
.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.
|
||||
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.
|
||||
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");
|
||||
cmd.args(args)
|
||||
.current_dir(dir)
|
||||
.env("GIT_AUTHOR_NAME", "Test User")
|
||||
.env("GIT_AUTHOR_EMAIL", "test@test.com")
|
||||
.env("GIT_COMMITTER_NAME", "Test User")
|
||||
.env("GIT_COMMITTER_EMAIL", "test@test.com")
|
||||
.env(
|
||||
"GIT_CONFIG_GLOBAL",
|
||||
if cfg!(windows) { "NUL" } else { "/dev/null" },
|
||||
)
|
||||
.env("GIT_CONFIG_NOSYSTEM", "1")
|
||||
.env("GIT_TERMINAL_PROMPT", "0");
|
||||
for (key, value) in envs {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
let output = cmd
|
||||
.output()
|
||||
.unwrap_or_else(|e| panic!("git {args:?} failed to spawn: {e}"));
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"git {:?} failed: {}",
|
||||
args,
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
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.
|
||||
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}"));
|
||||
std::fs::create_dir_all(&sub).expect("create populated dir");
|
||||
for f in 0..files_per_dir {
|
||||
std::fs::write(
|
||||
sub.join(format!("file_{f}.txt")),
|
||||
format!("content {d} {f}\n"),
|
||||
)
|
||||
.expect("write populated file");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
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"]);
|
||||
for k in 0..picks {
|
||||
let name = format!("pick_{k}.txt");
|
||||
std::fs::write(dir.join(&name), format!("pick {k}\n")).expect("write pick file");
|
||||
run_git(dir, &["add", &name]);
|
||||
run_git(dir, &["commit", "-m", &format!("pick {k}")]);
|
||||
}
|
||||
run_git(dir, &["checkout", &base]);
|
||||
std::fs::write(dir.join("base_advance.txt"), "advance\n").expect("write base advance file");
|
||||
run_git(dir, &["add", "base_advance.txt"]);
|
||||
run_git(dir, &["commit", "-m", "advance base"]);
|
||||
run_git(dir, &["checkout", "feature"]);
|
||||
base
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
//! 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.
|
||||
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(png);
|
||||
buf
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
//! Shared test utilities for xAI 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.
|
||||
|
||||
pub mod env;
|
||||
pub mod git;
|
||||
pub mod image;
|
||||
pub mod runfiles_util;
|
||||
pub mod tracing_capture;
|
||||
@@ -0,0 +1,46 @@
|
||||
//! Bazel runfiles helpers for locating 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.
|
||||
|
||||
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.
|
||||
pub fn try_resolve_runfiles(_path: &str) -> Option<PathBuf> {
|
||||
#[cfg(feature = "bazel")]
|
||||
{
|
||||
let r = runfiles::Runfiles::create().ok()?;
|
||||
runfiles::rlocation!(r, _path)
|
||||
}
|
||||
#[cfg(not(feature = "bazel"))]
|
||||
{
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// 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`.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```ignore
|
||||
/// use kigi_test_utils::crate_root;
|
||||
///
|
||||
/// fn test_data_dir() -> std::path::PathBuf {
|
||||
/// crate_root!("_main/crates/common/kigi-test-utils").join("testdata")
|
||||
/// }
|
||||
/// ```
|
||||
#[macro_export]
|
||||
macro_rules! crate_root {
|
||||
($runfiles_path:expr) => {
|
||||
$crate::runfiles_util::try_resolve_runfiles($runfiles_path)
|
||||
.unwrap_or_else(|| ::std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")))
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
//! Test-only tracing capture: count 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.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
/// Extracts the formatted `message` field of one event.
|
||||
#[derive(Default)]
|
||||
struct MessageVisitor(String);
|
||||
|
||||
impl tracing::field::Visit for MessageVisitor {
|
||||
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
|
||||
if field.name() == "message" {
|
||||
self.0 = format!("{value:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A `tracing_subscriber::Layer` counting, per registered prefix, the events
|
||||
/// whose `message` starts with it. Clones share the counts.
|
||||
#[derive(Clone)]
|
||||
pub struct MessagePrefixCounter {
|
||||
counters: Arc<Vec<(&'static str, AtomicUsize)>>,
|
||||
}
|
||||
|
||||
impl MessagePrefixCounter {
|
||||
pub fn new(prefixes: &[&'static str]) -> Self {
|
||||
Self {
|
||||
counters: Arc::new(prefixes.iter().map(|p| (*p, AtomicUsize::new(0))).collect()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn count(&self, prefix: &str) -> usize {
|
||||
self.counters
|
||||
.iter()
|
||||
.find(|(p, _)| *p == prefix)
|
||||
.unwrap_or_else(|| panic!("prefix not registered with this counter: {prefix:?}"))
|
||||
.1
|
||||
.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: tracing::Subscriber> tracing_subscriber::Layer<S> for MessagePrefixCounter {
|
||||
fn on_event(
|
||||
&self,
|
||||
event: &tracing::Event<'_>,
|
||||
_ctx: tracing_subscriber::layer::Context<'_, S>,
|
||||
) {
|
||||
let mut visitor = MessageVisitor::default();
|
||||
event.record(&mut visitor);
|
||||
for (prefix, count) in self.counters.iter() {
|
||||
if visitor.0.starts_with(prefix) {
|
||||
count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn install_prefix_counter_thread(
|
||||
prefixes: &[&'static str],
|
||||
) -> (tracing::subscriber::DefaultGuard, MessagePrefixCounter) {
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
let counter = MessagePrefixCounter::new(prefixes);
|
||||
let subscriber = tracing_subscriber::registry().with(counter.clone());
|
||||
(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.
|
||||
///
|
||||
/// `stderr_env_filter` additionally tees formatted logs matching the given
|
||||
/// `EnvFilter` directive to stderr (local debugging).
|
||||
pub fn install_prefix_counter_global(
|
||||
prefixes: &[&'static str],
|
||||
stderr_env_filter: Option<&str>,
|
||||
) -> MessagePrefixCounter {
|
||||
use tracing_subscriber::layer::{Layer as _, SubscriberExt as _};
|
||||
use tracing_subscriber::util::SubscriberInitExt as _;
|
||||
let counter = MessagePrefixCounter::new(prefixes);
|
||||
let fmt = stderr_env_filter.map(|filter| {
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_filter(tracing_subscriber::EnvFilter::new(filter))
|
||||
});
|
||||
tracing_subscriber::registry()
|
||||
.with(counter.clone())
|
||||
.with(fmt)
|
||||
.try_init()
|
||||
.expect("this test binary must own the global subscriber");
|
||||
counter
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-tool-protocol"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Wire-protocol types for the xAI Computer Hub"
|
||||
|
||||
[dependencies]
|
||||
kigi-tool-types = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
uuid = { workspace = true, features = ["v7"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,106 @@
|
||||
//! Per-tool capabilities and notification schemas.
|
||||
|
||||
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.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub streaming: Option<StreamingSpec>,
|
||||
|
||||
/// Tool honours `hook { Cancel }`.
|
||||
#[serde(default)]
|
||||
pub supports_cancel: bool,
|
||||
|
||||
/// Maximum concurrent invocations the tool will accept. `None` is
|
||||
/// unlimited.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_concurrency: Option<u32>,
|
||||
|
||||
/// Mirrors `Tool::is_read_only`; used by doom-loop detection.
|
||||
#[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>,
|
||||
|
||||
/// Opaque per-tool behaviour version. Bytewise-compared (NOT semver).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub behavior_version: Option<String>,
|
||||
|
||||
/// Per-tool override for the per-frame size cap. Service clamps to the
|
||||
/// 16 MiB hard ceiling.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub max_frame_bytes: Option<u32>,
|
||||
|
||||
/// Per-call timeout override (defaults to 60_000ms when omitted).
|
||||
#[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`.
|
||||
#[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.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct StreamingSpec {
|
||||
/// Stable snake_case discriminator the tool stamps on its
|
||||
/// `ToolProgress::Custom.subkind` (e.g. `"bash_output_chunk"`).
|
||||
pub subkind: String,
|
||||
|
||||
/// Per-frame `delta` byte cap (UTF-8-safe). Unset falls back to the
|
||||
/// runtime's 16 KiB default. Independent of
|
||||
/// [`ToolCapabilities::max_frame_bytes`], which caps whole frames.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
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 {
|
||||
OnSessionOpen,
|
||||
OnSessionClose,
|
||||
OnToolCallStart,
|
||||
OnToolCallResult,
|
||||
OnCancel,
|
||||
OnNotification,
|
||||
}
|
||||
|
||||
/// 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`.
|
||||
#[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.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct NotificationSchemas {
|
||||
/// Schemas for 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.
|
||||
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
|
||||
pub inbound: HashMap<String, serde_json::Value>,
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
//! Connection-shape and tool-definition-mode enums.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Role of a WebSocket connection. The computer hub uses this to decide
|
||||
/// 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 {
|
||||
Harness,
|
||||
ToolServer,
|
||||
}
|
||||
|
||||
/// 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": "..."}`.
|
||||
///
|
||||
/// `Copy` is intentionally NOT derived: `Concise`'s [`crate::ToolId`]
|
||||
/// fields wrap heap strings.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(tag = "mode", rename_all = "snake_case")]
|
||||
pub enum ToolDefinitionMode {
|
||||
/// Every `ToolDescription` is sent to the model directly.
|
||||
Full,
|
||||
/// Only the meta-tool pair is sent; everything else is discoverable
|
||||
/// through the search meta-tool.
|
||||
Concise {
|
||||
/// Model-facing name of the search/discovery meta-tool.
|
||||
meta_search: crate::ToolId,
|
||||
/// Model-facing name of the call/invoke meta-tool.
|
||||
meta_call: crate::ToolId,
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
//! JSON-RPC 2.0 envelope types with the Grok `session_id` / `seq`
|
||||
//! extensions.
|
||||
//!
|
||||
//! Two distinct id concepts coexist in this crate:
|
||||
//!
|
||||
//! - [`JsonRpcId`] (this module) is the JSON-RPC envelope `id` field —
|
||||
//! string OR number on the wire, per-connection, sender-allocated.
|
||||
//! - [`crate::RequestId`] is an opaque newtype wrapping a string, used
|
||||
//! internally as a correlator (e.g. to key in-flight maps). Convert
|
||||
//! between them via [`JsonRpcId::from_request_id`] /
|
||||
//! [`JsonRpcId::as_request_id`].
|
||||
|
||||
use std::fmt;
|
||||
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer, de};
|
||||
|
||||
use crate::{FrameSeq, IdError, RequestId, SessionId};
|
||||
|
||||
/// JSON-RPC 2.0 protocol version marker.
|
||||
///
|
||||
/// Serializes as the literal string `"2.0"` and rejects any other value on
|
||||
/// deserialize.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct JsonRpcVersion;
|
||||
|
||||
impl JsonRpcVersion {
|
||||
pub const VERSION: &'static str = "2.0";
|
||||
}
|
||||
|
||||
impl fmt::Display for JsonRpcVersion {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(Self::VERSION)
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for JsonRpcVersion {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.serialize_str(Self::VERSION)
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for JsonRpcVersion {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
struct V;
|
||||
impl de::Visitor<'_> for V {
|
||||
type Value = JsonRpcVersion;
|
||||
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "the literal string \"{}\"", JsonRpcVersion::VERSION)
|
||||
}
|
||||
fn visit_str<E: de::Error>(self, v: &str) -> Result<Self::Value, E> {
|
||||
if v == JsonRpcVersion::VERSION {
|
||||
Ok(JsonRpcVersion)
|
||||
} else {
|
||||
Err(E::custom(format!(
|
||||
"expected jsonrpc \"{}\", got {v:?}",
|
||||
JsonRpcVersion::VERSION
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
deserializer.deserialize_str(V)
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 envelope `id` field.
|
||||
///
|
||||
/// Per the spec the `id` MAY be a string, a number, or null. We accept the
|
||||
/// first two on deserialize and emit a string ourselves. Null ids are not
|
||||
/// produced and not modelled on the receive path.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum JsonRpcId {
|
||||
String(String),
|
||||
Number(i64),
|
||||
}
|
||||
|
||||
impl JsonRpcId {
|
||||
pub fn new_string(s: impl Into<String>) -> Self {
|
||||
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())
|
||||
}
|
||||
|
||||
pub fn from_request_id(id: &RequestId) -> Self {
|
||||
Self::String(id.as_str().to_owned())
|
||||
}
|
||||
|
||||
/// Project to a [`RequestId`]. Numeric ids are stringified. Returns
|
||||
/// an error if the resulting string would be empty.
|
||||
pub fn as_request_id(&self) -> Result<RequestId, IdError> {
|
||||
match self {
|
||||
Self::String(s) => RequestId::new(s.as_str()),
|
||||
Self::Number(n) => RequestId::new(n.to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for JsonRpcId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::String(s) => f.write_str(s),
|
||||
Self::Number(n) => write!(f, "{n}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 request envelope.
|
||||
///
|
||||
/// Generic over `params` so callers can pin a concrete schema (e.g.
|
||||
/// [`crate::frames::ToolCallParams`]) without losing the envelope's
|
||||
/// invariants.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct JsonRpcRequest<P = serde_json::Value> {
|
||||
pub jsonrpc: JsonRpcVersion,
|
||||
pub id: JsonRpcId,
|
||||
/// Grok extension: routing/sanity-check session id.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub session_id: Option<SessionId>,
|
||||
pub method: String,
|
||||
pub params: P,
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 notification envelope.
|
||||
///
|
||||
/// No `id` (notifications do not produce a response). `seq` is an
|
||||
/// optional per-connection monotonic counter so receivers can dedup and
|
||||
/// detect drops.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct JsonRpcNotification<P = serde_json::Value> {
|
||||
pub jsonrpc: JsonRpcVersion,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub session_id: Option<SessionId>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub seq: Option<FrameSeq>,
|
||||
pub method: String,
|
||||
pub params: P,
|
||||
}
|
||||
|
||||
/// JSON-RPC error object.
|
||||
///
|
||||
/// `code` is the numeric envelope code; `data` typically carries a
|
||||
/// serialized [`crate::error_wire::ToolErrorWire`] so receivers can switch
|
||||
/// on the stable string code rather than the numeric.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct JsonRpcError {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 response envelope.
|
||||
///
|
||||
/// Per the spec exactly one of `result` / `error` is present. The custom
|
||||
/// `Serialize` / `Deserialize` impls enforce that invariant: a payload
|
||||
/// containing both keys, or neither, fails to deserialize.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct JsonRpcResponse<R = serde_json::Value> {
|
||||
pub jsonrpc: JsonRpcVersion,
|
||||
pub id: JsonRpcId,
|
||||
pub session_id: Option<SessionId>,
|
||||
pub outcome: ResponseOutcome<R>,
|
||||
}
|
||||
|
||||
/// Either a `result` payload (success) or a [`JsonRpcError`] (failure).
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ResponseOutcome<R> {
|
||||
Result(R),
|
||||
Error(JsonRpcError),
|
||||
}
|
||||
|
||||
impl<R: Serialize> Serialize for JsonRpcResponse<R> {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
use serde::ser::SerializeMap;
|
||||
let mut len = 3;
|
||||
if self.session_id.is_some() {
|
||||
len += 1;
|
||||
}
|
||||
let mut map = serializer.serialize_map(Some(len))?;
|
||||
map.serialize_entry("jsonrpc", &self.jsonrpc)?;
|
||||
map.serialize_entry("id", &self.id)?;
|
||||
if let Some(sid) = &self.session_id {
|
||||
map.serialize_entry("session_id", sid)?;
|
||||
}
|
||||
match &self.outcome {
|
||||
ResponseOutcome::Result(r) => map.serialize_entry("result", r)?,
|
||||
ResponseOutcome::Error(e) => map.serialize_entry("error", e)?,
|
||||
}
|
||||
map.end()
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de, R: Deserialize<'de>> Deserialize<'de> for JsonRpcResponse<R> {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
// `Option<...>` deserialises to `None` when missing without
|
||||
// `#[serde(default)]`, avoiding a `R: Default` bound on the
|
||||
// result type parameter.
|
||||
#[derive(Deserialize)]
|
||||
struct Flat<R> {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id: JsonRpcId,
|
||||
session_id: Option<SessionId>,
|
||||
result: Option<R>,
|
||||
error: Option<JsonRpcError>,
|
||||
}
|
||||
|
||||
let flat = Flat::<R>::deserialize(deserializer)?;
|
||||
let outcome = match (flat.result, flat.error) {
|
||||
(Some(r), None) => ResponseOutcome::Result(r),
|
||||
(None, Some(e)) => ResponseOutcome::Error(e),
|
||||
(Some(_), Some(_)) => {
|
||||
return Err(de::Error::custom(
|
||||
"JSON-RPC response must contain `result` XOR `error`, got both",
|
||||
));
|
||||
}
|
||||
(None, None) => {
|
||||
return Err(de::Error::custom(
|
||||
"JSON-RPC response must contain `result` or `error`",
|
||||
));
|
||||
}
|
||||
};
|
||||
Ok(Self {
|
||||
jsonrpc: flat.jsonrpc,
|
||||
id: flat.id,
|
||||
session_id: flat.session_id,
|
||||
outcome,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<R> JsonRpcResponse<R> {
|
||||
pub fn ok(id: JsonRpcId, result: R) -> Self {
|
||||
Self {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id,
|
||||
session_id: None,
|
||||
outcome: ResponseOutcome::Result(result),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn err(id: JsonRpcId, error: JsonRpcError) -> Self {
|
||||
Self {
|
||||
jsonrpc: JsonRpcVersion,
|
||||
id,
|
||||
session_id: None,
|
||||
outcome: ResponseOutcome::Error(error),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_session(mut self, sid: SessionId) -> Self {
|
||||
self.session_id = Some(sid);
|
||||
self
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,300 @@
|
||||
//! Numeric ↔ string error-code mapping.
|
||||
//!
|
||||
//! Receivers SHOULD switch on `data.code` (the snake_case string) rather
|
||||
//! than the numeric JSON-RPC `error.code`. The numeric is the JSON-RPC
|
||||
//! envelope code; the string is the Grok 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.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error_wire::ToolErrorWire;
|
||||
|
||||
/// `(numeric_code, string_code)` pairs. Both columns are unique.
|
||||
pub const ERROR_CODES: &[(i32, &str)] = &[
|
||||
(-32700, "parse_error"),
|
||||
(-32600, "invalid_request"),
|
||||
(-32601, "method_not_found"),
|
||||
(-32602, "invalid_params"),
|
||||
(-32603, "internal_error"),
|
||||
(-32605, "unsupported_protocol_version"),
|
||||
(-32001, "timeout"),
|
||||
(-32002, "unauthorized"),
|
||||
(-32003, "forbidden"),
|
||||
(-32004, "connection_lost"),
|
||||
(-32005, "tool_server_gone"),
|
||||
(-32006, "session_not_found"),
|
||||
(-32008, "session_draining"),
|
||||
(-32011, "tool_not_found"),
|
||||
(-32012, "tool_already_registered"),
|
||||
(-32013, "tool_unavailable"),
|
||||
(-32014, "stale_generation"),
|
||||
(-32015, "duplicate_client_name"),
|
||||
(-32016, "tool_busy"),
|
||||
(-32017, "notification_schema_violation"),
|
||||
(-32018, "frame_too_large"),
|
||||
(-32019, "schema_unknown_kind"),
|
||||
(-32020, "behavior_version_unsupported"),
|
||||
(-32021, "server_id_in_use"),
|
||||
(-32022, "invalid_description"),
|
||||
(-32023, "render_limited"),
|
||||
(-32024, "terminal_error"),
|
||||
(-32099, "rate_limited"),
|
||||
];
|
||||
|
||||
/// Returns `None` for strings not in the table. Receivers should fall
|
||||
/// back to `-32603 internal_error` for unknown strings.
|
||||
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.
|
||||
pub fn from_tool_error_wire(err: &ToolErrorWire) -> i32 {
|
||||
match err {
|
||||
ToolErrorWire::ToolNotFound { .. } => -32011,
|
||||
ToolErrorWire::SessionMismatch => -32600,
|
||||
ToolErrorWire::PermissionDenied { .. } => -32003,
|
||||
ToolErrorWire::TransportClosed { .. } => -32004,
|
||||
ToolErrorWire::Timeout { .. } => -32001,
|
||||
ToolErrorWire::Cancelled { .. } => -32603,
|
||||
ToolErrorWire::InvalidArguments { .. } => -32602,
|
||||
ToolErrorWire::Execution { .. } => -32603,
|
||||
ToolErrorWire::UnsupportedProtocolVersion { .. } => -32605,
|
||||
ToolErrorWire::PayloadTooLarge { .. } => -32018,
|
||||
ToolErrorWire::BehaviorVersionUnsupported { .. } => -32020,
|
||||
ToolErrorWire::Internal { .. } => -32603,
|
||||
ToolErrorWire::RenderLimited { .. } => -32023,
|
||||
ToolErrorWire::TerminalError { .. } => -32024,
|
||||
ToolErrorWire::Custom { .. } => -32603,
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable identifier for "this session's workspace (tool) server is gone;
|
||||
/// re-provision and retry", used as both the [`ToolErrorWire::Custom`] subcode
|
||||
/// and the `details["code"]` value. Reusing `Custom` (not a new variant) keeps
|
||||
/// the frame deserializable on older peers.
|
||||
pub const WORKSPACE_UNAVAILABLE_SUBCODE: &str = "workspace_unavailable";
|
||||
|
||||
/// Generic, tenant-data-free message paired with the workspace-gone error.
|
||||
pub const WORKSPACE_UNAVAILABLE_MESSAGE: &str = "workspace server gone; re-provision and retry";
|
||||
|
||||
/// JSON-RPC envelope code paired with the workspace-unavailable error. Shares
|
||||
/// the canonical `tool_server_gone` numeric; recognizers key on `data.subcode`,
|
||||
/// not this companion.
|
||||
pub const WORKSPACE_UNAVAILABLE_JSONRPC_CODE: i32 = -32005;
|
||||
|
||||
/// Why the workspace (tool) server went away. `Unknown` absorbs values a newer
|
||||
/// peer may add, so the typed parse never fails across independently-deployed
|
||||
/// hub/SDK versions.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkspaceGoneReason {
|
||||
IdleTimeout,
|
||||
Disconnect,
|
||||
Shutdown,
|
||||
/// No owner has bound a tool-server for the session yet (an attach-time
|
||||
/// miss), as opposed to a workspace that was bound and then lost.
|
||||
NotBound,
|
||||
/// Target hub liveness key absent (origin reaper or forward-time check).
|
||||
InstanceGone,
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// When, relative to the failing tool call, the loss was observed. `Unknown`
|
||||
/// absorbs values a newer peer may add.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum WorkspaceGonePhase {
|
||||
InFlightCancelled,
|
||||
RouteMissing,
|
||||
/// Observed while resolving a `session_attach_server` request.
|
||||
Attach,
|
||||
#[serde(other)]
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// Structured payload placed in the wire `details` object. `code` mirrors the
|
||||
/// `Custom` subcode (the `ToolError::custom` convention), so it survives a
|
||||
/// `Wire → ToolError → Wire` round-trip and is the field recognizers read.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct WorkspaceUnavailableDetails {
|
||||
pub code: String,
|
||||
pub reason: WorkspaceGoneReason,
|
||||
pub phase: WorkspaceGonePhase,
|
||||
pub retryable: bool,
|
||||
}
|
||||
|
||||
/// Build the recognizable "workspace gone" error as a [`ToolErrorWire::Custom`].
|
||||
pub fn workspace_unavailable_wire(
|
||||
reason: WorkspaceGoneReason,
|
||||
phase: WorkspaceGonePhase,
|
||||
) -> ToolErrorWire {
|
||||
let details = serde_json::to_value(WorkspaceUnavailableDetails {
|
||||
code: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
reason,
|
||||
phase,
|
||||
retryable: true,
|
||||
});
|
||||
// This plain struct serializes infallibly; a missing `details` would make
|
||||
// the error unrecognizable, so guard the invariant in debug builds.
|
||||
debug_assert!(details.is_ok(), "workspace details must serialize");
|
||||
ToolErrorWire::Custom {
|
||||
subcode: WORKSPACE_UNAVAILABLE_SUBCODE.to_owned(),
|
||||
message: WORKSPACE_UNAVAILABLE_MESSAGE.to_owned(),
|
||||
details: details.ok(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
const REASONS: [WorkspaceGoneReason; 5] = [
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGoneReason::Disconnect,
|
||||
WorkspaceGoneReason::Shutdown,
|
||||
WorkspaceGoneReason::NotBound,
|
||||
WorkspaceGoneReason::InstanceGone,
|
||||
];
|
||||
const PHASES: [WorkspaceGonePhase; 3] = [
|
||||
WorkspaceGonePhase::InFlightCancelled,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
WorkspaceGonePhase::Attach,
|
||||
];
|
||||
|
||||
// Exhaustive-match helpers pin the exact snake_case wire strings; adding a
|
||||
// variant forces an update here.
|
||||
fn reason_wire(r: WorkspaceGoneReason) -> &'static str {
|
||||
match r {
|
||||
WorkspaceGoneReason::IdleTimeout => "idle_timeout",
|
||||
WorkspaceGoneReason::Disconnect => "disconnect",
|
||||
WorkspaceGoneReason::Shutdown => "shutdown",
|
||||
WorkspaceGoneReason::NotBound => "not_bound",
|
||||
WorkspaceGoneReason::InstanceGone => "instance_gone",
|
||||
WorkspaceGoneReason::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
fn phase_wire(p: WorkspaceGonePhase) -> &'static str {
|
||||
match p {
|
||||
WorkspaceGonePhase::InFlightCancelled => "in_flight_cancelled",
|
||||
WorkspaceGonePhase::RouteMissing => "route_missing",
|
||||
WorkspaceGonePhase::Attach => "attach",
|
||||
WorkspaceGonePhase::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_emits_custom_with_code_in_details_for_every_reason_and_phase() {
|
||||
for reason in REASONS {
|
||||
for phase in PHASES {
|
||||
let v = serde_json::to_value(workspace_unavailable_wire(reason, phase)).unwrap();
|
||||
assert_eq!(v["code"], json!("custom"), "outer discriminator");
|
||||
assert_eq!(v["subcode"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
|
||||
// details.code mirrors the subcode (round-trip identity).
|
||||
assert_eq!(v["details"]["code"], json!(WORKSPACE_UNAVAILABLE_SUBCODE));
|
||||
assert_eq!(v["details"]["reason"], json!(reason_wire(reason)));
|
||||
assert_eq!(v["details"]["phase"], json!(phase_wire(phase)));
|
||||
assert_eq!(v["details"]["retryable"], json!(true));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builder_uses_the_pinned_generic_message() {
|
||||
let ToolErrorWire::Custom { message, .. } = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::IdleTimeout,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
) else {
|
||||
panic!("expected Custom variant");
|
||||
};
|
||||
// Exact, tenant-data-free contract.
|
||||
assert_eq!(message, WORKSPACE_UNAVAILABLE_MESSAGE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_reason_and_phase_deserialize_to_unknown() {
|
||||
// Independently-deployed peers may emit reason/phase values this build
|
||||
// does not know; the typed parse must absorb them, not fail.
|
||||
let details: WorkspaceUnavailableDetails = serde_json::from_value(json!({
|
||||
"code": WORKSPACE_UNAVAILABLE_SUBCODE,
|
||||
"reason": "reason_from_a_newer_hub",
|
||||
"phase": "phase_from_a_newer_hub",
|
||||
"retryable": true,
|
||||
}))
|
||||
.expect("typed parse tolerates unknown enum values");
|
||||
assert_eq!(details.reason, WorkspaceGoneReason::Unknown);
|
||||
assert_eq!(details.phase, WorkspaceGonePhase::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_reason_serializes_and_round_trips() {
|
||||
// The route-missing classifier emits `Unknown` ("cause not observed"),
|
||||
// so — despite `Unknown` being the `#[serde(other)]` deserialize
|
||||
// catch-all — it must serialize to a stable `"unknown"` label and parse
|
||||
// back, both in the wire payload and as the bare enum.
|
||||
assert_eq!(
|
||||
serde_json::to_value(WorkspaceGoneReason::Unknown).unwrap(),
|
||||
json!("unknown"),
|
||||
);
|
||||
let wire = workspace_unavailable_wire(
|
||||
WorkspaceGoneReason::Unknown,
|
||||
WorkspaceGonePhase::RouteMissing,
|
||||
);
|
||||
let v = serde_json::to_value(&wire).unwrap();
|
||||
assert_eq!(v["details"]["reason"], json!("unknown"));
|
||||
let parsed: WorkspaceUnavailableDetails =
|
||||
serde_json::from_value(v["details"].clone()).expect("details round-trip");
|
||||
assert_eq!(parsed.reason, WorkspaceGoneReason::Unknown);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn custom_variant_tolerates_unknown_future_details_shape() {
|
||||
// An unknown subcode + richer future details must still deserialize rather than failing the frame.
|
||||
let future = json!({
|
||||
"code": "custom",
|
||||
"subcode": "some_future_subcode",
|
||||
"message": "from a newer peer",
|
||||
"details": {
|
||||
"code": "some_future_subcode",
|
||||
"extra_new_field": {"nested": [1, 2, 3]},
|
||||
},
|
||||
});
|
||||
let wire: ToolErrorWire =
|
||||
serde_json::from_value(future).expect("custom variant deserializes");
|
||||
let ToolErrorWire::Custom {
|
||||
subcode, details, ..
|
||||
} = &wire
|
||||
else {
|
||||
panic!("expected Custom variant");
|
||||
};
|
||||
assert_eq!(subcode, "some_future_subcode");
|
||||
assert!(
|
||||
details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("extra_new_field"))
|
||||
.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"],
|
||||
json!([1, 2, 3])
|
||||
);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user