mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-08-09 08:27:00 -07:00
joneba-google
252 lines
9.5 KiB
Python
252 lines
9.5 KiB
Python
"""Google Antigravity SDK Agent Runner and Context Management.
|
|
|
|
Provides execution wrappers for executing Coding and Evaluator AI Agents using
|
|
the Google Antigravity SDK. Includes serialized working directory controls
|
|
and automatic local sandbox approvals.
|
|
"""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
from typing import Iterator
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def working_directory(path: str | os.PathLike) -> Iterator[None]:
|
|
"""Safely and temporarily changes the working directory.
|
|
|
|
Guarantees restoration of the original CWD even in the event of failures.
|
|
|
|
Args:
|
|
path: Directory path to switch to.
|
|
|
|
Yields:
|
|
None.
|
|
"""
|
|
original_cwd = os.getcwd()
|
|
logging.debug("Switching working directory from %s to %s", original_cwd, path)
|
|
os.chdir(path)
|
|
try:
|
|
yield
|
|
finally:
|
|
logging.debug("Restoring working directory to %s", original_cwd)
|
|
os.chdir(original_cwd)
|
|
|
|
|
|
# Permitted tool allowlist for headless sandbox operations
|
|
ALLOWED_SANDBOX_TOOLS = {
|
|
# Reading tools
|
|
"view_file",
|
|
"read_file",
|
|
# File writing & editing tools
|
|
"replace_file_content",
|
|
"multi_replace_file_content",
|
|
"write_file",
|
|
"write_to_file",
|
|
# Command execution
|
|
"run_command",
|
|
}
|
|
|
|
# Registering global agent hooks for local sandbox tool calls
|
|
try:
|
|
from google.antigravity import Agent, LocalAgentConfig, hooks, policy
|
|
except ImportError:
|
|
Agent, LocalAgentConfig, hooks, policy = None, None, None, None
|
|
|
|
if hooks is not None:
|
|
@hooks.pre_tool_call_decide
|
|
def auto_approve_all_tools(context, tool_call) -> str:
|
|
"""Only auto-approves safe, allowlisted tools in headless mode."""
|
|
if tool_call.name in ALLOWED_SANDBOX_TOOLS:
|
|
logging.debug("Auto-approving allowlisted sandbox tool call: %s", tool_call.name)
|
|
return "PROCEED"
|
|
logging.warning("Rejecting non-allowlisted tool call: %s", tool_call.name)
|
|
return "REJECT"
|
|
|
|
|
|
class AgentRunnerError(Exception):
|
|
"""Raised when the AI Agent fails to run or complete execution loops."""
|
|
|
|
|
|
class AgentRunner:
|
|
"""Manages AI Agent setups and coordinates conversation execution loops."""
|
|
|
|
_cwd_lock: asyncio.Lock | None = None
|
|
|
|
def __init__(
|
|
self,
|
|
project_id: str,
|
|
location: str = "global",
|
|
model_name: str = "gemini-3.5-flash",
|
|
script_dir: str | None = None,
|
|
) -> None:
|
|
"""Initializes the runner with target Vertex AI details.
|
|
|
|
Args:
|
|
project_id: Target Google Cloud Platform Project ID.
|
|
location: Global endpoint location of Vertex AI services (default: "global").
|
|
model_name: Base LLM version string.
|
|
script_dir: Directory containing system/prompt markdown files.
|
|
"""
|
|
self.project_id = project_id
|
|
self.location = location or "global"
|
|
self.model_name = model_name
|
|
self.script_dir = script_dir or os.path.dirname(
|
|
os.path.abspath(__file__)
|
|
)
|
|
|
|
def _load_prompt_file(self, filename: str) -> str | None:
|
|
"""Helper to read a localized system instruction prompt markdown file.
|
|
|
|
Args:
|
|
filename: Name of the prompt file inside the script directory.
|
|
|
|
Returns:
|
|
The text content if file exists, else None.
|
|
"""
|
|
path = os.path.abspath(os.path.join(self.script_dir, filename))
|
|
if not path.startswith(os.path.abspath(self.script_dir)):
|
|
logging.warning("Path traversal attempt detected in prompt loading: %s", filename)
|
|
return None
|
|
|
|
if os.path.exists(path):
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return f.read()
|
|
except IOError as e:
|
|
logging.warning(
|
|
"Failed to read prompt file '%s': %s", filename, e
|
|
)
|
|
return None
|
|
|
|
async def run_agent(
|
|
self,
|
|
role: str,
|
|
prompt: str,
|
|
repo_path: str,
|
|
system_prompt_file: str | None = None,
|
|
) -> str:
|
|
"""Launches and manages an asynchronous conversation with an Antigravity Agent.
|
|
|
|
Args:
|
|
role: Label representing the agent's role (e.g., 'Coding Agent').
|
|
prompt: User message prompt guiding the immediate task.
|
|
repo_path: Target directory root of the repository to execute in.
|
|
system_prompt_file: Optional filename of system prompt markdown.
|
|
|
|
Returns:
|
|
A reconstructed single text block combining thoughts and outputs.
|
|
|
|
Raises:
|
|
AgentRunnerError: If Agent fails to run or execution fails.
|
|
"""
|
|
if Agent is None:
|
|
raise AgentRunnerError("Google Antigravity SDK is not installed.")
|
|
|
|
logging.info("Initializing Agent '%s' inside %s", role, repo_path)
|
|
|
|
# Build fallback / configured system instructions
|
|
system_instructions = f"You are the {role}. You must complete the requested tasks in the workspace."
|
|
if system_prompt_file:
|
|
loaded_instructions = self._load_prompt_file(system_prompt_file)
|
|
if loaded_instructions:
|
|
system_instructions = loaded_instructions
|
|
logging.info(
|
|
"System prompt successfully loaded from %s",
|
|
system_prompt_file
|
|
)
|
|
else:
|
|
logging.warning(
|
|
"Requested system prompt file '%s' not found. Reverting to default instructions.",
|
|
system_prompt_file,
|
|
)
|
|
|
|
config = LocalAgentConfig(
|
|
vertex=True,
|
|
project=self.project_id,
|
|
location=self.location,
|
|
model=self.model_name,
|
|
system_instructions=system_instructions,
|
|
policies=[policy.allow_all()],
|
|
workspaces=[repo_path],
|
|
)
|
|
|
|
stdout_list: list[str] = []
|
|
thinking_list: list[str] = []
|
|
|
|
if AgentRunner._cwd_lock is None:
|
|
AgentRunner._cwd_lock = asyncio.Lock()
|
|
|
|
try:
|
|
# We change CWD to the repo workspace because the Antigravity SDK Agent
|
|
# interacts relative to the current working process directory.
|
|
# Since os.chdir is process-wide, we must serialize execution to prevent
|
|
# concurrent tasks from corrupting the CWD.
|
|
async with AgentRunner._cwd_lock:
|
|
with working_directory(repo_path):
|
|
async with Agent(config) as agent:
|
|
logging.info(
|
|
"[%s] Sending initial task prompt to conversation loop...",
|
|
role,
|
|
)
|
|
await agent.conversation.send(prompt)
|
|
|
|
step_contents: dict[int, str] = {}
|
|
step_thoughts: dict[int, str] = {}
|
|
printed_steps: set[tuple[int, str]] = set()
|
|
|
|
async for step in agent.conversation.receive_steps():
|
|
if step.content:
|
|
step_contents[step.step_index] = step.content
|
|
|
|
# Retrieve thoughts if available via standard properties
|
|
thinking = getattr(step, "thinking", None) or getattr(
|
|
step, "thinking_delta", None
|
|
)
|
|
if thinking:
|
|
step_thoughts[step.step_index] = str(thinking)
|
|
|
|
step_key = (step.step_index, str(step.status))
|
|
if step_key not in printed_steps:
|
|
printed_steps.add(step_key)
|
|
logging.info(
|
|
"[%s Step %s] Type: %s (Source: %s, Status: %s)",
|
|
role,
|
|
step.step_index,
|
|
step.type,
|
|
step.source,
|
|
step.status,
|
|
)
|
|
if step.content:
|
|
logging.info("[%s Content]: %s", role, step.content)
|
|
if thinking:
|
|
logging.debug("[%s Thinking]: %s", role, thinking)
|
|
if step.tool_calls:
|
|
for call in step.tool_calls:
|
|
logging.info(
|
|
"[%s Tool Call]: %s with args %s",
|
|
role,
|
|
call.name,
|
|
call.args,
|
|
)
|
|
|
|
# Accumulate outputs
|
|
for step_idx in sorted(step_contents.keys()):
|
|
stdout_list.append(step_contents[step_idx])
|
|
|
|
for step_idx in sorted(step_thoughts.keys()):
|
|
thinking_list.append(step_thoughts[step_idx])
|
|
|
|
full_output = "\n".join(stdout_list)
|
|
if thinking_list:
|
|
joined_thoughts = "\n".join(thinking_list)
|
|
full_output += f"\nThoughts:\n{joined_thoughts}"
|
|
|
|
logging.info("Agent '%s' execution completed successfully.", role)
|
|
return full_output
|
|
|
|
except Exception as e:
|
|
logging.exception("Failed to execute agent loop for role: %s", role)
|
|
raise AgentRunnerError(f"Agent '{role}' execution failed: {e}") from e
|