Add x-restate-server to the response headers

python/restate/server.py

@@ -19,12 +19,14 @@
 from restate.server_types import Receive, Scope, Send, binary_to_header, header_to_binary
 from restate.vm import VMWrapper
 from restate._internal import PyIdentityVerifier, IdentityVerificationException # pylint: disable=import-error,no-name-in-module
+from restate._internal import SDK_VERSION # pylint: disable=import-error,no-name-in-module
 from restate.aws_lambda import is_running_on_lambda, wrap_asgi_as_lambda_handler
 
+X_RESTATE_SERVER = header_to_binary([("x-restate-server", f"restate-sdk-python/{SDK_VERSION}")])
 
 async def send_status(send, receive, status_code: int):
     """respond with a status code"""
-    await send({'type': 'http.response.start', 'status': status_code})
+    await send({'type': 'http.response.start', 'status': status_code, "headers": X_RESTATE_SERVER})
     # For more info on why this loop, see ServerInvocationContext.leave()
     # pylint: disable=R0801
     while True:
@@ -49,10 +51,12 @@ async def send_discovery(scope: Scope, send: Send, endpoint: Endpoint):
     else:
         discovered_as = "bidi"
     headers, js = compute_discovery_json(endpoint, 1, discovered_as)
+    bin_headers = header_to_binary(headers.items())
+    bin_headers.extend(X_RESTATE_SERVER)
     await send({
         'type': 'http.response.start',
         'status': 200,
-        'headers': header_to_binary(headers.items()),
+        'headers': bin_headers,
         'trailers': False
         })
     await send({
@@ -68,10 +72,12 @@ async def process_invocation_to_completion(vm: VMWrapper,
                                            send: Send):
     """Invoke the user code."""
     status, res_headers = vm.get_response_head()
+    res_bin_headers = header_to_binary(res_headers)
+    res_bin_headers.extend(X_RESTATE_SERVER)
     await send({
         'type': 'http.response.start',
         'status': status,
-        'headers': header_to_binary(res_headers),
+        'headers': res_bin_headers,
         'trailers': False
     })
     assert status == 200

python/restate/vm.py

@@ -15,7 +15,7 @@
 
 from dataclasses import dataclass
 import typing
-from restate._internal import PyVM, PyFailure, PySuspended, PyVoid, PyStateKeys # pylint: disable=import-error,no-name-in-module
+from restate._internal import PyVM, PyFailure, PySuspended, PyVoid, PyStateKeys  # pylint: disable=import-error,no-name-in-module
 
 @dataclass
 class Invocation:
@@ -59,7 +59,7 @@ def __init__(self, *args: object) -> None:
 class VMWrapper:
     """
     A wrapper class for the restate_sdk._internal.PyVM class.
-    It provides a type-friendly interface to our shared vm. 
+    It provides a type-friendly interface to our shared vm.
     """
 
     def __init__(self, headers: typing.List[typing.Tuple[str, str]]):

src/lib.rs

@@ -1,10 +1,16 @@
 use pyo3::create_exception;
 use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyNone};
-use restate_sdk_shared_core::{AsyncResultHandle, CoreVM, Failure, Header, IdentityVerifier, Input, NonEmptyValue, ResponseHead, RunEnterResult, SuspendedOrVMError, TakeOutputResult, Target, VMError, Value, VM};
+use restate_sdk_shared_core::{
+    AsyncResultHandle, CoreVM, Failure, Header, IdentityVerifier, Input, NonEmptyValue,
+    ResponseHead, RunEnterResult, SuspendedOrVMError, TakeOutputResult, Target, VMError, Value, VM,
+};
 use std::borrow::Cow;
 use std::time::Duration;
 
+// Current crate version
+const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
+
 // Data model
 
 #[pyclass]
@@ -103,7 +109,7 @@ impl From<PyFailure> for Failure {
 #[derive(Clone)]
 struct PyStateKeys {
     #[pyo3(get, set)]
-    keys: Vec<String>
+    keys: Vec<String>,
 }
 
 #[pyclass]
@@ -237,7 +243,7 @@ impl PyVM {
                 Ok(PyFailure::from(f).into_py(py).into_bound(py).into_any())
             }
             Ok(Some(Value::StateKeys(keys))) => {
-                Ok(PyStateKeys {keys}.into_py(py).into_bound(py).into_any())
+                Ok(PyStateKeys { keys }.into_py(py).into_bound(py).into_any())
             }
         }
     }
@@ -259,9 +265,7 @@ impl PyVM {
             .map_err(Into::into)
     }
 
-    fn sys_get_state_keys(
-        mut self_: PyRefMut<'_, Self>,
-    ) -> Result<PyAsyncResultHandle, PyVMError> {
+    fn sys_get_state_keys(mut self_: PyRefMut<'_, Self>) -> Result<PyAsyncResultHandle, PyVMError> {
         self_
             .vm
             .sys_state_get_keys()
@@ -563,5 +567,6 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
         "IdentityVerificationException",
         m.py().get_type_bound::<IdentityVerificationException>(),
     )?;
+    m.add("SDK_VERSION", CURRENT_VERSION)?;
     Ok(())
 }