Skip to content

Commit 50d9d90

Browse files
authored
fix(sleep): thread-safe backend cache + redact exports (#251)
* fix(sleep): thread-safe backend cache + redact exports - Guard CliBackend _cache/_tokens with a lock on the opt-in parallel replay path (SKILLOPT_SLEEP_WORKERS>1) so a concurrent miss cannot corrupt state or lose the token cost metric; the model call stays outside the lock so parallel workers still overlap. - Redact report.json (redact_secrets) before staging. - Redact harvest --output and --json exports (_redact_deep). - Add tests for CliBackend caching/thread-safety. * fix(sleep): mapping-key-aware redaction at export boundaries Address maintainer review on #251: - Use redact_secrets (mapping-key aware) instead of _redact_deep for harvest --output/--json and handoff exports, so nested api_key/token mappings are redacted (not just bare string leaves). - Redact report_md before writing it alongside report.json. - Add boundary tests (nested api_key/token + report.md). * fix(sleep): thread-safe cache access through locked helpers Address maintainer review on #251 (thread-safety): - Add _cache_get/_cache_pop/_cache_pop_if locked helpers and route the Pi and OpenCode _cached_call overrides through them (they previously read/pop the cache outside the lock). - tokens_used() now reads _tokens under the lock. - Popping a failed entry is conditional (_cache_pop_if): a failed caller only drops its own empty value, never another worker's just-stored success. - Add tests: barrier-forced overlapping misses stay consistent, and pop-if does not delete a successful entry. * fix(sleep): call-local token accounting for parallel replay Address maintainer review on #251 (last thread-safety item): - Record each model call's token delta on the calling thread (thread-local), so parallel replay_one() charges its own cost instead of a before/after global total that an overlapping worker inflates. - replay_one reads backend.token_delta() (falling back to the text-length heuristic for backends that don't track tokens). - Add tests for call-local and thread-isolated token deltas. * fix(sleep): never cache empty results; add Pi concurrency test Address maintainer review on #251 (deepen thread-safety): - _cached_call no longer caches empty (transient-failure) results and prefers a concurrently cached success, so an empty/duplicate cannot clobber or delete another worker's successful entry. - Add a Pi subclass-level concurrency test (barrier-forced empty-vs-success) asserting the success survives. * fix(sleep): key-aware redaction + correct concurrent token accounting - Redefine _redact_deep to delegate to the key-aware redact_secrets walker so {"api_key": "x"} is scrubbed at every boundary (--json, digests/snapshot files, gate_trials, extra, display), not just the --json link. - Reset _thread_local.delta on cache hit so a later hit doesn't reuse the previous call's delta. - Charge every real call's tokens on a concurrent miss (the dedup worker used to be free, undercounting). - Add DualBackend.token_delta() so replay_one() reads the target's call cost. - Regressions: cache-hit delta reset, barrier-forced concurrent charge, DualBackend token_delta, key-aware _redact_deep. * test(sleep): raise barrier wait timeout to avoid slow-CI flakiness - The barrier-forced concurrency tests waited 5s for all workers to reach the barrier; under a slow/loaded CI that can break the barrier mid-test and turn a pass into a spurious failure. Raise the wait to 15s (no semantic change). * docs(sleep): clarify DualBackend.token_delta intent Document that token_delta() is target-only by design (replay drives the target), that the optimizer only appears in replay via the rare model-judge fallback (rule/exact/answer tasks are scored locally, 0 tokens), and that the aggregate tokens_used() still counts both sub-backends so the total is not undercounted. * fix(sleep): call-local accounting for attempt_with_tools + refresh redaction test - Set the thread-local delta in every attempt_with_tools override that charged _tokens directly (Claude CLI, OpenCode, Codex, Cursor), so replay_one() reads real call-local usage instead of falling back to a response-length estimate. - Update the stale test_redact_deep_loses_mapping_key_context to assert redaction (it was asserting the old leak bug). - Add tool-replay regressions: attempt_with_tools sets call-local delta, and the dual-backend path surfaces the target's delta. * refactor(sleep): centralize token accounting in _record_cost - Add CliBackend._record_cost(prompt, response) as the single path to charge an inference's token cost (aggregate _tokens under _lock + call-local _thread_local.delta), so no path under- or over-counts. - Route _cached_call miss, all attempt_with_tools overrides (Claude/OpenCode/Codex/Cursor), and reflect through it; the OpenCode error path keeps its prompt-only charge. - No behavior change: identical delta model, just centralized — removing the duplicated len//4 accounting that caused the #251-class bugs to recur. * fix(sleep): reset call-local delta on no-call paths + one locked accounting helper Address the review: all token charging now routes through one locked helper (_record_cost), and the call-local delta is reset on every no-call / early-return path (_reset_call_delta), so a reused worker never reports a previous call's token count — covering the OpenCode disabled-tool-replay early return, cache hits, and the start of every tool-aware path. Add regressions: barrier-forced concurrent _record_cost (no lost updates) and same-thread prior-call -> disabled-tool-replay delta reset. * fix(sleep): record call-local delta on real-usage (Azure/OpenCode) paths - Add _record_delta(delta) and route the Azure/OpenCode real-usage accounting through it, so those backends also set the call-local _thread_local.delta — the last accounting path that did not (replay_one() saw 0/stale and fell back to a length estimate for these backends). * fix(sleep): single-owner token accounting, no double charge A backend that self-reports provider usage (AzureOpenAI/AzureResponses) must not be charged twice: _cached_call/_reflect now skip the len//4 estimate when the call already recorded its own exact usage (via a _charged_in_call marker), so a 30-token provider usage is charged once as 30, not as the ~110 estimate. - AzureOpenAI/AzureResponses _call accumulate usage across every paid attempt and record _record_delta(total) once; the call-local delta equals that increment. - OpenCodeCliBackend.attempt_with_tools error path routes prompt-only cost through _record_delta (no manual _tokens / _thread_local.delta update). - The _call str return contract is preserved, so llm_miner/rollout/slow_update keep working and their Azure cost stays accounted. - Added fake-client regressions for both Azure backends (single call, empty-retry + success accumulation, cache-hit no-charge) + OpenCode error-path. * refactor(sleep): thread-localize the self-charge marker Move _charged_in_call from a shared instance attribute to _thread_local.charged_in_call so parallel workers can never read another worker's marker; the marker already gets reset to False before every _call on the current thread. * fix(sleep): finalize accumulated usage on every exhausted-retry exit AzureOpenAI._call recorded usage_total only when last_exc was 'empty-response'; if the final attempt raised (after a paid empty response), it returned without _record_delta, so _cached_call substituted its len//4 estimate and dropped the known paid usage. Now always record the accumulated usage + set the marker on any exhausted-retry exit (usage_total 0 = genuinely no paid usage). Regression: empty-paid-response -> terminal-error keeps the exact 7 tokens. * fix(sleep): backward-compatible call-local accounting for tokens_used()-only backends replay_one no longer substitutes a text-length estimate when a backend implements only the older tokens_used() contract (no token_delta): it reports the same-thread before/after total difference instead, and DualBackend does the same for a target that lacks token_delta (snapshotted in attempt/attempt_with_tools). CliBackend and DualBackend-with-modern-target keep the thread-safe per-call delta, so the shared total is still never reduced by subtraction for concurrent CliBackend calls. Added existing-backend regressions (direct + DualBackend, tokens_used-only). * refactor(sleep): drop unused _cache_pop; note DualBackend snapshot not thread-safe * fix(sleep): report exact zero cost for cached replay results A replay served from the backend cache has no new model call: _cached_call() resets the per-thread delta to zero and the aggregate ledger is unchanged. replay_one() was still treating any zero as 'backend does not track tokens' and substituting a text-length estimate, so ReplayResult.tokens disagreed with the backend ledger (102 vs 0 in reproduction). With token_delta() present, the returned cost is authoritative - including a known zero cache hit. Only backends without the per-call interface fall back to before/after totals and the length estimate. * docs(sleep): spell out the parallel-replay accounting contract The cumulative-only fallback in DualBackend.token_delta() is sequential-only: `_target_tokens_before` is shared instance state, so two workers overlapping in one DualBackend overwrite each other's baseline. That is not the same guarantee as a target that reports a thread-local per-call delta, and it is not the same as a bare legacy backend (whose before/after snapshot replay_one takes on the worker's own thread). replay_batch()'s docstring now lists the four supported combinations, the DualBackend constructor note and its token_delta() fallback say which branch is sequential-only, and replay_one()'s comment no longer implies that any token_delta()-bearing backend is thread-safe. Tests: a regression for the modern path - replay_batch(workers=4) through a DualBackend over a target with a thread-local delta must produce exactly the same per-task costs as the sequential run, with the tasks spread far enough apart in size that the check is not vacuous. * docs(sleep): correct the parallel-replay claim for cumulative-only backends The first version of this contract filed a bare backend with only the older tokens_used() contract under "safe to share". It is not: replay_one differences a counter that the other workers are also spending, so an overlapping worker's tokens land in this task. Two workers against a backend where every attempt costs 37 report [74, 37]. Only a thread-local per-call delta is safe to share; every cumulative-total shape over-counts when workers overlap. The docstring now says that and names the failure direction - budgets see more spend than actually happened - rather than implying that legacy backends are fine. The parallel regression did not bite either: with no real overlap, a backend that had lost its thread-locality still produced costs identical to the sequential run. The fake target now blocks on a barrier until every worker is in flight, so the assertion fails when that locality goes away. Verified by substituting a cumulative delta (test fails) and by five consecutive runs with the real one (test passes). * fix(sleep): 区分用量缺失与权威零,避免误记为 0 provider 的 usage 为 None 时旧代码仍把 charged_in_call 置为 True, _cached_call 因此跳过 len//4 回退,聚合用量与 call-local 增量双双归零, 一次真实调用被记成免费。现在把是否上报用量与上报的数值分开判断: usage 缺失落到长度估算,上报 0 仍按权威零计费,多轮重试中已知的部分 用量继续保留。两个 backend(chat completions 与 responses)各自补齐 缺失、权威零、部分重试三类回归。 * fix(sleep): keep unknown usage distinct from zero on exhausted Azure retries A run of empty responses that all omit usage metadata exhausted the retry budget without any attempt reporting usage; the exit charged usage_total (0) and set charged_in_call, so unknown usage was recorded as an authoritative zero. Charge the accumulated provider usage only when at least one attempt actually reported it; otherwise leave the marker clear so _cached_call's len//4 estimate stands in for the unknown cost, as it does on main. Applies to AzureOpenAIBackend (Chat Completions) and AzureResponsesBackend (Responses); a reported zero still stays an authoritative zero, known partial retry usage is preserved, and cache hits remain unchanged. Regressions: exhausted-empty-without-usage and exhausted-empty-with-reported- zero for both API variants. Also drops an unused local in the OpenCode error path and tidies the new test files' imports (no behavior change). --------- Co-authored-by: WODE25500 <[email protected]> Co-authored-by: WODE25500 <[email protected]>
1 parent 343db22 commit 50d9d90

7 files changed

Lines changed: 1053 additions & 43 deletions

File tree

‎skillopt_sleep/__main__.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
json_safe,
4545
latest_staging,
4646
pending_staged_skills,
47+
redact_secrets,
4748
staged_skills,
4849
)
4950
from skillopt_sleep.staging import adopt as adopt_staging
@@ -333,15 +334,16 @@ def _handoff_dir_for(cfg) -> str:
333334

