vtcode_core/commands/
ask.rs

1//! Ask command implementation - single prompt without tools
2
3use crate::cli::input_hardening::validate_agent_safe_text;
4use crate::config::types::AgentConfig;
5use crate::llm::collect_single_response;
6use crate::llm::factory::{ProviderConfig, create_provider_with_config, infer_provider_from_model};
7use crate::llm::provider::{LLMRequest, Message};
8use crate::prompts::system::lightweight_instruction_text;
9use anyhow::Result;
10use std::io::IsTerminal;
11use std::sync::Arc;
12
13/// Handle the ask command - single prompt without tools
14#[cfg_attr(feature = "profiling", hotpath::measure)]
15pub async fn handle_ask_command(
16    config: AgentConfig,
17    prompt: Vec<String>,
18    options: crate::cli::AskCommandOptions,
19) -> Result<()> {
20    let prompt_text = prompt.join(" ");
21    validate_agent_safe_text("prompt", &prompt_text)?;
22
23    if config.verbose {
24        tracing::debug!(model = %config.model, prompt = %prompt_text, "Sending prompt");
25    }
26
27    let request = LLMRequest {
28        messages: Arc::new(vec![Message::user(prompt_text)]),
29        system_prompt: Some(Arc::from(lightweight_instruction_text())),
30        model: config.model.clone(),
31        ..Default::default()
32    };
33    let provider_name = if config.provider.trim().is_empty() {
34        infer_provider_from_model(&request.model)
35            .map(|provider| provider.to_string())
36            .ok_or_else(|| anyhow::anyhow!("Cannot determine provider for model: {}", request.model))?
37    } else {
38        config.provider.to_lowercase()
39    };
40    let provider = create_provider_with_config(
41        &provider_name,
42        ProviderConfig {
43            api_key: Some(config.api_key.clone()),
44            openai_chatgpt_auth: config.openai_chatgpt_auth.clone(),
45            copilot_auth: None,
46            base_url: None,
47            model: Some(request.model.clone()),
48            prompt_cache: None,
49            timeouts: None,
50            openai: None,
51            anthropic: None,
52            model_behavior: config.model_behavior.clone(),
53            workspace_root: Some(config.workspace.clone()),
54        },
55    )?;
56    let backend_kind = provider.name().to_string();
57    let response = collect_single_response(provider.as_ref(), request).await?;
58    let response_model = if response.model.is_empty() {
59        config.model.clone()
60    } else {
61        response.model.clone()
62    };
63
64    // Handle output based on format preference
65    if let Some(crate::cli::args::AskOutputFormat::Json) = options.output_format {
66        // Build a comprehensive JSON structure
67        let output = serde_json::json!({
68            "response": response,
69            "provider": {
70                "kind": backend_kind,
71                "model": response_model,
72            }
73        });
74        use std::io::Write;
75        let mut stdout = std::io::stdout().lock();
76        serde_json::to_writer_pretty(&mut stdout, &output)?;
77        writeln!(stdout)?;
78    } else {
79        use std::io::Write;
80        let mut stdout = std::io::stdout().lock();
81        if is_pipe_output() {
82            if let Some(code_only) = extract_code_only(response.content_text()) {
83                write!(stdout, "{code_only}")?;
84            } else {
85                writeln!(stdout, "{}", response.content_text())?;
86            }
87        } else {
88            // Print the response content directly (default behavior)
89            writeln!(stdout, "{}", response.content_text())?;
90        }
91    }
92
93    Ok(())
94}
95
96fn is_pipe_output() -> bool {
97    !std::io::stdout().is_terminal()
98}
99
100fn extract_code_only(text: &str) -> Option<String> {
101    let blocks = extract_code_fence_blocks(text);
102    let block = select_best_code_block(&blocks)?;
103    let mut output = block.lines.join("\n");
104    if !output.ends_with('\n') {
105        output.push('\n');
106    }
107    Some(output)
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use crate::llm::collect_single_response;
114    use crate::llm::provider::{
115        FinishReason, LLMError, LLMNormalizedStream, LLMProvider, LLMResponse, LLMStream, LLMStreamEvent,
116        NormalizedStreamEvent, ReasoningSource, Usage,
117    };
118    use async_trait::async_trait;
119    use futures::stream;
120
121    #[derive(Clone)]
122    struct StreamingOnlyProvider;
123
124    #[async_trait]
125    impl LLMProvider for StreamingOnlyProvider {
126        fn name(&self) -> &str {
127            "test"
128        }
129
130        fn supports_streaming(&self) -> bool {
131            true
132        }
133
134        fn supports_non_streaming(&self, _model: &str) -> bool {
135            false
136        }
137
138        async fn generate(&self, _request: LLMRequest) -> Result<LLMResponse, LLMError> {
139            panic!("generate should not be called for streaming-only provider")
140        }
141
142        async fn stream(&self, _request: LLMRequest) -> Result<LLMStream, LLMError> {
143            Ok(Box::pin(stream::iter(vec![
144                Ok(LLMStreamEvent::Token { delta: "hello ".to_string() }),
145                Ok(LLMStreamEvent::Token { delta: "world".to_string() }),
146                Ok(LLMStreamEvent::Completed {
147                    response: Box::new(LLMResponse {
148                        content: None,
149                        model: "gpt-5.2-codex".to_string(),
150                        tool_calls: None,
151                        usage: None,
152                        finish_reason: FinishReason::Stop,
153                        reasoning: None,
154                        reasoning_details: None,
155                        organization_id: None,
156                        request_id: None,
157                        tool_references: Vec::new(),
158                        compaction: None,
159                    }),
160                }),
161            ])))
162        }
163
164        fn supported_models(&self) -> Vec<String> {
165            vec!["gpt-5.2-codex".to_string()]
166        }
167
168        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
169            Ok(())
170        }
171    }
172
173    #[tokio::test]
174    async fn collect_single_response_uses_stream_when_non_streaming_is_unsupported() {
175        let provider = StreamingOnlyProvider;
176        let response = collect_single_response(
177            &provider,
178            LLMRequest {
179                model: "gpt-5.2-codex".to_string(),
180                ..Default::default()
181            },
182        )
183        .await
184        .expect("stream collection should succeed");
185
186        assert_eq!(response.content.as_deref(), Some("hello world"));
187    }
188
189    #[derive(Clone)]
190    struct NormalizedOnlyProvider;
191
192    #[async_trait]
193    impl LLMProvider for NormalizedOnlyProvider {
194        fn name(&self) -> &str {
195            "test"
196        }
197
198        fn supports_streaming(&self) -> bool {
199            true
200        }
201
202        fn supports_non_streaming(&self, _model: &str) -> bool {
203            false
204        }
205
206        async fn generate(&self, _request: LLMRequest) -> Result<LLMResponse, LLMError> {
207            panic!("generate should not be called for streaming-only provider")
208        }
209
210        async fn stream(&self, _request: LLMRequest) -> Result<LLMStream, LLMError> {
211            panic!("legacy stream should not be used when normalized stream is available")
212        }
213
214        async fn stream_normalized(&self, _request: LLMRequest) -> Result<LLMNormalizedStream, LLMError> {
215            Ok(Box::pin(stream::iter(vec![
216                Ok(NormalizedStreamEvent::TextDelta { delta: "hello ".to_string() }),
217                Ok(NormalizedStreamEvent::ReasoningDelta {
218                    delta: "thinking ".to_string(),
219                    source: ReasoningSource::ProviderSummary,
220                }),
221                Ok(NormalizedStreamEvent::Usage {
222                    usage: Usage {
223                        prompt_tokens: 10,
224                        completion_tokens: 2,
225                        total_tokens: 12,
226                        cached_prompt_tokens: None,
227                        cache_creation_tokens: None,
228                        cache_read_tokens: None,
229                        iterations: None,
230                    },
231                }),
232                Ok(NormalizedStreamEvent::Done {
233                    response: Box::new(LLMResponse {
234                        content: None,
235                        model: "gpt-5.2-codex".to_string(),
236                        tool_calls: None,
237                        usage: None,
238                        finish_reason: FinishReason::Stop,
239                        reasoning: None,
240                        reasoning_details: None,
241                        organization_id: None,
242                        request_id: None,
243                        tool_references: Vec::new(),
244                        compaction: None,
245                    }),
246                }),
247            ])))
248        }
249
250        fn supported_models(&self) -> Vec<String> {
251            vec!["gpt-5.2-codex".to_string()]
252        }
253
254        fn validate_request(&self, _request: &LLMRequest) -> Result<(), LLMError> {
255            Ok(())
256        }
257    }
258
259    #[tokio::test]
260    async fn collect_single_response_prefers_normalized_stream() {
261        let provider = NormalizedOnlyProvider;
262        let response = collect_single_response(
263            &provider,
264            LLMRequest {
265                model: "gpt-5.2-codex".to_string(),
266                ..Default::default()
267            },
268        )
269        .await
270        .expect("normalized stream collection should succeed");
271
272        assert_eq!(response.content.as_deref(), Some("hello "));
273        assert_eq!(response.reasoning.as_deref(), Some("thinking "));
274        assert_eq!(response.usage.as_ref().map(|usage| usage.total_tokens), Some(12));
275    }
276}
277
278fn extract_code_fence_blocks(text: &str) -> Vec<CodeFenceBlock> {
279    let mut blocks = Vec::new();
280    let mut current_language: Option<String> = None;
281    let mut current_lines: Vec<String> = Vec::new();
282
283    for raw_line in text.lines() {
284        let trimmed_start = raw_line.trim_start();
285        if let Some(rest) = trimmed_start.strip_prefix("```") {
286            let rest_clean = rest.trim_matches('\r');
287            let rest_trimmed = rest_clean.trim();
288            if current_language.is_some() {
289                if rest_trimmed.is_empty() {
290                    let language = current_language.take().and_then(|lang| {
291                        let cleaned = lang.trim_matches(|ch| matches!(ch, '"' | '\'' | '`'));
292                        let cleaned = cleaned.trim();
293                        if cleaned.is_empty() {
294                            None
295                        } else {
296                            Some(cleaned.to_string())
297                        }
298                    });
299                    let block_lines = std::mem::take(&mut current_lines);
300                    blocks.push(CodeFenceBlock { language, lines: block_lines });
301                    continue;
302                }
303            } else {
304                let token = rest_trimmed.split_whitespace().next().unwrap_or_default();
305                let normalized = token.trim_matches(|ch| matches!(ch, '"' | '\'' | '`')).trim();
306                current_language = Some(normalized.to_ascii_lowercase());
307                current_lines.clear();
308                continue;
309            }
310        }
311
312        if current_language.is_some() {
313            current_lines.push(raw_line.trim_end_matches('\r').to_string());
314        }
315    }
316
317    blocks
318}
319
320fn select_best_code_block(blocks: &[CodeFenceBlock]) -> Option<&CodeFenceBlock> {
321    let mut best = None;
322    let mut best_score = (0usize, 0u8);
323    for block in blocks {
324        let score = score_code_block(block);
325        if score > best_score {
326            best_score = score;
327            best = Some(block);
328        }
329    }
330    best
331}
332
333fn score_code_block(block: &CodeFenceBlock) -> (usize, u8) {
334    let line_count = block.lines.iter().filter(|line| !line.trim().is_empty()).count();
335    let has_language = block.language.as_ref().is_some_and(|lang| !lang.trim().is_empty());
336    (line_count, if has_language { 1 } else { 0 })
337}
338
339#[derive(Debug, Clone, PartialEq, Eq)]
340struct CodeFenceBlock {
341    language: Option<String>,
342    lines: Vec<String>,
343}
vtcode_core/commands/ask.rs

vtcode_core/commands/
ask.rs
