bingshuo/grok-build-upstream-mirror
Watch
1
0
Fork
You've already forked grok-build-upstream-mirror
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Synced from monorepo

Browse source 
Synced from monorepo

Changes:
- Shell: accept target response id on rewind execute
- Shell: stamp response id on chat user message chunks
- Worktree: optional rebuild and stale git registration cleanup in auto-GC
- Worktree: kind-aware auto-GC TTLs and config knobs
- Worktree: macOS process CWD scan and Unix PID liveness for GC guards
- Worktree: automatic throttled GC on startup (Linux age-based; non-Linux dead-only)
- Pager: add `[ui].combine_queued_prompts` to batch queued follow-ups
- Shell: stop overwriting user skills
- Tools: read markdown in `skills/` directories untruncated
- `/usage` shows per-session token and dollar usage in the TUI
- Security: prompt on environment-dumping `ps` variants
- Security: always-safe `kubectl` no longer runs arbitrary kubeconfig credential plugins without permission
- Tools: make scheduler deletion durable
- Shell: add relocation storage primitives
- Shell: give side model calls their own conversation ids
- Fix five workflow-runtime bugs (budget, pause, cancel, reconnect)
- Security: peel `env -S` / `--split-string` operands in the Bash permission gate (managed deny/ask)
- Pager: expose doctor in the TUI
- Security: block unauthorized RCE via abused safe commands
- Pager idle watcher cue: "1 subagent still running" instead of "watching · 1 subagent"
- Security: block `rg --pre` arbitrary code execution in auto-mode
- Voice: diagnose silent-mic failures (macOS permission) and add doctor/terminal-setup Voice section
- App builder deployer: `allow_forking` and `show_built_with_grok`
- Pager: stop stacking duplicate "Worked for" markers on parked turns
- Shell: support `max` as a distinct reasoning effort tier
- Tools: serialize background `/loop` fires on the whole work unit
- Shell: add working-directory relocation state primitives
- Proto: `ClientToolResult` and `ChatConfig` client-side tools
- Shell: model providers
- Chat: select App Builder product on the Build path
- Shell: attach author identity to feedback when the deployment opts in
- Doctor: fix for SSH wrap setup
- Workflow authoring skills: create-workflow and import-claude-workflow docs
- Add read-only grok doctor
- Sandbox: apply Landlock without a controlling TTY
- Pager: recover image paste over grok wrap on headless remotes
- Pager: make actions screen-mode aware
- Shell: resume sessions when the working directory moves
- Pager: centralize terminal diagnostics
- Workspace: gate inline shell file access
- Pager: centralize terminal probes
- Pager: edit minimal prompts in an external editor
- Pager: standardize backgrounding on Ctrl+B
- Shell: recap rides the parent turn's prompt cache
- Tools: add scheduler lifecycle version clock

Source-Revision: 0f4d7c91b8b2b408333f6de1e8a76cb8eaa71899
This commit is contained in:
grokkybara[bot] 2026-07-21 18:10:23 +00:00
commit 3af4d5d398
556 changed files with 56609 additions and 21892 deletions
Download patch file Download diff file Expand all files Collapse all files

266
crates/codegen/xai-grok-shell/src/tools/notification_bridge.rs
View file