334335

335336
def _redact_deep(obj):
336-
"""Redact secret-looking substrings in every string of a JSON-like tree."""
337+
"""Redact secrets key-aware across the whole structure (see redact_secrets).
338+
339+
This used to recurse values and only scrub string leaves, losing the
340+
mapping-key context — so ``{"api_key": "x"}`` leaked. Delegating to the
341+
key-aware ``redact_secrets`` walker fixes every output boundary that routes
342+
through this helper (--json, digests/snapshot files, gate_trials, extra,
343+
display) at once, keeping them all consistent.
344+
"""
337345
from skillopt_sleep.staging import redact_secrets
338-
if isinstance(obj, str):
339-
return redact_secrets(obj)
340-
if isinstance(obj, list):
341-
return [_redact_deep(x) for x in obj]
342-
if isinstance(obj, dict):
343-
return {k: _redact_deep(v) for k, v in obj.items()}
344-
return obj
346+
return redact_secrets(obj)
345347

346348

347349
def _display_error(exc: object) -> str:
@@ -483,7 +485,7 @@ def _handoff_mine_and_pin(cfg, args, backend, snapshot: str, dry: bool):
483485
# NOT marked reviewed: feeding this snapshot back through --tasks-file
484486
# with a real backend must still hit the human-review gate above. The
485487
# driver itself loads it directly, with the same trust as in-cycle mining.
486-
write_tasks_file(snapshot, _redact_deep(payload))
488+
write_tasks_file(snapshot, redact_secrets(payload))
487489
print(
488490
f"[sleep] handoff: pinned {len(tasks)} tasks -> {snapshot}",
489491
file=sys.stderr if args.json else sys.stdout,
@@ -807,9 +809,9 @@ def cmd_harvest(args) -> int:
807809
)
808810
output_path = ""
809811
if getattr(args, "output", ""):
810-
output_path = write_tasks_file(args.output, payload)
812+
output_path = write_tasks_file(args.output, redact_secrets(payload))
811813
if args.json:
812-
json_payload = dict(payload)
814+
json_payload = redact_secrets(payload)
813815
if output_path:
814816
json_payload["output"] = output_path
815817
print(json.dumps(json_payload, ensure_ascii=False, indent=2))

