yee88 0.3.0__py3-none-any.whl

yee88/runtime_loader.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+from collections.abc import Iterable, Mapping
+
+from .backends import EngineBackend
+from .config import ConfigError, ProjectsConfig
+from .engines import get_backend, list_backend_ids
+from .ids import RESERVED_CHAT_COMMANDS
+from .logging import get_logger
+from .router import AutoRouter, EngineStatus, RunnerEntry
+from .settings import TakopiSettings
+from .transport_runtime import TransportRuntime
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeSpec:
+    router: AutoRouter
+    projects: ProjectsConfig
+    allowlist: list[str] | None
+    plugin_configs: Mapping[str, Any] | None
+    watch_config: bool = False
+
+    def to_runtime(self, *, config_path: Path | None) -> TransportRuntime:
+        return TransportRuntime(
+            router=self.router,
+            projects=self.projects,
+            allowlist=self.allowlist,
+            config_path=config_path,
+            plugin_configs=self.plugin_configs,
+            watch_config=self.watch_config,
+        )
+
+    def apply(self, runtime: TransportRuntime, *, config_path: Path | None) -> None:
+        runtime.update(
+            router=self.router,
+            projects=self.projects,
+            allowlist=self.allowlist,
+            config_path=config_path,
+            plugin_configs=self.plugin_configs,
+            watch_config=self.watch_config,
+        )
+
+
+def resolve_plugins_allowlist(
+    settings: TakopiSettings | None,
+) -> list[str] | None:
+    if settings is None:
+        return None
+    enabled = list(settings.plugins.enabled)
+    return enabled or None
+
+
+def resolve_default_engine(
+    *,
+    override: str | None,
+    settings: TakopiSettings,
+    config_path: Path,
+    engine_ids: list[str],
+) -> str:
+    default_engine = override or settings.default_engine or "codex"
+    if default_engine not in engine_ids:
+        available = ", ".join(sorted(engine_ids))
+        raise ConfigError(
+            f"Unknown default engine {default_engine!r}. Available: {available}."
+        )
+    return default_engine
+
+
+def build_router(
+    *,
+    settings: TakopiSettings,
+    config_path: Path,
+    backends: list[EngineBackend],
+    default_engine: str,
+) -> AutoRouter:
+    entries: list[RunnerEntry] = []
+    warnings: list[str] = []
+
+    for backend in backends:
+        engine_id = backend.id
+        issue: str | None = None
+        status: EngineStatus = "ok"
+        engine_cfg: dict
+        try:
+            engine_cfg = settings.engine_config(engine_id, config_path=config_path)
+        except ConfigError as exc:
+            if engine_id == default_engine:
+                raise
+            issue = str(exc)
+            status = "bad_config"
+            engine_cfg = {}
+
+        try:
+            runner = backend.build_runner(engine_cfg, config_path)
+        except Exception as exc:
+            if engine_id == default_engine:
+                raise
+            issue = issue or str(exc)
+            if engine_cfg:
+                try:
+                    runner = backend.build_runner({}, config_path)
+                except Exception as fallback_exc:  # noqa: BLE001
+                    warnings.append(f"{engine_id}: {issue or str(fallback_exc)}")
+                    continue
+                status = "bad_config"
+            else:
+                status = "load_error"
+                warnings.append(f"{engine_id}: {issue}")
+                continue
+
+        cmd = backend.cli_cmd or backend.id
+        if shutil.which(cmd) is None:
+            status = "missing_cli"
+            if issue:
+                issue = f"{issue}; {cmd} not found on PATH"
+            else:
+                issue = f"{cmd} not found on PATH"
+
+        if status != "ok" and engine_id == default_engine:
+            raise ConfigError(f"Default engine {engine_id!r} unavailable: {issue}")
+
+        if status != "ok" and engine_id != default_engine:
+            warnings.append(f"{engine_id}: {issue}")
+
+        entries.append(
+            RunnerEntry(
+                engine=engine_id,
+                runner=runner,
+                status=status,
+                issue=issue,
+            )
+        )
+
+    for warning in warnings:
+        logger.warning("setup.warning", issue=warning)
+
+    return AutoRouter(entries=entries, default_engine=default_engine)
+
+
+def load_backends(
+    *,
+    engine_ids: list[str],
+    allowlist: list[str] | None,
+    default_engine: str,
+) -> list[EngineBackend]:
+    backends: list[EngineBackend] = []
+    load_issues: list[str] = []
+    for engine_id in engine_ids:
+        try:
+            backend = get_backend(engine_id, allowlist=allowlist)
+        except ConfigError as exc:
+            if engine_id == default_engine:
+                raise
+            load_issues.append(f"{engine_id}: {exc}")
+            continue
+        backends.append(backend)
+    if not backends:
+        raise ConfigError("No engine backends are available.")
+    for issue in load_issues:
+        logger.warning("setup.warning", issue=issue)
+    return backends
+
+
+def build_runtime_spec(
+    *,
+    settings: TakopiSettings,
+    config_path: Path,
+    default_engine_override: str | None = None,
+    reserved: Iterable[str] = RESERVED_CHAT_COMMANDS,
+) -> RuntimeSpec:
+    allowlist = resolve_plugins_allowlist(settings)
+    engine_ids = list_backend_ids(allowlist=allowlist)
+    projects = settings.to_projects_config(
+        config_path=config_path,
+        engine_ids=engine_ids,
+        reserved=reserved,
+    )
+    default_engine = resolve_default_engine(
+        override=default_engine_override,
+        settings=settings,
+        config_path=config_path,
+        engine_ids=engine_ids,
+    )
+    backends = load_backends(
+        engine_ids=engine_ids,
+        allowlist=allowlist,
+        default_engine=default_engine,
+    )
+    router = build_router(
+        settings=settings,
+        config_path=config_path,
+        backends=backends,
+        default_engine=default_engine,
+    )
+    return RuntimeSpec(
+        router=router,
+        projects=projects,
+        allowlist=allowlist,
+        plugin_configs=settings.plugins.model_extra,
+        watch_config=settings.watch_config,
+    )

yee88/scheduler.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass
+from typing import Any, Protocol
+from collections.abc import Awaitable, Callable
+
+import anyio
+
+from .context import RunContext
+from .logging import get_logger
+from .model import ResumeToken
+from .transport import ChannelId, MessageId, MessageRef, ThreadId
+
+logger = get_logger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class ThreadJob:
+    chat_id: ChannelId
+    user_msg_id: MessageId
+    text: str
+    resume_token: ResumeToken
+    context: RunContext | None = None
+    thread_id: ThreadId | None = None
+    session_key: tuple[int, int | None] | None = None
+    progress_ref: MessageRef | None = None
+
+
+RunJob = Callable[[ThreadJob], Awaitable[None]]
+
+
+class TaskGroup(Protocol):
+    def start_soon(
+        self, func: Callable[..., Awaitable[object]], *args: Any
+    ) -> None: ...
+
+
+class ThreadScheduler:
+    def __init__(self, *, task_group: TaskGroup, run_job: RunJob) -> None:
+        self._task_group = task_group
+        self._run_job = run_job
+        self._lock = anyio.Lock()
+        self._pending_by_thread: dict[str, deque[ThreadJob]] = {}
+        self._queued_by_progress: dict[tuple[ChannelId, MessageId], ThreadJob] = {}
+        self._active_threads: set[str] = set()
+        self._busy_until: dict[str, anyio.Event] = {}
+
+    @staticmethod
+    def thread_key(token: ResumeToken) -> str:
+        return f"{token.engine}:{token.value}"
+
+    async def note_thread_known(self, token: ResumeToken, done: anyio.Event) -> None:
+        key = self.thread_key(token)
+        async with self._lock:
+            current = self._busy_until.get(key)
+            if current is None or current.is_set():
+                self._busy_until[key] = done
+        self._task_group.start_soon(self._clear_busy, key, done)
+
+    async def enqueue(self, job: ThreadJob) -> None:
+        key = self.thread_key(job.resume_token)
+        async with self._lock:
+            queue = self._pending_by_thread.get(key)
+            if queue is None:
+                queue = deque()
+                self._pending_by_thread[key] = queue
+            queue.append(job)
+            if job.progress_ref is not None:
+                progress_key = (job.chat_id, job.progress_ref.message_id)
+                self._queued_by_progress[progress_key] = job
+            if key in self._active_threads:
+                return
+            self._active_threads.add(key)
+        self._task_group.start_soon(self._thread_worker, key)
+
+    async def enqueue_resume(
+        self,
+        chat_id: ChannelId,
+        user_msg_id: MessageId,
+        text: str,
+        resume_token: ResumeToken,
+        context: RunContext | None = None,
+        thread_id: ThreadId | None = None,
+        session_key: tuple[int, int | None] | None = None,
+        progress_ref: MessageRef | None = None,
+    ) -> None:
+        await self.enqueue(
+            ThreadJob(
+                chat_id=chat_id,
+                user_msg_id=user_msg_id,
+                text=text,
+                resume_token=resume_token,
+                context=context,
+                thread_id=thread_id,
+                session_key=session_key,
+                progress_ref=progress_ref,
+            )
+        )
+
+    async def cancel_queued(
+        self, chat_id: ChannelId, progress_msg_id: MessageId
+    ) -> ThreadJob | None:
+        progress_key = (chat_id, progress_msg_id)
+        async with self._lock:
+            job = self._queued_by_progress.pop(progress_key, None)
+            if job is None:
+                return None
+            thread_key = self.thread_key(job.resume_token)
+            queue = self._pending_by_thread.get(thread_key)
+            if queue is None:
+                return None
+            try:
+                queue.remove(job)
+            except ValueError:
+                return None
+            if not queue:
+                self._pending_by_thread.pop(thread_key, None)
+            return job
+
+    async def _clear_busy(self, key: str, done: anyio.Event) -> None:
+        await done.wait()
+        async with self._lock:
+            if self._busy_until.get(key) is done:
+                self._busy_until.pop(key, None)
+
+    async def _thread_worker(self, key: str) -> None:
+        try:
+            while True:
+                async with self._lock:
+                    done = self._busy_until.get(key)
+                    queue = self._pending_by_thread.get(key)
+                    if not queue:
+                        self._pending_by_thread.pop(key, None)
+                        self._active_threads.discard(key)
+                        return
+                    job = queue.popleft()
+                    if job.progress_ref is not None:
+                        progress_key = (job.chat_id, job.progress_ref.message_id)
+                        self._queued_by_progress.pop(progress_key, None)
+
+                if done is not None and not done.is_set():
+                    await done.wait()
+
+                try:
+                    await self._run_job(job)
+                except Exception as exc:  # noqa: BLE001
+                    logger.exception(
+                        "scheduler.job_failed",
+                        key=key,
+                        tag=job.resume_token.engine,
+                        chat_id=job.chat_id,
+                        user_msg_id=job.user_msg_id,
+                        error=str(exc),
+                        error_type=exc.__class__.__name__,
+                    )
+        finally:
+            async with self._lock:
+                self._active_threads.discard(key)

