Skip to content

fix(sleep): thread-safe backend cache + redact exports - #251

Open
WODE25500 wants to merge 14 commits into
microsoft:mainfrom
WODE25500:fix/sleep-hardening
Open

fix(sleep): thread-safe backend cache + redact exports#251
WODE25500 wants to merge 14 commits into
microsoft:mainfrom
WODE25500:fix/sleep-hardening

Conversation

@WODE25500

Copy link
Copy Markdown
Contributor

Sleep-cycle hardening.

  • Guard CliBackend _cache/_tokens with a lock on the opt-in parallel replay path (SKILLOPT_SLEEP_WORKERS>1); the model call stays outside the lock so parallel workers still overlap.
  • Redact report.json before staging (redact_secrets).
  • Redact harvest --output and --json exports (_redact_deep).
  • Add tests for CliBackend caching/thread-safety.

- 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.
@Yif-Yang

Copy link
Copy Markdown
Contributor

report.json is now correctly passed through the mapping-aware redactor, but the export-redaction and thread-safety fixes are still incomplete in three places.

  1. Harvest output uses _redact_deep(payload), while _redact_deep() only recurses into values and loses the mapping-key context. As a result, _redact_deep({"api_key": "plain-secret", "nested": {"token": "other"}}) returns both secrets unchanged. Please use the existing mapping-key-aware redact_secrets(payload) at every structured output boundary.
  2. write_staging() redacts report.json but writes caller-provided report_md verbatim. Edit content/rationale can therefore still expose credentials. Please redact the Markdown before writing it as well.
  3. Not all cache/token access uses the new lock: the Pi and OpenCode overrides still inspect and pop _cache directly, and tokens_used() reads _tokens without the lock. More importantly, parallel replay_one() attributes one call's tokens from a shared global before/after total, so overlapping workers can charge another worker's tokens to the wrong result. Please route cache/token access through locked helpers, prevent a failed caller from deleting another caller's successful cache entry, and use call-local accounting for per-result tokens.

Please add boundary-level tests for both harvest output/file and report.md, including nested api_key/token mappings. The concurrency tests should use barriers/events to force overlapping misses and assert exact call/cache/token outcomes, including a concurrent Pi/OpenCode empty-result versus successful-result case. The current immediate echo test does not guarantee overlap and can pass without exercising these races.

Address maintainer review on microsoft#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).
Address maintainer review on microsoft#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.
Address maintainer review on microsoft#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.
Address maintainer review on microsoft#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.
@WODE25500

WODE25500 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

Yifan Yang (@Yif-Yang) — thank you for the thorough review and guidance! I've addressed the feedback:

  • Structured outputs now use the mapping-key-aware redact_secrets (not _redact_deep, which loses the key context); report_md is also redacted before writing.
  • Thread-safety hardened: cache/token access goes through locked helpers (_cache_get/_cache_pop/_cache_pop_if); tokens_used() is locked; a failed caller only conditionally pops its own empty value (_cache_pop_if), never another worker's success; the base no longer caches empty (transient-failure) results.
  • Switched to call-local token accounting (thread-local delta) so parallel replay no longer misattributes tokens through a global before/after total.
  • Added tests: nested api_key/token redaction, report.md, barrier-forced concurrency, pop-if not deleting a success, Pi subclass empty-vs-success, and thread-isolated token deltas — all pass.

Thanks again for the detailed review!

@WODE25500

WODE25500 commented Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

Yifan Yang (@Yif-Yang) — thank you for the careful review and guidance. In fact I've done almost all of these submissions through DSH, which is exactly why I have a bold idea. Across your recent PR reviews I noticed that you consistently apply a set of quality baselines: fail-closed handling, structure-aware and boundary-consistent redaction, thread-safety with call-local accounting, validating against real contracts, PR hygiene, and so on. I'd like to distill that into a reusable review-standards: a general quality standard plus a pre-PR self-check CLI that contributors run before submitting, so they spend less time on rework and you receive fewer low-quality PRs — a win-win. I'd credit you as the original source of these standards. Looking forward to hearing from you whenever you get a chance.

