claude_wrapper/command/
query.rs

1use crate::Claude;
2use crate::command::ClaudeCommand;
3use crate::error::Result;
4use crate::exec::{self, CommandOutput};
5use crate::types::{Effort, InputFormat, OutputFormat, PermissionMode};
6
7/// Builder for `claude -p <prompt>` (oneshot print-mode queries).
8///
9/// This is the primary command for programmatic use. It runs a single
10/// prompt through Claude and returns the result.
11///
12/// # Example
13///
14/// ```no_run
15/// use claude_wrapper::{Claude, ClaudeCommand, QueryCommand, OutputFormat};
16///
17/// # async fn example() -> claude_wrapper::Result<()> {
18/// let claude = Claude::builder().build()?;
19///
20/// let output = QueryCommand::new("explain this error: file not found")
21///     .model("sonnet")
22///     .output_format(OutputFormat::Json)
23///     .max_turns(1)
24///     .execute(&claude)
25///     .await?;
26/// # Ok(())
27/// # }
28/// ```
29#[derive(Debug, Clone)]
30pub struct QueryCommand {
31    prompt: String,
32    model: Option<String>,
33    system_prompt: Option<String>,
34    append_system_prompt: Option<String>,
35    output_format: Option<OutputFormat>,
36    max_budget_usd: Option<f64>,
37    permission_mode: Option<PermissionMode>,
38    allowed_tools: Vec<String>,
39    disallowed_tools: Vec<String>,
40    mcp_config: Vec<String>,
41    add_dir: Vec<String>,
42    effort: Option<Effort>,
43    max_turns: Option<u32>,
44    json_schema: Option<String>,
45    continue_session: bool,
46    resume: Option<String>,
47    session_id: Option<String>,
48    fallback_model: Option<String>,
49    no_session_persistence: bool,
50    dangerously_skip_permissions: bool,
51    agent: Option<String>,
52    agents_json: Option<String>,
53    tools: Vec<String>,
54    file: Vec<String>,
55    include_partial_messages: bool,
56    input_format: Option<InputFormat>,
57    strict_mcp_config: bool,
58    settings: Option<String>,
59    fork_session: bool,
60    retry_policy: Option<crate::retry::RetryPolicy>,
61}
62
63impl QueryCommand {
64    /// Create a new query command with the given prompt.
65    #[must_use]
66    pub fn new(prompt: impl Into<String>) -> Self {
67        Self {
68            prompt: prompt.into(),
69            model: None,
70            system_prompt: None,
71            append_system_prompt: None,
72            output_format: None,
73            max_budget_usd: None,
74            permission_mode: None,
75            allowed_tools: Vec::new(),
76            disallowed_tools: Vec::new(),
77            mcp_config: Vec::new(),
78            add_dir: Vec::new(),
79            effort: None,
80            max_turns: None,
81            json_schema: None,
82            continue_session: false,
83            resume: None,
84            session_id: None,
85            fallback_model: None,
86            no_session_persistence: false,
87            dangerously_skip_permissions: false,
88            agent: None,
89            agents_json: None,
90            tools: Vec::new(),
91            file: Vec::new(),
92            include_partial_messages: false,
93            input_format: None,
94            strict_mcp_config: false,
95            settings: None,
96            fork_session: false,
97            retry_policy: None,
98        }
99    }
100
101    /// Set the model to use (e.g. "sonnet", "opus", or a full model ID).
102    #[must_use]
103    pub fn model(mut self, model: impl Into<String>) -> Self {
104        self.model = Some(model.into());
105        self
106    }
107
108    /// Set a custom system prompt (replaces the default).
109    #[must_use]
110    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
111        self.system_prompt = Some(prompt.into());
112        self
113    }
114
115    /// Append to the default system prompt.
116    #[must_use]
117    pub fn append_system_prompt(mut self, prompt: impl Into<String>) -> Self {
118        self.append_system_prompt = Some(prompt.into());
119        self
120    }
121
122    /// Set the output format.
123    #[must_use]
124    pub fn output_format(mut self, format: OutputFormat) -> Self {
125        self.output_format = Some(format);
126        self
127    }
128
129    /// Set the maximum budget in USD.
130    #[must_use]
131    pub fn max_budget_usd(mut self, budget: f64) -> Self {
132        self.max_budget_usd = Some(budget);
133        self
134    }
135
136    /// Set the permission mode.
137    #[must_use]
138    pub fn permission_mode(mut self, mode: PermissionMode) -> Self {
139        self.permission_mode = Some(mode);
140        self
141    }
142
143    /// Add allowed tools (e.g. "Bash", "Read", "mcp__my-server__*").
144    #[must_use]
145    pub fn allowed_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
146        self.allowed_tools.extend(tools.into_iter().map(Into::into));
147        self
148    }
149
150    /// Add a single allowed tool.
151    #[must_use]
152    pub fn allowed_tool(mut self, tool: impl Into<String>) -> Self {
153        self.allowed_tools.push(tool.into());
154        self
155    }
156
157    /// Add disallowed tools.
158    #[must_use]
159    pub fn disallowed_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
160        self.disallowed_tools
161            .extend(tools.into_iter().map(Into::into));
162        self
163    }
164
165    /// Add an MCP config file path.
166    #[must_use]
167    pub fn mcp_config(mut self, path: impl Into<String>) -> Self {
168        self.mcp_config.push(path.into());
169        self
170    }
171
172    /// Add an additional directory for tool access.
173    #[must_use]
174    pub fn add_dir(mut self, dir: impl Into<String>) -> Self {
175        self.add_dir.push(dir.into());
176        self
177    }
178
179    /// Set the effort level.
180    #[must_use]
181    pub fn effort(mut self, effort: Effort) -> Self {
182        self.effort = Some(effort);
183        self
184    }
185
186    /// Set the maximum number of turns.
187    #[must_use]
188    pub fn max_turns(mut self, turns: u32) -> Self {
189        self.max_turns = Some(turns);
190        self
191    }
192
193    /// Set a JSON schema for structured output validation.
194    #[must_use]
195    pub fn json_schema(mut self, schema: impl Into<String>) -> Self {
196        self.json_schema = Some(schema.into());
197        self
198    }
199
200    /// Continue the most recent conversation.
201    #[must_use]
202    pub fn continue_session(mut self) -> Self {
203        self.continue_session = true;
204        self
205    }
206
207    /// Resume a specific session by ID.
208    #[must_use]
209    pub fn resume(mut self, session_id: impl Into<String>) -> Self {
210        self.resume = Some(session_id.into());
211        self
212    }
213
214    /// Use a specific session ID.
215    #[must_use]
216    pub fn session_id(mut self, id: impl Into<String>) -> Self {
217        self.session_id = Some(id.into());
218        self
219    }
220
221    /// Set a fallback model for when the primary model is overloaded.
222    #[must_use]
223    pub fn fallback_model(mut self, model: impl Into<String>) -> Self {
224        self.fallback_model = Some(model.into());
225        self
226    }
227
228    /// Disable session persistence (sessions won't be saved to disk).
229    #[must_use]
230    pub fn no_session_persistence(mut self) -> Self {
231        self.no_session_persistence = true;
232        self
233    }
234
235    /// Bypass all permission checks. Only use in sandboxed environments.
236    #[must_use]
237    pub fn dangerously_skip_permissions(mut self) -> Self {
238        self.dangerously_skip_permissions = true;
239        self
240    }
241
242    /// Set the agent for the session.
243    #[must_use]
244    pub fn agent(mut self, agent: impl Into<String>) -> Self {
245        self.agent = Some(agent.into());
246        self
247    }
248
249    /// Set custom agents as a JSON object.
250    ///
251    /// Example: `{"reviewer": {"description": "Reviews code", "prompt": "You are a code reviewer"}}`
252    #[must_use]
253    pub fn agents_json(mut self, json: impl Into<String>) -> Self {
254        self.agents_json = Some(json.into());
255        self
256    }
257
258    /// Set the list of available built-in tools.
259    ///
260    /// Use `""` to disable all tools, `"default"` for all tools, or
261    /// specific tool names like `["Bash", "Edit", "Read"]`.
262    /// This is different from `allowed_tools` which controls MCP tool permissions.
263    #[must_use]
264    pub fn tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
265        self.tools.extend(tools.into_iter().map(Into::into));
266        self
267    }
268
269    /// Add a file resource to download at startup.
270    ///
271    /// Format: `file_id:relative_path` (e.g. `file_abc:doc.txt`).
272    #[must_use]
273    pub fn file(mut self, spec: impl Into<String>) -> Self {
274        self.file.push(spec.into());
275        self
276    }
277
278    /// Include partial message chunks as they arrive.
279    ///
280    /// Only works with `--output-format stream-json`.
281    #[must_use]
282    pub fn include_partial_messages(mut self) -> Self {
283        self.include_partial_messages = true;
284        self
285    }
286
287    /// Set the input format.
288    #[must_use]
289    pub fn input_format(mut self, format: InputFormat) -> Self {
290        self.input_format = Some(format);
291        self
292    }
293
294    /// Only use MCP servers from `--mcp-config`, ignoring all other MCP configurations.
295    #[must_use]
296    pub fn strict_mcp_config(mut self) -> Self {
297        self.strict_mcp_config = true;
298        self
299    }
300
301    /// Path to a settings JSON file or a JSON string.
302    #[must_use]
303    pub fn settings(mut self, settings: impl Into<String>) -> Self {
304        self.settings = Some(settings.into());
305        self
306    }
307
308    /// When resuming, create a new session ID instead of reusing the original.
309    #[must_use]
310    pub fn fork_session(mut self) -> Self {
311        self.fork_session = true;
312        self
313    }
314
315    /// Set a per-command retry policy, overriding the client default.
316    ///
317    /// # Example
318    ///
319    /// ```no_run
320    /// use claude_wrapper::{Claude, ClaudeCommand, QueryCommand, RetryPolicy};
321    /// use std::time::Duration;
322    ///
323    /// # async fn example() -> claude_wrapper::Result<()> {
324    /// let claude = Claude::builder().build()?;
325    ///
326    /// let output = QueryCommand::new("explain quicksort")
327    ///     .retry(RetryPolicy::new()
328    ///         .max_attempts(5)
329    ///         .initial_backoff(Duration::from_secs(2))
330    ///         .exponential()
331    ///         .retry_on_timeout(true))
332    ///     .execute(&claude)
333    ///     .await?;
334    /// # Ok(())
335    /// # }
336    /// ```
337    #[must_use]
338    pub fn retry(mut self, policy: crate::retry::RetryPolicy) -> Self {
339        self.retry_policy = Some(policy);
340        self
341    }
342
343    /// Return the full command as a string that could be run in a shell.
344    ///
345    /// Constructs a command string using the binary path from the Claude instance
346    /// and the arguments from this query. Arguments containing spaces or special
347    /// shell characters are shell-quoted to be safe for shell execution.
348    ///
349    /// # Example
350    ///
351    /// ```no_run
352    /// use claude_wrapper::{Claude, QueryCommand};
353    ///
354    /// # async fn example() -> claude_wrapper::Result<()> {
355    /// let claude = Claude::builder().build()?;
356    ///
357    /// let cmd = QueryCommand::new("explain quicksort")
358    ///     .model("sonnet");
359    ///
360    /// let command_str = cmd.to_command_string(&claude);
361    /// println!("Would run: {}", command_str);
362    /// # Ok(())
363    /// # }
364    /// ```
365    pub fn to_command_string(&self, claude: &Claude) -> String {
366        let args = self.build_args();
367        let quoted_args = args.iter().map(|arg| shell_quote(arg)).collect::<Vec<_>>();
368        format!("{} {}", claude.binary().display(), quoted_args.join(" "))
369    }
370
371    /// Execute the query and parse the JSON result.
372    ///
373    /// This is a convenience method that sets `OutputFormat::Json` and
374    /// deserializes the response into a [`QueryResult`](crate::types::QueryResult).
375    #[cfg(feature = "json")]
376    pub async fn execute_json(&self, claude: &Claude) -> Result<crate::types::QueryResult> {
377        // Build args with JSON output format forced
378        let mut args = self.build_args();
379
380        // Override output format to json if not already set
381        if self.output_format.is_none() {
382            args.push("--output-format".to_string());
383            args.push("json".to_string());
384        }
385
386        let output = exec::run_claude_with_retry(claude, args, self.retry_policy.as_ref()).await?;
387
388        serde_json::from_str(&output.stdout).map_err(|e| crate::error::Error::Json {
389            message: format!("failed to parse query result: {e}"),
390            source: e,
391        })
392    }
393
394    fn build_args(&self) -> Vec<String> {
395        let mut args = vec!["--print".to_string()];
396
397        if let Some(ref model) = self.model {
398            args.push("--model".to_string());
399            args.push(model.clone());
400        }
401
402        if let Some(ref prompt) = self.system_prompt {
403            args.push("--system-prompt".to_string());
404            args.push(prompt.clone());
405        }
406
407        if let Some(ref prompt) = self.append_system_prompt {
408            args.push("--append-system-prompt".to_string());
409            args.push(prompt.clone());
410        }
411
412        if let Some(ref format) = self.output_format {
413            args.push("--output-format".to_string());
414            args.push(format.as_arg().to_string());
415            // CLI v2.1.72+ requires --verbose when using stream-json with --print
416            if matches!(format, OutputFormat::StreamJson) {
417                args.push("--verbose".to_string());
418            }
419        }
420
421        if let Some(budget) = self.max_budget_usd {
422            args.push("--max-budget-usd".to_string());
423            args.push(budget.to_string());
424        }
425
426        if let Some(ref mode) = self.permission_mode {
427            args.push("--permission-mode".to_string());
428            args.push(mode.as_arg().to_string());
429        }
430
431        if !self.allowed_tools.is_empty() {
432            args.push("--allowed-tools".to_string());
433            args.push(self.allowed_tools.join(","));
434        }
435
436        if !self.disallowed_tools.is_empty() {
437            args.push("--disallowed-tools".to_string());
438            args.push(self.disallowed_tools.join(","));
439        }
440
441        for config in &self.mcp_config {
442            args.push("--mcp-config".to_string());
443            args.push(config.clone());
444        }
445
446        for dir in &self.add_dir {
447            args.push("--add-dir".to_string());
448            args.push(dir.clone());
449        }
450
451        if let Some(ref effort) = self.effort {
452            args.push("--effort".to_string());
453            args.push(effort.as_arg().to_string());
454        }
455
456        if let Some(turns) = self.max_turns {
457            args.push("--max-turns".to_string());
458            args.push(turns.to_string());
459        }
460
461        if let Some(ref schema) = self.json_schema {
462            args.push("--json-schema".to_string());
463            args.push(schema.clone());
464        }
465
466        if self.continue_session {
467            args.push("--continue".to_string());
468        }
469
470        if let Some(ref session_id) = self.resume {
471            args.push("--resume".to_string());
472            args.push(session_id.clone());
473        }
474
475        if let Some(ref id) = self.session_id {
476            args.push("--session-id".to_string());
477            args.push(id.clone());
478        }
479
480        if let Some(ref model) = self.fallback_model {
481            args.push("--fallback-model".to_string());
482            args.push(model.clone());
483        }
484
485        if self.no_session_persistence {
486            args.push("--no-session-persistence".to_string());
487        }
488
489        if self.dangerously_skip_permissions {
490            args.push("--dangerously-skip-permissions".to_string());
491        }
492
493        if let Some(ref agent) = self.agent {
494            args.push("--agent".to_string());
495            args.push(agent.clone());
496        }
497
498        if let Some(ref agents) = self.agents_json {
499            args.push("--agents".to_string());
500            args.push(agents.clone());
501        }
502
503        if !self.tools.is_empty() {
504            args.push("--tools".to_string());
505            args.push(self.tools.join(","));
506        }
507
508        for spec in &self.file {
509            args.push("--file".to_string());
510            args.push(spec.clone());
511        }
512
513        if self.include_partial_messages {
514            args.push("--include-partial-messages".to_string());
515        }
516
517        if let Some(ref format) = self.input_format {
518            args.push("--input-format".to_string());
519            args.push(format.as_arg().to_string());
520        }
521
522        if self.strict_mcp_config {
523            args.push("--strict-mcp-config".to_string());
524        }
525
526        if let Some(ref settings) = self.settings {
527            args.push("--settings".to_string());
528            args.push(settings.clone());
529        }
530
531        if self.fork_session {
532            args.push("--fork-session".to_string());
533        }
534
535        // Prompt is the positional argument at the end
536        args.push(self.prompt.clone());
537
538        args
539    }
540}
541
542impl ClaudeCommand for QueryCommand {
543    type Output = CommandOutput;
544
545    fn args(&self) -> Vec<String> {
546        self.build_args()
547    }
548
549    async fn execute(&self, claude: &Claude) -> Result<CommandOutput> {
550        exec::run_claude_with_retry(claude, self.args(), self.retry_policy.as_ref()).await
551    }
552}
553
554/// Shell-quote an argument if it contains spaces or special characters.
555fn shell_quote(arg: &str) -> String {
556    // Check if the argument needs quoting (contains whitespace or shell metacharacters)
557    if arg.contains(|c: char| c.is_whitespace() || "\"'$\\`|;<>&()[]{}".contains(c)) {
558        // Use single quotes and escape any existing single quotes
559        format!("'{}'", arg.replace("'", "'\\''"))
560    } else {
561        arg.to_string()
562    }
563}
564
565#[cfg(test)]
566mod tests {
567    use super::*;
568
569    #[test]
570    fn test_basic_query_args() {
571        let cmd = QueryCommand::new("hello world");
572        let args = cmd.args();
573        assert_eq!(args, vec!["--print", "hello world"]);
574    }
575
576    #[test]
577    fn test_full_query_args() {
578        let cmd = QueryCommand::new("explain this")
579            .model("sonnet")
580            .system_prompt("be concise")
581            .output_format(OutputFormat::Json)
582            .max_budget_usd(0.50)
583            .permission_mode(PermissionMode::BypassPermissions)
584            .allowed_tools(["Bash", "Read"])
585            .mcp_config("/tmp/mcp.json")
586            .effort(Effort::High)
587            .max_turns(3)
588            .no_session_persistence();
589
590        let args = cmd.args();
591        assert!(args.contains(&"--print".to_string()));
592        assert!(args.contains(&"--model".to_string()));
593        assert!(args.contains(&"sonnet".to_string()));
594        assert!(args.contains(&"--system-prompt".to_string()));
595        assert!(args.contains(&"--output-format".to_string()));
596        assert!(args.contains(&"json".to_string()));
597        // json format should NOT include --verbose (only stream-json needs it)
598        assert!(!args.contains(&"--verbose".to_string()));
599        assert!(args.contains(&"--max-budget-usd".to_string()));
600        assert!(args.contains(&"--permission-mode".to_string()));
601        assert!(args.contains(&"bypassPermissions".to_string()));
602        assert!(args.contains(&"--allowed-tools".to_string()));
603        assert!(args.contains(&"Bash,Read".to_string()));
604        assert!(args.contains(&"--effort".to_string()));
605        assert!(args.contains(&"high".to_string()));
606        assert!(args.contains(&"--max-turns".to_string()));
607        assert!(args.contains(&"--no-session-persistence".to_string()));
608        // Prompt is last
609        assert_eq!(args.last().unwrap(), "explain this");
610    }
611
612    #[test]
613    fn test_stream_json_includes_verbose() {
614        let cmd = QueryCommand::new("test").output_format(OutputFormat::StreamJson);
615        let args = cmd.args();
616        assert!(args.contains(&"--output-format".to_string()));
617        assert!(args.contains(&"stream-json".to_string()));
618        assert!(args.contains(&"--verbose".to_string()));
619    }
620
621    #[test]
622    fn test_to_command_string_simple() {
623        let claude = Claude::builder()
624            .binary("/usr/local/bin/claude")
625            .build()
626            .unwrap();
627
628        let cmd = QueryCommand::new("hello");
629        let command_str = cmd.to_command_string(&claude);
630
631        assert!(command_str.starts_with("/usr/local/bin/claude"));
632        assert!(command_str.contains("--print"));
633        assert!(command_str.contains("hello"));
634    }
635
636    #[test]
637    fn test_to_command_string_with_spaces() {
638        let claude = Claude::builder()
639            .binary("/usr/local/bin/claude")
640            .build()
641            .unwrap();
642
643        let cmd = QueryCommand::new("hello world").model("sonnet");
644        let command_str = cmd.to_command_string(&claude);
645
646        assert!(command_str.starts_with("/usr/local/bin/claude"));
647        assert!(command_str.contains("--print"));
648        // Prompt with spaces should be quoted
649        assert!(command_str.contains("'hello world'"));
650        assert!(command_str.contains("--model"));
651        assert!(command_str.contains("sonnet"));
652    }
653
654    #[test]
655    fn test_to_command_string_with_special_chars() {
656        let claude = Claude::builder()
657            .binary("/usr/local/bin/claude")
658            .build()
659            .unwrap();
660
661        let cmd = QueryCommand::new("test $VAR and `cmd`");
662        let command_str = cmd.to_command_string(&claude);
663
664        // Arguments with special shell characters should be quoted
665        assert!(command_str.contains("'test $VAR and `cmd`'"));
666    }
667
668    #[test]
669    fn test_to_command_string_with_single_quotes() {
670        let claude = Claude::builder()
671            .binary("/usr/local/bin/claude")
672            .build()
673            .unwrap();
674
675        let cmd = QueryCommand::new("it's");
676        let command_str = cmd.to_command_string(&claude);
677
678        // Single quotes should be escaped in shell
679        assert!(command_str.contains("'it'\\''s'"));
680    }
681}
claude_wrapper/command/query.rs

claude_wrapper/command/
query.rs
