chore: import upstream snapshot with attribution

python/assistant-stream/src/assistant_stream/create_run.py

@@ -0,0 +1,275 @@
+import asyncio
+import logging
+from typing import Any, AsyncGenerator, Callable, Coroutine, List, Optional, Sequence, Union
+from assistant_stream.assistant_stream_chunk import (
+    AssistantStreamChunk,
+    TextDeltaChunk,
+    ReasoningDeltaChunk,
+    ToolResultChunk,
+    DataChunk,
+    ErrorChunk,
+    SourceChunk,
+    ToolCallBeginChunk,
+)
+from assistant_stream.modules.tool_call import (
+    create_tool_call,
+    ToolCallController,
+    generate_openai_style_tool_call_id,
+)
+from assistant_stream.state_manager import StateManager
+
+logger = logging.getLogger(__name__)
+
+
+class ReadOnlyCancellationSignal:
+    """Read-only view over an asyncio.Event used for cancellation."""
+
+    def __init__(self, event: asyncio.Event):
+        self._event = event
+
+    def is_set(self) -> bool:
+        return self._event.is_set()
+
+    async def wait(self) -> bool:
+        return await self._event.wait()
+
+
+class RunController:
+    def __init__(self, queue, state_data, parent_id: Optional[str] = None):
+        self._queue = queue
+        self._loop = asyncio.get_running_loop()
+        self._dispose_callbacks = []
+        self._stream_tasks = []
+        self._state_manager = StateManager(self._put_chunk_nowait, state_data)
+        self._parent_id = parent_id
+        self._cancelled_event = asyncio.Event()
+        self._cancelled_signal = ReadOnlyCancellationSignal(self._cancelled_event)
+
+    def with_parent_id(self, parent_id: str) -> 'RunController':
+        """Create a new RunController instance with the specified parent_id."""
+        controller = RunController(self._queue, self._state_manager._state_data, parent_id)
+        controller._loop = self._loop
+        controller._dispose_callbacks = self._dispose_callbacks
+        controller._stream_tasks = self._stream_tasks
+        controller._state_manager = self._state_manager
+        controller._cancelled_event = self._cancelled_event
+        controller._cancelled_signal = self._cancelled_signal
+        return controller
+
+    def append_text(self, text_delta: str) -> None:
+        """Append a text delta to the stream."""
+        chunk = TextDeltaChunk(text_delta=text_delta, parent_id=self._parent_id)
+        self._flush_and_put_chunk(chunk)
+
+    def append_reasoning(self, reasoning_delta: str) -> None:
+        """Append a reasoning delta to the stream."""
+        chunk = ReasoningDeltaChunk(reasoning_delta=reasoning_delta, parent_id=self._parent_id)
+        self._flush_and_put_chunk(chunk)
+
+    def append_state_text(
+        self, path: Sequence[Union[str, int]], text_delta: str
+    ) -> None:
+        """Append a text delta at a state path using an append-text operation."""
+        self._state_manager.append_text(path, text_delta)
+
+    async def add_tool_call(
+        self, tool_name: str, tool_call_id: str = None
+    ) -> ToolCallController:
+        """Add a tool call to the stream."""
+        if tool_call_id is None:
+            tool_call_id = generate_openai_style_tool_call_id()
+
+        stream, controller = await create_tool_call(tool_name, tool_call_id, self._parent_id)
+        self._dispose_callbacks.append(controller.close)
+
+        self.add_stream(stream)
+        return controller
+
+    def add_tool_result(self, tool_call_id: str, result: Any) -> None:
+        """Add a tool result to the stream."""
+        chunk = ToolResultChunk(
+            tool_call_id=tool_call_id,
+            result=result,
+        )
+        self._flush_and_put_chunk(chunk)
+
+    def add_stream(self, stream: AsyncGenerator[AssistantStreamChunk, None]) -> None:
+        """Append a substream to the main stream."""
+
+        async def reader():
+            async for chunk in stream:
+                self._flush_and_put_chunk(chunk)
+
+        task = asyncio.create_task(reader())
+        self._stream_tasks.append(task)
+
+    def add_data(self, data: Any) -> None:
+        """Emit an event to the main stream."""
+        chunk = DataChunk(data=data)
+        self._flush_and_put_chunk(chunk)
+
+    def add_error(self, error: str) -> None:
+        """Emit an error to the main stream."""
+        chunk = ErrorChunk(error=error)
+        self._flush_and_put_chunk(chunk)
+    
+    def add_source(self, id: str, url: str, title: Optional[str] = None) -> None:
+        """Add a source to the stream."""
+        chunk = SourceChunk(
+            id=id,
+            url=url,
+            title=title,
+            parent_id=self._parent_id
+        )
+        self._flush_and_put_chunk(chunk)
+
+    def _put_chunk_nowait(self, chunk):
+        """Helper method to put a chunk in the queue without waiting.
+
+        This is used as a callback for the StateManager.
+        """
+        self._loop.call_soon_threadsafe(self._queue.put_nowait, chunk)
+
+    def _flush_and_put_chunk(self, chunk):
+        """Helper method to flush state operations and put a chunk in the queue.
+
+        This ensures state operations are sent before other operations.
+        """
+        # Flush any pending state operations first
+        self._state_manager.flush()
+        # Add the chunk to the queue
+        self._loop.call_soon_threadsafe(self._queue.put_nowait, chunk)
+
+    @property
+    def state(self):
+        """Access the state proxy object for making state updates.
+
+        This property provides a proxy object that allows navigating to any path
+        in the state, reading values, and setting values, which will trigger the
+        appropriate state update operation.
+
+        If the state is None, this property returns None directly.
+        You can set the root state directly by assigning to this property.
+
+        Example:
+            controller.state = {"user": {"name": "John"}, "messages": [{"text": "Hi"}]}
+            controller.state["user"]["name"] = "Bob"
+            name = controller.state["user"]["name"]
+            controller.state["messages"][0]["text"] += " chunk"  # emits append-text once non-empty
+
+            # Explicit equivalent:
+            controller.append_state_text(["messages", 0, "text"], " chunk")
+        """
+        return self._state_manager.state
+
+    @state.setter
+    def state(self, value):
+        """Set the entire state object.
+
+        Args:
+            value: The new state value to set
+        """
+        self._state_manager.add_operations(
+            [{"type": "set", "path": [], "value": value}]
+        )
+
+    @property
+    def cancelled_event(self) -> ReadOnlyCancellationSignal:
+        """Expose cancellation signal for cooperative cancellation."""
+        return self._cancelled_signal
+
+    @property
+    def is_cancelled(self) -> bool:
+        """Return whether this run has been cancelled."""
+        return self._cancelled_event.is_set()
+
+    def _mark_cancelled(self) -> None:
+        """Set cancellation signal once."""
+        if not self._cancelled_event.is_set():
+            self._cancelled_event.set()
+
+
+async def create_run(
+    callback: Callable[[RunController], Coroutine[Any, Any, None]],
+    *,
+    state: Any | None = None,
+) -> AsyncGenerator[AssistantStreamChunk, None]:
+    queue = asyncio.Queue()
+    controller = RunController(queue, state_data=state)
+
+    async def background_task():
+        try:
+            await callback(controller)
+        except Exception as e:
+            controller.add_error(str(e))
+            raise
+        finally:
+            # Flush any pending state updates before disposing
+            controller._state_manager.flush()
+
+            for dispose in controller._dispose_callbacks:
+                dispose()
+            try:
+                for task in controller._stream_tasks:
+                    await task
+            finally:
+                asyncio.get_running_loop().call_soon_threadsafe(queue.put_nowait, None)
+
+    task = asyncio.create_task(background_task())
+    ended_normally = False
+
+    try:
+        while True:
+            chunk = await controller._queue.get()
+            if chunk is None:
+                ended_normally = True
+                break
+            yield chunk
+            controller._queue.task_done()
+    finally:
+        if ended_normally:
+            # The `None` sentinel is queued at the end of `background_task`, so
+            # normal stream completion implies `task` is already done here.
+            # `result()` preserves normal-path error propagation.
+            task.result()
+        else:
+            controller._mark_cancelled()
+            # Yield to the event loop to allow the cancel signal to propagate.
+            await asyncio.sleep(0)
+            if not task.done():
+                # Give callbacks a brief chance to observe `is_cancelled`
+                # and exit cooperatively before forcing cancellation.
+                # 50ms keeps disconnect cleanup responsive without immediately
+                # interrupting callbacks that can stop themselves quickly.
+                try:
+                    await asyncio.wait_for(asyncio.shield(task), timeout=0.05)
+                except asyncio.TimeoutError:
+                    # Timeout means cooperative shutdown did not finish in time.
+                    pass
+                except Exception:
+                    # The stream consumer already disconnected, so suppress callback errors
+                    # but keep a log signal for postmortem debugging.
+                    logger.warning(
+                        "Suppressed callback exception during early-close grace period",
+                        exc_info=True,
+                    )
+            if not task.done():
+                task.cancel()
+            try:
+                # `shield()` lets caller-initiated cancellation interrupt `aclose()`
+                # without conflating it with our own forced `task.cancel()`.
+                await asyncio.shield(task)
+            except asyncio.CancelledError:
+                if task.cancelled():
+                    # Expected forced cancellation for early-close cleanup.
+                    pass
+                else:
+                    # Preserve caller-initiated cancellation (e.g. wait_for timeout).
+                    raise
+            except Exception:
+                # The stream consumer already disconnected, so suppress callback errors
+                # but keep a log signal for postmortem debugging.
+                logger.warning(
+                    "Suppressed callback exception after forced early-close cancellation",
+                    exc_info=True,
+                )