From 1773779dac2d84bc5c41c6ce5cd7912c78e09ce0 Mon Sep 17 00:00:00 2001 From: hailey Date: Tue, 14 Jul 2026 22:13:18 +0000 Subject: [PATCH] feat(ui-api): Ask SSE blueprint factory Add osprey_worker/ui_api/lib/ask: format_sse (versioned SSE framing that raises SerializationError on a non-serializable payload) and create_ask_blueprint -- a host-configurable Flask factory that validates the JSON body (strict Pydantic v1), derives the principal from a host-supplied provider (never the body), drives AskService.run_turn behind a safe pre-flight boundary (unexpected host-adapter failures become a safe JSON 500), and streams events with a terminal-seen guard so a started stream always ends in exactly one done/error. Not registered in create_app(); exercised via a dedicated test app. Also widen AskService.run_turn's return annotation to Generator so the transport can close() it for prompt lock release. Tests (16 passing under run-tests.sh) cover content-type/framing/versioning, malformed/ wrong-type/missing-provider/ownership/model/lock pre-flight errors, throwing factory/ principal safe 500s, safe provider-failure error frames, serialization-failure termination, and the terminal guard. Implements osprey-ask-ai AC4.1-AC4.5. --- .../src/osprey/worker/lib/ask/service.py | 4 +- .../worker/ui_api/osprey/lib/ask/__init__.py | 6 + .../worker/ui_api/osprey/lib/ask/blueprint.py | 133 +++++++++ .../worker/ui_api/osprey/lib/ask/sse.py | 30 ++ .../ui_api/osprey/lib/ask/tests/__init__.py | 0 .../osprey/lib/ask/tests/test_blueprint.py | 267 ++++++++++++++++++ .../ui_api/osprey/lib/ask/tests/test_sse.py | 29 ++ 7 files changed, 467 insertions(+), 2 deletions(-) create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/__init__.py create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/blueprint.py create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/sse.py create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/__init__.py create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_blueprint.py create mode 100644 osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_sse.py diff --git a/osprey_worker/src/osprey/worker/lib/ask/service.py b/osprey_worker/src/osprey/worker/lib/ask/service.py index facf3474..0c500bae 100644 --- a/osprey_worker/src/osprey/worker/lib/ask/service.py +++ b/osprey_worker/src/osprey/worker/lib/ask/service.py @@ -10,7 +10,7 @@ from __future__ import annotations import uuid -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional from osprey.worker.lib.ask.contracts import ( AskEvent, @@ -67,7 +67,7 @@ def run_turn( principal: Principal, *, cancel: Optional[Callable[[], bool]] = None, - ) -> Iterator[AskEvent]: + ) -> Generator[AskEvent, None, None]: cfg = self._cfg # ---- PRE-FLIGHT: raises before any yield; the transport maps to HTTP ---- diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/__init__.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/__init__.py new file mode 100644 index 00000000..267306f7 --- /dev/null +++ b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/__init__.py @@ -0,0 +1,6 @@ +"""Reusable Flask/SSE transport for the host-neutral Ask service.""" + +from osprey.worker.ui_api.osprey.lib.ask.blueprint import AskRequestModel, create_ask_blueprint +from osprey.worker.ui_api.osprey.lib.ask.sse import format_sse + +__all__ = ['AskRequestModel', 'create_ask_blueprint', 'format_sse'] diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/blueprint.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/blueprint.py new file mode 100644 index 00000000..848c3405 --- /dev/null +++ b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/blueprint.py @@ -0,0 +1,133 @@ +"""Reusable Flask blueprint factory exposing an AskService as versioned SSE. + +Host-configurable: the host supplies a ``service_factory`` (which builds an +:class:`AskService` from its concrete adapters) and a ``principal_provider`` (which +derives the server-authenticated principal -- never from the request body), and +chooses the URL. All pre-flight work runs inside a safe boundary so an unexpected +host-adapter failure becomes a safe JSON 500 rather than a raw stack trace. Once +streaming begins, exactly one terminal event is emitted. Nothing here imports +Discord/Smite/SML/Druid. +""" + +from __future__ import annotations + +from typing import Any, Callable, Generator, Iterator, Optional, Tuple + +import pydantic +from flask import Blueprint, Response, jsonify, request, stream_with_context +from osprey.worker.lib.ask import AskEvent, AskRequest, AskService, Principal +from osprey.worker.lib.ask.errors import ( + AskError, + InternalError, + InvalidRequest, + SerializationError, + to_public_payload, +) +from osprey.worker.ui_api.osprey.lib.ask.sse import format_sse + +ServiceFactory = Callable[[], AskService] +PrincipalProvider = Callable[[], Principal] + +_TERMINAL = ('done', 'error') + + +class AskRequestModel(pydantic.BaseModel): + """HTTP request body schema. Strict strings so a JSON number/bool is rejected.""" + + message: pydantic.StrictStr + conversation_id: Optional[pydantic.StrictStr] = None + model: Optional[pydantic.StrictStr] = None + context_ref: Optional[pydantic.StrictStr] = None + + class Config: + extra = 'forbid' + + +def _http_error(err: AskError) -> Tuple[Any, int]: + return jsonify({'error': to_public_payload(err)}), err.http_status + + +def create_ask_blueprint( + *, + name: str, + service_factory: ServiceFactory, + principal_provider: PrincipalProvider, + url: str = '/ask', +) -> Blueprint: + """Build a blueprint with a single POST route that streams Ask events as SSE.""" + bp = Blueprint(name, __name__) + + @bp.route(url, methods=['POST']) + def ask() -> Any: + body = request.get_json(silent=True) + if not isinstance(body, dict): + return _http_error(InvalidRequest()) + try: + model_in = AskRequestModel(**body) + except pydantic.ValidationError: + return _http_error(InvalidRequest()) + if not model_in.message.strip(): + return _http_error(InvalidRequest()) + + ask_request = AskRequest( + message=model_in.message, + conversation_id=model_in.conversation_id, + model=model_in.model, + context_ref=model_in.context_ref, + ) + + # Safe pre-flight boundary: principal derivation, service construction, generator + # creation, and the pre-flight next() must never leak a raw 500 / HTML page. + cancelled = {'v': False} + gen: Optional[Generator[AskEvent, None, None]] = None + try: + principal = principal_provider() + service = service_factory() + gen = service.run_turn(ask_request, principal, cancel=lambda: cancelled['v']) + first = next(gen) + except AskError as err: + if gen is not None: + gen.close() + return _http_error(err) + except StopIteration: + if gen is not None: + gen.close() + return _http_error(InternalError()) + except Exception: + if gen is not None: + gen.close() + return _http_error(InternalError()) + + active_gen = gen + assert active_gen is not None # reached only when next(gen) succeeded + + def stream() -> Iterator[str]: + terminal_seen = first.type in _TERMINAL + try: + yield format_sse(first) # conversation_started + for event in active_gen: + terminal_seen = terminal_seen or event.type in _TERMINAL + yield format_sse(event) + except SerializationError as err: + if not terminal_seen: + terminal_seen = True + yield format_sse(AskEvent('error', to_public_payload(err), first.conversation_id, first.turn_id)) + except GeneratorExit: + cancelled['v'] = True + active_gen.close() # service finally releases the lock + raise + except Exception: + if not terminal_seen: # terminal guard: never a second terminal after done + terminal_seen = True + yield format_sse( + AskEvent('error', to_public_payload(InternalError()), first.conversation_id, first.turn_id) + ) + finally: + active_gen.close() + + resp = Response(stream_with_context(stream()), mimetype='text/event-stream') + resp.headers['Cache-Control'] = 'no-cache' + resp.headers['X-Accel-Buffering'] = 'no' + return resp + + return bp diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/sse.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/sse.py new file mode 100644 index 00000000..19a72222 --- /dev/null +++ b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/sse.py @@ -0,0 +1,30 @@ +"""Versioned Server-Sent Events serialization for Ask events.""" + +from __future__ import annotations + +import json + +from osprey.worker.lib.ask import AskEvent +from osprey.worker.lib.ask.errors import SerializationError + + +def format_sse(event: AskEvent) -> str: + """Serialize an :class:`AskEvent` to a single SSE frame. + + The event name appears on the ``event:`` line and again as ``type`` inside the + JSON ``data`` payload, alongside the schema ``version``. A non-serializable payload + raises :class:`SerializationError` so the caller can emit a terminal error frame + (whose payload -- code + message strings -- is always serializable). + """ + data = { + 'version': event.version, + 'type': event.type, + 'conversation_id': event.conversation_id, + 'turn_id': event.turn_id, + 'payload': event.payload, + } + try: + body = json.dumps(data) + except (TypeError, ValueError) as exc: + raise SerializationError() from exc + return f'event: {event.type}\ndata: {body}\n\n' diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/__init__.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_blueprint.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_blueprint.py new file mode 100644 index 00000000..d25074ad --- /dev/null +++ b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_blueprint.py @@ -0,0 +1,267 @@ +"""Ask SSE blueprint: framing, versioning, pre-flight errors, safe failures, termination. + +Covers osprey-ask-ai AC4.1-AC4.5. A minimal Flask app registers the blueprint with a +fake ``service_factory`` (an AskService built from the Phase 2 port fakes) and a fake +``principal_provider``. These live under ui_api, whose conftest applies the autouse +Postgres fixture, so they run through the Docker harness (run-tests.sh). +""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Tuple + +from flask import Flask +from osprey.worker.lib.ask import AskConfig, AskLimits, AskService, Principal +from osprey.worker.lib.ask.contracts import Conversation +from osprey.worker.lib.ask.tests.fakes import ( + AllowAllModelPolicy, + BlockedLock, + InMemoryConversationStore, + InMemoryLock, + NoProviderPolicy, + RaisingProvider, + RejectingModelPolicy, + ScriptedProvider, +) +from osprey.worker.lib.llm import LLMResponse, Tool, ToolCall, ToolRegistry +from osprey.worker.ui_api.osprey.lib.ask.blueprint import create_ask_blueprint + +PRINCIPAL = Principal(id='u1', email='u1@example.com') + + +class RecordingRegistryFactory: + def __init__(self, registry: ToolRegistry = None) -> None: # type: ignore[assignment] + self.builds = 0 + self._reg = registry if registry is not None else ToolRegistry() + + def build(self, principal: Principal, snapshot: Any) -> ToolRegistry: + self.builds += 1 + return self._reg + + +class BadEvidence: + def normalize(self, tool_call: ToolCall, tool_result: Any) -> Dict[str, Any]: + return {'bad': object()} # not JSON-serializable -> forces a serialization error + + +class RaisingExitLock: + def __init__(self) -> None: + self.acquired = 0 + + @contextmanager + def _cm(self) -> Iterator[None]: + self.acquired += 1 + yield + raise RuntimeError('lock __exit__ blew up after the turn completed') + + def acquire(self, conversation_id: str, principal: Principal) -> Any: + return self._cm() + + +def _text(t: str) -> LLMResponse: + return LLMResponse(text=t) + + +def _tools(*calls: ToolCall) -> LLMResponse: + return LLMResponse(text='', tool_calls=list(calls)) + + +def _factory(provider: Any, *, store=None, lock=None, tools=None, policy=None, limits=None, evidence=None): + store = store or InMemoryConversationStore() + lock = lock or InMemoryLock() + tools = tools or RecordingRegistryFactory() + policy = policy or AllowAllModelPolicy(provider) + + def make() -> AskService: + return AskService( + AskConfig( + policy=policy, + conversations=store, + lock=lock, + tools=tools, + limits=limits or AskLimits(), + evidence=evidence, + ) + ) + + return make, store, lock, tools + + +def _app(service_factory: Any, principal_provider: Any = lambda: PRINCIPAL) -> Flask: + app = Flask(__name__) + app.register_blueprint( + create_ask_blueprint( + name='ask', service_factory=service_factory, principal_provider=principal_provider, url='/ask' + ) + ) + return app + + +def _post(app: Flask, payload: Any) -> Any: + return app.test_client().post('/ask', json=payload) + + +def _parse_sse(body: str) -> List[Tuple[str, Dict[str, Any]]]: + frames: List[Tuple[str, Dict[str, Any]]] = [] + for block in body.split('\n\n'): + if not block.strip(): + continue + name = '' + data = '' + for line in block.split('\n'): + if line.startswith('event: '): + name = line[len('event: ') :] + elif line.startswith('data: '): + data += line[len('data: ') :] + frames.append((name, json.loads(data))) + return frames + + +def _names(frames: List[Tuple[str, Dict[str, Any]]]) -> List[str]: + return [n for n, _ in frames] + + +# --- AC4.1 / AC4.2 --------------------------------------------------------- + + +def test_happy_path_content_type_framing_versioning(): + make, store, lock, factory = _factory(ScriptedProvider([_text('hello')])) + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 200 + assert resp.mimetype == 'text/event-stream' + frames = _parse_sse(resp.get_data(as_text=True)) + assert _names(frames) == ['conversation_started', 'assistant_message', 'done'] + for name, data in frames: + assert data['version'] == 1 + assert data['type'] == name + assert len(store.appended) == 1 + assert lock.acquired == 1 and lock.released == 1 # resources released through the route + + +# --- AC4.3 (appropriate errors, no stream, no tools) ----------------------- + + +def test_missing_message_400_no_tools(): + make, _, _, factory = _factory(ScriptedProvider([_text('x')])) + resp = _post(_app(make), {}) + assert resp.status_code == 400 + assert resp.get_json()['error']['code'] == 'invalid_request' + assert factory.builds == 0 + + +def test_wrong_type_message_400(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')])) + resp = _post(_app(make), {'message': 123}) + assert resp.status_code == 400 + assert resp.get_json()['error']['code'] == 'invalid_request' + + +def test_extra_field_400(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')])) + resp = _post(_app(make), {'message': 'hi', 'nope': 1}) + assert resp.status_code == 400 + + +def test_non_json_body_400(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')])) + resp = _app(make).test_client().post('/ask', data='not json', content_type='text/plain') + assert resp.status_code == 400 + + +def test_unavailable_provider_503(): + make, _, _, factory = _factory(ScriptedProvider([_text('x')]), policy=NoProviderPolicy()) + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 503 + assert resp.get_json()['error']['code'] == 'unavailable_provider' + assert factory.builds == 0 + + +def test_invalid_model_400(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')]), policy=RejectingModelPolicy()) + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 400 + assert resp.get_json()['error']['code'] == 'invalid_model' + + +def test_foreign_conversation_403(): + store = InMemoryConversationStore() + store.add(Conversation(id='c1', principal_id='someone-else')) + make, _, _, _ = _factory(ScriptedProvider([_text('x')]), store=store) + resp = _post(_app(make), {'message': 'hi', 'conversation_id': 'c1'}) + assert resp.status_code == 403 + assert resp.get_json()['error']['code'] == 'forbidden' + + +def test_lock_unavailable_409(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')]), lock=BlockedLock()) + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 409 + assert resp.get_json()['error']['code'] == 'lock_unavailable' + + +def test_throwing_service_factory_safe_500(): + def make() -> AskService: + raise RuntimeError('boom secret sk-XYZ') + + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 500 + assert resp.get_json()['error']['code'] == 'internal' + assert 'sk-XYZ' not in resp.get_data(as_text=True) + + +def test_throwing_principal_provider_safe_500(): + make, _, _, _ = _factory(ScriptedProvider([_text('x')])) + + def principal() -> Principal: + raise RuntimeError('identity secret') + + resp = _app(make, principal_provider=principal).test_client().post('/ask', json={'message': 'hi'}) + assert resp.status_code == 500 + assert resp.get_json()['error']['code'] == 'internal' + + +# --- AC4.4 (safe failures) + AC4.5 (exactly one terminal) ------------------ + + +def test_provider_failure_single_safe_error_frame(): + make, store, lock, _ = _factory(RaisingProvider()) # message carries 'sk-DEADBEEF' + resp = _post(_app(make), {'message': 'hi'}) + assert resp.status_code == 200 # the stream started + body = resp.get_data(as_text=True) + assert 'sk-DEADBEEF' not in body + frames = _parse_sse(body) + assert _names(frames) == ['conversation_started', 'error'] + assert frames[-1][1]['payload']['code'] == 'provider_error' + assert store.appended == [] + assert lock.acquired == 1 and lock.released == 1 + + +def test_serialization_failure_yields_terminal_error(): + reg = ToolRegistry() + reg.register(Tool(name='lookup', description='d', handler=lambda **k: 'ok')) + call = ToolCall(id='c', name='lookup', arguments={}) + make, _, lock, _ = _factory( + ScriptedProvider([_tools(call), _text('final')]), + tools=RecordingRegistryFactory(reg), + evidence=BadEvidence(), + ) + frames = _parse_sse(_post(_app(make), {'message': 'hi'}).get_data(as_text=True)) + names = _names(frames) + assert names[0] == 'conversation_started' + assert names[-1] == 'error' + assert frames[-1][1]['payload']['code'] == 'serialization_error' + assert names.count('error') == 1 + assert lock.acquired == 1 and lock.released == 1 + + +def test_terminal_guard_no_second_terminal_after_done(): + lock = RaisingExitLock() + make, store, _, _ = _factory(ScriptedProvider([_text('final')]), lock=lock) + frames = _parse_sse(_post(_app(make), {'message': 'hi'}).get_data(as_text=True)) + names = _names(frames) + assert names == ['conversation_started', 'assistant_message', 'done'] + assert names.count('done') == 1 + assert 'error' not in names + assert lock.acquired == 1 diff --git a/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_sse.py b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_sse.py new file mode 100644 index 00000000..3c944576 --- /dev/null +++ b/osprey_worker/src/osprey/worker/ui_api/osprey/lib/ask/tests/test_sse.py @@ -0,0 +1,29 @@ +"""SSE serialization: versioned framing and serialization-error behavior.""" + +import json + +import pytest +from osprey.worker.lib.ask import AskEvent +from osprey.worker.lib.ask.errors import SerializationError +from osprey.worker.ui_api.osprey.lib.ask.sse import format_sse + + +def test_frame_shape_and_versioned_payload(): + ev = AskEvent('assistant_message', {'text': 'hi'}, conversation_id='c1', turn_id='t1') + frame = format_sse(ev) + lines = frame.split('\n') + assert lines[0] == 'event: assistant_message' + assert lines[1].startswith('data: ') + assert frame.endswith('\n\n') + data = json.loads(lines[1][len('data: ') :]) + assert data['version'] == 1 + assert data['type'] == 'assistant_message' + assert data['conversation_id'] == 'c1' + assert data['turn_id'] == 't1' + assert data['payload'] == {'text': 'hi'} + + +def test_non_serializable_payload_raises_serialization_error(): + ev = AskEvent('query_result', {'obj': object()}) + with pytest.raises(SerializationError): + format_sse(ev)