1331 lines
52 KiB
Python
1331 lines
52 KiB
Python
"""
|
|
Unit tests for :class:`omnigent.inner.antigravity_executor.AntigravityExecutor`.
|
|
|
|
The fakes here mirror the real ``google.antigravity`` streaming surface the
|
|
executor depends on: ``agent.conversation`` yields :class:`Step` objects from
|
|
``receive_steps()`` (text / reasoning deltas, tool calls, status, usage) as the
|
|
turn runs, a registered ``PostToolCallHook`` fires per tool completion with a
|
|
``ToolResult``, and ``conversation.cancel()`` aborts a running turn. They let
|
|
the streaming / tool-pairing / cancellation logic be tested without the SDK
|
|
package or network.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import collections
|
|
import enum
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from omnigent.inner import antigravity_executor as ag
|
|
from omnigent.inner.antigravity_executor import AntigravityExecutor, _latest_user_text
|
|
from omnigent.inner.executor import (
|
|
ExecutorConfig,
|
|
ExecutorError,
|
|
ReasoningChunk,
|
|
TextChunk,
|
|
ToolCallComplete,
|
|
ToolCallRequest,
|
|
ToolCallStatus,
|
|
TurnCancelled,
|
|
TurnComplete,
|
|
)
|
|
from omnigent.llms._usage_observer import add_observer
|
|
|
|
# ── Fakes mirroring the real SDK streaming shapes ───────────────────────
|
|
|
|
|
|
class _StepType(enum.Enum):
|
|
"""Subset of ``google.antigravity.types.StepType`` the executor reads."""
|
|
|
|
TEXT_RESPONSE = "TEXT_RESPONSE"
|
|
TOOL_CALL = "TOOL_CALL"
|
|
FINISH = "FINISH"
|
|
|
|
|
|
class _StepStatus(enum.Enum):
|
|
"""Subset of ``google.antigravity.types.StepStatus`` the executor reads."""
|
|
|
|
ACTIVE = "ACTIVE"
|
|
DONE = "DONE"
|
|
CANCELED = "CANCELED"
|
|
ERROR = "ERROR"
|
|
TERMINAL_ERROR = "TERMINAL_ERROR"
|
|
|
|
|
|
class _StepSource(enum.Enum):
|
|
"""Subset of ``google.antigravity.types.StepSource``."""
|
|
|
|
SYSTEM = "SYSTEM"
|
|
USER = "USER"
|
|
MODEL = "MODEL"
|
|
|
|
|
|
class _StepTarget(enum.Enum):
|
|
"""Subset of ``google.antigravity.types.StepTarget``."""
|
|
|
|
USER = "USER"
|
|
ENVIRONMENT = "ENVIRONMENT"
|
|
|
|
|
|
class _AntigravityCancelledError(Exception):
|
|
"""Stand-in for ``google.antigravity.types.AntigravityCancelledError``."""
|
|
|
|
|
|
class _FakeToolCall:
|
|
def __init__(self, name: str, args: dict[str, Any], call_id: str | None = None) -> None:
|
|
self.name = name
|
|
self.args = args
|
|
self.id = call_id
|
|
|
|
|
|
class _FakeToolResult:
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
result: Any = None,
|
|
error: str | None = None,
|
|
call_id: str | None = None,
|
|
) -> None:
|
|
self.name = name
|
|
self.result = result
|
|
self.error = error
|
|
self.id = call_id
|
|
self.exception = None
|
|
|
|
|
|
class _FakeUsage:
|
|
def __init__(self) -> None:
|
|
self.prompt_token_count = 11
|
|
self.candidates_token_count = 7
|
|
self.total_token_count = 18
|
|
self.cached_content_token_count = 2
|
|
|
|
|
|
class _FakeStep:
|
|
"""Mirror of ``google.antigravity.types.Step`` (the fields the executor reads)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
step_type: _StepType | None = None,
|
|
status: _StepStatus | None = None,
|
|
content_delta: str = "",
|
|
thinking_delta: str = "",
|
|
tool_calls: list[_FakeToolCall] | None = None,
|
|
error: str = "",
|
|
usage_metadata: _FakeUsage | None = None,
|
|
source: _StepSource = _StepSource.MODEL,
|
|
target: _StepTarget = _StepTarget.USER,
|
|
) -> None:
|
|
self.type = step_type
|
|
self.status = status
|
|
self.content_delta = content_delta
|
|
self.thinking_delta = thinking_delta
|
|
self.tool_calls = tool_calls or []
|
|
self.error = error
|
|
self.usage_metadata = usage_metadata
|
|
# Default MODEL->USER (assistant-facing); set source=USER to model the
|
|
# SDK echoing the user's own input back in the step stream.
|
|
self.source = source
|
|
self.target = target
|
|
|
|
|
|
@dataclass
|
|
class _YieldStep:
|
|
"""Turn-script action: ``receive_steps`` yields this step."""
|
|
|
|
step: _FakeStep
|
|
|
|
|
|
@dataclass
|
|
class _FireToolResult:
|
|
"""Turn-script action: the SDK invokes each PostToolCallHook with this result."""
|
|
|
|
tool_result: _FakeToolResult
|
|
|
|
|
|
@dataclass
|
|
class _RaiseCancelled:
|
|
"""Turn-script action: ``receive_steps`` raises the SDK's cancellation error."""
|
|
|
|
|
|
@dataclass
|
|
class _RaiseGeneric:
|
|
"""Turn-script action: ``receive_steps`` raises a generic (non-cancel) error."""
|
|
|
|
message: str = "boom"
|
|
|
|
|
|
# A turn script is the ordered list of actions one ``receive_steps()`` replays.
|
|
_TurnAction = _YieldStep | _FireToolResult | _RaiseCancelled | _RaiseGeneric
|
|
|
|
|
|
class _FakeConversation:
|
|
"""Mirror of ``google.antigravity.conversation.Conversation`` (read paths)."""
|
|
|
|
def __init__(self, hooks: list[Any], scripts: collections.deque[list[_TurnAction]]) -> None:
|
|
self._hooks = hooks
|
|
self._scripts = scripts
|
|
self.sends: list[str] = []
|
|
self.cancel_called = 0
|
|
|
|
async def send(self, prompt: Any, **_kwargs: Any) -> None:
|
|
self.sends.append(prompt)
|
|
|
|
async def receive_steps(self) -> Any:
|
|
script = self._scripts.popleft() if self._scripts else []
|
|
for action in script:
|
|
if isinstance(action, _YieldStep):
|
|
yield action.step
|
|
elif isinstance(action, _FireToolResult):
|
|
for hook in self._hooks:
|
|
await hook.run(SimpleNamespace(), action.tool_result)
|
|
elif isinstance(action, _RaiseCancelled):
|
|
raise _AntigravityCancelledError("cancelled")
|
|
elif isinstance(action, _RaiseGeneric):
|
|
raise RuntimeError(action.message)
|
|
|
|
async def cancel(self) -> None:
|
|
self.cancel_called += 1
|
|
|
|
|
|
class _FakeAgent:
|
|
def __init__(self, config: Any, scripts: collections.deque[list[_TurnAction]]) -> None:
|
|
self.config = config
|
|
self._conversation = _FakeConversation(list(getattr(config, "hooks", []) or []), scripts)
|
|
self.closed = False
|
|
|
|
@property
|
|
def conversation(self) -> _FakeConversation:
|
|
return self._conversation
|
|
|
|
async def __aenter__(self) -> _FakeAgent:
|
|
return self
|
|
|
|
async def __aexit__(self, *_args: object) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class _FakeLocalAgentConfig:
|
|
"""Mirror of ``LocalAgentConfig`` — accepts exactly the fields the executor sets."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
system_instructions: str | None = None,
|
|
model: str | None = None,
|
|
api_key: str | None = None,
|
|
vertex: bool | None = None,
|
|
project: str | None = None,
|
|
location: str | None = None,
|
|
tools: Any = None,
|
|
hooks: Any = None,
|
|
) -> None:
|
|
self.system_instructions = system_instructions
|
|
self.model = model
|
|
self.api_key = api_key
|
|
self.vertex = vertex
|
|
self.project = project
|
|
self.location = location
|
|
self.tools = tools
|
|
self.hooks = hooks
|
|
|
|
|
|
class _FakePostToolCallHook:
|
|
"""Sub-classable stand-in for ``google.antigravity.hooks.PostToolCallHook``."""
|
|
|
|
async def run(self, context: Any, data: Any) -> None:
|
|
return None
|
|
|
|
|
|
def _install_fake_sdk(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
*,
|
|
scripts: list[list[_TurnAction]],
|
|
) -> dict[str, Any]:
|
|
"""Patch ``_ensure_antigravity_sdk`` to return a fake module.
|
|
|
|
:param monkeypatch: pytest monkeypatch fixture.
|
|
:param scripts: One turn-script (list of actions) per ``receive_steps`` call,
|
|
consumed front-to-back across turns / agent rebuilds.
|
|
:returns: A ``captured`` dict exposing the agents / configs built, so tests
|
|
can assert on what the executor passed to the SDK.
|
|
"""
|
|
queue: collections.deque[list[_TurnAction]] = collections.deque(scripts)
|
|
captured: dict[str, Any] = {"agents": [], "configs": []}
|
|
|
|
class _FakeHooks:
|
|
PostToolCallHook = _FakePostToolCallHook
|
|
|
|
class _FakeTypes:
|
|
AntigravityCancelledError = _AntigravityCancelledError
|
|
|
|
class _FakeModule:
|
|
LocalAgentConfig = _FakeLocalAgentConfig
|
|
hooks = _FakeHooks
|
|
types = _FakeTypes
|
|
|
|
@staticmethod
|
|
def Agent(config: Any) -> _FakeAgent:
|
|
agent = _FakeAgent(config, queue)
|
|
captured["agents"].append(agent)
|
|
captured["configs"].append(config)
|
|
return agent
|
|
|
|
monkeypatch.setattr(ag, "_ensure_antigravity_sdk", lambda: _FakeModule())
|
|
return captured
|
|
|
|
|
|
async def _drain(
|
|
executor: AntigravityExecutor,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None = None,
|
|
config: ExecutorConfig | None = None,
|
|
system_prompt: str = "sys",
|
|
) -> list[Any]:
|
|
events: list[Any] = []
|
|
async for event in executor.run_turn(
|
|
messages, tools=tools or [], system_prompt=system_prompt, config=config
|
|
):
|
|
events.append(event)
|
|
return events
|
|
|
|
|
|
def _text_step(delta: str) -> _YieldStep:
|
|
return _YieldStep(
|
|
_FakeStep(
|
|
step_type=_StepType.TEXT_RESPONSE, status=_StepStatus.ACTIVE, content_delta=delta
|
|
)
|
|
)
|
|
|
|
|
|
def _tool_call_step(call: _FakeToolCall, status: _StepStatus = _StepStatus.ACTIVE) -> _YieldStep:
|
|
return _YieldStep(_FakeStep(step_type=_StepType.TOOL_CALL, status=status, tool_calls=[call]))
|
|
|
|
|
|
# ── Tests ───────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_latest_user_text_prefers_last_user_message() -> None:
|
|
messages: list[dict[str, Any]] = [
|
|
{"role": "user", "content": "first"},
|
|
{"role": "assistant", "content": "reply"},
|
|
{"role": "user", "content": [{"type": "text", "text": "second"}]},
|
|
]
|
|
assert _latest_user_text(messages) == "second"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_maps_text_reasoning_and_usage(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Text/reasoning stream as separate deltas; usage + final text land on TurnComplete."""
|
|
script: list[_TurnAction] = [
|
|
_YieldStep(_FakeStep(status=_StepStatus.ACTIVE, thinking_delta="thinking...")),
|
|
_text_step("Hello "),
|
|
_text_step("world"),
|
|
_YieldStep(
|
|
_FakeStep(
|
|
step_type=_StepType.FINISH, status=_StepStatus.DONE, usage_metadata=_FakeUsage()
|
|
)
|
|
),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor(model="gemini-3-pro", api_key="k")
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
|
|
# Two TextChunks prove deltas stream incrementally rather than as one blob —
|
|
# if the executor reverted to a one-shot agent.chat() this would be 1 (or 0).
|
|
texts = [e.text for e in events if isinstance(e, TextChunk)]
|
|
assert texts == ["Hello ", "world"]
|
|
reasoning = [e for e in events if isinstance(e, ReasoningChunk)]
|
|
assert len(reasoning) == 1 and reasoning[0].delta == "thinking..."
|
|
assert reasoning[0].event_type == "reasoning_text"
|
|
|
|
completes = [e for e in events if isinstance(e, TurnComplete)]
|
|
assert len(completes) == 1
|
|
# Final text is the accumulation of the streamed deltas.
|
|
assert completes[0].response == "Hello world"
|
|
# Usage maps the SDK's UsageMetadata field names onto Omnigent's keys and
|
|
# stamps the resolved model so the scaffold can price the turn.
|
|
assert completes[0].usage == {
|
|
"input_tokens": 11,
|
|
"output_tokens": 7,
|
|
"total_tokens": 18,
|
|
"cache_read_input_tokens": 2,
|
|
"model": "gemini-3-pro",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_user_echoed_step_not_surfaced(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A USER-source step (the SDK echoing the prompt) must not leak into the output.
|
|
|
|
Regression guard for a real bug a live turn surfaced: the SDK streams the
|
|
user's own input back as a ``source=USER`` step; mapping its content_delta
|
|
to a TextChunk put the prompt into the assistant's response.
|
|
"""
|
|
script: list[_TurnAction] = [
|
|
_YieldStep(
|
|
_FakeStep(
|
|
step_type=_StepType.TEXT_RESPONSE,
|
|
status=_StepStatus.ACTIVE,
|
|
content_delta="echoed user prompt",
|
|
source=_StepSource.USER,
|
|
target=_StepTarget.USER,
|
|
)
|
|
),
|
|
_text_step("the real reply"), # MODEL->USER by default
|
|
_YieldStep(_FakeStep(step_type=_StepType.FINISH, status=_StepStatus.DONE)),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
texts = [e.text for e in events if isinstance(e, TextChunk)]
|
|
# Only the MODEL->USER reply — the USER-source echo is filtered out.
|
|
assert texts == ["the real reply"]
|
|
completes = [e for e in events if isinstance(e, TurnComplete)]
|
|
assert completes[0].response == "the real reply"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_request_and_completion_paired(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A tool call yields a request, then the PostToolCallHook yields a paired completion."""
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(_FakeToolCall("sys_shell", {"cmd": "ls"}, call_id="t1")),
|
|
_FireToolResult(_FakeToolResult("sys_shell", result={"ok": True}, call_id="t1")),
|
|
_text_step("done"),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
requests = [e for e in events if isinstance(e, ToolCallRequest)]
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
assert len(requests) == 1 and len(completes) == 1
|
|
assert requests[0].name == "sys_shell"
|
|
assert requests[0].args == {"cmd": "ls"}
|
|
assert requests[0].metadata == {"call_id": "t1"}
|
|
# Completion is paired to the request by call_id, carries the real result,
|
|
# and is classified SUCCESS (no error on the ToolResult).
|
|
assert completes[0].metadata == {"call_id": "t1"}
|
|
assert completes[0].name == "sys_shell"
|
|
assert completes[0].result == {"ok": True}
|
|
assert completes[0].status == ToolCallStatus.SUCCESS
|
|
# duration_ms is computed from the recorded request start; >= 0 proves the
|
|
# pending-tool table was populated by the request and read by the hook.
|
|
assert completes[0].duration_ms >= 0.0
|
|
# Request precedes completion in the stream.
|
|
assert events.index(requests[0]) < events.index(completes[0])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_completion_error_status(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A ToolResult carrying an error maps to a ToolCallComplete with ERROR status."""
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(_FakeToolCall("sys_shell", {}, call_id="t1")),
|
|
_FireToolResult(_FakeToolResult("sys_shell", error="permission denied", call_id="t1")),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
assert len(completes) == 1
|
|
# ERROR (not SUCCESS) because the ToolResult.error was set; the message is
|
|
# surfaced so the transcript shows why the tool failed.
|
|
assert completes[0].status == ToolCallStatus.ERROR
|
|
assert completes[0].error == "permission denied"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_result_payload_error_classified_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A ToolResult whose *payload* carries an error (not ToolResult.error) → ERROR."""
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(_FakeToolCall("sys_shell", {}, call_id="t1")),
|
|
# error=None, but the result payload self-describes as an error — this
|
|
# exercises classify_tool_result's payload branch, not the .error path.
|
|
_FireToolResult(_FakeToolResult("sys_shell", result={"error": "boom"}, call_id="t1")),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
assert len(completes) == 1
|
|
assert completes[0].status == ToolCallStatus.ERROR
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_call_without_id_still_completes(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""An id-less tool call still emits one request and one (unpaired) completion."""
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(_FakeToolCall("sys_shell", {"cmd": "ls"}, call_id=None)),
|
|
_FireToolResult(_FakeToolResult("sys_shell", result={"ok": True}, call_id=None)),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
requests = [e for e in events if isinstance(e, ToolCallRequest)]
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
# Exactly one of each: the request gets a synthetic id (so it's still shown);
|
|
# the id-less completion can't pair back, so its metadata is empty and it
|
|
# falls back to the ToolResult's own name. The tool must still "close".
|
|
assert len(requests) == 1
|
|
assert len(completes) == 1
|
|
assert completes[0].name == "sys_shell"
|
|
assert completes[0].metadata == {}
|
|
assert completes[0].duration_ms == 0.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_error_step_completes_without_hook(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""If a TOOL_CALL step errors and the hook never fires, the step closes the tool."""
|
|
call = _FakeToolCall("sys_shell", {"cmd": "ls"}, call_id="t1")
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(call, status=_StepStatus.ACTIVE),
|
|
# No _FireToolResult: simulate the SDK surfacing the tool error outside
|
|
# PostToolCallHook. The terminal TOOL_CALL ERROR step must still close it.
|
|
_tool_call_step(call, status=_StepStatus.ERROR),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
# Without the step-stream fallback the tool would stay "open" (0 completions);
|
|
# the fallback emits exactly one ERROR completion paired by call_id.
|
|
assert len(completes) == 1
|
|
assert completes[0].status == ToolCallStatus.ERROR
|
|
assert completes[0].metadata == {"call_id": "t1"}
|
|
# The turn itself is not failed — a tool error is not a turn-level error.
|
|
assert any(isinstance(e, TurnComplete) for e in events)
|
|
assert not any(isinstance(e, ExecutorError) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_completion_not_double_emitted(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""When both the hook and a terminal step fire, the tool completes exactly once."""
|
|
call = _FakeToolCall("sys_shell", {"cmd": "ls"}, call_id="t1")
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(call, status=_StepStatus.ACTIVE),
|
|
_FireToolResult(_FakeToolResult("sys_shell", result={"ok": True}, call_id="t1")),
|
|
# A trailing DONE step for the same call — the fallback must see it as
|
|
# already-completed (popped by the hook) and NOT emit a second event.
|
|
_tool_call_step(call, status=_StepStatus.DONE),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
completes = [e for e in events if isinstance(e, ToolCallComplete)]
|
|
# 1, not 2: the hook completed it (with the real result) and popped the
|
|
# pending entry, so the DONE-step fallback no-ops.
|
|
assert len(completes) == 1
|
|
assert completes[0].result == {"ok": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_request_deduped_across_steps(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""The same tool-call id appearing in multiple steps yields exactly one request."""
|
|
call = _FakeToolCall("sys_shell", {"cmd": "ls"}, call_id="dup")
|
|
script: list[_TurnAction] = [
|
|
_tool_call_step(call, status=_StepStatus.ACTIVE),
|
|
_tool_call_step(call, status=_StepStatus.DONE),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
requests = [e for e in events if isinstance(e, ToolCallRequest)]
|
|
# 1, not 2: the SDK re-emits the same ToolCall across dispatch/execution
|
|
# step transitions; the seen-id set must suppress the duplicate request.
|
|
assert len(requests) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminal_error_step_yields_executor_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A TERMINAL_ERROR step surfaces an ExecutorError and suppresses TurnComplete."""
|
|
script: list[_TurnAction] = [
|
|
_text_step("partial"),
|
|
_YieldStep(
|
|
_FakeStep(
|
|
step_type=_StepType.FINISH,
|
|
status=_StepStatus.TERMINAL_ERROR,
|
|
error="model exploded",
|
|
)
|
|
),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
errors = [e for e in events if isinstance(e, ExecutorError)]
|
|
assert len(errors) == 1
|
|
assert errors[0].message == "model exploded"
|
|
# TERMINAL_ERROR is non-retryable (a plain ERROR would be retryable).
|
|
assert errors[0].retryable is False
|
|
# No TurnComplete after a turn-level error — the workflow treats it as failed.
|
|
assert not any(isinstance(e, TurnComplete) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_error_step_without_message_still_fails(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""An ERROR step with no error text still yields an ExecutorError (not a silent success)."""
|
|
script: list[_TurnAction] = [
|
|
_YieldStep(_FakeStep(step_type=_StepType.FINISH, status=_StepStatus.ERROR, error="")),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
errors = [e for e in events if isinstance(e, ExecutorError)]
|
|
# An empty error string must not be reported as a successful (empty) turn;
|
|
# the executor substitutes a generic message and a plain ERROR is retryable.
|
|
assert len(errors) == 1
|
|
assert errors[0].message
|
|
assert errors[0].retryable is True
|
|
assert not any(isinstance(e, TurnComplete) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_turn_yields_turn_complete_none(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A turn that streams no text ends as TurnComplete(response=None), not ''."""
|
|
script: list[_TurnAction] = [
|
|
_YieldStep(_FakeStep(step_type=_StepType.FINISH, status=_StepStatus.DONE))
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
completes = [e for e in events if isinstance(e, TurnComplete)]
|
|
assert len(completes) == 1
|
|
# None (not "") is the load-bearing "produced nothing" signal documented on
|
|
# TurnComplete; a regression to "" would change how the empty turn renders.
|
|
assert completes[0].response is None
|
|
assert not any(isinstance(e, TextChunk) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_canceled_step_yields_turn_cancelled(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A CANCELED step surfaces TurnCancelled and no TurnComplete."""
|
|
script: list[_TurnAction] = [
|
|
_text_step("starting"),
|
|
_YieldStep(_FakeStep(status=_StepStatus.CANCELED)),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
assert any(isinstance(e, TurnCancelled) for e in events)
|
|
assert not any(isinstance(e, TurnComplete) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sdk_cancelled_error_yields_turn_cancelled(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""``AntigravityCancelledError`` from the SDK maps to TurnCancelled, not ExecutorError."""
|
|
script: list[_TurnAction] = [_text_step("starting"), _RaiseCancelled()]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
# The cancellation exception is caught specifically (via _cancelled_error_type
|
|
# resolving the SDK's type) and reported as a clean cancel, not a failure.
|
|
assert any(isinstance(e, TurnCancelled) for e in events)
|
|
assert not any(isinstance(e, ExecutorError) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generic_turn_failure_yields_retryable_error(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A non-cancel exception from the SDK becomes a retryable ExecutorError."""
|
|
script: list[_TurnAction] = [_text_step("partial"), _RaiseGeneric("kaboom")]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
|
|
errors = [e for e in events if isinstance(e, ExecutorError)]
|
|
assert len(errors) == 1
|
|
assert "kaboom" in errors[0].message
|
|
# retryable=True (unlike TERMINAL_ERROR) so the workflow picks RetryableLLMError;
|
|
# also distinguishes a generic failure from a clean cancel (TurnCancelled).
|
|
assert errors[0].retryable is True
|
|
assert not any(isinstance(e, (TurnComplete, TurnCancelled)) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_sdk_yields_executor_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
def _raise() -> Any:
|
|
raise ImportError("no google-antigravity")
|
|
|
|
monkeypatch.setattr(ag, "_ensure_antigravity_sdk", _raise)
|
|
executor = AntigravityExecutor()
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "q"}])
|
|
|
|
assert len(events) == 1
|
|
assert isinstance(events[0], ExecutorError)
|
|
assert "google-antigravity" in events[0].message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sys_tools_exposed_as_callables_routing_through_executor(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Omnigent tools become callable SDK tools whose calls hit ``_tool_executor``.
|
|
|
|
This is what lets an Antigravity agent drive Omnigent's sys / sub-agent
|
|
tools under policy (needed to run Polly / Debby).
|
|
"""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("done")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
calls: list[dict[str, Any]] = []
|
|
|
|
async def _fake_tool_executor(name: str, args: dict[str, Any]) -> dict[str, Any]:
|
|
calls.append({"name": name, "args": args})
|
|
return {"ok": True}
|
|
|
|
# The harness ExecutorAdapter assigns this in production; set it directly.
|
|
executor._tool_executor = _fake_tool_executor
|
|
|
|
tool_specs = [
|
|
{
|
|
"name": "sys_shell",
|
|
"description": "Run a shell command",
|
|
"parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}},
|
|
}
|
|
]
|
|
|
|
await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}], tool_specs)
|
|
|
|
sdk_tools = captured["configs"][0].tools
|
|
assert sdk_tools is not None and len(sdk_tools) == 1
|
|
sdk_tool = sdk_tools[0]
|
|
# LocalAgentConfig.tools is list[Callable]; the SDK reads __name__/__doc__.
|
|
assert callable(sdk_tool)
|
|
assert sdk_tool.__name__ == "sys_shell"
|
|
assert sdk_tool.__doc__ == "Run a shell command"
|
|
|
|
# Invoking the callable (kwargs form) routes back through the bridge.
|
|
assert await sdk_tool(cmd="ls") == {"ok": True}
|
|
# Single-dict argument form also works (SDK arg-shape tolerance).
|
|
assert await sdk_tool({"cmd": "pwd"}) == {"ok": True}
|
|
assert calls == [
|
|
{"name": "sys_shell", "args": {"cmd": "ls"}},
|
|
{"name": "sys_shell", "args": {"cmd": "pwd"}},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_tool_executor_means_no_sdk_tools(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Without a tool-executor bridge, no SDK tools are built (agent uses native)."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("done")]])
|
|
executor = AntigravityExecutor() # _tool_executor stays None
|
|
|
|
await _drain(
|
|
executor,
|
|
[{"role": "user", "content": "go"}],
|
|
[{"name": "sys_shell", "description": "", "parameters": {}}],
|
|
)
|
|
|
|
assert captured["configs"][0].tools is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_agent_reused_across_turns_same_session(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A second turn on the same session reuses the cached agent + conversation."""
|
|
captured = _install_fake_sdk(
|
|
monkeypatch, scripts=[[_text_step("one-reply")], [_text_step("two-reply")]]
|
|
)
|
|
executor = AntigravityExecutor()
|
|
|
|
await _drain(executor, [{"role": "user", "content": "one", "session_id": "s1"}])
|
|
await _drain(executor, [{"role": "user", "content": "two", "session_id": "s1"}])
|
|
|
|
# Exactly one agent built across two turns — the signature was unchanged so
|
|
# the cached agent (and its SDK conversation state) was reused.
|
|
assert len(captured["agents"]) == 1
|
|
assert captured["agents"][0].conversation.sends == ["one", "two"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fresh_session_replays_prior_history(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A FRESH agent seeds prior user/assistant turns into its first send().
|
|
|
|
Models a rebuilt/restarted session: the turn arrives with prior history but
|
|
the SDK conversation is brand new (and the SDK has no history-injection
|
|
API). The prior turns must ride into the single send() as a context prefix,
|
|
so the agent doesn't lose them. Without the seeding fix the agent would only
|
|
ever see the latest user text.
|
|
"""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("reply")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
messages = [
|
|
{"role": "user", "content": "what is 2+2?", "session_id": "s1"},
|
|
{"role": "assistant", "content": "4", "session_id": "s1"},
|
|
{"role": "user", "content": "and times 3?", "session_id": "s1"},
|
|
]
|
|
await _drain(executor, messages)
|
|
|
|
# One agent, one send. That send must carry the prior turns AND the latest
|
|
# user text — not just the latest text (the pre-fix behavior).
|
|
sends = captured["agents"][0].conversation.sends
|
|
assert len(sends) == 1
|
|
seeded = sends[0]
|
|
assert "what is 2+2?" in seeded # prior user turn replayed
|
|
assert "assistant: 4" in seeded # prior assistant turn replayed
|
|
assert "and times 3?" in seeded # latest user input still present
|
|
# The latest input is not the whole prompt — proves a prefix was prepended.
|
|
assert seeded != "and times 3?"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rebuilt_session_replays_history_after_signature_change(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A model switch rebuilds the agent; the rebuild re-seeds prior history.
|
|
|
|
A signature change (model/system-prompt/tools) discards the live SDK
|
|
conversation, so the *rebuilt* agent is fresh and must be re-seeded with the
|
|
history it just lost — the same context-loss bug a server restart causes.
|
|
"""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("a")], [_text_step("b")]])
|
|
executor = AntigravityExecutor(model="gemini-3-pro")
|
|
|
|
await _drain(executor, [{"role": "user", "content": "first", "session_id": "s1"}])
|
|
# Second turn carries the accumulated history AND switches model, forcing a
|
|
# rebuild of the (now fresh) agent.
|
|
second = [
|
|
{"role": "user", "content": "first", "session_id": "s1"},
|
|
{"role": "assistant", "content": "answer one", "session_id": "s1"},
|
|
{"role": "user", "content": "second", "session_id": "s1"},
|
|
]
|
|
await _drain(executor, second, config=ExecutorConfig(model="gemini-3-flash"))
|
|
|
|
# Two agents (model changed). The rebuilt agent's first send must replay the
|
|
# prior turns, not just the latest "second".
|
|
assert len(captured["agents"]) == 2
|
|
rebuilt_sends = captured["agents"][1].conversation.sends
|
|
assert len(rebuilt_sends) == 1
|
|
assert "first" in rebuilt_sends[0]
|
|
assert "assistant: answer one" in rebuilt_sends[0]
|
|
assert "second" in rebuilt_sends[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reused_session_does_not_reseed_history(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A REUSED agent must NOT re-seed history (it already holds it).
|
|
|
|
The live SDK conversation accumulated the prior turns itself, so re-seeding
|
|
on reuse would duplicate them. The second turn's send must be exactly the
|
|
latest user text, with no transcript prefix.
|
|
"""
|
|
captured = _install_fake_sdk(
|
|
monkeypatch, scripts=[[_text_step("one-reply")], [_text_step("two-reply")]]
|
|
)
|
|
executor = AntigravityExecutor()
|
|
|
|
await _drain(executor, [{"role": "user", "content": "one", "session_id": "s1"}])
|
|
# Second turn on the same session/signature reuses the agent. It carries the
|
|
# prior turn in its messages, but the reused conversation already has it.
|
|
second = [
|
|
{"role": "user", "content": "one", "session_id": "s1"},
|
|
{"role": "assistant", "content": "one-reply", "session_id": "s1"},
|
|
{"role": "user", "content": "two", "session_id": "s1"},
|
|
]
|
|
await _drain(executor, second)
|
|
|
|
assert len(captured["agents"]) == 1 # reused, not rebuilt
|
|
# The reused turn's send is the bare latest text — no "Conversation so far:"
|
|
# prefix and no duplicated prior turn.
|
|
sends = captured["agents"][0].conversation.sends
|
|
assert sends == ["one", "two"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_usage_observer_notified_on_turn_complete(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""The usage observer is notified with the turn's tokens before TurnComplete.
|
|
|
|
Peers notify in-process usage subscribers on every turn; without the fix,
|
|
antigravity turns fire nothing, so observers see no usage for them.
|
|
"""
|
|
script: list[_TurnAction] = [
|
|
_text_step("hi"),
|
|
_YieldStep(
|
|
_FakeStep(
|
|
step_type=_StepType.FINISH, status=_StepStatus.DONE, usage_metadata=_FakeUsage()
|
|
)
|
|
),
|
|
]
|
|
_install_fake_sdk(monkeypatch, scripts=[script])
|
|
executor = AntigravityExecutor(model="gemini-3-pro")
|
|
|
|
seen: list[dict[str, Any]] = []
|
|
|
|
def _observer(
|
|
*, model: str | None, input_tokens: int, output_tokens: int, total_tokens: int
|
|
) -> None:
|
|
seen.append(
|
|
{
|
|
"model": model,
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"total_tokens": total_tokens,
|
|
}
|
|
)
|
|
|
|
remove = add_observer(_observer)
|
|
try:
|
|
events = await _drain(executor, [{"role": "user", "content": "go", "session_id": "s1"}])
|
|
finally:
|
|
remove()
|
|
|
|
# Exactly one notification, carrying the model and the mapped token counts
|
|
# from the turn's UsageMetadata (_FakeUsage: prompt=11, candidates=7, total=18).
|
|
assert len(seen) == 1
|
|
assert seen[0] == {
|
|
"model": "gemini-3-pro",
|
|
"input_tokens": 11,
|
|
"output_tokens": 7,
|
|
"total_tokens": 18,
|
|
}
|
|
# The notification accompanies a normal TurnComplete.
|
|
assert any(isinstance(e, TurnComplete) for e in events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_switch_rebuilds_agent(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A per-turn model override changes the signature and rebuilds the agent."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("a")], [_text_step("b")]])
|
|
executor = AntigravityExecutor(model="gemini-3-pro")
|
|
|
|
await _drain(executor, [{"role": "user", "content": "one", "session_id": "s1"}])
|
|
await _drain(
|
|
executor,
|
|
[{"role": "user", "content": "two", "session_id": "s1"}],
|
|
config=ExecutorConfig(model="gemini-3-flash"),
|
|
)
|
|
|
|
# Two agents: the model changed (gemini-3-pro -> gemini-3-flash), which is
|
|
# part of the agent signature, so the executor rebuilt rather than reused.
|
|
assert len(captured["agents"]) == 2
|
|
assert captured["configs"][0].model == "gemini-3-pro"
|
|
assert captured["configs"][1].model == "gemini-3-flash"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_model_when_unset(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""With no model on the executor or per-turn config, the built-in default is pinned."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("ok")]])
|
|
executor = AntigravityExecutor() # no model anywhere
|
|
|
|
await _drain(executor, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
|
|
# Pins _ANTIGRAVITY_DEFAULT_MODEL; changing the default must update this.
|
|
assert captured["configs"][0].model == "gemini-3.5-flash"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_system_prompt_change_rebuilds_agent(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""A changed system_prompt is part of the agent signature, forcing a rebuild."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("a")], [_text_step("b")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
await _drain(
|
|
executor, [{"role": "user", "content": "one", "session_id": "s1"}], system_prompt="first"
|
|
)
|
|
await _drain(
|
|
executor, [{"role": "user", "content": "two", "session_id": "s1"}], system_prompt="second"
|
|
)
|
|
|
|
# Two agents: system_prompt is in the (model, system_prompt, tools) signature,
|
|
# so changing it rebuilds. A regression dropping system_prompt would be 1.
|
|
assert len(captured["agents"]) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_api_key_and_vertex_threaded_to_config(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""api_key and Vertex (project/location) reach LocalAgentConfig; base_url never does."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("ok")], [_text_step("ok")]])
|
|
|
|
key_exec = AntigravityExecutor(api_key="gem-key")
|
|
await _drain(key_exec, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
cfg = captured["configs"][0]
|
|
assert cfg.api_key == "gem-key"
|
|
# Vertex left unset on the API-key path.
|
|
assert cfg.vertex is None
|
|
|
|
vertex_exec = AntigravityExecutor(vertex=True, project="my-proj", location="us-central1")
|
|
await _drain(vertex_exec, [{"role": "user", "content": "hi", "session_id": "s2"}])
|
|
vcfg = captured["configs"][1]
|
|
assert vcfg.vertex is True
|
|
assert vcfg.project == "my-proj"
|
|
assert vcfg.location == "us-central1"
|
|
# The SDK config has no base_url field — the executor must never set one.
|
|
assert not hasattr(vcfg, "base_url")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_session_closes_agent(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""close_session() tears down the cached SDK agent for that session."""
|
|
captured = _install_fake_sdk(monkeypatch, scripts=[[_text_step("ok")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
await _drain(executor, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
agent = captured["agents"][0]
|
|
assert agent.closed is False # still open after the turn
|
|
|
|
await executor.close_session("s1")
|
|
# _close_agent awaited the agent's __aexit__, releasing the SDK connection.
|
|
assert agent.closed is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_open_agent_failure_leaves_no_session_state(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A failed agent build registers no dead, agent-less ``_session_states`` row.
|
|
|
|
Otherwise every turn on an un-buildable host (bad creds / missing glibc /
|
|
SDK drift) would accumulate a permanent empty entry that close_session never
|
|
reaps.
|
|
"""
|
|
_install_fake_sdk(monkeypatch, scripts=[[_text_step("ok")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
async def _boom(*_args: Any, **_kwargs: Any) -> Any:
|
|
raise RuntimeError("agent construction failed")
|
|
|
|
monkeypatch.setattr(executor, "_open_agent", _boom)
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
assert any(isinstance(e, ExecutorError) for e in events)
|
|
assert executor._session_states == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_conversation_access_failure_reaps_agent(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""If ``agent.conversation`` raises after ``__aenter__``, the entered agent
|
|
(and its native subprocess) is torn down rather than orphaned, and no
|
|
session state is registered."""
|
|
_install_fake_sdk(monkeypatch, scripts=[[_text_step("ok")]])
|
|
executor = AntigravityExecutor()
|
|
|
|
class _BadConversationAgent:
|
|
def __init__(self) -> None:
|
|
self.closed = False
|
|
|
|
@property
|
|
def conversation(self) -> Any:
|
|
raise RuntimeError("conversation unavailable")
|
|
|
|
async def __aexit__(self, *_args: object) -> None:
|
|
self.closed = True
|
|
|
|
bad_agent = _BadConversationAgent()
|
|
|
|
async def _open(*_args: Any, **_kwargs: Any) -> Any:
|
|
return bad_agent
|
|
|
|
monkeypatch.setattr(executor, "_open_agent", _open)
|
|
|
|
events = await _drain(executor, [{"role": "user", "content": "hi", "session_id": "s1"}])
|
|
assert any(isinstance(e, ExecutorError) for e in events)
|
|
assert bad_agent.closed is True # _close_agent reaped the entered agent
|
|
assert executor._session_states == {}
|
|
|
|
|
|
# ── Interrupt (cancellation) tests — real deterministic sync gates ──────
|
|
|
|
|
|
class _BlockingConversation:
|
|
"""Conversation that streams one delta, then blocks until cancel() releases it.
|
|
|
|
:param raise_on_release: when True, ``receive_steps`` raises the SDK
|
|
cancellation error after the gate opens (the "SDK reports a cancel"
|
|
path); when False it simply ends the stream cleanly (the "cancel ended
|
|
the turn quietly" path that exercises the ``interrupt_requested`` gate).
|
|
"""
|
|
|
|
def __init__(self, gate: asyncio.Event, raise_on_release: bool) -> None:
|
|
self._gate = gate
|
|
self._raise_on_release = raise_on_release
|
|
self.sends: list[str] = []
|
|
self.cancel_called = 0
|
|
|
|
async def send(self, prompt: Any, **_kw: Any) -> None:
|
|
self.sends.append(prompt)
|
|
|
|
async def receive_steps(self) -> Any:
|
|
yield _FakeStep(
|
|
step_type=_StepType.TEXT_RESPONSE, status=_StepStatus.ACTIVE, content_delta="streaming"
|
|
)
|
|
await self._gate.wait() # blocked until cancel() releases us
|
|
if self._raise_on_release:
|
|
raise _AntigravityCancelledError("cancelled")
|
|
|
|
async def cancel(self) -> None:
|
|
self.cancel_called += 1
|
|
self._gate.set()
|
|
|
|
|
|
def _install_blocking_sdk(
|
|
monkeypatch: pytest.MonkeyPatch, gate: asyncio.Event, *, raise_on_release: bool
|
|
) -> dict[str, Any]:
|
|
"""Install a fake SDK whose conversation blocks mid-turn until cancelled."""
|
|
captured: dict[str, Any] = {}
|
|
|
|
class _BlockingAgent:
|
|
def __init__(self, config: Any) -> None:
|
|
self.config = config
|
|
self._conversation = _BlockingConversation(gate, raise_on_release)
|
|
captured["conversation"] = self._conversation
|
|
|
|
@property
|
|
def conversation(self) -> _BlockingConversation:
|
|
return self._conversation
|
|
|
|
async def __aenter__(self) -> _BlockingAgent:
|
|
return self
|
|
|
|
async def __aexit__(self, *_a: object) -> None:
|
|
return None
|
|
|
|
class _FakeHooks:
|
|
PostToolCallHook = _FakePostToolCallHook
|
|
|
|
class _FakeTypes:
|
|
AntigravityCancelledError = _AntigravityCancelledError
|
|
|
|
class _FakeModule:
|
|
LocalAgentConfig = _FakeLocalAgentConfig
|
|
hooks = _FakeHooks
|
|
types = _FakeTypes
|
|
|
|
@staticmethod
|
|
def Agent(config: Any) -> _BlockingAgent:
|
|
return _BlockingAgent(config)
|
|
|
|
monkeypatch.setattr(ag, "_ensure_antigravity_sdk", lambda: _FakeModule())
|
|
return captured
|
|
|
|
|
|
async def _drive_until_first_text(
|
|
executor: AntigravityExecutor, collected: list[Any], first_text: asyncio.Event
|
|
) -> None:
|
|
async for event in executor.run_turn(
|
|
[{"role": "user", "content": "go", "session_id": "s1"}], tools=[], system_prompt="sys"
|
|
):
|
|
collected.append(event)
|
|
if isinstance(event, TextChunk):
|
|
first_text.set()
|
|
|
|
|
|
@pytest.mark.parametrize("raise_on_release", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_interrupt_session_cancels_running_turn(
|
|
monkeypatch: pytest.MonkeyPatch, raise_on_release: bool
|
|
) -> None:
|
|
"""interrupt_session cancels an in-flight turn -> TurnCancelled, no TurnComplete.
|
|
|
|
Deterministic race: the conversation blocks inside ``receive_steps`` after
|
|
streaming one delta; we interrupt only after observing that delta (so the
|
|
turn is provably mid-flight). The two parametrized cases cover both ways the
|
|
SDK can react to ``cancel()``: raising ``AntigravityCancelledError``
|
|
(raise_on_release=True), or ending the stream cleanly so the
|
|
``interrupt_requested`` gate in run_turn must convert it to TurnCancelled
|
|
(raise_on_release=False).
|
|
"""
|
|
gate = asyncio.Event()
|
|
first_text = asyncio.Event()
|
|
captured = _install_blocking_sdk(monkeypatch, gate, raise_on_release=raise_on_release)
|
|
executor = AntigravityExecutor()
|
|
|
|
collected: list[Any] = []
|
|
task = asyncio.create_task(_drive_until_first_text(executor, collected, first_text))
|
|
# Wait until the turn has streamed its first delta (provably mid-flight)
|
|
# before interrupting — this is the deterministic race window.
|
|
await asyncio.wait_for(first_text.wait(), timeout=5)
|
|
|
|
interrupted = await executor.interrupt_session("s1")
|
|
# Assert the cancel landed BEFORE awaiting the task: a broken interrupt that
|
|
# skips conversation.cancel() leaves the producer parked on the gate forever,
|
|
# so checking here fails crisply ("cancel never called") instead of as an
|
|
# opaque 5s task timeout below.
|
|
assert interrupted is True # a live conversation was found and asked to cancel
|
|
assert captured["conversation"].cancel_called == 1 # cancel reached the SDK boundary
|
|
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
# Either path must surface a clean cancel and never a TurnComplete.
|
|
assert any(isinstance(e, TurnCancelled) for e in collected)
|
|
assert not any(isinstance(e, TurnComplete) for e in collected)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_interrupt_session_unknown_returns_false(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""interrupt_session on a session with no open conversation returns False."""
|
|
_install_fake_sdk(monkeypatch, scripts=[])
|
|
executor = AntigravityExecutor()
|
|
assert await executor.interrupt_session("never-started") is False
|
|
|
|
|
|
class _RebuildConversation:
|
|
"""Conversation that blocks-then-cancels on turn 1, then runs clean on turn 2.
|
|
|
|
The first conversation (``blocking=True``) streams one delta and parks on a
|
|
gate until ``cancel()`` releases it, then raises the SDK cancellation error
|
|
— modelling an interrupted in-flight turn. The second conversation
|
|
(``blocking=False``) is the fresh one a post-interrupt rebuild must open and
|
|
runs a normal turn to completion. Each records its own ``sends`` so a test
|
|
can prove the second turn went to the *new* conversation, not the cancelled
|
|
one.
|
|
"""
|
|
|
|
def __init__(self, gate: asyncio.Event, *, blocking: bool) -> None:
|
|
self._gate = gate
|
|
self._blocking = blocking
|
|
self.sends: list[str] = []
|
|
self.cancel_called = 0
|
|
|
|
async def send(self, prompt: Any, **_kw: Any) -> None:
|
|
self.sends.append(prompt)
|
|
|
|
async def receive_steps(self) -> Any:
|
|
if self._blocking:
|
|
yield _FakeStep(
|
|
step_type=_StepType.TEXT_RESPONSE,
|
|
status=_StepStatus.ACTIVE,
|
|
content_delta="streaming",
|
|
)
|
|
await self._gate.wait() # parked until cancel() releases us
|
|
raise _AntigravityCancelledError("cancelled")
|
|
yield _FakeStep(
|
|
step_type=_StepType.TEXT_RESPONSE,
|
|
status=_StepStatus.ACTIVE,
|
|
content_delta="second-reply",
|
|
)
|
|
yield _FakeStep(step_type=_StepType.FINISH, status=_StepStatus.DONE)
|
|
|
|
async def cancel(self) -> None:
|
|
self.cancel_called += 1
|
|
self._gate.set()
|
|
|
|
|
|
def _install_rebuild_sdk(monkeypatch: pytest.MonkeyPatch, gate: asyncio.Event) -> dict[str, Any]:
|
|
"""Install a fake SDK: first agent blocks-then-cancels, later agents run clean."""
|
|
captured: dict[str, Any] = {"agents": [], "conversations": []}
|
|
|
|
class _RebuildAgent:
|
|
def __init__(self, config: Any, *, blocking: bool) -> None:
|
|
self.config = config
|
|
self._conversation = _RebuildConversation(gate, blocking=blocking)
|
|
self.closed = False
|
|
captured["conversations"].append(self._conversation)
|
|
|
|
@property
|
|
def conversation(self) -> _RebuildConversation:
|
|
return self._conversation
|
|
|
|
async def __aenter__(self) -> _RebuildAgent:
|
|
return self
|
|
|
|
async def __aexit__(self, *_a: object) -> None:
|
|
self.closed = True
|
|
|
|
class _FakeHooks:
|
|
PostToolCallHook = _FakePostToolCallHook
|
|
|
|
class _FakeTypes:
|
|
AntigravityCancelledError = _AntigravityCancelledError
|
|
|
|
class _FakeModule:
|
|
LocalAgentConfig = _FakeLocalAgentConfig
|
|
hooks = _FakeHooks
|
|
types = _FakeTypes
|
|
|
|
@staticmethod
|
|
def Agent(config: Any) -> _RebuildAgent:
|
|
# Only the very first agent blocks (the turn we interrupt); the
|
|
# rebuilt agent runs a normal turn.
|
|
agent = _RebuildAgent(config, blocking=not captured["agents"])
|
|
captured["agents"].append(agent)
|
|
return agent
|
|
|
|
monkeypatch.setattr(ag, "_ensure_antigravity_sdk", lambda: _FakeModule())
|
|
return captured
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_next_turn_rebuilds_after_interrupt(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""After an interrupt, the next turn rebuilds rather than reusing the cancelled convo.
|
|
|
|
Regression guard for F12: ``interrupt_session`` -> ``conversation.cancel()``
|
|
left the cancelled SDK conversation cached, so the next turn resumed from
|
|
aborted state. The fix invalidates the cached agent signature on interrupt
|
|
so the next ``run_turn`` closes the stale agent and opens a fresh
|
|
agent + conversation. Same model / system_prompt / tools across both turns,
|
|
so without the fix the agent would be reused (1 agent) and the second send
|
|
would land on the cancelled conversation.
|
|
"""
|
|
gate = asyncio.Event()
|
|
first_text = asyncio.Event()
|
|
captured = _install_rebuild_sdk(monkeypatch, gate)
|
|
executor = AntigravityExecutor()
|
|
|
|
# Turn 1: drive until the first delta (provably mid-flight), then interrupt.
|
|
collected: list[Any] = []
|
|
task = asyncio.create_task(_drive_until_first_text(executor, collected, first_text))
|
|
await asyncio.wait_for(first_text.wait(), timeout=5)
|
|
assert await executor.interrupt_session("s1") is True
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
# The interrupt produced a clean cancel (not an error) and cancelled the
|
|
# live conversation.
|
|
assert any(isinstance(e, TurnCancelled) for e in collected)
|
|
assert not any(isinstance(e, ExecutorError) for e in collected)
|
|
cancelled_conv = captured["conversations"][0]
|
|
assert cancelled_conv.cancel_called == 1
|
|
|
|
# Turn 2 on the SAME session/signature must NOT reuse the cancelled agent.
|
|
events = await _drain(executor, [{"role": "user", "content": "next", "session_id": "s1"}])
|
|
|
|
# A fresh agent was built (2 total) and the stale one was torn down — the
|
|
# exact rebuild path a signature change uses. Reuse (the bug) would be 1.
|
|
assert len(captured["agents"]) == 2
|
|
assert captured["agents"][0].closed is True
|
|
# The second turn went to the NEW conversation; the cancelled one never saw
|
|
# the follow-up prompt (no stale-state reuse).
|
|
new_conv = captured["conversations"][1]
|
|
assert new_conv.sends == ["next"]
|
|
assert "next" not in cancelled_conv.sends
|
|
# And the fresh turn completed normally.
|
|
completes = [e for e in events if isinstance(e, TurnComplete)]
|
|
assert len(completes) == 1
|
|
assert completes[0].response == "second-reply"
|