Handle each session initialization separately #60
Conversation
WalkthroughThis pull request refactors session initialization across several server components. The changes add two new methods— Changes
Possibly related PRs
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
server/server_test.go (1)
1095-1100: Implement ClientSession interface methods in fakeSession.The Initialize method is empty since the test sets the initialized state directly, and the Initialized method correctly returns the state. Consider making Initialize set the flag to true for consistency with other implementations.
-func (f fakeSession) Initialize() { +func (f fakeSession) Initialize() { + // Since fakeSession is a value type, this wouldn't update the receiver + // This is fine for tests, but in real implementations we use a pointer receiver + // or atomic operations to update the state }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
server/server.go(6 hunks)server/server_test.go(8 hunks)server/sse.go(3 hunks)server/stdio.go(3 hunks)
🧰 Additional context used
🧬 Code Definitions (2)
server/server_test.go (2)
mcp/types.go (3) (3)
t(15-17)t(19-30)JSONRPCNotification(152-155)server/server.go (2) (2)
MCPServer(75-88)ClientSessionFromContext(63-68)
server/server.go (2)
server/stdio.go (9) (9)
s(58-60)s(62-64)s(66-68)s(70-72)s(95-97)s(102-104)s(109-187)s(192-215)s(219-234)server/sse.go (7) (7)
s(33-35)s(37-39)s(41-43)s(45-47)s(146-153)s(157-170)s(174-255)
🔇 Additional comments (16)
server/sse.go (3)
11-11: Add sync/atomic import for atomic.Bool usage.The import of "sync/atomic" is added to support the new atomic.Bool field for thread-safe access to the initialization state.
25-25: Add initialized state flag to sseSession.Adding an atomic.Bool field ensures thread-safe operations on the initialization state flag, which is essential for concurrent access in an SSE server context.
41-47: Implement session initialization interface methods.These methods properly implement the ClientSession interface, allowing for per-session initialization tracking. The atomic operations ensure thread-safety when accessing the initialization state.
server/stdio.go (3)
12-12: Add sync/atomic import for atomic.Bool usage.The import of "sync/atomic" is required for the atomic.Bool field to ensure thread-safe operations on the initialization state.
55-55: Add initialized state flag to stdioSession.Adding an atomic.Bool field provides thread-safe access to the initialization state, consistent with the implementation in sseSession.
66-72: Implement ClientSession interface methods.These methods properly implement the required interface methods for session initialization, using atomic operations for thread safety.
server/server.go (6)
49-52: Add initialization methods to ClientSession interface.These new interface methods formalize the contract for session initialization, allowing each session implementation to manage its own initialization state instead of using a global flag.
142-143: Check session initialization before sending notifications.This change ensures that notifications are only sent to fully initialized sessions. This is important to prevent sending notifications to sessions that might not be ready to handle them yet.
160-162: Verify session initialization in SendNotificationToClient.This change adds protection against sending notifications to uninitialized sessions, returning an appropriate error message. This helps clients understand when notifications can't be sent and why.
533-534: Update comment to reflect session initialization check.The comment is updated to clarify that notifications are only sent to initialized sessions, which is consistent with the implementation changes.
553-554: Update comment for DeleteTools notification.Similar to the previous comment update, this clarifies that DeleteTools notifications are only sent to initialized sessions.
617-619: Initialize client session when handling initialize request.The server now initializes the client session directly rather than setting a global flag. This is the core change that enables per-session initialization tracking.
server/server_test.go (4)
183-183: Set session as initialized in test.The test now properly sets the initialized state to true for the test session, ensuring that notifications can be sent to it.
210-226: Test both initialized and uninitialized sessions.This test case now clearly distinguishes between initialized and uninitialized sessions, ensuring that the server correctly handles both types and only sends notifications to initialized sessions.
505-518: Add test case for uninitialized sessions.This new test case verifies that attempting to send notifications to uninitialized sessions results in an error, which is important for validating the new behavior.
1084-1084: Add initialized field to fakeSession struct.The fakeSession struct now includes the initialized flag to match the updated ClientSession interface requirements.
torkve
force-pushed
the
per-session-initialize
branch
from
47011cd to
1d5b53c
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
server/server_test.go (1)
1180-1185: Simple implementation of the ClientSession interfaceThe implementation of the required
Initialize()andInitialized()methods for thefakeSessionstruct allows for proper testing of the initialization logic, althoughInitialize()could be improved to actually update the initialized field state.Consider updating the
Initialize()method to set the initialized field to true:func (f fakeSession) Initialize() { + f.initialized = true }However, since
fakeSessionis being used with value semantics in the tests rather than pointers, this change might not have any effect without also updating how the sessions are created and stored in the tests.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
server/server.go(6 hunks)server/server_test.go(8 hunks)server/sse.go(3 hunks)server/stdio.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- server/stdio.go
- server/sse.go
🧰 Additional context used
🧬 Code Definitions (1)
server/server.go (2)
server/stdio.go (9)
s(58-60)s(62-64)s(66-68)s(70-72)s(95-97)s(102-104)s(109-187)s(192-215)s(219-234)server/sse.go (7)
s(34-36)s(38-40)s(42-44)s(46-48)s(161-168)s(172-185)s(189-265)
🔇 Additional comments (10)
server/server.go (6)
50-53: Appropriate interface extension to support per-session initializationThe addition of
Initialize()andInitialized() boolmethods to theClientSessioninterface is a well-designed approach that decouples session state management from the server, allowing each session type to handle its initialization state independently.
206-207: Good safeguard for notification deliveryThe condition to check
session.Initialized()before sending notifications prevents sending data to sessions that aren't ready to process it, which is a robust defensive programming practice.
225-225: Enhanced error handling for client notificationsThe additional check for
!session.Initialized()aligns with the PR objective of handling each session independently and ensures notifications are only sent to properly initialized sessions.
413-414: Clear notification comment improvementThe updated comment clarifies that notifications are only sent to initialized sessions, which matches the implementation and improves code readability.
433-434: Consistent notification comment improvementThis comment mirrors the one on line 413-414, maintaining consistency throughout the codebase when referring to the same behavior.
497-499: Per-session initialization implementationThis change correctly replaces the previous global initialization flag with per-session initialization, ensuring that each session can be managed independently as intended by the PR objective.
server/server_test.go (4)
184-184: Tests properly initialize active sessionsThe test cases appropriately set the
initializedfield totruefor active sessions, ensuring they receive notifications as expected.Also applies to: 215-215, 223-224, 257-257
219-225: Good test coverage for uninitialized sessionsAdding test cases for uninitialized sessions is important for validating that the new initialization check works correctly by preventing notification delivery to sessions that aren't ready.
505-519: Comprehensive test for uninitialized session notificationsThis test properly verifies that attempting to send a notification to an uninitialized session results in an error, ensuring the defensive check in
SendNotificationToClientworks as intended.
1169-1169: Proper field addition to the fakeSession structThe addition of the
initializedboolean field to thefakeSessionstruct enables proper testing of the initialization state management.
This removes global initialized flag, because each session can be initialized separately.
Summary by CodeRabbit