Skip to content

Commit c8ff7f8

Browse files
committed
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.
1 parent f3ef661 commit c8ff7f8

2 files changed

Lines changed: 52 additions & 1 deletion

File tree

‎skillopt_sleep/backend.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -360,14 +360,23 @@ def _record_cost(self, prompt: str, response: str) -> int:
360360
self._thread_local.delta = delta
361361
return delta
362362

363+
def _reset_call_delta(self) -> None:
364+
"""Zero the call-local delta for a NO-CALL path.
365+
366+
Called at the start of every tool-aware path (and used by cache hits) so
367+
a reused worker never reports a previous call's token count — every
368+
no-call / early-return path leaves the delta at 0.
369+
"""
370+
self._thread_local.delta = 0
371+
363372
def _cached_call(self, key: str, prompt: str, *, max_tokens: int = 1024) -> str:
364373
kind = key.split(":", 1)[0]
365374
ev = getattr(self, "evidence", None)
366375
with self._lock:
367376
cached = self._cache.get(key)
368377
if cached is not None:
369378
# cover: later cache hit must not report the previous call's delta.
370-
self._thread_local.delta = 0
379+
self._reset_call_delta()
371380
# cache hits log key-only (the full text is on the original miss event)
372381
if ev is not None:
373382
ev.log("replay", "model_call", kind=kind, cache_hit=True, key=key,
@@ -865,6 +874,7 @@ def _call(self, prompt: str, *, max_tokens: int = 1024) -> str:
865874
return out
866875

867876
def attempt_with_tools(self, task, skill, memory, tools):
877+
self._reset_call_delta()
868878
# Expose a REAL, callable `search` tool (a shell shim that logs each
869879
# call) so the gbrain quick-answerer judge (tool_called=search) is
870880
# validated honestly: we detect the call from the shim's log, not from
@@ -1356,6 +1366,7 @@ def attempt_with_tools(
13561366
tools: List[str],
13571367
) -> Tuple[str, List[str]]:
13581368
self.last_call_error = ""
1369+
self._reset_call_delta()
13591370
if not self.tool_replay:
13601371
self.last_call_error = (
13611372
"OpenCode CLI tool replay is not supported without explicit "
@@ -1616,6 +1627,7 @@ def _call(self, prompt: str, *, max_tokens: int = 1024, retries: int = 3) -> str
16161627
return out
16171628

16181629
def attempt_with_tools(self, task, skill, memory, tools):
1630+
self._reset_call_delta()
16191631
# Codex exec runs in a sandbox with shell access; expose the same real
16201632
# `search` shim and let it run (workspace-write so the shim can log).
16211633
import tempfile, shutil, stat
@@ -1880,6 +1892,7 @@ def _parse_jsonl_response(raw: str) -> str:
18801892
return "\n".join(parts).strip()
18811893

18821894
def attempt_with_tools(self, task, skill, memory, tools):
1895+
self._reset_call_delta()
18831896
# Expose REAL, callable tool shims in the working directory so the
18841897
# gbrain quick-answerer judge (tool_called=search) is validated
18851898
# honestly: we detect each call from the shim's log, not from a

‎tests/test_cli_backend_cache.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,44 @@ def fake_run(*args, **kwargs):
223223
assert db.token_delta() > 0, "dual-backend tool replay did not set target delta"
224224

225225

226+
def test_record_cost_concurrent_no_lost_updates():
227+
"""Barrier-forced concurrent _record_cost calls must not lose updates."""
228+
n = 30
229+
barrier = threading.Barrier(n)
230+
231+
class _BarrierEcho(_EchoBackend):
232+
def _cost_worker(self):
233+
barrier.wait(timeout=15)
234+
self._record_cost("hello", "world")
235+
236+
b = _BarrierEcho()
237+
threads = [threading.Thread(target=b._cost_worker) for _ in range(n)]
238+
for t in threads:
239+
t.start()
240+
for t in threads:
241+
t.join()
242+
per = len("hello") // 4 + len("world") // 4
243+
assert b._tokens == per * n, f"concurrent _record_cost lost updates: {b._tokens} != {per * n}"
244+
245+
246+
def test_disabled_tool_replay_resets_call_delta():
247+
"""A prior call's delta must not leak through a disabled-tool-replay path."""
248+
from types import SimpleNamespace
249+
250+
from skillopt_sleep.backend import OpenCodeCliBackend
251+
252+
b = OpenCodeCliBackend(model="", opencode_path="opencode", tool_replay=False)
253+
# A prior call on this thread set a nonzero call-local delta.
254+
b._record_cost("hello", "world")
255+
assert b.token_delta() > 0 # stale from the prior call
256+
# Disabled-tool-replay attempt_with_tools returns early; the delta must be
257+
# reset to 0 so a reused worker does not report the previous call's cost.
258+
task = SimpleNamespace(intent="intent", context_excerpt="ctx")
259+
out, called = b.attempt_with_tools(task, skill="s", memory="m", tools=["search"])
260+
assert out == "" and called == []
261+
assert b.token_delta() == 0, "disabled-tool-replay leaked the prior call's delta"
262+
263+
226264
def test_cmd_harvest_redact_deep_is_key_aware():
227265
"""_redact_deep must key-aware redact — `{"api_key": "x"}` used to leak."""
228266
from skillopt_sleep.__main__ import _redact_deep

0 commit comments

Comments
 (0)