oh-my-pi/python/omp-rpc/src/omp_rpc/client.py

from __future__ import annotations

import json
import os
import queue
import subprocess
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Generic, Mapping, Sequence, TypeVar, cast

from .host_tools import HostTool, HostToolContext
from .protocol import (
    AgentStartEvent,
    AgentEndEvent,
    AgentMessage,
    AssistantMessage,
    AutoCompactionEndEvent,
    AutoCompactionStartEvent,
    AutoRetryEndEvent,
    AutoRetryStartEvent,
    BashResult,
    BranchMessage,
    BranchResult,
    CancellationResult,
    CompactionResult,
    ExtensionError,
    ExtensionUiRequest,
    ImageContent,
    InterruptMode,
    JsonObject,
    JsonValue,
    MessageEndEvent,
    MessageStartEvent,
    MessageUpdateEvent,
    ModelCycleResult,
    ModelInfo,
    ReadyEvent,
    RetryFallbackAppliedEvent,
    RetryFallbackSucceededEvent,
    RpcAgentEvent,
    RpcNotification,
    SessionState,
    SessionStats,
    SteeringMode,
    StreamingBehavior,
    ThinkingLevel,
    ThinkingLevelCycleResult,
    TodoItem,
    TodoPhase,
    TodoStatus,
    TodoAutoClearEvent,
    TodoReminderEvent,
    ToolExecutionEndEvent,
    ToolExecutionStartEvent,
    ToolExecutionUpdateEvent,
    TtsrTriggeredEvent,
    TurnEndEvent,
    TurnStartEvent,
    UnknownNotification,
    assistant_text,
    parse_agent_messages,
    parse_bash_result,
    parse_branch_messages,
    parse_branch_result,
    parse_cancellation_result,
    parse_compaction_result,
    parse_model_cycle_result,
    parse_model_info,
    parse_notification,
    parse_session_state,
    parse_session_stats,
    parse_thinking_level_cycle_result,
    parse_todo_phases,
)

AgentEventListener = Callable[[RpcAgentEvent], None]
NotificationListener = Callable[[RpcNotification], None]
UiRequestListener = Callable[[ExtensionUiRequest], None]
ExtensionErrorListener = Callable[[ExtensionError], None]
ReadyListener = Callable[[ReadyEvent], None]
UnknownNotificationListener = Callable[[UnknownNotification], None]
AgentStartListener = Callable[[AgentStartEvent], None]
AgentEndListener = Callable[[AgentEndEvent], None]
TurnStartListener = Callable[[TurnStartEvent], None]
TurnEndListener = Callable[[TurnEndEvent], None]
MessageStartListener = Callable[[MessageStartEvent], None]
MessageUpdateListener = Callable[[MessageUpdateEvent], None]
MessageEndListener = Callable[[MessageEndEvent], None]
ToolExecutionStartListener = Callable[[ToolExecutionStartEvent], None]
ToolExecutionUpdateListener = Callable[[ToolExecutionUpdateEvent], None]
ToolExecutionEndListener = Callable[[ToolExecutionEndEvent], None]
AutoCompactionStartListener = Callable[[AutoCompactionStartEvent], None]
AutoCompactionEndListener = Callable[[AutoCompactionEndEvent], None]
AutoRetryStartListener = Callable[[AutoRetryStartEvent], None]
AutoRetryEndListener = Callable[[AutoRetryEndEvent], None]
RetryFallbackAppliedListener = Callable[[RetryFallbackAppliedEvent], None]
RetryFallbackSucceededListener = Callable[[RetryFallbackSucceededEvent], None]
TtsrTriggeredListener = Callable[[TtsrTriggeredEvent], None]
TodoReminderListener = Callable[[TodoReminderEvent], None]
TodoAutoClearListener = Callable[[TodoAutoClearEvent], None]
ProtocolErrorListener = Callable[["RpcProtocolError"], None]
ListenerErrorListener = Callable[["ListenerErrorEvent"], None]
TListener = TypeVar("TListener")
TEventListener = TypeVar("TEventListener", bound=Callable[..., None])
THistoryItem = TypeVar("THistoryItem")

_ASYNC_COMMANDS = frozenset({"prompt", "abort_and_prompt"})
_DEFAULT_ERROR_HISTORY_LIMIT = 128
_TODO_STATUS_VALUES = frozenset({"pending", "in_progress", "completed", "abandoned"})


def _clone_json_value(value: object) -> JsonValue:
    if value is None or isinstance(value, (str, int, float, bool)):
        return cast(JsonValue, value)
    if isinstance(value, list):
        return [_clone_json_value(item) for item in value]
    if isinstance(value, dict):
        cloned: JsonObject = {}
        for key, item in value.items():
            if not isinstance(key, str):
                raise RpcError("RPC payload objects must use string keys")
            cloned[key] = _clone_json_value(item)
        return cloned
    raise RpcError("RPC payload must be JSON-serializable")


def _clone_json_object(value: object) -> JsonObject:
    if not isinstance(value, dict):
        raise RpcError("RPC response payload must be an object")
    return cast(JsonObject, _clone_json_value(value))


class RpcError(RuntimeError):
    """Base exception for the Python RPC client."""


class RpcTimeoutError(RpcError):
    """Raised when the server does not respond before a timeout."""


class RpcProcessExitError(RpcError):
    """Raised when the RPC process exits while a request is pending."""


class RpcConcurrencyError(RpcError):
    """Raised when overlapping prompt lifecycle collectors would be ambiguous."""


class RpcCommandError(RpcError):
    """Raised when the RPC server returns `success: false`."""

    def __init__(self, command: str, error: str):
        super().__init__(f"{command}: {error}")
        self.command = command
        self.error = error


class RpcProtocolError(RpcError):
    """Raised or reported when the transport receives an unmatched RPC error response."""

    def __init__(self, payload: JsonObject):
        self.payload = dict(payload)
        command = payload.get("command")
        request_id = payload.get("id")
        error = payload.get("error")
        self.command = str(command) if isinstance(command, str) else None
        self.request_id = str(request_id) if isinstance(request_id, str) else None
        self.remote_error = str(error) if isinstance(error, str) else None

        fragments = ["Received unmatched RPC error response"]
        if self.command:
            fragments.append(f"for {self.command}")
        if self.request_id:
            fragments.append(f"(id={self.request_id})")
        if self.remote_error:
            fragments.append(f": {self.remote_error}")
        super().__init__(" ".join(fragments))


