refactor: reorganize message handling for better type safety and clarity (#239)

* refactor: improve typing with memory stream type aliases

Move memory stream type definitions to models.py and use them throughout
the codebase for better type safety and maintainability.

GitHub-Issue:#201

* refactor: move streams to ParsedMessage

* refactor: update test files to use ParsedMessage

Updates test files to work with the ParsedMessage stream type aliases
and fixes a line length issue in test_201_client_hangs_on_logging.py.

Github-Issue:#201

* refactor: rename ParsedMessage to MessageFrame for clarity

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

* refactor: move MessageFrame class to types.py for better code organization

🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>

* fix pyright

* refactor: update websocket client to use MessageFrame

Modified the websocket client to work with the new MessageFrame type,
preserving raw message text and properly extracting the root JSON-RPC
message when sending.

Github-Issue:#204

* fix: use NoneType instead of None for type parameters in MessageFrame

🤖 Generated with [Claude Code](https://claude.ai/code)

* refactor: rename root to message
This commit is contained in:
David Soria Parra
2025-03-13 13:44:55 +00:00
committed by GitHub
parent ad7f7a5473
commit 9d0f2daddb
17 changed files with 283 additions and 151 deletions

View File

@@ -180,6 +180,49 @@ class JSONRPCMessage(
pass
RawT = TypeVar("RawT")
class MessageFrame(BaseModel, Generic[RawT]):
"""
A wrapper around the general message received that contains both the parsed message
and the raw message.
This class serves as an encapsulation for JSON-RPC messages, providing access to
both the parsed structure (root) and the original raw data. This design is
particularly useful for Server-Sent Events (SSE) consumers who may need to access
additional metadata or headers associated with the message.
The 'root' attribute contains the parsed JSONRPCMessage, which could be a request,
notification, response, or error. The 'raw' attribute preserves the original
message as received, allowing access to any additional context or metadata that
might be lost in parsing.
This dual representation allows for flexible handling of messages, where consumers
can work with the structured data for standard operations, but still have the
option to examine or utilize the raw data when needed, such as for debugging,
logging, or accessing transport-specific information.
"""
message: JSONRPCMessage
raw: RawT | None = None
model_config = ConfigDict(extra="allow")
def model_dump(self, *args, **kwargs):
"""
Dumps the model to a dictionary, delegating to the root JSONRPCMessage.
This method allows for consistent serialization of the parsed message.
"""
return self.message.model_dump(*args, **kwargs)
def model_dump_json(self, *args, **kwargs):
"""
Dumps the model to a JSON string, delegating to the root JSONRPCMessage.
This method provides a convenient way to serialize the parsed message to JSON.
"""
return self.message.model_dump_json(*args, **kwargs)
class EmptyResult(Result):
"""A response that indicates success but carries no data."""