fix(workbuddy): keep prompt hooks from stalling on scope reads

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

@@ -20,6 +20,6 @@
 _SCRIPTS_ROOT = Path(__file__).resolve().parents[1] / "scripts"
 sys.path.insert(0, str(_SCRIPTS_ROOT))
 
-from workspace_scope import resolve_scope_id  # noqa: E402
+from workspace_scope import bind_response_deadline, open_bounded, resolve_scope_id  # noqa: E402
 
-__all__ = ["resolve_scope_id"]
+__all__ = ["bind_response_deadline", "open_bounded", "resolve_scope_id"]

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

@@ -27,18 +27,9 @@
 from hashlib import sha256
 from pathlib import Path
 from time import monotonic
-from typing import TYPE_CHECKING, Any, Protocol, TypeVar, cast
+from typing import Any, Protocol, cast
 from urllib.error import HTTPError
-from urllib.request import HTTPRedirectHandler, Request, build_opener
-
-if TYPE_CHECKING:
-    from typing_extensions import override
-else:
-    _MethodT = TypeVar("_MethodT")
-
-    def override(method: _MethodT, /) -> _MethodT:
-        return method
-
+from urllib.request import Request
 
 _HOOKS_ROOT = Path(__file__).resolve().parent
 _PLUGIN_ROOT = _HOOKS_ROOT.parent
@@ -51,9 +42,9 @@ def override(method: _MethodT, /) -> _MethodT:
 from workbuddy_settings import WorkBuddyPluginSettings  # noqa: E402
 
 if (_HOOKS_ROOT / "powercontext_scope_binding.py").is_file():
-    from powercontext_scope_binding import resolve_scope_id
+    from powercontext_scope_binding import bind_response_deadline, open_bounded, resolve_scope_id
 else:
-    from workspace_scope import resolve_scope_id
+    from workspace_scope import bind_response_deadline, open_bounded, resolve_scope_id
 
 _MAX_CONTEXT_BYTES = _prepared_context.MAX_CONTEXT_BYTES
 _InvalidResponseError = _prepared_context.InvalidPreparedContextResponse
@@ -69,7 +60,6 @@ def override(method: _MethodT, /) -> _MethodT:
 
 
 class _Response(Protocol):
-    fp: object
     status: int
 
     def __enter__(self) -> _Response: ...
@@ -79,25 +69,6 @@ def __exit__(self, *args: object) -> object: ...
     def read(self, amount: int = -1) -> bytes: ...
 
 
-class _RejectRedirects(HTTPRedirectHandler):
-    """Leave every 3xx response to urllib's default HTTP error handler."""
-
-    @override
-    def redirect_request(
-        self,
-        req: Request,
-        fp: object,
-        code: int,
-        msg: str,
-        headers: object,
-        newurl: str,
-    ) -> Request | None:
-        return None
-
-
-_URL_OPENER = build_opener(_RejectRedirects)
-
-
 class _HttpStatusError(RuntimeError):
     def __init__(self, status: int) -> None:
         self.status = status
