Skip to content

Commit 17869d2

Browse files
antigravity: forward GOOGLE_CLOUD_{PROJECT,LOCATION} env to AGY config (#223)
The AGY SDK requires vertex/project/location on its AgentConfig but does not read these env vars itself. Setting GOOGLE_GENAI_USE_VERTEXAI=True plus GOOGLE_CLOUD_PROJECT/LOCATION per the README's auth instructions therefore reaches the credential check, then dies later inside AGY with a confusing 'project and location, or an API key' error. Have the sidecar forward the env vars to gemini_config at startup, with programmatic config taking precedence. Validate fail-fast: if vertex is requested but project/location are missing, raise ValueError at startup naming the missing env var, instead of failing per-request later.
1 parent 9a4a581 commit 17869d2

2 files changed

Lines changed: 270 additions & 113 deletions

File tree

python/antigravity/harness_server.py

Lines changed: 116 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import logging
2323
import os
2424
import sys
25+
from typing import TypedDict
2526
import grpc
2627
from grpc_health.v1 import health, health_pb2, health_pb2_grpc
2728
from google.protobuf.struct_pb2 import Struct
@@ -52,59 +53,121 @@ def get_weather(city: str) -> str:
5253
else:
5354
return f"Weather information for '{city}' is not available."
5455

55-
# 2. Define the static agent config
56-
loaded_config = LocalAgentConfig(
57-
system_instructions="You are a helpful agent. Use the get_weather tool to answer weather questions.",
58-
tools=[get_weather]
59-
)
56+
class VertexKwargs(TypedDict, total=False):
57+
"""Typed subset of LocalAgentConfig kwargs needed to enable Vertex AI.
6058
61-
def _has_credentials(config: AgentConfig | None) -> bool:
62-
"""Checks if Gemini credentials are set either in env or config."""
63-
# Check environment variables
64-
has_api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
65-
use_vertex = (
59+
`total=False` so {} is a valid value (returned when env does not request
60+
Vertex). When populated, all three keys are present.
61+
"""
62+
vertex: bool
63+
project: str
64+
location: str
65+
66+
67+
def _env_use_vertex() -> bool:
68+
"""True if env requests the Vertex AI backend (vs. AI Studio API key)."""
69+
return (
6670
os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "").lower() in ("true", "1") or
6771
os.environ.get("GOOGLE_GENAI_USE_ENTERPRISE", "").lower() in ("true", "1")
6872
)
69-
if has_api_key or use_vertex:
73+
74+
def _vertex_kwargs_from_env() -> VertexKwargs:
75+
"""Returns LocalAgentConfig kwargs from GOOGLE_CLOUD_{PROJECT,LOCATION} env.
76+
77+
Temporary override until AGY supports reading these env vars natively.
78+
Returns {} when env does not request Vertex (caller's programmatic config
79+
stands as-is). When env requests Vertex, returns VertexKwargs populated
80+
for passing to LocalAgentConfig.__init__ so AGY's @model_validator picks
81+
them up.
82+
83+
Raises ValueError if env requests Vertex but project/location are missing.
84+
85+
TODO: remove once AGY reads these env vars natively.
86+
"""
87+
if not _env_use_vertex():
88+
return {}
89+
90+
project = os.environ.get("GOOGLE_CLOUD_PROJECT", "")
91+
location = os.environ.get("GOOGLE_CLOUD_LOCATION", "")
92+
93+
missing = [
94+
name for name, value in (
95+
("project (set GOOGLE_CLOUD_PROJECT)", project),
96+
("location (set GOOGLE_CLOUD_LOCATION)", location),
97+
) if not value
98+
]
99+
if missing:
100+
raise ValueError(
101+
"Vertex AI backend requested but missing required config: "
102+
+ ", ".join(missing)
103+
)
104+
105+
print(f"Vertex AI backend configured: project={project} location={location}")
106+
return {"vertex": True, "project": project, "location": location}
107+
108+
def _build_default_config() -> LocalAgentConfig:
109+
"""Builds the default agent config the sidecar serves on startup.
110+
111+
Vertex configuration is read from env via `_vertex_kwargs_from_env`.
112+
113+
TODO(#194): per-request `harness_config` will override fields of this
114+
default on a per-conversation basis. Until then, every conversation uses
115+
this config.
116+
"""
117+
return LocalAgentConfig(
118+
system_instructions="You are a helpful agent. Use the get_weather tool to answer weather questions.",
119+
tools=[get_weather],
120+
**_vertex_kwargs_from_env(),
121+
)
122+
123+
def _has_credentials(config: AgentConfig | None) -> bool:
124+
"""Checks if Gemini credentials are set per AGY's accepted sources.
125+
126+
Mirrors AGY's own validation. AGY accepts exactly these sources:
127+
1. GEMINI_API_KEY environment variable (read directly by AGY).
128+
2. config.api_key set programmatically (AI Studio path).
129+
3. config.vertex=True + config.{project,location} set (Vertex path).
130+
4. config.vertex=True + config.api_key set (Vertex Express Mode;
131+
covered by case 2).
132+
133+
Anything else (e.g. vertex=True without project/location) is rejected
134+
by AGY at request time, so we reject it here at startup too.
135+
"""
136+
# Check env - AGY reads GEMINI_API_KEY directly from os.environ.
137+
if os.environ.get("GEMINI_API_KEY"):
70138
return True
71-
72-
# Check configuration
139+
140+
# Check passed in config
73141
if config:
74-
# Check nested gemini_config
75-
gemini_config = getattr(config, "gemini_config", None)
76-
if gemini_config:
77-
# 1. Direct configuration
78-
if getattr(gemini_config, "api_key", None) or getattr(gemini_config, "vertex", False):
79-
return True
80-
# 2. Per-model configuration
81-
models = getattr(gemini_config, "models", None)
82-
default_model = getattr(models, "default", None) if models else None
83-
if default_model and getattr(default_model, "api_key", None):
84-
return True
85-
86-
# Check top-level config shorthands
87-
if getattr(config, "api_key", None) or getattr(config, "vertex", False):
142+
if getattr(config, "api_key", None):
88143
return True
89-
144+
if getattr(config, "vertex", False):
145+
# Vertex requires project + location, unless an api_key (Express
146+
# Mode) is set - but Express Mode would have returned True above.
147+
if getattr(config, "project", None) and getattr(config, "location", None):
148+
return True
149+
90150
return False
91151

