feat: add stream timing test

Author: brian14708

examples/agent.py

@@ -137,12 +137,14 @@ async def run_round(user_input: str) -> None:
     # Main conversation loop (max 100 rounds)
     for i in range(100):
         # Collect any queued messages without waiting
-        msgs: list[str] = list(await input_.next_nowait())
+        it = await input_.next_nowait()
+        msgs: list[str] = list(it) if it is not None else []
 
         # If no queued messages, wait for next input
         if not msgs:
+            await output.send(("assistant", "[Waiting on user input...]"))
             m = await input_.next()
-            msgs = [m]
+            msgs.extend(m)
 
         # Execute round with tracing and interruption support
         with span(f"Round #{i + 1}"):

src/duron/_core/session.py

@@ -339,10 +339,10 @@ async def _resume(self) -> bool:
             _ = await self._step()
             if is_entry(entry):
                 if entry["source"] == "task":
-                    recvd_msgs.add(entry["id"])
                     if not self._handle_message(o, entry):
-                        msg = "Extra messages found in log"
-                        raise RuntimeError(msg)
+                        e = "Extra messages found in log"
+                        raise RuntimeError(e)
+                    recvd_msgs.add(entry["id"])
                 else:
                     _ = self._handle_message(o, entry)
                 _ = await self._step()
@@ -353,10 +353,19 @@ async def _resume(self) -> bool:
                 self._pending_msg.pop()
                 recvd_msgs.remove(id_)
 
+        pending: list[Entry] = []
+        for msg in self._pending_msg:
+            if msg["id"] not in recvd_msgs:
+                pending.append(msg)
+            else:
+                recvd_msgs.remove(msg["id"])
+        self._pending_msg = pending
+
         if len(recvd_msgs) > 0:
-            msg = "Extra messages found in log"
-            raise RuntimeError(msg)
-        return self._main.done()
+            e = "Extra messages found in log"
+            raise RuntimeError(e)
+
+        return self._main.done() and len(self._pending_msg) == 0
 
     async def _start(self) -> None:
         if await self._resume():
@@ -599,6 +608,7 @@ def done(f: OpFuture) -> None:
                     tracer.end_op_span(op.future_id, promise_complete_entry)
                 await self._enqueue_log(promise_complete_entry)
                 self._loop.post_completion(id_, result=None)
+
             case _:
                 assert_never(op)
 

src/duron/_core/signal.py

@@ -2,11 +2,11 @@
 
 import asyncio
 import sys
-from collections import deque
-from typing import TYPE_CHECKING, Final, Generic, cast
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Generic, Literal, cast
 from typing_extensions import Any, TypeVar, final, override
 
-from duron._core.ops import Barrier, StreamCreate, create_op
+from duron._core.ops import StreamCreate, create_op
 from duron._core.stream import OpWriter
 
 if TYPE_CHECKING:
@@ -36,7 +36,10 @@ def __repr__(self) -> str:
         return f"SignalInterrupt(value={self.value!r})"
 
 
-_SIGNAL_TRIGGER: Final = object()
+@dataclass(slots=True)
+class _SignalState:
+    depth: int
+    triggered: Literal[False] | SignalInterrupt
 
 
 @final
@@ -58,22 +61,19 @@ class Signal(Generic[_T]):
 
     def __init__(self, loop: EventLoop) -> None:
         self._loop = loop
-        # task -> [offset, stack depth]
-        self._tasks: dict[asyncio.Task[Any], tuple[int, int]] = {}
-        self._trigger: deque[tuple[int, _T]] = deque()
+        self._tasks: dict[asyncio.Task[Any], _SignalState] = {}
 
     async def __aenter__(self) -> None:
         task = asyncio.current_task()
         if task is None:
             return
         assert task.get_loop() == self._loop
-        offset, _ = await create_op(self._loop, Barrier())
-        for toffset, value in self._trigger:
-            if toffset > offset:
-                raise SignalInterrupt(value=value)
-        _, depth = self._tasks.get(task, (0, -1))
-        self._tasks[task] = (offset, depth + 1)
-        self._flush()
+        if task not in self._tasks:
+            val = _SignalState(depth=0, triggered=False)
+            self._tasks[task] = val
+        else:
+            val = self._tasks[task]
+            val.depth += 1
 
     async def __aexit__(
         self,
@@ -84,40 +84,27 @@ async def __aexit__(
         task = asyncio.current_task()
         if task is None:
             return
-        offset_end, _ = await create_op(self._loop, Barrier())
-
-        offset_start, depth = self._tasks.pop(task)
-        if depth > 0:
-            self._tasks[task] = (offset_end, depth - 1)
-        for toffset, value in self._trigger:
-            if (
-                offset_start < toffset < offset_end
-                and exc_type is asyncio.CancelledError
-                and (args := cast("asyncio.CancelledError", exc_value).args)
-                and args[0] is _SIGNAL_TRIGGER
-            ):
-                if sys.version_info >= (3, 11):
-                    _ = task.uncancel()
-                self._flush()
-                raise SignalInterrupt(value=value)
-
-    def on_next(self, offset: int, value: _T) -> None:
-        self._trigger.append((offset, value))
-        for t, (toffset, _depth) in self._tasks.items():
-            if toffset < offset:
-                _ = self._loop.call_soon(t.cancel, _SIGNAL_TRIGGER)
+        state = self._tasks.pop(task)
+        _ = self._loop.generate_op_scope()
+        triggered = state.triggered
+        if state.depth > 0:
+            state.triggered = False
+            state.depth -= 1
+            self._tasks[task] = state
+        if triggered is not False:
+            if sys.version_info >= (3, 11):
+                _ = task.uncancel()
+            raise triggered from None
+
+    def on_next(self, _offset: int, value: _T) -> None:
+        for t, state in self._tasks.items():
+            if state.triggered is False:
+                state.triggered = SignalInterrupt(value=value)
+                _ = self._loop.call_soon(t.cancel, state.triggered)
 
     def on_close(self, _offset: int, _exc: Exception | None) -> None:
         pass
 
-    def _flush(self) -> None:
-        if not self._tasks:
-            self._trigger.clear()
-            return
-        min_offset = min((offset for offset, _ in self._tasks.values()))
-        while self._trigger and self._trigger[0][0] < min_offset:
-            _ = self._trigger.popleft()
-
 
 async def create_signal(
     loop: EventLoop, dtype: TypeHint[_T], name: str | None, metadata: OpMetadata