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:
- grok-shell: send an expired external-provider credential to the sign-in flow, not a 401 loop
- pager: clickable ▲ jumps to the top of the response being read
- grok-shell: keep a large task log from making the completion message too long
- Plan viewer scrollbar: widen grab zone to the border column; fix striped thumb in Terminal.app
- pager: poll the tmux probe teardown grace instead of sleeping it
- security: vendor-compat MCP kill switch is now actually enforced when reported as on
- grok-shell: restore session eviction when a leader client disconnects
- Bump rust-toolchain to 1.93.0
- workspace: lexical-normalize permission path patterns before glob matching
- pager: reject garbage Enter in the /resume picker
- pager: show Mermaid affordances in plan mode preview
- pager: drop manage-account link from /session-info
- workspace: auto-approve read-only git queries; defer write floor to auto classifier
- Add free-form pattern editor to the "Always allow" command prompt
- grok-shell: fix /btw caching
- pager: Tab walks answers in the ask_user_question card
- External-provider auth refresh: single 7s attempt instead of 3×5s
- pager: don't resurrect finished background tasks as Running when completion arrives first
- pager: report tmux truecolor clamping in Doctor
- Fix plan viewer scrollbar click+drag hijacked by comment gutter
- pager/shell: stop double Recap after the same last turn
- sampler: preserve x-should-retry through stream collection
- pager: clear plan-mode indicator immediately when the user approves a plan
- pager: tmux does not re-read its config on reattach

Source-Revision: 64c4de99cc822b25ce9c54ab5a4f372093d0885d
This commit is contained in:
grokkybara[bot] 2026-08-03 08:17:57 +00:00
commit 780d1388ff
323 changed files with 12258 additions and 7226 deletions
Download patch file Download diff file Expand all files Collapse all files

25
crates/codegen/xai-grok-shell/CHANGELOG.md
View file

@ -1,5 +1,26 @@
# Changelog
# 0.2.118 — 2026-07-31
## Features
- **Sessions** can now be permanently deleted from the dashboard by pressing Ctrl+X twice on an idle row, or from the welcome list with d then y.
- **Keyboard shortcuts help** (Ctrl+.) now shows how to browse prompt history and search the conversation.
- **grok doctor** now warns when tmux is reducing colors and can fix the config.
## Bug Fixes
- **`/btw`** now retries on temporary model overload instead of failing immediately.
- **Session sharing** is temporarily disabled.
- **`[stop]`** / Ctrl+C during `/compact` now cancels instead of no-opping.
- **Automatic recaps** no longer appear twice after the same turn.
- **Background task wait timeout** descriptions and limits now match the client's actual configured ceiling.
- **Background tasks** no longer stay stuck as 'Running' in the tasks pane when they finish quickly.
- **Plan mode indicator** now disappears right after approving a plan instead of lingering.
- **Dragging the scrollbar** in the plan preview now works as expected.
- **Compaction** now correctly handles certain context-length errors from the inference API.
# 0.2.117 — 2026-07-30
## Features
@ -37,6 +58,10 @@
# 0.2.115 — 2026-07-29
## Features
- **Delete sessions from the dashboard and welcome list.** On the dashboard, press `Ctrl+X` twice (or hover a settled row and click `[✗]` twice); in the welcome and `/resume` lists, press `d` then `y`.
## Bug Fixes
- **Fixed chat history corruption** that could duplicate tool results or cause later 400 errors after repeated identical tool calls.

3
crates/codegen/xai-grok-shell/Cargo.toml
View file

@ -1,12 +1,11 @@
[package]
license = "Apache-2.0"
name = "xai-grok-shell"
version = "0.2.117"
version = "0.2.118"
edition.workspace = true
[features]
default = []
unstable = []
dhat-heap = ["dep:dhat"]
# Session synthesis + in-process e2e harness (`session::testkit`) for soak,
# load, and bench tests. Off by default; the tests/benches that use it declare

11
crates/codegen/xai-grok-shell/README.md
View file

