2222import logging
2323import os
2424import sys
25+ from typing import TypedDict
2526import grpc
2627from grpc_health .v1 import health , health_pb2 , health_pb2_grpc
2728from 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
92152class 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
369425if __name__ == "__main__" :
370426 main ()
0 commit comments