Support eager state whitelist

README.md

@@ -64,11 +64,14 @@ The compatibility with Restate is described in the following table. sdk-python v
 | 1.3 - 1.4                 | ✅ <sup>(1)</sup> | ✅ <sup>(1)(2)</sup> |
 | 1.5 - 1.6                 | ✅                | ✅ <sup>(2)</sup>    |
 | 1.7                       | ✅                | ✅                   |
+| 1.8                       | ✅                | ✅ <sup>(3)</sup>    |
 
 <sup>(1)</sup> **Note** The new Service/Object/Workflow constructor field and the decorator field `invocation_retry_policy` works only from Restate 1.4 onward. Check the in-code documentation for more details.
 
 <sup>(2)</sup> **Note** Scopes and limit keys work only from Restate 1.7 onward. Check the in-code documentation for more details.
 
+<sup>(3)</sup> **Note** The list form of `state_preload` (eager state keys whitelist) works only from Restate 1.8 onward. Check the in-code documentation for more details.
+
 ## Contributing
 
 We’re excited if you join the Restate community and start contributing!

python/restate/discovery.py

@@ -27,7 +27,7 @@
 import json
 import typing
 from enum import Enum
-from typing import Dict, Optional, Any, List, get_args, get_origin
+from typing import Dict, Optional, Any, List, Tuple, Union, get_args, get_origin
 
 
 from restate.endpoint import Endpoint as RestateEndpoint