@Yif-Yang

Copy link
Copy Markdown
Contributor

Thanks — several original races and the report.md boundary are fixed, but four correctness gaps remain on this head.

  1. cmd_harvest --json still calls _redact_deep(payload). That function recurses into values and loses mapping-key context, so {"api_key":"top-secret"} is still emitted unchanged on JSON stdout. Please use the mapping-aware redact_secrets(payload) at this final boundary too.
  2. A cache hit returns before resetting self._thread_local.delta. Reusing a worker thread after a real call therefore makes a later cache hit report the previous call's token delta.
  3. Concurrent misses for the same key both make paid model calls, but when one worker finds the other's cached success it sets its own delta = 0 and does not add that real call to _tokens. The cache contents are safe, but actual usage is undercounted unless in-flight calls are coalesced or every real call is charged.
  4. DualBackend has no token_delta(). Consequently replay_one() falls back to a response-length estimate and loses the target backend's real call-local cost.

Please add boundary-level stdout coverage for mapping-key secrets, a same-thread miss→hit regression, a barrier-forced same-key concurrent-miss test that checks paid-call accounting, and a DualBackend replay accounting test. These are the remaining blockers; the other changes from the previous review look addressed.

- 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.
- 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).
@WODE25500

WODE25500 commented Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Already addressed the review feedback and updated this branch (#251):

  • Commits 37fe5b6 / d555992: _redact_deep now delegates to the key-aware redact_secrets walker (fixing every output boundary: --json, digests/snapshot files, gate_trials, extra, display); cache hits reset the thread-local delta; every real call on a concurrent miss is charged; DualBackend gained token_delta(); barrier tests added (with a longer timeout to avoid slow-CI flakiness).
    Please re-review, thanks.

Also added a comment on DualBackend.token_delta() documenting that target-only is intentional: replay drives the target, and the optimizer only appears via the rare model-judge fallback (rule/exact/answer tasks are scored locally); the aggregate tokens_used() still counts both, so the total is not undercounted.

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.
@Yif-Yang

Copy link
Copy Markdown
Contributor

Thanks — the original cache-hit, concurrent-call, boundary-redaction, and DualBackend changes are mostly addressed. Two blockers remain. First, the branch’s own full suite fails at tests/test_export_redaction.py::test_redact_deep_loses_mapping_key_context; that stale regression still asserts that the secret leaks and must be updated to assert redaction. Second, real tool-aware replay overrides update _tokens directly but do not set the new thread-local delta, so replay_one() receives zero or stale call-local usage and falls back to a response-length estimate. Please make call-local accounting cover attempt_with_tools() as well and add a tool-replay regression, including the dual-backend path.

…daction 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.
@WODE25500

Copy link
Copy Markdown
Contributor Author

Thanks for the re-review. Addressed both blockers: the stale est_redact_deep_loses_mapping_key_context now asserts redaction (was asserting the leak), and every �ttempt_with_tools override that charged _tokens now sets the call-local thread delta (Claude CLI, OpenCode, Codex, Cursor), so
eplay_one() sees real call-local usage. Added tool-replay regressions incl. the dual-backend path. Commit 3c6f95.

- 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 microsoft#251-class bugs to recur.
@Yif-Yang

Copy link
Copy Markdown
Contributor

Thanks — the tool-aware paths now set a call-local delta, but two accounting issues remain. OpenCodeCliBackend.attempt_with_tools() returns early when tool replay is disabled without resetting the thread-local delta, so a reused worker can report the previous call token count. Also, the Claude, OpenCode, Codex, and Cursor tool-aware paths still update _tokens directly outside _lock, unlike _cached_call(). Please route all token charging and per-call delta updates through one locked helper, reset the delta on every no-call/early-return path, and add a same-thread prior-call → disabled-tool-replay regression plus barrier-forced concurrent tool-call accounting coverage.

@WODE25500

Copy link
Copy Markdown
Contributor Author

Thanks to Teacher Yifan for the careful review, please wait a moment.

…unting 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.
- 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).
@WODE25500

