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,391 @@
|
||||
//! ChatStateActor — runs in a dedicated tokio task and owns all chat state.
|
||||
//!
|
||||
//! This module is organized into submodules by responsibility:
|
||||
//! - `state`: Internal state types (ChatState)
|
||||
//! - `mutations`: State mutation handlers (push_user_message, replace_conversation, etc.)
|
||||
//! - `queries`: Read-only query handlers (get_conversation, snapshot, etc.)
|
||||
|
||||
mod mutations;
|
||||
mod queries;
|
||||
pub(crate) mod request_builder;
|
||||
pub mod state;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::commands::ChatStateCommand;
|
||||
use crate::events::ChatStateEvent;
|
||||
use crate::handle::ChatStateHandle;
|
||||
use crate::persistence::ChatPersistence;
|
||||
use crate::types::{PruningConfig, TurnCapture};
|
||||
|
||||
use kigi_sampling_types::{ConversationItem, SamplingConfig};
|
||||
use state::ChatState;
|
||||
|
||||
/// The actor that owns all chat state.
|
||||
/// Runs in a dedicated tokio task and processes commands sequentially.
|
||||
pub struct ChatStateActor {
|
||||
/// Internal state — conversation, tokens, config, etc.
|
||||
state: ChatState,
|
||||
/// Pruning configuration for tool-result trimming.
|
||||
pruning_config: PruningConfig,
|
||||
/// Persistence implementation — owned exclusively, called with `&mut self`.
|
||||
persistence: Box<dyn ChatPersistence>,
|
||||
/// Channel to receive commands from handles.
|
||||
cmd_rx: mpsc::UnboundedReceiver<ChatStateCommand>,
|
||||
/// Channel to send events to the session main loop.
|
||||
event_tx: mpsc::UnboundedSender<ChatStateEvent>,
|
||||
/// Cancellation token for graceful shutdown.
|
||||
cancellation_token: tokio_util::sync::CancellationToken,
|
||||
}
|
||||
|
||||
impl ChatStateActor {
|
||||
/// Send an event to subscribers, logging if the channel is closed.
|
||||
fn send_event(&self, event: ChatStateEvent) {
|
||||
if self.event_tx.send(event).is_err() {
|
||||
debug!("ChatState event channel closed, event dropped");
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn the actor and return a handle to communicate with it.
|
||||
pub fn spawn(
|
||||
initial_conversation: Vec<ConversationItem>,
|
||||
sampling_config: SamplingConfig,
|
||||
persistence: Box<dyn ChatPersistence>,
|
||||
event_tx: mpsc::UnboundedSender<ChatStateEvent>,
|
||||
cancellation_token: tokio_util::sync::CancellationToken,
|
||||
) -> ChatStateHandle {
|
||||
Self::spawn_with_pruning(
|
||||
initial_conversation,
|
||||
sampling_config,
|
||||
PruningConfig::default(),
|
||||
persistence,
|
||||
event_tx,
|
||||
cancellation_token,
|
||||
)
|
||||
}
|
||||
|
||||
/// Spawn the actor with a custom pruning config.
|
||||
pub fn spawn_with_pruning(
|
||||
initial_conversation: Vec<ConversationItem>,
|
||||
sampling_config: SamplingConfig,
|
||||
pruning_config: PruningConfig,
|
||||
persistence: Box<dyn ChatPersistence>,
|
||||
event_tx: mpsc::UnboundedSender<ChatStateEvent>,
|
||||
cancellation_token: tokio_util::sync::CancellationToken,
|
||||
) -> ChatStateHandle {
|
||||
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
|
||||
|
||||
let actor = ChatStateActor {
|
||||
state: ChatState::new(initial_conversation, sampling_config),
|
||||
pruning_config,
|
||||
persistence,
|
||||
cmd_rx,
|
||||
event_tx,
|
||||
cancellation_token,
|
||||
};
|
||||
|
||||
tokio::spawn(actor.run());
|
||||
|
||||
ChatStateHandle::new(cmd_tx)
|
||||
}
|
||||
|
||||
/// Main actor loop — processes commands until shutdown or cancellation.
|
||||
async fn run(mut self) {
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = self.cancellation_token.cancelled() => {
|
||||
debug!("ChatStateActor shutting down via cancellation");
|
||||
break;
|
||||
}
|
||||
cmd = self.cmd_rx.recv() => {
|
||||
let Some(cmd) = cmd else {
|
||||
debug!("ChatStateActor shutting down: all handles dropped");
|
||||
break;
|
||||
};
|
||||
self.handle_command(cmd);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Dispatch a command to the appropriate mutation or query handler.
|
||||
fn handle_command(&mut self, cmd: ChatStateCommand) {
|
||||
match cmd {
|
||||
// ═══ Mutations ═══
|
||||
ChatStateCommand::PushUserMessage { item } => {
|
||||
self.push_user_message(item);
|
||||
}
|
||||
ChatStateCommand::PushUserMessageAndAck { item, reply } => {
|
||||
self.push_user_message(item);
|
||||
let _ = reply.send(());
|
||||
}
|
||||
ChatStateCommand::PushUserMessageWithRepairReason { item, reason } => {
|
||||
self.push_user_message_with_repair_reason(item, reason);
|
||||
}
|
||||
ChatStateCommand::PushAssistantResponse { item } => {
|
||||
self.push_message(item);
|
||||
}
|
||||
ChatStateCommand::PushToolResult { item } => {
|
||||
self.push_message(item);
|
||||
}
|
||||
ChatStateCommand::RecordTokenUsage { total_tokens } => {
|
||||
self.record_token_usage(total_tokens);
|
||||
}
|
||||
ChatStateCommand::RecordLastTurnUsage { usage } => {
|
||||
self.record_last_turn_usage(usage);
|
||||
}
|
||||
ChatStateCommand::RecordModelCallUsage {
|
||||
model_id,
|
||||
usage,
|
||||
api_duration_ms,
|
||||
cost_usd_ticks,
|
||||
} => {
|
||||
self.record_model_call_usage(model_id, &usage, api_duration_ms, cost_usd_ticks);
|
||||
}
|
||||
ChatStateCommand::RecordSubagentUsage {
|
||||
by_model,
|
||||
attribute_to_prompt,
|
||||
incomplete,
|
||||
reply,
|
||||
} => {
|
||||
self.record_subagent_usage(&by_model, attribute_to_prompt, incomplete);
|
||||
let _ = reply.send(());
|
||||
}
|
||||
ChatStateCommand::MarkUsageIncomplete {
|
||||
prompt,
|
||||
session,
|
||||
reply,
|
||||
} => {
|
||||
self.mark_usage_incomplete(prompt, session);
|
||||
let _ = reply.send(());
|
||||
}
|
||||
ChatStateCommand::IncrementPromptIndex => {
|
||||
self.increment_prompt_index();
|
||||
}
|
||||
ChatStateCommand::UpdateSamplingConfig { config } => {
|
||||
self.state.sampling_config = config;
|
||||
}
|
||||
ChatStateCommand::RecordAgentEditedPath { path } => {
|
||||
self.state.agent_edited_paths.insert(path);
|
||||
}
|
||||
ChatStateCommand::RecordStreamStart { timestamp_ms } => {
|
||||
self.state.stream_start_ms = Some(timestamp_ms);
|
||||
}
|
||||
ChatStateCommand::RecordTurnStart { timestamp_ms } => {
|
||||
self.state.turn_start_ms = Some(timestamp_ms);
|
||||
}
|
||||
ChatStateCommand::ReplaceConversation {
|
||||
items,
|
||||
is_compaction,
|
||||
} => {
|
||||
self.replace_conversation(items, is_compaction);
|
||||
}
|
||||
ChatStateCommand::RepairHistory {
|
||||
dry_run,
|
||||
turn_active,
|
||||
reply,
|
||||
} => {
|
||||
// Checked here so refusal and mutation are serialized; a
|
||||
// `false` at processing time means pre-turn state (see the
|
||||
// command's doc).
|
||||
let blocked = turn_active
|
||||
.as_ref()
|
||||
.map(|f| f.load(std::sync::atomic::Ordering::SeqCst))
|
||||
.unwrap_or(false);
|
||||
let result = if blocked {
|
||||
Err(crate::commands::RepairHistoryBlocked)
|
||||
} else {
|
||||
Ok(self.repair_history(dry_run))
|
||||
};
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
ChatStateCommand::ReplaceSystemHead { prompt, reply } => {
|
||||
let changed = self.replace_system_head(&prompt);
|
||||
let _ = reply.send(changed);
|
||||
}
|
||||
ChatStateCommand::CachePromptText { text } => {
|
||||
self.state.prompt_texts.push(text);
|
||||
}
|
||||
ChatStateCommand::RecordCompactionAt { prompt_index } => {
|
||||
self.state.last_compaction_prompt_index = Some(prompt_index);
|
||||
}
|
||||
ChatStateCommand::Flush => {
|
||||
self.persistence.flush();
|
||||
}
|
||||
ChatStateCommand::UpdateCredentials { credentials } => {
|
||||
self.state.credentials = credentials;
|
||||
}
|
||||
ChatStateCommand::RestoreSnapshot(snapshot) => {
|
||||
self.restore_snapshot(*snapshot);
|
||||
}
|
||||
ChatStateCommand::BeginTurnCapture => {
|
||||
self.state.turn_capture = Some(state::TurnCaptureState {
|
||||
turn_start_offset: self.state.conversation.len(),
|
||||
pre_replacement_messages: Vec::new(),
|
||||
compaction_occurred: false,
|
||||
});
|
||||
}
|
||||
ChatStateCommand::AppendHarnessTraceItems { items } => {
|
||||
self.state.harness_trace_buffer.extend(items);
|
||||
}
|
||||
ChatStateCommand::FlushHarnessTraceTurn => {
|
||||
self.state.seal_harness_trace_turn();
|
||||
}
|
||||
ChatStateCommand::RepairDanglingAfterHarnessHalt { class } => {
|
||||
self.repair_dangling_after_harness_halt(class);
|
||||
}
|
||||
|
||||
// ═══ Queries ═══
|
||||
//
|
||||
// Read queries are pure reads — repair only at write boundaries:
|
||||
// `ChatState::new()` (startup) and `push_user_message()` (new turn).
|
||||
// `BuildConversationRequest` retains the guard because it is only
|
||||
// ever issued by the agent loop between turns, never by background tasks.
|
||||
ChatStateCommand::BuildConversationRequest {
|
||||
tool_definitions,
|
||||
memory_reminder,
|
||||
persist_memory_reminder,
|
||||
trace,
|
||||
conv_id,
|
||||
req_id,
|
||||
reply,
|
||||
} => {
|
||||
self.ensure_conversation_integrity();
|
||||
let request = self.build_conversation_request(
|
||||
tool_definitions,
|
||||
memory_reminder,
|
||||
persist_memory_reminder,
|
||||
trace,
|
||||
conv_id,
|
||||
req_id,
|
||||
);
|
||||
let _ = reply.send(request);
|
||||
}
|
||||
ChatStateCommand::GetConversation { reply } => {
|
||||
tracing::debug!(
|
||||
conversation_len = self.state.conversation.len(),
|
||||
"ChatState: cloning full conversation for GetConversation"
|
||||
);
|
||||
let _ = reply.send(self.state.conversation.clone());
|
||||
}
|
||||
ChatStateCommand::GetPromptIndex { reply } => {
|
||||
let _ = reply.send(self.state.prompt_index);
|
||||
}
|
||||
ChatStateCommand::GetLastCompactionPromptIndex { reply } => {
|
||||
let _ = reply.send(self.state.last_compaction_prompt_index);
|
||||
}
|
||||
ChatStateCommand::GetTotalTokens { reply } => {
|
||||
let _ = reply.send(self.state.total_tokens);
|
||||
}
|
||||
ChatStateCommand::GetLastTurnUsage { reply } => {
|
||||
let _ = reply.send(self.state.last_turn_usage.clone());
|
||||
}
|
||||
ChatStateCommand::GetPromptUsage { reply } => {
|
||||
let _ = reply.send(self.state.prompt_usage.clone());
|
||||
}
|
||||
ChatStateCommand::GetSessionUsage { reply } => {
|
||||
let _ = reply.send(self.state.session_usage.clone());
|
||||
}
|
||||
ChatStateCommand::GetEstimatedTotalTokens { reply } => {
|
||||
let _ =
|
||||
reply.send(self.state.total_tokens + self.state.estimated_tokens_since_model);
|
||||
}
|
||||
ChatStateCommand::GetSamplingConfig { reply } => {
|
||||
let _ = reply.send(self.state.sampling_config.clone());
|
||||
}
|
||||
ChatStateCommand::GetAgentEditedPaths { reply } => {
|
||||
let _ = reply.send(self.state.agent_edited_paths.clone());
|
||||
}
|
||||
ChatStateCommand::GetNotificationMeta { reply } => {
|
||||
let _ = reply.send(self.get_notification_meta());
|
||||
}
|
||||
ChatStateCommand::Snapshot { reply } => {
|
||||
tracing::debug!(
|
||||
conversation_len = self.state.conversation.len(),
|
||||
"ChatState: cloning full state for Snapshot"
|
||||
);
|
||||
let _ = reply.send(self.snapshot());
|
||||
}
|
||||
ChatStateCommand::TruncateToPromptIndex {
|
||||
target_prompt_index,
|
||||
reply,
|
||||
} => {
|
||||
self.truncate_to_prompt_index(target_prompt_index);
|
||||
self.state.turn_capture = None;
|
||||
self.state.prompt_usage = None;
|
||||
// `harness_trace_buffer` / `harness_trace_turns` intentionally
|
||||
// survive a rewind: the goal planner / verifier subagents
|
||||
// genuinely ran, so their sealed trace turns stay uploadable as
|
||||
// siblings even when the live turn that triggered them is undone.
|
||||
let _ = reply.send(());
|
||||
}
|
||||
ChatStateCommand::CheckAutoCompactNeeded {
|
||||
threshold_percent,
|
||||
reply,
|
||||
} => {
|
||||
let _ = reply.send(self.check_auto_compact_needed(threshold_percent));
|
||||
}
|
||||
ChatStateCommand::GetCredentials { reply } => {
|
||||
let _ = reply.send(self.state.credentials.clone());
|
||||
}
|
||||
ChatStateCommand::GetLastModelMetadata { reply } => {
|
||||
let _ = reply.send(self.get_last_model_metadata());
|
||||
}
|
||||
ChatStateCommand::TakeTurnMessages { reply } => {
|
||||
let result = self.state.turn_capture.take().map(|cap| {
|
||||
let mut messages = cap.pre_replacement_messages;
|
||||
messages.extend(
|
||||
Self::turn_tail(&self.state.conversation, cap.turn_start_offset)
|
||||
.iter()
|
||||
.cloned(),
|
||||
);
|
||||
TurnCapture {
|
||||
messages,
|
||||
compaction_occurred: cap.compaction_occurred,
|
||||
}
|
||||
});
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
ChatStateCommand::TakeHarnessTraceTurns { reply } => {
|
||||
// Defensive seal: a phase that recorded items but never flushed
|
||||
// still rides its own turn rather than stranding.
|
||||
self.state.seal_harness_trace_turn();
|
||||
let _ = reply.send(std::mem::take(&mut self.state.harness_trace_turns));
|
||||
}
|
||||
|
||||
// ─── Narrow targeted queries ──────────────────────────────────
|
||||
ChatStateCommand::GetConversationLen { reply } => {
|
||||
let _ = reply.send(self.get_conversation_len());
|
||||
}
|
||||
ChatStateCommand::HasDanglingToolCalls { reply } => {
|
||||
let _ = reply.send(self.has_dangling_tool_calls());
|
||||
}
|
||||
ChatStateCommand::GetLastAssistantText { reply } => {
|
||||
let _ = reply.send(self.get_last_assistant_text());
|
||||
}
|
||||
ChatStateCommand::GetFirstUserText { reply } => {
|
||||
let _ = reply.send(self.get_first_user_text());
|
||||
}
|
||||
ChatStateCommand::GetConversationItemAt { index, reply } => {
|
||||
let _ = reply.send(self.get_conversation_item_at(index));
|
||||
}
|
||||
ChatStateCommand::GetLastUserQueryText { reply } => {
|
||||
let _ = reply.send(self.get_last_user_query_text());
|
||||
}
|
||||
ChatStateCommand::GetConversationCounts { reply } => {
|
||||
let _ = reply.send(self.get_conversation_counts());
|
||||
}
|
||||
ChatStateCommand::GetSystemMessage { reply } => {
|
||||
let _ = reply.send(self.get_system_message());
|
||||
}
|
||||
ChatStateCommand::GetEstimatedMessagesTokens { reply } => {
|
||||
let _ = reply.send(state::estimate_messages_tokens(&self.state.conversation));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,532 @@
|
||||
//! Mutation handlers for the ChatStateActor.
|
||||
|
||||
use kigi_sampling_types::{
|
||||
ContentPart, ConversationItem, DanglingToolCallReason, dedup_duplicate_tool_results,
|
||||
repair_dangling_tool_calls,
|
||||
};
|
||||
|
||||
use super::ChatStateActor;
|
||||
use super::request_builder::HARD_CLEAR_PLACEHOLDER;
|
||||
use crate::events::ChatStateEvent;
|
||||
use crate::types::ChatStateSnapshot;
|
||||
|
||||
/// Static string label for tracing on `ConversationItem` (avoids pulling
|
||||
/// the `Role` enum into the format string).
|
||||
fn item_kind_str(item: &ConversationItem) -> &'static str {
|
||||
match item {
|
||||
ConversationItem::System(_) => "system",
|
||||
ConversationItem::User(_) => "user",
|
||||
ConversationItem::Assistant(_) => "assistant",
|
||||
ConversationItem::ToolResult(_) => "tool_result",
|
||||
ConversationItem::BackendToolCall(_) => "backend_tool_call",
|
||||
ConversationItem::Reasoning(_) => "reasoning",
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatStateActor {
|
||||
/// Repair any dangling tool calls in the conversation and persist the fix.
|
||||
///
|
||||
/// A "dangling" tool call is an assistant message with tool call IDs that
|
||||
/// lack matching `ToolResult` entries. This can happen when:
|
||||
/// - The user cancels (Ctrl+C) mid-tool-execution in a live session
|
||||
/// - The process crashes between pushing the assistant and tool results
|
||||
/// - The tokio task is aborted at an `.await` point
|
||||
///
|
||||
/// This method repairs the state in-place and persists the fix to disk.
|
||||
/// It is idempotent — calling it on a clean conversation is a cheap no-op
|
||||
/// (single forward scan, no allocations).
|
||||
///
|
||||
/// Only call at write boundaries where the previous turn is definitively
|
||||
/// over (`ChatState::new()`, `push_user_message()`, `BuildConversationRequest`).
|
||||
/// Do NOT call from read handlers — background tasks run concurrently with
|
||||
/// tool execution and would misidentify in-flight calls as dangling.
|
||||
pub(super) fn ensure_conversation_integrity(&mut self) {
|
||||
self.ensure_conversation_integrity_with_reason(DanglingToolCallReason::UserCancelled);
|
||||
}
|
||||
|
||||
/// Like [`Self::ensure_conversation_integrity`] but takes an explicit reason.
|
||||
pub(super) fn ensure_conversation_integrity_with_reason(
|
||||
&mut self,
|
||||
reason: DanglingToolCallReason,
|
||||
) {
|
||||
// In-place integrity repair can add/remove items ahead of an active capture's
|
||||
// boundary, so snapshot + rebase the offset like the replace/restore paths.
|
||||
self.snapshot_turn_slice();
|
||||
let deduped = dedup_duplicate_tool_results(&mut self.state.conversation);
|
||||
if deduped > 0 {
|
||||
tracing::info!(
|
||||
deduped_count = deduped,
|
||||
"Removed duplicate tool results in conversation"
|
||||
);
|
||||
}
|
||||
let repaired = repair_dangling_tool_calls(&mut self.state.conversation, reason);
|
||||
if repaired > 0 || deduped > 0 {
|
||||
tracing::info!(
|
||||
repaired_count = repaired,
|
||||
"Repaired dangling tool calls in conversation"
|
||||
);
|
||||
self.persistence.replace_history(&self.state.conversation);
|
||||
}
|
||||
self.rebase_turn_capture_offset();
|
||||
}
|
||||
|
||||
/// Repair dangling tool calls after a harness-initiated halt.
|
||||
pub(super) fn repair_dangling_after_harness_halt(&mut self, class: &'static str) {
|
||||
self.ensure_conversation_integrity_with_reason(DanglingToolCallReason::HarnessHalted {
|
||||
class,
|
||||
});
|
||||
}
|
||||
|
||||
/// Out-of-band history repair (`x.ai/session/repair`): run
|
||||
/// [`crate::compaction_utils::repair_history`] and persist changes via
|
||||
/// [`Self::replace_conversation`]. Unlike
|
||||
/// [`Self::ensure_conversation_integrity`], this also removes orphaned
|
||||
/// `ToolResult`s — the shape that bricks a session with provider 400s.
|
||||
/// `dry_run` only reports.
|
||||
pub(super) fn repair_history(
|
||||
&mut self,
|
||||
dry_run: bool,
|
||||
) -> crate::compaction_utils::HistoryRepairReport {
|
||||
if dry_run {
|
||||
let mut copy = self.state.conversation.clone();
|
||||
return crate::compaction_utils::repair_history(&mut copy);
|
||||
}
|
||||
let mut items = std::mem::take(&mut self.state.conversation);
|
||||
let report = crate::compaction_utils::repair_history(&mut items);
|
||||
if report.changed() {
|
||||
tracing::warn!(
|
||||
duplicates_removed = report.duplicates_removed,
|
||||
stripped_tool_result_ids = ?report.stripped_tool_result_ids,
|
||||
synthetic_results_inserted = report.synthetic_results_inserted,
|
||||
"History repair modified the conversation"
|
||||
);
|
||||
// Full replace: persists atomically and re-bases token estimates.
|
||||
self.replace_conversation(items, false);
|
||||
} else {
|
||||
// Nothing changed — put the conversation back untouched.
|
||||
self.state.conversation = items;
|
||||
}
|
||||
report
|
||||
}
|
||||
|
||||
/// Push any conversation item (user, assistant, or tool result) and persist it.
|
||||
pub(super) fn push_message(&mut self, item: ConversationItem) {
|
||||
let count_in_delta = !matches!(item, ConversationItem::Assistant(_));
|
||||
if count_in_delta {
|
||||
let estimated_tokens = super::state::estimate_item_tokens(&item);
|
||||
self.state.estimated_tokens_since_model += estimated_tokens;
|
||||
tracing::debug!(
|
||||
item_kind = item_kind_str(&item),
|
||||
estimated_tokens_delta = estimated_tokens,
|
||||
estimated_total = self.state.total_tokens + self.state.estimated_tokens_since_model,
|
||||
model_reported_total = self.state.total_tokens,
|
||||
"ChatState: push_message updated estimated_tokens_since_model"
|
||||
);
|
||||
}
|
||||
self.persistence.persist_message(&item);
|
||||
self.state.conversation.push(item);
|
||||
}
|
||||
|
||||
/// Push a user message, ensuring conversation integrity first.
|
||||
///
|
||||
/// When the user cancels a turn while the model was executing parallel
|
||||
/// tool calls, the conversation may have dangling tool call IDs. This
|
||||
/// method repairs them before appending the new message so the on-disk
|
||||
/// and in-memory state stay consistent.
|
||||
///
|
||||
/// Also runs [`prune_retained_conversation`] to eagerly hard-clear very
|
||||
/// old tool results from the in-memory state, bounding long-session
|
||||
/// retained memory without waiting for the context-window threshold.
|
||||
pub(super) fn push_user_message(&mut self, item: ConversationItem) {
|
||||
self.push_user_message_with_repair_reason(item, DanglingToolCallReason::UserCancelled);
|
||||
}
|
||||
|
||||
/// Like [`Self::push_user_message`] but takes an explicit repair reason.
|
||||
pub(super) fn push_user_message_with_repair_reason(
|
||||
&mut self,
|
||||
item: ConversationItem,
|
||||
reason: DanglingToolCallReason,
|
||||
) {
|
||||
self.ensure_conversation_integrity_with_reason(reason);
|
||||
let estimated_tokens = super::state::estimate_item_tokens(&item);
|
||||
self.state.estimated_tokens_since_model += estimated_tokens;
|
||||
tracing::debug!(
|
||||
item_kind = item_kind_str(&item),
|
||||
estimated_tokens_delta = estimated_tokens,
|
||||
estimated_total = self.state.total_tokens + self.state.estimated_tokens_since_model,
|
||||
model_reported_total = self.state.total_tokens,
|
||||
"ChatState: push_user_message updated estimated_tokens_since_model"
|
||||
);
|
||||
self.persistence.persist_message(&item);
|
||||
self.state.conversation.push(item);
|
||||
self.prune_retained_conversation();
|
||||
}
|
||||
|
||||
/// Eagerly hard-clear tool results from very old turns in the retained
|
||||
/// in-memory conversation, freeing the actual string bytes.
|
||||
///
|
||||
/// Unlike the API-copy pruning in `build_conversation_request` (which runs
|
||||
/// on a *clone* only when context > 50% full), this operates on
|
||||
/// `self.state.conversation` directly and runs after every user turn.
|
||||
///
|
||||
/// # What this does
|
||||
///
|
||||
/// Only **hard-clears** are applied (no soft-trim). Soft-trimming is a
|
||||
/// context-management operation that changes what the model sees;
|
||||
/// hard-clearing is a memory-management operation that replaces content
|
||||
/// that is so old the model should not need it again. The threshold is
|
||||
/// controlled by `PruningConfig::hard_clear_age_turns`.
|
||||
///
|
||||
/// # Retained-memory measurement
|
||||
///
|
||||
/// When any clearing occurs, a `tracing::debug!` event reports:
|
||||
/// - `hard_cleared` — number of tool results cleared
|
||||
/// - `bytes_freed` — approximate bytes recovered (sum of content lengths)
|
||||
/// - `conversation_len` — total item count after the pass
|
||||
///
|
||||
/// # Synthetic User items and turn-age accuracy
|
||||
///
|
||||
/// The shell can inject synthetic `User` items mid-turn (e.g. system
|
||||
/// corrective warnings) without calling `increment_prompt_index`. These
|
||||
/// do not represent real user turns. The backward scan here counts every
|
||||
/// `User` item as a turn boundary, so synthetic items would normally cause
|
||||
/// old tool results to appear older than they really are.
|
||||
///
|
||||
/// This is compensated by raising the effective clearing threshold by the
|
||||
/// number of synthetic User items (`total_user_items - prompt_index`).
|
||||
/// The result: a tool result is never cleared before `hard_clear_age_turns`
|
||||
/// REAL turns have elapsed, even in sessions with many synthetic messages.
|
||||
///
|
||||
/// # Replay / rewind correctness
|
||||
///
|
||||
/// `updates.jsonl` is **never touched**, so cross-compaction
|
||||
/// `replay_to_prompt` is unaffected. The pruned `chat_history.jsonl`
|
||||
/// on disk mirrors the in-memory state — both lose old bulk content but
|
||||
/// `updates.jsonl` retains the original data for replay.
|
||||
pub(super) fn prune_retained_conversation(&mut self) -> usize {
|
||||
if !self.pruning_config.enabled {
|
||||
return 0;
|
||||
}
|
||||
// Fast exit: not enough turns have elapsed for any hard-clear to apply.
|
||||
if self.state.prompt_index < self.pruning_config.hard_clear_age_turns {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Compute how many synthetic User items exist (system reminders, etc.).
|
||||
// Synthetic User items are NOT real user turns — they are injected by the
|
||||
// shell mid-turn and do not increment `prompt_index`. The naive backward
|
||||
// scan counts every User item as a turn boundary, so synthetic items make
|
||||
// old tool results appear older than they really are and can cause
|
||||
// premature hard-clears.
|
||||
//
|
||||
// Fix: raise the effective clearing threshold by the number of synthetic
|
||||
// User items. This guarantees a tool result is never cleared before
|
||||
// `hard_clear_age_turns` REAL turns have elapsed, regardless of how many
|
||||
// synthetic messages the session contains.
|
||||
let total_user_items = self
|
||||
.state
|
||||
.conversation
|
||||
.iter()
|
||||
.filter(|i| matches!(i, ConversationItem::User(_)))
|
||||
.count();
|
||||
let synthetic_count = total_user_items.saturating_sub(self.state.prompt_index);
|
||||
let effective_threshold = self
|
||||
.pruning_config
|
||||
.hard_clear_age_turns
|
||||
.saturating_add(synthetic_count);
|
||||
|
||||
let before_bytes = self.conversation_content_bytes();
|
||||
let mut cleared = 0usize;
|
||||
let mut turn_from_end: usize = 0;
|
||||
let mut seen_first_user = false;
|
||||
|
||||
for i in (0..self.state.conversation.len()).rev() {
|
||||
if matches!(&self.state.conversation[i], ConversationItem::User(_)) {
|
||||
if seen_first_user {
|
||||
turn_from_end += 1;
|
||||
}
|
||||
seen_first_user = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
let ConversationItem::ToolResult(tr) = &mut self.state.conversation[i] else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if turn_from_end < effective_threshold {
|
||||
continue;
|
||||
}
|
||||
|
||||
if tr.content.as_ref() != HARD_CLEAR_PLACEHOLDER {
|
||||
tr.content = std::sync::Arc::<str>::from(HARD_CLEAR_PLACEHOLDER);
|
||||
cleared += 1;
|
||||
}
|
||||
}
|
||||
|
||||
if cleared > 0 {
|
||||
let after_bytes = self.conversation_content_bytes();
|
||||
tracing::debug!(
|
||||
hard_cleared = cleared,
|
||||
bytes_freed = before_bytes.saturating_sub(after_bytes),
|
||||
conversation_len = self.state.conversation.len(),
|
||||
"ChatState: in-memory tool-result prune"
|
||||
);
|
||||
self.persistence.replace_history(&self.state.conversation);
|
||||
}
|
||||
|
||||
cleared
|
||||
}
|
||||
|
||||
/// Approximate byte footprint of all string content in the conversation.
|
||||
///
|
||||
/// Used for before/after measurement logging when pruning runs.
|
||||
/// Sums the byte lengths of all string fields; does not allocate.
|
||||
fn conversation_content_bytes(&self) -> usize {
|
||||
self.state
|
||||
.conversation
|
||||
.iter()
|
||||
.map(|item| match item {
|
||||
ConversationItem::System(s) => s.content.len(),
|
||||
ConversationItem::User(u) => u
|
||||
.content
|
||||
.iter()
|
||||
.map(|p| match p {
|
||||
ContentPart::Text { text } => text.len(),
|
||||
ContentPart::Image { url } => url.len(),
|
||||
})
|
||||
.sum::<usize>(),
|
||||
ConversationItem::Assistant(a) => a.content.len(),
|
||||
ConversationItem::ToolResult(tr) => tr.content.len(),
|
||||
ConversationItem::BackendToolCall(b) => b.text_summary().len(),
|
||||
ConversationItem::Reasoning(r) => {
|
||||
kigi_sampling_types::reasoning_item_text(r).len()
|
||||
+ r.encrypted_content.as_deref().map(str::len).unwrap_or(0)
|
||||
}
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Record accumulated token usage and emit an event.
|
||||
pub(super) fn record_token_usage(&mut self, total_tokens: u64) {
|
||||
self.state.estimated_tokens_since_model = 0;
|
||||
self.state.estimate_at_last_response =
|
||||
super::state::estimate_conversation_tokens(&self.state.conversation);
|
||||
self.state.total_tokens = total_tokens;
|
||||
self.send_event(ChatStateEvent::TokensUpdated { total_tokens });
|
||||
}
|
||||
|
||||
/// Stash the per-turn `TokenUsage` from the most recent model response.
|
||||
/// No event is emitted — this slot is read on demand at `PromptResponse`
|
||||
/// construction time, not pushed to subscribers.
|
||||
pub(super) fn record_last_turn_usage(&mut self, usage: kigi_sampling_types::TokenUsage) {
|
||||
self.state.last_turn_usage = Some(usage);
|
||||
}
|
||||
|
||||
pub(super) fn record_model_call_usage(
|
||||
&mut self,
|
||||
model_id: Option<String>,
|
||||
usage: &kigi_sampling_types::TokenUsage,
|
||||
api_duration_ms: Option<u64>,
|
||||
cost_usd_ticks: Option<i64>,
|
||||
) {
|
||||
let model_key = match model_id.as_deref() {
|
||||
Some(id) if !id.is_empty() => id,
|
||||
_ => self.state.sampling_config.model.as_str(),
|
||||
}
|
||||
.to_owned();
|
||||
self.state
|
||||
.prompt_usage
|
||||
.get_or_insert_default()
|
||||
.record_main_loop_call(&model_key, usage, api_duration_ms, cost_usd_ticks);
|
||||
self.state.session_usage.record_main_loop_call(
|
||||
&model_key,
|
||||
usage,
|
||||
api_duration_ms,
|
||||
cost_usd_ticks,
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn record_subagent_usage(
|
||||
&mut self,
|
||||
by_model: &[(String, crate::usage::UsageTotals)],
|
||||
attribute_to_prompt: bool,
|
||||
incomplete: bool,
|
||||
) {
|
||||
if by_model.is_empty() && !incomplete {
|
||||
return;
|
||||
}
|
||||
if attribute_to_prompt {
|
||||
self.state
|
||||
.prompt_usage
|
||||
.get_or_insert_default()
|
||||
.record_subagent(by_model, incomplete);
|
||||
}
|
||||
// The session ledger always folds, even when the usage is not
|
||||
// attributable to the open prompt (its pin may belong to an earlier
|
||||
// prompt). Reporting that gap is the coordinator's sticky flag's job —
|
||||
// never mark a different live prompt's ledger.
|
||||
self.state
|
||||
.session_usage
|
||||
.record_subagent(by_model, incomplete);
|
||||
}
|
||||
|
||||
pub(super) fn mark_usage_incomplete(&mut self, prompt: bool, session: bool) {
|
||||
if prompt {
|
||||
self.state
|
||||
.prompt_usage
|
||||
.get_or_insert_default()
|
||||
.mark_incomplete();
|
||||
}
|
||||
if session {
|
||||
self.state.session_usage.mark_incomplete();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn increment_prompt_index(&mut self) {
|
||||
self.state.prompt_usage = None;
|
||||
self.state.prompt_index += 1;
|
||||
self.send_event(ChatStateEvent::PromptIndexChanged {
|
||||
new_index: self.state.prompt_index,
|
||||
});
|
||||
}
|
||||
|
||||
/// Replace the entire conversation, persist, re-estimate `total_tokens`,
|
||||
/// and emit reset + token-update events.
|
||||
///
|
||||
/// Compaction replaces carry the provider-side overhead forward as a
|
||||
/// *ratio* (`base_estimate × provider_total ÷ estimate_at_last_response`,
|
||||
/// capped at the pre-compaction total; `base_estimate` when that estimate is
|
||||
/// 0) so the reseed neither springs back nor over-counts (see
|
||||
/// `COMPACTION.md`).
|
||||
pub(super) fn replace_conversation(
|
||||
&mut self,
|
||||
items: Vec<ConversationItem>,
|
||||
is_compaction: bool,
|
||||
) {
|
||||
self.snapshot_turn_slice();
|
||||
if is_compaction && let Some(cap) = &mut self.state.turn_capture {
|
||||
cap.compaction_occurred = true;
|
||||
}
|
||||
let pre_replace_total = self.state.total_tokens;
|
||||
// `harness_trace_buffer` / `harness_trace_turns` intentionally untouched:
|
||||
// the planner/verifier subagents ran, so their sealed trace turns survive
|
||||
// a conversation replace (same intent as the `TruncateToPromptIndex` arm).
|
||||
self.persistence.replace_history(&items);
|
||||
let base_estimate = super::state::estimate_conversation_tokens(&items);
|
||||
let mut estimated_tokens =
|
||||
if is_compaction && pre_replace_total > 0 && self.state.estimate_at_last_response > 0 {
|
||||
let ratio = pre_replace_total as f64 / self.state.estimate_at_last_response as f64;
|
||||
(base_estimate as f64 * ratio).round() as u64
|
||||
} else {
|
||||
base_estimate
|
||||
};
|
||||
// Compaction must never appear to increase usage.
|
||||
if is_compaction && pre_replace_total > 0 {
|
||||
estimated_tokens = estimated_tokens.min(pre_replace_total);
|
||||
}
|
||||
self.state.conversation = items;
|
||||
self.state.estimated_tokens_since_model = 0;
|
||||
self.state.total_tokens = estimated_tokens;
|
||||
self.state.estimate_at_last_response =
|
||||
super::state::estimate_conversation_tokens(&self.state.conversation);
|
||||
self.rebase_turn_capture_offset();
|
||||
self.send_event(ChatStateEvent::ConversationReset {
|
||||
new_len: self.state.conversation.len(),
|
||||
});
|
||||
self.send_event(ChatStateEvent::TokensUpdated {
|
||||
total_tokens: estimated_tokens,
|
||||
});
|
||||
}
|
||||
|
||||
/// Atomically swap the leading `System` message with `prompt` (or insert one
|
||||
/// if absent), persisting when changed. Runs inside the actor's command loop
|
||||
/// so it serializes with turn pushes — no lost-update race on a mid-turn
|
||||
/// reconnect. Returns whether the conversation changed.
|
||||
///
|
||||
/// The conversation is cloned (items are `Arc`-backed, so the clone is
|
||||
/// shallow) rather than `mem::take`n: `replace_conversation` snapshots the
|
||||
/// in-flight turn-capture tail from `state.conversation` before swapping,
|
||||
/// so the state must stay intact until then.
|
||||
pub(super) fn replace_system_head(&mut self, prompt: &str) -> bool {
|
||||
if let Some(ConversationItem::System(sys)) = self.state.conversation.first()
|
||||
&& crate::conversation_util::canonical_system_prompt_eq(sys.content.as_ref(), prompt)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let mut conversation = self.state.conversation.clone();
|
||||
let changed =
|
||||
crate::conversation_util::replace_or_insert_system_head(&mut conversation, prompt);
|
||||
debug_assert!(changed, "head mismatch must produce a change");
|
||||
self.replace_conversation(conversation, false);
|
||||
changed
|
||||
}
|
||||
|
||||
/// Restore all state fields from a snapshot.
|
||||
pub(super) fn restore_snapshot(&mut self, snap: ChatStateSnapshot) {
|
||||
self.snapshot_turn_slice();
|
||||
// Harness trace buffers are transient (not part of the snapshot) and
|
||||
// intentionally survive a restore — see `replace_conversation`.
|
||||
self.state.conversation = snap.conversation;
|
||||
self.rebase_turn_capture_offset();
|
||||
self.state.sampling_config = snap.sampling_config;
|
||||
self.state.prompt_index = snap.prompt_index;
|
||||
self.state.total_tokens = snap.total_tokens;
|
||||
self.state.estimated_tokens_since_model = 0;
|
||||
self.state.estimate_at_last_response = if snap.estimate_at_last_response > 0 {
|
||||
snap.estimate_at_last_response
|
||||
} else {
|
||||
super::state::estimate_conversation_tokens(&self.state.conversation)
|
||||
};
|
||||
self.state.agent_edited_paths = snap.agent_edited_paths;
|
||||
self.state.prompt_texts = snap.prompt_texts;
|
||||
self.state.stream_start_ms = snap.stream_start_ms;
|
||||
self.state.turn_start_ms = snap.turn_start_ms;
|
||||
self.state.last_compaction_prompt_index = snap.last_compaction_prompt_index;
|
||||
self.state.credentials = snap.credentials;
|
||||
// Drop abandoned prompt billing; session ledger is lifetime.
|
||||
self.state.prompt_usage = None;
|
||||
}
|
||||
|
||||
/// If turn capture is active, append the current turn's tail items into
|
||||
/// `pre_replacement_messages` before an in-place mutation shifts or drops them.
|
||||
pub(super) fn snapshot_turn_slice(&mut self) {
|
||||
if let Some(cap) = &mut self.state.turn_capture {
|
||||
cap.pre_replacement_messages
|
||||
.extend_from_slice(Self::turn_tail(
|
||||
&self.state.conversation,
|
||||
cap.turn_start_offset,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-base an active turn capture's start offset to the current conversation
|
||||
/// length after an in-place mutation, keeping the tail slice valid.
|
||||
pub(super) fn rebase_turn_capture_offset(&mut self) {
|
||||
if let Some(cap) = &mut self.state.turn_capture {
|
||||
cap.turn_start_offset = self.state.conversation.len();
|
||||
}
|
||||
}
|
||||
|
||||
/// Fail-safe `conversation[offset..]` for turn capture: a capture accounting
|
||||
/// slip must never abort the user's session (a raw index here SIGABRT-crashed
|
||||
/// a live CLI), so an out-of-range offset yields an empty slice — loud in dev
|
||||
/// via `debug_assert!`, with a prod breadcrumb via `error!`.
|
||||
pub(super) fn turn_tail(
|
||||
conversation: &[ConversationItem],
|
||||
offset: usize,
|
||||
) -> &[ConversationItem] {
|
||||
debug_assert!(
|
||||
offset <= conversation.len(),
|
||||
"turn_start_offset {offset} > len {}",
|
||||
conversation.len()
|
||||
);
|
||||
conversation.get(offset..).unwrap_or_else(|| {
|
||||
tracing::error!(
|
||||
offset,
|
||||
len = conversation.len(),
|
||||
"turn-capture offset past conversation end; trace tail dropped"
|
||||
);
|
||||
&[]
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
//! Query handlers for the ChatStateActor.
|
||||
|
||||
use super::ChatStateActor;
|
||||
use crate::compaction_utils::extract_last_user_query;
|
||||
use crate::events::ChatStateEvent;
|
||||
use crate::types::{AutoCompactTrigger, ChatStateSnapshot, ConversationCounts, NotificationMeta};
|
||||
|
||||
impl ChatStateActor {
|
||||
/// Build a notification meta from current timing state.
|
||||
pub(super) fn get_notification_meta(&self) -> NotificationMeta {
|
||||
NotificationMeta {
|
||||
stream_start_ms: self.state.stream_start_ms,
|
||||
turn_start_ms: self.state.turn_start_ms,
|
||||
}
|
||||
}
|
||||
|
||||
/// Take a full snapshot of the actor's state.
|
||||
pub(super) fn snapshot(&self) -> ChatStateSnapshot {
|
||||
ChatStateSnapshot {
|
||||
conversation: self.state.conversation.clone(),
|
||||
sampling_config: self.state.sampling_config.clone(),
|
||||
prompt_index: self.state.prompt_index,
|
||||
total_tokens: self.state.total_tokens,
|
||||
estimate_at_last_response: self.state.estimate_at_last_response,
|
||||
agent_edited_paths: self.state.agent_edited_paths.clone(),
|
||||
prompt_texts: self.state.prompt_texts.clone(),
|
||||
stream_start_ms: self.state.stream_start_ms,
|
||||
turn_start_ms: self.state.turn_start_ms,
|
||||
last_compaction_prompt_index: self.state.last_compaction_prompt_index,
|
||||
credentials: self.state.credentials.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Truncate conversation to a target prompt index (rewind).
|
||||
///
|
||||
/// Walks the conversation to find the Nth `User` item (where N =
|
||||
/// `target_prompt_index`), truncates everything from that point onward,
|
||||
/// truncates `prompt_texts` to match, persists, and emits `ConversationReset`.
|
||||
///
|
||||
/// Prompt index semantics:
|
||||
/// - 0 = no user turns have started (only system message, if any)
|
||||
/// - 1 = one user turn completed
|
||||
/// - N = N user turns completed
|
||||
///
|
||||
/// Truncating to `target_prompt_index = 1` keeps only items up to (but not
|
||||
/// including) the 2nd `User` message.
|
||||
pub(super) fn truncate_to_prompt_index(&mut self, target_prompt_index: usize) {
|
||||
if target_prompt_index >= self.state.prompt_index {
|
||||
// Nothing to truncate — already at or before the target.
|
||||
return;
|
||||
}
|
||||
|
||||
// Find the conversation position of the Nth User item.
|
||||
// Items before that position are kept; from that position onward removed.
|
||||
let mut user_count = 0;
|
||||
let mut truncate_at = self.state.conversation.len();
|
||||
|
||||
for (i, item) in self.state.conversation.iter().enumerate() {
|
||||
if matches!(item, kigi_sampling_types::ConversationItem::User(_)) {
|
||||
if user_count == target_prompt_index {
|
||||
truncate_at = i;
|
||||
break;
|
||||
}
|
||||
user_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
self.state.conversation.truncate(truncate_at);
|
||||
self.state.prompt_texts.truncate(target_prompt_index);
|
||||
self.state.prompt_index = target_prompt_index;
|
||||
self.state.total_tokens =
|
||||
super::state::estimate_conversation_tokens(&self.state.conversation);
|
||||
self.state.estimated_tokens_since_model = 0;
|
||||
self.state.estimate_at_last_response = self.state.total_tokens;
|
||||
|
||||
self.persistence.replace_history(&self.state.conversation);
|
||||
|
||||
self.send_event(ChatStateEvent::ConversationReset {
|
||||
new_len: self.state.conversation.len(),
|
||||
});
|
||||
}
|
||||
|
||||
/// Check if auto-compact is needed based on token utilization.
|
||||
///
|
||||
/// Returns `Some(AutoCompactTrigger)` if `total_tokens` exceeds
|
||||
/// `context_window * threshold_percent / 100`, otherwise `None`.
|
||||
pub(super) fn check_auto_compact_needed(
|
||||
&self,
|
||||
threshold_percent: u8,
|
||||
) -> Option<AutoCompactTrigger> {
|
||||
let context_window = self.state.sampling_config.context_window;
|
||||
let cw = context_window.get();
|
||||
|
||||
if kigi_token_estimation::exceeds_threshold(self.state.total_tokens, cw, threshold_percent)
|
||||
{
|
||||
let utilization_percent =
|
||||
kigi_token_estimation::usage_percentage_truncated_u8(self.state.total_tokens, cw);
|
||||
Some(AutoCompactTrigger {
|
||||
total_tokens: self.state.total_tokens,
|
||||
context_window,
|
||||
utilization_percent,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn get_last_model_metadata(&self) -> crate::commands::ModelMetadata {
|
||||
self.state
|
||||
.conversation
|
||||
.iter()
|
||||
.rev()
|
||||
.find_map(|item| {
|
||||
if let kigi_sampling_types::ConversationItem::Assistant(a) = item {
|
||||
Some(crate::commands::ModelMetadata {
|
||||
resolved_model_id: a.model_id.clone(),
|
||||
model_fingerprint: a.model_fingerprint.clone(),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
// ─── Narrow targeted queries ─────────────────────────────────────────────
|
||||
|
||||
/// Return the number of items in the conversation.
|
||||
pub(super) fn get_conversation_len(&self) -> usize {
|
||||
self.state.conversation.len()
|
||||
}
|
||||
|
||||
/// Whether the conversation has any assistant tool call without a matching
|
||||
/// `ToolResult` (the dangling-tool-call repair would fire on the next build).
|
||||
pub(super) fn has_dangling_tool_calls(&self) -> bool {
|
||||
kigi_sampling_types::has_dangling_tool_calls(&self.state.conversation)
|
||||
}
|
||||
|
||||
/// Return the text content of the last assistant message with non-empty text.
|
||||
///
|
||||
/// Walks the conversation backwards and returns the first `Assistant` item
|
||||
/// whose `content` field is non-empty after trimming. Returns `None` when
|
||||
/// no such item exists.
|
||||
pub(super) fn get_last_assistant_text(&self) -> Option<String> {
|
||||
self.state.conversation.iter().rev().find_map(|item| {
|
||||
if let kigi_sampling_types::ConversationItem::Assistant(a) = item
|
||||
&& !a.content.trim().is_empty()
|
||||
{
|
||||
return Some(a.content.as_ref().to_owned());
|
||||
}
|
||||
None
|
||||
})
|
||||
}
|
||||
|
||||
/// Return the text of the **first content part** of the first `User` message,
|
||||
/// if and only if that part is `ContentPart::Text`.
|
||||
///
|
||||
/// Matches the original call-site semantics exactly: if the first user
|
||||
/// message leads with a non-text part (e.g. an image in a multimodal
|
||||
/// prompt), this returns `None` rather than scanning further parts.
|
||||
/// Callers that need "any text part" rather than "first-part-is-text"
|
||||
/// should use `get_conversation()` directly.
|
||||
pub(super) fn get_first_user_text(&self) -> Option<String> {
|
||||
self.state.conversation.iter().find_map(|item| {
|
||||
if let kigi_sampling_types::ConversationItem::User(u) = item {
|
||||
// Only return text if the first part is Text — behaviour-preserving
|
||||
// w.r.t. the original `content.first().and_then(|p| if Text { … })`.
|
||||
u.content.first().and_then(|part| {
|
||||
if let kigi_sampling_types::ContentPart::Text { text } = part {
|
||||
Some(text.as_ref().to_owned())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Return the conversation item at `index`, or `None` if out of bounds.
|
||||
pub(super) fn get_conversation_item_at(
|
||||
&self,
|
||||
index: usize,
|
||||
) -> Option<kigi_sampling_types::ConversationItem> {
|
||||
self.state.conversation.get(index).cloned()
|
||||
}
|
||||
|
||||
/// Return the processed text of the last user query (metadata tags stripped).
|
||||
///
|
||||
/// Delegates to [`extract_last_user_query`] so the caller does not need a
|
||||
/// full conversation clone.
|
||||
pub(super) fn get_last_user_query_text(&self) -> Option<String> {
|
||||
extract_last_user_query(&self.state.conversation)
|
||||
}
|
||||
|
||||
/// Return conversation item counts by role without cloning any items.
|
||||
pub(super) fn get_conversation_counts(&self) -> ConversationCounts {
|
||||
let mut counts = ConversationCounts {
|
||||
total: self.state.conversation.len(),
|
||||
..Default::default()
|
||||
};
|
||||
for item in &self.state.conversation {
|
||||
match item {
|
||||
kigi_sampling_types::ConversationItem::User(_) => counts.user += 1,
|
||||
kigi_sampling_types::ConversationItem::Assistant(_) => {
|
||||
counts.assistant += 1;
|
||||
}
|
||||
kigi_sampling_types::ConversationItem::ToolResult(_) => {
|
||||
counts.tool_result += 1;
|
||||
}
|
||||
kigi_sampling_types::ConversationItem::System(_) => {}
|
||||
kigi_sampling_types::ConversationItem::BackendToolCall(_) => {}
|
||||
kigi_sampling_types::ConversationItem::Reasoning(_) => {}
|
||||
}
|
||||
}
|
||||
counts
|
||||
}
|
||||
|
||||
/// Return the first `System` message in the conversation, or `None`.
|
||||
pub(super) fn get_system_message(&self) -> Option<kigi_sampling_types::ConversationItem> {
|
||||
self.state
|
||||
.conversation
|
||||
.iter()
|
||||
.find(|item| matches!(item, kigi_sampling_types::ConversationItem::System(_)))
|
||||
.cloned()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,865 @@
|
||||
//! ConversationRequest assembly — image compaction, pruning, repair, memory injection.
|
||||
|
||||
use kigi_sampling_types::{
|
||||
ContentPart, ConversationItem, ConversationRequest, ToolSpec, TraceContext,
|
||||
};
|
||||
|
||||
use super::ChatStateActor;
|
||||
use crate::events::ChatStateEvent;
|
||||
use crate::types::PruningConfig;
|
||||
|
||||
/// Placeholder inserted when a tool result is hard-cleared.
|
||||
///
|
||||
/// `pub(super)` so that `mutations.rs` can use the same string when it
|
||||
/// hard-clears tool results in the retained in-memory conversation.
|
||||
pub(super) const HARD_CLEAR_PLACEHOLDER: &str = "[Tool result omitted — too old]";
|
||||
|
||||
/// Separator inserted between head and tail in soft-trimmed results.
|
||||
const SOFT_TRIM_SEPARATOR: &str = "\n\n[…trimmed…]\n\n";
|
||||
|
||||
impl ChatStateActor {
|
||||
/// Build a `ConversationRequest` from the current actor state.
|
||||
///
|
||||
/// 1. Evict oldest inline images when the inline-image bytes near 50 MB
|
||||
/// 2. Prune old tool results if over 50% context utilization
|
||||
/// 3. Optionally persist the memory reminder into actor state
|
||||
/// 4. Inject memory reminder into the request clone (if needed)
|
||||
/// 5. Assemble and return the `ConversationRequest`
|
||||
///
|
||||
/// # Repair invariant
|
||||
///
|
||||
/// The `BuildConversationRequest` command handler calls
|
||||
/// `ensure_conversation_integrity()` on the actor's own conversation
|
||||
/// **before** this function runs. The clone therefore starts from an
|
||||
/// already-repaired state, so there is no need to run
|
||||
/// `dedup_duplicate_tool_results` / `repair_dangling_tool_calls` on the
|
||||
/// clone — those would be O(n) no-ops.
|
||||
pub(super) fn build_conversation_request(
|
||||
&mut self,
|
||||
tool_definitions: Vec<ToolSpec>,
|
||||
memory_reminder: Option<String>,
|
||||
persist_memory_reminder: bool,
|
||||
trace: Option<Box<dyn TraceContext>>,
|
||||
conv_id: String,
|
||||
req_id: String,
|
||||
) -> ConversationRequest {
|
||||
let needs_prune = should_prune(
|
||||
self.state.total_tokens,
|
||||
self.state.sampling_config.context_window,
|
||||
);
|
||||
let mut memory_reminder = memory_reminder;
|
||||
if let Some(reminder) = memory_reminder.as_deref()
|
||||
&& persist_memory_reminder
|
||||
{
|
||||
// A live in-place inject can prepend a `System` item, shifting indices
|
||||
// under an active capture; snapshot + rebase like the other mutators.
|
||||
self.snapshot_turn_slice();
|
||||
let injected = inject_memory_reminder(&mut self.state.conversation, reminder);
|
||||
if injected {
|
||||
self.persistence.replace_history(&self.state.conversation);
|
||||
memory_reminder = None;
|
||||
}
|
||||
self.rebase_turn_capture_offset();
|
||||
}
|
||||
// Measure the exact serialized body and evict only once it approaches
|
||||
// the 50 MB ceiling. `conversation_body_bytes` is wire-accurate yet
|
||||
// cheap — it skips the multi-MB base64 escape scan (see its docs) — so
|
||||
// it runs inline on every turn with no blocking-thread offload.
|
||||
// Eviction rewrites earlier turns and busts the KV-cache prefix, so we
|
||||
// only pay it when the body is actually near the limit (the original
|
||||
// behavior — evicting every turn — caused chronic cache misses).
|
||||
let body_bytes = conversation_body_bytes(&self.state.conversation);
|
||||
let inline_images = inline_image_count(&self.state.conversation);
|
||||
let needs_image_compaction = body_bytes >= IMAGE_COMPACT_TRIGGER_BYTES;
|
||||
let needs_mutation = needs_prune || memory_reminder.is_some() || needs_image_compaction;
|
||||
|
||||
// Only allocate the mutable working copy when a mutation path is taken.
|
||||
let mut eviction: Option<ImageEvictionOutcome> = None;
|
||||
let items = if needs_mutation {
|
||||
let mut items = self.state.conversation.clone();
|
||||
|
||||
// Step 1: When the body nears the 50 MB ceiling, evict oldest
|
||||
// images down to the low-water mark (not just under the trigger).
|
||||
// Reclaiming a batch frees headroom for many subsequent image
|
||||
// turns, so the prefix is rewritten once and then stays cache-warm
|
||||
// — instead of re-triggering and re-busting the cache every turn.
|
||||
if needs_image_compaction {
|
||||
eviction = Some(compact_images_to_byte_budget(
|
||||
&mut items,
|
||||
body_bytes,
|
||||
IMAGE_COMPACT_RECLAIM_TARGET_BYTES,
|
||||
));
|
||||
}
|
||||
|
||||
// Step 2: Prune old tool results if context is > 50% utilized
|
||||
if needs_prune {
|
||||
prune_conversation(&mut items, &self.pruning_config);
|
||||
}
|
||||
|
||||
// Step 3: Inject memory reminder into the system message
|
||||
if let Some(reminder) = memory_reminder {
|
||||
inject_memory_reminder(&mut items, &reminder);
|
||||
}
|
||||
|
||||
items
|
||||
} else {
|
||||
// Hot path: no pruning, no memory reminder, no old images —
|
||||
// clone directly into the request without any intermediate mutation passes.
|
||||
self.state.conversation.clone()
|
||||
};
|
||||
|
||||
// Per-turn image-budget record for local verification. Emitted on the
|
||||
// ChatState event channel (chat-state can't reach the shell's unified
|
||||
// log directly); the session consumer writes it to the local log file.
|
||||
// Only on image-bearing turns to avoid noise.
|
||||
if inline_images > 0 {
|
||||
self.send_event(ChatStateEvent::ImageBudget {
|
||||
body_bytes,
|
||||
trigger_bytes: IMAGE_COMPACT_TRIGGER_BYTES,
|
||||
reclaim_target_bytes: IMAGE_COMPACT_RECLAIM_TARGET_BYTES,
|
||||
inline_images,
|
||||
needs_image_compaction,
|
||||
evicted: eviction.as_ref().map_or(0, |o| o.evicted),
|
||||
body_bytes_after: eviction.as_ref().map_or(body_bytes, |o| o.body_bytes_after),
|
||||
});
|
||||
}
|
||||
|
||||
// Step 4: Assemble request
|
||||
ConversationRequest {
|
||||
items,
|
||||
tools: tool_definitions,
|
||||
hosted_tools: vec![],
|
||||
tool_choice: None,
|
||||
model: Some(self.state.sampling_config.model.clone()),
|
||||
temperature: self.state.sampling_config.temperature,
|
||||
max_output_tokens: self.state.sampling_config.max_completion_tokens,
|
||||
top_p: self.state.sampling_config.top_p,
|
||||
x_grok_conv_id: Some(conv_id),
|
||||
x_grok_req_id: Some(req_id),
|
||||
x_grok_session_id: None,
|
||||
x_grok_turn_idx: None,
|
||||
x_grok_agent_id: None,
|
||||
x_grok_deployment_id: None,
|
||||
x_grok_user_id: None,
|
||||
trace,
|
||||
reasoning_effort: self.state.sampling_config.reasoning_effort,
|
||||
json_schema: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Pruning (standalone functions, no actor state needed)
|
||||
// ============================================================================
|
||||
|
||||
/// Check whether pruning should run based on context utilization.
|
||||
///
|
||||
/// Returns `true` when `total_tokens` exceeds 50% of `context_window`.
|
||||
pub(crate) fn should_prune(total_tokens: u64, context_window: std::num::NonZeroU64) -> bool {
|
||||
total_tokens > context_window.get() / 2
|
||||
}
|
||||
|
||||
/// Prune old, large tool results from the conversation in place.
|
||||
///
|
||||
/// Turn age is estimated by walking backward through the conversation and
|
||||
/// counting `User` items to determine which "turn" each tool result belongs to.
|
||||
pub(crate) fn prune_conversation(conversation: &mut [ConversationItem], config: &PruningConfig) {
|
||||
if !config.enabled {
|
||||
return;
|
||||
}
|
||||
|
||||
let mut turn_from_end: usize = 0;
|
||||
let mut seen_first_user = false;
|
||||
|
||||
for i in (0..conversation.len()).rev() {
|
||||
if matches!(&conversation[i], ConversationItem::User(_)) {
|
||||
if seen_first_user {
|
||||
turn_from_end += 1;
|
||||
}
|
||||
seen_first_user = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
let ConversationItem::ToolResult(tool_result) = &mut conversation[i] else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Never prune recent turns.
|
||||
if turn_from_end < config.keep_last_n_turns {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Hard clear: very old tool results → replace entirely.
|
||||
if turn_from_end >= config.hard_clear_age_turns {
|
||||
if tool_result.content.as_ref() != HARD_CLEAR_PLACEHOLDER {
|
||||
tool_result.content = std::sync::Arc::<str>::from(HARD_CLEAR_PLACEHOLDER);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// Soft trim: large tool results → keep head + tail.
|
||||
let content_len = tool_result.content.chars().count();
|
||||
if content_len > config.soft_trim_threshold {
|
||||
let head = safe_char_slice(&tool_result.content, 0, config.soft_trim_head);
|
||||
let tail = safe_char_slice_tail(&tool_result.content, config.soft_trim_tail);
|
||||
tool_result.content =
|
||||
std::sync::Arc::<str>::from(format!("{head}{SOFT_TRIM_SEPARATOR}{tail}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Image size-gated compaction (request-copy only)
|
||||
// ============================================================================
|
||||
|
||||
/// Replaces an inline image evicted to keep the request body under the proxy's
|
||||
/// 50 MB limit. Phrased so the model treats the image as gone rather than
|
||||
/// describing it from memory — a silently-stripped image otherwise induces
|
||||
/// confident hallucination of its contents.
|
||||
const IMAGE_COMPACT_PLACEHOLDER: &str = "[An earlier image was removed to keep the request within its size limit and is no longer visible. Do not describe or reason about its contents from memory; ask the user to re-share it if you need to see it again.]";
|
||||
|
||||
/// Hard request-body ceiling enforced by the inference proxy
|
||||
/// (nginx `proxy-body-size`). Bodies larger than this are rejected with HTTP
|
||||
/// 413 — or a connection reset before the response is written. Inline image
|
||||
/// `data:` URLs (base64) are the dominant term in this size.
|
||||
const MAX_REQUEST_BYTES: usize = 50 * 1024 * 1024;
|
||||
|
||||
/// Evict old images once the serialized body reaches this size.
|
||||
///
|
||||
/// We gate on the exact body (see [`conversation_body_bytes`]) — system prompt,
|
||||
/// all message text, tool results, and image `data:` URLs are all counted
|
||||
/// precisely. This sits 3 MB below [`MAX_REQUEST_BYTES`] as headroom for the
|
||||
/// only parts of the wire request the body measurement does **not** include:
|
||||
/// - **tool definitions** — sent alongside the conversation but not part of it
|
||||
/// (tool JSON schemas + MCP tools); this is the bulk of the gap.
|
||||
/// - the request envelope and sampling params.
|
||||
/// - the small delta between our internal `ContentPart` JSON and the public-API
|
||||
/// wire format (the dominant base64 image bytes are identical in both).
|
||||
///
|
||||
/// The uncounted remainder is only sub-MB to low-MB in practice, so 3 MB covers
|
||||
/// it without needlessly sacrificing image capacity. The sampler's reactive 413
|
||||
/// image-strip is the final backstop if this is ever under-estimated.
|
||||
///
|
||||
/// Below this threshold every image stays in place so the KV-cache prefix is
|
||||
/// byte-stable across turns; eviction rewrites earlier turns and busts the
|
||||
/// prefix cache, so we only pay that cost when a 413 is actually near.
|
||||
pub(crate) const IMAGE_COMPACT_TRIGGER_BYTES: usize = MAX_REQUEST_BYTES - 3 * 1024 * 1024;
|
||||
|
||||
/// Low-water mark that eviction reclaims down to once it fires (hysteresis).
|
||||
///
|
||||
/// Eviction is **gated** at [`IMAGE_COMPACT_TRIGGER_BYTES`] but **reclaims** to
|
||||
/// this strictly lower mark. Evicting only enough to clear the trigger means
|
||||
/// the next image-bearing turn re-crosses it and evicts again — rewriting the
|
||||
/// prefix and busting the KV cache on essentially every turn once the body sits
|
||||
/// at the ceiling. Dropping to half the hard limit instead frees ~25 MB of
|
||||
/// headroom, so the prefix is rewritten once and then stays stable (cache-warm)
|
||||
/// across many turns until the headroom is consumed again. The oldest images
|
||||
/// (least useful) are sacrificed in a batch rather than one-per-turn — a
|
||||
/// high-water trigger paired with a lower reclaim mark (classic hysteresis).
|
||||
pub(crate) const IMAGE_COMPACT_RECLAIM_TARGET_BYTES: usize = MAX_REQUEST_BYTES / 2;
|
||||
|
||||
// Hysteresis invariant: eviction is gated at the trigger but reclaims to a
|
||||
// strictly lower mark, so one batch eviction buys many cache-warm turns rather
|
||||
// than re-triggering (and re-busting the prompt cache) every turn at the
|
||||
// ceiling. Enforced at compile time so the two constants can't drift together.
|
||||
const _: () = assert!(IMAGE_COMPACT_RECLAIM_TARGET_BYTES < IMAGE_COMPACT_TRIGGER_BYTES);
|
||||
|
||||
/// An [`std::io::Write`] sink that counts bytes instead of storing them. Lets
|
||||
/// us measure a `serde_json` encoding's length without allocating the full
|
||||
/// (potentially tens-of-MB) output buffer.
|
||||
#[derive(Default)]
|
||||
struct ByteCounter(usize);
|
||||
|
||||
impl std::io::Write for ByteCounter {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.0 += buf.len();
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Exact JSON-serialized byte length of any value, measured through a
|
||||
/// [`ByteCounter`] so no encoded buffer is allocated. JSON quoting and string
|
||||
/// escaping are captured precisely (not estimated from field lengths).
|
||||
fn serialized_json_bytes<T: serde::Serialize + ?Sized>(value: &T) -> usize {
|
||||
let mut counter = ByteCounter::default();
|
||||
if let Err(err) = serde_json::to_writer(&mut counter, value) {
|
||||
// Serializing in-memory state to a byte sink is infallible in
|
||||
// practice; if it ever fails, fall back to the bytes counted so far
|
||||
// (a lower bound) rather than forcing a needless compaction.
|
||||
tracing::warn!(%err, "failed to measure serialized size");
|
||||
}
|
||||
counter.0
|
||||
}
|
||||
|
||||
/// Serialized JSON frame of one image content part with an empty URL —
|
||||
/// `{"type":"image","url":""}`. The real payload adds exactly `url.len()` on
|
||||
/// top: an inline base64 `data:` URL contains no JSON-escaped characters, so
|
||||
/// its encoded length equals its raw length. Identical in our internal JSON and
|
||||
/// on the public-API wire (the base64 bytes are the same in both).
|
||||
const IMAGE_PART_FRAME_BYTES: usize = r#"{"type":"image","url":""}"#.len();
|
||||
|
||||
/// Exact serialized size of a single inline image part (frame + raw URL bytes).
|
||||
fn image_part_bytes(url: &str) -> usize {
|
||||
IMAGE_PART_FRAME_BYTES + url.len()
|
||||
}
|
||||
|
||||
/// Count of inline images in the conversation — for observability only.
|
||||
fn inline_image_count(conversation: &[ConversationItem]) -> usize {
|
||||
conversation
|
||||
.iter()
|
||||
.filter_map(|item| match item {
|
||||
ConversationItem::User(u) => Some(u),
|
||||
_ => None,
|
||||
})
|
||||
.flat_map(|u| u.content.iter())
|
||||
.filter(|p| matches!(p, ContentPart::Image { .. }))
|
||||
.count()
|
||||
}
|
||||
|
||||
/// Outcome of [`compact_images_to_byte_budget`], surfaced for logging and
|
||||
/// local verification.
|
||||
pub(crate) struct ImageEvictionOutcome {
|
||||
/// Number of inline images replaced with the placeholder.
|
||||
pub evicted: usize,
|
||||
/// Estimated serialized body size after eviction (`current_bytes` minus the
|
||||
/// net bytes freed) — at or below `target_bytes` once enough images go.
|
||||
pub body_bytes_after: usize,
|
||||
}
|
||||
|
||||
/// Exact serialized size of the conversation body — the figure the inference
|
||||
/// proxy weighs against its 50 MB limit — computed **without** scanning the
|
||||
/// multi-MB base64 image payloads.
|
||||
///
|
||||
/// `serde_json` escape-scans every byte of every string, so encoding the real
|
||||
/// conversation would walk tens of MB of base64 on every turn. Instead we
|
||||
/// serialize a copy with image URLs blanked (cheap: only the small non-image
|
||||
/// content — system prompt, message text, tool results — is scanned, and it is
|
||||
/// measured *exactly*, escaping included) and add back each URL's raw length.
|
||||
/// Because base64 never escapes, that length is its exact serialized
|
||||
/// contribution, so the result is byte-for-byte the true body size.
|
||||
///
|
||||
/// The blanking copy is cheap: image data lives behind `Arc<str>`, so cloning
|
||||
/// only bumps refcounts and the blanked clone drops them without copying bytes.
|
||||
fn conversation_body_bytes(conversation: &[ConversationItem]) -> usize {
|
||||
let mut blanked = conversation.to_vec();
|
||||
let mut image_url_bytes = 0usize;
|
||||
for item in &mut blanked {
|
||||
if let ConversationItem::User(user) = item {
|
||||
for part in &mut user.content {
|
||||
if let ContentPart::Image { url } = part {
|
||||
image_url_bytes += url.len();
|
||||
*url = std::sync::Arc::<str>::from("");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
serialized_json_bytes(&blanked) + image_url_bytes
|
||||
}
|
||||
|
||||
/// Replace the oldest inline images with [`IMAGE_COMPACT_PLACEHOLDER`] until
|
||||
/// the serialized request body drops back to `target_bytes`, keeping the
|
||||
/// newest images. `current_bytes` is the already-measured whole-body size (see
|
||||
/// [`conversation_body_bytes`]); each eviction drops `running` by the image
|
||||
/// part's exact serialized size minus the placeholder that replaces it, so it
|
||||
/// tracks the true body byte-for-byte as images are removed.
|
||||
///
|
||||
/// Operates on a mutable slice — intended for the request *copy* so the stored
|
||||
/// conversation is never modified.
|
||||
///
|
||||
/// ## Cache behavior
|
||||
///
|
||||
/// Eviction is **oldest-first**, which is sticky by construction: because we
|
||||
/// always retain the newest images, an image only transitions image →
|
||||
/// placeholder as *newer/larger* payloads push the body past the limit, never
|
||||
/// placeholder → image within a stable prefix. (Token compaction removes old
|
||||
/// turns wholesale and can free room to restore a previously-evicted image,
|
||||
/// but that already rewrites the prefix and invalidates the server-side prompt
|
||||
/// cache, so the restore is free.)
|
||||
///
|
||||
/// The caller gates eviction at [`IMAGE_COMPACT_TRIGGER_BYTES`] but passes the
|
||||
/// lower [`IMAGE_COMPACT_RECLAIM_TARGET_BYTES`] as `target_bytes`, so one
|
||||
/// eviction reclaims a batch of the oldest images and frees headroom for many
|
||||
/// later image turns. This turns "rewrite the prefix on essentially every turn
|
||||
/// once at the ceiling" into one larger, rare rewrite followed by a long
|
||||
/// cache-warm stretch — the prefix-cache cost of dropping the oldest (least
|
||||
/// useful) image is paid infrequently instead of per turn.
|
||||
///
|
||||
/// This replaces the previous policy — strip every image older than the most
|
||||
/// recent user turn on *every* request — which (a) busted the prompt-cache
|
||||
/// prefix on the turn after any image, and (b) dropped images the model still
|
||||
/// needed one turn later, causing it to hallucinate their contents.
|
||||
pub(crate) fn compact_images_to_byte_budget(
|
||||
conversation: &mut [ConversationItem],
|
||||
current_bytes: usize,
|
||||
target_bytes: usize,
|
||||
) -> ImageEvictionOutcome {
|
||||
if current_bytes <= target_bytes {
|
||||
return ImageEvictionOutcome {
|
||||
evicted: 0,
|
||||
body_bytes_after: current_bytes,
|
||||
};
|
||||
}
|
||||
|
||||
// The text part each evicted image is replaced with. Measured once: every
|
||||
// eviction shrinks the body by the image part's bytes and grows it back by
|
||||
// this placeholder's bytes, so the net saving is `image - placeholder`.
|
||||
let placeholder = ContentPart::Text {
|
||||
text: std::sync::Arc::<str>::from(IMAGE_COMPACT_PLACEHOLDER),
|
||||
};
|
||||
let placeholder_bytes = serialized_json_bytes(&placeholder);
|
||||
|
||||
// (item_idx, part_idx, exact serialized image-part bytes) for every inline
|
||||
// image, oldest-first.
|
||||
let mut images: Vec<(usize, usize, usize)> = Vec::new();
|
||||
for (i, item) in conversation.iter().enumerate() {
|
||||
if let ConversationItem::User(user) = item {
|
||||
for (j, part) in user.content.iter().enumerate() {
|
||||
if let ContentPart::Image { url } = part {
|
||||
images.push((i, j, image_part_bytes(url)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Evict oldest-first until the body fits again, keeping the newest images.
|
||||
let mut running = current_bytes;
|
||||
let mut evicted = 0usize;
|
||||
for &(i, j, image_bytes) in &images {
|
||||
if running <= target_bytes {
|
||||
break;
|
||||
}
|
||||
if let ConversationItem::User(user) = &mut conversation[i]
|
||||
&& let Some(part) = user.content.get_mut(j)
|
||||
{
|
||||
*part = placeholder.clone();
|
||||
// Net body saving: the image part leaves, the placeholder takes its
|
||||
// slot. Everything else (siblings, commas, brackets) is untouched,
|
||||
// so this is the exact change in the serialized body size.
|
||||
running = running.saturating_sub(image_bytes.saturating_sub(placeholder_bytes));
|
||||
evicted += 1;
|
||||
}
|
||||
}
|
||||
|
||||
ImageEvictionOutcome {
|
||||
evicted,
|
||||
body_bytes_after: running,
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Memory reminder injection
|
||||
// ============================================================================
|
||||
|
||||
use crate::types::MEMORY_CONTEXT_OPEN_TAG;
|
||||
|
||||
/// Upsert a memory reminder into the conversation's system message.
|
||||
///
|
||||
/// If the first item is a `System` message, any previously injected memory
|
||||
/// reminder section is replaced in-place; otherwise the reminder is appended.
|
||||
/// If no system message exists, a new `System` item is prepended.
|
||||
///
|
||||
/// Returns `true` when the conversation was changed.
|
||||
pub(super) fn inject_memory_reminder(items: &mut Vec<ConversationItem>, reminder: &str) -> bool {
|
||||
let reminder = reminder.trim();
|
||||
if reminder.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
if let Some(ConversationItem::System(sys)) = items.first_mut() {
|
||||
upsert_memory_reminder_text(&mut sys.content, reminder)
|
||||
} else {
|
||||
items.insert(0, ConversationItem::system(reminder));
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn upsert_memory_reminder_text(system_prompt: &mut std::sync::Arc<str>, reminder: &str) -> bool {
|
||||
let existing_start = system_prompt
|
||||
.find(MEMORY_CONTEXT_OPEN_TAG)
|
||||
.map(|idx| system_prompt[..idx].trim_end_matches('\n').len());
|
||||
|
||||
let updated: String = if let Some(prefix_len) = existing_start {
|
||||
let prefix = system_prompt[..prefix_len].trim_end_matches('\n');
|
||||
if prefix.is_empty() {
|
||||
reminder.to_string()
|
||||
} else {
|
||||
format!("{prefix}\n\n{reminder}")
|
||||
}
|
||||
} else if system_prompt.trim_end() == reminder {
|
||||
system_prompt.as_ref().to_owned()
|
||||
} else if system_prompt.is_empty() {
|
||||
reminder.to_string()
|
||||
} else {
|
||||
format!("{}\n\n{reminder}", system_prompt.trim_end_matches('\n'))
|
||||
};
|
||||
|
||||
if system_prompt.as_ref() == updated.as_str() {
|
||||
false
|
||||
} else {
|
||||
*system_prompt = std::sync::Arc::<str>::from(updated);
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// String helpers
|
||||
// ============================================================================
|
||||
|
||||
fn safe_char_slice(s: &str, start: usize, count: usize) -> String {
|
||||
s.chars().skip(start).take(count).collect()
|
||||
}
|
||||
|
||||
fn safe_char_slice_tail(s: &str, count: usize) -> String {
|
||||
let total = s.chars().count();
|
||||
if count >= total {
|
||||
return s.to_string();
|
||||
}
|
||||
s.chars().skip(total - count).collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn should_prune_gating() {
|
||||
use std::num::NonZeroU64;
|
||||
let cw = NonZeroU64::new(10000).unwrap();
|
||||
assert!(!should_prune(1000, cw)); // 10%
|
||||
assert!(should_prune(6000, cw)); // 60%
|
||||
assert!(!should_prune(5000, cw)); // 50% exact (> not >=)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_disabled_is_noop() {
|
||||
let mut conv = vec![ConversationItem::tool_result("c1", "x".repeat(10_000))];
|
||||
let config = PruningConfig {
|
||||
enabled: false,
|
||||
..Default::default()
|
||||
};
|
||||
prune_conversation(&mut conv, &config);
|
||||
if let ConversationItem::ToolResult(ref tr) = conv[0] {
|
||||
assert_eq!(tr.content.len(), 10_000);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_memory_into_existing_system() {
|
||||
let mut items = vec![
|
||||
ConversationItem::system("You are helpful."),
|
||||
ConversationItem::user("hi"),
|
||||
];
|
||||
inject_memory_reminder(&mut items, "Remember: user likes rust");
|
||||
if let ConversationItem::System(ref sys) = items[0] {
|
||||
assert!(sys.content.contains("Remember: user likes rust"));
|
||||
assert!(sys.content.starts_with("You are helpful."));
|
||||
}
|
||||
assert_eq!(items.len(), 2); // no new item added
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_memory_prepends_when_no_system() {
|
||||
let mut items = vec![ConversationItem::user("hi")];
|
||||
inject_memory_reminder(&mut items, "Remember: user likes rust");
|
||||
assert_eq!(items.len(), 2);
|
||||
assert!(matches!(&items[0], ConversationItem::System(_)));
|
||||
}
|
||||
|
||||
// -- image size-gated compaction tests --
|
||||
|
||||
/// A user message with a small fixed inline image.
|
||||
fn user_with_image(text: &str) -> ConversationItem {
|
||||
let mut item = ConversationItem::user(text);
|
||||
item.add_image("data:image/png;base64,iVBORw0KGgo=");
|
||||
item
|
||||
}
|
||||
|
||||
/// A user message carrying an inline image whose `data:` URL is exactly
|
||||
/// `url_bytes` long (must be >= the data-URL prefix length).
|
||||
fn user_with_image_of_bytes(text: &str, url_bytes: usize) -> ConversationItem {
|
||||
const PREFIX: &str = "data:image/png;base64,";
|
||||
let pad = url_bytes.saturating_sub(PREFIX.len());
|
||||
let mut item = ConversationItem::user(text);
|
||||
item.add_image(format!("{PREFIX}{}", "A".repeat(pad)));
|
||||
item
|
||||
}
|
||||
|
||||
fn has_image(item: &ConversationItem) -> bool {
|
||||
matches!(
|
||||
item,
|
||||
ConversationItem::User(u)
|
||||
if u.content.iter().any(|p| matches!(p, ContentPart::Image { .. }))
|
||||
)
|
||||
}
|
||||
|
||||
fn has_placeholder(item: &ConversationItem) -> bool {
|
||||
matches!(
|
||||
item,
|
||||
ConversationItem::User(u) if u.content.iter().any(|p| matches!(
|
||||
p,
|
||||
ContentPart::Text { text } if text.as_ref() == IMAGE_COMPACT_PLACEHOLDER
|
||||
))
|
||||
)
|
||||
}
|
||||
|
||||
// Images are sized ~100 KB so the ~235 B placeholder that replaces an
|
||||
// evicted image is negligible: each eviction frees ~one image's bytes.
|
||||
const TEST_IMG_BYTES: usize = 100_000;
|
||||
|
||||
#[test]
|
||||
fn no_eviction_when_at_or_below_target() {
|
||||
// Multiple old image turns are *retained* when the body already fits —
|
||||
// the key behavior change from the old "strip everything but newest".
|
||||
let mut conv = vec![
|
||||
ConversationItem::system("sys"),
|
||||
user_with_image_of_bytes("first", TEST_IMG_BYTES),
|
||||
ConversationItem::assistant("a"),
|
||||
user_with_image_of_bytes("second", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("third", TEST_IMG_BYTES),
|
||||
];
|
||||
// current < target: nothing to do.
|
||||
compact_images_to_byte_budget(&mut conv, 300_000, 400_000);
|
||||
assert_eq!(conv.iter().filter(|i| has_image(i)).count(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicts_oldest_until_under_target() {
|
||||
let mut conv = vec![
|
||||
user_with_image_of_bytes("oldest", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("middle", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("newest", TEST_IMG_BYTES),
|
||||
];
|
||||
// current 300k, target 250k: evicting the oldest (~100 KB) fits.
|
||||
compact_images_to_byte_budget(&mut conv, 300_000, 250_000);
|
||||
assert!(has_placeholder(&conv[0]), "oldest evicted");
|
||||
assert!(has_image(&conv[1]), "middle kept");
|
||||
assert!(has_image(&conv[2]), "newest kept");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicts_more_oldest_for_lower_target() {
|
||||
let mut conv = vec![
|
||||
user_with_image_of_bytes("oldest", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("middle", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("newest", TEST_IMG_BYTES),
|
||||
];
|
||||
// current 300k, target 150k: must drop the two oldest to fit.
|
||||
compact_images_to_byte_budget(&mut conv, 300_000, 150_000);
|
||||
assert!(has_placeholder(&conv[0]));
|
||||
assert!(has_placeholder(&conv[1]));
|
||||
assert!(has_image(&conv[2]), "newest kept");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_reclaims_batch_to_low_water_mark() {
|
||||
// Mirror production: a body sitting just over the trigger, made of many
|
||||
// equal images, is reclaimed in one pass down to the low-water mark —
|
||||
// dropping a *batch* of the oldest, not just the one image needed to
|
||||
// clear the trigger. This is the hysteresis that keeps the prefix
|
||||
// cache-warm for the following turns.
|
||||
let img_bytes = 1_000_000usize; // ~1 MB url each
|
||||
let n = (IMAGE_COMPACT_TRIGGER_BYTES / img_bytes) + 2; // body just over trigger
|
||||
let mut conv: Vec<ConversationItem> = (0..n)
|
||||
.map(|i| user_with_image_of_bytes(&format!("i{i}"), img_bytes))
|
||||
.collect();
|
||||
let current = n * img_bytes;
|
||||
assert!(current > IMAGE_COMPACT_TRIGGER_BYTES);
|
||||
|
||||
compact_images_to_byte_budget(&mut conv, current, IMAGE_COMPACT_RECLAIM_TARGET_BYTES);
|
||||
|
||||
let kept = conv.iter().filter(|i| has_image(i)).count();
|
||||
let evicted = conv.iter().filter(|i| has_placeholder(i)).count();
|
||||
|
||||
// Clearing only the trigger would evict ~3 images; reclaiming to the
|
||||
// low-water mark (~half the ceiling) must evict far more.
|
||||
assert!(
|
||||
evicted > n / 4,
|
||||
"expected batch eviction to the low-water mark, only {evicted}/{n} evicted"
|
||||
);
|
||||
// Oldest-first stops at the mark, so the most recent image survives.
|
||||
assert!(kept > 0);
|
||||
assert!(
|
||||
has_image(conv.last().unwrap()),
|
||||
"most recent image must be retained"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicts_all_when_target_below_one_image() {
|
||||
let mut conv = vec![
|
||||
user_with_image_of_bytes("a", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("b", TEST_IMG_BYTES),
|
||||
];
|
||||
compact_images_to_byte_budget(&mut conv, 200_000, 50_000);
|
||||
assert!(has_placeholder(&conv[0]));
|
||||
assert!(has_placeholder(&conv[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_keeps_newest_and_is_idempotent() {
|
||||
let mut conv = vec![
|
||||
user_with_image_of_bytes("i0", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("i1", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("i2", TEST_IMG_BYTES),
|
||||
user_with_image_of_bytes("i3", TEST_IMG_BYTES),
|
||||
];
|
||||
// current 400k, target 250k: drop the two oldest, keep the newest two.
|
||||
compact_images_to_byte_budget(&mut conv, 400_000, 250_000);
|
||||
assert!(has_placeholder(&conv[0]) && has_placeholder(&conv[1]));
|
||||
assert!(has_image(&conv[2]) && has_image(&conv[3]));
|
||||
|
||||
// Re-running with the now-smaller body is a no-op (sticky): the two
|
||||
// surviving images already fit.
|
||||
compact_images_to_byte_budget(&mut conv, 200_000, 250_000);
|
||||
assert!(has_placeholder(&conv[0]) && has_placeholder(&conv[1]));
|
||||
assert!(has_image(&conv[2]) && has_image(&conv[3]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicted_image_uses_honest_placeholder() {
|
||||
let mut conv = vec![user_with_image_of_bytes("x", TEST_IMG_BYTES)];
|
||||
compact_images_to_byte_budget(&mut conv, 100_000, 10);
|
||||
assert!(has_placeholder(&conv[0]));
|
||||
}
|
||||
|
||||
// -- conversation_body_bytes tests --
|
||||
|
||||
#[test]
|
||||
fn conversation_body_bytes_empty_is_json_array() {
|
||||
// serde encodes an empty slice as "[]" (2 bytes).
|
||||
assert_eq!(conversation_body_bytes(&[]), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_body_bytes_matches_serde_json_exactly() {
|
||||
// The blank-and-add-URLs measurement must equal a full serde_json
|
||||
// encode byte-for-byte — including non-image content and string
|
||||
// escaping. The `"` in the system text is escaped by serde; the
|
||||
// measurement must account for it.
|
||||
let conv = vec![
|
||||
ConversationItem::system("system \"quoted\" prompt"),
|
||||
user_with_image("look"),
|
||||
ConversationItem::assistant("a longer assistant reply with text"),
|
||||
ConversationItem::user("plain follow-up turn"),
|
||||
];
|
||||
let expected = serde_json::to_vec(&conv).unwrap().len();
|
||||
assert_eq!(conversation_body_bytes(&conv), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_body_bytes_matches_serde_json_with_large_image() {
|
||||
// Exact even for a multi-KB base64 payload — the scan we deliberately
|
||||
// skip still lands on the same byte count.
|
||||
let conv = vec![user_with_image_of_bytes("big", 50_000)];
|
||||
let expected = serde_json::to_vec(&conv).unwrap().len();
|
||||
assert_eq!(conversation_body_bytes(&conv), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_body_bytes_small_image_is_below_trigger() {
|
||||
// A normal small inline image must not trip the 50 MB gate — the case
|
||||
// the cache-miss fix preserves.
|
||||
let conv = vec![
|
||||
user_with_image("old"),
|
||||
ConversationItem::assistant("reply"),
|
||||
ConversationItem::user("current"),
|
||||
];
|
||||
assert!(conversation_body_bytes(&conv) < IMAGE_COMPACT_TRIGGER_BYTES);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn conversation_body_bytes_large_image_reaches_trigger() {
|
||||
let conv = vec![user_with_image_of_bytes("big", IMAGE_COMPACT_TRIGGER_BYTES)];
|
||||
assert!(conversation_body_bytes(&conv) >= IMAGE_COMPACT_TRIGGER_BYTES);
|
||||
}
|
||||
|
||||
// -- edge cases: exactness, boundaries, ordering --
|
||||
|
||||
#[test]
|
||||
fn body_bytes_parity_multi_image_unicode_escaping() {
|
||||
// The gate is only as correct as this equality. Exercise multiple
|
||||
// images in one turn, multibyte unicode (passed through, not escaped),
|
||||
// and chars serde *does* escape (`"`, `\`, control).
|
||||
let mut turn = ConversationItem::user("two pics 🚀 with \"quotes\" and \\ slash");
|
||||
turn.add_image("data:image/png;base64,AAAA");
|
||||
turn.add_image("data:image/png;base64,BBBBBB");
|
||||
let conv = vec![
|
||||
ConversationItem::system("sys 日本語 \t control"),
|
||||
turn,
|
||||
ConversationItem::assistant("reply"),
|
||||
ConversationItem::user("plain follow-up"),
|
||||
];
|
||||
assert_eq!(
|
||||
conversation_body_bytes(&conv),
|
||||
serde_json::to_vec(&conv).unwrap().len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_eviction_when_exactly_at_target() {
|
||||
// The no-op guard is `current <= target`; pin the inclusive boundary.
|
||||
let mut conv = vec![user_with_image_of_bytes("a", TEST_IMG_BYTES)];
|
||||
compact_images_to_byte_budget(&mut conv, 250_000, 250_000);
|
||||
assert!(has_image(&conv[0]), "exactly at target must not evict");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminates_when_placeholder_exceeds_image() {
|
||||
// Tiny images: each "saving" saturates to 0, but the loop must still
|
||||
// terminate and replace every image when the target is unreachable.
|
||||
let mut conv = vec![
|
||||
user_with_image_of_bytes("a", 40),
|
||||
user_with_image_of_bytes("b", 40),
|
||||
];
|
||||
compact_images_to_byte_budget(&mut conv, 1_000, 10);
|
||||
assert!(has_placeholder(&conv[0]) && has_placeholder(&conv[1]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn evicts_oldest_image_parts_first() {
|
||||
// `has_image`/`has_placeholder` are per-item, so count actual image
|
||||
// parts to verify oldest-first ordering across parts within a turn.
|
||||
fn image_parts(conv: &[ConversationItem]) -> usize {
|
||||
conv.iter()
|
||||
.filter_map(|i| match i {
|
||||
ConversationItem::User(u) => Some(u),
|
||||
_ => None,
|
||||
})
|
||||
.flat_map(|u| u.content.iter())
|
||||
.filter(|p| matches!(p, ContentPart::Image { .. }))
|
||||
.count()
|
||||
}
|
||||
let mut newest = ConversationItem::user("newest turn");
|
||||
newest.add_image(format!(
|
||||
"data:image/png;base64,{}",
|
||||
"A".repeat(TEST_IMG_BYTES)
|
||||
));
|
||||
newest.add_image(format!(
|
||||
"data:image/png;base64,{}",
|
||||
"B".repeat(TEST_IMG_BYTES)
|
||||
));
|
||||
let mut conv = vec![user_with_image_of_bytes("oldest", TEST_IMG_BYTES), newest];
|
||||
assert_eq!(image_parts(&conv), 3);
|
||||
|
||||
// ~300k body, reclaim to 150k: drop the two oldest, keep the newest.
|
||||
compact_images_to_byte_budget(&mut conv, 300_000, 150_000);
|
||||
assert_eq!(image_parts(&conv), 1, "newest image survives");
|
||||
assert!(has_placeholder(&conv[0]), "oldest turn evicted");
|
||||
assert!(has_image(&conv[1]), "newest turn keeps an image");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escaped_remote_url_is_a_lower_bound_only() {
|
||||
// base64 `data:` URLs are exact; a remote URL with a JSON-escaped char
|
||||
// under-counts by the escape bytes. Pin that documented bound so the
|
||||
// measurement can't silently drift past it.
|
||||
let mut item = ConversationItem::user("");
|
||||
item.add_image(r#"https://example.com/a"b"#);
|
||||
let conv = vec![item];
|
||||
assert!(conversation_body_bytes(&conv) <= serde_json::to_vec(&conv).unwrap().len());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,396 @@
|
||||
//! Internal state types for the ChatStateActor.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use kigi_sampling_types::{
|
||||
ConversationItem, DanglingToolCallReason, SamplingConfig, TokenUsage,
|
||||
dedup_duplicate_tool_results, repair_dangling_tool_calls,
|
||||
};
|
||||
|
||||
use crate::types::Credentials;
|
||||
use crate::usage::UsageLedger;
|
||||
|
||||
/// Bytes/4 estimate of the system prompt portion of a [`ConversationItem`].
|
||||
/// Returns 0 for non-system items so callers can pipe through whatever they
|
||||
/// have without unwrapping.
|
||||
pub fn estimate_system_message_tokens(item: &ConversationItem) -> u64 {
|
||||
match item {
|
||||
ConversationItem::System(s) => kigi_token_estimation::estimate_tokens(&s.content),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Bytes/4 estimate of one tool definition (name + description + the
|
||||
/// JSON-serialized parameters).
|
||||
pub fn estimate_tool_definition_tokens(td: &kigi_sampling_types::ToolDefinition) -> u64 {
|
||||
let name_len = td.function.name.len();
|
||||
let desc_len = td.function.description.as_deref().map_or(0, |d| d.len());
|
||||
let params_len = td.function.parameters.to_string().len();
|
||||
((name_len + desc_len + params_len) as u64) / kigi_token_estimation::BYTES_PER_TOKEN
|
||||
}
|
||||
|
||||
/// Sum [`estimate_tool_definition_tokens`] across a slice.
|
||||
pub fn estimate_tool_definitions_tokens(tds: &[kigi_sampling_types::ToolDefinition]) -> u64 {
|
||||
tds.iter().map(estimate_tool_definition_tokens).sum()
|
||||
}
|
||||
|
||||
/// Bytes/4 estimate for a single [`ConversationItem`].
|
||||
///
|
||||
/// Images are counted at [`kigi_token_estimation::IMAGE_TOKEN_ESTIMATE`] each.
|
||||
/// Shared by [`estimate_conversation_tokens`] and [`estimate_messages_tokens`]
|
||||
/// so the per-variant arithmetic stays in one place.
|
||||
pub fn estimate_item_tokens(item: &ConversationItem) -> u64 {
|
||||
use kigi_sampling_types::ContentPart;
|
||||
match item {
|
||||
ConversationItem::System(s) => kigi_token_estimation::estimate_tokens(&s.content),
|
||||
ConversationItem::User(u) => {
|
||||
let mut bytes: usize = 0;
|
||||
let mut images: u64 = 0;
|
||||
for p in &u.content {
|
||||
match p {
|
||||
ContentPart::Text { text } => bytes += text.len(),
|
||||
ContentPart::Image { .. } => images += 1,
|
||||
}
|
||||
}
|
||||
(bytes as u64) / kigi_token_estimation::BYTES_PER_TOKEN
|
||||
+ kigi_token_estimation::estimate_image_tokens(images)
|
||||
}
|
||||
ConversationItem::Assistant(a) => {
|
||||
let bytes = a.content.len()
|
||||
+ a.tool_calls
|
||||
.iter()
|
||||
.map(|tc| tc.arguments.len())
|
||||
.sum::<usize>();
|
||||
(bytes as u64) / kigi_token_estimation::BYTES_PER_TOKEN
|
||||
}
|
||||
ConversationItem::ToolResult(tr) => kigi_token_estimation::estimate_tokens(&tr.content),
|
||||
ConversationItem::BackendToolCall(b) => {
|
||||
kigi_token_estimation::estimate_tokens(&b.text_summary())
|
||||
}
|
||||
ConversationItem::Reasoning(r) => {
|
||||
// Summary + content text follow the standard bytes-per-token
|
||||
// estimate; encrypted blobs are base64 and don't survive
|
||||
// tokenization 1:1, so estimate at len/4 as well.
|
||||
let text_bytes = kigi_sampling_types::reasoning_item_text(r).len();
|
||||
let enc_bytes = r.encrypted_content.as_deref().map(str::len).unwrap_or(0);
|
||||
((text_bytes + enc_bytes) as u64) / kigi_token_estimation::BYTES_PER_TOKEN
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Estimate token footprint: text bytes / 4, images at the per-image
|
||||
/// constant defined by [`kigi_token_estimation::IMAGE_TOKEN_ESTIMATE`].
|
||||
pub fn estimate_conversation_tokens(items: &[ConversationItem]) -> u64 {
|
||||
items.iter().map(estimate_item_tokens).sum()
|
||||
}
|
||||
|
||||
/// grok-build's [`ItemTokenCounter`](kigi_compaction::ItemTokenCounter)
|
||||
/// for the shared compaction engine: the bytes/4 estimate grok-build already
|
||||
/// uses to drive its compaction triggers, exposed through the seam so the
|
||||
/// shared budgeting math gets the *same* trusted count.
|
||||
///
|
||||
/// Where another host plugs a real BPE tokenizer into the same seam,
|
||||
/// grok-build estimates instead, reusing [`estimate_item_tokens`] so the
|
||||
/// per-variant arithmetic (images, reasoning blobs, tool-call args) stays in
|
||||
/// one place.
|
||||
pub struct EstimatedItemTokenCounter;
|
||||
|
||||
impl kigi_compaction::ItemTokenCounter<ConversationItem> for EstimatedItemTokenCounter {
|
||||
fn count_item_tokens(&self, item: &ConversationItem) -> u32 {
|
||||
// The estimate is a `u64`; a single item never approaches `u32::MAX`
|
||||
// tokens, but saturate rather than wrap if one somehow does.
|
||||
estimate_item_tokens(item).try_into().unwrap_or(u32::MAX)
|
||||
}
|
||||
}
|
||||
|
||||
/// Bytes/4 estimate of every non-system item in `items`.
|
||||
pub fn estimate_messages_tokens(items: &[ConversationItem]) -> u64 {
|
||||
items
|
||||
.iter()
|
||||
.filter(|i| !matches!(i, ConversationItem::System(_)))
|
||||
.map(estimate_item_tokens)
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Internal mutable state for the ChatStateActor.
|
||||
///
|
||||
/// All fields are owned exclusively by the actor task — no locks needed.
|
||||
pub(crate) struct ChatState {
|
||||
/// The full conversation history.
|
||||
pub conversation: Vec<ConversationItem>,
|
||||
/// Current sampling configuration (model, context window, etc.).
|
||||
pub sampling_config: SamplingConfig,
|
||||
/// Current prompt index (incremented per user turn).
|
||||
pub prompt_index: usize,
|
||||
/// Cached prompt texts for rewind preview.
|
||||
pub prompt_texts: Vec<String>,
|
||||
/// Accumulated token usage.
|
||||
pub total_tokens: u64,
|
||||
/// Timestamp when the current stream started (epoch ms).
|
||||
pub stream_start_ms: Option<i64>,
|
||||
/// Timestamp when the current turn started (epoch ms).
|
||||
pub turn_start_ms: Option<i64>,
|
||||
/// File paths the agent has edited.
|
||||
pub agent_edited_paths: BTreeSet<String>,
|
||||
/// Prompt index at which the last compaction occurred.
|
||||
pub last_compaction_prompt_index: Option<usize>,
|
||||
/// Opaque credential secrets (api key, optional extra auth, client version).
|
||||
/// Stored opaquely — the actor never interprets them.
|
||||
pub credentials: Credentials,
|
||||
/// Bytes/4 estimate of tokens added since the last `record_token_usage`.
|
||||
/// Used by `check_preflight_overflow` to detect context window overflows
|
||||
/// between model responses.
|
||||
pub estimated_tokens_since_model: u64,
|
||||
/// Bytes/4 estimate of the conversation as of the last `record_token_usage`
|
||||
/// (or last reseed). `total_tokens − estimate_at_last_response` is the
|
||||
/// provider-side overhead carried across compaction.
|
||||
pub estimate_at_last_response: u64,
|
||||
/// Per-turn token usage from the most recent model response.
|
||||
/// Stashed by `record_last_turn_usage()` and read at `PromptResponse`
|
||||
/// construction to enrich `_meta` with `inputTokens` / `outputTokens` /
|
||||
/// `cachedReadTokens`. `None` means no model turn has completed yet
|
||||
/// in this session (or this is a freshly restored session that did not
|
||||
/// persist last_turn_usage). Always overwritten by the most recent turn —
|
||||
/// historical turns are not retained here.
|
||||
pub last_turn_usage: Option<TokenUsage>,
|
||||
/// Billing for the open prompt (cleared on next prompt; not persisted).
|
||||
pub prompt_usage: Option<UsageLedger>,
|
||||
/// Lifetime session billing (not persisted).
|
||||
pub session_usage: UsageLedger,
|
||||
/// Offset-based turn capture state. `Some` = capture active, `None` = inactive.
|
||||
/// Cleared on `TakeTurnMessages` (consumed), `BeginTurnCapture` (new turn),
|
||||
/// and `TruncateToPromptIndex` (rewind abandons the turn).
|
||||
pub(super) turn_capture: Option<TurnCaptureState>,
|
||||
/// Accumulator for the in-progress harness-subagent trace phase (the goal
|
||||
/// planner at `setup_goal`, or one verifier skeptic panel). Synthetic
|
||||
/// `task` pairs recorded via `AppendHarnessTraceItems` land here;
|
||||
/// `FlushHarnessTraceTurn` seals the accumulated items into one entry of
|
||||
/// `harness_trace_turns`. Independent of `turn_capture` (the planner runs
|
||||
/// ahead of `BeginTurnCapture`) and never enters the live `conversation`.
|
||||
pub(super) harness_trace_buffer: Vec<ConversationItem>,
|
||||
/// Sealed harness trace turns awaiting drain by the agent, which uploads
|
||||
/// each as its own sibling `turn_{N}` artifact so orchestrators can
|
||||
/// discover harness subagents via their `<subagent_result>` footer.
|
||||
/// Drained by `TakeHarnessTraceTurns` at the end of the user-facing turn.
|
||||
pub(super) harness_trace_turns: Vec<Vec<ConversationItem>>,
|
||||
}
|
||||
|
||||
/// Tracks which conversation items belong to the current turn without
|
||||
/// cloning every pushed item into a side buffer.
|
||||
///
|
||||
/// Instead of duplicating each `ConversationItem` on push, we record the
|
||||
/// conversation length at capture start (`turn_start_offset`). At take
|
||||
/// time, `conversation[turn_start_offset..]` gives us the turn's items
|
||||
/// with a single bulk clone.
|
||||
///
|
||||
/// When `replace_conversation` or `restore_snapshot` replaces the vec
|
||||
/// mid-turn, we snapshot `conversation[turn_start_offset..]` into
|
||||
/// `pre_replacement_messages` before the old vec is dropped, and reset
|
||||
/// the offset to the new vec's length.
|
||||
pub(super) struct TurnCaptureState {
|
||||
/// Index into `conversation` where this turn's messages start.
|
||||
pub turn_start_offset: usize,
|
||||
/// Messages saved from before a conversation replacement (compaction,
|
||||
/// snapshot restore). Extended (not replaced) if multiple replacements
|
||||
/// occur in one turn.
|
||||
pub pre_replacement_messages: Vec<ConversationItem>,
|
||||
/// Whether compaction occurred during this capture.
|
||||
pub compaction_occurred: bool,
|
||||
}
|
||||
|
||||
impl ChatState {
|
||||
/// Create a new `ChatState` with the given conversation and sampling config,
|
||||
/// all other fields defaulted.
|
||||
///
|
||||
/// Repairs any dangling tool calls in the initial conversation. This handles
|
||||
/// the race condition where the process was killed mid-tool-execution and
|
||||
/// `chat_history.jsonl` has an assistant message with tool call IDs that
|
||||
/// lack matching `ToolResult` entries. Without this, the in-memory state
|
||||
/// would carry broken conversation history until the next `build_request`.
|
||||
pub fn new(mut conversation: Vec<ConversationItem>, sampling_config: SamplingConfig) -> Self {
|
||||
let deduped = dedup_duplicate_tool_results(&mut conversation);
|
||||
if deduped > 0 {
|
||||
tracing::info!(
|
||||
deduped_count = deduped,
|
||||
"Removed duplicate tool results in initial conversation"
|
||||
);
|
||||
}
|
||||
let repaired =
|
||||
repair_dangling_tool_calls(&mut conversation, DanglingToolCallReason::UserCancelled);
|
||||
if repaired > 0 {
|
||||
tracing::info!(
|
||||
repaired_count = repaired,
|
||||
"Repaired dangling tool calls in initial conversation (likely from a previous crash)"
|
||||
);
|
||||
}
|
||||
|
||||
let initial_tokens = estimate_conversation_tokens(&conversation);
|
||||
|
||||
Self {
|
||||
conversation,
|
||||
sampling_config,
|
||||
prompt_index: 0,
|
||||
prompt_texts: Vec::new(),
|
||||
total_tokens: initial_tokens,
|
||||
stream_start_ms: None,
|
||||
turn_start_ms: None,
|
||||
agent_edited_paths: BTreeSet::new(),
|
||||
last_compaction_prompt_index: None,
|
||||
credentials: Credentials::default(),
|
||||
estimated_tokens_since_model: 0,
|
||||
estimate_at_last_response: initial_tokens,
|
||||
last_turn_usage: None,
|
||||
prompt_usage: None,
|
||||
session_usage: UsageLedger::default(),
|
||||
turn_capture: None,
|
||||
harness_trace_buffer: Vec::new(),
|
||||
harness_trace_turns: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Seal the items accumulated since the last flush into one harness trace
|
||||
/// turn. No-op when nothing was recorded since the last seal. Shared by the
|
||||
/// explicit `FlushHarnessTraceTurn` (one call per harness phase) and the
|
||||
/// defensive seal in `TakeHarnessTraceTurns`.
|
||||
pub(super) fn seal_harness_trace_turn(&mut self) {
|
||||
if !self.harness_trace_buffer.is_empty() {
|
||||
let turn = std::mem::take(&mut self.harness_trace_buffer);
|
||||
self.harness_trace_turns.push(turn);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_sampling_config() -> SamplingConfig {
|
||||
SamplingConfig {
|
||||
base_url: "https://api.example.com".to_string(),
|
||||
model: "test-model".to_string(),
|
||||
max_completion_tokens: None,
|
||||
temperature: None,
|
||||
top_p: None,
|
||||
api_backend: Default::default(),
|
||||
extra_headers: Default::default(),
|
||||
context_window: std::num::NonZeroU64::new(128_000).unwrap(),
|
||||
reasoning_effort: None,
|
||||
stream_tool_calls: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimated_item_token_counter_matches_estimate_item_tokens() {
|
||||
use kigi_compaction::ItemTokenCounter;
|
||||
|
||||
let counter = EstimatedItemTokenCounter;
|
||||
let items = vec![
|
||||
ConversationItem::system("you are a helpful assistant"),
|
||||
ConversationItem::user("fix the login bug in auth.rs"),
|
||||
ConversationItem::assistant("let me look at the file"),
|
||||
ConversationItem::tool_result("tc1", "fn login() {}"),
|
||||
];
|
||||
for item in &items {
|
||||
assert_eq!(
|
||||
u64::from(counter.count_item_tokens(item)),
|
||||
estimate_item_tokens(item),
|
||||
"counter must report the same trusted count as estimate_item_tokens"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_state_has_correct_defaults() {
|
||||
let state = ChatState::new(vec![], test_sampling_config());
|
||||
assert_eq!(state.prompt_index, 0);
|
||||
assert_eq!(state.total_tokens, 0); // empty conversation → 0
|
||||
assert!(state.conversation.is_empty());
|
||||
assert!(state.agent_edited_paths.is_empty());
|
||||
assert!(state.prompt_texts.is_empty());
|
||||
assert!(state.stream_start_ms.is_none());
|
||||
assert!(state.turn_start_ms.is_none());
|
||||
assert!(state.last_compaction_prompt_index.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_state_preserves_initial_conversation() {
|
||||
let items = vec![
|
||||
ConversationItem::system("sys"),
|
||||
ConversationItem::user("hello"),
|
||||
];
|
||||
let state = ChatState::new(items, test_sampling_config());
|
||||
assert_eq!(state.conversation.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_state_estimates_tokens_from_conversation() {
|
||||
// 4000 bytes of text per item, bytes / 4 = 1000 tokens each
|
||||
let items = vec![
|
||||
ConversationItem::system("x".repeat(4000).as_str()),
|
||||
ConversationItem::user("y".repeat(4000).as_str()),
|
||||
ConversationItem::assistant("z".repeat(4000).as_str()),
|
||||
ConversationItem::tool_result("call-1", "w".repeat(4000).as_str()),
|
||||
];
|
||||
let state = ChatState::new(items, test_sampling_config());
|
||||
assert_eq!(state.total_tokens, 4000); // 4 * (4000/4)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_system_message_tokens_only_counts_system_items() {
|
||||
let sys = ConversationItem::system("a".repeat(400));
|
||||
assert_eq!(estimate_system_message_tokens(&sys), 100);
|
||||
let user = ConversationItem::user("hello");
|
||||
assert_eq!(estimate_system_message_tokens(&user), 0);
|
||||
let asst = ConversationItem::assistant("hi");
|
||||
assert_eq!(estimate_system_message_tokens(&asst), 0);
|
||||
let tr = ConversationItem::tool_result("call-1", "x".repeat(4000).as_str());
|
||||
assert_eq!(estimate_system_message_tokens(&tr), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_tool_definition_tokens_counts_name_desc_params() {
|
||||
// Empty parameters serialize to "null" (4 bytes) in the JSON-string len
|
||||
let td = kigi_sampling_types::ToolDefinition::function(
|
||||
"search",
|
||||
Some("find a file"),
|
||||
serde_json::json!({}),
|
||||
);
|
||||
// name=6 + desc=11 + params=`{}`.len()=2 = 19, /4 = 4
|
||||
assert_eq!(estimate_tool_definition_tokens(&td), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_messages_tokens_excludes_system_and_sums_rest() {
|
||||
// 4000 bytes per item -> 1000 tokens each.
|
||||
let items = vec![
|
||||
ConversationItem::system("x".repeat(4000).as_str()),
|
||||
ConversationItem::user("y".repeat(4000).as_str()),
|
||||
ConversationItem::assistant("z".repeat(4000).as_str()),
|
||||
ConversationItem::tool_result("call-1", "w".repeat(4000).as_str()),
|
||||
];
|
||||
// Total = 4000 (4 items * 1000), system = 1000, messages = 3000.
|
||||
assert_eq!(estimate_conversation_tokens(&items), 4000);
|
||||
assert_eq!(estimate_messages_tokens(&items), 3000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_messages_tokens_zero_when_only_system() {
|
||||
let items = vec![ConversationItem::system("x".repeat(4000).as_str())];
|
||||
assert_eq!(estimate_messages_tokens(&items), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_messages_tokens_zero_for_empty() {
|
||||
assert_eq!(estimate_messages_tokens(&[]), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimate_tool_definitions_tokens_sums_across_slice() {
|
||||
let a =
|
||||
kigi_sampling_types::ToolDefinition::function("a", None::<&str>, serde_json::json!({}));
|
||||
let b =
|
||||
kigi_sampling_types::ToolDefinition::function("b", None::<&str>, serde_json::json!({}));
|
||||
let single = estimate_tool_definition_tokens(&a);
|
||||
assert_eq!(estimate_tool_definitions_tokens(&[a, b]), single * 2);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user