92152
class AntigravityHarnessServiceServicer(ax_pb2_grpc.HarnessServiceServicer):
93153
"""Implements the ax.HarnessService protocol over gRPC."""
94154

95-
def __init__(self):
155+
def __init__(self, default_config: AgentConfig):
96156
# TODO: Implement an eviction/idle-timeout policy to prevent unbounded memory growth in production.
157+
self._default_config = default_config
97158
self._agents = {}
98159
self._lock = asyncio.Lock()
99160

100161
async def _get_or_create_agent(self, conversation_id: str) -> Agent:
101162
async with self._lock:
102163
if conversation_id not in self._agents:
103-
global loaded_config
104-
if not loaded_config:
164+
# TODO(#194): derive a per-conversation AgentConfig by layering
165+
# request.start.harness_config on top of self._default_config,
166+
# instead of using the default verbatim for every conversation.
167+
if not self._default_config:
105168
raise ValueError("Agent config is not loaded on the server")
106169
print(f"[gRPC] Creating new Agent instance for conv_id={conversation_id}")
107-
agent = Agent(loaded_config)
170+
agent = Agent(self._default_config)
108171
await agent.__aenter__()
109172
self._agents[conversation_id] = agent
110173
return self._agents[conversation_id]
@@ -162,9 +225,8 @@ async def _run_turn(self, request):
162225
return
163226
latest_query_text = latest_message.content.text.text
164227

