Add end_session_confirm_message handshake

Previously the Python client sent end_session_message and immediately
closed the WS. Any artifact_upload_request_message the runtime produced
after that (from chokidar events that fired post-terminal, e.g. files
created by post-agent commands) was sent on a half-closed socket and
the upload_response could never make it back.

New protocol:
- Runtime receives end_session_message, drains the artifact watcher,
  sends end_session_confirm_message, then closes.
- Python client sends end_session_message, then pumps the receive
  loop — handling late artifact_upload_request_messages with the
  session's most-recent callback — until confirm arrives (or timeout).

Breaking change: older clients that close without waiting for confirm
will no longer get late artifacts uploaded. Both packages bumped to
0.12.0.

Author: jackcbrown89

packages/runtimeuse-client-python/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "runtimeuse-client"
-version = "0.11.0"
+version = "0.12.0"
 description = "Client library for AI agent runtime communication over WebSocket"
 readme = "README.md"
 license = {"text" = "FSL"}

packages/runtimeuse-client-python/src/runtimeuse_client/__init__.py

@@ -28,6 +28,7 @@
     ErrorMessageInterface,
     CancelMessage,
     EndSessionMessage,
+    EndSessionConfirmMessage,
     ArtifactUploadResult,
     OnAssistantMessageCallback,
     OnArtifactUploadRequestCallback,
@@ -63,6 +64,7 @@
     "ErrorMessageInterface",
     "CancelMessage",
     "EndSessionMessage",
+    "EndSessionConfirmMessage",
     "ArtifactUploadResult",
     "OnAssistantMessageCallback",
     "OnArtifactUploadRequestCallback",

packages/runtimeuse-client-python/src/runtimeuse_client/client.py

@@ -20,6 +20,7 @@
     AssistantMessageInterface,
     ArtifactUploadRequestMessageInterface,
     ArtifactUploadResponseMessageInterface,
+    OnArtifactUploadRequestCallback,
     QueryOptions,
     ExecuteCommandsOptions,
 )
@@ -183,6 +184,11 @@ def __init__(self, connected: ConnectedTransport):
         self._abort_event = asyncio.Event()
         self._send_queue: asyncio.Queue[dict] | None = None
         self._lock = asyncio.Lock()
+        # Track the most recent artifact upload callback so we can keep
+        # handling late artifact requests during session close (after the
+        # last query's request loop has exited).
+        self._last_artifact_cb: OnArtifactUploadRequestCallback | None = None
+        self._end_session_timeout_s: float = 60.0
 
     def abort(self) -> None:
         """Signal the in-flight request to cancel.
@@ -202,6 +208,8 @@ async def query(self, prompt: str, options: QueryOptions) -> QueryResult:
         async with self._lock:
             logger = options.logger or _default_logger
             self._abort_event = asyncio.Event()
+            if options.on_artifact_upload_request is not None:
+                self._last_artifact_cb = options.on_artifact_upload_request
 
             invocation = _build_invocation(prompt, options)
             send_queue: asyncio.Queue[dict] = asyncio.Queue()
@@ -238,6 +246,8 @@ async def execute_commands(
         async with self._lock:
             logger = options.logger or _default_logger
             self._abort_event = asyncio.Event()
+            if options.on_artifact_upload_request is not None:
+                self._last_artifact_cb = options.on_artifact_upload_request
 
             message = _build_command_execution(commands, options)
             send_queue: asyncio.Queue[dict] = asyncio.Queue()
@@ -266,6 +276,47 @@ async def execute_commands(
 
             return CommandExecutionResult(results=wire.results)
 
+    async def _end_session(self) -> None:
+        """Send end_session_message and drain late artifact uploads until the
+        runtime confirms. Called by the session context manager on exit.
+        """
+        end_session = getattr(self._connected, "end_session", None)
+        if end_session is None:
+            # Older/alternate transports: nothing to drain.
+            return
+
+        async def on_message(msg: dict) -> dict | None:
+            if msg.get("message_type") != "artifact_upload_request_message":
+                return None
+            if self._last_artifact_cb is None:
+                return None
+            try:
+                req = ArtifactUploadRequestMessageInterface.model_validate(msg)
+            except pydantic.ValidationError:
+                _default_logger.error(
+                    f"Malformed artifact upload request during end-of-session drain: {msg}"
+                )
+                return None
+            try:
+                result = await self._last_artifact_cb(req)
+            except Exception:
+                _default_logger.exception(
+                    "Artifact upload callback failed during end-of-session drain"
+                )
+                return None
+            response = ArtifactUploadResponseMessageInterface(
+                message_type="artifact_upload_response_message",
+                filename=req.filename,
+                filepath=req.filepath,
+                presigned_url=result.presigned_url,
+                content_type=result.content_type,
+            )
+            return response.model_dump(mode="json")
+
+        await end_session(
+            on_message=on_message, timeout_s=self._end_session_timeout_s
+        )
+
 
 class RuntimeUseClient:
     """Client for communicating with a runtimeuse agent runtime.
@@ -337,7 +388,11 @@ async def session(self) -> AsyncIterator[RuntimeUseSession]:
                 "Use a transport that implements PersistentTransport (e.g. WebSocketTransport)."
             )
         async with connect() as connected:
