Skip to main content

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}
61
62impl QueryCommand {
63    /// Create a new query command with the given prompt.
64    #[must_use]
65    pub fn new(prompt: impl Into<String>) -> Self {
66        Self {
67            prompt: prompt.into(),
68            model: None,
69            system_prompt: None,
70            append_system_prompt: None,
71            output_format: None,
72            max_budget_usd: None,
73            permission_mode: None,
74            allowed_tools: Vec::new(),
75            disallowed_tools: Vec::new(),
76            mcp_config: Vec::new(),
77            add_dir: Vec::new(),
78            effort: None,
79            max_turns: None,
80            json_schema: None,
81            continue_session: false,
82            resume: None,
83            session_id: None,
84            fallback_model: None,
85            no_session_persistence: false,
86            dangerously_skip_permissions: false,
87            agent: None,
88            agents_json: None,
89            tools: Vec::new(),
90            file: Vec::new(),
91            include_partial_messages: false,
92            input_format: None,
93            strict_mcp_config: false,
94            settings: None,
95            fork_session: false,
96        }
97    }
98
99    /// Set the model to use (e.g. "sonnet", "opus", or a full model ID).
100    #[must_use]
101    pub fn model(mut self, model: impl Into<String>) -> Self {
102        self.model = Some(model.into());
103        self
104    }
105
106    /// Set a custom system prompt (replaces the default).
107    #[must_use]
108    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
109        self.system_prompt = Some(prompt.into());
110        self
111    }
112
113    /// Append to the default system prompt.
114    #[must_use]
115    pub fn append_system_prompt(mut self, prompt: impl Into<String>) -> Self {
116        self.append_system_prompt = Some(prompt.into());
117        self
118    }
119
120    /// Set the output format.
121    #[must_use]
122    pub fn output_format(mut self, format: OutputFormat) -> Self {
123        self.output_format = Some(format);
124        self
125    }
126
127    /// Set the maximum budget in USD.
128    #[must_use]
129    pub fn max_budget_usd(mut self, budget: f64) -> Self {
130        self.max_budget_usd = Some(budget);
131        self
132    }
133
134    /// Set the permission mode.
135    #[must_use]
136    pub fn permission_mode(mut self, mode: PermissionMode) -> Self {
137        self.permission_mode = Some(mode);
138        self
139    }
140
141    /// Add allowed tools (e.g. "Bash", "Read", "mcp__my-server__*").
142    #[must_use]
143    pub fn allowed_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
144        self.allowed_tools.extend(tools.into_iter().map(Into::into));
145        self
146    }
147
148    /// Add a single allowed tool.
149    #[must_use]
150    pub fn allowed_tool(mut self, tool: impl Into<String>) -> Self {
151        self.allowed_tools.push(tool.into());
152        self
153    }
154
155    /// Add disallowed tools.
156    #[must_use]
157    pub fn disallowed_tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
158        self.disallowed_tools
159            .extend(tools.into_iter().map(Into::into));
160        self
161    }
162
163    /// Add an MCP config file path.
164    #[must_use]
165    pub fn mcp_config(mut self, path: impl Into<String>) -> Self {
166        self.mcp_config.push(path.into());
167        self
168    }
169
170    /// Add an additional directory for tool access.
171    #[must_use]
172    pub fn add_dir(mut self, dir: impl Into<String>) -> Self {
173        self.add_dir.push(dir.into());
174        self
175    }
176
177    /// Set the effort level.
178    #[must_use]
179    pub fn effort(mut self, effort: Effort) -> Self {
180        self.effort = Some(effort);
181        self
182    }
183
184    /// Set the maximum number of turns.
185    #[must_use]
186    pub fn max_turns(mut self, turns: u32) -> Self {
187        self.max_turns = Some(turns);
188        self
189    }
190
191    /// Set a JSON schema for structured output validation.
192    #[must_use]
193    pub fn json_schema(mut self, schema: impl Into<String>) -> Self {
194        self.json_schema = Some(schema.into());
195        self
196    }
197
198    /// Continue the most recent conversation.
199    #[must_use]
200    pub fn continue_session(mut self) -> Self {
201        self.continue_session = true;
202        self
203    }
204
205    /// Resume a specific session by ID.
206    #[must_use]
207    pub fn resume(mut self, session_id: impl Into<String>) -> Self {
208        self.resume = Some(session_id.into());
209        self
210    }
211
212    /// Use a specific session ID.
213    #[must_use]
214    pub fn session_id(mut self, id: impl Into<String>) -> Self {
215        self.session_id = Some(id.into());
216        self
217    }
218
219    /// Set a fallback model for when the primary model is overloaded.
220    #[must_use]
221    pub fn fallback_model(mut self, model: impl Into<String>) -> Self {
222        self.fallback_model = Some(model.into());
223        self
224    }
225
226    /// Disable session persistence (sessions won't be saved to disk).
227    #[must_use]
228    pub fn no_session_persistence(mut self) -> Self {
229        self.no_session_persistence = true;
230        self
231    }
232
233    /// Bypass all permission checks. Only use in sandboxed environments.
234    #[must_use]
235    pub fn dangerously_skip_permissions(mut self) -> Self {
236        self.dangerously_skip_permissions = true;
237        self
238    }
239
240    /// Set the agent for the session.
241    #[must_use]
242    pub fn agent(mut self, agent: impl Into<String>) -> Self {
243        self.agent = Some(agent.into());
244        self
245    }
246
247    /// Set custom agents as a JSON object.
248    ///
249    /// Example: `{"reviewer": {"description": "Reviews code", "prompt": "You are a code reviewer"}}`
250    #[must_use]
251    pub fn agents_json(mut self, json: impl Into<String>) -> Self {
252        self.agents_json = Some(json.into());
253        self
254    }
255
256    /// Set the list of available built-in tools.
257    ///
258    /// Use `""` to disable all tools, `"default"` for all tools, or
259    /// specific tool names like `["Bash", "Edit", "Read"]`.
260    /// This is different from `allowed_tools` which controls MCP tool permissions.
261    #[must_use]
262    pub fn tools(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
263        self.tools.extend(tools.into_iter().map(Into::into));
264        self
265    }
266
267    /// Add a file resource to download at startup.
268    ///
269    /// Format: `file_id:relative_path` (e.g. `file_abc:doc.txt`).
270    #[must_use]
271    pub fn file(mut self, spec: impl Into<String>) -> Self {
272        self.file.push(spec.into());
273        self
274    }
275
276    /// Include partial message chunks as they arrive.
277    ///
278    /// Only works with `--output-format stream-json`.
279    #[must_use]
280    pub fn include_partial_messages(mut self) -> Self {
281        self.include_partial_messages = true;
282        self
283    }
284
285    /// Set the input format.
286    #[must_use]
287    pub fn input_format(mut self, format: InputFormat) -> Self {
288        self.input_format = Some(format);
289        self
290    }
291
292    /// Only use MCP servers from `--mcp-config`, ignoring all other MCP configurations.
293    #[must_use]
294    pub fn strict_mcp_config(mut self) -> Self {
295        self.strict_mcp_config = true;
296        self
297    }
298
299    /// Path to a settings JSON file or a JSON string.
300    #[must_use]
301    pub fn settings(mut self, settings: impl Into<String>) -> Self {
302        self.settings = Some(settings.into());
303        self
304    }
305
306    /// When resuming, create a new session ID instead of reusing the original.
307    #[must_use]
308    pub fn fork_session(mut self) -> Self {
309        self.fork_session = true;
310        self
311    }
312
313    /// Execute the query and parse the JSON result.
314    ///
315    /// This is a convenience method that sets `OutputFormat::Json` and
316    /// deserializes the response into a [`QueryResult`](crate::types::QueryResult).
317    #[cfg(feature = "json")]
318    pub async fn execute_json(&self, claude: &Claude) -> Result<crate::types::QueryResult> {
319        // Build args with JSON output format forced
320        let mut args = self.build_args();
321
322        // Override output format to json if not already set
323        if self.output_format.is_none() {
324            args.push("--output-format".to_string());
325            args.push("json".to_string());
326        }
327
328        let output = exec::run_claude(claude, args).await?;
329
330        serde_json::from_str(&output.stdout).map_err(|e| crate::error::Error::Json {
331            message: format!("failed to parse query result: {e}"),
332            source: e,
333        })
334    }
335
336    fn build_args(&self) -> Vec<String> {
337        let mut args = vec!["--print".to_string()];
338
339        if let Some(ref model) = self.model {
340            args.push("--model".to_string());
341            args.push(model.clone());
342        }
343
344        if let Some(ref prompt) = self.system_prompt {
345            args.push("--system-prompt".to_string());
346            args.push(prompt.clone());
347        }
348
349        if let Some(ref prompt) = self.append_system_prompt {
350            args.push("--append-system-prompt".to_string());
351            args.push(prompt.clone());
352        }
353
354        if let Some(ref format) = self.output_format {
355            args.push("--output-format".to_string());
356            args.push(format.as_arg().to_string());
357        }
358
359        if let Some(budget) = self.max_budget_usd {
360            args.push("--max-budget-usd".to_string());
361            args.push(budget.to_string());
362        }
363
364        if let Some(ref mode) = self.permission_mode {
365            args.push("--permission-mode".to_string());
366            args.push(mode.as_arg().to_string());
367        }
368
369        if !self.allowed_tools.is_empty() {
370            args.push("--allowed-tools".to_string());
371            args.push(self.allowed_tools.join(","));
372        }
373
374        if !self.disallowed_tools.is_empty() {
375            args.push("--disallowed-tools".to_string());
376            args.push(self.disallowed_tools.join(","));
377        }
378
379        for config in &self.mcp_config {
380            args.push("--mcp-config".to_string());
381            args.push(config.clone());
382        }
383
384        for dir in &self.add_dir {
385            args.push("--add-dir".to_string());
386            args.push(dir.clone());
387        }
388
389        if let Some(ref effort) = self.effort {
390            args.push("--effort".to_string());
391            args.push(effort.as_arg().to_string());
392        }
393
394        if let Some(turns) = self.max_turns {
395            args.push("--max-turns".to_string());
396            args.push(turns.to_string());
397        }
398
399        if let Some(ref schema) = self.json_schema {
400            args.push("--json-schema".to_string());
401            args.push(schema.clone());
402        }
403
404        if self.continue_session {
405            args.push("--continue".to_string());
406        }
407
408        if let Some(ref session_id) = self.resume {
409            args.push("--resume".to_string());
410            args.push(session_id.clone());
411        }
412
413        if let Some(ref id) = self.session_id {
414            args.push("--session-id".to_string());
415            args.push(id.clone());
416        }
417
418        if let Some(ref model) = self.fallback_model {
419            args.push("--fallback-model".to_string());
420            args.push(model.clone());
421        }
422
423        if self.no_session_persistence {
424            args.push("--no-session-persistence".to_string());
425        }
426
427        if self.dangerously_skip_permissions {
428            args.push("--dangerously-skip-permissions".to_string());
429        }
430
431        if let Some(ref agent) = self.agent {
432            args.push("--agent".to_string());
433            args.push(agent.clone());
434        }
435
436        if let Some(ref agents) = self.agents_json {
437            args.push("--agents".to_string());
438            args.push(agents.clone());
439        }
440
441        if !self.tools.is_empty() {
442            args.push("--tools".to_string());
443            args.push(self.tools.join(","));
444        }
445
446        for spec in &self.file {
447            args.push("--file".to_string());
448            args.push(spec.clone());
449        }
450
451        if self.include_partial_messages {
452            args.push("--include-partial-messages".to_string());
453        }
454
455        if let Some(ref format) = self.input_format {
456            args.push("--input-format".to_string());
457            args.push(format.as_arg().to_string());
458        }
459
460        if self.strict_mcp_config {
461            args.push("--strict-mcp-config".to_string());
462        }
463
464        if let Some(ref settings) = self.settings {
465            args.push("--settings".to_string());
466            args.push(settings.clone());
467        }
468
469        if self.fork_session {
470            args.push("--fork-session".to_string());
471        }
472
473        // Prompt is the positional argument at the end
474        args.push(self.prompt.clone());
475
476        args
477    }
478}
479
480impl ClaudeCommand for QueryCommand {
481    type Output = CommandOutput;
482
483    fn args(&self) -> Vec<String> {
484        self.build_args()
485    }
486
487    async fn execute(&self, claude: &Claude) -> Result<CommandOutput> {
488        exec::run_claude(claude, self.args()).await
489    }
490}
491
492#[cfg(test)]
493mod tests {
494    use super::*;
495
496    #[test]
497    fn test_basic_query_args() {
498        let cmd = QueryCommand::new("hello world");
499        let args = cmd.args();
500        assert_eq!(args, vec!["--print", "hello world"]);
501    }
502
503    #[test]
504    fn test_full_query_args() {
505        let cmd = QueryCommand::new("explain this")
506            .model("sonnet")
507            .system_prompt("be concise")
508            .output_format(OutputFormat::Json)
509            .max_budget_usd(0.50)
510            .permission_mode(PermissionMode::BypassPermissions)
511            .allowed_tools(["Bash", "Read"])
512            .mcp_config("/tmp/mcp.json")
513            .effort(Effort::High)
514            .max_turns(3)
515            .no_session_persistence();
516
517        let args = cmd.args();
518        assert!(args.contains(&"--print".to_string()));
519        assert!(args.contains(&"--model".to_string()));
520        assert!(args.contains(&"sonnet".to_string()));
521        assert!(args.contains(&"--system-prompt".to_string()));
522        assert!(args.contains(&"--output-format".to_string()));
523        assert!(args.contains(&"json".to_string()));
524        assert!(args.contains(&"--max-budget-usd".to_string()));
525        assert!(args.contains(&"--permission-mode".to_string()));
526        assert!(args.contains(&"bypassPermissions".to_string()));
527        assert!(args.contains(&"--allowed-tools".to_string()));
528        assert!(args.contains(&"Bash,Read".to_string()));
529        assert!(args.contains(&"--effort".to_string()));
530        assert!(args.contains(&"high".to_string()));
531        assert!(args.contains(&"--max-turns".to_string()));
532        assert!(args.contains(&"--no-session-persistence".to_string()));
533        // Prompt is last
534        assert_eq!(args.last().unwrap(), "explain this");
535    }
536}