‎skillopt_sleep/backend.py‎

Lines changed: 223 additions & 25 deletions
Large diffs are not rendered by default.

‎skillopt_sleep/replay.py‎

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,16 +36,35 @@ def replay_one(backend: Backend, task: TaskRecord, skill: str, memory: str,
3636
tools = _required_tools(task)
3737
tools_called: List[str] = []
3838
t0 = time.time()
39-
tok_before = backend.tokens_used()
39+
# Backends that only expose the older tokens_used() contract (no per-call
40+
# token_delta) need a before/after difference on the same thread to report
41+
# the call-local cost; snapshot the total before the attempt for them.
42+
token_delta_fn = getattr(backend, "token_delta", None)
43+
tokens_before = None if token_delta_fn is not None else backend.tokens_used()
4044
if tools:
4145
response, tools_called = backend.attempt_with_tools(task, skill, memory, tools)
4246
else:
4347
response = backend.attempt(task, skill, memory, sample_id=sample_id)
4448
latency_ms = (time.time() - t0) * 1000.0
45-
tokens = max(0, backend.tokens_used() - tok_before)
46-
# if the backend doesn't track tokens (e.g. mock), approximate from text length
47-
if tokens == 0:
48-
tokens = (len(skill) + len(memory) + len(task.intent) + len(response)) // 4
49+
# Call-local token accounting: prefer the backend's per-call delta. For
50+
# CliBackend - and for a DualBackend whose target has one - that delta lives
51+
# in thread-local storage and is exact, including a known zero for a cache
52+
# hit, so no text-length estimate is substituted. A backend that only
53+
# implements the older tokens_used() contract has no per-call delta, so use
54+
# its before/after difference. That difference is only call-local while the
55+
# backend is used sequentially: the counter is shared, so an overlapping
56+
# worker's spend lands in it too. Fall back to a length estimate only when
57+
# the backend reports no tracking (zero) at all. See replay_batch() for which
58+
# combinations are safe to share across workers.
59+
if token_delta_fn is not None:
60+
tokens = token_delta_fn()
61+
else:
62+
tokens = max(0, backend.tokens_used() - (tokens_before or 0))
63+
# A backend without token_delta() that still changed tokens_used()
64+
# reports that real difference above. Only an unchanged total means the
65+
# backend does not track tokens (e.g. a mock), so approximate then.
66+
if tokens == 0:
67+
tokens = (len(skill) + len(memory) + len(task.intent) + len(response)) // 4
4968

5069
# rule judges may need the detected tool calls; score locally when possible
5170
if task.reference_kind == "rule" and task.judge:
@@ -109,6 +128,29 @@ def replay_batch(
109128
big test sets (like the research harness's --workers). ``workers`` defaults
110129
to env SKILLOPT_SLEEP_WORKERS or 1 (sequential). Mock stays sequential
111130
(deterministic) unless asked otherwise.
131+
132+
Every worker shares the one ``backend`` passed in, so per-call token
133+
accounting has to be call-local. Which combinations support that:
134+
135+
* ``CliBackend`` and subclasses - per-call delta in thread-local storage,
136+
safe to share across workers.
137+
* ``DualBackend`` whose target defines ``token_delta()`` - reports that
138+
target's thread-local delta, safe to share across workers.
139+
* a bare backend with only the older ``tokens_used()`` contract - sequential
140+
only: ``replay_one`` differences a counter the other workers are also
141+
spending, so an overlapping worker's tokens are counted into this task.
142+
* ``DualBackend`` whose target only has ``tokens_used()`` - sequential only,
143+
for that reason plus a baseline that lives on the shared DualBackend.
144+
145+
Only a thread-local per-call delta is safe to share. Everything built on a
146+
cumulative total - the bare legacy backend as much as the DualBackend
147+
fallback - over-counts when workers overlap: a 37-token attempt reported 74
148+
with two workers, because one worker's spend landed in the other's
149+
difference. Over-counting is the failure direction, so token and cost
150+
budgets see more spend than actually happened. Both shapes are supported
151+
sequentially and covered by tests there; neither is claimed here. Use
152+
``workers=1``, or a backend that reports a per-call delta, when the
153+
accounting has to be exact.
112154
"""
113155
if workers <= 0:
114156
workers = int(os.environ.get("SKILLOPT_SLEEP_WORKERS", "1") or "1")

‎skillopt_sleep/staging.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1130,13 +1130,13 @@ def write_staging(
11301130
(
11311131
os.path.join(out, "report.json"),
11321132
json.dumps(
1133-
json_safe(report.to_dict()),
1133+
json_safe(redact_secrets(report.to_dict())),
11341134
ensure_ascii=False,
11351135
indent=2,
11361136
allow_nan=False,
11371137
),
11381138
),
1139-
(os.path.join(out, "report.md"), report_md),
1139+
(os.path.join(out, "report.md"), redact_secrets(report_md)),
11401140
# The manifest is the publication marker and must always be last.
11411141
(
11421142
os.path.join(out, "manifest.json"),

0 commit comments

Comments
 (0)