165-
# 2. Initialize or get the Antigravity Agent session
166-
global loaded_config
167-
if not loaded_config:
228+
# TODO(#194): parse and validate request.start.harness_config here.
229+
if not self._default_config:
168230
yield ax_pb2.HarnessResponse(
169231
conversation_id=request.conversation_id,
170232
end=ax_pb2.HarnessEnd(
@@ -176,23 +238,6 @@ async def _run_turn(self, request):
176238
),
177239
)
178240
return
179-
180-
# Check credentials
181-
if not _has_credentials(loaded_config):
182-
yield ax_pb2.HarnessResponse(
183-
conversation_id=request.conversation_id,
184-
end=ax_pb2.HarnessEnd(
185-
state=ax_pb2.STATE_FAILED,
186-
error=ax_pb2.Error(
187-
code=9, # FAILED_PRECONDITION
188-
description=(
189-
"No Gemini credentials configured. Please set the GEMINI_API_KEY environment variable "
190-
"(AI Studio) or GOOGLE_GENAI_USE_VERTEXAI=True (Vertex AI) before starting the harness server."
191-
),
192-
),
193-
),
194-
)
195-
return
196241
try:
197242
agent = await self._get_or_create_agent(request.conversation_id)
198243
conversation = agent.conversation
@@ -302,9 +347,9 @@ def flush_thought():
302347
)
303348
return
304349

305-
async def serve(host: str, port: int):
350+
async def _serve(host: str, port: int, default_config: AgentConfig):
306351
server = grpc.aio.server()
307-
servicer = AntigravityHarnessServiceServicer()
352+
servicer = AntigravityHarnessServiceServicer(default_config)
308353
ax_pb2_grpc.add_HarnessServiceServicer_to_server(servicer, server)
309354

310355
# Serve the standard gRPC health protocol.
@@ -321,7 +366,7 @@ async def serve(host: str, port: int):
321366
finally:
322367
await servicer.cleanup()
323368

324-
def enhance_config_from_env(config) -> None:
369+
def _enhance_config_from_env(config) -> None:
325370
skills_dir = os.environ.get("SKILLS_DIR")
326371
if skills_dir and os.path.isdir(skills_dir):
327372
print(f"Adding preinstalled skills directory to agent config: {skills_dir}")
@@ -331,7 +376,7 @@ def enhance_config_from_env(config) -> None:
331376
if skills_dir not in config.skills_paths:
332377
config.skills_paths.append(skills_dir)
333378

334-
def resolve_localhost():
379+
def _resolve_localhost():
335380
"""Ensure `localhost` resolves to 127.0.0.1.
336381
337382
Substrate actors run under gVisor with no runtime-injected /etc/hosts.
@@ -357,14 +402,25 @@ def main():
357402
parser.add_argument("--host", default="localhost", help="Host to bind the server to")
358403
args = parser.parse_args()
359404

360-
global loaded_config
361-
enhance_config_from_env(loaded_config)
405+
try:
406+
default_config = _build_default_config()
407+
_enhance_config_from_env(default_config)
408+
if not _has_credentials(default_config):
409+
raise ValueError(
410+
"No Gemini credentials configured. Set GEMINI_API_KEY "
411+
"(AI Studio) or GOOGLE_GENAI_USE_VERTEXAI=True + "
412+
"GOOGLE_CLOUD_{PROJECT,LOCATION} (Vertex AI)."
413+
)
414+
except ValueError as e:
415+
# Single startup-config exit point.
416+
print(f"ERROR: {e}", file=sys.stderr)
417+
sys.exit(1)
362418

363419
# This is a hack, on Agent Substrate /etc/hosts end up not
364420
# having this entry even if it's the OCI image.
365-
resolve_localhost()
421+
_resolve_localhost()
366422

367-
asyncio.run(serve(args.host, args.port))
423+
asyncio.run(_serve(args.host, args.port, default_config))
368424

369425
if __name__ == "__main__":
370426
main()

0 commit comments

Comments
 (0)