@dataclass(slots=True, frozen=True)
class ListenerErrorEvent:
    listener_kind: str
    source_type: str | None
    listener: Callable[..., None]
    error: BaseException


@dataclass(slots=True, frozen=True)
class PromptTurn:
    events: tuple[RpcAgentEvent, ...]
    messages: tuple[AgentMessage, ...]
    assistant_message: AssistantMessage | None
    assistant_text: str | None

    def require_assistant_text(self) -> str:
        if self.assistant_text is None:
            raise RpcError("Prompt completed without a text assistant message")
        return self.assistant_text


TodoSeed = str | TodoItem | Mapping[str, object]
TodoPhaseSeed = TodoPhase | Mapping[str, object]


@dataclass(slots=True)
class _PendingRequest:
    command: str
    response_queue: queue.Queue[JsonObject | BaseException]


@dataclass(slots=True)
class _PendingHostToolCall:
    cancel_event: threading.Event


@dataclass(slots=True)
class _BoundedHistory(Generic[THistoryItem]):
    limit: int | None
    items: list[THistoryItem] = field(default_factory=list)
    offset: int = 0

    def clear(self) -> None:
        self.items.clear()
        self.offset = 0

    def append(self, item: THistoryItem) -> None:
        self.items.append(item)
        if self.limit is not None and len(self.items) > self.limit:
            trim = len(self.items) - self.limit
            del self.items[:trim]
            self.offset += trim

    def current_index(self) -> int:
        return self.offset + len(self.items)

    def snapshot(self) -> tuple[THistoryItem, ...]:
        return tuple(self.items)

    def snapshot_from(self, start_index: int) -> tuple[THistoryItem, ...]:
        return tuple(self.items[start_index - self.offset :])


@dataclass(slots=True)
class _PromptLifecycleCoordinator:
    lock: threading.Lock = field(default_factory=threading.Lock)
    active_operation: str | None = None

    def acquire(self, operation: str) -> None:
        with self.lock:
            if self.active_operation is not None:
                raise RpcConcurrencyError(
                    f"Cannot start {operation} while {self.active_operation} is already collecting prompt lifecycle events"
                )
            self.active_operation = operation

    def release(self, operation: str) -> None:
        with self.lock:
            if self.active_operation == operation:
                self.active_operation = None


