make context available in hooks, add OnRegisterSession hook

[zahmadsaleem]

Hi,

I wanted to make use of the sessionId generated by the McpServer to manage my own session data. The issue I was facing was, I could not get the sessionId from the /sse route and correlate to the message's sessionId. Also, I had no way of accessing the session information in the hooks. SSEContextFunc combined with OnRegisterSession will allow us to connect the initial /sse session to all messages outside the library.
I think context should be available on all hooks to trace it back to the request which may contain context information added by upstream http middlewares.

Summary by CodeRabbit

• New Features

  • Introduced an enhanced client session hook to improve session registration and handling.

• Refactor

  • Standardized context propagation across core server operations, improving request cancellation, deadline management, and error reporting.
  • Enhanced the message processing and session management workflows for increased stability and responsiveness.

[coderabbitai]

Walkthrough

Several files have been updated to add a context.Context parameter to function signatures. Hook callback functions in the MCPServer, its tests, and the hooks package were modified accordingly. In addition, a new session registration hook type and its associated methods were introduced. Session management methods in the server, SSE, and Stdio components have been updated to propagate context, ensuring consistent context awareness.

Changes

File(s)	Change Summary
examples/everything/main.go, server/hooks.go, server/request_handler.go, server/server_test.go	Updated hook callback signatures to include context.Context; added a new hook type (OnRegisterSessionHookFunc) and corresponding hook registration methods.
server/server.go, server/sse.go, server/stdio.go	Modified session registration methods to accept a context.Context parameter, ensuring context propagation in session management calls.

Possibly related PRs

• Fix tool calls #38: Introduced similar function signature changes involving the addition of context.Context to hook callbacks.
• feat: allow customising context via callback on transport servers #32: Focused on integrating context handling in the SSEServer and StdioServer structs, aligning with the current modifications.
• Add per-session notifications handling #46: Made related modifications to function signatures to incorporate context.Context in hook functions and session management methods.

Suggested reviewers

• ezynda3

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 051cda5 and 5d2c4e4.

⛔ Files ignored due to path filters (2)

