Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Mar 13, 2025

Conversation

dsp-ant
Copy link
Member

@dsp-ant dsp-ant commented Mar 3, 2025

Summary

  • Reorganizes message handling in the codebase by introducing MessageFrame (previously ParsedMessage) and improving type definitions
  • Enhances code organization by moving stream type definitions to more appropriate locations

Test plan

  • Ensure all existing tests pass
  • Verify type checking completes without errors
  • Check that renamed components maintain backward compatibility

GitHub-Issue:#201

@dsp-ant dsp-ant force-pushed the davidsp/raw-request branch from 2e3388b to 0d651d5 Compare March 3, 2025 20:09
@dsp-ant dsp-ant requested review from jerome3o-anthropic, jspahrsummers and bhosmer-ant and removed request for jerome3o-anthropic March 4, 2025 10:44
bhosmer-ant
bhosmer-ant previously approved these changes Mar 4, 2025
Copy link
Contributor

@bhosmer-ant bhosmer-ant left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Couple minor observations inline but LGTM!

RawT = TypeVar("RawT")


class MessageFrame(BaseModel, Generic[RawT]):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't know how fastidious we are about this in general, but fwiw the other classes here are docstringed up pretty thoroughly (and the point of raw in particular might be worth noting here)

dsp-ant added 9 commits March 13, 2025 13:23
Move memory stream type definitions to models.py and use them throughout
the codebase for better type safety and maintainability.

GitHub-Issue:#201
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
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
@dsp-ant dsp-ant force-pushed the davidsp/raw-request branch from 1d57efa to 1c53fc2 Compare March 13, 2025 13:40
@dsp-ant dsp-ant enabled auto-merge (squash) March 13, 2025 13:40
@dsp-ant dsp-ant disabled auto-merge March 13, 2025 13:44
@dsp-ant dsp-ant merged commit 9d0f2da into main Mar 13, 2025
6 checks passed
@dsp-ant dsp-ant deleted the davidsp/raw-request branch March 13, 2025 13:44
@Kludex
Copy link
Member

Kludex commented Mar 14, 2025

This broke a lot types 😢

Comment on lines +32 to +35
ReadStream = MemoryObjectReceiveStream[MessageFrame | Exception]
ReadStreamWriter = MemoryObjectSendStream[MessageFrame | Exception]
WriteStream = MemoryObjectSendStream[MessageFrame]
WriteStreamReader = MemoryObjectReceiveStream[MessageFrame]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can't do this if MessageFrame is a generic. It means that all those new types have unkown parameters to the generic type.

Comment on lines +211 to +223
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)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not type safe. You are hiding the parameters with args and kwargs.

@Kludex
Copy link
Member

Kludex commented Mar 14, 2025

What's this ParsedMessage you mention?

read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception],
write_stream: MemoryObjectSendStream[types.JSONRPCMessage],
read_stream: ReadStream,
write_stream: WriteStream,
Copy link
Member

@Kludex Kludex Mar 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not really the end type, since ReadStream and WriteStream are generics, something is missing here.

I believe it's easier to understand the code if we don't hide that those streams are anyio objects. Makes it easier to go around the code.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants