vtcode_llm/providers/
xai.rs

1use serde_json::{Map, Value};
2use vtcode_config::constants::{env_vars, models, urls};
3
4use super::extract_reasoning_trace;
5use super::openai_compat::{OpenAiCompatCore, OpenAiCompatSpec, SystemPromptPlacement, impl_openai_compat_provider};
6
7pub struct XaiSpec;
8
9fn xai_reasoning(message: &Value, _choice: &Value) -> Option<String> {
10    message.get("reasoning_content").and_then(extract_reasoning_trace)
11}
12
13impl OpenAiCompatSpec for XaiSpec {
14    const NAME: &'static str = "xAI";
15    const KEY: &'static str = "xai";
16    const API_KEY_ENV: &'static str = "XAI_API_KEY";
17    const DEFAULT_MODEL: &'static str = models::xai::DEFAULT_MODEL;
18    const DEFAULT_BASE_URL: &'static str = urls::XAI_API_BASE;
19    const BASE_URL_ENV: Option<&'static str> = Some(env_vars::XAI_BASE_URL);
20    const LISTED_MODELS: &'static [&'static str] = models::xai::SUPPORTED_MODELS;
21    const VALIDATION_ALLOWLIST: Option<&'static [&'static str]> = Some(models::xai::SUPPORTED_MODELS);
22
23    const SYSTEM_PROMPT: SystemPromptPlacement = SystemPromptPlacement::FirstMessage;
24    const STREAM_OPTIONS_INCLUDE_USAGE: bool = true;
25    const INCLUDE_USER_ID: bool = true;
26    const RESPONSE_REASONING_EXTRACTOR: Option<super::openai_compat::ReasoningExtractor> = Some(xai_reasoning);
27
28    fn insert_reasoning(
29        _core: &OpenAiCompatCore<Self>,
30        request: &crate::provider::LLMRequest,
31        payload: &mut Map<String, Value>,
32    ) -> Result<(), crate::provider::LLMError> {
33        if let Some(effort) = request.reasoning_effort {
34            if !matches!(
35                effort,
36                vtcode_config::types::ReasoningEffortLevel::None | vtcode_config::types::ReasoningEffortLevel::Unknown
37            ) {
38                // xAI natively supports `low`/`medium`/`high`/`xhigh` only:
39                // no native `minimal` (clamp to `low`) or `max` (clamp to `xhigh`).
40                // Older models treat `xhigh` as `high`.
41                let value = match effort {
42                    vtcode_config::types::ReasoningEffortLevel::Minimal
43                    | vtcode_config::types::ReasoningEffortLevel::Low => "low",
44                    vtcode_config::types::ReasoningEffortLevel::Max => "xhigh",
45                    other => other.as_str(),
46                };
47                payload.insert("reasoning_effort".to_owned(), serde_json::json!(value));
48            }
49        }
50        Ok(())
51    }
52
53    fn response_cache_metrics(core: &OpenAiCompatCore<Self>) -> bool {
54        core.prompt_cache_enabled
55    }
56
57    fn stream_cache_metrics(_core: &OpenAiCompatCore<Self>) -> bool {
58        true
59    }
60
61    /// xAI cache affinity: stable conversation identity from VT Code lineage.
62    /// Body `prompt_cache_key` covers Responses-style hosts; Chat Completions
63    /// sticky routing uses `x-grok-conv-id` via [`session_affinity_header`].
64    fn finish_payload(
65        _core: &OpenAiCompatCore<Self>,
66        request: &crate::provider::LLMRequest,
67        payload: &mut Map<String, Value>,
68    ) -> Result<(), crate::provider::LLMError> {
69        if let Some(lineage) =
70            crate::providers::shared::session_lineage_from_prompt_cache_key(request.prompt_cache_key.as_deref())
71        {
72            payload.insert("prompt_cache_key".to_owned(), Value::String(lineage.clone()));
73        }
74        Ok(())
75    }
76
77    fn session_affinity_header() -> Option<&'static str> {
78        Some("x-grok-conv-id")
79    }
80}
81
82impl XAIProvider {
83    /// xAI serves an OpenAI-compatible standalone compaction endpoint
84    /// (`POST /v1/responses/compact`), but only for curated Grok models on the
85    /// xAI API itself. Anything else stays on the universal local
86    /// summarization fallback.
87    fn xai_compact_model(&self, model: &str) -> bool {
88        let resolved = if model.trim().is_empty() {
89            self.core.model.as_str()
90        } else {
91            model
92        };
93        models::xai::SUPPORTED_MODELS.contains(&resolved) && self.core.base_url.contains("api.x.ai")
94    }
95
96    fn compact_client(&self, model: &str) -> crate::providers::openresponses::OpenResponsesProvider {
97        crate::providers::openresponses::OpenResponsesProvider::compact_endpoint_client(
98            &self.core.model,
99            &self.core.base_url,
100            &self.core.api_key,
101            model,
102        )
103    }
104}
105
106impl_openai_compat_provider!(XAIProvider, XaiSpec, {
107    fn supports_reasoning(&self, model: &str) -> bool {
108        let requested = if model.trim().is_empty() {
109            &self.core.model
110        } else {
111            model
112        };
113        self.core
114            .model_behavior
115            .as_ref()
116            .and_then(|b| b.model_supports_reasoning)
117            .unwrap_or(false)
118            || models::xai::REASONING_MODELS.contains(&requested)
119    }
120
121    fn supports_reasoning_effort(&self, model: &str) -> bool {
122        let requested = if model.trim().is_empty() {
123            &self.core.model
124        } else {
125            model
126        };
127        self.core
128            .model_behavior
129            .as_ref()
130            .and_then(|b| b.model_supports_reasoning_effort)
131            .unwrap_or_else(|| {
132                vtcode_config::models::model_catalog_entry("xai", requested)
133                    .is_some_and(|entry| !entry.reasoning_efforts.is_empty())
134            })
135    }
136
137    fn supports_responses_compaction(&self, model: &str) -> bool {
138        self.xai_compact_model(model)
139    }
140
141    fn supports_manual_openai_compaction(&self, model: &str) -> bool {
142        self.xai_compact_model(model)
143    }
144
145    async fn compact_history(
146        &self,
147        model: &str,
148        history: &[crate::provider::Message],
149    ) -> Result<Vec<crate::provider::Message>, crate::provider::LLMError> {
150        if !self.xai_compact_model(model) {
151            return Err(crate::provider::LLMError::Provider {
152                message: "xAI compaction is only supported for curated Grok models on the xAI API".to_string(),
153                metadata: None,
154            });
155        }
156        self.compact_client(model).compact_history_request(model, history).await
157    }
158
159    async fn compact_history_with_options(
160        &self,
161        model: &str,
162        history: &[crate::provider::Message],
163        _options: &crate::provider::ResponsesCompactionOptions,
164    ) -> Result<Vec<crate::provider::Message>, crate::provider::LLMError> {
165        self.compact_history(model, history).await
166    }
167});
168
169#[cfg(test)]
170mod tests {
171    #[test]
172    fn xai_advertises_grok_session_affinity_header() {
173        use super::XaiSpec;
174        use crate::providers::openai_compat::OpenAiCompatSpec;
175        assert_eq!(XaiSpec::session_affinity_header(), Some("x-grok-conv-id"));
176    }
177
178    #[test]
179    fn xai_finish_payload_injects_lineage_prompt_cache_key() {
180        let provider = XAIProvider::new("test-key".to_string());
181
182        let mut request = base_request();
183        request.prompt_cache_key = Some("vtcode:xai:session-lineage-1".to_string());
184        let payload = provider.core.convert_request(&request).unwrap();
185        assert_eq!(payload["prompt_cache_key"].as_str(), Some("session-lineage-1"));
186
187        let mut blank = base_request();
188        blank.prompt_cache_key = Some("   ".to_string());
189        let payload = provider.core.convert_request(&blank).unwrap();
190        assert!(payload.get("prompt_cache_key").is_none(), "blank lineage must omit prompt_cache_key");
191    }
192
193    #[test]
194    fn xai_dispatch_header_name_is_pinned_for_lineage_routing() {
195        use super::XaiSpec;
196        use crate::providers::openai_compat::OpenAiCompatSpec;
197        // Production dispatch attaches this header when request.prompt_cache_key
198        // carries session lineage (see openai_compat::dispatch).
199        assert_eq!(XaiSpec::session_affinity_header(), Some("x-grok-conv-id"));
200    }
201
202    use super::XAIProvider;
203    use crate::provider::{LLMRequest, Message, ToolChoice};
204    use std::sync::Arc;
205    use vtcode_config::constants::models;
206    use vtcode_config::types::ReasoningEffortLevel;
207
208    fn base_request() -> LLMRequest {
209        LLMRequest {
210            messages: vec![Message::user("hello".to_string())].into(),
211            system_prompt: Some(Arc::from("system guidance")),
212            model: models::xai::DEFAULT_MODEL.to_string(),
213            max_tokens: Some(512),
214            temperature: Some(0.5),
215            top_p: Some(0.25),
216            stream: true,
217            tool_choice: Some(ToolChoice::Auto),
218            ..Default::default()
219        }
220    }
221
222    #[test]
223    fn golden_payload_basic_shape() {
224        let provider = XAIProvider::new("test-key".to_string());
225        let payload = provider.core.convert_request(&base_request()).unwrap();
226
227        assert_eq!(payload["model"], models::xai::DEFAULT_MODEL);
228        let messages = payload["messages"].as_array().unwrap();
229        assert_eq!(messages.len(), 2);
230        assert_eq!(messages[0]["role"], "system");
231        assert_eq!(messages[0]["content"], "system guidance");
232        assert_eq!(messages[1]["role"], "user");
233        assert_eq!(payload["max_tokens"], 512);
234        assert_eq!(payload["temperature"], 0.5);
235        assert_eq!(payload["top_p"], 0.25);
236        assert_eq!(payload["stream"], true);
237        assert_eq!(payload["stream_options"]["include_usage"], true);
238        assert_eq!(payload["tool_choice"], "auto");
239    }
240
241    #[test]
242    fn golden_payload_with_reasoning_effort() {
243        let provider = XAIProvider::new("test-key".to_string());
244
245        let mut request = base_request();
246        request.reasoning_effort = Some(ReasoningEffortLevel::High);
247        let payload = provider.core.convert_request(&request).unwrap();
248        assert_eq!(payload["reasoning_effort"], "high");
249    }
250
251    #[test]
252    fn max_effort_clamps_to_xhigh() {
253        let provider = XAIProvider::new("test-key".to_string());
254
255        let mut request = base_request();
256        request.reasoning_effort = Some(ReasoningEffortLevel::Max);
257        let payload = provider.core.convert_request(&request).unwrap();
258        // No native `max`; closest supported level is `xhigh`.
259        assert_eq!(payload["reasoning_effort"], "xhigh");
260
261        // No native `minimal` either; closest supported level is `low`.
262        request.reasoning_effort = Some(ReasoningEffortLevel::Minimal);
263        let payload = provider.core.convert_request(&request).unwrap();
264        assert_eq!(payload["reasoning_effort"], "low");
265    }
266
267    #[test]
268    fn compat_payload_forwards_ultrafast_service_tier() {
269        let provider = XAIProvider::new("test-key".to_string());
270
271        let mut request = base_request();
272        request.service_tier = Some("ultrafast".to_string());
273        let payload = provider.core.convert_request(&request).unwrap();
274        assert_eq!(payload["service_tier"], "ultrafast");
275
276        let mut blank = base_request();
277        blank.service_tier = Some("   ".to_string());
278        let payload = provider.core.convert_request(&blank).unwrap();
279        assert!(payload.get("service_tier").is_none(), "blank tier must be omitted");
280    }
281
282    #[test]
283    fn compaction_support_is_curated_grok_models_on_xai_api_only() {
284        use crate::provider::LLMProvider;
285
286        // Default xAI endpoint: curated Grok models compact natively.
287        let provider = XAIProvider::new("test-key".to_string());
288        assert!(provider.supports_responses_compaction(models::xai::DEFAULT_MODEL));
289        assert!(provider.supports_manual_openai_compaction(models::xai::DEFAULT_MODEL));
290        for model in models::xai::SUPPORTED_MODELS {
291            assert!(provider.supports_manual_openai_compaction(model), "compact support for {model}");
292        }
293        assert!(!provider.supports_native_inline_compaction(models::xai::DEFAULT_MODEL));
294    }
295
296    #[tokio::test]
297    async fn compact_history_posts_to_xai_compact_endpoint() {
298        use wiremock::matchers::{method, path};
299        use wiremock::{Mock, MockServer, ResponseTemplate};
300
301        let server = MockServer::start().await;
302        Mock::given(method("POST"))
303            .and(path("/v1/responses/compact"))
304            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
305                "id": "cmp_01HZ9P0V8M2YQK3F7C4G6N5R2A",
306                "object": "response.compaction",
307                "created_at": 1748895600,
308                "model": "grok-4.6",
309                "output": [
310                    {
311                        "id": "msg_000",
312                        "type": "message",
313                        "status": "completed",
314                        "role": "user",
315                        "content": [{ "type": "input_text", "text": "Summarize the auth work." }]
316                    },
317                    {
318                        "id": "cmp_001",
319                        "type": "compaction",
320                        "encrypted_content": "opaque-blob"
321                    }
322                ]
323            })))
324            .mount(&server)
325            .await;
326
327        // The host gate only passes on the xAI API, so the transport is
328        // exercised through the compact client directly against the mock.
329        let provider = XAIProvider::new_with_client(
330            "test-key".to_string(),
331            models::xai::DEFAULT_MODEL.to_string(),
332            reqwest::Client::builder().no_proxy().build().expect("test client should build"),
333            format!("{}/v1", server.uri()),
334            vtcode_config::TimeoutsConfig::default(),
335        );
336        let history = vec![Message::user("Summarize the auth work.".to_string())];
337        let compacted = provider
338            .compact_client(models::xai::DEFAULT_MODEL)
339            .compact_history_request(models::xai::DEFAULT_MODEL, &history)
340            .await
341            .expect("xAI compaction should succeed");
342        assert!(!compacted.is_empty());
343        assert!(
344            compacted
345                .iter()
346                .any(|message| message.content.as_text().contains("Summarize the auth work.")),
347            "retained xAI input must survive compaction"
348        );
349    }
350
351    #[tokio::test]
352    async fn compact_history_rejects_unlisted_models() {
353        use crate::provider::LLMProvider;
354
355        let provider = XAIProvider::new("test-key".to_string());
356        assert!(!provider.supports_manual_openai_compaction("gpt-5"));
357        let history = vec![Message::user("hello".to_string())];
358        provider
359            .compact_history("gpt-5", &history)
360            .await
361            .expect_err("unlisted models must stay on local compaction");
362    }
363}
vtcode_llm/providers/xai.rs

vtcode_llm/providers/
xai.rs
