Files
omnigent-ai--omnigent/tests/e2e/omnigent/test_antigravity_streaming_e2e.py
2026-07-13 13:12:00 +08:00

474 lines
21 KiB
Python

"""Per-harness live characterization — antigravity (Gemini) SDK streaming fidelity.
Drives the ``antigravity`` harness end-to-end through the production
``/v1/sessions`` flow — register an inline Omnigent agent
(``executor.harness: antigravity``) → create + runner-bind a session →
``POST /v1/sessions/{id}/events`` → poll the session to idle — then asserts over
the *persisted transcript items* (``GET /v1/sessions/{id}/items``), NOT stdout
chrome. This exercises the full producer/consumer path in
:mod:`omnigent.inner.antigravity_executor`: ``conversation.send`` →
``receive_steps()`` → the MODEL→USER ``content_delta`` mapping onto
:class:`TextChunk` → the terminal :class:`TurnComplete` whose accumulated text
the runner persists as an assistant ``message`` item.
Scope is deliberately narrow: **streaming / output fidelity** of stable behavior
already on ``main`` (this branch is cut from ``main``). It covers three things a
streamed text channel can silently corrupt:
1. **Long output** is not truncated or duplicated — the per-delta accumulation
in ``run_turn`` (``final_text_parts.append`` then ``"".join(...)``) must
neither drop a tail nor replay a block.
2. **Unicode fidelity** — emoji / CJK / accented text round-trips byte-exact
through ``content_delta`` concatenation and JSON (de)serialization, with no
mojibake (``é``), replacement char (````), or double-encoding.
3. **Reasoning-heavy task** — a step-by-step prompt still yields a coherent,
non-empty final answer (we assert the answer sentinel, not the reasoning
internals, which the harness maps separately as :class:`ReasoningChunk`).
**Gating (mirrors the cursor per-harness test):** the antigravity harness is
Gemini-native — it authenticates with a Gemini / Antigravity API key and has NO
Databricks-gateway path — so this test SKIPS (rather than fails) when its
prerequisites are absent, keeping the e2e shards green where no Gemini key is
provisioned. It runs for real wherever the prerequisites are present:
- ``google.antigravity`` (the ``google-antigravity`` package) is importable, and
- a Gemini key is configured (``antigravity_api_key_configured()``) or present
ambiently as ``GEMINI_API_KEY`` / ``ANTIGRAVITY_API_KEY``.
.. note::
**glibc caveat.** The Antigravity SDK launches a bundled native
``localharness`` binary linked against a recent glibc (it needs
``GLIBC_ABI_DT_RELR``), so a turn only *completes* on a host with
glibc ≳ 2.36. On an older host (e.g. glibc 2.31) the turn fails at native
setup with ``version 'GLIBC_ABI_DT_RELR' not found`` and the session goes
``failed`` — that is an environment gap, not a regression in the code under
test. The gate above mirrors cursor's (package + key) and deliberately does
NOT probe glibc; on a glibc-2.31 dev box point the SDK at a loader-shim via
``ANTIGRAVITY_HARNESS_PATH`` (dev-only) to run the untouched binary under a
newer loader. CI runs this on a glibc-≥2.36 host.
**What breaks if this fails (with prerequisites present):**
- ``AntigravityExecutor`` regresses (the ``Step`` → :class:`TextChunk` mapping,
the per-turn text accumulation in ``run_turn``, or the MODEL→USER source/target
filter that keeps the echoed prompt out of the assistant reply).
- ``_build_antigravity_spawn_env`` stops resolving the Gemini key into
``HARNESS_ANTIGRAVITY_API_KEY`` for a spec with no ``executor.auth``.
- The sessions dispatch path stops persisting the streamed assistant text into
``conversation_items`` faithfully (truncation, duplication, or re-encoding).
"""
from __future__ import annotations
import importlib.util
import os
from pathlib import Path
from typing import Any
import httpx
import pytest
from omnigent.onboarding.antigravity_auth import antigravity_api_key_configured
from tests.e2e.conftest import (
create_runner_bound_session,
poll_session_until_terminal,
send_user_message_to_session,
upload_agent,
)
# Gemini id that runs on a plain AI-Studio key. ``gemini-3-pro`` 404s without
# Pro access, so the suite pins a Flash model; left un-rewritten (the harness is
# Gemini-native — a Databricks model map / profile stamp would break it).
_MODEL = "gemini-3.5-flash"
_HARNESS = "antigravity"
# A turn cold-starts the native binary, round-trips to Google's backend, and
# (for the long-output case) streams a multi-paragraph reply. 240s matches the
# headroom the other live session smokes allow without letting a hung run pin
# the suite forever.
_RUN_TIMEOUT_SEC = 240.0
def _antigravity_skip_reason() -> str | None:
"""Return why the antigravity harness can't run for real, or ``None``.
Mirrors the cursor per-harness gate's spirit: the harness is Gemini-native
(no Databricks-gateway fallback), so when its prerequisites are absent the
suite SKIPS rather than fails — CI shards without a provisioned Gemini key
stay green, and the test runs for real wherever a key is present.
Two prerequisites, both probed against the *running* interpreter (the e2e
server + runner the ``live_server`` fixture spawns with ``sys.executable``
are this same interpreter):
1. ``google.antigravity`` is importable (the ``google-antigravity`` extra).
2. A Gemini key is configured (the dedicated ``antigravity:`` config block)
or present ambiently as ``GEMINI_API_KEY`` / ``ANTIGRAVITY_API_KEY`` —
which ``_build_antigravity_spawn_env`` threads into the harness as
``HARNESS_ANTIGRAVITY_API_KEY`` for a spec with no ``executor.auth``.
This gate intentionally does NOT probe glibc (see the module docstring's
note): on a glibc-<2.36 host the prerequisites read as present and the turn
fails at native-binary setup — surfaced by the per-turn terminal-status
assertion as an environment gap, not masked as a silent skip.
:returns: A skip-reason string when a prerequisite is missing, else ``None``.
"""
if importlib.util.find_spec("google.antigravity") is None:
return (
"antigravity prerequisite missing: the 'google-antigravity' package "
"is not importable. Install it with: pip install 'omnigent[antigravity]'."
)
key_present = antigravity_api_key_configured() or bool(
os.environ.get("GEMINI_API_KEY") or os.environ.get("ANTIGRAVITY_API_KEY")
)
if not key_present:
return (
"antigravity prerequisite missing: no Gemini key configured. The "
"Antigravity SDK requires one (there is no login flow); configure it "
"via 'omni setup' (Antigravity) or export GEMINI_API_KEY / "
"ANTIGRAVITY_API_KEY. Skipped (not failed) so CI shards without a key "
"stay green."
)
return None
# Skip the whole module at COLLECTION time when the harness can't run. This must
# precede fixture setup: the live-server stack here is session-scoped (and the
# e2e suite's own ``--llm-api-key`` gate lives on a session fixture), and
# higher-scoped fixtures resolve before any function-scoped skip-guard could
# fire — so a function-level gate would surface as a fixture ERROR instead of a
# clean SKIP when a prerequisite is absent. ``allow_module_level`` runs the gate
# before any of that, yielding a clean skip with no server/runner spawned.
_SKIP_REASON = _antigravity_skip_reason()
if _SKIP_REASON is not None:
pytest.skip(_SKIP_REASON, allow_module_level=True)
def _write_antigravity_agent_yaml(tmp_path: Path, *, prompt: str) -> Path:
"""Write a minimal single-file ``harness: antigravity`` Omnigent bundle.
Deliberately carries NO ``executor.auth`` and NO Databricks ``profile``: the
harness resolves its Gemini key Gemini-natively (configured ``antigravity:``
block, then ambient ``GEMINI_API_KEY`` / ``ANTIGRAVITY_API_KEY``) — a global
``auth:`` / profile would be the OpenAI-gateway key the SDK can't use. The
bundle is uploaded with ``rewrite_model_for_databricks=False`` so the Gemini
model id and ``antigravity`` harness pass through unmangled.
:param tmp_path: Per-test temp dir to materialize the bundle directory in.
:param prompt: The agent's system prompt (per-test, so each scenario can
steer the model toward a deterministic, assertable shape).
:returns: The bundle directory to hand to :func:`upload_agent`.
"""
agent_dir = tmp_path / "antigravity-streaming-agent"
agent_dir.mkdir()
(agent_dir / "antigravity-streaming-agent.yaml").write_text(
"\n".join(
[
"name: antigravity-streaming-agent",
"description: Live streaming/output-fidelity probe for the antigravity harness.",
"executor:",
f" harness: {_HARNESS}",
f" model: {_MODEL}",
"prompt: |",
*(f" {line}" for line in prompt.splitlines()),
"",
]
),
encoding="utf-8",
)
return agent_dir
def _run_one_shot(
http_client: httpx.Client,
*,
runner_id: str,
agent_name: str,
user_text: str,
) -> str:
"""Run one antigravity turn through a fresh runner-bound session.
Creates a session, posts *user_text*, and polls to terminal — asserting the
turn completed (a ``failed`` status surfaces the error, e.g. the glibc
native-binary failure on an old host, instead of masquerading as empty
output).
:param http_client: HTTP client pointed at the live server.
:param runner_id: The live runner id to bind the session to.
:param agent_name: The registered antigravity agent's name.
:param user_text: The user prompt for this turn.
:returns: The session id (so the caller can read ``/items``).
"""
session_id = create_runner_bound_session(
http_client, agent_name=agent_name, runner_id=runner_id
)
response_id = send_user_message_to_session(
http_client, session_id=session_id, content=user_text
)
body = poll_session_until_terminal(
http_client,
session_id=session_id,
response_id=response_id,
timeout=_RUN_TIMEOUT_SEC,
)
assert body["status"] == "completed", (
f"antigravity turn did not complete: status={body['status']!r}, "
f"error={body.get('error')!r}. A 'failed' status with a "
f"GLIBC_ABI_DT_RELR error means the host's glibc is < 2.36 (the SDK's "
f"native binary can't load) — an environment gap, not a code "
f"regression. output={body.get('output')!r}"
)
return session_id
def _assistant_text_from_items(http_client: httpx.Client, session_id: str) -> str:
"""Concatenate assistant ``output_text`` from ``GET /v1/sessions/{id}/items``.
Reads the persisted transcript (the durable record the production session
flow commits), NOT stdout, so the assertions characterize what the harness
actually streamed and stored. The items endpoint returns the flat
Responses-style item shape — assistant text lives in ``message`` items with
``role == "assistant"`` whose ``content`` carries ``output_text`` blocks
(see ``omnigent/server/API.md`` § List Conversation Items).
:param http_client: HTTP client pointed at the live server.
:param session_id: The session/conversation id to read.
:returns: The assistant text, ``"\\n"``-joined across any assistant messages
in the turn (empty string if none).
"""
resp = http_client.get(f"/v1/sessions/{session_id}/items", params={"limit": 1000})
resp.raise_for_status()
items: list[dict[str, Any]] = resp.json()["data"]
parts: list[str] = []
for item in items:
if item.get("type") != "message" or item.get("role") != "assistant":
continue
for block in item.get("content", []):
if block.get("type") == "output_text":
text = block.get("text")
if text:
parts.append(text)
return "\n".join(parts)
def _max_repeated_run(text: str, *, block: int) -> int:
"""Return the largest count of any *block*-char substring repeated in *text*.
A cheap duplication detector: a streamed reply that double-appended a chunk
shows up as a long substring occurring far more than the prose would warrant.
Samples non-overlapping windows so it stays linear on long replies.
:param text: The text to scan.
:param block: Window size in characters; large enough that natural prose
rarely repeats a window verbatim.
:returns: The max occurrence count of any sampled window (1 when nothing
repeats; >= 2 indicates a verbatim repeat of a *block*-sized span).
"""
if len(text) < block:
return 1
counts: dict[str, int] = {}
best = 1
for start in range(0, len(text) - block + 1, block):
window = text[start : start + block]
counts[window] = counts.get(window, 0) + 1
best = max(best, counts[window])
return best
# ─── Tests ───────────────────────────────────────────────────
def test_antigravity_long_output_not_truncated_or_duplicated(
http_client: httpx.Client,
live_runner_id: str,
tmp_path: Path,
) -> None:
"""A long multi-paragraph reply streams in full — not truncated, not duplicated.
The model is steered to emit six numbered sections, each ending in a unique
``[SECT-k-END]`` sentinel. Asserting every sentinel survives proves nothing
was dropped from the head or tail of the streamed text; a length floor
proves the reply is genuinely substantial; and a verbatim-repeat check
proves the per-delta accumulation didn't replay a block. Together these
characterize ``run_turn``'s ``final_text_parts`` accumulation over a long
``content_delta`` stream.
"""
sentinels = [f"[SECT-{k}-END]" for k in range(1, 7)]
section_lines = "\n".join(
f"- Section {k}: at least four full sentences, then end the section with "
f"the exact marker {sent}"
for k, sent in enumerate(sentinels, start=1)
)
agent_name = upload_agent(
http_client,
_write_antigravity_agent_yaml(
tmp_path,
prompt=(
"You are a thorough technical writer. When asked for a multi-section "
"explanation, write every requested section in full, each as its own "
"paragraph, and place each section's exact end-marker on its own line. "
"Never abbreviate, summarize, or skip a section."
),
),
rewrite_model_for_databricks=False,
)
session_id = _run_one_shot(
http_client,
runner_id=live_runner_id,
agent_name=agent_name,
user_text=(
"Write a detailed explanation of how a CPU cache hierarchy works, "
"organized into exactly these six sections (keep them in order):\n"
f"{section_lines}\n"
"Write each section fully before moving to the next."
),
)
text = _assistant_text_from_items(http_client, session_id)
# Not truncated: every section's end-marker survived the stream. A dropped
# tail (the classic truncation bug) loses the later markers first.
missing = [s for s in sentinels if s not in text]
assert not missing, (
f"long-output reply is missing section markers {missing} — the streamed "
f"text was truncated before the end. Got {len(text)} chars:\n{text!r}"
)
# Substantial: six full sections clear this floor comfortably; an empty /
# one-line reply (a streaming or accumulation regression) does not.
assert len(text) >= 800, (
f"long-output reply is only {len(text)} chars; expected a substantial "
f"multi-section answer (>= 800). Reply:\n{text!r}"
)
# Not duplicated: no 80-char window recurs verbatim. A double-appended delta
# block would push some window's count to >= 2.
repeats = _max_repeated_run(text, block=80)
assert repeats < 2, (
f"long-output reply repeats an 80-char block {repeats}x — the stream "
f"likely double-appended a chunk. Reply:\n{text!r}"
)
def test_antigravity_unicode_fidelity_round_trips(
http_client: httpx.Client,
live_runner_id: str,
tmp_path: Path,
) -> None:
"""Emoji + CJK + accented text round-trips byte-exact into the transcript.
The model is asked to echo a fixed payload verbatim on its own line. We
assert each non-ASCII fragment appears EXACTLY in the persisted item text —
catching mojibake (``café`` → ``café``), the replacement char ``\\ufffd``,
and double-encoding, any of which would corrupt the ``content_delta``
concatenation or the JSON (de)serialization on the way into
``conversation_items``.
"""
# One emoji (incl. a ZWJ sequence + skin-tone modifier), CJK, and accents.
fragments = ["café", "naïve façade", "日本語のテキスト", "中文字符", "🚀", "👩‍🚀", "👍🏽"]
payload = " | ".join(fragments)
agent_name = upload_agent(
http_client,
_write_antigravity_agent_yaml(
tmp_path,
prompt=(
"You echo text exactly as given, preserving every character including "
"emoji, CJK, and accents. Never transliterate, escape, or normalize."
),
),
rewrite_model_for_databricks=False,
)
session_id = _run_one_shot(
http_client,
runner_id=live_runner_id,
agent_name=agent_name,
user_text=(
"Repeat the following text back to me EXACTLY, character for character, "
"on a single line, with no additional commentary, quotes, or code "
f"fences:\n{payload}"
),
)
text = _assistant_text_from_items(http_client, session_id)
# No corruption signatures anywhere in the persisted text.
assert "" not in text, (
f"transcript contains the Unicode replacement char (\\ufffd) — bytes were "
f"lost or mis-decoded in the stream. Reply:\n{text!r}"
)
# ``Ã`` is the tell-tale lead byte of UTF-8 mis-decoded as Latin-1 (e.g.
# ``é`` → ``é``); none of our intended fragments contain it.
assert "Ã" not in text and "â€" not in text, (
f"transcript shows mojibake / double-encoding (e.g. 'é' for 'é') — a "
f"decode/encode mismatch corrupted the non-ASCII text. Reply:\n{text!r}"
)
# Each intended fragment must survive verbatim.
missing = [frag for frag in fragments if frag not in text]
assert not missing, (
f"these unicode fragments did not round-trip into the transcript: "
f"{missing}. The harness must preserve emoji/CJK/accents byte-exact. "
f"Reply:\n{text!r}"
)
def test_antigravity_reasoning_heavy_task_final_answer(
http_client: httpx.Client,
live_runner_id: str,
tmp_path: Path,
) -> None:
"""A step-by-step task yields a coherent, non-empty final answer.
A small arithmetic word problem elicits multi-step reasoning; the model is
told to end with ``FINAL ANSWER: <n>``. We assert the answer sentinel with
the correct value is present and the reply is non-trivial — characterizing
that the harness streams a usable final answer for a reasoning-heavy turn.
We deliberately do NOT assert on the reasoning internals (the harness maps
``thinking_delta`` onto :class:`ReasoningChunk` separately, and whether/how
much reasoning surfaces is model- and config-dependent — over-specifying it
would make the test flaky without testing output fidelity).
"""
agent_name = upload_agent(
http_client,
_write_antigravity_agent_yaml(
tmp_path,
prompt=(
"You are a careful problem solver. Work through the problem step by "
"step, then state the result on a final line in the exact form "
"'FINAL ANSWER: <number>'."
),
),
rewrite_model_for_databricks=False,
)
# 3 crates * 4 boxes * 6 widgets = 72; +5 spares = 77. A unique integer the
# model is very unlikely to emit except as the correct result.
session_id = _run_one_shot(
http_client,
runner_id=live_runner_id,
agent_name=agent_name,
user_text=(
"A warehouse has 3 crates. Each crate holds 4 boxes, and each box holds "
"6 widgets. There are also 5 loose spare widgets on a shelf. How many "
"widgets are in the warehouse in total? Show your reasoning step by "
"step, then end with a line 'FINAL ANSWER: <number>'."
),
)
text = _assistant_text_from_items(http_client, session_id)
# Coherent + non-empty: a reasoning-heavy reply is more than a bare token.
assert len(text.strip()) >= 20, (
f"reasoning-task reply is too short to be a coherent answer "
f"({len(text.strip())} chars): {text!r}"
)
# The final answer surfaced with the correct value. ``77`` (not ``72``) also
# confirms the model accounted for the spares — i.e. the streamed reply
# carries the genuine end-of-reasoning result, not a truncated mid-step.
assert "FINAL ANSWER" in text, (
f"reply is missing the requested 'FINAL ANSWER:' line — the final-answer "
f"text didn't stream into the transcript. Reply:\n{text!r}"
)
assert "77" in text, (
f"reply's final answer is not 77 (3*4*6 + 5); the reasoning result is "
f"wrong or its tail was dropped from the stream. Reply:\n{text!r}"
)