Copy link
Copy Markdown
Contributor Author

Thanks for the re-review - both accounting issues are resolved (commits 3ef661, c8ff7f8, 8a17aa6):

  • All token charging now goes through one locked helper _record_cost (aggregate _tokens under _lock + call-local _thread_local.delta), so no path updates _tokens outside the lock.
  • The call-local delta is reset on every no-call / early-return path (_reset_call_delta), covering the OpenCode disabled-tool-replay early return and cache hits - a reused worker can no longer report the previous call's cost.
  • The Azure/OpenCode real-usage paths also record call-local delta via _record_delta.
  • Added regressions: barrier-forced concurrent _record_cost (no lost updates) and same-thread prior-call ? disabled-tool-replay delta reset.

@Yif-Yang

Copy link
Copy Markdown
Contributor

Thanks — the previous early-return and locking issues are fixed, but one accounting blocker remains on 8a17aa6.

CliBackend._cached_call() in skillopt_sleep/backend.py always calls _record_cost(prompt, out) after _call() (lines 393–395). However, AzureOpenAIBackend._call() already records provider usage through _record_delta() (lines 2476–2483), and AzureResponsesBackend._call() does the same (lines 2584–2595). Each successful Azure call is therefore charged twice, and the exact call-local usage is overwritten by the length estimate.

Minimal reproduction with a fake chat-completions response reporting 10 prompt + 20 completion tokens, a 400-character prompt, and a 40-character response:

provider_usage=30
tokens_used=140
token_delta=110

A two-attempt reproduction (7 tokens on an empty response, then 11 on success) similarly reports:

provider_usage=18
tokens_used=38
token_delta=20

Thus both aggregate cost and ReplayResult.tokens are incorrect. The newly added focused suites still pass (19 passed), because they only exercise length-estimated CliBackend implementations.

Please establish one owner for accounting: preferably have _call() return its provider-reported usage and let _cached_call() record it once, falling back to the length estimate only when usage is unavailable. Retry usage should be accumulated across every paid attempt, and the final call-local delta should equal that same accumulated increment.

Please add fake-client regressions for both AzureOpenAIBackend and AzureResponsesBackend, covering a single successful call, an empty-response retry followed by success, and a cache hit (delta == 0, aggregate unchanged).

Also, OpenCodeCliBackend.attempt_with_tools() still updates _tokens and _thread_local.delta manually on its error path (lines 1440–1445). Please route that path through _record_delta() as well and cover it with an error-path regression.

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.
@WODE25500
WODE25500 force-pushed the fix/sleep-hardening branch from 0712e44 to c72b1b4 Compare August 30, 2026 21:59
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.
@WODE25500

Copy link
Copy Markdown
Contributor Author

Yifan Yang (@Yif-Yang) — addressed the remaining accounting blocker on ef7b806.

Summary:

  • Single owner, no double charge: a backend that self-reports provider usage (AzureOpenAI / AzureResponses) records its exact, accumulated usage once inside _call and flips a per-thread marker; _cached_call and _reflect then skip their len//4 estimate. Previously the estimate ran on top of the provider usage, so a 30-token usage was charged as ~110 and the exact call-local delta was overwritten.
  • Retry accumulation: Azure _call sums usage across every paid attempt; the final call-local delta equals that accumulated increment (e.g. empty-then-success → 7 + 30 = 37).
  • Cache hit: no charge, delta reset to 0, aggregate unchanged.
  • OpenCode error path: attempt_with_tools now routes the prompt-only cost through _record_delta (no more manual _tokens / _thread_local.delta update).
  • Thread-safety: the marker is thread-local (_thread_local.charged_in_call), so parallel workers never read another worker's marker.

Added fake-client regressions (tests/test_azure_usage_accounting.py) for AzureOpenAI and AzureResponses covering a single successful call, an empty-response retry followed by success, and a cache hit, plus an OpenCode error-path regression — 7/7 pass; the full suite is green (1372 passed).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants