diff --git a/osprey_async_worker/src/osprey/async_worker/engine.py b/osprey_async_worker/src/osprey/async_worker/engine.py index 9bf7217d..c8286e6e 100644 --- a/osprey_async_worker/src/osprey/async_worker/engine.py +++ b/osprey_async_worker/src/osprey/async_worker/engine.py @@ -348,6 +348,7 @@ def _load_and_register_schemas(self) -> None: self._shadow_filter, self.get_known_action_names, self.register_specialized_graph, + schemas=self._execution_graph.validated_sources.sources.schemas(), ) async def execute( diff --git a/osprey_async_worker/src/osprey/async_worker/sinks/sink/rules_sink.py b/osprey_async_worker/src/osprey/async_worker/sinks/sink/rules_sink.py index 075161bd..9b29b9e6 100644 --- a/osprey_async_worker/src/osprey/async_worker/sinks/sink/rules_sink.py +++ b/osprey_async_worker/src/osprey/async_worker/sinks/sink/rules_sink.py @@ -19,7 +19,6 @@ from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink from osprey.async_worker.engine import AsyncOspreyEngine -from osprey.async_worker.executor import execute as async_execute from osprey.async_worker.sinks.sink.input_stream import AsyncBaseInputStream logger = logging.getLogger(__name__) @@ -92,8 +91,11 @@ async def classify_one( result: Optional[ExecutionResult] = None try: with metrics.timed('handled_message', tags=tags, use_ms=True): - result = await async_execute( - self._engine.execution_graph, + # Route through engine.execute() rather than async_execute() on the full graph + # directly, so typed-action-contract dispatch runs: an allowlisted action is served + # its specialized (pruned) graph and/or shadow-diffed. For non-allowlisted actions + # execute() serves the full execution graph — identical to the prior direct call. + result = await self._engine.execute( self._udf_helpers, action, max_concurrent=self._max_concurrent_udfs, diff --git a/osprey_async_worker/src/osprey/async_worker/tests/test_async_sinks.py b/osprey_async_worker/src/osprey/async_worker/tests/test_async_sinks.py index 5c6e7eab..5b1a0f7b 100644 --- a/osprey_async_worker/src/osprey/async_worker/tests/test_async_sinks.py +++ b/osprey_async_worker/src/osprey/async_worker/tests/test_async_sinks.py @@ -3,14 +3,16 @@ import asyncio from datetime import datetime from typing import List -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from osprey.engine.executor.execution_context import Action, ExecutionResult from osprey.async_worker.adaptor.interfaces import AsyncBaseOutputSink +from osprey.async_worker.sinks.sink import rules_sink as rules_sink_module from osprey.async_worker.sinks.sink.input_stream import AsyncStaticInputStream from osprey.async_worker.sinks.sink.output_sink import AsyncMultiOutputSink, AsyncStdoutOutputSink +from osprey.async_worker.sinks.sink.rules_sink import AsyncRulesRunner def _make_result(action_id: int = 1, action_name: str = 'test') -> ExecutionResult: @@ -189,3 +191,36 @@ async def test_multi_sink_stop(): async def test_stdout_sink_will_do_work(): sink = AsyncStdoutOutputSink() assert sink.will_do_work(_make_result()) is True + + +# --- AsyncRulesRunner dispatch wiring --- + + +@pytest.mark.asyncio +async def test_classify_one_routes_through_engine_execute_for_dispatch(): + """Regression: classify_one MUST execute via engine.execute() (which runs the + typed-action-contract resolve_dispatch + shadow), NOT via async_execute() against the + full execution_graph directly. The direct-graph path made the specializer a runtime + no-op on the asyncio worker (the sole prod path) — specialized graphs registered at + init but never served. This pins the sink to the dispatch-aware engine method.""" + served = _make_result(action_id=123, action_name='guild_invite_created') + engine = MagicMock() + engine.execute = AsyncMock(return_value=served) + # No per-action sample config -> ActionSampler returns _SAMPLE_NEVER (action not dropped). + engine.get_config_subkey.return_value.get_action_config.return_value = None + + output_sink = MagicMock() + output_sink.push = AsyncMock() + + runner = AsyncRulesRunner(engine, output_sink, MagicMock(), max_concurrent_udfs=1) + action = Action(action_id=123, action_name='guild_invite_created', data={}, timestamp=datetime.utcnow()) + + with patch.object(rules_sink_module, 'metrics', MagicMock()): + result = await runner.classify_one(action, tag='test') + + # The dispatch-aware engine method was used (the bug bypassed it for a direct full-graph run). + engine.execute.assert_awaited_once() + assert action in engine.execute.call_args.args, 'engine.execute must receive the action' + # The served result flows to the output sink and is returned. + output_sink.push.assert_awaited_once_with(served) + assert result is served diff --git a/osprey_worker/src/osprey/engine/ast/sources.py b/osprey_worker/src/osprey/engine/ast/sources.py index 289b0d95..d06872f4 100644 --- a/osprey_worker/src/osprey/engine/ast/sources.py +++ b/osprey_worker/src/osprey/engine/ast/sources.py @@ -13,24 +13,46 @@ SOURCE_ENTRY_POINT_PATH = 'main.sml' CONFIG_PATH = 'config.yaml' +# Typed-action-contract schema JSON rides on the same Sources payload as the rules, keyed +# under this prefix (e.g. 'schemas/guild_joined.json', 'schemas/types/user.json'). These keys +# never collide with the *.sml / config.yaml namespaces, so they round-trip through +# to_dict/from_dict (and thus etcd) without touching the deployer or provider. +SCHEMAS_PREFIX = 'schemas/' class Sources: """A collection of sources, and an arbitrary configuration which describes a set of imported rules and perhaps additional configuration that will be executed by the engine.""" - def __init__(self, sources: Dict[str, Source], config: Optional['SourcesConfig'] = None): + def __init__( + self, + sources: Dict[str, Source], + config: Optional['SourcesConfig'] = None, + schemas: Optional[Dict[str, str]] = None, + ): assert SOURCE_ENTRY_POINT_PATH in sources, ( "Sources requires a file with the `path` 'main.sml' to be present as the entry-point" ) self._sources = sources self._config = config if config is not None else SourcesConfig(Source(path=CONFIG_PATH, contents='')) + # Typed-action-contract schema JSON, keyed by repo-relative posix path under + # SCHEMAS_PREFIX (path -> raw JSON contents). Carried separately from `sources` so it + # never flows through the .sml-only `add_source` path. + self._schemas = dict(schemas or {}) self._hash: Optional[str] = None @property def config(self) -> 'SourcesConfig': return self._config + def schemas(self) -> Dict[str, str]: + """Returns a copy of the typed-action-contract schema map (path -> raw JSON).""" + return dict(self._schemas) + + def get_schema(self, path: str) -> Optional[str]: + """Returns the raw JSON contents of a schema by its path, if present.""" + return self._schemas.get(path) + def get_by_path(self, path: str) -> Optional[Source]: """Gets a source that belongs to a given path, if it exists.""" return self._sources.get(path) @@ -60,21 +82,31 @@ def to_dict(self) -> Dict[str, str]: if self._config.source.contents: sources.append(self._config.source) - return {source.path: source.contents for source in sources} + result = {source.path: source.contents for source in sources} + # Schema keys live under SCHEMAS_PREFIX and never collide with *.sml / config.yaml, so + # this is a no-op (byte-identical) when there are no schemas. + result.update(self._schemas) + return result @staticmethod def from_dict(sources_dict: Dict[str, str]) -> 'Sources': """Creates a Sources object from a dict of path -> contents.""" builder = SourcesBuilder() + schemas: Dict[str, str] = {} for path, contents in sources_dict.items(): + # Partition schema keys BEFORE the config/add_source dispatch — they must not flow + # through `add_source`, which asserts a `.sml` suffix. + if path.startswith(SCHEMAS_PREFIX): + schemas[path] = contents + continue source = Source(path=path, contents=contents) if source.path == CONFIG_PATH: builder.add_config(source) else: builder.add_source(source) - return builder.build() + return builder.build(schemas=schemas) @staticmethod def from_path(root: Path) -> 'Sources': @@ -94,7 +126,16 @@ def from_path(root: Path) -> 'Sources': ] builder.add_config(*config_sources) - return builder.build() + # Typed-action-contract schemas: keyed by repo-relative posix path under SCHEMAS_PREFIX + # (e.g. 'schemas/guild_joined.json', 'schemas/types/user.json'). Anchored at the repo + # root with glob (NOT rglob) so every key is rooted at 'schemas/' and round-trips + # through from_dict's prefix partition; rglob would also match a nested 'schemas/' dir, + # whose non-prefixed key would crash from_dict's .sml assert. + schemas = { + '/'.join(path.relative_to(root).parts): path.read_text() for path in root.glob('schemas/**/*.json') + } + + return builder.build(schemas=schemas) def hash(self) -> str: """Returns the hash of the sources - a good way to quickly identify what sources are being executed.""" @@ -113,6 +154,14 @@ def hash(self) -> str: hasher.update(b'|') hasher.update(source.contents.encode('utf-8')) + # Fold the schema map into the hash (sorted by path for determinism) so a + # schema-only edit invalidates the hash and the worker reloads to pick up the new + # specialized graphs. Empty schemas add nothing, preserving back-compat hashes. + for path in sorted(self._schemas): + hasher.update(path.encode('utf-8')) + hasher.update(b'|') + hasher.update(self._schemas[path].encode('utf-8')) + self._hash = hasher.hexdigest() return self._hash @@ -146,8 +195,8 @@ def add_config(self, *sources: Source) -> 'SourcesBuilder': self._config = SourcesConfig(*self._config_sources.values()) return self - def build(self) -> Sources: - return Sources(self._sources, config=self._config) + def build(self, schemas: Optional[Dict[str, str]] = None) -> Sources: + return Sources(self._sources, config=self._config, schemas=schemas) class SourcesConfig(Dict[str, Any]): diff --git a/osprey_worker/src/osprey/engine/ast/tests/__init__.py b/osprey_worker/src/osprey/engine/ast/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/osprey_worker/src/osprey/engine/ast/tests/test_sources.py b/osprey_worker/src/osprey/engine/ast/tests/test_sources.py new file mode 100644 index 00000000..f252d685 --- /dev/null +++ b/osprey_worker/src/osprey/engine/ast/tests/test_sources.py @@ -0,0 +1,142 @@ +"""Tests for carrying typed-action-contract schemas on the Sources payload. + +The schema JSON rides on the same etcd ``Sources`` payload as the rules: a +``schemas/...json`` key space that survives ``from_path`` / ``to_dict`` / +``from_dict`` and so reaches both engines with no deployer/publisher/provider +change. These tests pin that round-trip and the hash behavior. +""" + +import json + +from osprey.engine.ast.grammar import Source +from osprey.engine.ast.sources import SOURCE_ENTRY_POINT_PATH, Sources + +_MAIN_SML = '# main entry point\n' +_GUILD_SCHEMA = { + 'action': 'guild_joined', + 'version': 1, + 'provides': {'user': {'id': 'int'}}, + 'absent': ['target_user'], +} +_USER_TYPE = {'id': 'int', 'username': 'str'} + + +def _write_tree(root) -> None: + """Lay down main.sml + config.yaml + a couple schema files under ``root``.""" + (root / 'main.sml').write_text(_MAIN_SML) + (root / 'config.yaml').write_text('sample_rate: 100\n') + schemas_dir = root / 'schemas' + (schemas_dir / 'types').mkdir(parents=True) + (schemas_dir / 'guild_joined.json').write_text(json.dumps(_GUILD_SCHEMA)) + (schemas_dir / 'types' / 'user.json').write_text(json.dumps(_USER_TYPE)) + + +class TestSourcesSchemas: + def test_from_path_collects_schemas(self, tmp_path) -> None: + _write_tree(tmp_path) + sources = Sources.from_path(tmp_path) + + schemas = sources.schemas() + assert 'schemas/guild_joined.json' in schemas + assert 'schemas/types/user.json' in schemas + assert json.loads(schemas['schemas/guild_joined.json']) == _GUILD_SCHEMA + # Schema files must NOT leak into the .sml source collection. + assert 'schemas/guild_joined.json' not in sources.paths() + # get_schema accessor + assert sources.get_schema('schemas/types/user.json') == json.dumps(_USER_TYPE) + assert sources.get_schema('schemas/missing.json') is None + + def test_to_dict_from_dict_round_trips_schemas(self, tmp_path) -> None: + _write_tree(tmp_path) + sources = Sources.from_path(tmp_path) + + as_dict = sources.to_dict() + # Schema keys present alongside main.sml / config.yaml + assert 'schemas/guild_joined.json' in as_dict + assert 'schemas/types/user.json' in as_dict + assert SOURCE_ENTRY_POINT_PATH in as_dict + + rebuilt = Sources.from_dict(as_dict) + assert rebuilt.schemas() == sources.schemas() + # .sml sources preserved too + assert rebuilt.paths() == sources.paths() + + def test_schema_keys_coexist_with_main_sml_and_config(self, tmp_path) -> None: + _write_tree(tmp_path) + sources = Sources.from_path(tmp_path) + as_dict = sources.to_dict() + # main.sml + config.yaml + 2 schema keys + assert as_dict[SOURCE_ENTRY_POINT_PATH] == _MAIN_SML + assert 'config.yaml' in as_dict + assert sorted(k for k in as_dict if k.startswith('schemas/')) == [ + 'schemas/guild_joined.json', + 'schemas/types/user.json', + ] + + def test_back_compat_payload_without_schemas(self) -> None: + """A legacy payload of only {'main.sml': ...} builds with no schemas and no raise.""" + sources = Sources.from_dict({SOURCE_ENTRY_POINT_PATH: _MAIN_SML}) + assert sources.schemas() == {} + + def test_payload_with_schema_key_does_not_hit_sml_assert(self) -> None: + """A payload WITH a schemas/x.json key must build fine — schema keys never + flow through ``add_source`` (which asserts ``.sml``).""" + sources = Sources.from_dict( + { + SOURCE_ENTRY_POINT_PATH: _MAIN_SML, + 'schemas/guild_joined.json': json.dumps(_GUILD_SCHEMA), + } + ) + assert sources.schemas() == {'schemas/guild_joined.json': json.dumps(_GUILD_SCHEMA)} + + def test_to_dict_byte_identical_when_no_schemas(self) -> None: + """Back-compat: to_dict is byte-identical with vs without the schemas map when + ``_schemas`` is empty (schema keys only ever ADD to the dict via result.update). + + Use a payload with config.yaml so to_dict exercises its full body (a config-less + Sources hits an unrelated pre-existing path in SourcesConfig.source). + """ + base_payload = {SOURCE_ENTRY_POINT_PATH: _MAIN_SML, 'config.yaml': 'sample_rate: 100\n'} + no_schemas = Sources.from_dict(dict(base_payload)) + empty_schemas = Sources.from_dict(dict(base_payload)) + # Force an explicit empty schemas map on the second one to prove parity. + empty_schemas._schemas = {} + assert no_schemas.to_dict() == empty_schemas.to_dict() + assert no_schemas.to_dict()[SOURCE_ENTRY_POINT_PATH] == _MAIN_SML + assert all(not k.startswith('schemas/') for k in no_schemas.to_dict()) + + def test_hash_changes_when_schema_content_changes(self) -> None: + """A schema-only edit (identical .sml) must invalidate the hash so the worker + reloads and picks up new specialized graphs.""" + base = Sources.from_dict( + { + SOURCE_ENTRY_POINT_PATH: _MAIN_SML, + 'schemas/guild_joined.json': json.dumps(_GUILD_SCHEMA), + } + ) + changed_schema = dict(_GUILD_SCHEMA) + changed_schema['absent'] = ['target_user', 'captcha_response'] + changed = Sources.from_dict( + { + SOURCE_ENTRY_POINT_PATH: _MAIN_SML, + 'schemas/guild_joined.json': json.dumps(changed_schema), + } + ) + assert base.hash() != changed.hash() + + def test_hash_stable_for_identical_schemas(self) -> None: + payload = { + SOURCE_ENTRY_POINT_PATH: _MAIN_SML, + 'schemas/guild_joined.json': json.dumps(_GUILD_SCHEMA), + } + assert Sources.from_dict(payload).hash() == Sources.from_dict(dict(payload)).hash() + + def test_hash_unchanged_when_no_schemas_vs_baseline(self) -> None: + """Back-compat: adding the schemas map but leaving it empty must not change the + hash of a schema-less Sources.""" + no_schemas = Sources({SOURCE_ENTRY_POINT_PATH: Source(path=SOURCE_ENTRY_POINT_PATH, contents=_MAIN_SML)}) + explicit_empty = Sources( + {SOURCE_ENTRY_POINT_PATH: Source(path=SOURCE_ENTRY_POINT_PATH, contents=_MAIN_SML)}, + schemas={}, + ) + assert no_schemas.hash() == explicit_empty.hash() diff --git a/osprey_worker/src/osprey/engine/executor/tests/test_typed_contract_dispatch.py b/osprey_worker/src/osprey/engine/executor/tests/test_typed_contract_dispatch.py new file mode 100644 index 00000000..e950be7f --- /dev/null +++ b/osprey_worker/src/osprey/engine/executor/tests/test_typed_contract_dispatch.py @@ -0,0 +1,122 @@ +"""Tests for the schema-source seam in load_and_register_specialized_graphs. + +These exercise WHERE schemas come from (in-memory ``schemas=`` map vs the disk +``resolve_schemas_dir``), not the specialization itself — ``specialize_graph`` and +the loaders are patched so the routing decision is isolated. +""" + +from __future__ import annotations + +from typing import List, Tuple + +import osprey.engine.executor.typed_contract_dispatch as dispatch +from osprey.engine.executor.typed_contract_dispatch import load_and_register_specialized_graphs + +_ACTIONS = ['guild_joined', 'message_sent'] + + +def _run(monkeypatch, prune, shadow, schemas, *, resolve_should_be_called: bool): + """Drive load_and_register_specialized_graphs with the loaders patched. + + Returns the list of (action_name, source) tuples for which a schema was 'loaded'. + """ + loaded_from: List[Tuple[str, str]] = [] + registered: List[str] = [] + + def fake_specialize(full_graph, schema): + # schema is the sentinel we returned below; pass it straight through. + return schema + + def fake_from_sources(action_name, schemas_map): + loaded_from.append((action_name, 'sources')) + return object() # truthy sentinel -> a "specialized graph" + + def fake_from_disk(action_name, schemas_dir): + loaded_from.append((action_name, 'disk')) + return object() + + resolve_calls = {'n': 0} + + def fake_resolve(): + resolve_calls['n'] += 1 + from pathlib import Path + + return Path('/fake/schemas') + + monkeypatch.setattr(dispatch, 'specialize_graph', fake_specialize) + monkeypatch.setattr(dispatch, 'load_schema_for_action_from_sources', fake_from_sources) + monkeypatch.setattr(dispatch, 'load_schema_for_action', fake_from_disk) + monkeypatch.setattr(dispatch, 'resolve_schemas_dir', fake_resolve) + + count = load_and_register_specialized_graphs( + full_graph=object(), + prune_filter=prune, + shadow_filter=shadow, + get_action_names=lambda: _ACTIONS, + register=lambda name, graph: registered.append(name), + schemas=schemas, + ) + + if resolve_should_be_called: + assert resolve_calls['n'] >= 1, 'expected disk resolution to be used' + else: + assert resolve_calls['n'] == 0, 'resolve_schemas_dir must NOT be called when schemas= provided' + + return count, loaded_from, registered + + +class TestSchemaSourceSeam: + def test_loads_from_sources_map_when_provided(self, monkeypatch) -> None: + schemas = {'schemas/guild_joined.json': '{}'} + count, loaded_from, registered = _run( + monkeypatch, + prune=frozenset({'*'}), + shadow=frozenset(), + schemas=schemas, + resolve_should_be_called=False, + ) + assert count == len(_ACTIONS) + assert all(src == 'sources' for _, src in loaded_from) + assert set(registered) == set(_ACTIONS) + + def test_falls_back_to_disk_when_schemas_none(self, monkeypatch) -> None: + count, loaded_from, registered = _run( + monkeypatch, + prune=frozenset({'*'}), + shadow=frozenset(), + schemas=None, + resolve_should_be_called=True, + ) + assert count == len(_ACTIONS) + assert all(src == 'disk' for _, src in loaded_from) + + def test_empty_schemas_map_falls_back_to_disk(self, monkeypatch) -> None: + # An empty (falsy) schemas map means "no etcd schemas" -> use disk. + count, loaded_from, _ = _run( + monkeypatch, + prune=frozenset({'guild_joined'}), + shadow=frozenset(), + schemas={}, + resolve_should_be_called=True, + ) + assert count == 1 + assert loaded_from == [('guild_joined', 'disk')] + + def test_returns_zero_when_both_filters_empty(self, monkeypatch) -> None: + # Neither gate set: returns 0 and reads NOTHING (no resolve, no action names). + resolve_called = {'n': 0} + monkeypatch.setattr(dispatch, 'resolve_schemas_dir', lambda: resolve_called.__setitem__('n', 1)) + + def _boom(): + raise AssertionError('get_action_names must not be called when both filters empty') + + count = load_and_register_specialized_graphs( + full_graph=object(), + prune_filter=frozenset(), + shadow_filter=frozenset(), + get_action_names=_boom, + register=lambda name, graph: None, + schemas={'schemas/guild_joined.json': '{}'}, + ) + assert count == 0 + assert resolve_called['n'] == 0 diff --git a/osprey_worker/src/osprey/engine/executor/typed_contract_dispatch.py b/osprey_worker/src/osprey/engine/executor/typed_contract_dispatch.py index efc9d6e3..ece99b5a 100644 --- a/osprey_worker/src/osprey/engine/executor/typed_contract_dispatch.py +++ b/osprey_worker/src/osprey/engine/executor/typed_contract_dispatch.py @@ -25,6 +25,7 @@ SchemaLoadError, filter_includes, load_schema_for_action, + load_schema_for_action_from_sources, resolve_schemas_dir, ) from osprey.worker.lib.instruments import metrics @@ -38,28 +39,39 @@ def load_and_register_specialized_graphs( shadow_filter: FrozenSet[str], get_action_names: Callable[[], Iterable[str]], register: Callable[[str, ExecutionGraph], None], + schemas: Optional[Mapping[str, str]] = None, ) -> int: """Load schemas for allowlisted actions, specialize them against ``full_graph``, and register each via ``register(action_name, specialized_graph)``. Returns the count registered. - No-op (returns 0) when neither gate is set or no schemas dir resolves — so shipping - schema files on the rules path cannot change behavior until an action is explicitly - listed in ``OSPREY_TYPED_CONTRACT_PRUNING`` / ``_SHADOW``. ``get_action_names`` is - called lazily (only past those gate checks) so the disabled-by-default path does no work. + Schemas come from one of two sources: the in-memory ``schemas`` map carried on the etcd + Sources payload (when non-empty), else the on-disk schemas directory resolved via + ``resolve_schemas_dir``. The Sources path lets the specializer activate on the + etcd-sourced prod worker, which has no schemas directory on disk. + + No-op (returns 0) when neither gate is set, or when no schemas are provided AND no + schemas dir resolves — so shipping schema files cannot change behavior until an action is + explicitly listed in ``OSPREY_TYPED_CONTRACT_PRUNING`` / ``_SHADOW``. ``get_action_names`` + is called lazily (only past those gate checks) so the disabled-by-default path does no work. """ register_filter = prune_filter | shadow_filter if not register_filter: return 0 - schemas_dir = resolve_schemas_dir() - if schemas_dir is None: + use_sources = bool(schemas) + schemas_dir = None if use_sources else resolve_schemas_dir() + if not use_sources and schemas_dir is None: return 0 loaded = 0 for action_name in get_action_names(): if not filter_includes(register_filter, action_name): continue try: - schema = load_schema_for_action(action_name, schemas_dir) + if schemas: # in-memory etcd Sources map (the truthiness narrows Optional for mypy) + schema = load_schema_for_action_from_sources(action_name, schemas) + else: + assert schemas_dir is not None # guaranteed by the gate above + schema = load_schema_for_action(action_name, schemas_dir) except SchemaLoadError as e: log.warning("Failed to load schema for %s: %s", action_name, e) continue @@ -68,8 +80,9 @@ def load_and_register_specialized_graphs( register(action_name, specialize_graph(full_graph, schema)) loaded += 1 if loaded: + source_desc = "Sources" if use_sources else schemas_dir log.info("Loaded %d specialized graphs from %s (prune=%r shadow=%r)", - loaded, schemas_dir, sorted(prune_filter), sorted(shadow_filter)) + loaded, source_desc, sorted(prune_filter), sorted(shadow_filter)) return loaded diff --git a/osprey_worker/src/osprey/engine/schema/schema_loader.py b/osprey_worker/src/osprey/engine/schema/schema_loader.py index 49e19dc8..7e2cedb7 100644 --- a/osprey_worker/src/osprey/engine/schema/schema_loader.py +++ b/osprey_worker/src/osprey/engine/schema/schema_loader.py @@ -8,7 +8,7 @@ import os from dataclasses import dataclass from pathlib import Path -from typing import Dict, FrozenSet, List, Optional, Set +from typing import Callable, Dict, FrozenSet, List, Mapping, Optional, Set log = logging.getLogger(__name__) @@ -62,40 +62,33 @@ class SchemaLoadError(Exception): """Raised when a schema file cannot be parsed or is invalid.""" -def load_schema(schema_path: Path, schemas_dir: Optional[Path] = None) -> ActionSchema: - """Load and parse a single action schema JSON file. +def parse_schema(raw: dict, ref_reader: Callable[[str], dict], where: str) -> ActionSchema: + """Parse an already-decoded schema dict into an ActionSchema. - Args: - schema_path: Path to the .json schema file. - schemas_dir: Base directory for resolving $ref: paths. Defaults to - the parent directory of schema_path. + This is the source-agnostic core shared by the disk loader (``load_schema``) and the + in-memory loader (``load_schema_for_action_from_sources``). The two differ only in HOW + they read raw JSON and resolve ``$ref:`` references. - Returns: - Parsed ActionSchema. + Args: + raw: The decoded top-level schema JSON. + ref_reader: Resolves a ``$ref:`` string (e.g. ``"$ref:types/user.json"``) to the + decoded contents of the referenced type. Implementations must enforce any + path-traversal guards and raise SchemaLoadError on escape / not-found. + where: A human-readable location label (file path or sources key) for error messages. Raises: - SchemaLoadError: if the file is missing, malformed, or violates constraints. + SchemaLoadError: if the schema is malformed or violates constraints. """ - if schemas_dir is None: - schemas_dir = schema_path.parent - - try: - raw = json.loads(schema_path.read_text()) - except FileNotFoundError: - raise SchemaLoadError(f"Schema file not found: {schema_path}") - except json.JSONDecodeError as e: - raise SchemaLoadError(f"Invalid JSON in {schema_path}: {e}") - version = raw.get("version") if version != _SUPPORTED_VERSION: raise SchemaLoadError( - f"Unsupported schema version in {schema_path}: {version!r}. " + f"Unsupported schema version in {where}: {version!r}. " f"Expected {_SUPPORTED_VERSION!r}." ) action = raw.get("action", "") if not action: - raise SchemaLoadError(f"Missing 'action' field in {schema_path}") + raise SchemaLoadError(f"Missing 'action' field in {where}") raw_provides: Dict[str, object] = raw.get("provides", {}) absent_list: List[str] = raw.get("absent", []) @@ -109,34 +102,14 @@ def load_schema(schema_path: Path, schemas_dir: Optional[Path] = None) -> Action overlap = provides_groups & absent_groups if overlap: raise SchemaLoadError( - f"Schema {schema_path}: groups {overlap!r} appear in both 'provides' and 'absent'." + f"Schema {where}: groups {overlap!r} appear in both 'provides' and 'absent'." ) - # Resolve $ref: references — load referenced type files and merge into provides - # $ref values point to types/.json relative to schemas_dir - _schemas_dir_resolved = schemas_dir.resolve() - - def _resolve_ref_path(ref_str: str) -> Path: - """Resolve a $ref: path and assert it stays within schemas_dir.""" - ref_rel = ref_str[len("$ref:"):] - ref_path = (schemas_dir / ref_rel).resolve() - if not ref_path.is_relative_to(_schemas_dir_resolved): - raise SchemaLoadError( - f"$ref path escapes schemas directory: {ref_rel!r} resolves to {ref_path}" - ) - return ref_path - + # Resolve $ref: references — load referenced type definitions and merge into provides. resolved_provides: Dict[str, object] = {} for group, value in raw_provides.items(): if isinstance(value, str) and value.startswith("$ref:"): - ref_path = _resolve_ref_path(value) - try: - ref_data = json.loads(ref_path.read_text()) - except FileNotFoundError: - raise SchemaLoadError(f"Referenced type file not found: {ref_path} (from {schema_path})") - except json.JSONDecodeError as e: - raise SchemaLoadError(f"Invalid JSON in referenced file {ref_path}: {e}") - resolved_provides[group] = ref_data + resolved_provides[group] = ref_reader(value) else: resolved_provides[group] = value @@ -144,14 +117,11 @@ def _resolve_ref_path(ref_str: str) -> Path: for group, ref_str in types_used.items(): if isinstance(ref_str, str) and ref_str.startswith("$ref:") and group not in resolved_provides: try: - ref_path = _resolve_ref_path(ref_str) - ref_data = json.loads(ref_path.read_text()) - resolved_provides[group] = ref_data + resolved_provides[group] = ref_reader(ref_str) provides_groups.add(group) except SchemaLoadError: - raise - except (FileNotFoundError, json.JSONDecodeError): - # types_used refs are informational; skip if the file doesn't exist yet + # types_used refs are informational; skip if the type isn't available yet (a + # missing/escaping ref here must not break loading the schema). pass # Flatten provides to dot-notation field types: "user.id" -> "int" @@ -180,6 +150,51 @@ def _resolve_ref_path(ref_str: str) -> Path: ) +def load_schema(schema_path: Path, schemas_dir: Optional[Path] = None) -> ActionSchema: + """Load and parse a single action schema JSON file from disk. + + Args: + schema_path: Path to the .json schema file. + schemas_dir: Base directory for resolving $ref: paths. Defaults to + the parent directory of schema_path. + + Returns: + Parsed ActionSchema. + + Raises: + SchemaLoadError: if the file is missing, malformed, or violates constraints. + """ + if schemas_dir is None: + schemas_dir = schema_path.parent + + try: + raw = json.loads(schema_path.read_text()) + except FileNotFoundError: + raise SchemaLoadError(f"Schema file not found: {schema_path}") + except json.JSONDecodeError as e: + raise SchemaLoadError(f"Invalid JSON in {schema_path}: {e}") + + # $ref values point to types/.json relative to schemas_dir. + _schemas_dir_resolved = schemas_dir.resolve() + + def _disk_ref_reader(ref_str: str) -> dict: + """Resolve a $ref: path on disk, asserting it stays within schemas_dir.""" + ref_rel = ref_str[len("$ref:"):] + ref_path = (schemas_dir / ref_rel).resolve() + if not ref_path.is_relative_to(_schemas_dir_resolved): + raise SchemaLoadError( + f"$ref path escapes schemas directory: {ref_rel!r} resolves to {ref_path}" + ) + try: + return json.loads(ref_path.read_text()) + except FileNotFoundError: + raise SchemaLoadError(f"Referenced type file not found: {ref_path} (from {schema_path})") + except json.JSONDecodeError as e: + raise SchemaLoadError(f"Invalid JSON in referenced file {ref_path}: {e}") + + return parse_schema(raw, _disk_ref_reader, where=str(schema_path)) + + def resolve_schemas_dir() -> Optional[Path]: """Discover the schemas directory from the runtime environment. @@ -224,3 +239,48 @@ def load_schema_for_action(action_name: str, schemas_dir: Path) -> Optional[Acti except SchemaLoadError: log.exception("Failed to load schema for action %r from %s", action_name, schema_path) return None + + +def load_schema_for_action_from_sources( + action_name: str, schemas: Mapping[str, str] +) -> Optional[ActionSchema]: + """Load the schema for an action from the in-memory ``Sources`` schema map. + + ``schemas`` maps repo-relative posix paths (``schemas/.json``, + ``schemas/types/.json``) to raw JSON contents — the same map that rides on the + etcd Sources payload. This is the etcd-sourced counterpart to + :func:`load_schema_for_action`: it lets #55's specializer activate on the prod worker, + which has no schemas directory on disk. + + Returns None if no schema exists for ``action_name``. + + Raises: + SchemaLoadError: if the schema (or a referenced type) is malformed, missing, or a + ``$ref:`` attempts to escape the ``schemas/`` key space. + """ + key = f"schemas/{action_name}.json" + raw_text = schemas.get(key) + if raw_text is None: + return None + + def _sources_ref_reader(ref_str: str) -> dict: + """Resolve a $ref: against the schemas map, rejecting absolute / traversing paths.""" + ref_rel = ref_str[len("$ref:"):] + # Reject absolute paths or any parent-traversal segment — refs must stay within the + # `schemas/` key space, mirroring the disk loader's path-traversal guard. + if ref_rel.startswith("/") or any(part == ".." for part in ref_rel.split("/")): + raise SchemaLoadError(f"$ref path escapes schemas directory: {ref_rel!r}") + ref_key = f"schemas/{ref_rel}" + ref_text = schemas.get(ref_key) + if ref_text is None: + raise SchemaLoadError(f"Referenced type not found: {ref_key} (from {key})") + try: + return json.loads(ref_text) + except json.JSONDecodeError as e: + raise SchemaLoadError(f"Invalid JSON in referenced source {ref_key}: {e}") + + try: + raw = json.loads(raw_text) + except json.JSONDecodeError as e: + raise SchemaLoadError(f"Invalid JSON in {key}: {e}") + return parse_schema(raw, _sources_ref_reader, where=key) diff --git a/osprey_worker/src/osprey/engine/schema/tests/test_schema_loader.py b/osprey_worker/src/osprey/engine/schema/tests/test_schema_loader.py index d2f0af92..c7749eb6 100644 --- a/osprey_worker/src/osprey/engine/schema/tests/test_schema_loader.py +++ b/osprey_worker/src/osprey/engine/schema/tests/test_schema_loader.py @@ -10,6 +10,7 @@ SchemaLoadError, load_schema, load_schema_for_action, + load_schema_for_action_from_sources, resolve_schemas_dir, ) @@ -203,6 +204,82 @@ def test_schema_with_no_provides_and_no_absent(self, tmp_path: Path) -> None: assert len(schema.provides_field_types) == 0 +class TestLoadSchemaFromSources: + """``load_schema_for_action_from_sources`` reads schemas from the in-memory + ``Sources._schemas`` map (etcd payload) rather than disk — this is what activates typed + contracts on the etcd-sourced prod worker. + """ + + def test_finds_by_name(self) -> None: + schemas = {"schemas/guild_joined.json": json.dumps(_VALID_SCHEMA)} + schema = load_schema_for_action_from_sources("guild_joined", schemas) + assert schema is not None + assert schema.action == "guild_joined" + assert "user" in schema.provides_groups + + def test_returns_none_if_missing(self) -> None: + schemas = {"schemas/guild_joined.json": json.dumps(_VALID_SCHEMA)} + assert load_schema_for_action_from_sources("other_action", schemas) is None + assert load_schema_for_action_from_sources("guild_joined", {}) is None + + def test_resolves_ref_from_map(self) -> None: + schema_data = dict(_VALID_SCHEMA) + schema_data["provides"] = {"user": "$ref:types/user.json", "guild": {"id": "int"}} + schemas = { + "schemas/guild_joined.json": json.dumps(schema_data), + "schemas/types/user.json": json.dumps({"id": "int", "username": "str"}), + } + schema = load_schema_for_action_from_sources("guild_joined", schemas) + assert schema is not None + assert schema.provides_field_types["user.id"] == "int" + assert schema.provides_field_types["user.username"] == "str" + assert schema.provides_field_types["guild.id"] == "int" + + def test_missing_ref_raises(self) -> None: + schema_data = dict(_VALID_SCHEMA) + schema_data["provides"] = {"user": "$ref:types/nonexistent.json"} + schemas = {"schemas/guild_joined.json": json.dumps(schema_data)} + with pytest.raises(SchemaLoadError, match="not found"): + load_schema_for_action_from_sources("guild_joined", schemas) + + def test_relative_traversal_ref_raises(self) -> None: + schema_data = dict(_VALID_SCHEMA) + schema_data["provides"] = {"user": "$ref:../../etc/passwd"} + schemas = {"schemas/guild_joined.json": json.dumps(schema_data)} + with pytest.raises(SchemaLoadError, match="escapes schemas directory"): + load_schema_for_action_from_sources("guild_joined", schemas) + + def test_absolute_ref_raises(self) -> None: + schema_data = dict(_VALID_SCHEMA) + schema_data["provides"] = {"user": "$ref:/etc/hostname"} + schemas = {"schemas/guild_joined.json": json.dumps(schema_data)} + with pytest.raises(SchemaLoadError, match="escapes schemas directory"): + load_schema_for_action_from_sources("guild_joined", schemas) + + def test_parity_disk_vs_from_sources(self, tmp_path: Path) -> None: + """The same JSON loaded via disk ``load_schema`` and via the from-sources path must + yield equal ActionSchema objects (incl. a resolved $ref).""" + schema_data = dict(_VALID_SCHEMA) + schema_data["provides"] = {"user": "$ref:types/user.json", "guild": {"id": "int"}} + type_data = {"id": "int", "username": "str"} + + # Disk + types_dir = tmp_path / "types" + types_dir.mkdir() + (types_dir / "user.json").write_text(json.dumps(type_data)) + disk_path = _write_schema(tmp_path, schema_data) + disk_schema = load_schema(disk_path, schemas_dir=tmp_path) + + # From sources + schemas = { + "schemas/guild_joined.json": json.dumps(schema_data), + "schemas/types/user.json": json.dumps(type_data), + } + sources_schema = load_schema_for_action_from_sources("guild_joined", schemas) + + assert sources_schema == disk_schema + + class TestResolveSchemasDir: """resolve_schemas_dir picks up schemas without an explicit env var when OSPREY_RULES_PATH already points at a smite-rules checkout that has a diff --git a/osprey_worker/src/osprey/worker/lib/osprey_engine.py b/osprey_worker/src/osprey/worker/lib/osprey_engine.py index ccfd79cf..06489479 100644 --- a/osprey_worker/src/osprey/worker/lib/osprey_engine.py +++ b/osprey_worker/src/osprey/worker/lib/osprey_engine.py @@ -161,6 +161,7 @@ def _load_and_register_schemas(self) -> None: self._shadow_filter, self.get_known_action_names, self.register_specialized_graph, + schemas=self._execution_graph.validated_sources.sources.schemas(), ) def _handle_updated_sources(self) -> None: