Observe Antigravity progress with fewer duplicate reads

loopx/capabilities/native_chat/antigravity_adapter.py

@@ -101,6 +101,45 @@ def _observed_turn_steps(steps: list[Any], start: int, total: int) -> tuple[list
     return current, base + len(steps)
 
 
+_RUNNING_STEP_BODIES = {
+    "CORTEX_STEP_TYPE_PLANNER_RESPONSE": "plannerResponse",
+    "CORTEX_STEP_TYPE_RUN_COMMAND": "runCommand",
+    "CORTEX_STEP_TYPE_VIEW_FILE": "viewFile",
+    "CORTEX_STEP_TYPE_CODE_ACTION": "codeAction",
+    "CORTEX_STEP_TYPE_ERROR_MESSAGE": "errorMessage",
+}
+
+
+def _running_snapshot_steps(snapshot: dict[str, Any], *, cascade_id: str, start: int,
+                            total: int, sent_prompt: str) -> list[tuple[int, dict[str, Any]]] | None:
+    """Reuse a complete running snapshot only for the input this adapter sent."""
+    trajectory = snapshot.get("trajectory")
+    if (snapshot.get("status") != "CASCADE_RUN_STATUS_RUNNING" or not isinstance(trajectory, dict)
+            or trajectory.get("cascadeId") != cascade_id):
+        return None
+    raw = trajectory.get("steps")
+    if not isinstance(raw, list) or len(raw) != total:
+        return None
+    try:
+        indexed, covered = _observed_turn_steps(raw, start, total)
+        if covered != total or not indexed:
+            return None
+        first = indexed[0][1]
+        if (first.get("type") != "CORTEX_STEP_TYPE_USER_INPUT"
+                or _native_user_input_text(first) != sent_prompt):
+            return None
+        for _, step in indexed:
+            kind = step.get("type")
+            if not isinstance(kind, str) or not isinstance(step.get("status"), str):
+                return None
+            body = _RUNNING_STEP_BODIES.get(kind)
+            if body and (not isinstance(step.get(body), dict) or not step[body]):
+                return None
+    except (AttributeError, TypeError, ValueError):
+        return None
+    return indexed
+
+
 def _check_hard_timeout(value: float | None) -> None:
     if value is not None and (type(value) not in (int, float) or not math.isfinite(value) or value <= 0):
         raise ValueError("Antigravity hard timeout must be finite and positive")
@@ -1065,30 +1104,51 @@ def timeout(kind: str) -> None:
             error.native_error_info = {"timeout": detail}
             raise error
 
-        while True:
-            now = time.monotonic()
-            if deadline is not None and now >= deadline:
-                timeout("hard")
-            if now - last_progress_at >= ANTIGRAVITY_IDLE_TIMEOUT_SEC:
-                timeout("idle")
-            self._check_interrupted()
-            self._interrupted.wait(0.25)
-            self._check_interrupted()
+        observation_limit = 10.0
+
+        def observe(method: str, args: dict[str, Any],
+                    final_read_deadline: float | None) -> dict[str, Any] | None:
+            remaining = observation_limit
+            if final_read_deadline is not None:
+                now = time.monotonic()
+                if deadline is not None and now >= deadline:
+                    timeout("hard")
+                remaining = final_read_deadline - now
+                if remaining <= 0:
+                    timeout("idle")
             try:
-                traj = _rpc_post(
-                    self.ls_url,
-                    self.csrf_token,
-                    "GetCascadeTrajectory",
-                    {"cascadeId": self.cascade_id},
-                    timeout=10.0,
-                )
+                return _rpc_post(self.ls_url, self.csrf_token, method,
+                                 {"cascadeId": self.cascade_id, **args}, timeout=remaining)
             except CodexChatAgentError as exc:
-                # Losing an idempotent observation does not mean that the
-                # native invocation failed. Never resend its input, refresh its
-                # progress clock, or bypass the existing stop/deadline checks.
                 if exc.error_code != "antigravity_transport_error":
                     raise
+                # A missed observation neither resends the input nor refreshes
+                # the clock. The last observation has no transport retry.
+                if final_read_deadline is not None:
+                    self._check_interrupted()
+                    if deadline is not None and time.monotonic() >= deadline:
+                        timeout("hard")
+                    timeout("idle")
                 self._interrupted.wait(1.0)
+                return None
+
+        while True:
+            now = time.monotonic()
+            if deadline is not None and now >= deadline:
+                timeout("hard")
+            self._check_interrupted()
+            final_read_deadline = None
+            if now - last_progress_at >= ANTIGRAVITY_IDLE_TIMEOUT_SEC:
+                # Native work can advance between the last poll and the idle
+                # deadline. Check the same invocation once before cancelling.
+                final_read_deadline = now + observation_limit
+                if deadline is not None:
+                    final_read_deadline = min(final_read_deadline, deadline)
+            else:
+                self._interrupted.wait(0.25)
+                self._check_interrupted()
+            traj = observe("GetCascadeTrajectory", {}, final_read_deadline)
+            if traj is None:
                 continue
             status = str(traj.get("status") or "")
             total_steps = int(traj.get("numTotalSteps") or 0)
@@ -1099,30 +1159,29 @@ def timeout(kind: str) -> None:
             observed_total = total_steps
             if total_steps > start_step_index:
                 self._check_interrupted()
-                try:
-                    steps_resp = _rpc_post(
-                        self.ls_url,
-                        self.csrf_token,
-                        "GetCascadeTrajectorySteps",
-                        {
-                            "cascadeId": self.cascade_id,
-                            "startIndex": start_step_index,
-                            "endIndex": total_steps,
-                        },
-                        timeout=10.0,
+                snapshot_steps = None
+                if final_read_deadline is None:
+                    snapshot_steps = _running_snapshot_steps(
+                        traj, cascade_id=self.cascade_id, start=start_step_index,
+                        total=total_steps, sent_prompt=full_prompt,
                     )
-                except CodexChatAgentError as exc:
-                    if exc.error_code != "antigravity_transport_error":
-                        raise
-                    self._interrupted.wait(1.0)
-                    continue
-                steps = steps_resp.get("steps")
-                try:
-                    if not isinstance(steps, list):
-                        raise ValueError("Antigravity returned no trajectory coverage")
-                    current_steps, observed_total = _observed_turn_steps(steps, start_step_index, total_steps)
-                except (AttributeError, TypeError, ValueError) as exc:
-                    raise CodexChatAgentError(str(exc), gate=None, error_code="antigravity_rpc_error") from exc
+                if snapshot_steps is not None:
+                    current_steps = snapshot_steps
+                else:
+                    # Older and incomplete snapshots use the detailed read.
+                    # Keep that proof for the idle boundary and final answer.
+                    steps_resp = observe("GetCascadeTrajectorySteps", {
+                        "startIndex": start_step_index, "endIndex": total_steps,
+                    }, final_read_deadline)
+                    if steps_resp is None:
+                        continue
+                    steps = steps_resp.get("steps")
+                    try:
+                        if not isinstance(steps, list):
+                            raise ValueError("Antigravity returned no trajectory coverage")
+                        current_steps, observed_total = _observed_turn_steps(steps, start_step_index, total_steps)
+                    except (AttributeError, TypeError, ValueError) as exc:
+                        raise CodexChatAgentError(str(exc), gate=None, error_code="antigravity_rpc_error") from exc
                 if current_steps:
                     if sum(step.get("type") == "CORTEX_STEP_TYPE_USER_INPUT" for _, step in current_steps) > 1:
                         # Another native input is not this invocation's result or
@@ -1236,6 +1295,8 @@ def timeout(kind: str) -> None:
             self._check_interrupted()
             if deadline is not None and now >= deadline:
                 timeout("hard")
+            idle_expired = (final_read_deadline is not None
+                            and now - last_progress_at >= ANTIGRAVITY_IDLE_TIMEOUT_SEC)
             if saw_running_or_new_step and status == "CASCADE_RUN_STATUS_IDLE" and observed_total == total_steps:
                 final = current_steps[-1][1] if current_steps else {}
                 if final.get("type") == "CORTEX_STEP_TYPE_ERROR_MESSAGE":
@@ -1270,6 +1331,8 @@ def timeout(kind: str) -> None:
                         or (planner.get("stopReason") is None and not self._read_only_review)
                         or any(step.get("status") in {"CORTEX_STEP_STATUS_GENERATING", "CORTEX_STEP_STATUS_RUNNING"}
                                for _, step in current_steps))):
+                    if idle_expired:
+                        timeout("idle")
                     continue
                 if self._read_only_review:
                     inputs = [step for _, step in current_steps if step.get("type") == "CORTEX_STEP_TYPE_USER_INPUT"]
@@ -1291,6 +1354,8 @@ def timeout(kind: str) -> None:
                     raise CodexChatAgentError("Antigravity ended without a verified final response; execution may have occurred.",
                                               gate=None, error_code="upstream_transport_failure")
                 break
+            if idle_expired:
+                timeout("idle")
 
         if event_sink is not None:
             # Keep only completion facts and the original input/final text,

tests/test_steward_interrupted_answer.py

@@ -244,6 +244,61 @@ def test_original_native_code_action_result_survives_later_timeout_without_repla
     assert [method for method, _ in language_server["calls"]].count("SendUserCascadeMessage") == 1
 
 
+@pytest.mark.parametrize("embedded", ["complete", "short", "foreign", "wrong_input", "missing_payload"])
+def test_native_running_poll_reuses_only_complete_steps_for_the_current_input(
+    tmp_path, language_server, monkeypatch, embedded,
+):
+    from test_antigravity_steward_adapter import adapter_for, _native_clock, _native_frames
+
+    adapter = adapter_for(tmp_path)
+    older_input = {"type": "CORTEX_STEP_TYPE_USER_INPUT", "status": "CORTEX_STEP_STATUS_DONE",
+        "metadata": {"sourceTrajectoryStepInfo": {"stepIndex": 0}},
+        "userInput": {"items": [{"text": "An earlier request"}]}}
+    frames = _native_frames()
+    for frame in frames:
+        for index, step in enumerate(frame["steps"], start=1):
+            step["metadata"]["sourceTrajectoryStepInfo"]["stepIndex"] = index
+        frame["steps"].insert(0, deepcopy(older_input))
+    _native_clock(monkeypatch, language_server, frames)
+    rpc = native._rpc_post
+    detail_reads = []
+
+    def read(url, csrf, method, payload, timeout):
+        result = rpc(url, csrf, method, payload, timeout=timeout)
+        if method == "GetCascadeTrajectorySteps":
+            detail_reads.append(language_server["frame"]["status"])
+        if method != "GetCascadeTrajectory":
+            return result
+        if not language_server["sent_count"]:
+            return {**result, "numTotalSteps": 1}
+        steps = deepcopy(language_server["frame"]["steps"])
+        steps[1]["userInput"] = {"items": [{"text": language_server["inputs"][0]["text"]}]}
+        trajectory = {**result["trajectory"], "steps": steps}
+        if result["status"] == "CASCADE_RUN_STATUS_IDLE":
+            steps[-1]["plannerResponse"]["response"] = "Stale embedded final"
+        elif embedded == "short":
+            steps.pop(0)
+        elif embedded == "foreign":
+            trajectory["cascadeId"] = "unrelated-cascade"
+        elif embedded == "wrong_input":
+            steps[1]["userInput"]["items"][0]["text"] = "Another invocation"
+        elif embedded == "missing_payload":
+            steps[-1].pop("plannerResponse")
+        return {**result, "trajectory": trajectory}
+
+    monkeypatch.setattr(native, "_rpc_post", read)
+    events = []
+    result = adapter.start_turn("Return the original result", event_sink=lambda kind, data: events.append((kind, data)))
+    assert result["message"] == "native final"
+    assert any(kind == "answer.delta" and "progress" in data["text"] for kind, data in events)
+    assert "Stale embedded final" not in json.dumps(events)
+    assert detail_reads.count("CASCADE_RUN_STATUS_IDLE") == 1
+    assert detail_reads.count("CASCADE_RUN_STATUS_RUNNING") == (0 if embedded == "complete" else 4)
+    methods = [method for method, _ in language_server["calls"]]
+    assert methods.count("SendUserCascadeMessage") == 1
+    assert "CancelCascadeInvocation" not in methods
+
+
 @pytest.mark.parametrize("hard_limit, native_error", [(None, True), (150.0, True), (None, False)])
 def test_native_timeout_reports_its_clock_without_extending_native_retry(
     tmp_path, language_server, monkeypatch, hard_limit, native_error,
@@ -282,6 +337,49 @@ def test_native_timeout_reports_its_clock_without_extending_native_retry(
     assert len(errors) == diagnostic["native_error_count"]
 
 
+def test_original_steward_observes_native_work_at_idle_deadline_and_replies_once(
+    steward, language_server, monkeypatch,
+):
+    from test_antigravity_steward_adapter import _native_clock
+
+    monkeypatch.setenv("LOOPX_MANAGER_MODEL", "gemini-4-argon-high")
+    monkeypatch.setenv("LOOPX_MANAGER_REASONING_EFFORT", "high")
+    _, runtime, provider, transport, binding, _, _ = steward
+    binding = transport.bindings.configure(transport_ref="steward-app", project_ref=binding["project_ref"],
+        executor_endpoint_id="antigravity", context_kind="steward")
+    runtime.hard_timeout_sec = None
+
+    user = {"type": "CORTEX_STEP_TYPE_USER_INPUT", "status": "CORTEX_STEP_STATUS_DONE"}
+    previous = {"type": "CORTEX_STEP_TYPE_RUN_COMMAND", "status": "CORTEX_STEP_STATUS_DONE"}
+    resumed = {"type": "CORTEX_STEP_TYPE_RUN_COMMAND", "status": "CORTEX_STEP_STATUS_RUNNING"}
+    finished = {**resumed, "status": "CORTEX_STEP_STATUS_DONE"}
+    final = {"type": "CORTEX_STEP_TYPE_PLANNER_RESPONSE", "status": "CORTEX_STEP_STATUS_DONE",
+             "plannerResponse": {"response": "Original answer", "stopReason": "STOP_REASON_STOP_PATTERN"}}
+    _native_clock(monkeypatch, language_server, [
+        {"at": 1, "steps": [user, previous], "status": "CASCADE_RUN_STATUS_RUNNING"},
+        {"at": 302, "steps": [user, previous], "status": "CASCADE_RUN_STATUS_RUNNING"},
+        {"at": 302.01, "steps": [user, previous, resumed], "status": "CASCADE_RUN_STATUS_RUNNING"},
+        {"at": 304, "steps": [user, previous, finished, final], "status": "CASCADE_RUN_STATUS_IDLE"},
+    ])
+
+    source = provider.event("steward-app", "idle-deadline-source", "Finish this request")
+    transport.admit("steward-app", source)
+    ref = identity_ref(binding["provider_ref"], source["message_id"])
+    row = transport.core.read_request(ref)
+    done = runtime.wait_for_turn(session_id=row["session_id"], turn_id=row["turn_id"], timeout_sec=8)
+
+    assert done["status"] == "completed" and done["response"]["message"] == "Original answer"
+    assert transport.reconcile() == 1
+    delivery = json.loads((transport.root / f"{ref}.json").read_text())["deliveries"]["terminal"]
+    assert delivery["verified"] and delivery["text"] == done["response"]["message"]
+    returned = provider.messages[delivery["attempt"]["message_ref"]]
+    assert returned["chat_id"] == source["chat_id"]
+    assert "Original answer" in json.dumps(returned["body"])
+    assert transport.reconcile() == 0
+    assert sum(method == "SendUserCascadeMessage" for method, _ in language_server["calls"]) == 1
+    assert not language_server["cancelled"].is_set()
+
+
 @pytest.mark.parametrize("kind, expected", [
     ("idle", "等待原生助手有效进展超时"), ("hard", "原生助手达到本次执行时限"),
     ("unknown-host-value", "等待助手响应超时"), (None, "等待助手响应超时"),