@@ -300,7 +271,7 @@ def _post_json(
     request_timeout = min(settings.request_timeout_seconds, _remaining_time(deadline))
     request_deadline = min(deadline, monotonic() + request_timeout)
     try:
-        with _URL_OPENER.open(request, timeout=request_timeout) as response:
+        with open_bounded(request, timeout=request_timeout) as response:
             if expected_status is not None and response.status != expected_status:
                 raise _HttpStatusError(response.status)
             result = json.loads(_read_response(response, deadline=request_deadline))
@@ -325,6 +296,7 @@ def _request_headers(settings: WorkBuddyPluginSettings) -> dict[str, str]:
 def _read_response(response: _Response, *, deadline: float) -> bytes:
     """Read one response under a wall-clock deadline and a hard size bound."""
 
+    bind_response_deadline(response, deadline)
     content = bytearray()
     while True:
         _set_response_timeout(response, _remaining_time(deadline))
@@ -344,10 +316,10 @@ def _remaining_time(deadline: float) -> float:
     return remaining
 
 
-def _set_response_timeout(response: _Response, timeout: float) -> None:
+def _set_response_timeout(response: object, timeout: float) -> None:
     """Tighten urllib's socket timeout before each bounded read."""
 
-    raw = getattr(response.fp, "raw", None)
+    raw = getattr(getattr(response, "fp", None), "raw", None)
     sock = getattr(raw, "_sock", None)
     settimeout = getattr(sock, "settimeout", None)
     if settimeout is not None:

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

@@ -22,6 +22,7 @@
 import os
 import subprocess
 import sys
+import threading
 from collections.abc import Mapping, Sequence
 from hashlib import sha256
 from pathlib import Path
@@ -89,6 +90,34 @@ def redirect_request(
 _URL_OPENER = build_opener(_RejectRedirects)
 
 
+def open_bounded(request: Request, *, timeout: float) -> Any:
+    """Open one request under a hard wall-clock bound, response headers included.
+
+    urllib applies its timeout to each individual socket read, so a server that
+    trickles headers can outlive the caller's deadline. Running the open in a
+    daemon worker and abandoning it on expiry keeps the status line inside its
+    budget.
+    """
+
+    outcome: list[Any] = []
+
+    def _open() -> None:
+        try:
+            outcome.append(_URL_OPENER.open(request, timeout=timeout))
+        except BaseException as error:
+            outcome.append(error)
+
+    worker = threading.Thread(target=_open, name="powercontext-http", daemon=True)
+    worker.start()
+    worker.join(timeout)
+    if worker.is_alive():
+        raise TimeoutError
+    result = outcome[0] if outcome else TimeoutError()
+    if isinstance(result, BaseException):
+        raise result
+    return result
+
+
 def resolve_scope_id(
     cwd: str,
     *,
@@ -203,10 +232,12 @@ def _request_json(
         method=method,
     )
     try:
-        with _URL_OPENER.open(request, timeout=min(settings.request_timeout_seconds, remaining)) as response:
+        request_timeout = min(settings.request_timeout_seconds, remaining)
+        request_deadline = min(deadline, monotonic() + request_timeout)
+        with open_bounded(request, timeout=request_timeout) as response:
             if response.status < 200 or response.status >= 300:
                 raise ScopeBindingError
-            raw = _read_bounded(response)
+            raw = _read_bounded(response, deadline=request_deadline)
     except (HTTPError, OSError, TimeoutError) as error:
         raise ScopeBindingError from error
     try:
@@ -218,15 +249,82 @@ def _request_json(
     return value
 
 
-def _read_bounded(response: _Response) -> bytes:
+def _remaining_time(deadline: float) -> float:
+    remaining = deadline - monotonic()
+    if remaining <= 0:
+        raise TimeoutError
+    return remaining
+
+
+class _DeadlineSocket:
+    """Enforce one absolute deadline on every response socket read.
+
+    ``http.client`` can consume many socket reads inside a single ``read`` call
+    while it parses chunk framing, so tightening the socket timeout once per
+    bounded read cannot stop a server that trickles chunk extensions. Recomputing
+    the timeout before every receive keeps the caller's absolute deadline,
+    framing included.
+    """
+
+    def __init__(self, sock: Any, deadline: float) -> None:
+        self._sock = sock
+        self._deadline = deadline
+
+    def _remaining_time(self) -> float:
+        remaining = self._deadline - monotonic()
+        if remaining <= 0:
+            raise TimeoutError
+        return remaining
+
+    def recv(self, *args: Any) -> Any:
+        self._sock.settimeout(self._remaining_time())
+        return self._sock.recv(*args)
+
+    def recv_into(self, *args: Any) -> Any:
+        self._sock.settimeout(self._remaining_time())
+        return self._sock.recv_into(*args)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._sock, name)
+
+
+def bind_response_deadline(response: object, deadline: float) -> None:
+    """Keep every socket read of one open response inside the absolute deadline."""
+
+    if isinstance(response, HTTPError):
+        response = response.fp
+    raw: Any = getattr(getattr(response, "fp", None), "raw", None)
+    if raw is None:
+        return
+    sock = getattr(raw, "_sock", None)
+    if sock is None or isinstance(sock, _DeadlineSocket) or not hasattr(sock, "recv_into"):
+        return
+    raw._sock = _DeadlineSocket(sock, deadline)
+
+
+def _set_response_timeout(response: object, timeout: float) -> None:
+    """Tighten urllib's socket timeout before each bounded read."""
+
+    raw = getattr(getattr(response, "fp", None), "raw", None)
+    sock = getattr(raw, "_sock", None)
+    settimeout = getattr(sock, "settimeout", None)
+    if settimeout is not None:
+        settimeout(timeout)
+
+
+def _read_bounded(response: _Response, *, deadline: float) -> bytes:
+    bind_response_deadline(response, deadline)
     chunks: list[bytes] = []
     size = 0
-    while chunk := response.read(_READ_CHUNK_BYTES):
+    while True:
+        _set_response_timeout(response, _remaining_time(deadline))
+        chunk = response.read(_READ_CHUNK_BYTES)
+        if not chunk:
+            return b"".join(chunks)
         size += len(chunk)
         if size > _MAX_RESPONSE_BYTES:
             raise ScopeBindingError
         chunks.append(chunk)
-    return b"".join(chunks)
 
 
 def _git_value(cwd: str, *arguments: str) -> str | None:

tests/workbuddy_plugin/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2026 OceanBase.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the PowerContext WorkBuddy plugin."""

tests/workbuddy_plugin/conftest.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2026 OceanBase.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from collections.abc import Iterator
+from pathlib import Path
+from types import ModuleType
+
+import pytest
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
+WORKBUDDY_ROOT = REPOSITORY_ROOT / "integrations" / "workbuddy"
+PLUGIN_ROOT = WORKBUDDY_ROOT / "plugins" / "powercontext"
+_PLUGIN_MODULE_NAMES = (
+    "powercontext_client_config",
+    "powercontext_scope_binding",
+    "prepared_context",
+    "workbuddy_settings",
+    "workspace_scope",
+)
+
+
+def _load_module(name: str, path: Path) -> ModuleType:
+    spec = importlib.util.spec_from_file_location(name, path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception:
+        sys.modules.pop(name, None)
+        raise
+    return module
+
+
+@pytest.fixture
+def plugin_imports() -> Iterator[None]:
+    previous_path = list(sys.path)
+    previous_modules = {name: sys.modules.get(name) for name in _PLUGIN_MODULE_NAMES}
+    for name in _PLUGIN_MODULE_NAMES:
+        sys.modules.pop(name, None)
+    sys.path.insert(0, str(PLUGIN_ROOT))
+    try:
+        yield
+    finally:
+        sys.path[:] = previous_path
+        for name in _PLUGIN_MODULE_NAMES:
+            sys.modules.pop(name, None)
+        for name, module in previous_modules.items():
+            if module is not None:
+                sys.modules[name] = module
+
+
+@pytest.fixture
+def scope_module(plugin_imports: None) -> ModuleType:
+    return _load_module("powercontext_workbuddy_scope", PLUGIN_ROOT / "scripts" / "workspace_scope.py")
+
+
+@pytest.fixture
+def hook_module(plugin_imports: None) -> ModuleType:
+    return _load_module(
+        "powercontext_workbuddy_hook",
+        PLUGIN_ROOT / "hooks" / "workbuddy_powercontext_hook.py",
+    )