yee88/schemas/__init__.py

@@ -0,0 +1 @@
+"""Event schemas for runner JSONL streams."""

yee88/schemas/claude.py

@@ -0,0 +1,238 @@
+"""Msgspec models and decoder for Claude Code stream-json output."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import msgspec
+
+
+class StreamTextBlock(
+    msgspec.Struct, tag="text", tag_field="type", forbid_unknown_fields=False
+):
+    text: str
+
+
+class StreamThinkingBlock(
+    msgspec.Struct, tag="thinking", tag_field="type", forbid_unknown_fields=False
+):
+    thinking: str
+    signature: str
+
+
+class StreamToolUseBlock(
+    msgspec.Struct, tag="tool_use", tag_field="type", forbid_unknown_fields=False
+):
+    id: str
+    name: str
+    input: dict[str, Any]
+
+
+class StreamToolResultBlock(
+    msgspec.Struct, tag="tool_result", tag_field="type", forbid_unknown_fields=False
+):
+    tool_use_id: str
+    content: str | list[dict[str, Any]] | None = None
+    is_error: bool | None = None
+
+
+type StreamContentBlock = (
+    StreamTextBlock | StreamThinkingBlock | StreamToolUseBlock | StreamToolResultBlock
+)
+
+
+class StreamUserMessageBody(msgspec.Struct, forbid_unknown_fields=False):
+    role: Literal["user"]
+    content: str | list[StreamContentBlock]
+
+
+class StreamAssistantMessageBody(msgspec.Struct, forbid_unknown_fields=False):
+    role: Literal["assistant"]
+    content: list[StreamContentBlock]
+    model: str
+    error: str | None = None
+
+
+class StreamUserMessage(
+    msgspec.Struct, tag="user", tag_field="type", forbid_unknown_fields=False
+):
+    message: StreamUserMessageBody
+    uuid: str | None = None
+    parent_tool_use_id: str | None = None
+    session_id: str | None = None
+
+
+class StreamAssistantMessage(
+    msgspec.Struct, tag="assistant", tag_field="type", forbid_unknown_fields=False
+):
+    message: StreamAssistantMessageBody
+    parent_tool_use_id: str | None = None
+    uuid: str | None = None
+    session_id: str | None = None
+
+
+class StreamSystemMessage(
+    msgspec.Struct, tag="system", tag_field="type", forbid_unknown_fields=False
+):
+    subtype: str
+    session_id: str | None = None
+    uuid: str | None = None
+    cwd: str | None = None
+    tools: list[str] | None = None
+    mcp_servers: list[Any] | None = None
+    model: str | None = None
+    permissionMode: str | None = None
+    output_style: str | None = None
+    apiKeySource: str | None = None
+
+
+class StreamResultMessage(
+    msgspec.Struct, tag="result", tag_field="type", forbid_unknown_fields=False
+):
+    subtype: str
+    duration_ms: int
+    duration_api_ms: int
+    is_error: bool
+    num_turns: int
+    session_id: str
+    total_cost_usd: float | None = None
+    usage: dict[str, Any] | None = None
+    result: str | None = None
+    structured_output: Any = None
+
+
+class StreamEventMessage(
+    msgspec.Struct, tag="stream_event", tag_field="type", forbid_unknown_fields=False
+):
+    uuid: str
+    session_id: str
+    event: dict[str, Any]
+    parent_tool_use_id: str | None = None
+
+
+class ControlInterruptRequest(
+    msgspec.Struct, tag="interrupt", tag_field="subtype", forbid_unknown_fields=False
+):
+    pass
+
+
+class ControlCanUseToolRequest(
+    msgspec.Struct, tag="can_use_tool", tag_field="subtype", forbid_unknown_fields=False
+):
+    tool_name: str
+    input: dict[str, Any]
+    permission_suggestions: list[Any] | None = None
+    blocked_path: str | None = None
+
+
+class ControlInitializeRequest(
+    msgspec.Struct, tag="initialize", tag_field="subtype", forbid_unknown_fields=False
+):
+    hooks: dict[str, Any] | None = None
+
+
+class ControlSetPermissionModeRequest(
+    msgspec.Struct,
+    tag="set_permission_mode",
+    tag_field="subtype",
+    forbid_unknown_fields=False,
+):
+    mode: str
+
+
+class ControlHookCallbackRequest(
+    msgspec.Struct,
+    tag="hook_callback",
+    tag_field="subtype",
+    forbid_unknown_fields=False,
+):
+    callback_id: str
+    input: Any
+    tool_use_id: str | None = None
+
+
+class ControlMcpMessageRequest(
+    msgspec.Struct, tag="mcp_message", tag_field="subtype", forbid_unknown_fields=False
+):
+    server_name: str
+    message: Any
+
+
+class ControlRewindFilesRequest(
+    msgspec.Struct, tag="rewind_files", tag_field="subtype", forbid_unknown_fields=False
+):
+    user_message_id: str
+
+
+type ControlRequest = (
+    ControlInterruptRequest
+    | ControlCanUseToolRequest
+    | ControlInitializeRequest
+    | ControlSetPermissionModeRequest
+    | ControlHookCallbackRequest
+    | ControlMcpMessageRequest
+    | ControlRewindFilesRequest
+)
+
+
+class StreamControlRequest(
+    msgspec.Struct, tag="control_request", tag_field="type", forbid_unknown_fields=False
+):
+    request_id: str
+    request: ControlRequest
+
+
+class ControlSuccessResponse(
+    msgspec.Struct, tag="success", tag_field="subtype", forbid_unknown_fields=False
+):
+    request_id: str
+    response: dict[str, Any] | None = None
+
+
+class ControlErrorResponse(
+    msgspec.Struct, tag="error", tag_field="subtype", forbid_unknown_fields=False
+):
+    request_id: str
+    error: str
+
+
+type ControlResponse = ControlSuccessResponse | ControlErrorResponse
+
+
+class StreamControlResponse(
+    msgspec.Struct,
+    tag="control_response",
+    tag_field="type",
+    forbid_unknown_fields=False,
+):
+    response: ControlResponse
+
+
+class StreamControlCancelRequest(
+    msgspec.Struct,
+    tag="control_cancel_request",
+    tag_field="type",
+    forbid_unknown_fields=False,
+):
+    request_id: str | None = None
+
+
+type StreamJsonMessage = (
+    StreamUserMessage
+    | StreamAssistantMessage
+    | StreamSystemMessage
+    | StreamResultMessage
+    | StreamEventMessage
+    | StreamControlRequest
+    | StreamControlResponse
+    | StreamControlCancelRequest
+)
+
+
+STREAM_JSON_SCHEMA = msgspec.json.schema(StreamJsonMessage)
+
+_DECODER = msgspec.json.Decoder(StreamJsonMessage)
+
+
+def decode_stream_json_line(line: str | bytes) -> StreamJsonMessage:
+    return _DECODER.decode(line)