Skip to content
Open
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
6 changes: 3 additions & 3 deletions src/ucode/agents/pi.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Pi coding agent: writes ~/.pi/agent/models.json with Databricks-backed providers.
"""Pi coding agent: writes a ucode-private models.json with Databricks-backed providers.

Pi (https://pi.dev) is a multi-provider coding agent. We register three
providers in its `models.json`, each speaking the API dialect best suited to
Expand Down Expand Up @@ -108,7 +108,7 @@ def render_overlay(
codex_models: list[str],
gemini_models: list[str],
) -> tuple[dict, list[list[str]]]:
"""Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
"""Return (overlay, managed_key_paths) for Pi's private agent config."""
providers: dict = {}
keys: list[list[str]] = [["model"]]
# Pi expands header values that match an env var name. Our UA contains
Expand Down Expand Up @@ -237,7 +237,7 @@ def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
def build_runtime_env(token: str) -> dict[str, str]:
env = os.environ.copy()
env["OAUTH_TOKEN"] = token
env["HOME"] = str(PI_UCODE_HOME)
env["PI_CODING_AGENT_DIR"] = str(PI_CONFIG_DIR)
return env


Expand Down
8 changes: 6 additions & 2 deletions tests/test_agent_pi.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,9 +240,13 @@ def test_sets_oauth_token(self):
env = pi.build_runtime_env("tok")
assert env["OAUTH_TOKEN"] == "tok"

def test_sets_ucode_home(self):
def test_sets_private_agent_dir_without_replacing_home(self, monkeypatch):
monkeypatch.setenv("HOME", "/real-user-home")

env = pi.build_runtime_env("tok")
assert env["HOME"] == str(pi.PI_UCODE_HOME)

assert env["PI_CODING_AGENT_DIR"] == str(pi.PI_CONFIG_DIR)
assert env["HOME"] == "/real-user-home"


class TestPiValidateCmd:
Expand Down
7 changes: 5 additions & 2 deletions tests/test_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -911,14 +911,17 @@ def test_launch_pi_per_model(self, tmp_path, monkeypatch, e2e_state, e2e_workspa
pytest.skip("No Pi-compatible models available on this workspace")

monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path)
# Pi reads models.json below HOME/.pi/agent. Point both pi's runtime
# HOME and our writer at the same isolated tmp home.
# Point PI_CODING_AGENT_DIR and ucode's config writer at the same
# isolated directory without changing the process HOME.
pi_home = tmp_path / "pi-home"
pi_dir = pi_home / ".pi" / "agent"
config_path = pi_dir / "models.json"
backup_path = tmp_path / "pi-models.backup.json"
monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home)
monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir)
monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path)
monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json")
monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json")
monkeypatch.setattr(pi, "PI_BACKUP_PATH", backup_path)

failures = []
Expand Down
3 changes: 3 additions & 0 deletions tests/test_e2e_user_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -325,7 +325,10 @@ def test_user_agent_arrives_at_gateway(self, tmp_path, monkeypatch, capture_serv

monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path)
monkeypatch.setattr(pi, "PI_UCODE_HOME", pi_home)
monkeypatch.setattr(pi, "PI_CONFIG_DIR", pi_dir)
monkeypatch.setattr(pi, "PI_CONFIG_PATH", config_path)
monkeypatch.setattr(pi, "PI_SETTINGS_PATH", pi_dir / "settings.json")
monkeypatch.setattr(pi, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings.backup.json")
monkeypatch.setattr(pi, "PI_BACKUP_PATH", tmp_path / "pi.backup.json")

state = {
Expand Down