@@ -86,6 +86,7 @@ def __init__(
         idempotencyRetention: Optional[int] = None,
         workflowCompletionRetention: Optional[int] = None,
         enableLazyState: Optional[bool] = None,
+        eagerStateKeysWhitelist: Optional[List[str]] = None,
         ingressPrivate: Optional[bool] = None,
         retryPolicyInitialInterval: Optional[int] = None,
         retryPolicyMaxInterval: Optional[int] = None,
@@ -105,6 +106,7 @@ def __init__(
         self.idempotencyRetention = idempotencyRetention
         self.workflowCompletionRetention = workflowCompletionRetention
         self.enableLazyState = enableLazyState
+        self.eagerStateKeysWhitelist = eagerStateKeysWhitelist
         self.ingressPrivate = ingressPrivate
         self.retryPolicyInitialInterval = retryPolicyInitialInterval
         self.retryPolicyMaxInterval = retryPolicyMaxInterval
@@ -127,6 +129,7 @@ def __init__(
         journalRetention: Optional[int] = None,
         idempotencyRetention: Optional[int] = None,
         enableLazyState: Optional[bool] = None,
+        eagerStateKeysWhitelist: Optional[List[str]] = None,
         ingressPrivate: Optional[bool] = None,
         retryPolicyInitialInterval: Optional[int] = None,
         retryPolicyMaxInterval: Optional[int] = None,
@@ -144,6 +147,7 @@ def __init__(
         self.journalRetention = journalRetention
         self.idempotencyRetention = idempotencyRetention
         self.enableLazyState = enableLazyState
+        self.eagerStateKeysWhitelist = eagerStateKeysWhitelist
         self.ingressPrivate = ingressPrivate
         self.retryPolicyInitialInterval = retryPolicyInitialInterval
         self.retryPolicyMaxInterval = retryPolicyMaxInterval
@@ -225,6 +229,27 @@ def json_schema_from_type_hint(type_hint: Optional[TypeHint[Any]]) -> Any:
     return type_hint_to_json_schema(type_hint.annotation)
 
 
+def resolve_state_preload(
+    state_preload: Optional[Union[bool, List[str]]], enable_lazy_state: Optional[bool]
+) -> Tuple[Optional[bool], Optional[List[str]]]:
+    """
+    Resolve the user facing state preload options into the manifest pair (enableLazyState, eagerStateKeysWhitelist).
+    """
+    if state_preload is None:
+        # Legacy passthrough (may itself be None).
+        return enable_lazy_state, None
+
+    if isinstance(state_preload, bool):
+        # True  => eager (preload all)  => not lazy
+        # False => lazy  (preload none) => lazy
+        return not state_preload, None
+
+    # state_preload is a whitelist of keys to preload eagerly (selective preloading).
+    # An empty whitelist is equivalent to fully lazy state.
+    # Copy, so later mutations of the user's list don't affect discovery.
+    return True, list(state_preload) if len(state_preload) > 0 else None
+
+
 # pylint: disable=R0912,R0915
 def compute_discovery_json(
     endpoint: RestateEndpoint, version: int, discovered_as: typing.Literal["bidi", "request_response"]
@@ -236,6 +261,15 @@ def compute_discovery_json(
     ep = compute_discovery(endpoint, discovered_as)
 
     # Validate that new discovery fields aren't used with older protocol versions
+    if version <= 4:
+        # Note: enableLazyState itself is available since v3; only the eager state whitelist requires v5.
+        for service in ep.services:
+            if service.eagerStateKeysWhitelist is not None:
+                raise ValueError("eagerStateKeysWhitelist is only supported in discovery protocol version 5")
+            for handler in service.handlers:
+                if handler.eagerStateKeysWhitelist is not None:
+                    raise ValueError("eagerStateKeysWhitelist is only supported in discovery protocol version 5")
+
     if version <= 3:
         # Strip lambdaCompression for older discovery versions
         ep.lambdaCompression = None
@@ -336,6 +370,9 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal["
                     contentType=handler.handler_io.content_type,
                     jsonSchema=json_schema_from_type_hint(handler.handler_io.output_type),
                 )
+            handler_enable_lazy_state, handler_eager_state_keys_whitelist = resolve_state_preload(
+                handler.state_preload, handler.enable_lazy_state
+            )
             # add the handler
             service_handlers.append(
                 Handler(
@@ -360,7 +397,8 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal["
                     workflowCompletionRetention=int(handler.workflow_retention.total_seconds() * 1000)
                     if handler.workflow_retention is not None
                     else None,
-                    enableLazyState=handler.enable_lazy_state,
+                    enableLazyState=handler_enable_lazy_state,
+                    eagerStateKeysWhitelist=handler_eager_state_keys_whitelist,
                     ingressPrivate=handler.ingress_private,
                     retryPolicyInitialInterval=int(
                         handler.invocation_retry_policy.initial_interval.total_seconds() * 1000
@@ -386,6 +424,13 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal["
                 )
             )
         # add the service
+        if isinstance(service, (Workflow, VirtualObject)):
+            service_enable_lazy_state, service_eager_state_keys_whitelist = resolve_state_preload(
+                service.state_preload, service.enable_lazy_state
+            )
+        else:
+            # Stateless services must not advertise state options.
+            service_enable_lazy_state, service_eager_state_keys_whitelist = None, None
         description = service.service_tag.description
         metadata = service.service_tag.metadata
         services.append(
@@ -407,7 +452,8 @@ def compute_discovery(endpoint: RestateEndpoint, discovered_as: typing.Literal["
                 idempotencyRetention=int(service.idempotency_retention.total_seconds() * 1000)
                 if service.idempotency_retention is not None
                 else None,
-                enableLazyState=service.enable_lazy_state if isinstance(service, (Workflow, VirtualObject)) else None,
+                enableLazyState=service_enable_lazy_state,
+                eagerStateKeysWhitelist=service_eager_state_keys_whitelist,
                 ingressPrivate=service.ingress_private,
                 retryPolicyInitialInterval=int(service.invocation_retry_policy.initial_interval.total_seconds() * 1000)
                 if service.invocation_retry_policy and service.invocation_retry_policy.initial_interval

python/restate/handler.py

@@ -29,6 +29,7 @@
     Literal,
     Optional,
     TypeVar,
+    Union,
 )
 
 from restate.retry_policy import InvocationRetryPolicy
@@ -216,6 +217,19 @@ def update_handler_io_with_type_hints(handler_io: HandlerIO[I, O], signature: Si
     update_handler_io_with_return_type_hints(handler_io, signature)
 
 
+def check_state_preload_options(
+    state_preload: Optional[Union[bool, List[str]]], enable_lazy_state: Optional[bool]
+) -> None:
+    """
+    Validate that the new `state_preload` option and the deprecated `enable_lazy_state` option are not both set.
+    """
+    if state_preload is not None and enable_lazy_state is not None:
+        raise ValueError(
+            "Cannot set both the 'state_preload' option and the deprecated 'enable_lazy_state' option at the same time. "
+            "Use only 'state_preload'."
+        )
+
+
 # pylint: disable=R0902
 @dataclass
 class Handler(Generic[I, O]):
@@ -236,7 +250,9 @@ class Handler(Generic[I, O]):
         journal_retention: Journal retention duration.
         idempotency_retention: Idempotency retention duration.
         workflow_retention: Workflow completion retention duration.
-        enable_lazy_state: If true, lazy state is enabled.
+        enable_lazy_state: Deprecated, use `state_preload` instead.
+        state_preload: Configures which state keys are preloaded ("eager state") at the start of each invocation.
+            True preloads all state, False preloads nothing, a list preloads only the listed keys.
         ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress.
         invocation_retry_policy: Optional retry policy configuration applied to this handler.
         explicit_recovery: If true, the handler runs with explicit recovery, see restate._preview.explicit_recovery.
@@ -256,6 +272,7 @@ class Handler(Generic[I, O]):
     idempotency_retention: Optional[timedelta] = None
     workflow_retention: Optional[timedelta] = None
     enable_lazy_state: Optional[bool] = None
+    state_preload: Optional[Union[bool, List[str]]] = None
     ingress_private: Optional[bool] = None
     invocation_retry_policy: Optional[InvocationRetryPolicy] = None
     context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None
@@ -280,6 +297,7 @@ def make_handler(
     idempotency_retention: Optional[timedelta] = None,
     workflow_retention: Optional[timedelta] = None,
     enable_lazy_state: Optional[bool] = None,
+    state_preload: Optional[Union[bool, List[str]]] = None,
     ingress_private: Optional[bool] = None,
     invocation_retry_policy: Optional[InvocationRetryPolicy] = None,
     context_managers: Optional[List[Callable[[], AsyncContextManager[None]]]] = None,
@@ -301,7 +319,9 @@ def make_handler(
         journal_retention: Journal retention duration.
         idempotency_retention: Idempotency retention duration.
         workflow_retention: Workflow completion retention duration.
-        enable_lazy_state: If true, lazy state is enabled.
+        enable_lazy_state: Deprecated, use `state_preload` instead.
+        state_preload: Configures which state keys are preloaded ("eager state") at the start of each invocation.
+            True preloads all state, False preloads nothing, a list preloads only the listed keys.
         ingress_private: If true, the handler cannot be invoked from the HTTP nor Kafka ingress.
         invocation_retry_policy: Retry policy used by Restate when invoking this handler.
             NOTE: only supported with restate-server >= 1.5.
@@ -316,6 +336,8 @@ def make_handler(
     if len(signature.parameters) == 0:
         raise ValueError("Handler must have at least one parameter")
 
+    check_state_preload_options(state_preload, enable_lazy_state)
+
     arity = len(signature.parameters)
     update_handler_io_with_type_hints(handler_io, signature)  # mutates handler_io
 
@@ -334,6 +356,7 @@ def make_handler(
         idempotency_retention=idempotency_retention,
         workflow_retention=workflow_retention,
         enable_lazy_state=enable_lazy_state,
+        state_preload=state_preload,
         ingress_private=ingress_private,
         invocation_retry_policy=invocation_retry_policy,
         context_managers=context_managers,

python/restate/object.py

@@ -21,7 +21,7 @@
 
 from restate.serde import Serde, DefaultSerde
 from restate.retry_policy import InvocationRetryPolicy
-from restate.handler import Handler, HandlerIO, ServiceTag, make_handler
+from restate.handler import Handler, HandlerIO, ServiceTag, check_state_preload_options, make_handler
 
 I = typing.TypeVar("I")
 O = typing.TypeVar("O")
@@ -63,9 +63,13 @@ class VirtualObject:
         idempotency_retention (timedelta, optional): The retention duration of idempotent requests to this service.
             NOTE: You can set this field only if you register this service against restate-server >= 1.4,
             otherwise the service discovery will fail.
-        enable_lazy_state (bool, optional): When set to True, lazy state will be enabled for all invocations to this service.
-            NOTE: You can set this field only if you register this service against restate-server >= 1.4,
-            otherwise the service discovery will fail.
+        enable_lazy_state (bool, optional): Deprecated, use `state_preload` instead.
+        state_preload (bool | List[str], optional): Configures which state keys are preloaded ("eager state")
+            at the start of each invocation to this virtual object. Can be overridden per handler.
+            - `True`: preload all state (this is the default behavior).
+            - `False`: preload nothing (lazy state).
+            - `["key1", "key2"]`: preload only the listed keys eagerly.
+            NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail.
         ingress_private (bool, optional): When set to True this service, with all its handlers, cannot be invoked from the restate-server
             HTTP and Kafka ingress, but only from other services.
             NOTE: You can set this field only if you register this service against restate-server >= 1.4,
@@ -85,6 +89,7 @@ def __init__(
         journal_retention: typing.Optional[timedelta] = None,
         idempotency_retention: typing.Optional[timedelta] = None,
         enable_lazy_state: typing.Optional[bool] = None,
+        state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None,
         ingress_private: typing.Optional[bool] = None,
         invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None,
         invocation_context_managers: typing.Optional[
@@ -97,7 +102,9 @@ def __init__(
         self.abort_timeout = abort_timeout
         self.journal_retention = journal_retention
         self.idempotency_retention = idempotency_retention
+        check_state_preload_options(state_preload, enable_lazy_state)
         self.enable_lazy_state = enable_lazy_state
+        self.state_preload = state_preload
         self.ingress_private = ingress_private
         self.invocation_retry_policy = invocation_retry_policy
         self.context_managers = invocation_context_managers
@@ -125,6 +132,7 @@ def handler(
         journal_retention: typing.Optional[timedelta] = None,
         idempotency_retention: typing.Optional[timedelta] = None,
         enable_lazy_state: typing.Optional[bool] = None,
+        state_preload: typing.Optional[typing.Union[bool, typing.List[str]]] = None,
         ingress_private: typing.Optional[bool] = None,
         invocation_retry_policy: typing.Optional[InvocationRetryPolicy] = None,
         invocation_context_managers: typing.Optional[
@@ -163,9 +171,13 @@ def handler(
             idempotency_retention: The retention duration of idempotent requests to this service.
                 NOTE: You can set this field only if you register this service against restate-server >= 1.4,
                 otherwise the service discovery will fail.
-            enable_lazy_state: When set to True, lazy state will be enabled for all invocations to this handler.
-                NOTE: You can set this field only if you register this service against restate-server >= 1.4,
-                otherwise the service discovery will fail.
+            enable_lazy_state: Deprecated, use `state_preload` instead.
+            state_preload: Configures which state keys are preloaded ("eager state") at the start of each
+                invocation to this handler. Overrides the value set on the service.
+                - `True`: preload all state (this is the default behavior).
+                - `False`: preload nothing (lazy state).
+                - `["key1", "key2"]`: preload only the listed keys eagerly.
+                NOTE: The list form is supported only by restate-server >= 1.8, otherwise the service discovery will fail.
             ingress_private: When set to True this handler cannot be invoked from the restate-server HTTP and Kafka ingress,
                 but only from other services.
                 NOTE: You can set this field only if you register this service against restate-server >= 1.4,
@@ -212,6 +224,7 @@ def wrapped(*args, **kwargs):
                 idempotency_retention,
                 None,
                 enable_lazy_state,
+                state_preload,
                 ingress_private,
                 invocation_retry_policy,
                 combined_context_managers,

python/restate/server.py

@@ -105,14 +105,16 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint):
     # Negotiate discovery protocol version
     version = 2
     if accept_header:
-        if "application/vnd.restate.endpointmanifest.v4+json" in accept_header:
+        if "application/vnd.restate.endpointmanifest.v5+json" in accept_header:
+            version = 5
+        elif "application/vnd.restate.endpointmanifest.v4+json" in accept_header:
             version = 4
         elif "application/vnd.restate.endpointmanifest.v3+json" in accept_header:
             version = 3
         elif "application/vnd.restate.endpointmanifest.v2+json" in accept_header:
             version = 2
         else:
-            await send_status_with_error_text(send, 415, f"Unsupported discovery version ${accept_header}")
+            await send_status_with_error_text(send, 415, f"Unsupported discovery version {accept_header}")
             return
 
     try:
@@ -128,7 +130,7 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint):
             }
         )
     except ValueError as e:
-        await send_status_with_error_text(send, 500, f"Error when computing discovery ${e}")
+        await send_status_with_error_text(send, 500, f"Error when computing discovery {e}")
         return
 
 

python/restate/service.py

@@ -203,6 +203,7 @@ def wrapped(*args, **kwargs):
                 idempotency_retention,
                 None,
                 None,
+                None,
                 ingress_private,
                 invocation_retry_policy,
                 combined_context_managers,