• server/internal/gen/hooks.go.tmpl is excluded by !**/gen/**
• server/internal/gen/request_handler.go.tmpl is excluded by !**/gen/**

📒 Files selected for processing (7)

• examples/everything/main.go (1 hunks)
• server/hooks.go (14 hunks)
• server/request_handler.go (9 hunks)
• server/server.go (1 hunks)
• server/server_test.go (10 hunks)
• server/sse.go (1 hunks)
• server/stdio.go (1 hunks)

🧰 Additional context used

🧬 Code Definitions (2)

examples/everything/main.go (2)

mcp/types.go (2)

• InitializeRequest (266-275)
• InitializeResult (279-293)

mcp/tools.go (2)

• CallToolRequest (44-59)
• CallToolResult (34-41)

server/hooks.go (2)

server/server.go (1)

• ClientSession (49-58)

mcp/types.go (20)

• MCPMethod (11-11)
• InitializeRequest (266-275)
• InitializeResult (279-293)
• PingRequest (355-357)
• EmptyResult (232-232)
• ListResourcesRequest (400-402)
• ListResourcesResult (406-409)
• ListResourceTemplatesRequest (413-415)
• ListResourceTemplatesResult (419-422)
• ReadResourceRequest (426-435)
• ReadResourceResult (439-442)
• MethodInitialize (16-16)
• MethodPing (20-20)
• MethodResourcesList (24-24)
• MethodResourcesTemplatesList (28-28)
• MethodResourcesRead (32-32)
• MethodPromptsList (36-36)
• MethodPromptsGet (40-40)
• MethodToolsList (44-44)
• MethodToolsCall (48-48)

🔇 Additional comments (15)

server/sse.go (1)

219-219: Good addition of context propagation to session registration.

This change properly propagates the HTTP request context to the RegisterSession method, enabling access to the context in hooks and allowing for tracing back to the original request. This aligns well with the PR objective of making context available in hooks.

server/stdio.go (1)

194-194: Consistent context propagation implementation.

The change ensures that the stdio server correctly passes the context to session registration, maintaining parity with the SSE implementation. This is a good practice for context propagation throughout the system.

server/server.go (2)

174-176: Appropriate signature update to support context propagation.

Adding the context parameter to the RegisterSession method is a good enhancement that enables context-aware operations during session registration. This change is fundamental to the PR's objective of making context available in hooks.

---

182-182: Well implemented session registration hook.

The addition of this hook call enables custom processing during session registration with access to both the context and session information. This is crucial for implementing the OnRegisterSession hook functionality mentioned in the PR objectives.

server/request_handler.go (1)

69-69: Comprehensive context propagation to all hook calls.

The changes consistently pass the context to all hook method calls throughout the request handling process. This systematic approach ensures that context information is available at every stage of request processing, which is essential for tracing and context-aware operations.

Also applies to: 73-73, 76-76, 88-89, 92-92, 95-96, 113-114, 117-117, 120-121, 138-139, 142-142, 145-146, 163-164, 167-167, 170-171, 188-189, 192-192, 195-196, 213-214, 217-217, 220-221, 238-239, 242-242, 245-246, 263-264, 267-267, 270-271

server/server_test.go (2)

181-181: Properly updated RegisterSession call with context parameter.

The addition of context.TODO() parameter to RegisterSession is consistent with the PR objectives to make context available in hooks. This change ensures that the test cases align with the updated method signature.

---

212-212: Consistent context parameter usage throughout test cases.

All RegisterSession calls are consistently updated to include the context.TODO() parameter across multiple test scenarios, maintaining compatibility with the new method signature. This ensures that all test cases properly exercise the updated functionality.

Also applies to: 221-221, 254-254, 282-282

examples/everything/main.go (1)

35-55: Successfully updated hook function signatures with context parameter.

The hook function signatures have been properly updated to include a context.Context parameter. This change is aligned with the PR objectives of making context available in all hooks, enabling better tracing of requests and access to context information added by upstream HTTP middlewares.

The changes are consistent across all hook types:

• BeforeAny
• OnSuccess
• OnError
• BeforeInitialize
• AfterInitialize
• AfterCallTool
• BeforeCallTool

server/hooks.go (7)

11-13: Added new OnRegisterSession hook type.

This addition implements the core functionality described in the PR objectives - creating an OnRegisterSessionHookFunc that will allow developers to link the initial /sse session to all subsequent messages outside the library. This hook is triggered when a new session is registered, providing access to both the context and the client session.

---

16-16: Updated existing hook function signatures with context parameter.

All hook function signatures have been properly updated to include a context.Context parameter. This change enables hooks to access the context associated with the request, which may contain important information added by upstream HTTP middlewares.

Also applies to: 20-20, 28-28, 55-55

---

57-82: Consistently applied context parameter to all method-specific hooks.

The context parameter has been systematically added to all method-specific hook functions, ensuring that context information is consistently available throughout the hook chain. This provides improved traceability and access to request-scoped values across the entire request lifecycle.

---

85-85: Added OnRegisterSession slice to Hooks struct.

The OnRegisterSession field has been added to the Hooks struct to store registered session hooks. This complements the new hook type and enables the hook registration mechanism.

---

166-205: Updated hook execution methods with context parameter.

The internal hook execution methods (beforeAny, onSuccess, onError) have been updated to accept and pass the context parameter to registered hooks. This ensures that context is propagated consistently throughout the hook execution chain.

---

207-218: Implemented new session registration hook functionality.

These new methods implement the core functionality needed for the OnRegisterSession hook:

1. AddOnRegisterSession - Allows registering new session hooks
2. RegisterSession - Executes all registered session hooks with the provided context and session

This implementation is critical for fulfilling the PR objective of making session IDs accessible for managing session data and linking sessions to messages.

---

227-460: Consistently updated all hook execution methods with context parameter.

All method-specific hook execution functions (like beforeInitialize, afterInitialize, etc.) have been updated to include the context parameter and pass it to both general hooks (via beforeAny/onSuccess) and method-specific hooks. This ensures that context is consistently available throughout all hook execution chains, maintaining the integrity of the request context flow.

✨ Finishing Touches

• 📝 Generate Docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

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)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai plan to trigger planning for file edits and PR creation.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.