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
3 changes: 3 additions & 0 deletions src/copilot_experiments/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
preflight_pier_backend,
prepare_pier_job_for_run,
run_pier_job,
sync_saved_run_config,
)
from .pier_results import (
describe_missing_pier_analysis_source,
Expand Down Expand Up @@ -253,6 +254,8 @@ def run(
if verbose:
prepared.config.debug = True
inject_copilot_token(prepared.config, auth.token)
if prepared.resumed:
sync_saved_run_config(prepared.config)
console.print(f"[bold]Running Pier job[/bold] {prepared.label}")
if prepared.resumed:
console.print(f"[dim]resume:[/dim] reusing existing Pier run {prepared.label}")
Expand Down
81 changes: 78 additions & 3 deletions src/copilot_experiments/pier_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,9 @@ def _nearest_existing_ancestor(path: Path) -> Path:
return Path(path.anchor) if path.anchor else Path(".")


_TOKEN_ENV_KEYS: frozenset[str] = frozenset({"COPILOT_GITHUB_TOKEN", "GITHUB_TOKEN", "GH_TOKEN"})


def inject_copilot_token(config: Any, token: str) -> None:
"""Inject a GitHub token into local Copilot agents without persisting it to config."""

Expand All @@ -195,6 +198,38 @@ def inject_copilot_token(config: Any, token: str) -> None:
agent.env.setdefault("GH_TOKEN", token)


def sync_saved_run_config(config: Any) -> None:
"""Overwrite the Copilot token env vars in ``config.json`` of an existing run dir.

Call this after ``inject_copilot_token`` and before ``run_pier_job`` when resuming.
Pier reads ``config.json`` on ``Job.create`` and raises ``FileExistsError`` when the
stored config differs from the one passed to it — most commonly because the auth token
has rotated between runs. Patching only the token fields in the raw JSON preserves
all other config fields exactly as Pier wrote them, so the comparison always passes.
"""
run_dir = Path(config.jobs_dir) / str(config.job_name)
config_path = run_dir / "config.json"
if not config_path.exists():
return
try:
data = json.loads(config_path.read_text(encoding="utf-8"))
# Update token env vars in-place in the raw dict so no secret-masking occurs.
for saved_agent, live_agent in zip(data.get("agents", []), config.agents, strict=False):
env: dict[str, str] = saved_agent.setdefault("env", {})
for key in _TOKEN_ENV_KEYS:
live_val = live_agent.env.get(key)
if live_val is not None:
env[key] = live_val
else:
env.pop(key, None)
config_path.write_text(
json.dumps(data, ensure_ascii=False, indent=2),
encoding="utf-8",
)
except Exception:
pass


def prepare_pier_job_for_run(
config: Any,
*,
Expand All @@ -215,9 +250,8 @@ def prepare_pier_job_for_run(
if resume:
existing = _latest_existing_run_dir(prepared)
if existing is not None:
prepared.jobs_dir = existing.parent
prepared.job_name = existing.name
return PreparedPierJob(prepared, requested_name, existing.name, resumed=True)
resumed_config = _config_for_resume(prepared, existing)
return PreparedPierJob(resumed_config, requested_name, existing.name, resumed=True)

base_run_name = (now or datetime.now()).strftime("%Y%m%d-%H%M%S")
run_name = base_run_name
Expand Down Expand Up @@ -254,6 +288,47 @@ def _job_dir(config: Any) -> Path:
return Path(config.jobs_dir) / str(config.job_name)


def _config_for_resume(current_config: Any, run_dir: Path) -> Any:
"""Reconstruct the Pier config from a saved run directory for resuming.

Pier's ``Job.create`` raises ``FileExistsError`` when the config passed to it
differs from the one stored in ``config.json``. This helper rebuilds the config
from the saved file (so all structural fields match), then clears the token env
vars so that ``inject_copilot_token`` can inject fresh values. Falls back to
the current in-memory config (pointing at the run dir) when the file cannot be
loaded or parsed.
"""
config_path = run_dir / "config.json"
if config_path.exists():
try:
data = json.loads(config_path.read_text(encoding="utf-8"))
loaded = type(current_config).model_validate(data)
_clear_copilot_token_env(loaded)
# Always anchor jobs_dir/job_name to the actual run directory — the saved
# values may be relative or point to a different machine layout.
loaded.jobs_dir = run_dir.parent
loaded.job_name = run_dir.name
return loaded
except Exception:
pass
current_config.jobs_dir = run_dir.parent
current_config.job_name = run_dir.name
return current_config


def _clear_copilot_token_env(config: Any) -> None:
"""Remove Copilot token env vars from agents so fresh values can be injected."""
for agent in config.agents:
is_copilot = (
agent.import_path == COPILOT_CLI_AGENT_IMPORT_PATH
or agent.name == COPILOT_CLI_AGENT_NAME
)
if not is_copilot:
continue
for key in _TOKEN_ENV_KEYS:
agent.env.pop(key, None)


def _latest_existing_run_dir(config: Any) -> Path | None:
"""Return the latest resumable run directory for a stable job config."""

Expand Down
75 changes: 75 additions & 0 deletions tests/test_pier_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
load_pier_job_config,
preflight_pier_backend,
prepare_pier_job_for_run,
sync_saved_run_config,
)


Expand Down Expand Up @@ -140,6 +141,80 @@ def test_prepare_pier_job_for_run_resume_uses_latest_nested_run(tmp_path: Path):
assert prepared.resumed


def test_prepare_pier_job_for_run_resume_loads_saved_config_and_clears_tokens(tmp_path: Path):
"""Resume rebuilds config from config.json so Pier's equality check passes."""
import json

config_path = tmp_path / "job.yaml"
config_path.write_text(
"job_name: smoke\njobs_dir: jobs\nagents:\n - name: copilot-cli\n",
encoding="utf-8",
)
config = load_pier_job_config(config_path, root=tmp_path)
run_dir = tmp_path / "jobs" / "smoke" / "20260620-160000"
run_dir.mkdir(parents=True)

# Simulate a previous run: inject a (now-stale) token and save config.json
inject_copilot_token(config, "old-token")
saved_config = config.model_copy(deep=True)
saved_config.jobs_dir = run_dir.parent
saved_config.job_name = run_dir.name
(run_dir / "config.json").write_text(
json.dumps(saved_config.model_dump(mode="json"), ensure_ascii=False),
encoding="utf-8",
)
(run_dir / "copilot-experiments-run.json").write_text("{}", encoding="utf-8")

# Load fresh config (token not yet injected) and resume
fresh_config = load_pier_job_config(config_path, root=tmp_path)
prepared = prepare_pier_job_for_run(fresh_config, resume=True)

assert prepared.resumed
assert prepared.run_name == "20260620-160000"
# Token env vars must be cleared so inject_copilot_token can set fresh values
copilot_agent = prepared.config.agents[0]
assert "COPILOT_GITHUB_TOKEN" not in copilot_agent.env
assert "GITHUB_TOKEN" not in copilot_agent.env
assert "GH_TOKEN" not in copilot_agent.env
# After injecting a new token it should be present
inject_copilot_token(prepared.config, "new-token")
assert copilot_agent.env["COPILOT_GITHUB_TOKEN"] == "new-token"


def test_sync_saved_run_config_updates_config_json(tmp_path: Path):
"""sync_saved_run_config overwrites config.json with the current in-memory config."""
import json

config_path = tmp_path / "job.yaml"
config_path.write_text(
"job_name: smoke\njobs_dir: jobs\nagents:\n - name: copilot-cli\n",
encoding="utf-8",
)
config = load_pier_job_config(config_path, root=tmp_path)
run_dir = tmp_path / "jobs" / "smoke" / "20260620-160000"
run_dir.mkdir(parents=True)

# Write a stale config.json with an old token
config.jobs_dir = run_dir.parent
config.job_name = run_dir.name
inject_copilot_token(config, "old-token")
(run_dir / "config.json").write_text(
json.dumps(config.model_dump(mode="json"), ensure_ascii=False),
encoding="utf-8",
)

# Update the in-memory token and sync to disk
config.agents[0].env["COPILOT_GITHUB_TOKEN"] = "new-token"
config.agents[0].env["GITHUB_TOKEN"] = "new-token"
config.agents[0].env["GH_TOKEN"] = "new-token"
sync_saved_run_config(config)

saved = json.loads((run_dir / "config.json").read_text(encoding="utf-8"))
agent_env = saved["agents"][0]["env"]
assert agent_env["COPILOT_GITHUB_TOKEN"] == "new-token"
assert agent_env["GITHUB_TOKEN"] == "new-token"


def test_preflight_pier_backend_reports_missing_docker(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
Expand Down