@ -302,6 +302,7 @@ This is transparent — you don't need to do anything. Grok handles it in the ba
- **Before expiry:** If your binary returned `expires_in` in its JSON output, or you set `auth_token_ttl` in config, Grok re-runs the binary ~5 minutes before the token expires, so you never see an auth error.
- **On auth error:** If the server rejects a request with 401/403 (e.g. token was revoked or expired), Grok re-runs the binary and retries the request once.
- **When the refresh run can't mint:** refreshes are headless (no stdin, short timeout), so a binary that needs you to complete an SSO flow cannot succeed there. Grok then stops treating the stored credential as usable and runs your binary in its interactive mode instead — at startup that is the same sign-in flow a machine with no credentials gets; mid-session the turn fails with a re-auth prompt and `/login` re-runs the binary.
- **OIDC:** If you're using OIDC and have a `refresh_token`, Grok silently refreshes via your IdP without re-opening the browser.
**Tuning the refresh buffer:**
@ -333,7 +334,7 @@ Common log messages:
| `auth: running external auth provider` | Your binary is being called (includes the command and whether it's a refresh) |
| `auth: external auth provider returned fresh token` | Success — token was parsed and stored |
| `auth: external auth provider failed` | Binary exited non-zero, or exited 0 but stdout was empty/unparseable (the `error` field has details) |
| `auth: external auth provider timed out (likely needs interactive auth), killing` | Binary didn't exit before the timeout (60s initial, 5s mid-session refresh) and was killed |
| `auth: external auth provider timed out (likely needs interactive auth), killing` | Binary didn't exit before the timeout (60s initial, 7s mid-session refresh) and was killed |
| `auth: failed to start external auth provider` | The command couldn't be spawned (e.g. binary not found) |
### Per-Model Auth Providers
@ -1350,12 +1351,6 @@ output_byte_limit = 65536 # max output size (64KB)
[toolset.web_fetch]
proxy_endpoint = "https://proxy.example.com" # egress proxy URL (all requests routed through it)
allowed_domains = ["docs.rs", "x.ai"] # override the built-in ~84-domain allowlist
[shortcuts]
send = ["Enter"]
newline = ["Shift+Enter", "Alt+Enter"]
quit = ["Ctrl+D", "Ctrl+Q"]
confirm_quit = true
```
### Telemetry
@ -2009,7 +2004,7 @@ args = ["-y", "mcp-remote", "https://mcp.linear.app/mcp"]
If you also have a `linear` server in `~/.grok/config.toml`, the project version replaces it entirely.
> **Note:** Only `[mcp_servers]` is supported in project-scoped `.grok/config.toml`. Other config sections (models, shortcuts, etc.) are only read from `~/.grok/config.toml`.
> **Note:** Only `[mcp_servers]` is supported in project-scoped `.grok/config.toml`. Other config sections (models, etc.) are only read from `~/.grok/config.toml`.
### Tool Naming

62
crates/codegen/xai-grok-shell/changelogs/0.2.118.json Normal file
View file

@ -0,0 +1,62 @@
[
{
"category": "fixes",
"description": "**`/btw`** now retries on temporary model overload instead of failing immediately.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Session sharing** is temporarily disabled.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**`[stop]`** / Ctrl+C during `/compact` now cancels instead of no-opping.",
"breaking_change": false
},
{
"category": "features",
"description": "**Sessions** can now be permanently deleted from the dashboard by pressing Ctrl+X twice on an idle row, or from the welcome list with d then y.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Automatic recaps** no longer appear twice after the same turn.",
"breaking_change": false
},
{
"category": "features",
"description": "**Keyboard shortcuts help** (Ctrl+.) now shows how to browse prompt history and search the conversation.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Background task wait timeout** descriptions and limits now match the client's actual configured ceiling.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Background tasks** no longer stay stuck as 'Running' in the tasks pane when they finish quickly.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Plan mode indicator** now disappears right after approving a plan instead of lingering.",
"breaking_change": false
},
{
"category": "features",
"description": "**grok doctor** now warns when tmux is reducing colors and can fix the config.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Dragging the scrollbar** in the plan preview now works as expected.",
"breaking_change": false
},
{
"category": "fixes",
"description": "**Compaction** now correctly handles certain context-length errors from the inference API.",
"breaking_change": false
}
]

20
crates/codegen/xai-grok-shell/changelogs/0.2.118.md Normal file
View file

@ -0,0 +1,20 @@
# 0.2.118 — 2026-07-31
## Features
- **Sessions** can now be permanently deleted from the dashboard by pressing Ctrl+X twice on an idle row, or from the welcome list with d then y.
- **Keyboard shortcuts help** (Ctrl+.) now shows how to browse prompt history and search the conversation.
- **grok doctor** now warns when tmux is reducing colors and can fix the config.
## Bug Fixes
- **`/btw`** now retries on temporary model overload instead of failing immediately.
- **Session sharing** is temporarily disabled.
- **`[stop]`** / Ctrl+C during `/compact` now cancels instead of no-opping.
- **Automatic recaps** no longer appear twice after the same turn.
- **Background task wait timeout** descriptions and limits now match the client's actual configured ceiling.
- **Background tasks** no longer stay stuck as 'Running' in the tasks pane when they finish quickly.
- **Plan mode indicator** now disappears right after approving a plan instead of lingering.
- **Dragging the scrollbar** in the plan preview now works as expected.
- **Compaction** now correctly handles certain context-length errors from the inference API.

5
crates/codegen/xai-grok-shell/src/active_sessions.rs
View file

@ -30,11 +30,6 @@ pub fn register(session: ActiveSession) -> io::Result<()> {
register_in(&crate::util::grok_home::grok_home(), session)
}
/// Unregister a session (clean exit). No-op if not found.
pub fn unregister(session_id: &acp::SessionId) -> io::Result<()> {
unregister_in(&crate::util::grok_home::grok_home(), session_id)
}
/// Non-blocking unregister for signal handlers. Returns `Ok(false)` on
/// lock contention; the orphan is cleaned up by `collect_crashed` next launch.
pub fn try_unregister(session_id: &acp::SessionId) -> io::Result<bool> {

152
crates/codegen/xai-grok-shell/src/agent/app.rs
View file

@ -21,6 +21,7 @@ use crate::agent::init::{bootstrap, exit_on_config_error};
use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
use crate::agent::mvp_agent::MvpAgent;
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig, run_auth_flow};
use crate::leader::protocol::InternalMethod;
use crate::util::grok_home;
use dirs;
@ -206,27 +207,12 @@ fn spawn_agent_local(
handle_io
}
/// Build a newline-terminated JSON-RPC request line for an internal
/// `x.ai/...` extension method, for injection into the agent's inbound ACP
/// stream by the leader's own watcher tasks (config hot-reload, skills).
///
/// The wire method is written **`_`-prefixed** (`_x.ai/internal/...`):
/// `agent-client-protocol`'s inbound decoder routes a non-built-in method to
/// `ext_method` only when it carries the `_` extension prefix and rejects
/// bare custom methods with `-32601 method_not_found`. These injections were
/// historically sent un-prefixed, so every watcher-driven hot-reload
/// (models, skills, MCP servers) was silently rejected at decode — the
/// watcher-side "change detected" logs fired but the reload handlers never
/// ran. Keep `method` here as the un-prefixed name; the prefix is a wire
/// detail added in one place.
fn internal_reload_request_line(id: &str, method: &str, params: serde_json::Value) -> String {
let msg = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": format!("_{method}"),
"params": params,
});
format!("{}\n", msg)
fn internal_reload_request_line(
id: &str,
method: InternalMethod,
params: serde_json::Value,
) -> String {
crate::leader::protocol::internal_request_line(id, method, params)
}
/// Start a skills file watcher and wire it to inject `x.ai/internal/reload_skills`
@ -254,17 +240,17 @@ where
let (id, method) = match change {
crate::config::watcher::DiscoveryChange::Skills if !created_discovery_dir => {
info!("Skill directory changed on disk, reloading skills for all sessions");
("skills-reload", "x.ai/internal/reload_skills")
("skills-reload", InternalMethod::ReloadSkills)
}
crate::config::watcher::DiscoveryChange::Skills => {
info!("Discovery directory created on disk, reloading skills and workflows");
("skills-reload", "x.ai/internal/reload_skills")
("skills-reload", InternalMethod::ReloadSkills)
}
crate::config::watcher::DiscoveryChange::Workflows => {
info!(
"Workflow directory changed on disk, re-advertising commands for all sessions"
);
("workflows-reload", "x.ai/internal/reload_workflows")
("workflows-reload", InternalMethod::ReloadWorkflows)
}
};
let line = internal_reload_request_line(id, method, serde_json::json!({}));
@ -395,8 +381,6 @@ pub async fn run_stdio_agent(
// Restore managed policy right before bootstrap reads it (no stale window after prefetch).
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
// Fail-closed external-OTEL gate: suppress until settings resolve,
// opening now only for a pure env-API-key user (no remote policy).
apply_otel_config(&auth_manager, &agent_config.grok_com_config);
let handle_io = spawn_agent_local(
agent_config,
@ -426,24 +410,6 @@ pub async fn run_headless(
agent_config: &AgentConfig,
reauthenticate: bool,
memory_config: Option<crate::config::MemoryConfig>,
) -> anyhow::Result<()> {
run_headless_inner(agent_config, reauthenticate, false, memory_config).await
}
/// Run the headless agent without opening any browser windows.
/// If no cached credentials exist, returns an error instead of starting OAuth flow.
pub async fn run_headless_no_browser(
agent_config: &AgentConfig,
memory_config: Option<crate::config::MemoryConfig>,
) -> anyhow::Result<()> {
run_headless_inner(agent_config, false, true, memory_config).await
}
async fn run_headless_inner(
agent_config: &AgentConfig,
reauthenticate: bool,
no_browser: bool,
memory_config: Option<crate::config::MemoryConfig>,
) -> anyhow::Result<()> {
register_fs_watch_runtime();
xai_grok_telemetry::unified_log::set_version(xai_grok_version::VERSION);
@ -469,17 +435,7 @@ async fn run_headless_inner(
agent_config.mode = crate::agent::config::AgentMode::Headless;
let ctx = &agent_config.grok_com_config;
let (mut auth, did_browser_flow) = if no_browser {
// No-browser mode: only use cached credentials, skip OAuth flow
let auth_manager = agent_config.create_auth_manager();
match auth_manager.current() {
Some(auth) => (auth, false),
None if auth_manager.is_expired() => {
anyhow::bail!("Session expired. Please run 'grok login' to re-authenticate.")
}
None => anyhow::bail!("No cached credentials found. Run `grok login`."),
}
} else if reauthenticate {
let (mut auth, did_browser_flow) = if reauthenticate {
let auth_manager = Arc::new(AuthManager::new(&grok_home::grok_home(), ctx.clone()));
run_auth_flow(
&auth_manager,
@ -565,7 +521,7 @@ async fn run_headless_inner(
// Create first-connection callback for headless-specific behavior
let on_first_connect: Box<dyn FnOnce() + Send + 'static> = Box::new(move || {
if !did_browser_flow && !no_browser {
if !did_browser_flow {
// Print to stderr (not logger) so user sees it
eprintln!();
eprintln!(
@ -971,10 +927,8 @@ pub fn suppress_otel() {
}
/// Startup external-OTEL gate for an in-process (embedded) agent. Mirrors the
/// leader startup gate so the pager process is fail-closed by construction at the
/// agent boundary: suppress until the agent's first settings outcome, except a
/// pure env-API-key user (no session now, none minting) whose stream has no
/// remote policy and may emit immediately.
/// leader startup gate so the pager process is fail-closed by construction at
/// the agent boundary.
pub fn apply_otel_config(auth_manager: &AuthManager, grok_com_config: &GrokComConfig) {
suppress_otel();
// Session presence is disk-based (valid or expired), not refresh success: an
@ -982,10 +936,9 @@ pub fn apply_otel_config(auth_manager: &AuthManager, grok_com_config: &GrokComCo
// must keep the gate closed.
let has_session = auth_manager.current().is_some() || auth_manager.read_disk_auth().is_some();
if crate::agent::otel_gate::should_open_at_startup(crate::agent::otel_gate::StartupGate {
channel: crate::agent::otel_gate::resolved_policy_channel(),
has_session,
has_api_key_env: crate::agent::auth_method::has_xai_api_key_env(),
session_pending: crate::agent::otel_gate::is_session_pending(has_session, grok_com_config),
remote_fetch_enabled: crate::util::config::resolve_remote_fetch_enabled(),
}) {
crate::agent::otel_gate::open_at_startup();
}
@ -1251,13 +1204,19 @@ pub async fn run_leader(
.is_some();
let session_pending =
crate::agent::otel_gate::is_session_pending(has_session, &agent_config.grok_com_config);
let policy_channel =
crate::agent::otel_gate::policy_channel_for(&agent_config.endpoints.proxy_url());
if crate::agent::otel_gate::should_open_at_startup(crate::agent::otel_gate::StartupGate {
channel: policy_channel,
has_session,
has_api_key_env: crate::agent::auth_method::has_xai_api_key_env(),
session_pending,
remote_fetch_enabled: crate::util::config::resolve_remote_fetch_enabled(),
}) {
info!("Pure env-API-key leader; opening external-OTEL gate (no remote policy applies)");
info!(
channel = ?policy_channel,
has_session,
session_pending,
"Opening external-OTEL gate at startup: no fleet policy is pending for this leader"
);
crate::agent::otel_gate::open_at_startup();
}
@ -1676,7 +1635,7 @@ pub async fn run_leader(
models_manager_for_config.on_auth_changed().await;
let line = internal_reload_request_line(
"config-auth-reloaded",
"x.ai/internal/reload_all_mcp_servers",
InternalMethod::ReloadAllMcpServers,
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1688,7 +1647,7 @@ pub async fn run_leader(
auth_manager_for_config.clear_in_memory();
let line = internal_reload_request_line(
"config-auth-cleared",
"x.ai/internal/auth_cleared",
InternalMethod::AuthCleared,
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1707,7 +1666,7 @@ pub async fn run_leader(
info!("MCP server config change detected — reloading active sessions");
let line = internal_reload_request_line(
"config-reload-mcp",
"x.ai/internal/reload_all_mcp_servers",
InternalMethod::ReloadAllMcpServers,
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1730,7 +1689,7 @@ pub async fn run_leader(
);
let line = internal_reload_request_line(
"config-reload-project-mcp",
"x.ai/internal/reload_project_mcp_servers",
InternalMethod::ReloadProjectMcpServers,
serde_json::json!({ "cwd": cwd.to_string_lossy() }),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1745,7 +1704,7 @@ pub async fn run_leader(
info!("Model config change detected — reloading agent model list");
let line = internal_reload_request_line(
"config-reload-models",
"x.ai/internal/reload_models",
InternalMethod::ReloadModels,
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1771,7 +1730,7 @@ pub async fn run_leader(
info!("Models cache change detected — reloading agent model catalog");
let line = internal_reload_request_line(
"config-reload-models-cache",
"x.ai/internal/reload_models_cache",
InternalMethod::ReloadModelsCache,
serde_json::json!({}),
);
let mut tx = acp_tx_for_config.lock().await;
@ -1982,40 +1941,10 @@ mod tests {
.expect("x.ai OIDC session must be relay-eligible")
}
/// The external-OTEL gate opens at startup only for a pure env-API-key leader:
/// env key set, no session, no pending mint. Any session (resolved of any
/// credential type, or about to be minted) makes it wait for the fetch.
#[test]
fn otel_gate_opens_only_for_pure_env_api_key_leader() {
use crate::agent::otel_gate::{StartupGate, should_open_at_startup};
let opens = |has_session, has_api_key_env, session_pending| {
should_open_at_startup(StartupGate {
has_session,
has_api_key_env,
session_pending,
remote_fetch_enabled: true,
})
};
// (has_session, has_api_key_env, session_pending)
assert!(opens(false, true, false), "pure env API key → opens");
assert!(!opens(true, true, false), "any resolved session → waits");
assert!(!opens(true, false, false), "session, no env key → waits");
assert!(
!opens(false, true, true),
"pending mint → session coming, waits"
);
assert!(
!opens(false, false, false),
"no env key, no session → waits"
);
}
/// The embedded startup gate (every pager `--no-leader` / fallback path) must be
/// fail-closed by construction: a session user stays closed until the agent
/// resolves settings, even when an env API key is also present (the key must
/// not bypass the session's remote policy). The pure env-API-key open path
/// is covered by `otel_gate_opens_only_for_pure_env_api_key_leader`, since
/// `is_session_pending` is environment-dependent (true in a devbox/CI pod).
/// not bypass the session's remote policy).
#[test]
#[serial_test::serial]
fn embedded_otel_gate_keeps_a_session_user_fail_closed() {
@ -2035,6 +1964,7 @@ mod tests {
struct Restore {
key: Option<std::ffi::OsString>,
legacy: Option<std::ffi::OsString>,
proxy: Option<std::ffi::OsString>,
}
impl Drop for Restore {
fn drop(&mut self) {
@ -2042,14 +1972,17 @@ mod tests {
unsafe {
set_or_clear(XAI_API_KEY_ENV_VAR, self.key.take());
set_or_clear(LEGACY_XAI_API_KEY_ENV_VAR, self.legacy.take());
set_or_clear(PROXY_ENV_VAR, self.proxy.take());
}
mark_external_otel_settings_resolved();
}
}
const PROXY_ENV_VAR: &str = "GROK_CLI_CHAT_PROXY_BASE_URL";
let _restore = Restore {
key: std::env::var_os(XAI_API_KEY_ENV_VAR),
legacy: std::env::var_os(LEGACY_XAI_API_KEY_ENV_VAR),
proxy: std::env::var_os(PROXY_ENV_VAR),
};
let cfg = GrokComConfig::default();
@ -2057,6 +1990,7 @@ mod tests {
unsafe {
std::env::set_var(XAI_API_KEY_ENV_VAR, "test-key");
std::env::remove_var(LEGACY_XAI_API_KEY_ENV_VAR);
std::env::remove_var(PROXY_ENV_VAR);
}
let session = GrokAuth {
@ -2345,17 +2279,11 @@ mod tests {
cancel.cancel();
}
/// The watcher-injected internal reload requests must carry the ACP
/// wire-level `_` extension prefix. `agent-client-protocol`'s inbound
/// decoder routes non-built-in methods to `ext_method` only when
/// `_`-prefixed and rejects bare custom methods with `-32601`, so an
/// un-prefixed injection means every config-driven hot-reload silently
/// dies at decode (watcher logs fire, handlers never run).
#[test]
fn internal_reload_request_line_uses_wire_ext_prefix() {
fn internal_reload_request_line_carries_id_params_and_newline() {
let line = internal_reload_request_line(
"config-reload-models",
"x.ai/internal/reload_models",
InternalMethod::ReloadModels,
serde_json::json!({}),
);
assert!(line.ends_with('\n'), "must be a newline-terminated line");
@ -2371,7 +2299,7 @@ mod tests {
// Params must pass through verbatim (project-MCP reload carries cwd).
let line = internal_reload_request_line(
"config-reload-project-mcp",
"x.ai/internal/reload_project_mcp_servers",
InternalMethod::ReloadProjectMcpServers,
serde_json::json!({ "cwd": "/repo/x" }),
);
let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap();
@ -2379,7 +2307,7 @@ mod tests {
let line = internal_reload_request_line(
"config-auth-cleared",
"x.ai/internal/auth_cleared",
InternalMethod::AuthCleared,
serde_json::json!({}),
);
let msg: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap();

30
crates/codegen/xai-grok-shell/src/agent/auth_method.rs
View file

@ -33,7 +33,7 @@ pub const LEGACY_XAI_API_KEY_ENV_VAR: &str = "GROK_CODE_XAI_API_KEY";
///
/// Checks `XAI_API_KEY` first, then falls back to the legacy
/// `GROK_CODE_XAI_API_KEY` for backward compatibility.
pub fn read_xai_api_key_env() -> Result<String, std::env::VarError> {
pub(crate) fn read_xai_api_key_env() -> Result<String, std::env::VarError> {
std::env::var(XAI_API_KEY_ENV_VAR).or_else(|_| std::env::var(LEGACY_XAI_API_KEY_ENV_VAR))
}
@ -60,7 +60,7 @@ pub fn has_xai_api_key_env() -> bool {
/// `GROK_DISABLE_API_KEY_AUTH`) is the admin kill switch: when true the
/// method is never advertised, regardless of available credentials, so
/// `XAI_API_KEY` can't bypass a deployment's forced IdP login.
pub fn should_advertise_xai_api_key<'a, I>(disable_api_key_auth: bool, models: I) -> bool
pub(crate) fn should_advertise_xai_api_key<'a, I>(disable_api_key_auth: bool, models: I) -> bool
where
I: IntoIterator<Item = &'a ModelEntry>,
{
@ -311,7 +311,7 @@ impl AuthMethodKind {
}
/// `true` for session-based methods (cached_token, grok.com, oidc).
pub fn is_session_based(self) -> bool {
pub(crate) fn is_session_based(self) -> bool {
matches!(self, Self::CachedToken | Self::GrokCom | Self::Oidc)
}
@ -319,25 +319,17 @@ impl AuthMethodKind {
pub fn needs_interactive_login(self) -> bool {
matches!(self, Self::GrokCom | Self::Oidc)
}
pub fn auth_error_message(self) -> &'static str {
if self.is_session_based() {
AUTH_ERROR_SESSION_EXPIRED
} else {
AUTH_ERROR_API_KEY
}
}
}
/// `true` for session-based ACP methods (cached_token, grok.com, oidc).
pub fn is_session_based_method(method_id: &acp::AuthMethodId) -> bool {
pub(crate) fn is_session_based_method(method_id: &acp::AuthMethodId) -> bool {
AuthMethodKind::from_id(method_id).is_session_based()
}
/// Per-model BYOK status: whether the selected model carries its own
/// `[model.*]` `api_key`/`env_key`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ModelByok {
pub(crate) enum ModelByok {
/// Model has its own per-model key (not refreshable).
Byok,
/// Model has no per-model key (session auth governs).
@ -374,7 +366,7 @@ impl ModelByok {
/// where sending the session token cannot leak to a third-party BYOK
/// endpoint. A definite `NotByok` always refreshes (it only ever routes to
/// the session endpoint); a definite `Byok` never does.
pub fn session_token_auth_gate(
pub(crate) fn session_token_auth_gate(
is_session_based_method: bool,
model_byok: ModelByok,
endpoint_is_first_party: bool,
@ -401,7 +393,7 @@ pub const AUTH_ERROR_API_KEY: &str = "Authentication failed. Run `grok login`, s
/// Pinned `oidc`: **no** fallthrough to api_key — return `None` so the caller
/// fails auth. Pinned `api_key` should not reach this path (cached_token is
/// not advertised).
pub fn method_id_after_cached_token_unavailable(
pub(crate) fn method_id_after_cached_token_unavailable(
has_external_api_key: bool,
preferred_method: Option<PreferredAuthMethod>,
) -> Option<&'static str> {
@ -423,7 +415,7 @@ pub const PREFERRED_OIDC_UNAVAILABLE: &str =
"preferred_method=oidc but no session is available. Run `grok login` to authenticate.";
pub const XAI_API_KEY_METHOD_ID: &str = "xai.api_key";
pub fn xai_api_key_auth_method() -> acp::AuthMethod {
pub(crate) fn xai_api_key_auth_method() -> acp::AuthMethod {
acp::AuthMethod::Agent(
acp::AuthMethodAgent::new(
acp::AuthMethodId::new(XAI_API_KEY_METHOD_ID),
@ -436,7 +428,7 @@ pub fn xai_api_key_auth_method() -> acp::AuthMethod {
}
pub const CACHED_TOKEN_AUTH_METHOD_ID: &str = "cached_token";
pub fn cached_token_auth_method() -> acp::AuthMethod {
pub(crate) fn cached_token_auth_method() -> acp::AuthMethod {
acp::AuthMethod::Agent(
acp::AuthMethodAgent::new(
acp::AuthMethodId::new(CACHED_TOKEN_AUTH_METHOD_ID),
@ -449,7 +441,7 @@ pub fn cached_token_auth_method() -> acp::AuthMethod {
pub const GROK_COM_METHOD_ID: &str = "grok.com";
/// xAI OAuth2/OIDC auth. Method id `"grok.com"` kept for ACP wire-compat.
pub fn grok_com_auth_method(
pub(crate) fn grok_com_auth_method(
label: Option<&str>,
has_auth_provider_command: bool,
) -> acp::AuthMethod {
@ -469,7 +461,7 @@ pub fn grok_com_auth_method(
}
pub const OIDC_METHOD_ID: &str = "oidc";
pub fn oidc_auth_method(issuer: &str, label: Option<&str>) -> acp::AuthMethod {
pub(crate) fn oidc_auth_method(issuer: &str, label: Option<&str>) -> acp::AuthMethod {
let name = label
.map(|l| l.to_string())
.unwrap_or_else(|| format!("Single sign-on ({})", issuer));

6
crates/codegen/xai-grok-shell/src/agent/chat_modes.rs
View file

@ -33,7 +33,7 @@ struct CachedModes {
}
/// Thread-safe, cheaply-cloneable manager. Cloning bumps the inner `Arc`.
#[derive(Clone)]
pub struct ChatModesManager {
pub(crate) struct ChatModesManager {
inner: Arc<Inner>,
}
struct Inner {
@ -151,7 +151,7 @@ impl ChatModesManager {
}
/// Kick a background `/rest/modes` fill when auth is already present so
/// `--chat` initialize / first `session/new` hit a warm cache.
pub fn warm_in_background(&self) {
pub(crate) fn warm_in_background(&self) {
let Some(user_id) = self.current_user_id() else {
return;
};
@ -164,7 +164,7 @@ fn empty_state() -> acp::SessionModelState {
/// Maps grok.com modes → `SessionModelState`: keeps only `available` modes,
/// reconciles `current_model_id` (default → first available → empty, never
/// out-of-set), and stashes `badgeText`/`iconHint`/`tags` in `_meta`.
pub fn modes_to_model_state(resp: &ListModesResponse) -> acp::SessionModelState {
pub(crate) fn modes_to_model_state(resp: &ListModesResponse) -> acp::SessionModelState {
let available_models: Vec<acp::ModelInfo> = resp
.modes
.iter()

299
crates/codegen/xai-grok-shell/src/agent/config.rs
View file

@ -42,15 +42,13 @@ pub enum AgentMode {
/// Default agent type when the server or user config doesn't specify one.
pub const DEFAULT_AGENT_TYPE: &str = "grok-build-plan";
/// Serde default for `ModelInfo.agent_type` and `ModelEntryConfig.agent_type`.
pub fn default_agent_type() -> String {
pub(crate) fn default_agent_type() -> String {
DEFAULT_AGENT_TYPE.to_owned()
}
/// Default base URL for the cli chat proxy.
pub const CLI_CHAT_PROXY_BASE_URL_DEFAULT: &str = "https://cli-chat-proxy.grok.com/v1";
/// Default base URL for the public xAI API.
pub const XAI_API_BASE_URL_DEFAULT: &str = "https://api.x.ai/v1";
/// Default base URL for the asset server (profile images, etc.).
pub const ASSET_SERVER_URL_DEFAULT: &str = "https://assets.grok.com";
/// One or more environment variable names that may hold a model API key.
///
/// Serde `untagged`: accepts a string or an array in TOML/JSON.
@ -109,11 +107,11 @@ impl EnvKeys {
}
}
/// Resolve the first set, non-blank process env value among configured names.
pub fn resolve_value(&self) -> Option<String> {
pub(crate) fn resolve_value(&self) -> Option<String> {
self.resolve_value_with(|name| std::env::var(name).ok())
}
/// Testable resolve with an injected getenv.
pub fn resolve_value_with(
pub(crate) fn resolve_value_with(
&self,
mut getenv: impl FnMut(&str) -> Option<String>,
) -> Option<String> {
@ -239,10 +237,6 @@ pub struct EndpointsConfig {
/// Env: `OTEL_EXPORTER_OTLP_TIMEOUT`. Export HTTP timeout (ms).
#[serde(skip_serializing_if = "Option::is_none")]
pub otel_exporter_otlp_timeout: Option<u64>,
/// Base URL for the asset server (profile images, etc.).
/// Env: `GROK_ASSET_SERVER_URL`.
#[serde(default = "default_asset_server_url")]
pub asset_server_url: String,
/// Read by `load_management_api_key_sync()`. Declared for `serde_ignored`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub management_api_key: Option<String>,
@ -250,9 +244,6 @@ pub struct EndpointsConfig {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub gcs_service_account_key: Option<String>,
}
pub(crate) fn default_asset_server_url() -> String {
std::env::var("GROK_ASSET_SERVER_URL").unwrap_or_else(|_| ASSET_SERVER_URL_DEFAULT.to_owned())
}
/// A blank or whitespace-only override counts as unset. Single source of truth
/// for the "empty value = not configured" rule shared by the endpoint resolvers.
fn blank_as_unset(opt: &Option<String>) -> Option<String> {
@ -280,7 +271,7 @@ impl EndpointsConfig {
/// startup fetches use the configured (not public) endpoints. Only merges
/// layers — never derives one endpoint from another. Falls back to
/// `default()` on load failure.
pub fn from_effective_config() -> Self {
pub(crate) fn from_effective_config() -> Self {
match crate::config::load_effective_config() {
Ok(cfg) => Self::from_config_value(&cfg),
Err(_) => Self::default(),
@ -311,25 +302,25 @@ impl EndpointsConfig {
blank_as_unset(&self.cli_chat_proxy_base_url)
.unwrap_or_else(|| CLI_CHAT_PROXY_BASE_URL_DEFAULT.to_owned())
}
pub fn resolve_inference_base_url(&self) -> String {
pub(crate) fn resolve_inference_base_url(&self) -> String {
self.models_base_url
.clone()
.unwrap_or_else(|| self.proxy_url())
}
/// Feedback endpoint — an auxiliary service, so it defaults to the
/// cli-chat-proxy, never `xai_api_base_url`.
pub fn resolve_feedback_base_url(&self) -> String {
pub(crate) fn resolve_feedback_base_url(&self) -> String {
blank_as_unset(&self.feedback_base_url).unwrap_or_else(|| self.proxy_url())
}
/// Trace upload endpoint — an auxiliary service, so it defaults to the
/// cli-chat-proxy, never `xai_api_base_url`.
pub fn resolve_trace_upload_url(&self) -> String {
pub(crate) fn resolve_trace_upload_url(&self) -> String {
blank_as_unset(&self.trace_upload_url).unwrap_or_else(|| self.proxy_url())
}
/// Managed deployment-config URL (`grok setup`): explicit `managed_config_url`,
/// else `proxy_url` + `/deployment/config`. Never `xai_api_base_url`, so the
/// deployment key reaches the proxy, not the inference host.
pub fn resolve_managed_config_url(&self) -> String {
pub(crate) fn resolve_managed_config_url(&self) -> String {
blank_as_unset(&self.managed_config_url).unwrap_or_else(|| {
format!(
"{}/deployment/config",
@ -348,7 +339,7 @@ impl EndpointsConfig {
/// master switch IS set, the standard `OTEL_EXPORTER_OTLP_*` values are
/// completely ignored here so the internally-authed firehose never lands
/// at an external collector.
pub fn resolve_otlp_traces_endpoint(&self) -> String {
pub(crate) fn resolve_otlp_traces_endpoint(&self) -> String {
if let Some(full) = blank_as_unset(&self.grok_internal_otlp_traces_endpoint) {
return full.trim_end_matches('/').to_string();
}
@ -378,7 +369,7 @@ impl EndpointsConfig {
/// Extra headers for the INTERNAL export: `grok_internal_otlp_headers`
/// first; legacy fallback to `otel_exporter_otlp_headers` ONLY when the
/// external-OTEL master switch is unset (back-compat for existing users).
pub fn resolve_otlp_headers(&self) -> Vec<(String, String)> {
pub(crate) fn resolve_otlp_headers(&self) -> Vec<(String, String)> {
if let Some(headers) = blank_as_unset(&self.grok_internal_otlp_headers) {
return parse_otlp_header_list(&headers);
}
@ -399,7 +390,7 @@ impl EndpointsConfig {
/// CONTRACT: this flag is passed to the external OTEL stream's init, which
/// MUST refuse to activate when it is true — the same standard vars cannot
/// feed both pipelines (no-double-send invariant, enforced in code).
pub fn internal_otlp_consumed_standard_vars(&self) -> bool {
pub(crate) fn internal_otlp_consumed_standard_vars(&self) -> bool {
if self.external_otel_master_switch {
return false;
}
@ -412,7 +403,7 @@ impl EndpointsConfig {
/// Trace export enabled unless `OTEL_TRACES_EXPORTER=none`. Deliberately
/// still honored by the internal pipeline even with `GROK_EXTERNAL_OTEL`
/// set: disabling internal span export is the safe direction.
pub fn resolve_traces_export_enabled(&self) -> bool {
pub(crate) fn resolve_traces_export_enabled(&self) -> bool {
!matches!(
self.otel_traces_exporter.as_deref().map(str::trim),
Some("none")
@ -420,18 +411,18 @@ impl EndpointsConfig {
}
/// `OTEL_BSP_SCHEDULE_DELAY` / `OTEL_TRACES_EXPORT_INTERVAL` — tuning-only,
/// deliberately shared between the internal and external pipelines.
pub fn resolve_otlp_export_interval(&self) -> Option<std::time::Duration> {
pub(crate) fn resolve_otlp_export_interval(&self) -> Option<std::time::Duration> {
self.otel_traces_export_interval
.map(std::time::Duration::from_millis)
}
/// `OTEL_EXPORTER_OTLP_TIMEOUT` — tuning-only, deliberately shared between
/// the internal and external pipelines.
pub fn resolve_otlp_timeout(&self) -> Option<std::time::Duration> {
pub(crate) fn resolve_otlp_timeout(&self) -> Option<std::time::Duration> {
self.otel_exporter_otlp_timeout
.map(std::time::Duration::from_millis)
}
/// Resolve trace upload credentials: inline > file > `None` (ambient).
pub fn resolve_trace_credentials(&self) -> Option<String> {
pub(crate) fn resolve_trace_credentials(&self) -> Option<String> {
if let Some(ref inline) = self.trace_upload_credentials {
let trimmed = inline.trim();
if !trimmed.is_empty() {
@ -531,7 +522,7 @@ impl EndpointsConfig {
})
}
/// `models_list_url` > `{models_base_url}/models` > `{proxy_base_url}/models`.
pub fn resolve_models_list_url(&self) -> String {
pub(crate) fn resolve_models_list_url(&self) -> String {
if let Some(ref url) = self.models_list_url {
return url.clone();
}
@ -572,7 +563,6 @@ impl Default for EndpointsConfig {
.and_then(|s| s.parse().ok()),
otel_exporter_otlp_timeout: env_string("OTEL_EXPORTER_OTLP_TIMEOUT")
.and_then(|s| s.parse().ok()),
asset_server_url: default_asset_server_url(),
management_api_key: None,
gcs_service_account_key: None,
}
@ -613,7 +603,6 @@ pub struct Requirements {
pub trace_upload: Constrained<bool>,
pub feedback: Constrained<bool>,
pub lsp_tools: Constrained<bool>,
pub tool_search: Constrained<bool>,
pub web_fetch: Constrained<bool>,
pub ask_user_question: Constrained<bool>,
pub image_gen: Constrained<bool>,
@ -915,7 +904,7 @@ impl PluginsConfig {
/// enabling attacker-controlled hooks (e.g. SessionStart → RCE).
/// Native `.grok/config.toml` entries already present take precedence:
/// a name is only added if it isn't already in the opposite list.
pub fn merge_claude_enabled_plugins(&mut self, _cwd: Option<&std::path::Path>) {
pub(crate) fn merge_claude_enabled_plugins(&mut self, _cwd: Option<&std::path::Path>) {
if crate::claude_import::is_claude_import_marked_with_log("merge_claude_enabled_plugins") {
return;
}
@ -939,7 +928,9 @@ impl PluginsConfig {
}
}
/// Build a `DiscoveryConfig` from this plugins config.
pub fn to_discovery_config(&self) -> xai_grok_agent::plugins::discovery::DiscoveryConfig {
pub(crate) fn to_discovery_config(
&self,
) -> xai_grok_agent::plugins::discovery::DiscoveryConfig {
xai_grok_agent::plugins::discovery::DiscoveryConfig {
cli_plugin_dirs: self.cli_plugin_dirs.clone(),
config_paths: self.paths.iter().map(std::path::PathBuf::from).collect(),
@ -1117,12 +1108,6 @@ pub struct RelayConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub enabled: Option<bool>,
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct RemoteConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub secret: Option<String>,
}
/// `[hub]` section from config.toml.
///
/// Optional default Computer Hub URL for **workspace provider** exposure
@ -1178,7 +1163,7 @@ pub struct SandboxSettingsConfig {
pub auto_allow_bash: Option<bool>,
}
impl SandboxSettingsConfig {
pub fn from_effective_config() -> Self {
pub(crate) fn from_effective_config() -> Self {
crate::config::load_effective_config()
.ok()
.and_then(|v| v.get("sandbox")?.clone().try_into().ok())
@ -1197,7 +1182,7 @@ impl SandboxSettingsConfig {
.unwrap_or_else(|| Resolved::new("off".to_owned(), ConfigSource::Default))
}
/// Resolve auto_allow_bash: requirement > env > config > default (false).
pub fn resolve_auto_allow_bash(&self, requirement: Option<bool>) -> Resolved<bool> {
pub(crate) fn resolve_auto_allow_bash(&self, requirement: Option<bool>) -> Resolved<bool> {
BoolFlag::env("GROK_SANDBOX_AUTO_ALLOW_BASH")
.requirement(requirement)
.config(self.auto_allow_bash)
@ -1371,8 +1356,6 @@ pub struct Config {
pub auth_providers: IndexMap<String, crate::auth::AuthProviderConfig>,
#[serde(skip)]
pub model_providers: IndexMap<String, ModelProviderConfig>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub shortcuts: Option<toml::Value>,
/// Written by the client via `config_toml_edit`; absorbed so it isn't
/// flagged as an unrecognized key.
#[serde(default, skip_serializing_if = "Option::is_none")]
@ -1423,8 +1406,6 @@ pub struct Config {
pub harness: HarnessConfig,
#[serde(default, skip_serializing)]
pub relay: RelayConfig,
#[serde(default, skip_serializing)]
pub remote: RemoteConfig,
/// Computer Hub configuration (`[hub]` in config.toml).
#[serde(default, skip_serializing)]
pub hub: HubConfig,
@ -1657,7 +1638,7 @@ impl CliAgentOverrides {
/// the flags are authoritative, so they replace the agent's own fields.
/// Spawned subagents instead layer these on top of an author's definition —
/// see [`Self::apply_to_subagent_definition`].
pub fn apply_to_definition(&self, def: &mut xai_grok_agent::config::AgentDefinition) {
pub(crate) fn apply_to_definition(&self, def: &mut xai_grok_agent::config::AgentDefinition) {
if let Some(ref tools) = self.tools {
def.tools = tools.clone();
}
@ -1671,7 +1652,10 @@ impl CliAgentOverrides {
/// Subagent variant of [`Self::apply_to_definition`]: records the flags as
/// session-clamp state (see [`AgentDefinition::session_tools_allowlist`])
/// instead of overwriting the agent author's own fields.
pub fn apply_to_subagent_definition(&self, def: &mut xai_grok_agent::config::AgentDefinition) {
pub(crate) fn apply_to_subagent_definition(
&self,
def: &mut xai_grok_agent::config::AgentDefinition,
) {
def.session_tools_allowlist = self.tools.clone();
def.session_tools_denylist = self.disallowed_tools.clone();
if let Some(ref parent_mode) = self.permission_mode
@ -1681,7 +1665,7 @@ impl CliAgentOverrides {
resolve_subagent_permission_mode(def.permission_mode.clone(), parent_mode);
}
}
pub fn has_definition_overrides(&self) -> bool {
pub(crate) fn has_definition_overrides(&self) -> bool {
self.tools.is_some() || self.disallowed_tools.is_some() || self.permission_mode.is_some()
}
}
@ -1808,7 +1792,6 @@ impl Default for Config {
grok_com_config: GrokComConfig::default(),
auth_providers: IndexMap::new(),
model_providers: IndexMap::new(),
shortcuts: None,
hints: None,
ui: UiConfig::default(),
toolset: ShellToolsetConfig::default(),
@ -1827,7 +1810,6 @@ impl Default for Config {
models: ModelsConfig::default(),
harness: HarnessConfig::default(),
relay: RelayConfig::default(),
remote: RemoteConfig::default(),
hub: HubConfig::default(),
worktree_pool: WorktreePoolConfig::default(),
sandbox: SandboxSettingsConfig::default(),
@ -1967,7 +1949,7 @@ fn parse_auth_providers(
impl Config {
/// Reject invalid glob patterns in the model-filter lists at config load, so
/// a typo fails loudly instead of silently changing availability.
pub fn validate_model_filters(&self) -> Result<(), String> {
pub(crate) fn validate_model_filters(&self) -> Result<(), String> {
for (field, list) in [
("allowed_models", &self.models.allowed_models),
("disabled_models", &self.models.disabled_models),
@ -2170,7 +2152,7 @@ impl Config {
/// Must be called after `new_from_toml_cfg` on the **primary startup path**
/// before the config is handed to `MvpAgent`. Project definitions are overlaid
/// per cwd after that cwd's authoritative folder-trust resolve.
pub fn resolve_subagents(&mut self, cli_flag: bool, raw_config: &toml::Value) {
pub(crate) fn resolve_subagents(&mut self, cli_flag: bool, raw_config: &toml::Value) {
let sa = crate::config::SubagentsConfig::resolve(cli_flag, raw_config);
self.subagents_enabled = sa.enabled;
self.subagent_model_overrides = sa.models;
@ -2270,7 +2252,7 @@ impl Config {
/// the CLI flags already stored on this `Config`.
///
/// Integration test coverage: `tests/test_settings_refresh.rs`.
pub fn re_resolve_runtime_fields(&mut self, raw_config: &toml::Value) {
pub(crate) fn re_resolve_runtime_fields(&mut self, raw_config: &toml::Value) {
let remote_settings = self.remote_settings.clone();
let cli_web_search_model = self.web_search_model_override.clone();
let cli_session_summary_model = self.session_summary_model_override.clone();
@ -2321,25 +2303,25 @@ impl Config {
self.features.telemetry = Some(mode);
}
}
pub fn is_telemetry_enabled(&self) -> bool {
pub(crate) fn is_telemetry_enabled(&self) -> bool {
self.resolve_telemetry_mode().value.is_enabled()
}
pub fn is_trace_upload_enabled(&self) -> bool {
self.resolve_trace_upload().value
}
pub fn is_feedback_enabled(&self) -> bool {
pub(crate) fn is_feedback_enabled(&self) -> bool {
self.resolve_feedback().value
}
pub fn is_session_recap_enabled(&self) -> bool {
pub(crate) fn is_session_recap_enabled(&self) -> bool {
self.resolve_session_recap().value
}
pub fn is_voice_mode_enabled(&self) -> bool {
pub(crate) fn is_voice_mode_enabled(&self) -> bool {
self.resolve_voice_mode().value
}
/// Two-pass (prefire) compaction gate. Default OFF (opt-in) — enable via
/// remote settings `two_pass_compaction_enabled`, the `[features] two_pass_compaction`
/// config.toml key, or `GROK_TWO_PASS_COMPACTION` env.
pub fn is_two_pass_compaction_enabled(&self) -> bool {
pub(crate) fn is_two_pass_compaction_enabled(&self) -> bool {
self.resolve_two_pass_compaction().value
}
pub(crate) fn resolve_telemetry_mode(&self) -> Resolved<TelemetryMode> {
@ -2396,7 +2378,7 @@ impl Config {
)
}
/// K12 scoped resolve: fresh jemalloc fields + current gates (no remote rewrite).
pub fn resolve_jemalloc_heap_profile_from_partial(
pub(crate) fn resolve_jemalloc_heap_profile_from_partial(
&self,
jemalloc_enabled: Option<bool>,
jemalloc_thresholds: Option<&[u64]>,
@ -2501,7 +2483,7 @@ impl Config {
/// `[worktree.auto_gc]` TOML > remote `worktree_auto_gc` > defaults.
/// Platform age-expiry (non-Linux dead-only) is enforced inside
/// `xai_fast_worktree::maybe_auto_gc`, not here.
pub fn resolve_worktree_auto_gc(&self) -> xai_fast_worktree::ResolvedWorktreeAutoGc {
pub(crate) fn resolve_worktree_auto_gc(&self) -> xai_fast_worktree::ResolvedWorktreeAutoGc {
crate::util::config::resolve_worktree_auto_gc_from_settings(
Some(&self.worktree.auto_gc),
self.remote_settings
@ -2997,7 +2979,7 @@ impl Config {
.and_then(|r| r.compaction_detail.as_deref()),
)
}
pub fn resolve_cancel_rewind(&self) -> Resolved<bool> {
pub(crate) fn resolve_cancel_rewind(&self) -> Resolved<bool> {
let ff = self
.remote_settings
.as_ref()
@ -3014,74 +2996,12 @@ impl Config {
/// the default xAI OAuth2 fallback when no enterprise OIDC is configured.
///
/// Priority: `--oauth` > GROK_OAUTH_ENABLED env > default (true = OAuth).
pub fn resolve_grok_oauth(&self, cli_oidc: Option<bool>) -> Resolved<bool> {
pub(crate) fn resolve_grok_oauth(&self, cli_oidc: Option<bool>) -> Resolved<bool> {
BoolFlag::env("GROK_OAUTH_ENABLED")
.cli(cli_oidc)
.default(true)
.resolve()
}
/// Resolve whether to spawn the per-`Ready`-client transport
/// liveness pollers and the session-actor `StatusDispatcher`.
///
/// Thin delegate to the canonical
/// [`resolve_mcp_liveness_watchers`] free function, which unifies
/// the two previous implementations so they can't drift. CLI / managed / feature-flag inputs are
/// `None` here because the `Config` method only has visibility
/// into the embedded `Features` table; richer call sites (e.g.
/// the session-actor spawn path) go through
/// [`crate::util::config::resolve_mcp_liveness_watchers`] which
/// stacks all 7 layers.
pub fn resolve_mcp_liveness_watchers(&self) -> Resolved<bool> {
resolve_mcp_liveness_watchers(None, None, self.features.mcp_liveness_watchers, None, None)
}
/// Resolve whether the bounded stdio auto-restart task is allowed
/// to fire. Thin delegate to
/// [`resolve_mcp_auto_restart`]; mirrors
/// [`Self::resolve_mcp_liveness_watchers`]. The 7-step precedence
/// stack lives in the canonical free function. CLI / managed /
/// feature-flag inputs are `None` here because the `Config`
/// method only has visibility into the embedded `Features`
/// table; richer call sites go through
/// [`crate::util::config::resolve_mcp_auto_restart`] which stacks
/// all 7 layers.
pub fn resolve_mcp_auto_restart(&self) -> Resolved<bool> {
resolve_mcp_auto_restart(None, None, self.features.mcp_auto_restart, None, None)
}
/// Resolve whether the pager subscribes to the per-server
/// `x.ai/mcp/server_status` push.
///
/// Thin delegate to the canonical
/// [`resolve_mcp_push_server_status`] free function — mirrors the
/// `resolve_mcp_liveness_watchers` pattern so the two
/// implementations can't drift. CLI / managed / feature-flag
/// inputs are `None` here because the `Config` method only has
/// visibility into the embedded `Features` table; richer call
/// sites go through
/// [`crate::util::config::resolve_mcp_push_server_status`] which
/// stacks all 7 layers.
pub fn resolve_mcp_push_server_status(&self) -> Resolved<bool> {
resolve_mcp_push_server_status(None, None, self.features.mcp_push_server_status, None, None)
}
/// Resolve whether the leader's `ConfigFileWatcher` adds the two
/// narrow non-recursive watches for `<cwd>/` and `<cwd>/.grok/`.
///
/// Thin delegate to the canonical
/// [`resolve_mcp_recursive_config_watch`] free function — mirrors
/// the same delegation pattern. CLI / managed /
/// feature-flag inputs are `None` here because the `Config`
/// method only sees the embedded `Features` table; richer call
/// sites (notably the leader's watcher spawn path) go through
/// [`crate::util::config::resolve_mcp_recursive_config_watch`]
/// which stacks all 7 layers.
pub fn resolve_mcp_recursive_config_watch(&self) -> Resolved<bool> {
resolve_mcp_recursive_config_watch(
None,
None,
self.features.mcp_recursive_config_watch,
None,
None,
)
}
}
/// Canonical resolver for `mcp.liveness_watchers`. Stacks the full
/// 7-step `BoolFlag` precedence:
@ -3089,14 +3009,13 @@ impl Config {
/// `requirement > cli > env (GROK_MCP_LIVENESS_WATCHERS) > config >
/// managed > feature_flag > default (true)`.
///
/// Both `Config::resolve_mcp_liveness_watchers` and
/// `util::config::resolve_mcp_liveness_watchers` delegate here so the
/// `util::config::resolve_mcp_liveness_watchers` delegates here so the
/// precedence is single-sourced.
///
/// The default is `true` — it gates the watcher + dispatcher
/// default-on, with this flag existing primarily as a kill switch
/// during the rollout.
pub fn resolve_mcp_liveness_watchers(
pub(crate) fn resolve_mcp_liveness_watchers(
requirement: Option<bool>,
cli: Option<bool>,
config: Option<bool>,
@ -3119,13 +3038,12 @@ pub fn resolve_mcp_liveness_watchers(
/// managed > feature_flag > default (true)`.
///
/// Mirrors [`resolve_mcp_liveness_watchers`]. Both
/// `Config::resolve_mcp_auto_restart` and
/// `util::config::resolve_mcp_auto_restart` delegate here so the
/// `util::config::resolve_mcp_auto_restart` delegates here so the
/// precedence is single-sourced.
///
/// Recovery is on by default; opt out via `GROK_MCP_AUTO_RESTART=false`,
/// `[features] mcp_auto_restart`, or `requirements.toml`.
pub fn resolve_mcp_auto_restart(
pub(crate) fn resolve_mcp_auto_restart(
requirement: Option<bool>,
cli: Option<bool>,
config: Option<bool>,
@ -3148,8 +3066,7 @@ pub fn resolve_mcp_auto_restart(
/// `requirement > cli > env (GROK_MCP_PUSH_SERVER_STATUS) > config >
/// managed > feature_flag > default (true)`.
///
/// Both `Config::resolve_mcp_push_server_status` and
/// `util::config::resolve_mcp_push_server_status` delegate here so
/// `util::config::resolve_mcp_push_server_status` delegates here so
/// the precedence is single-sourced.
///
/// The default is `true` — the pager's subscription to
@ -3178,8 +3095,7 @@ pub fn resolve_mcp_push_server_status(
/// `requirement > cli > env (GROK_MCP_RECURSIVE_CONFIG_WATCH) >
/// config > managed > feature_flag > default (true)`.
///
/// Both `Config::resolve_mcp_recursive_config_watch` and
/// `util::config::resolve_mcp_recursive_config_watch` delegate here
/// `util::config::resolve_mcp_recursive_config_watch` delegates here
/// so the precedence is single-sourced.
///
/// The default is `true`. It enables the two narrow
@ -3195,7 +3111,7 @@ pub fn resolve_mcp_push_server_status(
/// are non-recursive (by design, to avoid blowing through
/// `fs.inotify.max_user_watches` on large repos). The flag name
/// follows the rollout-gate naming convention.
pub fn resolve_mcp_recursive_config_watch(
pub(crate) fn resolve_mcp_recursive_config_watch(
requirement: Option<bool>,
cli: Option<bool>,
config: Option<bool>,
@ -3226,7 +3142,7 @@ pub fn resolve_mcp_recursive_config_watch(
/// 4. process env via `enable_env` (either direction)
/// 5. merged config (user/managed defaults)
/// 6. `inherit`, then `default`
pub struct SyncBoolFlag {
pub(crate) struct SyncBoolFlag {
extract_toml: fn(&toml::Value) -> Option<bool>,
disable_env: Option<&'static str>,
enable_env: Option<fn() -> Option<bool>>,
@ -3245,13 +3161,13 @@ impl SyncBoolFlag {
}
/// Force-off env name (e.g. `"DISABLE_TELEMETRY"`). Truthy at this name
/// in `managed_settings.json` or process env disables the flag.
pub const fn disable_env(mut self, name: &'static str) -> Self {
pub(crate) const fn disable_env(mut self, name: &'static str) -> Self {
self.disable_env = Some(name);
self
}
/// Either-direction env resolver (typically `GROK_*`). Returns
/// `Some(enabled)` for an explicit signal, `None` to fall through.
pub const fn enable_env(mut self, resolver: fn() -> Option<bool>) -> Self {
pub(crate) const fn enable_env(mut self, resolver: fn() -> Option<bool>) -> Self {
self.enable_env = Some(resolver);
self
}
@ -3298,7 +3214,7 @@ impl SyncBoolFlag {
}
/// Sync slice of [`Config::resolve_telemetry_mode`] for use before the tokio
/// runtime (e.g. `init_sentry`). `true` only when explicitly off.
pub fn is_telemetry_disabled_sync() -> bool {
pub(crate) fn is_telemetry_disabled_sync() -> bool {
!SyncBoolFlag::new(telemetry_enabled_from_toml)
.disable_env("DISABLE_TELEMETRY")
.enable_env(grok_telemetry_env_enabled)
@ -3307,7 +3223,7 @@ pub fn is_telemetry_disabled_sync() -> bool {
/// Like [`is_telemetry_disabled_sync`] but only `true` when telemetry is
/// *explicitly* off; absence is not disabled (`.default(true)`) so remote-only
/// enablement still builds the OTLP exporter (the runtime gate then governs it).
pub fn is_telemetry_explicitly_disabled_sync() -> bool {
pub(crate) fn is_telemetry_explicitly_disabled_sync() -> bool {
!SyncBoolFlag::new(telemetry_enabled_from_toml)
.disable_env("DISABLE_TELEMETRY")
.enable_env(grok_telemetry_env_enabled)
@ -3475,7 +3391,9 @@ pub(crate) fn resolve_external_otel_config_with(
/// stream (fleet kill switch + content-gate lock). Tighten-only by
/// construction — there is no remote enable direction — so it is safe to
/// call on every settings refresh.
pub fn apply_external_otel_remote_policy(settings: Option<&crate::util::config::RemoteSettings>) {
pub(crate) fn apply_external_otel_remote_policy(
settings: Option<&crate::util::config::RemoteSettings>,
) {
let Some(settings) = settings else { return };
let policy = xai_grok_telemetry::external::ExternalOtelRemotePolicy {
force_disable: settings.external_otel_disabled.unwrap_or(false),
@ -3537,7 +3455,7 @@ fn managed_settings_env_flag(key: &str) -> Option<bool> {
}
/// Assemble the final model map. Priority (highest wins):
/// config.toml `[model.*]` > prefetched (remote) > hardcoded defaults.
pub fn resolve_model_list(
pub(crate) fn resolve_model_list(
cfg: &Config,
prefetched: Option<IndexMap<String, ModelEntry>>,
) -> IndexMap<String, ModelEntry> {
@ -3756,7 +3674,7 @@ fn apply_global_scalar_defaults(
}
}
/// Built-in default models. Prefer `resolve_model_list()`.
pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, ModelEntry> {
pub(crate) fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, ModelEntry> {
default_models(endpoints)
.into_iter()
.map(|(key, entry)| (key, ModelEntry::from_config_entry(&entry)))
@ -3764,7 +3682,7 @@ pub fn default_model_entries(endpoints: &EndpointsConfig) -> IndexMap<String, Mo
}
/// Resolve a model against the available model map.
/// Checks the map key (id) first, then falls back to a slug scan.
pub fn find_model_by_id<'a>(
pub(crate) fn find_model_by_id<'a>(
models: &'a IndexMap<String, ModelEntry>,
model_id: &str,
) -> Option<&'a ModelEntry> {
@ -3777,7 +3695,7 @@ pub fn find_model_by_id<'a>(
/// the session model the worker falls back to. Not-found-in-catalog ⇒ `false`
/// (conservative; also covers the Tier-2 synthetic proxy entry). Drives the
/// built-in `low` effort default.
pub fn effective_classifier_supports_re(
pub(crate) fn effective_classifier_supports_re(
aux_model: Option<&str>,
session_model: &str,
models: &IndexMap<String, ModelEntry>,
@ -4300,7 +4218,7 @@ impl ModelInfo {
}
}
/// Extract shared model metadata from a flat config entry.
pub fn from_config(entry: &ModelEntryConfig) -> Self {
pub(crate) fn from_config(entry: &ModelEntryConfig) -> Self {
ModelInfo {
user_selectable: true,
id: entry.id.clone(),
@ -4363,7 +4281,7 @@ impl ModelInfo {
/// | true | _ | hidden | hidden |
/// | false | true | visible | visible |
/// | false | false | visible | **hidden** |
pub fn visible_for_auth(&self, is_session_auth: bool) -> bool {
pub(crate) fn visible_for_auth(&self, is_session_auth: bool) -> bool {
!self.hidden && (is_session_auth || self.supported_in_api)
}
}
@ -4398,7 +4316,7 @@ impl ModelEntry {
pub fn info(&self) -> &ModelInfo {
&self.info
}
pub fn from_config_entry(entry: &ModelEntryConfig) -> Self {
pub(crate) fn from_config_entry(entry: &ModelEntryConfig) -> Self {
Self {
info: ModelInfo::from_config(entry),
api_key: entry.api_key.clone(),
@ -4426,7 +4344,7 @@ impl ModelEntry {
/// resolves to a non-empty value, or a named auth provider.
/// Probes `std::env::var` at call time: result is not stable across env
/// changes. Never executes a provider command.
pub fn has_own_credentials(&self) -> bool {
pub(crate) fn has_own_credentials(&self) -> bool {
self.own_credential().is_some() || self.auth_provider.is_some()
}
}
@ -4466,7 +4384,7 @@ impl Default for CodebaseIndexingSetting {
impl CodebaseIndexingSetting {
/// Should `path` be indexed? For `Enabled(true)`, always yes (caller gates on git-root).
/// For `Patterns`, path must match an include and not match any `!exclude`.
pub fn should_index(&self, path: &std::path::Path) -> bool {
pub(crate) fn should_index(&self, path: &std::path::Path) -> bool {
match self {
Self::Enabled(b) => *b,
Self::Patterns(patterns) => {
@ -4639,9 +4557,6 @@ pub struct Features {
pub managed_config: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub lsp_tools: Option<bool>,
/// MCP tool search/discovery. `None` = defer to remote settings / env / default (true).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_search: Option<bool>,
/// Web fetch tool. `None` = defer to remote settings / env / default (false).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub web_fetch: Option<bool>,
@ -4719,7 +4634,10 @@ pub struct Features {
/// dispatcher are spawned — useful as an emergency kill switch
/// for the rollout. `None` = defer to env / default (true).
///
/// Resolved via [`Config::resolve_mcp_liveness_watchers`].
/// Not read through this struct: the live resolver re-reads the
/// `[features]` key out-of-band from raw TOML in
/// `util::config::resolve::mcp`. Declared so `serde_ignored`
/// does not report it as an unrecognized key.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mcp_liveness_watchers: Option<bool>,
/// Bounded stdio auto-restart task.
@ -4733,7 +4651,10 @@ pub struct Features {
/// (recovery is on by default; set `false` here / via
/// `GROK_MCP_AUTO_RESTART` to opt out).
///
/// Resolved via [`Config::resolve_mcp_auto_restart`].
/// Not read through this struct: the live resolver re-reads the
/// `[features]` key out-of-band from raw TOML in
/// `util::config::resolve::mcp`. Declared so `serde_ignored`
/// does not report it as an unrecognized key.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mcp_auto_restart: Option<bool>,
/// Pager-side subscription to the `x.ai/mcp/server_status` push.
@ -4745,14 +4666,15 @@ pub struct Features {
/// back to the legacy `x.ai/mcp/tools_changed` debounced refetch
/// path. `None` = defer to env / default (true).
///
/// The pager-side gate
/// Not read through this struct. The pager-side gate
/// (`acp_handler::push_server_status_enabled`) uses an
/// **env-only** OnceLock cache via
/// [`crate::util::config::resolve_mcp_push_server_status(None, None, None)`].
/// That function consults `BoolFlag::env` and the default `true`
/// — it does NOT read this `Features` field. The shell-side
/// `Config::resolve_mcp_push_server_status` does delegate
/// through this field, but the pager never holds a `Config`.
/// [`crate::util::config::resolve_mcp_push_server_status(None, None, None)`],
/// which consults `BoolFlag::env` and the default `true`. The
/// `[features]` key itself is honoured out-of-band, re-read from
/// raw TOML in `util::config::resolve::mcp`. This field is
/// declared so `serde_ignored` does not report the key as
/// unrecognized.
///
/// Practical consequence: setting
/// `[features] mcp_push_server_status = false` in
@ -4785,13 +4707,16 @@ pub struct Features {
/// name and behavior; deferred to a follow-up to avoid widening
/// the config surface across requirements.toml / managed configs.
///
/// Resolved via [`Config::resolve_mcp_recursive_config_watch`].
/// Not read through this struct: the live resolver re-reads the
/// `[features]` key out-of-band from raw TOML in
/// `util::config::resolve::mcp`. Declared so `serde_ignored`
/// does not report it as an unrecognized key.
/// `None` = defer to env / default (true).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mcp_recursive_config_watch: Option<bool>,
}
/// Resolved credentials for a model session.
pub struct ResolvedCredentials {
pub(crate) struct ResolvedCredentials {
pub api_key: Option<String>,
pub base_url: String,
pub auth_type: xai_chat_state::AuthType,
@ -4811,7 +4736,10 @@ pub(crate) fn first_own_credential(
}
/// Priority: model api_key/env_key > cached auth-provider token > session
/// token > XAI_API_KEY.
pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> ResolvedCredentials {
pub(crate) fn resolve_credentials(
model: &ModelEntry,
session_key: Option<&str>,
) -> ResolvedCredentials {
let info = model.info();
let (api_key, base_url, auth_type) = if let Some(key) = model.own_credential() {
(
@ -4871,7 +4799,7 @@ pub fn resolve_credentials(model: &ModelEntry, session_key: Option<&str>) -> Res
/// `disable_api_key_auth` at the credential seam: swap a first-party xAI API
/// key for the IdP session (absent => request fails => forces login). BYOK
/// (non-xAI `base_url`) is untouched; no-op when the switch is off.
pub fn enforce_disable_api_key_auth(
pub(crate) fn enforce_disable_api_key_auth(
creds: &mut ResolvedCredentials,
disable_api_key_auth: bool,
session_key: Option<&str>,
@ -4909,7 +4837,7 @@ pub use xai_grok_telemetry::config::deployment_id_from_key;
/// Returns `None` (with a warning) if config loading, parsing, or model
/// lookup fails. `session_key` should only be passed when `auth_type` is
/// `SessionToken` — callers must guard this.
pub fn try_resolve_model_credentials(
pub(crate) fn try_resolve_model_credentials(
model_id: &str,
session_key: Option<&str>,
) -> Option<ResolvedCredentials> {
@ -4932,7 +4860,7 @@ pub fn try_resolve_model_credentials(
/// Per-model auth facts (BYOK status + auth scheme) from one effective-config
/// load, memoized by the session actor.
#[derive(Clone, Copy)]
pub struct ModelAuthFacts {
pub(crate) struct ModelAuthFacts {
pub byok: ModelByok,
pub auth_scheme: AuthScheme,
}
@ -4942,7 +4870,7 @@ pub struct ModelAuthFacts {
/// model absent from the catalog → `NotByok`. An empty `model_id` (no sampling
/// config yet) → `Unknown`, not `NotByok`, so the gate isn't activated for an
/// unidentified model.
pub fn resolve_model_auth_facts_and_provider(
pub(crate) fn resolve_model_auth_facts_and_provider(
model_id: &str,
) -> (ModelAuthFacts, Option<crate::auth::AuthProviderRef>) {
if model_id.is_empty() {
@ -5004,7 +4932,7 @@ fn with_resolved_model<T>(model_id: &str, f: impl FnOnce(ModelLookup) -> T) -> T
/// description, session summary, ...), resolved through the catalog so a
/// `[model.*]` override redirects it to its own endpoint, credentials, and
/// routing `model`. `None` → caller falls back to the active session's model.
pub fn resolve_aux_model_sampling_config(
pub(crate) fn resolve_aux_model_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
endpoints: &EndpointsConfig,
@ -5107,7 +5035,7 @@ pub fn resolve_aux_model_sampling_config(
/// The resolver gate is host-based, stricter than `session_token_auth_gate`:
/// a session-token deployment on a custom `models_base_url` loses aux-sampler
/// refresh, rather than risk the session bearer on a third-party endpoint.
pub fn stamp_session_local_sampler_fields(
pub(crate) fn stamp_session_local_sampler_fields(
cfg: &mut SamplerConfig,
active_session_config: &SamplerConfig,
client_identifier: Option<String>,
@ -5128,7 +5056,7 @@ pub fn stamp_session_local_sampler_fields(
/// On `None`, fall back to the active session model and full config (not
/// forcing `image_description_model` onto the agent endpoint, which 404s on
/// BYOK / non-proxy routes for internal slugs like `grok-build`).
pub fn finalize_image_describe_sampler_config(
pub(crate) fn finalize_image_describe_sampler_config(
resolved_aux: Option<SamplerConfig>,
active_session_config: &SamplerConfig,
client_identifier: Option<String>,
@ -5154,7 +5082,7 @@ pub fn finalize_image_describe_sampler_config(
/// Re-derive `auth_type` from the model's own credentials so BYOK env-key
/// models stay on `ApiKey` even when a session token is present. Falls
/// back to `fallback` when the model isn't in the on-disk catalog.
pub fn resolve_chat_state_auth_type(
pub(crate) fn resolve_chat_state_auth_type(
model_id: &str,
session_key: Option<&str>,
fallback: xai_chat_state::AuthType,
@ -5163,7 +5091,7 @@ pub fn resolve_chat_state_auth_type(
.map(|r| r.auth_type)
.unwrap_or(fallback)
}
pub fn sampling_config_for_model(
pub(crate) fn sampling_config_for_model(
model: &ModelEntry,
credentials: ResolvedCredentials,
alpha_test_key: Option<String>,
@ -5229,7 +5157,7 @@ pub fn sampling_config_for_model(
/// get an extra access header from the corresponding key argument.
///
/// Existing entries are never overwritten so callers can pre-set a value.
pub fn inject_url_derived_headers(
pub(crate) fn inject_url_derived_headers(
headers: &mut IndexMap<String, String>,
alpha_test_key: Option<&str>,
base_url: &str,
@ -5247,27 +5175,6 @@ pub fn inject_url_derived_headers(
}
let _ = (alpha_test_key, base_url);
}
pub fn resolve_model_to_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
alpha_test_key: Option<String>,
client_version: Option<String>,
fallback_entry: Option<ModelEntry>,
) -> Option<SamplerConfig> {
let entry = find_model_by_id(models, model_id)
.cloned()
.or(fallback_entry)?;
let credentials = resolve_credentials(&entry, session_key);
Some(sampling_config_for_model(
&entry,
credentials,
alpha_test_key,
client_version,
None,
None,
))
}
fn resolve_hidden_default_web_search_sampling_config(
model_id: &str,
session_key: Option<&str>,
@ -5326,7 +5233,7 @@ fn resolve_hidden_default_web_search_sampling_config(
None,
)
}
pub fn resolve_web_search_sampling_config(
pub(crate) fn resolve_web_search_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
@ -5372,7 +5279,7 @@ pub fn resolve_web_search_sampling_config(
}
resolved.map(crate::tools::config::web_search_sampling_config)
}
pub fn to_acp_model_info(
pub(crate) fn to_acp_model_info(
models: &IndexMap<String, ModelEntry>,
) -> IndexMap<acp::ModelId, acp::ModelInfo> {
models
@ -5449,7 +5356,7 @@ pub struct ModelSwitchIncompatibleAgentError {
}
impl ModelSwitchIncompatibleAgentError {
/// Build an `acp::Error` with this structured payload.
pub fn into_acp_error(self) -> acp::Error {
pub(crate) fn into_acp_error(self) -> acp::Error {
let message = format!(
"Cannot switch to model '{}': it requires agent '{}' but the active agent is '{}'. \
Start a new session to use this model.",
@ -10586,8 +10493,6 @@ agent_type = "cursor"
enabled = false
[relay]
enabled = false
[remote]
secret = "value"
[worktree_pool]
pool_size = 4
[managed_mcps]
@ -10597,8 +10502,6 @@ agent_type = "cursor"
[toolset.bash]
timeout_secs = 120
login_shell_capture = true
[shortcuts]
ctrl_k = "search"
[grok_com_config]
token_header = "test"
[auth.oidc]

18
crates/codegen/xai-grok-shell/src/agent/feedback_client.rs
View file

@ -42,7 +42,7 @@ const CLIENT_VERSION_HEADER: &str = "x-grok-client-version";
/// measurement (e.g. IDs, timestamps, client type).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionTurnDelta {
pub(crate) struct SessionTurnDelta {
// ── Context fields ──────────────────────────────────────────────────
/// **[context]** Which client surface produced this record (e.g. CLI, TUI).
pub client_type: ClientType,
@ -284,7 +284,7 @@ pub struct SessionTurnDelta {
/// Response from the turn-deltas endpoint.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionTurnDeltaResponse {
pub(crate) struct SessionTurnDeltaResponse {
pub session_id: String,
pub turn_number: i64,
pub recorded_at: chrono::DateTime<chrono::Utc>,
@ -296,7 +296,7 @@ pub struct SessionTurnDeltaResponse {
/// without fragile string matching on error messages.
#[derive(Debug, thiserror::Error)]
#[error("{context} failed with status {status}: {body}")]
pub struct FeedbackApiError {
pub(crate) struct FeedbackApiError {
pub status: reqwest::StatusCode,
pub context: &'static str,
pub body: String,
@ -304,12 +304,12 @@ pub struct FeedbackApiError {
impl FeedbackApiError {
/// Returns `true` if this is a 401 Unauthorized response.
pub fn is_unauthorized(&self) -> bool {
pub(crate) fn is_unauthorized(&self) -> bool {
self.status == reqwest::StatusCode::UNAUTHORIZED
}
/// Returns `true` if this is a 403 Forbidden response.
pub fn is_forbidden(&self) -> bool {
pub(crate) fn is_forbidden(&self) -> bool {
self.status == reqwest::StatusCode::FORBIDDEN
}
}
@ -384,7 +384,7 @@ impl FeedbackClient {
/// Whether this client can refresh credentials on a 401: requires both an
/// attached `AuthManager` and a wired `TokenRefresher` (e.g. static
/// deployment-key sessions return false).
pub fn has_token_refresher(&self) -> bool {
pub(crate) fn has_token_refresher(&self) -> bool {
self.credentials
.auth_manager()
.is_some_and(|am| am.has_refresher_attached())
@ -459,7 +459,7 @@ impl FeedbackClient {
}
}
pub async fn try_refresh_credentials(&self) -> bool {
pub(crate) async fn try_refresh_credentials(&self) -> bool {
let Some(manager) = self.credentials.auth_manager() else {
return false;
};
@ -672,7 +672,7 @@ impl FeedbackClient {
///
/// Called at the end of every turn to stream time-series data for
/// regression tracking and session analytics.
pub async fn send_turn_delta(
pub(crate) async fn send_turn_delta(
&self,
session_id: &str,
delta: &SessionTurnDelta,
@ -779,7 +779,7 @@ pub fn signals_to_update(
/// `loc_tracking_enabled` indicates whether the LOC attribution hunk tracker
/// was active for this session. When `false`, LOC delta fields are zeros
/// because the tracker was never spawned — not because no code changed.
pub fn snapshot_to_turn_delta(
pub(crate) fn snapshot_to_turn_delta(
snapshot: &crate::session::signals::TurnDeltaSnapshot,
client_type: ClientType,
request_id: Option<String>,

18
crates/codegen/xai-grok-shell/src/agent/folder_trust.rs
View file

@ -76,7 +76,7 @@ static DECISIONS: LazyLock<Mutex<HashMap<PathBuf, bool>>> =
/// their grants and denies, so a cache deny would be the one verdict nothing
/// (grant, store, prompt) could ever lift. Returns whether the folder had been
/// trusted. Symmetric with [`grant_folder_trust`].
pub fn revoke_folder_trust(cwd: &Path) -> bool {
pub(crate) fn revoke_folder_trust(cwd: &Path) -> bool {
// Local/dev builds are fully inert: nothing was trusted-via-gate to revoke,
// and recording `false` here would make `project_scope_allowed` wrongly gate.
if folder_trust_inert() {
@ -123,7 +123,7 @@ pub fn revoke_folder_trust(cwd: &Path) -> bool {
///
/// `DECISIONS` uses `parking_lot::Mutex` (no poisoning), so this gate cannot
/// fail OPEN on a poisoned lock.
pub fn project_scope_allowed(cwd: &Path) -> bool {
pub(crate) fn project_scope_allowed(cwd: &Path) -> bool {
let key = workspace_key(cwd);
// Copy out of the lock so the Some(false) reconcile can re-acquire it
// (parking_lot mutexes are not re-entrant).
@ -231,7 +231,11 @@ pub(crate) fn record_for_test(cwd: &Path, allowed: bool) {
/// whose cwd differs from the launch dir, `grok mcp doctor`) passes `false`, so
/// an unresolved interactive-but-untrusted workspace resolves **fail-closed**
/// (untrusted, no prompt) — only the launch dir is ever prompted for.
pub fn resolve_and_record(cwd: &Path, remote: Option<&RemoteSettings>, allow_prompt: bool) -> bool {
pub(crate) fn resolve_and_record(
cwd: &Path,
remote: Option<&RemoteSettings>,
allow_prompt: bool,
) -> bool {
// Local/dev builds are fully inert: project scope is always allowed, so skip
// the `trusted_folders.toml` read entirely.
if folder_trust_inert() {
@ -265,7 +269,7 @@ pub fn resolve_and_record(cwd: &Path, remote: Option<&RemoteSettings>, allow_pro
/// no-configs case (config added post-startup via git pull / agent write is
/// caught). The init-time dedup belongs to the one-shot caller (a `OnceCell` on
/// `MvpAgent`), NOT to any new shared-cache entry.
pub fn resolve_launch_dir_trust(cwd: &Path, remote: Option<&RemoteSettings>) -> bool {
pub(crate) fn resolve_launch_dir_trust(cwd: &Path, remote: Option<&RemoteSettings>) -> bool {
// Local/dev builds are fully inert: project scope is always allowed, skipping
// the store read + repo scan entirely.
if folder_trust_inert() {
@ -408,7 +412,7 @@ fn compute_from_inputs(
/// Edge case: a name declared in BOTH a project config and the global
/// `~/.grok/config.toml` is dropped when untrusted. This is intended — untrusted
/// project content must not influence the command spawned for a shared name.
pub fn project_scoped_mcp_names(cwd: &Path) -> HashSet<String> {
pub(crate) fn project_scoped_mcp_names(cwd: &Path) -> HashSet<String> {
let mut names = HashSet::new();
// `.grok/config.toml [mcp_servers]` entries tagged project (the loader's key
@ -450,7 +454,7 @@ pub fn project_scoped_mcp_names(cwd: &Path) -> HashSet<String> {
/// with a project-declared name is ALSO dropped when untrusted: an untrusted repo
/// must not influence the command spawned for that name (see
/// [`project_scoped_mcp_names`]). Servers with no project-name collision are kept.
pub fn filter_untrusted_project_mcp(
pub(crate) fn filter_untrusted_project_mcp(
cwd: &Path,
merged: Vec<acp::McpServer>,
) -> Vec<acp::McpServer> {
@ -482,7 +486,7 @@ pub fn filter_untrusted_project_mcp(
/// Thin `cwd`→verdict wrapper over the shared
/// [`xai_grok_tools::implementations::lsp::config::filter_project_lsp_when_untrusted`]
/// predicate, so Site B and the workspace build path share one gate.
pub fn filter_untrusted_project_lsp(
pub(crate) fn filter_untrusted_project_lsp(
cwd: &Path,
sourced: std::collections::BTreeMap<
String,

2
crates/codegen/xai-grok-shell/src/agent/handlers/session.rs
View file

@ -293,7 +293,7 @@ async fn handle_session_list(
/// Build sessions in exactly the requested directory. A page walk cannot reach
/// past `over_fetch(limit)` rows per cwd: the local lane re-scans that window
/// each page instead of seeking to the cursor.
pub async fn handle_list_sessions(
pub(crate) async fn handle_list_sessions(
agent: &MvpAgent,
args: acp::ListSessionsRequest,
) -> Result<acp::ListSessionsResponse, acp::Error> {

5
crates/codegen/xai-grok-shell/src/agent/init.rs
View file

@ -152,7 +152,8 @@ fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
// agent) passes through here, so diagnostic uploads always carry
// the version stamp and the resource ceilings in effect.
xai_grok_telemetry::unified_log::set_version(xai_grok_version::VERSION);
crate::util::limits::log_effective_limits();
let limits = crate::util::limits::ProcessLimits::read();
limits.log();
if !cfg!(test) {
// Clear a logged-out team's files before the background sync runs.
@ -196,6 +197,8 @@ fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
);
}
update_telemetry_config(cfg, auth_manager);
// Emitted here: the event needs the client update_telemetry_config installs.
xai_grok_telemetry::session_ctx::log_event(limits.into_event());
});
}

2
crates/codegen/xai-grok-shell/src/agent/mod.rs
View file

@ -3,7 +3,7 @@ pub mod app;
pub mod auth_method;
pub mod chat_modes;
pub mod config;
pub mod config_model_override_parse;
pub(crate) mod config_model_override_parse;
mod ext_parsers;
pub mod feedback_client;
pub mod folder_trust;

43
crates/codegen/xai-grok-shell/src/agent/models.rs
View file

@ -242,17 +242,17 @@ impl ModelsManager {
}
/// Subscribe to model-switch events. Returns a `watch::Receiver`
pub fn subscribe_model_switch(&self) -> tokio::sync::watch::Receiver<u64> {
pub(crate) fn subscribe_model_switch(&self) -> tokio::sync::watch::Receiver<u64> {
self.inner.model_switch_watch.subscribe()
}
/// Cheap snapshot of the current model-switch generation, for the laziness-check poll loop.
pub fn model_switch_generation(&self) -> u64 {
pub(crate) fn model_switch_generation(&self) -> u64 {
*self.inner.model_switch_watch.borrow()
}
/// Build from a resolved config. Falls back to bundled default if no models available.
pub fn from_config(
pub(crate) fn from_config(
cfg: &config::Config,
prefetched_models: Option<IndexMap<String, ModelEntry>>,
auth_manager: Arc<AuthManager>,
@ -307,7 +307,7 @@ impl ModelsManager {
}
/// Swap config, rebuild catalog, and reselect the model.
pub fn apply_config(&self, new_config: config::Config) {
pub(crate) fn apply_config(&self, new_config: config::Config) {
if let Err(e) = new_config.validate_model_filters() {
tracing::error!(error = %e, "ignoring config reload: invalid model filters");
return;
@ -370,7 +370,7 @@ impl ModelsManager {
}
/// [`Self::apply_config`] plus an unconditional default re-resolve, for remote-settings arrival while no session exists.
pub fn apply_config_reselecting_default(&self, new_config: config::Config) {
pub(crate) fn apply_config_reselecting_default(&self, new_config: config::Config) {
self.apply_config(new_config.clone());
self.reselect_default_model(&new_config);
self.notify_models_updated();
@ -421,7 +421,7 @@ impl ModelsManager {
self.inner.current_model_id.read().clone()
}
pub fn set_current_model_id(&self, id: acp::ModelId) {
pub(crate) fn set_current_model_id(&self, id: acp::ModelId) {
self.inner
.user_selected_model
.store(true, Ordering::Relaxed);
@ -443,7 +443,10 @@ impl ModelsManager {
}
/// Per-model Layer-3 LazinessDetector config for `model_id` (disabled default when absent).
pub fn laziness_detector_for(&self, model_id: &str) -> config::LazinessDetectorPerModelConfig {
pub(crate) fn laziness_detector_for(
&self,
model_id: &str,
) -> config::LazinessDetectorPerModelConfig {
self.inner
.catalog
.read()
@ -459,16 +462,16 @@ impl ModelsManager {
self.inner.catalog.write().models.insert(id.into(), entry);
}
pub fn current_reasoning_effort(&self) -> Option<ReasoningEffort> {
pub(crate) fn current_reasoning_effort(&self) -> Option<ReasoningEffort> {
*self.inner.current_reasoning_effort.read()
}
pub fn set_current_reasoning_effort(&self, effort: Option<ReasoningEffort>) {
pub(crate) fn set_current_reasoning_effort(&self, effort: Option<ReasoningEffort>) {
*self.inner.current_reasoning_effort.write() = effort;
}
/// Whether the given model supports reasoning effort according to the catalog.
pub fn model_supports_reasoning_effort(&self, model_id: &str) -> bool {
pub(crate) fn model_supports_reasoning_effort(&self, model_id: &str) -> bool {
self.inner
.catalog
.read()
@ -478,7 +481,7 @@ impl ModelsManager {
.unwrap_or(false)
}
pub fn model_default_reasoning_effort(&self, model_id: &str) -> Option<ReasoningEffort> {
pub(crate) fn model_default_reasoning_effort(&self, model_id: &str) -> Option<ReasoningEffort> {
self.inner
.catalog
.read()
@ -488,7 +491,7 @@ impl ModelsManager {
}
/// The raw catalog `reasoning_efforts` list for `model_id` with no fallback,
pub fn model_reasoning_efforts(&self, model_id: &str) -> Vec<ReasoningEffortOption> {
pub(crate) fn model_reasoning_efforts(&self, model_id: &str) -> Vec<ReasoningEffortOption> {
self.inner
.catalog
.read()
@ -498,7 +501,7 @@ impl ModelsManager {
.unwrap_or_default()
}
pub fn model_supports_backend_search(&self, model_id: &str) -> bool {
pub(crate) fn model_supports_backend_search(&self, model_id: &str) -> bool {
self.inner
.catalog
.read()
@ -508,7 +511,7 @@ impl ModelsManager {
.unwrap_or(false)
}
pub fn model_compactions_remaining(
pub(crate) fn model_compactions_remaining(
&self,
model_id: &str,
) -> Option<xai_grok_sampling_types::CompactionsRemaining> {
@ -520,7 +523,7 @@ impl ModelsManager {
.and_then(|e| e.info().compactions_remaining)
}
pub fn model_compaction_at_tokens(
pub(crate) fn model_compaction_at_tokens(
&self,
model_id: &str,
) -> Option<xai_grok_sampling_types::CompactionAtTokens> {
@ -533,7 +536,7 @@ impl ModelsManager {
}
/// Catalog opt-in to display the served-checkpoint fingerprint for this model.
pub fn model_show_model_fingerprint(&self, model_id: &str) -> bool {
pub(crate) fn model_show_model_fingerprint(&self, model_id: &str) -> bool {
let cat = self.inner.catalog.read();
let models = &cat.models;
resolve_catalog_key(models, &acp::ModelId::new(model_id))
@ -543,12 +546,12 @@ impl ModelsManager {
}
/// Resolved next-prompt-suggestion model pin from the live config
pub fn prompt_suggest_model_pin(&self) -> crate::config::PromptSuggestModelPin {
pub(crate) fn prompt_suggest_model_pin(&self) -> crate::config::PromptSuggestModelPin {
self.inner.cfg.read().prompt_suggest_model_pin.clone()
}
/// Whether `model_id` resolves in the current catalog — as a config key
pub fn model_in_catalog(&self, model_id: &str) -> bool {
pub(crate) fn model_in_catalog(&self, model_id: &str) -> bool {
let cat = self.inner.catalog.read();
let models = &cat.models;
resolve_catalog_key(models, &acp::ModelId::new(model_id)).is_some()
@ -571,7 +574,7 @@ impl ModelsManager {
}
/// Refresh models when the etag changes.
pub async fn refresh_if_new_etag(&self, etag: String) {
pub(crate) async fn refresh_if_new_etag(&self, etag: String) {
let same_etag = {
let cat = self.inner.catalog.read();
cat.etag.as_deref() == Some(etag.as_str())
@ -659,7 +662,7 @@ impl ModelsManager {
}
/// Hot-reload the catalog from `~/.grok/models_cache.json` after an external write (config-watcher detected).
pub fn reload_from_disk_cache(&self) {
pub(crate) fn reload_from_disk_cache(&self) {
self.reload_from_cache_manager(&self.inner.cache);
}

2
crates/codegen/xai-grok-shell/src/agent/models/resolution.rs
View file

@ -195,7 +195,7 @@ impl ModelGlobSet {
}
/// Single source of truth for the catalog. Applies, in order: `disabled_models`
pub fn resolve_model_catalog(
pub(crate) fn resolve_model_catalog(
cfg: &config::Config,
prefetched: Option<IndexMap<String, ModelEntry>>,
) -> IndexMap<String, ModelEntry> {

113
crates/codegen/xai-grok-shell/src/agent/mvp_agent/acp_agent.rs
View file

@ -3,6 +3,8 @@
//! [`acp::Agent`] trait implementation for [`MvpAgent`].
//! Co-located child of `mvp_agent` (`use super::*`).
use super::*;
use crate::auth::SilentRefresh;
use crate::leader::protocol::InternalMethod;
/// Which `x_search` sub-tools enforce the date cutoff, sent in `initialize`. `x_user_search` and
/// `x_thread_fetch` are `false`: they don't honor it yet.
#[derive(serde::Serialize)]
@ -318,45 +320,10 @@ impl acp::Agent for MvpAgent {
);
let mut has_cached_token = init_has_current;
if !init_has_current && init_is_expired {
let refreshed = matches!(
tokio::time::timeout(
crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
self.auth_manager.auth(),
)
.await,
Ok(Ok(_))
);
if refreshed {
tracing::debug!(
auth_type = ?self.auth_type(),
"auth: initialize() silent refresh succeeded",
);
xai_grok_telemetry::unified_log::info(
"auth: initialize() silent refresh succeeded",
None,
Some(
serde_json::json!({ "auth_type": format!("{:?}", self.auth_type()) }),
),
);
has_cached_token = true;
} else if !self.auth_manager.requires_manual_reauth() {
tracing::info!("auth: silent refresh failed transiently; advertising cached_token");
xai_grok_telemetry::unified_log::info(
"auth: initialize() silent refresh failed transiently, keeping cached_token",
None,
None,
);
has_cached_token = true;
} else {
tracing::warn!(
"auth: token expired, silent refresh failed - re-authentication required"
);
xai_grok_telemetry::unified_log::warn(
"auth: token expired, silent refresh failed - re-authentication required",
None,
None,
);
}
has_cached_token = match self.auth_manager.silent_refresh().await {
SilentRefresh::Renewed(_) => true,
SilentRefresh::Failed(remedy) => remedy.is_self_healing(),
};
}
let (
login_label,
@ -729,43 +696,19 @@ impl acp::Agent for MvpAgent {
}
}
}
if self.auth_manager.current().is_none()
&& self.auth_manager.is_expired()
{
let am = self.auth_manager.clone();
let refresh = tokio::spawn(async move { am.auth().await });
match tokio::time::timeout(
crate::http::STARTUP_AUTH_REFRESH_TIMEOUT,
refresh,
)
.await
{
Ok(Ok(Ok(_))) => {}
outcome => {
tracing::debug!(
timed_out = outcome.is_err(),
"auth: cached_token pre-check refresh did not produce a token (yet)"
)
let resolved = match self.auth_manager.current() {
Some(auth) => Some(auth),
None if !self.auth_manager.is_expired() => None,
None => {
match self.auth_manager.silent_refresh().await {
SilentRefresh::Renewed(auth) => Some(*auth),
SilentRefresh::Failed(remedy) if remedy.is_self_healing() => {
self.auth_manager.current_or_expired()
}
SilentRefresh::Failed(_) => None,
}
}
}
let resolved = self
.auth_manager
.current()
.or_else(|| {
if self.auth_manager.is_expired()
&& !self.auth_manager.requires_manual_reauth()
{
xai_grok_telemetry::unified_log::info(
"auth cached_token: accepting expired-but-refreshable session",
None,
None,
);
self.auth_manager.current_or_expired()
} else {
None
}
});
};
let Some(auth) = resolved else {
let message = if self.auth_manager.is_expired() {
"Session expired, re-authentication required"
@ -1005,8 +948,7 @@ impl acp::Agent for MvpAgent {
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let remote_settings = self.cfg.borrow().remote_settings.clone();
folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false);
let initial_client_mcp_servers = arguments.mcp_servers.clone();
let (mcp_servers, managed_mcp_expires_at) = self
let (initial_client_mcp_servers, mcp_servers, managed_mcp_expires_at) = self
.resolve_mcp_servers(arguments.mcp_servers, cwd.as_path())
.await;
let mcp_meta_config_map = parse_mcp_meta_config(arguments.meta.as_ref());
@ -1490,8 +1432,7 @@ impl acp::Agent for MvpAgent {
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let remote_settings = self.cfg.borrow().remote_settings.clone();
folder_trust::resolve_and_record(cwd.as_path(), remote_settings.as_ref(), false);
let initial_client_mcp_servers = client_mcp_servers.clone();
let (mcp_servers, managed_mcp_expires_at) = self
let (initial_client_mcp_servers, mcp_servers, managed_mcp_expires_at) = self
.resolve_mcp_servers(client_mcp_servers, cwd.as_path())
.await;
let mcp_meta_config_map = parse_mcp_meta_config(request_meta.as_ref());
@ -1874,7 +1815,7 @@ impl acp::Agent for MvpAgent {
)
.await?;
drop(spawn_timer);
} else if !mcp_servers.is_empty() {
} else {
tracing::info!(
session_id = %session_id.0,
mcp_server_count = mcp_servers.len(),
@ -1890,11 +1831,6 @@ impl acp::Agent for MvpAgent {
respond_to: tx,
});
}
} else {
tracing::info!(
session_id = %session_id.0,
"load_session: reconnecting to existing session (feedback manager already initialized)"
);
}
{
let init_meta = self
@ -3525,13 +3461,12 @@ impl acp::Agent for MvpAgent {
}
"x.ai/session/rename" | "x.ai/session/delete"
| "x.ai/session/update_mcp_servers" | "x.ai/session/fork"
| "x.ai/internal/reload_all_mcp_servers"
| "x.ai/internal/reload_project_mcp_servers" | "x.ai/internal/reload_skills"
| "x.ai/internal/reload_workflows" | "x.ai/internal/reload_models"
| "x.ai/internal/reload_models_cache" | "x.ai/internal/auth_cleared"
| "x.ai/plugins/reload" | "x.ai/commands/list" => {
crate::extensions::session_admin::handle(self, &args).await
}
m if InternalMethod::from_name(m).is_some() => {
crate::extensions::session_admin::handle(self, &args).await
}
"x.ai/session/repair" => crate::extensions::repair::handle(self, &args).await,
"x.ai/session/usage" => crate::extensions::usage::handle(self, &args).await,
"x.ai/memory/flush" | "x.ai/memory/rewrite" => {
@ -3950,7 +3885,7 @@ impl acp::Agent for MvpAgent {
"Permission state reset for matching sessions"
);
}
if args.method.as_ref() == "x.ai/internal/evict_sessions" {
if args.method.as_ref() == InternalMethod::EvictSessions.name() {
self.handle_evict_sessions(&args.params).await;
}
if args.method.as_ref() == "x.ai/toggle_plan_mode"

73
crates/codegen/xai-grok-shell/src/agent/mvp_agent/agent_ops.rs
View file

@ -208,7 +208,7 @@ impl MvpAgent {
self.cfg.borrow().managed_mcp_gateway_tools_enabled
&& self.has_managed_mcp_auth()
}
pub async fn get_managed_mcp_configs(
pub(crate) async fn get_managed_mcp_configs(
&self,
) -> Vec<crate::session::managed_mcp::ManagedMcpConfig> {
if !self.can_fetch_managed_mcps() {
@ -222,7 +222,7 @@ impl MvpAgent {
)
.await
}
pub async fn get_managed_mcp_gateway_tool_catalog(
pub(crate) async fn get_managed_mcp_gateway_tool_catalog(
&self,
) -> Option<crate::session::managed_mcp::GatewayToolCatalog> {
if !self.can_fetch_managed_mcp_gateway_tools() {
@ -244,7 +244,7 @@ impl MvpAgent {
)
.await
}
pub fn managed_mcp_cache(
pub(crate) fn managed_mcp_cache(
&self,
) -> &crate::session::managed_mcp::ManagedMcpStateHandle {
&self.managed_mcp_cache
@ -465,7 +465,7 @@ impl MvpAgent {
}
});
}
pub fn agent_mcp_state(
pub(crate) fn agent_mcp_state(
&self,
) -> std::sync::Arc<tokio::sync::Mutex<crate::session::mcp_servers::McpState>> {
self.agent_mcp_state.clone()
@ -496,23 +496,38 @@ impl MvpAgent {
"lazily populated plugin registry snapshot"
);
}
/// Fetch managed configs, merge with client servers, return merged list + earliest expiry.
/// Fetch managed configs, admit client servers under a post-await compat
/// snapshot, merge, and return `(admitted_seed, merged, earliest_expiry)`.
///
/// Compat is read **after** the managed-config await so admit + merge share
/// one snapshot; a settings reapply during the await cannot make the
/// retained seed and the spawned set disagree.
pub(super) async fn resolve_mcp_servers(
&self,
client_servers: Vec<acp::McpServer>,
cwd: &std::path::Path,
) -> (Vec<acp::McpServer>, Option<chrono::DateTime<chrono::Utc>>) {
) -> (
Vec<acp::McpServer>,
Vec<acp::McpServer>,
Option<chrono::DateTime<chrono::Utc>>,
) {
self.ensure_plugin_registry();
let managed = self.get_managed_mcp_configs().await;
let expires_at = managed.iter().filter_map(|c| c.token_expires_at).min();
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
let compat = self.cfg.borrow().compat_resolved;
let admitted = crate::session::managed_mcp::admit_client_mcp_servers(
client_servers,
cwd,
&compat,
);
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
admitted.clone(),
cwd,
&managed,
self.plugin_registry_handle.snapshot().as_deref(),
&self.cfg.borrow().compat_resolved,
&compat,
);
(merged, expires_at)
(admitted, merged, expires_at)
}
/// Set the memory configuration (called from TUI after config resolution).
pub fn set_memory_config(&mut self, config: crate::config::MemoryConfig) {
@ -524,7 +539,10 @@ impl MvpAgent {
///
/// Must be called right after construction: entries registered on the
/// constructor-created default instance are NOT migrated.
pub fn set_activity(&mut self, activity: crate::agent::activity::AgentActivity) {
pub(crate) fn set_activity(
&mut self,
activity: crate::agent::activity::AgentActivity,
) {
self.activity = activity;
}
/// Send [`SessionCommand::Shutdown`] to every live session actor and wait
@ -542,7 +560,7 @@ impl MvpAgent {
/// the watcher is constructed in `agent/app.rs`. In simple /
/// non-leader mode the channel is never wired and
/// `notify_session_cwd_for_watch` is a no-op.
pub fn set_config_watcher_path_tx(
pub(crate) fn set_config_watcher_path_tx(
&mut self,
tx: tokio::sync::mpsc::UnboundedSender<std::path::PathBuf>,
) {
@ -1627,7 +1645,11 @@ impl MvpAgent {
auth: &crate::auth::GrokAuth,
) -> Option<crate::util::config::RemoteSettings> {
let identity = auth.user_id.clone();
self.otel_gate.rearm_on_switch(&identity);
let channel = {
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
crate::agent::otel_gate::policy_channel_for(&proxy_url)
};
self.otel_gate.rearm_on_switch(&identity, channel);
let outcome = self.fetch_settings_self_healing_401(auth).await;
let live = self.auth_manager.current_or_expired().map(|a| a.user_id);
self.otel_gate.resolve(&identity, outcome, live.as_deref())
@ -2816,12 +2838,12 @@ impl MvpAgent {
self.storage_mode.get()
}
/// Returns the background copy context for managing background file copy tasks.
pub fn background_copy_context(&self) -> BackgroundCopyContext {
pub(crate) fn background_copy_context(&self) -> BackgroundCopyContext {
self.background_copy_context.clone()
}
/// Move a foreground bash command to background.
/// Routes through the session's tool bridge to unblock the agent loop.
pub async fn background_foreground_command(
pub(crate) async fn background_foreground_command(
&self,
session_id: &str,
tool_call_id: &str,
@ -2835,7 +2857,7 @@ impl MvpAgent {
}
/// Kill a background task by task_id.
/// Routes through the session's tool bridge to the TerminalBackend.
pub async fn kill_background_task(
pub(crate) async fn kill_background_task(
&self,
session_id: &str,
task_id: &str,
@ -2847,7 +2869,7 @@ impl MvpAgent {
Err("session not found".to_string())
}
}
pub async fn delete_scheduled_task(
pub(crate) async fn delete_scheduled_task(
&self,
session_id: &str,
task_id: &str,
@ -2862,7 +2884,7 @@ impl MvpAgent {
/// Cancel a subagent by id, returning a typed outcome that backs the pager's
/// `x.ai/subagent/cancel`. Active/pending → cancelled (a finish follows);
/// already-finished → its terminal status; unknown id → `NotFound`.
pub async fn cancel_subagent(
pub(crate) async fn cancel_subagent(
&self,
subagent_id: &str,
) -> xai_grok_tools::implementations::grok_build::task::types::SubagentCancelOutcome {
@ -3063,13 +3085,16 @@ impl MvpAgent {
}
/// Get a session's cwd by session_id.
/// Returns None if the session is not found.
pub fn get_session_cwd(&self, session_id: &acp::SessionId) -> Option<PathBuf> {
pub(crate) fn get_session_cwd(
&self,
session_id: &acp::SessionId,
) -> Option<PathBuf> {
let sessions = self.sessions.borrow();
sessions.get(session_id).map(|handle| PathBuf::from(&handle.info.cwd))
}
/// Get a session handle by session_id.
/// Returns None if the session is not found.
pub fn get_session_handle(
pub(crate) fn get_session_handle(
&self,
session_id: &acp::SessionId,
) -> Option<crate::session::SessionHandle> {
@ -3077,7 +3102,7 @@ impl MvpAgent {
sessions.get(session_id).cloned()
}
/// Get hooks list for a session (for `x.ai/hooks/list` extension).
pub async fn list_hooks(
pub(crate) async fn list_hooks(
&self,
session_id: &acp::SessionId,
) -> Option<xai_hooks_plugins_types::HooksListResponse> {
@ -3085,7 +3110,7 @@ impl MvpAgent {
handle.get_hooks_list().await
}
/// Execute a hooks management action (for `x.ai/hooks/action`).
pub async fn execute_hooks_action(
pub(crate) async fn execute_hooks_action(
&self,
session_id: &acp::SessionId,
action: xai_hooks_plugins_types::HooksAction,
@ -3101,7 +3126,7 @@ impl MvpAgent {
handle.execute_hooks_action(action).await
}
/// Execute a plugins management action (for `x.ai/plugins/action`).
pub async fn execute_plugins_action(
pub(crate) async fn execute_plugins_action(
&self,
session_id: &acp::SessionId,
action: xai_hooks_plugins_types::PluginsAction,
@ -3119,7 +3144,7 @@ impl MvpAgent {
outcome
}
/// Get a snapshot of the shared plugin registry (for `x.ai/plugins/list`).
pub fn plugin_registry_snapshot(
pub(crate) fn plugin_registry_snapshot(
&self,
) -> Option<std::sync::Arc<xai_grok_agent::plugins::PluginRegistry>> {
self.plugin_registry_handle.snapshot()
@ -3127,7 +3152,7 @@ impl MvpAgent {
/// Run content search at agent level.
/// This allows content search to work with just a cwd, without requiring a session.
/// Returns an upload method, or `None` when trace uploads are disabled.
pub async fn trace_upload_config(
pub(crate) async fn trace_upload_config(
&self,
) -> Option<crate::session::repo_changes::UploadMethod> {
let (method, _reason) = self.trace_upload_config_with_reason().await;

15
crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs
View file

@ -128,7 +128,7 @@ impl MvpAgent {
/// only the **last** client to call `initialize()`.
///
/// Falls back to global agent state when no session_id is given.
pub fn code_nav_eligibility_for_request(
pub(crate) fn code_nav_eligibility_for_request(
&self,
session_id: Option<&acp::SessionId>,
cwd: &std::path::Path,
@ -155,17 +155,6 @@ impl MvpAgent {
self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
}
/// Check eligibility using the stored initialize_request context.
///
/// **Not safe in leader mode** — reads the last `initialize()` call's
/// client_type and capability. Prefer [`code_nav_eligibility_for_request`]
/// when a session_id is available.
pub fn code_nav_eligibility(&self, cwd: &std::path::Path) -> Result<(), CodeNavEligibility> {
let client_type = *self.client_type.borrow();
let code_nav_enabled = self.code_nav_enabled.get();
self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
}
/// Resolve and get-or-create the codebase index for `cwd`, applying config
/// and git-root eligibility checks.
///
@ -251,7 +240,7 @@ impl MvpAgent {
/// Get an existing codebase index for the given cwd.
/// Returns None if no index exists for this cwd.
pub fn get_codebase_index(
pub(crate) fn get_codebase_index(
&self,
cwd: &std::path::Path,
) -> Option<std::sync::Arc<xai_codebase_graph::IndexManagerHandle>> {

7
crates/codegen/xai-grok-shell/src/agent/mvp_agent/heap_profile.rs
View file

@ -69,6 +69,13 @@ impl MvpAgent {
};
tokio::time::sleep(poll_interval).await;
let reported = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
agent_ref.get().report_resource_usage_if_due();
}));
if reported.is_err() {
tracing::error!("resource telemetry: report tick panicked; continuing");
}
let enabled = agent_ref
.get()
.heap_profile_monitor

164
crates/codegen/xai-grok-shell/src/agent/mvp_agent/mod.rs
View file

@ -660,7 +660,7 @@ fn announcements_refresh_interval() -> std::time::Duration {
/// gates fails. Used in `x.ai/code/status` responses and to generate
/// clear error messages on code-nav requests from ineligible clients.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodeNavEligibility {
pub(crate) enum CodeNavEligibility {
/// Client type is not web (web-only for initial rollout).
ClientNotWeb,
/// Client did not advertise `x.ai/codeNavigation.enabled`.
@ -1342,6 +1342,7 @@ impl Drop for SessionLoadGuard<'_> {
mod code_nav;
mod folder_trust_prompt;
mod heap_profile;
mod resource_telemetry;
mod session_registry;
mod session_lifecycle;
mod subagent_coordinator;
@ -1384,150 +1385,6 @@ pub(crate) struct OrphanedTask {
cwd: String,
}
impl MvpAgent {
/// Forward one raw JSONL replay line and collect its completion receiver.
///
/// Dispatches by on-disk method name:
/// - ACP updates (`"session/update"`) → typed `SessionNotification` for correct
/// TUI dispatch (direct dispatch preserves Rust types, not method strings).
/// - xAI updates (`"_x.ai/session/update"`) → `ExtNotification`.
///
/// When `mark_replay` is true, the notification is tagged with
/// `_meta.isReplay: true` so the client knows it's historical data.
/// Cursor-based reconnects set this to false for events after the cursor
/// so the client processes them as live updates.
fn forward_raw_replay_line(
&self,
line: &str,
persist_data: Option<&serde_json::Value>,
target_client_id: Option<&serde_json::Value>,
completions: &mut Vec<
tokio::sync::oneshot::Receiver<xai_acp_lib::AcpResult<()>>,
>,
mark_replay: bool,
pending_tool_calls: &mut std::collections::HashMap<
acp::ToolCallId,
acp::ToolCall,
>,
) {
use crate::session::storage::RawLinePeek;
let env = match serde_json::from_str::<RawLinePeek<'_>>(line) {
Ok(e) => e,
Err(e) => {
tracing::debug!(?e, "replay: skipping unparseable JSONL line");
return;
}
};
let method = env.method.unwrap_or("session/update");
let Some(raw_params) = env.params else {
tracing::debug!("replay: skipping JSONL line with no params");
return;
};
let is_xai = method == "_x.ai/session/update";
if is_xai {
if target_client_id.is_none() && !mark_replay {
if let Ok(owned) = serde_json::value::RawValue::from_string(
raw_params.get().to_owned(),
) {
completions
.push(
self
.gateway
.forward_with_completion(
acp::ExtNotification::new(
"x.ai/session/update",
std::sync::Arc::from(owned),
),
),
);
}
} else {
let Ok(mut params) = serde_json::from_str::<
serde_json::Value,
>(raw_params.get()) else {
tracing::debug!("replay: skipping xAI update with unparseable params");
return;
};
if let Some(obj) = params.as_object_mut() {
let meta = obj
.entry("_meta")
.or_insert_with(|| serde_json::json!({}));
if let Some(m) = meta.as_object_mut() {
if mark_replay {
m.insert("isReplay".to_string(), serde_json::json!(true));
}
if let Some(pd) = persist_data {
m.insert("x.ai/persist".to_string(), pd.clone());
}
if let Some(tid) = target_client_id {
m.insert("x.ai/leaderClientId".to_string(), tid.clone());
}
}
}
if let Ok(raw_val) = serde_json::value::to_raw_value(&params) {
completions
.push(
self
.gateway
.forward_with_completion(
acp::ExtNotification::new(
"x.ai/session/update",
std::sync::Arc::from(raw_val),
),
),
);
}
}
} else {
let Ok(mut notification) = serde_json::from_str::<
acp::SessionNotification,
>(raw_params.get()) else {
tracing::debug!("replay: skipping ACP update with unparseable params");
return;
};
match &mut notification.update {
acp::SessionUpdate::ToolCall(tc) => {
let is_pre_completed = matches!(
tc.status,
acp::ToolCallStatus::Completed | acp::ToolCallStatus::Failed
);
if is_pre_completed {} else {
pending_tool_calls.insert(tc.tool_call_id.clone(), tc.clone());
return;
}
}
acp::SessionUpdate::ToolCallUpdate(u) => {
match u.fields.status {
Some(acp::ToolCallStatus::Completed)
| Some(acp::ToolCallStatus::Failed) => {
if let Some(mut base) = pending_tool_calls
.remove(&u.tool_call_id)
{
base.update(std::mem::take(&mut u.fields));
notification.update = acp::SessionUpdate::ToolCall(base);
}
}
None => {
if let Some(base) = pending_tool_calls
.get_mut(&u.tool_call_id)
{
base.update(std::mem::take(&mut u.fields));
}
return;
}
_ => return,
}
}
_ => {}
}
if mark_replay {
mark_as_replay(&mut notification.meta, persist_data);
}
if let Some(tid) = target_client_id {
stamp_meta_value(&mut notification.meta, "x.ai/leaderClientId", tid);
}
completions.push(self.gateway.forward_with_completion(notification));
}
}
/// Replay updates from disk and drain completions.
/// Returns `(initial_total_tokens, end_offset)`.
pub(super) async fn replay_session_updates(
@ -1774,8 +1631,9 @@ impl MvpAgent {
owner_session_id: None,
description: None,
is_backgrounded: true,
output_total_bytes: 0,
};
let notification = crate::extensions::notification::SessionNotification {
let mut notification = crate::extensions::notification::SessionNotification {
session_id: session_id.clone(),
update: crate::extensions::notification::SessionUpdate::TaskCompleted {
task_snapshot: snapshot,
@ -1783,9 +1641,9 @@ impl MvpAgent {
},
meta: None,
};
if let Ok(params) = serde_json::to_value(&notification)
.and_then(|v| serde_json::value::to_raw_value(&v))
{
if let Some(params) = crate::tools::task_completed_frame::encode(
&mut notification,
) {
completions
.push(
self
@ -1793,7 +1651,7 @@ impl MvpAgent {
.forward_with_completion(
acp::ExtNotification::new(
"x.ai/task_completed",
params.into(),
params.into_inner().into(),
),
),
);
@ -1930,7 +1788,8 @@ impl MvpAgent {
),
);
let remote_was_absent = self.cfg.borrow().remote_settings.is_none();
if let Some(auth) = self.auth_manager.current()
if crate::util::config::resolve_remote_fetch_enabled()
&& let Some(auth) = self.auth_manager.current()
&& let Some(settings) = self.fetch_settings_resolving_gate(&auth).await
{
self.install_remote_settings(settings);
@ -2708,6 +2567,9 @@ pub(crate) fn settings_allow_access(
) -> bool {
!matches!(rs.and_then(|s| s.allow_access), Some(false))
}
mod replay;
#[cfg(test)]
mod replay_tests;
#[cfg(test)]
mod tests;
#[cfg(test)]

186
crates/codegen/xai-grok-shell/src/agent/mvp_agent/replay.rs Normal file
View file

@ -0,0 +1,186 @@
//! Forwards recorded session updates back to a loading client, fitting
//! completion records written before the size limit existed.
use agent_client_protocol as acp;
use super::{MvpAgent, mark_as_replay, stamp_meta_value};
impl MvpAgent {
/// Records written before completions were bounded can still be too long
/// for a client to read. `None` drops one that cannot be shrunk, which
/// costs a completion event but keeps the connection.
fn fitted_replay_params(
params: Box<serde_json::value::RawValue>,
) -> Option<Box<serde_json::value::RawValue>> {
use crate::tools::task_completed_frame::{Refit, refit_recorded};
match refit_recorded(&params) {
Refit::Unchanged => Some(params),
Refit::Fitted(fitted) => Some(fitted.into_inner()),
Refit::Unfittable => {
tracing::warn!(
bytes = params.get().len(),
"replay: dropping a completion too long to send"
);
None
}
}
}
/// Forward one raw JSONL replay line and collect its completion receiver.
///
/// Dispatches by on-disk method name:
/// - ACP updates (`"session/update"`) → typed `SessionNotification` for correct
/// TUI dispatch (direct dispatch preserves Rust types, not method strings).
/// - xAI updates (`"_x.ai/session/update"`) → `ExtNotification`.
///
/// When `mark_replay` is true, the notification is tagged with
/// `_meta.isReplay: true` so the client knows it's historical data.
/// Cursor-based reconnects set this to false for events after the cursor
/// so the client processes them as live updates.
pub(super) fn forward_raw_replay_line(
&self,
line: &str,
persist_data: Option<&serde_json::Value>,
target_client_id: Option<&serde_json::Value>,
completions: &mut Vec<tokio::sync::oneshot::Receiver<xai_acp_lib::AcpResult<()>>>,
mark_replay: bool,
pending_tool_calls: &mut std::collections::HashMap<acp::ToolCallId, acp::ToolCall>,
) {
use crate::session::storage::RawLinePeek;
let env = match serde_json::from_str::<RawLinePeek<'_>>(line) {
Ok(e) => e,
Err(e) => {
tracing::debug!(?e, "replay: skipping unparseable JSONL line");
return;
}
};
// updates.jsonl only persists `_x.ai/session/update` and `session/update`.
// Unknown methods fall through to the ACP parse below and are dropped on error.
let method = env.method.unwrap_or("session/update");
let Some(raw_params) = env.params else {
tracing::debug!("replay: skipping JSONL line with no params");
return;
};
let is_xai = method == "_x.ai/session/update";
if is_xai {
// The fast-path forwards raw params with no `_meta` round-trip, so it
// can stamp nothing. When a `target_client_id` is present we MUST take
// the injection path instead, otherwise the replay would lose the
// target and the leader would broadcast it to every subscriber.
if target_client_id.is_none() && !mark_replay {
// Fast-path: forward raw params without Value round-trip.
if let Ok(owned) =
serde_json::value::RawValue::from_string(raw_params.get().to_owned())
&& let Some(owned) = Self::fitted_replay_params(owned)
{
completions.push(self.gateway.forward_with_completion(
acp::ExtNotification::new(
"x.ai/session/update",
std::sync::Arc::from(owned),
),
));
}
} else {
// Inject _meta — requires parse + re-serialize.
let Ok(mut params) = serde_json::from_str::<serde_json::Value>(raw_params.get())
else {
tracing::debug!("replay: skipping xAI update with unparseable params");
return;
};
if let Some(obj) = params.as_object_mut() {
let meta = obj.entry("_meta").or_insert_with(|| serde_json::json!({}));
if let Some(m) = meta.as_object_mut() {
// `isReplay` only applies to historical replay events, not the
// post-cursor live deltas that reach this path when a target is set.
if mark_replay {
m.insert("isReplay".to_string(), serde_json::json!(true));
}
if let Some(pd) = persist_data {
m.insert("x.ai/persist".to_string(), pd.clone());
}
if let Some(tid) = target_client_id {
m.insert("x.ai/leaderClientId".to_string(), tid.clone());
}
}
}
// Fit after `_meta` is added, so what is measured is what is sent.
if let Ok(raw_val) = serde_json::value::to_raw_value(&params)
&& let Some(raw_val) = Self::fitted_replay_params(raw_val)
{
completions.push(self.gateway.forward_with_completion(
acp::ExtNotification::new(
"x.ai/session/update",
std::sync::Arc::from(raw_val),
),
));
}
}
} else {
// ACP — forward as typed SessionNotification for correct TUI dispatch.
let Ok(mut notification) =
serde_json::from_str::<acp::SessionNotification>(raw_params.get())
else {
tracing::debug!("replay: skipping ACP update with unparseable params");
return;
};
// Collapse ToolCall + all ToolCallUpdates into a single
// pre-completed ToolCall during replay. This gives the pager
// one push() per tool call instead of 2-4.
//
// - ToolCall (registration): buffer by ID, don't forward yet.
// - ToolCallUpdate status=None (start metadata): merge into buffer.
// - ToolCallUpdate InProgress/Pending (streaming): drop.
// - ToolCallUpdate Completed/Failed: merge into buffer, forward
// as a single SessionUpdate::ToolCall with final status.
match &mut notification.update {
acp::SessionUpdate::ToolCall(tc) => {
let is_pre_completed = matches!(
tc.status,
acp::ToolCallStatus::Completed | acp::ToolCallStatus::Failed
);
if is_pre_completed {
// Already complete — forward as-is (no updates will follow).
} else {
pending_tool_calls.insert(tc.tool_call_id.clone(), tc.clone());
return;
}
}
acp::SessionUpdate::ToolCallUpdate(u) => {
match u.fields.status {
Some(acp::ToolCallStatus::Completed)
| Some(acp::ToolCallStatus::Failed) => {
if let Some(mut base) = pending_tool_calls.remove(&u.tool_call_id) {
base.update(std::mem::take(&mut u.fields));
notification.update = acp::SessionUpdate::ToolCall(base);
}
// If no buffered base, forward the ToolCallUpdate as-is.
}
None => {
// Start metadata (title, kind, rawInput, locations).
// Merge into the buffered ToolCall.
if let Some(base) = pending_tool_calls.get_mut(&u.tool_call_id) {
base.update(std::mem::take(&mut u.fields));
}
return;
}
_ => return, // InProgress / Pending — drop
}
}
_ => {}
}
if mark_replay {
mark_as_replay(&mut notification.meta, persist_data);
}
// Stamp the leader unicast target regardless of mark_replay so the
// leader routes both historical and post-cursor live deltas only to
// the loading client.
if let Some(tid) = target_client_id {
stamp_meta_value(&mut notification.meta, "x.ai/leaderClientId", tid);
}
completions.push(self.gateway.forward_with_completion(notification));
}
}
}

230
crates/codegen/xai-grok-shell/src/agent/mvp_agent/replay_tests.rs Normal file
View file

@ -0,0 +1,230 @@
//! Replay of records earlier builds wrote, asserted on what reaches the client.
use std::collections::HashMap;
use agent_client_protocol as acp;
use serde_json::Value;
use super::MvpAgent;
use crate::tools::task_completed_frame::FRAME_MAX_BYTES;
use xai_acp_lib::{AcpAgentGatewaySender as GatewaySender, AcpClientMessage};
/// The record is built from the real notification type, so renaming a field
/// the refit looks up by name fails this test rather than silently putting
/// oversized lines back on the wire.
#[tokio::test]
async fn replay_shrinks_an_oversized_completion() {
let (agent, mut rx) = build_agent_with_gateway();
let line = replay_line(&recorded_completion("Z".repeat(2 * 1024 * 1024)));
agent.forward_raw_replay_line(
&line,
/*persist_data*/ None,
/*target_client_id*/ None,
&mut Vec::new(),
/*mark_replay*/ false,
&mut HashMap::new(),
);
let params = next_ext_notification_params(&mut rx).expect("the record must still be sent");
assert!(
params.len() <= FRAME_MAX_BYTES,
"replayed {} bytes",
params.len()
);
assert!(
params.contains("/tmp/bg-old.log"),
"the log pointer is kept"
);
assert!(
params.contains("keep me"),
"a field this build does not model must survive the refit"
);
}
fn recorded_completion(output: String) -> Value {
use crate::extensions::notification::{SessionNotification, SessionUpdate};
let notification = SessionNotification {
session_id: acp::SessionId::new("s"),
update: SessionUpdate::TaskCompleted {
task_snapshot: xai_grok_tools::types::TaskSnapshot {
task_id: "bg-old".to_string(),
command: "grep -r pattern .".to_string(),
display_command: None,
cwd: "/workspace".to_string(),
start_time: std::time::SystemTime::now(),
end_time: Some(std::time::SystemTime::now()),
output,
output_file: std::path::PathBuf::from("/tmp/bg-old.log"),
truncated: false,
output_total_bytes: 0,
exit_code: Some(0),
signal: None,
completed: true,
block_waited: false,
explicitly_killed: false,
kind: Default::default(),
owner_session_id: None,
description: None,
is_backgrounded: true,
},
will_wake: false,
},
meta: None,
};
let mut record = serde_json::to_value(&notification).expect("serialize");
record["update"]["task_snapshot"]["a_field_from_another_build"] =
Value::String("keep me".to_string());
record
}
fn replay_line(record: &Value) -> String {
serde_json::json!({ "method": "_x.ai/session/update", "params": record }).to_string()
}
/// The branch a plain resume takes: `_meta` is added to the record, so the fit
/// has to happen after that and not before.
#[tokio::test]
async fn a_marked_replay_is_fitted_after_its_metadata_is_added() {
let (agent, mut rx) = build_agent_with_gateway();
let line = replay_line(&recorded_completion("Z".repeat(2 * 1024 * 1024)));
let persist = serde_json::json!({ "padding": "p".repeat(8 * 1024) });
agent.forward_raw_replay_line(
&line,
Some(&persist),
/*target_client_id*/ None,
&mut Vec::new(),
/*mark_replay*/ true,
&mut HashMap::new(),
);
let params = next_ext_notification_params(&mut rx).expect("the record must still be sent");
assert!(
params.len() <= FRAME_MAX_BYTES,
"replayed {} bytes once the metadata was added",
params.len()
);
assert!(params.contains("isReplay"));
}
/// A completion nothing can shrink is dropped, because sending it is what
/// closes the connection.
#[tokio::test]
async fn replay_drops_a_completion_nothing_can_shrink() {
let (agent, mut rx) = build_agent_with_gateway();
let mut record = recorded_completion(String::new());
record["update"]["task_snapshot"]["task_id"] = Value::String("t".repeat(2 * FRAME_MAX_BYTES));
let line = replay_line(&record);
agent.forward_raw_replay_line(
&line,
/*persist_data*/ None,
/*target_client_id*/ None,
&mut Vec::new(),
/*mark_replay*/ false,
&mut HashMap::new(),
);
assert!(next_ext_notification_params(&mut rx).is_none());
}
#[tokio::test]
async fn replay_forwards_records_within_the_limit_untouched() {
let (agent, mut rx) = build_agent_with_gateway();
let record = recorded_completion("hi\n".to_string());
let line = replay_line(&record);
agent.forward_raw_replay_line(
&line,
/*persist_data*/ None,
/*target_client_id*/ None,
&mut Vec::new(),
/*mark_replay*/ false,
&mut HashMap::new(),
);
let params = next_ext_notification_params(&mut rx).expect("forwarded");
assert_eq!(params, serde_json::to_string(&record).unwrap());
}
#[tokio::test]
async fn replay_leaves_other_oversized_records_alone() {
let (agent, mut rx) = build_agent_with_gateway();
let recorded = format!(
r#"{{"sessionId":"s","update":{{"sessionUpdate":"subagent_spawned","subagent_id":"{}"}}}}"#,
"x".repeat(64 * 1024)
);
let line = format!(r#"{{"method":"_x.ai/session/update","params":{recorded}}}"#);
agent.forward_raw_replay_line(
&line,
/*persist_data*/ None,
/*target_client_id*/ None,
&mut Vec::new(),
/*mark_replay*/ false,
&mut HashMap::new(),
);
let params = next_ext_notification_params(&mut rx).expect("forwarded");
assert_eq!(params, recorded);
}
/// Stale-task reconciliation on a cold load builds its completion from a
/// recorded command of any size, so its line is measured like the rest.
#[tokio::test]
async fn a_stale_task_completion_is_frame_bounded() {
let (agent, mut rx) = build_agent_with_gateway();
let dir = tempfile::tempdir().unwrap();
let line = format!(
r#"{{"timestamp":1,"method":"_x.ai/session/update","params":{{"sessionId":"s","update":{{"sessionUpdate":"task_backgrounded","task_id":"stale-1","command":"{}","cwd":"/tmp"}}}}}}"#,
"c".repeat(64 * 1024)
);
let path = dir.path().join("updates.jsonl");
std::fs::write(&path, line).unwrap();
agent.reconcile_stale_background_tasks(&acp::SessionId::new("s"), &Some(path));
let params = next_ext_notification_params(&mut rx).expect("the stale task must be reported");
assert!(
params.len() <= FRAME_MAX_BYTES,
"reconciliation emitted {} bytes",
params.len()
);
assert!(params.contains("session_restart"));
}
fn build_agent_with_gateway() -> (
MvpAgent,
tokio::sync::mpsc::UnboundedReceiver<AcpClientMessage>,
) {
use crate::agent::config::Config as AgentConfig;
use crate::auth::{AuthManager, GrokComConfig};
let temp_dir = tempfile::tempdir().unwrap();
let auth_manager =
std::sync::Arc::new(AuthManager::new(temp_dir.path(), GrokComConfig::default()));
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let agent = MvpAgent::new(
GatewaySender::new(tx),
&AgentConfig::default(),
auth_manager,
None,
)
.expect("valid test config");
(agent, rx)
}
fn next_ext_notification_params(
rx: &mut tokio::sync::mpsc::UnboundedReceiver<AcpClientMessage>,
) -> Option<String> {
let mut params = None;
while let Ok(msg) = rx.try_recv() {
if let AcpClientMessage::ExtNotification(args) = msg {
params.get_or_insert_with(|| args.request.params.get().to_string());
let _ = args.response_tx.send(Ok(()));
}
}
params
}

152
crates/codegen/xai-grok-shell/src/agent/mvp_agent/resource_telemetry.rs Normal file
View file

@ -0,0 +1,152 @@
//! What the process is holding, at session close and on a bounded cadence.
use std::cell::RefCell;
use std::time::{Duration, Instant};
use xai_grok_telemetry::events::ResourceReportTrigger;
use super::*;
/// Floor between growth-driven reports, so a fast climb costs a bounded number
/// of events.
const MIN_REPORT_INTERVAL: Duration = Duration::from_secs(5 * 60);
/// A quiet process still reports this often, so a flat line is evidence rather
/// than absence.
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(60 * 60);
/// Fraction of the last reading that counts as movement worth an event.
const MATERIAL_CHANGE_DIVISOR: u64 = 10;
/// Keeps the periodic reports rare while a process is behaving, and frequent
/// while it is not.
#[derive(Default)]
struct ReportCadence {
last: Option<Report>,
}
// The gauges describe the process, so the cadence that rate-limits them belongs
// to the process, not to any one agent. `//` comments: `///` above
// `thread_local!` trips `clippy::unused_doc_comments`.
thread_local! {
static CADENCE: RefCell<ReportCadence> = RefCell::new(ReportCadence::default());
}
struct Report {
at: Instant,
rss_bytes: Option<u64>,
}
impl ReportCadence {
fn is_due(&self, now: Instant, rss_bytes: Option<u64>) -> bool {
let Some(last) = &self.last else {
return true;
};
let elapsed = now.saturating_duration_since(last.at);
if elapsed >= HEARTBEAT_INTERVAL {
return true;
}
if elapsed < MIN_REPORT_INTERVAL {
return false;
}
match (rss_bytes, last.rss_bytes) {
(Some(current), Some(previous)) => {
current.abs_diff(previous) >= previous / MATERIAL_CHANGE_DIVISOR
}
_ => false,
}
}
fn record(&mut self, at: Instant, rss_bytes: Option<u64>) {
self.last = Some(Report { at, rss_bytes });
}
}
impl MvpAgent {
/// Sampled after removal, so a leak reads as a rising tail across releases.
pub(super) fn log_resource_usage(&self, trigger: ResourceReportTrigger) {
let usage = xai_tty_utils::sample_process_resources();
CADENCE.with_borrow_mut(|cadence| cadence.record(Instant::now(), usage.rss_bytes));
xai_grok_telemetry::session_ctx::log_event(
xai_grok_telemetry::events::ProcessResourceUsage {
trigger,
rss_bytes: usage.rss_bytes,
peak_rss_bytes: usage.peak_rss_bytes,
footprint_bytes: usage.footprint_bytes,
threads: usage.threads,
open_files: usage.open_files,
resident_sessions: self.sessions.borrow().len(),
session_threads: self.session_registry.counts().session_threads,
},
);
}
/// Called from the heap monitor's poll loop, which runs whether or not
/// profiling is on. A session that never closes reports through this.
///
/// Every tick pays one memory read; only a tick that reports pays for the
/// thread and descriptor counts.
pub(super) fn report_resource_usage_if_due(&self) {
let memory = xai_tty_utils::sample_process_memory();
let due = CADENCE.with_borrow(|cadence| cadence.is_due(Instant::now(), memory.rss_bytes));
if due {
self.log_resource_usage(ResourceReportTrigger::Periodic);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn at(base: Instant, secs: u64) -> Instant {
base + Duration::from_secs(secs)
}
fn after(rss_bytes: u64) -> (Instant, ReportCadence) {
let base = Instant::now();
let mut cadence = ReportCadence::default();
cadence.record(base, Some(rss_bytes));
(base, cadence)
}
#[test]
fn the_first_report_is_always_due() {
assert!(ReportCadence::default().is_due(Instant::now(), Some(100)));
}
#[test]
fn a_quiet_process_still_reports_hourly() {
let (base, cadence) = after(1_000);
assert!(cadence.is_due(at(base, 60 * 60), Some(1_000)));
}
#[test]
fn movement_reports_once_past_the_floor() {
let (base, cadence) = after(1_000);
assert!(
!cadence.is_due(at(base, 60), Some(2_000)),
"the floor bounds what a fast climb can cost"
);
assert!(cadence.is_due(at(base, 6 * 60), Some(1_200)));
assert!(
!cadence.is_due(at(base, 6 * 60), Some(1_050)),
"under a tenth is noise"
);
assert!(
cadence.is_due(at(base, 6 * 60), Some(800)),
"a purge is as worth reporting as a leak"
);
}
#[test]
fn an_unreadable_gauge_falls_back_to_the_heartbeat() {
let base = Instant::now();
let mut cadence = ReportCadence::default();
cadence.record(base, None);
assert!(!cadence.is_due(at(base, 30 * 60), None));
assert!(cadence.is_due(at(base, 60 * 60), None));
}
}

1
crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs
View file

@ -78,6 +78,7 @@ impl MvpAgent {
if let Some(ops) = self.workspace_ops.borrow().as_ref() {
ops.end_local_session(id.0.as_ref());
}
self.log_resource_usage(xai_grok_telemetry::events::ResourceReportTrigger::SessionClose);
}
/// Get-or-create the per-session dispatch lock. `prompt` holds it across
/// intake so a cancel cannot overtake the prompt it targets.

18
crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests.rs
View file

@ -4395,13 +4395,9 @@ async fn post_auth_settings_non_xai_keeps_local_but_still_emits() {
"settings arrival must push x.ai/settings/update for non-xai auth too"
);
}
/// A failed post-auth fetch must re-close the gate and leave it closed. Guards
/// two behaviors a passing-on-`Fetched` test can't: the account-switch
/// re-suppress fires (gate was open, identity not yet resolved), and a
/// transient/4xx outcome (`Retry`) does not reopen it.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn post_auth_settings_retry_re_suppresses_and_stays_closed() {
async fn post_auth_settings_failure_resolves_gate_onto_local_policy() {
use crate::agent::config::AgentMode;
use crate::auth::{GrokAuth, XAI_OAUTH2_ISSUER};
let _restore = RestoreOtelGate;
@ -4413,12 +4409,16 @@ async fn post_auth_settings_retry_re_suppresses_and_stays_closed() {
..GrokAuth::test_default()
};
let (agent, _rx) = build_agent_with_auth_and_proxy(xai_auth, server.url(), AgentMode::Leader);
xai_grok_telemetry::external::mark_external_otel_settings_resolved();
assert!(xai_grok_telemetry::external::is_settings_gate_open());
xai_grok_telemetry::external::suppress_external_otel_until_settings();
assert!(!xai_grok_telemetry::external::is_settings_gate_open());
agent.maybe_fetch_post_auth_settings().await;
assert!(
!xai_grok_telemetry::external::is_settings_gate_open(),
"a Retry (failed) post-auth fetch must re-close the gate and keep it closed"
xai_grok_telemetry::external::is_settings_gate_open(),
"an exhausted fetch is a definitive answer: open on local policy"
);
assert!(
agent.cfg.borrow().remote_settings.is_none(),
"opening the gate must not fabricate settings; none were fetched"
);
}
/// A same-credential refresh must NOT re-suppress a gate already resolved for

184
crates/codegen/xai-grok-shell/src/agent/otel_gate.rs
View file

@ -5,39 +5,83 @@
//! [`xai_grok_telemetry::external`]:
//!
//! 1. Startup (no leader instance yet): [`suppress`] closes the gate before
//! telemetry init; [`open_at_startup`] re-opens it only for a pure
//! env-API-key leader ([`should_open_at_startup`]), which has no remote
//! policy to fetch.
//! telemetry init; [`open_at_startup`] re-opens it when nothing will deliver
//! a fleet policy to this process ([`should_open_at_startup`]).
//! 2. Post-auth/refresh (per-leader): [`OtelGate::resolve`] drives the gate
//! from the [`SettingsFetch`] outcome for the still-live identity.
//!
//! A leader that never authenticates keeps the gate closed for life: the gate
//! fails safe by dropping telemetry, never by shipping it early.
use std::time::Duration;
use crate::remote::SettingsFetch;
use crate::util::config::RemoteSettings;
pub(crate) const SETTINGS_GATE_MAX_WAIT: Duration = crate::http::SETTINGS_REAPPLY_TIMEOUT;
/// Closes the gate. Process-global and idempotent; callable before any
/// `AgentConfig` exists.
pub(crate) fn suppress() {
xai_grok_telemetry::external::set_settings_gate_max_wait(SETTINGS_GATE_MAX_WAIT);
xai_grok_telemetry::external::suppress_external_otel_until_settings();
}
/// Whether an xAI fleet policy can govern this process at all.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PolicyChannel {
Applies,
Unavailable(NoPolicy),
}
impl PolicyChannel {
pub(crate) fn is_unavailable(self) -> bool {
matches!(self, Self::Unavailable(_))
}
}
/// Why no fleet policy can reach this process.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum NoPolicy {
RemoteFetchDisabled,
ProxyRepointed,
}
pub(crate) fn policy_channel(remote_fetch_enabled: bool, proxy_is_xai: bool) -> PolicyChannel {
if !remote_fetch_enabled {
return PolicyChannel::Unavailable(NoPolicy::RemoteFetchDisabled);
}
if !proxy_is_xai {
return PolicyChannel::Unavailable(NoPolicy::ProxyRepointed);
}
PolicyChannel::Applies
}
/// [`policy_channel`] resolved against live config for the proxy actually in
/// use.
pub(crate) fn policy_channel_for(proxy_url: &str) -> PolicyChannel {
policy_channel(
crate::util::config::resolve_remote_fetch_enabled(),
crate::util::is_cli_chat_proxy_url(proxy_url),
)
}
/// [`policy_channel_for`] against the effective config, for startup call sites
/// that run before an `AgentConfig` exists.
pub(crate) fn resolved_policy_channel() -> PolicyChannel {
policy_channel_for(&crate::agent::config::EndpointsConfig::from_effective_config().proxy_url())
}
/// Inputs to [`should_open_at_startup`]. Named fields prevent transposed
pub(crate) struct StartupGate {
pub(crate) channel: PolicyChannel,
pub(crate) has_session: bool,
pub(crate) has_api_key_env: bool,
pub(crate) session_pending: bool,
/// When false, no remote fleet policy can arrive, so the gate fails open to the leader's local telemetry decision.
pub(crate) remote_fetch_enabled: bool,
}
/// Returns whether a leader opens the gate at startup: only a pure
/// Returns whether a leader opens the gate at startup.
pub(crate) fn should_open_at_startup(gate: StartupGate) -> bool {
// No remote policy will arrive with `remote_fetch` off, so fail open to
if !gate.remote_fetch_enabled {
if gate.channel.is_unavailable() {
return true;
}
!gate.has_session && gate.has_api_key_env && !gate.session_pending
!gate.has_session && !gate.session_pending
}
/// Returns whether a session-less startup is about to mint a grok.com session
@ -62,14 +106,20 @@ pub(crate) struct OtelGate {
}
impl OtelGate {
/// Re-closes the gate before fetching a different identity's policy, so a stale open can't leak across an account switch.
pub(crate) fn rearm_on_switch(&self, identity: &str) {
/// Re-closes the gate before fetching a different identity's policy, so a
/// stale open can't leak across an account switch.
pub(crate) fn rearm_on_switch(&self, identity: &str, channel: PolicyChannel) {
if channel.is_unavailable() {
return;
}
if identity.is_empty() || self.resolved_for.borrow().as_deref() != Some(identity) {
xai_grok_telemetry::external::suppress_external_otel_until_settings();
suppress();
}
}
/// Drives the gate from a settings-fetch `outcome` for `identity`: fail-closed on transient outcomes, opens on a definitive one. Returns settings only when fetched.
/// Drives the gate from a settings-fetch `outcome` for `identity`. Every
/// outcome for the live identity is definitive and opens the gate; only the
/// `Fetched` one carries a policy (and settings) to apply.
pub(crate) fn resolve(
&self,
identity: &str,
@ -84,11 +134,10 @@ impl OtelGate {
self.apply_and_open(identity, Some(&settings));
Some(*settings)
}
SettingsFetch::Rejected => {
SettingsFetch::Rejected | SettingsFetch::Retry => {
self.apply_and_open(identity, None);
None
}
SettingsFetch::Retry => None,
}
}
@ -126,38 +175,81 @@ mod tests {
}
#[test]
fn startup_gate_fails_open_when_remote_fetch_disabled() {
// remote_fetch off => no remote policy will ever arrive => fail open,
assert!(should_open_at_startup(StartupGate {
has_session: true,
has_api_key_env: false,
session_pending: false,
remote_fetch_enabled: false,
}));
assert!(!should_open_at_startup(StartupGate {
has_session: true,
has_api_key_env: false,
session_pending: false,
remote_fetch_enabled: true,
}));
fn policy_channel_reports_every_structural_reason() {
assert_eq!(
policy_channel(false, true),
PolicyChannel::Unavailable(NoPolicy::RemoteFetchDisabled),
"remote_fetch off: the deployment declared it never calls xAI"
);
assert_eq!(
policy_channel(true, false),
PolicyChannel::Unavailable(NoPolicy::ProxyRepointed),
"a non-xAI proxy is not governed by xAI fleet policy"
);
assert_eq!(
policy_channel(false, false),
PolicyChannel::Unavailable(NoPolicy::RemoteFetchDisabled),
"the explicit config decision is reported ahead of the endpoint"
);
assert_eq!(
policy_channel(true, true),
PolicyChannel::Applies,
"xAI proxy + fetches allowed: a policy can arrive, so wait for it"
);
}
#[test]
fn startup_gate_opens_whenever_no_policy_will_arrive() {
let opens = |channel, has_session, session_pending| {
should_open_at_startup(StartupGate {
channel,
has_session,
session_pending,
})
};
let applies = PolicyChannel::Applies;
for reason in [NoPolicy::RemoteFetchDisabled, NoPolicy::ProxyRepointed] {
let none = PolicyChannel::Unavailable(reason);
assert!(
opens(none, true, false),
"{reason:?}: no policy can arrive, so a session must not wait"
);
assert!(opens(none, false, true), "{reason:?}: nor a pending mint");
}
assert!(
!opens(applies, true, false),
"a session with a reachable policy waits for it"
);
assert!(
!opens(applies, false, true),
"a pending mint is a session about to exist; wait for its policy"
);
assert!(
opens(applies, false, false),
"no session and none pending: nothing will query the channel yet"
);
}
#[test]
#[serial_test::serial]
fn resolve_opens_only_on_definitive_outcome_for_live_identity() {
fn resolve_opens_on_every_definitive_outcome_for_the_live_identity() {
let _restore = RestoreGate;
let gate = OtelGate::default();
suppress_external_otel_until_settings();
assert!(
gate.resolve("alice", SettingsFetch::Retry, Some("alice"))
.is_none()
.is_none(),
"a failed fetch yields no settings"
);
assert!(
!is_settings_gate_open(),
"a transient outcome stays fail-closed"
is_settings_gate_open(),
"an exhausted fetch must open the gate rather than mute the stream"
);
suppress_external_otel_until_settings();
assert!(
gate.resolve("alice", SettingsFetch::Rejected, Some("alice"))
.is_none()
@ -200,10 +292,26 @@ mod tests {
gate.set_resolved_for("");
mark_external_otel_settings_resolved();
gate.rearm_on_switch("");
gate.rearm_on_switch("", PolicyChannel::Applies);
assert!(
!is_settings_gate_open(),
"an empty identity must always re-close (cannot prove same credential)"
);
}
#[test]
#[serial_test::serial]
fn rearm_never_re_closes_when_no_policy_can_arrive() {
let _restore = RestoreGate;
let gate = OtelGate::default();
for reason in [NoPolicy::RemoteFetchDisabled, NoPolicy::ProxyRepointed] {
mark_external_otel_settings_resolved();
gate.rearm_on_switch("alice", PolicyChannel::Unavailable(reason));
assert!(
is_settings_gate_open(),
"{reason:?}: re-closing would wait on a policy that cannot arrive"
);
}
}
}

4
crates/codegen/xai-grok-shell/src/agent/proxy.rs
View file

@ -32,7 +32,7 @@ use tracing::debug;
/// 2. If `HTTPS_PROXY` (or `https_proxy`) is set, return its value.
/// 3. If `HTTP_PROXY` (or `http_proxy`) is set, return its value.
/// 4. Otherwise return `None`.
pub fn resolve_proxy_for_host(target_host: &str) -> Option<String> {
pub(crate) fn resolve_proxy_for_host(target_host: &str) -> Option<String> {
resolve_proxy_for_host_with(target_host, |key| std::env::var(key))
}
@ -120,7 +120,7 @@ fn is_host_bypassed(host: &str, no_proxy: &str) -> bool {
/// 3. Wrap the tunnel in TLS (using rustls with native root certificates).
/// 4. Return the stream as `MaybeTlsStream<TcpStream>` so it is compatible
/// with `tokio_tungstenite::client_async`.
pub async fn connect_via_proxy(
pub(crate) async fn connect_via_proxy(
proxy_url: &str,
target_host: &str,
target_port: u16,

4
crates/codegen/xai-grok-shell/src/agent/relay.rs
View file

@ -82,7 +82,7 @@ impl RelayConfig {
}
}
/// Callback type for first connection event.
pub type FirstConnectCallback = Box<dyn FnOnce() + Send + 'static>;
pub(crate) type FirstConnectCallback = Box<dyn FnOnce() + Send + 'static>;
/// Handle to a running relay connection.
///
/// The relay maintains a persistent WebSocket connection to grok.com with
@ -130,7 +130,7 @@ pub fn spawn_relay_connection(
///
/// Same as `spawn_relay_connection` but allows providing a callback that will be
/// called once when the first successful connection is established.
pub fn spawn_relay_connection_with_callback(
pub(crate) fn spawn_relay_connection_with_callback(
config: RelayConfig,
to_agent_tx: mpsc::UnboundedSender<String>,
parent_cancel: Option<CancellationToken>,

5
crates/codegen/xai-grok-shell/src/agent/server.rs
View file

@ -83,7 +83,7 @@ struct NewConnectionChannels {
/// Query parameters for WebSocket connection.
#[derive(Debug, serde::Deserialize, Default)]
pub struct WsQueryParams {
pub(crate) struct WsQueryParams {
#[serde(rename = "server-key")]
pub server_key: Option<String>,
}
@ -307,9 +307,6 @@ async fn run_persistent_agent(
// Restore managed policy right before bootstrap reads it — the agent is created lazily here,
// so an earlier restore could go stale before the gate.
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
// Fail-closed external-OTEL gate: suppress until settings resolve, opening
// now only for a pure env-API-key user (no remote policy). Matches the
// stdio/leader boot; per-connection settings reopen it via `initialize`.
crate::agent::app::apply_otel_config(&auth_manager, &agent_config.grok_com_config);
let agent = Rc::new(
MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)

4
crates/codegen/xai-grok-shell/src/agent/session_config.rs
View file

@ -14,7 +14,7 @@ pub(crate) const SELECTABLE_REASONING_EFFORTS: [ReasoningEffort; 5] = [
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionConfigOption {
pub(crate) struct SessionConfigOption {
pub id: String,
pub category: String,
pub label: String,
@ -25,7 +25,7 @@ pub struct SessionConfigOption {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GrokSessionDetail {
pub(crate) struct GrokSessionDetail {
pub session_id: String,
pub kind: String,
pub cwd: String,

2
crates/codegen/xai-grok-shell/src/agent/session_registry_client.rs
View file

@ -323,7 +323,7 @@ impl SessionRegistryClient {
}
/// GET /v1/sessions/{id}/download — returns a signed GCS URL without downloading.
pub async fn get_download_url(
pub(crate) async fn get_download_url(
&self,
session_id: &str,
file: &str,

16
crates/codegen/xai-grok-shell/src/agent/subagent/mod.rs
View file

@ -340,7 +340,7 @@ impl SubagentSpawnContext {
/// catalog used to pick the subagent's `SamplerConfig`); user TOML and
/// GB global tiers are sourced from the parent's snapshot captured at
/// spawn-context build time.
pub fn resolve_auto_compact_threshold_percent(&self, subagent_model_id: &str) -> u8 {
pub(crate) fn resolve_auto_compact_threshold_percent(&self, subagent_model_id: &str) -> u8 {
let gb_per_model =
crate::agent::config::find_model_by_id(&self.available_models, subagent_model_id)
.and_then(|e| e.info.auto_compact_threshold_percent);
@ -362,7 +362,7 @@ impl SubagentSpawnContext {
}
}
/// Subagent verbatim-input flag, mirroring `Config::resolve_compaction_verbatim_input` (env > config > remote settings > default `true`).
pub fn resolve_compaction_verbatim_input(&self) -> bool {
pub(crate) fn resolve_compaction_verbatim_input(&self) -> bool {
crate::agent::config::BoolFlag::env("GROK_COMPACTION_VERBATIM_INPUT")
.config(
self.agent_config
@ -378,7 +378,9 @@ impl SubagentSpawnContext {
.resolve()
.value
}
pub fn resolve_compaction_tool_choice(&self) -> crate::util::config::CompactionToolChoice {
pub(crate) fn resolve_compaction_tool_choice(
&self,
) -> crate::util::config::CompactionToolChoice {
crate::util::config::resolve_compaction_tool_choice_from(
crate::agent::config::env_string(crate::util::config::ENV_COMPACTION_TOOL_CHOICE)
.as_deref(),
@ -395,7 +397,7 @@ impl SubagentSpawnContext {
/// (env > config > remote settings > default). Default `false` so it ships dark;
/// `managed_config.toml` `[features] subagent_worktree_snapshot` is the
/// per-deployment rollout lever.
pub fn resolve_subagent_worktree_snapshot_enabled(&self) -> bool {
pub(crate) fn resolve_subagent_worktree_snapshot_enabled(&self) -> bool {
crate::agent::config::BoolFlag::env("GROK_SUBAGENT_WORKTREE_SNAPSHOT")
.config(
self.agent_config
@ -416,7 +418,7 @@ impl SubagentSpawnContext {
/// parent (requirements/env/user/managed from disk; remote from the
/// parent's snapshot) and follows the session into subagents. Bash stays
/// on tool defaults, as before that knob existed.
pub fn resolve_tool_params_json(
pub(crate) fn resolve_tool_params_json(
&self,
) -> crate::session::agent_rebuild::ResolvedToolParamsJson {
let params = crate::util::config::resolve_ask_user_question_params_from_disk(
@ -2268,7 +2270,7 @@ pub(crate) struct SubagentMeta {
/// locally. Schema is versioned for forward compatibility.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SubagentSessionMetadata {
pub(crate) struct SubagentSessionMetadata {
pub schema_version: u32,
pub session_id: String,
pub session_kind: String,
@ -2327,7 +2329,7 @@ impl SubagentSessionMetadata {
/// Current schema version.
pub const SCHEMA_VERSION: u32 = 1;
/// Build from a `SubagentMeta` + additional runtime context.
pub fn from_meta(
pub(crate) fn from_meta(
meta: &SubagentMeta,
model_id: Option<&str>,
cwd: Option<&str>,

6
crates/codegen/xai-grok-shell/src/auth/attribution.rs
View file

@ -132,7 +132,7 @@ impl ShellAttribution {
/// The two callbacks share the same underlying impl and emit the
/// same `auth_401_attribution` event format -- only the trait
/// signature differs (`SamplingConsumer` vs. `ToolConsumer`).
pub fn new_tool_callback(
pub(crate) fn new_tool_callback(
auth_manager: Arc<AuthManager>,
session_id: Option<String>,
) -> Arc<dyn ToolAuth401AttributionCallback> {
@ -797,14 +797,14 @@ mod tests {
use tracing_subscriber::registry::LookupSpan;
#[derive(Debug, Default, Clone)]
pub struct CapturedSpan {
pub(crate) struct CapturedSpan {
pub name: String,
pub fields_str: std::collections::BTreeMap<String, String>,
pub fields_i64: std::collections::BTreeMap<String, i64>,
pub fields_bool: std::collections::BTreeMap<String, bool>,
}
pub struct SpanCollector {
pub(crate) struct SpanCollector {
pub spans: std::sync::Arc<Mutex<Vec<CapturedSpan>>>,
}

14
crates/codegen/xai-grok-shell/src/auth/config.rs
View file

@ -137,7 +137,7 @@ pub const XAI_OAUTH2_ISSUER: &str = "https://auth.x.ai";
const PROD_ACCOUNTS_APP_ORIGINS: &[&str] = &["https://accounts.x.ai"];
/// See the opt-in non-production feature variant above — builds without
/// the feature accept only the production accounts app.
pub fn allowed_accounts_app_origins() -> Vec<String> {
pub(crate) fn allowed_accounts_app_origins() -> Vec<String> {
PROD_ACCOUNTS_APP_ORIGINS
.iter()
.map(|o| o.to_string())
@ -148,7 +148,7 @@ pub fn allowed_accounts_app_origins() -> Vec<String> {
///
/// Callers can chain additional configuration (e.g. `.allow_headers(...)` or
/// `.allow_private_network(true)`) onto the returned layer.
pub fn accounts_app_cors_layer(method: axum::http::Method) -> tower_http::cors::CorsLayer {
pub(crate) fn accounts_app_cors_layer(method: axum::http::Method) -> tower_http::cors::CorsLayer {
tower_http::cors::CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::list(
allowed_accounts_app_origins()
@ -168,7 +168,7 @@ const XAI_OAUTH2_LOCAL_ISSUER: &str = "http://localhost:22255";
const DEFAULT_OAUTH2_REFERRER: &str = "grok-build";
/// Returns `true` when `GROK_LOCAL_AUTH=1` is set,
/// indicating the local accounts-app should be used as the OAuth2 issuer.
pub fn use_local_auth() -> bool {
pub(crate) fn use_local_auth() -> bool {
std::env::var("GROK_LOCAL_AUTH")
.map(|v| !v.is_empty() && v != "0")
.unwrap_or(false)
@ -201,7 +201,7 @@ impl GrokComConfig {
/// otherwise set `disable_api_key_auth = false` and override it — so the env
/// is OR-ed in here and cannot be turned back off by a user layer. Trusted
/// `requirements.toml` already wins over `config.toml` via layer precedence.
pub fn api_key_auth_disabled(&self) -> bool {
pub(crate) fn api_key_auth_disabled(&self) -> bool {
self.disable_api_key_auth == Some(true)
|| self.force_login_team_uuid.is_some()
|| env_lockdown_forced()
@ -210,7 +210,7 @@ impl GrokComConfig {
/// interactive browser login, external auth provider) must not run — the
/// pin is fail-closed. Explicit `grok login --devbox` / `--api-key` bypass
/// this by not consulting automatic flow helpers.
pub fn blocks_automatic_oidc(&self) -> bool {
pub(crate) fn blocks_automatic_oidc(&self) -> bool {
matches!(self.preferred_method, Some(PreferredAuthMethod::ApiKey))
}
/// The auth.json scope key for this config.
@ -252,7 +252,7 @@ impl OAuth2ProviderConfig {
})
}
/// Convert to [`OidcAuthConfig`] to reuse the OIDC login flow.
pub fn as_oidc(&self) -> OidcAuthConfig {
pub(crate) fn as_oidc(&self) -> OidcAuthConfig {
OidcAuthConfig {
issuer: self.issuer.clone(),
client_id: self.client_id.clone(),
@ -260,7 +260,7 @@ impl OAuth2ProviderConfig {
audience: None,
}
}
pub fn base_auth_scope(&self) -> String {
pub(crate) fn base_auth_scope(&self) -> String {
format!("{}::{}", self.issuer.trim_end_matches('/'), self.client_id)
}
pub fn auth_scope(&self) -> String {

14
crates/codegen/xai-grok-shell/src/auth/credential_provider.rs
View file

@ -35,7 +35,7 @@ impl xai_grok_sampler::BearerResolver for WireValidBearerResolver {
}
/// Production impl: wraps the live `AuthManager`. 401 recovery
/// delegates to `AuthManager::unauthorized_recovery`.
pub struct ShellAuthCredentialProvider {
pub(crate) struct ShellAuthCredentialProvider {
auth_manager: Arc<AuthManager>,
static_credentials: GrokAuthCredentials,
}
@ -250,7 +250,7 @@ impl xai_file_utils::storage_client::Auth401AttributionCallback for StorageClien
///
/// Before upgrade, the bootstrap manager provides disk-read-only
/// behavior (equivalent to the pre-consolidation OTel path).
pub struct OtelAuthCredentialProvider {
pub(crate) struct OtelAuthCredentialProvider {
/// Bootstrap manager used before the live one is available.
bootstrap: Arc<AuthManager>,
/// Swapped to the agent's live `AuthManager` via `set_live()`.
@ -271,10 +271,10 @@ impl OtelAuthCredentialProvider {
/// `snapshot()` reads from the live manager (proactive refresh
/// keeps it hot) and `refresh_after_unauthorized()` drives the
/// full recovery state machine.
pub fn set_live(&self, auth_manager: Arc<AuthManager>) {
pub(crate) fn set_live(&self, auth_manager: Arc<AuthManager>) {
self.live.store(Arc::new(Some(auth_manager)));
}
pub fn set_deployment_key(&self, key: String) {
pub(crate) fn set_deployment_key(&self, key: String) {
self.deployment_key.store(Arc::new(Some(key)));
}
/// Single-load snapshot of the live/bootstrap state.
@ -380,7 +380,7 @@ static OTEL_PROVIDER: std::sync::OnceLock<Arc<OtelAuthCredentialProvider>> =
/// `AuthManager`. Call this once after the main `AuthManager` is
/// constructed and has its refresher configured. No-ops if the OTel
/// layer was never initialized (e.g. `InstrumentationMode::Disabled`).
pub fn wire_otel_auth_manager(auth_manager: Arc<AuthManager>) {
pub(crate) fn wire_otel_auth_manager(auth_manager: Arc<AuthManager>) {
if let Some(provider) = OTEL_PROVIDER.get() {
provider.set_live(auth_manager);
tracing::debug!("otel: upgraded credential provider to live AuthManager");
@ -392,7 +392,7 @@ pub fn wire_otel_auth_manager(auth_manager: Arc<AuthManager>) {
/// stamps per-export, so both pipelines attribute identically. No-op when
/// the OTel provider was never initialized or the external stream is
/// dormant.
pub fn sync_external_otel_identity() {
pub(crate) fn sync_external_otel_identity() {
if let Some(provider) = OTEL_PROVIDER.get() {
let snapshot = provider.snapshot();
xai_grok_telemetry::external::set_identity(
@ -401,7 +401,7 @@ pub fn sync_external_otel_identity() {
}
}
/// No-ops if the OTel layer was never initialized.
pub fn wire_otel_deployment_key(key: String) {
pub(crate) fn wire_otel_deployment_key(key: String) {
if let Some(provider) = OTEL_PROVIDER.get() {
provider.set_deployment_key(key);
tracing::debug!("otel: set deployment key on credential provider");

4
crates/codegen/xai-grok-shell/src/auth/devbox_login_stub.rs
View file

@ -32,6 +32,8 @@ pub(super) async fn mint_devbox_auth_raw() -> anyhow::Result<GrokAuth> {
}
/// `grok login --devbox` entry point: always errors in this build.
pub async fn run_devbox_login(_config: &crate::agent::config::Config) -> anyhow::Result<GrokAuth> {
pub(crate) async fn run_devbox_login(
_config: &crate::agent::config::Config,
) -> anyhow::Result<GrokAuth> {
anyhow::bail!(UNAVAILABLE)
}

12
crates/codegen/xai-grok-shell/src/auth/device_code.rs
View file

@ -27,7 +27,7 @@ const MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS: i64 = 10 * 60;
/// variant hides the `reqwest::Error` the login funnel classifies, because
/// transparent forwards `source()` past the error it wraps.
#[derive(Debug, Error)]
pub enum DeviceCodeError {
pub(crate) enum DeviceCodeError {
#[error(
"Device-code login is not available for this deployment. \
Try `grok login` or set XAI_API_KEY instead."
@ -44,7 +44,7 @@ pub enum DeviceCodeError {
/// consent page — the traffic that otherwise pollutes the device-flow
/// conversion denominator.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientSurface {
pub(crate) enum ClientSurface {
/// An interactive front-end (TUI / IDE) renders the URL + code to a human.
Ui,
/// CLI attached to an interactive terminal (stderr is a TTY).
@ -79,7 +79,7 @@ fn detect_cli_surface() -> ClientSurface {
/// Callers display `verification_uri` + `user_code` to the user,
/// then pass this struct to `complete_device_code_login`.
#[derive(Debug, Clone)]
pub struct DeviceCode {
pub(crate) struct DeviceCode {
pub verification_uri: String,
pub verification_uri_complete: Option<String>,
pub user_code: String,
@ -129,7 +129,7 @@ struct IdTokenClaims {
/// This is a single HTTP POST. The caller is responsible for displaying
/// `DeviceCode::verification_uri` and `DeviceCode::user_code` to the user
/// before calling `complete_device_code_login`.
pub async fn request_device_code(
pub(crate) async fn request_device_code(
issuer: &str,
client_id: &str,
scopes: &[String],
@ -203,7 +203,7 @@ pub async fn request_device_code(
///
/// Callers should have already displayed `device_code.verification_uri`
/// and `device_code.user_code` to the user before calling this.
pub async fn complete_device_code_login(
pub(crate) async fn complete_device_code_login(
issuer: &str,
client_id: &str,
device_code: DeviceCode,
@ -290,7 +290,7 @@ pub async fn complete_device_code_login(
///
/// Takes `channels` by `&mut`, consuming it only after the device code is
/// obtained, so callers can reuse it for a loopback fallback on `NotEnabled`.
pub async fn run_device_code_login_channels(
pub(crate) async fn run_device_code_login_channels(
issuer: &str,
client_id: &str,
scopes: &[String],

43
crates/codegen/xai-grok-shell/src/auth/error.rs
View file

@ -1,3 +1,5 @@
use std::borrow::Cow;
use thiserror::Error;
#[derive(Debug, Error)]
@ -84,9 +86,12 @@ pub enum RefreshTokenFailedReason {
RefreshTokenRejected,
/// `invalid_client` — the client/app credential was rejected.
ClientRejected,
/// Escalation from repeated transient failures (OIDC) or a single
/// external-binary failure. Never a raw IdP code: an unrecognized terminal
/// code is classified transient, not `Other` (see `classify_terminal`).
/// The operator's `auth_provider_command` could not mint a credential in a
/// headless run (`GROK_AUTH_EXPIRED=1`).
ProviderInteractiveRequired,
/// Escalation from repeated transient failures (OIDC). Never a raw IdP
/// code: an unrecognized terminal code is classified transient, not
/// `Other` (see `classify_terminal`).
Other,
}
@ -97,27 +102,53 @@ impl RefreshTokenFailedReason {
pub(crate) fn is_sticky(self) -> bool {
match self {
Self::RefreshTokenRejected => true,
Self::ClientRejected | Self::ProviderInteractiveRequired | Self::Other => false,
}
}
/// Whether the verdict rules out an unattended retry for as long as it
/// stands. Orthogonal to [`Self::is_sticky`], which is about whether the
/// verdict ever ages out.
pub(crate) fn blocks_unattended_retry(self) -> bool {
match self {
Self::RefreshTokenRejected | Self::ProviderInteractiveRequired => true,
Self::ClientRejected | Self::Other => false,
}
}
/// User-facing copy for a terminal refresh failure; the raw IdP code stays
/// in logs.
pub(crate) fn user_message(self) -> &'static str {
pub(crate) fn user_message(self) -> Cow<'static, str> {
match self {
Self::RefreshTokenRejected => {
"Your session has expired. Run `grok login` to sign in again."
"Your session has expired. Run `grok login` to sign in again.".into()
}
Self::ClientRejected => {
"Authentication is temporarily unavailable. Run `grok login` if this persists."
.into()
}
Self::ProviderInteractiveRequired => provider_login_message(None),
Self::Other => {
"Authentication could not be refreshed. Run `grok login` to sign in again."
"Authentication could not be refreshed. Run `grok login` to sign in again.".into()
}
}
}
}
/// `label` is the operator's `auth_provider_label`, where the surface has one.
pub(crate) fn provider_login_message(label: Option<&str>) -> Cow<'static, str> {
match label {
Some(label) => format!(
"Your session expired and {label} could not renew it in the background. \
Run /login to sign in again."
)
.into(),
None => "Your session expired and your sign-in helper could not renew it in the \
background. Run /login to sign in again."
.into(),
}
}
impl AuthError {
/// A retryable refresh failure with a message-only cause, for the genuinely
/// message-only sites (lock timeout, sleep/dark-wake defer, no refresher);

3
crates/codegen/xai-grok-shell/src/auth/external_auth.rs
View file

@ -40,7 +40,8 @@ pub(crate) fn parse_output(output: &std::process::Output) -> anyhow::Result<Grok
}
/// Short timeout for a mid-session refresh: it must not hang the session.
const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
/// The single run in `ExternalBinaryRefresher` gets this whole budget.
const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(7);
/// Runs the external auth binary for a headless mid-session refresh. Initial,
/// interactive sign-in takes a separate path (`flow::run_external_auth_provider`,

8
crates/codegen/xai-grok-shell/src/auth/flow.rs
View file

@ -11,7 +11,7 @@ use crate::http::TransportFailureKind;
use crate::util::grok_home;
use xai_grok_telemetry::events::{LoginFailed, LoginFailureKind};
pub type StderrCallback = Box<dyn Fn(&str)>;
pub(crate) type StderrCallback = Box<dyn Fn(&str)>;
/// Reject a cached credential for reuse if it lacks `oidc_issuer`, has a
/// mismatched issuer, or its team principal violates the `force_login_team_uuid`
@ -186,7 +186,7 @@ impl AuthUrlMode {
}
/// Back-compat flag for older clients that only read `external_provider`.
pub fn is_external_provider(self) -> bool {
pub(crate) fn is_external_provider(self) -> bool {
matches!(self, Self::Command)
}
}
@ -319,7 +319,7 @@ async fn run_external_auth_provider(
}
/// GUI auth: bridges external provider stderr to `url_tx`, pipes code submission via `code_rx`.
pub async fn run_auth_flow_with_stderr_bridge(
pub(crate) async fn run_auth_flow_with_stderr_bridge(
auth_manager: &Arc<AuthManager>,
grok_com_config: &GrokComConfig,
channels: AuthChannels,
@ -432,7 +432,7 @@ pub async fn run_auth_flow(
/// Like [`run_auth_flow`] but with `force_interactive`: skip cached
/// credentials without clearing them. Used by `/login` for mid-session
/// re-auth where abandoning the flow must not disrupt the session.
pub async fn run_auth_flow_interactive(
pub(crate) async fn run_auth_flow_interactive(
auth_manager: &Arc<AuthManager>,
grok_com_config: &GrokComConfig,
on_stderr: Option<StderrCallback>,

50
crates/codegen/xai-grok-shell/src/auth/manager.rs
View file

@ -14,6 +14,9 @@ use tokio_util::sync::CancellationToken;
mod enrichment;
#[path = "manager/lock.rs"]
pub(super) mod lock;
#[path = "manager/remedy.rs"]
mod remedy;
pub(crate) use remedy::{AuthRemedy, SilentRefresh};
#[path = "manager/sleep_gate.rs"]
mod sleep_gate;
@ -64,8 +67,8 @@ pub(crate) enum RefreshReason {
pub(crate) const AUTH_LOCK_TIMEOUT: StdDuration = StdDuration::from_secs(10);
/// Lock timeout for `refresh_chain`, held across the IdP call to prevent
/// refresh-token reuse. Must exceed the external-auth refresh timeout
/// (`EXTERNAL_AUTH_REFRESH_TIMEOUT`, 5 s) so followers wait rather than retry.
/// refresh-token reuse. Must exceed the external-auth refresh budget
/// (a single 7s run) so followers wait rather than retry.
const REFRESH_LOCK_TIMEOUT: StdDuration = StdDuration::from_secs(45);
/// Long poll interval used by the proactive refresh task when no
@ -977,7 +980,7 @@ impl AuthManager {
/// Used by [`ModelsManager`] to trigger model catalog recovery
/// after sleep/wake, bypassing the FSEvents file watcher which
/// can silently die on macOS after resume.
pub fn refresh_notifier(&self) -> Arc<tokio::sync::Notify> {
pub(crate) fn refresh_notifier(&self) -> Arc<tokio::sync::Notify> {
self.refresh_notify.clone()
}
@ -997,7 +1000,7 @@ impl AuthManager {
/// to the primary refresh path instead of driving their own
/// `ServerRejected` recovery, avoiding concurrent refresh storms that
/// amplify 401 bursts at CCP.
pub async fn wait_for_token_refresh(&self, timeout: std::time::Duration) -> bool {
pub(crate) async fn wait_for_token_refresh(&self, timeout: std::time::Duration) -> bool {
let pre_key = self.current().map(|a| a.key.clone());
tokio::select! {
_ = self.refresh_notify.notified() => {}
@ -1362,6 +1365,9 @@ impl AuthManager {
// Snapshot inner ONCE for dispatch atomicity (closes a TOCTOU
// where a concurrent `clear()` raced `token_type()` + `inner.read()`).
let snapshot: Option<GrokAuth> = self.with_inner_read(|inner| inner.cloned());
// Kept alongside `snapshot`, which the grace arm below consumes: the
// devbox arms still need to name the credential they gave up on.
let snapshot_key: Option<String> = snapshot.as_ref().map(|a| a.key.clone());
let token_type = TokenType::from_auth(snapshot.as_ref());
tracing::Span::current().record("token_type", tracing::field::debug(token_type));
@ -1394,7 +1400,7 @@ impl AuthManager {
// preferred_method=api_key forbids automatic OIDC mint.
if !self.grok_com_config.blocks_automatic_oidc()
&& self.is_devbox_environment()
&& let Ok(auth) = self.try_devbox_recovery().await
&& let Ok(auth) = self.try_devbox_recovery(snapshot_key.as_deref()).await
{
return Ok(auth);
}
@ -1469,7 +1475,7 @@ impl AuthManager {
if result.is_err()
&& !self.grok_com_config.blocks_automatic_oidc()
&& self.is_devbox_environment()
&& let Ok(auth) = self.try_devbox_recovery().await
&& let Ok(auth) = self.try_devbox_recovery(snapshot_key.as_deref()).await
{
return Ok(auth);
}
@ -1500,7 +1506,15 @@ impl AuthManager {
///
/// Fail-closed under `preferred_method=api_key` (no automatic OIDC mint),
/// including direct callers such as sampler 401 recovery.
pub(crate) async fn try_devbox_recovery(self: &Arc<Self>) -> Result<GrokAuth, AuthError> {
///
/// `unusable` is the credential the caller has already established cannot
/// work — the bearer the server rejected, or the snapshot that failed to
/// refresh. It is what makes the wait-on-the-lock double-check below mean
/// "somebody else fixed this" instead of "the dead token is still here".
pub(crate) async fn try_devbox_recovery(
self: &Arc<Self>,
unusable: Option<&str>,
) -> Result<GrokAuth, AuthError> {
if self.grok_com_config.blocks_automatic_oidc() {
tracing::debug!(
"auth: devbox recovery skipped (preferred_method=api_key blocks automatic OIDC)"
@ -1515,8 +1529,15 @@ impl AuthManager {
let _guard = self.refresh_lock.lock().await;
// Double-check: another task may have recovered while we waited.
if let Some(auth) = self.current() {
// Double-check: another task may have recovered while we waited. Only
// a credential that is not the caller's `unusable` one counts. Without
// that filter a 401 on a still-locally-valid bearer reports recovery
// with the very token the server just rejected, and the caller
// resubmits it until its retry budget runs out.
if let Some(auth) = self
.current()
.filter(|auth| unusable != Some(auth.key.as_str()))
{
return Ok(auth);
}
@ -2169,10 +2190,8 @@ impl AuthManager {
/// decision.
pub(crate) fn requires_manual_reauth(&self) -> bool {
use crate::auth::error::RefreshTokenError;
// Sticky IdP rejection of the credential a refresh would send:
// no retry can fix it.
if let Some(AuthError::Refresh(RefreshTokenError::Permanent(e))) = self.permanent_failure()
&& e.reason.is_sticky()
&& e.reason.blocks_unattended_retry()
{
return true;
}
@ -2190,6 +2209,11 @@ impl AuthManager {
!(mem_refreshable || disk_refreshable)
}
fn is_external_provider_refresh_authority(&self) -> bool {
self.grok_com_config.auth_provider_command.is_some()
&& self.token_type() == TokenType::ExternalBinary
}
/// `true` iff a [`TokenRefresher`] is wired in. `false` for static-key
/// or pre-`configure_refresher` managers.
pub(crate) fn has_refresher_attached(&self) -> bool {
@ -2678,7 +2702,7 @@ impl AuthManager {
}
/// Set the process model key (empty clears). Not for session tokens.
pub fn set_process_static_api_key(&self, key: Option<String>) {
pub(crate) fn set_process_static_api_key(&self, key: Option<String>) {
let key = key.map(|k| k.trim().to_string()).filter(|k| !k.is_empty());
*self.process_static_api_key.write() = key;
}

322
crates/codegen/xai-grok-shell/src/auth/manager/remedy.rs Normal file
View file

@ -0,0 +1,322 @@
//! What it takes to get a session back to a usable credential, and the one
//! bounded unattended attempt the startup paths make before asking the user.
use std::sync::Arc;
use super::AuthManager;
use crate::auth::model::GrokAuth;
/// The way back to a usable credential, as of right now.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum AuthRemedy {
/// A later unattended refresh can still succeed — in the field, almost
/// always a launch seconds after wake, before the network is up.
SelfHealing,
/// Only an interactive run of the operator's auth provider can mint one.
ProviderLogin { label: Option<String> },
/// Only a user-driven login can.
ManualLogin,
}
impl AuthRemedy {
pub(crate) fn is_self_healing(&self) -> bool {
matches!(self, Self::SelfHealing)
}
/// `error_type` for a turn that died on this credential.
pub(crate) fn turn_error_type(&self) -> &'static str {
match self {
Self::SelfHealing => "auth_transient",
Self::ProviderLogin { .. } | Self::ManualLogin => "auth",
}
}
/// The same remedy, for a turn that has already spent its automatic
/// retries. [`Self::SelfHealing`] cannot survive that: its whole message
/// is "retry in a few seconds", which is exactly what just failed several
/// times over. What is left is a plain re-authentication — no advice of
/// our own, and classified so the client offers its own way back.
pub(crate) fn after_retries_exhausted(self) -> Self {
match self {
Self::SelfHealing => Self::ManualLogin,
provider_or_manual => provider_or_manual,
}
}
/// What to tell the user beyond the failure itself.
pub(crate) fn advice(&self) -> Option<String> {
match self {
Self::SelfHealing => Some(
"Authentication is temporarily unavailable (often a network blip right \
after wake). Your session is still signed in and will recover \
automatically — retry in a few seconds; no need to run /login."
.to_owned(),
),
Self::ProviderLogin { label } => {
Some(crate::auth::error::provider_login_message(label.as_deref()).into_owned())
}
Self::ManualLogin => None,
}
}
}
/// What a [`AuthManager::silent_refresh`] attempt leaves the caller holding.
#[derive(Debug, Clone)]
pub(crate) enum SilentRefresh {
/// The credential [`AuthManager::auth`] vouched for — the one the next
/// request would carry.
///
/// Carried, not re-read: `auth()` also succeeds on its grace arm, serving a
/// token that is still wire-valid but inside the early-invalidation buffer,
/// and [`AuthManager::current`] hides exactly that token. A caller that
/// answered `Renewed` with `current()` would reject the session this
/// outcome just accepted — and disagree with the `Failed(SelfHealing)` arm
/// on the very same credential.
Renewed(Box<GrokAuth>),
Failed(AuthRemedy),
}
impl AuthManager {
/// Attempt one unattended refresh, bounded because the caller's response
/// gates the client's first draw.
///
/// Spawned rather than awaited inline: dropping the future at the deadline
/// abandons an IdP exchange whose rotated refresh token the server may
/// already have burned, which is how a suspend mid-refresh revoked whole
/// token families in the field.
pub(crate) async fn silent_refresh(self: &Arc<Self>) -> SilentRefresh {
let manager = Arc::clone(self);
let attempt = tokio::spawn(async move { manager.auth().await });
let outcome =
match tokio::time::timeout(crate::http::STARTUP_AUTH_REFRESH_TIMEOUT, attempt).await {
Ok(Ok(Ok(auth))) => SilentRefresh::Renewed(Box::new(auth)),
_ => SilentRefresh::Failed(self.auth_remedy()),
};
// The variant, not the outcome: the `Renewed` payload is a credential.
let logged = match &outcome {
SilentRefresh::Renewed(_) => "Renewed".to_owned(),
SilentRefresh::Failed(remedy) => format!("Failed({remedy:?})"),
};
xai_grok_telemetry::unified_log::info(
"auth: silent refresh",
None,
Some(serde_json::json!({ "outcome": logged })),
);
outcome
}
/// Classify the current credential's way back.
///
/// The provider arm deliberately ignores the recorded verdict: real
/// interactive-only binaries block until something kills them, so their
/// run routinely ends with nothing recorded at all.
pub(crate) fn auth_remedy(&self) -> AuthRemedy {
let provider_mints_sessions = self.is_external_provider_refresh_authority();
let user_must_act = self.requires_manual_reauth()
|| (provider_mints_sessions && self.current_wire_valid().is_none());
match (user_must_act, provider_mints_sessions) {
(false, _) => AuthRemedy::SelfHealing,
(true, true) => AuthRemedy::ProviderLogin {
label: self.grok_com_config().auth_provider_label.clone(),
},
(true, false) => AuthRemedy::ManualLogin,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::model::AuthMode;
use crate::auth::refresh::{RefreshOutcome, TokenRefresher};
use crate::auth::{GrokComConfig, error::RefreshTokenFailedReason};
use chrono::{Duration, Utc};
fn external_provider_config() -> GrokComConfig {
GrokComConfig {
auth_provider_command: Some("acme-auth".to_owned()),
auth_provider_label: Some("Acme SSO".to_owned()),
..GrokComConfig::default()
}
}
fn external_credential(expires_at: chrono::DateTime<Utc>) -> GrokAuth {
GrokAuth {
key: "external".into(),
auth_mode: AuthMode::External,
expires_at: Some(expires_at),
..GrokAuth::test_default()
}
}
/// Wired as production wires it, so nothing here passes on the
/// "no refresh authority" arm.
fn provider_manager(dir: &std::path::Path, credential: GrokAuth) -> Arc<AuthManager> {
let config = external_provider_config();
let command = config.auth_provider_command.clone();
let manager = Arc::new(AuthManager::new(dir, config));
manager.hot_swap(credential);
manager.configure_refresher(command, None);
manager
}
/// The verdict-free arm.
#[test]
fn hard_expired_external_credential_needs_the_provider_without_a_verdict() {
let dir = tempfile::tempdir().unwrap();
let manager = provider_manager(
dir.path(),
external_credential(Utc::now() - Duration::hours(1)),
);
assert!(!manager.has_permanent_failure());
assert_eq!(
manager.auth_remedy(),
AuthRemedy::ProviderLogin {
label: Some("Acme SSO".to_owned())
}
);
}
/// A bare-token credential the backend rejects: it never expires locally,
/// so only the verdict from the failed run says the user has to act.
#[test]
fn wire_valid_external_credential_needs_the_provider_once_its_run_failed() {
let dir = tempfile::tempdir().unwrap();
let manager = provider_manager(
dir.path(),
external_credential(Utc::now() + Duration::hours(1)),
);
assert_eq!(manager.auth_remedy(), AuthRemedy::SelfHealing);
manager.record_permanent_failure(
"external".to_owned(),
RefreshTokenFailedReason::ProviderInteractiveRequired.into(),
);
assert_eq!(
manager.auth_remedy(),
AuthRemedy::ProviderLogin {
label: Some("Acme SSO".to_owned())
}
);
}
/// Inside the early-invalidation buffer the proxy still accepts the token,
/// so a failed refresh must not cost the user a login.
#[test]
fn buffer_window_external_credential_is_self_healing() {
let dir = tempfile::tempdir().unwrap();
let manager = provider_manager(
dir.path(),
external_credential(Utc::now() + Duration::minutes(1)),
);
assert!(manager.current().is_none(), "buffer window hides the token");
assert_eq!(manager.auth_remedy(), AuthRemedy::SelfHealing);
}
/// A provider command configured alongside OIDC must not capture OIDC's
/// own refresh path.
#[test]
fn expired_oidc_credential_with_a_refresh_token_is_self_healing() {
let dir = tempfile::tempdir().unwrap();
let manager = provider_manager(
dir.path(),
GrokAuth {
key: "expired-oidc".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-live".into()),
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
},
);
assert_eq!(manager.auth_remedy(), AuthRemedy::SelfHealing);
}
/// With no provider command there is no binary to escalate to.
#[test]
fn expired_credential_without_a_provider_command_needs_a_manual_login() {
let dir = tempfile::tempdir().unwrap();
let manager = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
manager.hot_swap(external_credential(Utc::now() - Duration::hours(1)));
manager.configure_refresher(None, None);
assert_eq!(manager.auth_remedy(), AuthRemedy::SelfHealing);
manager.record_permanent_failure(
"external".to_owned(),
RefreshTokenFailedReason::ProviderInteractiveRequired.into(),
);
assert_eq!(manager.auth_remedy(), AuthRemedy::ManualLogin);
}
/// A refresh that fails over a token still inside the early-invalidation
/// buffer is a *success* for [`AuthManager::silent_refresh`]: `auth()`
/// serves the cached bearer the proxy still accepts. `current()` hides that
/// token, so `Renewed` must carry the credential — a caller re-reading
/// `current()` here would reject a session that `Failed(SelfHealing)`, on
/// this very credential, would have accepted via `current_or_expired()`.
#[tokio::test]
async fn renewed_carries_the_wire_valid_bearer_the_buffer_hides() {
struct OfflineRefresher;
#[async_trait::async_trait]
impl TokenRefresher for OfflineRefresher {
async fn refresh(
&self,
_reason: crate::auth::manager::RefreshReason,
) -> RefreshOutcome {
RefreshOutcome::transient("network unreachable")
}
}
let dir = tempfile::tempdir().unwrap();
let manager = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
// CI runs in pods where `is_devbox_environment()` is true; a mint would
// resolve the credential for the wrong reason.
manager.set_devbox_env_for_test(false);
// A minute from real expiry: inside the 5-min buffer, still on the wire.
manager.hot_swap(GrokAuth {
key: "wire-valid".into(),
auth_mode: AuthMode::Oidc,
refresh_token: Some("rt-live".into()),
expires_at: Some(Utc::now() + Duration::minutes(1)),
..GrokAuth::test_default()
});
manager.set_refresher(Arc::new(OfflineRefresher));
assert!(manager.current().is_none(), "the buffer hides the token");
assert!(manager.is_expired(), "and reports the session expired");
let SilentRefresh::Renewed(auth) = manager.silent_refresh().await else {
panic!("auth()'s grace arm serves the wire-valid bearer");
};
assert_eq!(auth.key, "wire-valid");
assert!(
manager.current().is_none(),
"current() still hides it — which is why the outcome carries it",
);
assert!(
manager.auth_remedy().is_self_healing(),
"and the other arm would have accepted the same credential",
);
}
#[test]
fn turn_surface_matches_the_remedy() {
assert_eq!(AuthRemedy::SelfHealing.turn_error_type(), "auth_transient");
assert!(
AuthRemedy::SelfHealing
.advice()
.is_some_and(|a| a.contains("no need to run /login"))
);
assert_eq!(AuthRemedy::ManualLogin.turn_error_type(), "auth");
assert_eq!(
AuthRemedy::ManualLogin.advice(),
None,
"the client's own banner already tells the user to run /login"
);
let provider = AuthRemedy::ProviderLogin {
label: Some("Acme SSO".to_owned()),
};
assert_eq!(provider.turn_error_type(), "auth");
let advice = provider.advice().expect("provider advice");
assert!(advice.contains("Acme SSO") && advice.contains("/login"));
assert!(!advice.contains("no need to run /login"));
}
}

116
crates/codegen/xai-grok-shell/src/auth/manager_tests.rs
View file

@ -4706,6 +4706,12 @@ fn manual_auth_reason_maps_terminal_and_skips_non_forcing() {
}),
Some(R::WrongTeam)
);
// Before this reason existed these lockouts hid under the self-healing
// `Other` bucket and never surfaced in the KPI at all.
assert_eq!(
permanent(Reason::ProviderInteractiveRequired),
Some(R::ProviderInteractiveRequired)
);
// Self-healing (TTL) reasons, transient / no-credential, and API-key
// lockouts (out of scope for this KPI) don't count.
assert_eq!(permanent(Reason::ClientRejected), None);
@ -4739,6 +4745,10 @@ fn relay_should_cancel_gives_up_only_on_terminal_failures() {
// Cancelled even though it never emits the KPI (a kill-switched API key
// means rotate the key, not `/login`).
assert!(relay_should_cancel(&AuthError::ApiKeyAuthDisabled));
// Reconnecting would replay the same 401 until the user signs in.
assert!(relay_should_cancel(&AuthError::permanent(
Reason::ProviderInteractiveRequired
)));
// Recoverable: fall through and reconnect.
assert!(!relay_should_cancel(&AuthError::transient("network blip")));
@ -4942,6 +4952,52 @@ async fn requires_manual_reauth_true_for_sticky_verdict_and_no_refresher() {
);
}
/// Treating a failed provider run as self-healing is what let an expired
/// credential in and then 401'd every turn. The verdict still ages out, so a
/// later launch gets to retry the provider.
#[tokio::test]
async fn requires_manual_reauth_true_after_external_provider_refresh_failed() {
let dir = tempfile::tempdir().unwrap();
let mgr = Arc::new(AuthManager::new(dir.path(), external_provider_config()));
mgr.hot_swap(GrokAuth {
key: "expired-external".into(),
auth_mode: AuthMode::External,
expires_at: Some(Utc::now() - Duration::hours(1)),
..GrokAuth::test_default()
});
mgr.set_refresher(Arc::new(FailingRefresher {
call_count: Arc::new(AtomicU32::new(0)),
}));
assert!(
!mgr.requires_manual_reauth(),
"before any attempt the provider may still mint silently"
);
record_permanent_failure(
&mgr,
crate::auth::error::RefreshTokenFailedReason::ProviderInteractiveRequired,
);
assert!(
mgr.requires_manual_reauth(),
"a failed headless provider run leaves only the interactive flow"
);
mgr.force_permanent_failure_aged_out();
assert!(
!mgr.requires_manual_reauth(),
"the verdict is non-sticky: past its TTL the provider gets another chance"
);
}
/// Config for a deployment that mints sessions with an external binary.
fn external_provider_config() -> GrokComConfig {
GrokComConfig {
auth_provider_command: Some("acme-auth".to_owned()),
..GrokComConfig::default()
}
}
// ── proactive_failure_backoff ────────────────────────────────────────
/// The proactive loop's failure backoff: zero before any failure (schedule is
@ -4972,3 +5028,63 @@ fn proactive_failure_backoff_shape() {
"backoff must cap at BACKOFF_INTERVAL (+jitter), got {huge:?}"
);
}
// ── try_devbox_recovery: the wait-on-the-lock double-check ───────────
/// Seed a credential that is locally valid but that the caller has been told
/// the server rejects — the shape that made the double-check lie.
fn devbox_manager(dir: &std::path::Path, key: &str) -> Arc<AuthManager> {
let mgr = Arc::new(AuthManager::new(dir, GrokComConfig::default()));
mgr.set_devbox_env_for_test(true);
mgr.hot_swap(GrokAuth {
key: key.into(),
auth_mode: AuthMode::External,
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
mgr
}
/// The credential the caller already knows is dead can never be the answer.
///
/// `try_devbox_recovery` short-circuits on whatever `current()` holds, to
/// catch a sibling task that refreshed while we waited on `refresh_lock`.
/// Told nothing about the rejected bearer it used to return that bearer, so
/// on a devbox every 401 against a still-locally-valid token reported
/// "recovered" and the turn resubmitted it until its retry budget ran out.
#[tokio::test]
async fn devbox_recovery_never_re_serves_the_credential_it_was_given_up_on() {
let dir = tempfile::tempdir().unwrap();
let mgr = devbox_manager(dir.path(), "rejected-but-locally-valid");
assert!(
mgr.current().is_some(),
"precondition: the rejected bearer is still locally valid"
);
// Asserted as "not this credential" rather than as an error: on a real
// devbox the mint can genuinely succeed, and a *different* credential is
// exactly the outcome we want. Everywhere else there is no mint endpoint
// and this is an error.
let outcome = mgr
.try_devbox_recovery(Some("rejected-but-locally-valid"))
.await;
assert!(
!matches!(&outcome, Ok(auth) if auth.key == "rejected-but-locally-valid"),
"recovery must not report success with the rejected bearer, got {outcome:?}"
);
}
/// The double-check still does its job: a credential that is *not* the one
/// the caller gave up on means a sibling task refreshed, so take it and skip
/// the mint.
#[tokio::test]
async fn devbox_recovery_short_circuits_on_a_credential_someone_else_landed() {
let dir = tempfile::tempdir().unwrap();
let mgr = devbox_manager(dir.path(), "landed-by-a-sibling-task");
let auth = mgr
.try_devbox_recovery(Some("the-bearer-the-server-rejected"))
.await
.expect("a different live credential is a recovery");
assert_eq!(auth.key, "landed-by-a-sibling-task");
}

3
crates/codegen/xai-grok-shell/src/auth/mod.rs
View file

@ -4,7 +4,7 @@ mod config;
pub mod credential_provider;
#[path = "devbox_login_stub.rs"]
pub(crate) mod devbox_login;
pub mod device_code;
pub(crate) mod device_code;
pub mod error;
mod external_auth;
mod flow;
@ -43,6 +43,7 @@ pub use jwt::{is_jwt_expired_or_near, parse_jwt_expiration};
mod meta;
pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
pub use manager::{AuthManager, shared_api_key_provider};
pub(crate) use manager::{AuthRemedy, SilentRefresh};
pub use meta::{AuthMeta, GateInfo};
pub use model::{AuthMode, GrokAuth, lookup_auth};
pub(crate) use model::{

6
crates/codegen/xai-grok-shell/src/auth/model.rs
View file

@ -161,7 +161,7 @@ impl GrokAuth {
/// qualify; external-provider credentials qualify only when first-party
/// (`is_xai_auth`), matching the built-in devbox login they replace.
/// Plain API keys never do.
pub fn is_session_auth(&self) -> bool {
pub(crate) fn is_session_auth(&self) -> bool {
match self.auth_mode {
AuthMode::WebLogin | AuthMode::Oidc => true,
AuthMode::External => self.is_xai_auth(),
@ -184,7 +184,7 @@ impl GrokAuth {
/// retention. Use this for trace-upload and research-data gates.
/// Product analytics (`telemetry_enabled`) and user-facing sync
/// features should use `is_zdr_team()` directly.
pub fn is_data_collection_disabled(&self) -> bool {
pub(crate) fn is_data_collection_disabled(&self) -> bool {
self.is_zdr_team() || self.coding_data_retention_opt_out
}
@ -244,7 +244,7 @@ impl GrokAuth {
/// ```ignore
/// GrokAuth { key: "my-key".into(), ..GrokAuth::test_default() }
/// ```
pub fn test_default() -> Self {
pub(crate) fn test_default() -> Self {
Self {
key: "test-key".into(),
user_id: "test-user".into(),

2
crates/codegen/xai-grok-shell/src/auth/oidc/protocol.rs
View file

@ -94,7 +94,7 @@ pub(crate) fn with_alpha_test_key(
let _ = url;
builder
}
pub fn is_configured(config: &GrokComConfig) -> bool {
pub(crate) fn is_configured(config: &GrokComConfig) -> bool {
config.oidc.is_some()
}
/// Peek at the unverified access token JWT to extract the `principal_type`

8
crates/codegen/xai-grok-shell/src/auth/recovery.rs
View file

@ -29,6 +29,7 @@ pub(crate) fn manual_auth_reason(err: &AuthError) -> Option<ManualAuthReason> {
Some(match err {
AuthError::Refresh(RefreshTokenError::Permanent(e)) => match e.reason {
RefreshTokenFailedReason::RefreshTokenRejected => R::RefreshTokenRejected,
RefreshTokenFailedReason::ProviderInteractiveRequired => R::ProviderInteractiveRequired,
// Self-healing via the TTL, not a manual re-auth.
RefreshTokenFailedReason::ClientRejected | RefreshTokenFailedReason::Other => {
return None;
@ -217,7 +218,7 @@ enum RecoveryStep {
}
/// State machine that walks through recovery strategies after a 401.
pub struct UnauthorizedRecovery {
pub(crate) struct UnauthorizedRecovery {
auth_manager: Arc<AuthManager>,
/// The token that was rejected by the server.
rejected_token: String,
@ -324,7 +325,10 @@ impl UnauthorizedRecovery {
// preferred_method=api_key forbids automatic OIDC mint.
if !self.auth_manager.grok_com_config().blocks_automatic_oidc()
&& self.auth_manager.is_devbox_environment()
&& let Ok(auth) = self.auth_manager.try_devbox_recovery().await
&& let Ok(auth) = self
.auth_manager
.try_devbox_recovery(Some(&self.rejected_token))
.await
{
return Ok(auth);
}

68
crates/codegen/xai-grok-shell/src/auth/refresh/external_refresher.rs
View file

@ -20,15 +20,13 @@ impl ExternalBinaryRefresher {
Self { runner, command }
}
/// A failed or timed-out binary run is a single-strike `Other` permanent
/// failure. `Other` is non-sticky, so `PERMANENT_FAILURE_TTL` lets a flaky
/// or briefly slow binary self-heal without `/login`. The async runner
/// bounds every run and group-kills the child on timeout, so there is no
/// wedged-process case that would need a separate transient outcome.
/// A failed or timed-out binary run is a single-strike permanent failure;
/// the reason is non-sticky so a flaky or briefly slow binary still
/// recovers without the user.
fn record_failure(&self, message: &str) -> RefreshOutcome {
tracing::warn!(%message, "auth: external binary refresh failed permanently");
// No token key in the binary flow; the caller scopes the verdict.
RefreshOutcome::permanent(RefreshTokenFailedReason::Other, None)
RefreshOutcome::permanent(RefreshTokenFailedReason::ProviderInteractiveRequired, None)
}
}
@ -57,33 +55,49 @@ impl TokenRefresher for ExternalBinaryRefresher {
mod tests {
use super::*;
use crate::auth::GrokAuth;
use std::sync::Mutex;
use std::sync::atomic::{AtomicU32, Ordering};
/// Minimal runner whose external command returns a fixed result.
/// Runner that yields scripted results in order, then `None`.
struct FakeRunner {
external_result: Option<GrokAuth>,
results: Mutex<Vec<Option<GrokAuth>>>,
calls: AtomicU32,
}
impl FakeRunner {
fn new(results: Vec<Option<GrokAuth>>) -> Self {
Self {
results: Mutex::new(results),
calls: AtomicU32::new(0),
}
}
fn calls(&self) -> u32 {
self.calls.load(Ordering::SeqCst)
}
}
#[async_trait::async_trait]
impl ExternalCommandRunner for FakeRunner {
async fn run_external_command(&self, _command: &str) -> Option<GrokAuth> {
self.external_result.clone()
self.calls.fetch_add(1, Ordering::SeqCst);
let mut results = self.results.lock().unwrap();
if results.is_empty() {
return None;
}
results.remove(0)
}
}
/// A failed binary run is a single-strike `Other` permanent failure that is
/// NON-sticky: it must age out via the TTL, never lock an external-binary
/// user out forever. (Flipping this to a sticky reason would be a silent
/// lockout regression.)
/// A failed binary run must stay NON-sticky: it has to age out via the
/// TTL, never lock an external-binary user out forever.
#[tokio::test]
async fn external_binary_failure_is_single_strike_non_sticky_permanent() {
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: None,
}),
"auth-binary".into(),
);
let runner = Arc::new(FakeRunner::new(vec![None]));
let refresher = ExternalBinaryRefresher::new(runner.clone(), "auth-binary".into());
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::PermanentFailure { error, .. } => {
assert_eq!(error.reason, RefreshTokenFailedReason::Other);
assert_eq!(
error.reason,
RefreshTokenFailedReason::ProviderInteractiveRequired
);
assert!(
!error.reason.is_sticky(),
"external-binary failure must age out, not strand the user forever",
@ -91,6 +105,7 @@ mod tests {
}
other => panic!("a failed binary run must be a permanent Other failure, got {other:?}"),
}
assert_eq!(runner.calls(), 1, "the single run gets the whole 7s budget");
}
#[tokio::test]
@ -99,15 +114,16 @@ mod tests {
key: "ext-fresh".into(),
..GrokAuth::test_default()
};
let refresher = ExternalBinaryRefresher::new(
Arc::new(FakeRunner {
external_result: Some(token),
}),
"auth-binary".into(),
);
let runner = Arc::new(FakeRunner::new(vec![Some(token)]));
let refresher = ExternalBinaryRefresher::new(runner.clone(), "auth-binary".into());
match refresher.refresh(RefreshReason::ServerRejected).await {
RefreshOutcome::Success(auth) => assert_eq!(auth.key, "ext-fresh"),
other => panic!("a successful binary run must return Success, got {other:?}"),
}
assert_eq!(
runner.calls(),
1,
"a success must run the binary exactly once"
);
}
}

224
crates/codegen/xai-grok-shell/src/bin/trace_classify.rs
View file

@ -1,224 +0,0 @@
//! Replay an offline session trace against the Layer-2 TodoGate and
//! Layer-3 LazinessDetector classifier, emitting one JSONL line per
//! turn.
//!
//! Usage:
//! cargo run --bin trace_classify -- \
//! --trace /path/to/trace-<id>-all-turns.json \
//! [--output out.jsonl] \
//! [--model grok-4.5] \
//! [--api-base-url https://api.x.ai/v1] \
//! [--api-key <key> | $XAI_API_KEY | <grok-home>/auth.json] \
//! [--min-confidence 0.7] \
//! [--include-reasoning true] \
//! [--grok-home <path>]
//!
//! The binary name is `trace_classify` (underscore) — that's the file
//! name in `src/bin/`, which cargo's auto-discovery uses verbatim.
//! The task brief calls it `trace-classify` (hyphen) in prose; the
//! canonical CLI invocation is the underscore form.
//!
//! Each JSONL line carries the per-turn gate decision, the parsed
//! classifier verdict (or the abort/parse error if the call failed),
//! and the inputs that drove them.
use std::path::PathBuf;
use clap::Parser;
use xai_grok_shell::trace_classifier::{RunArgs, run, validate_min_confidence};
#[derive(Parser, Debug)]
#[command(
name = "trace_classify",
about = "Replay a session trace against the TodoGate + Laziness classifier"
)]
struct Cli {
/// Path to the offline trace JSON (a top-level array of turn records).
#[arg(long)]
trace: PathBuf,
/// Write JSONL output here (one line per turn). Defaults to stdout
/// when omitted.
#[arg(long)]
output: Option<PathBuf>,
/// Model the classifier sampler calls. Must be a model the API key
/// has access to.
#[arg(long, default_value = "grok-4.5")]
model: String,
/// Sampler base URL.
#[arg(long, default_value = "https://api.x.ai/v1")]
api_base_url: String,
/// API key. Overrides `$XAI_API_KEY` when set; falls back to
/// `$XAI_API_KEY`, then `<grok-home>/auth.json` (`xai::api_key`
/// scope) when absent or empty.
#[arg(long)]
api_key: Option<String>,
/// Override the LazinessDetector min-confidence threshold (default
/// matches production's `LAZINESS_DEFAULT_MIN_CONFIDENCE`). Must
/// be a finite float in `[0.0, 1.0]`. Use this to mirror a
/// per-model override from the production models catalog. (F6/N5)
#[arg(long, value_parser = validate_min_confidence)]
min_confidence: Option<f32>,
/// Override the harness `[assistant reasoning]` emission flag.
/// When absent (the default), the binary uses the harness default
/// `LAZINESS_INCLUDE_REASONING`. Accepts `true` / `false`. The
/// offline replay tool has no per-model config to consult, so
/// this is the only override surface here — production resolves
/// `LazinessDetectorPerModelConfig::include_reasoning` separately.
#[arg(long)]
include_reasoning: Option<bool>,
/// Override the directory containing `auth.json` for the
/// third-tier API-key fallback. Defaults to the same path the
/// shell uses (`$GROK_HOME` or `~/.grok`). Exposed primarily for
/// tests / sandboxed invocations.
#[arg(long)]
grok_home: Option<PathBuf>,
}
/// `current_thread` flavour: the replay is strictly sequential
/// (one turn at a time), and a multi-threaded runtime would force
/// every writer (including `StdoutLock`) to be `Send` — which it
/// isn't. The sequential nature also means we never schedule work in
/// parallel, so `current_thread` is the right cost shape too.
#[tokio::main(flavor = "current_thread")]
async fn main() -> anyhow::Result<()> {
let cli = Cli::parse();
let args = RunArgs {
trace: cli.trace,
output: cli.output,
model_id: cli.model,
api_base_url: cli.api_base_url,
api_key: cli.api_key,
min_confidence: cli.min_confidence,
include_reasoning: cli.include_reasoning,
grok_home: cli.grok_home,
};
let summary = run(args).await?;
eprintln!("{}", summary.render());
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use clap::CommandFactory;
#[test]
fn cli_parses_minimal_args() {
let cli = Cli::try_parse_from(["trace_classify", "--trace", "foo.json", "--model", "bar"])
.expect("parse");
assert_eq!(cli.trace, PathBuf::from("foo.json"));
assert_eq!(cli.model, "bar");
assert_eq!(cli.api_base_url, "https://api.x.ai/v1");
assert!(cli.output.is_none());
assert!(cli.api_key.is_none());
assert!(cli.min_confidence.is_none());
assert!(cli.include_reasoning.is_none());
assert!(cli.grok_home.is_none());
}
/// Per-model knob (mirrored as a CLI override on the offline tool):
/// `--include-reasoning true` and `--include-reasoning false` both
/// parse; absent → `None` so the harness default applies.
#[test]
fn cli_include_reasoning_override_parses() {
let cli_true = Cli::try_parse_from([
"trace_classify",
"--trace",
"foo.json",
"--include-reasoning",
"true",
])
.expect("parse true");
assert_eq!(cli_true.include_reasoning, Some(true));
let cli_false = Cli::try_parse_from([
"trace_classify",
"--trace",
"foo.json",
"--include-reasoning",
"false",
])
.expect("parse false");
assert_eq!(cli_false.include_reasoning, Some(false));
let cli_absent =
Cli::try_parse_from(["trace_classify", "--trace", "foo.json"]).expect("parse absent");
assert!(cli_absent.include_reasoning.is_none());
}
#[test]
fn cli_grok_home_override_parses() {
let cli = Cli::try_parse_from([
"trace_classify",
"--trace",
"foo.json",
"--grok-home",
"/tmp/scratch-grok",
])
.expect("parse");
assert_eq!(cli.grok_home, Some(PathBuf::from("/tmp/scratch-grok")));
}
#[test]
fn cli_requires_trace() {
let err = Cli::try_parse_from(["trace_classify"]).expect_err("missing --trace");
let msg = err.to_string();
assert!(msg.contains("--trace"), "error mentions --trace: {msg}");
}
/// F18 — assert the documented defaults actually take effect.
#[test]
fn cli_defaults_match_documented_values() {
let cmd = Cli::command();
let by_id = |id: &str| {
cmd.get_arguments()
.find(|a| a.get_id().as_str() == id)
.unwrap_or_else(|| panic!("arg {id} missing"))
.get_default_values()
.iter()
.map(|v| v.to_string_lossy().into_owned())
.collect::<Vec<_>>()
};
assert_eq!(by_id("model"), vec!["grok-4.5"]);
assert_eq!(by_id("api_base_url"), vec!["https://api.x.ai/v1"]);
assert!(by_id("min_confidence").is_empty(), "no default");
assert!(by_id("include_reasoning").is_empty(), "no default");
}
/// F6 — `--min-confidence 0.5` parses and lands in `RunArgs`.
#[test]
fn cli_min_confidence_override_parses() {
let cli = Cli::try_parse_from([
"trace_classify",
"--trace",
"foo.json",
"--min-confidence",
"0.42",
])
.expect("parse");
assert_eq!(cli.min_confidence, Some(0.42));
}
/// N5 — clap `value_parser` rejects out-of-range / non-finite
/// floats at parse time, before they reach `RunArgs`. Bad values
/// are passed via `--min-confidence=VALUE` syntax so negative
/// literals aren't mis-parsed as short flags.
#[test]
fn cli_min_confidence_rejects_bad_values() {
for bad in ["1.5", "-0.1", "nan", "inf", "not-a-float"] {
let arg = format!("--min-confidence={bad}");
let err = Cli::try_parse_from(["trace_classify", "--trace", "foo.json", arg.as_str()])
.expect_err(bad);
// Parsing failed — that's all we need. Exact error text
// is clap-version-dependent.
let _ = err.to_string();
}
}
}

23
crates/codegen/xai-grok-shell/src/bundle.rs
View file

@ -19,7 +19,7 @@ struct ArchiveBundleMetadata {
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct BundleManifest {
pub(crate) struct BundleManifest {
pub version: String,
pub checksums: HashMap<String, String>,
}
@ -83,11 +83,11 @@ struct BundleFile<'a> {
content: &'a str,
}
pub fn bundled_root() -> PathBuf {
pub(crate) fn bundled_root() -> PathBuf {
xai_grok_config::grok_home().join(BUNDLED_DIR_NAME)
}
pub fn read_cached_manifest(root: &Path) -> Result<Option<BundleManifest>> {
pub(crate) fn read_cached_manifest(root: &Path) -> Result<Option<BundleManifest>> {
let manifest_path = manifest_path(root);
let bytes = match std::fs::read(&manifest_path) {
Ok(bytes) => bytes,
@ -103,7 +103,10 @@ pub fn read_cached_manifest(root: &Path) -> Result<Option<BundleManifest>> {
.map(Some)
}
pub fn write_bundle_to_cache(root: &Path, bundle: &SubagentBundle) -> Result<BundleManifest> {
pub(crate) fn write_bundle_to_cache(
root: &Path,
bundle: &SubagentBundle,
) -> Result<BundleManifest> {
let old_manifest = read_cached_manifest(root)?.map(sanitize_manifest);
ensure_bundle_dirs(root)?;
@ -149,7 +152,7 @@ pub fn write_bundle_to_cache(root: &Path, bundle: &SubagentBundle) -> Result<Bun
Ok(next_manifest)
}
pub fn extract_bundle_archive(root: &Path, archive_bytes: &[u8]) -> Result<BundleManifest> {
pub(crate) fn extract_bundle_archive(root: &Path, archive_bytes: &[u8]) -> Result<BundleManifest> {
let decoder = flate2::read::GzDecoder::new(archive_bytes);
let mut archive = tar::Archive::new(decoder);
@ -257,7 +260,7 @@ pub fn extract_bundle_archive(root: &Path, archive_bytes: &[u8]) -> Result<Bundl
Ok(next_manifest)
}
pub fn checksum_bytes(bytes: &[u8]) -> String {
pub(crate) fn checksum_bytes(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
format!("{:x}", hasher.finalize())
@ -269,7 +272,7 @@ pub fn checksum_file(path: &Path) -> Result<String> {
Ok(checksum_bytes(&bytes))
}
pub fn prune_removed_files(
pub(crate) fn prune_removed_files(
root: &Path,
old_manifest: &BundleManifest,
retained_checksums: &mut HashMap<String, String>,
@ -411,7 +414,7 @@ fn map_archive_path_to_cache_path(archive_path: &str) -> Option<String> {
None
}
pub fn count_entries_by_prefix(manifest: &BundleManifest, prefix: &str) -> usize {
pub(crate) fn count_entries_by_prefix(manifest: &BundleManifest, prefix: &str) -> usize {
manifest
.checksums
.keys()
@ -467,7 +470,7 @@ fn validate_bundle_name(kind: BundleFileKind, name: &str) -> Result<()> {
#[cfg(test)]
pub(crate) mod test_helpers {
pub fn make_test_archive(entries: &[(&str, &[u8])]) -> Vec<u8> {
pub(crate) fn make_test_archive(entries: &[(&str, &[u8])]) -> Vec<u8> {
let encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
let mut builder = tar::Builder::new(encoder);
for &(path, content) in entries {
@ -481,7 +484,7 @@ pub(crate) mod test_helpers {
encoder.finish().unwrap()
}
pub fn bundle_json(version: &str) -> String {
pub(crate) fn bundle_json(version: &str) -> String {
format!(r#"{{"version":"{version}"}}"#)
}
}

8
crates/codegen/xai-grok-shell/src/claude_import.rs
View file

@ -536,7 +536,7 @@ static MARKER_CACHE: std::sync::RwLock<Option<bool>> = std::sync::RwLock::new(No
/// on the marker (so users see one log line indicating the cutoff fired).
/// Use the bare version for read-time display logic that already has its own
/// path (e.g. UI listings in `extensions/skills.rs` and `inspect.rs`).
pub fn is_claude_import_marked() -> bool {
pub(crate) fn is_claude_import_marked() -> bool {
if let Some(v) = *MARKER_CACHE.read().expect("MARKER_CACHE poisoned") {
return v;
}
@ -551,7 +551,7 @@ pub fn is_claude_import_marked() -> bool {
/// Called from the slash command after `apply_import` writes the marker so
/// that subsequent in-process gate checks reflect the new state without
/// waiting for restart.
pub fn refresh_marker_cache(value: bool) {
pub(crate) fn refresh_marker_cache(value: bool) {
*MARKER_CACHE.write().expect("MARKER_CACHE poisoned") = Some(value);
}
@ -570,7 +570,7 @@ pub(crate) fn reset_marker_cache_for_test() {
/// Call sites are runtime fallback paths in `claude_compat.rs`,
/// `util/config.rs`, `util/hooks.rs`, and `agent/config.rs` that previously
/// read `.claude/`.
pub fn is_claude_import_marked_with_log(gate_name: &'static str) -> bool {
pub(crate) fn is_claude_import_marked_with_log(gate_name: &'static str) -> bool {
static LOGGED: OnceLock<()> = OnceLock::new();
let marked = is_claude_import_marked();
if marked {
@ -585,7 +585,7 @@ pub fn is_claude_import_marked_with_log(gate_name: &'static str) -> bool {
}
/// Testable variant of [`is_claude_import_marked`] that reads from the given path.
pub fn is_claude_import_marked_at(config_path: &Path) -> bool {
pub(crate) fn is_claude_import_marked_at(config_path: &Path) -> bool {
let content = match std::fs::read_to_string(config_path) {
Ok(s) => s,
Err(_) => return false,

65
crates/codegen/xai-grok-shell/src/claude_import_state.rs
View file

@ -10,7 +10,7 @@ use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tracing::{debug, warn};
use tracing::warn;
use xai_grok_workspace::permission::claude_settings::find_claude_settings_paths;
@ -18,7 +18,7 @@ use xai_grok_workspace::permission::claude_settings::find_claude_settings_paths;
/// Persistent import state, loaded from / saved to `~/.grok/claude_import_state.json`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImportState {
pub(crate) struct ImportState {
/// Schema version for forward compatibility.
pub version: u32,
/// Hash of global Claude settings (`~/.claude/settings*.json`, `~/.claude.json`).
@ -31,7 +31,7 @@ pub struct ImportState {
/// Import state for a single scope (global or one project).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScopeState {
pub(crate) struct ScopeState {
/// SHA-256 hex digest of the concatenated Claude settings file contents.
pub last_hash: String,
/// RFC 3339 timestamp of when the hash was last recorded.
@ -56,7 +56,7 @@ fn state_path() -> PathBuf {
}
/// Load the import state from disk. Returns default if missing or unreadable.
pub fn load_import_state() -> ImportState {
pub(crate) fn load_import_state() -> ImportState {
let path = state_path();
match std::fs::read_to_string(&path) {
Ok(content) => serde_json::from_str(&content).unwrap_or_else(|e| {
@ -80,7 +80,7 @@ pub fn load_import_state() -> ImportState {
}
/// Save the import state to disk (atomic write via tmp + rename).
pub fn save_import_state(state: &ImportState) -> std::io::Result<()> {
pub(crate) fn save_import_state(state: &ImportState) -> std::io::Result<()> {
let path = state_path();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
@ -181,61 +181,6 @@ fn compute_project_hash(cwd: &Path) -> (String, Vec<PathBuf>) {
// Change Detection
/// Check if any Claude settings files have changed since the last import/dismiss.
///
/// Returns `true` if:
/// - Global settings exist and have a different hash than last recorded
/// - Project settings exist and have a different hash than last recorded
/// - No import state exists yet but Claude settings files are present
pub fn has_new_changes(cwd: &Path) -> bool {
let state = load_import_state();
// Check global scope.
let (global_hash, global_paths) = compute_global_hash();
let global_files_exist = global_paths.iter().any(|p| p.exists());
if global_files_exist {
match &state.global {
None => {
debug!("Claude import: global settings found, no previous import state");
return true;
}
Some(s) if s.last_hash != global_hash => {
debug!(
old = %s.last_hash,
new = %global_hash,
"Claude import: global settings changed since last import"
);
return true;
}
_ => {}
}
}
// Check project scope.
let (project_hash, project_paths) = compute_project_hash(cwd);
let project_files_exist = project_paths.iter().any(|p| p.exists());
if project_files_exist {
let cwd_key = cwd.to_string_lossy().to_string();
match state.projects.get(&cwd_key) {
None => {
debug!("Claude import: project settings found, no previous import state");
return true;
}
Some(s) if s.last_hash != project_hash => {
debug!(
old = %s.last_hash,
new = %project_hash,
"Claude import: project settings changed since last import"
);
return true;
}
_ => {}
}
}
false
}
// State Updates
fn now_rfc3339() -> String {

52
crates/codegen/xai-grok-shell/src/config/mod.rs
View file

@ -386,7 +386,7 @@ impl SubagentsConfig {
/// File-based personas are loaded from `{cwd}/.grok/personas/*.toml`.
/// Each file defines a single `SubagentPersona`. The file stem becomes
/// the persona name. Inline config takes precedence.
pub fn discover_personas(&mut self, cwd: &std::path::Path) {
pub(crate) fn discover_personas(&mut self, cwd: &std::path::Path) {
let dir = cwd.join(".grok").join("personas");
self.discover_personas_in_dir(&dir);
}
@ -429,7 +429,7 @@ impl SubagentsConfig {
/// The file stem becomes the role name.
///
/// Precedence: inline config roles override file-based roles with the same name.
pub fn discover_roles(&mut self, cwd: &std::path::Path) {
pub(crate) fn discover_roles(&mut self, cwd: &std::path::Path) {
let roles_dir = cwd.join(".grok").join("roles");
self.discover_roles_in_dir(&roles_dir);
}
@ -437,7 +437,7 @@ impl SubagentsConfig {
pub const DEFAULT_MAX_DEPTH: u32 = 1;
/// Clamp to `1..=u32::MAX`. Values below 1 (including 0 / negatives) warn
/// and become 1 so nesting is never accidentally disabled.
pub fn clamp_max_depth(raw: i64, source: &str) -> u32 {
pub(crate) fn clamp_max_depth(raw: i64, source: &str) -> u32 {
if raw < i64::from(Self::DEFAULT_MAX_DEPTH) {
tracing::warn!(
source,
@ -461,7 +461,11 @@ impl SubagentsConfig {
/// Depth 0 is the top-level session; a child is parent+1. Spawn is rejected
/// when `depth >= max`. So `max = 1` allows only top-level spawns; nested
/// spawns from a first-level subagent need `max >= 2`.
pub fn resolve_max_depth(env: Option<&str>, config: Option<i64>, remote: Option<u32>) -> u32 {
pub(crate) fn resolve_max_depth(
env: Option<&str>,
config: Option<i64>,
remote: Option<u32>,
) -> u32 {
if let Some(raw) = env {
match raw.trim().parse::<i64>() {
Ok(v) => return Self::clamp_max_depth(v, "env"),
@ -615,7 +619,7 @@ impl ManagedMcpsConfig {
/// Auxiliary model overrides under `[models]`.
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
#[serde(default)]
pub struct ModelOverrideConfig {
pub(crate) struct ModelOverrideConfig {
pub web_search: String,
/// `None` = current model.
pub session_summary: Option<String>,
@ -895,7 +899,7 @@ impl StorageMode {
/// requires grok.com auth (it syncs to grok-code-backend). This is the
/// single home for that gate, used at boot ([`crate::agent::init`]) and by
/// the post-readiness self-heal (`MvpAgent::reapply_storage_mode`).
pub fn from_remote_gated(
pub(crate) fn from_remote_gated(
remote: Option<&crate::util::config::RemoteSettings>,
has_xai_auth: bool,
) -> Self {
@ -904,10 +908,6 @@ impl StorageMode {
mode => mode,
}
}
/// Returns true if this mode syncs to the backend.
pub fn is_writeback(&self) -> bool {
matches!(self, Self::Writeback)
}
}
pub use xai_grok_config::ConfigLayers;
pub use xai_grok_config::{
@ -920,7 +920,7 @@ pub use xai_grok_config::{
normalize_identity, requirements_layers, system_config_dir, user_grok_home,
};
/// Map of "dotted.path" to which config file the value came from.
pub fn config_origins(
pub(crate) fn config_origins(
layers: &ConfigLayers,
) -> std::collections::HashMap<String, crate::agent::config::ConfigSource> {
use crate::agent::config::ConfigSource;
@ -1017,7 +1017,7 @@ pub struct Sourced<T> {
}
/// A config field clamped by requirements.
#[derive(Debug, Clone)]
pub struct EnforcedField {
pub(crate) struct EnforcedField {
pub path: &'static str,
pub value: String,
pub source: RequirementSource,
@ -1029,7 +1029,7 @@ impl std::fmt::Display for EnforcedField {
}
/// Apply overrides from external `managed-settings.json`.
/// Called before `apply_requirements()` so requirements.toml can override.
pub fn apply_managed_settings_features(
pub(crate) fn apply_managed_settings_features(
config: &mut crate::agent::config::Config,
) -> Vec<EnforcedField> {
let ms = xai_grok_workspace::permission::resolution::managed_settings();
@ -1064,7 +1064,7 @@ fn apply_managed_settings_features_inner(
}
/// Clamp `AgentConfig` fields per `requirements.toml`. No-op if absent.
/// System pins win over user pins on conflict.
pub fn apply_requirements(config: &mut crate::agent::config::Config) -> Vec<EnforcedField> {
pub(crate) fn apply_requirements(config: &mut crate::agent::config::Config) -> Vec<EnforcedField> {
requirements_layers()
.into_iter()
.flat_map(|layer| {
@ -1149,7 +1149,6 @@ fn apply_requirements_inner(
}
pin_feature!(feedback);
pin_feature!(lsp_tools);
pin_feature!(tool_search);
pin_feature!(web_fetch);
pin_feature!(ask_user_question);
pin_feature!(image_gen);
@ -1477,11 +1476,6 @@ pub fn apply_sandbox(
sandbox.install();
}
}
/// Load `<cwd>/.grok/config.toml` (with this layer's `[[version_overrides]]`
/// applied). Empty table if the file is missing.
pub fn load_project_config(cwd: &std::path::Path) -> std::io::Result<toml::Value> {
load_config_file(&cwd.join(".grok").join("config.toml"))
}
pub use xai_grok_workspace::project_config::find_project_configs;
/// Resolve the effective `[plugins]` config for a working directory the same
/// way a session does at reload time: global/user config
@ -1493,7 +1487,7 @@ pub use xai_grok_workspace::project_config::find_project_configs;
/// eager plugin-registry fan-out so all three discover the same plugins for a
/// given cwd. Centralizing it prevents the paths/disabled/discovered-command
/// drift those callers would otherwise accumulate.
pub fn resolve_effective_plugins_config(
pub(crate) fn resolve_effective_plugins_config(
cwd: &std::path::Path,
) -> crate::agent::config::PluginsConfig {
let extract = |toml_val: &toml::Value| -> Option<crate::agent::config::PluginsConfig> {
@ -1524,7 +1518,7 @@ pub use xai_grok_config::{deep_merge_toml, expand_env_vars_in_string, expand_env
///
/// Creates the `[plugins]` section and `paths` array if they don't exist.
/// Deduplicates: if the path is already present, this is a no-op.
pub fn add_plugin_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
pub(crate) fn add_plugin_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
let config_path = crate::util::grok_home::grok_home().join("config.toml");
let content = std::fs::read_to_string(&config_path).unwrap_or_default();
let mut config: toml::Value = if content.is_empty() {
@ -1565,7 +1559,7 @@ pub fn add_plugin_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
/// Remove a plugin path from `[plugins].paths` in `~/.grok/config.toml`.
///
/// If the path is not found, this is a no-op (returns Ok).
pub fn remove_plugin_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
pub(crate) fn remove_plugin_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
let config_path = crate::util::grok_home::grok_home().join("config.toml");
let content = match std::fs::read_to_string(&config_path) {
Ok(c) => c,
@ -1741,7 +1735,7 @@ pub fn dismissed_plugin_ctas_in_file(
/// CWE-427: Only paths under `~/.grok/` are allowed to prevent
/// arbitrary hook path injection that bypasses the project trust gate.
/// Paths are canonicalized (resolving symlinks and `..`) before checking.
pub fn validate_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
pub(crate) fn validate_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
let candidate = std::path::Path::new(path);
if !candidate.is_absolute() {
return Err("Hook path must be absolute.".into());
@ -1781,7 +1775,7 @@ pub fn validate_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>>
///
/// Auto-enables all plugins in the repo so they are active after the next reload.
/// Returns `(plugin_names, warnings)` for status messaging.
pub fn post_install_plugin(repo_key: &str) -> (Vec<String>, Vec<String>) {
pub(crate) fn post_install_plugin(repo_key: &str) -> (Vec<String>, Vec<String>) {
let registry = xai_grok_agent::plugins::InstallRegistry::load();
let Some(repo) = registry.get_repo(repo_key) else {
return (
@ -1867,7 +1861,7 @@ pub fn remove_enabled_plugin(plugin_id: &str) -> Result<(), Box<dyn std::error::
///
/// If the path is already present (exact string match), this is a no-op.
/// CWE-427: The path is validated to be under `~/.grok/` before writing.
pub fn add_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
pub(crate) fn add_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
validate_hooks_path(path)?;
add_hooks_path_to_file(
path,
@ -1875,7 +1869,7 @@ pub fn add_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
)
}
/// Add a hook path to a specific file (for tests).
pub fn add_hooks_path_to_file(
pub(crate) fn add_hooks_path_to_file(
path: &str,
paths_file: &std::path::Path,
) -> Result<(), Box<dyn std::error::Error>> {
@ -1898,14 +1892,14 @@ pub fn add_hooks_path_to_file(
///
/// If the path is not found (exact string match), this is a no-op.
/// Matches the same exact-string behavior as `add_hooks_path`.
pub fn remove_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
pub(crate) fn remove_hooks_path(path: &str) -> Result<(), Box<dyn std::error::Error>> {
remove_hooks_path_from_file(
path,
&crate::util::grok_home::grok_home().join("hooks-paths"),
)
}
/// Remove a hook path from a specific file (for tests).
pub fn remove_hooks_path_from_file(
pub(crate) fn remove_hooks_path_from_file(
path: &str,
paths_file: &std::path::Path,
) -> Result<(), Box<dyn std::error::Error>> {

2
crates/codegen/xai-grok-shell/src/config/reloader.rs
View file

@ -81,7 +81,7 @@ pub enum ConfigUpdate {
/// Runs on `tokio::spawn` (`Send`). Receives raw [`ConfigChangeEvent`]s from
/// the file watcher, diffs against last-known state, and sends [`ConfigUpdate`]
/// messages to the agent via an `mpsc` channel.
pub struct ConfigReloader {
pub(crate) struct ConfigReloader {
last_auth_key_hash: u64,
last_global_config: toml::Value,
/// Per-cwd content hash of the project MCP config files, used to

6
crates/codegen/xai-grok-shell/src/config/watcher.rs
View file

@ -16,7 +16,7 @@ const DEFAULT_DEBOUNCE: Duration = Duration::from_millis(1000);
/// reload's own reads schedule the next reload, a ~1/sec self-sustaining loop.
/// Dropping `Access` is safe: writes still emit `Modify`/`Create` and chmod
/// emits `Modify(Metadata)`; only reads are `Access`-only.
pub struct AccessFilteredWatcher(notify::RecommendedWatcher);
pub(crate) struct AccessFilteredWatcher(notify::RecommendedWatcher);
impl notify::Watcher for AccessFilteredWatcher {
fn new<F: notify::EventHandler>(
@ -585,7 +585,7 @@ fn plan_skills_watch_targets(
///
/// After a [`DiscoveryChange`], call [`Self::refresh_new_dirs`] so newly created
/// seed dirs get watches attached.
pub struct ProjectDiscoveryWatcher {
pub(crate) struct ProjectDiscoveryWatcher {
debouncer: Debouncer<AccessFilteredWatcher>,
refresh_dirs: Vec<(PathBuf, RecursiveMode)>,
refreshed_dirs: HashSet<PathBuf>,
@ -651,7 +651,7 @@ impl ProjectDiscoveryWatcher {
}
/// Attach watches for seed dirs that now exist (call after a discovery event).
pub fn refresh_new_dirs(&mut self) {
pub(crate) fn refresh_new_dirs(&mut self) {
attach_new_refresh_dirs(
&mut self.debouncer,
&self.refresh_dirs,

2
crates/codegen/xai-grok-shell/src/extensions/bundle.rs
View file

@ -51,7 +51,7 @@ struct BundleSyncRequest {
}
#[derive(Debug, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct BundleSyncResult {
pub(crate) struct BundleSyncResult {
pub updated: bool,
pub version: String,
pub personas_count: usize,

10
crates/codegen/xai-grok-shell/src/extensions/code_nav.rs
View file

@ -57,7 +57,7 @@ type ExtResult = Result<acp::ExtResponse, acp::Error>;
/// receive `reason: sessionRequired` in the error response.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GotoRequest {
pub(crate) struct GotoRequest {
/// Session ID — required for code navigation.
pub session_id: Option<acp::SessionId>,
/// Working directory (optional when session_id is provided).
@ -75,7 +75,7 @@ pub struct GotoRequest {
/// **`sessionId` is required** — same contract as [`GotoRequest`].
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FindSymbolRequest {
pub(crate) struct FindSymbolRequest {
/// Session ID — required for code navigation.
pub session_id: Option<acp::SessionId>,
/// Working directory (optional when session_id is provided).
@ -91,7 +91,7 @@ pub struct FindSymbolRequest {
/// **`sessionId` is required** — same contract as [`GotoRequest`].
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StatusRequest {
pub(crate) struct StatusRequest {
/// Session ID — required for code navigation.
pub session_id: Option<acp::SessionId>,
/// Working directory (optional when session_id is provided).
@ -134,7 +134,7 @@ pub struct SymbolLocation {
/// Serialised as a camelCase string so clients can pattern-match on it.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub enum IndexStatusReason {
pub(crate) enum IndexStatusReason {
/// Index is running and ready.
Active,
/// Index is eligible but has not been started yet (first code-nav request
@ -155,7 +155,7 @@ pub enum IndexStatusReason {
/// Response for status query.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct StatusResponse {
pub(crate) struct StatusResponse {
/// Whether an index is currently active for this cwd.
pub indexed: bool,
/// Whether this client is eligible to use codebase indexing.

10
crates/codegen/xai-grok-shell/src/extensions/fs.rs
View file

@ -35,7 +35,7 @@ fn default_include_hidden() -> bool {
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FsListRequest {
pub(crate) struct FsListRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,
@ -74,14 +74,14 @@ impl FsListRequest {
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FsExistsRequest {
pub(crate) struct FsExistsRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FsReadFileRequest {
pub(crate) struct FsReadFileRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,
@ -109,7 +109,7 @@ pub struct FsReadFileRequest {
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FsWriteFileRequest {
pub(crate) struct FsWriteFileRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,
@ -119,7 +119,7 @@ pub struct FsWriteFileRequest {
}
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FsDeleteFileRequest {
pub(crate) struct FsDeleteFileRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,

28
crates/codegen/xai-grok-shell/src/extensions/git.rs
View file

@ -60,7 +60,7 @@ fn default_working() -> String {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitStatusRequest {
pub(crate) struct GitStatusRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -72,7 +72,7 @@ pub struct GitStatusRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitFilesRequest {
pub(crate) struct GitFilesRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -83,7 +83,7 @@ pub struct GitFilesRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitDiffsRequest {
pub(crate) struct GitDiffsRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -107,7 +107,7 @@ pub struct GitDiffsRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitStageRequest {
pub(crate) struct GitStageRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -116,7 +116,7 @@ pub struct GitStageRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitStageContentRequest {
pub(crate) struct GitStageContentRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -126,7 +126,7 @@ pub struct GitStageContentRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitUnstageRequest {
pub(crate) struct GitUnstageRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -152,7 +152,7 @@ impl From<GitDiscardScope> for DiscardScope {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitDiscardRequest {
pub(crate) struct GitDiscardRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -165,7 +165,7 @@ pub struct GitDiscardRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitCommitRequest {
pub(crate) struct GitCommitRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -182,7 +182,7 @@ pub struct GitCommitRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitStashRequest {
pub(crate) struct GitStashRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -192,7 +192,7 @@ pub struct GitStashRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitCheckoutRequest {
pub(crate) struct GitCheckoutRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -212,7 +212,7 @@ struct CheckoutSessionHeadRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitInfoRequest {
pub(crate) struct GitInfoRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -220,7 +220,7 @@ pub struct GitInfoRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitBranchesRequest {
pub(crate) struct GitBranchesRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -228,7 +228,7 @@ pub struct GitBranchesRequest {
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitCurrentCommitRequest {
pub(crate) struct GitCurrentCommitRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]
@ -237,7 +237,7 @@ pub struct GitCurrentCommitRequest {
/// Request for x.ai/git/checkout_commit extension method.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GitCheckoutCommitRequest {
pub(crate) struct GitCheckoutCommitRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
#[serde(default)]

4
crates/codegen/xai-grok-shell/src/extensions/hooks.rs
View file

@ -21,7 +21,7 @@ struct ListRequest {
session_id: String,
}
pub fn hook_spec_to_info(spec: &xai_grok_hooks::config::HookSpec) -> HookInfo {
pub(crate) fn hook_spec_to_info(spec: &xai_grok_hooks::config::HookSpec) -> HookInfo {
use xai_grok_hooks::event::HookEventName;
let event = match spec.event {
@ -92,7 +92,7 @@ pub struct ClientHookGroup {
pub timeout: Option<std::time::Duration>,
}
pub type ClientHooks = HashMap<HookEventName, Vec<ClientHookGroup>>;
pub(crate) type ClientHooks = HashMap<HookEventName, Vec<ClientHookGroup>>;
/// One hook dispatched to a client callback: the shared [`HookEventEnvelope`]
/// (flattened, camelCase) plus the `hookCallbackId` it targets. The same shape is sent

20
crates/codegen/xai-grok-shell/src/extensions/hunk_tracker.rs
View file

@ -27,7 +27,7 @@ use xai_hunk_tracker::{
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetHunksRequest {
pub(crate) struct GetHunksRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
/// Filter by file path (optional)
@ -39,14 +39,14 @@ pub struct GetHunksRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetFilesRequest {
pub(crate) struct GetFilesRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HunkActionRequest {
pub(crate) struct HunkActionRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub hunk_id: String,
@ -55,7 +55,7 @@ pub struct HunkActionRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FileActionRequest {
pub(crate) struct FileActionRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub path: String,
@ -64,7 +64,7 @@ pub struct FileActionRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TurnActionRequest {
pub(crate) struct TurnActionRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub prompt_index: usize,
@ -73,7 +73,7 @@ pub struct TurnActionRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AllActionRequest {
pub(crate) struct AllActionRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
pub action: String, // "accept" | "reject"
@ -81,7 +81,7 @@ pub struct AllActionRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GetSummaryRequest {
pub(crate) struct GetSummaryRequest {
#[serde(default)]
pub session_id: Option<acp::SessionId>,
}
@ -125,19 +125,19 @@ pub struct FileSummary {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GetFilesResponse {
pub(crate) struct GetFilesResponse {
pub files: Vec<FileSummary>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GetAllFileContentsResponse {
pub(crate) struct GetAllFileContentsResponse {
pub files: Vec<FileContentEntry>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ActionResponse {
pub(crate) struct ActionResponse {
pub success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,

2
crates/codegen/xai-grok-shell/src/extensions/jj.rs
View file

@ -9,7 +9,7 @@ use xai_grok_workspace::session::jj;
/// Handle a `x.ai/git/*` method for a jj-colocated repo.
///
/// Returns `Some(result)` if handled, `None` to fall through to git.
pub async fn try_handle(
pub(crate) async fn try_handle(
method: &str,
git_root: &std::path::Path,
raw_params: &serde_json::value::RawValue,

4
crates/codegen/xai-grok-shell/src/extensions/marketplace.rs
View file

@ -1268,7 +1268,7 @@ fn read_default_skills_installs_purged(config_path: &std::path::Path) -> bool {
///
/// Gated by sticky `default_skills_installs_purged` in config.toml. Best-effort:
/// errors are logged and never block startup.
pub fn purge_default_skills_installs(grok_home: &std::path::Path) {
pub(crate) fn purge_default_skills_installs(grok_home: &std::path::Path) {
purge_default_skills_installs_impl(grok_home, || {
xai_grok_agent::plugins::install_registry::InstallRegistry::try_load_from(
xai_grok_agent::plugins::install_registry::InstallRegistry::resolve_install_dir(),
@ -1368,7 +1368,7 @@ fn purge_default_skills_installs_impl(
/// `official_marketplace_auto_installed` is set. Under a process-wide flock it
/// adds the source (or just sets the flag if it's already present in config.toml
/// or a JSON store). Best-effort: errors are logged and never block startup.
pub fn ensure_official_marketplace_source(grok_home: &std::path::Path) {
pub(crate) fn ensure_official_marketplace_source(grok_home: &std::path::Path) {
let config_path = grok_home.join("config.toml");
if read_official_marketplace_auto_installed(&config_path) {

19
crates/codegen/xai-grok-shell/src/extensions/mcp.rs
View file

@ -72,7 +72,7 @@ fn default_true() -> bool {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct McpListResponse {
pub(crate) struct McpListResponse {
pub servers: Vec<McpServerEntry>,
}
@ -179,7 +179,7 @@ pub struct McpToolEntry {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpCallRequest {
pub(crate) struct McpCallRequest {
/// When present: session pool. When absent: agent pool (config.toml only).
#[serde(default)]
pub session_id: Option<String>,
@ -287,7 +287,7 @@ pub use crate::session::mcp_dispatcher::{
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpReadResourceRequest {
pub(crate) struct McpReadResourceRequest {
#[serde(default)]
pub session_id: Option<String>,
pub server: String,
@ -411,7 +411,7 @@ pub fn build_mcp_catalog(
build_mcp_catalog_with_gateway_tools(managed_configs, local_servers, None, &Default::default())
}
pub fn build_mcp_catalog_with_gateway_tools(
pub(crate) fn build_mcp_catalog_with_gateway_tools(
managed_configs: &[crate::session::managed_mcp::ManagedMcpConfig],
local_servers: &[acp::McpServer],
gateway_catalog: Option<&crate::session::managed_mcp::GatewayToolCatalog>,
@ -608,7 +608,7 @@ fn disabled_server_placeholder_entry(name: &str) -> McpServerEntry {
/// Build session MCP status: which servers are enabled, healthy, and what tools they expose.
/// Clones state under lock then releases — does not hold lock across awaits.
pub async fn build_mcp_status(
pub(crate) async fn build_mcp_status(
mcp_state: &Arc<TokioMutex<McpState>>,
tool_bridge: &Arc<xai_grok_tools::bridge::ToolBridge>,
event_writer: Option<&xai_file_utils::events::EventWriter>,
@ -767,7 +767,10 @@ async fn ensure_agent_pool_initialized(mcp_state: &Arc<TokioMutex<McpState>>) {
/// Spawn config.toml MCP clients into the agent pool. Handshakes happen
/// lazily on first `CallMcpTool`.
pub async fn init_agent_mcp_pool(mcp_state: &Arc<TokioMutex<McpState>>, cwd: &std::path::Path) {
pub(crate) async fn init_agent_mcp_pool(
mcp_state: &Arc<TokioMutex<McpState>>,
cwd: &std::path::Path,
) {
use crate::session::mcp_servers::start_mcp_servers;
let configs = {
@ -1240,7 +1243,7 @@ async fn handle_read_resource(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
to_ext_response(Ok(result))
}
pub async fn read_mcp_resource(
pub(crate) async fn read_mcp_resource(
mcp_state: &Arc<TokioMutex<McpState>>,
server_name: &str,
uri: &str,
@ -1329,7 +1332,7 @@ pub async fn read_mcp_resource(
///
/// Injected into the agent's `SharedResources` via `tool_bridge.update_resource()`
/// at session startup so tools can enumerate and fetch MCP resources.
pub struct McpStateResourceProvider(pub Arc<TokioMutex<McpState>>);
pub(crate) struct McpStateResourceProvider(pub Arc<TokioMutex<McpState>>);
#[async_trait::async_trait]
impl xai_grok_tools::types::resources::McpResourceProvider for McpStateResourceProvider {

14
crates/codegen/xai-grok-shell/src/extensions/mod.rs
View file

@ -2,7 +2,7 @@ pub mod auth;
pub(crate) mod auth_gate;
pub mod billing;
pub mod bundle;
pub mod chat_conversation_history;
pub(crate) mod chat_conversation_history;
pub mod code_nav;
pub mod debug;
pub mod feedback;
@ -27,9 +27,9 @@ pub mod rewind;
pub mod rollout;
pub mod routing;
pub mod search;
pub mod session_admin;
pub(crate) mod session_admin;
pub mod session_search;
pub mod session_state;
pub(crate) mod session_state;
pub mod session_updates;
pub mod share;
pub mod skills;
@ -44,13 +44,13 @@ use serde::Serialize;
use serde::de::DeserializeOwned;
use std::sync::Arc;
pub type ExtResult = Result<acp::ExtResponse, acp::Error>;
pub fn parse_params<T: DeserializeOwned>(args: &acp::ExtRequest) -> Result<T, acp::Error> {
pub(crate) fn parse_params<T: DeserializeOwned>(args: &acp::ExtRequest) -> Result<T, acp::Error> {
parse_params_str(args.params.get())
}
/// Deserialize ACP params from their raw JSON string, mapping a parse failure
/// to `invalid_params`. Used by [`parse_params`] and the bridge `encode` hooks,
/// which hold the params `RawValue` directly.
pub fn parse_params_str<T: DeserializeOwned>(raw: &str) -> Result<T, acp::Error> {
pub(crate) fn parse_params_str<T: DeserializeOwned>(raw: &str) -> Result<T, acp::Error> {
serde_json::from_str(raw)
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {}", e)))
}
@ -66,13 +66,13 @@ pub fn to_ext_response<T: Serialize>(result: anyhow::Result<T>) -> ExtResult {
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
/// Wrap a serializable value as an `ExtResponse` without the `ExtMethodResult` envelope.
pub fn to_raw_response<T: Serialize>(v: &T) -> ExtResult {
pub(crate) fn to_raw_response<T: Serialize>(v: &T) -> ExtResult {
serde_json::value::to_raw_value(v)
.map(|raw| acp::ExtResponse::new(Arc::from(raw)))
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
/// Convert a result with optional warning to an ExtResponse.
pub fn to_ext_response_partial<T: Serialize>(
pub(crate) fn to_ext_response_partial<T: Serialize>(
result: anyhow::Result<T>,
warning: Option<String>,
) -> ExtResult {

10
crates/codegen/xai-grok-shell/src/extensions/notification.rs
View file

@ -92,7 +92,7 @@ impl PromptUsage {
/// Project a ledger snapshot for the wire. Returns `Some` whenever
/// `incomplete` is set — even if `ledger` is `None` — so the flag is never
/// dropped by omission. Always scrubs untrustworthy costs.
pub fn project_from_ledger(
pub(crate) fn project_from_ledger(
ledger: Option<&xai_chat_state::UsageLedger>,
incomplete: bool,
) -> Option<Self> {
@ -116,7 +116,7 @@ impl PromptUsage {
/// Error-path attach: any open ledger is always incomplete (may under-count
/// without a freeze drain). `may_undercount` only matters when the ledger is empty.
pub fn for_error_path(
pub(crate) fn for_error_path(
ledger: Option<&xai_chat_state::UsageLedger>,
may_undercount: bool,
) -> Option<Self> {
@ -129,7 +129,7 @@ impl PromptUsage {
/// Drop cost ticks when partial or incomplete so all wire surfaces fail closed.
/// Incomplete bills clear ticks even when `cost_is_partial` is false.
pub fn scrub_untrustworthy_costs(&mut self) {
pub(crate) fn scrub_untrustworthy_costs(&mut self) {
if !(self.usage_is_incomplete || self.totals.cost_is_partial) {
return;
}
@ -282,7 +282,7 @@ pub fn ticks_to_usd(ticks: i64) -> f64 {
}
/// Full ACP input → headless uncached input (`full − cache_read`).
pub fn uncached_input_tokens(full_input: u64, cached_read: u64) -> u64 {
pub(crate) fn uncached_input_tokens(full_input: u64, cached_read: u64) -> u64 {
full_input.saturating_sub(cached_read)
}
@ -294,7 +294,7 @@ pub fn uncached_input_tokens(full_input: u64, cached_read: u64) -> u64 {
/// - Omits all cost floats when partial or incomplete (absence ≠ free).
/// - Incomplete with no tokens emits only `usage_is_incomplete` (no zero usage object).
/// - `modelUsage` rows are a reduced external-compat schema (camelCase; no reasoning/duration).
pub fn project_result_usage(result: &mut serde_json::Value, usage: &PromptUsage) {
pub(crate) fn project_result_usage(result: &mut serde_json::Value, usage: &PromptUsage) {
if usage.usage_is_incomplete && usage.is_token_empty() {
result["usage_is_incomplete"] = true.into();
return;

2
crates/codegen/xai-grok-shell/src/extensions/plugins.rs
View file

@ -19,7 +19,7 @@ struct ListRequest {
}
/// Convert a `LoadedPlugin` to a `PluginInfo` DTO.
pub fn loaded_plugin_to_info(plugin: &xai_grok_agent::plugins::LoadedPlugin) -> PluginInfo {
pub(crate) fn loaded_plugin_to_info(plugin: &xai_grok_agent::plugins::LoadedPlugin) -> PluginInfo {
use xai_grok_agent::plugins::discovery::PluginScope as AgentScope;
let scope = match plugin.scope {

6
crates/codegen/xai-grok-shell/src/extensions/pr.rs
View file

@ -6,21 +6,21 @@ use crate::agent::MvpAgent;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PrStatusRequest {
pub(crate) struct PrStatusRequest {
pub cwd: String,
pub branch: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PrStatusResponse {
pub(crate) struct PrStatusResponse {
pub pr: Option<PrData>,
pub updated_session_ids: Vec<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PrData {
pub(crate) struct PrData {
pub url: String,
pub state: String,
pub is_in_merge_queue: bool,

2
crates/codegen/xai-grok-shell/src/extensions/repair.rs
View file

@ -35,7 +35,7 @@ struct RepairSessionRequest {
/// Response payload for `x.ai/session/repair`.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RepairSessionResponse {
pub(crate) struct RepairSessionResponse {
/// Whether the repair modified (or, for `dryRun`, would modify) the history.
pub repaired: bool,
/// Echo of the request's `dryRun` flag.

5
crates/codegen/xai-grok-shell/src/extensions/routing.rs
View file

@ -25,7 +25,10 @@ pub struct NotificationMeta {
/// Inject `targetClientId` into the `_meta` field of a JSON params object.
/// Merges with any existing `_meta` fields rather than replacing them.
pub fn inject_routing_meta(params: &mut serde_json::Value, target_client_id: &TargetClientId) {
pub(crate) fn inject_routing_meta(
params: &mut serde_json::Value,
target_client_id: &TargetClientId,
) {
if target_client_id.is_none() {
return;
}

10
crates/codegen/xai-grok-shell/src/extensions/search.rs
View file

@ -14,14 +14,14 @@ type ExtResult = Result<acp::ExtResponse, acp::Error>;
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct FuzzyOpenResponse {
pub(crate) struct FuzzyOpenResponse {
pub session_id: String,
pub search_id: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct FuzzyChangeResponse {
pub(crate) struct FuzzyChangeResponse {
pub session_id: String,
pub search_id: String,
}
@ -65,7 +65,7 @@ fn resolve_cwd(
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FuzzyOpenRequest {
pub(crate) struct FuzzyOpenRequest {
/// Optional session ID - used to lookup cwd if cwd not provided directly
#[serde(default)]
pub session_id: Option<acp::SessionId>,
@ -86,7 +86,7 @@ pub struct FuzzyOpenRequest {
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FuzzyChangeRequest {
pub(crate) struct FuzzyChangeRequest {
pub search_id: String,
pub query: String,
#[serde(default)]
@ -97,7 +97,7 @@ pub struct FuzzyChangeRequest {
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FuzzyCloseRequest {
pub(crate) struct FuzzyCloseRequest {
pub search_id: String,
}

53
crates/codegen/xai-grok-shell/src/extensions/session_admin.rs
View file

@ -27,6 +27,7 @@ use serde::Deserialize;
use super::{ExtResult, parse_params, to_raw_response};
use crate::agent::MvpAgent;
use crate::leader::protocol::InternalMethod;
use crate::session::persistence::list_summaries;
use crate::session::storage::StorageAdapter;
use crate::session::storage::jsonl::JsonlStorageAdapter;
@ -36,6 +37,9 @@ use xai_grok_telemetry::id::agent_id;
#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
if let Some(method) = InternalMethod::from_name(args.method.as_ref()) {
return handle_internal(agent, args, method).await;
}
match args.method.as_ref() {
"x.ai/session/rename" => handle_session_rename(agent, args).await,
"x.ai/session/delete" => handle_session_delete(agent, args).await,
@ -43,21 +47,33 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
#[cfg(feature = "local-workspace")]
"x.ai/session/add_local_workspace" => handle_add_local_workspace(agent, args).await,
"x.ai/session/fork" => handle_session_fork(agent, args).await,
"x.ai/internal/reload_all_mcp_servers" => handle_reload_all_mcp_servers(agent).await,
"x.ai/internal/reload_project_mcp_servers" => {
handle_reload_project_mcp_servers(agent, args).await
}
"x.ai/internal/reload_skills" => handle_reload_skills(agent),
"x.ai/internal/reload_workflows" => handle_reload_workflows(agent),
"x.ai/internal/reload_models" => handle_reload_models(agent),
"x.ai/internal/reload_models_cache" => handle_reload_models_cache(agent),
"x.ai/internal/auth_cleared" => handle_auth_cleared(agent),
"x.ai/plugins/reload" => handle_plugins_reload(agent).await,
"x.ai/commands/list" => handle_commands_list(agent, args).await,
_ => Err(acp::Error::method_not_found()),
}
}
/// Exhaustive, so a new [`InternalMethod`] cannot compile without a handler.
async fn handle_internal(
agent: &MvpAgent,
args: &acp::ExtRequest,
method: InternalMethod,
) -> ExtResult {
match method {
InternalMethod::ReloadAllMcpServers => handle_reload_all_mcp_servers(agent).await,
InternalMethod::ReloadProjectMcpServers => {
handle_reload_project_mcp_servers(agent, args).await
}
InternalMethod::ReloadSkills => handle_reload_skills(agent),
InternalMethod::ReloadWorkflows => handle_reload_workflows(agent),
InternalMethod::ReloadModels => handle_reload_models(agent),
InternalMethod::ReloadModelsCache => handle_reload_models_cache(agent),
InternalMethod::AuthCleared => handle_auth_cleared(agent),
// Arrives as a notification, so it never reaches this request path.
InternalMethod::EvictSessions => Err(acp::Error::method_not_found()),
}
}
// session/rename
/// Handles renaming a session.
@ -337,13 +353,18 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) ->
(h, cwd)
};
// Await managed first, then one compat snapshot for admit + merge so a
// settings reapply mid-await cannot make the seed and spawned set disagree.
let managed = agent.get_managed_mcp_configs().await;
let compat = agent.cfg.borrow().compat_resolved;
let admitted =
crate::session::managed_mcp::admit_client_mcp_servers(params.mcp_servers, &cwd, &compat);
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
params.mcp_servers.clone(),
admitted.clone(),
&cwd,
&managed,
agent.plugin_registry_handle().snapshot().as_deref(),
&agent.cfg.borrow().compat_resolved,
&compat,
);
let (tx, rx) = tokio::sync::oneshot::channel();
@ -360,13 +381,11 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) ->
.map_err(|_| acp::Error::internal_error().data("session closed"))?
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
// Persist the new client set on the handle so config hot-reloads
// (`reload_all_mcp_servers` / `reload_project_mcp_servers`) re-merge from
// the client's latest intent rather than the `session/new` snapshot —
// otherwise a reload would resurrect servers the client just removed
// (or drop ones it just added).
// Store the admitted (not raw) client set: hot-reloads re-merge from this
// seed, and a raw list would re-spawn a previously rejected vendor server
// once on-disk attribution vanishes.
if let Some(h) = agent.sessions.borrow_mut().get_mut(&params.session_id) {
h.initial_client_mcp_servers = params.mcp_servers;
h.initial_client_mcp_servers = admitted;
}
ExtMethodResult::success(serde_json::json!({ "ok": true }))

4
crates/codegen/xai-grok-shell/src/extensions/session_search.rs
View file

@ -21,7 +21,7 @@ use super::ExtResult;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchSessionsRequest {
pub(crate) struct SearchSessionsRequest {
/// The search query string.
pub query: String,
/// Optional workspace directory to scope results to.
@ -44,7 +44,7 @@ fn default_limit() -> usize {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchSessionsResponse {
pub(crate) struct SearchSessionsResponse {
pub results: Vec<SearchSessionHit>,
pub next_offset: Option<usize>,
pub total_estimate: Option<usize>,

4
crates/codegen/xai-grok-shell/src/extensions/session_state.rs
View file

@ -44,7 +44,7 @@ fn validate_session_uuid(session_id: &str) -> Result<(), acp::Error> {
/// `x.ai/session/state`: return metadata columns keyed by logical name. Errors when
/// the session isn't found on this host, since it reads a single record whose absence
/// is not an empty result (unlike the collection returned by `x.ai/session/updates`).
pub async fn handle_state(args: &acp::ExtRequest) -> ExtResult {
pub(crate) async fn handle_state(args: &acp::ExtRequest) -> ExtResult {
let request: StateRequest = super::parse_params(args)?;
validate_session_uuid(&request.session_id)?;
@ -76,7 +76,7 @@ struct ImportRequest {
/// `x.ai/session/import`: recreate a session on this host from mirrored columns and
/// transcript. A session that already exists locally is left unchanged.
pub async fn handle_import(args: &acp::ExtRequest) -> ExtResult {
pub(crate) async fn handle_import(args: &acp::ExtRequest) -> ExtResult {
let mut request: ImportRequest = super::parse_params(args)?;
validate_session_uuid(&request.session_id)?;

2
crates/codegen/xai-grok-shell/src/extensions/skills.rs
View file

@ -73,7 +73,7 @@ pub struct SkillsResetResponse {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SkillsToggleRequest {
pub(crate) struct SkillsToggleRequest {
/// Skill name to toggle.
pub name: String,
/// Whether to enable (`true`) or disable (`false`) the skill.

4
crates/codegen/xai-grok-shell/src/extensions/task.rs
View file

@ -387,7 +387,7 @@ struct DeleteScheduledTaskResponse {
}
/// Handle `x.ai/scheduler/*` extension methods.
pub async fn handle_scheduler(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
pub(crate) async fn handle_scheduler(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/scheduler/delete" => {
let req: DeleteScheduledTaskRequest = parse(args)?;
@ -405,7 +405,7 @@ pub async fn handle_scheduler(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRe
}
/// Handle `x.ai/subagent/*` extension methods.
pub async fn handle_subagent(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
pub(crate) async fn handle_subagent(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/subagent/cancel" => {
let req: CancelSubagentRequest = parse(args)?;

16
crates/codegen/xai-grok-shell/src/extensions/terminal.rs
View file

@ -30,7 +30,7 @@ pub struct CreateTerminalRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TerminalIdRequest {
pub(crate) struct TerminalIdRequest {
pub session_id: String,
pub terminal_id: String,
}
@ -44,7 +44,7 @@ pub struct CreateTerminalResponse {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PtyCreateRequest {
pub(crate) struct PtyCreateRequest {
pub shell: Option<String>,
pub cwd: Option<String>,
#[serde(default)]
@ -60,7 +60,7 @@ pub struct PtyCreateRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PtyLoadRequest {
pub(crate) struct PtyLoadRequest {
pub terminal_id: String,
#[serde(default, rename = "_meta")]
pub meta: Option<RequestMeta>,
@ -70,14 +70,14 @@ pub struct PtyLoadRequest {
/// ignored for PTY terminals (looked up by `terminal_id` alone).
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct KillTerminalRequest {
pub(crate) struct KillTerminalRequest {
pub terminal_id: String,
pub session_id: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PtyResizeRequest {
pub(crate) struct PtyResizeRequest {
pub terminal_id: String,
pub rows: u16,
pub cols: u16,
@ -85,14 +85,14 @@ pub struct PtyResizeRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PtyInputNotification {
pub(crate) struct PtyInputNotification {
pub terminal_id: String,
pub data: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TerminalListResponse {
pub(crate) struct TerminalListResponse {
pub terminals: Vec<terminal::TerminalInfo>,
}
@ -348,7 +348,7 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
}
}
pub async fn handle_pty_input(params: &serde_json::Value) {
pub(crate) async fn handle_pty_input(params: &serde_json::Value) {
use base64::Engine as _;
let Ok(input) = serde_json::from_value::<PtyInputNotification>(params.clone()) else {

8
crates/codegen/xai-grok-shell/src/extensions/worktree.rs
View file

@ -71,13 +71,13 @@ pub struct ListWorktreeRequest {
#[derive(Debug, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ShowWorktreeRequest {
pub(crate) struct ShowWorktreeRequest {
pub id_or_path: String,
}
#[derive(Debug, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GcWorktreeRequest {
pub(crate) struct GcWorktreeRequest {
#[serde(default)]
pub dry_run: bool,
/// Duration string like "7d", "24h", "30m", "60s".
@ -95,14 +95,14 @@ pub struct WorktreeDbPathResponse {
#[derive(Debug, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ResolveLocalForWorktreeResumeRequest {
pub(crate) struct ResolveLocalForWorktreeResumeRequest {
pub session_id: String,
pub cwd: String,
}
#[derive(Debug, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ResolveLocalForWorktreeResumeResponse {
pub(crate) struct ResolveLocalForWorktreeResumeResponse {
pub found: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub resolved_session_id: Option<String>,

18
crates/codegen/xai-grok-shell/src/heap_profile/monitor.rs
View file

@ -209,7 +209,7 @@ impl HeapProfileMonitor {
self.upload_in_flight
}
pub fn clear_upload_in_flight(&mut self) {
pub(crate) fn clear_upload_in_flight(&mut self) {
self.upload_in_flight = false;
}
@ -256,7 +256,7 @@ impl HeapProfileMonitor {
/// Start a dump when a threshold is crossed. Deferred paths return `None`
/// without latching.
pub fn begin_tick(&mut self) -> Option<PendingDump> {
pub(crate) fn begin_tick(&mut self) -> Option<PendingDump> {
if !self.config.enabled || self.upload_in_flight {
if self.upload_in_flight {
tracing::debug!(reason = "in_flight", "heap_profile: skipped");
@ -298,7 +298,7 @@ impl HeapProfileMonitor {
})
}
pub fn finish_tick(&mut self, threshold: u64, outcome: DumpAttemptOutcome) {
pub(crate) fn finish_tick(&mut self, threshold: u64, outcome: DumpAttemptOutcome) {
if should_latch(outcome) {
self.latched.insert(threshold);
}
@ -357,7 +357,7 @@ impl HeapProfileMonitor {
}
/// Work item produced by [`HeapProfileMonitor::begin_tick`].
pub struct PendingDump {
pub(crate) struct PendingDump {
pub threshold: u64,
stats: super::JemallocStats,
session_id: Arc<str>,
@ -391,6 +391,16 @@ impl PendingDump {
"heap_profile: threshold_crossed"
);
xai_grok_telemetry::session_ctx::log_event(
xai_grok_telemetry::events::HeapThresholdCrossed {
threshold_bytes: threshold,
resident_bytes: stats.resident,
allocated_bytes: stats.allocated,
// The sampler reports 0 when the platform has no cheap read.
rss_peak_bytes: (rss_peak > 0).then_some(rss_peak),
},
);
let temp_dir = match PrivateTempDir::create() {
Ok(d) => d,
Err(e) => {

26
crates/codegen/xai-grok-shell/src/inspect/mod.rs
View file

@ -53,7 +53,7 @@ impl std::fmt::Display for Scope {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct InspectReport {
pub(crate) struct InspectReport {
pub grok_version: String,
pub channel: String,
pub cwd: String,
@ -82,7 +82,7 @@ pub struct InspectReport {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct InstructionFile {
pub(crate) struct InstructionFile {
pub path: String,
pub scope: Scope,
pub file_type: String,
@ -100,7 +100,7 @@ pub struct InstructionFile {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PermissionsReport {
pub(crate) struct PermissionsReport {
pub sources: Vec<String>,
pub loaded: usize,
pub skipped: Vec<SkippedRule>,
@ -124,7 +124,7 @@ pub struct PermissionsReport {
/// derives its line from these fields (see `enforced_label`).
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EnforcedPolicy {
pub(crate) struct EnforcedPolicy {
/// Stable key: "alwaysApprove" | "telemetry" | "feedback".
pub setting: String,
/// The enforced value.
@ -135,7 +135,7 @@ pub struct EnforcedPolicy {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SkippedRule {
pub(crate) struct SkippedRule {
pub rule: String,
pub reason: String,
}
@ -145,7 +145,7 @@ pub struct SkippedRule {
/// The team pin is admin policy, not a secret, so it is shown verbatim.
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LoginPolicyReport {
pub(crate) struct LoginPolicyReport {
/// Raw `disable_api_key_auth` knob (env `GROK_DISABLE_API_KEY_AUTH`).
pub disable_api_key_auth: Option<bool>,
/// Configured team pin: single string, list, or null when unset.
@ -156,7 +156,7 @@ pub struct LoginPolicyReport {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct HookEntry {
pub(crate) struct HookEntry {
pub event: String,
pub hook_type: String,
pub target: String,
@ -173,7 +173,7 @@ pub struct HookEntry {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SkillEntry {
pub(crate) struct SkillEntry {
pub name: String,
pub description: String,
pub source: ConfigSource,
@ -190,7 +190,7 @@ pub struct SkillEntry {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AgentEntry {
pub(crate) struct AgentEntry {
pub name: String,
pub description: String,
pub source: ConfigSource,
@ -217,7 +217,7 @@ pub struct PluginProvides {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MarketplaceEntry {
pub(crate) struct MarketplaceEntry {
pub name: String,
pub path: String,
pub enabled_plugins: usize,
@ -243,7 +243,7 @@ pub struct McpServerEntry {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LspServerEntry {
pub(crate) struct LspServerEntry {
pub name: String,
pub command: String,
pub args: Vec<String>,
@ -256,7 +256,7 @@ pub struct LspServerEntry {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ConfigSources {
pub(crate) struct ConfigSources {
/// Config layers (system + user managed, user + system requirements, user
/// config.toml, the macOS MDM managed-preferences layer, and project
/// .grok/config.toml files). Driven from the same resolvers used at runtime
@ -269,7 +269,7 @@ pub struct ConfigSources {
/// A single config layer entry for `grok inspect`.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ConfigLayer {
pub(crate) struct ConfigLayer {
/// Logical role of the layer: "system-managed", "managed", "user",
/// "system-requirements", "requirements", "mdm", or "project".
pub role: String,

2
crates/codegen/xai-grok-shell/src/leader/client.rs
View file

@ -284,7 +284,7 @@ impl LeaderClient {
///
/// Like [`into_channels()`](Self::into_channels) but also returns the
/// disconnect watch so the caller can observe why the connection ended.
pub fn into_channels_with_disconnect(
pub(crate) fn into_channels_with_disconnect(
self,
) -> (
mpsc::UnboundedSender<String>,

21
crates/codegen/xai-grok-shell/src/leader/lock.rs
View file

@ -201,18 +201,6 @@ impl LeaderLock {
}
}
/// Acquire exclusive lock, blocking until available.
///
/// Used by the leader process on startup. Blocks until the lock is available.
/// After acquiring, call `write_pid()` to record the leader's PID.
pub fn acquire_blocking(&mut self) -> Result<(), LockError> {
let file = self.open_lock_file()?;
file.lock_exclusive()?;
self.mark_acquired(file);
Ok(())
}
/// Acquire exclusive lock with a bounded wait, re-opening the lock-file path
/// on every attempt.
///
@ -224,7 +212,10 @@ impl LeaderLock {
///
/// Async so the 200ms poll yields to the Tokio runtime instead of blocking a
/// worker thread — `run_leader` calls this on the multi-thread runtime.
pub async fn acquire_reopen_timeout(&mut self, timeout: Duration) -> Result<(), LockError> {
pub(crate) async fn acquire_reopen_timeout(
&mut self,
timeout: Duration,
) -> Result<(), LockError> {
let deadline = Instant::now() + timeout;
let poll_interval = Duration::from_millis(200);
@ -263,7 +254,7 @@ impl LeaderLock {
Self::read_pid_from_path(&self.lock_path)
}
pub fn read_pid_from_path(path: &Path) -> Option<u32> {
pub(crate) fn read_pid_from_path(path: &Path) -> Option<u32> {
let mut content = String::new();
File::open(path)
.and_then(|mut f| f.read_to_string(&mut content))
@ -272,7 +263,7 @@ impl LeaderLock {
}
/// Delete the socket file. Call while holding the lock.
pub fn cleanup_socket(&self) -> io::Result<()> {
pub(crate) fn cleanup_socket(&self) -> io::Result<()> {
match fs::remove_file(&self.sock_path) {
Ok(()) => Ok(()),
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),

2
crates/codegen/xai-grok-shell/src/leader/mod.rs
View file

@ -868,7 +868,7 @@ impl LeaderConnection {
/// Like [`into_channels()`](Self::into_channels) but also returns a
/// [`watch::Receiver<DisconnectReason>`] so the caller can observe
/// why the connection ended (e.g., `LeaderShutdown` vs `ConnectionLost`).
pub fn into_channels_with_disconnect(
pub(crate) fn into_channels_with_disconnect(
self,
) -> (
mpsc::UnboundedSender<String>,

90
crates/codegen/xai-grok-shell/src/leader/protocol.rs
View file

@ -19,7 +19,9 @@ pub enum ProtocolError {
ConnectionClosed,
}
pub async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, ProtocolError> {
pub(crate) async fn read_frame<R: AsyncRead + Unpin>(
reader: &mut R,
) -> Result<Vec<u8>, ProtocolError> {
let mut len_buf = [0u8; 4];
match reader.read_exact(&mut len_buf).await {
Ok(_) => {}
@ -39,7 +41,7 @@ pub async fn read_frame<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>,
Ok(buf)
}
pub async fn write_frame<W: AsyncWrite + Unpin>(
pub(crate) async fn write_frame<W: AsyncWrite + Unpin>(
writer: &mut W,
data: &[u8],
) -> Result<(), ProtocolError> {
@ -391,6 +393,71 @@ pub enum ServerMessage {
LeaderReady,
}
/// Extension methods injected into the agent, named as the agent matches them.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumIter)]
pub(crate) enum InternalMethod {
AuthCleared,
EvictSessions,
ReloadAllMcpServers,
ReloadModels,
ReloadModelsCache,
ReloadProjectMcpServers,
ReloadSkills,
ReloadWorkflows,
}
impl InternalMethod {
pub const fn name(self) -> &'static str {
match self {
Self::AuthCleared => "x.ai/internal/auth_cleared",
Self::EvictSessions => "x.ai/internal/evict_sessions",
Self::ReloadAllMcpServers => "x.ai/internal/reload_all_mcp_servers",
Self::ReloadModels => "x.ai/internal/reload_models",
Self::ReloadModelsCache => "x.ai/internal/reload_models_cache",
Self::ReloadProjectMcpServers => "x.ai/internal/reload_project_mcp_servers",
Self::ReloadSkills => "x.ai/internal/reload_skills",
Self::ReloadWorkflows => "x.ai/internal/reload_workflows",
}
}
pub fn from_name(name: &str) -> Option<Self> {
use strum::IntoEnumIterator;
Self::iter().find(|method| method.name() == name)
}
/// The decoder routes a custom method to `ext_method` / `ext_notification`
/// only when it carries the `_` prefix, and rejects the bare name.
fn wire_name(self) -> String {
format!("_{}", self.name())
}
}
/// Not newline-terminated: the `acp_tx` forwarding loop appends the terminator.
pub(crate) fn internal_notification(method: InternalMethod, params: serde_json::Value) -> String {
serde_json::json!({
"jsonrpc": "2.0",
"method": method.wire_name(),
"params": params,
})
.to_string()
}
/// Newline-terminated for direct injection.
pub(crate) fn internal_request_line(
id: &str,
method: InternalMethod,
params: serde_json::Value,
) -> String {
let msg = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": method.wire_name(),
"params": params,
});
format!("{msg}\n")
}
#[cfg(test)]
mod tests {
use super::*;
@ -716,6 +783,25 @@ mod tests {
}
}
#[test]
fn every_internal_method_carries_the_routable_prefix() {
use strum::IntoEnumIterator;
for method in InternalMethod::iter() {
for line in [
internal_notification(method, serde_json::json!({})),
internal_request_line("id", method, serde_json::json!({})),
] {
let json: serde_json::Value = serde_json::from_str(line.trim_end()).unwrap();
assert_eq!(
json["method"].as_str().and_then(|m| m.strip_prefix('_')),
Some(method.name()),
"unroutable wire method: {line}"
);
}
}
}
#[test]
fn shutdown_reason_variants_serialize_correctly() {
let auto = serde_json::to_string(&ShutdownReason::AutoUpdate).unwrap();

23
crates/codegen/xai-grok-shell/src/leader/server.rs
View file

@ -15,8 +15,8 @@ const LEADER_VERSION: &str = match option_env!("VERSION_WITH_COMMIT") {
};
use super::protocol::{
ClientCapabilities, ClientId, ClientMessage, ClientMode, ControlCommand, ControlPayload,
LEADER_PROTOCOL_VERSION, LeaderCapabilities, ProtocolError, ServerMessage, read_message,
write_message,
InternalMethod, LEADER_PROTOCOL_VERSION, LeaderCapabilities, ProtocolError, ServerMessage,
internal_notification, read_message, write_message,
};
use super::transport::{LeaderListener, LeaderStream};
use crate::agent::activity::AgentActivity;
@ -145,7 +145,7 @@ impl LeaderServerControlState {
workspace: Arc::new(WorkspaceControl::new(None)),
}
}
pub fn with_default_hub_url(mut self, default_hub_url: Option<String>) -> Self {
pub(crate) fn with_default_hub_url(mut self, default_hub_url: Option<String>) -> Self {
self.workspace = Arc::new(WorkspaceControl::new(default_hub_url));
self
}
@ -184,7 +184,7 @@ impl WorkspaceControl {
}
/// Wire the hub credential to the leader's shared `AuthManager` (sole
/// owner of refresh + persistence).
pub fn set_auth_manager(&self, auth_manager: Arc<AuthManager>) {
pub(crate) fn set_auth_manager(&self, auth_manager: Arc<AuthManager>) {
self.auth.send_replace(Some(Arc::new(LeaderAuthProvider {
auth_manager,
refresh_in_flight: Arc::new(std::sync::atomic::AtomicBool::new(false)),
@ -1730,12 +1730,10 @@ pub async fn run_leader_server(
last_active_client = None;
}
if !detached_sessions.is_empty() {
let evict_notification = serde_json::json!({
"jsonrpc": "2.0",
"method": "x.ai/internal/evict_sessions",
"params": { "sessionIds": detached_sessions }
});
let _ = acp_tx.send(evict_notification.to_string());
let _ = acp_tx.send(internal_notification(
InternalMethod::EvictSessions,
serde_json::json!({ "sessionIds": detached_sessions }),
));
info!(
client_id = id.0,
session_count = detached_sessions.len(),
@ -5260,7 +5258,10 @@ mod tests {
.expect("channel should not be closed");
let json: serde_json::Value =
serde_json::from_str(&eviction_msg).expect("should be valid JSON");
assert_eq!(json["method"], "x.ai/internal/evict_sessions");
assert_eq!(
json["method"].as_str().and_then(|m| m.strip_prefix('_')),
Some(InternalMethod::EvictSessions.name()),
);
let session_ids = json["params"]["sessionIds"]
.as_array()
.expect("sessionIds should be an array");

1
crates/codegen/xai-grok-shell/src/lib.rs
View file

@ -41,6 +41,5 @@ pub mod terminal;
pub(crate) mod test_support;
pub mod tier;
pub mod tools;
pub mod trace_classifier;
pub mod upload;
pub mod util;

10
crates/codegen/xai-grok-shell/src/plugin.rs
View file

@ -41,7 +41,7 @@ pub struct InstallOutcome {
/// Classify an install source as local (filesystem) vs git (remote) without
/// installing — used for telemetry `install_kind` on the failure path, where no
/// [`InstallOutcome`] is available.
pub fn install_source_is_local(source: &str, cwd: &Path) -> bool {
pub(crate) fn install_source_is_local(source: &str, cwd: &Path) -> bool {
matches!(
git_install::parse_install_source(source, cwd),
git_install::InstallSource::Local { .. }
@ -346,7 +346,7 @@ pub fn update_plugins(name: Option<&str>) -> Result<Vec<RepoUpdateOutcome>, Upda
update_plugins_by_selector(name.map(|name| PluginUpdateSelector::PluginName(name.to_string())))
}
pub fn update_plugins_by_selector(
pub(crate) fn update_plugins_by_selector(
selector: Option<PluginUpdateSelector>,
) -> Result<Vec<RepoUpdateOutcome>, UpdateError> {
let mut registry = InstallRegistry::load();
@ -706,14 +706,14 @@ fn bullet_list(items: &[String]) -> String {
/// The require-sha pin policy for remote plugin code. Disk-only config + env,
/// both tighten-only: a remote campaign overlay must not be able to relax a
/// local security policy, and an unreadable config falls back to the env knob.
pub fn marketplace_require_sha() -> bool {
pub(crate) fn marketplace_require_sha() -> bool {
xai_grok_config::load_effective_config_disk_only()
.map(|c| xai_grok_plugin_marketplace::load_require_sha(&c))
.unwrap_or_else(|_| xai_grok_plugin_marketplace::env_require_sha())
}
/// Marketplace sources from config.toml + settings JSON, unfiltered.
pub fn load_marketplace_sources() -> Vec<MarketplaceSource> {
pub(crate) fn load_marketplace_sources() -> Vec<MarketplaceSource> {
let config = crate::config::load_effective_config()
.ok()
.unwrap_or(toml::Value::Table(toml::map::Map::new()));
@ -725,7 +725,7 @@ pub fn load_marketplace_sources() -> Vec<MarketplaceSource> {
/// Like [`load_marketplace_sources`] but drops git sources blocked by the
/// managed `marketplace_allowlist`. Install paths must use this so policy
/// cannot be bypassed.
pub fn load_filtered_marketplace_sources() -> Vec<MarketplaceSource> {
pub(crate) fn load_filtered_marketplace_sources() -> Vec<MarketplaceSource> {
let allowlist =
&xai_grok_workspace::permission::resolution::managed_settings().marketplace_allowlist;
filter_sources_by_allowlist(load_marketplace_sources(), allowlist)

8
crates/codegen/xai-grok-shell/src/relay/sync.rs
View file

@ -26,7 +26,7 @@ const DROP_BATCH_SIZE: usize = 64;
/// Build the share URL for a session.
/// Format: https://grok.com/build/{sessionId}
pub fn build_share_url(session_id: &str) -> String {
pub(crate) fn build_share_url(session_id: &str) -> String {
let base_url =
std::env::var("GROK_CODE_WEB_URL").unwrap_or_else(|_| "https://grok.com".to_string());
format!("{}/build/{}", base_url, session_id)
@ -143,7 +143,7 @@ impl RelaySyncState {
}
/// Update the cursor after a successful sync.
pub fn update_cursor(&mut self, event_id: String) {
pub(crate) fn update_cursor(&mut self, event_id: String) {
self.last_synced_event_id = Some(event_id);
self.last_synced_at = Some(
std::time::SystemTime::now()
@ -309,7 +309,7 @@ impl RelaySync {
}
/// Get the current connection state.
pub fn connection_state(&self) -> ConnectionState {
pub(crate) fn connection_state(&self) -> ConnectionState {
*self.connection_state_rx.borrow()
}
@ -325,7 +325,7 @@ impl RelaySync {
}
/// Subscribe to connection state changes.
pub fn subscribe_state(&self) -> watch::Receiver<ConnectionState> {
pub(crate) fn subscribe_state(&self) -> watch::Receiver<ConnectionState> {
self.connection_state_rx.clone()
}
}

112
crates/codegen/xai-grok-shell/src/remote/agent.rs
View file

@ -125,7 +125,7 @@ impl SandboxClient {
}
/// Terminate a sandbox session.
pub async fn terminate_session(
pub(crate) async fn terminate_session(
&self,
session_id: &str,
request: &SandboxTerminateRequest,
@ -152,78 +152,6 @@ impl SandboxClient {
// Session Lifecycle
// ========================================================================
/// Start a sandbox session (non-TUI).
pub async fn start_session(
&self,
request: &SandboxStartRequest,
) -> Result<SandboxStartResponse> {
let url = format!("{}/sandbox/sessions/start", self.base_url);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send start session request")?;
Self::parse_response(response, "start session").await
}
/// Get sandbox session status.
pub async fn get_session_status(&self, session_id: &str) -> Result<SandboxStatusResponse> {
let url = format!("{}/sandbox/sessions/{}/status", self.base_url, session_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get session status request")?;
Self::parse_response(response, "get session status").await
}
/// Get sandbox session logs.
pub async fn get_session_logs(&self, session_id: &str) -> Result<SandboxLogsResponse> {
let url = format!("{}/sandbox/sessions/{}/logs", self.base_url, session_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get session logs request")?;
Self::parse_response(response, "get session logs").await
}
/// Hibernate a sandbox session (snapshot rootfs to GCS and terminate).
pub async fn hibernate_session(&self, session_id: &str) -> Result<SandboxHibernateResponse> {
let url = format!(
"{}/sandbox/sessions/{}/hibernate",
self.base_url, session_id
);
let response = self
.auth_headers(self.client.post(&url))
.await?
.send()
.await
.context("failed to send hibernate session request")?;
Self::parse_response(response, "hibernate session").await
}
/// Restore a previously hibernated sandbox session from its snapshot.
pub async fn restore_session(
&self,
session_id: &str,
request: &SandboxRestoreRequest,
) -> Result<SandboxRestoreResponse> {
let url = format!("{}/sandbox/sessions/{}/restore", self.base_url, session_id);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send restore session request")?;
Self::parse_response(response, "restore session").await
}
// ========================================================================
// Environment CRUD
// ========================================================================
@ -249,7 +177,7 @@ impl SandboxClient {
}
/// Create a new sandbox environment.
pub async fn create_environment(
pub(crate) async fn create_environment(
&self,
request: &SandboxCreateEnvironmentRequest,
) -> Result<SandboxEnvironmentResponse> {
@ -264,23 +192,8 @@ impl SandboxClient {
Self::parse_response(response, "create environment").await
}
/// Get a sandbox environment by ID.
pub async fn get_environment(
&self,
environment_id: &str,
) -> Result<SandboxEnvironmentResponse> {
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get environment request")?;
Self::parse_response(response, "get environment").await
}
/// Update a sandbox environment.
pub async fn update_environment(
pub(crate) async fn update_environment(
&self,
environment_id: &str,
request: &SandboxUpdateEnvironmentRequest,
@ -297,7 +210,7 @@ impl SandboxClient {
}
/// Delete a sandbox environment.
pub async fn delete_environment(&self, environment_id: &str) -> Result<()> {
pub(crate) async fn delete_environment(&self, environment_id: &str) -> Result<()> {
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
let response = self
.auth_headers(self.client.delete(&url))
@ -307,21 +220,4 @@ impl SandboxClient {
.context("failed to send delete environment request")?;
Self::check_response(response, "delete environment").await
}
/// List preinstalled packages available for sandbox environments.
pub async fn list_preinstalled_packages(
&self,
) -> Result<SandboxListPreinstalledPackagesResponse> {
let url = format!(
"{}/sandbox/environments/preinstalled-packages",
self.base_url
);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send list preinstalled packages request")?;
Self::parse_response(response, "list preinstalled packages").await
}
}

5
crates/codegen/xai-grok-shell/src/remote/chat_models_client.rs
View file

@ -107,7 +107,10 @@ impl ChatModelsClient {
/// Gated only on a valid grok.com bearer — deliberately NOT `is_xai_auth()`
/// (unlike workspaces/conversations), since `/rest/modes` is the public chat
/// endpoint and that gate would exclude API-key / cached-token chat users.
pub async fn list_modes(&self, locale: &str) -> Result<ListModesResponse, ChatModelsError> {
pub(crate) async fn list_modes(
&self,
locale: &str,
) -> Result<ListModesResponse, ChatModelsError> {
let auth = self
.auth
.auth()

44
crates/codegen/xai-grok-shell/src/remote/client.rs
View file

@ -209,16 +209,16 @@ async fn fetch_bundle_inner(
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ShareResponse {
pub(crate) struct ShareResponse {
pub permission_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadDataResponse {
pub(crate) struct LoadDataResponse {
pub messages: Option<Vec<LoadedMessage>>,
pub session: Option<SessionInfo>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadedMessage {
pub(crate) struct LoadedMessage {
pub id: String,
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
@ -237,14 +237,14 @@ pub struct SessionInfo {
pub metadata: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SaveDataRequest {
pub(crate) struct SaveDataRequest {
pub messages: Vec<ExportedMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UpsertSessionRequest {
pub(crate) struct UpsertSessionRequest {
pub session: SessionUpdate,
pub agent_id: String,
}
@ -322,7 +322,10 @@ impl BackendClient {
}
/// Attach a live `AuthManager` so every request resolves a fresh token
/// instead of requiring the caller to pass `&GrokAuth`.
pub fn with_auth_manager(mut self, manager: std::sync::Arc<crate::auth::AuthManager>) -> Self {
pub(crate) fn with_auth_manager(
mut self,
manager: std::sync::Arc<crate::auth::AuthManager>,
) -> Self {
let credentials: std::sync::Arc<dyn xai_grok_auth::AuthCredentialProvider> =
std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
@ -384,22 +387,6 @@ impl BackendClient {
let share_response = self.create_share_link(&session.session_id).await?;
Ok(share_url(&share_response.permission_id))
}
/// Sync session to backend without creating a share link.
pub async fn sync_session(
&self,
session: &ExportedSession,
agent_id: &str,
) -> Result<(), BackendError> {
self.upsert_session(&session.session_id, &session.metadata, agent_id)
.await?;
self.save_session_data(
&session.session_id,
&session.messages,
Some(&session.metadata),
)
.await?;
Ok(())
}
/// Build auth + identity headers.
/// Must include X-XAI-Token-Auth so nginx auth subrequest routes to authenticate_xai_grok_cli_token.
/// See: crates/codegen/xai-grok-shell/src/agent/app.rs:run_headless
@ -474,7 +461,7 @@ impl BackendClient {
}
Ok(())
}
pub async fn save_session_data(
pub(crate) async fn save_session_data(
&self,
session_id: &str,
messages: &[ExportedMessage],
@ -511,7 +498,7 @@ impl BackendClient {
let data: ListResponse = response.json().await?;
Ok(data.sessions)
}
pub async fn load_session_data(
pub(crate) async fn load_session_data(
&self,
session_id: &str,
) -> Result<LoadDataResponse, BackendError> {
@ -530,7 +517,10 @@ impl BackendClient {
let data: LoadDataResponse = response.json().await?;
Ok(data)
}
pub async fn create_share_link(&self, session_id: &str) -> Result<ShareResponse, BackendError> {
pub(crate) async fn create_share_link(
&self,
session_id: &str,
) -> Result<ShareResponse, BackendError> {
let url = format!("{}/sessions/{}/share", self.base_url, session_id);
let response = self.send_with_auth(self.reqwest_client.post(&url)).await?;
if !response.status().is_success() {
@ -541,7 +531,7 @@ impl BackendClient {
let share_response: ShareResponse = response.json().await?;
Ok(share_response)
}
pub async fn delete_session_data(&self, session_id: &str) -> Result<(), BackendError> {
pub(crate) async fn delete_session_data(&self, session_id: &str) -> Result<(), BackendError> {
let url = format!("{}/sessions/{}/data", self.base_url, session_id);
let response = self
.send_with_auth(self.reqwest_client.delete(&url))
@ -839,7 +829,7 @@ pub(crate) fn fetch_models_blocking(
}
/// Parse a single model entry from the /models-v2 response.
/// Used by both initial model fetch and session-resume metadata refresh.
pub fn parse_remote_model_value(
pub(crate) fn parse_remote_model_value(
value: &serde_json::Value,
default_base_url: &str,
) -> Option<crate::agent::config::ModelEntryConfig> {

5
crates/codegen/xai-grok-shell/src/remote/conversations_client.rs
View file

@ -223,7 +223,10 @@ impl ConversationsClient {
}
/// `DELETE /rest/app-chat/conversations/soft/{conversation_id}` — soft-delete.
pub async fn soft_delete_conversation(&self, conversation_id: &str) -> Result<(), ConvError> {
pub(crate) async fn soft_delete_conversation(
&self,
conversation_id: &str,
) -> Result<(), ConvError> {
let auth = self.require_xai_auth().await?;
let url = format!(
"{}/rest/app-chat/conversations/soft/{}",

4
crates/codegen/xai-grok-shell/src/remote/mod.rs
View file

@ -1,13 +1,13 @@
//! Remote storage client for the backend.
pub mod agent;
pub mod chat_models_client;
pub(crate) mod chat_models_client;
pub mod client;
pub mod conversations_client;
pub mod pull;
#[cfg(test)]
mod pull_smoke_test;
pub mod skills_client;
pub(crate) mod skills_client;
pub mod sync;
pub mod workspaces_client;

4
crates/codegen/xai-grok-shell/src/remote/skills_client.rs
View file

@ -105,7 +105,7 @@ impl ProductSkillsCatalog {
///
/// Mirrors grok-web: enabled user skills first (and hide same-named
/// bundled entries they override), then remaining bundled skills.
pub fn to_skill_infos(&self) -> Vec<SkillInfo> {
pub(crate) fn to_skill_infos(&self) -> Vec<SkillInfo> {
let enabled_user: Vec<(String, &UserSkill)> = self
.user
.iter()
@ -662,7 +662,7 @@ impl SkillsClient {
/// alt while primary is tenant-tagged. Callers still success-cache under
/// the **primary** identity (team/org of primary) so the same session
/// hits TTL; personal primaries cannot match that entry.
pub async fn try_list_catalog(
pub(crate) async fn try_list_catalog(
&self,
locale: &str,
) -> Result<(ProductSkillsCatalog, bool), SkillsError> {

2
crates/codegen/xai-grok-shell/src/remote/sync.rs
View file

@ -80,7 +80,7 @@ impl RemoteSync {
let _ = self.tx.send(SyncMsg::SetTitle(title));
}
pub fn set_model_id(&self, model_id: String) {
pub(crate) fn set_model_id(&self, model_id: String) {
let _ = self.tx.send(SyncMsg::SetModelId(model_id));
}
}

2
crates/codegen/xai-grok-shell/src/remote/workspaces_client.rs
View file

@ -75,7 +75,7 @@ impl WorkspacesClient {
}
}
pub async fn list_workspaces(&self, q: &WsQuery) -> Result<ListWorkspacesPage, WsError> {
pub(crate) async fn list_workspaces(&self, q: &WsQuery) -> Result<ListWorkspacesPage, WsError> {
let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?;
if !auth.is_xai_auth() {
return Err(WsError::NoOauth);

11
crates/codegen/xai-grok-shell/src/sampling/error.rs
View file

@ -102,7 +102,7 @@ pub const OVERLOADED_USER_MESSAGE: &str = "Model is temporarily overloaded. Try
/// Map a `SamplingError` to an ACP `Error` for client-facing responses.
/// This stays in xai-grok-shell because it depends on `agent_client_protocol::Error`.
pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
pub(crate) fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
use reqwest::StatusCode;
// Capacity/overload gets the same short copy on every surface. Message
// only, `data` deliberately unset: `Display` appends JSON-encoded `data`,
@ -187,7 +187,10 @@ pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
}
}
pub fn error_data_with_status(message: String, http_status: Option<u16>) -> serde_json::Value {
pub(crate) fn error_data_with_status(
message: String,
http_status: Option<u16>,
) -> serde_json::Value {
match http_status {
Some(sc) => serde_json::json!({ "message": message, "http_status": sc }),
None => serde_json::Value::String(message),
@ -195,7 +198,7 @@ pub fn error_data_with_status(message: String, http_status: Option<u16>) -> serd
}
/// Terminal-failure `acp::Error.data`: max-tokens truncation carries an `error_kind` marker (the kind's stable `as_str` name); other kinds keep the legacy shape.
pub fn terminal_error_data(
pub(crate) fn terminal_error_data(
message: String,
http_status: Option<u16>,
kind: xai_grok_sampler::SamplingErrorKind,
@ -297,7 +300,7 @@ pub fn prompt_usage_from_error(
/// notification from a prompt result. Rate-limit errors produce
/// `("rate_limit", null)` so the client shows its own upgrade message;
/// other errors produce `("error", <detail>)`.
pub fn prompt_complete_fields(
pub(crate) fn prompt_complete_fields(
result: &std::result::Result<acp::StopReason, acp::Error>,
) -> (serde_json::Value, serde_json::Value) {
match result {

12
crates/codegen/xai-grok-shell/src/session/acp_conversion.rs
View file

@ -24,7 +24,7 @@ use xai_tool_types::{KillTaskOutput, TaskOutputOutput};
/// directory (e.g., `/root/.grok/worktrees/project/fork-019cb252-...`). The
/// client UI should instead see the original project path (the `display_cwd`).
#[derive(Clone, Debug)]
pub struct PathRewriter {
pub(crate) struct PathRewriter {
/// The real worktree path (what tools actually see).
real_cwd: String,
/// The display path (what the client UI should see).
@ -67,7 +67,7 @@ impl PathRewriter {
}
/// Rewrite a `PathBuf` if it starts with the real worktree path.
pub fn rewrite_path(&self, path: &Path) -> PathBuf {
pub(crate) fn rewrite_path(&self, path: &Path) -> PathBuf {
match path.strip_prefix(&self.real_cwd) {
Ok(relative) => PathBuf::from(&self.display_cwd).join(relative),
Err(_) => path.to_path_buf(),
@ -80,7 +80,7 @@ impl PathRewriter {
/// paths embedded anywhere in the JSON tree without needing to walk the
/// structure. Reuses `rewrite()` so both plain and encoded replacements
/// are applied consistently.
pub fn rewrite_json(&self, value: serde_json::Value) -> serde_json::Value {
pub(crate) fn rewrite_json(&self, value: serde_json::Value) -> serde_json::Value {
let serialized = value.to_string();
let rewritten = self.rewrite(&serialized);
if rewritten == serialized {
@ -110,7 +110,7 @@ fn maybe_rewrite_path(rewriter: Option<&PathRewriter>, path: PathBuf) -> PathBuf
///
/// Uses serde directly — ToolOutput derives Serialize with `#[serde(tag = "type")]`,
/// so the JSON round-trips cleanly with the TUI's deserialization.
pub fn raw_output_json(
pub(crate) fn raw_output_json(
output: &ToolOutput,
rewriter: Option<&PathRewriter>,
) -> Option<serde_json::Value> {
@ -129,7 +129,7 @@ pub fn raw_output_json(
/// `tool_meta` is attached as `_meta` on the update for MCP tools that have
/// MCP Apps UI metadata (e.g., `_meta.ui.resourceUri`). This allows clients
/// to render interactive UIs without maintaining a separate metadata store.
pub fn acp_tool_update(
pub(crate) fn acp_tool_update(
output: &ToolOutput,
tool_call_id: &str,
rewriter: Option<&PathRewriter>,
@ -644,7 +644,7 @@ pub fn acp_tool_update(
/// This converts `xai-grok-tools`' TodoItem (which has `id`, `content: Option<String>`,
/// `status: Option<String>`) to `acp::PlanEntry` (which has `content`, `priority`, `status`).
/// The `id` is not directly represented in `PlanEntry` but the ordering is preserved.
pub fn acp_plan_update(output: &ToolOutput) -> Option<acp::Plan> {
pub(crate) fn acp_plan_update(output: &ToolOutput) -> Option<acp::Plan> {
use crate::tools::todo::plan_entry_from_todo_item;
use xai_grok_tools::types::output::TodoWriteOutput;
match output {

Some files were not shown because too many files have changed in this diff Show more