ZacharyZhang-NY
d6c20fc13f
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).
2203 lines
96 KiB
Rust
2203 lines
96 KiB
Rust
//! Turn-execution concern for `SessionActor` (`handle_prompt`, turn-end,
|
||
//! sampling loop).
|
||
use super::*;
|
||
/// Synthetic tool the model calls to return its schema-constrained final answer
|
||
/// on backends that can't constrain output natively (Messages API). Intercepted
|
||
/// in the loop, never executed as a real tool.
|
||
const STRUCTURED_OUTPUT_TOOL: &str = "StructuredOutput";
|
||
/// Max times the model may re-call `StructuredOutput` with non-conforming args
|
||
/// before the turn ends with the last validation error.
|
||
const STRUCTURED_OUTPUT_MAX_RETRIES: u32 = 3;
|
||
/// What a `StructuredOutput` tool call means for the turn (see
|
||
/// `handle_structured_output_tool_call`).
|
||
enum StructuredOutputStep {
|
||
/// Accepted, or retries exhausted: the carried result is the final output.
|
||
Complete(Result<serde_json::Value, String>),
|
||
/// Non-conforming args; a corrective tool_result was pushed — re-sample.
|
||
Retry,
|
||
/// No sole StructuredOutput call (absent, or co-emitted with real tools that
|
||
/// should run this round).
|
||
Proceed,
|
||
}
|
||
/// Parse `raw` as JSON and validate it against a `validator` compiled once per
|
||
/// turn. Returns the value on success, or a human-readable error (surfaced to
|
||
/// the model on retry and to the client as `structuredOutputError`). A `validator`
|
||
/// of `Err` means the user's schema itself was invalid.
|
||
fn validate_structured_output(
|
||
validator: &Result<jsonschema::Validator, String>,
|
||
raw: &str,
|
||
) -> Result<serde_json::Value, String> {
|
||
let validator = validator.as_ref().map_err(Clone::clone)?;
|
||
let value: serde_json::Value = serde_json::from_str(raw.trim())
|
||
.map_err(|e| format!("model output was not valid JSON: {e}"))?;
|
||
match validator.validate(&value) {
|
||
Ok(()) => Ok(value),
|
||
Err(e) => Err(format!("output does not match the required schema: {e}")),
|
||
}
|
||
}
|
||
/// Result of the turn-end usage drain (and cancel's no-drain snapshot).
|
||
///
|
||
/// **Ledger marks** only when [`Self::fail_closed`]. Sticky and background
|
||
/// live are **report-level only** (tokens still land on the session ledger).
|
||
pub(super) struct UsageDrainOutcome {
|
||
/// Query failure, FG still live after timeout/cancel. Marks both
|
||
/// the prompt and session bills incomplete. (True apply-miss stains
|
||
/// ledgers at fold time via `mark_apply_miss_incomplete`, not here.)
|
||
pub(super) fail_closed: bool,
|
||
/// A background child is still running: only this prompt's report is
|
||
/// incomplete; its spend reaches the session ledger at completion.
|
||
pub(super) background_live: bool,
|
||
/// Pin-scoped sticky (session-only attribution or apply-miss report).
|
||
/// Report incomplete only — does not stain ledgers by itself.
|
||
pub(super) sticky_report: bool,
|
||
}
|
||
impl UsageDrainOutcome {
|
||
/// Wire / attach incomplete: fail-closed ∪ background ∪ sticky.
|
||
pub(super) fn report_incomplete(&self) -> bool {
|
||
self.fail_closed || self.background_live || self.sticky_report
|
||
}
|
||
/// Map an outstanding reply without a multi-second drain (cancel path).
|
||
/// Same policy as freeze's terminal outcome: FG live → fail-closed;
|
||
/// sticky and background → report only.
|
||
pub(super) fn from_outstanding_reply(
|
||
reply: Option<
|
||
&kigi_tools::implementations::grok_build::task::types::SubagentOutstandingReply,
|
||
>,
|
||
) -> Self {
|
||
match reply {
|
||
None => Self {
|
||
fail_closed: true,
|
||
background_live: false,
|
||
sticky_report: false,
|
||
},
|
||
Some(r) => Self {
|
||
fail_closed: !r.live_ids.is_empty(),
|
||
background_live: r.background_live,
|
||
sticky_report: r.subagent_usage_not_applied,
|
||
},
|
||
}
|
||
}
|
||
}
|
||
/// Accumulates a turn's per-call token usage and tool-call presence across the
|
||
/// agentic loop's model calls, recording running totals on the turn span. Kept
|
||
/// out of the loop body so telemetry bookkeeping doesn't obscure control flow.
|
||
#[derive(Default)]
|
||
struct TurnSpanTotals {
|
||
input_tokens: i64,
|
||
output_tokens: i64,
|
||
cache_read_tokens: i64,
|
||
has_tool_call: bool,
|
||
}
|
||
impl TurnSpanTotals {
|
||
/// Fold one model response into the totals (tokens sum — each call is billed
|
||
/// its full prompt; has_tool_call OR-s — the final call has none) and update
|
||
/// the span. `stop_reason` is last-wins (the terminal reason), not summed.
|
||
fn record(&mut self, span: &tracing::Span, response: &ConversationResponse) {
|
||
if let Some(u) = response.usage.as_ref() {
|
||
self.input_tokens += i64::from(u.prompt_tokens);
|
||
self.output_tokens += i64::from(u.completion_tokens);
|
||
self.cache_read_tokens += i64::from(u.cached_prompt_tokens);
|
||
span.record("input_tokens", self.input_tokens);
|
||
span.record("output_tokens", self.output_tokens);
|
||
span.record("cache_read_tokens", self.cache_read_tokens);
|
||
}
|
||
if let Some(sr) = response.stop_reason {
|
||
span.record("stop_reason", sr.as_str());
|
||
}
|
||
self.has_tool_call |= !response.tool_calls().is_empty();
|
||
span.record("response.has_tool_call", self.has_tool_call);
|
||
}
|
||
}
|
||
/// How the turn's per-block user-message echo is published to clients /
|
||
/// `updates.jsonl`.
|
||
///
|
||
/// Every turn consumes a `prompt_index`, and rewind / fork truncation
|
||
/// (`replay_to_prompt`, `updates_truncate_for_prompt`) recover turn
|
||
/// boundaries by counting persisted `UserMessageChunk` runs — so every mode
|
||
/// persists the echo. Turns whose content must not render as a user prompt
|
||
/// (notification drain) are hidden by the *pager* via the
|
||
/// `hideFromScrollback` chunk meta, not by omitting the persisted line.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
enum UserEchoMode {
|
||
/// Live + persist (real user / cron / skill turns).
|
||
Broadcast,
|
||
/// Persist without live broadcast. Interject-fallback: panes already
|
||
/// rendered the text, so a live echo would duplicate it. Notification
|
||
/// drain: model-only content (the UI surfaces it via side channels:
|
||
/// monitor gutter, task pane) that no pane should render live.
|
||
PersistOnly,
|
||
}
|
||
fn user_echo_mode(prompt_id: &str) -> UserEchoMode {
|
||
if prompt_id.starts_with(super::interjection::INTERJECT_FALLBACK_PROMPT_PREFIX) {
|
||
return UserEchoMode::PersistOnly;
|
||
}
|
||
match super::super::PromptOrigin::from_prompt_id(prompt_id) {
|
||
super::super::PromptOrigin::NotificationDrain => UserEchoMode::PersistOnly,
|
||
_ => UserEchoMode::Broadcast,
|
||
}
|
||
}
|
||
impl SessionActor {
|
||
/// Run the image-normalization pipeline (re-encode caps, min-side and
|
||
/// integrity checks) and surface its outcomes: compression / re-encode
|
||
/// fallback / dropped notices are appended to `text_out` (TEXT only —
|
||
/// image data never enters a string) and mirrored as
|
||
/// `ImageCompressed`/`ImageDropped` notifications. Returns the surviving
|
||
/// images. Single owner of the notice/notify wiring, shared by the
|
||
/// prompt path and the interjection drain.
|
||
pub(crate) async fn normalize_images_with_notices(
|
||
&self,
|
||
text_out: &mut String,
|
||
images: Vec<acp::ImageContent>,
|
||
is_cursor: bool,
|
||
) -> Vec<acp::ImageContent> {
|
||
let mut norm_result =
|
||
crate::session::image_normalize::normalize_images(images, is_cursor).await;
|
||
let user_images = std::mem::take(&mut norm_result.images);
|
||
use crate::extensions::notification::ImageCompressedEntry;
|
||
if !norm_result.compressed.is_empty() {
|
||
text_out.push_str(&crate::session::image_normalize::render_compression_notice(
|
||
&norm_result.compressed,
|
||
is_cursor,
|
||
));
|
||
let message = norm_result
|
||
.compressed
|
||
.iter()
|
||
.map(|c| c.display())
|
||
.collect::<Vec<_>>()
|
||
.join("; ");
|
||
let images = norm_result
|
||
.compressed
|
||
.iter()
|
||
.map(ImageCompressedEntry::from)
|
||
.collect();
|
||
self.send_xai_notification(XaiSessionUpdate::ImageCompressed { images, message })
|
||
.await;
|
||
}
|
||
if !norm_result.re_encode_fallbacks.is_empty() {
|
||
text_out.push_str(
|
||
&crate::session::image_normalize::render_re_encode_fallback_notice(
|
||
&norm_result.re_encode_fallbacks,
|
||
is_cursor,
|
||
),
|
||
);
|
||
self.send_xai_notification(XaiSessionUpdate::ImageCompressed {
|
||
images: vec![],
|
||
message: norm_result.re_encode_fallbacks.join(" "),
|
||
})
|
||
.await;
|
||
}
|
||
if let Some((notice, notes)) = crate::session::image_normalize::dropped_to_envelope(
|
||
std::mem::take(&mut norm_result.dropped),
|
||
is_cursor,
|
||
) {
|
||
text_out.push_str(¬ice);
|
||
self.send_xai_notification(XaiSessionUpdate::ImageDropped { notes })
|
||
.await;
|
||
}
|
||
user_images
|
||
}
|
||
#[tracing::instrument(
|
||
name = "session.handle_prompt",
|
||
skip_all,
|
||
fields(
|
||
session_id = %self.session_info.id.0,
|
||
prompt_id = %prompt_id,
|
||
prompt_length = tracing::field::Empty,
|
||
command_name = tracing::field::Empty,
|
||
command_source = tracing::field::Empty,
|
||
)
|
||
)]
|
||
pub(super) async fn handle_prompt(
|
||
self: &Arc<Self>,
|
||
prompt_id: &str,
|
||
prompt_blocks: Vec<acp::ContentBlock>,
|
||
prompt_mode: PromptMode,
|
||
prompt_client_identifier: Option<String>,
|
||
prompt_screen_mode: Option<String>,
|
||
verbatim: bool,
|
||
json_schema: Option<serde_json::Value>,
|
||
persist_ack: Option<oneshot::Sender<()>>,
|
||
) -> PromptTurnResult {
|
||
let handle_prompt_start = std::time::Instant::now();
|
||
let prompt_length: usize = prompt_blocks
|
||
.iter()
|
||
.map(|b| match b {
|
||
acp::ContentBlock::Text(t) => t.text.len(),
|
||
_ => 0,
|
||
})
|
||
.sum();
|
||
tracing::Span::current().record("prompt_length", prompt_length as i64);
|
||
*self.active_skill.lock() = None;
|
||
kigi_log::unified_log::info(
|
||
"shell.handle_prompt.start",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "prompt_id" : prompt_id, "block_count" : prompt_blocks.len(), }
|
||
)),
|
||
);
|
||
if !super::super::PromptOrigin::from_prompt_id(prompt_id).is_synthetic() {
|
||
self.cancel_pending_recap_for_new_prompt();
|
||
}
|
||
*self.turn_start_prompt_mode.lock() = prompt_mode;
|
||
*self.turn_prompt_mode.lock() = prompt_mode;
|
||
self.signals_handle().increment_turn();
|
||
self.reconcile_plan_mode_with_prompt(prompt_mode);
|
||
let _turn_active_guard =
|
||
TurnActiveGuard::activate(self.tool_context.is_turn_active.as_ref());
|
||
let _session_turn_active_guard = TurnActiveGuard::activate(Some(&self.session_turn_active));
|
||
let turn_start_input = kigi_agent_lifecycle::TurnStartInput::new(
|
||
super::super::PromptOrigin::from_prompt_id(prompt_id).is_synthetic(),
|
||
);
|
||
for contributor in self.extension_registry.turn_lifecycle_contributors() {
|
||
contributor.on_turn_start(&turn_start_input).await;
|
||
}
|
||
if let Ok(mut pending) = self.rewind_pending_prompt.lock()
|
||
&& let Some(prev_text) = pending.take()
|
||
{
|
||
let new_text = prompt_blocks.iter().fold(String::new(), |mut acc, b| {
|
||
if let acp::ContentBlock::Text(t) = b {
|
||
acc.push_str(&t.text);
|
||
}
|
||
acc
|
||
});
|
||
if new_text.trim() == prev_text.trim() {
|
||
self.signals_handle().record_regeneration();
|
||
} else {
|
||
self.signals_handle().record_edit_and_retry();
|
||
}
|
||
}
|
||
if let Some(bash_command) = Self::extract_bash_command(&prompt_blocks) {
|
||
return self
|
||
.handle_direct_bash_command(prompt_id, bash_command, &prompt_blocks)
|
||
.await;
|
||
}
|
||
let slash_skills = self
|
||
.agent
|
||
.borrow()
|
||
.tool_bridge()
|
||
.clone()
|
||
.slash_skills()
|
||
.await;
|
||
let skill_rewrite = if crate::session::is_cursor_user_template(
|
||
&self.agent.borrow().definition().user_message_template,
|
||
) {
|
||
slash_commands::SkillSlashRewrite::Passthrough
|
||
} else {
|
||
slash_commands::SkillSlashRewrite::RewriteToRun
|
||
};
|
||
let availability = self.command_availability().await;
|
||
let mut pending_skill_information: Option<String> = None;
|
||
let prompt_blocks = match slash_commands::resolve(
|
||
prompt_blocks,
|
||
&slash_skills,
|
||
availability,
|
||
skill_rewrite,
|
||
) {
|
||
Ok(blocks) => blocks,
|
||
Err(SlashCommandOutcome::Builtin(action)) => {
|
||
let text_block =
|
||
|text: String| acp::ContentBlock::Text(acp::TextContent::new(text));
|
||
{
|
||
let span = tracing::Span::current();
|
||
span.record("command_name", action.command_name());
|
||
span.record("command_source", "builtin");
|
||
}
|
||
match action {
|
||
BuiltinAction::GoalSet {
|
||
objective,
|
||
token_budget,
|
||
} => {
|
||
let reminder = self.setup_goal(&objective, token_budget).await;
|
||
vec![text_block(reminder), text_block(objective)]
|
||
}
|
||
BuiltinAction::GoalResume => match self.resume_goal().await {
|
||
GoalResumeOutcome::Inference { reminder, user_msg } => {
|
||
self.send_slash_command_output(&user_msg).await;
|
||
vec![text_block(reminder)]
|
||
}
|
||
GoalResumeOutcome::Message(msg) => {
|
||
self.send_slash_command_output(&msg).await;
|
||
return ok_end_turn(0, None);
|
||
}
|
||
},
|
||
_ => return self.execute_builtin_slash_command(action).await,
|
||
}
|
||
}
|
||
Err(SlashCommandOutcome::InvokeSkill {
|
||
blocks: original_blocks,
|
||
skills: parsed_skills,
|
||
}) => {
|
||
if let Some(first) = parsed_skills.first() {
|
||
*self.active_skill.lock() = Some(first.name.clone());
|
||
let span = tracing::Span::current();
|
||
span.record("command_name", first.name.as_str());
|
||
span.record(
|
||
"command_source",
|
||
if first.plugin_name.is_some() {
|
||
"plugin"
|
||
} else {
|
||
"skill"
|
||
},
|
||
);
|
||
}
|
||
for sk in &parsed_skills {
|
||
let skill_source = if sk.plugin_name.is_some() {
|
||
"plugin"
|
||
} else {
|
||
crate::session::telemetry::skill_source_label(
|
||
&sk.skill_path,
|
||
self.session_info.cwd.as_str(),
|
||
)
|
||
};
|
||
tracing::info_span!(
|
||
"skill.activated", skill_name = % sk.name, invocation_trigger =
|
||
"slash_command", skill_source = skill_source,
|
||
)
|
||
.in_scope(|| {});
|
||
if let Some(ref pname) = sk.plugin_name {
|
||
tracing::info_span!(
|
||
"plugin.used", plugin_name = % pname, skill_name = % sk.name,
|
||
)
|
||
.in_scope(|| {});
|
||
}
|
||
}
|
||
pending_skill_information = slash_commands::build_skill_information_for_refs(
|
||
&parsed_skills,
|
||
&slash_skills,
|
||
&self.session_id_string(),
|
||
)
|
||
.await;
|
||
original_blocks
|
||
}
|
||
};
|
||
self.events.begin_turn();
|
||
let model_id = self.current_model_id().await;
|
||
let turn_number = self.chat_state_handle.get_prompt_index().await as u64;
|
||
self.current_turn_number.set(turn_number);
|
||
let yolo_mode = self.permissions.is_yolo_mode();
|
||
let msg_count = self.chat_state_handle.get_conversation_len().await;
|
||
let redirect_kind = if matches!(
|
||
super::super::PromptOrigin::from_prompt_id(prompt_id),
|
||
super::super::PromptOrigin::User
|
||
) {
|
||
self.events.take_prior_redirect_kind()
|
||
} else {
|
||
None
|
||
};
|
||
self.emit_event(crate::session::events::Event::TurnStarted {
|
||
session_id: self.session_id_string(),
|
||
turn_number,
|
||
model_id: model_id.clone(),
|
||
yolo_mode,
|
||
conversation_message_count: msg_count,
|
||
session_relationship: crate::session::events::SessionRelationship::Primary,
|
||
schema_version: crate::session::events::EVENT_SCHEMA_VERSION.into(),
|
||
redirect_kind,
|
||
});
|
||
self.observability_bridge
|
||
.emit(
|
||
kigi_tool_protocol::session_event::SessionEvent::TurnStarted {
|
||
turn_number,
|
||
model_id: model_id.clone(),
|
||
yolo_mode,
|
||
},
|
||
)
|
||
.await;
|
||
self.send_before_turn_event(kigi_tool_protocol::turn_hook::BeforeTurnPayload {
|
||
turn_number: self.chat_state_handle.get_prompt_index().await as u64,
|
||
model_id: model_id.clone(),
|
||
yolo_mode: self.permissions.is_yolo_mode(),
|
||
conversation_message_count: msg_count,
|
||
session_relationship: kigi_tool_protocol::turn_hook::DEFAULT_SESSION_RELATIONSHIP
|
||
.to_string(),
|
||
schema_version: crate::session::events::EVENT_SCHEMA_VERSION.to_string(),
|
||
})
|
||
.await;
|
||
let turn_idx = self.chat_state_handle.get_prompt_index().await as u64;
|
||
let current_prompt_index = self.chat_state_handle.get_prompt_index().await;
|
||
let origin = super::super::PromptOrigin::from_prompt_id(prompt_id);
|
||
let mut chunk_meta = serde_json::Map::new();
|
||
chunk_meta.insert("modelId".into(), serde_json::json!(model_id));
|
||
chunk_meta.insert(
|
||
"promptIndex".into(),
|
||
serde_json::json!(current_prompt_index),
|
||
);
|
||
if origin.hide_user_echo_from_scrollback() {
|
||
chunk_meta.insert("hideFromScrollback".into(), serde_json::json!(true));
|
||
}
|
||
let user_chunk_meta = Some(chunk_meta);
|
||
self.chat_state_handle.increment_prompt_index();
|
||
let text = prompt_blocks.iter().fold(String::new(), |mut acc, b| {
|
||
if let acp::ContentBlock::Text(t) = b {
|
||
acc.push_str(&t.text);
|
||
}
|
||
acc
|
||
});
|
||
let trimmed = text.trim().to_string();
|
||
if !trimmed.is_empty() {
|
||
self.chat_state_handle.cache_prompt_text(trimmed);
|
||
}
|
||
*self.tool_context.prompt_index.lock().await = current_prompt_index;
|
||
self.file_state_tracker
|
||
.begin_prompt(current_prompt_index)
|
||
.await;
|
||
let echo_mode = user_echo_mode(prompt_id);
|
||
for block in prompt_blocks.iter() {
|
||
let update = acp::SessionUpdate::UserMessageChunk(
|
||
acp::ContentChunk::new(block.clone()).meta(user_chunk_meta.clone()),
|
||
);
|
||
let notification_meta = self.build_notification_meta();
|
||
let notification = acp::SessionNotification::new(self.session_info.id.clone(), update)
|
||
.meta(notification_meta.as_object().cloned());
|
||
if echo_mode == UserEchoMode::PersistOnly {
|
||
let _ = self
|
||
.notifications
|
||
.persistence_tx
|
||
.send(PersistenceMsg::Update(
|
||
crate::session::storage::SessionUpdate::Acp(Box::new(notification)),
|
||
));
|
||
} else {
|
||
self.emit_notification_direct(notification).await;
|
||
}
|
||
}
|
||
let crate::session::prompt_parser::ParsedPrompt {
|
||
mut context,
|
||
query,
|
||
skill_information: skill_info,
|
||
images: mut raw_images,
|
||
is_cursor,
|
||
} = match parse_prompt_with_skills(
|
||
&prompt_blocks,
|
||
self.tool_context.cwd.to_path_buf(),
|
||
&self.session_info,
|
||
verbatim,
|
||
self.is_cursor_harness(),
|
||
pending_skill_information.take().unwrap_or_default(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(v) => v,
|
||
Err(err) => {
|
||
tracing::warn!("Invalid prompt: {}", err.message);
|
||
return Err(err);
|
||
}
|
||
};
|
||
let recovered = crate::session::placeholder_images::recover_orphan_placeholders(
|
||
&query,
|
||
&mut raw_images,
|
||
std::path::Path::new(&self.session_info.cwd),
|
||
);
|
||
if recovered > 0 {
|
||
tracing::info!(
|
||
session_id = % self.session_info.id, recovered,
|
||
"server-side placeholder fallback: loaded orphan image(s) from disk",
|
||
);
|
||
}
|
||
let query = crate::session::placeholder_images::strip_paths_from_image_placeholders(query);
|
||
let user_images = self
|
||
.normalize_images_with_notices(&mut context, raw_images, is_cursor)
|
||
.await;
|
||
let (query, extra_images) = if !self.is_cursor_harness() {
|
||
let extraction = kigi_tools::util::base64_images::extract_base64_images(query);
|
||
if extraction.images.is_empty() {
|
||
(extraction.text, Vec::new())
|
||
} else {
|
||
let cleaned_text = extraction.text;
|
||
let count = extraction.images.len();
|
||
tracing::info!(
|
||
session_id = % self.session_info.id, count,
|
||
"base64 images extracted from user query",
|
||
);
|
||
let acp_imgs: Vec<agent_client_protocol::ImageContent> = extraction
|
||
.images
|
||
.into_iter()
|
||
.map(|img| agent_client_protocol::ImageContent::new(img.data, img.mime_type))
|
||
.collect();
|
||
let nr = crate::session::image_normalize::normalize_images(acp_imgs, false).await;
|
||
if !nr.re_encode_fallbacks.is_empty() {
|
||
tracing::warn!(
|
||
session_id = % self.session_info.id, notes = % nr
|
||
.re_encode_fallbacks.join(" "),
|
||
"Extracted user query image kept original after re-encode failure",
|
||
);
|
||
}
|
||
(cleaned_text, nr.images)
|
||
}
|
||
} else {
|
||
(query, Vec::new())
|
||
};
|
||
let assembled = crate::session::prompt_parser::ParsedPrompt::assemble_parts_with_skills(
|
||
&context,
|
||
&query,
|
||
&skill_info,
|
||
is_cursor,
|
||
);
|
||
let (user_message, _truncated_local_path) = if verbatim {
|
||
(assembled, None)
|
||
} else {
|
||
self.maybe_truncate_large_prompt_with_skills(
|
||
context,
|
||
query,
|
||
skill_info,
|
||
is_cursor,
|
||
current_prompt_index,
|
||
)
|
||
.await
|
||
};
|
||
let _ = self
|
||
.notifications
|
||
.persistence_tx
|
||
.send(PersistenceMsg::ContentChunk(PersistenceContentChunk::new(
|
||
prompt_blocks.to_vec(),
|
||
)));
|
||
let model_id = self
|
||
.chat_state_handle
|
||
.get_sampling_config()
|
||
.await
|
||
.map(|c| c.model)
|
||
.unwrap_or_default();
|
||
self.maybe_inject_mcp_reminder().await;
|
||
self.maybe_inject_mcp_connecting_reminder().await;
|
||
self.maybe_inject_date_rollover_reminder().await;
|
||
self.inject_plan_mode_reminders().await;
|
||
self.inject_resumed_tasks_reminder();
|
||
self.drain_between_turn_completions().await;
|
||
let user_message = if user_images.is_empty() {
|
||
user_message
|
||
} else if self.is_cursor_harness() {
|
||
self.transcribe_user_images(user_message, &user_images)
|
||
.await?
|
||
} else {
|
||
let session_dir =
|
||
crate::session::persistence::session_dir(&crate::session::info::Info {
|
||
id: self.session_info.id.clone(),
|
||
cwd: self.session_info.cwd.clone(),
|
||
});
|
||
crate::session::image_describe::persist_and_prepend_image_files(
|
||
&session_dir,
|
||
&user_images,
|
||
&user_message,
|
||
)
|
||
.map_err(|e| {
|
||
acp::Error::internal_error()
|
||
.data(format!("failed to save user images to assets dir: {e}"))
|
||
})?
|
||
};
|
||
let attached_image_refs = if self.is_cursor_harness() {
|
||
Vec::new()
|
||
} else {
|
||
crate::session::placeholder_images::attached_image_references(&user_images)
|
||
};
|
||
self.tool_bridge_handle()
|
||
.update_resource(kigi_tools::types::resources::AttachedImages(
|
||
attached_image_refs,
|
||
))
|
||
.await;
|
||
let prompt_text_for_hook = user_message.clone();
|
||
{
|
||
let origin = super::super::PromptOrigin::from_prompt_id(prompt_id);
|
||
if matches!(origin, super::super::PromptOrigin::User) {
|
||
self.maybe_inject_interrupt_reminder().await;
|
||
}
|
||
let mut user_chat = match &origin {
|
||
super::super::PromptOrigin::TaskCompleted { .. } => {
|
||
ConversationItem::task_completed(user_message)
|
||
}
|
||
super::super::PromptOrigin::SubagentCompleted { .. } => {
|
||
ConversationItem::subagent_completed(user_message)
|
||
}
|
||
super::super::PromptOrigin::NotificationDrain => {
|
||
ConversationItem::notification_drain(user_message)
|
||
}
|
||
super::super::PromptOrigin::GoalSummary => {
|
||
ConversationItem::goal_summary(user_message)
|
||
}
|
||
super::super::PromptOrigin::GoalClassifierNudge => {
|
||
ConversationItem::goal_classifier_nudge(user_message)
|
||
}
|
||
super::super::PromptOrigin::SchedulerFired => {
|
||
ConversationItem::scheduler_fired(user_message)
|
||
}
|
||
super::super::PromptOrigin::PlanResume => ConversationItem::user(user_message),
|
||
super::super::PromptOrigin::User => {
|
||
let mut item = ConversationItem::user(user_message);
|
||
if let Some(interrupt) = self
|
||
.events
|
||
.take_prior_interrupt_category()
|
||
.and_then(crate::session::events::prior_turn_interrupt_from_cancellation)
|
||
{
|
||
item.set_prior_turn_interrupt(interrupt);
|
||
}
|
||
item
|
||
}
|
||
};
|
||
user_chat.set_prompt_index(current_prompt_index);
|
||
if !self.is_cursor_harness() {
|
||
for image in &user_images {
|
||
user_chat.add_image(pick_user_image_url(image));
|
||
}
|
||
for image in &extra_images {
|
||
user_chat.add_image(format!("data:{};base64,{}", image.mime_type, image.data));
|
||
}
|
||
}
|
||
if let Some(ack) = persist_ack {
|
||
if self
|
||
.chat_state_handle
|
||
.push_user_message_and_ack(user_chat)
|
||
.await
|
||
.is_some()
|
||
{
|
||
let (flush_tx, flush_rx) = oneshot::channel();
|
||
if self
|
||
.notifications
|
||
.persistence_tx
|
||
.send(PersistenceMsg::FlushAndAck {
|
||
respond_to: flush_tx,
|
||
})
|
||
.is_ok()
|
||
&& flush_rx.await.is_ok()
|
||
{
|
||
let _ = ack.send(());
|
||
} else {
|
||
tracing::error!(
|
||
session_id = % self.session_info.id.0, prompt_id = %
|
||
prompt_id, "persist_ack flush barrier failed"
|
||
);
|
||
}
|
||
} else {
|
||
tracing::error!(
|
||
session_id = % self.session_info.id.0, prompt_id = % prompt_id,
|
||
"persist_ack skipped: chat-state actor unavailable"
|
||
);
|
||
}
|
||
} else {
|
||
self.chat_state_handle.push_user_message(user_chat);
|
||
}
|
||
}
|
||
self.dispatch_hook(
|
||
kigi_hooks::event::HookEventName::UserPromptSubmit,
|
||
kigi_hooks::event::HookPayload::UserPromptSubmit {
|
||
prompt: Some(prompt_text_for_hook),
|
||
},
|
||
Some(prompt_id),
|
||
None,
|
||
)
|
||
.await;
|
||
let turn_scope_guard =
|
||
TurnSubagentScopeGuard::new(self.current_prompt_id.clone(), prompt_id.to_string());
|
||
let turn_model_id = self.current_model_id().await;
|
||
let doom_event_model = turn_model_id.clone();
|
||
let turn_timer = std::time::Instant::now();
|
||
let result = {
|
||
loop {
|
||
if self.goal_harness_enabled() {
|
||
let goal_loop_active = self.goal_tracker.lock().status()
|
||
== Some(crate::session::goal_tracker::GoalStatus::Active);
|
||
self.set_goal_loop_active_resource(goal_loop_active).await;
|
||
}
|
||
let round = self
|
||
.process_conversation_turn_with_recovery(prompt_id, json_schema.clone())
|
||
.await;
|
||
if !matches!(round, Ok(TurnOutcome::Completed { .. })) {
|
||
break round;
|
||
}
|
||
if matches!(round, Ok(TurnOutcome::Completed { refusal: true, .. })) {
|
||
break round;
|
||
}
|
||
let goal_active = laziness_injection_active(
|
||
self.goal_harness_enabled(),
|
||
self.goal_tracker.lock().status(),
|
||
);
|
||
if !goal_active {
|
||
break round;
|
||
}
|
||
match self.run_goal_round_end().await {
|
||
GoalRoundDecision::Continue(directive) => {
|
||
self.inject_goal_continuation_message(directive).await;
|
||
}
|
||
GoalRoundDecision::EndTurn => break round,
|
||
}
|
||
}
|
||
};
|
||
let turn_duration_ms = turn_timer.elapsed().as_millis() as u64;
|
||
let handle_prompt_elapsed_ms = handle_prompt_start.elapsed().as_millis() as u64;
|
||
kigi_log::unified_log::info(
|
||
"shell.handle_prompt.done",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "prompt_id" : prompt_id, "total_elapsed_ms" :
|
||
handle_prompt_elapsed_ms, "turn_elapsed_ms" : turn_duration_ms,
|
||
"pre_turn_ms" : handle_prompt_elapsed_ms
|
||
.saturating_sub(turn_duration_ms), "ok" : result.is_ok(), }
|
||
)),
|
||
);
|
||
let turn_tool_count = self.events.tool_count_this_turn();
|
||
let bridge_outcome = turn_result_to_hook_outcome(&result);
|
||
self.observability_bridge
|
||
.emit(kigi_tool_protocol::session_event::SessionEvent::TurnEnded {
|
||
turn_number: current_prompt_index as u64,
|
||
outcome: bridge_outcome,
|
||
duration_ms: turn_duration_ms,
|
||
tool_call_count: turn_tool_count,
|
||
model_id: turn_model_id.clone(),
|
||
})
|
||
.await;
|
||
match &result {
|
||
Ok(TurnOutcome::Completed { .. }) => {
|
||
self.emit_turn_ended(
|
||
crate::session::events::TurnOutcomeLabel::Completed,
|
||
None,
|
||
None,
|
||
);
|
||
self.send_after_turn_event(kigi_tool_protocol::turn_hook::AfterTurnPayload {
|
||
turn_number: current_prompt_index as u64,
|
||
outcome: kigi_tool_protocol::turn_hook::TurnHookOutcome::Completed,
|
||
duration_ms: turn_duration_ms,
|
||
tool_call_count: turn_tool_count,
|
||
model_id: turn_model_id.clone(),
|
||
written_repo_paths: Vec::new(),
|
||
cancellation_category: None,
|
||
cancellation_context: None,
|
||
})
|
||
.await;
|
||
}
|
||
Ok(TurnOutcome::Cancelled { category, context }) => {
|
||
self.emit_turn_ended(
|
||
crate::session::events::TurnOutcomeLabel::Cancelled,
|
||
*category,
|
||
context.clone(),
|
||
);
|
||
if let Some(cause) = category {
|
||
self.events.set_prior_interrupt_category(*cause);
|
||
}
|
||
self.send_after_turn_event(kigi_tool_protocol::turn_hook::AfterTurnPayload {
|
||
turn_number: current_prompt_index as u64,
|
||
outcome: kigi_tool_protocol::turn_hook::TurnHookOutcome::Cancelled,
|
||
duration_ms: turn_duration_ms,
|
||
tool_call_count: turn_tool_count,
|
||
model_id: turn_model_id.clone(),
|
||
written_repo_paths: Vec::new(),
|
||
cancellation_category: cancellation_category_wire_string(*category),
|
||
cancellation_context: context.clone(),
|
||
})
|
||
.await;
|
||
}
|
||
Ok(TurnOutcome::MaxTurnsReached { limit }) => {
|
||
tracing::info!(limit, "turn ended: max_turns reached");
|
||
self.emit_turn_ended(
|
||
crate::session::events::TurnOutcomeLabel::Cancelled,
|
||
None,
|
||
Some(serde_json::json!(
|
||
{ "reason" : "max_turns_reached", "limit" : limit, }
|
||
)),
|
||
);
|
||
self.send_after_turn_event(kigi_tool_protocol::turn_hook::AfterTurnPayload {
|
||
turn_number: current_prompt_index as u64,
|
||
outcome: kigi_tool_protocol::turn_hook::TurnHookOutcome::Cancelled,
|
||
duration_ms: turn_duration_ms,
|
||
tool_call_count: turn_tool_count,
|
||
model_id: turn_model_id.clone(),
|
||
written_repo_paths: Vec::new(),
|
||
cancellation_category: None,
|
||
cancellation_context: Some(serde_json::json!(
|
||
{ "reason" : "max_turns_reached", "limit" : limit, }
|
||
)),
|
||
})
|
||
.await;
|
||
}
|
||
Err(err) => {
|
||
self.emit_turn_ended(crate::session::events::TurnOutcomeLabel::Error, None, None);
|
||
self.send_after_turn_event(kigi_tool_protocol::turn_hook::AfterTurnPayload {
|
||
turn_number: current_prompt_index as u64,
|
||
outcome: kigi_tool_protocol::turn_hook::TurnHookOutcome::Error,
|
||
duration_ms: turn_duration_ms,
|
||
tool_call_count: turn_tool_count,
|
||
model_id: turn_model_id.clone(),
|
||
written_repo_paths: Vec::new(),
|
||
cancellation_category: None,
|
||
cancellation_context: None,
|
||
})
|
||
.await;
|
||
let error_category = Self::classify_turn_error(err);
|
||
self.dispatch_hook(
|
||
kigi_hooks::event::HookEventName::StopFailure,
|
||
kigi_hooks::event::HookPayload::StopFailure {
|
||
error: format!("{err}"),
|
||
},
|
||
Some(prompt_id),
|
||
None,
|
||
)
|
||
.await;
|
||
}
|
||
}
|
||
let doom_tally = std::mem::take(&mut *self.doom_loop_turn_tally.lock());
|
||
doom_tally.fired();
|
||
let stop_reason_str = match &result {
|
||
Ok(TurnOutcome::Completed { .. }) => "end_turn",
|
||
Ok(TurnOutcome::Cancelled { .. }) | Ok(TurnOutcome::MaxTurnsReached { .. }) => {
|
||
"cancelled"
|
||
}
|
||
Err(_) => "error",
|
||
};
|
||
self.dispatch_hook(
|
||
kigi_hooks::event::HookEventName::Stop,
|
||
kigi_hooks::event::HookPayload::Stop {
|
||
reason: stop_reason_str.to_string(),
|
||
},
|
||
Some(prompt_id),
|
||
None,
|
||
)
|
||
.await;
|
||
match &result {
|
||
Ok(TurnOutcome::Completed { .. }) => {
|
||
for contributor in self.extension_registry.turn_lifecycle_contributors() {
|
||
contributor
|
||
.on_turn_done(&kigi_agent_lifecycle::TurnDoneInput)
|
||
.await;
|
||
}
|
||
}
|
||
Ok(TurnOutcome::Cancelled { .. }) | Ok(TurnOutcome::MaxTurnsReached { .. }) => {
|
||
let input = kigi_agent_lifecycle::TurnAbortInput::new(
|
||
kigi_agent_lifecycle::TurnAbortReason::Interrupted,
|
||
);
|
||
for contributor in self.extension_registry.turn_lifecycle_contributors() {
|
||
contributor.on_turn_abort(&input).await;
|
||
}
|
||
}
|
||
Err(err) => {
|
||
let message = err.to_string();
|
||
let input = kigi_agent_lifecycle::TurnErrorInput { message: &message };
|
||
for contributor in self.extension_registry.turn_lifecycle_contributors() {
|
||
contributor.on_turn_error(&input).await;
|
||
}
|
||
}
|
||
}
|
||
if matches!(
|
||
result,
|
||
Ok(TurnOutcome::Cancelled { .. }) | Ok(TurnOutcome::MaxTurnsReached { .. })
|
||
) {
|
||
self.cancel_running_turn_subagents();
|
||
}
|
||
self.flush_to_disk().await;
|
||
self.file_state_tracker
|
||
.end_prompt(&self.tool_context.fs, current_prompt_index)
|
||
.await;
|
||
if let Some(mut rewind_point) = self
|
||
.file_state_tracker
|
||
.get_rewind_point(current_prompt_index)
|
||
.await
|
||
{
|
||
rewind_point.normalize_to_relative(self.tool_context.cwd.as_ref());
|
||
let _ = self
|
||
.notifications
|
||
.persistence_tx
|
||
.send(PersistenceMsg::RewindPoint(rewind_point));
|
||
}
|
||
match result {
|
||
Ok(outcome) => {
|
||
let usage = self.freeze_prompt_usage(prompt_id).await;
|
||
drop(turn_scope_guard);
|
||
self.chat_state_handle.flush();
|
||
let total_tokens = self.chat_state_handle.get_total_tokens().await;
|
||
let (stop_reason, mut snapshot, completion_kind, structured_output) = match outcome
|
||
{
|
||
TurnOutcome::Completed {
|
||
snapshot,
|
||
structured_output,
|
||
refusal,
|
||
..
|
||
} => (
|
||
if refusal {
|
||
acp::StopReason::Refusal
|
||
} else {
|
||
acp::StopReason::EndTurn
|
||
},
|
||
*snapshot,
|
||
PromptCompletionKind::Completed,
|
||
structured_output,
|
||
),
|
||
TurnOutcome::Cancelled { category, context } => {
|
||
let cancellation_ctx = context.and_then(|v| serde_json::from_value(v).ok());
|
||
(
|
||
acp::StopReason::Cancelled,
|
||
None,
|
||
PromptCompletionKind::Cancelled {
|
||
category,
|
||
context: cancellation_ctx,
|
||
},
|
||
None,
|
||
)
|
||
}
|
||
TurnOutcome::MaxTurnsReached { limit } => (
|
||
acp::StopReason::Cancelled,
|
||
None,
|
||
PromptCompletionKind::MaxTurnsReached { limit },
|
||
None,
|
||
),
|
||
};
|
||
if let Some(snapshot) = snapshot.as_mut() {
|
||
self.apply_prompt_modes_to_snapshot(snapshot);
|
||
}
|
||
Ok(crate::session::commands::PromptTurnOk {
|
||
stop_reason,
|
||
total_tokens,
|
||
turn_snapshot: snapshot,
|
||
completion_kind,
|
||
structured_output,
|
||
usage,
|
||
})
|
||
}
|
||
Err(e) => {
|
||
let usage = self.freeze_prompt_usage(prompt_id).await;
|
||
drop(turn_scope_guard);
|
||
Err(crate::sampling::error::attach_prompt_usage(e, usage))
|
||
}
|
||
}
|
||
}
|
||
/// Wait for turn-blocking subagents (up to 120s on the turn task),
|
||
/// snapshot, clear sticky. Background children never gate the drain: the
|
||
/// prompt report is marked incomplete immediately and their spend reaches
|
||
/// the session ledger when they finish.
|
||
/// Cancel intentionally skips this multi-second drain (actor-loop safety).
|
||
pub(super) async fn freeze_prompt_usage(
|
||
&self,
|
||
prompt_id: &str,
|
||
) -> Option<crate::extensions::notification::PromptUsage> {
|
||
const DRAIN: std::time::Duration = std::time::Duration::from_secs(120);
|
||
self.freeze_prompt_usage_bounded(prompt_id, DRAIN).await
|
||
}
|
||
/// [`freeze_prompt_usage`] with an explicit drain bound, for tests.
|
||
pub(super) async fn freeze_prompt_usage_bounded(
|
||
&self,
|
||
prompt_id: &str,
|
||
max_wait: std::time::Duration,
|
||
) -> Option<crate::extensions::notification::PromptUsage> {
|
||
let drain = self
|
||
.drain_subagent_usage_for_prompt_bounded(prompt_id, max_wait)
|
||
.await;
|
||
self.finalize_usage_from_outcome(prompt_id, drain).await
|
||
}
|
||
/// Waits for turn-blocking folds only.
|
||
/// `fail_closed` on timeout or query failure; sticky and `background_live`
|
||
/// are report-level only (no ledger mark). Must run on the turn task (not
|
||
/// the session actor loop) so folds can land.
|
||
pub(super) async fn drain_subagent_usage_for_prompt_bounded(
|
||
&self,
|
||
prompt_id: &str,
|
||
max_wait: std::time::Duration,
|
||
) -> UsageDrainOutcome {
|
||
const POLL: std::time::Duration = std::time::Duration::from_millis(50);
|
||
let deadline = std::time::Instant::now() + max_wait;
|
||
loop {
|
||
let reply = self.outstanding_reply_for_prompt(prompt_id).await;
|
||
match reply.as_ref() {
|
||
None => {
|
||
tracing::warn!(
|
||
prompt_id,
|
||
"outstanding subagent query failed; treating usage as incomplete"
|
||
);
|
||
return UsageDrainOutcome {
|
||
fail_closed: true,
|
||
background_live: false,
|
||
sticky_report: false,
|
||
};
|
||
}
|
||
Some(r) if r.live_ids.is_empty() => {
|
||
return UsageDrainOutcome {
|
||
fail_closed: false,
|
||
background_live: r.background_live,
|
||
sticky_report: r.subagent_usage_not_applied,
|
||
};
|
||
}
|
||
Some(r) => {
|
||
if std::time::Instant::now() >= deadline {
|
||
tracing::warn!(
|
||
prompt_id,
|
||
count = r.live_ids.len(),
|
||
max_wait_ms = max_wait.as_millis() as u64,
|
||
"subagent usage drain timed out; usage may under-count"
|
||
);
|
||
return UsageDrainOutcome {
|
||
fail_closed: true,
|
||
background_live: r.background_live,
|
||
sticky_report: r.subagent_usage_not_applied,
|
||
};
|
||
}
|
||
}
|
||
}
|
||
tokio::time::sleep(POLL).await;
|
||
}
|
||
}
|
||
pub(super) async fn snapshot_prompt_usage(
|
||
&self,
|
||
) -> Option<crate::extensions::notification::PromptUsage> {
|
||
self.snapshot_prompt_usage_marked(false).await
|
||
}
|
||
pub(super) async fn snapshot_prompt_usage_marked(
|
||
&self,
|
||
incomplete: bool,
|
||
) -> Option<crate::extensions::notification::PromptUsage> {
|
||
match self.chat_state_handle.try_get_prompt_usage().await {
|
||
Ok(ledger) => {
|
||
let incomplete = incomplete || ledger.as_ref().is_some_and(|l| l.incomplete);
|
||
crate::extensions::notification::PromptUsage::project_from_ledger(
|
||
ledger.as_ref(),
|
||
incomplete,
|
||
)
|
||
}
|
||
Err(()) => {
|
||
crate::extensions::notification::PromptUsage::project_from_ledger(None, true)
|
||
}
|
||
}
|
||
}
|
||
/// When freeze did not attach: incomplete if billed or may under-count; else omit.
|
||
pub(super) async fn error_path_usage_fallback(
|
||
&self,
|
||
prompt_id: &str,
|
||
) -> Option<crate::extensions::notification::PromptUsage> {
|
||
let may_undercount = Self::usage_incomplete_from_reply(
|
||
self.outstanding_reply_for_prompt(prompt_id).await.as_ref(),
|
||
);
|
||
match self.chat_state_handle.try_get_prompt_usage().await {
|
||
Ok(ledger) => crate::extensions::notification::PromptUsage::for_error_path(
|
||
ledger.as_ref(),
|
||
may_undercount,
|
||
),
|
||
Err(()) => crate::extensions::notification::PromptUsage::for_error_path(None, true),
|
||
}
|
||
}
|
||
/// Sticky incomplete for `prompt_id`, or the live pin when `None`.
|
||
/// Returns true only if the coordinator acked the mark.
|
||
pub(super) async fn mark_subagent_usage_not_applied(&self, prompt_id: Option<&str>) -> bool {
|
||
let resolved = prompt_id
|
||
.map(str::to_owned)
|
||
.or_else(|| self.current_prompt_id.lock().ok().and_then(|g| g.clone()));
|
||
let Some(pid) = resolved else {
|
||
return false;
|
||
};
|
||
let Some(tx) = &self.tool_context.subagent_event_tx else {
|
||
return false;
|
||
};
|
||
use kigi_tools::implementations::grok_build::task::types::{
|
||
SubagentEvent, SubagentMarkUsageNotAppliedRequest,
|
||
};
|
||
let (respond_to, ack) = tokio::sync::oneshot::channel();
|
||
if tx
|
||
.send(SubagentEvent::MarkUsageNotApplied(
|
||
SubagentMarkUsageNotAppliedRequest {
|
||
prompt_id: pid,
|
||
respond_to,
|
||
},
|
||
))
|
||
.is_err()
|
||
{
|
||
return false;
|
||
}
|
||
ack.await.is_ok()
|
||
}
|
||
/// Drain this session's buffered mid-turn monitor events
|
||
/// (`drain_owned` — leader mode shares the buffer) into ONE hidden
|
||
/// synthetic user message, tagged `SyntheticReason::SystemReminder` so
|
||
/// compaction/fork/pruning skip it. Deliberately a bare
|
||
/// `push_user_message`, NOT `inject_synthetic_user_message`: the latter
|
||
/// persists a `UserMessageChunk` to `updates.jsonl`, which resume
|
||
/// replays — the raw XML would render as a user prompt. Clients see
|
||
/// monitor events only via the structured `x.ai/monitor_event` channel.
|
||
pub(crate) async fn inject_pending_monitor_events(&self) {
|
||
let Some(buffer) = &self.tool_context.monitor_event_buffer else {
|
||
return;
|
||
};
|
||
let mine = kigi_tools::implementations::grok_build::task::types::drain_owned(
|
||
buffer,
|
||
Some(self.session_info.id.0.as_ref()),
|
||
);
|
||
if mine.is_empty() {
|
||
return;
|
||
}
|
||
let Some(body) = kigi_tools::reminders::task_completion::format_monitor_events(
|
||
&mine,
|
||
Some(&self.tool_context.task_output_tool_name),
|
||
) else {
|
||
return;
|
||
};
|
||
let wrapped = kigi_tools::reminders::wrap_reminder(&body);
|
||
self.chat_state_handle
|
||
.push_user_message(ConversationItem::system_reminder(wrapped));
|
||
tracing::info!(
|
||
session_id = % self.session_info.id.0, count = mine.len(),
|
||
"injected mid-turn monitor events as hidden synthetic user message"
|
||
);
|
||
}
|
||
/// Per-turn hook called from the event-loop completion handler
|
||
/// after every turn finishes. Two terminal branches when the
|
||
/// goal is `Active` (`goal_active_now == true`):
|
||
///
|
||
/// 1. **Success.** Reset `goal_continuation_streak` to 0, then call
|
||
/// `maybe_queue_goal_continuation`. That helper verifies any
|
||
/// pending completion via its turn-end drain, queues the
|
||
/// continuation reminder if the goal is still `Active`, and runs
|
||
/// the stop-detector to select the nudge flavor (generic vs.
|
||
/// bail-specific) and emit `Event::GoalPrematureStopDetected`.
|
||
/// 2. **Non-success.** Increment `goal_continuation_streak`. At
|
||
/// [`GOAL_CONTINUATION_BACKOFF_THRESHOLD`] consecutive hits,
|
||
/// reset the streak and auto-pause with
|
||
/// `GoalPauseReason::BackOff`. No continuation is queued on this path: an
|
||
/// infra-error / cancelled turn rarely carries a deliberate
|
||
/// turn-final message, and stop-detection lives on the success
|
||
/// path inside `maybe_queue_goal_continuation`.
|
||
///
|
||
/// When the goal is not `Active` (`goal_active_now == false` —
|
||
/// the doom-loop / infra-error branches in the event loop ran
|
||
/// before this method and already transitioned the goal out of
|
||
/// Active), both branches are skipped: neither streak moves and the
|
||
/// existing pause cause is preserved.
|
||
///
|
||
/// `goal_blocked_streak` is never touched here — blocked attempts end
|
||
/// their turns successfully, so a per-turn reset would make the
|
||
/// 3-streak pause unreachable. It resets on goal create / resume /
|
||
/// clear and on a `completed: true` `update_goal`.
|
||
pub(crate) async fn handle_turn_end(&self, turn_succeeded: bool) {
|
||
let goal_active_now = laziness_injection_active(
|
||
self.goal_harness_enabled(),
|
||
self.goal_tracker.lock().status(),
|
||
);
|
||
if turn_succeeded && goal_active_now {
|
||
self.goal_continuation_streak
|
||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
||
self.maybe_queue_goal_continuation().await;
|
||
return;
|
||
}
|
||
if !turn_succeeded && goal_active_now {
|
||
let current_tokens = self.chat_state_handle.get_total_tokens().await as i64;
|
||
if self.enforce_goal_token_budget(current_tokens).await {
|
||
return;
|
||
}
|
||
let streak = self
|
||
.goal_continuation_streak
|
||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
|
||
+ 1;
|
||
if streak >= GOAL_CONTINUATION_BACKOFF_THRESHOLD {
|
||
self.goal_continuation_streak
|
||
.store(0, std::sync::atomic::Ordering::Relaxed);
|
||
self.auto_pause_goal_if_active(
|
||
crate::session::goal_tracker::GoalPauseReason::BackOff,
|
||
)
|
||
.await;
|
||
self.send_slash_command_output(&format!(
|
||
"Goal auto-paused after {GOAL_CONTINUATION_BACKOFF_THRESHOLD} consecutive \
|
||
non-completing turns. The model is not making progress. \
|
||
Use /goal resume to retry or /goal clear to abandon."
|
||
))
|
||
.await;
|
||
}
|
||
}
|
||
}
|
||
/// Wraps `process_conversation_turn` with auto-recovery for agents that opt in.
|
||
///
|
||
/// Agents with a `completion_requirement` in their definition require the model
|
||
/// to call a specific tool before finishing. If a prompt turn ends without that
|
||
/// tool having been called, this method injects the recovery prompt and re-runs
|
||
/// the turn with exponential backoff.
|
||
///
|
||
/// Agents without `completion_requirement` bypass this entirely.
|
||
#[tracing::instrument(
|
||
name = "session.process_conversation_turn_with_recovery",
|
||
skip_all,
|
||
err,
|
||
fields(req_id = %req_id, session_id = %self.session_info.id.0)
|
||
)]
|
||
pub(super) async fn process_conversation_turn_with_recovery(
|
||
self: &Arc<Self>,
|
||
req_id: &str,
|
||
json_schema: Option<serde_json::Value>,
|
||
) -> Result<TurnOutcome, acp::Error> {
|
||
let _ = self.compaction.auto_compact_suppressed.compare_exchange(
|
||
crate::session::compaction_config::SUPPRESS_TURN,
|
||
crate::session::compaction_config::SUPPRESS_NONE,
|
||
std::sync::atomic::Ordering::Relaxed,
|
||
std::sync::atomic::Ordering::Relaxed,
|
||
);
|
||
let agent_ref = self.agent.borrow();
|
||
let completion_req = match agent_ref.completion_requirement() {
|
||
Some(req) => req,
|
||
None => {
|
||
return self.process_conversation_turn(req_id, json_schema).await;
|
||
}
|
||
};
|
||
let recovery = match &completion_req.recovery {
|
||
Some(r) => r.clone(),
|
||
None => {
|
||
return self.process_conversation_turn(req_id, json_schema).await;
|
||
}
|
||
};
|
||
let required_tool = completion_req.tool.clone();
|
||
let recovery_prompt = completion_req.reminder.clone();
|
||
let mut result = self
|
||
.process_conversation_turn(req_id, json_schema.clone())
|
||
.await;
|
||
if matches!(result, Ok(TurnOutcome::MaxTurnsReached { .. })) {
|
||
return result;
|
||
}
|
||
if let Ok(TurnOutcome::Completed {
|
||
ref tools_called, ..
|
||
}) = result
|
||
&& tools_called.iter().any(|name| name == &required_tool)
|
||
{
|
||
tracing::info!(
|
||
"Completion requirement satisfied (tool '{}' called) for session {}",
|
||
required_tool,
|
||
self.session_info.id.0,
|
||
);
|
||
return result;
|
||
}
|
||
let mut attempt = 0u32;
|
||
loop {
|
||
attempt += 1;
|
||
let error_desc = match &result {
|
||
Ok(_) => "Agent finished without completing required task".into(),
|
||
Err(e) => format!("{e:?}"),
|
||
};
|
||
if attempt > recovery.max_retries {
|
||
tracing::error!(
|
||
"Auto-recovery exhausted after {attempt} attempts for session {}: {error_desc}",
|
||
self.session_info.id.0,
|
||
);
|
||
self.send_xai_notification(XaiSessionUpdate::AutoRecoveryExhausted {
|
||
attempts: attempt,
|
||
error: error_desc,
|
||
})
|
||
.await;
|
||
return result;
|
||
}
|
||
let delay_ms = std::cmp::min(
|
||
recovery.base_delay_ms * 2u64.pow(attempt.saturating_sub(1)),
|
||
recovery.max_delay_ms,
|
||
);
|
||
let delay = std::time::Duration::from_millis(delay_ms);
|
||
tracing::warn!(
|
||
"Auto-recovery attempt {}/{} for session {}: {error_desc}. Retrying in {}ms",
|
||
attempt,
|
||
recovery.max_retries,
|
||
self.session_info.id.0,
|
||
delay.as_millis(),
|
||
);
|
||
self.send_xai_notification(XaiSessionUpdate::AutoRecoveryStarted {
|
||
attempt,
|
||
max_retries: recovery.max_retries,
|
||
error: error_desc,
|
||
delay_ms: delay.as_millis() as u64,
|
||
})
|
||
.await;
|
||
sleep(delay).await;
|
||
let recovery_message = ConversationItem::auto_recovery(recovery_prompt.clone());
|
||
self.chat_state_handle.push_user_message(recovery_message);
|
||
result = self.process_conversation_turn(req_id, None).await;
|
||
if matches!(result, Ok(TurnOutcome::MaxTurnsReached { .. })) {
|
||
return result;
|
||
}
|
||
if let Ok(TurnOutcome::Completed {
|
||
ref tools_called, ..
|
||
}) = result
|
||
&& tools_called.iter().any(|name| name == &required_tool)
|
||
{
|
||
tracing::info!(
|
||
"Completion requirement satisfied after {} recovery attempt(s) \
|
||
(tool '{}' called) for session {}",
|
||
attempt,
|
||
required_tool,
|
||
self.session_info.id.0,
|
||
);
|
||
return result;
|
||
}
|
||
}
|
||
}
|
||
/// Compute the first-turn memory reminder, if one should be injected.
|
||
///
|
||
/// A block persisted by an earlier session segment (a prior `--resume`
|
||
/// process, or a turn before a compaction) is reused verbatim — see
|
||
/// [`conversation_has_memory_context`] for why re-searching is harmful.
|
||
///
|
||
/// [`conversation_has_memory_context`]: crate::session::helpers::memory_context::conversation_has_memory_context
|
||
pub(crate) async fn first_turn_memory_reminder(&self) -> Option<String> {
|
||
if self
|
||
.memory
|
||
.context_injected
|
||
.load(std::sync::atomic::Ordering::Relaxed)
|
||
{
|
||
return None;
|
||
}
|
||
self.memory
|
||
.context_injected
|
||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||
if !self.memory.initial_injection_config.enabled {
|
||
tracing::info!(
|
||
target : kigi_log::memory_log::TARGET,
|
||
"MEMORY_INJECT: first-turn injection disabled by config"
|
||
);
|
||
return None;
|
||
}
|
||
let (Some(storage), Some(params)) =
|
||
(self.memory.storage(), self.memory.backend_params.as_ref())
|
||
else {
|
||
return None;
|
||
};
|
||
let conversation = self.chat_state_handle.get_conversation().await;
|
||
if crate::session::helpers::memory_context::conversation_has_memory_context(&conversation) {
|
||
tracing::info!(
|
||
target : kigi_log::memory_log::TARGET,
|
||
"MEMORY_INJECT: existing memory-context block present in system message -- skipping re-injection to preserve prompt cache"
|
||
);
|
||
return None;
|
||
}
|
||
use kigi_tools::types::memory_backend::MemoryBackend as _;
|
||
let (injection_params, configured_min_score) =
|
||
build_initial_injection_backend_params(params, &self.memory.initial_injection_config);
|
||
let backend = crate::session::memory::MemoryBackendImpl::from_session_params(
|
||
storage,
|
||
&injection_params,
|
||
);
|
||
let raw_query =
|
||
crate::session::helpers::session_compact::extract_last_real_user_query(&conversation)
|
||
.unwrap_or_default();
|
||
let was_greeting = raw_query.is_empty()
|
||
|| raw_query.len() < 20
|
||
|| crate::session::helpers::memory_context::is_greeting(&raw_query);
|
||
let query = if was_greeting {
|
||
"project conventions preferences architecture".to_string()
|
||
} else {
|
||
raw_query
|
||
};
|
||
let inject_start = std::time::Instant::now();
|
||
let inject_results = backend.search(&query, 6, configured_min_score).await.ok();
|
||
let result_count = inject_results.as_ref().map_or(0, |r| r.len());
|
||
let top_score = inject_results
|
||
.as_ref()
|
||
.and_then(|r| r.first())
|
||
.map_or(0.0, |r| r.score);
|
||
let total_snippet_chars: usize = inject_results
|
||
.as_ref()
|
||
.map_or(0, |r| r.iter().map(|s| s.snippet.len()).sum());
|
||
tracing::info!(
|
||
target : kigi_log::memory_log::TARGET, configured_min_score,
|
||
"MEMORY_INJECT_SEARCH: results={result_count}"
|
||
);
|
||
inject_results.and_then(|results| {
|
||
crate::session::helpers::memory_context::format_memory_reminder(&results)
|
||
})
|
||
}
|
||
/// Inspect `tool_calls` for a `StructuredOutput` call and decide the turn's
|
||
/// next step, pushing the call's `tool_result` (correction / retry error /
|
||
/// terminal) as a side effect. Validates the args against `validator` and
|
||
/// bumps `retries` on a non-conforming retry.
|
||
async fn handle_structured_output_tool_call(
|
||
&self,
|
||
tool_calls: &mut Vec<kigi_sampling_types::conversation::ToolCall>,
|
||
validator: &Result<jsonschema::Validator, String>,
|
||
retries: &mut u32,
|
||
) -> StructuredOutputStep {
|
||
let Some(pos) = tool_calls
|
||
.iter()
|
||
.position(|tc| tc.name == STRUCTURED_OUTPUT_TOOL)
|
||
else {
|
||
return StructuredOutputStep::Proceed;
|
||
};
|
||
if tool_calls.len() > 1 {
|
||
for tc in tool_calls
|
||
.iter()
|
||
.filter(|tc| tc.name == STRUCTURED_OUTPUT_TOOL)
|
||
{
|
||
self.chat_state_handle
|
||
.push_tool_result(ConversationItem::tool_result(
|
||
tc.id.as_ref().to_owned(),
|
||
"Call StructuredOutput alone, exactly once, after all other tools finish.",
|
||
));
|
||
}
|
||
tool_calls.retain(|tc| tc.name != STRUCTURED_OUTPUT_TOOL);
|
||
return StructuredOutputStep::Proceed;
|
||
}
|
||
let call_id = tool_calls[pos].id.as_ref().to_owned();
|
||
let validated = validate_structured_output(validator, &tool_calls[pos].arguments);
|
||
if let Err(err) = &validated
|
||
&& *retries < STRUCTURED_OUTPUT_MAX_RETRIES
|
||
{
|
||
*retries += 1;
|
||
self.chat_state_handle
|
||
.push_tool_result(ConversationItem::tool_result(
|
||
call_id,
|
||
format!("{err}\nFix the arguments and call StructuredOutput again."),
|
||
));
|
||
return StructuredOutputStep::Retry;
|
||
}
|
||
self.chat_state_handle
|
||
.push_tool_result(ConversationItem::tool_result(
|
||
call_id,
|
||
match &validated {
|
||
Ok(_) => "Structured output accepted.".to_string(),
|
||
Err(err) => err.clone(),
|
||
},
|
||
));
|
||
StructuredOutputStep::Complete(validated)
|
||
}
|
||
/// Shared turn-completion bookkeeping (plan cleanup, signals snapshot +
|
||
/// persistence, BigQuery turn delta, feedback prompt). Runs identically for
|
||
/// the native and StructuredOutput-tool completion paths. Returns the
|
||
/// turn-end snapshot for `TurnOutcome::Completed`.
|
||
async fn finalize_turn_bookkeeping(
|
||
&self,
|
||
req_id: &str,
|
||
conv_turn_start: std::time::Instant,
|
||
turn_span_totals: &TurnSpanTotals,
|
||
model_fingerprint: Option<String>,
|
||
) -> Option<TurnDeltaSnapshot> {
|
||
self.emit_turn_end_plan_cleanup().await;
|
||
self.signals_handle().record_turn_complete();
|
||
let mut snapshot = self.signals_handle().take_turn_end_snapshot().await;
|
||
if let Some(snap) = snapshot.as_mut() {
|
||
self.apply_prompt_modes_to_snapshot(snap);
|
||
snap.turn_input_tokens = turn_span_totals.input_tokens.max(0) as u64;
|
||
snap.turn_output_tokens = turn_span_totals.output_tokens.max(0) as u64;
|
||
snap.turn_cached_input_tokens = turn_span_totals.cache_read_tokens.max(0) as u64;
|
||
for pr in &snap.delta.prs_created_this_turn {}
|
||
}
|
||
if let Some(snap) = snapshot.as_ref() {
|
||
let _ = self
|
||
.notifications
|
||
.persistence_tx
|
||
.send(PersistenceMsg::Signals(snap.current.clone()));
|
||
}
|
||
if let Some(request) = self
|
||
.feedback_manager
|
||
.maybe_request_feedback(Some(req_id.to_string()))
|
||
.await
|
||
{
|
||
self.send_feedback_notification(request).await;
|
||
}
|
||
snapshot
|
||
}
|
||
#[tracing::instrument(
|
||
name = "session.process_conversation_turn",
|
||
skip_all,
|
||
err,
|
||
fields(
|
||
session_id = %self.session_info.id.0,
|
||
model_id,
|
||
turn_tool_count,
|
||
turn_model_calls,
|
||
input_tokens = tracing::field::Empty,
|
||
output_tokens = tracing::field::Empty,
|
||
cache_read_tokens = tracing::field::Empty,
|
||
stop_reason = tracing::field::Empty,
|
||
response.has_tool_call = tracing::field::Empty,
|
||
request_id = tracing::field::Empty,
|
||
ttft_ms = tracing::field::Empty,
|
||
mcp_server.name = tracing::field::Empty,
|
||
mcp_tool.name = tracing::field::Empty,
|
||
agent.name = tracing::field::Empty,
|
||
skill.name = tracing::field::Empty,
|
||
query_source = tracing::field::Empty,
|
||
effort = tracing::field::Empty,
|
||
attempt = tracing::field::Empty,
|
||
parent_agent_id = tracing::field::Empty,
|
||
)
|
||
)]
|
||
async fn process_conversation_turn(
|
||
self: &Arc<Self>,
|
||
req_id: &str,
|
||
json_schema: Option<serde_json::Value>,
|
||
) -> Result<TurnOutcome, acp::Error> {
|
||
let conv_turn_start = std::time::Instant::now();
|
||
self.maybe_refresh_model_metadata_on_resume().await;
|
||
self.maybe_compact_on_model_switch().await;
|
||
self.chat_state_handle
|
||
.record_turn_start(chrono::Utc::now().timestamp_millis());
|
||
{
|
||
let span = tracing::Span::current();
|
||
if let Some(agent) = self.active_agent_type.lock().clone() {
|
||
span.record("agent.name", agent.as_str());
|
||
}
|
||
if let Some(skill) = self.active_skill.lock().clone() {
|
||
span.record("skill.name", skill.as_str());
|
||
}
|
||
span.record(
|
||
"query_source",
|
||
if self.startup_hints.is_subagent {
|
||
"subagent"
|
||
} else {
|
||
"main"
|
||
},
|
||
);
|
||
if let Some(parent) = self.startup_hints.parent_session_id.as_deref() {
|
||
span.record("parent_agent_id", parent);
|
||
}
|
||
}
|
||
if let Some(cfg) = self.chat_state_handle.get_sampling_config().await {
|
||
let span = tracing::Span::current();
|
||
span.record("model_id", cfg.model.as_str());
|
||
if let Some(effort) = cfg.reasoning_effort {
|
||
span.record("effort", effort.as_str());
|
||
}
|
||
}
|
||
let tool_prep_start = std::time::Instant::now();
|
||
let (tool_definitions, mcp_wait_ms) = self.prepare_tool_definitions_timed().await;
|
||
let total_prep_ms = tool_prep_start.elapsed().as_millis() as u64;
|
||
kigi_log::unified_log::info(
|
||
"shell.turn.tool_prep_done",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "tool_count" : tool_definitions.len(), "mcp_wait_ms" : mcp_wait_ms,
|
||
"total_prep_ms" : total_prep_ms, "elapsed_since_turn_start_ms" :
|
||
conv_turn_start.elapsed().as_millis() as u64, }
|
||
)),
|
||
);
|
||
self.record_turn_model().await;
|
||
let mut metrics_drop_guard = TurnMetrics::new();
|
||
let mut turn_tools_called: Vec<String> = Vec::new();
|
||
let mut tool_turn_count: usize = 1;
|
||
let mut loop_index: u32 = 0;
|
||
let mut todo_gate_fires: u32 = 0;
|
||
let mut auth_retry_schedule = AuthRetrySchedule::new();
|
||
let mut turn_span_totals = TurnSpanTotals::default();
|
||
let mut model_fingerprint: Option<String> = None;
|
||
let mut structured_output_retries: u32 = 0;
|
||
let structured_output_validator = json_schema.as_ref().map(|schema| {
|
||
jsonschema::validator_for(schema).map_err(|e| format!("invalid output schema: {e}"))
|
||
});
|
||
let schema_ok = matches!(structured_output_validator, Some(Ok(_)));
|
||
let native_backend = if json_schema.is_some() {
|
||
match self.chat_state_handle.get_sampling_config().await {
|
||
Some(c) => c.api_backend.supports_native_schema(),
|
||
None => {
|
||
tracing::warn!(
|
||
"structured output: no sampling config; using StructuredOutput tool"
|
||
);
|
||
false
|
||
}
|
||
}
|
||
} else {
|
||
false
|
||
};
|
||
let structured_output_native = schema_ok && native_backend;
|
||
let structured_output_tool = schema_ok && !native_backend;
|
||
if structured_output_tool {
|
||
self.push_system_reminder(
|
||
"A response schema is required. After any tool use, call the \
|
||
`StructuredOutput` tool exactly once with your final answer as its \
|
||
arguments; do not return the answer as text.",
|
||
);
|
||
}
|
||
loop {
|
||
self.emit_event(crate::session::events::Event::LoopStarted { loop_index });
|
||
loop_index += 1;
|
||
self.drain_pending_interjections().await;
|
||
self.flush_pending_skill_reminders().await;
|
||
self.inject_pending_monitor_events().await;
|
||
let memory_reminder = self.first_turn_memory_reminder().await;
|
||
if memory_reminder.is_some() {
|
||
self.memory
|
||
.injection_count
|
||
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||
tracing::info!(
|
||
target : kigi_log::memory_log::TARGET,
|
||
"MEMORY_INJECT: first-turn memory context injected"
|
||
);
|
||
}
|
||
self.maybe_inject_mcp_reminder().await;
|
||
if self.two_pass_active()
|
||
&& !self.compaction.prefire.has_cache()
|
||
&& self.should_prefire_two_pass().await
|
||
&& self.compaction.prefire.try_begin()
|
||
{
|
||
let actor = std::sync::Arc::clone(self);
|
||
let handle = tokio::task::spawn_local(async move {
|
||
actor.run_prefire_pass1().await;
|
||
});
|
||
self.compaction.prefire.set_handle(handle);
|
||
}
|
||
if let Some(trigger_info) = self.check_auto_compact_needed().await
|
||
&& let Err(e) = self.run_compact_only(trigger_info).await
|
||
{
|
||
tracing::error!(error = % e, "Pre-sampling auto-compaction failed");
|
||
}
|
||
let use_backend_search =
|
||
self.agent.borrow().backend_search_enabled() && self.supports_backend_search.get();
|
||
tracing::debug!(use_backend_search, "backend_search: turn tool resolution");
|
||
let mut effective_tools: Vec<ToolSpec> =
|
||
if let Some(ref override_tools) = self.forked_tool_override {
|
||
override_tools.clone()
|
||
} else {
|
||
self.turn_base_tool_specs(&tool_definitions)
|
||
};
|
||
if structured_output_tool && let Some(schema) = json_schema.clone() {
|
||
effective_tools.push(ToolSpec {
|
||
name: STRUCTURED_OUTPUT_TOOL.to_string(),
|
||
description: Some(
|
||
"Return your final answer as JSON matching the required schema. \
|
||
Call this exactly once, at the end."
|
||
.to_string(),
|
||
),
|
||
parameters: schema,
|
||
});
|
||
}
|
||
let build_req_start = std::time::Instant::now();
|
||
let request = self
|
||
.chat_state_handle
|
||
.build_request(
|
||
effective_tools,
|
||
memory_reminder,
|
||
self.memory.is_enabled(),
|
||
None,
|
||
self.session_info.id.to_string(),
|
||
req_id.to_owned(),
|
||
)
|
||
.await
|
||
.expect("chat state actor should be alive");
|
||
kigi_log::unified_log::debug(
|
||
"shell.turn.build_request_done",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "build_request_ms" : build_req_start.elapsed().as_millis() as
|
||
u64, "loop_index" : loop_index, }
|
||
)),
|
||
);
|
||
let mut request = request;
|
||
request.x_grok_session_id = Some(self.session_info.id.to_string());
|
||
request.x_grok_turn_idx =
|
||
Some(self.chat_state_handle.get_prompt_index().await.to_string());
|
||
request.x_grok_agent_id = Some(crate::util::agent_id::agent_id());
|
||
if request.x_grok_deployment_id.is_none() {
|
||
request.x_grok_deployment_id = crate::managed_config::resolve_deployment_id(
|
||
crate::managed_config::resolve_deployment_key().as_deref(),
|
||
);
|
||
}
|
||
if structured_output_native {
|
||
request.json_schema = json_schema.clone();
|
||
}
|
||
if use_backend_search {
|
||
request.hosted_tools = self.agent.borrow().hosted_tools().to_vec();
|
||
}
|
||
self.emit_event(crate::session::events::Event::PhaseChanged {
|
||
phase: crate::session::events::Phase::WaitingForModel,
|
||
});
|
||
self.observability_bridge
|
||
.emit(
|
||
kigi_tool_protocol::session_event::SessionEvent::PhaseChanged {
|
||
phase: kigi_tool_protocol::session_event::SessionPhase::Sampling,
|
||
},
|
||
)
|
||
.await;
|
||
kigi_log::unified_log::info(
|
||
"shell.turn.inference_start",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "loop_index" : loop_index, "elapsed_since_turn_start_ms" :
|
||
conv_turn_start.elapsed().as_millis() as u64, }
|
||
)),
|
||
);
|
||
let model_timer = std::time::Instant::now();
|
||
let (response, latency) = match self.run_turn_via_sampler(request.clone()).await? {
|
||
SamplerTurnOutcome::Response(r, latency) => (r, latency),
|
||
SamplerTurnOutcome::CompactAndResubmit => {
|
||
auth_retry_schedule.reset();
|
||
continue;
|
||
}
|
||
SamplerTurnOutcome::RefreshAuthAndResubmit => {
|
||
if let Some((attempt, delay)) = auth_retry_schedule.next_delay() {
|
||
let delay_ms = delay.as_millis() as u64;
|
||
tracing::warn!(
|
||
attempt,
|
||
delay_ms,
|
||
"auth 401 retry: backing off before resubmit"
|
||
);
|
||
kigi_log::unified_log::warn(
|
||
"shell.turn.auth_retry_backoff",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "loop_index" : loop_index, "attempt" : attempt,
|
||
"max_retries" : AuthRetrySchedule::MAX_RETRIES, "delay_ms" :
|
||
delay_ms, }
|
||
)),
|
||
);
|
||
self.send_xai_notification(XaiSessionUpdate::RetryState(
|
||
crate::extensions::notification::RetryState::Retrying {
|
||
attempt,
|
||
max_retries: AuthRetrySchedule::MAX_RETRIES,
|
||
reason: "Re-authenticated after 401; retrying request".to_string(),
|
||
},
|
||
))
|
||
.await;
|
||
sleep(delay).await;
|
||
continue;
|
||
}
|
||
let msg = format!(
|
||
"Auth recovery succeeded but inference request was \
|
||
still rejected (401) after {} retries",
|
||
AuthRetrySchedule::MAX_RETRIES
|
||
);
|
||
tracing::error!(msg);
|
||
return Err(acp::Error::internal_error().data(
|
||
crate::sampling::error::error_data_with_status(msg, Some(401)),
|
||
));
|
||
}
|
||
};
|
||
auth_retry_schedule.reset();
|
||
let model_elapsed_ms = model_timer.elapsed().as_millis() as u64;
|
||
let usage = response.usage.as_ref();
|
||
let prompt_tokens = usage.map(|u| u.prompt_tokens);
|
||
let cached_prompt_tokens = usage.map(|u| u.cached_prompt_tokens);
|
||
let completion_tokens = usage.map(|u| u.completion_tokens);
|
||
let reasoning_tokens = usage.map(|u| u.reasoning_tokens);
|
||
let ttft_ms = latency.time_to_first_token_ms;
|
||
let tokens_per_sec = match completion_tokens {
|
||
Some(ct) if ct > 0 => {
|
||
let decode_ms = match ttft_ms {
|
||
Some(ttft) if model_elapsed_ms > ttft => model_elapsed_ms - ttft,
|
||
_ => model_elapsed_ms,
|
||
};
|
||
(decode_ms > 0).then(|| {
|
||
let tps = f64::from(ct) * 1000.0 / decode_ms as f64;
|
||
(tps * 10.0).round() / 10.0
|
||
})
|
||
}
|
||
_ => None,
|
||
};
|
||
kigi_log::unified_log::info(
|
||
"shell.turn.inference_done",
|
||
Some(self.session_info.id.0.as_ref()),
|
||
Some(serde_json::json!(
|
||
{ "loop_index" : loop_index, "model_elapsed_ms" :
|
||
model_elapsed_ms, "elapsed_since_turn_start_ms" : conv_turn_start
|
||
.elapsed().as_millis() as u64, "ttft_ms" : ttft_ms, "itl_p50_ms"
|
||
: latency.itl_p50_ms, "attempts" : latency.attempts,
|
||
"prompt_tokens" : prompt_tokens, "cached_prompt_tokens" :
|
||
cached_prompt_tokens, "completion_tokens" : completion_tokens,
|
||
"reasoning_tokens" : reasoning_tokens, "tokens_per_sec" :
|
||
tokens_per_sec, }
|
||
)),
|
||
);
|
||
if let Some(usage) = response.usage.as_ref() {
|
||
self.chat_state_handle
|
||
.record_token_usage(u64::from(usage.total_tokens));
|
||
self.send_available_commands_update().await;
|
||
}
|
||
turn_span_totals.record(&tracing::Span::current(), &response);
|
||
let _ = self.compaction.auto_compact_suppressed.compare_exchange(
|
||
crate::session::compaction_config::SUPPRESS_UNTIL_SUCCESS,
|
||
crate::session::compaction_config::SUPPRESS_NONE,
|
||
std::sync::atomic::Ordering::Relaxed,
|
||
std::sync::atomic::Ordering::Relaxed,
|
||
);
|
||
let model_duration_ms = model_timer.elapsed().as_millis() as u64;
|
||
{
|
||
let model_id = self.current_model_id().await;
|
||
}
|
||
self.record_response_token_usage(&response, Some(model_duration_ms));
|
||
let mut tool_calls = response.tool_calls().to_vec();
|
||
metrics_drop_guard.record_model_response(tool_calls.len());
|
||
if let Some(fp) = response
|
||
.assistant()
|
||
.and_then(|a| a.model_fingerprint.clone())
|
||
{
|
||
model_fingerprint = Some(fp);
|
||
}
|
||
let fallback_text = response.fallback_text();
|
||
let stop_reason = response.stop_reason;
|
||
let response_is_empty = response.is_empty();
|
||
let turn_refused = stop_reason == Some(kigi_sampling_types::StopReason::ContentFilter);
|
||
let refusal_explanation = response.stop_message.clone();
|
||
let final_answer_text = json_schema.is_some().then(|| response.assistant_text());
|
||
for item in response.items {
|
||
match item {
|
||
kigi_sampling_types::ConversationItem::Assistant(_) => {
|
||
self.record_assistant_response(item).await;
|
||
}
|
||
_ => {
|
||
self.chat_state_handle.push_tool_result(item);
|
||
}
|
||
}
|
||
}
|
||
if let Some(text) = fallback_text {
|
||
tracing::warn!(
|
||
text_len = text.len(),
|
||
"emitting fallback AgentMessageChunk — no text chunks were streamed"
|
||
);
|
||
self.send_update(
|
||
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(
|
||
acp::ContentBlock::Text(acp::TextContent::new(text)),
|
||
)),
|
||
None,
|
||
)
|
||
.await;
|
||
}
|
||
if turn_refused && response_is_empty {
|
||
let mut notice = "The model provider refused to generate a response \
|
||
for this turn (content filter)."
|
||
.to_string();
|
||
if let Some(explanation) = refusal_explanation.as_deref() {
|
||
notice.push_str("\n\nProvider explanation: ");
|
||
notice.push_str(explanation);
|
||
}
|
||
tracing::warn!(
|
||
has_explanation = refusal_explanation.is_some(),
|
||
"model response was a provider refusal — emitting notice chunk"
|
||
);
|
||
self.send_update(
|
||
acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(
|
||
acp::ContentBlock::Text(acp::TextContent::new(notice)),
|
||
)),
|
||
None,
|
||
)
|
||
.await;
|
||
}
|
||
if tool_calls.is_empty() {
|
||
if !schema_ok
|
||
&& !turn_refused
|
||
&& let Some(gate_cfg) = self.todo_gate_policy()
|
||
{
|
||
let collected = self.collect_todo_gate_input(req_id).await;
|
||
let input = collected.as_input();
|
||
if let TodoGateDecision::Nudge { reminder, reason } = evaluate_todo_gate(&input)
|
||
{
|
||
if todo_gate_fires < gate_cfg.max_fires_per_prompt {
|
||
todo_gate_fires += 1;
|
||
tracing::info!(
|
||
prompt_id = % req_id, pending = ? input.pending,
|
||
unbacked_in_progress = ? input.in_progress_unbacked,
|
||
backed_in_progress = ? input.in_progress_backed,
|
||
backing_task_count = input.backing_task_count,
|
||
todo_gate_fires, reason = reason.as_str(),
|
||
"turn-end TodoGate: nudging model to advance remaining todos"
|
||
);
|
||
self.events
|
||
.emit(crate::session::events::Event::TodoGateFired {
|
||
fires: todo_gate_fires,
|
||
pending: input.pending.len(),
|
||
in_progress: input.in_progress_unbacked.len()
|
||
+ input.in_progress_backed.len(),
|
||
reason: reason.as_str(),
|
||
});
|
||
let rendered = self
|
||
.tool_bridge_handle()
|
||
.render_prompt(&reminder, &serde_json::json!({}))
|
||
.await
|
||
.unwrap_or(reminder);
|
||
self.push_system_reminder(&rendered);
|
||
continue;
|
||
}
|
||
let cap = gate_cfg.max_fires_per_prompt;
|
||
tracing::warn!(
|
||
prompt_id = % req_id, todo_gate_cap = cap,
|
||
"turn-end TodoGate: exhausted retries, falling through"
|
||
);
|
||
self.events
|
||
.emit(crate::session::events::Event::TodoGateExhausted {
|
||
pending: input.pending.len(),
|
||
});
|
||
self.push_system_reminder(&format!(
|
||
"The agent attempted to end this turn {cap} times \
|
||
with todos still pending or in_progress. Falling through \
|
||
to user. If you want autonomous progress, prompt the agent \
|
||
to continue explicitly, or clean up the todo list."
|
||
));
|
||
}
|
||
}
|
||
if self.drain_pending_interjections().await {
|
||
tracing::info!("Drained interjection(s) before turn completion — continuing");
|
||
continue;
|
||
}
|
||
let snapshot = self
|
||
.finalize_turn_bookkeeping(
|
||
req_id,
|
||
conv_turn_start,
|
||
&turn_span_totals,
|
||
model_fingerprint.clone(),
|
||
)
|
||
.await;
|
||
if self.drain_pending_interjections().await {
|
||
tracing::info!(
|
||
"Drained late interjection(s) during turn-end bookkeeping — continuing"
|
||
);
|
||
continue;
|
||
}
|
||
let structured_output = match (
|
||
structured_output_validator.as_ref(),
|
||
final_answer_text.as_ref(),
|
||
) {
|
||
(Some(validator), Some(text)) => {
|
||
Some(validate_structured_output(validator, text))
|
||
}
|
||
_ => None,
|
||
};
|
||
return Ok(TurnOutcome::Completed {
|
||
snapshot: Box::new(snapshot),
|
||
tools_called: turn_tools_called,
|
||
structured_output,
|
||
refusal: turn_refused,
|
||
});
|
||
}
|
||
if structured_output_tool && let Some(validator) = structured_output_validator.as_ref()
|
||
{
|
||
match self
|
||
.handle_structured_output_tool_call(
|
||
&mut tool_calls,
|
||
validator,
|
||
&mut structured_output_retries,
|
||
)
|
||
.await
|
||
{
|
||
StructuredOutputStep::Complete(validated) => {
|
||
turn_tools_called.push(STRUCTURED_OUTPUT_TOOL.to_string());
|
||
let snapshot = self
|
||
.finalize_turn_bookkeeping(
|
||
req_id,
|
||
conv_turn_start,
|
||
&turn_span_totals,
|
||
model_fingerprint.clone(),
|
||
)
|
||
.await;
|
||
return Ok(TurnOutcome::Completed {
|
||
snapshot: Box::new(snapshot),
|
||
tools_called: turn_tools_called,
|
||
structured_output: Some(validated),
|
||
refusal: false,
|
||
});
|
||
}
|
||
StructuredOutputStep::Retry => continue,
|
||
StructuredOutputStep::Proceed => {}
|
||
}
|
||
}
|
||
for tc in &tool_calls {
|
||
if let Some((server, tool)) =
|
||
crate::session::mcp_servers::parse_mcp_tool_name(&tc.name)
|
||
{
|
||
let span = tracing::Span::current();
|
||
span.record("mcp_server.name", server.as_str());
|
||
span.record("mcp_tool.name", tool.as_str());
|
||
}
|
||
turn_tools_called.push(tc.name.clone());
|
||
}
|
||
let tool_call_responses: Vec<ToolCallResponse> = tool_calls
|
||
.into_iter()
|
||
.map(|tc| ToolCallResponse {
|
||
id: tc.id.as_ref().to_owned(),
|
||
kind: "function".to_string(),
|
||
function: crate::sampling::types::ToolCallFunction {
|
||
name: tc.name,
|
||
arguments: tc.arguments.as_ref().to_owned(),
|
||
},
|
||
})
|
||
.collect();
|
||
self.emit_event(crate::session::events::Event::PhaseChanged {
|
||
phase: crate::session::events::Phase::ToolExecution,
|
||
});
|
||
self.observability_bridge
|
||
.emit(
|
||
kigi_tool_protocol::session_event::SessionEvent::PhaseChanged {
|
||
phase: kigi_tool_protocol::session_event::SessionPhase::ToolExecution,
|
||
},
|
||
)
|
||
.await;
|
||
let execute_tool_calls_result = self.execute_tool_calls(tool_call_responses).await;
|
||
match execute_tool_calls_result {
|
||
Ok(ToolLoop::PermissionReject { tool_name, reason }) => {
|
||
return Ok(TurnOutcome::Cancelled {
|
||
category: Some(
|
||
crate::session::events::CancellationCategory::PermissionRejected,
|
||
),
|
||
context: Some(serde_json::json!(
|
||
{ "tool_name" : tool_name, "reason" : reason, }
|
||
)),
|
||
});
|
||
}
|
||
Ok(ToolLoop::HookDenied { .. }) => {}
|
||
Ok(ToolLoop::Cancelled) => {
|
||
return Ok(TurnOutcome::Cancelled {
|
||
category: Some(
|
||
crate::session::events::CancellationCategory::PermissionCancelled,
|
||
),
|
||
context: None,
|
||
});
|
||
}
|
||
Ok(ToolLoop::FollowupMessage(followup_message)) => {
|
||
self.add_followup_message_as_user_turn(&followup_message)
|
||
.await;
|
||
continue;
|
||
}
|
||
_ => {}
|
||
}
|
||
let next_turn = tool_turn_count + 1;
|
||
if let Some(limit) = self.max_turns
|
||
&& next_turn > limit
|
||
{
|
||
tracing::info!(
|
||
session_id = % self.session_info.id, tool_turn_count, limit,
|
||
"max-turns limit reached, stopping"
|
||
);
|
||
return Ok(TurnOutcome::MaxTurnsReached { limit });
|
||
}
|
||
tool_turn_count = next_turn;
|
||
if let Some(trigger_info) = self.check_preflight_overflow().await {
|
||
if let Err(e) = self.run_compact_only(trigger_info).await {
|
||
tracing::error!(error = % e, "Preflight overflow compaction failed");
|
||
}
|
||
continue;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
/// Backoff schedule for resubmits after a *successful* 401 auth recovery
|
||
/// (fresh token minted, request to be re-sent).
|
||
///
|
||
/// Two hard-won invariants, both regressions from the silent-hang incident
|
||
/// where a turn froze 16m40s and then 11.6 days (user-cancelled at 27min):
|
||
///
|
||
/// - **Delays must be 1s/2s/4s.** `tokio_retry::ExponentialBackoff::
|
||
/// from_millis(base)` raises `base` to the attempt number, so the base must
|
||
/// stay small: `from_millis(1000)` yields 1000ⁿ ms = 1s → 16m40s → 11.57
|
||
/// days. `from_millis(2).factor(500)` yields 2ⁿ × 500ms = 1s, 2s, 4s.
|
||
/// - **The schedule is per-incident, not per-turn.** A long turn can span
|
||
/// several hourly gateway token rotations; each rotation is an independent
|
||
/// 401→refresh→retry event. Without `reset()` after a successful response,
|
||
/// the third rotation of one turn would land on the last (largest) delay
|
||
/// and the fourth would fail the turn outright.
|
||
struct AuthRetrySchedule {
|
||
delays: std::iter::Take<ExponentialBackoff>,
|
||
attempt: u32,
|
||
}
|
||
impl AuthRetrySchedule {
|
||
/// Consecutive post-recovery 401s tolerated before the turn fails.
|
||
const MAX_RETRIES: u32 = 3;
|
||
fn new() -> Self {
|
||
Self {
|
||
delays: ExponentialBackoff::from_millis(2)
|
||
.factor(500)
|
||
.max_delay(std::time::Duration::from_secs(10))
|
||
.take(Self::MAX_RETRIES as usize),
|
||
attempt: 0,
|
||
}
|
||
}
|
||
/// Next `(attempt_number, delay)` (1-indexed), or `None` once exhausted.
|
||
fn next_delay(&mut self) -> Option<(u32, std::time::Duration)> {
|
||
let delay = self.delays.next()?;
|
||
self.attempt += 1;
|
||
Some((self.attempt, delay))
|
||
}
|
||
/// A successful model response closes the incident: restart the schedule
|
||
/// so the next token rotation starts back at the shortest delay.
|
||
fn reset(&mut self) {
|
||
*self = Self::new();
|
||
}
|
||
}
|
||
#[cfg(test)]
|
||
mod auth_retry_schedule_tests {
|
||
use super::AuthRetrySchedule;
|
||
use std::time::Duration;
|
||
/// Pins the exact schedule. Guards against the `from_millis(1000)`
|
||
/// footgun (baseⁿ semantics): that spelling produced sleeps of 1s,
|
||
/// 16m40s, and 11.57 days, observed in the field as a silent
|
||
/// ~27-minute hang in `waiting_model` that the user had to cancel.
|
||
#[test]
|
||
fn schedule_is_one_two_four_seconds_then_exhausted() {
|
||
let mut schedule = AuthRetrySchedule::new();
|
||
let steps: Vec<_> = std::iter::from_fn(|| schedule.next_delay()).collect();
|
||
assert_eq!(
|
||
steps,
|
||
vec![
|
||
(1, Duration::from_secs(1)),
|
||
(2, Duration::from_secs(2)),
|
||
(3, Duration::from_secs(4)),
|
||
],
|
||
);
|
||
assert_eq!(
|
||
schedule.next_delay(),
|
||
None,
|
||
"must exhaust after MAX_RETRIES"
|
||
);
|
||
}
|
||
/// Each successful response must restart the schedule: hourly token
|
||
/// rotations within one long turn are independent incidents, so they
|
||
/// must not escalate toward exhaustion (turn failure).
|
||
#[test]
|
||
fn reset_restarts_delays_and_attempt_numbering() {
|
||
let mut schedule = AuthRetrySchedule::new();
|
||
schedule.next_delay();
|
||
schedule.next_delay();
|
||
schedule.reset();
|
||
assert_eq!(schedule.next_delay(), Some((1, Duration::from_secs(1))));
|
||
}
|
||
}
|
||
#[cfg(test)]
|
||
mod user_echo_broadcast_tests {
|
||
use super::{UserEchoMode, user_echo_mode};
|
||
/// Notification-drain: persisted (rewind/fork count user-chunk runs as
|
||
/// turn boundaries) but never broadcast live; the pager hides it via the
|
||
/// `hideFromScrollback` chunk meta.
|
||
#[test]
|
||
fn notification_drain_turn_is_persist_only() {
|
||
assert_eq!(
|
||
user_echo_mode("notifications-019e0000-0000-7000-8000-0000000000aa"),
|
||
UserEchoMode::PersistOnly
|
||
);
|
||
}
|
||
/// Real user prompts, cron (`/loop`) fires, and other turns still broadcast
|
||
/// live so multi-client / dashboard viewers stay in sync.
|
||
#[test]
|
||
fn user_and_cron_turns_broadcast_live() {
|
||
assert_eq!(user_echo_mode("my-prompt"), UserEchoMode::Broadcast);
|
||
assert_eq!(
|
||
user_echo_mode("scheduler-fired-abc"),
|
||
UserEchoMode::Broadcast
|
||
);
|
||
assert_eq!(
|
||
user_echo_mode("task-completed-bg-1"),
|
||
UserEchoMode::Broadcast
|
||
);
|
||
assert_eq!(
|
||
user_echo_mode("subagent-completed-xyz"),
|
||
UserEchoMode::Broadcast
|
||
);
|
||
}
|
||
/// Interject-fallback turns are persist-only: every pane already rendered
|
||
/// the text from the `x.ai/session/interjection` broadcast, so a live
|
||
/// echo would duplicate the block.
|
||
#[test]
|
||
fn interject_fallback_turn_is_persist_only() {
|
||
assert_eq!(
|
||
user_echo_mode("interject-fallback-019e24b7"),
|
||
UserEchoMode::PersistOnly
|
||
);
|
||
}
|
||
}
|
||
#[cfg(test)]
|
||
mod structured_output_validation_tests {
|
||
use super::validate_structured_output;
|
||
fn validator() -> Result<jsonschema::Validator, String> {
|
||
let schema = serde_json::json!(
|
||
{ "type" : "object", "properties" : { "name" : { "type" : "string" }, "age" :
|
||
{ "type" : "integer" } }, "required" : ["name", "age"],
|
||
"additionalProperties" : false, }
|
||
);
|
||
jsonschema::validator_for(&schema).map_err(|e| e.to_string())
|
||
}
|
||
#[test]
|
||
fn accepts_conforming_json() {
|
||
let v = validate_structured_output(&validator(), r#"{"name":"alice","age":30}"#).unwrap();
|
||
assert_eq!(v["name"], "alice");
|
||
}
|
||
#[test]
|
||
fn rejects_non_json() {
|
||
let err = validate_structured_output(&validator(), "not json").unwrap_err();
|
||
assert!(err.starts_with("model output was not valid JSON: "));
|
||
}
|
||
#[test]
|
||
fn rejects_schema_violation() {
|
||
let err = validate_structured_output(&validator(), r#"{"name":"alice"}"#).unwrap_err();
|
||
assert!(err.starts_with("output does not match the required schema: "));
|
||
}
|
||
#[test]
|
||
fn surfaces_invalid_schema_error() {
|
||
let bad: Result<jsonschema::Validator, String> = Err("invalid output schema: boom".into());
|
||
let err = validate_structured_output(&bad, r#"{"name":"alice","age":1}"#).unwrap_err();
|
||
assert_eq!(err, "invalid output schema: boom");
|
||
}
|
||
}
|