diff --git a/integrations/workbuddy/plugins/powercontext/hooks/powercontext_scope_binding.py b/integrations/workbuddy/plugins/powercontext/hooks/powercontext_scope_binding.py index 91fe82bbc9..bfdb333562 100644 --- a/integrations/workbuddy/plugins/powercontext/hooks/powercontext_scope_binding.py +++ b/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"] diff --git a/integrations/workbuddy/plugins/powercontext/hooks/workbuddy_powercontext_hook.py b/integrations/workbuddy/plugins/powercontext/hooks/workbuddy_powercontext_hook.py index 8dc5701d2a..4510d38f9a 100644 --- a/integrations/workbuddy/plugins/powercontext/hooks/workbuddy_powercontext_hook.py +++ b/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: diff --git a/integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py b/integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py index 70532a45f8..ffd6991064 100644 --- a/integrations/workbuddy/plugins/powercontext/scripts/workspace_scope.py +++ b/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: diff --git a/tests/workbuddy_plugin/__init__.py b/tests/workbuddy_plugin/__init__.py new file mode 100644 index 0000000000..85a43f8d7d --- /dev/null +++ b/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.""" diff --git a/tests/workbuddy_plugin/conftest.py b/tests/workbuddy_plugin/conftest.py new file mode 100644 index 0000000000..379f8fce22 --- /dev/null +++ b/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", + ) diff --git a/tests/workbuddy_plugin/test_hook.py b/tests/workbuddy_plugin/test_hook.py new file mode 100644 index 0000000000..10233a3cb4 --- /dev/null +++ b/tests/workbuddy_plugin/test_hook.py @@ -0,0 +1,154 @@ +# 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 io +import json +import sys +import threading +import time +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +_SCOPE_BODY = json.dumps({"scope_id": "scope-1"}).encode() +_DRIP_INTERVAL_SECONDS = 0.05 +_DRIP_WRITES = 60 + + +@contextmanager +def _serve(handler: type[BaseHTTPRequestHandler]) -> Generator[str, None, None]: + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + server.shutdown() + thread.join(timeout=1) + server.server_close() + + +def test_prompt_hook_fails_open_within_the_http_budget_when_recall_reads_drip( + hook_module: ModuleType, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """The recall/capture reads must stay inside the hook's own HTTP budget, not the host hook timeout.""" + + paths: list[str] = [] + + class DripRecallHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + paths.append(self.path) + self.send_response(200) + self.send_header("Content-Type", "application/json") + if self.path == "/v1/scope-bindings/resolve": + self.send_header("Content-Length", str(len(_SCOPE_BODY))) + self.end_headers() + self.wfile.write(_SCOPE_BODY) + self.wfile.flush() + return + self.end_headers() + for _ in range(_DRIP_WRITES): + try: + self.wfile.write(b" ") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + time.sleep(_DRIP_INTERVAL_SECONDS) + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + monkeypatch.setattr( + sys, + "stdin", + io.StringIO( + json.dumps({ + "hook_event_name": "UserPromptSubmit", + "prompt": "what did we decide about the recall budget?", + "cwd": str(tmp_path), + "session_id": "session-1", + }) + ), + ) + output = io.StringIO() + monkeypatch.setattr(sys, "stdout", output) + + with _serve(DripRecallHandler) as server_url: + settings = hook_module.WorkBuddyPluginSettings( + server_url=server_url, + request_timeout_seconds=0.3, + http_budget_seconds=1.0, + ) + started = time.monotonic() + assert hook_module.main(settings) == 0 + elapsed = time.monotonic() - started + + assert elapsed < 1.5 + assert "/v1/scope-bindings/resolve" in paths + assert "/v1/context/prepare" in paths + assert json.loads(output.getvalue())["hookSpecificOutput"]["additionalContext"] == "" + + +def test_prompt_hook_refuses_redirects( + hook_module: ModuleType, +) -> None: + """Sharing the bounded opener must not let a redirect carry the hook's credentials elsewhere.""" + + target_headers: list[dict[str, str]] = [] + + class TargetHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + target_headers.append(dict(self.headers)) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(TargetHandler) as target_url: + + class RedirectHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + self.send_response(302) + self.send_header("Location", f"{target_url}/stolen") + self.end_headers() + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(RedirectHandler) as source_url: + settings = hook_module.WorkBuddyPluginSettings( + server_url=source_url, + authorization="Bearer secret-token", + ) + with pytest.raises(RuntimeError): + hook_module._post_json( + "/v1/context/prepare", + {"scope_id": "scope-1", "query": "what did we decide?"}, + settings=settings, + deadline=time.monotonic() + 1.0, + ) + + assert target_headers == [] diff --git a/tests/workbuddy_plugin/test_scope_binding.py b/tests/workbuddy_plugin/test_scope_binding.py new file mode 100644 index 0000000000..e32592fbf7 --- /dev/null +++ b/tests/workbuddy_plugin/test_scope_binding.py @@ -0,0 +1,202 @@ +# 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 json +import threading +import time +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from types import ModuleType +from typing import Any + +import pytest + +_SCOPE_BODY = json.dumps({"scope_id": "scope-1"}).encode() +_DRIP_INTERVAL_SECONDS = 0.05 +_DRIP_WRITES = 60 +_DRIP_HEADER_LINES = 60 +_LARGE_PAYLOAD_BYTES = 1_000_000 + + +@contextmanager +def _serve(handler: type[BaseHTTPRequestHandler]) -> Generator[str, None, None]: + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + server.shutdown() + thread.join(timeout=1) + server.server_close() + + +def _settings(scope_module: ModuleType, server_url: str, **overrides: Any) -> Any: + return scope_module.WorkBuddyPluginSettings(server_url=server_url, **overrides) + + +def _drain_request(request: BaseHTTPRequestHandler) -> None: + request.rfile.read(int(request.headers.get("Content-Length", "0"))) + + +def test_response_body_drip_aborts_at_the_request_timeout( + scope_module: ModuleType, + tmp_path: Path, +) -> None: + """A server that trickles one byte faster than the socket timeout must not outlive the deadline.""" + + class DripBodyHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + _drain_request(self) + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + for _ in range(_DRIP_WRITES): + try: + self.wfile.write(b" ") + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + time.sleep(_DRIP_INTERVAL_SECONDS) + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(DripBodyHandler) as server_url: + started = time.monotonic() + with pytest.raises(scope_module.ScopeBindingError): + scope_module.resolve_scope_id( + str(tmp_path), + session_id=None, + settings=_settings(scope_module, server_url, request_timeout_seconds=0.3), + deadline=time.monotonic() + 2.0, + ) + elapsed = time.monotonic() - started + + assert elapsed < 1.0 + + +def test_response_header_drip_aborts_at_the_request_timeout( + scope_module: ModuleType, + tmp_path: Path, +) -> None: + """A server that trickles response headers inside the socket timeout must not outlive the deadline.""" + + class DripHeaderHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + _drain_request(self) + try: + self.wfile.write(b"HTTP/1.1 200 OK\r\n") + self.wfile.flush() + for index in range(_DRIP_HEADER_LINES): + self.wfile.write(f"X-Pad-{index:02d}: {'p' * 16}\r\n".encode()) + self.wfile.flush() + time.sleep(_DRIP_INTERVAL_SECONDS) + self.wfile.write(b"Content-Type: application/json\r\n") + self.wfile.write(f"Content-Length: {len(_SCOPE_BODY)}\r\n".encode()) + self.wfile.write(b"\r\n") + self.wfile.write(_SCOPE_BODY) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(DripHeaderHandler) as server_url: + started = time.monotonic() + with pytest.raises(scope_module.ScopeBindingError): + scope_module.resolve_scope_id( + str(tmp_path), + session_id=None, + settings=_settings(scope_module, server_url, request_timeout_seconds=0.3), + deadline=time.monotonic() + 2.0, + ) + elapsed = time.monotonic() - started + + assert elapsed < 1.0 + + +def test_large_response_body_within_the_request_timeout_succeeds( + scope_module: ModuleType, + tmp_path: Path, +) -> None: + """Bounding the deadline must not turn a large but healthy response into a failure.""" + + payload = json.dumps({"scope_id": "scope-1", "pad": "x" * _LARGE_PAYLOAD_BYTES}).encode() + + class LargeBodyHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + _drain_request(self) + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + self.wfile.flush() + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(LargeBodyHandler) as server_url: + started = time.monotonic() + scope_id = scope_module.resolve_scope_id( + str(tmp_path), + session_id=None, + settings=_settings(scope_module, server_url, request_timeout_seconds=2.0), + deadline=time.monotonic() + 2.0, + ) + elapsed = time.monotonic() - started + + assert scope_id == "scope-1" + assert elapsed < 2.0 + + +def test_expired_deadline_fails_without_contacting_the_server( + scope_module: ModuleType, + tmp_path: Path, +) -> None: + requests: list[str] = [] + + class RecordingHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + _drain_request(self) + requests.append(self.path) + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(_SCOPE_BODY))) + self.end_headers() + self.wfile.write(_SCOPE_BODY) + self.wfile.flush() + + def log_message(self, format: str, *args: Any) -> None: # noqa: A002 + pass + + with _serve(RecordingHandler) as server_url: + started = time.monotonic() + with pytest.raises(scope_module.ScopeBindingError): + scope_module.resolve_scope_id( + str(tmp_path), + session_id=None, + settings=_settings(scope_module, server_url, request_timeout_seconds=0.3), + deadline=time.monotonic() - 1.0, + ) + elapsed = time.monotonic() - started + + assert requests == [] + assert elapsed < 1.0