Merge pull request #12 from onyx-dot-app/docker-exec-stream

Danelegend · web-flow · commit e6c4af89df62 · 2026-02-20T19:47:41.000-08:00
feat: stream code execution
diff --git a/code-interpreter/app/api/routes.py b/code-interpreter/app/api/routes.py
@@ -1,21 +1,25 @@
 from __future__ import annotations
 
+from collections.abc import Iterator
 from pathlib import Path
 
 from fastapi import APIRouter, File, HTTPException, UploadFile, status
-from fastapi.responses import Response
+from fastapi.responses import Response, StreamingResponse
 
 from app.app_configs import get_settings
 from app.models.schemas import (
     ExecuteRequest,
     ExecuteResponse,
     FileMetadataResponse,
     ListFilesResponse,
+    StreamErrorEvent,
+    StreamOutputEvent,
+    StreamResultEvent,
     UploadFileResponse,
     WorkspaceFile,
 )
-from app.services.executor_base import EntryKind, WorkspaceEntry
-from app.services.executor_factory import execute_python
+from app.services.executor_base import EntryKind, StreamChunk, StreamResult, WorkspaceEntry
+from app.services.executor_factory import execute_python, execute_python_streaming
 from app.services.file_storage import FileStorageService
 
 router = APIRouter()
@@ -118,6 +122,51 @@ def execute(req: ExecuteRequest) -> ExecuteResponse:
     )
 
 
