[engine] carry typed-action-contract schemas on the Sources payload

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(

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,

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

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]):

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()