diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index 4ce0c020..bd2ecd93 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -25,7 +25,6 @@ import pathlib import re import sys -from typing import TypedDict import grpc from grpc_health.v1 import health, health_pb2, health_pb2_grpc from google.protobuf.struct_pb2 import Struct @@ -74,18 +73,6 @@ def _validate_conversation_id(conversation_id: str) -> None: ) -class VertexKwargs(TypedDict, total=False): - """Typed subset of LocalAgentConfig kwargs needed to enable Vertex AI. - - `total=False` so {} is a valid value (returned when env does not request - Vertex). When populated, all three keys are present. - """ - - vertex: bool - project: str - location: str - - def _env_use_vertex() -> bool: """True if env requests the Vertex AI backend (vs. AI Studio API key).""" return os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ( @@ -94,56 +81,17 @@ def _env_use_vertex() -> bool: ) or os.environ.get("GOOGLE_GENAI_USE_ENTERPRISE", "").lower() in ("true", "1") -def _vertex_kwargs_from_env() -> VertexKwargs: - """Returns LocalAgentConfig kwargs from GOOGLE_CLOUD_{PROJECT,LOCATION} env. - - Temporary override until AGY supports reading these env vars natively. - Returns {} when env does not request Vertex (caller's programmatic config - stands as-is). When env requests Vertex, returns VertexKwargs populated - for passing to LocalAgentConfig.__init__ so AGY's @model_validator picks - them up. - - Raises ValueError if env requests Vertex but project/location are missing. - - TODO: remove once AGY reads these env vars natively. - """ - if not _env_use_vertex(): - return {} - - project = os.environ.get("GOOGLE_CLOUD_PROJECT", "") - location = os.environ.get("GOOGLE_CLOUD_LOCATION", "") - - missing = [ - name - for name, value in ( - ("project (set GOOGLE_CLOUD_PROJECT)", project), - ("location (set GOOGLE_CLOUD_LOCATION)", location), - ) - if not value - ] - if missing: - raise ValueError( - "Vertex AI backend requested but missing required config: " - + ", ".join(missing) - ) - - print(f"Vertex AI backend configured: project={project} location={location}") - return {"vertex": True, "project": project, "location": location} - - def _build_default_config() -> LocalAgentConfig: """Builds the default agent config the sidecar serves on startup. - Vertex configuration is read from env via `_vertex_kwargs_from_env`. + Credentials/backend come from the standard GenAI env vars, which the + AGY SDK reads natively as of google-antigravity 0.1.7. TODO(#194): per-request `harness_config` will override fields of this default on a per-conversation basis. Until then, every conversation uses this config. """ - return LocalAgentConfig( - system_instructions="You are a helpful agent.", - **_vertex_kwargs_from_env(), - ) + return LocalAgentConfig(system_instructions="You are a helpful agent.") def _has_credentials(config: AgentConfig | None) -> bool: @@ -152,26 +100,25 @@ def _has_credentials(config: AgentConfig | None) -> bool: Mirrors AGY's own validation. AGY accepts exactly these sources: 1. GEMINI_API_KEY environment variable (read directly by AGY). 2. config.api_key set programmatically (AI Studio path). - 3. config.vertex=True + config.{project,location} set (Vertex path). - 4. config.vertex=True + config.api_key set (Vertex Express Mode; - covered by case 2). + 3. Vertex requested (config.vertex or GOOGLE_GENAI_USE_VERTEXAI / + GOOGLE_GENAI_USE_ENTERPRISE) + GOOGLE_CLOUD_{PROJECT,LOCATION}. + 4. Vertex requested + config.api_key set (Express Mode; covered by 2). - Anything else (e.g. vertex=True without project/location) is rejected - by AGY at request time, so we reject it here at startup too. """ # Check env - AGY reads GEMINI_API_KEY directly from os.environ. if os.environ.get("GEMINI_API_KEY"): return True - # Check passed in config - if config: - if getattr(config, "api_key", None): + # AI Studio path via programmatic api_key. + if config and getattr(config, "api_key", None): + return True + + # Vertex path: requested via config.vertex or env, project+location via env. + if _env_use_vertex() or bool(getattr(config, "vertex", False)): + if os.environ.get("GOOGLE_CLOUD_PROJECT") and os.environ.get( + "GOOGLE_CLOUD_LOCATION" + ): return True - if getattr(config, "vertex", False): - # Vertex requires project + location, unless an api_key (Express - # Mode) is set - but Express Mode would have returned True above. - if getattr(config, "project", None) and getattr(config, "location", None): - return True return False diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py index 56c3c469..95916606 100644 --- a/python/antigravity/harness_server_test.py +++ b/python/antigravity/harness_server_test.py @@ -204,22 +204,27 @@ def test_has_credentials_missing(monkeypatch): def test_has_credentials_vertex_requires_project_and_location(monkeypatch): - """vertex=True alone is not enough; AGY requires project+location too.""" + """Vertex needs project+location; as of AGY 0.1.7 these come from env.""" from python.antigravity.harness_server import _has_credentials monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) - cfg_vertex_only = LocalAgentConfig(system_instructions="test", vertex=True) - assert _has_credentials(cfg_vertex_only) is False + cfg = LocalAgentConfig(system_instructions="test") + assert _has_credentials(cfg) is False - cfg_vertex_proj_only = LocalAgentConfig(system_instructions="test", vertex=True, project="p") - assert _has_credentials(cfg_vertex_proj_only) is False + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "p") + assert _has_credentials(cfg) is False - cfg_vertex_loc_only = LocalAgentConfig(system_instructions="test", vertex=True, location="us-central1") - assert _has_credentials(cfg_vertex_loc_only) is False + monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") + assert _has_credentials(cfg) is False - cfg_vertex_full = LocalAgentConfig(system_instructions="test", vertex=True, project="p", location="us-central1") - assert _has_credentials(cfg_vertex_full) is True + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "p") + monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") + assert _has_credentials(cfg) is True def test_has_credentials_vertex_express_mode(monkeypatch): @@ -400,73 +405,25 @@ async def request_iter(): asyncio.run(_run()) -def test_vertex_kwargs_from_env_returns_kwargs(monkeypatch): - """GOOGLE_GENAI_USE_VERTEXAI + GOOGLE_CLOUD_{PROJECT,LOCATION} -> kwargs dict.""" - from python.antigravity.harness_server import _vertex_kwargs_from_env - monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") - monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") - monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") - assert _vertex_kwargs_from_env() == { - "vertex": True, - "project": "env-project", - "location": "us-east1", - } - - -def test_vertex_kwargs_from_env_no_op_without_vertex(monkeypatch): - """Without GOOGLE_GENAI_USE_VERTEXAI, returns empty dict.""" - from python.antigravity.harness_server import _vertex_kwargs_from_env - monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) - monkeypatch.delenv("GOOGLE_GENAI_USE_ENTERPRISE", raising=False) - monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "should-be-ignored") - monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "should-be-ignored") - assert _vertex_kwargs_from_env() == {} +def test_build_default_config_routes_to_vertex_via_env(monkeypatch): + """Bare default config + Vertex env vars routes to Vertex (AGY 0.1.7). - -def test_vertex_kwargs_from_env_raises_when_project_missing(monkeypatch): - """Vertex requested with no project -> ValueError naming the env var.""" - from python.antigravity.harness_server import _vertex_kwargs_from_env - monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") - monkeypatch.delenv("GOOGLE_CLOUD_PROJECT", raising=False) - monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") - with pytest.raises(ValueError, match="GOOGLE_CLOUD_PROJECT"): - _vertex_kwargs_from_env() - - -def test_vertex_kwargs_from_env_raises_when_location_missing(monkeypatch): - """Vertex requested with no location -> ValueError naming the env var.""" - from python.antigravity.harness_server import _vertex_kwargs_from_env - monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") - monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") - monkeypatch.delenv("GOOGLE_CLOUD_LOCATION", raising=False) - with pytest.raises(ValueError, match="GOOGLE_CLOUD_LOCATION"): - _vertex_kwargs_from_env() - - -def test_vertex_kwargs_from_env_enterprise_alias(monkeypatch): - """GOOGLE_GENAI_USE_ENTERPRISE is an alias for VERTEXAI.""" - from python.antigravity.harness_server import _vertex_kwargs_from_env - monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) - monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true") - monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "ent-project") - monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") - assert _vertex_kwargs_from_env() == { - "vertex": True, - "project": "ent-project", - "location": "us-central1", - } - - -def test_build_default_config_picks_up_vertex_env(monkeypatch): - """End-to-end: env vars flow through _build_default_config into LocalAgentConfig.""" + The SDK hydrates vertex onto the config and project/location onto the + VertexEndpoint, so we assert the resolved endpoint rather than + config.{project,location}. + """ from python.antigravity.harness_server import _build_default_config + from google.antigravity import types + monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "True") monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "env-project") monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-east1") cfg = _build_default_config() assert cfg.vertex is True - assert cfg.project == "env-project" - assert cfg.location == "us-east1" + endpoint = cfg._build_shorthand_endpoint() + assert isinstance(endpoint, types.VertexEndpoint) + assert endpoint.project == "env-project" + assert endpoint.location == "us-east1" def test_servicer_requires_default_config():