No public description

PiperOrigin-RevId: 926283277
Change-Id: Ie6208176173b83e8a5d3593b2278e662396d078a

Author: Antigravity Team

google/antigravity/connections/local/local_connection.py

@@ -46,7 +46,6 @@
 from google.antigravity.hooks import hooks
 from google.antigravity.tools import tool_runner as t_runner
 
-
 _ANY_ADAPTER = pydantic.TypeAdapter(Any)
 
 
@@ -132,6 +131,7 @@ def mark_handled(self, request_type: str) -> bool:
 def _get_mcp_tool_name(server_name: str, tool_name: str) -> str:
   return f"{_MCP_TOOL_PREFIX}{server_name}_{tool_name}"
 
+
 _IDLE_SENTINEL = object()
 _CLOSE_SENTINEL = None
 
@@ -601,6 +601,9 @@ async def receive_steps(self) -> AsyncIterator[LocalConnectionStep]:
 
         step_obj = await self._step_queue.get()
 
+        if isinstance(step_obj, Exception):
+          raise step_obj
+
         if step_obj is _IDLE_SENTINEL:
           continue
         if step_obj is None:
@@ -858,9 +861,7 @@ async def _ws_reader_loop(self) -> None:
               and step_obj.source == types.StepSource.MODEL
               and step_obj.content
           ):
-            self._subagent_responses[step_obj.trajectory_id] = (
-                step_obj.content
-            )
+            self._subagent_responses[step_obj.trajectory_id] = step_obj.content
 
           # Dispatch post-tool-call or on-tool-error hooks for built-in tools
           # that were approved via ToolConfirmation. The harness executes them
@@ -945,16 +946,12 @@ async def _ws_reader_loop(self) -> None:
               self._active_subagent_ids.discard(tsu.trajectory_id)
               if self._hook_runner:
                 op_ctx = hooks.OperationContext(self._get_turn_context())
-                response = self._subagent_responses.pop(
-                    tsu.trajectory_id, ""
-                )
+                response = self._subagent_responses.pop(tsu.trajectory_id, "")
                 result = types.ToolResult(
                     name=types.BuiltinTools.START_SUBAGENT.value,
                     result=response or tsu.trajectory_id,
                 )
-                await self._hook_runner.dispatch_post_tool_call(
-                    op_ctx, result
-                )
+                await self._hook_runner.dispatch_post_tool_call(op_ctx, result)
             else:
               # Parent trajectory went idle.
               self._parent_idle = True
@@ -965,6 +962,20 @@ async def _ws_reader_loop(self) -> None:
               if not self._is_idle.is_set():
                 self._is_idle.set()
                 await self._step_queue.put(_IDLE_SENTINEL)
+
+            if tsu.HasField("error"):
+              if is_subagent:
+                # Subagent failures are only logged. If an issue is serious
+                # enough to be worth raising an exception then it should
+                # affect the main trajectory too.
+                logging.info(
+                    "Subagent trajectory failed with error: %s", tsu.error
+                )
+              else:
+                await self._step_queue.put(
+                    types.AntigravityExecutionError(tsu.error)
+                )
+
         elif event.HasField("tool_call"):
           self._run_in_background(self._handle_tool_call(event.tool_call))
     except websockets.ConnectionClosed as e:
