ai_agent/
query_engine.rs

1#![allow(dead_code)]
2
3//! QueryEngine - translates QueryEngine.ts from TypeScript
4//! Owns the query lifecycle and session state for a conversation.
5//!
6//! One QueryEngine per conversation. Each submitMessage() call starts a new
7//! turn within the same conversation. State (messages, file cache, usage, etc.)
8//! persists across turns.
9
10use crate::error::AgentError;
11use crate::services::model_cost::TokenUsage;
12use crate::types::*;
13use serde_json::Value;
14use std::collections::{HashMap, HashSet};
15use std::sync::{Arc, Mutex};
16use uuid::Uuid;
17
18/// Async generator result type - yields SDK messages
19pub type SDKMessage = Value;
20
21/// Query engine configuration (matches TypeScript QueryEngineConfig)
22pub struct QueryEngineConfig {
23    pub cwd: String,
24    pub tools: Vec<ToolDefinition>,
25    pub commands: Vec<Value>, // Command[]
26    pub mcp_clients: Vec<crate::mcp::McpConnection>,
27    pub agents: Vec<Value>, // AgentDefinition[]
28    pub initial_messages: Option<Vec<Message>>,
29    pub read_file_cache: Option<FileStateCache>,
30    pub custom_system_prompt: Option<String>,
31    pub append_system_prompt: Option<String>,
32    pub user_specified_model: Option<String>,
33    pub fallback_model: Option<String>,
34    pub thinking_config: Option<ThinkingConfig>,
35    pub max_turns: Option<u32>,
36    pub max_budget_usd: Option<f64>,
37    pub task_budget: Option<TaskBudget>,
38    pub json_schema: Option<Value>,
39    pub verbose: bool,
40    pub replay_user_messages: bool,
41    pub include_partial_messages: bool,
42    pub abort_controller: Option<AbortController>,
43    pub orphaned_permission: Option<OrphanedPermission>,
44}
45
46impl Default for QueryEngineConfig {
47    fn default() -> Self {
48        Self {
49            cwd: String::new(),
50            tools: vec![],
51            commands: vec![],
52            mcp_clients: vec![],
53            agents: vec![],
54            initial_messages: None,
55            read_file_cache: None,
56            custom_system_prompt: None,
57            append_system_prompt: None,
58            user_specified_model: None,
59            fallback_model: None,
60            thinking_config: None,
61            max_turns: None,
62            max_budget_usd: None,
63            task_budget: None,
64            json_schema: None,
65            verbose: false,
66            replay_user_messages: false,
67            include_partial_messages: false,
68            abort_controller: None,
69            orphaned_permission: None,
70        }
71    }
72}
73
74/// Task budget configuration
75#[derive(Debug, Clone)]
76pub struct TaskBudget {
77    pub total: f64,
78}
79
80/// Thinking configuration
81#[derive(Debug, Clone)]
82pub struct ThinkingConfig {
83    pub thinking_type: ThinkingType,
84}
85
86#[derive(Debug, Clone)]
87pub enum ThinkingType {
88    Adaptive,
89    Enabled,
90    Disabled,
91}
92
93impl Default for ThinkingConfig {
94    fn default() -> Self {
95        Self {
96            thinking_type: ThinkingType::Adaptive,
97        }
98    }
99}
100
101/// Permission mode enum
102#[derive(Debug, Clone, PartialEq, Default)]
103pub enum PermissionMode {
104    #[default]
105    Ask,
106    Allow,
107    Deny,
108    Bypass,
109}
110
111/// SDK permission denial
112#[derive(Debug, Clone)]
113pub struct SDKPermissionDenial {
114    pub tool_name: String,
115    pub tool_use_id: String,
116    pub tool_input: Value,
117}
118
119/// SDK status
120#[derive(Debug, Clone)]
121pub struct SDKStatus {
122    pub status: String,
123    pub message: Option<String>,
124}
125
126/// Abort controller
127#[derive(Debug, Clone)]
128pub struct AbortController {
129    aborted: Arc<Mutex<bool>>,
130}
131
132impl AbortController {
133    pub fn new() -> Self {
134        Self {
135            aborted: Arc::new(Mutex::new(false)),
136        }
137    }
138
139    pub fn abort(&self) {
140        *self.aborted.lock().unwrap() = true;
141    }
142
143    pub fn is_aborted(&self) -> bool {
144        *self.aborted.lock().unwrap()
145    }
146}
147
148impl Default for AbortController {
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154/// File state cache
155#[derive(Debug, Clone, Default)]
156pub struct FileStateCache {
157    pub cache: HashMap<String, Value>,
158}
159
160/// Orphaned permission
161#[derive(Debug, Clone)]
162pub struct OrphanedPermission {
163    pub tool_name: String,
164    pub tool_input: Value,
165    pub tool_use_id: String,
166}
167
168/// Elicitation request
169#[derive(Debug, Clone)]
170pub struct ElicitationRequest {
171    pub tool_name: String,
172    pub message: String,
173    pub url: Option<String>,
174}
175
176/// Elicitation response
177#[derive(Debug, Clone)]
178pub struct ElicitationResponse {
179    pub url: Option<String>,
180    pub selection: Option<String>,
181}
182
183/// Snip result from compaction
184#[derive(Debug, Clone)]
185pub struct SnipResult {
186    pub messages: Vec<SDKMessage>,
187    pub executed: bool,
188}
189
190/// Tool use context
191pub struct ToolUseContext {
192    pub cwd: String,
193    pub session_id: String,
194    pub agent_id: Option<String>,
195    pub query_tracking: Option<QueryTracking>,
196    pub options: ToolUseContextOptions,
197    pub abort_controller: AbortController,
198    pub read_file_state: FileStateCache,
199}
200
201pub struct ToolUseContextOptions {
202    pub commands: Vec<Value>,
203    pub debug: bool,
204    pub tools: Vec<ToolDefinition>,
205    pub verbose: bool,
206    pub main_loop_model: Option<String>,
207    pub thinking_config: Option<ThinkingConfig>,
208    pub mcp_clients: Vec<crate::mcp::McpConnection>,
209    pub mcp_resources: HashMap<String, Value>,
210    pub ide_installation_status: Option<Value>,
211    pub is_non_interactive_session: bool,
212    pub custom_system_prompt: Option<String>,
213    pub append_system_prompt: Option<String>,
214    pub agent_definitions: AgentDefinitions,
215    pub theme: Option<String>,
216    pub max_budget_usd: Option<f64>,
217}
218
219impl Default for ToolUseContextOptions {
220    fn default() -> Self {
221        Self {
222            commands: vec![],
223            debug: false,
224            tools: vec![],
225            verbose: false,
226            main_loop_model: None,
227            thinking_config: None,
228            mcp_clients: vec![],
229            mcp_resources: HashMap::new(),
230            ide_installation_status: None,
231            is_non_interactive_session: false,
232            custom_system_prompt: None,
233            append_system_prompt: None,
234            agent_definitions: AgentDefinitions::default(),
235            theme: None,
236            max_budget_usd: None,
237        }
238    }
239}
240
241#[derive(Debug, Clone, Default)]
242pub struct AgentDefinitions {
243    pub active_agents: Vec<Value>,
244    pub all_agents: Vec<Value>,
245    pub allowed_agent_types: Option<Vec<String>>,
246}
247
248/// Query tracking for analytics
249#[derive(Debug, Clone)]
250pub struct QueryTracking {
251    pub chain_id: String,
252    pub depth: u32,
253}
254
255/// CanUseTool function type
256pub type CanUseToolFn = dyn Fn(
257        &ToolDefinition,
258        &Value,
259        &ToolUseContext,
260        &Option<Message>,
261        &str,
262        Option<bool>,
263    )
264        -> std::pin::Pin<Box<dyn std::future::Future<Output = PermissionDecision> + Send + Sync>>
265    + Send
266    + Sync;
267
268/// Handle elicitation function type (boxed)
269pub type HandleElicitationFn = Box<
270    dyn Fn(
271            ElicitationRequest,
272        ) -> std::pin::Pin<
273            Box<dyn std::future::Future<Output = Option<ElicitationResponse>> + Send + Sync>,
274        > + Send
275        + Sync,
276>;
277
278/// Set SDK status function type (boxed)
279pub type SetSdkStatusFn = Box<dyn Fn(SDKStatus) + Send + Sync>;
280
281/// Snip replay function type (boxed)
282pub type SnipReplayFn = Box<dyn Fn(&Message, &[Message]) -> Option<SnipResult> + Send + Sync>;
283
284/// Permission decision
285#[derive(Debug, Clone)]
286pub enum PermissionDecision {
287    Allow,
288    Deny { reason: Option<String> },
289    Ask { expires_at: Option<u64> },
290}
291
292/// App state for SDK
293#[derive(Debug, Clone, Default)]
294pub struct AppState {
295    pub tool_permission_context: ToolPermissionContext,
296    pub fast_mode: bool,
297    pub file_history: Value,
298    pub attribution: Value,
299    pub mcp: McpState,
300    pub effort_value: Option<f64>,
301    pub advisor_model: Option<String>,
302}
303
304#[derive(Debug, Clone, Default)]
305pub struct ToolPermissionContext {
306    pub mode: PermissionMode,
307    pub always_allow_rules: AlwaysAllowRules,
308    pub additional_working_directories: HashMap<String, String>,
309}
310
311#[derive(Debug, Clone, Default)]
312pub struct AlwaysAllowRules {
313    pub command: Vec<String>,
314}
315
316#[derive(Debug, Clone, Default)]
317pub struct McpState {
318    pub tools: Vec<Value>,
319    pub clients: Vec<McpClient>,
320}
321
322#[derive(Debug, Clone)]
323pub struct McpClient {
324    pub name: String,
325    pub client_type: String, // "pending", "connected", etc.
326}
327
328/// Token usage (non-nullable)
329#[derive(Debug, Clone, Default)]
330pub struct NonNullableUsage {
331    pub input_tokens: u64,
332    pub output_tokens: u64,
333    pub cache_creation_input_tokens: Option<u64>,
334    pub cache_read_input_tokens: Option<u64>,
335}
336
337impl From<TokenUsage> for NonNullableUsage {
338    fn from(usage: TokenUsage) -> Self {
339        Self {
340            input_tokens: usage.input_tokens as u64,
341            output_tokens: usage.output_tokens as u64,
342            cache_creation_input_tokens: Some(usage.prompt_cache_write_tokens as u64),
343            cache_read_input_tokens: Some(usage.prompt_cache_read_tokens as u64),
344        }
345    }
346}
347
348/// Options for submit_message
349#[derive(Debug, Default)]
350pub struct SubmitOptions {
351    pub uuid: Option<String>,
352    pub is_meta: Option<bool>,
353}
354
355/// Main ask function - convenience wrapper around QueryEngine
356/// Returns a future that resolves to a vector of SDK messages
357pub async fn ask(config: AskConfig) -> Result<Vec<SDKMessage>, AgentError> {
358    // Convert Vec<Value> to Vec<Message> if needed
359    let initial_messages: Option<Vec<Message>> = config.mutable_messages.map(|msgs| {
360        msgs.into_iter()
361            .map(|v| Message {
362                role: MessageRole::User,
363                content: v.to_string(),
364                attachments: None,
365                tool_call_id: None,
366                tool_calls: None,
367                is_error: None,
368            })
369            .collect()
370    });
371
372    let engine = QueryEngine::new(QueryEngineConfig {
373        cwd: config.cwd,
374        tools: config.tools,
375        commands: vec![],
376        mcp_clients: config.mcp_clients.unwrap_or_default(),
377        agents: config.agents.unwrap_or_default(),
378        initial_messages,
379        read_file_cache: None,
380        custom_system_prompt: config.custom_system_prompt,
381        append_system_prompt: config.append_system_prompt,
382        user_specified_model: config.user_specified_model,
383        fallback_model: config.fallback_model,
384        thinking_config: config.thinking_config,
385        max_turns: config.max_turns,
386        max_budget_usd: config.max_budget_usd,
387        task_budget: config.task_budget,
388        json_schema: config.json_schema,
389        verbose: config.verbose.unwrap_or(false),
390        replay_user_messages: config.replay_user_messages.unwrap_or(false),
391        include_partial_messages: config.include_partial_messages.unwrap_or(false),
392        abort_controller: config.abort_controller,
393        orphaned_permission: config.orphaned_permission,
394    });
395
396    // For now, return empty vec - actual implementation would call submit_message
397    // Convert Vec<Message> to Vec<SDKMessage> (Value)
398    let messages: Vec<SDKMessage> = engine
399        .get_messages()
400        .iter()
401        .map(|m| {
402            serde_json::json!({
403                "role": format!("{:?}", m.role),
404                "content": m.content,
405            })
406        })
407        .collect();
408    Ok(messages)
409}
410
411/// Configuration for ask function
412pub struct AskConfig {
413    pub prompt: String,
414    pub prompt_uuid: Option<String>,
415    pub is_meta: Option<bool>,
416    pub cwd: String,
417    pub tools: Vec<ToolDefinition>,
418    pub mcp_clients: Option<Vec<crate::mcp::McpConnection>>,
419    pub verbose: Option<bool>,
420    pub thinking_config: Option<ThinkingConfig>,
421    pub max_turns: Option<u32>,
422    pub max_budget_usd: Option<f64>,
423    pub task_budget: Option<TaskBudget>,
424    pub mutable_messages: Option<Vec<SDKMessage>>,
425    pub custom_system_prompt: Option<String>,
426    pub append_system_prompt: Option<String>,
427    pub user_specified_model: Option<String>,
428    pub fallback_model: Option<String>,
429    pub json_schema: Option<Value>,
430    pub abort_controller: Option<AbortController>,
431    pub replay_user_messages: Option<bool>,
432    pub include_partial_messages: Option<bool>,
433    pub agents: Option<Vec<Value>>,
434    pub orphaned_permission: Option<OrphanedPermission>,
435}
436
437impl Default for AskConfig {
438    fn default() -> Self {
439        Self {
440            prompt: String::new(),
441            prompt_uuid: None,
442            is_meta: None,
443            cwd: String::new(),
444            tools: vec![],
445            mcp_clients: None,
446            verbose: None,
447            thinking_config: None,
448            max_turns: None,
449            max_budget_usd: None,
450            task_budget: None,
451            mutable_messages: None,
452            custom_system_prompt: None,
453            append_system_prompt: None,
454            user_specified_model: None,
455            fallback_model: None,
456            json_schema: None,
457            abort_controller: None,
458            replay_user_messages: None,
459            include_partial_messages: None,
460            agents: None,
461            orphaned_permission: None,
462        }
463    }
464}
465
466/// QueryEngine - owns the query lifecycle and session state
467pub struct QueryEngine {
468    config: QueryEngineConfig,
469    mutable_messages: Vec<Message>,
470    abort_controller: AbortController,
471    permission_denials: Vec<SDKPermissionDenial>,
472    total_usage: NonNullableUsage,
473    has_handled_orphaned_permission: bool,
474    read_file_state: FileStateCache,
475    discovered_skill_names: HashSet<String>,
476    loaded_nested_memory_paths: HashSet<String>,
477}
478
479impl QueryEngine {
480    pub fn new(config: QueryEngineConfig) -> Self {
481        Self {
482            config,
483            mutable_messages: vec![],
484            abort_controller: AbortController::new(),
485            permission_denials: vec![],
486            total_usage: NonNullableUsage::default(),
487            has_handled_orphaned_permission: false,
488            read_file_state: FileStateCache::default(),
489            discovered_skill_names: HashSet::new(),
490            loaded_nested_memory_paths: HashSet::new(),
491        }
492    }
493
494    pub fn interrupt(&mut self) {
495        self.abort_controller.abort();
496    }
497
498    pub fn get_messages(&self) -> &Vec<Message> {
499        &self.mutable_messages
500    }
501
502    pub fn get_read_file_state(&self) -> &FileStateCache {
503        &self.read_file_state
504    }
505
506    pub fn get_session_id(&self) -> String {
507        Uuid::new_v4().to_string()
508    }
509
510    pub fn set_model(&mut self, model: String) {
511        self.config.user_specified_model = Some(model);
512    }
513}
ai_agent/query_engine.rs

ai_agent/
query_engine.rs
