Skip to main content

antigravity_sdk_rust/
types.rs

1//! Common configuration structures, enums, and SDK data models.
2//!
3//! This module houses all the data types shared across the SDK, including Gemini configuration
4//! parameters, system instructions, capability filters, built-in tools list, and step execution progress structs.
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::HashMap;
9
10/// The default model name used when none is specified.
11pub const DEFAULT_MODEL: &str = "gemini-3.5-flash";
12
13/// The default image generation model name used.
14pub const DEFAULT_IMAGE_GENERATION_MODEL: &str = "gemini-3.1-flash-image-preview";
15
16/// Configures the intensity of the reasoning/thinking process for models that support it.
17#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
18#[serde(rename_all = "lowercase")]
19pub enum ThinkingLevel {
20    /// Minimal reasoning overhead.
21    Minimal,
22    /// Low reasoning.
23    Low,
24    /// Medium reasoning.
25    Medium,
26    /// High reasoning.
27    High,
28}
29
30/// Generation configuration parameters.
31#[derive(Debug, Clone, Serialize, Deserialize, Default)]
32pub struct GenerationConfig {
33    /// Desired thinking level for reasoning-based models.
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub thinking_level: Option<ThinkingLevel>,
36}
37
38/// Specific model entry defining the model name, key, and generation settings.
39#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct ModelEntry {
41    /// The name/identifier of the model.
42    pub name: String,
43    /// Model-specific API key (if overriding the global key).
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub api_key: Option<String>,
46    /// Generation settings (e.g. thinking configurations).
47    #[serde(default)]
48    pub generation: GenerationConfig,
49}
50
51impl Default for ModelEntry {
52    fn default() -> Self {
53        Self {
54            name: DEFAULT_MODEL.to_string(),
55            api_key: None,
56            generation: GenerationConfig::default(),
57        }
58    }
59}
60
61/// Mapping of models configured for different tasks in the agent's session.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct ModelConfig {
64    /// The primary text/chat model.
65    #[serde(default = "default_model_entry")]
66    pub default: ModelEntry,
67    /// The model used for image generation tasks.
68    #[serde(default = "default_image_generation_entry")]
69    pub image_generation: ModelEntry,
70}
71
72fn default_model_entry() -> ModelEntry {
73    ModelEntry {
74        name: DEFAULT_MODEL.to_string(),
75        api_key: None,
76        generation: GenerationConfig::default(),
77    }
78}
79
80fn default_image_generation_entry() -> ModelEntry {
81    ModelEntry {
82        name: DEFAULT_IMAGE_GENERATION_MODEL.to_string(),
83        api_key: None,
84        generation: GenerationConfig::default(),
85    }
86}
87
88impl Default for ModelConfig {
89    fn default() -> Self {
90        Self {
91            default: default_model_entry(),
92            image_generation: default_image_generation_entry(),
93        }
94    }
95}
96
97/// Root configurations for the Gemini AI model endpoints.
98#[derive(Debug, Clone, Serialize, Deserialize, Default)]
99pub struct GeminiConfig {
100    /// Global API key for Gemini endpoints.
101    #[serde(skip_serializing_if = "Option::is_none")]
102    pub api_key: Option<String>,
103    /// Model configurations.
104    #[serde(default)]
105    pub models: ModelConfig,
106    /// Option to enable Google Search grounding tool.
107    #[serde(skip_serializing_if = "Option::is_none")]
108    pub enable_google_search: Option<bool>,
109    /// Option to enable URL context resolution.
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub enable_url_context: Option<bool>,
112}
113
114/// A structured section appended to system instructions.
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct SystemInstructionSection {
117    /// Main markdown or text body of the section.
118    pub content: String,
119    /// Described title of the section.
120    pub title: String,
121}
122
123/// Directly supplied system instruction instructions text override.
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct CustomSystemInstructions {
126    /// Custom raw instructions text.
127    pub text: String,
128}
129
130/// Appended instructions format, maintaining identity overrides and section segments.
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct AppendedSystemInstructions {
133    /// Optional override for the agent's custom identity block.
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub custom_identity: Option<String>,
136    /// Sections to be appended to the standard system instructions.
137    #[serde(default)]
138    pub appended_sections: Vec<SystemInstructionSection>,
139}
140
141/// Represents the style or content source for system instructions.
142#[derive(Debug, Clone, Serialize, Deserialize)]
143#[serde(untagged)]
144pub enum SystemInstructions {
145    /// Completely custom text override.
146    Custom(CustomSystemInstructions),
147    /// Standard structured segments appended to the system identity.
148    Appended(AppendedSystemInstructions),
149}
150
151/// Enumeration of built-in tools supported by the agent system.
152#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
153pub enum BuiltinTools {
154    /// Tool to create a new file.
155    #[serde(rename = "CREATE_FILE")]
156    CreateFile,
157    /// Tool to edit an existing file.
158    #[serde(rename = "EDIT_FILE")]
159    EditFile,
160    /// Tool to query/find files in a directory.
161    #[serde(rename = "FIND_FILE")]
162    FindFile,
163    /// Tool to list files inside a directory.
164    #[serde(rename = "LIST_DIR")]
165    ListDir,
166    /// Tool to execute a shell command.
167    #[serde(rename = "RUN_COMMAND")]
168    RunCommand,
169    /// Tool to perform ripgrep searches.
170    #[serde(rename = "SEARCH_DIR")]
171    SearchDir,
172    /// Tool to view a file's content.
173    #[serde(rename = "VIEW_FILE")]
174    ViewFile,
175    /// Tool to instantiate a subagent.
176    #[serde(rename = "START_SUBAGENT")]
177    StartSubagent,
178    /// Tool to generate images from descriptions.
179    #[serde(rename = "GENERATE_IMAGE")]
180    GenerateImage,
181    /// Terminating signal indicating the task is completed.
182    #[serde(rename = "FINISH")]
183    Finish,
184}
185
186impl BuiltinTools {
187    /// Returns the static string slice mapping to the tool name.
188    pub const fn as_str(&self) -> &'static str {
189        match self {
190            Self::CreateFile => "CREATE_FILE",
191            Self::EditFile => "EDIT_FILE",
192            Self::FindFile => "FIND_FILE",
193            Self::ListDir => "LIST_DIR",
194            Self::RunCommand => "RUN_COMMAND",
195            Self::SearchDir => "SEARCH_DIR",
196            Self::ViewFile => "VIEW_FILE",
197            Self::StartSubagent => "START_SUBAGENT",
198            Self::GenerateImage => "GENERATE_IMAGE",
199            Self::Finish => "FINISH",
200        }
201    }
202
203    /// Returns a list of all safe, read-only tools.
204    pub fn read_only() -> Vec<Self> {
205        vec![
206            Self::FindFile,
207            Self::ListDir,
208            Self::ViewFile,
209            Self::SearchDir,
210        ]
211    }
212}
213
214/// Agent capabilities and tool restrictions configuration.
215#[derive(Debug, Clone, Serialize, Deserialize, Default)]
216pub struct CapabilitiesConfig {
217    /// List of explicitly enabled tools.
218    #[serde(skip_serializing_if = "Option::is_none")]
219    pub enabled_tools: Option<Vec<BuiltinTools>>,
220    /// List of explicitly disabled tools.
221    #[serde(skip_serializing_if = "Option::is_none")]
222    pub disabled_tools: Option<Vec<BuiltinTools>>,
223    /// Threshold at which the message history is compacted/summarized.
224    #[serde(skip_serializing_if = "Option::is_none")]
225    pub compaction_threshold: Option<u32>,
226    /// Custom schema override for the finish tool schema.
227    #[serde(skip_serializing_if = "Option::is_none")]
228    pub finish_tool_schema_json: Option<String>,
229    /// Model designated for processing images.
230    #[serde(skip_serializing_if = "Option::is_none")]
231    pub image_model: Option<String>,
232}
233
234/// Configuration settings for Model Context Protocol (MCP) servers.
235#[derive(Debug, Clone, Serialize, Deserialize)]
236#[serde(tag = "type")]
237pub enum McpServerConfig {
238    /// Launch the MCP server as a local stdio process.
239    #[serde(rename = "stdio")]
240    Stdio {
241        /// command binary.
242        command: String,
243        /// execution arguments.
244        args: Vec<String>,
245    },
246    /// Connect to the MCP server via Server-Sent Events (SSE).
247    #[serde(rename = "sse")]
248    Sse {
249        /// HTTP URL endpoint.
250        url: String,
251        /// Additional HTTP headers.
252        #[serde(skip_serializing_if = "Option::is_none")]
253        headers: Option<HashMap<String, String>>,
254    },
255    /// Connect to the MCP server via standard HTTP.
256    #[serde(rename = "http")]
257    Http {
258        /// HTTP URL endpoint.
259        url: String,
260        /// Additional HTTP headers.
261        #[serde(skip_serializing_if = "Option::is_none")]
262        headers: Option<HashMap<String, String>>,
263        /// General connection timeout in seconds.
264        #[serde(default = "default_mcp_timeout")]
265        timeout: f64,
266        /// Reading timeout for the SSE listener.
267        #[serde(default = "default_mcp_sse_timeout")]
268        sse_read_timeout: f64,
269        /// Flag whether to terminate the channel connection when closed.
270        #[serde(default = "default_true")]
271        terminate_on_close: bool,
272    },
273}
274
275const fn default_mcp_timeout() -> f64 {
276    30.0
277}
278const fn default_mcp_sse_timeout() -> f64 {
279    300.0
280}
281const fn default_true() -> bool {
282    true
283}
284
285/// Describes a model's request to execute a registered tool.
286#[derive(Debug, Clone, Serialize, Deserialize)]
287pub struct ToolCall {
288    /// Unique call ID generated for correlation.
289    pub id: String,
290    /// Name of the target tool.
291    pub name: String,
292    /// Arguments payload parsed as JSON.
293    pub args: Value,
294    /// Canonical file system path (if the tool targets a file/directory).
295    #[serde(skip_serializing_if = "Option::is_none")]
296    pub canonical_path: Option<String>,
297}
298
299/// The response outcome of executing a client-side tool.
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub struct ToolResult {
302    /// Name of the executed tool.
303    pub name: String,
304    /// Optional matching tool call ID.
305    #[serde(skip_serializing_if = "Option::is_none")]
306    pub id: Option<String>,
307    /// Output result of successful tool execution.
308    #[serde(skip_serializing_if = "Option::is_none")]
309    pub result: Option<Value>,
310    /// Error message string if tool execution failed.
311    #[serde(skip_serializing_if = "Option::is_none")]
312    pub error: Option<String>,
313}
314
315/// Consumption stats for API usage tracking.
316#[derive(Debug, Clone, Serialize, Deserialize, Default)]
317pub struct UsageMetadata {
318    /// Tokens included in the request prompt.
319    pub prompt_token_count: i32,
320    /// Tokens generated in candidates.
321    pub candidates_token_count: i32,
322    /// Total combined tokens.
323    pub total_token_count: i32,
324    /// Cache hit content tokens.
325    #[serde(default)]
326    pub cached_content_token_count: i32,
327    /// Tokens consumed during inner thinking/reasoning.
328    #[serde(default)]
329    pub thoughts_token_count: i32,
330}
331
332/// The classification type of a step in the trajectory.
333#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
334pub enum StepType {
335    /// Raw text content returned by the model.
336    #[serde(rename = "TEXT_RESPONSE")]
337    TextResponse,
338    /// Execution of a tool call.
339    #[serde(rename = "TOOL_CALL")]
340    ToolCall,
341    /// Logging or notification events from the system.
342    #[serde(rename = "SYSTEM_MESSAGE")]
343    SystemMessage,
344    /// A history compaction step summarizing context.
345    #[serde(rename = "COMPACTION")]
346    Compaction,
347    /// Terminating milestone indicator.
348    #[serde(rename = "FINISH")]
349    Finish,
350    /// Catch-all variant for unrecognized steps.
351    #[serde(rename = "UNKNOWN")]
352    Unknown,
353}
354
355/// The originating source component of a trajectory step.
356#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
357pub enum StepSource {
358    /// Internal orchestration environment.
359    #[serde(rename = "SYSTEM")]
360    System,
361    /// End-user input.
362    #[serde(rename = "USER")]
363    User,
364    /// Generative model prediction.
365    #[serde(rename = "MODEL")]
366    Model,
367    /// Unknown author.
368    #[serde(rename = "UNKNOWN")]
369    Unknown,
370}
371
372/// The target destination of a step event.
373#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
374pub enum StepTarget {
375    /// Event addressed to the user.
376    #[serde(rename = "TARGET_USER")]
377    User,
378    /// Event executing in the sandbox environment.
379    #[serde(rename = "TARGET_ENVIRONMENT")]
380    Environment,
381    /// Unspecified destination.
382    #[serde(rename = "TARGET_UNSPECIFIED")]
383    Unspecified,
384    /// Unknown destination.
385    #[serde(rename = "UNKNOWN")]
386    Unknown,
387}
388
389/// Lifecycle execution status of a trajectory step.
390#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
391pub enum StepStatus {
392    /// Active/running state.
393    #[serde(rename = "ACTIVE")]
394    Active,
395    /// Completed successfully.
396    #[serde(rename = "DONE")]
397    Done,
398    /// Waiting for user response/confirmation.
399    #[serde(rename = "WAITING_FOR_USER")]
400    WaitingForUser,
401    /// Finished with a fatal error.
402    #[serde(rename = "ERROR")]
403    Error,
404    /// Execution was canceled.
405    #[serde(rename = "CANCELED")]
406    Canceled,
407    /// Unknown status.
408    #[serde(rename = "UNKNOWN")]
409    Unknown,
410}
411
412/// Individual step event recording an action in the agent's history trajectory.
413#[derive(Debug, Clone, Serialize, Deserialize)]
414pub struct Step {
415    /// Unique identifier for this step.
416    pub id: String,
417    /// Positional index in the trajectory sequence.
418    pub step_index: u32,
419    /// Functional type of the step.
420    pub r#type: StepType,
421    /// Originating author.
422    pub source: StepSource,
423    /// Destination target.
424    pub target: StepTarget,
425    /// Execution status.
426    pub status: StepStatus,
427    /// Main text/markdown content associated with the step.
428    pub content: String,
429    /// Text difference delta relative to previous steps.
430    pub content_delta: String,
431    /// Reasoning thoughts generated for this step.
432    pub thinking: String,
433    /// Thinking reasoning difference delta relative to previous steps.
434    pub thinking_delta: String,
435    /// Custom tool executions registered in this step.
436    pub tool_calls: Vec<ToolCall>,
437    /// Captured execution errors.
438    pub error: String,
439    /// True if this represents the final response segment from the model.
440    pub is_complete_response: Option<bool>,
441    /// Parsed structured JSON output.
442    pub structured_output: Option<Value>,
443    /// Token usage details.
444    pub usage_metadata: Option<UsageMetadata>,
445    /// Unique identifier of the execution cascade grouping subagents.
446    #[serde(default)]
447    pub cascade_id: String,
448    /// Sub-agent trajectory grouping identifier.
449    #[serde(default)]
450    pub trajectory_id: String,
451    /// HTTP status code (if from a network action).
452    #[serde(default)]
453    pub http_code: u32,
454}
455
456impl Default for Step {
457    fn default() -> Self {
458        Self {
459            id: String::new(),
460            step_index: 0,
461            r#type: StepType::Unknown,
462            source: StepSource::Unknown,
463            target: StepTarget::Unknown,
464            status: StepStatus::Unknown,
465            content: String::new(),
466            content_delta: String::new(),
467            thinking: String::new(),
468            thinking_delta: String::new(),
469            tool_calls: Vec::new(),
470            error: String::new(),
471            is_complete_response: None,
472            structured_output: None,
473            usage_metadata: None,
474            cascade_id: String::new(),
475            trajectory_id: String::new(),
476            http_code: 0,
477        }
478    }
479}
480
481/// The result returned by a middleware hook.
482#[derive(Debug, Clone, Serialize, Deserialize, Default)]
483pub struct HookResult {
484    /// True if the operation is allowed to proceed.
485    pub allow: bool,
486    /// Diagnostic or error message details.
487    #[serde(default)]
488    pub message: String,
489}
490
491/// Individual multiple-choice or freeform answer to an interactive user question.
492#[derive(Debug, Clone, Serialize, Deserialize)]
493pub struct QuestionResponse {
494    /// Selected index choices (if multiple-choice).
495    pub selected_option_ids: Option<Vec<String>>,
496    /// Freeform response text.
497    #[serde(default)]
498    pub freeform_response: String,
499    /// True if the question was skipped.
500    #[serde(default)]
501    pub skipped: bool,
502}
503
504/// Complete collection of responses answered to a set of interactive questions.
505#[derive(Debug, Clone, Serialize, Deserialize)]
506pub struct QuestionHookResult {
507    /// List of user responses.
508    pub responses: Vec<QuestionResponse>,
509    /// True if the question panel dialogue was canceled.
510    #[serde(default)]
511    pub cancelled: bool,
512}
513
514/// Single choice option in a multiple-choice question.
515#[derive(Debug, Clone, Serialize, Deserialize)]
516pub struct AskQuestionOption {
517    /// Unique identifier for this choice option.
518    pub id: String,
519    /// Visual label text for the choice.
520    pub text: String,
521}
522
523/// Interactive user question entry structure.
524#[derive(Debug, Clone, Serialize, Deserialize)]
525pub struct AskQuestionEntry {
526    /// Main question prompt text.
527    pub question: String,
528    /// List of multiple-choice options.
529    pub options: Vec<AskQuestionOption>,
530    /// True if multiple selections are supported.
531    #[serde(default)]
532    pub is_multi_select: bool,
533}
534
535/// Final summary outcome of a chat interaction.
536#[derive(Debug, Clone, Serialize, Deserialize)]
537pub struct ChatResponse {
538    /// Combined text output returned.
539    pub text: String,
540    /// Combined reasoning thoughts.
541    pub thinking: String,
542    /// Sequence of intermediate execution steps.
543    pub steps: Vec<Step>,
544    /// Token usage metrics.
545    pub usage_metadata: UsageMetadata,
546}
547
548/// Streaming fragment sent over chunk-based event listeners.
549#[derive(Debug, Clone, Serialize, Deserialize)]
550#[serde(tag = "chunk_type")]
551pub enum StreamChunk {
552    /// Streaming thinking fragment.
553    Thought {
554        /// Step index identifier.
555        step_index: u32,
556        /// Thinking segment.
557        text: String,
558    },
559    /// Streaming text response fragment.
560    Text {
561        /// Step index identifier.
562        step_index: u32,
563        /// Text segment.
564        text: String,
565    },
566    /// Complete parsed tool call definition.
567    ToolCall(ToolCall),
568}
569
570#[cfg(test)]
571mod tests {
572    #![allow(
573        clippy::unwrap_used,
574        clippy::expect_used,
575        clippy::panic,
576        clippy::field_reassign_with_default
577    )]
578    use super::*;
579    use serde_json::json;
580
581    #[test]
582    fn test_tool_call_construction() {
583        let tc = ToolCall {
584            id: "call_1".to_string(),
585            name: "read_file".to_string(),
586            args: json!({"path": "/tmp/foo"}),
587            canonical_path: None,
588        };
589        assert_eq!(tc.name, "read_file");
590        assert_eq!(tc.args["path"], "/tmp/foo");
591        assert_eq!(tc.id, "call_1");
592        assert_eq!(tc.canonical_path, None);
593    }
594
595    #[test]
596    fn test_tool_call_serialization() {
597        let json_data = r#"{"id":"call_1","name":"read_file","args":{"path":"/tmp/foo"}}"#;
598        let tc: ToolCall = serde_json::from_str(json_data).unwrap();
599        assert_eq!(tc.name, "read_file");
600        assert_eq!(tc.args["path"], "/tmp/foo");
601        assert_eq!(tc.id, "call_1");
602        assert_eq!(tc.canonical_path, None);
603    }
604
605    #[test]
606    fn test_tool_result_success() {
607        let tr = ToolResult {
608            name: "sum_tool".to_string(),
609            id: Some("call_1".to_string()),
610            result: Some(json!(42)),
611            error: None,
612        };
613        assert_eq!(tr.name, "sum_tool");
614        assert_eq!(tr.result.unwrap(), 42);
615        assert!(tr.error.is_none());
616        assert_eq!(tr.id.unwrap(), "call_1");
617    }
618
619    #[test]
620    fn test_tool_result_error() {
621        let tr = ToolResult {
622            name: "bad_tool".to_string(),
623            id: None,
624            result: None,
625            error: Some("kaboom".to_string()),
626        };
627        assert_eq!(tr.name, "bad_tool");
628        assert!(tr.result.is_none());
629        assert_eq!(tr.error.unwrap(), "kaboom");
630        assert!(tr.id.is_none());
631    }
632
633    #[test]
634    fn test_tool_result_mutability() {
635        let mut tr = ToolResult {
636            name: "tool".to_string(),
637            id: None,
638            result: None,
639            error: None,
640        };
641        tr.result = Some(json!("updated"));
642        assert_eq!(tr.result.unwrap(), "updated");
643    }
644
645    #[test]
646    fn test_step_defaults() {
647        let step = Step::default();
648        assert_eq!(step.id, "");
649        assert_eq!(step.step_index, 0);
650        assert!(matches!(step.r#type, StepType::Unknown));
651        assert!(matches!(step.status, StepStatus::Unknown));
652        assert!(matches!(step.source, StepSource::Unknown));
653        assert_eq!(step.content, "");
654        assert!(step.tool_calls.is_empty());
655        assert_eq!(step.error, "");
656    }
657
658    #[test]
659    fn test_step_mutability() {
660        let mut step = Step::default();
661        step.content = "goodbye".to_string();
662        assert_eq!(step.content, "goodbye");
663    }
664
665    #[test]
666    fn test_hook_result_defaults() {
667        let hr = HookResult::default();
668        assert!(!hr.allow); // derived default for bool in Rust is false
669        assert_eq!(hr.message, "");
670    }
671
672    #[test]
673    fn test_question_response_defaults() {
674        let qr = QuestionResponse {
675            selected_option_ids: None,
676            freeform_response: String::new(),
677            skipped: false,
678        };
679        assert!(qr.selected_option_ids.is_none());
680        assert_eq!(qr.freeform_response, "");
681        assert!(!qr.skipped);
682    }
683
684    #[test]
685    fn test_question_response_skipped() {
686        let qr = QuestionResponse {
687            selected_option_ids: None,
688            freeform_response: String::new(),
689            skipped: true,
690        };
691        assert!(qr.skipped);
692    }
693
694    #[test]
695    fn test_gemini_config_defaults() {
696        let config = GeminiConfig::default();
697        assert!(config.api_key.is_none());
698        assert_eq!(config.models.default.name, DEFAULT_MODEL);
699        assert!(config.models.default.generation.thinking_level.is_none());
700        assert!(config.enable_google_search.is_none());
701        assert!(config.enable_url_context.is_none());
702    }
703
704    #[test]
705    fn test_thinking_level_serialization() {
706        let level = ThinkingLevel::Low;
707        let json_str = serde_json::to_string(&level).unwrap();
708        assert_eq!(json_str, "\"low\"");
709    }
710
711    #[test]
712    fn test_capabilities_config_defaults() {
713        let config = CapabilitiesConfig::default();
714        assert!(config.enabled_tools.is_none());
715        assert!(config.disabled_tools.is_none());
716        assert!(config.compaction_threshold.is_none());
717    }
718}