Agents: the workflow tool sends then attaches, so a crash no longer reports a run as failed; the CLI shows steered messages

python/restate/agents/__main__.py

@@ -52,17 +52,38 @@ 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:
+            # 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)
@@ -71,6 +92,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 +128,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 +151,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 +184,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:

python/restate/agents/_calls.py

@@ -20,7 +20,10 @@
 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 ``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

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_invocation, start
 
 T = typing.TypeVar("T")
 
@@ -326,6 +327,21 @@ async def _once(
     return result
 
 
+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.
+
+    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)
+    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:
     """Return ``ctx`` untyped, to call a handler whose input type is known only at run time."""
     return ctx
@@ -523,7 +539,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.
@@ -546,9 +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 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)
 
         return await _once(service, arguments["raw"], tool_call, made)
 

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
@@ -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``.
@@ -251,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)
         return _update(output, result)
 
     return node

python/restate/sessions/_calls.py

@@ -74,6 +74,8 @@ def generic_send(
         headers: dict[str, str] | None = None,
     ) -> Any: ...
 
+    def attach_invocation(self, invocation_id: str) -> Awaitable[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``.
@@ -245,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."""
 

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,6 +167,9 @@ def generic_call(
     ) -> Any:
         return self.context.generic_call(service, handler, arg, key, idempotency_key, headers)
 
+    def attach_invocation(self, invocation_id: str) -> Awaitable[bytes]:
+        return self.context.attach_invocation(invocation_id, serde=BytesSerde())
+
     def generic_send(
         self,
         service: str,
@@ -336,6 +339,37 @@ async def cancel_calls(ctx: SessionContext) -> list[str]:
     return await _preview.cancel_previous_calls(_explicit_of(ctx, "cancel_calls"))
 
 
+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)
+
+
+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.
+    """
+    return _host_of(ctx).attach_invocation(invocation_id)
+
+
 def definition_of(ctx: SessionContext | ViewContext) -> SessionObject | None:
     """Return the session object whose handler or view ``ctx`` belongs to, or ``None``.
 

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``.
+- ``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 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
@@ -35,10 +38,12 @@
     "Schema",
     "ScopeKey",
     "SqlHost",
+    "attach_invocation",
     "cancel_calls",
     "cancelled",
     "database",
     "definition_of",
     "invoke",
     "reset",
+    "start",
 ]

tests/agents_cli.py

@@ -0,0 +1,88 @@
+#
+#  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"]
+
+
+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"]