Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 31 additions & 12 deletions src/agent_env/agents/agents/gemini/gemini_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand All @@ -76,23 +83,35 @@ 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],
}
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"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Hash length lacks a name

The new [:32] hardcodes how much of the prompt hash goes into the filename. The repository requires magic numbers to use descriptive names rather than inline values. Name this length before merging so readers can tell what it controls.

Rule Used: Store magic numbers as class or instance variables with descriptive names rather than using them inline in the code. (source)

Learned From
scaleapi/scaleapi#126388

Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/agents/agents/gemini/gemini_config.py
Line: 105

Comment:
**Hash length lacks a name**

The new `[:32]` hardcodes how much of the prompt hash goes into the filename. The repository requires magic numbers to use descriptive names rather than inline values. Name this length before merging so readers can tell what it controls.

**Rule Used:** Store magic numbers as class or instance variables with descriptive names rather than using them inline in the code. ([source](https://app.greptile.com/scale-ai/-/custom-context?memory=002e0051-41ad-46c1-9098-47433c580150))

**Learned From**
[scaleapi/scaleapi#126388](https://github.com/scaleapi/scaleapi/pull/126388)

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Fix in Cursor Fix in Claude Code Fix in Codex

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)
22 changes: 18 additions & 4 deletions tst/unit/agents/gemini/test_extensions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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"] == {
Expand All @@ -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"}},
Expand All @@ -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"})

Expand Down
Loading