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

antigravity_sdk_rust/
local.rs
