deeper-bucket-sync / tests /test_share_trace.py
cmpatino's picture
cmpatino HF Staff
Upload folder using huggingface_hub
4879fc7 verified
Raw
History Blame Contribute Delete
15 kB
"""Env-anchored session detection in the share-trace client.
The bug this guards against: with multiple agents working from one directory,
`detect(cwd, "auto")` used to pick by cwd+recency, Claude-Code-first — so a Codex
agent would upload a co-located Claude Code log. Detection now keys on the
harness that INVOKES the script (its env), so agents are never cross-attributed.
"""
from __future__ import annotations
import json
import os
import re
import sys
import time
from pathlib import Path
import pytest
# The client is a standalone single file under clients/, not on the test path.
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "clients"))
import share_trace as st # noqa: E402
CWD = "/work/proj" # absolute; detect only uses it to compute slugs / cwd-match
_MARKERS = ("CLAUDE_CODE_SESSION_ID", "CLAUDECODE", "CODEX_SANDBOX", "CODEX_SANDBOX_NETWORK_DISABLED")
def _slug(cwd: str) -> str:
return re.sub(r"[/._]", "-", os.path.abspath(cwd))
@pytest.fixture
def home(tmp_path, monkeypatch):
monkeypatch.setenv("HOME", str(tmp_path)) # Path.home() → tmp
for v in _MARKERS:
monkeypatch.delenv(v, raising=False)
return tmp_path
def _cc(home: Path, sid: str, *, mtime: float | None = None) -> Path:
d = home / ".claude" / "projects" / _slug(CWD)
d.mkdir(parents=True, exist_ok=True)
f = d / f"{sid}.jsonl"
f.write_text('{"type":"user","message":{"role":"user","content":"hi"}}\n')
if mtime is not None:
os.utime(f, (mtime, mtime))
return f
def _codex(home: Path, *, cwd: str = CWD) -> Path:
d = home / ".codex" / "sessions" / "2026" / "06" / "26"
d.mkdir(parents=True, exist_ok=True)
f = d / "rollout-2026-06-26T00-00-00-abc.jsonl"
f.write_text(json.dumps({"type": "session_meta", "payload": {"cwd": cwd}}) + "\n")
return f
def test_cc_pins_invoking_session_not_newest(home, monkeypatch):
# CLAUDE_CODE_SESSION_ID must win over the newest-mtime heuristic.
mine = _cc(home, "mine-sid", mtime=time.time() - 100)
_cc(home, "other-sid", mtime=time.time()) # newer; would win by recency
monkeypatch.setenv("CLAUDE_CODE_SESSION_ID", "mine-sid")
monkeypatch.setenv("CLAUDECODE", "1")
harness, path, uncertain = st.detect(CWD, "auto")
assert harness == "claude-code"
assert path == mine # the invoking session, not "other-sid"
assert uncertain is False # exact pin → no confirmation needed
def test_codex_marker_never_grabs_a_claude_log(home, monkeypatch):
# The reported bug: a Codex agent in a dir that ALSO has a Claude Code session.
_cc(home, "cc-sid") # co-located CC log (the trap)
rollout = _codex(home)
monkeypatch.setenv("CODEX_SANDBOX", "seatbelt")
harness, path, _ = st.detect(CWD, "auto")
assert harness == "codex"
assert path == rollout # never the CC log
def test_ambiguous_without_markers_refuses(home):
# No env marker + both harnesses present → refuse rather than misattribute.
_cc(home, "cc-sid")
_codex(home)
with pytest.raises(SystemExit):
st.detect(CWD, "auto")
def test_no_marker_single_harness_is_used(home):
rollout = _codex(home) # only Codex present, no markers
harness, path, _ = st.detect(CWD, "auto")
assert harness == "codex" and path == rollout
def test_missing_session_id_refuses_to_guess(home, monkeypatch):
# If Claude exposes an exact session id, selecting any other log is unsafe.
_cc(home, "real-sid")
monkeypatch.setenv("CLAUDE_CODE_SESSION_ID", "ghost-sid")
monkeypatch.setenv("CLAUDECODE", "1")
with pytest.raises(SystemExit) as exc:
st.detect(CWD, "auto")
assert "ghost-sid" in str(exc.value)
assert "refusing to guess" in str(exc.value)
def test_explicit_transcript_dry_run_works(home, monkeypatch, tmp_path, capsys):
# Explicit transcript mode is the strongest way to pin a session; it must not
# require auto-detection's `uncertain` return value.
transcript = tmp_path / "rollout-2026-06-29T00-00-00-explicit.jsonl"
transcript.write_text(
"\n".join(
[
json.dumps(
{
"type": "session_meta",
"timestamp": "2026-06-29T00:00:00Z",
"payload": {"session_id": "explicit-sess", "model": "gpt-test"},
}
),
json.dumps(
{
"type": "event_msg",
"timestamp": "2026-06-29T00:01:00Z",
"payload": {
"type": "token_count",
"info": {
"total_token_usage": {
"input_tokens": 1,
"output_tokens": 2,
"cached_input_tokens": 0,
"reasoning_output_tokens": 0,
"total_tokens": 3,
}
},
},
}
),
json.dumps(
{
"type": "response_item",
"timestamp": "2026-06-29T00:02:00Z",
"payload": {"type": "local_shell_call", "call_id": "c1"},
}
),
]
)
+ "\n"
)
monkeypatch.setattr(
sys,
"argv",
[
"share_trace.py",
"--transcript",
str(transcript),
"--harness",
"codex",
"--dry-run",
],
)
assert st.main() == 0
out = capsys.readouterr().out
assert "session : explicit-sess" in out
assert "tokens : 3" in out
def test_redactor_preserves_trace_structure_and_task_context():
secrets = {
"hf": "hf_" + "A" * 24,
"github_classic": "ghp_" + "B" * 24,
"github_fine": "github_pat_" + "C" * 30,
"sk": "sk-" + "D" * 24,
"aws_long_lived": "AKIA" + "E" * 16,
"aws_temporary": "ASIA" + "F" * 16,
"slack": "xoxb-" + "1" * 12 + "-" + "G" * 24,
"gitlab": "glpat-" + "H" * 24,
"google": "AIza" + "I" * 35,
"npm": "npm_" + "J" * 36,
"pypi": "pypi-" + "K" * 24,
"jwt": "eyJ" + "L" * 12 + "." + "M" * 12 + "." + "N" * 12,
}
record = {
"type": "response_item",
"payload": {
"task": "Fix the payment retry while preserving the commit history.",
"commit": "7f4d3b2a" * 5,
"headers": {
"Authorization": "Bearer " + secrets["hf"],
"Cookie": "session=top-secret-cookie",
},
"password": "correct horse battery staple",
"aws_secret_access_key": "O" * 40,
"command": (
"curl -H 'Authorization: Basic dXNlcjpwYXNz' "
"'https://alice:db-pass@db.internal/app"
"?access_token=query-secret'"
),
"private_key": (
"-----BEGIN PRIVATE KEY-----\nsecret-material\n"
"-----END PRIVATE KEY-----"
),
"provider_values": list(secrets.values()),
},
}
redactor = st.TraceRedactor("secrets")
result = redactor.redact_jsonl(json.dumps(record) + "\n")
parsed = json.loads(result)
assert parsed["type"] == "response_item"
assert parsed["payload"]["task"] == record["payload"]["task"]
assert parsed["payload"]["commit"] == record["payload"]["commit"]
assert "curl -H" in parsed["payload"]["command"]
assert "db.internal/app" in parsed["payload"]["command"]
assert all(secret not in result for secret in secrets.values())
for sensitive in (
"top-secret-cookie",
"correct horse battery staple",
"O" * 40,
"dXNlcjpwYXNz",
"alice",
"db-pass",
"query-secret",
"secret-material",
):
assert sensitive not in result
assert "<REDACTED:BEARER_TOKEN_1>" in result
assert "<REDACTED:PRIVATE_KEY_" in result
assert redactor.summary()
def test_balanced_redaction_uses_stable_aliases_and_preserves_relative_paths():
text = json.dumps(
{
"message": (
"Ask alice@example.com to inspect /Users/alice/work/app.py; "
"alice@example.com owns src/app.py"
)
}
) + "\n"
redactor = st.TraceRedactor("balanced")
result = redactor.redact_jsonl(text)
assert result.count("<REDACTED:EMAIL_1>") == 2
assert "$HOME/work/app.py" in result
assert "src/app.py" in result
assert redactor.summary()["EMAIL"] == 2
assert redactor.summary()["HOME_PATH"] == 1
def test_privacy_levels_are_progressively_stricter():
text = (
"Contact alice@example.com under /home/alice/work, then call "
"https://api.internal.example/v1 from 10.20.30.40"
)
secrets = st.redact(text, privacy="secrets")
balanced = st.redact(text, privacy="balanced")
strict = st.redact(text, privacy="strict")
assert "alice@example.com" in secrets
assert "/home/alice/work" in secrets
assert "api.internal.example" in secrets
assert "10.20.30.40" in secrets
assert "alice@example.com" not in balanced
assert "$HOME/work" in balanced
assert "api.internal.example" in balanced
assert "10.20.30.40" in balanced
assert "api.internal.example" not in strict
assert "10.20.30.40" not in strict
assert "https://<REDACTED:HOST_1>/v1" in strict
def test_escaped_authorization_header_is_redacted_without_losing_command():
line = json.dumps(
{
"command": (
'curl -H \\"Authorization: Bearer escaped.token-value\\" '
"https://example.com/v1"
)
}
) + "\n"
result = st.redact(line, privacy="secrets")
assert "escaped.token-value" not in result
assert "Authorization: Bearer <REDACTED:BEARER_TOKEN_1>" in result
assert "curl -H" in result
assert "https://example.com/v1" in result
def test_redaction_is_idempotent_for_quoted_and_unquoted_assignments():
text = json.dumps(
{
"command": "run --password='two words' api_key=opaque-value",
"password": "structured value",
}
) + "\n"
first = st.redact(text, privacy="balanced")
second = st.redact(first, privacy="balanced")
assert second == first
assert "two words" not in first
assert "opaque-value" not in first
assert "structured value" not in first
def test_custom_patterns_are_stable_and_pattern_file_is_validated(tmp_path):
patterns = tmp_path / "redact-patterns.txt"
patterns.write_text("# customer identifiers\nAcme-(?:North|South)\n")
compiled = st._custom_patterns(str(patterns))
redactor = st.TraceRedactor("secrets", compiled)
result = redactor.redact_jsonl(
json.dumps({"task": "Compare Acme-North with Acme-North and Acme-South"}) + "\n"
)
assert result.count("<REDACTED:CUSTOM_1>") == 2
assert result.count("<REDACTED:CUSTOM_2>") == 1
assert "Compare " in result
patterns.write_text("(\n")
with pytest.raises(SystemExit, match="invalid regex"):
st._custom_patterns(str(patterns))
def test_full_upload_only_sends_scrubbed_content_and_neutral_filename(
home, monkeypatch, tmp_path
):
secret = "github_pat_" + "Z" * 30
transcript = tmp_path / "alice@example.com.jsonl"
transcript.write_text(
"\n".join(
[
json.dumps(
{
"type": "session_meta",
"timestamp": "2026-06-29T00:00:00Z",
"payload": {
"session_id": "scrubbed-session",
"model": "gpt-test",
},
}
),
json.dumps(
{
"type": "event_msg",
"timestamp": "2026-06-29T00:01:00Z",
"payload": {
"type": "token_count",
"info": {
"total_token_usage": {
"input_tokens": 1,
"output_tokens": 2,
"cached_input_tokens": 0,
"reasoning_output_tokens": 0,
"total_tokens": 3,
}
},
},
}
),
json.dumps(
{
"type": "response_item",
"timestamp": "2026-06-29T00:02:00Z",
"payload": {
"type": "local_shell_call",
"call_id": "c1",
"command": (
f"deploy with {secret} for alice@example.com "
"from /Users/alice/work/app.py"
),
},
}
),
]
)
+ "\n"
)
uploads = {}
def capture_upload(local, destination):
uploads[destination] = Path(local).read_text()
monkeypatch.setattr(st, "_hf_cp", capture_upload)
monkeypatch.setattr(st.shutil, "which", lambda _: "/usr/bin/hf")
monkeypatch.setattr(
sys,
"argv",
[
"share_trace.py",
"--transcript",
str(transcript),
"--harness",
"codex",
"--full",
"--yes",
"--upload-only",
"--agent-id",
"agent-1",
"--org",
"test-org",
"--slug",
"test-collab",
],
)
assert st.main() == 0
manifest_uri = next(uri for uri in uploads if uri.endswith("/manifest.md"))
trace_uri = next(uri for uri in uploads if uri.endswith("/trace.jsonl"))
assert "alice@example.com" not in uploads[manifest_uri]
assert '"privacy": "balanced"' in uploads[manifest_uri]
assert '"GITHUB_TOKEN": 1' in uploads[manifest_uri]
assert secret not in uploads[trace_uri]
assert "alice@example.com" not in uploads[trace_uri]
assert "/Users/alice" not in uploads[trace_uri]
assert "$HOME/work/app.py" in uploads[trace_uri]
assert all("alice@example.com.jsonl" not in uri for uri in uploads)
def test_session_id_must_be_a_safe_bucket_component():
assert st._safe_session_id("rollout-2026.06_29") == "rollout-2026.06_29"
with pytest.raises(SystemExit, match="safe --session-id"):
st._safe_session_id("../another-session")