+@router.post("/execute/stream")
+def execute_stream(req: ExecuteRequest) -> StreamingResponse:
+    """Execute Python code with streaming output via Server-Sent Events."""
+    _validate_timeout(req)
+    settings = get_settings()
+    storage = get_file_storage()
+    staged_files, input_files_map = _stage_request_files(req, storage)
+
+    def generate() -> Iterator[str]:
+        try:
+            for event in execute_python_streaming(
+                code=req.code,
+                stdin=req.stdin,
+                timeout_ms=req.timeout_ms,
+                max_output_bytes=settings.max_output_bytes,
+                cpu_time_limit_sec=settings.cpu_time_limit_sec,
+                memory_limit_mb=settings.memory_limit_mb,
+                files=staged_files,
+                last_line_interactive=req.last_line_interactive,
+            ):
+                if isinstance(event, StreamChunk):
+                    yield StreamOutputEvent(stream=event.stream, data=event.data).to_sse()
+
+                elif isinstance(event, StreamResult):
+                    yield StreamResultEvent(
+                        exit_code=event.exit_code,
+                        timed_out=event.timed_out,
+                        duration_ms=event.duration_ms,
+                        files=_save_workspace_files(event.files, input_files_map, storage),
+                    ).to_sse()
+
+        except Exception as exc:
+            yield StreamErrorEvent(message=str(exc)).to_sse()
+
+    return StreamingResponse(
+        generate(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+
 @router.post("/files", response_model=UploadFileResponse, status_code=status.HTTP_201_CREATED)
 async def upload_file(file: UploadFile = File(...)) -> UploadFileResponse:  # noqa: B008
     """Upload a file for later use in code execution."""
diff --git a/code-interpreter/app/models/schemas.py b/code-interpreter/app/models/schemas.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import ClassVar, Literal
+
 from pydantic import BaseModel, Field, StrictInt, StrictStr
 
 from app.services.executor_base import EntryKind
@@ -50,6 +52,51 @@ class ExecuteResponse(BaseModel):
     )
 
 
+class SSEModel(BaseModel):
+    """Base for Server-Sent Event payloads.
+
+    Subclasses declare ``sse_event`` as a ClassVar to set the SSE event type.
+    Call ``to_sse()`` to get a fully-formatted SSE frame.
+    """
+
+    sse_event: ClassVar[str]
+
+    def to_sse(self) -> str:
+        data = self.model_dump_json()
+        return f"event: {self.sse_event}\ndata: {data}\n\n"
+
+
+class StreamOutputEvent(SSEModel):
+    """Payload for 'output' SSE events."""
+
+    sse_event: ClassVar[str] = "output"
+
+    stream: Literal["stdout", "stderr"]
+    data: StrictStr
+
+
+class StreamResultEvent(SSEModel):
+    """Payload for the final 'result' SSE event."""
+
+    sse_event: ClassVar[str] = "result"
+
+    exit_code: int | None
+    timed_out: bool
+    duration_ms: StrictInt
+    files: list[WorkspaceFile] = Field(
+        default_factory=list,
+        description="Snapshot of the execution workspace after completion.",
+    )
+
+
+class StreamErrorEvent(SSEModel):
+    """Payload for 'error' SSE events."""
+
+    sse_event: ClassVar[str] = "error"
+
+    message: StrictStr
+
+
 class UploadFileResponse(BaseModel):
     file_id: StrictStr = Field(..., description="Unique identifier for the uploaded file.")
     filename: StrictStr = Field(..., description="Original filename as provided during upload.")
diff --git a/code-interpreter/app/services/executor_base.py b/code-interpreter/app/services/executor_base.py
@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Sequence
+from collections.abc import Generator, Sequence
 from dataclasses import dataclass
 from enum import StrEnum
-from typing import Protocol
+from typing import Literal, Protocol
 
 
 def wrap_last_line_interactive(code: str) -> str:
@@ -77,6 +77,27 @@ class WorkspaceEntry:
     content: bytes | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class StreamChunk:
+    """A chunk of output from the execution."""
+
+    stream: Literal["stdout", "stderr"]
+    data: str
+
+
+@dataclass(frozen=True, slots=True)
+class StreamResult:
+    """Final execution result emitted at end of stream."""
+
+    exit_code: int | None
+    timed_out: bool
+    duration_ms: int
+    files: tuple[WorkspaceEntry, ...]
+
+
+StreamEvent = StreamChunk | StreamResult
+
+
 class ExecutorProtocol(Protocol):
     def execute_python(
         self,
@@ -113,6 +134,25 @@ def execute_python(
                                    if it's a bare expression (only the last line is affected).
         """
 
+    def execute_python_streaming(
+        self,
+        *,
+        code: str,
+        stdin: str | None,
+        timeout_ms: int,
+        max_output_bytes: int,
+        cpu_time_limit_sec: int | None = None,
+        memory_limit_mb: int | None = None,
+        files: Sequence[tuple[str, bytes]] | None = None,
+        last_line_interactive: bool = True,
+    ) -> Generator[StreamEvent, None, None]:
+        """Execute Python code and yield output chunks as they arrive.
+
+        Yields StreamChunk events during execution, then a single StreamResult
+        at the end. Default implementation raises NotImplementedError.
+        """
+        raise NotImplementedError(f"{type(self).__name__} does not support streaming execution")
+
     @staticmethod
     def truncate_output(stream: bytes, max_bytes: int) -> str:
         if len(stream) <= max_bytes:
diff --git a/code-interpreter/app/services/executor_docker.py b/code-interpreter/app/services/executor_docker.py
@@ -1,5 +1,8 @@
+import codecs
 import io
 import logging
+import os
+import selectors
 import shlex
 import subprocess
 import tarfile
@@ -10,6 +13,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from shutil import which
+from typing import Literal
 
 from app.app_configs import (
     PYTHON_EXECUTOR_DOCKER_BIN,
@@ -20,6 +24,9 @@
     BaseExecutor,
     EntryKind,
     ExecutionResult,
+    StreamChunk,
+    StreamEvent,
+    StreamResult,
     WorkspaceEntry,
     wrap_last_line_interactive,
 )
@@ -406,3 +413,159 @@ def execute_python(
             duration_ms=duration_ms,
             files=workspace_snapshot,
         )
+
+    def _terminate_process(self, ctx: _ExecContext, timed_out: bool) -> None:
+        """Kill the process on timeout or wait for normal exit."""
+        if timed_out:
+            subprocess.run(
+                [
+                    self.docker_binary,
+                    "exec",
+                    ctx.container_name,
+                    "pkill",
+                    "-9",
+                    "python",
+                ],
+                capture_output=True,
+            )
+            ctx.proc.kill()
+        ctx.proc.wait()
+
+    def execute_python_streaming(
+        self,
+        *,
+        code: str,
+        stdin: str | None,
+        timeout_ms: int,
+        max_output_bytes: int,
+        cpu_time_limit_sec: int | None = None,
+        memory_limit_mb: int | None = None,
+        files: Sequence[tuple[str, bytes]] | None = None,
+        last_line_interactive: bool = True,
+    ) -> Generator[StreamEvent, None, None]:
+        """Execute Python code and yield output chunks as they arrive via SSE.
+
+        Yields StreamChunk events during execution, then a single StreamResult
+        at the end containing exit_code, timing, and workspace files.
+        """
+        with self._run_in_container(
+            code=code,
+            cpu_time_limit_sec=cpu_time_limit_sec,
+            memory_limit_mb=memory_limit_mb,
+            timeout_ms=timeout_ms,
+            files=files,
+            last_line_interactive=last_line_interactive,
+        ) as ctx:
+            _write_stdin(ctx.proc, stdin)
+
+            deadline = time.monotonic() + (timeout_ms / 1000.0)
+            timed_out = yield from _stream_process_output(ctx.proc, deadline, max_output_bytes)
+
+            self._terminate_process(ctx, timed_out)
+            workspace_snapshot = self._extract_workspace_snapshot(ctx.container_name)
+
+        duration_ms = int((time.perf_counter() - ctx.start) * 1000)
+        exit_code = None if timed_out else ctx.proc.returncode
+
+        yield StreamResult(
+            exit_code=exit_code,
+            timed_out=timed_out,
+            duration_ms=duration_ms,
+            files=workspace_snapshot,
+        )
+
+
+def _write_stdin(proc: subprocess.Popen[bytes], stdin: str | None) -> None:
+    """Write optional stdin data and close the pipe."""
+    if proc.stdin is None:
+        raise RuntimeError("Failed to open subprocess stdin pipe")
+    if stdin is not None:
+        proc.stdin.write(stdin.encode("utf-8"))
+    proc.stdin.close()
+
+
+class _StreamTracker:
+    """Per-stream state for incremental decoding with truncation."""
+
+    __slots__ = ("stream", "decoder", "bytes_sent", "max_bytes")
+
+    def __init__(self, stream: Literal["stdout", "stderr"], max_bytes: int) -> None:
+        self.stream = stream
+        self.decoder = codecs.getincrementaldecoder("utf-8")("replace")
+        self.bytes_sent = 0
+        self.max_bytes = max_bytes
+
+    def decode_chunk(self, data: bytes) -> StreamChunk | None:
+        """Decode a raw chunk and return a ``StreamChunk`` if within limits."""
+        chunk: StreamChunk | None = None
+        if self.bytes_sent < self.max_bytes:
+            allowed = self.max_bytes - self.bytes_sent
+            text = self.decoder.decode(data[:allowed], False)
+            if text:
+                chunk = StreamChunk(stream=self.stream, data=text)
+        self.bytes_sent += len(data)
+        return chunk
+
+    def flush(self) -> StreamChunk | None:
+        """Flush the decoder and return a final chunk if any bytes remain."""
+        text = self.decoder.decode(b"", True)
+        if text:
+            return StreamChunk(stream=self.stream, data=text)
+        return None
+
+
+def _stream_process_output(
+    proc: subprocess.Popen[bytes],
+    deadline: float,
+    max_output_bytes: int,
+) -> Generator[StreamChunk, None, bool]:
+    """Read stdout/stderr incrementally and yield ``StreamChunk`` events.
+
+    Returns ``True`` if the process timed out, ``False`` otherwise.
+    """
+    if proc.stdout is None or proc.stderr is None:
+        raise RuntimeError("Failed to open subprocess output pipes")
+
+    sel = selectors.DefaultSelector()
+    sel.register(proc.stdout, selectors.EVENT_READ, "stdout")
+    sel.register(proc.stderr, selectors.EVENT_READ, "stderr")
+
+    trackers: dict[str, _StreamTracker] = {
+        "stdout": _StreamTracker("stdout", max_output_bytes),
+        "stderr": _StreamTracker("stderr", max_output_bytes),
+    }
+    fds: dict[str, int] = {
+        "stdout": proc.stdout.fileno(),
+        "stderr": proc.stderr.fileno(),
+    }
+    timed_out = False
+    chunk_size = 4096
+
+    try:
+        while sel.get_map():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                timed_out = True
+                break
+
+            events = sel.select(timeout=min(remaining, 5.0))
+
+            for key, _ in events:
+                stream_name: str = key.data
+                data = os.read(fds[stream_name], chunk_size)
+                if not data:
+                    sel.unregister(key.fileobj)
+                    continue
+
+                chunk = trackers[stream_name].decode_chunk(data)
+                if chunk is not None:
+                    yield chunk
+    finally:
+        sel.close()
+
+    for tracker in trackers.values():
+        chunk = tracker.flush()
+        if chunk is not None:
+            yield chunk
+
+    return timed_out
diff --git a/code-interpreter/app/services/executor_factory.py b/code-interpreter/app/services/executor_factory.py
diff --git a/code-interpreter/tests/integration_tests/test_streaming.py b/code-interpreter/tests/integration_tests/test_streaming.py
