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        }
416
417        if let Some(budget) = self.max_budget_usd {
418            args.push("--max-budget-usd".to_string());
419            args.push(budget.to_string());
420        }
421
422        if let Some(ref mode) = self.permission_mode {
423            args.push("--permission-mode".to_string());
424            args.push(mode.as_arg().to_string());
425        }
426
427        if !self.allowed_tools.is_empty() {
428            args.push("--allowed-tools".to_string());
429            args.push(self.allowed_tools.join(","));
430        }
431
432        if !self.disallowed_tools.is_empty() {
433            args.push("--disallowed-tools".to_string());
434            args.push(self.disallowed_tools.join(","));
435        }
436
437        for config in &self.mcp_config {
438            args.push("--mcp-config".to_string());
439            args.push(config.clone());
440        }
441
442        for dir in &self.add_dir {
443            args.push("--add-dir".to_string());
444            args.push(dir.clone());
445        }
446
447        if let Some(ref effort) = self.effort {
448            args.push("--effort".to_string());
449            args.push(effort.as_arg().to_string());
450        }
451
452        if let Some(turns) = self.max_turns {
453            args.push("--max-turns".to_string());
454            args.push(turns.to_string());
455        }
456
457        if let Some(ref schema) = self.json_schema {
458            args.push("--json-schema".to_string());
459            args.push(schema.clone());
460        }
461
462        if self.continue_session {
463            args.push("--continue".to_string());
464        }
465
466        if let Some(ref session_id) = self.resume {
467            args.push("--resume".to_string());
468            args.push(session_id.clone());
469        }
470
471        if let Some(ref id) = self.session_id {
472            args.push("--session-id".to_string());
473            args.push(id.clone());
474        }
475
476        if let Some(ref model) = self.fallback_model {
477            args.push("--fallback-model".to_string());
478            args.push(model.clone());
479        }
480
481        if self.no_session_persistence {
482            args.push("--no-session-persistence".to_string());
483        }
484
485        if self.dangerously_skip_permissions {
486            args.push("--dangerously-skip-permissions".to_string());
487        }
488
489        if let Some(ref agent) = self.agent {
490            args.push("--agent".to_string());
491            args.push(agent.clone());
492        }
493
494        if let Some(ref agents) = self.agents_json {
495            args.push("--agents".to_string());
496            args.push(agents.clone());
497        }
498
499        if !self.tools.is_empty() {
500            args.push("--tools".to_string());
501            args.push(self.tools.join(","));
502        }
503
504        for spec in &self.file {
505            args.push("--file".to_string());
506            args.push(spec.clone());
507        }
508
509        if self.include_partial_messages {
510            args.push("--include-partial-messages".to_string());
511        }
512
513        if let Some(ref format) = self.input_format {
514            args.push("--input-format".to_string());
515            args.push(format.as_arg().to_string());
516        }
517
518        if self.strict_mcp_config {
519            args.push("--strict-mcp-config".to_string());
520        }
521
522        if let Some(ref settings) = self.settings {
523            args.push("--settings".to_string());
524            args.push(settings.clone());
525        }
526
527        if self.fork_session {
528            args.push("--fork-session".to_string());
529        }
530
531        // Prompt is the positional argument at the end
532        args.push(self.prompt.clone());
533
534        args
535    }
536}
537
538impl ClaudeCommand for QueryCommand {
539    type Output = CommandOutput;
540
541    fn args(&self) -> Vec<String> {
542        self.build_args()
543    }
544
545    async fn execute(&self, claude: &Claude) -> Result<CommandOutput> {
546        exec::run_claude_with_retry(claude, self.args(), self.retry_policy.as_ref()).await
547    }
548}
549
550/// Shell-quote an argument if it contains spaces or special characters.
551fn shell_quote(arg: &str) -> String {
552    // Check if the argument needs quoting (contains whitespace or shell metacharacters)
553    if arg.contains(|c: char| c.is_whitespace() || "\"'$\\`|;<>&()[]{}".contains(c)) {
554        // Use single quotes and escape any existing single quotes
555        format!("'{}'", arg.replace("'", "'\\''"))
556    } else {
557        arg.to_string()
558    }
559}
560
561#[cfg(test)]
562mod tests {
563    use super::*;
564
565    #[test]
566    fn test_basic_query_args() {
567        let cmd = QueryCommand::new("hello world");
568        let args = cmd.args();
569        assert_eq!(args, vec!["--print", "hello world"]);
570    }
571
572    #[test]
573    fn test_full_query_args() {
574        let cmd = QueryCommand::new("explain this")
575            .model("sonnet")
576            .system_prompt("be concise")
577            .output_format(OutputFormat::Json)
578            .max_budget_usd(0.50)
579            .permission_mode(PermissionMode::BypassPermissions)
580            .allowed_tools(["Bash", "Read"])
581            .mcp_config("/tmp/mcp.json")
582            .effort(Effort::High)
583            .max_turns(3)
584            .no_session_persistence();
585
586        let args = cmd.args();
587        assert!(args.contains(&"--print".to_string()));
588        assert!(args.contains(&"--model".to_string()));
589        assert!(args.contains(&"sonnet".to_string()));
590        assert!(args.contains(&"--system-prompt".to_string()));
591        assert!(args.contains(&"--output-format".to_string()));
592        assert!(args.contains(&"json".to_string()));
593        assert!(args.contains(&"--max-budget-usd".to_string()));
594        assert!(args.contains(&"--permission-mode".to_string()));
595        assert!(args.contains(&"bypassPermissions".to_string()));
596        assert!(args.contains(&"--allowed-tools".to_string()));
597        assert!(args.contains(&"Bash,Read".to_string()));
598        assert!(args.contains(&"--effort".to_string()));
599        assert!(args.contains(&"high".to_string()));
600        assert!(args.contains(&"--max-turns".to_string()));
601        assert!(args.contains(&"--no-session-persistence".to_string()));
602        // Prompt is last
603        assert_eq!(args.last().unwrap(), "explain this");
604    }
605
606    #[test]
607    fn test_to_command_string_simple() {
608        let claude = Claude::builder()
609            .binary("/usr/local/bin/claude")
610            .build()
611            .unwrap();
612
613        let cmd = QueryCommand::new("hello");
614        let command_str = cmd.to_command_string(&claude);
615
616        assert!(command_str.starts_with("/usr/local/bin/claude"));
617        assert!(command_str.contains("--print"));
618        assert!(command_str.contains("hello"));
619    }
620
621    #[test]
622    fn test_to_command_string_with_spaces() {
623        let claude = Claude::builder()
624            .binary("/usr/local/bin/claude")
625            .build()
626            .unwrap();
627
628        let cmd = QueryCommand::new("hello world").model("sonnet");
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        // Prompt with spaces should be quoted
634        assert!(command_str.contains("'hello world'"));
635        assert!(command_str.contains("--model"));
636        assert!(command_str.contains("sonnet"));
637    }
638
639    #[test]
640    fn test_to_command_string_with_special_chars() {
641        let claude = Claude::builder()
642            .binary("/usr/local/bin/claude")
643            .build()
644            .unwrap();
645
646        let cmd = QueryCommand::new("test $VAR and `cmd`");
647        let command_str = cmd.to_command_string(&claude);
648
649        // Arguments with special shell characters should be quoted
650        assert!(command_str.contains("'test $VAR and `cmd`'"));
651    }
652
653    #[test]
654    fn test_to_command_string_with_single_quotes() {
655        let claude = Claude::builder()
656            .binary("/usr/local/bin/claude")
657            .build()
658            .unwrap();
659
660        let cmd = QueryCommand::new("it's");
661        let command_str = cmd.to_command_string(&claude);
662
663        // Single quotes should be escaped in shell
664        assert!(command_str.contains("'it'\\''s'"));
665    }
666}
claude_wrapper/command/query.rs

claude_wrapper/command/
query.rs
