From c4b2bb1d4f08fbd1715940aa9e4e373eb4170ac0 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 16:21:08 +0000 Subject: [PATCH 1/7] Sessions: a call to a workflow made again after a crash waits for the run it started A session object runs again from its first line after a crash, so it calls again what the crashed attempt called. An idempotency key makes that the same call, except for a workflow's main handler: Restate takes no idempotency key for it, and refuses the second call (409, "the workflow method was already invoked"). The agents' `workflow` tool, and LangGraph's `nodes.workflow`, then told the model that the workflow failed, while it ran to the end. A call to a workflow's main handler (`ctx.call(wf.main, ..., key=id)`) that is refused this way now attaches to the workflow's run by its id, through the shared core's `AttachInvocationTarget::WorkflowId` (`attach_workflow` on the VM and the server context), and returns its result: the run of that id, whoever started it. Any other 409 (a cancellation), a workflow's shared handlers, and calls by name are unchanged. Tests: the rejoin and what isn't rejoined (fake host), the agents' workflow tool after a crash, and a session handler that crashes once its workflow call is made, against a real server. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/_calls.py | 4 +- python/restate/agents/_tools.py | 4 +- python/restate/agents/langgraph/nodes.py | 5 +- python/restate/server_context.py | 11 +++ python/restate/sessions/_api.py | 3 + python/restate/sessions/_calls.py | 50 ++++++++++--- python/restate/sessions/_object.py | 3 + python/restate/vm.py | 6 ++ src/lib.rs | 19 +++++ tests/agents_tools.py | 22 +++++- tests/sessions_harness.py | 31 +++++++- tests/sessions_object.py | 90 ++++++++++++++++++++++++ 12 files changed, 234 insertions(+), 14 deletions(-) diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index 8c2701ae..3e517097 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -20,7 +20,9 @@ interrupted, and the model makes it again under a new id. Either way, a call with the same tool and input as one an earlier attempt made reuses that call's id, and so its idempotency key: it is the same Restate call, which it -waits for again or whose result it gets. +waits for again or whose result it gets. (A workflow's main handler takes no +idempotency key: the call made again waits for the run of its workflow id +instead, as ``restate.sessions``' calls do.) A call is kept until the framework's transcript has its result under the latest id it was made with (``settle`` checks this at each publish and before diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index f6422a2f..34de4ca6 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -523,7 +523,9 @@ def workflow( """Create a tool that runs a workflow to completion and returns its result. Each tool call starts a workflow of its own, whose id is built from the - agent's key and the tool call's id, unless ``key`` is given. + agent's key and the tool call's id, unless ``key`` is given. A workflow + runs once per id: a tool call run again after a crash waits for the run + it started, and a ``key`` that already ran gets that run's result. Args: target: The workflow's main handler. diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 67609622..49c6f0b9 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -233,7 +233,10 @@ def workflow( takes one. Plain data (a dict, say) is validated into the handler's input type, such as a Pydantic model. key: The workflow id, from the state. Defaults to one of the node's - own, made of the agent's key and the node's task. + own, made of the agent's key and the node's task. A workflow + runs once per id: a node run again after a crash waits for the + run it started, and an id that already ran gets that run's + result. Returns: The node, for ``StateGraph.add_node``. diff --git a/python/restate/server_context.py b/python/restate/server_context.py index aa22cfc8..809d295e 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -1319,6 +1319,17 @@ def attach_invocation( update_restate_context_is_replaying(self.vm) return self.create_future(handle, serde) + def attach_workflow(self, name: str, key: str) -> RestateDurableFuture[bytes]: + """Attach to the run of the workflow ``name`` with the id ``key``: its result, as bytes, once it ends. + + For a call to a workflow's main handler that Restate refused because the + run had started (a call made again after a crash, from a handler that + runs again from its first line rather than replaying its journal). + """ + handle = self.vm.attach_workflow(name, key) + update_restate_context_is_replaying(self.vm) + return self.create_future(handle, BytesSerde()) + def get_extension_data(self, key: str) -> Any: """Get extension data by key.""" return self.extension_data.get(key) diff --git a/python/restate/sessions/_api.py b/python/restate/sessions/_api.py index e9c12b35..d15e39f2 100644 --- a/python/restate/sessions/_api.py +++ b/python/restate/sessions/_api.py @@ -128,6 +128,9 @@ def call( key: The object's or workflow's key; a service takes none. idempotency_key: The same on every attempt, so that a retried attempt gets the first call's result instead of calling again. + A workflow's main handler needs none: it runs once per + workflow id, and a call that Restate refuses because that + run started waits for the run, and returns its result. headers: Headers to send with the call. Raises: diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 4ecdb9cd..e8f0ed6c 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -18,6 +18,13 @@ count = await ctx.call(increment, 5, key="a", idempotency_key=f"{ctx.request.id}:inc") ctx.send(record, entry, key="audit", delay=timedelta(hours=1)) answer = await ctx.generic_call("Counter", "get", key="a") + +A workflow's main handler runs once per workflow id, and Restate takes no +idempotency key for it: a second call is refused (409, "already invoked"). +A session object runs again from its first line after a crash, so it would +make such a call again. ``call`` makes a call to a workflow's main handler +that is refused this way wait for that run instead, and return its result: +the run of that id, whoever started it. """ from __future__ import annotations @@ -28,7 +35,7 @@ from datetime import timedelta from typing import Any, Protocol -from restate.context import RestateDurableCallFuture +from restate.context import RestateDurableCallFuture, RestateDurableFuture from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable from restate.serde import JsonSerde, Serde @@ -74,9 +81,11 @@ def generic_send( headers: dict[str, str] | None = None, ) -> Any: ... + def attach_workflow(self, name: str, key: str) -> RestateDurableFuture[bytes]: ... + -def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Serde[Any], Serde[Any]]: - """Return the service name, handler name and serdes of ``handler``, checking ``key``. +def _target(handler: Callable[..., Any], key: str | None) -> tuple[Handler[Any, Any], str, str, Serde[Any], Serde[Any]]: + """Return ``handler``'s definition, its service name, handler name and serdes, checking ``key``. Raises ``TerminalError`` (since a retry would fail the same way) if ``handler`` is not a handler, if an object's or workflow's handler has no @@ -93,7 +102,12 @@ def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Ser raise TerminalError(f"{service} is {what}: call it with its key") if not keyed and key is not None: raise TerminalError(f"{service} is a service: call it without a key") - return service, target.name, target.handler_io.input_serde, target.handler_io.output_serde + return target, service, target.name, target.handler_io.input_serde, target.handler_io.output_serde + + +def already_invoked(error: TerminalError) -> bool: + """Whether ``error`` is Restate refusing a call to a workflow's main handler because that workflow's run started.""" + return error.status_code == 409 and "already invoked" in error.message class CallFuture(RestateDurableCallFuture[O]): @@ -101,15 +115,31 @@ class CallFuture(RestateDurableCallFuture[O]): Await it for the result. ``invocation_id()`` returns the id of the call's invocation, and ``cancel_invocation()`` cancels it. + + With ``rejoin`` (a call to a workflow's main handler), a call that Restate + refuses because the workflow's run started waits for that run instead + (``rejoin()``), and returns its result. """ - def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> None: + def __init__( + self, + call: RestateDurableCallFuture[bytes], + serde: Serde[O], + rejoin: Callable[[], Awaitable[bytes]] | None = None, + ) -> None: super().__init__() self._call: RestateDurableCallFuture[bytes] = call self._serde: Serde[O] = serde + self._rejoin = rejoin async def _result(self) -> O: - return typing.cast(O, self._serde.deserialize(await self._call)) + try: + result = await self._call + except TerminalError as error: + if self._rejoin is None or not already_invoked(error): + raise + result = await self._rejoin() + return typing.cast(O, self._serde.deserialize(result)) def __await__(self) -> Generator[Any, Any, O]: return self._result().__await__() @@ -159,8 +189,12 @@ def call( idempotency_key: str | None = None, headers: dict[str, str] | None = None, ) -> CallFuture[Any]: - service, name, input_serde, output_serde = _target(handler, key) + target, service, name, input_serde, output_serde = _target(handler, key) future = self._call_host.generic_call(service, name, input_serde.serialize(arg), key, idempotency_key, headers) + if target.kind == "workflow" and key is not None: + # The main handler of a workflow: a call made again (after a crash) joins the run it started. + host, workflow_id = self._call_host, key + return CallFuture(future, output_serde, rejoin=lambda: host.attach_workflow(service, workflow_id)) return CallFuture(future, output_serde) def generic_call( @@ -225,7 +259,7 @@ def send( idempotency_key: str | None = None, headers: dict[str, str] | None = None, ) -> None: - service, name, input_serde, _ = _target(handler, key) + _, service, name, input_serde, _ = _target(handler, key) parameter = input_serde.serialize(arg) self._send_sink(GenericSend(service, name, parameter, key, delay, idempotency_key, headers)) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 8ce7e9ae..28a1d670 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -167,6 +167,9 @@ def generic_call( ) -> Any: return self.context.generic_call(service, handler, arg, key, idempotency_key, headers) + def attach_workflow(self, name: str, key: str) -> Any: + return self.context.attach_workflow(name, key) + def generic_send( self, service: str, diff --git a/python/restate/vm.py b/python/restate/vm.py index 85be9830..3956708a 100644 --- a/python/restate/vm.py +++ b/python/restate/vm.py @@ -574,6 +574,12 @@ def attach_invocation(self, invocation_id: str) -> int: """ return self.vm.attach_invocation(invocation_id) + def attach_workflow(self, name: str, key: str) -> int: + """ + Attach to the run of the workflow ``name`` with the id ``key`` + """ + return self.vm.attach_workflow(name, key) + def is_replaying(self) -> bool: """Returns true if the state machine is replaying.""" return self.vm.is_replaying() diff --git a/src/lib.rs b/src/lib.rs index 1236cd53..5837ae71 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -959,6 +959,25 @@ impl PyVM { .map_err(Into::into) } + /// Attach to the run of the workflow `name` with the id `key` (unscoped): its result, once it ends. + fn attach_workflow( + mut self_: PyRefMut<'_, Self>, + name: String, + key: String, + ) -> Result { + self_ + .vm + .sys_attach_invocation( + restate_sdk_shared_core::AttachInvocationTarget::WorkflowId { + name, + key, + scope: None, + }, + ) + .map(Into::into) + .map_err(Into::into) + } + fn sys_end(mut self_: PyRefMut<'_, Self>) -> Result<(), PyVMError> { self_.vm.sys_end().map_err(Into::into) } diff --git a/tests/agents_tools.py b/tests/agents_tools.py index e238d835..ddc70f4a 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -24,10 +24,10 @@ from restate.sessions import internal from restate.agents._common import ToolError from restate.agents._object import AgentObject -from restate.agents._tools import Tool, ToolCall, call, subagent, tool, wait_for_signal +from restate.agents._tools import Tool, ToolCall, call, subagent, tool, wait_for_signal, workflow from restate.context import Request from sessions_database import Backend # pylint: disable=import-error -from sessions_object import Host, request # pylint: disable=import-error +from sessions_object import ALREADY_INVOKED, Host, pay, request # pylint: disable=import-error @pytest.fixture @@ -121,6 +121,24 @@ async def test_call_lets_the_model_choose_the_key_and_makes_the_call_once_per_to assert host.calls[-1][4] == f"{req.id}:t1" +# --- workflow ------------------------------------------------------------------------------------- + + +async def test_workflow_made_again_after_a_crash_waits_for_the_run_it_started(): + """A crash after the tool call started its workflow runs the tool call again: Restate refuses the second call to + the workflow's main handler (409, already invoked), and the tool answers with that run's result, not the error.""" + tools["pay"] = workflow(pay, description="Pay an amount out.") + host = Host( + Backend(), + service="runner", + key="alice", + refusals={"pay": ALREADY_INVOKED}, + answers={"payouts/alice-c1": b'"paid 30"'}, + ) + assert await run_tool(host, "pay", {"input": 30}, "c1") == "paid 30" + assert host.attached == [("payouts", "alice-c1")] + + # --- signal and subagent -------------------------------------------------------------------------- diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 2678f068..654c49a4 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -156,10 +156,32 @@ async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: return cancelled +payouts = restate.Workflow("payouts") + + +@payouts.main() +async def pay(ctx: restate.WorkflowContext, amount: int) -> str: + await ctx.sleep(timedelta(milliseconds=100)) + return f"paid {amount} for {ctx.key()}" + + +@notes.handler() +async def refund(ctx: sessions.SessionContext, order: str) -> str: + """Pays ``order`` out with a workflow, and crashes on its first attempt once the call is made: the retry calls the + workflow again, which Restate refuses (409, already invoked), and so waits for the run the first attempt started.""" + payout = ctx.call(pay, 30, key=order) + await payout.invocation_id() # the call is made + seen = await ctx.db.get("SELECT count(*) AS n FROM notes WHERE text = ?", [order]) + if seen is not None and seen["n"] == 0: + await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [order]) + raise ValueError("crash after the call was made") + return await payout + + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([inbox, pricing, sleeper, notes]), restate_image=RESTATE_IMAGE + restate.app([inbox, pricing, sleeper, payouts, notes]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -181,6 +203,13 @@ async def test_a_retry_runs_the_handler_from_the_start_against_what_was_committe assert await harness.client.object_call(texts, key=key, arg=None) == ["once"] +async def test_a_call_to_a_workflow_made_again_after_a_crash_returns_the_run_it_started(harness: HarnessEnvironment): + key, order = str(uuid.uuid4()), f"order-{uuid.uuid4()}" + assert await harness.client.object_call(refund, key=key, arg=order) == f"paid 30 for {order}" + # Another key's call for the same workflow id gets that run's result too: a workflow runs once per id. + assert await harness.client.object_call(refund, key=str(uuid.uuid4()), arg=order) == f"paid 30 for {order}" + + async def test_calls_and_sql_are_awaited_together(harness: HarnessEnvironment): key = str(uuid.uuid4()) await harness.client.object_call(add, key=key, arg="x") diff --git a/tests/sessions_object.py b/tests/sessions_object.py index fbc5cb29..dd7102e7 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -84,6 +84,26 @@ async def cancel_invocation(self) -> None: pass +class Refused(restate.RestateDurableCallFuture[bytes]): + """A call that Restate refused with ``error``.""" + + def __init__(self, error: TerminalError) -> None: + super().__init__() + self.error = error + + def __await__(self): + async def refused() -> bytes: + raise self.error + + return refused().__await__() + + async def invocation_id(self) -> str: + return "inv_refused" + + async def cancel_invocation(self) -> None: + pass + + @dataclass class Host: """An attempt's context over ``backend``, the object's state, as the SDK's gives it.""" @@ -95,6 +115,10 @@ class Host: sends: list[Send] = field(default_factory=list) calls: list[tuple[str, str, bytes, str | None, str | None]] = field(default_factory=list) answers: dict[str, bytes] = field(default_factory=dict) + # The calls Restate refuses, by handler: the error they fail with. + refusals: dict[str, TerminalError] = field(default_factory=dict) + # The workflow runs attached to, as (workflow, id); their results are the answers of "/". + attached: list[tuple[str, str]] = field(default_factory=list) cancelled: list[str] = field(default_factory=list) # How many more syncs go through before the process crashes: every sync after fails, losing its writes. _syncs_left: int | None = None @@ -127,8 +151,14 @@ async def sync(self) -> None: def generic_call(self, service, handler, arg, key=None, idempotency_key=None, headers=None): self.calls.append((service, handler, arg, key, idempotency_key)) + if handler in self.refusals: + return Refused(self.refusals[handler]) return Answer(self.answers[handler]) + def attach_workflow(self, name: str, key: str): + self.attached.append((name, key)) + return Answer(self.answers[f"{name}/{key}"]) + def generic_send(self, service, handler, arg, key=None, send_delay=None, idempotency_key=None, headers=None): self.sends.append(Send(service, handler, arg, key, send_delay, idempotency_key)) @@ -359,6 +389,66 @@ async def literal(ctx: sessions.SessionContext) -> Any: assert await run(caller, "literal", host) == 42 +payouts = restate.Workflow("payouts") + + +@payouts.main() +async def pay(ctx: restate.WorkflowContext, amount: int) -> str: + return f"paid {amount}" + + +@payouts.handler() +async def paid(ctx: restate.WorkflowSharedContext) -> str: + return "paid" + + +ALREADY_INVOKED = TerminalError("the workflow method was already invoked", status_code=409) + + +async def test_a_call_to_a_workflow_refused_as_already_invoked_waits_for_that_run(): + """A session object runs again from its first line, so it calls a workflow again after a crash: Restate refuses the + second call to its main handler (it takes no idempotency key), and the call waits for the run it started.""" + caller = sessions.SessionObject("caller") + + @caller.handler() + async def refund(ctx: sessions.SessionContext, amount: int) -> str: + return await ctx.call(pay, amount, key="order-7") + + host = Host( + Backend(), service="caller", refusals={"pay": ALREADY_INVOKED}, answers={"payouts/order-7": b'"paid 30"'} + ) + assert await run(caller, "refund", host, 30) == "paid 30" + assert host.attached == [("payouts", "order-7")] + + +async def test_only_a_workflows_main_handler_refused_as_already_invoked_is_rejoined(): + caller = sessions.SessionObject("caller") + + @caller.handler() + async def refund(ctx: sessions.SessionContext, amount: int) -> str: + return await ctx.call(pay, amount, key="order-7") + + @caller.handler() + async def receipt(ctx: sessions.SessionContext) -> str: + return await ctx.call(paid, key="order-7") + + @caller.handler() + async def by_name(ctx: sessions.SessionContext) -> str: + return await ctx.generic_call("payouts", "pay", 30, key="order-7") + + cancelled = TerminalError("cancelled", status_code=409) + host = Host(Backend(), service="caller", refusals={"pay": cancelled}) + with pytest.raises(TerminalError, match="cancelled"): + await run(caller, "refund", host, 30) + host = Host(Backend(), service="caller", refusals={"paid": ALREADY_INVOKED, "pay": ALREADY_INVOKED}) + with pytest.raises(TerminalError, match="already invoked"): + await run(caller, "receipt", host) + # By name, the call can't tell a workflow's main handler from another: it is not rejoined. + with pytest.raises(TerminalError, match="already invoked"): + await run(caller, "by_name", host) + assert host.attached == [] + + async def test_ctx_send_is_made_at_once_and_cancel_takes_an_invocation_id(): sender = sessions.SessionObject("sender") From fe7c57d71add643c21257d104c00abcea94a681b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 16:21:08 +0000 Subject: [PATCH 2/7] Agents CLI: show the user's messages it did not type, a steered one where the turn took it After the conversation so far, the terminal showed none of the user's messages, since it took them to be typed here. So a message sent with `/steer` never showed where the turn took it (between a tool's result and the next step), and messages from other clients didn't show at all. The terminal now notes the messages it sends to `chat`, and hides only those (once each, by text); every other message of the user shows in its place. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/__main__.py | 46 +++++++++++++++--- tests/agents_cli.py | 78 +++++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+), 6 deletions(-) create mode 100644 tests/agents_cli.py diff --git a/python/restate/agents/__main__.py b/python/restate/agents/__main__.py index d1b3e5e5..e0c24105 100644 --- a/python/restate/agents/__main__.py +++ b/python/restate/agents/__main__.py @@ -52,17 +52,35 @@ def _clipped(text: str, limit: int = 160) -> str: class Terminal: - """What the terminal showed of the conversation: each part of a message once, and a tool's output once.""" + """What the terminal showed of the conversation: each part of a message once, and a tool's output once. + + The user's messages are shown too, where they are in the conversation (a + message steered to a turn, say, among its steps), but for those typed here + as messages to ``chat``: the terminal shows them already. + """ def __init__(self, agent: str) -> None: self.agent = agent self.shown: set[str] = set() + self.sent: list[str] = [] + """The messages typed here and sent to ``chat`` that the conversation has not shown yet.""" + self.typed: set[int] = set() + """The ids of the user's messages that were typed here, and so not shown.""" self.history = True - """The first change is the conversation so far, your messages too; after it, you typed them.""" + """The first change is the conversation so far: all of it is shown, the messages typed before too.""" + + def sending(self, text: str) -> None: + """Note a message typed here and sent to ``chat``: when the conversation has it, it isn't shown again.""" + self.sent.append(text) + + def unsent(self, text: str) -> None: + """Forget a message ``sending`` noted, whose send failed.""" + with contextlib.suppress(ValueError): + self.sent.remove(text) def show(self, change: Any) -> None: for message in change.messages: - if message["role"] == "user" and not self.history: + if message["role"] == "user" and not self.history and self._typed(message): continue for i, part in enumerate(message["parts"]): self._part(f"{message['id']}.{i}", message["role"], part) @@ -71,6 +89,17 @@ def show(self, change: Any) -> None: self._once(wait["id"], f"⏸ {wait['name']}: {json.dumps(wait.get('input'))} ({hint})") self.history = False + def _typed(self, message: Any) -> bool: + """Whether the user's ``message`` was typed here: the first sent message of its text, and that one after.""" + if message["id"] in self.typed: + return True + text = "".join(part["text"] for part in message["parts"] if part["type"] == "text") + if text not in self.sent: + return False + self.sent.remove(text) + self.typed.add(message["id"]) + return True + def _part(self, key: str, role: str, part: Part) -> None: if part["type"] == "text": self._once(key, f"{'you' if role == 'user' else self.agent}: {part['text']}") @@ -96,7 +125,7 @@ def _value(words: list[str]) -> object: return text -async def _say(agent: Any, line: str) -> None: +async def _say(agent: Any, terminal: Terminal, line: str) -> None: """Run a command, or send a message to the agent.""" command, *words = line.split() if command == "/steer" and words: @@ -119,7 +148,12 @@ async def _say(agent: Any, line: str) -> None: elif command.startswith("/"): print(USAGE, end="", flush=True) else: - await agent.send(line) + terminal.sending(line) + try: + await agent.send(line) + except BaseException: + terminal.unsent(line) + raise async def chat(target: str, ingress: str | None) -> None: @@ -147,7 +181,7 @@ async def following() -> None: if not line.strip(): continue try: - await _say(agent, line.strip()) + await _say(agent, terminal, line.strip()) except Exception as error: # Restate's answer, shown rather than ending the chat print(f"Restate answered {error}", flush=True) finally: diff --git a/tests/agents_cli.py b/tests/agents_cli.py new file mode 100644 index 00000000..003b560a --- /dev/null +++ b/tests/agents_cli.py @@ -0,0 +1,78 @@ +# +# 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 +# +"""The agents' command line: what its terminal shows of the conversation as it follows it.""" + +# pylint: disable=missing-function-docstring + +from types import SimpleNamespace +from typing import Any + +import pytest + +from restate.agents.__main__ import Terminal + + +def change(*messages: dict[str, Any]) -> Any: + """A change of the conversation, as ``AgentClient.follow`` yields it.""" + return SimpleNamespace(messages=list(messages), waiting=[]) + + +def user(message_id: int, text: str) -> dict[str, Any]: + return {"id": message_id, "role": "user", "parts": [{"type": "text", "text": text}]} + + +def answer(message_id: int, text: str) -> dict[str, Any]: + return {"id": message_id, "role": "assistant", "parts": [{"type": "text", "text": text}]} + + +def tool(message_id: int, name: str, output: str | None) -> dict[str, Any]: + part = {"type": "tool", "id": f"call-{message_id}", "name": name, "input": {}} + if output is not None: + part["output"] = output + return {"id": message_id, "role": "assistant", "parts": [part]} + + +def test_the_conversation_so_far_is_shown_whole(capsys: pytest.CaptureFixture[str]): + terminal = Terminal("support") + terminal.show(change(user(0, "Hi"), answer(1, "Hello"))) + assert capsys.readouterr().out.splitlines() == ["you: Hi", "support: Hello"] + + +def test_a_message_typed_here_is_not_shown_again_but_a_steered_one_is_shown_where_the_turn_took_it( + capsys: pytest.CaptureFixture[str], +): + terminal = Terminal("support") + terminal.show(change()) + terminal.sending("Refund order 103") # typed: the terminal shows it already + terminal.show(change(user(0, "Refund order 103"), tool(1, "approval", None))) + # Steered with /steer: the turn took it after the approval, before the refund. + terminal.show(change(tool(1, "approval", "approved"), user(2, "And order 104?"), tool(3, "refund", "done"))) + terminal.show(change(user(0, "Refund order 103"), answer(4, "Refunded, and 104 ships today."))) + assert capsys.readouterr().out.splitlines() == [ + " → approval({})", + " ← approved", + "you: And order 104?", + " → refund({})", + " ← done", + "support: Refunded, and 104 ships today.", + ] + + +def test_a_message_from_elsewhere_is_shown_and_each_typed_message_is_hidden_once(capsys: pytest.CaptureFixture[str]): + terminal = Terminal("support") + terminal.show(change()) + terminal.sending("again") + terminal.show(change(user(0, "again"), answer(1, "ok"))) + terminal.show(change(user(2, "again"))) # the same text, from another client + terminal.sending("never sent") + terminal.unsent("never sent") # its send failed + terminal.show(change(user(3, "never sent"))) + assert capsys.readouterr().out.splitlines() == ["support: ok", "you: again", "you: never sent"] From df1de24ab3666bf279c524b4245b032e95ef38a1 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 16:39:29 +0000 Subject: [PATCH 3/7] Test: attach_workflow returns the result of the workflow's run A harness test of the new API: a workflow returns "hi", a service sends it a run (ignoring the result), attaches to it with `attach_workflow`, and answers with what the attach returned. Co-Authored-By: Claude Opus 5.5 --- tests/servercontext.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/servercontext.py b/tests/servercontext.py index bb1f53dc..844466b8 100644 --- a/tests/servercontext.py +++ b/tests/servercontext.py @@ -26,6 +26,8 @@ WorkflowContext, ) from restate.serde import DefaultSerde +from restate.server_context import ServerInvocationContext +import json import pytest import typing @@ -191,3 +193,23 @@ async def greet(ctx: Context, name: str) -> str | None: async with simple_harness(greeter) as client: res = await client.service_call(greet, arg="bob") assert res == "hi" + + +async def test_attach_workflow_returns_the_result_of_its_run(): + greeter = Workflow("greeter_workflow") + + @greeter.main() + async def hello(ctx: WorkflowContext) -> str: + return "hi" + + caller = Service("caller") + + @caller.handler() + async def send_and_attach(ctx: Context, workflow_id: str) -> str: + ctx.workflow_send(hello, key=workflow_id, arg=None) # starts the run, ignoring its result + # attach_workflow is the server context's: the run's result, as JSON bytes. + result = await typing.cast(ServerInvocationContext, ctx).attach_workflow("greeter_workflow", workflow_id) + return json.loads(result) + + async with restate.create_test_harness(restate.app([greeter, caller]), restate_image=RESTATE_IMAGE) as harness: + assert await harness.client.service_call(send_and_attach, arg="wf-1") == "hi" From 226910774b0be0b4aca615fdb68f82e308ccf04e Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 16:43:17 +0000 Subject: [PATCH 4/7] Agents CLI: a message typed before the conversation so far came uses up its text there The first change (the conversation so far) shows every message, so it never used up the text of a line typed before it came. A later user message of the same text (steered, or from another client) was then hidden in its place. The first change now uses it up too, while still showing the message. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/__main__.py | 5 ++++- tests/agents_cli.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/restate/agents/__main__.py b/python/restate/agents/__main__.py index e0c24105..ed1a5441 100644 --- a/python/restate/agents/__main__.py +++ b/python/restate/agents/__main__.py @@ -80,7 +80,10 @@ def unsent(self, text: str) -> None: def show(self, change: Any) -> None: for message in change.messages: - if message["role"] == "user" and not self.history and self._typed(message): + # The conversation so far shows every message, but one typed here still uses up its text there, so that a + # later message of the same text (steered, or from another client) isn't hidden in its place. + typed = message["role"] == "user" and self._typed(message) + if typed and not self.history: continue for i, part in enumerate(message["parts"]): self._part(f"{message['id']}.{i}", message["role"], part) diff --git a/tests/agents_cli.py b/tests/agents_cli.py index 003b560a..aae49a0f 100644 --- a/tests/agents_cli.py +++ b/tests/agents_cli.py @@ -76,3 +76,13 @@ def test_a_message_from_elsewhere_is_shown_and_each_typed_message_is_hidden_once terminal.unsent("never sent") # its send failed terminal.show(change(user(3, "never sent"))) assert capsys.readouterr().out.splitlines() == ["support: ok", "you: again", "you: never sent"] + + +def test_a_message_typed_before_the_conversation_so_far_came_is_shown_there_and_hides_no_later_one( + capsys: pytest.CaptureFixture[str], +): + terminal = Terminal("support") + terminal.sending("Hi") # typed before the first change came, which has it already + terminal.show(change(user(0, "Hi"), answer(1, "Hello"))) + terminal.show(change(user(2, "Hi"))) # the same text, from another client: shown + assert capsys.readouterr().out.splitlines() == ["you: Hi", "support: Hello", "you: Hi"] From 64a97379226441a72f42ec9a7e244acb549b82b9 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 16:58:26 +0000 Subject: [PATCH 5/7] Agents: a workflow run again after a crash waits for its run, in the agents' tools rather than in ctx.call The rejoin moves up, out of restate.sessions' `ctx.call`, which is back to what 1.5.0 does: a call to a workflow's main handler that Restate refuses is that error again, for every session object. The agents' `workflow` tool and LangGraph's `nodes.workflow` now run the workflow through `run_workflow`: they call its main handler, and if Restate refuses the call because the run of that id started (409, "already invoked": the tool call run again after a crash), they wait for that run with `attach_workflow` and return its result. A first call stays a call of the attempt, so stopping the turn still cancels it. `restate.sessions.internal.attach_workflow(ctx, name, key)` is how a library reaches the server context's `attach_workflow` from a session object's context. Tests: the helper (fake host); the workflow tool rejoins on "already invoked" and not on another 409; against a real server, a workflow tool call run again after a crash that followed its call answers with the run's result (it answers "Error: the workflow method was already invoked" without the rejoin). The sessions-level rejoin tests are gone with it. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/_calls.py | 5 +- python/restate/agents/_tools.py | 37 +++++++++++- python/restate/agents/langgraph/nodes.py | 6 +- python/restate/sessions/_api.py | 3 - python/restate/sessions/_calls.py | 50 +++------------- python/restate/sessions/_object.py | 16 ++++++ python/restate/sessions/internal.py | 6 +- tests/agents_tools.py | 8 +++ tests/agents_tools_server.py | 73 ++++++++++++++++++++++++ tests/sessions_harness.py | 31 +--------- tests/sessions_object.py | 51 +++-------------- 11 files changed, 160 insertions(+), 126 deletions(-) create mode 100644 tests/agents_tools_server.py diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index 3e517097..2b429318 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -21,8 +21,9 @@ with the same tool and input as one an earlier attempt made reuses that call's id, and so its idempotency key: it is the same Restate call, which it waits for again or whose result it gets. (A workflow's main handler takes no -idempotency key: the call made again waits for the run of its workflow id -instead, as ``restate.sessions``' calls do.) +idempotency key, and Restate refuses the call made again: the ``workflow`` +tool, and LangGraph's ``nodes.workflow``, then wait for the run of its +workflow id instead, ``run_workflow``.) A call is kept until the framework's transcript has its result under the latest id it was made with (``settle`` checks this at each publish and before diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 34de4ca6..f99b8f74 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -45,6 +45,7 @@ from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable from restate.sessions import SessionContext +from restate.sessions.internal import attach_workflow T = typing.TypeVar("T") @@ -326,6 +327,38 @@ async def _once( return result +def already_invoked(error: TerminalError) -> bool: + """Whether ``error`` is Restate refusing a call to a workflow's main handler because that workflow's run started. + + Its status code is a cancellation's too (409): the message tells them apart. + """ + return error.status_code == 409 and "already invoked" in error.message + + +async def run_workflow( + ctx: SessionContext, target: Callable[..., Any], value: Any, workflow_id: str, idempotency_key: str +) -> Any: + """Run the workflow ``target`` (its main handler) with the id ``workflow_id`` to its end, and return its result. + + A workflow runs once per id. Restate takes no idempotency key for its main + handler, and refuses a call to it made again (after a crash, which runs + the tool call again) with 409, "already invoked": the call then waits for + the run of that id instead (``attach_workflow``), whoever started it. + Waiting for the run this way is not a call of the attempt, so stopping the + turn doesn't cancel the run, as it cancels the first call's. + """ + handler = handler_of(target) + try: + return await untyped(ctx).call( + target, *inputs_of(handler, value), key=workflow_id, idempotency_key=idempotency_key + ) + except TerminalError as error: + if not already_invoked(error): + raise + result = await attach_workflow(ctx, handler.service_tag.name, workflow_id) + return handler.handler_io.output_serde.deserialize(result) + + def untyped(ctx: SessionContext) -> Any: """Return ``ctx`` untyped, to call a handler whose input type is known only at run time.""" return ctx @@ -548,9 +581,7 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: workflow_id = f"{c.ctx.key}-{c.call_id}" if key is None else await configured(key, c.ctx) - return await untyped(c.ctx).call( - target, *inputs_of(handler, value), key=workflow_id, idempotency_key=c.idempotency_key - ) + return await run_workflow(c.ctx, target, value, workflow_id, c.idempotency_key) return await _once(service, arguments["raw"], tool_call, made) diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 49c6f0b9..27524bf7 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -60,7 +60,7 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: from pydantic_core import to_jsonable_python from restate.agents import _signals -from restate.agents._tools import handler_of, input_type_of, inputs_of, keyed, untyped +from restate.agents._tools import handler_of, input_type_of, inputs_of, keyed, run_workflow, untyped from restate.agents.langgraph._context import handler_context from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -254,9 +254,7 @@ async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "workflow") workflow_id = key(state) if key is not None else step.idempotency_key value = called.validated(input(state)) if input is not None else None - result = await untyped(step.ctx).call( - target, *inputs_of(called.handler, value), key=workflow_id, idempotency_key=step.idempotency_key - ) + result = await run_workflow(step.ctx, target, value, workflow_id, step.idempotency_key) return _update(output, result) return node diff --git a/python/restate/sessions/_api.py b/python/restate/sessions/_api.py index d15e39f2..e9c12b35 100644 --- a/python/restate/sessions/_api.py +++ b/python/restate/sessions/_api.py @@ -128,9 +128,6 @@ def call( key: The object's or workflow's key; a service takes none. idempotency_key: The same on every attempt, so that a retried attempt gets the first call's result instead of calling again. - A workflow's main handler needs none: it runs once per - workflow id, and a call that Restate refuses because that - run started waits for the run, and returns its result. headers: Headers to send with the call. Raises: diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index e8f0ed6c..2e44ba21 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -18,13 +18,6 @@ count = await ctx.call(increment, 5, key="a", idempotency_key=f"{ctx.request.id}:inc") ctx.send(record, entry, key="audit", delay=timedelta(hours=1)) answer = await ctx.generic_call("Counter", "get", key="a") - -A workflow's main handler runs once per workflow id, and Restate takes no -idempotency key for it: a second call is refused (409, "already invoked"). -A session object runs again from its first line after a crash, so it would -make such a call again. ``call`` makes a call to a workflow's main handler -that is refused this way wait for that run instead, and return its result: -the run of that id, whoever started it. """ from __future__ import annotations @@ -35,7 +28,7 @@ from datetime import timedelta from typing import Any, Protocol -from restate.context import RestateDurableCallFuture, RestateDurableFuture +from restate.context import RestateDurableCallFuture from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable from restate.serde import JsonSerde, Serde @@ -81,11 +74,11 @@ def generic_send( headers: dict[str, str] | None = None, ) -> Any: ... - def attach_workflow(self, name: str, key: str) -> RestateDurableFuture[bytes]: ... + def attach_workflow(self, name: str, key: str) -> Awaitable[bytes]: ... -def _target(handler: Callable[..., Any], key: str | None) -> tuple[Handler[Any, Any], str, str, Serde[Any], Serde[Any]]: - """Return ``handler``'s definition, its service name, handler name and serdes, checking ``key``. +def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Serde[Any], Serde[Any]]: + """Return the service name, handler name and serdes of ``handler``, checking ``key``. Raises ``TerminalError`` (since a retry would fail the same way) if ``handler`` is not a handler, if an object's or workflow's handler has no @@ -102,12 +95,7 @@ def _target(handler: Callable[..., Any], key: str | None) -> tuple[Handler[Any, raise TerminalError(f"{service} is {what}: call it with its key") if not keyed and key is not None: raise TerminalError(f"{service} is a service: call it without a key") - return target, service, target.name, target.handler_io.input_serde, target.handler_io.output_serde - - -def already_invoked(error: TerminalError) -> bool: - """Whether ``error`` is Restate refusing a call to a workflow's main handler because that workflow's run started.""" - return error.status_code == 409 and "already invoked" in error.message + return service, target.name, target.handler_io.input_serde, target.handler_io.output_serde class CallFuture(RestateDurableCallFuture[O]): @@ -115,31 +103,15 @@ class CallFuture(RestateDurableCallFuture[O]): Await it for the result. ``invocation_id()`` returns the id of the call's invocation, and ``cancel_invocation()`` cancels it. - - With ``rejoin`` (a call to a workflow's main handler), a call that Restate - refuses because the workflow's run started waits for that run instead - (``rejoin()``), and returns its result. """ - def __init__( - self, - call: RestateDurableCallFuture[bytes], - serde: Serde[O], - rejoin: Callable[[], Awaitable[bytes]] | None = None, - ) -> None: + def __init__(self, call: RestateDurableCallFuture[bytes], serde: Serde[O]) -> None: super().__init__() self._call: RestateDurableCallFuture[bytes] = call self._serde: Serde[O] = serde - self._rejoin = rejoin async def _result(self) -> O: - try: - result = await self._call - except TerminalError as error: - if self._rejoin is None or not already_invoked(error): - raise - result = await self._rejoin() - return typing.cast(O, self._serde.deserialize(result)) + return typing.cast(O, self._serde.deserialize(await self._call)) def __await__(self) -> Generator[Any, Any, O]: return self._result().__await__() @@ -189,12 +161,8 @@ def call( idempotency_key: str | None = None, headers: dict[str, str] | None = None, ) -> CallFuture[Any]: - target, service, name, input_serde, output_serde = _target(handler, key) + service, name, input_serde, output_serde = _target(handler, key) future = self._call_host.generic_call(service, name, input_serde.serialize(arg), key, idempotency_key, headers) - if target.kind == "workflow" and key is not None: - # The main handler of a workflow: a call made again (after a crash) joins the run it started. - host, workflow_id = self._call_host, key - return CallFuture(future, output_serde, rejoin=lambda: host.attach_workflow(service, workflow_id)) return CallFuture(future, output_serde) def generic_call( @@ -259,7 +227,7 @@ def send( idempotency_key: str | None = None, headers: dict[str, str] | None = None, ) -> None: - _, service, name, input_serde, _ = _target(handler, key) + service, name, input_serde, _ = _target(handler, key) parameter = input_serde.serialize(arg) self._send_sink(GenericSend(service, name, parameter, key, delay, idempotency_key, headers)) diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 28a1d670..964f714d 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -339,6 +339,22 @@ async def cancel_calls(ctx: SessionContext) -> list[str]: return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls")) +def attach_workflow(ctx: SessionContext | ViewContext, name: str, key: str) -> Awaitable[bytes]: + """Wait for the run of the workflow ``name`` with the id ``key``, and return its result, serialized. + + For libraries that run workflows from session objects. Restate takes no + idempotency key for a workflow's main handler, and refuses a call to it + made again, after a crash, with 409 ("the workflow method was already + invoked"): this waits for the run the first call started instead. + + Raises: + TypeError: If ``ctx`` is not the context of a session object's handler or view. + """ + if not isinstance(ctx, (AttemptContext, AttemptViewContext)): + raise TypeError(f"{ctx!r} is not the context of a session object's handler or view") + return ctx._call_host.attach_workflow(name, key) # pyright: ignore[reportPrivateUsage] (this module's own class) + + def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None: """Return the session object whose handler or view ``ctx`` belongs to, or ``None``. diff --git a/python/restate/sessions/internal.py b/python/restate/sessions/internal.py index 7c1784bb..95042655 100644 --- a/python/restate/sessions/internal.py +++ b/python/restate/sessions/internal.py @@ -19,13 +19,16 @@ - ``reset``: empty the key's database inside a transaction. - ``cancelled`` and ``cancel_calls``: explicit cancellation, for handlers declared with ``explicit_cancellation=True``. +- ``attach_workflow``: wait for a workflow's run by its id, for a call to its + main handler that Restate refused as already invoked (made again after a + crash). - ``ScopeKey`` and ``definition_of``: state a library keeps per attempt, and the object an attempt belongs to. - ``invoke``, ``SqlHost`` and ``ObjectConfig``: run one attempt of a handler against a host of your own, without Restate, for tests. """ -from restate.sessions._object import cancel_calls, cancelled, definition_of, invoke +from restate.sessions._object import attach_workflow, cancel_calls, cancelled, definition_of, invoke from restate.sessions._schema import Schema from restate.sessions._scope import ScopeKey from restate.sessions._session import ObjectConfig, SqlHost, database, reset @@ -35,6 +38,7 @@ "Schema", "ScopeKey", "SqlHost", + "attach_workflow", "cancel_calls", "cancelled", "database", diff --git a/tests/agents_tools.py b/tests/agents_tools.py index ddc70f4a..48ffea53 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -139,6 +139,14 @@ async def test_workflow_made_again_after_a_crash_waits_for_the_run_it_started(): assert host.attached == [("payouts", "alice-c1")] +async def test_workflow_refused_otherwise_answers_with_the_error_and_waits_for_no_run(): + tools["pay"] = workflow(pay, description="Pay an amount out.") + cancelled = restate.TerminalError("cancelled", status_code=409) # a 409 that is no "already invoked" + host = Host(Backend(), service="runner", key="alice", refusals={"pay": cancelled}) + assert await run_tool(host, "pay", {"input": 30}, "c1") == "Error: cancelled" + assert host.attached == [] + + # --- signal and subagent -------------------------------------------------------------------------- diff --git a/tests/agents_tools_server.py b/tests/agents_tools_server.py new file mode 100644 index 00000000..1420d35f --- /dev/null +++ b/tests/agents_tools_server.py @@ -0,0 +1,73 @@ +# +# 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 +# +"""The agents' Restate tools against a real Restate server: a workflow tool call run again after a crash.""" + +# pylint: disable=missing-function-docstring,import-error + +import uuid +from datetime import timedelta + +import pytest +from support import RESTATE_IMAGE + +import restate +from restate import HarnessEnvironment, sessions +from restate.agents._object import AgentObject +from restate.agents._tools import ToolCall, workflow + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +payouts = restate.Workflow("tools_payouts") + + +@payouts.main() +async def pay(ctx: restate.WorkflowContext, amount: int) -> str: + await ctx.sleep(timedelta(milliseconds=100)) + return f"paid {amount} for {ctx.key()}" + + +pay_tool = workflow(pay, description="Pay an amount out.") + +runner = AgentObject("tools_runner", migrations=["CREATE TABLE crashed (call TEXT PRIMARY KEY)"]) + + +@runner.handler +async def refund(ctx: sessions.SessionContext, amount: int) -> str: + """Runs the workflow tool as tool call ``c1``. + + The first attempt starts the tool call's workflow, as the tool does, and crashes before its result. The retry runs + the tool call, whose call to the workflow Restate refuses (409, already invoked): it waits for that run instead. + """ + if await ctx.db.get("SELECT 1 FROM crashed WHERE call = 'c1'") is None: + started = ctx.call(pay, amount, key=f"{ctx.key}-c1") + await started.invocation_id() # the run started + await ctx.db.execute("INSERT INTO crashed (call) VALUES ('c1')") + raise ValueError("crash after the workflow's run started") + return await pay_tool.run({"input": amount}, ToolCall(ctx, "c1")) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness(restate.app([payouts, runner]), restate_image=RESTATE_IMAGE) as harness: + yield harness + + +async def test_a_workflow_tool_call_run_again_after_a_crash_answers_with_the_run_it_started( + harness: HarnessEnvironment, +): + key = str(uuid.uuid4()) + assert await harness.client.object_call(refund, key=key, arg=30) == f"paid 30 for {key}-c1" diff --git a/tests/sessions_harness.py b/tests/sessions_harness.py index 654c49a4..2678f068 100644 --- a/tests/sessions_harness.py +++ b/tests/sessions_harness.py @@ -156,32 +156,10 @@ async def stoppable(ctx: sessions.SessionContext, text: str) -> list[str]: return cancelled -payouts = restate.Workflow("payouts") - - -@payouts.main() -async def pay(ctx: restate.WorkflowContext, amount: int) -> str: - await ctx.sleep(timedelta(milliseconds=100)) - return f"paid {amount} for {ctx.key()}" - - -@notes.handler() -async def refund(ctx: sessions.SessionContext, order: str) -> str: - """Pays ``order`` out with a workflow, and crashes on its first attempt once the call is made: the retry calls the - workflow again, which Restate refuses (409, already invoked), and so waits for the run the first attempt started.""" - payout = ctx.call(pay, 30, key=order) - await payout.invocation_id() # the call is made - seen = await ctx.db.get("SELECT count(*) AS n FROM notes WHERE text = ?", [order]) - if seen is not None and seen["n"] == 0: - await ctx.db.execute("INSERT INTO notes (text) VALUES (?)", [order]) - raise ValueError("crash after the call was made") - return await payout - - @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([inbox, pricing, sleeper, payouts, notes]), restate_image=RESTATE_IMAGE + restate.app([inbox, pricing, sleeper, notes]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -203,13 +181,6 @@ async def test_a_retry_runs_the_handler_from_the_start_against_what_was_committe assert await harness.client.object_call(texts, key=key, arg=None) == ["once"] -async def test_a_call_to_a_workflow_made_again_after_a_crash_returns_the_run_it_started(harness: HarnessEnvironment): - key, order = str(uuid.uuid4()), f"order-{uuid.uuid4()}" - assert await harness.client.object_call(refund, key=key, arg=order) == f"paid 30 for {order}" - # Another key's call for the same workflow id gets that run's result too: a workflow runs once per id. - assert await harness.client.object_call(refund, key=str(uuid.uuid4()), arg=order) == f"paid 30 for {order}" - - async def test_calls_and_sql_are_awaited_together(harness: HarnessEnvironment): key = str(uuid.uuid4()) await harness.client.object_call(add, key=key, arg="x") diff --git a/tests/sessions_object.py b/tests/sessions_object.py index dd7102e7..edb5d6a0 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -14,6 +14,7 @@ import asyncio import itertools +import json import typing from dataclasses import dataclass, field from typing import Any @@ -397,56 +398,22 @@ async def pay(ctx: restate.WorkflowContext, amount: int) -> str: return f"paid {amount}" -@payouts.handler() -async def paid(ctx: restate.WorkflowSharedContext) -> str: - return "paid" - - ALREADY_INVOKED = TerminalError("the workflow method was already invoked", status_code=409) -async def test_a_call_to_a_workflow_refused_as_already_invoked_waits_for_that_run(): - """A session object runs again from its first line, so it calls a workflow again after a crash: Restate refuses the - second call to its main handler (it takes no idempotency key), and the call waits for the run it started.""" +async def test_attach_workflow_waits_for_a_workflows_run_by_its_id(): + """For libraries: the run of a workflow, by its id, which a call to its main handler made again can't get.""" caller = sessions.SessionObject("caller") @caller.handler() - async def refund(ctx: sessions.SessionContext, amount: int) -> str: - return await ctx.call(pay, amount, key="order-7") + async def receipt(ctx: sessions.SessionContext, order: str) -> str: + return json.loads(await internal.attach_workflow(ctx, "payouts", order)) - host = Host( - Backend(), service="caller", refusals={"pay": ALREADY_INVOKED}, answers={"payouts/order-7": b'"paid 30"'} - ) - assert await run(caller, "refund", host, 30) == "paid 30" + host = Host(Backend(), service="caller", answers={"payouts/order-7": b'"paid 30"'}) + assert await run(caller, "receipt", host, "order-7") == "paid 30" assert host.attached == [("payouts", "order-7")] - - -async def test_only_a_workflows_main_handler_refused_as_already_invoked_is_rejoined(): - caller = sessions.SessionObject("caller") - - @caller.handler() - async def refund(ctx: sessions.SessionContext, amount: int) -> str: - return await ctx.call(pay, amount, key="order-7") - - @caller.handler() - async def receipt(ctx: sessions.SessionContext) -> str: - return await ctx.call(paid, key="order-7") - - @caller.handler() - async def by_name(ctx: sessions.SessionContext) -> str: - return await ctx.generic_call("payouts", "pay", 30, key="order-7") - - cancelled = TerminalError("cancelled", status_code=409) - host = Host(Backend(), service="caller", refusals={"pay": cancelled}) - with pytest.raises(TerminalError, match="cancelled"): - await run(caller, "refund", host, 30) - host = Host(Backend(), service="caller", refusals={"paid": ALREADY_INVOKED, "pay": ALREADY_INVOKED}) - with pytest.raises(TerminalError, match="already invoked"): - await run(caller, "receipt", host) - # By name, the call can't tell a workflow's main handler from another: it is not rejoined. - with pytest.raises(TerminalError, match="already invoked"): - await run(caller, "by_name", host) - assert host.attached == [] + with pytest.raises(TypeError, match="is not the context of a session object"): + internal.attach_workflow(typing.cast(sessions.SessionContext, object()), "payouts", "order-7") async def test_ctx_send_is_made_at_once_and_cancel_takes_an_invocation_id(): From ac2a20d71810e21cef1e772efb63f2d6d262d6d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 17:07:38 +0000 Subject: [PATCH 6/7] Agents: the workflow tool sends to the workflow, then attaches to the send's invocation A workflow runs once per id, and Restate takes no idempotency key for its main handler. The agents' `workflow` tool, and LangGraph's `nodes.workflow`, now always send to the workflow, keep the send's invocation id, and attach to that invocation for the result (`run_workflow`). A send to a workflow whose run started already starts nothing, and its invocation id is that run's: a tool call made again after a crash, or a `key` that ran before, waits for that run. No call is refused, and no error is matched. The run is the workflow's own: stopping the turn doesn't cancel it. `restate.sessions.internal` gains `start` (a send right away that returns its invocation id) and `attach_invocation`, for libraries. The session-context `attach_workflow` helper of the previous commit is gone; the SDK's `attach_workflow` stays. Tests: against a real server, two sends to one workflow id give the same invocation id, which attach returns "hi" for; and the workflow tool run again after a crash that followed its send answers "hi". The tool and the helpers, against the fake host. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/_calls.py | 6 +-- python/restate/agents/_tools.py | 39 ++++----------- python/restate/agents/langgraph/nodes.py | 2 +- python/restate/sessions/_calls.py | 8 ++- python/restate/sessions/_object.py | 41 ++++++++++----- python/restate/sessions/internal.py | 11 ++-- tests/agents_tools.py | 32 ++++-------- tests/agents_tools_server.py | 64 +++++++++--------------- tests/servercontext.py | 20 ++++++++ tests/sessions_object.py | 59 +++++++++------------- 10 files changed, 133 insertions(+), 149 deletions(-) diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index 2b429318..c10369a0 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -21,9 +21,9 @@ with the same tool and input as one an earlier attempt made reuses that call's id, and so its idempotency key: it is the same Restate call, which it waits for again or whose result it gets. (A workflow's main handler takes no -idempotency key, and Restate refuses the call made again: the ``workflow`` -tool, and LangGraph's ``nodes.workflow``, then wait for the run of its -workflow id instead, ``run_workflow``.) +idempotency key: the ``workflow`` tool, and LangGraph's ``nodes.workflow``, +send to it, which starts its run once per workflow id, and wait for that run, +``run_workflow``.) A call is kept until the framework's transcript has its result under the latest id it was made with (``settle`` checks this at each publish and before diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index f99b8f74..7cd100c5 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -45,7 +45,7 @@ from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable from restate.sessions import SessionContext -from restate.sessions.internal import attach_workflow +from restate.sessions.internal import attach_invocation, start T = typing.TypeVar("T") @@ -327,36 +327,19 @@ async def _once( return result -def already_invoked(error: TerminalError) -> bool: - """Whether ``error`` is Restate refusing a call to a workflow's main handler because that workflow's run started. - - Its status code is a cancellation's too (409): the message tells them apart. - """ - return error.status_code == 409 and "already invoked" in error.message - - -async def run_workflow( - ctx: SessionContext, target: Callable[..., Any], value: Any, workflow_id: str, idempotency_key: str -) -> Any: +async def run_workflow(ctx: SessionContext, target: Callable[..., Any], value: Any, workflow_id: str) -> Any: """Run the workflow ``target`` (its main handler) with the id ``workflow_id`` to its end, and return its result. - A workflow runs once per id. Restate takes no idempotency key for its main - handler, and refuses a call to it made again (after a crash, which runs - the tool call again) with 409, "already invoked": the call then waits for - the run of that id instead (``attach_workflow``), whoever started it. - Waiting for the run this way is not a call of the attempt, so stopping the - turn doesn't cancel the run, as it cancels the first call's. + It sends to the workflow, and waits for the invocation the send started. + A workflow runs once per id: a send to one whose run started already + starts nothing, and its invocation id is that run's. So a tool call made + again after a crash, or a ``key`` that ran before, waits for that run. The + run is the workflow's own, not a call of the turn: stopping the turn + doesn't cancel it. """ handler = handler_of(target) - try: - return await untyped(ctx).call( - target, *inputs_of(handler, value), key=workflow_id, idempotency_key=idempotency_key - ) - except TerminalError as error: - if not already_invoked(error): - raise - result = await attach_workflow(ctx, handler.service_tag.name, workflow_id) - return handler.handler_io.output_serde.deserialize(result) + invocation_id = await start(ctx, target, *inputs_of(handler, value), key=workflow_id) + return handler.handler_io.output_serde.deserialize(await attach_invocation(ctx, invocation_id)) def untyped(ctx: SessionContext) -> Any: @@ -581,7 +564,7 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: workflow_id = f"{c.ctx.key}-{c.call_id}" if key is None else await configured(key, c.ctx) - return await run_workflow(c.ctx, target, value, workflow_id, c.idempotency_key) + return await run_workflow(c.ctx, target, value, workflow_id) return await _once(service, arguments["raw"], tool_call, made) diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 27524bf7..ac81d768 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -254,7 +254,7 @@ async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "workflow") workflow_id = key(state) if key is not None else step.idempotency_key value = called.validated(input(state)) if input is not None else None - result = await run_workflow(step.ctx, target, value, workflow_id, step.idempotency_key) + result = await run_workflow(step.ctx, target, value, workflow_id) return _update(output, result) return node diff --git a/python/restate/sessions/_calls.py b/python/restate/sessions/_calls.py index 2e44ba21..0ca72051 100644 --- a/python/restate/sessions/_calls.py +++ b/python/restate/sessions/_calls.py @@ -74,7 +74,7 @@ def generic_send( headers: dict[str, str] | None = None, ) -> Any: ... - def attach_workflow(self, name: str, key: str) -> Awaitable[bytes]: ... + def attach_invocation(self, invocation_id: str) -> Awaitable[bytes]: ... def _target(handler: Callable[..., Any], key: str | None) -> tuple[str, str, Serde[Any], Serde[Any]]: @@ -247,6 +247,12 @@ def generic_send( self._send_sink(GenericSend(service, handler, parameter, key, delay, idempotency_key, headers)) +def send_to(host: CallHost, handler: Callable[..., Awaitable[Any]], arg: Any, key: str | None) -> Any: + """Send ``arg`` to ``handler`` right away, as ``ctx.send`` does, and return the SDK's handle of the send.""" + service, name, input_serde, _ = _target(handler, key) + return host.generic_send(service, name, input_serde.serialize(arg), key) + + def send_now(host: CallHost) -> Callable[[GenericSend], None]: """Return the sink of ``ctx.send``, which sends each send right away.""" diff --git a/python/restate/sessions/_object.py b/python/restate/sessions/_object.py index 964f714d..3ade0dc8 100644 --- a/python/restate/sessions/_object.py +++ b/python/restate/sessions/_object.py @@ -48,7 +48,7 @@ async def list_notes(ctx: sessions.ViewContext) -> list[str]: from restate.serde import BytesSerde, DefaultSerde, Serde from restate.server_context import DisconnectedException, ServerInvocationContext from restate.sessions._api import Database, DatabaseReader, SessionContext, ViewContext -from restate.sessions._calls import CallHost, Calls, GenericSend, Sends, send_now +from restate.sessions._calls import CallHost, Calls, GenericSend, Sends, send_now, send_to from restate.sessions._link import Link from restate.sessions._schema import Migration, Schema, steps_of from restate.sessions._scope import SCOPE, Scope @@ -167,8 +167,8 @@ def generic_call( ) -> Any: return self.context.generic_call(service, handler, arg, key, idempotency_key, headers) - def attach_workflow(self, name: str, key: str) -> Any: - return self.context.attach_workflow(name, key) + def attach_invocation(self, invocation_id: str) -> Awaitable[bytes]: + return self.context.attach_invocation(invocation_id, serde=BytesSerde()) def generic_send( self, @@ -339,20 +339,35 @@ async def cancel_calls(ctx: SessionContext) -> list[str]: return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls")) -def attach_workflow(ctx: SessionContext | ViewContext, name: str, key: str) -> Awaitable[bytes]: - """Wait for the run of the workflow ``name`` with the id ``key``, and return its result, serialized. +def _host_of(ctx: SessionContext) -> CallHost: + if not isinstance(ctx, AttemptContext): + raise TypeError(f"{ctx!r} is not the context of a session object's handler") + return ctx._call_host # pyright: ignore[reportPrivateUsage] (this module's own class) - For libraries that run workflows from session objects. Restate takes no - idempotency key for a workflow's main handler, and refuses a call to it - made again, after a crash, with 409 ("the workflow method was already - invoked"): this waits for the run the first call started instead. + +async def start( + ctx: SessionContext, handler: Callable[..., Awaitable[Any]], arg: Any = None, *, key: str | None = None +) -> str: + """Start an invocation of ``handler`` with ``arg`` by a send, made right away as ``ctx.send`` does, and return its id. + + For libraries that wait for what they sent (``attach_invocation``). A + send to a workflow's main handler whose run of that id started already + starts nothing: its invocation id is that run's. + + Raises: + TypeError: If ``ctx`` is not the context of a session object's handler. + """ + sent = send_to(_host_of(ctx), handler, arg, key) + return typing.cast(str, await sent.invocation_id()) + + +def attach_invocation(ctx: SessionContext, invocation_id: str) -> Awaitable[bytes]: + """Wait for the invocation ``invocation_id`` to end, and return its result, serialized. Raises: - TypeError: If ``ctx`` is not the context of a session object's handler or view. + TypeError: If ``ctx`` is not the context of a session object's handler. """ - if not isinstance(ctx, (AttemptContext, AttemptViewContext)): - raise TypeError(f"{ctx!r} is not the context of a session object's handler or view") - return ctx._call_host.attach_workflow(name, key) # pyright: ignore[reportPrivateUsage] (this module's own class) + return _host_of(ctx).attach_invocation(invocation_id) def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None: diff --git a/python/restate/sessions/internal.py b/python/restate/sessions/internal.py index 95042655..c6b18dd2 100644 --- a/python/restate/sessions/internal.py +++ b/python/restate/sessions/internal.py @@ -19,16 +19,16 @@ - ``reset``: empty the key's database inside a transaction. - ``cancelled`` and ``cancel_calls``: explicit cancellation, for handlers declared with ``explicit_cancellation=True``. -- ``attach_workflow``: wait for a workflow's run by its id, for a call to its - main handler that Restate refused as already invoked (made again after a - crash). +- ``start`` and ``attach_invocation``: send, keeping the invocation id, + and wait for that invocation's result; a library runs a workflow this way + (a send to it again, after a crash, starts nothing: the same run). - ``ScopeKey`` and ``definition_of``: state a library keeps per attempt, and the object an attempt belongs to. - ``invoke``, ``SqlHost`` and ``ObjectConfig``: run one attempt of a handler against a host of your own, without Restate, for tests. """ -from restate.sessions._object import attach_workflow, cancel_calls, cancelled, definition_of, invoke +from restate.sessions._object import attach_invocation, cancel_calls, cancelled, definition_of, invoke, start from restate.sessions._schema import Schema from restate.sessions._scope import ScopeKey from restate.sessions._session import ObjectConfig, SqlHost, database, reset @@ -38,11 +38,12 @@ "Schema", "ScopeKey", "SqlHost", - "attach_workflow", + "attach_invocation", "cancel_calls", "cancelled", "database", "definition_of", "invoke", "reset", + "start", ] diff --git a/tests/agents_tools.py b/tests/agents_tools.py index 48ffea53..a0d73585 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -27,7 +27,7 @@ from restate.agents._tools import Tool, ToolCall, call, subagent, tool, wait_for_signal, workflow from restate.context import Request from sessions_database import Backend # pylint: disable=import-error -from sessions_object import ALREADY_INVOKED, Host, pay, request # pylint: disable=import-error +from sessions_object import Host, pay, request # pylint: disable=import-error @pytest.fixture @@ -124,27 +124,17 @@ async def test_call_lets_the_model_choose_the_key_and_makes_the_call_once_per_to # --- workflow ------------------------------------------------------------------------------------- -async def test_workflow_made_again_after_a_crash_waits_for_the_run_it_started(): - """A crash after the tool call started its workflow runs the tool call again: Restate refuses the second call to - the workflow's main handler (409, already invoked), and the tool answers with that run's result, not the error.""" +async def test_workflow_sends_to_the_workflow_and_waits_for_the_run_the_send_started(): + """The tool sends to the workflow, and attaches to the invocation of the send. Run again (after a crash), it sends + again, which Restate starts nothing for (the same run, of the same id), and waits for that run again.""" tools["pay"] = workflow(pay, description="Pay an amount out.") - host = Host( - Backend(), - service="runner", - key="alice", - refusals={"pay": ALREADY_INVOKED}, - answers={"payouts/alice-c1": b'"paid 30"'}, - ) - assert await run_tool(host, "pay", {"input": 30}, "c1") == "paid 30" - assert host.attached == [("payouts", "alice-c1")] - - -async def test_workflow_refused_otherwise_answers_with_the_error_and_waits_for_no_run(): - tools["pay"] = workflow(pay, description="Pay an amount out.") - cancelled = restate.TerminalError("cancelled", status_code=409) # a 409 that is no "already invoked" - host = Host(Backend(), service="runner", key="alice", refusals={"pay": cancelled}) - assert await run_tool(host, "pay", {"input": 30}, "c1") == "Error: cancelled" - assert host.attached == [] + host = Host(Backend(), service="runner", key="alice", answers={"inv_payouts/alice-c1": b'"paid 30"'}) + req = request() + assert await run_tool(host, "pay", {"input": 30}, "c1", req) == "paid 30" + assert await run_tool(host, "pay", {"input": 30}, "c1", req) == "paid 30" + assert [(s.service, s.handler, s.arg, s.key) for s in host.sends] == [("payouts", "pay", b"30", "alice-c1")] * 2 + assert host.attached == ["inv_payouts/alice-c1"] * 2 + assert host.calls == [] # --- signal and subagent -------------------------------------------------------------------------- diff --git a/tests/agents_tools_server.py b/tests/agents_tools_server.py index 1420d35f..a270d71e 100644 --- a/tests/agents_tools_server.py +++ b/tests/agents_tools_server.py @@ -8,20 +8,18 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""The agents' Restate tools against a real Restate server: a workflow tool call run again after a crash.""" +"""The agents' Restate tools against a real Restate server: the workflow tool sends, then attaches.""" # pylint: disable=missing-function-docstring,import-error -import uuid -from datetime import timedelta - import pytest from support import RESTATE_IMAGE import restate -from restate import HarnessEnvironment, sessions +from restate import sessions from restate.agents._object import AgentObject from restate.agents._tools import ToolCall, workflow +from restate.sessions import internal @pytest.fixture(scope="session") @@ -29,45 +27,29 @@ def anyio_backend(): return "asyncio" -pytestmark = [pytest.mark.anyio] - -payouts = restate.Workflow("tools_payouts") - - -@payouts.main() -async def pay(ctx: restate.WorkflowContext, amount: int) -> str: - await ctx.sleep(timedelta(milliseconds=100)) - return f"paid {amount} for {ctx.key()}" - +pytestmark = [pytest.mark.anyio, pytest.mark.server] -pay_tool = workflow(pay, description="Pay an amount out.") -runner = AgentObject("tools_runner", migrations=["CREATE TABLE crashed (call TEXT PRIMARY KEY)"]) +async def test_the_workflow_tool_answers_with_its_run_run_again_after_a_crash_too(): + greeter = restate.Workflow("tools_greeter") + @greeter.main() + async def hello(ctx: restate.WorkflowContext) -> str: + return "hi" -@runner.handler -async def refund(ctx: sessions.SessionContext, amount: int) -> str: - """Runs the workflow tool as tool call ``c1``. - - The first attempt starts the tool call's workflow, as the tool does, and crashes before its result. The retry runs - the tool call, whose call to the workflow Restate refuses (409, already invoked): it waits for that run instead. - """ - if await ctx.db.get("SELECT 1 FROM crashed WHERE call = 'c1'") is None: - started = ctx.call(pay, amount, key=f"{ctx.key}-c1") - await started.invocation_id() # the run started - await ctx.db.execute("INSERT INTO crashed (call) VALUES ('c1')") - raise ValueError("crash after the workflow's run started") - return await pay_tool.run({"input": amount}, ToolCall(ctx, "c1")) - - -@pytest.fixture(scope="session") -async def harness(): - async with restate.create_test_harness(restate.app([payouts, runner]), restate_image=RESTATE_IMAGE) as harness: - yield harness + say_hi = workflow(hello, description="Say hi.") + runner = AgentObject("tools_runner", migrations=["CREATE TABLE crashed (call TEXT PRIMARY KEY)"]) + @runner.handler + async def greet(ctx: sessions.SessionContext, call_id: str) -> str: + """Runs the workflow tool as tool call ``call_id``. Its first attempt starts the tool call's workflow, as the + tool does, and crashes before its result: the retry runs the tool call, whose send starts nothing (the run + started), and which waits for that run.""" + if await ctx.db.get("SELECT 1 FROM crashed WHERE call = ?", [call_id]) is None: + await internal.start(ctx, hello, key=f"{ctx.key}-{call_id}") + await ctx.db.execute("INSERT INTO crashed (call) VALUES (?)", [call_id]) + raise ValueError("crash after the workflow's run started") + return await say_hi.run({}, ToolCall(ctx, call_id)) -async def test_a_workflow_tool_call_run_again_after_a_crash_answers_with_the_run_it_started( - harness: HarnessEnvironment, -): - key = str(uuid.uuid4()) - assert await harness.client.object_call(refund, key=key, arg=30) == f"paid 30 for {key}-c1" + async with restate.create_test_harness(restate.app([greeter, runner]), restate_image=RESTATE_IMAGE) as harness: + assert await harness.client.object_call(greet, key="alice", arg="c1") == "hi" diff --git a/tests/servercontext.py b/tests/servercontext.py index 844466b8..5c652e5c 100644 --- a/tests/servercontext.py +++ b/tests/servercontext.py @@ -213,3 +213,23 @@ async def send_and_attach(ctx: Context, workflow_id: str) -> str: async with restate.create_test_harness(restate.app([greeter, caller]), restate_image=RESTATE_IMAGE) as harness: assert await harness.client.service_call(send_and_attach, arg="wf-1") == "hi" + + +async def test_a_send_to_a_workflow_gives_the_invocation_id_that_attach_waits_for(): + greeter = Workflow("greeter_workflow_send") + + @greeter.main() + async def hello(ctx: WorkflowContext) -> str: + return "hi" + + caller = Service("caller_send") + + @caller.handler() + async def send_and_attach(ctx: Context, workflow_id: str) -> list[str]: + first = await ctx.workflow_send(hello, key=workflow_id, arg=None).invocation_id() # starts the run + again = await ctx.workflow_send(hello, key=workflow_id, arg=None).invocation_id() # starts nothing: same run + assert again == first + return [await ctx.attach_invocation(again, type_hint=str), await ctx.attach_invocation(first, type_hint=str)] + + async with restate.create_test_harness(restate.app([greeter, caller]), restate_image=RESTATE_IMAGE) as harness: + assert await harness.client.service_call(send_and_attach, arg="wf-1") == ["hi", "hi"] diff --git a/tests/sessions_object.py b/tests/sessions_object.py index edb5d6a0..756c52da 100644 --- a/tests/sessions_object.py +++ b/tests/sessions_object.py @@ -85,24 +85,14 @@ async def cancel_invocation(self) -> None: pass -class Refused(restate.RestateDurableCallFuture[bytes]): - """A call that Restate refused with ``error``.""" +class Sent: + """The SDK's handle of a send: its invocation id is made of its target.""" - def __init__(self, error: TerminalError) -> None: - super().__init__() - self.error = error - - def __await__(self): - async def refused() -> bytes: - raise self.error - - return refused().__await__() + def __init__(self, invocation_id: str) -> None: + self.id = invocation_id async def invocation_id(self) -> str: - return "inv_refused" - - async def cancel_invocation(self) -> None: - pass + return self.id @dataclass @@ -116,10 +106,8 @@ class Host: sends: list[Send] = field(default_factory=list) calls: list[tuple[str, str, bytes, str | None, str | None]] = field(default_factory=list) answers: dict[str, bytes] = field(default_factory=dict) - # The calls Restate refuses, by handler: the error they fail with. - refusals: dict[str, TerminalError] = field(default_factory=dict) - # The workflow runs attached to, as (workflow, id); their results are the answers of "/". - attached: list[tuple[str, str]] = field(default_factory=list) + # The invocations attached to, by id; their results are the answers of their ids. + attached: list[str] = field(default_factory=list) cancelled: list[str] = field(default_factory=list) # How many more syncs go through before the process crashes: every sync after fails, losing its writes. _syncs_left: int | None = None @@ -152,16 +140,15 @@ async def sync(self) -> None: def generic_call(self, service, handler, arg, key=None, idempotency_key=None, headers=None): self.calls.append((service, handler, arg, key, idempotency_key)) - if handler in self.refusals: - return Refused(self.refusals[handler]) return Answer(self.answers[handler]) - def attach_workflow(self, name: str, key: str): - self.attached.append((name, key)) - return Answer(self.answers[f"{name}/{key}"]) + def attach_invocation(self, invocation_id: str): + self.attached.append(invocation_id) + return Answer(self.answers[invocation_id]) def generic_send(self, service, handler, arg, key=None, send_delay=None, idempotency_key=None, headers=None): self.sends.append(Send(service, handler, arg, key, send_delay, idempotency_key)) + return Sent(f"inv_{service}/{key}") def cancel_invocation(self, invocation_id: str) -> None: self.cancelled.append(invocation_id) @@ -398,22 +385,22 @@ async def pay(ctx: restate.WorkflowContext, amount: int) -> str: return f"paid {amount}" -ALREADY_INVOKED = TerminalError("the workflow method was already invoked", status_code=409) - - -async def test_attach_workflow_waits_for_a_workflows_run_by_its_id(): - """For libraries: the run of a workflow, by its id, which a call to its main handler made again can't get.""" +async def test_start_sends_and_returns_the_invocation_id_that_attach_invocation_waits_for(): + """For libraries: start an invocation by a send, and wait for its result by its id.""" caller = sessions.SessionObject("caller") @caller.handler() - async def receipt(ctx: sessions.SessionContext, order: str) -> str: - return json.loads(await internal.attach_workflow(ctx, "payouts", order)) - - host = Host(Backend(), service="caller", answers={"payouts/order-7": b'"paid 30"'}) - assert await run(caller, "receipt", host, "order-7") == "paid 30" - assert host.attached == [("payouts", "order-7")] + async def refund(ctx: sessions.SessionContext, amount: int) -> str: + invocation_id = await internal.start(ctx, pay, amount, key="order-7") + return json.loads(await internal.attach_invocation(ctx, invocation_id)) + + host = Host(Backend(), service="caller", answers={"inv_payouts/order-7": b'"paid 30"'}) + assert await run(caller, "refund", host, 30) == "paid 30" + [send] = host.sends + assert (send.service, send.handler, send.arg, send.key) == ("payouts", "pay", b"30", "order-7") + assert host.attached == ["inv_payouts/order-7"] with pytest.raises(TypeError, match="is not the context of a session object"): - internal.attach_workflow(typing.cast(sessions.SessionContext, object()), "payouts", "order-7") + await internal.start(typing.cast(sessions.SessionContext, object()), pay, 30, key="order-7") async def test_ctx_send_is_made_at_once_and_cancel_takes_an_invocation_id(): From bc6d1ee512310a90732a91ccc1910d335999aefb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 17:09:59 +0000 Subject: [PATCH 7/7] Remove attach_workflow: the send's invocation id is what the workflow tool attaches to The workflow tool gets the run's invocation id from its send, and attaches to it with `attach_invocation`, so attaching by workflow id is not needed: the binding (`PyVM.attach_workflow`), `VMWrapper.attach_workflow`, `ServerInvocationContext.attach_workflow` and its harness test go. This PR changes no Rust. Co-Authored-By: Claude Opus 5.5 --- python/restate/server_context.py | 11 ----------- python/restate/vm.py | 6 ------ src/lib.rs | 19 ------------------- tests/servercontext.py | 22 ---------------------- 4 files changed, 58 deletions(-) diff --git a/python/restate/server_context.py b/python/restate/server_context.py index 809d295e..aa22cfc8 100644 --- a/python/restate/server_context.py +++ b/python/restate/server_context.py @@ -1319,17 +1319,6 @@ def attach_invocation( update_restate_context_is_replaying(self.vm) return self.create_future(handle, serde) - def attach_workflow(self, name: str, key: str) -> RestateDurableFuture[bytes]: - """Attach to the run of the workflow ``name`` with the id ``key``: its result, as bytes, once it ends. - - For a call to a workflow's main handler that Restate refused because the - run had started (a call made again after a crash, from a handler that - runs again from its first line rather than replaying its journal). - """ - handle = self.vm.attach_workflow(name, key) - update_restate_context_is_replaying(self.vm) - return self.create_future(handle, BytesSerde()) - def get_extension_data(self, key: str) -> Any: """Get extension data by key.""" return self.extension_data.get(key) diff --git a/python/restate/vm.py b/python/restate/vm.py index 3956708a..85be9830 100644 --- a/python/restate/vm.py +++ b/python/restate/vm.py @@ -574,12 +574,6 @@ def attach_invocation(self, invocation_id: str) -> int: """ return self.vm.attach_invocation(invocation_id) - def attach_workflow(self, name: str, key: str) -> int: - """ - Attach to the run of the workflow ``name`` with the id ``key`` - """ - return self.vm.attach_workflow(name, key) - def is_replaying(self) -> bool: """Returns true if the state machine is replaying.""" return self.vm.is_replaying() diff --git a/src/lib.rs b/src/lib.rs index 5837ae71..1236cd53 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -959,25 +959,6 @@ impl PyVM { .map_err(Into::into) } - /// Attach to the run of the workflow `name` with the id `key` (unscoped): its result, once it ends. - fn attach_workflow( - mut self_: PyRefMut<'_, Self>, - name: String, - key: String, - ) -> Result { - self_ - .vm - .sys_attach_invocation( - restate_sdk_shared_core::AttachInvocationTarget::WorkflowId { - name, - key, - scope: None, - }, - ) - .map(Into::into) - .map_err(Into::into) - } - fn sys_end(mut self_: PyRefMut<'_, Self>) -> Result<(), PyVMError> { self_.vm.sys_end().map_err(Into::into) } diff --git a/tests/servercontext.py b/tests/servercontext.py index 5c652e5c..91429041 100644 --- a/tests/servercontext.py +++ b/tests/servercontext.py @@ -26,8 +26,6 @@ WorkflowContext, ) from restate.serde import DefaultSerde -from restate.server_context import ServerInvocationContext -import json import pytest import typing @@ -195,26 +193,6 @@ async def greet(ctx: Context, name: str) -> str | None: assert res == "hi" -async def test_attach_workflow_returns_the_result_of_its_run(): - greeter = Workflow("greeter_workflow") - - @greeter.main() - async def hello(ctx: WorkflowContext) -> str: - return "hi" - - caller = Service("caller") - - @caller.handler() - async def send_and_attach(ctx: Context, workflow_id: str) -> str: - ctx.workflow_send(hello, key=workflow_id, arg=None) # starts the run, ignoring its result - # attach_workflow is the server context's: the run's result, as JSON bytes. - result = await typing.cast(ServerInvocationContext, ctx).attach_workflow("greeter_workflow", workflow_id) - return json.loads(result) - - async with restate.create_test_harness(restate.app([greeter, caller]), restate_image=RESTATE_IMAGE) as harness: - assert await harness.client.service_call(send_and_attach, arg="wf-1") == "hi" - - async def test_a_send_to_a_workflow_gives_the_invocation_id_that_attach_waits_for(): greeter = Workflow("greeter_workflow_send")