harness-sdk-python 0.4.3__tar.gz → 0.4.4__tar.gz

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: harness-sdk-python
-Version: 0.4.3
+Version: 0.4.4
 Summary: RunManager: the harness-sdk runs subsystem for Python Statewire hosts
 Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
 License-Expression: MIT

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "harness-sdk-python"
-version = "0.4.3"
+version = "0.4.4"
 description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
 readme = "README.md"
 license = "MIT"

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/src/harness_sdk/run_manager.py

@@ -184,8 +184,15 @@ class RunManager:
         self._run_futures: list["asyncio.Future[Any]"] = []
         self._input_requests: list[dict[str, Any]] = []
         self._input_answers: dict[str, Any] = {}
+        self._idle = asyncio.Event()
+        self._idle.set()
         self._init_state()
 
+    @property
+    def idle(self) -> asyncio.Event:
+        """Set while no run is active and nothing is about to dispatch; cleared for the whole drain chain."""
+        return self._idle
+
     def _init_state(self) -> None:
         # Runs state is write-only and not durable: overwrite whatever is there.
         self._state["status"] = "ready"
@@ -247,6 +254,7 @@ class RunManager:
                 self._dispatched_ids = ()
             if self._staged_rewinds:
                 rewind = self._staged_rewinds.pop(0)
+                self._clear_input()
                 self._dispatch(
                     rewind.type,
                     rewind.messages,
@@ -276,6 +284,8 @@ class RunManager:
         for future in fx.continues:
             if not future.done():
                 future.set_result(None)
+        if self._task is None:
+            self._idle.set()
 
     def _apply(self, entry: Any, fx: _Effects) -> None:
         if isinstance(entry, _Stop):
@@ -400,6 +410,7 @@ class RunManager:
         )
         self._ctx = ctx
         self._sync_steering()
+        self._idle.clear()
         self._task = self._create_task(self._run(ctx))
 
     async def _run(self, ctx: "RunManager.StartContext") -> None:
@@ -424,6 +435,13 @@ class RunManager:
                 raise RuntimeError(
                     "run settled without acking its messages (call ctx.ack_messages())"
                 )
+        except asyncio.CancelledError:
+            self._settle(ctx)
+            self._state["status"] = "stopped"
+            self._settle_initiators(_reject("stopped", "run cancelled"))
+            self._revert_dispatching()
+            self._idle.set()
+            raise  # no drain, no freeze
         except Exception as exc:
             self._settle(ctx)
             message = str(exc) or type(exc).__name__
@@ -741,8 +759,6 @@ class RunManager:
             raise _reject("capability-missing", "the rewind capability is not enabled")
         if e.source_meta is None:
             raise _reject("unknown-id", f"message {e.message_id} is unknown")
