diff --git a/.devrefs/devrefs.yml b/.devrefs/devrefs.yml index 5d33098..2d7da1b 100644 --- a/.devrefs/devrefs.yml +++ b/.devrefs/devrefs.yml @@ -13,3 +13,7 @@ references: description: "a ratatui agent harness that we can use as reference, especially since it's by openai" remote: https://github.com/openai/codex.git branch: main + - id: xai-org/grok-build + description: "source code for xai grok" + remote: https://github.com/xai-org/grok-build.git + branch: main diff --git a/README.md b/README.md index 1480a2d..ac85978 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ A purely Rust-based AI CLI coding agent with a beautiful terminal UI for interac - **Multiple Model Support** - Works w/ the same models.dev support. - **Command System** - Intuitive commands: `/sessions`, `/new`, `/connect`, `/models`, `/exit` + custom commands. - **Session Management** - Create and manage multiple chat sessions - - **Streaming Responses** - Real-time streaming of AI responses (w/ [aisdk.rs](https://aisdk.rs)) + - **Streaming Responses** - Real-time streaming of AI responses + websockets using OpenAI. ## Installation @@ -118,9 +118,9 @@ Like any benchmark, please take this with a grain of salt. I have a cherry-picke | Agent | Score | Checks | Avg time | Est. tokens | Est. cost | | ----------- | ----: | -----: | -------: | ----------: | --------: | -| 🦀 crabcode | 100% | 19/19 | 21.3s | 2851 | $0.0103 | -| 🔲 opencode | 100% | 19/19 | 24.2s | 3837 | $0.0201 | -| ⚛️ codex | 100% | 19/19 | 24.9s | 56172 | $0.5435 | +| 🦀 crabcode | 100% | 19/19 | 29.8s | 2768 | $0.0094 | +| 🔲 opencode | 100% | 19/19 | 34.9s | 4612 | $0.0279 | +| ⚛️ codex | 100% | 19/19 | 33.7s | 36888 | $0.3506 | ## Contributing diff --git a/_docs/__PARITY.md b/_docs/__PARITY.md index 73a97f8..3ea1211 100644 --- a/_docs/__PARITY.md +++ b/_docs/__PARITY.md @@ -25,7 +25,7 @@ Recent implementation notes: | 1.7 | Configurable max steps per agent | Per-agent `max_steps` controls agent loop depth. | **Present.** Agent registry supports `steps`, `maxSteps`, and `max_steps`; app/print/task paths read max steps from registry definitions. | No material max-step config gap beyond the duplicated runner logic noted in 1.3. | | 2.1 | Provider-specific header and behavior instructions | Provider/model-specific prompt variants for Beast/OpenAI, Anthropic, Gemini, and Codex. | **Mostly present.** `src/prompt/mod.rs` has OpenAI, Anthropic, Gemini, Codex, and Generic prompt branches. | Selection is based on model-id string heuristics, not resolved provider kind/model metadata, so OpenAI-compatible or renamed models can get the wrong prompt. | | 2.2 | Environment context block | Includes workdir, git status, platform, and date. | **Present.** `SystemPromptComposer::get_environment_context` emits `` with working directory, git-repo flag, platform, and date. | No material harness gap. | -| 2.3 | Tool schemas block | System prompt lists registered tool schemas as JSON. | **Present.** `SystemPromptComposer::with_tool_registry` emits JSON schemas; app and print mode compose prompts with scoped dynamic registries. | Schemas are scoped to visible tools for the current mode, not literally every registered tool; registry ordering is not deterministic because it is backed by a `HashMap`. | +| 2.3 | Tool schemas block | System prompt lists registered tool schemas as JSON. | **Changed.** `SystemPromptComposer::with_tool_registry` now emits tool **names** + short guidance only; full schemas go solely via the provider `tools` field (OpenCode/Codex/Grok Build). | Avoids double-paying schema tokens every step; registry ordering is still non-deterministic because it is backed by a `HashMap`. | | 2.4 | Custom instructions discovery | Walk-up discovery for `AGENTS.md`/`CLAUDE.md` plus global fallback. | **Partial.** `src/prompt/rules.rs` walks upward for nearest local `AGENTS.md` or `CLAUDE.md`; global fallback checks `$XDG_CONFIG_HOME/crabcode/AGENTS.md` and `~/.claude/CLAUDE.md`. | Walk-up does not stop at the git/project root; it does not load OpenCode global instruction locations; it returns the first local match instead of a layered instruction stack. | | 2.5 | Available skills XML block | Prompt lists discovered skills as ``. | **Present.** `src/prompt/mod.rs` appends skills from `SkillStore`; `src/tools/skill.rs` repeats them in the tool description. | No material gap beyond discovery/permission gaps listed below. | | 2.6 | Available subagents XML block | Prompt lists subagent names/descriptions so the primary agent can use Task. | **Mostly present.** `src/prompt/mod.rs` emits visible subagent definitions from `AgentRegistry`, including config-defined agents. | Missing OpenCode built-in `scout` and `vlm-agent`; hidden agents are intentionally omitted from prompt listings. | diff --git a/npm/README.md b/npm/README.md index 1480a2d..ac85978 100644 --- a/npm/README.md +++ b/npm/README.md @@ -23,7 +23,7 @@ A purely Rust-based AI CLI coding agent with a beautiful terminal UI for interac - **Multiple Model Support** - Works w/ the same models.dev support. - **Command System** - Intuitive commands: `/sessions`, `/new`, `/connect`, `/models`, `/exit` + custom commands. - **Session Management** - Create and manage multiple chat sessions - - **Streaming Responses** - Real-time streaming of AI responses (w/ [aisdk.rs](https://aisdk.rs)) + - **Streaming Responses** - Real-time streaming of AI responses + websockets using OpenAI. ## Installation @@ -118,9 +118,9 @@ Like any benchmark, please take this with a grain of salt. I have a cherry-picke | Agent | Score | Checks | Avg time | Est. tokens | Est. cost | | ----------- | ----: | -----: | -------: | ----------: | --------: | -| 🦀 crabcode | 100% | 19/19 | 21.3s | 2851 | $0.0103 | -| 🔲 opencode | 100% | 19/19 | 24.2s | 3837 | $0.0201 | -| ⚛️ codex | 100% | 19/19 | 24.9s | 56172 | $0.5435 | +| 🦀 crabcode | 100% | 19/19 | 29.8s | 2768 | $0.0094 | +| 🔲 opencode | 100% | 19/19 | 34.9s | 4612 | $0.0279 | +| ⚛️ codex | 100% | 19/19 | 33.7s | 36888 | $0.3506 | ## Contributing diff --git a/src/agent/config.rs b/src/agent/config.rs index 1e4fe53..c599fed 100644 --- a/src/agent/config.rs +++ b/src/agent/config.rs @@ -10,6 +10,9 @@ pub struct OpenAIRequestOptions { pub default_instructions: Option, pub disallow_system_messages: bool, pub force_tool_strict_false: bool, + /// Sticky prompt-cache routing key (Responses / chat-completions). + /// Typically the crabcode session id. + pub prompt_cache_key: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -83,6 +86,8 @@ pub struct LlmSessionConfig { pub reasoning_effort: Option, pub supports_image_input: bool, pub openai_options: OpenAIRequestOptions, + /// Sticky prompt-cache key for this session (OpenAI/xAI/compatible). + pub prompt_cache_key: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -131,6 +136,7 @@ mod tests { reasoning_effort: None, supports_image_input: false, openai_options: OpenAIRequestOptions::default(), + prompt_cache_key: None, } } diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index 4e9bcae..0ee00eb 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -323,6 +323,13 @@ async fn start_subagent_stream( if let Some(effort) = session.reasoning_effort { builder = builder.reasoning_effort(effort.as_str()); } + if let Some(cache_key) = session + .prompt_cache_key + .as_deref() + .or(session.openai_options.prompt_cache_key.as_deref()) + { + builder = builder.prompt_cache_key(cache_key); + } let provider = builder .build() .map_err(|e| format!("Failed to build OpenAICompatible provider: {}", e))?; @@ -389,6 +396,13 @@ async fn start_subagent_stream( if session.openai_options.force_tool_strict_false { builder = builder.tool_strict_override(false); } + if let Some(cache_key) = session + .prompt_cache_key + .as_deref() + .or(session.openai_options.prompt_cache_key.as_deref()) + { + builder = builder.prompt_cache_key(cache_key); + } if !session.openai_options.additional_headers.is_empty() { builder = builder.headers(session.openai_options.additional_headers.clone()); } @@ -615,6 +629,7 @@ mod tests { reasoning_effort, supports_image_input: false, openai_options: crate::agent::config::OpenAIRequestOptions::default(), + prompt_cache_key: None, } } } diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index 71b268b..d252f88 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -96,7 +96,7 @@ impl Provider for Anthropic { ) -> Result { let url = format!("{}/v1/messages", self.base_url.trim_end_matches('/')); - let system_prompts: Vec = messages + let mut system_prompts: Vec = messages .iter() .filter_map(|m| match m { Message::System(s) => Some(serde_json::json!({ @@ -107,6 +107,17 @@ impl Provider for Anthropic { }) .collect(); + // Anthropic prompt caching: mark the last system block ephemeral so the + // stable prefix can be reused across tool steps in a session. + if let Some(last) = system_prompts.last_mut() { + if let Some(obj) = last.as_object_mut() { + obj.insert( + "cache_control".to_string(), + serde_json::json!({ "type": "ephemeral" }), + ); + } + } + let user_messages: Vec = messages .iter() .filter_map(|m| match m { diff --git a/src/aisdk/providers/compatible.rs b/src/aisdk/providers/compatible.rs index 6179ee5..8507b57 100644 --- a/src/aisdk/providers/compatible.rs +++ b/src/aisdk/providers/compatible.rs @@ -18,6 +18,7 @@ pub struct OpenAICompatible { model_name: String, provider_name: String, reasoning_effort: Option, + prompt_cache_key: Option, } impl OpenAICompatible { @@ -33,6 +34,7 @@ pub struct OpenAICompatibleBuilder { model_name: Option, provider_name: Option, reasoning_effort: Option, + prompt_cache_key: Option, } impl OpenAICompatibleBuilder { @@ -61,6 +63,11 @@ impl OpenAICompatibleBuilder { self } + pub fn prompt_cache_key(mut self, key: impl Into) -> Self { + self.prompt_cache_key = Some(key.into()); + self + } + pub fn build(self) -> Result { Ok(OpenAICompatible { base_url: self @@ -74,6 +81,7 @@ impl OpenAICompatibleBuilder { .provider_name .unwrap_or_else(|| "openai-compatible".to_string()), reasoning_effort: self.reasoning_effort, + prompt_cache_key: self.prompt_cache_key, }) } } @@ -134,6 +142,12 @@ impl Provider for OpenAICompatible { body["reasoning_effort"] = serde_json::Value::String(effort.clone()); } + if let Some(key) = &self.prompt_cache_key { + if !key.is_empty() { + body["prompt_cache_key"] = serde_json::Value::String(key.clone()); + } + } + let mut request_headers = reqwest::header::HeaderMap::new(); request_headers.insert( reqwest::header::CONTENT_TYPE, diff --git a/src/aisdk/providers/openai.rs b/src/aisdk/providers/openai.rs index dc8dc5e..d716065 100644 --- a/src/aisdk/providers/openai.rs +++ b/src/aisdk/providers/openai.rs @@ -28,6 +28,14 @@ const OPENAI_WEBSOCKET_IO_TIMEOUT: Duration = Duration::from_secs(300); const OPENAI_WEBSOCKET_STREAM_RETRIES: usize = 1; const OPENAI_WEBSOCKET_FAILURES_BEFORE_FALLBACK: usize = 5; +#[async_trait] +pub trait HttpResponseRetryPolicy: Send + Sync + std::fmt::Debug { + async fn retry_headers( + &self, + status: reqwest::StatusCode, + ) -> Option; +} + #[derive(Debug, Clone)] pub struct OpenAI { base_url: String, @@ -42,6 +50,8 @@ pub struct OpenAI { default_instructions: Option, reasoning_effort: Option, responses_websocket: bool, + prompt_cache_key: Option, + response_retry_policy: Option>, websocket_state: Arc>, } @@ -77,6 +87,8 @@ pub struct OpenAIBuilder { default_instructions: Option, reasoning_effort: Option, responses_websocket: bool, + prompt_cache_key: Option, + response_retry_policy: Option>, } impl OpenAIBuilder { @@ -140,6 +152,16 @@ impl OpenAIBuilder { self } + pub fn prompt_cache_key(mut self, key: impl Into) -> Self { + self.prompt_cache_key = Some(key.into()); + self + } + + pub fn response_retry_policy(mut self, policy: Arc) -> Self { + self.response_retry_policy = Some(policy); + self + } + pub fn build(self) -> Result { let base_url = self .base_url @@ -174,6 +196,8 @@ impl OpenAIBuilder { default_instructions: self.default_instructions, reasoning_effort: self.reasoning_effort, responses_websocket: self.responses_websocket, + prompt_cache_key: self.prompt_cache_key, + response_retry_policy: self.response_retry_policy, websocket_state: Arc::new(Mutex::new(OpenAIWebsocketState::default())), }) } @@ -364,9 +388,9 @@ impl Provider for OpenAI { )) .build() .map_err(|e| Error::Provider(format!("Failed to build client: {}", e)))?; - let response = client + let mut response = client .post(&url) - .headers(request_headers) + .headers(request_headers.clone()) .json(&body) .send() .await @@ -379,6 +403,28 @@ impl Provider for OpenAI { ))) })?; + if let Some(policy) = &self.response_retry_policy { + if let Some(retry_headers) = policy.retry_headers(response.status()).await { + request_headers.extend(retry_headers); + response = client + .post(&url) + .headers(request_headers) + .json(&body) + .send() + .await + .map_err(|err| { + Error::RetryableProvider(RetryError::from_message( + format_openai_request_error( + "send_after_response_retry_policy", + &url, + &err, + Some(&request_diagnostics), + ), + )) + })?; + } + } + if !response.status().is_success() { let status = response.status(); let response_url = sanitized_url(response.url()); @@ -530,6 +576,12 @@ impl OpenAI { body["reasoning"] = serde_json::json!({ "effort": effort }); } + if let Some(key) = &self.prompt_cache_key { + if !key.is_empty() { + body["prompt_cache_key"] = serde_json::Value::String(key.clone()); + } + } + body } @@ -2100,6 +2152,27 @@ mod tests { assert!(body.get("tools").is_none()); assert!(body.get("tool_choice").is_none()); assert!(body.get("parallel_tool_calls").is_none()); + assert!(body.get("prompt_cache_key").is_none()); + } + + #[test] + fn responses_body_includes_prompt_cache_key_and_store_false() { + let provider = OpenAI::builder() + .base_url("https://api.x.ai") + .api_key("test-key") + .model_name("grok-composer-2.5-fast") + .prompt_cache_key("session-abc") + .store_override(false) + .build() + .unwrap(); + + let body = provider.build_responses_body( + vec![serde_json::json!({"role": "user", "content": "hi"})], + &[], + ); + + assert_eq!(body["prompt_cache_key"], "session-abc"); + assert_eq!(body["store"], false); } #[test] diff --git a/src/aisdk/response.rs b/src/aisdk/response.rs index 7dfc3c4..1e0b953 100644 --- a/src/aisdk/response.rs +++ b/src/aisdk/response.rs @@ -14,6 +14,15 @@ use tokio::sync::mpsc; const PHASELESS_AMBIGUOUS_FOLLOW_UP_LIMIT: usize = 1; const PROVIDER_STEP_MAX_RETRIES: usize = 10; +/// Keep the newest N tool outputs intact for the model; older ones are pruned. +/// Matches Grok Build / OpenCode mid-session tool-result retention behavior. +const KEEP_RECENT_TOOL_OUTPUTS: usize = 6; +/// Soft-trim threshold for older-but-still-retained tool outputs (chars). +const TOOL_OUTPUT_SOFT_TRIM_CHARS: usize = 4_000; +const TOOL_OUTPUT_SOFT_TRIM_HEAD: usize = 1_500; +const TOOL_OUTPUT_SOFT_TRIM_TAIL: usize = 1_500; +const PRUNED_TOOL_OUTPUT_PLACEHOLDER: &str = "[Old tool result content cleared]"; + pub struct StreamTextResponse { pub stream: LanguageModelStream, stop_reason: Arc>>, @@ -127,6 +136,13 @@ pub async fn stream_with_tools( } let step_summary = provider_step_log_summary(¤t_messages, &tools); + let pruned = prune_stale_tool_outputs_in_place(&mut current_messages); + if pruned > 0 { + let _ = tx_loop.send(ChunkType::Metadata(format!( + "tool_outputs_pruned count={} keep_recent={}", + pruned, KEEP_RECENT_TOOL_OUTPUTS + ))); + } let _ = tx_loop.send(ChunkType::Metadata(format!( "provider_step_start step={} messages={} tools={} {}", step_idx, @@ -758,6 +774,86 @@ fn rollback_provider_attempt( *emitted_non_replayable_output = false; } +/// Prune older tool outputs in the live multi-step transcript before each +/// provider request. Durable UI/history copies are left intact — this only +/// mutates the request-facing message list. +/// +/// Strategy (Grok Build / OpenCode inspired): +/// - Keep the newest [`KEEP_RECENT_TOOL_OUTPUTS`] results full-size +/// - Soft-trim large older results to head+tail +/// - Hard-clear anything older than 2× the keep window to a placeholder +/// - Drop attached images on pruned outputs (base64 is extremely expensive) +fn prune_stale_tool_outputs_in_place(messages: &mut [Message]) -> usize { + let tool_output_indices: Vec = messages + .iter() + .enumerate() + .filter_map(|(idx, message)| matches!(message, Message::ToolOutput(_)).then_some(idx)) + .collect(); + + if tool_output_indices.is_empty() { + return 0; + } + + let keep_from = tool_output_indices + .len() + .saturating_sub(KEEP_RECENT_TOOL_OUTPUTS); + let hard_clear_before = tool_output_indices + .len() + .saturating_sub(KEEP_RECENT_TOOL_OUTPUTS.saturating_mul(2)); + let mut pruned = 0usize; + + for (rank, &idx) in tool_output_indices.iter().enumerate() { + if rank >= keep_from { + continue; + } + + let Message::ToolOutput(output) = &mut messages[idx] else { + continue; + }; + + let had_images = !output.images.is_empty(); + let original_len = output.output.len(); + if rank < hard_clear_before { + if original_len > PRUNED_TOOL_OUTPUT_PLACEHOLDER.len() || had_images { + output.output = PRUNED_TOOL_OUTPUT_PLACEHOLDER.to_string(); + output.images.clear(); + pruned += 1; + } + continue; + } + + let trimmed = soft_trim_tool_output(&output.output); + if trimmed.len() < original_len || had_images { + output.output = trimmed; + output.images.clear(); + pruned += 1; + } + } + + pruned +} + +fn soft_trim_tool_output(text: &str) -> String { + let char_count = text.chars().count(); + if char_count <= TOOL_OUTPUT_SOFT_TRIM_CHARS { + return text.to_string(); + } + + let head: String = text.chars().take(TOOL_OUTPUT_SOFT_TRIM_HEAD).collect(); + let tail: String = text + .chars() + .rev() + .take(TOOL_OUTPUT_SOFT_TRIM_TAIL) + .collect::() + .chars() + .rev() + .collect(); + let omitted = char_count + .saturating_sub(TOOL_OUTPUT_SOFT_TRIM_HEAD) + .saturating_sub(TOOL_OUTPUT_SOFT_TRIM_TAIL); + format!("{head}\n\n...[{omitted} chars truncated]...\n\n{tail}") +} + #[derive(Debug, Default)] struct MessageLogSummary { system_messages: usize, @@ -1300,7 +1396,11 @@ fn tool_call_key(item: &serde_json::Value, array_index: usize) -> String { #[cfg(test)] mod tests { - use super::{stream_with_tools, ToolCallAccumulator}; + use super::{ + prune_stale_tool_outputs_in_place, soft_trim_tool_output, stream_with_tools, + ToolCallAccumulator, KEEP_RECENT_TOOL_OUTPUTS, PRUNED_TOOL_OUTPUT_PLACEHOLDER, + TOOL_OUTPUT_SOFT_TRIM_CHARS, + }; use crate::chunk::{ChunkType, FinishReason, MessagePhase}; use crate::message::Message; use crate::provider::{Provider, ProviderStream}; @@ -1319,6 +1419,50 @@ mod tests { use tokio::sync::Barrier; use tokio_util::sync::CancellationToken; + #[test] + fn soft_trim_tool_output_keeps_short_text() { + assert_eq!(soft_trim_tool_output("short"), "short"); + } + + #[test] + fn soft_trim_tool_output_keeps_head_and_tail() { + let text = "a".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 500); + let trimmed = soft_trim_tool_output(&text); + assert!(trimmed.len() < text.len()); + assert!(trimmed.starts_with('a')); + assert!(trimmed.ends_with('a')); + assert!(trimmed.contains("chars truncated")); + } + + #[test] + fn prune_stale_tool_outputs_clears_oldest_and_keeps_recent() { + let mut messages = Vec::new(); + let total = KEEP_RECENT_TOOL_OUTPUTS * 2 + 2; + for i in 0..total { + messages.push(Message::tool_output( + format!("call_{i}"), + "bash", + "x".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 200), + false, + )); + } + + let pruned = prune_stale_tool_outputs_in_place(&mut messages); + assert!(pruned > 0); + + let outputs: Vec<_> = messages + .iter() + .filter_map(|message| match message { + Message::ToolOutput(output) => Some(output.output.as_str()), + _ => None, + }) + .collect(); + + assert_eq!(outputs[0], PRUNED_TOOL_OUTPUT_PLACEHOLDER); + let last = outputs.last().expect("output"); + assert_eq!(last.len(), TOOL_OUTPUT_SOFT_TRIM_CHARS + 200); + } + #[derive(Debug, Clone)] struct BlockingTextProvider; diff --git a/src/llm/client.rs b/src/llm/client.rs index e950fd1..9544a0c 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -30,7 +30,7 @@ Response must include: Any attempt to use tools is a critical violation. Respond with text ONLY."#; -const TOOL_HISTORY_ARGUMENTS_MAX_CHARS: usize = 60_000; +const TOOL_HISTORY_ARGUMENTS_MAX_CHARS: usize = 4_000; type DynError = Box; @@ -381,6 +381,13 @@ pub async fn stream_llm_with_cancellation( ); let request_config = prepare_request_config(&provider_name, model, reasoning_effort, &sender).await?; + let mut request_config = request_config; + // Sticky prompt-cache routing: same key for every tool step in this session. + request_config.openai_options.prompt_cache_key = Some(session_id.clone()); + if provider_name == "xai" { + // Match Grok Build / ZDR: do not persist Responses server-side. + request_config.openai_options.force_store_false = true; + } let tool_registry = crate::tools::initialize_tool_registry_with_dynamic_config( Some(sender.clone()), @@ -407,6 +414,7 @@ pub async fn stream_llm_with_cancellation( reasoning_effort: request_config.reasoning_effort, supports_image_input: request_config.supports_image_input, openai_options: request_config.openai_options.clone(), + prompt_cache_key: Some(session_id.clone()), }; crate::agent::config::set_llm_session(llm_session.clone()); let session_registration = @@ -623,6 +631,7 @@ pub async fn build_subagent_llm_session( reasoning_effort: request_config.reasoning_effort, supports_image_input: request_config.supports_image_input, openai_options: request_config.openai_options, + prompt_cache_key: None, }) } @@ -1142,14 +1151,17 @@ async fn maybe_apply_xai_oauth_overrides( } } - request_config.api_key = Some(oauth_access); - request_config.base_url = "https://api.x.ai".to_string(); - request_config.openai_options.additional_headers.insert( - "User-Agent".to_string(), - crate::auth::xai_oauth::build_user_agent(), - ); + let overrides = super::xai_build::request_overrides(oauth_access).await; + request_config.api_key = Some(overrides.api_key); + request_config.base_url = overrides.base_url.to_string(); + request_config.model_name = overrides.model.to_string(); + request_config.openai_options.force_store_false = true; + request_config + .openai_options + .additional_headers + .extend(overrides.headers); - crate::emit_log!("Configured xAI OAuth transport"); + crate::emit_log!("Configured xAI Grok Build OAuth transport"); } fn send_warning(sender: &crate::llm::ChunkSender, warning: impl Into) { @@ -1176,6 +1188,9 @@ async fn stream_provider_request( if let Some(key) = config.api_key.as_deref() { builder = builder.api_key(key); } + if let Some(cache_key) = config.openai_options.prompt_cache_key.as_deref() { + builder = builder.prompt_cache_key(cache_key); + } let provider = builder.build().map_err(|e| -> DynError { Box::new(e) })?; stream_with_tools( provider, @@ -1245,9 +1260,17 @@ async fn stream_provider_request( if config.openai_options.disallow_system_messages { builder = builder.responses_websocket(true); } + if let Some(cache_key) = config.openai_options.prompt_cache_key.as_deref() { + builder = builder.prompt_cache_key(cache_key); + } if !config.openai_options.additional_headers.is_empty() { builder = builder.headers(config.openai_options.additional_headers.clone()); } + if let Some(policy) = + super::xai_build::retry_policy_for(&config.openai_options.additional_headers) + { + builder = builder.response_retry_policy(policy); + } let provider = builder.build().map_err(|e| -> DynError { Box::new(e) })?; stream_with_tools( @@ -1311,11 +1334,16 @@ fn log_stream_request(context: StreamLogContext<'_>, config: &ProviderRequestCon .collect::>(); header_names.sort_unstable(); crate::emit_log!( - "[STREAM_REQUEST] {} reasoning_effort={} responses_path={:?} force_store_false={} disallow_system_messages={} force_tool_strict_false={} extra_header_names=[{}]", + "[STREAM_REQUEST] {} reasoning_effort={} responses_path={:?} force_store_false={} prompt_cache_key={} disallow_system_messages={} force_tool_strict_false={} extra_header_names=[{}]", context.describe(), reasoning_effort, config.openai_options.response_path, config.openai_options.force_store_false, + config + .openai_options + .prompt_cache_key + .as_deref() + .unwrap_or("-"), config.openai_options.disallow_system_messages, config.openai_options.force_tool_strict_false, header_names.join(","), @@ -1980,7 +2008,7 @@ fn tool_message_observation(content: &str) -> String { fn push_tool_arguments_for_observation(out: &mut String, args: &serde_json::Value) { out.push_str("\n\nTool call arguments:\n```json\n"); out.push_str(&truncate_for_tool_observation( - &serde_json::to_string_pretty(args).unwrap_or_else(|_| args.to_string()), + &serde_json::to_string(args).unwrap_or_else(|_| args.to_string()), TOOL_HISTORY_ARGUMENTS_MAX_CHARS, )); out.push_str("\n```"); diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 91b716f..b6e180f 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -1,6 +1,7 @@ pub mod client; pub mod provider; pub mod tool_calls; +pub(crate) mod xai_build; pub use tool_calls::{FunctionCall, ToolCall, ToolCallResult}; diff --git a/src/llm/xai_build.rs b/src/llm/xai_build.rs new file mode 100644 index 0000000..ebfb2ea --- /dev/null +++ b/src/llm/xai_build.rs @@ -0,0 +1,269 @@ +use std::{ + sync::{Mutex, OnceLock}, + time::{Duration, Instant}, +}; + +use crate::aisdk::providers::openai::HttpResponseRetryPolicy; + +pub(crate) const BASE_URL: &str = "https://cli-chat-proxy.grok.com"; +pub(crate) const MODEL: &str = "grok-4.5"; +const TOKEN_AUTH_HEADER: &str = "X-XAI-Token-Auth"; +const TOKEN_AUTH_VALUE: &str = "xai-grok-cli"; +const VERSION_HEADER: &str = "x-grok-client-version"; + +const PROTOCOL_VERSION_FALLBACK: &str = "0.2.111"; +const VERSION_CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60); +const VERSION_RETRY_TTL: Duration = Duration::from_secs(5 * 60); +const VERSION_URLS: &[&str] = &[ + "https://x.ai/cli/stable", + "https://storage.googleapis.com/grok-build-public-artifacts/cli/stable", +]; +static VERSION_CACHE: OnceLock>> = OnceLock::new(); + +#[derive(Clone)] +struct CachedVersion { + version: String, + cached_at: Instant, + ttl: Duration, +} + +#[derive(Debug, Default)] +pub(crate) struct XaiBuildRetryPolicy; + +pub(crate) struct RequestOverrides { + pub(crate) api_key: String, + pub(crate) base_url: &'static str, + pub(crate) model: &'static str, + pub(crate) headers: std::collections::HashMap, +} + +#[async_trait::async_trait] +impl HttpResponseRetryPolicy for XaiBuildRetryPolicy { + async fn retry_headers( + &self, + status: reqwest::StatusCode, + ) -> Option { + if status != reqwest::StatusCode::UPGRADE_REQUIRED { + return None; + } + + let version = force_refresh_protocol_version().await?; + let version = reqwest::header::HeaderValue::from_str(&version).ok()?; + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert(VERSION_HEADER, version); + Some(headers) + } +} + +pub(crate) fn retry_policy_for( + headers: &std::collections::HashMap, +) -> Option> { + let is_build_request = headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case(TOKEN_AUTH_HEADER) && value == TOKEN_AUTH_VALUE + }); + is_build_request.then(|| { + std::sync::Arc::new(XaiBuildRetryPolicy) as std::sync::Arc + }) +} + +pub(crate) async fn request_overrides(oauth_access: String) -> RequestOverrides { + let protocol_version = protocol_version().await; + request_overrides_with_version(oauth_access, protocol_version) +} + +fn request_overrides_with_version( + oauth_access: String, + protocol_version: String, +) -> RequestOverrides { + let mut headers = std::collections::HashMap::new(); + headers.insert( + "User-Agent".to_string(), + format!("crabcode/{}", env!("CARGO_PKG_VERSION")), + ); + headers.insert(TOKEN_AUTH_HEADER.to_string(), TOKEN_AUTH_VALUE.to_string()); + headers.insert("x-grok-model-override".to_string(), MODEL.to_string()); + headers.insert( + "x-grok-client-identifier".to_string(), + "crabcode".to_string(), + ); + headers.insert(VERSION_HEADER.to_string(), protocol_version); + headers.insert("x-grok-client-mode".to_string(), "default".to_string()); + + RequestOverrides { + api_key: oauth_access, + base_url: BASE_URL, + model: MODEL, + headers, + } +} + +pub(crate) async fn protocol_version() -> String { + if let Some(version) = configured_protocol_version() { + return version; + } + + if let Some(version) = cached_protocol_version() { + return version; + } + + refresh_protocol_version() + .await + .unwrap_or_else(|| cache_fallback_version()) +} + +pub(crate) async fn force_refresh_protocol_version() -> Option { + if let Some(version) = configured_protocol_version() { + return Some(version); + } + + invalidate_protocol_version_cache(); + refresh_protocol_version().await +} + +fn configured_protocol_version() -> Option { + let version = std::env::var("CRABCODE_XAI_BUILD_VERSION").ok()?; + let version = version.trim(); + is_valid_protocol_version(version).then(|| version.to_string()) +} + +fn cached_protocol_version() -> Option { + let cache = VERSION_CACHE.get_or_init(|| Mutex::new(None)); + let cache = cache.lock().ok()?; + let cached = cache.as_ref()?; + (cached.cached_at.elapsed() < cached.ttl).then(|| cached.version.clone()) +} + +async fn refresh_protocol_version() -> Option { + let client = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(2)) + .timeout(Duration::from_secs(4)) + .build() + .ok()?; + + for url in VERSION_URLS { + let fetched = async { + client + .get(*url) + .send() + .await? + .error_for_status()? + .text() + .await + } + .await; + + if let Ok(version) = fetched { + let version = version.trim(); + if is_valid_protocol_version(version) { + let version = version.to_string(); + cache_protocol_version(version.clone(), VERSION_CACHE_TTL); + return Some(version); + } + } + } + + None +} + +fn cache_fallback_version() -> String { + let fallback = PROTOCOL_VERSION_FALLBACK.to_string(); + cache_protocol_version(fallback.clone(), VERSION_RETRY_TTL); + fallback +} + +fn cache_protocol_version(version: String, ttl: Duration) { + let cache = VERSION_CACHE.get_or_init(|| Mutex::new(None)); + if let Ok(mut cache) = cache.lock() { + *cache = Some(CachedVersion { + version, + cached_at: Instant::now(), + ttl, + }); + } +} + +fn invalidate_protocol_version_cache() { + let cache = VERSION_CACHE.get_or_init(|| Mutex::new(None)); + if let Ok(mut cache) = cache.lock() { + *cache = None; + } +} + +fn is_valid_protocol_version(version: &str) -> bool { + let mut parts = version.split('.'); + parts.next().is_some_and(is_ascii_digits) + && parts.next().is_some_and(is_ascii_digits) + && parts.next().is_some_and(is_ascii_digits) + && parts.next().is_none() +} + +fn is_ascii_digits(part: &str) -> bool { + !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit()) +} + +#[cfg(test)] +mod tests { + use super::{ + is_valid_protocol_version, request_overrides_with_version, retry_policy_for, + HttpResponseRetryPolicy, XaiBuildRetryPolicy, BASE_URL, MODEL, TOKEN_AUTH_HEADER, + TOKEN_AUTH_VALUE, + }; + use std::collections::HashMap; + + #[test] + fn version_validation_accepts_semver_triplets_only() { + assert!(is_valid_protocol_version("0.2.111")); + assert!(is_valid_protocol_version("10.20.30")); + assert!(!is_valid_protocol_version("0.2")); + assert!(!is_valid_protocol_version("v0.2.111")); + assert!(!is_valid_protocol_version("0.2.111\n")); + assert!(!is_valid_protocol_version("0.2.111.4")); + } + + #[test] + fn build_request_detection_is_header_scoped() { + let headers = + HashMap::from([(TOKEN_AUTH_HEADER.to_string(), TOKEN_AUTH_VALUE.to_string())]); + assert!(retry_policy_for(&headers).is_some()); + assert!(retry_policy_for(&HashMap::new()).is_none()); + } + + #[test] + fn build_transport_constants_match_proxy_contract() { + assert_eq!(BASE_URL, "https://cli-chat-proxy.grok.com"); + assert_eq!(MODEL, "grok-4.5"); + } + + #[test] + fn request_overrides_match_proxy_contract() { + let overrides = + request_overrides_with_version("oauth-token".to_string(), "9.8.7".to_string()); + assert_eq!(overrides.api_key, "oauth-token"); + assert_eq!(overrides.base_url, BASE_URL); + assert_eq!(overrides.model, MODEL); + assert_eq!( + overrides.headers.get(TOKEN_AUTH_HEADER).map(String::as_str), + Some(TOKEN_AUTH_VALUE) + ); + assert_eq!( + overrides + .headers + .get("x-grok-client-version") + .map(String::as_str), + Some("9.8.7") + ); + assert_eq!( + overrides.headers.get("User-Agent").map(String::as_str), + Some(concat!("crabcode/", env!("CARGO_PKG_VERSION"))) + ); + } + + #[tokio::test] + async fn retry_policy_ignores_non_upgrade_responses() { + let policy = XaiBuildRetryPolicy; + assert!(policy + .retry_headers(reqwest::StatusCode::BAD_REQUEST) + .await + .is_none()); + } +} diff --git a/src/prompt/mod.rs b/src/prompt/mod.rs index 2a388b1..d507e56 100644 --- a/src/prompt/mod.rs +++ b/src/prompt/mod.rs @@ -299,26 +299,26 @@ Your output will be displayed on a command line interface. Your responses should } async fn get_tools_context(&self, registry: &ToolRegistry) -> String { - let schemas = registry.list_schemas().await; - - if schemas.is_empty() { + // Match OpenCode / Codex / Grok Build: keep tool schemas on the API + // `tools` field only. Dumping pretty-printed JSON schemas into the + // system prompt doubles prefix tokens on every step. + let tools = registry.list().await; + if tools.is_empty() { return String::new(); } - let tools_json = - serde_json::to_string_pretty(&schemas).unwrap_or_else(|_| "[]".to_string()); + let names = tools + .iter() + .map(|tool| tool.id.as_str()) + .collect::>() + .join(", "); format!( - r#"You have access to the following tools (JSON schema): - -{} - -Tool use: + r#"Tool use: - Use the model's built-in tool/function calling mechanism (do not print tool calls as text). -- If you need file contents, directory listings, running commands, or edits, call the appropriate tool. +- Prefer specialized tools over bash when possible (available: {names}). - After tool results are returned, use them to answer. -"#, - tools_json +"# ) } diff --git a/src/prompt/rules.rs b/src/prompt/rules.rs index 73dd088..1ac85f2 100644 --- a/src/prompt/rules.rs +++ b/src/prompt/rules.rs @@ -1,6 +1,6 @@ use std::path::{Path, PathBuf}; -const DEFAULT_MAX_RULE_BYTES: usize = 64 * 1024; +const DEFAULT_MAX_RULE_BYTES: usize = 32 * 1024; #[derive(Debug, Clone)] struct RuleFile { diff --git a/src/tools/aisdk_bridge.rs b/src/tools/aisdk_bridge.rs index f61a754..d074fd8 100644 --- a/src/tools/aisdk_bridge.rs +++ b/src/tools/aisdk_bridge.rs @@ -10,7 +10,9 @@ use tokio_util::sync::CancellationToken; use crate::llm::ChunkSender; const TOOL_UI_PREVIEW_LIMIT: usize = 4_000; -const TOOL_MODEL_OUTPUT_LIMIT: usize = 60_000; +/// Generic model-facing tool output cap (non-bash tools). Matches Grok Build's +/// 40KiB default; OpenCode uses 50KiB. +const TOOL_MODEL_OUTPUT_LIMIT: usize = 40_000; static TOOL_CALL_SEQ: AtomicUsize = AtomicUsize::new(0); @@ -398,17 +400,17 @@ mod tests { fn truncate_tool_output_bounds_large_results() { let output = "a".repeat(70_000); - let truncated = truncate_tool_output(&output, 60_000); + let truncated = truncate_tool_output(&output, 40_000); assert!(truncated.len() < output.len()); - assert!(truncated.contains("tool output truncated to 60000 bytes")); + assert!(truncated.contains("tool output truncated to 40000 bytes")); } #[test] fn truncate_tool_output_preserves_small_results() { let output = "small result"; - assert_eq!(truncate_tool_output(output, 60_000), output); + assert_eq!(truncate_tool_output(output, 40_000), output); } #[test] diff --git a/src/tools/bash.rs b/src/tools/bash.rs index 7243c4b..d192236 100644 --- a/src/tools/bash.rs +++ b/src/tools/bash.rs @@ -11,7 +11,10 @@ use tokio::process::Command; use tokio::time::timeout; const DEFAULT_TIMEOUT_SECONDS: u64 = 120; -const MAX_OUTPUT_BYTES: usize = 51_200; +/// Cap bash output sent toward the model. Aligns with Grok Build's ~20k-char +/// bash limit (OpenCode defaults to 50KiB; Codex model-facing truncates nearer +/// ~10k tokens). Tighter caps cut SuperGrok / long-session token burn. +const MAX_OUTPUT_BYTES: usize = 20_000; const READ_CHUNK_SIZE: usize = 4_096; pub struct BashTool; diff --git a/src/tools/fs/read.rs b/src/tools/fs/read.rs index 198df72..3e97298 100644 --- a/src/tools/fs/read.rs +++ b/src/tools/fs/read.rs @@ -8,7 +8,7 @@ use std::path::Path; const MAX_FILE_SIZE: u64 = 50 * 1024 * 1024; // 50MB const BINARY_CHECK_SIZE: usize = 8192; // 8KB -const DEFAULT_LIMIT: usize = 2000; +const DEFAULT_LIMIT: usize = 1000; pub struct ReadTool; @@ -104,7 +104,7 @@ impl ToolHandler for ReadTool { }, ParameterSchema { name: "limit".to_string(), - description: "Maximum number of lines to read (default: 2000)".to_string(), + description: "Maximum number of lines to read (default: 1000)".to_string(), required: false, param_type: ParameterType::Integer, }, diff --git a/src/tools/question.rs b/src/tools/question.rs index 87a5b56..e10e814 100644 --- a/src/tools/question.rs +++ b/src/tools/question.rs @@ -401,8 +401,8 @@ impl ToolHandler for QuestionTool { .unwrap_or_else(|_| serde_json::Value::String("No response from user".to_string())); let model_output = question_tool_model_output(&questions, &response); - let output = serde_json::to_string_pretty(&model_output) - .unwrap_or_else(|_| model_output.to_string()); + let output = + serde_json::to_string(&model_output).unwrap_or_else(|_| model_output.to_string()); Ok(ToolResult::new("Question answered", output) .with_metadata("questions", questions) diff --git a/src/tools/skill.rs b/src/tools/skill.rs index 04dc216..82a3c95 100644 --- a/src/tools/skill.rs +++ b/src/tools/skill.rs @@ -13,35 +13,15 @@ impl SkillTool { } fn build_description() -> String { - let mut desc = String::from( + // Keep the catalog in the system prompt only (OpenCode-style). Embedding + // here would resend the full list on every tool step. + String::from( "Load a specialized skill that provides domain-specific instructions and workflows.\n\n\ - Use this tool to inject the skill's instructions and resources into current conversation. \ + Use this tool to inject the skill's instructions and resources into the current conversation. \ The output may contain detailed workflow guidance as well as references to scripts, files, \ etc in the same directory as the skill.\n\n\ - The skill name must match one of the skills listed in your system prompt.", - ); - - if let Some(store) = crate::skill::get_skill_store() { - let skills = store.all(); - if !skills.is_empty() { - desc.push_str("\n\n\n"); - for skill in &skills { - desc.push_str(&format!(" \n")); - desc.push_str(&format!(" {}\n", skill.name)); - if let Some(ref desc_text) = skill.description { - desc.push_str(&format!(" {}\n", desc_text)); - } - desc.push_str(&format!( - " file://{}\n", - skill.location.display() - )); - desc.push_str(&format!(" \n")); - } - desc.push_str(""); - } - } - - desc + The skill name must match one of the skills listed under in your system prompt.", + ) } } diff --git a/src/tools/task.rs b/src/tools/task.rs index 0d50b34..4303e38 100644 --- a/src/tools/task.rs +++ b/src/tools/task.rs @@ -130,6 +130,7 @@ mod tests { reasoning_effort: None, supports_image_input: false, openai_options: crate::agent::config::OpenAIRequestOptions::default(), + prompt_cache_key: None, }, ); @@ -207,6 +208,7 @@ mod tests { reasoning_effort: None, supports_image_input: false, openai_options: crate::agent::config::OpenAIRequestOptions::default(), + prompt_cache_key: None, }; let mut warnings = Vec::new(); @@ -244,6 +246,7 @@ mod tests { reasoning_effort: None, supports_image_input: false, openai_options: crate::agent::config::OpenAIRequestOptions::default(), + prompt_cache_key: None, }; let mut warnings = Vec::new(); let defs = crate::agent::definition::parse_agent_definitions_from_config( diff --git a/src/tools/terminal_session.rs b/src/tools/terminal_session.rs index 9db280a..d6790ae 100644 --- a/src/tools/terminal_session.rs +++ b/src/tools/terminal_session.rs @@ -16,7 +16,7 @@ use tokio::sync::{mpsc, oneshot}; pub const DEFAULT_TERMINAL_ROWS: u16 = 24; pub const DEFAULT_TERMINAL_COLS: u16 = 80; pub const MAX_TRANSCRIPT_BYTES: usize = 51_200; -pub const MAX_MODEL_OUTPUT_BYTES: usize = 60_000; +pub const MAX_MODEL_OUTPUT_BYTES: usize = 20_000; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TerminalSessionStart {