Programmatic Control

pip install ag-ui-antigravity google-antigravity

def api_key() -> str | None:
    """The Gemini API key.

    ``GEMINI_API_KEY`` is the only name the Antigravity SDK reads. The showcase
    fleet exports ``GOOGLE_API_KEY`` for its Gemini integrations, so accept
    that too rather than fail on the first model call.
    """
    return os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")


def gemini_base_url() -> str | None:
    """The Gemini-compatible server to call instead of Google's API.

    Showcase compose sets ``GOOGLE_GEMINI_BASE_URL=http://aimock:4010``, the
    same variable the google-adk integration reads.
    """
    return os.environ.get("GOOGLE_GEMINI_BASE_URL", "").rstrip("/") or None


def aimock_headers() -> dict[str, str]:
    """Headers for a model call: the package's fixture context plus the
    current request's mock and diagnostic headers.

    ``X-AIMock-Context`` is fixed: it selects this package's fixtures whatever
    the request says. The rest comes from the inbound request (captured by
    ``HeaderForwardingHTTPMiddleware``) and is forwarded only when present, so
    ordinary demo traffic is never forced into strict mode.
    """
    forwarded = {
        k: v
        for k, v in get_forwarded_headers().items()
        if k.startswith(_FORWARDED_PREFIXES)
    }
    return {**forwarded, "x-aimock-context": SLUG}


def endpoint() -> GeminiAPIEndpoint | None:
    """Where every model call goes: aimock under compose, Google's API otherwise.

    Passed to the adapter as a callable, so it runs each time a session is
    built, inside that conversation's first request. The harness makes the
    model calls itself, so headers can only reach them on the endpoint: the
    request's ``X-AIMock-Strict`` and ``x-test-id``/``x-diag-*`` ride along
    from here (PNI-576). They are fixed per conversation, because the SDK
    sets a conversation's model configuration when it starts.
    """
    base_url = gemini_base_url()
    if base_url is None:
        return None
    return GeminiAPIEndpoint(
        base_url=base_url,
        api_key=api_key(),
        http_headers=aimock_headers(),
    )


def chat_only_capabilities() -> CapabilitiesConfig:
    """Only ``finish`` stays enabled.

    ``search_web`` returns an empty summary without Google credentials and the
    model retries it forever; ``ask_question`` would park on an interrupt the
    CopilotKit demos never answer. ``finish`` is how the harness ends a turn.
    """
    return CapabilitiesConfig(
        enabled_tools=[BuiltinTools.FINISH], enable_subagents=False
    )


_DIRS_READY = False


def _ensure_dirs() -> None:
    """Creates the workspace / save_dir once, on first agent construction.

    Deliberately NOT done at import time: the defaults live under ``/data``,
    which only exists inside the container image, so an import-time
    ``os.makedirs`` blows up on a developer machine (read-only ``/`` on macOS)
    and takes every ``import agent_server`` — including the package's pytest
    suite — down with it.
    """
    global _DIRS_READY
    if _DIRS_READY:
        return
    os.makedirs(WORKSPACE, exist_ok=True)
    os.makedirs(SAVE_DIR, exist_ok=True)
    _DIRS_READY = True


_POOL = None


def shared_pool():
    """One HarnessPool for the whole server: one Go process, not one per agent."""
    global _POOL
    if _POOL is None:
        from ag_ui_antigravity.harness_pool import HarnessPool

        _POOL = HarnessPool()
    return _POOL


def build(**kwargs):
    """Creates an AntigravityAgent with the package defaults applied."""
    from ag_ui_antigravity import AntigravityAgent

    _ensure_dirs()

    defaults = dict(
        model=MODEL,
        api_key=api_key(),
        endpoint=endpoint,
        workspaces=[WORKSPACE],
        save_dir=SAVE_DIR,
        harness_pool=shared_pool(),
        capabilities=chat_only_capabilities(),
        session_timeout_seconds=SESSION_TIMEOUT_SECONDS,
        max_sessions=MAX_SESSIONS,
        max_tool_calls_per_turn=MAX_TOOL_CALLS_PER_TURN,
    )
    defaults.update(kwargs)
    return AntigravityAgent(**defaults)