-        if self._task is not None and not e.source_meta["isLeaf"]:
-            raise _reject("not-leaf", "only the leaf may be edited during a run")
         self._check_anchor(e)
         self._staged_rewinds.append(
             _Rewind(
@@ -788,7 +804,6 @@ class RunManager:
         return None
 
     def _apply_edit(self, e: _Edit) -> Any:
-        self._check_leaf_lanes(e.source_meta, "run/edit")
         if (
             e.message["id"] != e.source_id
             and self._lane_of(e.message["id"]) is not None
@@ -808,7 +823,6 @@ class RunManager:
         return _PARKED
 
     def _apply_reload(self, e: _Reload) -> Any:
-        self._check_leaf_lanes(e.source_meta, "run/reload")
         self._staged_rewinds.append(
             _Rewind(
                 "message-reload",
@@ -821,15 +835,6 @@ class RunManager:
         )
         return _PARKED
 
-    def _check_leaf_lanes(self, meta: dict[str, Any], command: str) -> None:
-        if meta["isLeaf"]:
-            return
-        if self._lane_items("steerQueue") or self._lane_items("queue"):
-            raise _reject(
-                "not-leaf",
-                f"only the leaf accepts {command} while the queue is non-empty",
-            )
-
     # ─── Command handlers ───────────────────────────────────
 
     async def _send(
@@ -912,10 +917,7 @@ class RunManager:
         return await self._stage(_Dequeue(message_id))
 
     def _check_rewind_gate(self, command: str) -> None:
-        status = self._status()
-        if status == "input-required":
-            raise _reject("wrong-state", f"{command} is rejected in input-required")
-        if status == "running":
+        if self._status() == "running":
             if "rewind-during-run" not in self._capabilities:
                 raise _reject(
                     "capability-missing",
@@ -924,10 +926,6 @@ class RunManager:
         elif "rewind" not in self._capabilities:
             raise _reject("capability-missing", "the rewind capability is not enabled")
 
-    def _check_leaf_running(self, meta: dict[str, Any], command: str) -> None:
-        if not meta["isLeaf"] and self._status() == "running":
-            raise _reject("not-leaf", f"only the leaf accepts {command} during a run")
-
     async def edit(
         self, params: Any, *, meta: Any = None, ack: Callable[[], None]
     ) -> Any:
@@ -938,7 +936,6 @@ class RunManager:
         source_meta = await self._get_message_meta(source_id)
         if source_meta is None:
             raise _reject("unknown-id", f"message {source_id} is unknown")
-        self._check_leaf_running(source_meta, "run/edit")
         if source_meta["role"] != "user" and "assistant-edit" not in self._capabilities:
             raise _reject(
                 "capability-missing", "the assistant-edit capability is not enabled"
@@ -965,7 +962,6 @@ class RunManager:
             raise _reject("unknown-id", f"message {source_id} is unknown")
         if source_meta["role"] != "assistant":
             raise _reject("invalid-message", "sourceId must name an assistant message")
-        self._check_leaf_running(source_meta, "run/reload")
         if source_meta["parentId"] is not None:
             parent = await self._get_message_meta(source_meta["parentId"])
             if (

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_edit_dispatched.py

@@ -47,15 +47,24 @@ async def test_edit_of_dispatched_item_requires_rewind_during_run(capabilities):
 
 
 @pytest.mark.parametrize("command", ["run/enqueue", "run/steer"])
-async def test_edit_of_dispatched_non_leaf_rejects(command):
+async def test_edit_of_dispatched_non_leaf_stops_and_reruns(command):
     script = Script()
     async with run_host(script, capabilities=("rewind", "rewind-during-run")) as (drv, host):
         await drv.command("run/enqueue", add("m1"), terminal=False)
-        await script.next_call()
+        call = await script.next_call()
         script.thread["m1"] = {"parentId": None, "role": "user", "isLeaf": False}
-        assert_rejected(
-            await drv.command(command, {"message": msg("m1", "edited")}), "not-leaf"
+        call.ack()
+        pending = await drv.command(
+            command, {"message": msg("m1", "edited")}, terminal=False
         )
+        assert pending["type"] == "pending"
+        call.finish(RunManager.Stop(dispatch_queue=False))
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-edit"
+        assert [m["parts"][0]["text"] for m in rerun.ctx.messages] == ["edited"]
+        rerun.ack()
+        rerun.finish(RunManager.Stop(dispatch_queue=False))
+        assert (await drv.res(pending["seq"]))["type"] == "accepted"
 
 
 async def test_edit_of_dispatched_item_with_queued_items_stops_and_reruns():

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_edit_reload.py

@@ -49,7 +49,7 @@ async def test_edit_while_running_requires_rewind_during_run():
         )
 
 
-async def test_deep_edit_rejected_while_queue_non_empty():
+async def test_deep_edit_with_queue_non_empty_reruns_and_keeps_queue():
     script = Script()
     thread_with_turn(script)
     async with run_host(script, capabilities=("rewind",)) as (drv, host):
@@ -59,14 +59,12 @@ async def test_deep_edit_rejected_while_queue_non_empty():
         call.ack()
         call.fail(RuntimeError("boom"))
         await drv.wait_status("error")
-        assert_rejected(
-            await drv.command("run/edit", {"sourceId": "u1", "message": msg("u2")}, terminal=False),
-            "not-leaf",
-        )
-        await drv.command("run/dequeue", {"messageId": "m2"})
         res = await drv.command("run/edit", {"sourceId": "u1", "message": msg("u2")}, terminal=False)
         assert res["type"] == "pending"
-        assert (await script.next_call()).ctx.type == "message-edit"
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-edit"
+        assert rerun.ctx.rollback_to is None
+        assert [item["id"] for item in drv.replica["queue"]] == ["m2"]
 
 
 async def test_leaf_edit_in_error_accepted_with_queue_non_empty():
@@ -164,7 +162,7 @@ async def test_reload_without_rewind_rejects():
         )
 
 
-async def test_reload_leaf_only_while_queue_non_empty():
+async def test_deep_reload_with_queue_non_empty_reruns_and_keeps_queue():
     script = Script()
     script.thread.update(
         {
@@ -181,13 +179,12 @@ async def test_reload_leaf_only_while_queue_non_empty():
         call.ack()
         call.fail(RuntimeError("boom"))
         await drv.wait_status("error")
-        assert_rejected(
-            await drv.command("run/reload", {"sourceId": "a0"}, terminal=False), "not-leaf"
-        )
-        res = await drv.command("run/reload", {"sourceId": "a1"}, terminal=False)
+        res = await drv.command("run/reload", {"sourceId": "a0"}, terminal=False)
         assert res["type"] == "pending"
         call = await script.next_call()
-        assert call.ctx.rollback_to == "m1"
+        assert call.ctx.type == "message-reload"
+        assert call.ctx.rollback_to == "u0"
+        assert [item["id"] for item in drv.replica["queue"]] == ["m2"]
 
 
 async def test_reload_of_assistant_parented_response_needs_continuation():

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_input_required.py

@@ -257,24 +257,54 @@ async def test_queue_while_parked_dispatches_after_resume():
         await drv.wait_status("ready")
 
 
-async def test_lifecycle_commands_rejected_while_parked():
+async def test_continue_rejected_while_parked():
     script = Script()
+    async with run_host(script) as (drv, host):
+        await park(drv, script, [tool_call("r1")])
+        assert_rejected(await drv.command("run/continue"), "wrong-state")
+
+
+def parked_thread(script):
     script.thread.update(
         {
             "u1": {"parentId": None, "role": "user", "onActiveBranch": True},
             "a1": {"parentId": "u1", "role": "assistant", "isLeaf": True, "onActiveBranch": True},
         }
     )
+
+
+async def test_edit_while_parked_clears_requests_and_reruns():
+    script = Script()
+    parked_thread(script)
     async with run_host(script, capabilities=("rewind",)) as (drv, host):
         await park(drv, script, [tool_call("r1")], anchor="a1")
-        assert_rejected(await drv.command("run/continue"), "wrong-state")
-        assert_rejected(
-            await drv.command("run/edit", {"sourceId": "u1", "message": msg("u2")}, terminal=False),
-            "wrong-state",
-        )
-        assert_rejected(
-            await drv.command("run/reload", {"sourceId": "a1"}, terminal=False), "wrong-state"
+        res = await drv.command(
+            "run/edit", {"sourceId": "u1", "message": msg("u2")}, terminal=False
         )
+        assert res["type"] == "pending"
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-edit"
+        assert rerun.ctx.rollback_to is None
+        await drv.wait(lambda s: "inputRequests" not in s)
+        rerun.ack()
+        rerun.finish(RunManager.Complete())
+        await drv.wait_status("ready")
+
+
+async def test_reload_while_parked_clears_requests_and_reruns():
+    script = Script()
+    parked_thread(script)
+    async with run_host(script, capabilities=("rewind",)) as (drv, host):
+        await park(drv, script, [tool_call("r1")], anchor="a1")
+        res = await drv.command("run/reload", {"sourceId": "a1"}, terminal=False)
+        assert res["type"] == "pending"
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-reload"
+        assert rerun.ctx.rollback_to == "u1"
+        await drv.wait(lambda s: "inputRequests" not in s)
+        rerun.ack()
+        rerun.finish(RunManager.Complete())
+        await drv.wait_status("ready")
 
 
 async def test_steer_while_parked_stays_queued():

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_outcomes.py

@@ -138,3 +138,22 @@ async def test_unknown_capability_rejects_at_construction():
             schedule=lambda fn: None,
             capabilities=("time-travel",),
         )
+
+
+async def test_idle_spans_the_drain_chain():
+    script = Script()
+    async with run_host(script) as (drv, host):
+        runs = host.live.runs
+        assert runs.idle.is_set()
+        await drv.command("run/enqueue", add("m1"), terminal=False)
+        first = await script.next_call()
+        assert not runs.idle.is_set()
+        await drv.command("run/enqueue", add("m2", anchor="m1"), terminal=False)
+        first.ack()
+        first.finish(RunManager.Complete())
+        second = await script.next_call()
+        assert not runs.idle.is_set()
+        second.ack()
+        second.finish(RunManager.Complete())
+        await drv.wait_status("ready")
+        assert runs.idle.is_set()

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_rewind_during_run.py

@@ -89,17 +89,39 @@ async def test_running_edit_and_reload_reject_without_capability(capabilities):
         )
 
 
-async def test_running_deeper_than_leaf_edit_and_reload_reject():
+async def test_running_non_leaf_edit_stops_settles_then_rewinds():
     script = Script()
     async with run_host(script, capabilities=("rewind", "rewind-during-run")) as (drv, host):
-        await start_run(drv, script, leaf=False)
-        assert_rejected(
-            await drv.command("run/edit", {"sourceId": "m1", "message": msg("m2")}, terminal=False),
-            "not-leaf",
-        )
-        assert_rejected(
-            await drv.command("run/reload", {"sourceId": "a1"}, terminal=False), "not-leaf"
+        call = await start_run(drv, script, leaf=False)
+        pending = await drv.command(
+            "run/edit", {"sourceId": "m1", "message": msg("m2")}, terminal=False
         )
+        assert pending["type"] == "pending"
+        await asyncio.wait_for(call.ctx.stop_requested.wait(), 5)
+        call.finish(RunManager.Stop(dispatch_queue=False))
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-edit"
+        assert [m["id"] for m in rerun.ctx.messages] == ["m2"]
+        assert rerun.ctx.rollback_to is None
+        rerun.ack()
+        rerun.finish(RunManager.Stop(dispatch_queue=False))
+        assert (await drv.res(pending["seq"]))["type"] == "accepted"
+
+
+async def test_running_non_leaf_reload_stops_settles_then_rewinds():
+    script = Script()
+    async with run_host(script, capabilities=("rewind", "rewind-during-run")) as (drv, host):
+        call = await start_run(drv, script, leaf=False)
+        pending = await drv.command("run/reload", {"sourceId": "a1"}, terminal=False)
+        assert pending["type"] == "pending"
+        await asyncio.wait_for(call.ctx.stop_requested.wait(), 5)
+        call.finish(RunManager.Stop(dispatch_queue=False))
+        rerun = await script.next_call()
+        assert rerun.ctx.type == "message-reload"
+        assert rerun.ctx.rollback_to == "m1"
+        rerun.ack()
+        rerun.finish(RunManager.Stop(dispatch_queue=False))
+        assert (await drv.res(pending["seq"]))["type"] == "accepted"
 
 
 async def test_rewind_during_run_requires_rewind():

{harness_sdk_python-0.4.3 → harness_sdk_python-0.4.4}/tests/test_settle.py

@@ -215,3 +215,33 @@ async def test_rewind_edit_failure_settles_with_the_payload():
         settled = await drv.res(pending["seq"])
         assert settled["type"] == "rejected"
         assert settled["payload"] == {"code": 402}
+
+
+async def test_task_cancellation_reverts_entries_and_settles_sends():
+    script = Script()
+    async with run_host(script) as (drv, host):
+        pending = await drv.command("run/enqueue", add("m1"), terminal=False)
+        await script.next_call()
+        await drv.command("run/enqueue", add("m2"))
+        task = host.live.runs._task
+        task.cancel()
+        assert (await drv.res(pending["seq"]))["type"] == "accepted"
+        await drv.wait_status("stopped")
+        assert queue_ids(drv.replica, "queue") == ["m1", "m2"]
+        assert drv.replica["runLeafMessageId"] is None
+        assert task.cancelled()
+
+
+async def test_task_cancellation_rejects_a_rewind_initiator():
+    script = Script()
+    script.thread["u1"] = {"parentId": None, "role": "user", "isLeaf": True}
+    async with run_host(script, capabilities=("rewind",)) as (drv, host):
+        pending = await drv.command(
+            "run/edit", {"sourceId": "u1", "message": msg("u2")}, terminal=False
+        )
+        await script.next_call()
+        host.live.runs._task.cancel()
+        settled = await drv.res(pending["seq"])
+        assert settled["type"] == "rejected"
+        assert settled["payload"] == {"reason": "stopped"}
+        await drv.wait_status("stopped")