diff --git a/src/agent_env/agents/agents/gemini/gemini_config.py b/src/agent_env/agents/agents/gemini/gemini_config.py index f377b2343..a1f324c89 100644 --- a/src/agent_env/agents/agents/gemini/gemini_config.py +++ b/src/agent_env/agents/agents/gemini/gemini_config.py @@ -3,16 +3,20 @@ `~/.gemini/settings.json` holds what every run shares, the changelog's hook. Each run gets settings of its own as the CLI's system settings, which override the user's: the task's MCP servers with the turn's role header, and the turn's system prompt as a -context file. So turns running at once never share a role or a system prompt. +context file named by its content. So turns running at once never share a role, and a +system prompt keeps one path from turn to turn. """ from __future__ import annotations +import hashlib import json import logging import os +import threading import uuid +from collections import Counter from collections.abc import Iterator, Mapping -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from pathlib import Path from typing import Any @@ -24,6 +28,8 @@ _DEFAULT_SYSTEM_SETTINGS = Path("/etc/gemini-cli/settings.json") # The CLI's context file in a workspace, which a run with a system prompt still loads beside its own. _WORKSPACE_CONTEXT_FILE = "GEMINI.md" +_prompt_files: Counter[Path] = Counter() +_prompt_files_lock = threading.Lock() def gemini_home() -> Path: @@ -66,7 +72,8 @@ def run_settings(servers: Mapping[str, Mapping[str, Any]], role: str | None, system_prompt: str | None) -> Iterator[Path]: """One run's system settings, in an owner-only file of its own that names `servers` with `role`; with a `system_prompt`, also a context file the CLI layers on its own system prompt (the CLI has no flag for one, and - replacing its system prompt would lose its tool protocol). Both are removed when the run ends.""" + replacing its system prompt would lose its tool protocol). The settings go when the run ends, the context file + when the last run naming it does.""" home = gemini_home() runs = home / "agentenv-runs" runs.mkdir(mode=0o700, parents=True, exist_ok=True) @@ -76,13 +83,9 @@ def run_settings(servers: Mapping[str, Mapping[str, Any]], role: str | None, **(settings.get("mcpServers") if isinstance(settings.get("mcpServers"), dict) else {}), **{settings_name(name): _server_entry(info["url"], info.get("headers"), role) for name, info in servers.items()}, } - written: list[Path] = [] - try: + with ExitStack() as stack: if system_prompt: - # The CLI's global context files live in its home, and only this run's settings name this one. - context = home / f"agentenv-system-prompt-{token}.md" - context.write_text(system_prompt) - written.append(context) + context = stack.enter_context(_prompt_file(home, system_prompt)) settings["context"] = { **(settings.get("context") if isinstance(settings.get("context"), dict) else {}), "fileName": [_WORKSPACE_CONTEXT_FILE, context.name], @@ -90,9 +93,25 @@ def run_settings(servers: Mapping[str, Mapping[str, Any]], role: str | None, path = runs / f"{token}.json" # It carries the MCP servers' headers. with os.fdopen(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w") as handle: + stack.callback(path.unlink, missing_ok=True) json.dump(settings, handle) - written.append(path) + yield path + + +@contextmanager +def _prompt_file(home: Path, system_prompt: str) -> Iterator[Path]: + """The system prompt as a global context file named by its content: the CLI writes the file's path into the + model's system instruction, so a name new each turn would miss the provider's prompt cache on every turn.""" + path = home / f"agentenv-system-prompt-{hashlib.sha256(system_prompt.encode()).hexdigest()[:32]}.md" + with _prompt_files_lock: + if not _prompt_files[path]: + path.write_text(system_prompt) + _prompt_files[path] += 1 + try: yield path finally: - for path in written: - path.unlink(missing_ok=True) + with _prompt_files_lock: + _prompt_files[path] -= 1 + if not _prompt_files[path]: + del _prompt_files[path] + path.unlink(missing_ok=True) diff --git a/tst/unit/agents/gemini/test_extensions.py b/tst/unit/agents/gemini/test_extensions.py index b49d371f2..e10b24f12 100644 --- a/tst/unit/agents/gemini/test_extensions.py +++ b/tst/unit/agents/gemini/test_extensions.py @@ -11,7 +11,7 @@ from agentenv_protocol.a2a_agent import AGENT_CONFIG_V1 from starlette.testclient import TestClient -from tst.unit.agents.gemini.agent import a2a_server, gemini_bridge +from tst.unit.agents.gemini.agent import a2a_server, gemini_bridge, gemini_config from tst.unit.agents.gemini.support import card, namespace_grant, send AGENT_CONFIG_PATH = AGENT_CONFIG_V1.endpoint @@ -76,8 +76,9 @@ def test_a_server_is_registered_without_underscores(home, client): def test_each_run_gets_its_servers_role_and_system_prompt_in_settings_of_its_own(monkeypatch, home, client): - """Turns of different roles running at once share no file: each CLI's system settings name its turn's role, the - SDK's peer tools among the servers, and its system prompt as a context file, all removed once the turn ends.""" + """Turns of different roles running at once share no settings: each CLI's system settings name its turn's role, + the SDK's peer tools among the servers, and its system prompt as a context file named by its content, so the same + prompt keeps one path from turn to turn; all are removed once the turn ends.""" runs = [] async def fake_run_gemini(**kwargs): @@ -94,8 +95,10 @@ async def fake_run_gemini(**kwargs): send(client, "one", context_id="ctx-1") client.post(AGENT_CONFIG_PATH, json={"role": None, "system_prompt": None}) send(client, "two", context_id="ctx-2") + client.post(AGENT_CONFIG_PATH, json={"role": "solver", "system_prompt": "be brief"}) + send(client, "three", context_id="ctx-1") - (first_path, first, first_context), (second_path, second, second_context) = runs + (first_path, first, first_context), (second_path, second, second_context), (_, third, _) = runs peer = first["mcpServers"]["peeragents"]["httpUrl"] assert re.fullmatch(r"http://127\.0\.0\.1:\d+/ext/peer-agents/mcp", peer) assert first["mcpServers"] == { @@ -105,6 +108,7 @@ async def fake_run_gemini(**kwargs): } (prompt_file,) = first_context assert first["context"] == {"fileName": ["GEMINI.md", prompt_file]} and first_context[prompt_file] == "be brief" + assert third["context"] == first["context"] assert second["mcpServers"] == { "before": {"httpUrl": "https://gw/a/mcp", "headers": {"X-Key": "k"}}, "audited": {"httpUrl": "https://gw/b/mcp", "headers": {"agentenv-role": "auditor"}}, @@ -115,6 +119,16 @@ async def fake_run_gemini(**kwargs): assert not (home / ".gemini" / prompt_file).exists() + +def test_turns_sharing_a_system_prompt_share_its_file_until_the_last_ends(home): + with gemini_config.run_settings({}, "solver", "be brief") as first: + with gemini_config.run_settings({}, "auditor", "be brief") as second: + names = [json.loads(p.read_text())["context"]["fileName"][1] for p in (first, second)] + assert names[0] == names[1] + prompt = home / ".gemini" / names[0] + assert prompt.read_text() == "be brief" and not second.exists() + assert not prompt.exists() and not first.exists() + def test_a_duplicate_url_or_name_is_a_409(home, client): client.post(MCP_CONFIG_PATH, json={"url": "https://gw/a/mcp", "name": "env1"})