@ -2,7 +2,7 @@
//! into `xai-grok-shell`'s native systems (ACP gateway, hunk tracker, file state tracker).
use crate::session::commands::SessionCommand;
use crate::session::commands::{NotificationPriority, NotificationSource};
use crate::session::persistence::PersistenceMsg;
use crate::session::persistence::{DurableAppendError, PersistenceHandle, PersistenceMsg};
use agent_client_protocol::{self as acp, Client as _};
use std::collections::HashMap;
use std::path::PathBuf;
@ -31,9 +31,8 @@ pub struct NotificationBridgeConfig {
/// Shared gate: when false, suppress gateway forwarding.
/// Events are still processed for hunk tracking and file state.
pub gateway_enabled: std::sync::Arc<std::sync::atomic::AtomicBool>,
/// Persistence channel for durably storing notifications.
/// Used to persist bash output even when the gateway gate is closed.
pub persistence_tx: mpsc::UnboundedSender<PersistenceMsg>,
/// Persistence handle for FIFO ordinary writes and durable tombstone barriers.
pub persistence: PersistenceHandle,
/// When true, send incremental `output_delta` instead of full `output`
/// in bash streaming updates. The client must opt in via the
/// `x.ai/incrementalBashOutput` capability.
@ -104,14 +103,112 @@ pub(crate) fn resolved_tool_name(slot: &std::sync::OnceLock<Option<String>>) ->
fn stamp_event_id(config: &NotificationBridgeConfig, meta: &mut Option<acp::Meta>) {
crate::util::event_id::ensure_event_id_meta(&config.session_id.0, meta);
}
fn stamp_scheduler_meta(
config: &NotificationBridgeConfig,
meta: &mut Option<acp::Meta>,
generation: &str,
revision: u64,
) {
stamp_event_id(config, meta);
let meta = meta.get_or_insert_with(acp::Meta::new);
meta.insert("x.ai/schedulerGeneration".to_owned(), generation.into());
meta.insert("x.ai/schedulerRevision".to_owned(), revision.into());
}
fn durable_append_landed(result: Result<(), DurableAppendError>) -> Result<(), String> {
match result {
Ok(()) => Ok(()),
Err(DurableAppendError::Committed(error)) => {
tracing::warn!(
% error, "Scheduler tombstone committed with bookkeeping failure"
);
Ok(())
}
Err(DurableAppendError::NotCommitted(error)) => {
Err(format!("scheduler tombstone was not committed: {error}"))
}
Err(DurableAppendError::AcknowledgementLost(error)) => Err(format!(
"scheduler tombstone commit status is unknown: {error}"
)),
}
}
async fn handle_scheduled_task_removed(
config: &NotificationBridgeConfig,
removed: xai_grok_tools::notification::ScheduledTaskRemoved,
acknowledgement: Option<tokio::sync::oneshot::Sender<Result<(), String>>>,
) -> Result<(), String> {
tracing::info!(task_id = % removed.task_id, "Scheduled task removed");
let result: Result<Box<serde_json::value::RawValue>, String> = async {
let mut meta = None;
stamp_scheduler_meta(config, &mut meta, &removed.generation, removed.revision);
let notification = crate::extensions::notification::SessionNotification {
session_id: config.session_id.clone(),
update: crate::extensions::notification::SessionUpdate::ScheduledTaskDeleted {
task_id: removed.task_id,
},
meta: meta.map(serde_json::Value::Object),
};
let params = serde_json::to_value(&notification)
.and_then(|value| serde_json::value::to_raw_value(&value))
.map_err(|error| format!("failed to serialize scheduled task deletion: {error}"))?;
let update = crate::session::storage::SessionUpdate::Xai(Box::new(notification));
if acknowledgement.is_some() {
durable_append_landed(config.persistence.append_update_durably(update).await)?;
} else {
config
.persistence
.tx
.send(PersistenceMsg::Update(update))
.map_err(|_| "session persistence stopped".to_owned())?;
}
Ok(params)
}
.await;
match result {
Ok(params) => {
if let Some(acknowledgement) = acknowledgement {
let _ = acknowledgement.send(Ok(()));
}
config
.gateway
.forward_fire_and_forget(acp::ExtNotification::new(
"x.ai/scheduled_task_deleted",
params.into(),
));
Ok(())
}
Err(error) => {
if let Some(acknowledgement) = acknowledgement {
let _ = acknowledgement.send(Err(error.clone()));
}
Err(error)
}
}
}
/// Create a `ToolNotificationHandle` and spawn a bridge task that
/// translates notifications into shell-native systems.
pub fn spawn_notification_bridge(config: NotificationBridgeConfig) -> ToolNotificationHandle {
let (handle, mut rx) = ToolNotificationHandle::channel();
let (handle, mut rx) = ToolNotificationHandle::acknowledged_channel();
tokio::task::spawn_local(async move {
let mut offsets: HashMap<String, usize> = HashMap::new();
while let Some(notification) = rx.recv().await {
handle_notification(&config, notification, &mut offsets).await;
while let Some(delivery) = rx.recv().await {
let acknowledgement = delivery.acknowledgement;
match delivery.notification {
ToolNotification::ScheduledTaskRemoved(removed) => {
if let Err(error) =
handle_scheduled_task_removed(&config, removed, acknowledgement).await
{
tracing::warn!(
% error, "Failed to handle scheduled task removal"
);
}
}
notification => {
handle_notification(&config, notification, &mut offsets).await;
if let Some(acknowledgement) = acknowledgement {
let _ = acknowledgement.send(Ok(()));
}
}
}
}
tracing::debug!("Notification bridge task exiting (sender dropped)");
});
@ -131,7 +228,7 @@ async fn emit_current_mode_update(
)),
);
stamp_event_id(config, &mut notification.meta);
let _ = config.persistence_tx.send(PersistenceMsg::Update(
let _ = config.persistence.tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Acp(Box::new(notification.clone())),
));
config.gateway.forward_fire_and_forget(notification);
@ -187,7 +284,7 @@ async fn handle_notification(
));
let mut notification = acp::SessionNotification::new(config.session_id.clone(), update);
stamp_event_id(config, &mut notification.meta);
let _ = config.persistence_tx.send(PersistenceMsg::Update(
let _ = config.persistence.tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Acp(Box::new(notification.clone())),
));
if config
@ -240,7 +337,7 @@ async fn handle_notification(
stamp_event_id(config, &mut meta_map);
notification.meta = meta_map.map(serde_json::Value::Object);
}
let _ = config.persistence_tx.send(PersistenceMsg::Update(
let _ = config.persistence.tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())),
));
let params = serde_json::to_value(&notification)
@ -476,7 +573,7 @@ async fn handle_notification(
stamp_event_id(config, &mut meta_map);
notification.meta = meta_map.map(serde_json::Value::Object);
}
let _ = config.persistence_tx.send(PersistenceMsg::Update(
let _ = config.persistence.tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())),
));
let params = serde_json::to_value(&notification)
@ -503,7 +600,8 @@ async fn handle_notification(
*config.turn_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Plan;
let snapshot = config.plan_mode.lock().snapshot();
let _ = config
.persistence_tx
.persistence
.tx
.send(PersistenceMsg::PlanModeState(snapshot));
emit_current_mode_update(config, xai_grok_tools::types::SessionMode::Plan).await;
}
@ -530,7 +628,8 @@ async fn handle_notification(
*config.turn_prompt_mode.lock() = crate::session::plan_mode::PromptMode::Agent;
let snapshot = config.plan_mode.lock().snapshot();
let _ = config
.persistence_tx
.persistence
.tx
.send(PersistenceMsg::PlanModeState(snapshot));
emit_current_mode_update(config, xai_grok_tools::types::SessionMode::Default).await;
}
@ -584,6 +683,8 @@ async fn handle_notification(
));
}
}
let mut meta = None;
stamp_scheduler_meta(config, &mut meta, &fired.generation, fired.revision);
let fired_notif = crate::extensions::notification::SessionNotification {
session_id: config.session_id.clone(),
update: crate::extensions::notification::SessionUpdate::ScheduledTaskFired {
@ -593,7 +694,7 @@ async fn handle_notification(
next_fire_at: fired.next_fire_at,
subagent_id: fired.subagent_id,
},
meta: None,
meta: meta.map(serde_json::Value::Object),
};
if let Ok(params) =
serde_json::to_value(&fired_notif).and_then(|v| serde_json::value::to_raw_value(&v))
@ -665,36 +766,15 @@ async fn handle_notification(
});
}
ToolNotification::ScheduledTaskRemoved(removed) => {
tracing::info!(task_id = % removed.task_id, "Scheduled task removed");
let mut notification = crate::extensions::notification::SessionNotification {
session_id: config.session_id.clone(),
update: crate::extensions::notification::SessionUpdate::ScheduledTaskDeleted {
task_id: removed.task_id,
},
meta: None,
};
{
let mut meta_map = None;
stamp_event_id(config, &mut meta_map);
notification.meta = meta_map.map(serde_json::Value::Object);
}
let _ = config.persistence_tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())),
));
if let Ok(params) = serde_json::to_value(&notification)
.and_then(|v| serde_json::value::to_raw_value(&v))
{
config
.gateway
.forward_fire_and_forget(acp::ExtNotification::new(
"x.ai/scheduled_task_deleted",
params.into(),
));
if let Err(error) = handle_scheduled_task_removed(config, removed, None).await {
tracing::warn!(% error, "Failed to handle scheduled task removal");
}
}
ToolNotification::ScheduledTaskCreated(created) => {
tracing::info!(task_id = % created.task_id, "Scheduled task created");
let mut notification = crate::extensions::notification::SessionNotification {
let mut meta = None;
stamp_scheduler_meta(config, &mut meta, &created.generation, created.revision);
let notification = crate::extensions::notification::SessionNotification {
session_id: config.session_id.clone(),
update: crate::extensions::notification::SessionUpdate::ScheduledTaskCreated {
task_id: created.task_id,
@ -702,14 +782,9 @@ async fn handle_notification(
human_schedule: created.human_schedule,
next_fire_at: created.next_fire_at,
},
meta: None,
meta: meta.map(serde_json::Value::Object),
};
{
let mut meta_map = None;
stamp_event_id(config, &mut meta_map);
notification.meta = meta_map.map(serde_json::Value::Object);
}
let _ = config.persistence_tx.send(PersistenceMsg::Update(
let _ = config.persistence.tx.send(PersistenceMsg::Update(
crate::session::storage::SessionUpdate::Xai(Box::new(notification.clone())),
));
if let Ok(params) = serde_json::to_value(&notification)
@ -796,7 +871,7 @@ mod tests {
prompt_index: Arc::new(TokioMutex::new(0)),
cwd: PathBuf::from("/tmp"),
gateway_enabled: Arc::new(std::sync::atomic::AtomicBool::new(true)),
persistence_tx,
persistence: PersistenceHandle::from_sender_for_test(persistence_tx),
incremental_bash_output: false,
plan_mode: Arc::new(parking_lot::Mutex::new(
crate::session::plan_mode::PlanModeTracker::new(PathBuf::from("/tmp/test-session")),
@ -1164,8 +1239,7 @@ mod tests {
};
assert!(
matches!(admission.fallback.source, NotificationSource::MonitorCompleted {
ref task_id }
if task_id == "mon-timeout")
ref task_id } if task_id == "mon-timeout")
);
assert!(admission.respond_to.send(true).is_err());
let _ = respond_to.send(Ok(crate::session::commands::PromptTurnOk {
@ -1463,6 +1537,8 @@ if task_id == "mon-timeout")
prompt: "check deploy".into(),
human_schedule: "every 5 minutes".into(),
next_fire_at: Some("2026-01-01T00:00:00Z".into()),
generation: "generation-a".into(),
revision: 1,
},
);
let mut offsets = HashMap::new();
@ -1476,6 +1552,9 @@ if task_id == "mon-timeout")
&notif.update,
crate::extensions::notification::SessionUpdate::ScheduledTaskCreated { .. }
));
let meta = notif.meta.as_ref().expect("scheduler metadata");
assert_eq!(meta["x.ai/schedulerGeneration"], "generation-a");
assert_eq!(meta["x.ai/schedulerRevision"], 1);
assert!(
notif
.meta
@ -1537,13 +1616,14 @@ if task_id == "mon-timeout")
#[tokio::test]
async fn scheduled_task_removed_is_persisted() {
let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full();
let notification = ToolNotification::ScheduledTaskRemoved(
xai_grok_tools::notification::types::ScheduledTaskRemoved {
task_id: "loop-1".into(),
},
);
let mut offsets = HashMap::new();
handle_notification(&config, notification, &mut offsets).await;
let removed = xai_grok_tools::notification::ScheduledTaskRemoved {
task_id: "loop-1".into(),
generation: "generation-a".into(),
revision: 2,
};
handle_scheduled_task_removed(&config, removed, None)
.await
.unwrap();
let msg = persistence_rx
.try_recv()
.expect("scheduled_task_removed must be persisted");
@ -1557,10 +1637,49 @@ if task_id == "mon-timeout")
xai_persisted_event_id(&notif).is_some(),
"the persisted deletion line must be stamped"
);
let meta = notif.meta.as_ref().expect("scheduler metadata");
assert_eq!(meta["x.ai/schedulerGeneration"], "generation-a");
assert_eq!(meta["x.ai/schedulerRevision"], 2);
}
_ => panic!("expected PersistenceMsg::Update(Xai(ScheduledTaskDeleted))"),
}
}
#[tokio::test]
async fn acknowledged_scheduler_removal_appends_before_ack_and_broadcast() {
let (config, mut gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full();
let removed = xai_grok_tools::notification::ScheduledTaskRemoved {
task_id: "loop-ack".into(),
generation: "generation-a".into(),
revision: 17,
};
let (acknowledgement, mut receipt) = tokio::sync::oneshot::channel::<Result<(), String>>();
let persistence = async {
let PersistenceMsg::AppendUpdateDurablyAndAck {
update: crate::session::storage::SessionUpdate::Xai(notification),
respond_to,
} = persistence_rx.recv().await.expect("durable append")
else {
panic!("expected durable scheduler tombstone");
};
assert_eq!(notification.meta.unwrap()["x.ai/schedulerRevision"], 17);
assert!(gateway_rx.try_recv().is_err());
assert!(matches!(
receipt.try_recv(),
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
));
respond_to.send(Ok(())).unwrap();
receipt.await.unwrap().unwrap();
};
let (result, ()) = tokio::join!(
handle_scheduled_task_removed(&config, removed, Some(acknowledgement)),
persistence,
);
result.unwrap();
assert!(matches!(
gateway_rx.try_recv(),
Ok(xai_acp_lib::AcpClientMessage::ExtNotification(_))
));
}
fn xai_persisted_event_id(
notif: &crate::extensions::notification::SessionNotification,
) -> Option<String> {
@ -1644,16 +1763,33 @@ if task_id == "mon-timeout")
_ => panic!("expected Acp update"),
}
}
#[test]
fn durable_append_mapping_respects_commit_disposition() {
assert!(
durable_append_landed(Err(DurableAppendError::Committed(std::io::Error::other(
"summary failed"
),)))
.is_ok()
);
for failure in [
DurableAppendError::NotCommitted(std::io::Error::other("append failed")),
DurableAppendError::AcknowledgementLost(std::io::Error::other("lost")),
] {
assert!(durable_append_landed(Err(failure)).is_err());
}
}
#[tokio::test]
async fn scheduled_task_fired_is_not_persisted() {
let (config, _gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full();
let (config, mut gateway_rx, mut persistence_rx, _cmd_rx) = make_test_config_full();
let notification = ToolNotification::ScheduledTaskFired(
xai_grok_tools::notification::types::ScheduledTaskFired {
task_id: "loop-1".into(),
prompt: "check deploy".into(),
human_schedule: "every 5 minutes".into(),
next_fire_at: Some("2026-01-01T00:00:00Z".into()),
subagent_id: None,
subagent_id: Some("subagent-1".into()),
generation: "generation-a".into(),
revision: 3,
},
);
let mut offsets = HashMap::new();
@ -1662,6 +1798,15 @@ if task_id == "mon-timeout")
persistence_rx.try_recv().is_err(),
"scheduled_task_fired must NOT be persisted (recurring \u{2192} unbounded log growth)"
);
let fired = gateway_rx
.try_recv()
.expect("scheduled fire must be broadcast");
let xai_acp_lib::AcpClientMessage::ExtNotification(fired) = fired else {
panic!("expected scheduler fire notification");
};
let value: serde_json::Value = serde_json::from_str(fired.request.params.get()).unwrap();
assert_eq!(value["_meta"]["x.ai/schedulerGeneration"], "generation-a");
assert_eq!(value["_meta"]["x.ai/schedulerRevision"], 3);
}
fn make_monitor_event_notification(task_id: &str, owner: Option<&str>) -> ToolNotification {
ToolNotification::MonitorEvent(xai_grok_tools::notification::types::MonitorEvent {
@ -1827,8 +1972,7 @@ if task_id == "mon-timeout")
assert_eq!(priority, NotificationPriority::Later);
assert!(
matches!(source, NotificationSource::BashTaskCompleted { ref task_id
}
if task_id == "bg-disabled")
} if task_id == "bg-disabled")
);
let text = match &prompt_blocks[0] {
acp::ContentBlock::Text(t) => &t.text,

193
crates/codegen/xai-grok-shell/src/tools/tool_context.rs
View file

@ -16,23 +16,117 @@ use xai_grok_paths::AbsPathBuf;
use xai_grok_workspace::file_system::{AsyncFileSystem, AsyncFsWrapper};
use xai_grok_workspace::session::file_state::FileStateHandle;
use xai_hunk_tracker::HunkTrackerHandle;
/// RAII marker: the turn is blocked inside an interruptible wait. Increments
/// [`ToolContext::blocking_wait_depth`] for its lifetime; `Drop` decrements
/// (a cancelled turn can't leak the count).
pub(crate) struct BlockingWaitGuard(Arc<std::sync::atomic::AtomicUsize>);
#[derive(Debug, Clone, Default)]
pub struct TaskOutputTokenBudget {
inner: Arc<parking_lot::Mutex<TaskOutputTokenBudgetState>>,
}
#[derive(Debug, Default)]
struct TaskOutputTokenBudgetState {
total: Option<u64>,
spent: u64,
incomplete: bool,
}
impl TaskOutputTokenBudget {
pub fn limited(total: u64) -> Self {
debug_assert!(total > 0, "task output grant must be positive");
Self {
inner: Arc::new(parking_lot::Mutex::new(TaskOutputTokenBudgetState {
total: Some(total),
spent: 0,
incomplete: false,
})),
}
}
pub fn remaining(&self) -> Option<u64> {
let state = self.inner.lock();
state.total.map(|total| total.saturating_sub(state.spent))
}
pub fn clamp_request(&self, configured: Option<u32>) -> Option<u32> {
let remaining = self.remaining()?;
if remaining == 0 {
return Some(0);
}
let remaining = u32::try_from(remaining).unwrap_or(u32::MAX);
Some(configured.map_or(remaining, |configured| configured.min(remaining)))
}
pub fn record_reported_output(&self, output_tokens: u64) {
let mut state = self.inner.lock();
state.spent = state.spent.saturating_add(output_tokens);
if let Some(total) = state.total
&& state.spent > total
{
state.spent = total;
state.incomplete = true;
}
}
pub fn mark_incomplete_and_exhaust(&self) {
let mut state = self.inner.lock();
state.incomplete = true;
if let Some(total) = state.total {
state.spent = state.spent.max(total);
}
}
pub fn usage(&self) -> (u64, bool) {
let state = self.inner.lock();
(state.spent, state.incomplete)
}
pub fn is_limited(&self) -> bool {
self.inner.lock().total.is_some()
}
}
pub struct BlockingWaitState(std::sync::Mutex<BlockingWaitInner>);
#[derive(Default)]
struct BlockingWaitInner {
depth: usize,
generation: u64,
}
impl BlockingWaitState {
pub(crate) fn new() -> Self {
Self(std::sync::Mutex::new(BlockingWaitInner::default()))
}
pub(crate) fn depth(&self) -> usize {
self.0
.lock()
.expect("blocking wait state mutex poisoned")
.depth
}
#[cfg(test)]
pub(crate) fn set_depth_for_test(&self, depth: usize) {
self.0
.lock()
.expect("blocking wait state mutex poisoned")
.depth = depth;
}
pub(crate) fn reset(&self) {
let mut state = self.0.lock().expect("blocking wait state mutex poisoned");
state.generation = state.generation.wrapping_add(1);
state.depth = 0;
}
}
pub(crate) struct BlockingWaitGuard {
state: Arc<BlockingWaitState>,
generation: u64,
}
impl BlockingWaitGuard {
pub(crate) fn enter(depth: Arc<std::sync::atomic::AtomicUsize>) -> Self {
depth.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Self(depth)
pub(crate) fn enter(state: Arc<BlockingWaitState>) -> Self {
let generation = {
let mut inner = state.0.lock().expect("blocking wait state mutex poisoned");
inner.depth = inner.depth.saturating_add(1);
inner.generation
};
Self { state, generation }
}
}
impl Drop for BlockingWaitGuard {
fn drop(&mut self) {
let _ = self.0.fetch_update(
std::sync::atomic::Ordering::SeqCst,
std::sync::atomic::Ordering::SeqCst,
|depth| Some(depth.saturating_sub(1)),
);
let mut inner = self
.state
.0
.lock()
.expect("blocking wait state mutex poisoned");
if inner.generation == self.generation {
inner.depth = inner.depth.saturating_sub(1);
}
}
}
/// Session-level context. NOT used for tool execution (bridge handles that).
@ -71,6 +165,7 @@ pub struct ToolContext {
/// Shared turn-active flag — set `true` at turn start, `false` at turn end.
/// Used by the between-turn completion drain in `handle_prompt`.
pub is_turn_active: Option<Arc<std::sync::atomic::AtomicBool>>,
pub(crate) unattributed_background_usage: Arc<std::sync::atomic::AtomicBool>,
/// Shared buffer for mid-turn monitor event notifications.
/// Events pushed here are drained by the session turn loop
/// (`inject_pending_monitor_events`) and surfaced as ONE hidden
@ -111,9 +206,33 @@ pub struct ToolContext {
/// Count of interruptible blocking waits the running turn is parked in (via
/// [`BlockingWaitGuard`]). `queue_input` reads it: a prompt arriving while
/// non-zero takes the send-now path.
pub blocking_wait_depth: Arc<std::sync::atomic::AtomicUsize>,
pub blocking_wait_depth: Arc<BlockingWaitState>,
pub task_output_token_budget: Option<TaskOutputTokenBudget>,
pub(crate) sampler_retry_only_before_output: bool,
}
impl ToolContext {
pub(crate) fn clamp_task_model_request(
&self,
configured: Option<u32>,
) -> Result<Option<u32>, &'static str> {
match self.task_output_token_budget.as_ref() {
Some(budget) => match budget.clamp_request(configured) {
Some(0) => Err("workflow child output-token budget exhausted"),
clamped => Ok(clamped),
},
None => Ok(configured),
}
}
pub(crate) fn record_task_model_output(&self, output_tokens: u64) {
if let Some(budget) = self.task_output_token_budget.as_ref() {
budget.record_reported_output(output_tokens);
}
}
pub(crate) fn fail_task_output_usage_closed(&self) {
if let Some(budget) = self.task_output_token_budget.as_ref() {
budget.mark_incomplete_and_exhaust();
}
}
pub fn new(
cwd: AbsPathBuf,
gateway: Option<GatewaySender>,
@ -142,6 +261,7 @@ impl ToolContext {
lsp: None,
lsp_server_names: Vec::new(),
is_turn_active: None,
unattributed_background_usage: Arc::new(std::sync::atomic::AtomicBool::new(false)),
monitor_event_buffer: None,
task_completion_reservations: None,
task_wake_suppressed: None,
@ -151,7 +271,9 @@ impl ToolContext {
xai_grok_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL.to_string(),
auto_wake_enabled: true,
goal_loop_active_gate: Arc::new(std::sync::atomic::AtomicBool::new(false)),
blocking_wait_depth: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
task_output_token_budget: None,
sampler_retry_only_before_output: false,
}
}
pub fn with_preloaded_env(
@ -179,6 +301,7 @@ impl ToolContext {
lsp: None,
lsp_server_names: Vec::new(),
is_turn_active: None,
unattributed_background_usage: Arc::new(std::sync::atomic::AtomicBool::new(false)),
monitor_event_buffer: None,
task_completion_reservations: None,
task_wake_suppressed: None,
@ -188,7 +311,9 @@ impl ToolContext {
xai_grok_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL.to_string(),
auto_wake_enabled: true,
goal_loop_active_gate: Arc::new(std::sync::atomic::AtomicBool::new(false)),
blocking_wait_depth: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
task_output_token_budget: None,
sampler_retry_only_before_output: false,
}
}
pub fn with_file_state_handle(mut self, handle: FileStateHandle) -> Self {
@ -207,7 +332,40 @@ impl ToolContext {
}
}
#[cfg(test)]
mod output_budget_tests {
use super::TaskOutputTokenBudget;
#[test]
fn clamps_every_request_to_remaining_and_stops_at_zero() {
let budget = TaskOutputTokenBudget::limited(10);
assert_eq!(budget.clamp_request(None), Some(10));
assert_eq!(budget.clamp_request(Some(7)), Some(7));
budget.record_reported_output(6);
assert_eq!(budget.clamp_request(None), Some(4));
assert_eq!(budget.clamp_request(Some(9)), Some(4));
budget.record_reported_output(4);
assert_eq!(budget.clamp_request(None), Some(0));
}
#[test]
fn provider_output_not_context_drives_spend() {
let budget = TaskOutputTokenBudget::limited(100);
let provider_prompt_tokens = 90_000u64;
budget.record_reported_output(25);
assert_eq!(budget.usage(), (25, false));
assert_eq!(provider_prompt_tokens, 90_000);
assert_eq!(budget.remaining(), Some(75));
}
#[test]
fn unknown_usage_exhausts_grant_pessimistically() {
let budget = TaskOutputTokenBudget::limited(50);
budget.record_reported_output(7);
budget.mark_incomplete_and_exhaust();
assert_eq!(budget.usage(), (50, true));
assert_eq!(budget.clamp_request(None), Some(0));
}
}
#[cfg(test)]
mod tests {
use super::BlockingWaitState;
use crate::{terminal::AsyncTerminalRunner, tools::ToolContext};
use std::collections::HashMap;
use std::sync::Arc;
@ -236,6 +394,7 @@ mod tests {
lsp: None,
lsp_server_names: Vec::new(),
is_turn_active: None,
unattributed_background_usage: Arc::new(std::sync::atomic::AtomicBool::new(false)),
monitor_event_buffer: None,
task_completion_reservations: None,
task_wake_suppressed: None,
@ -245,7 +404,9 @@ mod tests {
xai_grok_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL.to_string(),
auto_wake_enabled: true,
goal_loop_active_gate: Arc::new(std::sync::atomic::AtomicBool::new(false)),
blocking_wait_depth: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
blocking_wait_depth: Arc::new(BlockingWaitState::new()),
task_output_token_budget: None,
sampler_retry_only_before_output: false,
}
}
}