support step context from Python Workflows SDK

Author: Caio-Nogueira

src/pyodide/internal/workers-api/src/workers/_workers.py

@@ -1163,8 +1163,20 @@ async def _resolve_dependency(self, dep):
 
 
 async def _do_call(entrypoint, name, config, callback, *results):
-    async def _callback():
-        result = callback(*results)
+    async def _callback(ctx=None):
+        sig = inspect.signature(callback)
+        params = [
+            p
+            for p in sig.parameters.values()
+            if p.kind
+            not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+        ]
+        wants_ctx = len(params) > len(results)
+
+        if wants_ctx and ctx is not None:
+            result = callback(python_from_rpc(ctx), *results)
+        else:
+            result = callback(*results)
 
         if inspect.iscoroutine(result):
             result = await result

src/workerd/server/tests/python/workflow-entrypoint/worker.js

@@ -4,7 +4,8 @@ import * as assert from 'node:assert';
 class Context extends RpcTarget {
   async do(name, fn) {
     try {
-      const result = await fn();
+      const ctx = { attempt: '1', metadata: 'expected_return_metadata' };
+      const result = await fn(ctx);
       return result;
     } catch (e) {
       console.log(`Error received: ${e.name} Message: ${e.message}`);
@@ -17,14 +18,22 @@ class Context extends RpcTarget {
 export default {
   async test(ctrl, env, ctx) {
     // JS types
-    const stubStep = new Context();
+    const stubStep= new Context();
 
-    const resp = await env.PythonWorkflow.run(
+    let resp = await env.PythonWorkflow.run(
       {
         foo: 'bar',
       },
       stubStep
     );
     assert.deepStrictEqual(resp, 'foobar');
+
+    resp = await env.PythonWorkflowWithCtx.run(
+      {
+        foo: 'bar',
+      },
+      stubStep
+    );
+    assert.deepStrictEqual(resp, 'expected_return_metadata');
   },
 };

src/workerd/server/tests/python/workflow-entrypoint/workflow-entrypoint.wd-test

@@ -18,6 +18,7 @@ const pyWorker :Workerd.Worker = (
 
   bindings = [
     (name = "PythonWorkflow", service = (name = "py", entrypoint = "WorkflowEntrypointExample")),
+    (name = "PythonWorkflowWithCtx", service = (name = "py", entrypoint = "WorkflowEntrypointWithCtx")),
   ],
 );
 
@@ -32,5 +33,6 @@ const jsWorker :Workerd.Worker = (
 
   bindings = [
     (name = "PythonWorkflow", service = (name = "py", entrypoint = "WorkflowEntrypointExample")),
+    (name = "PythonWorkflowWithCtx", service = (name = "py", entrypoint = "WorkflowEntrypointWithCtx")),
   ],
 );

src/workerd/server/tests/python/workflow-entrypoint/workflow.py

@@ -44,23 +44,48 @@ async def step_3(_result1, _result2):
         await await_step(step_3)
 
         # `step_1` and `step_2` run serially
-        @step.do("step_4", depends=[step_1, step_2], concurrent=False)
+        @step.do("step_4", depends=[step_1, step_2], concurrent=True)
         async def step_4(result1, result2):
             print("Executing step 4 (depends on step 1 and step 2)")
             assert result1["foo"] == "foo"
             assert result2["bar"] == "bar"
 
         await await_step(step_4)
 
+        # tests step memoization - steps 1 and 2 are already resolved
         @step.do("step_5", depends=[step_1, step_2], concurrent=False)
-        async def step_5(result1, result2):
+        async def step_5(result1=(), result2=()):
             print("Executing step 5 (depends on step 1 and step 2)")
             assert result1["foo"] == "foo"
             assert result2["bar"] == "bar"
 
             return result1["foo"] + result2["bar"]
 
-        return await await_step(step_5)
+        return await step_5()
+
+class WorkflowEntrypointWithCtx(WorkflowEntrypoint):
+    async def run(self, event, step):
+        @step.do("step_1")
+        async def step_1():
+            # tests backwards compat with workflows that don't have ctx in the step callback
+            print("Executing step 1")
+            return {"foo": "foo"}
+
+        @step.do("step_2")
+        async def step_2(ctx):
+            print("Executing step 2")
+            return {"bar": "bar"}
+
+        @step.do("step_3", depends=[step_1, step_2], concurrent=False)
+        async def step_3(context, result1=(), result2=()):
+            print(f"context exists and is {context}")
+            print("Executing step 5 (depends on step 1 and step 2)")
+            assert result1["foo"] == "foo"
+            assert result2["bar"] == "bar"
+            assert context["attempt"] == "1"
+            return context["metadata"]
+
+        return await step_3()
 
 
 async def test(ctrl, env, ctx):