class RpcClient:
    def __init__(
        self,
        *,
        command: Sequence[str] | None = None,
        executable: str = "omp",
        provider: str | None = None,
        model: str | None = None,
        session_dir: str | Path | None = None,
        cwd: str | Path | None = None,
        env: Mapping[str, str] | None = None,
        thinking: ThinkingLevel | None = None,
        append_system_prompt: str | None = None,
        provider_session_id: str | None = None,
        tools: Sequence[str] | None = None,
        custom_tools: Sequence[HostTool[Any, Any]] | None = None,
        no_session: bool = False,
        no_skills: bool = False,
        no_rules: bool = False,
        no_title: bool | None = None,
        rpc_defaults: bool = True,
        extra_args: Sequence[str] = (),
        startup_timeout: float = 30.0,
        request_timeout: float = 30.0,
        max_event_history: int | None = 10_000,
        max_stderr_chunks: int | None = 512,
    ) -> None:
        self._command = tuple(command) if command is not None else None
        self._executable = executable
        self._provider = provider
        self._model = model
        self._session_dir = Path(session_dir) if session_dir is not None else None
        self._cwd = Path(cwd) if cwd is not None else None
        self._env = dict(env or {})
        self._thinking = thinking
        self._append_system_prompt = append_system_prompt
        self._provider_session_id = provider_session_id
        self._tools = tuple(tools) if tools is not None else None
        self._custom_tools = tuple(custom_tools) if custom_tools is not None else ()
        self._no_session = no_session
        self._no_skills = no_skills
        self._no_rules = no_rules
        self._no_title = no_title
        self._rpc_defaults = rpc_defaults
        self._extra_args = tuple(extra_args)
        self._startup_timeout = startup_timeout
        self._request_timeout = request_timeout
        self._max_event_history = self._validate_history_limit("max_event_history", max_event_history)
        self._max_stderr_chunks = self._validate_history_limit("max_stderr_chunks", max_stderr_chunks)

        self._process: subprocess.Popen[str] | None = None
        self._stdout_thread: threading.Thread | None = None
        self._stderr_thread: threading.Thread | None = None
        self._ready = threading.Event()
        self._write_lock = threading.Lock()
        self._state_lock = threading.Lock()
        self._event_condition = threading.Condition()
        self._pending: dict[str, _PendingRequest] = {}
        self._pending_host_tool_calls: dict[str, _PendingHostToolCall] = {}
        self._request_id = 0
        self._events = _BoundedHistory[JsonObject](self._max_event_history)
        self._async_errors = _BoundedHistory[BaseException](_DEFAULT_ERROR_HISTORY_LIMIT)
        self._scheduled_agent_runs = 0
        self._completed_agent_runs = 0
        self._last_schedule_async_error_index = 0
        self._ui_requests: queue.Queue[ExtensionUiRequest] = queue.Queue()
        self._stderr_chunks = _BoundedHistory[str](self._max_stderr_chunks)
        self._closed_error: BaseException | None = None
        self._stopping = False
        self._ready_received = False
        self._protocol_errors = _BoundedHistory[RpcProtocolError](_DEFAULT_ERROR_HISTORY_LIMIT)
        self._listener_errors = _BoundedHistory[ListenerErrorEvent](_DEFAULT_ERROR_HISTORY_LIMIT)
        self._prompt_lifecycle = _PromptLifecycleCoordinator()

        self._notification_listeners: list[NotificationListener] = []
        self._event_listeners: list[AgentEventListener] = []
        self._typed_event_listeners: dict[str, list[AgentEventListener]] = {}
        self._ready_listeners: list[ReadyListener] = []
        self._unknown_notification_listeners: list[UnknownNotificationListener] = []
        self._ui_request_listeners: list[UiRequestListener] = []
        self._extension_error_listeners: list[ExtensionErrorListener] = []
        self._protocol_error_listeners: list[ProtocolErrorListener] = []
        self._listener_error_listeners: list[ListenerErrorListener] = []

    def __enter__(self) -> RpcClient:
        return self.start()

    def __exit__(self, _exc_type: object, _exc: object, _tb: object) -> None:
        self.stop()

    @property
    def stderr(self) -> str:
        with self._state_lock:
            return "".join(self._stderr_chunks.snapshot())

    @property
    def command(self) -> tuple[str, ...]:
        return self._build_command()

    @property
    def protocol_errors(self) -> tuple[RpcProtocolError, ...]:
        with self._state_lock:
            return self._protocol_errors.snapshot()

    @property
    def listener_errors(self) -> tuple[ListenerErrorEvent, ...]:
        with self._state_lock:
            return self._listener_errors.snapshot()

    def start(self) -> RpcClient:
        if self._process is not None:
            raise RpcError("RPC client is already started")

        self._ready.clear()
        self._stopping = False
        self._closed_error = None
        self._ready_received = False
        self._events.clear()
        self._async_errors.clear()
        self._scheduled_agent_runs = 0
        self._completed_agent_runs = 0
        self._last_schedule_async_error_index = 0
        self._ui_requests = queue.Queue()
        with self._state_lock:
            self._stderr_chunks.clear()
        with self._state_lock:
            self._protocol_errors.clear()
            self._listener_errors.clear()

        process = subprocess.Popen(
            list(self._build_command()),
            cwd=str(self._cwd) if self._cwd is not None else None,
            env={**os.environ, **self._env},
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
            bufsize=1,
        )
        self._process = process

        self._stdout_thread = threading.Thread(target=self._read_stdout_loop, name="omp-rpc-stdout", daemon=True)
        self._stderr_thread = threading.Thread(target=self._read_stderr_loop, name="omp-rpc-stderr", daemon=True)
        self._stdout_thread.start()
        self._stderr_thread.start()

        if not self._ready.wait(self._startup_timeout):
            stderr = self.stderr
            self.stop()
            raise RpcTimeoutError(f"Timed out waiting for RPC ready signal. Stderr: {stderr}")

        if not self._ready_received:
            error = self._closed_error
            stderr = self.stderr
            self.stop()
            if isinstance(error, RpcError):
                raise error
            if error is not None:
                raise RpcProcessExitError(f"RPC process stopped before ready: {error}. Stderr: {stderr}") from error
            raise RpcTimeoutError(f"Timed out waiting for RPC ready signal. Stderr: {stderr}")

        if self._custom_tools:
            self.set_custom_tools(self._custom_tools)
        return self

    def stop(self) -> None:
        process = self._process
        if process is None:
            return

        self._stopping = True
        for pending_call in self._pending_host_tool_calls.values():
            pending_call.cancel_event.set()

        try:
            if process.stdin is not None:
                try:
                    process.stdin.close()
                except OSError:
                    pass

            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=1.0)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait(timeout=1.0)
        finally:
            if process.stdout is not None:
                try:
                    process.stdout.close()
                except OSError:
                    pass
            if process.stderr is not None:
                try:
                    process.stderr.close()
                except OSError:
                    pass
            self._fail_pending(RpcProcessExitError("RPC process stopped"))
            self._pending_host_tool_calls.clear()
            self._process = None
            self._ready.set()
            with self._event_condition:
                self._event_condition.notify_all()
            if self._stdout_thread is not None:
                self._stdout_thread.join(timeout=1.0)
            if self._stderr_thread is not None:
                self._stderr_thread.join(timeout=1.0)
            self._stdout_thread = None
            self._stderr_thread = None

    def on_event(self, listener: AgentEventListener) -> Callable[[], None]:
        self._event_listeners.append(listener)
        return lambda: self._remove_listener(self._event_listeners, listener)

    def on_notification(self, listener: NotificationListener) -> Callable[[], None]:
        self._notification_listeners.append(listener)
        return lambda: self._remove_listener(self._notification_listeners, listener)

    def on_ready(self, listener: ReadyListener) -> Callable[[], None]:
        self._ready_listeners.append(listener)
        return lambda: self._remove_listener(self._ready_listeners, listener)

    def on_agent_start(self, listener: AgentStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("agent_start", listener)

    def on_agent_end(self, listener: AgentEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("agent_end", listener)

    def on_turn_start(self, listener: TurnStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("turn_start", listener)

    def on_turn_end(self, listener: TurnEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("turn_end", listener)

    def on_message_start(self, listener: MessageStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("message_start", listener)

    def on_message_update(self, listener: MessageUpdateListener) -> Callable[[], None]:
        return self._add_typed_event_listener("message_update", listener)

    def on_message_end(self, listener: MessageEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("message_end", listener)

    def on_tool_execution_start(self, listener: ToolExecutionStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("tool_execution_start", listener)

    def on_tool_execution_update(self, listener: ToolExecutionUpdateListener) -> Callable[[], None]:
        return self._add_typed_event_listener("tool_execution_update", listener)

    def on_tool_execution_end(self, listener: ToolExecutionEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("tool_execution_end", listener)

    def on_auto_compaction_start(self, listener: AutoCompactionStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("auto_compaction_start", listener)

    def on_auto_compaction_end(self, listener: AutoCompactionEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("auto_compaction_end", listener)

    def on_auto_retry_start(self, listener: AutoRetryStartListener) -> Callable[[], None]:
        return self._add_typed_event_listener("auto_retry_start", listener)

    def on_auto_retry_end(self, listener: AutoRetryEndListener) -> Callable[[], None]:
        return self._add_typed_event_listener("auto_retry_end", listener)

    def on_retry_fallback_applied(self, listener: RetryFallbackAppliedListener) -> Callable[[], None]:
        return self._add_typed_event_listener("retry_fallback_applied", listener)

    def on_retry_fallback_succeeded(self, listener: RetryFallbackSucceededListener) -> Callable[[], None]:
        return self._add_typed_event_listener("retry_fallback_succeeded", listener)

    def on_ttsr_triggered(self, listener: TtsrTriggeredListener) -> Callable[[], None]:
        return self._add_typed_event_listener("ttsr_triggered", listener)

    def on_todo_reminder(self, listener: TodoReminderListener) -> Callable[[], None]:
        return self._add_typed_event_listener("todo_reminder", listener)

    def on_todo_auto_clear(self, listener: TodoAutoClearListener) -> Callable[[], None]:
        return self._add_typed_event_listener("todo_auto_clear", listener)

    def on_ui_request(self, listener: UiRequestListener) -> Callable[[], None]:
        self._ui_request_listeners.append(listener)
        return lambda: self._remove_listener(self._ui_request_listeners, listener)

    def on_extension_error(self, listener: ExtensionErrorListener) -> Callable[[], None]:
        self._extension_error_listeners.append(listener)
        return lambda: self._remove_listener(self._extension_error_listeners, listener)

    def on_protocol_error(self, listener: ProtocolErrorListener) -> Callable[[], None]:
        self._protocol_error_listeners.append(listener)
        return lambda: self._remove_listener(self._protocol_error_listeners, listener)

    def on_listener_error(self, listener: ListenerErrorListener) -> Callable[[], None]:
        self._listener_error_listeners.append(listener)
        return lambda: self._remove_listener(self._listener_error_listeners, listener)

    def on_unknown_notification(self, listener: UnknownNotificationListener) -> Callable[[], None]:
        self._unknown_notification_listeners.append(listener)
        return lambda: self._remove_listener(self._unknown_notification_listeners, listener)

    def install_headless_ui(
        self,
        *,
        on_request: UiRequestListener | None = None,
        confirm: bool = False,
        select_value: str | None = None,
        input_value: str | None = None,
        editor_value: str | None = None,
    ) -> Callable[[], None]:
        """Auto-handle RPC UI requests for non-interactive hosts.

        Passive UI methods such as notifications and status updates are ignored.
        Confirm dialogs default to `False`. Select, input, and editor requests
        are cancelled unless an explicit value is provided.
        """

        def handle(request: ExtensionUiRequest) -> None:
            if on_request is not None:
                try:
                    on_request(request)
                except Exception as exc:
                    self._record_listener_error(
                        ListenerErrorEvent(
                            listener_kind="headless_ui_request",
                            source_type=request.type,
                            listener=on_request,
                            error=exc,
                        )
                    )

            if request.method == "cancel" or request.is_passive():
                return
            if request.method == "confirm":
                self.send_ui_confirmation(request.id, confirm)
                return
            if request.method == "select":
                if select_value is not None:
                    self.send_ui_value(request.id, select_value)
                else:
                    self.cancel_ui_request(request.id)
                return
            if request.method == "input":
                if input_value is not None:
                    self.send_ui_value(request.id, input_value)
                else:
                    self.cancel_ui_request(request.id)
                return
            if request.method == "editor":
                if editor_value is not None:
                    self.send_ui_value(request.id, editor_value)
                else:
                    self.cancel_ui_request(request.id)

        return self.on_ui_request(handle)

    def next_ui_request(self, timeout: float | None = None) -> ExtensionUiRequest:
        try:
            return self._ui_requests.get(timeout=timeout)
        except queue.Empty as exc:
            raise RpcTimeoutError("Timed out waiting for an extension UI request") from exc

    def send_ui_value(self, request_id: str, value: str) -> None:
        self._send_notification({"type": "extension_ui_response", "id": request_id, "value": value})

    def send_ui_confirmation(self, request_id: str, confirmed: bool) -> None:
        self._send_notification({"type": "extension_ui_response", "id": request_id, "confirmed": confirmed})

    def cancel_ui_request(self, request_id: str, *, timed_out: bool = False) -> None:
        payload: JsonObject = {"type": "extension_ui_response", "id": request_id, "cancelled": True}
        if timed_out:
            payload["timedOut"] = True
        self._send_notification(payload)

    def get_state(self) -> SessionState:
        payload = self._request("get_state")
        return parse_session_state(payload)

    def set_model(self, provider: str, model_id: str) -> ModelInfo:
        payload = self._request("set_model", provider=provider, modelId=model_id)
        model = parse_model_info(payload)
        if model is None:
            raise RpcError("set_model returned an empty payload")
        return model

    def cycle_model(self) -> ModelCycleResult | None:
        return parse_model_cycle_result(self._request("cycle_model"))

    def get_available_models(self) -> tuple[ModelInfo, ...]:
        payload = self._request("get_available_models")
        models = cast(list[JsonObject], payload.get("models") or [])
        return tuple(filter(None, (parse_model_info(model) for model in models)))

    def set_thinking_level(self, level: ThinkingLevel) -> None:
        self._request("set_thinking_level", level=level)

    def cycle_thinking_level(self) -> ThinkingLevelCycleResult | None:
        return parse_thinking_level_cycle_result(self._request("cycle_thinking_level"))

    def set_steering_mode(self, mode: SteeringMode) -> None:
        self._request("set_steering_mode", mode=mode)

    def set_follow_up_mode(self, mode: SteeringMode) -> None:
        self._request("set_follow_up_mode", mode=mode)

    def set_interrupt_mode(self, mode: InterruptMode) -> None:
        self._request("set_interrupt_mode", mode=mode)

    def compact(self, custom_instructions: str | None = None) -> CompactionResult:
        payload = self._request("compact", customInstructions=custom_instructions)
        return parse_compaction_result(payload)

    def set_auto_compaction(self, enabled: bool) -> None:
        self._request("set_auto_compaction", enabled=enabled)

    def set_auto_retry(self, enabled: bool) -> None:
        self._request("set_auto_retry", enabled=enabled)

    def abort_retry(self) -> None:
        self._request("abort_retry")

    def bash(self, command: str) -> BashResult:
        payload = self._request("bash", command=command)
        return parse_bash_result(payload)

    def abort_bash(self) -> None:
        self._request("abort_bash")

    def get_session_stats(self) -> SessionStats:
        payload = self._request("get_session_stats")
        return parse_session_stats(payload)

    def export_html(self, output_path: str | Path | None = None) -> Path:
        payload = self._request("export_html", outputPath=str(output_path) if output_path is not None else None)
        return Path(str(payload["path"]))

    def new_session(self, parent_session: str | None = None) -> CancellationResult:
        return parse_cancellation_result(self._request("new_session", parentSession=parent_session))

    def switch_session(self, session_path: str | Path) -> CancellationResult:
        return parse_cancellation_result(self._request("switch_session", sessionPath=str(session_path)))

    def branch(self, entry_id: str) -> BranchResult:
        return parse_branch_result(self._request("branch", entryId=entry_id))

    def get_branch_messages(self) -> tuple[BranchMessage, ...]:
        return parse_branch_messages(self._request("get_branch_messages"))

    def get_last_assistant_text(self) -> str | None:
        payload = self._request("get_last_assistant_text")
        value = payload.get("text")
        return str(value) if isinstance(value, str) else None

    def set_session_name(self, name: str) -> None:
        self._request("set_session_name", name=name)

    def get_todos(self) -> tuple[TodoPhase, ...]:
        return self.get_state().todo_phases

    def set_todos(self, todos: Sequence[TodoSeed | TodoPhaseSeed]) -> tuple[TodoPhase, ...]:
        phases = self._normalize_todo_phases(todos)
        payload = self._request("set_todos", phases=cast(JsonValue, phases))
        return parse_todo_phases(payload.get("todoPhases"))

    def clear_todos(self) -> tuple[TodoPhase, ...]:
        return self.set_todos(())

    def get_messages(self) -> tuple[AgentMessage, ...]:
        payload = self._request("get_messages")
        return parse_agent_messages(cast(JsonValue | None, payload.get("messages")))

    def set_custom_tools(self, tools: Sequence[HostTool[Any, Any]]) -> tuple[str, ...]:
        self._custom_tools = tuple(tools)
        if self._process is None:
            return tuple(tool.name for tool in self._custom_tools)

        payload = self._request(
            "set_host_tools",
            tools=cast(
                JsonValue,
                [
                    {
                        "name": tool.name,
                        "label": tool.label,
                        "description": tool.description,
                        "parameters": tool.parameters,
                        "hidden": tool.hidden,
                    }
                    for tool in self._custom_tools
                ],
            ),
        )
        tool_names = payload.get("toolNames") or []
        if not isinstance(tool_names, list):
            raise RpcError("set_host_tools response did not include toolNames")
        return tuple(str(name) for name in tool_names)

    def prompt(
        self,
        message: str,
        *,
        images: Sequence[ImageContent] | None = None,
        streaming_behavior: StreamingBehavior | None = None,
    ) -> None:
        self._request(
            "prompt",
            message=message,
            images=list(images) if images is not None else None,
            streamingBehavior=streaming_behavior,
        )
        self._mark_agent_run_scheduled()

    def steer(self, message: str, *, images: Sequence[ImageContent] | None = None) -> None:
        self._request("steer", message=message, images=list(images) if images is not None else None)

    def follow_up(self, message: str, *, images: Sequence[ImageContent] | None = None) -> None:
        self._request("follow_up", message=message, images=list(images) if images is not None else None)

    def abort(self) -> None:
        self._request("abort")

    def abort_and_prompt(self, message: str, *, images: Sequence[ImageContent] | None = None) -> None:
        self._request("abort_and_prompt", message=message, images=list(images) if images is not None else None)
        self._mark_agent_run_scheduled()

    def prompt_and_wait(
        self,
        message: str,
        *,
        images: Sequence[ImageContent] | None = None,
        streaming_behavior: StreamingBehavior | None = None,
        timeout: float | None = None,
    ) -> PromptTurn:
        operation = "prompt_and_wait"
        self._prompt_lifecycle.acquire(operation)
        try:
            start_index = self._current_event_index()
            start_async_error_index = self._current_async_error_index()
            self.prompt(message, images=images, streaming_behavior=streaming_behavior)
            events = self._wait_for_agent_end(start_index, start_async_error_index, timeout=timeout)
            return self._build_prompt_turn(events)
        finally:
            self._prompt_lifecycle.release(operation)

    def wait_for_idle(self, timeout: float | None = None) -> None:
        operation = "wait_for_idle"
        self._prompt_lifecycle.acquire(operation)
        try:
            if self._is_agent_idle():
                self._check_async_errors()
                return
            start_index = self._current_event_index()
            start_async_error_index = self._current_async_error_index()
            self._wait_for_agent_end(start_index, start_async_error_index, timeout=timeout)
        finally:
            self._prompt_lifecycle.release(operation)

    def collect_events(self, timeout: float | None = None) -> tuple[RpcAgentEvent, ...]:
        operation = "collect_events"
        self._prompt_lifecycle.acquire(operation)
        try:
            start_index = self._current_event_index()
            start_async_error_index = self._current_async_error_index()
            return self._wait_for_agent_end(start_index, start_async_error_index, timeout=timeout)
        finally:
            self._prompt_lifecycle.release(operation)

    def request_raw(self, command_type: str, **payload: JsonValue) -> JsonObject:
        return self._request(command_type, **payload)

    def _current_event_index(self) -> int:
        with self._event_condition:
            return self._events.current_index()

    def _current_async_error_index(self) -> int:
        with self._event_condition:
            return self._async_errors.current_index()

    def _mark_agent_run_scheduled(self) -> None:
        with self._event_condition:
            self._scheduled_agent_runs += 1
            self._last_schedule_async_error_index = self._async_errors.current_index()
    def _mark_agent_run_completed(self) -> None:
        with self._event_condition:
            self._completed_agent_runs += 1
            self._event_condition.notify_all()

    def _is_agent_idle(self) -> bool:
        with self._event_condition:
            return self._scheduled_agent_runs == self._completed_agent_runs

    def _check_async_errors(self) -> None:
        with self._event_condition:
            errors = self._async_errors.snapshot_from(self._last_schedule_async_error_index)
        if errors:
            raise errors[0]

    def _build_prompt_turn(self, events: tuple[RpcAgentEvent, ...]) -> PromptTurn:
        final_messages: tuple[AgentMessage, ...] = ()
        for event in reversed(events):
            if isinstance(event, AgentEndEvent):
                final_messages = event.messages
                break

        assistant_message: AssistantMessage | None = None
        for message in reversed(final_messages):
            if message.get("role") == "assistant":
                assistant_message = cast(AssistantMessage, message)
                break

        if assistant_message is None:
            for event in reversed(events):
                if hasattr(event, "message"):
                    message = cast(AgentMessage | None, getattr(event, "message", None))
                    if isinstance(message, dict) and message.get("role") == "assistant":
                        assistant_message = cast(AssistantMessage, message)
                        break

        return PromptTurn(
            events=events,
            messages=final_messages,
            assistant_message=assistant_message,
            assistant_text=assistant_text(assistant_message) if assistant_message is not None else None,
        )

    def _wait_for_agent_end(
        self,
        start_index: int,
        start_async_error_index: int,
        timeout: float | None = None,
    ) -> tuple[RpcAgentEvent, ...]:
        deadline = time.monotonic() + (timeout if timeout is not None else 60.0)
        with self._event_condition:
            while True:
                if self._closed_error is not None:
                    raise RpcProcessExitError(str(self._closed_error))

                if start_index < self._events.offset:
                    raise RpcError(
                        "Event history limit was exceeded while waiting for agent_end. "
                        "Increase max_event_history to retain more streamed events."
                    )

                if start_async_error_index < self._async_errors.offset:
                    raise RpcError(
                        "Async error history limit was exceeded while waiting for agent_end. "
                        "Increase max_event_history if your host needs to retain more background failures."
                    )

                async_errors = self._async_errors.snapshot_from(start_async_error_index)
                if len(async_errors) > 0:
                    raise async_errors[0]

                event_payloads = self._events.snapshot_from(start_index)
                if any(payload.get("type") == "agent_end" for payload in event_payloads):
                    events = tuple(cast(RpcAgentEvent, parse_notification(payload)) for payload in event_payloads)
                    return events

                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise RpcTimeoutError(f"Timed out waiting for agent_end. Stderr: {self.stderr}")
                self._event_condition.wait(remaining)

    def _request(self, command_type: str, **payload: JsonValue) -> JsonObject:
        process = self._require_process()
        request_id = self._next_request_id()
        envelope: JsonObject = {"id": request_id, "type": command_type}
        for key, value in payload.items():
            if value is not None:
                envelope[key] = value

        response_queue: queue.Queue[JsonObject | BaseException] = queue.Queue(maxsize=1)
        with self._state_lock:
            self._pending[request_id] = _PendingRequest(command=command_type, response_queue=response_queue)

        try:
            self._write_json(process, envelope)
        except BaseException:
            with self._state_lock:
                self._pending.pop(request_id, None)
            raise

        try:
            response = response_queue.get(timeout=self._request_timeout)
        except queue.Empty as exc:
            with self._state_lock:
                self._pending.pop(request_id, None)
            raise RpcTimeoutError(f"Timed out waiting for response to {command_type}. Stderr: {self.stderr}") from exc

        if isinstance(response, BaseException):
            raise response

        if not bool(response.get("success", False)):
            raise RpcCommandError(command=str(response.get("command", command_type)), error=str(response.get("error", "")))

        data = response.get("data")
        if data is None:
            return {}
        return _clone_json_object(data)

    def _send_notification(self, payload: JsonObject) -> None:
        process = self._require_process()
        self._write_json(process, payload)

    def _normalize_host_tool_result(self, result: object) -> JsonObject:
        if isinstance(result, str):
            return {"content": [{"type": "text", "text": result}]}
        if isinstance(result, Mapping):
            return cast(JsonObject, dict(result))
        raise RpcError("Host tool handlers must return a string or a result mapping")

    def _handle_host_tool_call(self, payload: JsonObject) -> None:
        request_id = payload.get("id")
        tool_name = payload.get("toolName")
        tool_call_id = payload.get("toolCallId")
        raw_arguments = payload.get("arguments")
        if not isinstance(request_id, str) or not isinstance(tool_name, str) or not isinstance(tool_call_id, str):
            return
        if not isinstance(raw_arguments, Mapping):
            self._send_notification(
                {
                    "type": "host_tool_result",
                    "id": request_id,
                    "result": {"content": [{"type": "text", "text": "Host tool arguments must be an object"}], "details": {}},
                    "isError": True,
                }
            )
            return

        tool = next((candidate for candidate in self._custom_tools if candidate.name == tool_name), None)
        if tool is None:
            self._send_notification(
                {
                    "type": "host_tool_result",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": f'Host tool "{tool_name}" is not registered'}],
                        "details": {},
                    },
                    "isError": True,
                }
            )
            return

        pending_call = _PendingHostToolCall(cancel_event=threading.Event())
        self._pending_host_tool_calls[request_id] = pending_call

        def run_tool() -> None:
            try:
                params = tool.parse_params(cast(JsonObject, dict(raw_arguments)))
                context = HostToolContext(
                    tool_call_id=tool_call_id,
                    _cancel_event=pending_call.cancel_event,
                    _send_update=lambda result: self._send_notification(
                        {"type": "host_tool_update", "id": request_id, "partialResult": result}
                    ),
                )
                result = tool.execute(params, context)
                if pending_call.cancel_event.is_set():
                    return
                self._send_notification(
                    {
                        "type": "host_tool_result",
                        "id": request_id,
                        "result": self._normalize_host_tool_result(result),
                    }
                )
            except Exception as exc:
                if pending_call.cancel_event.is_set():
                    return
                self._send_notification(
                    {
                        "type": "host_tool_result",
                        "id": request_id,
                        "result": {"content": [{"type": "text", "text": str(exc)}], "details": {}},
                        "isError": True,
                    }
                )
            finally:
                self._pending_host_tool_calls.pop(request_id, None)

        threading.Thread(target=run_tool, name=f"omp-rpc-host-tool:{tool_name}", daemon=True).start()

    def _handle_host_tool_cancel(self, payload: JsonObject) -> None:
        target_id = payload.get("targetId")
        if not isinstance(target_id, str):
            return
        pending_call = self._pending_host_tool_calls.get(target_id)
        if pending_call is not None:
            pending_call.cancel_event.set()

    def _add_typed_event_listener(self, event_type: str, listener: TEventListener) -> Callable[[], None]:
        listeners = self._typed_event_listeners.setdefault(event_type, [])
        typed_listener = cast(AgentEventListener, listener)
        listeners.append(typed_listener)
        return lambda: self._remove_listener(listeners, typed_listener)

    @staticmethod
    def _normalize_todo_phases(todos: Sequence[TodoSeed | TodoPhaseSeed]) -> list[JsonObject]:
        if len(todos) == 0:
            return []

        next_task_id = 1

        def next_task() -> str:
            nonlocal next_task_id
            task_id = f"task-{next_task_id}"
            next_task_id += 1
            return task_id

        def normalize_todo_item(seed: TodoSeed) -> JsonObject:
            if isinstance(seed, str):
                return {"id": next_task(), "content": seed, "status": cast(JsonValue, "pending")}

            if isinstance(seed, TodoItem):
                if seed.status not in _TODO_STATUS_VALUES:
                    raise RpcError(f"Unsupported todo status: {seed.status}")
                return {
                    "id": seed.id or next_task(),
                    "content": seed.content,
                    "status": cast(JsonValue, seed.status),
                    "notes": seed.notes,
                    "details": seed.details,
                }

            content = seed.get("content")
            if not isinstance(content, str) or not content.strip():
                raise RpcError("Todo items must provide a non-empty 'content' value")

            raw_id = seed.get("id")
            raw_status = seed.get("status")
            raw_notes = seed.get("notes")
            raw_details = seed.get("details")
            if isinstance(raw_status, str):
                if raw_status not in _TODO_STATUS_VALUES:
                    raise RpcError(f"Unsupported todo status: {raw_status}")
                status: TodoStatus = cast(TodoStatus, raw_status)
            else:
                status = "pending"
            return {
                "id": str(raw_id) if isinstance(raw_id, str) and raw_id else next_task(),
                "content": content,
                "status": cast(JsonValue, status),
                "notes": raw_notes if isinstance(raw_notes, str) else None,
                "details": raw_details if isinstance(raw_details, str) else None,
            }

        def is_phase_seed(seed: TodoSeed | TodoPhaseSeed) -> bool:
            if isinstance(seed, TodoPhase):
                return True
            if not isinstance(seed, Mapping):
                return False
            return "tasks" in seed or ("name" in seed and "content" not in seed)

        def normalize_phase(seed: TodoPhaseSeed, index: int) -> JsonObject:
            if isinstance(seed, TodoPhase):
                phase_id = seed.id or f"phase-{index}"
                name = seed.name
                tasks = [normalize_todo_item(task) for task in seed.tasks]
            else:
                raw_name = seed.get("name")
                if not isinstance(raw_name, str) or not raw_name.strip():
                    raise RpcError("Todo phases must provide a non-empty 'name' value")
                phase_id_value = seed.get("id")
                raw_tasks = seed.get("tasks") or ()
                if not isinstance(raw_tasks, Sequence) or isinstance(raw_tasks, (str, bytes)):
                    raise RpcError("Todo phase 'tasks' must be a sequence")
                phase_id = str(phase_id_value) if isinstance(phase_id_value, str) and phase_id_value else f"phase-{index}"
                name = raw_name
                tasks = [normalize_todo_item(cast(TodoSeed, task)) for task in raw_tasks]

            return {"id": phase_id, "name": name, "tasks": tasks}

        if any(is_phase_seed(todo) for todo in todos):
            phases: list[JsonObject] = []
            for index, seed in enumerate(todos, start=1):
                if not is_phase_seed(seed):
                    raise RpcError("Cannot mix flat todo items with todo phases in one set_todos() call")
                phases.append(normalize_phase(cast(TodoPhaseSeed, seed), index))
            return phases

        return [{"id": "phase-1", "name": "Todos", "tasks": [normalize_todo_item(cast(TodoSeed, todo)) for todo in todos]}]

    def _build_command(self) -> tuple[str, ...]:
        if self._command is not None:
            return self._command

        command: list[str] = [self._executable, "--mode", "rpc"]
        if self._provider:
            command.extend(["--provider", self._provider])
        if self._model:
            command.extend(["--model", self._model])
        if self._session_dir is not None:
            command.extend(["--session-dir", str(self._session_dir)])
        if self._thinking is not None:
            command.extend(["--thinking", self._thinking])
        if self._append_system_prompt is not None:
            command.extend(["--append-system-prompt", self._append_system_prompt])
        if self._provider_session_id is not None:
            command.extend(["--provider-session-id", self._provider_session_id])
        if self._tools is not None:
            if len(self._tools) == 0:
                command.append("--no-tools")
            else:
                command.extend(["--tools", ",".join(self._tools)])
        if self._no_session:
            command.append("--no-session")
        if self._no_skills:
            command.append("--no-skills")
        if self._no_rules:
            command.append("--no-rules")
        emit_no_title = self._no_title if self._no_title is not None else self._rpc_defaults
        if emit_no_title:
            command.append("--no-title")
        command.extend(self._extra_args)
        return tuple(command)

    def _next_request_id(self) -> str:
        with self._state_lock:
            self._request_id += 1
            return f"req_{self._request_id}"

    def _require_process(self) -> subprocess.Popen[str]:
        if self._process is None:
            raise RpcError("RPC client is not started")
        return self._process

    def _write_json(self, process: subprocess.Popen[str], payload: JsonObject) -> None:
        if process.stdin is None:
            raise RpcProcessExitError("RPC process stdin is unavailable")
        with self._write_lock:
            try:
                process.stdin.write(json.dumps(payload))
                process.stdin.write("\n")
                process.stdin.flush()
            except (BrokenPipeError, OSError) as exc:
                raise RpcProcessExitError(f"Failed to write RPC command: {exc}") from exc

    def _read_stdout_loop(self) -> None:
        process = self._process
        if process is None or process.stdout is None:
            return

        line_number = 0
        try:
            for line in process.stdout:
                line_number += 1
                stripped = line.strip()
                if not stripped:
                    continue

                try:
                    payload = cast(JsonObject, json.loads(stripped))
                except json.JSONDecodeError as exc:
                    snippet = stripped
                    if len(snippet) > 240:
                        snippet = f"{snippet[:237]}..."
                    raise RpcError(
                        f"Failed to decode RPC output on line {line_number}: {exc}. Frame: {snippet!r}"
                    ) from exc
                if payload.get("type") == "response":
                    self._handle_response(payload)
                    continue
                if payload.get("type") == "host_tool_call":
                    self._handle_host_tool_call(payload)
                    continue
                if payload.get("type") == "host_tool_cancel":
                    self._handle_host_tool_cancel(payload)
                    continue

                notification = parse_notification(payload)
                listener_notification = parse_notification(payload)
                self._dispatch_listeners(
                    "notification",
                    listener_notification.type,
                    self._notification_listeners,
                    listener_notification,
                )

                if isinstance(notification, ReadyEvent):
                    self._ready_received = True
                    self._ready.set()
                    self._dispatch_listeners("ready", listener_notification.type, self._ready_listeners, listener_notification)
                    continue

                if isinstance(notification, ExtensionUiRequest):
                    self._ui_requests.put(notification)
                    self._dispatch_listeners(
                        "ui_request",
                        listener_notification.type,
                        self._ui_request_listeners,
                        cast(ExtensionUiRequest, listener_notification),
                    )
                    continue

                if isinstance(notification, ExtensionError):
                    self._dispatch_listeners(
                        "extension_error",
                        listener_notification.type,
                        self._extension_error_listeners,
                        cast(ExtensionError, listener_notification),
                    )
                    continue

                if isinstance(notification, UnknownNotification):
                    self._dispatch_listeners(
                        "unknown_notification",
                        listener_notification.type,
                        self._unknown_notification_listeners,
                        cast(UnknownNotification, listener_notification),
                    )
                    continue

                listener_event = cast(RpcAgentEvent, listener_notification)
                self._append_event(payload)
                if listener_event.type == "agent_end":
                    self._mark_agent_run_completed()
                self._dispatch_listeners("event", listener_event.type, self._event_listeners, listener_event)
                self._dispatch_listeners(
                    "typed_event", listener_event.type, self._typed_event_listeners.get(listener_event.type, []), listener_event
                )
        except Exception as exc:
            self._mark_closed(exc)
        else:
            if not self._stopping:
                exit_code = process.poll()
                if exit_code is None:
                    try:
                        exit_code = process.wait(timeout=1.0)
                    except subprocess.TimeoutExpired:
                        self._mark_closed(RpcProcessExitError("RPC process stdout closed before the process exited"))
                        return
                self._mark_closed(RpcProcessExitError(f"RPC process exited with code {exit_code}. Stderr: {self.stderr}"))

    def _read_stderr_loop(self) -> None:
        process = self._process
        if process is None or process.stderr is None:
            return
        try:
            for chunk in process.stderr:
                with self._state_lock:
                    self._stderr_chunks.append(chunk)
        except Exception as exc:
            if not self._stopping:
                self._mark_closed(RpcError(f"Failed to read RPC stderr: {exc}"))

    def _mark_closed(self, error: BaseException) -> None:
        if self._closed_error is not None:
            return
        self._closed_error = error
        self._ready.set()
        self._fail_pending(error)
        with self._event_condition:
            self._event_condition.notify_all()

    def _fail_pending(self, error: BaseException) -> None:
        with self._state_lock:
            pending = [pending.response_queue for pending in self._pending.values()]
            self._pending.clear()
        for response_queue in pending:
            response_queue.put(error)

    def _handle_response(self, payload: JsonObject) -> None:
        request_id = payload.get("id")
        if isinstance(request_id, str):
            with self._state_lock:
                pending = self._pending.pop(request_id, None)
            if pending is not None:
                pending.response_queue.put(payload)
                return

        if self._deliver_correlated_error_response(payload):
            return

        protocol_error = self._build_protocol_error(payload)
        if protocol_error is None:
            return

        if protocol_error.command in _ASYNC_COMMANDS and protocol_error.remote_error is not None:
            self._append_async_error(RpcCommandError(protocol_error.command, protocol_error.remote_error))
            self._mark_agent_run_completed()

        self._record_protocol_error(protocol_error)

    def _deliver_correlated_error_response(self, payload: JsonObject) -> bool:
        if bool(payload.get("success", False)):
            return False

        command = payload.get("command")
        if not isinstance(command, str):
            return False

        with self._state_lock:
            matching_ids = [request_id for request_id, pending in self._pending.items() if pending.command == command]
            target_id: str | None = None
            if len(matching_ids) == 1:
                target_id = matching_ids[0]
            elif command == "parse" and len(self._pending) == 1:
                target_id = next(iter(self._pending))

            if target_id is None:
                return False

            pending = self._pending.pop(target_id)

        pending.response_queue.put(payload)
        return True

    def _build_protocol_error(self, payload: JsonObject) -> RpcProtocolError | None:
        if payload.get("type") != "response":
            return None
        if bool(payload.get("success", False)):
            return None
        return RpcProtocolError(_clone_json_object(payload))

    def _append_event(self, payload: JsonObject) -> None:
        with self._event_condition:
            self._events.append(_clone_json_object(payload))
            self._event_condition.notify_all()

    def _append_async_error(self, error: BaseException) -> None:
        with self._event_condition:
            self._async_errors.append(error)
            self._event_condition.notify_all()

    def _record_protocol_error(self, error: RpcProtocolError) -> None:
        with self._state_lock:
            self._protocol_errors.append(error)
        self._dispatch_listeners("protocol_error", error.command, self._protocol_error_listeners, error)

    def _record_listener_error(self, event: ListenerErrorEvent) -> None:
        with self._state_lock:
            self._listener_errors.append(event)

        for listener in list(self._listener_error_listeners):
            try:
                listener(event)
            except Exception:
                continue

    def _dispatch_listeners(
        self,
        listener_kind: str,
        source_type: str | None,
        listeners: Sequence[Callable[[Any], None]],
        payload: Any,
    ) -> None:
        for listener in list(listeners):
            try:
                listener(payload)
            except Exception as exc:
                self._record_listener_error(
                    ListenerErrorEvent(
                        listener_kind=listener_kind,
                        source_type=source_type,
                        listener=listener,
                        error=exc,
                    )
                )

    @staticmethod
    def _validate_history_limit(name: str, limit: int | None) -> int | None:
        if limit is None:
            return None
        if limit <= 0:
            raise ValueError(f"{name} must be greater than zero")
        return limit

    @staticmethod
    def _remove_listener(listeners: list[TListener], listener: TListener) -> None:
        try:
            listeners.remove(listener)
        except ValueError:
            pass