codex-autorunner 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

codex_autorunner/integrations/agents/opencode_backend.py

@@ -1,18 +1,41 @@
+import asyncio
+import contextlib
 import json
 import logging
+from pathlib import Path
 from typing import Any, AsyncGenerator, Dict, Optional
 
 from ...agents.opencode.client import OpenCodeClient
 from ...agents.opencode.events import SSEEvent
-from .agent_backend import AgentBackend, AgentEvent, AgentEventType, now_iso
-from .run_event import (
+from ...agents.opencode.logging import OpenCodeEventFormatter
+from ...agents.opencode.runtime import (
+    OpenCodeTurnOutput,
+    build_turn_id,
+    collect_opencode_output,
+    extract_session_id,
+    map_approval_policy_to_permission,
+    opencode_missing_env,
+    parse_message_response,
+    split_model_id,
+)
+from ...agents.opencode.supervisor import OpenCodeSupervisor
+from ...core.ports.agent_backend import (
+    AgentBackend,
+    AgentEvent,
+    AgentEventType,
+    now_iso,
+)
+from ...core.ports.run_event import (
     Completed,
     Failed,
     OutputDelta,
     RunEvent,
+    RunNotice,
     Started,
+    TokenUsage,
     ToolCall,
 )
+from ...core.text_delta_coalescer import StreamingTextCoalescer
 
 _logger = logging.getLogger(__name__)
 
@@ -20,32 +43,74 @@ _logger = logging.getLogger(__name__)
 class OpenCodeBackend(AgentBackend):
     def __init__(
         self,
-        base_url: str,
         *,
+        base_url: Optional[str] = None,
+        supervisor: Optional[OpenCodeSupervisor] = None,
+        workspace_root: Optional[Path] = None,
         auth: Optional[tuple[str, str]] = None,
         timeout: Optional[float] = None,
         agent: Optional[str] = None,
-        model: Optional[dict[str, str]] = None,
+        model: Optional[str] = None,
+        reasoning: Optional[str] = None,
+        approval_policy: Optional[str] = None,
+        session_stall_timeout_seconds: Optional[float] = None,
+        logger: Optional[logging.Logger] = None,
     ):
-        self._client = OpenCodeClient(
-            base_url=base_url,
-            auth=auth,
-            timeout=timeout,
-        )
+        self._supervisor = supervisor
+        self._workspace_root = Path(workspace_root) if workspace_root else None
+        self._client: Optional[OpenCodeClient]
+        if base_url:
+            self._client = OpenCodeClient(
+                base_url=base_url,
+                auth=auth,
+                timeout=timeout,
+                logger=logger,
+            )
+        else:
+            self._client = None
         self._agent = agent
         self._model = model
+        self._reasoning = reasoning
+        self._approval_policy = approval_policy
+        self._session_stall_timeout_seconds = session_stall_timeout_seconds
+        self._logger = logger or _logger
 
         self._session_id: Optional[str] = None
         self._message_count: int = 0
         self._final_messages: list[str] = []
+        self._last_turn_id: Optional[str] = None
+        self._last_token_total: Optional[dict[str, Any]] = None
+        self._event_formatter = OpenCodeEventFormatter()
+
+    def configure(
+        self,
+        *,
+        model: Optional[str],
+        reasoning: Optional[str],
+        approval_policy: Optional[str],
+    ) -> None:
+        self._model = model
+        self._reasoning = reasoning
+        self._approval_policy = approval_policy
 
     async def start_session(self, target: dict, context: dict) -> str:
-        result = await self._client.create_session(
-            title=f"Flow session {self._message_count}",
-            directory=None,
-        )
+        client = await self._ensure_client()
+        workspace_root = self._workspace_root or Path(context.get("workspace") or ".")
+        resume_session = context.get("session_id") or context.get("thread_id")
+        if isinstance(resume_session, str) and resume_session:
+            try:
+                await client.get_session(resume_session)
+                self._session_id = resume_session
+            except Exception:
+                self._session_id = None
+
+        if not self._session_id:
+            result = await client.create_session(
+                title=f"Flow session {self._message_count}",
+                directory=str(workspace_root),
+            )
+            self._session_id = extract_session_id(result, allow_fallback_id=True)
 
-        self._session_id = result.get("id")
         if not self._session_id:
             raise RuntimeError("Failed to create OpenCode session: missing session ID")
 
@@ -56,6 +121,7 @@ class OpenCodeBackend(AgentBackend):
     async def run_turn(
         self, session_id: str, message: str
     ) -> AsyncGenerator[AgentEvent, None]:
+        client = await self._ensure_client()
         if session_id:
             self._session_id = session_id
         if not self._session_id:
@@ -65,11 +131,11 @@ class OpenCodeBackend(AgentBackend):
 
         yield AgentEvent.stream_delta(content=message, delta_type="user_message")
 
-        await self._client.send_message(
+        await client.send_message(
             self._session_id,
             message=message,
             agent=self._agent,
-            model=self._model,
+            model=split_model_id(self._model) if self._model else None,
         )
 
         self._message_count += 1
@@ -79,34 +145,192 @@ class OpenCodeBackend(AgentBackend):
     async def run_turn_events(
         self, session_id: str, message: str
     ) -> AsyncGenerator[RunEvent, None]:
+        client = await self._ensure_client()
+        workspace_root = self._workspace_root or Path(".")
+
         if session_id:
             self._session_id = session_id
         if not self._session_id:
-            self._session_id = await self.start_session(target={}, context={})
+            self._session_id = await self.start_session(
+                target={},
+                context={"workspace": str(workspace_root)},
+            )
 
         _logger.info("Running turn events on session %s", self._session_id)
 
+        self._last_turn_id = build_turn_id(self._session_id)
+
         yield Started(timestamp=now_iso(), session_id=self._session_id)
 
         yield OutputDelta(
             timestamp=now_iso(), content=message, delta_type="user_message"
         )
 
-        await self._client.send_message(
-            self._session_id,
-            message=message,
-            agent=self._agent,
-            model=self._model,
+        model_payload = split_model_id(self._model) if self._model else None
+        missing_env = await opencode_missing_env(
+            client, str(workspace_root), model_payload
         )
+        if missing_env:
+            provider_id = model_payload.get("providerID") if model_payload else None
+            missing_label = ", ".join(missing_env)
+            yield Failed(
+                timestamp=now_iso(),
+                error_message=(
+                    "OpenCode provider "
+                    f"{provider_id or 'selected'} requires env vars: {missing_label}"
+                ),
+            )
+            return
 
-        self._message_count += 1
+        permission_policy = map_approval_policy_to_permission(
+            self._approval_policy, default="allow"
+        )
+
+        event_queue: asyncio.Queue[RunEvent] = asyncio.Queue()
+        self._event_formatter.reset()
+        assistant_stream_coalescer = StreamingTextCoalescer()
+
+        async def _enqueue_lines(lines: list[str]) -> None:
+            for line in lines:
+                await event_queue.put(
+                    OutputDelta(
+                        timestamp=now_iso(), content=line, delta_type="log_line"
+                    )
+                )
+
+        async def _part_handler(
+            part_type: str, part: dict[str, Any], delta_text: Optional[str]
+        ) -> None:
+            if part_type == "usage" and isinstance(part, dict):
+                self._last_token_total = _usage_to_token_total(part)
+                await event_queue.put(TokenUsage(timestamp=now_iso(), usage=dict(part)))
+                await _enqueue_lines(self._event_formatter.format_usage(part))
+            else:
+                await _enqueue_lines(
+                    self._event_formatter.format_part(part_type, part, delta_text)
+                )
+            if part_type == "text" and isinstance(delta_text, str) and delta_text:
+                for chunk in assistant_stream_coalescer.add(delta_text):
+                    await event_queue.put(
+                        OutputDelta(
+                            timestamp=now_iso(),
+                            content=chunk,
+                            delta_type="assistant_stream",
+                        )
+                    )
+
+        ready_event = asyncio.Event()
+        output_task = asyncio.create_task(
+            collect_opencode_output(
+                client,
+                session_id=self._session_id,
+                workspace_path=str(workspace_root),
+                model_payload=model_payload,
+                permission_policy=permission_policy,
+                part_handler=_part_handler,
+                ready_event=ready_event,
+                stall_timeout_seconds=self._session_stall_timeout_seconds,
+            )
+        )
+        try:
+            await asyncio.wait_for(ready_event.wait(), timeout=2.0)
+        except asyncio.TimeoutError:
+            await event_queue.put(
+                RunNotice(
+                    timestamp=now_iso(),
+                    kind="ready_timeout",
+                    message="OpenCode stream readiness wait timed out",
+                    data={"timeout_seconds": 2.0},
+                )
+            )
+
+        prompt_response: Any = None
+        prompt_task: Optional[asyncio.Task[Any]] = asyncio.create_task(
+            client.prompt_async(
+                self._session_id,
+                message=message,
+                agent=self._agent,
+                model=model_payload,
+                variant=self._reasoning,
+            )
+        )
 
+        output_result = None
         try:
-            async for run_event in self._yield_run_events_until_completion():
-                yield run_event
-        except Exception as e:
-            _logger.error("Error during turn execution: %s", e)
-            yield Failed(timestamp=now_iso(), error_message=str(e))
+            while True:
+                queue_task = asyncio.create_task(event_queue.get())
+                tasks = {output_task, queue_task}
+                if prompt_task is not None:
+                    tasks.add(prompt_task)
+                done, pending = await asyncio.wait(
+                    tasks, return_when=asyncio.FIRST_COMPLETED
+                )
+
+                if queue_task in done:
+                    yield queue_task.result()
+                else:
+                    queue_task.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await queue_task
+
+                if prompt_task is not None and prompt_task in done:
+                    try:
+                        prompt_response = prompt_task.result()
+                    except Exception as exc:
+                        output_task.cancel()
+                        with contextlib.suppress(asyncio.CancelledError):
+                            await output_task
+                        yield Failed(timestamp=now_iso(), error_message=str(exc))
+                        return
+                    prompt_task = None
+
+                if output_task in done:
+                    output_result = await output_task
+                    break
+
+        finally:
+            if prompt_task is not None:
+                with contextlib.suppress(asyncio.CancelledError):
+                    await prompt_task
+            for line in self._event_formatter.flush_all_reasoning():
+                await event_queue.put(
+                    OutputDelta(
+                        timestamp=now_iso(), content=line, delta_type="log_line"
+                    )
+                )
+            for chunk in assistant_stream_coalescer.flush():
+                await event_queue.put(
+                    OutputDelta(
+                        timestamp=now_iso(),
+                        content=chunk,
+                        delta_type="assistant_stream",
+                    )
+                )
+
+        while not event_queue.empty():
+            yield event_queue.get_nowait()
+
+        if output_result is None:
+            yield Failed(timestamp=now_iso(), error_message="OpenCode output failed")
+            return
+
+        if prompt_response is not None and not output_result.text:
+            fallback = parse_message_response(prompt_response)
+            if fallback.text:
+                output_result = OpenCodeTurnOutput(
+                    text=fallback.text, error=output_result.error
+                )
+            if fallback.error and not output_result.error:
+                output_result = OpenCodeTurnOutput(
+                    text=output_result.text, error=fallback.error
+                )
+
+        if output_result.text:
+            yield Completed(timestamp=now_iso(), final_message=output_result.text)
+        elif output_result.error:
+            yield Failed(timestamp=now_iso(), error_message=output_result.error)
+        else:
+            yield Completed(timestamp=now_iso(), final_message="")
 
     async def stream_events(self, session_id: str) -> AsyncGenerator[AgentEvent, None]:
         if session_id:
@@ -114,15 +338,17 @@ class OpenCodeBackend(AgentBackend):
         if not self._session_id:
             raise RuntimeError("Session not started. Call start_session() first.")
 
-        async for sse in self._client.stream_events(directory=None):
+        client = await self._ensure_client()
+        async for sse in client.stream_events(directory=None):
             for agent_event in self._convert_sse_to_agent_event(sse):
                 yield agent_event
 
     async def interrupt(self, session_id: str) -> None:
         target_session = session_id or self._session_id
         if target_session:
+            client = await self._ensure_client()
             try:
-                await self._client.abort(target_session)
+                await client.abort(target_session)
                 _logger.info("Interrupted OpenCode session %s", target_session)
             except Exception as e:
                 _logger.warning("Failed to interrupt session: %s", e)
@@ -140,7 +366,8 @@ class OpenCodeBackend(AgentBackend):
         if self._session_id:
             paths.insert(0, f"/session/{self._session_id}/event")
         try:
-            async for sse in self._client.stream_events(
+            client = await self._ensure_client()
+            async for sse in client.stream_events(
                 directory=None,
                 paths=paths,
             ):
@@ -168,7 +395,8 @@ class OpenCodeBackend(AgentBackend):
         if self._session_id:
             paths.insert(0, f"/session/{self._session_id}/event")
         try:
-            async for sse in self._client.stream_events(
+            client = await self._ensure_client()
+            async for sse in client.stream_events(
                 directory=None,
                 paths=paths,
             ):
@@ -323,3 +551,48 @@ class OpenCodeBackend(AgentBackend):
             # If server does not tag events, do not drop them.
             return True
         return session_id == self._session_id
+
+    async def _ensure_client(self) -> OpenCodeClient:
+        if self._client is not None:
+            return self._client
+        if self._supervisor is None or self._workspace_root is None:
+            raise RuntimeError("OpenCode client unavailable: supervisor not configured")
+        client = await self._supervisor.get_client(self._workspace_root)
+        self._client = client
+        return client
+
+    @property
+    def last_turn_id(self) -> Optional[str]:
+        return self._last_turn_id
+
+    @property
+    def last_token_total(self) -> Optional[dict[str, Any]]:
+        return self._last_token_total
+
+
+def _usage_to_token_total(usage: dict[str, Any]) -> Optional[dict[str, int]]:
+    if not isinstance(usage, dict):
+        return None
+
+    def _int(key: str) -> int:
+        value = usage.get(key)
+        return int(value) if isinstance(value, (int, float)) else 0
+
+    total = usage.get("totalTokens")
+    total_tokens = int(total) if isinstance(total, (int, float)) else None
+    input_tokens = _int("inputTokens")
+    cached_tokens = _int("cachedInputTokens")
+    output_tokens = _int("outputTokens")
+    reasoning_tokens = _int("reasoningTokens")
+    if total_tokens is None:
+        total_tokens = input_tokens + cached_tokens + output_tokens + reasoning_tokens
+    return {
+        "total": total_tokens,
+        "input_tokens": input_tokens,
+        "prompt_tokens": input_tokens,
+        "cached_input_tokens": cached_tokens,
+        "output_tokens": output_tokens,
+        "completion_tokens": output_tokens,
+        "reasoning_tokens": reasoning_tokens,
+        "reasoning_output_tokens": reasoning_tokens,
+    }

codex_autorunner/integrations/agents/runner.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import logging
+from typing import AsyncGenerator, Callable, Optional
+
+from ...core.ports.agent_backend import AgentBackend, AgentEvent
+from ...core.ports.run_event import (
+    Failed,
+    OutputDelta,
+    RunEvent,
+    Started,
+    now_iso,
+)
+
+_logger = logging.getLogger(__name__)
+
+LogHandler = Callable[[str], None]
+EventCallback = Callable[[RunEvent], None]
+
+
+async def run_turn_with_backend(
+    backend: AgentBackend,
+    message: str,
+    session_id: Optional[str],
+    *,
+    log_handler: Optional[LogHandler] = None,
+    event_callback: Optional[EventCallback] = None,
+) -> int:
+    """
+    Execute a turn using the AgentBackend protocol.
+
+    Returns:
+        Exit code (0 for success, non-zero for failure)
+    """
+    try:
+        if not session_id:
+            session_id = await backend.start_session(target={}, context={})
+
+        if event_callback:
+            event_callback(Started(timestamp=now_iso(), session_id=session_id))
+
+        if log_handler:
+            log_handler(message)
+
+        events_consumed = False
+        if hasattr(backend, "run_turn_events"):
+            async for run_event in backend.run_turn_events(session_id, message):
+                events_consumed = True
+                if event_callback:
+                    event_callback(run_event)
+                if log_handler and isinstance(run_event, OutputDelta):
+                    log_handler(run_event.content)
+
+        if not events_consumed:
+            async for agent_event in backend.run_turn(session_id, message):
+                if isinstance(agent_event, AgentEvent):
+                    if log_handler:
+                        if agent_event.data.get("content"):
+                            log_handler(agent_event.data["content"])
+                elif isinstance(agent_event, str):
+                    if log_handler:
+                        log_handler(agent_event)
+
+        return 0
+    except Exception as exc:
+        _logger.error("Turn execution failed: %s", exc)
+        if event_callback:
+            event_callback(Failed(timestamp=now_iso(), error_message=str(exc)))
+        return 1
+
+
+async def stream_turn_events(
+    backend: AgentBackend,
+    session_id: str,
+) -> AsyncGenerator[AgentEvent, None]:
+    """
+    Stream events from a backend for an existing session.
+
+    This is used for external streaming (e.g., WebSocket UI) where the turn
+    has already been initiated and we want to stream events as they arrive.
+    """
+    if hasattr(backend, "stream_events"):
+        async for event in backend.stream_events(session_id):
+            yield event
+    else:
+        yield AgentEvent.stream_delta(content="", delta_type="noop")