@@ -1061,13 +1072,12 @@ async def _handle_question_request(
       logging.exception("_handle_question_request failed; sending error")
       error_answer = localharness_pb2.UserQuestionAnswer(
           multiple_choice_answer=localharness_pb2.MultipleChoiceAnswer(
-              freeform_response=(
-                  f"SDK error processing question: {e!r}"
-              ),
+              freeform_response=f"SDK error processing question: {e!r}",
           ),
       )
       await self._send_question_response(
-          step_update, [error_answer],
+          step_update,
+          [error_answer],
       )
 
   async def _send_question_response(
@@ -1170,9 +1180,7 @@ async def _handle_tool_confirmation_request(
       # ToolConfirmation only has a bool field (no error/reason field), so
       # rejecting is the only option. The harness transitions the step to
       # STATE_ERROR, which the model does see.
-      logging.exception(
-          "_handle_tool_confirmation_request failed; rejecting"
-      )
+      logging.exception("_handle_tool_confirmation_request failed; rejecting")
       await self._send_tool_confirmation(step_update, False)
 
   async def _send_tool_confirmation(
@@ -1388,12 +1396,10 @@ def _get_default_binary_path_external() -> str:
     if dist.files:
       for f in dist.files:
         normalized_path = str(f).replace("\\", "/")
-        if normalized_path.endswith(
-            (
-                "google/antigravity/bin/localharness",
-                "google/antigravity/bin/localharness.exe",
-            )
-        ):
+        if normalized_path.endswith((
+            "google/antigravity/bin/localharness",
+            "google/antigravity/bin/localharness.exe",
+        )):
           binary_path = os.path.abspath(str(f.locate()))
           if os.path.exists(binary_path):
             return binary_path
@@ -1750,7 +1756,7 @@ async def __aenter__(self) -> None:
               f"Failed to connect to WebSocket at {ws_url} after"
               f" {max_retries} attempts. Stderr: {stderr_output}"
           ) from e
-        await asyncio.sleep(0.1 * (2 ** attempt))
+        await asyncio.sleep(0.1 * (2**attempt))
 
     assert ws is not None
     try:

google/antigravity/connections/local/local_connection_test.py

@@ -137,6 +137,47 @@ async def test_receive_steps_system_error_401(self):
       async for _ in harness.conn.receive_steps():
         pass
 
+  async def test_receive_steps_trajectory_error(self):
+    harness = self._make_harness()
+
+    await harness.conn.send("Hello")
+    init_data = await harness.wait_for_response()
+    self.assertEqual(init_data.get("userInput"), "Hello")
+
+    # Set the cascade ID.
+    event1 = localharness_pb2.OutputEvent(
+        step_update=localharness_pb2.StepUpdate(
+            cascade_id="my_cascade",
+            trajectory_id="my_cascade",
+            step_index=1,
+            state=localharness_pb2.StepUpdate.STATE_ACTIVE,
+            source=localharness_pb2.StepUpdate.SOURCE_MODEL,
+            text="I'm working",
+        )
+    )
+    await harness.send_event(event1)
+
+    # Send an error.
+    event2 = localharness_pb2.OutputEvent(
+        trajectory_state_update=localharness_pb2.TrajectoryStateUpdate(
+            trajectory_id="my_cascade",
+            state=localharness_pb2.TrajectoryStateUpdate.STATE_IDLE,
+            error="Trajectory execution failed",
+        )
+    )
+    await harness.send_event(event2)
+    await harness.close_from_harness_side()
+
+    steps = []
+    with self.assertRaisesRegex(
+        types.AntigravityExecutionError, "Trajectory execution failed"
+    ):
+      async for step in harness.conn.receive_steps():
+        steps.append(step)
+
+    self.assertEqual(len(steps), 1)
+    self.assertEqual(steps[0].content, "I'm working")
+
   def test_local_connection_step_from_dict(self):
     """Tests that LocalConnectionStep maps fields correctly."""
     step_dict = {

google/antigravity/connections/local/localharness_pb2.py

@@ -193,9 +193,10 @@
     b' \x01(\x0b\x32\x30.antigravity.localharness.UserInput.SlashCommandH\x00\x42\x06\n\x04part"O\n\x10ToolConfirmation\x12\x15\n\rtrajectory_id\x18\x01'
     b' \x01(\t\x12\x12\n\nstep_index\x18\x02'
     b' \x01(\r\x12\x10\n\x08\x61\x63\x63\x65pted\x18\x03'
-    b' \x01(\x08"\xb7\x01\n\x15TrajectoryStateUpdate\x12\x15\n\rtrajectory_id\x18\x02'
+    b' \x01(\x08"\xc6\x01\n\x15TrajectoryStateUpdate\x12\x15\n\rtrajectory_id\x18\x02'
     b' \x01(\t\x12\x44\n\x05state\x18\x03'
-    b' \x01(\x0e\x32\x35.antigravity.localharness.TrajectoryStateUpdate.State"A\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_RUNNING\x10\x01\x12\x0e\n\nSTATE_IDLE\x10\x02"9\n\x06Struct\x12/\n\x06\x66ields\x18\x01'
+    b' \x01(\x0e\x32\x35.antigravity.localharness.TrajectoryStateUpdate.State\x12\r\n\x05\x65rror\x18\x04'
+    b' \x01(\t"A\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_RUNNING\x10\x01\x12\x0e\n\nSTATE_IDLE\x10\x02"9\n\x06Struct\x12/\n\x06\x66ields\x18\x01'
     b' \x03(\x0b\x32\x1f.antigravity.localharness.Field"E\n\x05\x46ield\x12\x0c\n\x04name\x18\x01'
     b' \x01(\t\x12.\n\x05value\x18\x02'
     b' \x01(\x0b\x32\x1f.antigravity.localharness.Value"\xb7\x02\n\x05Value\x12\x39\n\nnull_value\x18\x01'
@@ -263,8 +264,8 @@
   _globals['_MCPSTDIOTRANSPORT_ENVENTRY']._serialized_options = b'8\001'
   _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._loaded_options = None
   _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._serialized_options = b'8\001'
-  _globals['_NULLVALUE']._serialized_start = 10385
-  _globals['_NULLVALUE']._serialized_end = 10412
+  _globals['_NULLVALUE']._serialized_start = 10400
+  _globals['_NULLVALUE']._serialized_end = 10427
   _globals['_INPUTCONFIG']._serialized_start = 49
   _globals['_INPUTCONFIG']._serialized_end = 195
   _globals['_INITIALIZECONVERSATIONEVENT']._serialized_start = 197
@@ -384,45 +385,45 @@
   _globals['_TOOLCONFIRMATION']._serialized_start = 7802
   _globals['_TOOLCONFIRMATION']._serialized_end = 7881
   _globals['_TRAJECTORYSTATEUPDATE']._serialized_start = 7884
-  _globals['_TRAJECTORYSTATEUPDATE']._serialized_end = 8067
-  _globals['_TRAJECTORYSTATEUPDATE_STATE']._serialized_start = 8002
-  _globals['_TRAJECTORYSTATEUPDATE_STATE']._serialized_end = 8067
-  _globals['_STRUCT']._serialized_start = 8069
-  _globals['_STRUCT']._serialized_end = 8126
-  _globals['_FIELD']._serialized_start = 8128
-  _globals['_FIELD']._serialized_end = 8197
-  _globals['_VALUE']._serialized_start = 8200
-  _globals['_VALUE']._serialized_end = 8511
-  _globals['_LISTVALUE']._serialized_start = 8513
-  _globals['_LISTVALUE']._serialized_end = 8573
-  _globals['_TOOLCALL']._serialized_start = 8575
-  _globals['_TOOLCALL']._serialized_end = 8688
-  _globals['_TOOLRESPONSE']._serialized_start = 8691
-  _globals['_TOOLRESPONSE']._serialized_end = 8853
-  _globals['_USERQUESTIONSRESPONSE']._serialized_start = 8856
-  _globals['_USERQUESTIONSRESPONSE']._serialized_end = 9124
-  _globals['_USERQUESTIONSRESPONSE_QUESTIONSRESPONSE']._serialized_start = 9032
-  _globals['_USERQUESTIONSRESPONSE_QUESTIONSRESPONSE']._serialized_end = 9114
-  _globals['_USERQUESTIONANSWER']._serialized_start = 9127
-  _globals['_USERQUESTIONANSWER']._serialized_end = 9261
-  _globals['_MULTIPLECHOICEANSWER']._serialized_start = 9263
-  _globals['_MULTIPLECHOICEANSWER']._serialized_end = 9345
+  _globals['_TRAJECTORYSTATEUPDATE']._serialized_end = 8082
+  _globals['_TRAJECTORYSTATEUPDATE_STATE']._serialized_start = 8017
+  _globals['_TRAJECTORYSTATEUPDATE_STATE']._serialized_end = 8082
+  _globals['_STRUCT']._serialized_start = 8084
+  _globals['_STRUCT']._serialized_end = 8141
+  _globals['_FIELD']._serialized_start = 8143
+  _globals['_FIELD']._serialized_end = 8212
+  _globals['_VALUE']._serialized_start = 8215
+  _globals['_VALUE']._serialized_end = 8526
+  _globals['_LISTVALUE']._serialized_start = 8528
+  _globals['_LISTVALUE']._serialized_end = 8588
+  _globals['_TOOLCALL']._serialized_start = 8590
+  _globals['_TOOLCALL']._serialized_end = 8703
+  _globals['_TOOLRESPONSE']._serialized_start = 8706
+  _globals['_TOOLRESPONSE']._serialized_end = 8868
+  _globals['_USERQUESTIONSRESPONSE']._serialized_start = 8871
+  _globals['_USERQUESTIONSRESPONSE']._serialized_end = 9139
+  _globals['_USERQUESTIONSRESPONSE_QUESTIONSRESPONSE']._serialized_start = 9047
+  _globals['_USERQUESTIONSRESPONSE_QUESTIONSRESPONSE']._serialized_end = 9129
+  _globals['_USERQUESTIONANSWER']._serialized_start = 9142
+  _globals['_USERQUESTIONANSWER']._serialized_end = 9276
+  _globals['_MULTIPLECHOICEANSWER']._serialized_start = 9278
+  _globals['_MULTIPLECHOICEANSWER']._serialized_end = 9360
   _globals['_MEDIA']._serialized_start = 7541
   _globals['_MEDIA']._serialized_end = 7602
-  _globals['_USAGEMETADATA']._serialized_start = 9411
-  _globals['_USAGEMETADATA']._serialized_end = 9579
-  _globals['_MCPSERVERCONFIG']._serialized_start = 9582
-  _globals['_MCPSERVERCONFIG']._serialized_end = 9982
-  _globals['_MCPSERVERCONFIG_AUTHPROVIDERTYPE']._serialized_start = 9872
-  _globals['_MCPSERVERCONFIG_AUTHPROVIDERTYPE']._serialized_end = 9969
-  _globals['_MCPSTDIOTRANSPORT']._serialized_start = 9985
-  _globals['_MCPSTDIOTRANSPORT']._serialized_end = 10146
-  _globals['_MCPSTDIOTRANSPORT_ENVENTRY']._serialized_start = 10104
-  _globals['_MCPSTDIOTRANSPORT_ENVENTRY']._serialized_end = 10146
-  _globals['_MCPHTTPTRANSPORT']._serialized_start = 10149
-  _globals['_MCPHTTPTRANSPORT']._serialized_end = 10302
-  _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._serialized_start = 10256
-  _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._serialized_end = 10302
-  _globals['_ACTIONMCPTOOL']._serialized_start = 10304
-  _globals['_ACTIONMCPTOOL']._serialized_end = 10383
+  _globals['_USAGEMETADATA']._serialized_start = 9426
+  _globals['_USAGEMETADATA']._serialized_end = 9594
+  _globals['_MCPSERVERCONFIG']._serialized_start = 9597
+  _globals['_MCPSERVERCONFIG']._serialized_end = 9997
+  _globals['_MCPSERVERCONFIG_AUTHPROVIDERTYPE']._serialized_start = 9887
+  _globals['_MCPSERVERCONFIG_AUTHPROVIDERTYPE']._serialized_end = 9984
+  _globals['_MCPSTDIOTRANSPORT']._serialized_start = 10000
+  _globals['_MCPSTDIOTRANSPORT']._serialized_end = 10161
+  _globals['_MCPSTDIOTRANSPORT_ENVENTRY']._serialized_start = 10119
+  _globals['_MCPSTDIOTRANSPORT_ENVENTRY']._serialized_end = 10161
+  _globals['_MCPHTTPTRANSPORT']._serialized_start = 10164
+  _globals['_MCPHTTPTRANSPORT']._serialized_end = 10317
+  _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._serialized_start = 10271
+  _globals['_MCPHTTPTRANSPORT_HEADERSENTRY']._serialized_end = 10317
+  _globals['_ACTIONMCPTOOL']._serialized_start = 10319
+  _globals['_ACTIONMCPTOOL']._serialized_end = 10398
 # @@protoc_insertion_point(module_scope)