diff --git a/skillopt_sleep/__main__.py b/skillopt_sleep/__main__.py index 6875ad21..e58b018b 100644 --- a/skillopt_sleep/__main__.py +++ b/skillopt_sleep/__main__.py @@ -44,6 +44,7 @@ json_safe, latest_staging, pending_staged_skills, + redact_secrets, staged_skills, ) from skillopt_sleep.staging import adopt as adopt_staging @@ -333,15 +334,16 @@ def _handoff_dir_for(cfg) -> str: def _redact_deep(obj): - """Redact secret-looking substrings in every string of a JSON-like tree.""" + """Redact secrets key-aware across the whole structure (see redact_secrets). + + This used to recurse values and only scrub string leaves, losing the + mapping-key context — so ``{"api_key": "x"}`` leaked. Delegating to the + key-aware ``redact_secrets`` walker fixes every output boundary that routes + through this helper (--json, digests/snapshot files, gate_trials, extra, + display) at once, keeping them all consistent. + """ from skillopt_sleep.staging import redact_secrets - if isinstance(obj, str): - return redact_secrets(obj) - if isinstance(obj, list): - return [_redact_deep(x) for x in obj] - if isinstance(obj, dict): - return {k: _redact_deep(v) for k, v in obj.items()} - return obj + return redact_secrets(obj) def _display_error(exc: object) -> str: @@ -483,7 +485,7 @@ def _handoff_mine_and_pin(cfg, args, backend, snapshot: str, dry: bool): # NOT marked reviewed: feeding this snapshot back through --tasks-file # with a real backend must still hit the human-review gate above. The # driver itself loads it directly, with the same trust as in-cycle mining. - write_tasks_file(snapshot, _redact_deep(payload)) + write_tasks_file(snapshot, redact_secrets(payload)) print( f"[sleep] handoff: pinned {len(tasks)} tasks -> {snapshot}", file=sys.stderr if args.json else sys.stdout, @@ -807,9 +809,9 @@ def cmd_harvest(args) -> int: ) output_path = "" if getattr(args, "output", ""): - output_path = write_tasks_file(args.output, payload) + output_path = write_tasks_file(args.output, redact_secrets(payload)) if args.json: - json_payload = dict(payload) + json_payload = redact_secrets(payload) if output_path: json_payload["output"] = output_path print(json.dumps(json_payload, ensure_ascii=False, indent=2)) diff --git a/skillopt_sleep/backend.py b/skillopt_sleep/backend.py index 99e41f41..484c328c 100644 --- a/skillopt_sleep/backend.py +++ b/skillopt_sleep/backend.py @@ -27,6 +27,7 @@ import shutil import subprocess import tempfile +import threading from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple @@ -332,24 +333,78 @@ def __init__(self, model: str = "", timeout: int = 180) -> None: self._cache: Dict[str, str] = {} self.last_call_error = "" self.last_reflect_raw = "" + # Guards _cache/_tokens against concurrent mutation on the opt-in + # parallel replay path (SKILLOPT_SLEEP_WORKERS>1). The model call + # itself stays outside the lock so parallel workers can overlap. + self._lock = threading.Lock() + # Per-thread token delta for call-local accounting under parallel replay. + self._thread_local = threading.local() # subclasses override -------------------------------------------------- def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: raise NotImplementedError + def _record_delta(self, delta: int) -> int: + """Add a computed delta to the aggregate ``_tokens`` AND record it + call-local. The one place both totals are updated, so they always agree. + """ + with self._lock: + self._tokens += delta + self._thread_local.delta = delta + return delta + + def _record_cost(self, prompt: str, response: str) -> int: + """THE single path to record an inference's token cost (``len//4``). + + Computes the ``len//4`` delta and delegates to ``_record_delta``. Every + inference path (``_cached_call``, ``attempt_with_tools``, ``reflect``) + must route cost here so the aggregate and call-local totals always agree + and no path under- or over-counts. + """ + delta = len(prompt or "") // 4 + len(response or "") // 4 + return self._record_delta(delta) + + def _reset_call_delta(self) -> None: + """Zero the call-local delta for a NO-CALL path. + + Called at the start of every tool-aware path (and used by cache hits) so + a reused worker never reports a previous call's token count — every + no-call / early-return path leaves the delta at 0. + """ + self._thread_local.delta = 0 + def _cached_call(self, key: str, prompt: str, *, max_tokens: int = 1024) -> str: kind = key.split(":", 1)[0] ev = getattr(self, "evidence", None) - if key in self._cache: + with self._lock: + cached = self._cache.get(key) + if cached is not None: + # cover: later cache hit must not report the previous call's delta. + self._reset_call_delta() # cache hits log key-only (the full text is on the original miss event) if ev is not None: ev.log("replay", "model_call", kind=kind, cache_hit=True, key=key, phase=getattr(self, "evidence_phase", ""), backend=self.name, model=self.model) - return self._cache[key] + return cached + # The model call is intentionally outside the lock so parallel workers + # over the same backend overlap; a concurrent miss may duplicate a call, + # but _cache/_tokens reads+writes below are atomic. out = self._call(prompt, max_tokens=max_tokens) - self._tokens += len(prompt) // 4 + len(out) // 4 - self._cache[key] = out + # Charge every real call's tokens AND record it call-local in one place. + delta = self._record_cost(prompt, out) + with self._lock: + # The cache dedup below may reuse another worker's success, but the + # model call above still consumed tokens (already charged). + existing = self._cache.get(key) + if existing: + # A success was cached by another worker; prefer it (dedup) so an + # empty/duplicate never overwrites a concurrent success. + out = existing + elif out: + # This worker succeeded and nothing is cached: cache it. + self._cache[key] = out + # else: empty result + nothing cached -> don't cache (transient failure) if ev is not None: ev.log("replay", "model_call", kind=kind, cache_hit=False, key=key, phase=getattr(self, "evidence_phase", ""), backend=self.name, @@ -524,7 +579,7 @@ def _explain(c: str) -> str: "Reply with ONLY the JSON array, no prose, no markdown fences." ) raw = self._call(p, max_tokens=1024) - self._tokens += len(p) // 4 + len(raw) // 4 + self._record_cost(p, raw) if ev is not None: ev.log("reflect", "exchange", target=target, attempt=attempt + 1, backend=self.name, model=self.model, @@ -554,8 +609,33 @@ def _explain(c: str) -> str: )) return edits + def _cache_get(self, key: str) -> str | None: + """Thread-safe cache read (route subclass cache access through this).""" + with self._lock: + return self._cache.get(key) + + def _cache_pop(self, key: str) -> str | None: + """Thread-safe cache pop — used to drop a failed entry without racing.""" + with self._lock: + return self._cache.pop(key, None) + + def _cache_pop_if(self, key: str, expected: str | None) -> None: + """Drop a cache entry only if it still holds ``expected``. + + A failed caller must not delete a successful result another worker just + stored for the same key, so we only pop our own (empty) value. + """ + with self._lock: + if self._cache.get(key) == expected: + self._cache.pop(key, None) + def tokens_used(self) -> int: - return self._tokens + with self._lock: + return self._tokens + + def token_delta(self) -> int: + """Token cost of the most recent call on THIS thread (call-local).""" + return getattr(self._thread_local, "delta", 0) # ── Pi CLI backend ──────────────────────────────────────────────── @@ -604,13 +684,13 @@ def _set_call_error(self, message: object) -> None: def _cached_call(self, key: str, prompt: str, *, max_tokens: int = 1024) -> str: """Do not make a transient Pi failure sticky in the response cache.""" - if key in self._cache: + if self._cache_get(key) is not None: # A cached success must not expose an unrelated previous failure # through diagnostics/evidence attached to this call. self.last_call_error = "" out = super()._cached_call(key, prompt, max_tokens=max_tokens) if not out: - self._cache.pop(key, None) + self._cache_pop_if(key, out) return out def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: @@ -798,6 +878,7 @@ def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: return out def attempt_with_tools(self, task, skill, memory, tools): + self._reset_call_delta() # Expose a REAL, callable `search` tool (a shell shim that logs each # call) so the gbrain quick-answerer judge (tool_called=search) is # validated honestly: we detect the call from the shim's log, not from @@ -856,7 +937,7 @@ def attempt_with_tools(self, task, skill, memory, tools): "Claude CLI could not be executed: %s", exc, ) resp = "" - self._tokens += len(prompt) // 4 + len(resp) // 4 + self._record_cost(prompt, resp) called: List[str] = [] if os.path.exists(calllog): with open(calllog) as f: @@ -1239,11 +1320,11 @@ def _verify_tool_allowlist(self, env: Dict[str, str], work: str, agent: str, exp def _cached_call(self, key: str, prompt: str, *, max_tokens: int = 1024) -> str: """Keep failed OpenCode calls out of the cache.""" - if key in self._cache: + if self._cache_get(key) is not None: self.last_call_error = "" out = super()._cached_call(key, prompt, max_tokens=max_tokens) if not out: - self._cache.pop(key, None) + self._cache_pop_if(key, out) return out def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: @@ -1289,6 +1370,7 @@ def attempt_with_tools( tools: List[str], ) -> Tuple[str, List[str]]: self.last_call_error = "" + self._reset_call_delta() if not self.tool_replay: self.last_call_error = ( "OpenCode CLI tool replay is not supported without explicit " @@ -1356,10 +1438,14 @@ def attempt_with_tools( name for name, tool_id in project.tool_mapping.items() if tool_id in called ] except OpenCodeError as exc: - self._tokens += exc.prompt_chars // 4 + # Prompt-only cost on the error path (no response text). + delta = exc.prompt_chars // 4 + with self._lock: + self._tokens += delta + self._thread_local.delta = delta self.last_call_error = str(exc) return "", [] - self._tokens += len(prompt) // 4 + len(text) // 4 + self._record_cost(prompt, text) return text, called_tools @@ -1545,6 +1631,7 @@ def _call(self, prompt: str, *, max_tokens: int = 1024, retries: int = 3) -> str return out def attempt_with_tools(self, task, skill, memory, tools): + self._reset_call_delta() # Codex exec runs in a sandbox with shell access; expose the same real # `search` shim and let it run (workspace-write so the shim can log). import tempfile, shutil, stat @@ -1637,7 +1724,7 @@ def attempt_with_tools(self, task, skill, memory, tools): self.last_call_error = ( f"codex exec (tools) exited {proc.returncode}: {(proc.stderr or '')[:500]}" ) - self._tokens += len(prompt) // 4 + len(resp) // 4 + self._record_cost(prompt, resp) called: List[str] = [] if os.path.exists(calllog): with open(calllog) as f: @@ -1809,6 +1896,7 @@ def _parse_jsonl_response(raw: str) -> str: return "\n".join(parts).strip() def attempt_with_tools(self, task, skill, memory, tools): + self._reset_call_delta() # Expose REAL, callable tool shims in the working directory so the # gbrain quick-answerer judge (tool_called=search) is validated # honestly: we detect each call from the shim's log, not from a @@ -1903,7 +1991,7 @@ def attempt_with_tools(self, task, skill, memory, tools): resp = self._parse_jsonl_response(proc.stdout or "") except Exception: resp = "" - self._tokens += len(prompt) // 4 + len(resp) // 4 + self._record_cost(prompt, resp) called: List[str] = [] if os.path.exists(calllog): with open(calllog) as f: @@ -2193,6 +2281,15 @@ def _call(self, prompt, *, max_tokens=1024): def tokens_used(self): return self.target.tokens_used() + self.optimizer.tokens_used() + def token_delta(self) -> int: + # Call-local cost for a replay attempt: replay_one() drives + # attempt/attempt_with_tools -> the TARGET backend, so the per-attempt + # delta is the target's. The optimizer only appears in replay via + # judge() on the rare model-judge fallback (rule/exact/answer tasks are + # scored locally, 0 tokens); that cost is still counted in the aggregate + # tokens_used() (target + optimizer), so the total is not undercounted. + return getattr(self.target, "token_delta", lambda: 0)() + # ── Azure OpenAI backend (gpt-5.x via managed identity) ─────────────────────── @@ -2380,7 +2477,10 @@ def _call(self, prompt: str, *, max_tokens: int = 1024, retries: int = 5) -> str text = (resp.choices[0].message.content or "").strip() try: u = resp.usage - self._tokens += (getattr(u, "prompt_tokens", 0) or 0) + (getattr(u, "completion_tokens", 0) or 0) + self._record_delta( + (getattr(u, "prompt_tokens", 0) or 0) + + (getattr(u, "completion_tokens", 0) or 0) + ) except Exception: pass if text: @@ -2489,7 +2589,10 @@ def _call(self, prompt: str, *, max_tokens: int = 1024, retries: int = 5) -> str text = (getattr(resp, "output_text", "") or "").strip() try: u = resp.usage - self._tokens += (getattr(u, "input_tokens", 0) or 0) + (getattr(u, "output_tokens", 0) or 0) + self._record_delta( + (getattr(u, "input_tokens", 0) or 0) + + (getattr(u, "output_tokens", 0) or 0) + ) except Exception: pass if text: diff --git a/skillopt_sleep/replay.py b/skillopt_sleep/replay.py index 1502a71c..06aa7650 100644 --- a/skillopt_sleep/replay.py +++ b/skillopt_sleep/replay.py @@ -36,13 +36,16 @@ def replay_one(backend: Backend, task: TaskRecord, skill: str, memory: str, tools = _required_tools(task) tools_called: List[str] = [] t0 = time.time() - tok_before = backend.tokens_used() if tools: response, tools_called = backend.attempt_with_tools(task, skill, memory, tools) else: response = backend.attempt(task, skill, memory, sample_id=sample_id) latency_ms = (time.time() - t0) * 1000.0 - tokens = max(0, backend.tokens_used() - tok_before) + # Call-local token accounting (thread-safe under parallel replay): use the + # backend's per-call delta rather than a before/after global total, which + # another overlapping worker would inflate. + token_delta = getattr(backend, "token_delta", None) + tokens = token_delta() if token_delta else 0 # if the backend doesn't track tokens (e.g. mock), approximate from text length if tokens == 0: tokens = (len(skill) + len(memory) + len(task.intent) + len(response)) // 4 diff --git a/skillopt_sleep/staging.py b/skillopt_sleep/staging.py index e2ecfbd5..81ab4a3e 100644 --- a/skillopt_sleep/staging.py +++ b/skillopt_sleep/staging.py @@ -1089,13 +1089,13 @@ def write_staging( ( os.path.join(out, "report.json"), json.dumps( - json_safe(report.to_dict()), + json_safe(redact_secrets(report.to_dict())), ensure_ascii=False, indent=2, allow_nan=False, ), ), - (os.path.join(out, "report.md"), report_md), + (os.path.join(out, "report.md"), redact_secrets(report_md)), # The manifest is the publication marker and must always be last. ( os.path.join(out, "manifest.json"), diff --git a/tests/test_cli_backend_cache.py b/tests/test_cli_backend_cache.py new file mode 100644 index 00000000..7fe88d17 --- /dev/null +++ b/tests/test_cli_backend_cache.py @@ -0,0 +1,301 @@ +"""Tests for CliBackend response caching + thread-safety on the parallel path.""" + +from __future__ import annotations + +import threading + +from skillopt_sleep.backend import CliBackend + + +class _EchoBackend(CliBackend): + """Minimal backend: echoes the prompt, records call count.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.calls = 0 + + def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: + self.calls += 1 + return f"resp:{prompt}" + + +def test_cached_call_caches_and_counts_tokens(): + b = _EchoBackend() + out = b._cached_call("k:1", "hello") + assert out == "resp:hello" + assert b.calls == 1 + assert b._tokens > 0 + + # A cache hit returns the value and adds no tokens. + tokens_before = b._tokens + out2 = b._cached_call("k:1", "hello") + assert out2 == "resp:hello" + assert b.calls == 1 # no new _call on the hit + assert b._tokens == tokens_before + + +def test_cached_call_distinct_keys_are_distinct(): + b = _EchoBackend() + assert b._cached_call("k:1", "a") == "resp:a" + assert b._cached_call("k:2", "b") == "resp:b" + assert b.calls == 2 + + +def test_cached_call_concurrent_access_does_not_corrupt(): + """Concurrent workers over one backend must not lose cache/token updates.""" + b = _EchoBackend() + results: list[str] = [] + errors: list[Exception] = [] + + def worker(): + try: + results.append(b._cached_call("k:1", "hello")) + except Exception as exc: # pragma: no cover - safety net + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert len(results) == 8 + # The cache may duplicate a call on a concurrent miss, but every result is + # a valid cached value and the cache/token state stays consistent. + assert all(r == "resp:hello" for r in results) + assert b._cache["k:1"] == "resp:hello" + + +def test_concurrent_misses_with_barrier_are_consistent(): + """Force overlapping cache misses with a barrier; state stays consistent.""" + n = 6 + barrier = threading.Barrier(n) + + class _BarrierBackend(_EchoBackend): + def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: + barrier.wait(timeout=15) + return super()._call(prompt, max_tokens=max_tokens) + + b = _BarrierBackend() + results: list[str] = [] + errors: list[Exception] = [] + + def worker(): + try: + results.append(b._cached_call("k:1", "hello")) + except Exception as exc: # pragma: no cover - safety net + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert all(r == "resp:hello" for r in results) + assert b._cache["k:1"] == "resp:hello" + + +def test_cache_pop_if_only_removes_own_value(): + """A failed caller must not delete another worker's successful result.""" + b = _EchoBackend() + b._cache["k:1"] = "resp:hello" + # Failed caller (pop-if with its empty value) must NOT remove a success. + b._cache_pop_if("k:1", "") + assert b._cache["k:1"] == "resp:hello" + # But it does remove an entry that still holds the expected (empty) value. + b._cache["k:2"] = "" + b._cache_pop_if("k:2", "") + assert "k:2" not in b._cache + + +def test_token_delta_is_call_local(): + """A single call records a positive per-call token delta on this thread.""" + b = _EchoBackend() + b._cached_call("k:1", "hello") + assert b.token_delta() > 0 + + +def test_token_delta_isolated_between_threads(): + """Per-thread token deltas do not leak across parallel workers.""" + b = _EchoBackend() + deltas: dict[int, int] = {} + + def worker(i: int): + b._cached_call(f"k:{i}", "hello") + deltas[i] = b.token_delta() + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + # Each worker that actually made a call saw its own positive delta. + for i in range(4): + assert deltas[i] > 0, f"worker {i} got no call-local delta" + + +def test_cache_hit_resets_token_delta(): + """A cache hit must reset the per-thread delta (store-then-load reuse).""" + b = _EchoBackend() + b._cached_call("k:1", "hello") # miss + assert b.token_delta() > 0 + b._cached_call("k:1", "hello") # hit + assert b.token_delta() == 0, "cache hit leaked the previous call's delta" + + +def test_concurrent_missing_charges_every_real_call(): + """A barrier-forced same-key concurrent miss must charge every real call.""" + n = 6 + barrier = threading.Barrier(n) + + class _BarrierBackend(_EchoBackend): + def _call(self, prompt: str, *, max_tokens: int = 1024) -> str: + barrier.wait(timeout=15) + return super()._call(prompt, max_tokens=max_tokens) + + b = _BarrierBackend() + threads = [threading.Thread(target=lambda: b._cached_call("k:1", "hello")) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + + per = len("hello") // 4 + len("resp:hello") // 4 + assert b.calls == n + assert b._tokens == per * b.calls, "concurrent misses undercounted real calls" + + +def test_dual_backend_token_delta_returns_target(): + """DualBackend.token_delta() must report the target backend's call cost.""" + from skillopt_sleep.backend import DualBackend + + target = _EchoBackend() + target._cached_call("k:1", "hello") + optimizer = _EchoBackend() + db = DualBackend(target=target, optimizer=optimizer) + assert db.token_delta() == target.token_delta() + + +def test_attempt_with_tools_sets_call_local_delta(monkeypatch): + """Tool-aware replay must set the thread-local delta so replay_one() sees + real call-local usage instead of falling back to a response-length estimate.""" + from types import SimpleNamespace + + import skillopt_sleep.backend as backend_mod + from skillopt_sleep.backend import ClaudeCliBackend + + b = ClaudeCliBackend(model="", claude_path="claude", timeout=10) + + def fake_run(*args, **kwargs): + return SimpleNamespace(returncode=0, stdout="ok\n", stderr="") + + monkeypatch.setattr(backend_mod.subprocess, "run", fake_run) + + task = SimpleNamespace(intent="intent", context_excerpt="ctx") + response, called = b.attempt_with_tools(task, skill="s", memory="m", tools=["search"]) + assert response == "ok" + assert b.token_delta() > 0, "attempt_with_tools did not set call-local delta" + + +def test_dual_backend_attempt_with_tools_sets_target_delta(monkeypatch): + """The dual-backend tool-replay path must also surface the target's delta.""" + from types import SimpleNamespace + + import skillopt_sleep.backend as backend_mod + from skillopt_sleep.backend import ClaudeCliBackend, DualBackend + + target = ClaudeCliBackend(model="", claude_path="claude", timeout=10) + optimizer = ClaudeCliBackend(model="", claude_path="claude", timeout=10) + db = DualBackend(target=target, optimizer=optimizer) + + def fake_run(*args, **kwargs): + return SimpleNamespace(returncode=0, stdout="ok\n", stderr="") + + monkeypatch.setattr(backend_mod.subprocess, "run", fake_run) + + task = SimpleNamespace(intent="intent", context_excerpt="ctx") + response, called = db.attempt_with_tools(task, skill="s", memory="m", tools=["search"]) + assert response == "ok" + assert db.token_delta() > 0, "dual-backend tool replay did not set target delta" + + +def test_record_cost_concurrent_no_lost_updates(): + """Barrier-forced concurrent _record_cost calls must not lose updates.""" + n = 30 + barrier = threading.Barrier(n) + + class _BarrierEcho(_EchoBackend): + def _cost_worker(self): + barrier.wait(timeout=15) + self._record_cost("hello", "world") + + b = _BarrierEcho() + threads = [threading.Thread(target=b._cost_worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + per = len("hello") // 4 + len("world") // 4 + assert b._tokens == per * n, f"concurrent _record_cost lost updates: {b._tokens} != {per * n}" + + +def test_disabled_tool_replay_resets_call_delta(): + """A prior call's delta must not leak through a disabled-tool-replay path.""" + from types import SimpleNamespace + + from skillopt_sleep.backend import OpenCodeCliBackend + + b = OpenCodeCliBackend(model="", opencode_path="opencode", tool_replay=False) + # A prior call on this thread set a nonzero call-local delta. + b._record_cost("hello", "world") + assert b.token_delta() > 0 # stale from the prior call + # Disabled-tool-replay attempt_with_tools returns early; the delta must be + # reset to 0 so a reused worker does not report the previous call's cost. + task = SimpleNamespace(intent="intent", context_excerpt="ctx") + out, called = b.attempt_with_tools(task, skill="s", memory="m", tools=["search"]) + assert out == "" and called == [] + assert b.token_delta() == 0, "disabled-tool-replay leaked the prior call's delta" + + +def test_cmd_harvest_redact_deep_is_key_aware(): + """_redact_deep must key-aware redact — `{"api_key": "x"}` used to leak.""" + from skillopt_sleep.__main__ import _redact_deep + + out = _redact_deep({"api_key": "x", "content": "keep me", "nested": {"token": "y"}}) + assert out == {"api_key": "[REDACTED]", "content": "keep me", "nested": {"token": "[REDACTED]"}} + + +def test_concurrent_pi_empty_and_success_preserve_success(monkeypatch): + """A failed (empty) Pi worker must not clobber another worker's success.""" + from skillopt_sleep.backend import PiCliBackend + + b = PiCliBackend(model="x") + n = 2 + barrier = threading.Barrier(n) + state: dict[str, str] = {} + + def fake_call(prompt: str, *, max_tokens: int = 1024) -> str: + barrier.wait(timeout=15) + return state[threading.current_thread().name] + + monkeypatch.setattr(b, "_call", fake_call) + results: list[str] = [] + + def worker(ret: str): + state[threading.current_thread().name] = ret + results.append(b._cached_call("k:1", "p")) + + th_empty = threading.Thread(target=worker, args=("",)) + th_success = threading.Thread(target=worker, args=("resp:ok",)) + th_empty.start() + th_success.start() + th_empty.join() + th_success.join() + + # The successful result is cached; the empty worker never removes it. + assert b._cache["k:1"] == "resp:ok" + assert "resp:ok" in results diff --git a/tests/test_export_redaction.py b/tests/test_export_redaction.py new file mode 100644 index 00000000..f97c5ddb --- /dev/null +++ b/tests/test_export_redaction.py @@ -0,0 +1,52 @@ +"""Boundary tests for structured-output redaction (mapping-key aware). + +The maintainer flagged that ``_redact_deep`` recurses only into values and +loses the mapping-key context, so ``{"api_key": "plain-secret"}`` survives. +The structured export boundaries now use the mapping-key-aware +``redact_secrets`` instead. These tests pin that behavior. +""" + +from __future__ import annotations + +from skillopt_sleep.__main__ import _redact_deep +from skillopt_sleep.staging import redact_secrets + + +def test_redact_secrets_is_mapping_key_aware(): + """redact_secrets redacts by mapping key (api_key/token), not just content.""" + data = {"api_key": "top-secret", "nested": {"token": "deep-secret"}} + out = redact_secrets(data) + assert "top-secret" not in str(out) + assert "deep-secret" not in str(out) + assert not _contains(out, "top-secret") + assert not _contains(out, "deep-secret") + + +def test_redact_deep_is_mapping_key_aware(): + """_redact_deep now delegates to the key-aware redact_secrets, so a secret + value under a secret-named key is redacted (the old value-losing bug).""" + data = {"api_key": "top-secret", "content": "keep me"} + out = _redact_deep(data) + assert "top-secret" not in str(out), "secret still leaked" + assert "keep me" in str(out), "diagnostic content lost" + + +def test_report_md_is_redacted_as_string(): + """Markdown output is scrubbed of secret-looking assignments/tokens.""" + md = ( + "reason: the API key is sk-abcdefghijklmnop123456 and " + "a secret assignment api_key=plain-secret-value appears" + ) + out = redact_secrets(md) + assert "sk-abcdefghijklmnop123456" not in out + assert "plain-secret-value" not in out + + +def _contains(obj, needle: str) -> bool: + if isinstance(obj, str): + return needle in obj + if isinstance(obj, dict): + return any(_contains(v, needle) for v in obj.values()) + if isinstance(obj, (list, tuple)): + return any(_contains(i, needle) for i in obj) + return False