diff --git a/e2e/bub/tests/test_bub_capture.py b/e2e/bub/tests/test_bub_capture.py new file mode 100644 index 000000000..fd602ec5a --- /dev/null +++ b/e2e/bub/tests/test_bub_capture.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import asyncio +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +from bub.hooks.interception import ToolCall, ToolCallResult +from powercontext_bub import plugin as plugin_module +from powercontext_bub.plugin import PowerContextPlugin, PowerContextSettings + + +def test_tool_capture_redacts_credentials_before_crossing_the_client_boundary( + monkeypatch, + tmp_path: Path, +) -> None: + sensitive_value = "provider-secret-sentinel" + captured_requests: list[Any] = [] + + class RecordingClient: + def __init__(self, base_url: str, *, timeout: float) -> None: + del base_url, timeout + + async def __aenter__(self) -> RecordingClient: + return self + + async def __aexit__(self, *exc_info: object) -> None: + del exc_info + + async def capture_content_source(self, request: Any) -> SimpleNamespace: + captured_requests.append(request) + return SimpleNamespace(position=1) + + settings = PowerContextSettings( + base_url="http://127.0.0.1:8000", + scope_id="test:scope", + capture_events=True, + capture_checkpoint_every=100, + ) + monkeypatch.setenv("BUB_API_KEY", sensitive_value) + monkeypatch.setattr(plugin_module, "ensure_config", lambda _: settings) + monkeypatch.setattr(plugin_module, "PowerContextClient", RecordingClient) + plugin = PowerContextPlugin(SimpleNamespace(workspace=tmp_path)) + state = plugin.load_state(message=None, session_id="session-1") + state["session_id"] = "session-1" + + asyncio.run( + plugin.after_tool_call( + ToolCall(run_id="run-1", tool="provider.request", arguments={"api_key": sensitive_value}), + ToolCallResult( + run_id="run-1", + tool="provider.request", + arguments={"api_key": sensitive_value}, + result=f"response contained {sensitive_value}", + ), + state, + ) + ) + + assert len(captured_requests) == 1 + request = captured_requests[0] + assert request.metadata["event"] == "tool_result" + assert sensitive_value not in request.content + assert "[REDACTED]" in request.content diff --git a/e2e/bub/uv.lock b/e2e/bub/uv.lock index ae2e3bc71..d05276f1b 100644 --- a/e2e/bub/uv.lock +++ b/e2e/bub/uv.lock @@ -1045,12 +1045,14 @@ source = { editable = "../../integrations/bub" } dependencies = [ { name = "bub" }, { name = "powercontext", extra = ["client"] }, + { name = "pydantic-settings" }, ] [package.metadata] requires-dist = [ { name = "bub", specifier = ">=0.4.0,<0.5.0" }, { name = "powercontext", extras = ["client"], specifier = ">=0.0.1" }, + { name = "pydantic-settings", specifier = ">=2.7,<3" }, ] [[package]] diff --git a/integrations/bub/README.md b/integrations/bub/README.md index f2c2d00b5..d7a7ca869 100644 --- a/integrations/bub/README.md +++ b/integrations/bub/README.md @@ -7,17 +7,41 @@ This package connects Bub to a running PowerContext Server through the public Py - `powercontext.context` prepares bounded context for a question. Before each model call, the plugin also prepares relevant context and adds it as host-supplied historical evidence. -The plugin does not persist Bub conversation history. A new Bub session can observe an earlier session only through the -configured PowerContext scope. +Automatic trajectory capture is opt-in. When enabled, the plugin captures the initial task and completed LLM and tool +events as bounded Content Sources. It periodically flushes those Sources through the Memory pipeline so later model +steps in the same Bub run can recall earlier findings. Provider-hidden reasoning is never available to the hook and is +not captured. ## Configuration +The plugin uses Bub's Pydantic settings extension. Configuration can live in the `powercontext` section of Bub's +configuration file: + +```yaml +powercontext: + base_url: http://127.0.0.1:8000 + scope_id: project:example + capture_events: true + capture_checkpoint_every: 5 +``` + +Environment variables use the `POWERCONTEXT_BUB_` prefix and take precedence over file values. Values are parsed and +validated by Pydantic before the plugin starts. + | Variable | Default | Purpose | | --- | --- | --- | | `POWERCONTEXT_BUB_BASE_URL` | `http://127.0.0.1:8000` | PowerContext Server URL | | `POWERCONTEXT_BUB_SCOPE_ID` | workspace-derived | Durable scope shared by Bub sessions | | `POWERCONTEXT_BUB_TIMEOUT` | `10` | Client timeout in seconds | | `POWERCONTEXT_BUB_MAX_BYTES` | `8000` | Maximum prepared-context size | +| `POWERCONTEXT_BUB_CAPTURE_EVENTS` | `false` | Capture completed Bub events as Content Sources | +| `POWERCONTEXT_BUB_CAPTURE_CHECKPOINT_EVERY` | `5` | Flush Memory after this many captured events | +| `POWERCONTEXT_BUB_CAPTURE_MAX_BYTES` | `8192` | Maximum UTF-8 bytes stored for one captured event | +| `POWERCONTEXT_BUB_CAPTURE_LOG` | unset | Optional JSONL evidence path; records metadata but not event content | + +Captured tool arguments redact values under credential-like keys. Known credential environment values are also +removed from serialized event content. Keep the PowerContext scope and optional capture log protected because normal +tool output can still contain sensitive project data. Install the package together with PowerContext and Bub: diff --git a/integrations/bub/pyproject.toml b/integrations/bub/pyproject.toml index a1223880e..88ce05fa2 100644 --- a/integrations/bub/pyproject.toml +++ b/integrations/bub/pyproject.toml @@ -6,6 +6,7 @@ requires-python = ">=3.12,<4.0" dependencies = [ "bub>=0.4.0,<0.5.0", "powercontext[client]>=0.0.1", + "pydantic-settings>=2.7,<3", ] [project.entry-points."bub"] diff --git a/integrations/bub/src/powercontext_bub/__init__.py b/integrations/bub/src/powercontext_bub/__init__.py index 9f885dfd1..226ab7c99 100644 --- a/integrations/bub/src/powercontext_bub/__init__.py +++ b/integrations/bub/src/powercontext_bub/__init__.py @@ -1,6 +1,6 @@ """Bub integration for PowerContext.""" from powercontext_bub import tools as _tools # noqa: F401 -from powercontext_bub.plugin import PowerContextPlugin +from powercontext_bub.plugin import PowerContextPlugin, PowerContextSettings -__all__ = ["PowerContextPlugin"] +__all__ = ["PowerContextPlugin", "PowerContextSettings"] diff --git a/integrations/bub/src/powercontext_bub/plugin.py b/integrations/bub/src/powercontext_bub/plugin.py index 2dac9fd2a..fad06ba59 100644 --- a/integrations/bub/src/powercontext_bub/plugin.py +++ b/integrations/bub/src/powercontext_bub/plugin.py @@ -2,18 +2,23 @@ from __future__ import annotations +import asyncio import hashlib +import json import os -from dataclasses import dataclass, replace +from dataclasses import replace +from datetime import UTC, datetime from pathlib import Path from typing import Any -from bub import hookimpl -from bub.hooks.interception import LlmCallRequest +from bub import Settings, config, ensure_config, hookimpl +from bub.hooks.interception import LlmCallRequest, LlmCallResult, ToolCall, ToolCallResult from bub.turn import TurnState +from pydantic import Field, HttpUrl +from pydantic_settings import SettingsConfigDict from powercontext.client import InvalidResponseError, PowerContextClient, ServerResponseError, TransportError -from powercontext.http import PrepareContextRequest +from powercontext.http import CaptureContentSourceRequest, FlushMemoryRequest, PrepareContextRequest STATE_KEY = "_powercontext" CONTEXT_MARKER = "PowerContext host-supplied context" @@ -23,57 +28,53 @@ Relevant host-supplied context is injected automatically before each model call. Use powercontext.search for follow-up recall beyond the injected context. Use powercontext.remember when the user establishes a durable decision, preference, constraint, or procedure.""" +CAPTURE_SCHEMA = "powercontext.bub-capture-event/v1" +SENSITIVE_KEY_PARTS = ("api_key", "authorization", "cookie", "password", "secret", "token") -class ConfigurationError(ValueError): - """Report invalid PowerContext settings during Bub startup.""" +@config(name="powercontext") +class PowerContextSettings(Settings): + """Validated Bub configuration for the PowerContext plugin.""" - def __init__(self) -> None: - super().__init__("POWERCONTEXT_BUB_MAX_BYTES must be between 512 and 32768") + model_config = SettingsConfigDict( + env_prefix="POWERCONTEXT_BUB_", + env_ignore_empty=True, + extra="ignore", + frozen=True, + ) - -@dataclass(frozen=True, slots=True) -class Settings: - """PowerContext settings for one Bub process.""" - - base_url: str - scope_id: str - timeout: float - max_bytes: int - - @classmethod - def from_environment(cls, workspace: Path) -> Settings: - scope_id = os.getenv("POWERCONTEXT_BUB_SCOPE_ID", "").strip() - if not scope_id: - digest = hashlib.sha256(str(workspace.resolve()).encode()).hexdigest()[:20] - scope_id = f"bub:{digest}" - - max_bytes = int(os.getenv("POWERCONTEXT_BUB_MAX_BYTES", "8000")) - if not 512 <= max_bytes <= 32768: - raise ConfigurationError - - return cls( - base_url=os.getenv("POWERCONTEXT_BUB_BASE_URL", "http://127.0.0.1:8000").strip(), - scope_id=scope_id, - timeout=float(os.getenv("POWERCONTEXT_BUB_TIMEOUT", "10")), - max_bytes=max_bytes, - ) + base_url: HttpUrl = HttpUrl("http://127.0.0.1:8000") + scope_id: str | None = Field(default=None, min_length=1) + timeout: float = Field(default=10, gt=0) + max_bytes: int = Field(default=8000, ge=512, le=32768) + capture_events: bool = False + capture_checkpoint_every: int = Field(default=5, ge=1, le=100) + capture_max_bytes: int = Field(default=8192, ge=512, le=32768) + capture_log: Path | None = None class PowerContextPlugin: """Bub hooks backed by the public PowerContext client.""" def __init__(self, framework: Any) -> None: - self.settings = Settings.from_environment(Path(framework.workspace)) + self.settings = ensure_config(PowerContextSettings) + self.base_url = str(self.settings.base_url).rstrip("/") + self.scope_id = self.settings.scope_id or _workspace_scope(Path(framework.workspace)) + self._capture_lock = asyncio.Lock() @hookimpl def load_state(self, message: Any, session_id: str) -> TurnState: del message, session_id return { STATE_KEY: { - "base_url": self.settings.base_url, - "scope_id": self.settings.scope_id, + "base_url": self.base_url, + "scope_id": self.scope_id, "timeout": self.settings.timeout, + "capture_sequence": 0, + "captured_events": 0, + "captured_position": 0, + "flushed_position": 0, + "prompt_captured": False, } } @@ -84,10 +85,20 @@ def system_prompt(self, prompt: str | list[dict[str, Any]], state: TurnState) -> @hookimpl async def before_llm_call(self, request: LlmCallRequest, state: TurnState) -> LlmCallRequest | None: + query = _latest_user_text(request.messages) + capture_state = state[STATE_KEY] + if self.settings.capture_events and query and not capture_state["prompt_captured"]: + capture_state["prompt_captured"] = True + await self._capture_event( + event="user_prompt", + run_id=request.run_id, + payload={"text": query}, + state=state, + ) + if any(_contains_context_marker(message) for message in request.messages): return None - query = _latest_user_text(request.messages) if not query: return None @@ -101,20 +112,173 @@ async def before_llm_call(self, request: LlmCallRequest, state: TurnState) -> Ll } return replace(request, messages=[context_message, *request.messages]) + @hookimpl + async def after_llm_call(self, request: LlmCallRequest, result: LlmCallResult, state: TurnState) -> None: + if not self.settings.capture_events: + return + await self._capture_event( + event="llm_result", + run_id=request.run_id, + payload={ + "text": result.text, + "tool_calls": result.tool_calls, + "error": _error_name(result.error), + "duration_ms": result.duration_ms, + }, + state=state, + ) + + @hookimpl + async def after_tool_call(self, call: ToolCall, result: ToolCallResult, state: TurnState) -> None: + if not self.settings.capture_events: + return + await self._capture_event( + event="tool_result", + run_id=call.run_id, + payload={ + "tool": call.tool, + "arguments": call.arguments, + "result": result.result, + "error": _error_name(result.error), + "duration_ms": result.duration_ms, + }, + state=state, + ) + + @hookimpl + async def save_state(self, session_id: str, state: TurnState, message: Any, model_output: str) -> None: + del session_id, message, model_output + if not self.settings.capture_events: + return + async with self._capture_lock: + await self._flush_captured_sources(state, final=True) + async def _prepare_context(self, query: str, state: TurnState) -> str | None: request = PrepareContextRequest( - scope_id=self.settings.scope_id, + scope_id=self.scope_id, query=query, max_bytes=self.settings.max_bytes, ) try: - async with PowerContextClient(self.settings.base_url, timeout=self.settings.timeout) as client: + async with PowerContextClient(self.base_url, timeout=self.settings.timeout) as client: prepared = await client.prepare_context(request) except CLIENT_ERRORS as exc: state[STATE_KEY]["prepare_error"] = type(exc).__name__ + self._write_capture_record( + event="context", + status="failed", + error=type(exc).__name__, + captured_events=state[STATE_KEY]["captured_events"], + ) return None + self._write_capture_record( + event="context", + status=prepared.status.value, + content_bytes=prepared.content_bytes, + captured_events=state[STATE_KEY]["captured_events"], + flushed_position=state[STATE_KEY]["flushed_position"], + ) return prepared.content + async def _capture_event( + self, + *, + event: str, + run_id: str, + payload: dict[str, Any], + state: TurnState, + ) -> None: + async with self._capture_lock: + capture_state = state[STATE_KEY] + capture_state["capture_sequence"] += 1 + sequence = capture_state["capture_sequence"] + session_id = str(state.get("session_id", "unknown")) + source_id = _source_id(self.scope_id, session_id, sequence, event, run_id) + content = _capture_content(event, sequence, payload, self.settings.capture_max_bytes) + request = CaptureContentSourceRequest( + scope_id=self.scope_id, + source_id=source_id, + content=content, + metadata={ + "origin": "bub", + "kind": "agent-trajectory", + "event": event, + "sequence": sequence, + "session_id": session_id, + "run_id": run_id, + }, + ) + try: + async with PowerContextClient(self.base_url, timeout=self.settings.timeout) as client: + response = await client.capture_content_source(request) + except CLIENT_ERRORS as exc: + self._write_capture_record( + event=event, + sequence=sequence, + status="failed", + source_id=source_id, + error=type(exc).__name__, + ) + return + + capture_state["captured_events"] += 1 + capture_state["captured_position"] = max(capture_state["captured_position"], response.position) + self._write_capture_record( + event=event, + sequence=sequence, + status="captured", + source_id=source_id, + source_position=response.position, + ) + if capture_state["captured_events"] % self.settings.capture_checkpoint_every == 0: + await self._flush_captured_sources(state, final=False) + + async def _flush_captured_sources(self, state: TurnState, *, final: bool) -> None: + capture_state = state[STATE_KEY] + target_position = capture_state["captured_position"] + if target_position <= capture_state["flushed_position"]: + return + + try: + async with PowerContextClient(self.base_url, timeout=self.settings.timeout) as client: + response = await client.flush_memory(FlushMemoryRequest(scope_id=self.scope_id)) + except CLIENT_ERRORS as exc: + self._write_capture_record( + event="checkpoint", + status="failed", + final=final, + target_position=target_position, + error=type(exc).__name__, + ) + return + + capture_state["flushed_position"] = response.current_cursor + self._write_capture_record( + event="checkpoint", + status=response.status.value, + final=final, + target_position=target_position, + previous_cursor=response.previous_cursor, + current_cursor=response.current_cursor, + high_watermark=response.high_watermark, + processed_source_count=response.processed_source_count, + memory_created=response.memory is not None, + ) + + def _write_capture_record(self, *, event: str, status: str, **values: Any) -> None: + if self.settings.capture_log is None: + return + record = { + "schema": CAPTURE_SCHEMA, + "recorded_at": datetime.now(UTC).isoformat(), + "event": event, + "status": status, + **values, + } + self.settings.capture_log.parent.mkdir(parents=True, exist_ok=True) + with self.settings.capture_log.open("a", encoding="utf-8") as capture_log: + capture_log.write(json.dumps(record, ensure_ascii=True, sort_keys=True) + "\n") + def _latest_user_text(messages: list[dict[str, Any]]) -> str: for message in reversed(messages): @@ -135,3 +299,95 @@ def _latest_user_text(messages: list[dict[str, Any]]) -> str: def _contains_context_marker(message: dict[str, Any]) -> bool: content = message.get("content") return isinstance(content, str) and CONTEXT_MARKER in content + + +def _workspace_scope(workspace: Path) -> str: + digest = hashlib.sha256(str(workspace.resolve()).encode()).hexdigest()[:20] + return f"bub:{digest}" + + +def _source_id(scope_id: str, session_id: str, sequence: int, event: str, run_id: str) -> str: + identity = "\0".join((scope_id, session_id, str(sequence), event, run_id)) + return f"bub-event:{hashlib.sha256(identity.encode()).hexdigest()}" + + +def _capture_content(event: str, sequence: int, payload: dict[str, Any], max_bytes: int) -> str: + safe_payload = _sanitize(payload) + content = _redact_known_secrets( + json.dumps( + {"event": event, "sequence": sequence, "payload": safe_payload}, + ensure_ascii=True, + sort_keys=True, + default=str, + ) + ) + encoded = content.encode("utf-8") + if len(encoded) <= max_bytes: + return content + + envelope = {"event": event, "sequence": sequence, "payload_excerpt": "", "truncated": True} + lower_bound = 0 + upper_bound = len(content) + rendered = json.dumps(envelope, ensure_ascii=True, sort_keys=True) + while lower_bound <= upper_bound: + candidate_length = (lower_bound + upper_bound) // 2 + envelope["payload_excerpt"] = content[:candidate_length] + candidate = json.dumps(envelope, ensure_ascii=True, sort_keys=True) + if len(candidate.encode("utf-8")) <= max_bytes: + rendered = candidate + lower_bound = candidate_length + 1 + else: + upper_bound = candidate_length - 1 + return rendered + + +def _sanitize(value: Any) -> Any: + if isinstance(value, dict): + return { + str(key): "[REDACTED]" if _is_sensitive_key(str(key)) else _sanitize(item) for key, item in value.items() + } + if isinstance(value, list | tuple): + return [_sanitize(item) for item in value] + return value + + +def _is_sensitive_key(key: str) -> bool: + folded = key.casefold().replace("-", "_") + return any(part in folded for part in SENSITIVE_KEY_PARTS) + + +def _error_name(error: Exception | None) -> str | None: + return None if error is None else type(error).__name__ + + +def _redact_known_secrets(value: str) -> str: + secrets = {secret for name, secret in os.environ.items() if secret and len(secret) >= 8 and _is_sensitive_key(name)} + secrets.update(_codex_auth_secrets()) + for secret in secrets: + value = value.replace(secret, "[REDACTED]") + return value + + +def _codex_auth_secrets() -> set[str]: + codex_home = Path(os.getenv("CODEX_HOME", str(Path.home() / ".codex"))).expanduser() + try: + auth = json.loads((codex_home / "auth.json").read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return set() + return _sensitive_values(auth) + + +def _sensitive_values(value: Any, *, sensitive: bool = False) -> set[str]: + if isinstance(value, dict): + secrets: set[str] = set() + for key, item in value.items(): + secrets.update(_sensitive_values(item, sensitive=sensitive or _is_sensitive_key(str(key)))) + return secrets + if isinstance(value, list | tuple): + secrets = set() + for item in value: + secrets.update(_sensitive_values(item, sensitive=sensitive)) + return secrets + if sensitive and isinstance(value, str) and len(value) >= 8: + return {value} + return set()