-            yield RuntimeUseSession(connected)
+            runtime_session = RuntimeUseSession(connected)
+            try:
+                yield runtime_session
+            finally:
+                await runtime_session._end_session()
 
     async def query(
         self,

packages/runtimeuse-client-python/src/runtimeuse_client/transports/transport.py

@@ -1,5 +1,8 @@
 import asyncio
-from typing import Any, AsyncGenerator, AsyncContextManager, Protocol
+from typing import Any, AsyncGenerator, AsyncContextManager, Awaitable, Callable, Protocol
+
+
+EndSessionMessageHandler = Callable[[dict[str, Any]], Awaitable[dict[str, Any] | None]]
 
 
 class Transport(Protocol):
@@ -30,6 +33,12 @@ def request(
         self, send_queue: asyncio.Queue[dict]
     ) -> AsyncGenerator[dict[str, Any], None]: ...
 
+    async def end_session(
+        self,
+        on_message: EndSessionMessageHandler | None = None,
+        timeout_s: float = 60.0,
+    ) -> None: ...
+
     async def close(self) -> None: ...
 
 

packages/runtimeuse-client-python/src/runtimeuse_client/transports/websocket_transport.py

@@ -6,6 +6,8 @@
 
 import websockets
 
+from .transport import EndSessionMessageHandler
+
 _logger = logging.getLogger(__name__)
 
 
@@ -41,15 +43,59 @@ async def request(
             except asyncio.CancelledError:
                 pass
 
-    async def close(self, send_end_message: bool = True) -> None:
-        """Close the connection, optionally sending end_session_message first."""
-        if send_end_message:
-            try:
-                await self._ws.send(
-                    json.dumps({"message_type": "end_session_message"})
-                )
-            except websockets.exceptions.ConnectionClosed:
-                pass
+    async def end_session(
+        self,
+        on_message: EndSessionMessageHandler | None = None,
+        timeout_s: float = 60.0,
+    ) -> None:
+        """Send ``end_session_message`` and pump the receive loop until the
+        server's ``end_session_confirm_message`` arrives (or timeout).
+
+        The runtime drains its artifact watcher before confirming, so late
+        ``artifact_upload_request_message``s may arrive in this window. For
+        each incoming message, ``on_message`` is awaited and may return a
+        dict response to send back over the socket.
+        """
+        try:
+            await self._ws.send(
+                json.dumps({"message_type": "end_session_message"})
+            )
+        except websockets.exceptions.ConnectionClosed:
+            return
+
+        try:
+            async with asyncio.timeout(timeout_s):
+                async for raw in self._ws:
+                    try:
+                        msg = json.loads(raw)
+                    except json.JSONDecodeError:
+                        continue
+                    if msg.get("message_type") == "end_session_confirm_message":
+                        return
+                    if on_message is None:
+                        continue
+                    try:
+                        response = await on_message(msg)
+                    except Exception:
+                        _logger.exception(
+                            "Error handling message during end_session drain"
+                        )
+                        continue
+                    if response is None:
+                        continue
+                    try:
+                        await self._ws.send(json.dumps(response))
+                    except websockets.exceptions.ConnectionClosed:
+                        return
+        except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed):
+            return
+
+    async def close(self) -> None:
+        """Close the underlying socket without sending end_session_message.
+
+        Callers that need a graceful session end should call
+        :meth:`end_session` first.
+        """
         await self._ws.close()
 
     async def _queue_sender(self, send_queue: asyncio.Queue[dict]) -> None:
@@ -89,13 +135,15 @@ async def __call__(
     async def connect(self) -> AsyncIterator[ConnectedWebSocketTransport]:
         """Open a persistent WebSocket connection for use with a session.
 
-        Sends ``end_session_message`` and closes the socket on exit.
+        The connection is closed on exit. Callers that want a graceful session
+        end (draining late artifacts) should call
+        :meth:`ConnectedWebSocketTransport.end_session` before exiting the
+        context.
         """
         _logger.info("Connecting persistent WebSocket to %s", self.ws_url)
         async with websockets.connect(self.ws_url, open_timeout=60) as ws:
             connected = ConnectedWebSocketTransport(ws)
             try:
                 yield connected
             finally:
-                await connected.close(send_end_message=True)
                 _logger.info("Persistent agent runtime connection closed")

packages/runtimeuse-client-python/src/runtimeuse_client/types.py

@@ -106,6 +106,10 @@ class EndSessionMessage(BaseModel):
     message_type: Literal["end_session_message"]
 
 
+class EndSessionConfirmMessage(BaseModel):
+    message_type: Literal["end_session_confirm_message"]
+
+
 class CommandExecutionMessage(BaseModel):
     message_type: Literal["command_execution_message"]
     source_id: str | None = None

packages/runtimeuse-client-python/test/conftest.py

@@ -67,6 +67,9 @@ async def request(
         finally:
             self._drain(send_queue)
 
+    async def end_session(self, on_message=None, timeout_s: float = 60.0) -> None:
+        return None
+
     async def close(self) -> None:
         self.closed = True
 

packages/runtimeuse-client-python/test/test_client.py

@@ -982,6 +982,9 @@ async def request(self, send_queue: asyncio.Queue[dict]):
         finally:
             self._drain_send(send_queue)
 
+    async def end_session(self, on_message=None, timeout_s: float = 60.0) -> None:
+        return None
+
     async def close(self) -> None:
         self.closed = True