From 44cfecaef7013adf5c8bc62d35b063746cc4fe0e Mon Sep 17 00:00:00 2001 From: hailey Date: Tue, 14 Jul 2026 22:13:18 +0000 Subject: [PATCH] feat(ask): host-neutral Ask turn service Add osprey_worker/lib/ask: vendor-neutral contracts (AskRequest/AskEvent/Principal/ AskLimits/Conversation/ContextSnapshot/ResolvedModel), a stable error taxonomy, host adapter Protocols (ModelPolicy, ConversationStore, ConversationLock, ContextSnapshotProvider, ToolRegistryFactory, Redactor, AuditSink, EvidenceNormalizer) + AskConfig, and AskService.run_turn -- a generator that pre-flights (raising before any event), streams a deterministic bounded tool loop, redacts structured payloads, enforces history/context/output budgets, checks cancellation, persists only successful turns, and releases the conversation lock on every terminal path. Tests (28 passing under run-tests.sh) cover the deterministic event sequence, tool feed-back, multi-tool ordering, iteration/output budgets, ownership + pre-flight rejections, provider/tool/cancellation/lock failures, structured redaction, and audit-failure isolation. Implements osprey-ask-ai AC2.1-AC2.5, AC3.1-AC3.5. --- .../src/osprey/worker/lib/ask/__init__.py | 87 ++++ .../src/osprey/worker/lib/ask/contracts.py | 110 +++++ .../src/osprey/worker/lib/ask/errors.py | 124 +++++ .../src/osprey/worker/lib/ask/ports.py | 136 ++++++ .../src/osprey/worker/lib/ask/service.py | 256 +++++++++++ .../src/osprey/worker/lib/ask/tests/fakes.py | 240 ++++++++++ .../worker/lib/ask/tests/test_errors.py | 66 +++ .../worker/lib/ask/tests/test_service.py | 430 ++++++++++++++++++ 8 files changed, 1449 insertions(+) create mode 100644 osprey_worker/src/osprey/worker/lib/ask/__init__.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/contracts.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/errors.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/ports.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/service.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/tests/fakes.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/tests/test_errors.py create mode 100644 osprey_worker/src/osprey/worker/lib/ask/tests/test_service.py diff --git a/osprey_worker/src/osprey/worker/lib/ask/__init__.py b/osprey_worker/src/osprey/worker/lib/ask/__init__.py new file mode 100644 index 00000000..cab88e86 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/__init__.py @@ -0,0 +1,87 @@ +"""Host-neutral Ask AI turn service. + +This package drives a validated user turn through an LLM provider, a principal-scoped +tool registry, host persistence/locking, redaction, audit, and evidence adapters, +emitting a deterministic, bounded stream of versioned events. All product- and +vendor-specific behavior is supplied by the host through the Protocols in +:mod:`osprey.worker.lib.ask.ports`; nothing here imports Discord/Smite/SML/Druid or a +concrete vendor SDK. +""" + +from osprey.worker.lib.ask.contracts import ( + ASK_EVENT_VERSION, + AskEvent, + AskEventType, + AskLimits, + AskRequest, + ContextSnapshot, + Conversation, + Principal, + ResolvedModel, +) +from osprey.worker.lib.ask.errors import ( + AskError, + AskErrorCode, + BudgetExceeded, + Cancelled, + Forbidden, + InternalError, + InvalidContext, + InvalidModel, + InvalidRequest, + LockUnavailable, + ProviderError, + SerializationError, + ToolExecutionError, + UnavailableProvider, + to_public_payload, +) +from osprey.worker.lib.ask.ports import ( + AskConfig, + AuditSink, + ContextSnapshotProvider, + ConversationLock, + ConversationStore, + EvidenceNormalizer, + ModelPolicy, + Redactor, + ToolRegistryFactory, +) +from osprey.worker.lib.ask.service import AskService + +__all__ = [ + 'ASK_EVENT_VERSION', + 'AskEvent', + 'AskEventType', + 'AskLimits', + 'AskRequest', + 'ContextSnapshot', + 'Conversation', + 'Principal', + 'ResolvedModel', + 'AskError', + 'AskErrorCode', + 'BudgetExceeded', + 'Cancelled', + 'Forbidden', + 'InternalError', + 'InvalidContext', + 'InvalidModel', + 'InvalidRequest', + 'LockUnavailable', + 'ProviderError', + 'SerializationError', + 'ToolExecutionError', + 'UnavailableProvider', + 'to_public_payload', + 'AskConfig', + 'AuditSink', + 'ContextSnapshotProvider', + 'ConversationLock', + 'ConversationStore', + 'EvidenceNormalizer', + 'ModelPolicy', + 'Redactor', + 'ToolRegistryFactory', + 'AskService', +] diff --git a/osprey_worker/src/osprey/worker/lib/ask/contracts.py b/osprey_worker/src/osprey/worker/lib/ask/contracts.py new file mode 100644 index 00000000..1bdea33e --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/contracts.py @@ -0,0 +1,110 @@ +"""Vendor-neutral Ask AI domain contracts. + +Request, event, principal, conversation, and limit types shared across the Ask +service and its transports. No host, product, or vendor specifics live here. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, List, Literal, Mapping, Optional + +from osprey.worker.lib.llm import BaseLLMProvider, LLMMessage + +ASK_EVENT_VERSION = 1 + +AskEventType = Literal[ + 'conversation_started', + 'tool_call', + 'query_result', + 'assistant_message', + 'done', + 'error', +] + + +@dataclass(frozen=True) +class Principal: + """The server-authenticated actor a turn is scoped to. + + Derived by the transport from server-side identity; it is never read from the + client request body. Ownership, permissions, and tool authorization are all + keyed off this principal. + """ + + id: str + email: str + display_name: Optional[str] = None + attributes: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass +class AskRequest: + """A single user turn request. ``message`` is required and non-empty.""" + + message: str + conversation_id: Optional[str] = None + model: Optional[str] = None + context_ref: Optional[str] = None + + +@dataclass +class AskEvent: + """A versioned event emitted while a turn streams. + + ``type`` selects the event; ``payload`` carries type-specific data. Every event + carries the schema ``version`` so clients and hosts can evolve independently. + """ + + type: AskEventType + payload: dict[str, Any] = field(default_factory=dict) + conversation_id: Optional[str] = None + turn_id: Optional[str] = None + version: int = ASK_EVENT_VERSION + + +@dataclass(frozen=True) +class AskLimits: + """Host-configurable bounds the service enforces on every turn.""" + + max_history_messages: int = 40 + max_context_chars: int = 20_000 + max_tool_iterations: int = 8 + max_output_tokens: int = 1024 + max_output_chars: int = 40_000 + max_evidence_chars: int = 4_000 + + +@dataclass +class Conversation: + """A conversation owned by a principal. ``messages`` are oldest-first.""" + + id: str + principal_id: str + messages: List[LLMMessage] = field(default_factory=list) + + +@dataclass +class ContextSnapshot: + """Validated grounding state referenced by a request's ``context_ref``. + + Grounding input only: it never grants authorization. The host renders it to a + string via :meth:`ContextSnapshotProvider.render`. + """ + + ref: str + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass +class ResolvedModel: + """The outcome of :meth:`ModelPolicy.resolve`. + + ``provider`` is a *ready* provider (SDK/credentials validated) or ``None`` when + no provider is available (mapped to ``unavailable_provider``). ``limits`` + optionally overrides the service default for this turn. + """ + + model: Optional[str] + provider: Optional[BaseLLMProvider] + limits: Optional[AskLimits] = None diff --git a/osprey_worker/src/osprey/worker/lib/ask/errors.py b/osprey_worker/src/osprey/worker/lib/ask/errors.py new file mode 100644 index 00000000..63b9e4e3 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/errors.py @@ -0,0 +1,124 @@ +"""Ask error taxonomy mapped to stable, host-safe codes. + +Every error carries a stable ``code`` and a ``public_message`` that is always safe to +show a client -- it never contains a raw exception, provider payload, or credential. +``preflight`` errors are raised before any event is emitted (the transport maps them +to an HTTP status); non-preflight errors occur mid-stream and become a terminal +``error`` event. +""" + +from __future__ import annotations + +from typing import Dict, Literal, Optional + +AskErrorCode = Literal[ + 'unavailable_provider', + 'invalid_request', + 'invalid_model', + 'invalid_context', + 'forbidden', + 'provider_error', + 'tool_error', + 'budget_exceeded', + 'lock_unavailable', + 'serialization_error', + 'cancelled', + 'internal', +] + + +class AskError(Exception): + """Base class for all Ask errors.""" + + code: AskErrorCode = 'internal' + http_status: int = 500 + preflight: bool = False + public_message: str = 'An internal error occurred.' + + def __init__(self, public_message: Optional[str] = None) -> None: + if public_message is not None: + self.public_message = public_message + super().__init__(self.public_message) + + +# --- Pre-flight errors: raised before streaming; transport returns HTTP --- + + +class InvalidRequest(AskError): + code: AskErrorCode = 'invalid_request' + http_status = 400 + preflight = True + public_message = 'The request was invalid.' + + +class InvalidModel(AskError): + code: AskErrorCode = 'invalid_model' + http_status = 400 + preflight = True + public_message = 'The requested model is not available.' + + +class InvalidContext(AskError): + code: AskErrorCode = 'invalid_context' + http_status = 400 + preflight = True + public_message = 'The request context was invalid.' + + +class Forbidden(AskError): + code: AskErrorCode = 'forbidden' + http_status = 403 + preflight = True + public_message = 'You do not have access to this conversation.' + + +class UnavailableProvider(AskError): + code: AskErrorCode = 'unavailable_provider' + http_status = 503 + preflight = True + public_message = 'No language model provider is available.' + + +class LockUnavailable(AskError): + code: AskErrorCode = 'lock_unavailable' + http_status = 409 + preflight = True + public_message = 'Another turn is already in progress for this conversation.' + + +# --- In-stream errors: surfaced as a terminal ``error`` event --- + + +class ProviderError(AskError): + code: AskErrorCode = 'provider_error' + public_message = 'The language model provider failed to respond.' + + +class ToolExecutionError(AskError): + code: AskErrorCode = 'tool_error' + public_message = 'A tool failed to execute.' + + +class BudgetExceeded(AskError): + code: AskErrorCode = 'budget_exceeded' + public_message = 'The response exceeded the configured limits.' + + +class Cancelled(AskError): + code: AskErrorCode = 'cancelled' + public_message = 'The request was cancelled.' + + +class SerializationError(AskError): + code: AskErrorCode = 'serialization_error' + public_message = 'The response could not be serialized.' + + +class InternalError(AskError): + code: AskErrorCode = 'internal' + public_message = 'An internal error occurred.' + + +def to_public_payload(err: AskError) -> Dict[str, str]: + """Return the safe, client-facing ``{code, message}`` for an error.""" + return {'code': err.code, 'message': err.public_message} diff --git a/osprey_worker/src/osprey/worker/lib/ask/ports.py b/osprey_worker/src/osprey/worker/lib/ask/ports.py new file mode 100644 index 00000000..1eab4822 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/ports.py @@ -0,0 +1,136 @@ +"""Host adapter Protocols and the AskConfig ports bundle. + +The host supplies concrete implementations of these Protocols; the generic package +defines only the interfaces (and test fakes). See README section 3.4 for the invariant +each port must uphold. Nothing here imports Discord/Smite/SML/Druid or a vendor SDK. +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Any, Mapping, Optional, Protocol, Sequence, runtime_checkable + +from osprey.worker.lib.ask.contracts import ( + AskLimits, + AskRequest, + ContextSnapshot, + Conversation, + Principal, + ResolvedModel, +) +from osprey.worker.lib.llm import LLMResponse, ToolCall, ToolRegistry, ToolResult + + +@runtime_checkable +class ModelPolicy(Protocol): + def resolve(self, requested_model: Optional[str], principal: Principal) -> ResolvedModel: + """Validate the requested model and return a ready provider + limits. + + Raise ``InvalidModel`` for a rejected model. Return ``provider=None`` (or + raise ``UnavailableProvider``) when no provider is available or it cannot + initialize -- both map to a pre-flight ``unavailable_provider``. + """ + ... + + +@runtime_checkable +class ConversationStore(Protocol): + def load(self, conversation_id: str, principal: Principal) -> Optional[Conversation]: + """Load a conversation, enforcing principal ownership. + + Return ``None`` when the id is unknown. Raise ``Forbidden`` when the + conversation is owned by another principal. + """ + ... + + def create(self, principal: Principal) -> Conversation: + """Create and return a new conversation owned by ``principal``.""" + ... + + def append_assistant_turn( + self, + conversation: Conversation, + request: AskRequest, + response: LLMResponse, + evidence: Sequence[Mapping[str, Any]], + ) -> None: + """Persist a successful assistant turn. Called only on the success path.""" + ... + + +@runtime_checkable +class ConversationLock(Protocol): + def acquire(self, conversation_id: str, principal: Principal) -> AbstractContextManager[None]: + """Return a context manager serializing turns for this conversation. + + Raise ``LockUnavailable`` when a turn is already in flight. The manager is + released on every terminal path. + """ + ... + + +@runtime_checkable +class ContextSnapshotProvider(Protocol): + def resolve(self, context_ref: str, principal: Principal) -> Optional[ContextSnapshot]: + """Validate a ``context_ref`` and return a snapshot, or ``None`` if invalid.""" + ... + + def render(self, snapshot: ContextSnapshot) -> str: + """Render a snapshot to grounding text (truncated to limits by the service).""" + ... + + +@runtime_checkable +class ToolRegistryFactory(Protocol): + def build(self, principal: Principal, snapshot: Optional[ContextSnapshot]) -> ToolRegistry: + """Build a principal-scoped tool registry for a turn.""" + ... + + +@runtime_checkable +class Redactor(Protocol): + def redact_text(self, text: str) -> str: + """Redact a plain string before it crosses the transport boundary.""" + ... + + def redact_payload(self, value: Any) -> Any: + """Recursively redact a JSON-compatible payload (dict/list/str/scalar).""" + ... + + +@runtime_checkable +class AuditSink(Protocol): + def record( + self, + principal: Principal, + conversation_id: str, + turn_id: str, + outcome: str, + detail: Mapping[str, Any], + ) -> None: + """Record a terminal outcome. Best-effort; failures must not break the stream.""" + ... + + +@runtime_checkable +class EvidenceNormalizer(Protocol): + def normalize(self, tool_call: ToolCall, tool_result: ToolResult) -> dict[str, Any]: + """Map a (tool_call, tool_result) pair to a ``query_result`` payload.""" + ... + + +@dataclass +class AskConfig: + """The ports + limits bundle an :class:`AskService` is constructed with.""" + + policy: ModelPolicy + conversations: ConversationStore + lock: ConversationLock + tools: ToolRegistryFactory + limits: AskLimits = AskLimits() + context: Optional[ContextSnapshotProvider] = None + redactor: Optional[Redactor] = None + audit: Optional[AuditSink] = None + evidence: Optional[EvidenceNormalizer] = None + system_prompt: Optional[str] = None diff --git a/osprey_worker/src/osprey/worker/lib/ask/service.py b/osprey_worker/src/osprey/worker/lib/ask/service.py new file mode 100644 index 00000000..facf3474 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/service.py @@ -0,0 +1,256 @@ +"""AskService: the host-neutral turn orchestrator. + +``run_turn`` is a generator. On the first ``next()`` it performs all pre-flight checks +(raising typed pre-flight errors before yielding anything), then streams a +deterministic, bounded sequence of events over the provider and a principal-scoped +tool registry, persists a successful turn, and yields exactly one terminal event. +The conversation lock is released on every terminal path, including client disconnect. +""" + +from __future__ import annotations + +import uuid +from typing import Any, Callable, Dict, Iterator, List, Optional + +from osprey.worker.lib.ask.contracts import ( + AskEvent, + AskLimits, + AskRequest, + ContextSnapshot, + Conversation, + Principal, + ResolvedModel, +) +from osprey.worker.lib.ask.errors import ( + AskError, + BudgetExceeded, + Cancelled, + Forbidden, + InternalError, + InvalidContext, + InvalidRequest, + ProviderError, + ToolExecutionError, + UnavailableProvider, + to_public_payload, +) +from osprey.worker.lib.ask.ports import AskConfig +from osprey.worker.lib.llm import LLMMessage, LLMResponse, LLMUsage, ToolCall, ToolResult + +_TRUNCATION_MARK = '…[truncated]' + + +def _new_id() -> str: + return uuid.uuid4().hex + + +def _usage_payload(usage: Optional[LLMUsage]) -> Optional[Dict[str, int]]: + if usage is None: + return None + return { + 'input_tokens': usage.input_tokens, + 'output_tokens': usage.output_tokens, + 'cache_read_tokens': usage.cache_read_tokens, + 'cache_write_tokens': usage.cache_write_tokens, + } + + +class AskService: + """Drives a single validated user turn to a bounded stream of :class:`AskEvent`.""" + + def __init__(self, config: AskConfig) -> None: + self._cfg = config + + def run_turn( + self, + request: AskRequest, + principal: Principal, + *, + cancel: Optional[Callable[[], bool]] = None, + ) -> Iterator[AskEvent]: + cfg = self._cfg + + # ---- PRE-FLIGHT: raises before any yield; the transport maps to HTTP ---- + if not request.message or not request.message.strip(): + raise InvalidRequest() + + resolved: ResolvedModel = cfg.policy.resolve(request.model, principal) + if resolved.provider is None: + raise UnavailableProvider() + provider = resolved.provider + limits: AskLimits = resolved.limits or cfg.limits + + snapshot: Optional[ContextSnapshot] = None + grounding: Optional[str] = None + if request.context_ref is not None: + context_provider = cfg.context + if context_provider is None: + raise InvalidContext() + snapshot = context_provider.resolve(request.context_ref, principal) + if snapshot is None: + raise InvalidContext() + try: + grounding = context_provider.render(snapshot)[: limits.max_context_chars] + except AskError: + raise + except Exception as exc: # render failure is a pre-flight invalid-context, not mid-stream + raise InvalidContext() from exc + + conversation = self._load_or_create(request.conversation_id, principal) + lock_cm = cfg.lock.acquire(conversation.id, principal) # may raise LockUnavailable + + # ---- STREAM: lock held; released on success, error, and GeneratorExit ---- + turn_id = _new_id() + cid = conversation.id + with lock_cm: + try: + yield AskEvent('conversation_started', {'conversation_id': cid, 'turn_id': turn_id}, cid, turn_id) + try: + registry = cfg.tools.build(principal, snapshot) # principal-scoped tools + except AskError: + raise # a factory may reject with e.g. Forbidden; keep its code + except Exception as exc: + raise ToolExecutionError() from exc # framework tool failure => terminal tool_error + messages = self._build_messages(conversation, request, grounding, limits) + evidence_log: List[Dict[str, Any]] = [] + final: Optional[LLMResponse] = None + + for i in range(limits.max_tool_iterations): + self._raise_if_cancelled(cancel) + try: + resp = provider.chat( + messages=messages, + system=cfg.system_prompt, + tools=registry.definitions(), + model=resolved.model, + max_tokens=limits.max_output_tokens, + ) + except Exception as exc: + raise ProviderError() from exc # never leak vendor text + self._raise_if_cancelled(cancel) # after provider return, before side effects + + if not resp.tool_calls: + final = resp + break + if i == limits.max_tool_iterations - 1: + raise BudgetExceeded() # do not dispatch tools we can't feed back + + messages.append( + LLMMessage(role='assistant', content=resp.text or None, tool_calls=list(resp.tool_calls)) + ) + results: List[ToolResult] = [] + for tc in resp.tool_calls: # deterministic order + self._raise_if_cancelled(cancel) + yield AskEvent('tool_call', self._tool_call_payload(tc), cid, turn_id) + try: + tr = registry.dispatch(tc) # handler errors are bounded is_error results + except Exception as exc: + raise ToolExecutionError() from exc # framework tool failure is terminal + results.append(tr) + evidence = self._evidence_payload(tc, tr, limits) + evidence_log.append(evidence) + yield AskEvent('query_result', evidence, cid, turn_id) + messages.append(LLMMessage(role='tool', tool_results=results)) + else: + raise BudgetExceeded() # exhausted iterations without a final answer + + assert final is not None # guaranteed: loop either breaks with final or raises + self._raise_if_cancelled(cancel) + # Build + validate all terminal data BEFORE committing, so a budget failure + # never persists a "successful" turn whose terminal event cannot ship. + text = self._finalize_text(final.text, limits) + done_payload = {'conversation_id': cid, 'turn_id': turn_id, 'usage': _usage_payload(final.usage)} + yield AskEvent('assistant_message', {'text': text}, cid, turn_id) + cfg.conversations.append_assistant_turn(conversation, request, final, evidence_log) + self._safe_audit(principal, cid, turn_id, 'done', {}) + yield AskEvent('done', done_payload, cid, turn_id) + except AskError as err: + self._safe_audit(principal, cid, turn_id, 'error', {'code': err.code}) + yield self._error_event(err, cid, turn_id) # streaming began => terminal error event + except Exception: + self._safe_audit(principal, cid, turn_id, 'error', {'code': 'internal'}) + yield self._error_event(InternalError(), cid, turn_id) + + # --- helpers ------------------------------------------------------------- + + def _load_or_create(self, conversation_id: Optional[str], principal: Principal) -> Conversation: + if conversation_id is None: + return self._cfg.conversations.create(principal) + conv = self._cfg.conversations.load(conversation_id, principal) + if conv is None: + # A provided id that is unknown or not owned is treated as forbidden, so we + # never leak the existence of another principal's conversation. + raise Forbidden() + return conv + + def _build_messages( + self, + conversation: Conversation, + request: AskRequest, + grounding: Optional[str], + limits: AskLimits, + ) -> List[LLMMessage]: + history = list(conversation.messages)[-limits.max_history_messages :] + messages: List[LLMMessage] = [] + if grounding: + messages.append(LLMMessage(role='system', content=grounding)) + messages.extend(history) + messages.append(LLMMessage(role='user', content=request.message)) + return messages + + def _finalize_text(self, text: str, limits: AskLimits) -> str: + redactor = self._cfg.redactor + if redactor is not None: + text = redactor.redact_text(text) + if len(text) > limits.max_output_chars: + raise BudgetExceeded() + return text + + def _tool_call_payload(self, tc: ToolCall) -> Dict[str, Any]: + return {'id': tc.id, 'name': tc.name, 'arguments': self._redact_value(tc.arguments)} + + def _evidence_payload(self, tc: ToolCall, tr: ToolResult, limits: AskLimits) -> Dict[str, Any]: + cfg = self._cfg + if cfg.evidence is not None: + payload: Dict[str, Any] = dict(cfg.evidence.normalize(tc, tr)) + else: + payload = { + 'tool_call_id': tr.tool_call_id, + 'name': tc.name, + 'is_error': tr.is_error, + 'content': tr.content, + } + payload = self._redact_value(payload) + return self._truncate_strings(payload, limits.max_evidence_chars) + + def _redact_value(self, value: Any) -> Any: + redactor = self._cfg.redactor + if redactor is None: + return value + return redactor.redact_payload(value) + + @staticmethod + def _truncate_strings(value: Any, limit: int) -> Any: + if isinstance(value, str): + return value if len(value) <= limit else value[:limit] + _TRUNCATION_MARK + if isinstance(value, dict): + return {k: AskService._truncate_strings(v, limit) for k, v in value.items()} + if isinstance(value, list): + return [AskService._truncate_strings(v, limit) for v in value] + return value + + def _error_event(self, err: AskError, cid: str, turn_id: str) -> AskEvent: + return AskEvent('error', to_public_payload(err), cid, turn_id) + + def _raise_if_cancelled(self, cancel: Optional[Callable[[], bool]]) -> None: + if cancel is not None and cancel(): + raise Cancelled() + + def _safe_audit(self, principal: Principal, cid: str, turn_id: str, outcome: str, detail: Dict[str, Any]) -> None: + audit = self._cfg.audit + if audit is None: + return + try: + audit.record(principal, cid, turn_id, outcome, detail) + except Exception: # audit is best-effort and must never corrupt the stream + pass diff --git a/osprey_worker/src/osprey/worker/lib/ask/tests/fakes.py b/osprey_worker/src/osprey/worker/lib/ask/tests/fakes.py new file mode 100644 index 00000000..d6517f08 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/tests/fakes.py @@ -0,0 +1,240 @@ +"""In-memory fakes implementing the Ask host ports, for service tests. + +None of these import product or vendor code; they are built from the generic Ask + +LLM types alone, which is itself the demonstration that a host can wire the service +without importing Discord/Smite/SML/Druid or a concrete vendor SDK (AC2.1/AC2.2). +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager, contextmanager +from typing import Any, Dict, Iterator, List, Mapping, Optional, Sequence + +from osprey.worker.lib.ask.contracts import ( + AskLimits, + AskRequest, + ContextSnapshot, + Conversation, + Principal, + ResolvedModel, +) +from osprey.worker.lib.ask.errors import Forbidden, InvalidModel, LockUnavailable +from osprey.worker.lib.llm import ( + BaseLLMProvider, + LLMMessage, + LLMResponse, + ToolCall, + ToolDefinition, + ToolRegistry, + ToolResult, +) + + +class ScriptedProvider(BaseLLMProvider): + """Returns a queued list of responses; records the messages seen on each call.""" + + def __init__(self, responses: Sequence[LLMResponse]) -> None: + self._responses = list(responses) + self.calls: List[List[LLMMessage]] = [] + + def chat( + self, + *, + messages: Sequence[LLMMessage], + system: Optional[str] = None, + tools: Optional[Sequence[ToolDefinition]] = None, + model: Optional[str] = None, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + **params: Any, + ) -> LLMResponse: + self.calls.append(list(messages)) + return self._responses[len(self.calls) - 1] + + +class RaisingProvider(BaseLLMProvider): + """Raises on chat(); the message carries a secret to prove it never leaks.""" + + def __init__(self, exc: Optional[BaseException] = None) -> None: + self._exc = exc or RuntimeError('vendor stack trace secret: sk-DEADBEEF') + + def chat( + self, + *, + messages: Sequence[LLMMessage], + system: Optional[str] = None, + tools: Optional[Sequence[ToolDefinition]] = None, + model: Optional[str] = None, + max_tokens: Optional[int] = None, + temperature: Optional[float] = None, + **params: Any, + ) -> LLMResponse: + raise self._exc + + +class AllowAllModelPolicy: + def __init__( + self, + provider: BaseLLMProvider, + model: Optional[str] = 'test-model', + limits: Optional[AskLimits] = None, + ) -> None: + self._provider = provider + self._model = model + self._limits = limits + + def resolve(self, requested_model: Optional[str], principal: Principal) -> ResolvedModel: + return ResolvedModel(model=requested_model or self._model, provider=self._provider, limits=self._limits) + + +class RejectingModelPolicy: + def resolve(self, requested_model: Optional[str], principal: Principal) -> ResolvedModel: + raise InvalidModel() + + +class NoProviderPolicy: + def resolve(self, requested_model: Optional[str], principal: Principal) -> ResolvedModel: + return ResolvedModel(model=requested_model, provider=None, limits=None) + + +class InMemoryConversationStore: + def __init__(self) -> None: + self._by_id: Dict[str, Conversation] = {} + self._counter = 0 + self.appended: List[Dict[str, Any]] = [] + + def add(self, conversation: Conversation) -> Conversation: + self._by_id[conversation.id] = conversation + return conversation + + def create(self, principal: Principal) -> Conversation: + self._counter += 1 + conv = Conversation(id=f'conv-{self._counter}', principal_id=principal.id) + self._by_id[conv.id] = conv + return conv + + def load(self, conversation_id: str, principal: Principal) -> Optional[Conversation]: + conv = self._by_id.get(conversation_id) + if conv is None: + return None + if conv.principal_id != principal.id: + raise Forbidden() + return conv + + def append_assistant_turn( + self, + conversation: Conversation, + request: AskRequest, + response: LLMResponse, + evidence: Sequence[Mapping[str, Any]], + ) -> None: + self.appended.append( + { + 'conversation_id': conversation.id, + 'request': request, + 'response': response, + 'evidence': [dict(e) for e in evidence], + } + ) + + +class InMemoryLock: + def __init__(self) -> None: + self.acquired = 0 + self.released = 0 + + @contextmanager + def _cm(self) -> Iterator[None]: + self.acquired += 1 + try: + yield + finally: + self.released += 1 + + def acquire(self, conversation_id: str, principal: Principal) -> AbstractContextManager[None]: + return self._cm() + + +class BlockedLock: + def acquire(self, conversation_id: str, principal: Principal) -> AbstractContextManager[None]: + raise LockUnavailable() + + +class StaticContextProvider: + def __init__( + self, + snapshots: Optional[Mapping[str, ContextSnapshot]] = None, + render_text: str = 'GROUNDING', + render_error: bool = False, + ) -> None: + self._snapshots = dict(snapshots or {}) + self._render_text = render_text + self._render_error = render_error + + def resolve(self, context_ref: str, principal: Principal) -> Optional[ContextSnapshot]: + return self._snapshots.get(context_ref) + + def render(self, snapshot: ContextSnapshot) -> str: + if self._render_error: + raise RuntimeError('render blew up') + return self._render_text + + +class RecordingToolRegistryFactory: + def __init__(self, registry: ToolRegistry) -> None: + self._registry = registry + self.built_with_principal: Optional[Principal] = None + self.built_with_snapshot: Optional[ContextSnapshot] = None + + def build(self, principal: Principal, snapshot: Optional[ContextSnapshot]) -> ToolRegistry: + self.built_with_principal = principal + self.built_with_snapshot = snapshot + return self._registry + + +class RaisingToolRegistryFactory: + def build(self, principal: Principal, snapshot: Optional[ContextSnapshot]) -> ToolRegistry: + raise RuntimeError('tool factory blew up') + + +class TokenRedactor: + """Replaces a secret token in text and (recursively) in structured payloads.""" + + def __init__(self, secret: str, replacement: str = '[REDACTED]') -> None: + self._secret = secret + self._replacement = replacement + + def redact_text(self, text: str) -> str: + return text.replace(self._secret, self._replacement) + + def redact_payload(self, value: Any) -> Any: + if isinstance(value, str): + return value.replace(self._secret, self._replacement) + if isinstance(value, dict): + return {k: self.redact_payload(v) for k, v in value.items()} + if isinstance(value, list): + return [self.redact_payload(v) for v in value] + return value + + +class RecordingAudit: + def __init__(self, fail: bool = False) -> None: + self.records: List[Dict[str, Any]] = [] + self._fail = fail + + def record( + self, + principal: Principal, + conversation_id: str, + turn_id: str, + outcome: str, + detail: Mapping[str, Any], + ) -> None: + if self._fail: + raise RuntimeError('audit sink down') + self.records.append({'outcome': outcome, 'detail': dict(detail), 'turn_id': turn_id}) + + +class CustomEvidence: + def normalize(self, tool_call: ToolCall, tool_result: ToolResult) -> Dict[str, Any]: + return {'kind': 'custom', 'tool': tool_call.name, 'ok': not tool_result.is_error} diff --git a/osprey_worker/src/osprey/worker/lib/ask/tests/test_errors.py b/osprey_worker/src/osprey/worker/lib/ask/tests/test_errors.py new file mode 100644 index 00000000..14c275ef --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/tests/test_errors.py @@ -0,0 +1,66 @@ +"""Error taxonomy: stable codes/statuses, preflight flags, and safe public payloads.""" + +from osprey.worker.lib.ask.errors import ( + AskError, + BudgetExceeded, + Cancelled, + Forbidden, + InternalError, + InvalidContext, + InvalidModel, + InvalidRequest, + LockUnavailable, + ProviderError, + SerializationError, + ToolExecutionError, + UnavailableProvider, + to_public_payload, +) + +PREFLIGHT = [ + (InvalidRequest, 'invalid_request', 400), + (InvalidModel, 'invalid_model', 400), + (InvalidContext, 'invalid_context', 400), + (Forbidden, 'forbidden', 403), + (UnavailableProvider, 'unavailable_provider', 503), + (LockUnavailable, 'lock_unavailable', 409), +] + +IN_STREAM = [ + (ProviderError, 'provider_error'), + (ToolExecutionError, 'tool_error'), + (BudgetExceeded, 'budget_exceeded'), + (Cancelled, 'cancelled'), + (SerializationError, 'serialization_error'), + (InternalError, 'internal'), +] + + +def test_preflight_errors_have_code_status_and_flag(): + for cls, code, status in PREFLIGHT: + err = cls() + assert isinstance(err, AskError) + assert err.code == code + assert err.http_status == status + assert err.preflight is True + + +def test_in_stream_errors_are_not_preflight(): + for cls, code in IN_STREAM: + err = cls() + assert err.code == code + assert err.preflight is False + + +def test_public_payload_is_safe_and_hides_wrapped_exception(): + secret = 'sk-DEADBEEF top secret' + try: + raise ValueError(secret) + except ValueError as exc: + err = ProviderError() + err.__cause__ = exc + payload = to_public_payload(err) + assert set(payload) == {'code', 'message'} + assert payload['code'] == 'provider_error' + assert secret not in payload['message'] + assert 'ValueError' not in payload['message'] diff --git a/osprey_worker/src/osprey/worker/lib/ask/tests/test_service.py b/osprey_worker/src/osprey/worker/lib/ask/tests/test_service.py new file mode 100644 index 00000000..1b169321 --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/ask/tests/test_service.py @@ -0,0 +1,430 @@ +"""AskService: deterministic, bounded, safe turns. + +Covers osprey-ask-ai AC2.1-AC2.5 and AC3.1-AC3.5. Tests drain ``run_turn`` and assert +on the event-type sequence and payloads. These live under ``lib/`` and therefore run +through the Docker harness (``run-tests.sh``); the autouse Postgres fixture applies but +is not otherwise exercised here. +""" + +from __future__ import annotations + +from typing import Any, List, Optional + +import pytest +from osprey.worker.lib.ask import ( + AskConfig, + AskEvent, + AskLimits, + AskRequest, + AskService, + Forbidden, + InvalidContext, + InvalidModel, + InvalidRequest, + LockUnavailable, + Principal, + UnavailableProvider, +) +from osprey.worker.lib.ask.contracts import ContextSnapshot, Conversation +from osprey.worker.lib.ask.tests.fakes import ( + AllowAllModelPolicy, + BlockedLock, + CustomEvidence, + InMemoryConversationStore, + InMemoryLock, + NoProviderPolicy, + RaisingProvider, + RaisingToolRegistryFactory, + RecordingAudit, + RecordingToolRegistryFactory, + RejectingModelPolicy, + ScriptedProvider, + StaticContextProvider, + TokenRedactor, +) +from osprey.worker.lib.llm import LLMMessage, LLMResponse, Tool, ToolCall, ToolRegistry + +PRINCIPAL = Principal(id='u1', email='u1@example.com') +OTHER = Principal(id='u2', email='u2@example.com') + + +def _text(text: str) -> LLMResponse: + return LLMResponse(text=text) + + +def _tools(*calls: ToolCall) -> LLMResponse: + return LLMResponse(text='', tool_calls=list(calls)) + + +def _registry(*tools: Tool) -> ToolRegistry: + reg = ToolRegistry() + for t in tools: + reg.register(t) + return reg + + +def _ok_tool(name: str = 'lookup', result: Any = None) -> Tool: + def handler(**kwargs: Any) -> Any: + return {'result': 'ok'} if result is None else result + + return Tool(name=name, description='d', handler=handler) + + +def _raising_tool(name: str = 'boom', message: str = 'kaboom') -> Tool: + def handler(**kwargs: Any) -> Any: + raise ValueError(message) + + return Tool(name=name, description='d', handler=handler) + + +def _build( + provider: Any, + *, + registry: Optional[ToolRegistry] = None, + store: Optional[InMemoryConversationStore] = None, + lock: Optional[Any] = None, + context: Optional[Any] = None, + redactor: Optional[Any] = None, + audit: Optional[Any] = None, + evidence: Optional[Any] = None, + limits: Optional[AskLimits] = None, + tool_factory: Optional[Any] = None, +) -> Any: + store = store or InMemoryConversationStore() + lock = lock or InMemoryLock() + factory = tool_factory or RecordingToolRegistryFactory(registry or ToolRegistry()) + cfg = AskConfig( + policy=AllowAllModelPolicy(provider), + conversations=store, + lock=lock, + tools=factory, + limits=limits or AskLimits(), + context=context, + redactor=redactor, + audit=audit, + evidence=evidence, + ) + return AskService(cfg), store, lock, factory + + +def _drain( + service: AskService, request: AskRequest, principal: Principal = PRINCIPAL, cancel: Any = None +) -> List[AskEvent]: + return list(service.run_turn(request, principal, cancel=cancel)) + + +def _types(events: List[AskEvent]) -> List[str]: + return [e.type for e in events] + + +# --- AC3.1 / AC3.2 --------------------------------------------------------- + + +def test_happy_path_one_tool_call_sequence_and_persist(): + call = ToolCall(id='c1', name='lookup', arguments={'q': 7}) + provider = ScriptedProvider([_tools(call), _text('final answer')]) + service, store, lock, _ = _build(provider, registry=_registry(_ok_tool('lookup', result={'v': 1}))) + events = _drain(service, AskRequest(message='hi')) + assert _types(events) == ['conversation_started', 'tool_call', 'query_result', 'assistant_message', 'done'] + assert _types(events).count('done') == 1 + assert not any(e.type == 'error' for e in events) + assert events[3].payload['text'] == 'final answer' + # tool result fed back: the 2nd provider call ends with a 'tool' message + assert len(provider.calls) == 2 + assert provider.calls[1][-1].role == 'tool' + # persisted exactly once with the final response + assert len(store.appended) == 1 + assert store.appended[0]['response'].text == 'final answer' + assert lock.acquired == 1 and lock.released == 1 + + +def test_new_conversation_id_flows_through_events_and_persist(): + provider = ScriptedProvider([_text('ok')]) + service, store, _, _ = _build(provider) + events = _drain(service, AskRequest(message='hi')) + cid = events[0].conversation_id + assert events[0].type == 'conversation_started' + assert cid and cid.startswith('conv-') + assert all(e.conversation_id == cid for e in events) + assert store.appended[0]['conversation_id'] == cid + + +# --- AC3.5 (multiple tool calls; iteration budget) ------------------------- + + +def test_multiple_tool_calls_ordered_deterministically(): + c1 = ToolCall(id='c1', name='a', arguments={}) + c2 = ToolCall(id='c2', name='b', arguments={}) + provider = ScriptedProvider([_tools(c1, c2), _text('done')]) + service, _, _, _ = _build(provider, registry=_registry(_ok_tool('a', 'ra'), _ok_tool('b', 'rb'))) + events = _drain(service, AskRequest(message='hi')) + assert _types(events) == [ + 'conversation_started', + 'tool_call', + 'query_result', + 'tool_call', + 'query_result', + 'assistant_message', + 'done', + ] + assert [e.payload['name'] for e in events if e.type == 'tool_call'] == ['a', 'b'] + + +def test_tool_iteration_budget_does_not_dispatch_on_final_iteration(): + call = ToolCall(id='c', name='counter', arguments={}) + provider = ScriptedProvider([_tools(call)] * 10) + dispatched = {'n': 0} + + def handler(**kwargs: Any) -> Any: + dispatched['n'] += 1 + return 'x' + + reg = _registry(Tool(name='counter', description='d', handler=handler)) + service, store, lock, _ = _build(provider, registry=reg, limits=AskLimits(max_tool_iterations=3)) + events = _drain(service, AskRequest(message='hi')) + assert events[-1].type == 'error' + assert events[-1].payload['code'] == 'budget_exceeded' + # iters 0 and 1 dispatch; iter 2 is the final allowed call -> raise before dispatch + assert dispatched['n'] == 2 + assert store.appended == [] + assert lock.acquired == 1 and lock.released == 1 + + +# --- AC2.5 (bounds) -------------------------------------------------------- + + +def test_output_char_budget_exceeded_blocks_persist(): + provider = ScriptedProvider([_text('x' * 100)]) + service, store, _, _ = _build(provider, limits=AskLimits(max_output_chars=10)) + events = _drain(service, AskRequest(message='hi')) + assert _types(events) == ['conversation_started', 'error'] + assert events[-1].payload['code'] == 'budget_exceeded' + assert store.appended == [] + + +def test_history_truncated_to_limit(): + provider = ScriptedProvider([_text('ok')]) + store = InMemoryConversationStore() + store.add( + Conversation( + id='c1', + principal_id=PRINCIPAL.id, + messages=[LLMMessage(role='user', content=f'm{i}') for i in range(50)], + ) + ) + service, _, _, _ = _build(provider, store=store, limits=AskLimits(max_history_messages=5)) + _drain(service, AskRequest(message='new', conversation_id='c1')) + seen = provider.calls[0] + assert len(seen) == 6 # 5 history + 1 new user message + assert seen[-1].content == 'new' + + +def test_context_grounding_truncated_and_tools_scoped_to_principal(): + snap = ContextSnapshot(ref='r1') + context = StaticContextProvider(snapshots={'r1': snap}, render_text='G' * 100) + provider = ScriptedProvider([_text('ok')]) + factory = RecordingToolRegistryFactory(ToolRegistry()) + service, _, _, _ = _build(provider, context=context, limits=AskLimits(max_context_chars=10), tool_factory=factory) + _drain(service, AskRequest(message='hi', context_ref='r1')) + seen = provider.calls[0] + assert seen[0].role == 'system' + assert seen[0].content == 'G' * 10 # truncated grounding + assert factory.built_with_principal is PRINCIPAL # server principal, not context + assert factory.built_with_snapshot is snap + + +def test_evidence_content_truncated(): + call = ToolCall(id='c', name='lookup', arguments={}) + provider = ScriptedProvider([_tools(call), _text('ok')]) + service, _, _, _ = _build( + provider, registry=_registry(_ok_tool('lookup', result='y' * 100)), limits=AskLimits(max_evidence_chars=10) + ) + events = _drain(service, AskRequest(message='hi')) + content = [e for e in events if e.type == 'query_result'][0].payload['content'] + assert content.startswith('y' * 10) + assert content.endswith('…[truncated]') + + +# --- AC2.4 (pre-flight rejections; raised before any event) ---------------- + + +def test_empty_message_rejected_preflight(): + provider = ScriptedProvider([_text('x')]) + service, _, lock, _ = _build(provider) + with pytest.raises(InvalidRequest): + _drain(service, AskRequest(message=' ')) + assert provider.calls == [] + assert lock.acquired == 0 + + +def test_invalid_model_rejected_preflight(): + cfg = AskConfig( + policy=RejectingModelPolicy(), + conversations=InMemoryConversationStore(), + lock=InMemoryLock(), + tools=RecordingToolRegistryFactory(ToolRegistry()), + ) + with pytest.raises(InvalidModel): + list(AskService(cfg).run_turn(AskRequest(message='hi'), PRINCIPAL)) + + +def test_unavailable_provider_rejected_preflight(): + cfg = AskConfig( + policy=NoProviderPolicy(), + conversations=InMemoryConversationStore(), + lock=InMemoryLock(), + tools=RecordingToolRegistryFactory(ToolRegistry()), + ) + with pytest.raises(UnavailableProvider): + list(AskService(cfg).run_turn(AskRequest(message='hi'), PRINCIPAL)) + + +def test_unknown_context_ref_rejected_preflight(): + provider = ScriptedProvider([_text('x')]) + service, _, lock, _ = _build(provider, context=StaticContextProvider(snapshots={})) + with pytest.raises(InvalidContext): + _drain(service, AskRequest(message='hi', context_ref='missing')) + assert lock.acquired == 0 + + +def test_context_ref_without_provider_rejected_preflight(): + provider = ScriptedProvider([_text('x')]) + service, _, _, _ = _build(provider, context=None) + with pytest.raises(InvalidContext): + _drain(service, AskRequest(message='hi', context_ref='r1')) + + +def test_context_render_failure_is_preflight_invalid_context(): + provider = ScriptedProvider([_text('x')]) + context = StaticContextProvider(snapshots={'r1': ContextSnapshot(ref='r1')}, render_error=True) + service, _, lock, _ = _build(provider, context=context) + with pytest.raises(InvalidContext): + _drain(service, AskRequest(message='hi', context_ref='r1')) + assert provider.calls == [] + assert lock.acquired == 0 + + +def test_foreign_conversation_forbidden_preflight(): + provider = ScriptedProvider([_text('x')]) + store = InMemoryConversationStore() + store.add(Conversation(id='c1', principal_id=OTHER.id)) + service, _, lock, _ = _build(provider, store=store) + with pytest.raises(Forbidden): + _drain(service, AskRequest(message='hi', conversation_id='c1')) + assert lock.acquired == 0 + + +def test_unknown_conversation_id_forbidden_preflight(): + provider = ScriptedProvider([_text('x')]) + service, _, _, _ = _build(provider) + with pytest.raises(Forbidden): + _drain(service, AskRequest(message='hi', conversation_id='does-not-exist')) + + +def test_lock_unavailable_rejected_preflight(): + provider = ScriptedProvider([_text('x')]) + service, _, _, _ = _build(provider, lock=BlockedLock()) + with pytest.raises(LockUnavailable): + _drain(service, AskRequest(message='hi')) + assert provider.calls == [] + + +# --- AC3.3 / AC3.4 (failures are terminal + safe; lock released; no persist) -- + + +def test_provider_error_is_safe_terminal_and_releases_lock(): + provider = RaisingProvider() # message contains 'sk-DEADBEEF' + service, store, lock, _ = _build(provider) + events = _drain(service, AskRequest(message='hi')) + assert _types(events) == ['conversation_started', 'error'] + assert events[-1].payload['code'] == 'provider_error' + assert 'sk-DEADBEEF' not in events[-1].payload['message'] + assert 'RuntimeError' not in events[-1].payload['message'] + assert store.appended == [] + assert lock.acquired == 1 and lock.released == 1 + + +def test_tool_handler_exception_is_nonterminal_query_result(): + call = ToolCall(id='c', name='boom', arguments={}) + provider = ScriptedProvider([_tools(call), _text('recovered')]) + service, store, _, _ = _build(provider, registry=_registry(_raising_tool('boom', 'kaboom'))) + events = _drain(service, AskRequest(message='hi')) + assert _types(events) == ['conversation_started', 'tool_call', 'query_result', 'assistant_message', 'done'] + qr = [e for e in events if e.type == 'query_result'][0] + assert qr.payload['is_error'] is True + assert 'kaboom' in qr.payload['content'] + assert len(store.appended) == 1 + + +def test_tool_factory_failure_is_terminal_tool_error(): + provider = ScriptedProvider([_text('never')]) + service, store, lock, _ = _build(provider, tool_factory=RaisingToolRegistryFactory()) + events = _drain(service, AskRequest(message='hi')) + assert events[-1].type == 'error' + assert events[-1].payload['code'] == 'tool_error' + assert store.appended == [] + assert lock.acquired == 1 and lock.released == 1 + + +def test_cancellation_is_terminal_and_skips_persist(): + call = ToolCall(id='c', name='lookup', arguments={}) + provider = ScriptedProvider([_tools(call), _text('final')]) + state = {'n': 0} + + def cancel() -> bool: + state['n'] += 1 + return state['n'] >= 2 # allow the first check, cancel on the next + + service, store, lock, _ = _build(provider, registry=_registry(_ok_tool('lookup'))) + events = _drain(service, AskRequest(message='hi'), cancel=cancel) + assert events[-1].type == 'error' + assert events[-1].payload['code'] == 'cancelled' + assert store.appended == [] + assert lock.acquired == 1 and lock.released == 1 + + +def test_generator_exit_releases_lock_without_persist(): + call = ToolCall(id='c', name='lookup', arguments={}) + provider = ScriptedProvider([_tools(call), _text('final')]) + service, store, lock, _ = _build(provider, registry=_registry(_ok_tool('lookup'))) + gen = service.run_turn(AskRequest(message='hi'), PRINCIPAL) + assert next(gen).type == 'conversation_started' + gen.close() # simulate client disconnect + assert lock.acquired == 1 and lock.released == 1 + assert store.appended == [] + + +# --- AC4.4 boundary (structured redaction) + AC2.1 (adapters) -------------- + + +def test_structured_redaction_of_args_evidence_and_text(): + secret = 'sk-SECRET' + call = ToolCall(id='c', name='lookup', arguments={'token': secret, 'nested': {'k': secret}}) + provider = ScriptedProvider([_tools(call), _text(f'answer {secret}')]) + service, _, _, _ = _build( + provider, registry=_registry(_raising_tool('lookup', f'failed with {secret}')), redactor=TokenRedactor(secret) + ) + events = _drain(service, AskRequest(message='hi')) + tc = [e for e in events if e.type == 'tool_call'][0] + qr = [e for e in events if e.type == 'query_result'][0] + am = [e for e in events if e.type == 'assistant_message'][0] + assert secret not in str(tc.payload) and '[REDACTED]' in str(tc.payload) + assert secret not in str(qr.payload) + assert secret not in am.payload['text'] + + +def test_custom_evidence_normalizer_used(): + call = ToolCall(id='c', name='lookup', arguments={}) + provider = ScriptedProvider([_tools(call), _text('ok')]) + service, _, _, _ = _build(provider, registry=_registry(_ok_tool('lookup')), evidence=CustomEvidence()) + events = _drain(service, AskRequest(message='hi')) + qr = [e for e in events if e.type == 'query_result'][0] + assert qr.payload == {'kind': 'custom', 'tool': 'lookup', 'ok': True} + + +def test_audit_failure_does_not_break_stream(): + provider = ScriptedProvider([_text('ok')]) + service, store, _, _ = _build(provider, audit=RecordingAudit(fail=True)) + events = _drain(service, AskRequest(message='hi')) + assert events[-1].type == 'done' + assert len(store.appended) == 1