wehub-resource-sync
367 lines
14 KiB
Python
367 lines
14 KiB
Python
"""SSE frame parser — converts raw byte chunks into typed events.
|
|
|
|
Handles the ``event:`` / ``data:`` / ``[DONE]`` framing from the
|
|
server's ``text/event-stream`` responses.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from omnigent.server import schemas as _srv_events
|
|
|
|
from ._events import (
|
|
NATIVE_TOOL_TYPES,
|
|
ClientTaskCancel,
|
|
CompactionInProgress,
|
|
ElicitationRequest,
|
|
ErrorEvent,
|
|
MessageDone,
|
|
NativeToolCall,
|
|
OutputFileDone,
|
|
ReasoningDelta,
|
|
ReasoningStarted,
|
|
ReasoningSummaryDelta,
|
|
ResponseCancelled,
|
|
ResponseCompleted,
|
|
ResponseCreated,
|
|
ResponseFailed,
|
|
ResponseIncomplete,
|
|
ResponseInProgress,
|
|
ResponseQueued,
|
|
RetryEvent,
|
|
StreamEvent,
|
|
TextDelta,
|
|
ToolCall,
|
|
ToolResult,
|
|
)
|
|
from ._types import ErrorInfo, Response
|
|
|
|
_log = logging.getLogger("omnigent_client.sse")
|
|
|
|
|
|
def _wire_type(cls: type) -> str:
|
|
"""
|
|
Extract the wire ``type`` literal from a server event class.
|
|
|
|
Each :mod:`omnigent.server.schemas` event class pins its
|
|
``type`` field as ``Literal["..."]``; this helper unwraps that
|
|
to a plain string so the SDK's ``str == str`` dispatch table
|
|
stays a ``str`` comparison rather than introducing a class-side
|
|
isinstance check.
|
|
|
|
:param cls: A subclass of the server's ``_SSEEventBase``,
|
|
e.g. :class:`omnigent.server.schemas.OutputTextDeltaEvent`.
|
|
:returns: The wire ``type`` literal, e.g.
|
|
``"response.output_text.delta"``.
|
|
"""
|
|
return cls.model_fields["type"].annotation.__args__[0] # type: ignore[union-attr]
|
|
|
|
|
|
# Wire type literals — pulled from the server's typed source of
|
|
# truth so a rename there is a one-edit change here as well.
|
|
_T_RESPONSE_CREATED = _wire_type(_srv_events.CreatedEvent)
|
|
_T_RESPONSE_QUEUED = _wire_type(_srv_events.QueuedEvent)
|
|
_T_RESPONSE_IN_PROGRESS = _wire_type(_srv_events.InProgressEvent)
|
|
_T_RESPONSE_COMPLETED = _wire_type(_srv_events.CompletedEvent)
|
|
_T_RESPONSE_FAILED = _wire_type(_srv_events.FailedEvent)
|
|
_T_RESPONSE_INCOMPLETE = _wire_type(_srv_events.IncompleteEvent)
|
|
_T_RESPONSE_CANCELLED = _wire_type(_srv_events.CancelledEvent)
|
|
_T_RESPONSE_OUTPUT_TEXT_DELTA = _wire_type(_srv_events.OutputTextDeltaEvent)
|
|
_T_RESPONSE_REASONING_STARTED = _wire_type(_srv_events.ReasoningStartedEvent)
|
|
_T_RESPONSE_REASONING_TEXT_DELTA = _wire_type(_srv_events.ReasoningTextDeltaEvent)
|
|
_T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA = _wire_type(_srv_events.ReasoningSummaryTextDeltaEvent)
|
|
_T_RESPONSE_OUTPUT_ITEM_DONE = _wire_type(_srv_events.OutputItemDoneEvent)
|
|
_T_RESPONSE_OUTPUT_FILE_DONE = _wire_type(_srv_events.OutputFileDoneEvent)
|
|
_T_RESPONSE_RETRY = _wire_type(_srv_events.RetryEvent)
|
|
_T_RESPONSE_ERROR = _wire_type(_srv_events.ErrorEvent)
|
|
_T_RESPONSE_COMPACTION_IN_PROGRESS = _wire_type(_srv_events.CompactionInProgressEvent)
|
|
_T_RESPONSE_CLIENT_TASK_CANCEL = _wire_type(_srv_events.ClientTaskCancelEvent)
|
|
_T_RESPONSE_ELICITATION_REQUEST = _wire_type(_srv_events.ElicitationRequestEvent)
|
|
|
|
|
|
async def parse_sse_stream(
|
|
byte_stream: AsyncIterator[bytes],
|
|
) -> AsyncIterator[StreamEvent]:
|
|
"""Parse an SSE byte stream into typed events.
|
|
|
|
:param byte_stream: Raw bytes from ``httpx.Response.aiter_bytes()``.
|
|
:yields: Typed :class:`StreamEvent` instances.
|
|
"""
|
|
buf = ""
|
|
current_event: str | None = None
|
|
|
|
async for chunk in byte_stream:
|
|
buf += chunk.decode("utf-8", errors="replace")
|
|
while "\n" in buf:
|
|
line, buf = buf.split("\n", 1)
|
|
line = line.rstrip("\r")
|
|
|
|
if line.startswith("event: "):
|
|
current_event = line[7:]
|
|
elif line.startswith("data: "):
|
|
data_str = line[6:]
|
|
if data_str.strip() == "[DONE]":
|
|
# The server's terminal sentinel is a bare ``data: [DONE]``
|
|
# with no preceding ``event:`` line, so detect it regardless
|
|
# of ``current_event``.
|
|
return
|
|
# A ``data:`` line is only a parseable event when an ``event:``
|
|
# line preceded it; a lone data line (other than [DONE]) is
|
|
# ignored.
|
|
if current_event is not None:
|
|
try:
|
|
data = json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
_log.warning("Failed to parse SSE data: %s", data_str[:200])
|
|
current_event = None
|
|
continue
|
|
event = _parse_event(current_event, data)
|
|
if event is not None:
|
|
yield event
|
|
current_event = None
|
|
elif line == "":
|
|
current_event = None
|
|
|
|
|
|
def _normalize_event_type(event_type: str) -> str:
|
|
"""Normalize event type to handle server enum rendering.
|
|
|
|
The server builds terminal events as ``f"response.{task.status}"``
|
|
where ``task.status`` may be a Python enum (rendering as
|
|
``response.TaskStatus.COMPLETED``) instead of the expected
|
|
``response.completed``. Normalize by extracting and lowercasing
|
|
the enum value.
|
|
"""
|
|
if ".TaskStatus." in event_type:
|
|
# "response.TaskStatus.COMPLETED" → "response.completed"
|
|
parts = event_type.split(".")
|
|
status = parts[-1].lower()
|
|
return f"response.{status}"
|
|
return event_type
|
|
|
|
|
|
def _parse_event(event_type: str, data: dict[str, Any]) -> StreamEvent | None:
|
|
"""Convert a raw SSE event type + JSON data into a typed event.
|
|
|
|
:param event_type: Wire name of the SSE ``event:`` field, e.g.
|
|
``"response.output_text.delta"``. Compared against the
|
|
``_T_RESPONSE_*`` wire-type constants (sourced from
|
|
:mod:`omnigent.server.schemas`) to dispatch.
|
|
:param data: Decoded JSON payload from the SSE ``data:`` field.
|
|
:returns: A typed :class:`StreamEvent` for known event names, or
|
|
``None`` when the payload is missing required fields or the
|
|
event type is unrecognized (forward-compatible skip).
|
|
"""
|
|
event_type = _normalize_event_type(event_type)
|
|
|
|
# Response lifecycle
|
|
if event_type == _T_RESPONSE_CREATED:
|
|
return ResponseCreated(response=_parse_response(data))
|
|
if event_type == _T_RESPONSE_QUEUED:
|
|
return ResponseQueued(response=_parse_response(data))
|
|
if event_type == _T_RESPONSE_IN_PROGRESS:
|
|
return ResponseInProgress(response=_parse_response(data))
|
|
if event_type == _T_RESPONSE_COMPLETED:
|
|
return ResponseCompleted(response=_parse_response(data))
|
|
if event_type == _T_RESPONSE_FAILED:
|
|
return ResponseFailed(response=_parse_response(data))
|
|
if event_type == _T_RESPONSE_INCOMPLETE:
|
|
resp = _parse_response(data)
|
|
reason = ""
|
|
if resp.incomplete_details is not None:
|
|
reason = resp.incomplete_details.reason
|
|
return ResponseIncomplete(response=resp, reason=reason)
|
|
if event_type == _T_RESPONSE_CANCELLED:
|
|
return ResponseCancelled(response=_parse_response(data))
|
|
|
|
# Text streaming
|
|
if event_type == _T_RESPONSE_OUTPUT_TEXT_DELTA:
|
|
delta = data.get("delta")
|
|
if isinstance(delta, str):
|
|
return TextDelta(delta=delta)
|
|
return None
|
|
|
|
# Reasoning
|
|
if event_type == _T_RESPONSE_REASONING_STARTED:
|
|
return ReasoningStarted()
|
|
if event_type == _T_RESPONSE_REASONING_TEXT_DELTA:
|
|
delta = data.get("delta")
|
|
if isinstance(delta, str):
|
|
return ReasoningDelta(delta=delta)
|
|
return None
|
|
if event_type == _T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA:
|
|
delta = data.get("delta")
|
|
if isinstance(delta, str):
|
|
return ReasoningSummaryDelta(delta=delta)
|
|
return None
|
|
|
|
# Output items
|
|
if event_type == _T_RESPONSE_OUTPUT_ITEM_DONE:
|
|
return _parse_output_item(data)
|
|
|
|
# File output
|
|
if event_type == _T_RESPONSE_OUTPUT_FILE_DONE:
|
|
return OutputFileDone(
|
|
file_id=str(data.get("file_id", "")),
|
|
filename=str(data["filename"]) if data.get("filename") is not None else None,
|
|
content_type=str(data["content_type"])
|
|
if data.get("content_type") is not None
|
|
else None,
|
|
)
|
|
|
|
# Retry
|
|
if event_type == _T_RESPONSE_RETRY:
|
|
return RetryEvent(
|
|
source=str(data.get("source", "")),
|
|
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
|
attempt=int(data.get("attempt", 0)),
|
|
max_attempts=int(data.get("max_attempts", 0)),
|
|
delay_seconds=float(data.get("delay_seconds", 0.0)),
|
|
error=_parse_error_info(data.get("error", {})),
|
|
)
|
|
|
|
# Error
|
|
if event_type == _T_RESPONSE_ERROR:
|
|
return ErrorEvent(
|
|
source=str(data.get("source", "")),
|
|
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
|
error=_parse_error_info(data.get("error", {})),
|
|
)
|
|
|
|
# Compaction
|
|
if event_type == _T_RESPONSE_COMPACTION_IN_PROGRESS:
|
|
return CompactionInProgress()
|
|
|
|
# Async client-tool cancel notification
|
|
if event_type == _T_RESPONSE_CLIENT_TASK_CANCEL:
|
|
task_id = data.get("task_id")
|
|
if isinstance(task_id, str) and task_id:
|
|
raw_call_id = data.get("call_id")
|
|
call_id = raw_call_id if isinstance(raw_call_id, str) and raw_call_id else None
|
|
return ClientTaskCancel(task_id=task_id, call_id=call_id)
|
|
_log.warning("response.client_task.cancel missing task_id: %r", data)
|
|
return None
|
|
|
|
# MCP-shape elicitation request (POLICIES.md §7 + the
|
|
# universal API additions in
|
|
# designs/SERVER_HARNESS_CONTRACT.md). The ``params`` block
|
|
# mirrors MCP's ``ElicitRequestFormParams`` field-for-field;
|
|
# we surface it as a typed event so the stream consumer can
|
|
# route it to the elicitation hook (not into a ToolHandler).
|
|
if event_type == "response.elicitation_request":
|
|
elicitation_id = data.get("elicitation_id")
|
|
if not isinstance(elicitation_id, str) or not elicitation_id:
|
|
_log.warning(
|
|
"response.elicitation_request missing elicitation_id: %r",
|
|
data,
|
|
)
|
|
return None
|
|
params = data.get("params")
|
|
if not isinstance(params, dict):
|
|
_log.warning(
|
|
"response.elicitation_request missing params dict: %r",
|
|
data,
|
|
)
|
|
return None
|
|
requested_schema = params.get("requestedSchema")
|
|
return ElicitationRequest(
|
|
elicitation_id=elicitation_id,
|
|
message=str(params.get("message") or ""),
|
|
# MCP spec restricts requestedSchema to a JSON-Schema
|
|
# dict — we accept anything dict-shaped and pass it
|
|
# through; non-dict inputs (defensive against
|
|
# malformed producers) become an empty schema.
|
|
requested_schema=requested_schema if isinstance(requested_schema, dict) else {},
|
|
mode=str(params.get("mode") or "form"),
|
|
phase=str(params.get("phase") or ""),
|
|
policy_name=str(params.get("policy_name") or ""),
|
|
content_preview=str(params.get("content_preview") or ""),
|
|
url=str(params["url"]) if isinstance(params.get("url"), str) else None,
|
|
)
|
|
|
|
# Unknown event — skip gracefully for forward-compatibility
|
|
_log.debug("Skipping unknown SSE event type: %s", event_type)
|
|
return None
|
|
|
|
|
|
def _parse_output_item(data: dict[str, Any]) -> StreamEvent | None:
|
|
"""Parse a ``response.output_item.done`` event into a typed event."""
|
|
item = data.get("item")
|
|
if not isinstance(item, dict):
|
|
return None
|
|
|
|
item_type = item.get("type", "")
|
|
|
|
if item_type == "function_call":
|
|
args_str = str(item.get("arguments", "{}"))
|
|
try:
|
|
arguments = json.loads(args_str)
|
|
except json.JSONDecodeError:
|
|
arguments = {}
|
|
name = str(item.get("name", ""))
|
|
call_id = str(item.get("call_id", ""))
|
|
return ToolCall(
|
|
name=name,
|
|
arguments=arguments,
|
|
call_id=call_id,
|
|
status=str(item.get("status", "")),
|
|
agent_name=str(item.get("model", "")),
|
|
)
|
|
|
|
if item_type == "function_call_output":
|
|
raw_arguments = item.get("arguments")
|
|
arguments: dict[str, object] = {}
|
|
if isinstance(raw_arguments, dict):
|
|
arguments = raw_arguments
|
|
elif isinstance(raw_arguments, str) and raw_arguments:
|
|
try:
|
|
decoded_arguments = json.loads(raw_arguments)
|
|
except json.JSONDecodeError:
|
|
decoded_arguments = {}
|
|
if isinstance(decoded_arguments, dict):
|
|
arguments = decoded_arguments
|
|
return ToolResult(
|
|
call_id=str(item.get("call_id", "")),
|
|
output=str(item.get("output", "")),
|
|
arguments=arguments,
|
|
)
|
|
|
|
if item_type == "message":
|
|
content = item.get("content", [])
|
|
return MessageDone(
|
|
content=content if isinstance(content, list) else [],
|
|
)
|
|
|
|
if item_type in NATIVE_TOOL_TYPES:
|
|
return NativeToolCall(
|
|
tool_type=item_type,
|
|
data=item,
|
|
)
|
|
|
|
# Compaction items, reasoning items, etc. — skip
|
|
_log.debug("Skipping output item type: %s", item_type)
|
|
return None
|
|
|
|
|
|
def _parse_response(data: dict[str, Any]) -> Response:
|
|
"""Extract the Response object from an SSE event payload."""
|
|
resp_data = data.get("response")
|
|
if isinstance(resp_data, dict):
|
|
return Response.from_dict(resp_data)
|
|
# Some events put fields at the top level
|
|
return Response.from_dict(data)
|
|
|
|
|
|
def _parse_error_info(raw: Any) -> ErrorInfo:
|
|
"""Parse an ErrorInfo from a nested dict."""
|
|
if isinstance(raw, dict):
|
|
return ErrorInfo(
|
|
code=str(raw.get("code", "")),
|
|
message=str(raw.get("message", "")),
|
|
)
|
|
return ErrorInfo(code="", message=str(raw))
|