antigravity_sdk_rust/
local.rs

1//! Local subprocess connection and negotiation.
2//!
3//! This module provides [`LocalConnectionStrategy`] to initialize and spawn the local
4//! agent subprocess, perform the initial handshake, and transition to a WebSocket session
5//! wrapped by [`LocalConnection`].
6
7use crate::connection::Connection;
8use crate::hooks::HookRunner;
9use crate::proto::localharness::{
10    ClientInfo as ProtoClientInfo, FileEditToolConfig, FilesystemWorkspace, FindToolConfig,
11    GeminiConfig as ProtoGeminiConfig, GenerateImageToolConfig, GrepSearchToolConfig,
12    HarnessConfig, HarnessSideTools, InitializeConversationEvent, InputConfig, InputEvent,
13    ListDirToolConfig, MultipleChoiceAnswer, OutputConfig, OutputEvent, RunCommandToolConfig,
14    SubagentsConfig, SystemInstructions as ProtoSystemInstructions, Tool as ProtoTool,
15    ToolConfirmation, ToolResponse, UserQuestionAnswer, UserQuestionsConfig, UserQuestionsResponse,
16    ViewFileToolConfig, Workspace as ProtoWorkspace, WriteToFileToolConfig,
17    appended_system_instructions::Section, custom_system_instructions::Part,
18    user_questions_response::QuestionsResponse, workspace::WorkspaceType,
19};
20use crate::tools::ToolRunner;
21use crate::types::{
22    AntigravityExecutionError, AskQuestionEntry, AskQuestionOption, BuiltinTools,
23    CapabilitiesConfig, GeminiConfig, McpServerConfig, QuestionHookResult, Step, StepSource,
24    StepStatus, StepTarget, StepType, SystemInstructions, ToolCall, ToolResult, UsageMetadata,
25};
26
27use anyhow::anyhow;
28use futures_util::stream::{self, BoxStream};
29use futures_util::{SinkExt, StreamExt};
30use prost::Message;
31use serde_json::Value;
32use std::collections::{HashMap, HashSet};
33use std::process::Stdio;
34use std::sync::Arc;
35use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
36use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
37use tokio::process::Command;
38use tokio::sync::Mutex;
39use tokio::sync::mpsc::{self, UnboundedSender};
40use tokio_tungstenite::connect_async;
41use tokio_tungstenite::tungstenite::Message as WsMessage;
42use tokio_tungstenite::tungstenite::client::IntoClientRequest;
43
44/// Connection strategy implementation communicating with a local subprocess harness.
45///
46/// `LocalConnection` handles launching and lifecycle tracking of the `localharness` binary,
47/// reading standard error output, upgrading communication to standard `WebSockets`, and managing
48/// subagent and tool invocation state.
49#[allow(dead_code)]
50pub struct LocalConnection {
51    conversation_id: String,
52    learned_id: Arc<std::sync::OnceLock<String>>,
53    process: Arc<Mutex<tokio::process::Child>>,
54    child_stdin: Arc<Mutex<Option<tokio::process::ChildStdin>>>,
55    is_idle: Arc<AtomicBool>,
56    step_rx: Arc<Mutex<Option<mpsc::UnboundedReceiver<Result<Step, anyhow::Error>>>>>,
57    ws_tx: UnboundedSender<String>,
58    tool_runner: Option<ToolRunner>,
59    hook_runner: Option<HookRunner>,
60    parent_idle: Arc<Mutex<bool>>,
61    active_subagent_ids: Arc<Mutex<HashSet<String>>>,
62    step_trackers: Arc<Mutex<HashMap<(String, u32), StepTracker>>>,
63}
64
65impl std::fmt::Debug for LocalConnection {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.debug_struct("LocalConnection")
68            .field("conversation_id", &self.conversation_id())
69            .field("is_idle", &self.is_idle)
70            .field("tool_runner", &self.tool_runner)
71            .field("hook_runner", &self.hook_runner)
72            .field("parent_idle", &self.parent_idle)
73            .field("active_subagent_ids", &self.active_subagent_ids)
74            .field("step_trackers", &self.step_trackers)
75            .finish_non_exhaustive()
76    }
77}
78
79impl Connection for LocalConnection {
80    fn conversation_id(&self) -> &str {
81        if self.conversation_id.is_empty() {
82            self.learned_id.get().map_or("", String::as_str)
83        } else {
84            &self.conversation_id
85        }
86    }
87
88    fn is_idle(&self) -> bool {
89        self.is_idle.load(Ordering::SeqCst)
90    }
91
92    fn receive_steps(&self) -> BoxStream<'static, Result<Step, anyhow::Error>> {
93        let step_rx = self.step_rx.clone();
94        let is_idle = self.is_idle.clone();
95        stream::unfold(false, move |mut checked_initial_idle| {
96            let step_rx = step_rx.clone();
97            let is_idle = is_idle.clone();
98            async move {
99                // If the connection is already idle on the first poll and the queue is empty, terminate.
100                if !checked_initial_idle {
101                    checked_initial_idle = true;
102                    let mut guard = step_rx.lock().await;
103                    if guard
104                        .as_mut()
105                        .is_some_and(|rx| rx.is_empty() && is_idle.load(Ordering::SeqCst))
106                    {
107                        return None;
108                    }
109                }
110
111                loop {
112                    let mut guard = step_rx.lock().await;
113                    let Some(rx) = &mut *guard else {
114                        return None;
115                    };
116                    match rx.try_recv() {
117                        Ok(step_res) => match &step_res {
118                            Ok(step) if step.id == "IDLE_SENTINEL" => {
119                                if is_idle.load(Ordering::SeqCst) {
120                                    return None;
121                                }
122                            }
123                            _ => {
124                                return Some((step_res, checked_initial_idle));
125                            }
126                        },
127                        Err(mpsc::error::TryRecvError::Empty) => {
128                            drop(guard);
129                            let mut guard2 = step_rx.lock().await;
130                            let Some(rx2) = &mut *guard2 else {
131                                return None;
132                            };
133                            let step_res = rx2.recv().await;
134                            drop(guard2);
135                            match step_res {
136                                Some(res) => match &res {
137                                    Ok(step) if step.id == "IDLE_SENTINEL" => {
138                                        if is_idle.load(Ordering::SeqCst) {
139                                            return None;
140                                        }
141                                    }
142                                    _ => {
143                                        return Some((res, checked_initial_idle));
144                                    }
145                                },
146                                None => return None,
147                            }
148                        }
149                        Err(mpsc::error::TryRecvError::Disconnected) => {
150                            return None;
151                        }
152                    }
153                }
154            }
155        })
156        .boxed()
157    }
158
159    async fn send(&self, content: &str) -> Result<(), anyhow::Error> {
160        self.is_idle.store(false, Ordering::SeqCst);
161        {
162            let mut p_idle = self.parent_idle.lock().await;
163            *p_idle = false;
164        }
165        {
166            let mut active = self.active_subagent_ids.lock().await;
167            active.clear();
168        }
169        {
170            let mut guard = self.step_rx.lock().await;
171            if let Some(rx) = &mut *guard {
172                while rx.try_recv().is_ok() {}
173            }
174        }
175
176        let input_event = InputEvent {
177            event: Some(crate::proto::localharness::input_event::Event::UserInput(
178                content.to_string(),
179            )),
180        };
181        let raw_json = serde_json::to_string(&input_event)?;
182        self.ws_tx.send(raw_json)?;
183        Ok(())
184    }
185
186    async fn send_trigger_notification(&self, content: &str) -> Result<(), anyhow::Error> {
187        let input_event = InputEvent {
188            event: Some(
189                crate::proto::localharness::input_event::Event::AutomatedTrigger(
190                    content.to_string(),
191                ),
192            ),
193        };
194        let raw_json = serde_json::to_string(&input_event)?;
195        self.ws_tx.send(raw_json)?;
196        Ok(())
197    }
198
199    async fn send_halt_request(&self) -> Result<(), anyhow::Error> {
200        let input_event = InputEvent {
201            event: Some(crate::proto::localharness::input_event::Event::HaltRequest(
202                true,
203            )),
204        };
205        let raw_json = serde_json::to_string(&input_event)?;
206        self.ws_tx.send(raw_json)?;
207        Ok(())
208    }
209
210    async fn send_tool_confirmation(
211        &self,
212        trajectory_id: &str,
213        step_index: u32,
214        accepted: bool,
215    ) -> Result<(), anyhow::Error> {
216        let conf = ToolConfirmation {
217            trajectory_id: Some(trajectory_id.to_string()),
218            step_index: Some(step_index),
219            accepted: Some(accepted),
220        };
221        let input_event = InputEvent {
222            event: Some(crate::proto::localharness::input_event::Event::ToolConfirmation(conf)),
223        };
224        let raw_json = serde_json::to_string(&input_event)?;
225        self.ws_tx.send(raw_json)?;
226        Ok(())
227    }
228
229    async fn send_tool_response(&self, id: &str, result: ToolResult) -> Result<(), anyhow::Error> {
230        let resp_json = if let Some(ref val) = result.result {
231            // The Go harness expects responseJson to always be a JSON object.
232            // If the tool returned a non-object value (string, number, array, etc.),
233            // wrap it under a "result" key to match the Python SDK's behavior.
234            if val.is_object() {
235                serde_json::to_string(val)?
236            } else {
237                serde_json::to_string(&serde_json::json!({ "result": val }))?
238            }
239        } else if let Some(ref err) = result.error {
240            serde_json::to_string(&serde_json::json!({ "error": err }))?
241        } else {
242            "{}".to_string()
243        };
244
245        let resp = ToolResponse {
246            id: Some(id.to_string()),
247            response_json: Some(resp_json),
248            supplemental_media: Vec::new(),
249            response: None,
250        };
251        let input_event = InputEvent {
252            event: Some(crate::proto::localharness::input_event::Event::ToolResponse(resp)),
253        };
254        let raw_json = serde_json::to_string(&input_event)?;
255        self.ws_tx.send(raw_json)?;
256        Ok(())
257    }
258
259    async fn send_question_response(
260        &self,
261        trajectory_id: &str,
262        step_index: u32,
263        answers: QuestionHookResult,
264    ) -> Result<(), anyhow::Error> {
265        let mut proto_answers = Vec::new();
266        for r in answers.responses {
267            if r.skipped {
268                proto_answers.push(UserQuestionAnswer {
269                    answer: Some(
270                        crate::proto::localharness::user_question_answer::Answer::Unanswered(true),
271                    ),
272                });
273            } else {
274                let mut mc_ans = MultipleChoiceAnswer {
275                    selected_choice_indices: Vec::new(),
276                    freeform_response: Some(r.freeform_response.clone()),
277                };
278                if let Some(ref opts) = r.selected_option_ids {
279                    for opt in opts {
280                        if let Ok(idx) = opt.parse::<i32>() {
281                            mc_ans.selected_choice_indices.push(idx - 1);
282                        }
283                    }
284                }
285                proto_answers.push(UserQuestionAnswer {
286                    answer: Some(crate::proto::localharness::user_question_answer::Answer::MultipleChoiceAnswer(mc_ans)),
287                });
288            }
289        }
290
291        let resp = UserQuestionsResponse {
292            trajectory_id: Some(trajectory_id.to_string()),
293            step_index: Some(step_index),
294            result: Some(if answers.cancelled {
295                crate::proto::localharness::user_questions_response::Result::Cancelled(true)
296            } else {
297                crate::proto::localharness::user_questions_response::Result::Response(
298                    QuestionsResponse {
299                        answers: proto_answers,
300                    },
301                )
302            }),
303        };
304
305        let input_event = InputEvent {
306            event: Some(crate::proto::localharness::input_event::Event::QuestionResponse(resp)),
307        };
308        let raw_json = serde_json::to_string(&input_event)?;
309        self.ws_tx.send(raw_json)?;
310        Ok(())
311    }
312
313    async fn disconnect(&self) -> Result<(), anyhow::Error> {
314        let mut proc = self.process.lock().await;
315        let _ = proc.kill().await;
316        drop(proc);
317        Ok(())
318    }
319}
320
321/// Configurator and builder to spawn a local helper subprocess and build a connection.
322#[derive(Debug)]
323pub struct LocalConnectionStrategy {
324    /// Path to the `localharness` binary.
325    pub binary_path: String,
326    /// Gemini configuration parameters.
327    pub gemini_config: GeminiConfig,
328    /// Capability config specifying enabled tools.
329    pub capabilities_config: CapabilitiesConfig,
330    /// System instruction content or templates.
331    pub system_instructions: Option<SystemInstructions>,
332    /// Optional directory to store state and execution logs.
333    pub save_dir: Option<String>,
334    /// Workspace paths.
335    pub workspaces: Vec<String>,
336    /// Folders containing custom skill modules.
337    pub skills_paths: Vec<String>,
338    /// Optional coordinator runner to handle custom tools.
339    pub tool_runner: Option<ToolRunner>,
340    /// Optional coordinator runner to handle lifecycle hooks.
341    pub hook_runner: Option<HookRunner>,
342    /// Conversation ID for standard session resuming or tracking.
343    pub conversation_id: String,
344    /// MCP server configurations.
345    pub mcp_servers: Vec<McpServerConfig>,
346}
347
348impl LocalConnectionStrategy {
349    /// Creates a new `LocalConnectionStrategy`.
350    #[allow(clippy::too_many_arguments)]
351    pub const fn new(
352        binary_path: String,
353        gemini_config: GeminiConfig,
354        capabilities_config: CapabilitiesConfig,
355        system_instructions: Option<SystemInstructions>,
356        save_dir: Option<String>,
357        workspaces: Vec<String>,
358        skills_paths: Vec<String>,
359        tool_runner: Option<ToolRunner>,
360        hook_runner: Option<HookRunner>,
361        conversation_id: String,
362        mcp_servers: Vec<McpServerConfig>,
363    ) -> Self {
364        Self {
365            binary_path,
366            gemini_config,
367            capabilities_config,
368            system_instructions,
369            save_dir,
370            workspaces,
371            skills_paths,
372            tool_runner,
373            hook_runner,
374            conversation_id,
375            mcp_servers,
376        }
377    }
378
379    /// Spawns the subprocess helper binary, performs length-prefixed initialization,
380    /// upgrades connection to WebSocket, and builds the stateful [`LocalConnection`].
381    ///
382    /// # Errors
383    ///
384    /// Returns an error if the subprocess cannot be launched, the handshake fails,
385    /// or the WebSocket upgrade fails.
386    #[allow(clippy::too_many_lines)]
387    pub async fn connect(&self) -> Result<LocalConnection, anyhow::Error> {
388        let use_vertex = self.gemini_config.vertex;
389        let api_key = self
390            .gemini_config
391            .models
392            .default
393            .api_key
394            .clone()
395            .or_else(|| self.gemini_config.api_key.clone())
396            .or_else(|| std::env::var("GEMINI_API_KEY").ok());
397
398        if !use_vertex && api_key.is_none() {
399            return Err(anyhow!(
400                "A Gemini API key is required. Set it via GeminiConfig or GEMINI_API_KEY env var."
401            ));
402        }
403
404        if use_vertex {
405            let has_project = self.gemini_config.project.is_some();
406            let has_location = self.gemini_config.location.is_some();
407            if api_key.is_none() && !(has_project && has_location) {
408                return Err(anyhow!(
409                    "For Vertex AI, either a GCP project and location, or an API key \
410                     (Express Mode) must be set."
411                ));
412            }
413        }
414
415        let api_key = api_key.unwrap_or_default();
416
417        // 1. Spawning localharness subprocess
418        // Explicitly forward SHELL and PATH so the harness can fork /bin/sh for
419        // RUN_COMMAND steps regardless of what working directory is used.
420        // Without this, the harness may fail with "fork/exec /bin/sh: no such file
421        // or directory" when running commands in absolute-path subdirectories.
422        let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/sh".to_string());
423        let path = std::env::var("PATH")
424            .unwrap_or_else(|_| "/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin".to_string());
425        let mut child = Command::new(&self.binary_path)
426            .stdin(Stdio::piped())
427            .stdout(Stdio::piped())
428            .stderr(Stdio::piped())
429            .env("SHELL", &shell)
430            .env("PATH", &path)
431            .spawn()?;
432
433        let mut child_stdin = child
434            .stdin
435            .take()
436            .ok_or_else(|| anyhow!("Failed to open child stdin"))?;
437        let mut child_stdout = child
438            .stdout
439            .take()
440            .ok_or_else(|| anyhow!("Failed to open child stdout"))?;
441        let child_stderr = child
442            .stderr
443            .take()
444            .ok_or_else(|| anyhow!("Failed to open child stderr"))?;
445
446        // 2. Perform Handshake via length-prefixed protocol buffer over stdin/stdout
447        let client_info = ProtoClientInfo {
448            language: Some("rust".to_string()),
449            version: Some(env!("CARGO_PKG_VERSION").to_string()),
450            language_version: Some(rustc_version()),
451        };
452
453        let input_config = InputConfig {
454            storage_directory: self.save_dir.clone(),
455            port: None,
456            bind_address: None,
457            client_info: Some(client_info),
458        };
459
460        let mut input_buf = Vec::new();
461        input_config.encode(&mut input_buf)?;
462
463        let size = input_buf.len() as u32;
464        child_stdin.write_all(&size.to_le_bytes()).await?;
465        child_stdin.write_all(&input_buf).await?;
466        child_stdin.flush().await?;
467
468        let mut size_bytes = [0u8; 4];
469        child_stdout.read_exact(&mut size_bytes).await?;
470        let length = u32::from_le_bytes(size_bytes) as usize;
471        let mut output_buf = vec![0u8; length];
472        child_stdout.read_exact(&mut output_buf).await?;
473        let output_config = OutputConfig::decode(&output_buf[..])?;
474
475        let port = output_config
476            .port
477            .ok_or_else(|| anyhow!("Harness OutputConfig missing port"))?;
478        let harness_api_key = output_config
479            .api_key
480            .ok_or_else(|| anyhow!("Harness OutputConfig missing api_key"))?;
481
482        // 3. Setup WebSocket connection
483        let ws_url = format!("ws://localhost:{port}/");
484        let mut req = ws_url.clone().into_client_request()?;
485        req.headers_mut()
486            .insert("x-goog-api-key", harness_api_key.parse()?);
487
488        // Connect with retry/backoff
489        let mut ws_stream = None;
490        let mut delay = std::time::Duration::from_millis(100);
491        for attempt in 0..5 {
492            match connect_async(req.clone()).await {
493                Ok((stream, _)) => {
494                    ws_stream = Some(stream);
495                    break;
496                }
497                Err(e) => {
498                    if attempt == 4 {
499                        let _ = child.kill().await;
500                        return Err(anyhow!("Failed to connect to WS at {ws_url}: {e:?}"));
501                    }
502                    tokio::time::sleep(delay).await;
503                    delay *= 2;
504                }
505            }
506        }
507
508        let ws = ws_stream.ok_or_else(|| anyhow!("Failed to connect to WS"))?;
509        let (mut ws_write, mut ws_read) = ws.split();
510
511        // 4. Build HarnessConfig proto
512        let mut proto_tools = Vec::new();
513        if let Some(ref runner) = self.tool_runner {
514            let tools = runner.tools.read().await;
515            for t in tools.values() {
516                proto_tools.push(ProtoTool {
517                    name: Some(t.name().to_string()),
518                    description: Some(t.description().to_string()),
519                    parameters_json_schema: Some(t.parameters_json_schema().to_string()),
520                    response_json_schema: None,
521                });
522            }
523        }
524
525        let proto_sys = self.system_instructions.as_ref().map(|sys| match sys {
526            SystemInstructions::Custom(custom) => {
527                let instr_part = Part {
528                    part: Some(
529                        crate::proto::localharness::custom_system_instructions::part::Part::Text(
530                            custom.text.clone(),
531                        ),
532                    ),
533                };
534                ProtoSystemInstructions {
535                    r#type: Some(
536                        crate::proto::localharness::system_instructions::Type::Custom(
537                            crate::proto::localharness::CustomSystemInstructions {
538                                part: vec![instr_part],
539                            },
540                        ),
541                    ),
542                }
543            }
544            SystemInstructions::Appended(appended) => {
545                let mut sections = Vec::new();
546                for sec in &appended.appended_sections {
547                    sections.push(Section {
548                        title: Some(sec.title.clone()),
549                        content: Some(sec.content.clone()),
550                    });
551                }
552                ProtoSystemInstructions {
553                    r#type: Some(
554                        crate::proto::localharness::system_instructions::Type::Appended(
555                            crate::proto::localharness::AppendedSystemInstructions {
556                                custom_identity: appended.custom_identity.clone(),
557                                appended_sections: sections,
558                            },
559                        ),
560                    ),
561                }
562            }
563        });
564
565        let proto_gemini = ProtoGeminiConfig {
566            api_key: Some(api_key),
567            base_url: None,
568            model_name: Some(self.gemini_config.models.default.name.clone()),
569            thinking_level: self
570                .gemini_config
571                .models
572                .default
573                .generation
574                .thinking_level
575                .map(|l| match l {
576                    crate::types::ThinkingLevel::Minimal => "minimal".to_string(),
577                    crate::types::ThinkingLevel::Low => "low".to_string(),
578                    crate::types::ThinkingLevel::Medium => "medium".to_string(),
579                    crate::types::ThinkingLevel::High => "high".to_string(),
580                }),
581            enable_url_context: self.gemini_config.enable_url_context,
582            enable_google_search: self.gemini_config.enable_google_search,
583            use_vertex: Some(self.gemini_config.vertex),
584            project: self.gemini_config.project.clone(),
585            location: self.gemini_config.location.clone(),
586        };
587
588        let mut proto_workspaces = Vec::new();
589        for w in &self.workspaces {
590            proto_workspaces.push(ProtoWorkspace {
591                workspace_type: Some(WorkspaceType::FilesystemWorkspace(FilesystemWorkspace {
592                    directory: Some(w.clone()),
593                })),
594            });
595        }
596
597        let all_tools = vec![
598            BuiltinTools::CreateFile,
599            BuiltinTools::EditFile,
600            BuiltinTools::FindFile,
601            BuiltinTools::ListDir,
602            BuiltinTools::RunCommand,
603            BuiltinTools::SearchDir,
604            BuiltinTools::ViewFile,
605            BuiltinTools::StartSubagent,
606            BuiltinTools::GenerateImage,
607            BuiltinTools::Finish,
608        ];
609        let active_tools: HashSet<BuiltinTools> =
610            self.capabilities_config.enabled_tools.as_ref().map_or_else(
611                || {
612                    if let Some(ref disabled) = self.capabilities_config.disabled_tools {
613                        let disabled_set: HashSet<BuiltinTools> =
614                            disabled.iter().copied().collect();
615                        all_tools
616                            .into_iter()
617                            .filter(|t| !disabled_set.contains(t))
618                            .collect()
619                    } else {
620                        all_tools.into_iter().collect()
621                    }
622                },
623                |enabled| enabled.iter().copied().collect(),
624            );
625
626        let side_tools = HarnessSideTools {
627            find: Some(FindToolConfig {
628                enabled: Some(active_tools.contains(&BuiltinTools::FindFile)),
629            }),
630            run_command: Some(RunCommandToolConfig {
631                enabled: Some(active_tools.contains(&BuiltinTools::RunCommand)),
632            }),
633            subagents: Some(SubagentsConfig {
634                enabled: Some(active_tools.contains(&BuiltinTools::StartSubagent)),
635            }),
636            user_questions: Some(UserQuestionsConfig {
637                enabled: Some(true),
638            }),
639            file_edit: Some(FileEditToolConfig {
640                enabled: Some(active_tools.contains(&BuiltinTools::EditFile)),
641            }),
642            view_file: Some(ViewFileToolConfig {
643                enabled: Some(active_tools.contains(&BuiltinTools::ViewFile)),
644            }),
645            write_to_file: Some(WriteToFileToolConfig {
646                enabled: Some(active_tools.contains(&BuiltinTools::CreateFile)),
647            }),
648            grep_search: Some(GrepSearchToolConfig {
649                enabled: Some(active_tools.contains(&BuiltinTools::SearchDir)),
650            }),
651            list_dir: Some(ListDirToolConfig {
652                enabled: Some(active_tools.contains(&BuiltinTools::ListDir)),
653            }),
654            permissions: None,
655            generate_image: Some(GenerateImageToolConfig {
656                enabled: Some(active_tools.contains(&BuiltinTools::GenerateImage)),
657                model_name: self.capabilities_config.image_model.clone(),
658            }),
659        };
660
661        let harness_config = HarnessConfig {
662            cascade_id: Some(self.conversation_id.clone()),
663            model_config: Some(
664                crate::proto::localharness::harness_config::ModelConfig::GeminiConfig(proto_gemini),
665            ),
666            system_instructions: proto_sys,
667            tools: proto_tools,
668            harness_side_tools: Some(side_tools),
669            compaction_threshold: self.capabilities_config.compaction_threshold,
670            workspaces: proto_workspaces,
671            skills_paths: self.skills_paths.clone(),
672            finish_tool_schema_json: self.capabilities_config.finish_tool_schema_json.clone(),
673            initial_trajectory: None,
674            app_data_dir: self.save_dir.clone(),
675        };
676
677        // 5. Send InitializeConversationEvent
678        let init_event = InitializeConversationEvent {
679            config: Some(harness_config),
680        };
681        let init_json = serde_json::to_string(&init_event)?;
682        ws_write.send(WsMessage::Text(init_json)).await?;
683
684        // 6. Spawn Background WS Sender Loop
685        let (ws_tx, mut ws_rx) = mpsc::unbounded_channel::<String>();
686        tokio::spawn(async move {
687            while let Some(msg) = ws_rx.recv().await {
688                if let Err(e) = ws_write.send(WsMessage::Text(msg)).await {
689                    tracing::error!("WS Write Loop Error: {:?}", e);
690                    break;
691                }
692            }
693        });
694
695        // 7. Setup channels for step stream
696        let (step_tx, step_rx) = mpsc::unbounded_channel::<Result<Step, anyhow::Error>>();
697        let client_tool_step_counter = Arc::new(AtomicU32::new(50_000));
698
699        let is_idle = Arc::new(AtomicBool::new(false));
700        let parent_idle = Arc::new(Mutex::new(false));
701        let active_subagent_ids = Arc::new(Mutex::new(HashSet::new()));
702        let step_trackers = Arc::new(Mutex::new(HashMap::new()));
703
704        let conn_ws_tx = ws_tx.clone();
705        let conn_is_idle = is_idle.clone();
706        let conn_parent_idle = parent_idle.clone();
707        let conn_active_subagents = active_subagent_ids.clone();
708        let conn_step_trackers = step_trackers.clone();
709
710        let tool_runner = self.tool_runner.clone();
711        let hook_runner = self.hook_runner.clone();
712        let conversation_id = self.conversation_id.clone();
713        let learned_id = Arc::new(std::sync::OnceLock::new());
714        let conn_learned_id = learned_id.clone();
715        let conn_cascade_id: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
716        let conn_cascade_id_for_ws = conn_cascade_id.clone();
717
718        // 8. Spawn Stderr Reader
719        let mut reader = tokio::io::BufReader::new(child_stderr);
720        tokio::spawn(async move {
721            let mut line = String::new();
722            while let Ok(n) = reader.read_line(&mut line).await {
723                if n == 0 {
724                    break;
725                }
726                tracing::info!("Harness stderr: {}", line.trim_end());
727                line.clear();
728            }
729        });
730
731        // 9. Spawn WS Reader Loop
732        let pending_builtin_tool_calls =
733            Arc::new(Mutex::new(HashMap::<(String, u32), ToolCall>::new()));
734        tokio::spawn(async move {
735            while let Some(msg_res) = ws_read.next().await {
736                match msg_res {
737                    Ok(WsMessage::Text(raw_text)) => {
738                        tracing::debug!("WS raw message: {}", &raw_text[..raw_text.len().min(500)]);
739                        match serde_json::from_str::<OutputEvent>(&raw_text) {
740                            Ok(output_event) => {
741                                if let Some(event) = output_event.event {
742                                    match event {
743                                        crate::proto::localharness::output_event::Event::StepUpdate(step_update) => {
744                                            let traj_id = step_update.trajectory_id.clone().unwrap_or_default();
745                                            let step_idx = step_update.step_index.unwrap_or(0);
746                                            let key = (traj_id.clone(), step_idx);
747
748                                            // Learn the cascade_id from the first StepUpdate
749                                            // where cascade_id == trajectory_id (Python parity)
750                                            {
751                                                let cascade_id_val = step_update.cascade_id.clone().unwrap_or_default();
752                                                if !cascade_id_val.is_empty() && cascade_id_val == traj_id {
753                                                    let _ = conn_learned_id.set(cascade_id_val.clone());
754                                                    let mut cid = conn_cascade_id_for_ws.lock().await;
755                                                    if cid.is_none() {
756                                                        tracing::debug!("Learned cascade_id from StepUpdate: {}", cascade_id_val);
757                                                        *cid = Some(cascade_id_val);
758                                                    }
759                                                }
760                                            }
761
762                                            let (is_questions_new, is_tool_conf_new) = {
763                                                let mut trackers = conn_step_trackers.lock().await;
764                                                let tracker = trackers.entry(key.clone()).or_insert_with(|| StepTracker {
765                                                    state: 0,
766                                                    handled_requests: HashSet::new(),
767                                                });
768                                                if let Some(st) = step_update.state {
769                                                    tracker.update_state(st);
770                                                }
771                                                let is_q = if step_update.questions_request.is_some() {
772                                                    tracker.mark_handled("questions_request")
773                                                } else {
774                                                    false
775                                                };
776                                                let is_tc = if step_update.tool_confirmation_request.is_some() {
777                                                    tracker.mark_handled("tool_confirmation_request")
778                                                } else {
779                                                    false
780                                                };
781                                                drop(trackers);
782                                                (is_q, is_tc)
783                                            };
784
785                                            // Map StepUpdate proto to Step domain model
786                                            let step_type = if step_update.compaction.is_some() {
787                                                StepType::Compaction
788                                            } else if step_update.finish.is_some() {
789                                                StepType::Finish
790                                            } else if step_update.list_directory.is_some()
791                                                || step_update.find_file.is_some()
792                                                || step_update.search_directory.is_some()
793                                                || step_update.view_file.is_some()
794                                                || step_update.create_file.is_some()
795                                                || step_update.edit_file.is_some()
796                                                || step_update.run_command.is_some()
797                                                || step_update.invoke_subagent.is_some()
798                                                || step_update.generate_image.is_some()
799                                            {
800                                                StepType::ToolCall
801                                            } else if step_update.text.is_some() {
802                                                StepType::TextResponse
803                                            } else {
804                                                StepType::Unknown
805                                            };
806
807                                            let mut tool_calls = Vec::new();
808                                            if let Some(tc) = extract_builtin_tool_call(&step_update) {
809                                                tool_calls.push(tc);
810                                            }
811
812                                            let source = match step_update.source {
813                                                Some(1) => StepSource::System,
814                                                Some(2) => StepSource::User,
815                                                Some(3) => StepSource::Model,
816                                                _ => StepSource::Unknown,
817                                            };
818
819                                            let status = match step_update.state {
820                                                Some(1) => StepStatus::Active,
821                                                Some(2) => StepStatus::Done,
822                                                Some(3) => StepStatus::WaitingForUser,
823                                                Some(4) => StepStatus::Error,
824                                                Some(5) => StepStatus::TerminalError,
825                                                _ => StepStatus::Unknown,
826                                            };
827
828                                            let target = match step_update.target {
829                                                Some(1) => StepTarget::User,
830                                                Some(2 | 3) => StepTarget::Environment,
831                                                _ => StepTarget::Unknown,
832                                            };
833
834                                            let usage = output_event.usage_metadata.map(|u| UsageMetadata {
835                                                prompt_token_count: u.prompt_token_count.unwrap_or(0),
836                                                candidates_token_count: u.candidates_token_count.unwrap_or(0),
837                                                total_token_count: u.total_token_count.unwrap_or(0),
838                                                cached_content_token_count: u.cached_content_token_count.unwrap_or(0),
839                                                thoughts_token_count: u.thoughts_token_count.unwrap_or(0),
840                                            });
841
842                                            let is_complete = Some(
843                                                source == StepSource::Model
844                                                    && status == StepStatus::Done
845                                                    && step_update.text.is_some()
846                                                    && target == StepTarget::User
847                                            );
848
849                                            let structured = step_update.finish.as_ref().and_then(|f| {
850                                                f.output_string.as_ref().and_then(|s| serde_json::from_str(s).ok())
851                                            });
852
853                                            let error_msg = step_update.error_message.clone().unwrap_or_default();
854                                            let http_code = step_update.error.as_ref().and_then(|e| e.http_code).unwrap_or(0);
855
856                                            let step = Step {
857                                                id: format!("{traj_id}_{step_idx}"),
858                                                step_index: step_idx,
859                                                r#type: step_type,
860                                                source,
861                                                target,
862                                                status,
863                                                content: step_update.text.clone().unwrap_or_default(),
864                                                content_delta: step_update.text_delta.clone().unwrap_or_default(),
865                                                thinking: step_update.thinking.clone().unwrap_or_default(),
866                                                thinking_delta: step_update.thinking_delta.clone().unwrap_or_default(),
867                                                tool_calls,
868                                                error: error_msg,
869                                                is_complete_response: is_complete,
870                                                structured_output: structured,
871                                                usage_metadata: usage,
872                                                cascade_id: step_update.cascade_id.clone().unwrap_or_default(),
873                                                trajectory_id: traj_id.clone(),
874                                                http_code,
875                                            };
876
877                                            let _ = step_tx.send(Ok(step));
878
879                                            // Detect platform-level errors (source=SYSTEM) and propagate them.
880                                            if source == StepSource::System
881                                                && status == StepStatus::Error
882                                                && (http_code == 400 || http_code == 401 || http_code == 403)
883                                            {
884                                                let err_str = step_update.error.as_ref().and_then(|e| e.error_message.clone()).unwrap_or_else(|| "System error occurred.".to_string());
885                                                let _ = step_tx.send(Err(anyhow!("System step error (HTTP {}): {}", http_code, err_str)));
886                                                break;
887                                             }
888
889                                            // Handle terminal errors — non-recoverable agent execution failure.
890                                            if status == StepStatus::TerminalError {
891                                                let err_msg = step_update.error_message.clone()
892                                                    .unwrap_or_else(|| "Terminal error occurred during execution".to_string());
893                                                let _ = step_tx.send(Err(
894                                                    AntigravityExecutionError { message: err_msg }.into()
895                                                ));
896                                                break;
897                                            }
898
899                                            // Dispatch post-tool-call or on-tool-error hooks for built-in tools
900                                            let state_val = step_update.state.unwrap_or(0);
901                                            if state_val == 2 || state_val == 4 || state_val == 5 {
902                                                let mut pending = pending_builtin_tool_calls.lock().await;
903                                                if let (Some(tc), Some(runner)) = (pending.remove(&key), hook_runner.as_ref()) {
904                                                    if state_val == 2 {
905                                                        let extracted = extract_tool_result(&step_update);
906                                                        let tr = ToolResult {
907                                                            name: tc.name.clone(),
908                                                            id: Some(tc.id.clone()),
909                                                            result: extracted.and_then(|r| r.result).or_else(|| step_update.text.clone().map(Value::String)),
910                                                            error: None,
911                                                        };
912                                                        let runner_clone = runner.clone();
913                                                        tokio::spawn(async move {
914                                                            let _ = runner_clone.dispatch_post_tool_call(&tr).await;
915                                                        });
916                                                    } else {
917                                                        let err_msg = step_update.error_message.clone().unwrap_or_else(|| "Built-in tool failed".to_string());
918                                                        let err = anyhow!(err_msg);
919                                                        let runner_clone = runner.clone();
920                                                        tokio::spawn(async move {
921                                                            let _ = runner_clone.dispatch_on_tool_error(&err).await;
922                                                        });
923                                                    }
924                                                }
925                                            }
926
927                                            if let (true, Some(q_req)) = (is_questions_new, &step_update.questions_request) {
928                                                let conn_ws_tx = conn_ws_tx.clone();
929                                                let hook_runner = hook_runner.clone();
930                                                let q_req_clone = q_req.clone();
931                                                let trajectory_id = step_update.trajectory_id.clone();
932                                                let step_index = step_update.step_index;
933                                                tokio::spawn(async move {
934                                                    let mut questions_list = Vec::new();
935                                                    for uq in &q_req_clone.questions {
936                                                        if let Some(crate::proto::localharness::user_question::QuestionType::MultipleChoice(ref mc)) = uq.question_type {
937                                                            let mut opts = Vec::new();
938                                                            for (j, choice) in mc.choices.iter().enumerate() {
939                                                                opts.push(AskQuestionOption {
940                                                                    id: (j + 1).to_string(),
941                                                                    text: choice.clone(),
942                                                                });
943                                                            }
944                                                            questions_list.push(AskQuestionEntry {
945                                                                question: mc.question.clone().unwrap_or_default(),
946                                                                options: opts,
947                                                                is_multi_select: mc.is_multi_select.unwrap_or(false),
948                                                            });
949                                                        }
950                                                    }
951
952                                                    let mut proto_answers = vec![
953                                                        UserQuestionAnswer {
954                                                            answer: Some(crate::proto::localharness::user_question_answer::Answer::Unanswered(true)),
955                                                        };
956                                                        q_req_clone.questions.len()
957                                                    ];
958
959                                                    if let Some(runner) = hook_runner.as_ref().filter(|_| !questions_list.is_empty()) {
960                                                        let res = runner.dispatch_interaction(&questions_list).await;
961                                                        if let Ok(Some(q_res)) = res {
962                                                            for (orig_idx, r) in q_res.responses.iter().enumerate() {
963                                                                if !r.skipped {
964                                                                    let mut mc_ans = MultipleChoiceAnswer {
965                                                                        selected_choice_indices: Vec::new(),
966                                                                        freeform_response: Some(r.freeform_response.clone()),
967                                                                    };
968                                                                    if let Some(ref opts) = r.selected_option_ids {
969                                                                        for opt in opts {
970                                                                            if let Ok(idx) = opt.parse::<i32>() {
971                                                                                mc_ans.selected_choice_indices.push(idx - 1);
972                                                                            }
973                                                                        }
974                                                                    }
975                                                                    proto_answers[orig_idx] = UserQuestionAnswer {
976                                                                        answer: Some(crate::proto::localharness::user_question_answer::Answer::MultipleChoiceAnswer(mc_ans)),
977                                                                    };
978                                                                }
979                                                            }
980                                                        }
981                                                    }
982
983                                                    let resp = UserQuestionsResponse {
984                                                        trajectory_id,
985                                                        step_index,
986                                                        result: Some(crate::proto::localharness::user_questions_response::Result::Response(QuestionsResponse {
987                                                            answers: proto_answers,
988                                                        })),
989                                                    };
990                                                    let input_event = InputEvent {
991                                                        event: Some(crate::proto::localharness::input_event::Event::QuestionResponse(resp)),
992                                                    };
993                                                    if let Ok(raw_json) = serde_json::to_string(&input_event) {
994                                                        let _ = conn_ws_tx.send(raw_json);
995                                                    }
996                                                });
997                                            }
998
999                                            if is_tool_conf_new {
1000                                                let conn_ws_tx = conn_ws_tx.clone();
1001                                                let hook_runner = hook_runner.clone();
1002                                                let step_update_clone = step_update.clone();
1003                                                let pending_calls = pending_builtin_tool_calls.clone();
1004                                                tokio::spawn(async move {
1005                                                    let mut allow = true;
1006                                                    let tool_call = extract_builtin_tool_call(&step_update_clone);
1007                                                    if let Some(ref tc) = tool_call {
1008                                                        if let Some(ref runner) = hook_runner {
1009                                                            let pre_call = runner.dispatch_pre_tool_call(tc).await;
1010                                                            if let Ok(res) = pre_call {
1011                                                                allow = res.allow;
1012                                                            }
1013                                                        }
1014                                                        if allow {
1015                                                            let key = (step_update_clone.trajectory_id.clone().unwrap_or_default(), step_update_clone.step_index.unwrap_or(0));
1016                                                            pending_calls.lock().await.insert(key, tc.clone());
1017                                                        }
1018                                                    }
1019
1020                                                    let conf = ToolConfirmation {
1021                                                        trajectory_id: step_update_clone.trajectory_id.clone(),
1022                                                        step_index: step_update_clone.step_index,
1023                                                        accepted: Some(allow),
1024                                                    };
1025                                                    let input_event = InputEvent {
1026                                                        event: Some(crate::proto::localharness::input_event::Event::ToolConfirmation(conf)),
1027                                                    };
1028                                                    if let Ok(raw_json) = serde_json::to_string(&input_event) {
1029                                                        let _ = conn_ws_tx.send(raw_json);
1030                                                    }
1031                                                });
1032                                            }
1033                                        }
1034                                        crate::proto::localharness::output_event::Event::TrajectoryStateUpdate(tsu) => {
1035                                            let sub_id = tsu.trajectory_id.clone().unwrap_or_default();
1036                                            let learned_cascade = conn_cascade_id_for_ws.lock().await;
1037                                            let is_subagent = learned_cascade.as_ref().is_some_and(|cid| !sub_id.is_empty() && sub_id != *cid);
1038                                            tracing::debug!("TrajectoryStateUpdate: trajectory_id={:?}, state={:?}, is_subagent={}, learned_cascade_id={:?}", sub_id, tsu.state, is_subagent, *learned_cascade);
1039                                            drop(learned_cascade);
1040
1041                                            let mut active_subs = conn_active_subagents.lock().await;
1042                                            let mut p_idle = conn_parent_idle.lock().await;
1043
1044                                            if tsu.state == Some(1) { // STATE_RUNNING
1045                                                if is_subagent {
1046                                                    active_subs.insert(sub_id);
1047                                                }
1048                                            } else if tsu.state == Some(2) { // STATE_IDLE
1049                                                if is_subagent {
1050                                                    active_subs.remove(&sub_id);
1051                                                } else {
1052                                                    *p_idle = true;
1053                                                }
1054                                            }
1055
1056                                            tracing::debug!("TrajectoryStateUpdate: p_idle={}, active_subs_empty={}", *p_idle, active_subs.is_empty());
1057                                            if *p_idle && active_subs.is_empty() && !conn_is_idle.swap(true, Ordering::SeqCst) {
1058                                                tracing::debug!("Connection transitioned to IDLE, sending sentinel");
1059                                                let sentinel = Step {
1060                                                    id: "IDLE_SENTINEL".to_string(),
1061                                                    ..Default::default()
1062                                                };
1063                                                let _ = step_tx.send(Ok(sentinel));
1064                                            }
1065                                        }
1066                                        crate::proto::localharness::output_event::Event::ToolCall(tool_call) => {
1067                                            let conn_ws_tx = conn_ws_tx.clone();
1068                                            let tool_runner = tool_runner.clone();
1069                                            let hook_runner = hook_runner.clone();
1070                                            let step_tx_clone = step_tx.clone();
1071                                            let learned_id_clone = conn_learned_id.clone();
1072                                            let counter = client_tool_step_counter.clone();
1073                                            tokio::spawn(async move {
1074                                                let args: Value = serde_json::from_str(&tool_call.arguments_json.clone().unwrap_or_default()).unwrap_or(Value::Null);
1075                                                let tc = ToolCall {
1076                                                    id: tool_call.id.clone().unwrap_or_default(),
1077                                                    name: tool_call.name.clone().unwrap_or_default(),
1078                                                    args: args.clone(),
1079                                                    canonical_path: None,
1080                                                };
1081                                                tracing::debug!("ToolCall event received: id={}, name={}", tc.id, tc.name);
1082
1083                                                // Emit ACTIVE step so the UI can show a tool card
1084                                                let synth_idx = counter.fetch_add(1, Ordering::SeqCst);
1085                                                let traj_id = learned_id_clone.get()
1086                                                    .cloned()
1087                                                    .unwrap_or_default();
1088                                                let active_step = Step {
1089                                                    id: tc.id.clone(),
1090                                                    step_index: synth_idx,
1091                                                    r#type: StepType::ToolCall,
1092                                                    source: StepSource::Model,
1093                                                    target: StepTarget::Environment,
1094                                                    status: StepStatus::Active,
1095                                                    content: tc.name.clone(),
1096                                                    tool_calls: vec![tc.clone()],
1097                                                    trajectory_id: traj_id.clone(),
1098                                                    ..Default::default()
1099                                                };
1100                                                let _ = step_tx_clone.send(Ok(active_step));
1101
1102                                                let allow = if let Some(runner) = hook_runner.as_ref() {
1103                                                    let res = runner.dispatch_pre_tool_call(&tc).await.map_or(true, |res| res.allow);
1104                                                    tracing::debug!("Policy decision for tool {}: allow={}", tc.name, res);
1105                                                    res
1106                                                } else {
1107                                                    true
1108                                                };
1109
1110                                                if !allow {
1111                                                    // Emit ERROR step for denied tool call
1112                                                    let denied_step = Step {
1113                                                        id: tc.id.clone(),
1114                                                        step_index: synth_idx,
1115                                                        r#type: StepType::ToolCall,
1116                                                        source: StepSource::Model,
1117                                                        target: StepTarget::Environment,
1118                                                        status: StepStatus::Error,
1119                                                        content: tc.name.clone(),
1120                                                        error: "Execution denied by hook policy".to_string(),
1121                                                        tool_calls: vec![tc.clone()],
1122                                                        trajectory_id: traj_id,
1123                                                        ..Default::default()
1124                                                    };
1125                                                    let _ = step_tx_clone.send(Ok(denied_step));
1126
1127                                                    let resp = ToolResponse {
1128                                                        id: tool_call.id.clone(),
1129                                                        response_json: Some("{\"error\": \"Execution denied by hook policy\"}".to_string()),
1130                                                        supplemental_media: Vec::new(),
1131                                                        response: None,
1132                                                    };
1133                                                    let input_event = InputEvent {
1134                                                        event: Some(crate::proto::localharness::input_event::Event::ToolResponse(resp)),
1135                                                    };
1136                                                    if let Ok(raw_json) = serde_json::to_string(&input_event) {
1137                                                        let _ = conn_ws_tx.send(raw_json);
1138                                                    }
1139                                                    return;
1140                                                }
1141
1142                                                let mut result = ToolResult {
1143                                                    id: tool_call.id.clone(),
1144                                                    name: tc.name.clone(),
1145                                                    result: None,
1146                                                    error: None,
1147                                                };
1148
1149                                                if let Some(ref runner) = tool_runner {
1150                                                    tracing::debug!("Executing tool {} with args: {:?}", tc.name, tc.args);
1151                                                    let results = runner.process_tool_calls(vec![tc.clone()]).await;
1152                                                    if let Some(r) = results.into_iter().next() {
1153                                                        result = r;
1154                                                    }
1155                                                } else {
1156                                                    result.error = Some("No tool runner registered".to_string());
1157                                                }
1158
1159                                                if let (Some(err_str), Some(runner)) = (result.error.as_ref(), hook_runner.as_ref()) {
1160                                                    if let Ok((res, val)) = runner.dispatch_on_tool_error(&anyhow!(err_str.clone())).await {
1161                                                        let allow_error = res.allow;
1162                                                        if allow_error {
1163                                                            result.result = val;
1164                                                            result.error = None;
1165                                                        }
1166                                                    }
1167                                                } else if let Some(runner) = hook_runner.as_ref() {
1168                                                    let _ = runner.dispatch_post_tool_call(&result).await;
1169                                                }
1170
1171                                                // Emit DONE or ERROR step with the execution result
1172                                                let (final_status, error_msg, result_args) = if result.error.is_some() {
1173                                                    (StepStatus::Error, result.error.clone().unwrap_or_default(), Value::Null)
1174                                                } else {
1175                                                    (StepStatus::Done, String::new(), result.result.clone().unwrap_or(Value::Null))
1176                                                };
1177                                                let done_step = Step {
1178                                                    id: tc.id.clone(),
1179                                                    step_index: synth_idx,
1180                                                    r#type: StepType::ToolCall,
1181                                                    source: StepSource::Model,
1182                                                    target: StepTarget::Environment,
1183                                                    status: final_status,
1184                                                    content: tc.name.clone(),
1185                                                    error: error_msg,
1186                                                    tool_calls: vec![ToolCall {
1187                                                        id: tc.id.clone(),
1188                                                        name: tc.name.clone(),
1189                                                        args: result_args,
1190                                                        canonical_path: None,
1191                                                    }],
1192                                                    trajectory_id: traj_id,
1193                                                    ..Default::default()
1194                                                };
1195                                                let _ = step_tx_clone.send(Ok(done_step));
1196
1197                                                // The Go harness expects responseJson to always be a JSON object.
1198                                                // Wrap non-object values (string, number, array, etc.) under "result".
1199                                                let resp_json = if let Some(ref val) = result.result {
1200                                                    if val.is_object() {
1201                                                        serde_json::to_string(val).unwrap_or_default()
1202                                                    } else {
1203                                                        serde_json::to_string(&serde_json::json!({ "result": val })).unwrap_or_default()
1204                                                    }
1205                                                } else if let Some(ref err) = result.error {
1206                                                    serde_json::to_string(&serde_json::json!({ "error": err })).unwrap_or_default()
1207                                                } else {
1208                                                    "{}".to_string()
1209                                                };
1210
1211                                                let resp = ToolResponse {
1212                                                    id: tool_call.id.clone(),
1213                                                    response_json: Some(resp_json),
1214                                                    supplemental_media: Vec::new(),
1215                                                    response: None,
1216                                                };
1217                                                let input_event = InputEvent {
1218                                                    event: Some(crate::proto::localharness::input_event::Event::ToolResponse(resp)),
1219                                                };
1220                                                if let Ok(raw_json) = serde_json::to_string(&input_event) {
1221                                                    tracing::debug!("Sending ToolResponse input_event: {}", raw_json);
1222                                                    let _ = conn_ws_tx.send(raw_json);
1223                                                }
1224                                            });
1225                                        }
1226                                    }
1227                                }
1228                            }
1229                            Err(e) => {
1230                                tracing::error!(
1231                                    "Failed to deserialize OutputEvent: {:?}. Raw: {}",
1232                                    e,
1233                                    &raw_text[..raw_text.len().min(300)]
1234                                );
1235                            }
1236                        }
1237                    }
1238                    Ok(_) => {}
1239                    Err(e) => {
1240                        let _ = step_tx.send(Err(anyhow!("WS read error: {e:?}")));
1241                        break;
1242                    }
1243                }
1244            }
1245        });
1246
1247        // 10. Hook runners dispatch session start
1248        if let Some(ref runner) = self.hook_runner {
1249            runner.dispatch_session_start().await?;
1250        }
1251
1252        Ok(LocalConnection {
1253            conversation_id,
1254            learned_id,
1255            process: Arc::new(Mutex::new(child)),
1256            child_stdin: Arc::new(Mutex::new(Some(child_stdin))),
1257            is_idle,
1258            step_rx: Arc::new(Mutex::new(Some(step_rx))),
1259            ws_tx,
1260            tool_runner: self.tool_runner.clone(),
1261            hook_runner: self.hook_runner.clone(),
1262            parent_idle,
1263            active_subagent_ids,
1264            step_trackers,
1265        })
1266    }
1267}
1268
1269/// Internal state tracker for matching `StepUpdate` payloads with active handshakes.
1270#[derive(Debug)]
1271pub struct StepTracker {
1272    state: i32,
1273    handled_requests: HashSet<String>,
1274}
1275
1276impl StepTracker {
1277    /// Updates the tracked step status state.
1278    pub const fn update_state(&mut self, state: i32) {
1279        self.state = state;
1280    }
1281
1282    /// Marks a specific request payload (e.g. "`tool_confirmation_request`") as handled.
1283    /// Returns true if the tracker transitioned to active for the request.
1284    pub fn mark_handled(&mut self, request_name: &str) -> bool {
1285        if self.state == 3 && !self.handled_requests.contains(request_name) {
1286            self.handled_requests.insert(request_name.to_string());
1287            return true;
1288        }
1289        false
1290    }
1291}
1292
1293#[allow(clippy::too_many_lines)]
1294fn extract_builtin_tool_call(
1295    step_update: &crate::proto::localharness::StepUpdate,
1296) -> Option<ToolCall> {
1297    let traj_id = step_update.trajectory_id.clone().unwrap_or_default();
1298    let step_idx = step_update.step_index.unwrap_or(0);
1299    let id = format!("{traj_id}_{step_idx}");
1300
1301    if step_update.invoke_subagent.is_some() {
1302        return Some(ToolCall {
1303            id,
1304            name: "START_SUBAGENT".to_string(),
1305            args: serde_json::json!({
1306                "prompt": step_update.request_text.clone().unwrap_or_default()
1307            }),
1308            canonical_path: None,
1309        });
1310    }
1311
1312    if let Some(ref fd) = step_update.find_file {
1313        return Some(ToolCall {
1314            id,
1315            name: "FIND_FILE".to_string(),
1316            args: serde_json::json!({
1317                "directory_path": fd.directory_path,
1318                "query": fd.query,
1319            }),
1320            canonical_path: fd.directory_path.clone(),
1321        });
1322    }
1323    if let Some(ref run) = step_update.run_command {
1324        return Some(ToolCall {
1325            id,
1326            name: "RUN_COMMAND".to_string(),
1327            args: serde_json::json!({
1328                "command_line":    run.command_line,
1329                "working_dir":     run.working_dir,
1330                // Include the execution result fields so the frontend
1331                // can display stdout/stderr instead of "(no output)".
1332                "combined_output": run.combined_output,
1333                "exit_code":       run.exit_code,
1334            }),
1335            canonical_path: None,
1336        });
1337    }
1338    if let Some(ref view) = step_update.view_file {
1339        return Some(ToolCall {
1340            id,
1341            name: "VIEW_FILE".to_string(),
1342            args: serde_json::json!({
1343                "file_path": view.file_path,
1344                "start_line": view.start_line,
1345                "end_line": view.end_line,
1346            }),
1347            canonical_path: view.file_path.clone(),
1348        });
1349    }
1350    if let Some(ref write) = step_update.create_file {
1351        return Some(ToolCall {
1352            id,
1353            name: "CREATE_FILE".to_string(),
1354            args: serde_json::json!({
1355                "file_path": write.file_path,
1356                "contents": write.contents,
1357            }),
1358            canonical_path: write.file_path.clone(),
1359        });
1360    }
1361    if let Some(ref edit) = step_update.edit_file {
1362        return Some(ToolCall {
1363            id,
1364            name: "EDIT_FILE".to_string(),
1365            args: serde_json::json!({
1366                "file_path": edit.file_path,
1367            }),
1368            canonical_path: edit.file_path.clone(),
1369        });
1370    }
1371    if let Some(ref search) = step_update.search_directory {
1372        // The harness puts grep/search results into `step_update.text`.
1373        // Pack them into `args.output` so the frontend can display them,
1374        // mirroring how RUN_COMMAND packs `combined_output`.
1375        return Some(ToolCall {
1376            id,
1377            name: "SEARCH_DIR".to_string(),
1378            args: serde_json::json!({
1379                "directory_path": search.directory_path,
1380                "query": search.query,
1381                "num_results": search.num_results,
1382                // Actual grep results from the harness
1383                "output": step_update.text,
1384            }),
1385            canonical_path: search.directory_path.clone(),
1386        });
1387    }
1388    if let Some(ref list) = step_update.list_directory {
1389        return Some(ToolCall {
1390            id,
1391            name: "LIST_DIR".to_string(),
1392            args: serde_json::json!({
1393                "directory_path": list.directory_path,
1394            }),
1395            canonical_path: list.directory_path.clone(),
1396        });
1397    }
1398    if let Some(ref img_gen) = step_update.generate_image {
1399        return Some(ToolCall {
1400            id,
1401            name: "GENERATE_IMAGE".to_string(),
1402            args: serde_json::json!({
1403                "prompt": img_gen.prompt,
1404                "image_paths": img_gen.image_paths,
1405                "image_name": img_gen.image_name,
1406            }),
1407            canonical_path: None,
1408        });
1409    }
1410    None
1411}
1412
1413fn extract_tool_result(step_update: &crate::proto::localharness::StepUpdate) -> Option<ToolResult> {
1414    let traj_id = step_update.trajectory_id.clone().unwrap_or_default();
1415    let step_idx = step_update.step_index.unwrap_or(0);
1416    let id = format!("{traj_id}_{step_idx}");
1417
1418    let tool_call = extract_builtin_tool_call(step_update)?;
1419    let result = step_update.text.clone().map(Value::String);
1420    let error = step_update.error_message.clone();
1421
1422    Some(ToolResult {
1423        id: Some(id),
1424        name: tool_call.name,
1425        result,
1426        error,
1427    })
1428}
1429
1430/// Returns the Rust compiler version used to build this crate.
1431fn rustc_version() -> String {
1432    option_env!("RUSTC_VERSION")
1433        .or(option_env!("CARGO_PKG_RUST_VERSION"))
1434        .unwrap_or("unknown")
1435        .to_string()
1436}
antigravity_sdk_rust/local.rs

antigravity_sdk_rust/
local.rs
