smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(rpc): harden Python protocol compatibility

Browse Source
This commit is contained in:
Wolfgang Schoenberger
2026-07-30 15:16:54 -07:00
parent 4df68d6043
commit 0f033ed706
6 changed files with 208 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/omp-rpc/src/omp_rpc/__init__.py
+2
View File
@@ -49,6 +49,7 @@ from .protocol import (
CancellationResult,
CompactionResult,
CompactionSummaryMessage,
ContextUsage,
CustomMessage,
DeveloperMessage,
Effort,
@@ -116,6 +117,7 @@ __all__ = [
"CancellationResult",
"CompactionResult",
"CompactionSummaryMessage",
"ContextUsage",
"CustomMessage",
"DeveloperMessage",
"Effort",
python/omp-rpc/src/omp_rpc/client.py
+17 -4
View File
@@ -1353,7 +1353,9 @@ class RpcClient:
event_payloads = self._events.snapshot_from(start_index)
if any(
payload.get("type") == "agent_end" for payload in event_payloads
payload.get("type") == "agent_end"
and payload.get("isTerminal") is not False
for payload in event_payloads
):
events = tuple(
cast(RpcAgentEvent, parse_notification(payload))
@@ -1898,8 +1900,16 @@ class RpcClient:
payload_type = payload.get("type")
if payload_type in ("tool_execution_update", "tool_execution_end"):
self._normalize_host_tool_event(payload)
notification = parse_notification(payload)
listener_notification = parse_notification(payload)
try:
notification = parse_notification(payload)
listener_notification = parse_notification(payload)
except (TypeError, ValueError) as exc:
notification = UnknownNotification(
_clone_json_object(payload), parse_error=str(exc)
)
listener_notification = UnknownNotification(
_clone_json_object(payload), parse_error=str(exc)
)
self._dispatch_listeners(
"notification",
listener_notification.type,
@@ -1949,7 +1959,10 @@ class RpcClient:
listener_event = cast(RpcAgentEvent, listener_notification)
self._append_event(payload)
if listener_event.type == "agent_end":
if (
isinstance(listener_event, AgentEndEvent)
and listener_event.is_terminal is not False
):
self._mark_agent_run_completed()
self._dispatch_listeners(
"event", listener_event.type, self._event_listeners, listener_event
python/omp-rpc/src/omp_rpc/protocol.py
+56 -10
View File
@@ -35,17 +35,30 @@ ExtensionUiMethod: TypeAlias = Literal[
"setWidget",
"setTitle",
"set_editor_text",
"open_url",
]
InteractiveExtensionUiMethod: TypeAlias = Literal[
"select", "confirm", "input", "editor"
]
PassiveExtensionUiMethod: TypeAlias = Literal[
"notify", "setStatus", "setWidget", "setTitle", "set_editor_text"
"notify",
"setStatus",
"setWidget",
"setTitle",
"set_editor_text",
"open_url",
]
ValueExtensionUiMethod: TypeAlias = Literal["select", "input", "editor"]
PASSIVE_EXTENSION_UI_METHODS: Final[frozenset[PassiveExtensionUiMethod]] = frozenset(
{"notify", "setStatus", "setWidget", "setTitle", "set_editor_text"}
{
"notify",
"setStatus",
"setWidget",
"setTitle",
"set_editor_text",
"open_url",
}
)
INTERACTIVE_EXTENSION_UI_METHODS: Final[frozenset[InteractiveExtensionUiMethod]] = (
frozenset({"select", "confirm", "input", "editor"})
@@ -81,6 +94,7 @@ _EXTENSION_UI_METHOD_VALUES: Final[frozenset[str]] = frozenset(
"setWidget",
"setTitle",
"set_editor_text",
"open_url",
}
)
_AGENT_MESSAGE_ROLE_VALUES: Final[frozenset[str]] = frozenset(
@@ -119,10 +133,10 @@ _ASSISTANT_DONE_REASON_VALUES: Final[frozenset[str]] = frozenset(
)
_ASSISTANT_ERROR_REASON_VALUES: Final[frozenset[str]] = frozenset({"aborted", "error"})
_AUTO_COMPACTION_REASON_VALUES: Final[frozenset[str]] = frozenset(
{"threshold", "overflow", "idle"}
{"threshold", "overflow", "idle", "incomplete"}
)
_AUTO_COMPACTION_ACTION_VALUES: Final[frozenset[str]] = frozenset(
{"context-full", "handoff"}
{"context-full", "handoff", "shake", "snapcompact"}
)
@@ -774,6 +788,13 @@ class TodoPhase:
tasks: tuple[TodoItem, ...]
@dataclass(slots=True, frozen=True)
class ContextUsage:
tokens: int
context_window: int
percent: float
@dataclass(slots=True, frozen=True)
class SessionState:
model: ModelInfo | None
@@ -795,6 +816,7 @@ class SessionState:
fast_mode_enabled: bool = False
fast_mode_active: bool = False
tokens_per_second: float | None = None
context_usage: ContextUsage | None = None
@dataclass(slots=True, frozen=True)
@@ -913,6 +935,9 @@ class ExtensionUiRequest:
widget_lines: tuple[str, ...] | None = None
widget_placement: WidgetPlacement | None = None
text: str | None = None
url: str | None = None
launch_url: str | None = None
instructions: str | None = None
type: Literal["extension_ui_request"] = "extension_ui_request"
def is_passive(self) -> bool:
@@ -946,6 +971,7 @@ class AgentEndEvent:
messages: tuple[AgentMessage, ...]
type: Literal["agent_end"] = "agent_end"
message_count: int | None = field(default=None, kw_only=True)
is_terminal: bool | None = field(default=None, kw_only=True)
@dataclass(slots=True, frozen=True)
@@ -1008,14 +1034,14 @@ class ToolExecutionEndEvent:
@dataclass(slots=True, frozen=True)
class AutoCompactionStartEvent:
reason: Literal["threshold", "overflow", "idle"]
action: Literal["context-full", "handoff"]
reason: Literal["threshold", "overflow", "idle", "incomplete"]
action: Literal["context-full", "handoff", "shake", "snapcompact"]
type: Literal["auto_compaction_start"] = "auto_compaction_start"
@dataclass(slots=True, frozen=True)
class AutoCompactionEndEvent:
action: Literal["context-full", "handoff"]
action: Literal["context-full", "handoff", "shake", "snapcompact"]
result: CompactionResult | None
aborted: bool
will_retry: bool
@@ -1079,6 +1105,7 @@ class TodoAutoClearEvent:
class UnknownNotification:
payload: JsonObject
type: Literal["unknown"] = "unknown"
parse_error: str | None = field(default=None, kw_only=True)
RpcAgentEvent: TypeAlias = (
@@ -1372,6 +1399,11 @@ def parse_session_state(payload: JsonObject) -> SessionState:
fast_mode_enabled=bool(payload.get("fastModeEnabled", False)),
fast_mode_active=bool(payload.get("fastModeActive", False)),
tokens_per_second=_optional_float(payload, "tokensPerSecond"),
context_usage=parse_context_usage(
_optional_json_object(
payload.get("contextUsage"), field="sessionState.contextUsage"
)
),
)
@@ -1483,6 +1515,16 @@ def parse_session_stats(payload: JsonObject) -> SessionStats:
)
def parse_context_usage(payload: JsonObject | None) -> ContextUsage | None:
if payload is None:
return None
return ContextUsage(
tokens=int(payload.get("tokens", 0)),
context_window=int(payload.get("contextWindow", 0)),
percent=float(payload.get("percent", 0.0)),
)
def parse_extension_ui_request(payload: JsonObject) -> ExtensionUiRequest:
return ExtensionUiRequest(
id=_require_str(payload, "id"),
@@ -1527,6 +1569,9 @@ def parse_extension_ui_request(payload: JsonObject) -> ExtensionUiRequest:
),
),
text=_optional_str(payload, "text"),
url=_optional_str(payload, "url"),
launch_url=_optional_str(payload, "launchUrl"),
instructions=_optional_str(payload, "instructions"),
)
@@ -1570,6 +1615,7 @@ def parse_notification(payload: JsonObject) -> RpcNotification:
cast(JsonValue | None, payload.get("messages"))
),
message_count=_optional_int(payload, "messageCount"),
is_terminal=_optional_bool(payload, "isTerminal"),
)
if event_type == "turn_start":
return TurnStartEvent()
@@ -1661,7 +1707,7 @@ def parse_notification(payload: JsonObject) -> RpcNotification:
if event_type == "auto_compaction_start":
return AutoCompactionStartEvent(
reason=cast(
Literal["threshold", "overflow", "idle"],
Literal["threshold", "overflow", "idle", "incomplete"],
_require_literal(
payload.get("reason", "threshold"),
_AUTO_COMPACTION_REASON_VALUES,
@@ -1669,7 +1715,7 @@ def parse_notification(payload: JsonObject) -> RpcNotification:
),
),
action=cast(
Literal["context-full", "handoff"],
Literal["context-full", "handoff", "shake", "snapcompact"],
_require_literal(
payload.get("action", "context-full"),
_AUTO_COMPACTION_ACTION_VALUES,
@@ -1681,7 +1727,7 @@ def parse_notification(payload: JsonObject) -> RpcNotification:
result_payload = payload.get("result")
return AutoCompactionEndEvent(
action=cast(
Literal["context-full", "handoff"],
Literal["context-full", "handoff", "shake", "snapcompact"],
_require_literal(
payload.get("action", "context-full"),
_AUTO_COMPACTION_ACTION_VALUES,