diff --git a/python/restate/__init__.py b/python/restate/__init__.py index 16935e3..3b27a44 100644 --- a/python/restate/__init__.py +++ b/python/restate/__init__.py @@ -38,6 +38,7 @@ SendHandle, RunOptions, ) +from .entry_codec import JournalValueCodec, JournalValueCodecProvider from .exceptions import TerminalError, SdkInternalBaseException, is_internal_exception from .asyncio import as_completed, gather, wait_completed, select @@ -58,6 +59,7 @@ def create_test_harness( restate_image: str = "docker.io/restatedev/restate:latest", always_replay: bool = False, disable_retries: bool = False, + journal_value_codec: typing.Optional[JournalValueCodec] = None, ) -> typing.AsyncGenerator[HarnessEnvironment, None]: """a dummy harness constructor to raise ImportError. Install restate-sdk[harness] to use this feature""" raise ImportError("Install restate-sdk[harness] to use this feature") @@ -81,7 +83,9 @@ def test_harness( @asynccontextmanager async def create_client( - ingress: str, headers: typing.Optional[dict] = None + ingress: str, + headers: typing.Optional[dict] = None, + journal_value_codec: typing.Optional[JournalValueCodec] = None, ) -> typing.AsyncGenerator[RestateClient, None]: """a dummy client constructor to raise ImportError. Install restate-sdk[client] to use this feature""" raise ImportError("Install restate-sdk[client] to use this feature") @@ -105,6 +109,8 @@ async def create_client( "ScopedContext", "RunOptions", "TerminalError", + "JournalValueCodec", + "JournalValueCodecProvider", "app", "create_test_harness", "test_harness", diff --git a/python/restate/client.py b/python/restate/client.py index dc39213..35ed238 100644 --- a/python/restate/client.py +++ b/python/restate/client.py @@ -20,6 +20,7 @@ from .client_types import RestateClient, RestateClientSendHandle, RestateScopedClient, HttpError from .context import HandlerType +from .entry_codec import JournalValueCodec from .serde import BytesSerde, JsonSerde, Serde from .handler import handler_from_callable @@ -32,9 +33,15 @@ class Client(RestateClient): A basic client for connecting to the Restate service. """ - def __init__(self, client: httpx.AsyncClient, headers: typing.Optional[dict] = None): + def __init__( + self, + client: httpx.AsyncClient, + headers: typing.Optional[dict] = None, + journal_codec: typing.Optional[JournalValueCodec] = None, + ): self.headers = headers or {} self.client = client + self.journal_codec = journal_codec def scope(self, scope: str) -> RestateScopedClient: return ScopedClient(self, scope) @@ -103,6 +110,8 @@ async def do_raw_call( ) -> O: """Make an RPC call to the given handler""" parameter = input_serde.serialize(input_param) + if self.journal_codec is not None: + parameter = self.journal_codec.encode(parameter) if headers is not None: headers_kvs = list(headers.items()) else: @@ -124,6 +133,9 @@ async def do_raw_call( scope=scope, limit_key=limit_key, ) + # A send returns the invocation-id envelope, not a codec'd payload, so skip decode for it. + if not send and self.journal_codec is not None: + res = await self.journal_codec.decode(res) return output_serde.deserialize(res) # type: ignore async def post( @@ -471,10 +483,19 @@ async def workflow_send( @asynccontextmanager async def create_client( - ingress: str, headers: typing.Optional[dict] = None + ingress: str, + headers: typing.Optional[dict] = None, + journal_value_codec: typing.Optional[JournalValueCodec] = None, ) -> typing.AsyncGenerator[RestateClient, None]: """ Create a new Restate client. + + Args: + ingress: The base URL of the Restate ingress. + headers: Optional default headers to send with every request. + journal_value_codec: Optional journal value codec. When set, request bodies are encoded and + call responses decoded through it. It must match the codec configured on the endpoint. + NOTE: This is experimental and may change in future releases. """ async with httpx.AsyncClient(base_url=ingress, headers=headers, http2=True) as http_client: - yield Client(http_client, headers) + yield Client(http_client, headers, journal_codec=journal_value_codec) diff --git a/python/restate/endpoint.py b/python/restate/endpoint.py index f8c3a83..0c6fc32 100644 --- a/python/restate/endpoint.py +++ b/python/restate/endpoint.py @@ -17,6 +17,7 @@ from restate.service import Service from restate.object import VirtualObject from restate.workflow import Workflow +from restate.entry_codec import JournalValueCodec, JournalValueCodecProvider # disable too few methods in a class @@ -31,6 +32,7 @@ class Endpoint: services: typing.Dict[str, typing.Union[Service, VirtualObject, Workflow]] protocol: typing.Optional[typing.Literal["bidi", "request_response"]] identity_keys: typing.List[str] + journal_value_codec: typing.Optional[typing.Union[JournalValueCodec, JournalValueCodecProvider]] def __init__(self): """ @@ -44,6 +46,10 @@ def __init__(self): self.identity_keys = [] + # An optional journal value codec (or async provider building one). When set, the SDK + # transforms serialized journal values through it. None means no codec is configured. + self.journal_value_codec = None + def bind(self, *services: typing.Union[Service, VirtualObject, Workflow]): """ Bind a service to the endpoint @@ -79,6 +85,23 @@ def identity_key(self, identity_key: str): """Add an identity key to this endpoint.""" self.identity_keys.append(identity_key) + def set_journal_value_codec(self, codec_or_provider: typing.Union[JournalValueCodec, JournalValueCodecProvider]): + """ + Set the journal value codec for this endpoint. + + NOTE: This is experimental and may change in future releases. + + Args: + codec_or_provider: Either a :class:`JournalValueCodec` instance, or an async provider + (a zero-arg callable returning an awaitable of a codec) that is invoked once and + cached for the lifetime of the endpoint. + + Returns: + The updated Endpoint instance. + """ + self.journal_value_codec = codec_or_provider + return self + def app(self): """ Returns the ASGI application for this endpoint. @@ -101,6 +124,7 @@ def app( services: typing.Iterable[typing.Union[Service, VirtualObject, Workflow]], protocol: typing.Optional[typing.Literal["bidi", "request_response"]] = None, identity_keys: typing.Optional[typing.List[str]] = None, + journal_value_codec: typing.Optional[typing.Union[JournalValueCodec, JournalValueCodecProvider]] = None, ): """A restate ASGI application that hosts the given services.""" endpoint = Endpoint() @@ -113,4 +137,6 @@ def app( if identity_keys: for key in identity_keys: endpoint.identity_key(key) + if journal_value_codec is not None: + endpoint.set_journal_value_codec(journal_value_codec) return endpoint.app() diff --git a/python/restate/entry_codec.py b/python/restate/entry_codec.py new file mode 100644 index 0000000..95bfa0e --- /dev/null +++ b/python/restate/entry_codec.py @@ -0,0 +1,71 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""This module contains the journal value codec interface.""" + +import abc +import typing + +# disable too few public methods +# pylint: disable=R0903 + + +class JournalValueCodec(abc.ABC): + """ + Journal values codec. + + This allows to transform journal values after being serialized, before writing them to the + wire, and vice-versa. It sits *between* the ``Serde`` layer and the journal/wire. + The canonical use case is to encrypt or compress everything the SDK persists to the journal. + + Values that are passed through the codec: + + * Handlers input and success output + * ``ctx.run`` success results + * Awakeables/Promise success results + * State values + * Call/send request parameters and call responses + + Failures are never passed through the codec. + + NOTE: This is preview and may change in future releases. + """ + + @abc.abstractmethod + def encode(self, buf: bytes) -> bytes: + """ + Encodes the given buffer. This will be applied *after* serialization. + + Args: + buf: The buffer to encode. Empty byte buffers should be appropriately handled as well. + + Returns: + The encoded buffer. + """ + + @abc.abstractmethod + async def decode(self, buf: bytes) -> bytes: + """ + Decodes the given buffer. This will be applied *before* deserialization. + + Args: + buf: The buffer to decode. + + Returns: + The decoded buffer. + """ + + +JournalValueCodecProvider = typing.Callable[[], typing.Awaitable[JournalValueCodec]] +"""A provider that asynchronously builds a :class:`JournalValueCodec`. + +It is invoked once, and the resulting codec is reused for the lifetime of the endpoint. +This is useful to perform async setup at startup, e.g. loading an encryption key. +""" diff --git a/python/restate/handler.py b/python/restate/handler.py index b3371b2..91a55a7 100644 --- a/python/restate/handler.py +++ b/python/restate/handler.py @@ -34,6 +34,7 @@ from restate.retry_policy import InvocationRetryPolicy from restate.context import HandlerType +from restate.entry_codec import JournalValueCodec from restate.exceptions import TerminalError from restate.serde import DefaultSerde, PydanticJsonSerde, MsgspecJsonSerde, Serde, is_pydantic, Msgspec from restate.types import extract_core_type @@ -350,11 +351,18 @@ def handler_from_callable(wrapper: HandlerType[I, O]) -> Handler[I, O]: raise ValueError("Handler not found") # pylint: disable=raise-missing-from -async def invoke_handler(handler: Handler[I, O], ctx: Any, in_buffer: bytes) -> bytes: +async def invoke_handler( + handler: Handler[I, O], ctx: Any, in_buffer: bytes, journal_codec: Optional[JournalValueCodec] = None +) -> bytes: """ Invoke the handler with the given context and input. """ if handler.arity == 2: + if journal_codec is not None: + try: + in_buffer = await journal_codec.decode(in_buffer) + except Exception as e: + raise TerminalError(message="Failed to decode input using journal value codec", status_code=400) from e try: in_arg = handler.handler_io.input_serde.deserialize(in_buffer) except Exception as e: @@ -363,4 +371,6 @@ async def invoke_handler(handler: Handler[I, O], ctx: Any, in_buffer: bytes) -> else: out_arg = await handler.fn(ctx) # type: ignore [call-arg] out_buffer = handler.handler_io.output_serde.serialize(out_arg) + if journal_codec is not None: + out_buffer = journal_codec.encode(bytes(out_buffer)) return bytes(out_buffer) diff --git a/python/restate/harness.py b/python/restate/harness.py index 5861124..7a2c739 100644 --- a/python/restate/harness.py +++ b/python/restate/harness.py @@ -21,6 +21,7 @@ from hypercorn.config import Config from hypercorn.asyncio import serve from restate.client import create_client +from restate.entry_codec import JournalValueCodec from restate.server_types import RestateAppT from restate.types import HarnessEnvironment from testcontainers.core.container import DockerContainer # type: ignore @@ -338,6 +339,7 @@ async def create_test_harness( restate_image: str = "docker.io/restatedev/restate:latest", always_replay: bool = False, disable_retries: bool = False, + journal_value_codec: typing.Optional[JournalValueCodec] = None, ) -> typing.AsyncGenerator[HarnessEnvironment, None]: """ Creates a test harness for running Restate services together with restate-server. @@ -360,6 +362,8 @@ async def create_test_harness( on a suspension point. This is useful to hunt non-deterministic bugs that might prevent your code to replay correctly (default is False). :param disable_retries: When True, retries are disabled (default is False). + :param journal_value_codec: Optional journal value codec to configure on the ingress client, so + it matches a codec configured on the endpoint under test (default is None). """ with ( create_restate_container( @@ -377,7 +381,7 @@ async def create_test_harness( msg = f"unable to register the services at {bind_address} - {res.status_code} {res.text}" raise AssertionError(msg) - async with create_client(runtime.ingress_url()) as client: + async with create_client(runtime.ingress_url(), journal_value_codec=journal_value_codec) as client: yield HarnessEnvironment( ingress_url=runtime.ingress_url(), admin_api_url=runtime.admin_url(), client=client ) diff --git a/python/restate/server.py b/python/restate/server.py index 195ea5f..c5e233f 100644 --- a/python/restate/server.py +++ b/python/restate/server.py @@ -11,12 +11,14 @@ """This module contains the ASGI server for the restate framework.""" import asyncio +import functools import logging import signal -from typing import Dict, Set, TypedDict, Literal +from typing import Dict, Optional, Set, TypedDict, Literal from restate.discovery import compute_discovery_json from restate.endpoint import Endpoint +from restate.entry_codec import JournalValueCodec from restate.server_context import ServerInvocationContext, DisconnectedException from restate.server_types import Receive, ReceiveChannel, RestateAppT, Scope, Send, binary_to_header, header_to_binary # pylint: disable=line-too-long from restate.vm import VMWrapper @@ -125,7 +127,12 @@ async def send_health_check(send: Send): async def process_invocation_to_completion( - vm: VMWrapper, handler, attempt_headers: Dict[str, str], receive: ReceiveChannel, send: Send + vm: VMWrapper, + handler, + attempt_headers: Dict[str, str], + receive: ReceiveChannel, + send: Send, + journal_codec: Optional[JournalValueCodec] = None, ): """Invoke the user code.""" status, res_headers = vm.get_response_head() @@ -155,7 +162,13 @@ async def process_invocation_to_completion( # ======================================== invocation = vm.sys_input() context = ServerInvocationContext( - vm=vm, handler=handler, invocation=invocation, attempt_headers=attempt_headers, send=send, receive=receive + vm=vm, + handler=handler, + invocation=invocation, + attempt_headers=attempt_headers, + send=send, + receive=receive, + journal_codec=journal_codec, ) try: await context.enter() @@ -217,6 +230,27 @@ def asgi_app(endpoint: Endpoint) -> RestateAppT: active_channels: Set[ReceiveChannel] = set() sigterm_installed = False + # Journal value codec resolution. + # A codec (or an async provider building one) may be configured on the endpoint. Whether one is + # configured is known synchronously and drives the VM's payload-checks flag. An async provider + # is resolved once, lazily on the first invocation, and the resulting codec is reused. + codec_config = endpoint.journal_value_codec + disable_payload_checks = codec_config is not None + + # Memoize the provider's resolution. functools.cache can't wrap the async provider directly + # (it would cache a coroutine, which can't be awaited twice), so we cache a sync factory that + # returns the Task instead: the provider runs once and every caller awaits the same Task. + @functools.cache + def _codec_task() -> "asyncio.Task[JournalValueCodec]": + assert callable(codec_config) + return asyncio.ensure_future(codec_config()) + + async def get_codec() -> Optional[JournalValueCodec]: + if codec_config is None or isinstance(codec_config, JournalValueCodec): + # Nothing configured, or an already-built instance: nothing to await. + return codec_config + return await _codec_task() + def _on_sigterm() -> None: """Notify all active receive channels of graceful shutdown.""" for ch in active_channels: @@ -284,8 +318,14 @@ async def app(scope: Scope, receive: Receive, send: Send): receive_channel = ReceiveChannel(receive) active_channels.add(receive_channel) try: + journal_codec = await get_codec() await process_invocation_to_completion( - VMWrapper(request_headers), handler, dict(request_headers), receive_channel, send + VMWrapper(request_headers, disable_payload_checks=disable_payload_checks), + handler, + dict(request_headers), + receive_channel, + send, + journal_codec=journal_codec, ) finally: active_channels.discard(receive_channel) diff --git a/python/restate/server_context.py b/python/restate/server_context.py index 4b1b60b..80fbec3 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -54,6 +54,7 @@ SuspendedException, RetryableError, ) +from restate.entry_codec import JournalValueCodec from restate.handler import Handler, handler_from_callable, invoke_handler from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_types import ReceiveChannel, Send @@ -371,7 +372,7 @@ def value(self) -> RestateDurableFuture[Any]: def resolve(self, value: Any) -> Awaitable[None]: vm: VMWrapper = self.server_context.vm assert self.serde is not None - value_buffer = self.serde.serialize(value) + value_buffer = self.server_context._encode(self.serde.serialize(value)) # pylint: disable=protected-access handle = vm.sys_complete_promise_success(self.name, value_buffer) update_restate_context_is_replaying(self.server_context.vm) @@ -518,6 +519,7 @@ def __init__( attempt_headers: Dict[str, str], send: Send, receive: ReceiveChannel, + journal_codec: Optional[JournalValueCodec] = None, ) -> None: super().__init__() self.vm = vm @@ -527,11 +529,18 @@ def __init__( self.send = send self.random_instance = Random(invocation.random_seed) self.receive = receive + self.journal_codec = journal_codec self.run_coros_to_execute: dict[int, Callable[[], Awaitable[None]]] = {} self.request_finished_event = asyncio.Event() self.tasks = Tasks() self.extension_data: Dict[str, Any] = {} + def _encode(self, buffer: bytes) -> bytes: + """Apply the journal value codec encode step, if a codec is configured.""" + if self.journal_codec is None: + return buffer + return self.journal_codec.encode(buffer) + async def enter(self): """Invoke the user code.""" update_restate_context_is_replaying(self.vm) @@ -543,7 +552,9 @@ async def enter(self): await stack.enter_async_context(manager()) await stack.enter_async_context(auto_close_extension_data(self.extension_data)) - out_buffer = await invoke_handler(handler=self.handler, ctx=self, in_buffer=in_buffer) + out_buffer = await invoke_handler( + handler=self.handler, ctx=self, in_buffer=in_buffer, journal_codec=self.journal_codec + ) restate_context_is_replaying.set(False) self.vm.sys_write_output_success(bytes(out_buffer)) self.vm.sys_end() @@ -721,6 +732,8 @@ async def fetch_result(): if res is None or serde is None: return res if isinstance(res, bytes): + if self.journal_codec is not None: + res = await self.journal_codec.decode(res) return serde.deserialize(res) return res @@ -770,7 +783,7 @@ def set(self, name: str, value: T, serde: Serde[T] = DefaultSerde()) -> None: """Set the value associated with the given name.""" if isinstance(serde, DefaultSerde): serde = serde.with_maybe_type(type(value)) - buffer = serde.serialize(value) + buffer = self._encode(serde.serialize(value)) self.vm.sys_set_state(name, bytes(buffer)) update_restate_context_is_replaying(self.vm) @@ -810,7 +823,7 @@ def resolve_signal(self, invocation_id: str, name: str, value: I, serde: Serde[I """Resolve a named signal on a target invocation.""" if isinstance(serde, DefaultSerde): serde = serde.with_maybe_type(type(value)) - buf = serde.serialize(value) + buf = self._encode(serde.serialize(value)) self.vm.sys_resolve_signal(invocation_id, name, buf) update_restate_context_is_replaying(self.vm) @@ -854,7 +867,7 @@ async def create_run_coroutine( self.tasks.add(action_result_future) action_result = typing.cast(T, await action_result_future) - buffer = serde.serialize(action_result) + buffer = self._encode(serde.serialize(action_result)) self.vm.propose_run_completion_success(handle, buffer) except TerminalError as t: failure = Failure(code=t.status_code, message=t.message, metadata=t.metadata) @@ -1026,7 +1039,7 @@ def do_raw_call( limit_key: str | None = None, ) -> RestateDurableCallFuture[O] | SendHandle: """Make an RPC call to the given handler""" - parameter = input_serde.serialize(input_param) + parameter = self._encode(input_serde.serialize(input_param)) if headers is not None: headers_kvs = list(headers.items()) else: @@ -1226,7 +1239,7 @@ def awakeable( def resolve_awakeable(self, name: str, value: I, serde: Serde[I] = DefaultSerde()) -> None: if isinstance(serde, DefaultSerde): serde = serde.with_maybe_type(type(value)) - buf = serde.serialize(value) + buf = self._encode(serde.serialize(value)) self.vm.sys_resolve_awakeable(name, buf) update_restate_context_is_replaying(self.vm) diff --git a/python/restate/vm.py b/python/restate/vm.py index f1e91af..7b7b91f 100644 --- a/python/restate/vm.py +++ b/python/restate/vm.py @@ -205,8 +205,8 @@ class VMWrapper: It provides a type-friendly interface to our shared vm. """ - def __init__(self, headers: typing.List[typing.Tuple[str, str]]): - self.vm = PyVM(headers) + def __init__(self, headers: typing.List[typing.Tuple[str, str]], disable_payload_checks: bool = False): + self.vm = PyVM(headers, disable_payload_checks) def get_response_head(self) -> typing.Tuple[int, typing.Iterable[typing.Tuple[str, str]]]: """ diff --git a/src/lib.rs b/src/lib.rs index c128bda..319281a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,9 +4,9 @@ use pyo3::types::{PyBytes, PyNone, PyString}; use restate_sdk_shared_core::fmt::{set_error_formatter, ErrorFormatter}; use restate_sdk_shared_core::{ AwaitResponse, AwakeableHandle, CallHandle, CoreVM, Error, Header, IdentityVerifier, Input, - NonEmptyValue, NotificationHandle, OnMaxAttempts, ResponseHead, RetryPolicy, RunExitResult, - RunHandle, Target, TerminalFailure, UnresolvedFuture, VMOptions, Value, - CANCEL_NOTIFICATION_HANDLE, VM, + NonDeterministicChecksOption, NonEmptyValue, NotificationHandle, OnMaxAttempts, ResponseHead, + RetryPolicy, RunExitResult, RunHandle, Target, TerminalFailure, UnresolvedFuture, VMOptions, + Value, CANCEL_NOTIFICATION_HANDLE, VM, }; use std::fmt; use std::time::{Duration, SystemTime}; @@ -371,9 +371,24 @@ struct PyVM { #[pymethods] impl PyVM { #[new] - fn new(headers: Vec<(String, String)>) -> Result { + #[pyo3(signature = (headers, disable_payload_checks=false))] + fn new( + headers: Vec<(String, String)>, + disable_payload_checks: bool, + ) -> Result { + // When a journal value codec is configured, the serialized payloads written to the journal + // may legitimately differ between attempts (e.g. encryption with a random nonce), so we + // must disable the VM's payload determinism checks. + let options = VMOptions { + non_determinism_checks: if disable_payload_checks { + NonDeterministicChecksOption::PayloadChecksDisabled + } else { + NonDeterministicChecksOption::Enabled + }, + ..Default::default() + }; Ok(Self { - vm: CoreVM::new(headers, VMOptions::default())?, + vm: CoreVM::new(headers, options)?, }) } diff --git a/tests/entry_codec.py b/tests/entry_codec.py new file mode 100644 index 0000000..5c25d3b --- /dev/null +++ b/tests/entry_codec.py @@ -0,0 +1,235 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# + +"""Tests for the JournalValueCodec feature.""" + +import typing +from contextlib import asynccontextmanager + +import pytest + +import restate +from restate import Context, JournalValueCodec, Service, TerminalError +from restate.client import Client +from restate.handler import handler_from_callable, invoke_handler +from restate.serde import JsonSerde + +# ----- Asyncio fixtures + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [ + pytest.mark.anyio, +] + + +# ----- Codec used across the tests + +_MAGIC = b"\xca" + + +class MagicCodec(JournalValueCodec): + """A tiny symmetric codec: prepends a magic byte on encode, strips+validates it on decode. + + Non-trivial enough that, if the codec were NOT applied on both sides, decoding would fail — + which lets us prove the codec actually runs. Also exercises the empty-buffer contract: + ``encode(b"")`` yields the (non-empty) magic byte, and ``decode`` strips it back to ``b""``. + """ + + def encode(self, buf: bytes) -> bytes: + return _MAGIC + buf + + async def decode(self, buf: bytes) -> bytes: + if not buf.startswith(_MAGIC): + raise ValueError("missing magic prefix") + return buf[len(_MAGIC) :] + + +# ============================================================ +# Fast, docker-free unit tests of the encode/decode seams +# ============================================================ + + +async def test_codec_roundtrip_and_empty_buffer(): + codec = MagicCodec() + assert await codec.decode(codec.encode(b"hello")) == b"hello" + # empty buffer must be handled gracefully + assert await codec.decode(codec.encode(b"")) == b"" + + +async def test_invoke_handler_decodes_input_and_encodes_output(): + svc = Service("greeter") + + @svc.handler() + async def greet(ctx: Context, name: str) -> str: # pylint: disable=unused-argument + return f"hi {name}" + + handler = handler_from_callable(greet) + codec = MagicCodec() + + # The buffer handed to invoke_handler is what the VM stores: the codec-encoded input. + encoded_input = codec.encode(handler.handler_io.input_serde.serialize("bob")) + + out = await invoke_handler(handler=handler, ctx=None, in_buffer=encoded_input, journal_codec=codec) + + # Output must be codec-encoded; decoding it back must yield the serialized "hi bob". + assert out.startswith(_MAGIC) + decoded_out = await codec.decode(out) + assert handler.handler_io.output_serde.deserialize(decoded_out) == "hi bob" + + +async def test_invoke_handler_bad_input_raises_terminal_400(): + svc = Service("greeter") + + @svc.handler() + async def greet(ctx: Context, name: str) -> str: # pylint: disable=unused-argument + return f"hi {name}" + + handler = handler_from_callable(greet) + codec = MagicCodec() + + # Raw (unencoded) input has no magic prefix -> decode must fail as a terminal 400. + raw_input = handler.handler_io.input_serde.serialize("bob") + with pytest.raises(TerminalError) as exc: + await invoke_handler(handler=handler, ctx=None, in_buffer=raw_input, journal_codec=codec) + assert exc.value.status_code == 400 + + +class _RecordingClient(Client): + """A Client that records the last request content and returns a canned response body.""" + + def __init__(self, journal_codec, canned_response: bytes): + super().__init__(client=None, journal_codec=journal_codec) # type: ignore[arg-type] + self.last_content: typing.Optional[bytes] = None + self.canned_response = canned_response + + async def post(self, /, service, handler, send, content, **kwargs): # type: ignore[override] + self.last_content = content + return self.canned_response + + +async def test_client_encodes_request_and_decodes_response(): + codec = MagicCodec() + # Server would have stored an encoded success value; simulate that as the response body. + response_body = codec.encode(JsonSerde[str]().serialize("pong")) + client = _RecordingClient(codec, response_body) + + result: str = await client.do_raw_call( + service="s", + handler="h", + input_param="ping", + input_serde=JsonSerde[str](), + output_serde=JsonSerde[str](), + ) + + # Request body was encoded by the codec... + assert client.last_content is not None and client.last_content.startswith(_MAGIC) + assert await codec.decode(client.last_content) == JsonSerde[str]().serialize("ping") + # ...and the response was decoded before deserialization. + assert result == "pong" + + +async def test_client_send_skips_response_decode(): + codec = MagicCodec() + # A send returns the invocation-id envelope (plain JSON), which must NOT be codec-decoded. + envelope = JsonSerde[dict]().serialize({"invocationId": "inv_123"}) + client = _RecordingClient(codec, envelope) + + result: dict = await client.do_raw_call( + service="s", + handler="h", + input_param="ping", + input_serde=JsonSerde[str](), + output_serde=JsonSerde[dict](), + send=True, + ) + + # Request is still encoded, but the response envelope is returned verbatim (no decode attempted). + assert client.last_content is not None and client.last_content.startswith(_MAGIC) + assert result == {"invocationId": "inv_123"} + + +# ============================================================ +# End-to-end test against a real restate-server (needs docker) +# ============================================================ + + +@asynccontextmanager +async def codec_harness( + service: typing.Union[Service, restate.VirtualObject, restate.Workflow], codec: JournalValueCodec +) -> typing.AsyncIterator[restate.RestateClient]: + """Spin up a harness where BOTH the endpoint and the ingress client share the same codec.""" + async with restate.create_test_harness( + restate.app([service], journal_value_codec=codec), + journal_value_codec=codec, + restate_image="ghcr.io/restatedev/restate:latest", + ) as harness: + yield harness.client + + +async def test_codec_end_to_end(): + codec = MagicCodec() + obj = restate.VirtualObject("codec_obj") + + @obj.handler() + async def exercise(ctx: restate.ObjectContext, name: str) -> str: + # state set + get round-trips through the codec + ctx.set("who", name) + stored = await ctx.get("who", type_hint=str) + assert stored == name + + # ctx.run success result round-trips through the codec + ran = await ctx.run_typed("compute", lambda: name.upper()) + assert ran == name.upper() + + # awakeable resolve + await round-trips through the codec + awk_id, awk = ctx.awakeable(type_hint=str) + ctx.resolve_awakeable(awk_id, "signal-value") + assert await awk == "signal-value" + + return f"hi {name}" + + async with codec_harness(obj, codec) as client: + # handler input + output round-trip through the codec on both client and server + result = await client.object_call(exercise, key="k1", arg="bob") + assert result == "hi bob" + + +async def test_codec_end_to_end_with_async_provider(): + codec = MagicCodec() + provider_calls = 0 + + async def provider() -> JournalValueCodec: + nonlocal provider_calls + provider_calls += 1 + return codec + + svc = Service("codec_svc") + + @svc.handler() + async def greet(ctx: Context, name: str) -> str: # pylint: disable=unused-argument + return f"hi {name}" + + # The endpoint is configured with an ASYNC PROVIDER; the ingress client gets the built instance. + async with restate.create_test_harness( + restate.app([svc], journal_value_codec=provider), + journal_value_codec=codec, + restate_image="ghcr.io/restatedev/restate:latest", + ) as harness: + assert await harness.client.service_call(greet, arg="bob") == "hi bob" + assert await harness.client.service_call(greet, arg="alice") == "hi alice" + + # The provider must have been resolved exactly once across multiple invocations. + assert provider_calls == 1