diff --git a/pyproject.toml b/pyproject.toml index 7169eb7..09f9f4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,7 @@ dev = [ "pytest-asyncio>=0.23", "build>=1.2", "respx>=0.21", - "ruff>=0.6", + "ruff>=0.6,<0.16", ] [tool.hatch.build.targets.wheel] diff --git a/src/openlinker/runtime/__init__.py b/src/openlinker/runtime/__init__.py index e153618..3e32b5a 100644 --- a/src/openlinker/runtime/__init__.py +++ b/src/openlinker/runtime/__init__.py @@ -6,6 +6,7 @@ RUNTIME_REQUIRED_FEATURES, RuntimeAssignment, RuntimeAttemptIdentity, + RuntimeAuthority, RuntimeCallOptions, RuntimeDrainTimeoutError, RuntimeEvent, @@ -33,6 +34,7 @@ "MemoryRuntimeStore", "RuntimeAssignment", "RuntimeAttemptIdentity", + "RuntimeAuthority", "RuntimeCallOptions", "RuntimeContext", "RuntimeDrainTimeoutError", diff --git a/src/openlinker/runtime/types.py b/src/openlinker/runtime/types.py index 8f889fe..0d3b15b 100644 --- a/src/openlinker/runtime/types.py +++ b/src/openlinker/runtime/types.py @@ -114,6 +114,14 @@ class RuntimeMTLS: server_name: str = "" +@dataclass(frozen=True) +class RuntimeAuthority: + principal_scope_id: str + runtime_session_id: str + runtime_session_epoch: int + runtime_attachment_id: str + + @dataclass(frozen=True) class RuntimeAttemptIdentity: run_id: str diff --git a/src/openlinker/runtime/worker.py b/src/openlinker/runtime/worker.py index ed15306..e274739 100644 --- a/src/openlinker/runtime/worker.py +++ b/src/openlinker/runtime/worker.py @@ -48,6 +48,7 @@ from .types import ( RUNTIME_MAX_CAPACITY, RuntimeAttemptIdentity, + RuntimeAuthority, RuntimeDrainTimeoutError, RuntimeEvent, RuntimeHandlerError, @@ -98,6 +99,10 @@ def __init__(self) -> None: super().__init__("RuntimeWorker stopped before its durable drain completed") +class _RuntimeAssignmentAuthorityError(ValueError): + pass + + class RuntimeHandler(Protocol): async def handle(self, context: RuntimeContext) -> RuntimeResult | dict[str, Any] | Any: ... @@ -128,7 +133,12 @@ def __init__(self, worker: RuntimeWorker, active: _ActiveAttempt) -> None: self.run_id = attempt.run_id self.agent_id = attempt.agent_id self.input = dict(active.assignment.input) - self.metadata = dict(active.assignment.metadata) + self.metadata, self.authority = _runtime_authority_from_metadata( + active.assignment.metadata, + active.assignment.identity.session_epoch, + attempt.runtime_session_id, + worker._ready, + ) @property def cancelled(self) -> bool: @@ -924,8 +934,9 @@ async def _start_confirmed_attempt( async def _execute_attempt(self, active: _ActiveAttempt) -> None: started = time.monotonic() - context = RuntimeContext(self, active) + context: RuntimeContext | None = None try: + context = RuntimeContext(self, active) raw = await _invoke_handler(self.handler, context) context._close() result = _normalize_result(raw) @@ -937,12 +948,18 @@ async def _execute_attempt(self, active: _ActiveAttempt) -> None: result = RuntimeResult.failed( "HANDLER_CANCELLED", "handler stopped without a Runtime cancellation" ) + except _RuntimeAssignmentAuthorityError: + result = RuntimeResult.failed( + "ASSIGNMENT_AUTHORITY_INVALID", + "assignment Runtime authority is invalid", + ) except Exception as exc: result = RuntimeResult.failed( "HANDLER_ERROR", _bounded(str(exc), 500, "handler failed") ) finally: - context._close() + if context is not None: + context._close() if active.cancel_event.is_set() or self._force_cancel.is_set(): return duration_ms = result.duration_ms or max(0, int((time.monotonic() - started) * 1000)) @@ -2624,6 +2641,59 @@ def _canonical_uuid(value: str, label: str) -> None: raise ValueError(f"{label} must be a lowercase non-zero UUID") +_RUNTIME_AUTHORITY_METADATA_KEY = "_openlinker_runtime_authority" + + +def _runtime_authority_from_metadata( + raw_metadata: dict[str, Any], + session_epoch: int, + runtime_session_id: str, + ready: RuntimeReady | None, +) -> tuple[dict[str, Any], RuntimeAuthority | None]: + metadata = dict(raw_metadata) + missing = object() + raw = metadata.pop(_RUNTIME_AUTHORITY_METADATA_KEY, missing) + if raw is missing: + return metadata, None + if ( + ready is None + or not isinstance(raw, dict) + or set(raw) != {"principal_scope_id", "source"} + or raw.get("source") != "core" + or not isinstance(raw.get("principal_scope_id"), str) + or session_epoch < 1 + ): + raise _RuntimeAssignmentAuthorityError( + "assignment Runtime authority is invalid" + ) + principal_scope_id = raw["principal_scope_id"] + try: + _canonical_runtime_principal_scope(principal_scope_id) + _canonical_uuid(runtime_session_id, "runtime_session_id") + _canonical_uuid(ready.attachment_id, "runtime_attachment_id") + except ValueError as exc: + raise _RuntimeAssignmentAuthorityError( + "assignment Runtime authority is invalid" + ) from exc + return metadata, RuntimeAuthority( + principal_scope_id=principal_scope_id, + runtime_session_id=runtime_session_id, + runtime_session_epoch=session_epoch, + runtime_attachment_id=ready.attachment_id, + ) + + +def _canonical_runtime_principal_scope(value: str) -> None: + allowed = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789._:-" + if ( + not value + or len(value) > 256 + or value.strip() != value + or any(character not in allowed for character in value) + ): + raise ValueError("principal_scope_id must be an opaque Runtime identifier") + + def _token_scoped_runtime_node_id(agent_token: str) -> str: digest = hashlib.sha256( b"openlinker/runtime-worker/token-scoped-node/v1\x00" diff --git a/tests/test_runtime.py b/tests/test_runtime.py index 1aa8691..195a06f 100644 --- a/tests/test_runtime.py +++ b/tests/test_runtime.py @@ -446,10 +446,25 @@ async def test_assignment_is_durable_and_confirmed_before_handler_runs(): store = runtime.MemoryRuntimeStore() transport = FakeTransport() transport.assignment = assignment(store) + transport.assignment.metadata.update( + { + "_openlinker_runtime_authority": { + "principal_scope_id": "ps1_" + ("A" * 43), + "source": "core", + } + } + ) transport.ack_release.clear() handler_started = asyncio.Event() async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: + assert context.metadata == {"source": "test"} + assert context.authority == runtime.RuntimeAuthority( + principal_scope_id="ps1_" + ("A" * 43), + runtime_session_id=store.identity.runtime_session_id, + runtime_session_epoch=store.identity.session_epoch, + runtime_attachment_id=ATTACHMENT_ID, + ) handler_started.set() await context.emit("run.progress", {"step": 1}) return {"answer": "ok"} @@ -469,6 +484,42 @@ async def handler(context: runtime.RuntimeContext) -> dict[str, Any]: assert transport.session_closed +@pytest.mark.asyncio +async def test_malformed_runtime_authority_is_rejected_before_handler(): + store = runtime.MemoryRuntimeStore() + transport = FakeTransport() + transport.assignment = assignment(store) + transport.assignment.metadata.update( + { + "_openlinker_runtime_authority": { + "principal_scope_id": "not/an/opaque-id", + "source": "core", + } + } + ) + handler_calls = 0 + + async def handler(_context: runtime.RuntimeContext) -> dict[str, Any]: + nonlocal handler_calls + handler_calls += 1 + return {} + + worker = make_worker(store, transport, handler) + running = asyncio.create_task(worker.run()) + try: + await asyncio.wait_for(transport.result_acked.wait(), timeout=1) + assert handler_calls == 0 + assert transport.result_attempts[0]["status"] == "failed" + assert transport.result_attempts[0]["error"] == { + "error_code": "ASSIGNMENT_AUTHORITY_INVALID", + "message": "assignment Runtime authority is invalid", + "retryable_hint": False, + } + finally: + await worker.stop() + await running + + @pytest.mark.asyncio async def test_lost_acks_replay_the_same_assignment_event_and_result_ids(): store = runtime.MemoryRuntimeStore()