Skip to content

Commit daa20a6

Browse files
AlexStocksOmX
andauthored
fix(workbuddy): keep prompt hooks from stalling on scope reads (#1729)
* fix(workbuddy): keep prompt hooks from stalling on scope reads WorkBuddy prompt submission should fail open within the scope binding HTTP budget even when a server response stalls after headers arrive. Constraint: WorkBuddy enforces an outer UserPromptSubmit hook timeout Scope-risk: narrow Tested: uv run pytest tests/workbuddy_plugin/test_scope_binding.py -q Tested: uv run ruff check integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py tests/workbuddy_plugin/conftest.py tests/workbuddy_plugin/test_scope_binding.py Tested: uv run ruff format --check integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py tests/workbuddy_plugin/conftest.py tests/workbuddy_plugin/test_scope_binding.py Co-authored-by: OmX <omx@oh-my-codex.dev> * fix(workbuddy): bound hook response reads at every socket receive f42a2bf added an absolute deadline around the scope resolver's response body, but it still tightened urllib's socket timeout only once per 64 KiB block read. http.client consumes many socket reads inside a single read call, so a server that trickles a byte faster than that timeout keeps the caller inside one read and past its deadline; the observed worst case is the host's 30s UserPromptSubmit timeout, which is exactly what the change set out to remove. Reuse the primitives the claude-code and codex scope modules already ship: - open_bounded runs the request in a daemon worker, so connect and header exchange stay inside the caller's budget. - bind_response_deadline recomputes the remaining time before every receive, framing included. Both now cover the scope resolver and the hook's own recall, capture, and flush reads. The hook shares them through the installed powercontext_scope_binding shim instead of carrying its own opener, so the shim forwards them too. _set_response_timeout reads the socket through a guarded getattr chain, so a missing fp can no longer raise AttributeError outside the scope client's error contract. The deadline test replaces its fake socket with real loopback properties: body drip, header drip, a large healthy response, an expired deadline, and the hook's own recall read. The three drip cases fail on f42a2bf (3.19s, 3.17s, 3.20s against a 1.0s assertion) and pass here. --------- Co-authored-by: OmX <omx@oh-my-codex.dev>
1 parent 2f1a76b commit daa20a6

7 files changed

Lines changed: 562 additions & 43 deletions

File tree

‎integrations/workbuddy/plugins/powercontext/hooks/powercontext_scope_binding.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,6 @@
2020
_SCRIPTS_ROOT = Path(__file__).resolve().parents[1] / "scripts"
2121
sys.path.insert(0, str(_SCRIPTS_ROOT))
2222

23-
from workspace_scope import resolve_scope_id # noqa: E402
23+
from workspace_scope import bind_response_deadline, open_bounded, resolve_scope_id # noqa: E402
2424

25-
__all__ = ["resolve_scope_id"]
25+
__all__ = ["bind_response_deadline", "open_bounded", "resolve_scope_id"]

‎integrations/workbuddy/plugins/powercontext/hooks/workbuddy_powercontext_hook.py‎

Lines changed: 8 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,9 @@
2727
from hashlib import sha256
2828
from pathlib import Path
2929
from time import monotonic
30-
from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast
30+
from typing import Any, Protocol, cast
3131
from urllib.error import HTTPError
32-
from urllib.request import HTTPRedirectHandler, Request, build_opener
33-
34-
if TYPE_CHECKING:
35-
from typing_extensions import override
36-
else:
37-
_MethodT = TypeVar("_MethodT")
38-
39-
def override(method: _MethodT, /) -> _MethodT:
40-
return method
41-
32+
from urllib.request import Request
4233

4334
_HOOKS_ROOT = Path(__file__).resolve().parent
4435
_PLUGIN_ROOT = _HOOKS_ROOT.parent
@@ -51,9 +42,9 @@ def override(method: _MethodT, /) -> _MethodT:
5142
from workbuddy_settings import WorkBuddyPluginSettings # noqa: E402
5243

5344
if (_HOOKS_ROOT / "powercontext_scope_binding.py").is_file():
54-
from powercontext_scope_binding import resolve_scope_id
45+
from powercontext_scope_binding import bind_response_deadline, open_bounded, resolve_scope_id
5546
else:
56-
from workspace_scope import resolve_scope_id
47+
from workspace_scope import bind_response_deadline, open_bounded, resolve_scope_id
5748

5849
_MAX_CONTEXT_BYTES = _prepared_context.MAX_CONTEXT_BYTES
5950
_InvalidResponseError = _prepared_context.InvalidPreparedContextResponse
@@ -69,7 +60,6 @@ def override(method: _MethodT, /) -> _MethodT:
6960

7061

7162
class _Response(Protocol):
72-
fp: object
7363
status: int
7464

7565
def __enter__(self) -> _Response: ...
@@ -79,25 +69,6 @@ def __exit__(self, *args: object) -> object: ...
7969
def read(self, amount: int = -1) -> bytes: ...
8070

8171

82-
class _RejectRedirects(HTTPRedirectHandler):
83-
"""Leave every 3xx response to urllib's default HTTP error handler."""
84-
85-
@override
86-
def redirect_request(
87-
self,
88-
req: Request,
89-
fp: object,
90-
code: int,
91-
msg: str,
92-
headers: object,
93-
newurl: str,
94-
) -> Request | None:
95-
return None
96-
97-
98-
_URL_OPENER = build_opener(_RejectRedirects)
99-
100-
10172
class _HttpStatusError(RuntimeError):
10273
def __init__(self, status: int) -> None:
10374
self.status = status
@@ -300,7 +271,7 @@ def _post_json(
300271
request_timeout = min(settings.request_timeout_seconds, _remaining_time(deadline))
301272
request_deadline = min(deadline, monotonic() + request_timeout)
302273
try:
303-
with _URL_OPENER.open(request, timeout=request_timeout) as response:
274+
with open_bounded(request, timeout=request_timeout) as response:
304275
if expected_status is not None and response.status != expected_status:
305276
raise _HttpStatusError(response.status)
306277
result = json.loads(_read_response(response, deadline=request_deadline))
@@ -325,6 +296,7 @@ def _request_headers(settings: WorkBuddyPluginSettings) -> dict[str, str]:
325296
def _read_response(response: _Response, *, deadline: float) -> bytes:
326297
"""Read one response under a wall-clock deadline and a hard size bound."""
327298

299+
bind_response_deadline(response, deadline)
328300
content = bytearray()
329301
while True:
330302
_set_response_timeout(response, _remaining_time(deadline))
@@ -344,10 +316,10 @@ def _remaining_time(deadline: float) -> float:
344316
return remaining
345317

346318

347-
def _set_response_timeout(response: _Response, timeout: float) -> None:
319+
def _set_response_timeout(response: object, timeout: float) -> None:
348320
"""Tighten urllib's socket timeout before each bounded read."""
349321

350-
raw = getattr(response.fp, "raw", None)
322+
raw = getattr(getattr(response, "fp", None), "raw", None)
351323
sock = getattr(raw, "_sock", None)
352324
settimeout = getattr(sock, "settimeout", None)
353325
if settimeout is not None:

‎integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py‎

Lines changed: 103 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import os
2323
import subprocess
2424
import sys
25+
import threading
2526
from collections.abc import Mapping, Sequence
2627
from hashlib import sha256
2728
from pathlib import Path
@@ -89,6 +90,34 @@ def redirect_request(
8990
_URL_OPENER = build_opener(_RejectRedirects)
9091

9192

93+
def open_bounded(request: Request, *, timeout: float) -> Any:
94+
"""Open one request under a hard wall-clock bound, response headers included.
95+
96+
urllib applies its timeout to each individual socket read, so a server that
97+
trickles headers can outlive the caller's deadline. Running the open in a
98+
daemon worker and abandoning it on expiry keeps the status line inside its
99+
budget.
100+
"""
101+
102+
outcome: list[Any] = []
103+
104+
def _open() -> None:
105+
try:
106+
outcome.append(_URL_OPENER.open(request, timeout=timeout))
107+
except BaseException as error:
108+
outcome.append(error)
109+
110+
worker = threading.Thread(target=_open, name="powercontext-http", daemon=True)
111+
worker.start()
112+
worker.join(timeout)
113+
if worker.is_alive():
114+
raise TimeoutError
115+
result = outcome[0] if outcome else TimeoutError()
116+
if isinstance(result, BaseException):
117+
raise result
118+
return result
119+
120+
92121
def resolve_scope_id(
93122
cwd: str,
94123
*,
@@ -203,10 +232,12 @@ def _request_json(
203232
method=method,
204233
)
205234
try:
206-
with _URL_OPENER.open(request, timeout=min(settings.request_timeout_seconds, remaining)) as response:
235+
request_timeout = min(settings.request_timeout_seconds, remaining)
236+
request_deadline = min(deadline, monotonic() + request_timeout)
237+
with open_bounded(request, timeout=request_timeout) as response:
207238
if response.status < 200 or response.status >= 300:
208239
raise ScopeBindingError
209-
raw = _read_bounded(response)
240+
raw = _read_bounded(response, deadline=request_deadline)
210241
except (HTTPError, OSError, TimeoutError) as error:
211242
raise ScopeBindingError from error
212243
try:
@@ -218,15 +249,82 @@ def _request_json(
218249
return value
219250

220251

221-
def _read_bounded(response: _Response) -> bytes:
252+
def _remaining_time(deadline: float) -> float:
253+
remaining = deadline - monotonic()
254+
if remaining <= 0:
255+
raise TimeoutError
256+
return remaining
257+
258+
259+
class _DeadlineSocket:
260+
"""Enforce one absolute deadline on every response socket read.
261+
262+
``http.client`` can consume many socket reads inside a single ``read`` call
263+
while it parses chunk framing, so tightening the socket timeout once per
264+
bounded read cannot stop a server that trickles chunk extensions. Recomputing
265+
the timeout before every receive keeps the caller's absolute deadline,
266+
framing included.
267+
"""
268+
269+
def __init__(self, sock: Any, deadline: float) -> None:
270+
self._sock = sock
271+
self._deadline = deadline
272+
273+
def _remaining_time(self) -> float:
274+
remaining = self._deadline - monotonic()
275+
if remaining <= 0:
276+
raise TimeoutError
277+
return remaining
278+
279+
def recv(self, *args: Any) -> Any:
280+
self._sock.settimeout(self._remaining_time())
281+
return self._sock.recv(*args)
282+
283+
def recv_into(self, *args: Any) -> Any:
284+
self._sock.settimeout(self._remaining_time())
285+
return self._sock.recv_into(*args)
286+
287+
def __getattr__(self, name: str) -> Any:
288+
return getattr(self._sock, name)
289+
290+
291+
def bind_response_deadline(response: object, deadline: float) -> None:
292+
"""Keep every socket read of one open response inside the absolute deadline."""
293+
294+
if isinstance(response, HTTPError):
295+
response = response.fp
296+
raw: Any = getattr(getattr(response, "fp", None), "raw", None)
297+
if raw is None:
298+
return
299+
sock = getattr(raw, "_sock", None)
300+
if sock is None or isinstance(sock, _DeadlineSocket) or not hasattr(sock, "recv_into"):
301+
return
302+
raw._sock = _DeadlineSocket(sock, deadline)
303+
304+
305+
def _set_response_timeout(response: object, timeout: float) -> None:
306+
"""Tighten urllib's socket timeout before each bounded read."""
307+
308+
raw = getattr(getattr(response, "fp", None), "raw", None)
309+
sock = getattr(raw, "_sock", None)
310+
settimeout = getattr(sock, "settimeout", None)
311+
if settimeout is not None:
312+
settimeout(timeout)
313+
314+
315+
def _read_bounded(response: _Response, *, deadline: float) -> bytes:
316+
bind_response_deadline(response, deadline)
222317
chunks: list[bytes] = []
223318
size = 0
224-
while chunk := response.read(_READ_CHUNK_BYTES):
319+
while True:
320+
_set_response_timeout(response, _remaining_time(deadline))
321+
chunk = response.read(_READ_CHUNK_BYTES)
322+
if not chunk:
323+
return b"".join(chunks)
225324
size += len(chunk)
226325
if size > _MAX_RESPONSE_BYTES:
227326
raise ScopeBindingError
228327
chunks.append(chunk)
229-
return b"".join(chunks)
230328

231329

232330
def _git_value(cwd: str, *arguments: str) -> str | None:

‎tests/workbuddy_plugin/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright (c) 2026 OceanBase.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Tests for the PowerContext WorkBuddy plugin."""

‎tests/workbuddy_plugin/conftest.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
# Copyright (c) 2026 OceanBase.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import importlib.util
18+
import sys
19+
from collections.abc import Iterator
20+
from pathlib import Path
21+
from types import ModuleType
22+
23+
import pytest
24+
25+
REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
26+
WORKBUDDY_ROOT = REPOSITORY_ROOT / "integrations" / "workbuddy"
27+
PLUGIN_ROOT = WORKBUDDY_ROOT / "plugins" / "powercontext"
28+
_PLUGIN_MODULE_NAMES = (
29+
"powercontext_client_config",
30+
"powercontext_scope_binding",
31+
"prepared_context",
32+
"workbuddy_settings",
33+
"workspace_scope",
34+
)
35+
36+
37+
def _load_module(name: str, path: Path) -> ModuleType:
38+
spec = importlib.util.spec_from_file_location(name, path)
39+
assert spec is not None and spec.loader is not None
40+
module = importlib.util.module_from_spec(spec)
41+
sys.modules[name] = module
42+
try:
43+
spec.loader.exec_module(module)
44+
except Exception:
45+
sys.modules.pop(name, None)
46+
raise
47+
return module
48+
49+
50+
@pytest.fixture
51+
def plugin_imports() -> Iterator[None]:
52+
previous_path = list(sys.path)
53+
previous_modules = {name: sys.modules.get(name) for name in _PLUGIN_MODULE_NAMES}
54+
for name in _PLUGIN_MODULE_NAMES:
55+
sys.modules.pop(name, None)
56+
sys.path.insert(0, str(PLUGIN_ROOT))
57+
try:
58+
yield
59+
finally:
60+
sys.path[:] = previous_path
61+
for name in _PLUGIN_MODULE_NAMES:
62+
sys.modules.pop(name, None)
63+
for name, module in previous_modules.items():
64+
if module is not None:
65+
sys.modules[name] = module
66+
67+
68+
@pytest.fixture
69+
def scope_module(plugin_imports: None) -> ModuleType:
70+
return _load_module("powercontext_workbuddy_scope", PLUGIN_ROOT / "scripts" / "workspace_scope.py")
71+
72+
73+
@pytest.fixture
74+
def hook_module(plugin_imports: None) -> ModuleType:
75+
return _load_module(
76+
"powercontext_workbuddy_hook",
77+
PLUGIN_ROOT / "hooks" / "workbuddy_powercontext_hook.py",
78+
)

0 commit comments

Comments
 (0)