Protocol V6 + random_seed support (#542)

slinkydeveloper · web-flow · commit 665f394d28b6 · 2025-09-09T14:36:20.000+02:00
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/InvocationIdImpl.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/InvocationIdImpl.java
@@ -13,15 +13,17 @@
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Objects;
+import org.jspecify.annotations.Nullable;
 
 final class InvocationIdImpl implements InvocationId {
 
   private final String id;
   private Long seed;
 
-  InvocationIdImpl(String debugId) {
+  InvocationIdImpl(String debugId, @Nullable Long seed) {
     this.id = debugId;
-    this.seed = null;
+    // If random seed null, it will be computed
+    this.seed = seed;
   }
 
   @Override
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/ReplayingState.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/ReplayingState.java
@@ -101,7 +101,8 @@ public StateMachine.Input processInputCommand(StateContext stateContext) {
 
     //noinspection unchecked
     return new StateMachine.Input(
-        new InvocationIdImpl(stateContext.getStartInfo().debugId()),
+        new InvocationIdImpl(
+            stateContext.getStartInfo().debugId(), stateContext.getStartInfo().randomSeed()),
         byteStringToSlice(inputCommandMessage.getValue().getContent()),
         Map.ofEntries(
             inputCommandMessage.getHeadersList().stream()
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/ServiceProtocol.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/ServiceProtocol.java
@@ -15,7 +15,7 @@ public class ServiceProtocol {
   public static final Protocol.ServiceProtocolVersion MIN_SERVICE_PROTOCOL_VERSION =
       Protocol.ServiceProtocolVersion.V5;
   public static final Protocol.ServiceProtocolVersion MAX_SERVICE_PROTOCOL_VERSION =
-      Protocol.ServiceProtocolVersion.V5;
+      Protocol.ServiceProtocolVersion.V6;
 
   static final String CONTENT_TYPE = "content-type";
 
@@ -40,6 +40,9 @@ static Protocol.ServiceProtocolVersion parseServiceProtocolVersion(String versio
     if (version.equals("application/vnd.restate.invocation.v5")) {
       return Protocol.ServiceProtocolVersion.V5;
     }
+    if (version.equals("application/vnd.restate.invocation.v6")) {
+      return Protocol.ServiceProtocolVersion.V6;
+    }
     return Protocol.ServiceProtocolVersion.SERVICE_PROTOCOL_VERSION_UNSPECIFIED;
   }
 
@@ -59,6 +62,9 @@ static String serviceProtocolVersionToHeaderValue(Protocol.ServiceProtocolVersio
     if (Objects.requireNonNull(version) == Protocol.ServiceProtocolVersion.V5) {
       return "application/vnd.restate.invocation.v5";
     }
+    if (Objects.requireNonNull(version) == Protocol.ServiceProtocolVersion.V6) {
+      return "application/vnd.restate.invocation.v6";
+    }
     throw new IllegalArgumentException(
         String.format("Service protocol version '%s' has no header value", version.getNumber()));
   }
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StartInfo.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StartInfo.java
@@ -17,4 +17,5 @@ record StartInfo(
     String objectKey,
     int entriesToReplay,
     int retryCountSinceLastStoredEntry,
-    Duration durationSinceLastStoredEntry) {}
+    Duration durationSinceLastStoredEntry,
+    Long randomSeed) {}
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateContext.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateContext.java
@@ -10,24 +10,33 @@
 
 import com.google.protobuf.MessageLite;
 import dev.restate.sdk.core.EndpointRequestHandler;
+import dev.restate.sdk.core.generated.protocol.Protocol;
 import java.util.Objects;
 import java.util.concurrent.Flow;
 
 final class StateContext {
 
+  private final Protocol.ServiceProtocolVersion negotiatedProtocolVersion;
   private final StateHolder stateHolder;
   private final Journal journal;
   private EagerState eagerState;
   private transient StartInfo startInfo;
   private boolean inputClosed;
   private Flow.Subscriber<MessageLite> outputSubscriber;
 
-  StateContext(EndpointRequestHandler.LoggingContextSetter loggingContextSetter) {
+  StateContext(
+      EndpointRequestHandler.LoggingContextSetter loggingContextSetter,
+      Protocol.ServiceProtocolVersion negotiatedProtocolVersion) {
     this.stateHolder = new StateHolder(loggingContextSetter);
+    this.negotiatedProtocolVersion = negotiatedProtocolVersion;
     this.journal = new Journal();
     this.inputClosed = false;
   }
 
+  public Protocol.ServiceProtocolVersion getNegotiatedProtocolVersion() {
+    return negotiatedProtocolVersion;
+  }
+
   public State getCurrentState() {
     return stateHolder.getState();
   }
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachineImpl.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/StateMachineImpl.java
@@ -36,8 +36,6 @@ class StateMachineImpl implements StateMachine {
   private static final String AWAKEABLE_IDENTIFIER_PREFIX = "sign_1";
   private static final int CANCEL_SIGNAL_ID = 1;
 
-  private final Protocol.ServiceProtocolVersion serviceProtocolVersion;
-
   // Callbacks
   private final CompletableFuture<Void> waitForReadyFuture = new CompletableFuture<>();
   private CompletableFuture<Void> waitNextProcessedInput;
@@ -54,17 +52,16 @@ class StateMachineImpl implements StateMachine {
       EndpointRequestHandler.LoggingContextSetter loggingContextSetter) {
     String contentTypeHeader = headersAccessor.get(ServiceProtocol.CONTENT_TYPE);
 
-    this.serviceProtocolVersion = ServiceProtocol.parseServiceProtocolVersion(contentTypeHeader);
-
-    if (!ServiceProtocol.isSupported(this.serviceProtocolVersion)) {
+    var serviceProtocolVersion = ServiceProtocol.parseServiceProtocolVersion(contentTypeHeader);
+    if (!ServiceProtocol.isSupported(serviceProtocolVersion)) {
       throw new ProtocolException(
           String.format(
               "Service endpoint does not support the service protocol version '%s'.",
               contentTypeHeader),
           ProtocolException.UNSUPPORTED_MEDIA_TYPE_CODE);
     }
 
-    this.stateContext = new StateContext(loggingContextSetter);
+    this.stateContext = new StateContext(loggingContextSetter, serviceProtocolVersion);
   }
 
   // -- Few callbacks
@@ -176,7 +173,8 @@ public void onComplete() {
 
   @Override
   public String getResponseContentType() {
-    return ServiceProtocol.serviceProtocolVersionToHeaderValue(serviceProtocolVersion);
+    return ServiceProtocol.serviceProtocolVersionToHeaderValue(
+        stateContext.getNegotiatedProtocolVersion());
   }
 
   @Override
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/WaitingStartState.java b/sdk-core/src/main/java/dev/restate/sdk/core/statemachine/WaitingStartState.java
@@ -41,7 +41,12 @@ public void onNewMessage(
             startMessage.getKey(),
             startMessage.getKnownEntries(),
             startMessage.getRetryCountSinceLastStoredEntry(),
-            Duration.ofMillis(startMessage.getDurationSinceLastStoredEntry())));
+            Duration.ofMillis(startMessage.getDurationSinceLastStoredEntry()),
+            // Random seed from start message will be set only if protocol >= 6
+            stateContext.getNegotiatedProtocolVersion().getNumber()
+                    >= Protocol.ServiceProtocolVersion.V6_VALUE
+                ? startMessage.getRandomSeed()
+                : null));
     stateContext.setEagerState(new EagerState(startMessage));
 
     // Tracing and logging setup
diff --git a/sdk-core/src/main/service-protocol/dev/restate/service/protocol.proto b/sdk-core/src/main/service-protocol/dev/restate/service/protocol.proto
@@ -32,6 +32,9 @@ enum ServiceProtocolVersion {
   // * New command to attach to existing invocation
   // * New command to get output of existing invocation
   V5 = 5;
+  // Added:
+  // * StartMessage.random_seed
+  V6 = 6;
 }
 
 // --- Core frames ---
@@ -73,6 +76,10 @@ message StartMessage {
   // Please note this duration might not be accurate,
   // and might change depending on which Restate replica executes the request.
   uint64 duration_since_last_stored_entry = 8;
+
+  // Random seed to use to seed the deterministic RNG exposed in the context API.
+  // This will be stable across restarts.
+  uint64 random_seed = 9;
 }
 
 // Type: 0x0000 + 1
diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/MockBidiStream.java b/sdk-core/src/test/java/dev/restate/sdk/core/MockBidiStream.java
@@ -60,7 +60,10 @@ public void executeTest(TestDefinitions.TestDefinition definition) {
         server.processorForRequest(
             "/" + serviceDefinition.getServiceName() + "/" + definition.getMethod(),
             HeadersAccessor.wrap(
-                Map.of("content-type", ProtoUtils.serviceProtocolContentTypeHeader())),
+                Map.of(
+                    "content-type",
+                    ProtoUtils.serviceProtocolContentTypeHeader(
+                        definition.isEnablePreviewContext()))),
             EndpointRequestHandler.LoggingContextSetter.THREAD_LOCAL_INSTANCE,
             coreExecutor,
             true);
diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/MockRequestResponse.java b/sdk-core/src/test/java/dev/restate/sdk/core/MockRequestResponse.java
@@ -58,7 +58,10 @@ public void executeTest(TestDefinition definition) {
         server.processorForRequest(
             "/" + serviceDefinition.getServiceName() + "/" + definition.getMethod(),
             HeadersAccessor.wrap(
-                Map.of("content-type", ProtoUtils.serviceProtocolContentTypeHeader())),
+                Map.of(
+                    "content-type",
+                    ProtoUtils.serviceProtocolContentTypeHeader(
+                        definition.isEnablePreviewContext()))),
             EndpointRequestHandler.LoggingContextSetter.THREAD_LOCAL_INSTANCE,
             syscallsExecutor,
             false);
diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/RandomTestSuite.java b/sdk-core/src/test/java/dev/restate/sdk/core/RandomTestSuite.java
@@ -8,8 +8,8 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.sdk.core;
 
-import static dev.restate.sdk.core.AssertUtils.*;
 import static dev.restate.sdk.core.statemachine.ProtoUtils.*;
+import static dev.restate.sdk.core.statemachine.ProtoUtils.invocationIdToRandomSeed;
 
 import dev.restate.sdk.core.TestDefinitions.TestDefinition;
 import dev.restate.sdk.core.TestDefinitions.TestInvocationBuilder;
@@ -26,12 +26,21 @@ public abstract class RandomTestSuite implements TestSuite {
   @Override
   public Stream<TestDefinition> definitions() {
     String debugId = "my-id";
+    long startMessageSeed = System.currentTimeMillis();
 
     return Stream.of(
+        this.randomShouldBeDeterministic()
+            .withInput(
+                startMessage(1).setDebugId(debugId).setRandomSeed(startMessageSeed),
+                ProtoUtils.inputCmd())
+            // This enables protocol v6
+            .enablePreviewContext()
+            .expectingOutput(outputCmd(getExpectedInt(startMessageSeed)), END_MESSAGE)
+            .named("Using StartMessage.random_seed"),
         this.randomShouldBeDeterministic()
             .withInput(startMessage(1).setDebugId(debugId), ProtoUtils.inputCmd())
             .expectingOutput(
-                outputCmd(getExpectedInt(ProtoUtils.invocationIdToRandomSeed(debugId))),
-                END_MESSAGE));
+                outputCmd(getExpectedInt(invocationIdToRandomSeed(debugId))), END_MESSAGE)
+            .named("Using invocation id"));
   }
 }
diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/lambda/LambdaHandlerTest.java b/sdk-core/src/test/java/dev/restate/sdk/core/lambda/LambdaHandlerTest.java
@@ -43,7 +43,7 @@ public void testInvoke() throws IOException {
 
     // Mock request
     APIGatewayProxyRequestEvent request = new APIGatewayProxyRequestEvent();
-    request.setHeaders(Map.of("content-type", ProtoUtils.serviceProtocolContentTypeHeader()));
+    request.setHeaders(Map.of("content-type", ProtoUtils.serviceProtocolContentTypeHeader(false)));
     request.setPath(
         "/a/path/prefix/invoke/" + JavaCounterServiceHandlers.Metadata.SERVICE_NAME + "/get");
     request.setHttpMethod("POST");
@@ -66,7 +66,7 @@ public void testInvoke() throws IOException {
     // Assert response
     assertThat(response.getStatusCode()).isEqualTo(200);
     assertThat(response.getHeaders())
-        .containsEntry("content-type", ProtoUtils.serviceProtocolContentTypeHeader());
+        .containsEntry("content-type", ProtoUtils.serviceProtocolContentTypeHeader(false));
     assertThat(response.getIsBase64Encoded()).isTrue();
     assertThat(response.getBody())
         .asBase64Decoded()
diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/statemachine/ProtoUtils.java b/sdk-core/src/test/java/dev/restate/sdk/core/statemachine/ProtoUtils.java
@@ -29,17 +29,14 @@
 public class ProtoUtils {
 
   public static long invocationIdToRandomSeed(String invocationId) {
-    return new InvocationIdImpl(invocationId).toRandomSeed();
-  }
-
-  public static String serviceProtocolContentTypeHeader() {
-    return ServiceProtocol.serviceProtocolVersionToHeaderValue(
-        ServiceProtocol.MIN_SERVICE_PROTOCOL_VERSION);
+    return new InvocationIdImpl(invocationId, null).toRandomSeed();
   }
 
   public static String serviceProtocolContentTypeHeader(boolean enableContextPreview) {
     return ServiceProtocol.serviceProtocolVersionToHeaderValue(
-        ServiceProtocol.MAX_SERVICE_PROTOCOL_VERSION);
+        enableContextPreview
+            ? ServiceProtocol.MAX_SERVICE_PROTOCOL_VERSION
+            : ServiceProtocol.MIN_SERVICE_PROTOCOL_VERSION);
   }
 
   public static String serviceProtocolDiscoveryContentTypeHeader() {
diff --git a/sdk-core/src/test/kotlin/dev/restate/sdk/core/vertx/RestateHttpServerTest.kt b/sdk-core/src/test/kotlin/dev/restate/sdk/core/vertx/RestateHttpServerTest.kt
@@ -102,8 +102,8 @@ internal class RestateHttpServerTest {
         // Prepare request header
         request
             .setChunked(true)
-            .putHeader(HttpHeaders.CONTENT_TYPE, serviceProtocolContentTypeHeader())
-            .putHeader(HttpHeaders.ACCEPT, serviceProtocolContentTypeHeader())
+            .putHeader(HttpHeaders.CONTENT_TYPE, serviceProtocolContentTypeHeader(false))
+            .putHeader(HttpHeaders.ACCEPT, serviceProtocolContentTypeHeader(false))
         request.write(encode(startMessage(0).build()))
 
         val response = request.response().coAwait()

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ public class ServiceProtocol {`
`15`	`15`	`public static final Protocol.ServiceProtocolVersion MIN_SERVICE_PROTOCOL_VERSION =`
`16`	`16`	`Protocol.ServiceProtocolVersion.V5;`
`17`	`17`	`public static final Protocol.ServiceProtocolVersion MAX_SERVICE_PROTOCOL_VERSION =`
`18`		`- Protocol.ServiceProtocolVersion.V5;`
	`18`	`+ Protocol.ServiceProtocolVersion.V6;`
`19`	`19`
`20`	`20`	`static final String CONTENT_TYPE = "content-type";`
`21`	`21`
`@@ -40,6 +40,9 @@ static Protocol.ServiceProtocolVersion parseServiceProtocolVersion(String versio`
`40`	`40`	`if (version.equals("application/vnd.restate.invocation.v5")) {`
`41`	`41`	`return Protocol.ServiceProtocolVersion.V5;`
`42`	`42`	`}`
	`43`	`+ if (version.equals("application/vnd.restate.invocation.v6")) {`
	`44`	`+ return Protocol.ServiceProtocolVersion.V6;`
	`45`	`+ }`
`43`	`46`	`return Protocol.ServiceProtocolVersion.SERVICE_PROTOCOL_VERSION_UNSPECIFIED;`
`44`	`47`	`}`
`45`	`48`
`@@ -59,6 +62,9 @@ static String serviceProtocolVersionToHeaderValue(Protocol.ServiceProtocolVersio`
`59`	`62`	`if (Objects.requireNonNull(version) == Protocol.ServiceProtocolVersion.V5) {`
`60`	`63`	`return "application/vnd.restate.invocation.v5";`
`61`	`64`	`}`
	`65`	`+ if (Objects.requireNonNull(version) == Protocol.ServiceProtocolVersion.V6) {`
	`66`	`+ return "application/vnd.restate.invocation.v6";`
	`67`	`+ }`
`62`	`68`	`throw new IllegalArgumentException(`
`63`	`69`	`String.format("Service protocol version '%s' has no header value", version.getNumber()));`
`64`	`70`	`}`