fix session WIP

[pavelkraleu]

Summary by CodeRabbit

• New Features

  • Improved session management with explicit session scoping and cleanup capabilities.
  • Enhanced logging for session and server operations, including session IDs and server creation details.

• Bug Fixes

  • Unified session lifecycle handling and removed reliance on a default session, reducing potential session inconsistencies.

• Refactor

  • Server management now uses a session-aware map for more precise tracking and removal.
  • Session management logic simplified for clarity and consistency.
  • Internal context handling updated for immutability and clearer session ID access.

• Chores

  • Removed support for stdio transport; HTTP mode is now required.
  • Updated .gitignore to exclude .idea directory.
  • Removed session management tools and related exports no longer in use.

[coderabbitai]

Walkthrough

This set of changes refactors session and server management throughout the codebase, removing the concept of a default browser session and introducing explicit session scoping via unique MCP session IDs. The Context class and related functions now require and propagate these session IDs, with enhanced logging for session operations. The ServerList class was reworked to use a map keyed by session IDs, supporting more precise server tracking and removal. Session management functions were simplified, with the removal of default session logic and the addition of a closeSession function. The stdio transport mode was removed, and HTTP transport now manages sessions using a new TransportSession interface. The session management tools in src/tools/session.ts were deleted, and related exports were removed from the tools index. Logging and error handling were standardized and expanded across the updated modules.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

• Complexity: The changes are broad and touch multiple core modules (context, server management, session management, transport, and tools), including removal of features, API changes, and internal refactoring.
• Review time: The reviewer will need to verify the correctness of session handling, ensure proper propagation of session IDs, check for regressions due to the removal of default session logic and stdio transport, and confirm that all logging and error handling is consistent and informative. The removal of session management tools and changes to exported entities also require careful attention to potential downstream impacts.

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

✨ Finishing Touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch pavel-session

---

🪧 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.
  • Explain this complex logic.
  • 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 explain this code block.
• 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 explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

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 generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai generate unit tests to generate unit tests for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @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.

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (2)

src/transport.ts (1)

39-50: Consider adding collision check for MCP session ID

While UUID collision is extremely unlikely, consider adding a check to ensure the generated mcpSessionId doesn't already exist in your session tracking system for complete robustness.

src/sessionManager.ts (1)

20-27: Consider consistent logging approach

The function mixes console.log for info messages and process.stderr.write for warnings. Consider using a consistent logging approach throughout the module.

📜 Review details

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

📥 Commits

Reviewing files that changed from the base of the PR and between 1a0c827 and 1261a87.

⛔ Files ignored due to path filters (1)

• package-lock.json is excluded by !**/package-lock.json

📒 Files selected for processing (10)

• .gitignore (1 hunks)
• src/context.ts (4 hunks)
• src/index.ts (2 hunks)
• src/program.ts (2 hunks)
• src/server.ts (1 hunks)
• src/sessionManager.ts (8 hunks)
• src/tools/index.ts (0 hunks)
• src/tools/multiSession.ts (1 hunks)
• src/tools/session.ts (0 hunks)
• src/transport.ts (2 hunks)

💤 Files with no reviewable changes (2)

• src/tools/index.ts
• src/tools/session.ts

🧰 Additional context used

🧬 Code Graph Analysis (4)

src/program.ts (2)

src/server.ts (1)

• ServerList (3-55)

src/transport.ts (1)

• startHttpTransport (76-127)

src/transport.ts (1)

src/server.ts (1)

• ServerList (3-55)

src/context.ts (3)

config.d.ts (1)

• Config (4-101)

src/sessionManager.ts (2)

• getSession (257-305)
• closeSession (360-377)

src/types/types.ts (2)

• BrowserSession (25-30)
• MCPTool (38-38)

src/index.ts (1)

src/context.ts (1)

• Context (9-162)

🔇 Additional comments (26)

.gitignore (1)

296-297: Addition of .idea directory looks good

Ignoring IDE-specific metadata keeps the repository clean. No further action required.

src/program.ts (3)

8-8: LGTM: Stdio transport removal is consistent with the refactoring.

The removal of startStdioTransport from the import aligns with the broader architectural change to HTTP-only transport.

---

71-76: LGTM: Port validation enforces HTTP-only mode correctly.

The mandatory port requirement with clear error messaging properly enforces the new HTTP-only transport mode.

---

82-93: LGTM: Server factory integrates well with session-aware architecture.

The server factory properly accepts and logs the mcpSessionId parameter, which integrates correctly with the updated ServerList class that manages servers by session ID. The unconditional call to startHttpTransport is appropriate given the removal of stdio transport.

src/index.ts (3)

118-124: LGTM: Function signature correctly supports session-aware architecture.

Adding the optional mcpSessionId parameter maintains flexibility while enabling the session-aware server creation pattern.

---

132-136: LGTM: Session ID validation is necessary and well-implemented.

The validation correctly enforces the requirement for mcpSessionId in HTTP mode, which is necessary for the session-aware Context class integration.

---

157-158: LGTM: Context integration with session ID is correct.

The Context constructor call properly passes the mcpSessionId, which aligns with the updated Context class that requires this parameter for session-scoped browser management.

src/tools/multiSession.ts (1)

67-81: LGTM: Context wrapper refactoring improves clarity and aligns with Context class changes.

The shift from Object.create() to object spreading with explicit method overrides is cleaner and more maintainable. The currentSessionId getter correctly aligns with the updated Context class design where this property is read-only.

src/server.ts (4)

4-9: LGTM: Map-based server storage enables proper session management.

Changing from array to Map<string, Server> allows for efficient session-keyed server management, and the updated factory signature supports session-aware server creation.

---

11-28: LGTM: Create method properly handles session-aware server management.

The method correctly uses the mcpSessionId as a key when provided and falls back to generating a unique timestamp-based key for legacy scenarios. The enhanced logging provides good visibility into server lifecycle management.

---

30-44: LGTM: Close method correctly handles Map-based server removal.

The method properly finds the server by reference in the Map and removes it by key. The enhanced logging provides good traceability for server lifecycle management.

---

46-55: LGTM: CloseAll method efficiently handles bulk server closure.

The method correctly converts Map values to an array for concurrent closure and properly clears the Map afterward. The logging provides good visibility into the bulk closure process.

src/context.ts (6)

6-6: LGTM: Import change aligns with explicit session management.

Replacing defaultSessionId with closeSession is consistent with the removal of default session concepts in favor of explicit session lifecycle management.

---

11-24: LGTM: Constructor properly establishes session-aware context.

The required mcpSessionId parameter and generated browserSessionId create a clear hierarchy between MCP sessions and browser sessions. The logging provides excellent traceability for debugging session-related issues.

---

26-29: LGTM: Read-only getter prevents accidental session ID modification.

Making currentSessionId a read-only getter that returns browserSessionId provides clear access to the session identifier while preventing accidental modifications.

---

38-48: LGTM: Enhanced getStagehand method supports flexible session targeting.

The optional sessionId parameter enables both default and explicit session targeting, while the enhanced logging and improved error messages provide better debugging support.

---

52-68: LGTM: Consistent logging pattern improves debugging capability.

The enhanced logging with both MCP and browser session IDs provides excellent context for debugging session-related issues across all browser access methods.

---

80-96: LGTM: Cleanup method provides proper session lifecycle management.

The cleanup method correctly uses the closeSession function for browser session cleanup, with appropriate error handling and comprehensive logging to ensure reliable resource management during MCP session termination.

src/transport.ts (4)

8-12: Good interface design for session tracking

The TransportSession interface effectively groups related session data, improving code organization and type safety.

---

22-36: Comprehensive session lookup with appropriate logging

The session lookup logic properly handles missing sessions and provides detailed logging for debugging.

---

52-63: Robust cleanup logic with detailed logging

The cleanup handler properly checks for transport.sessionId existence before attempting deletion and provides comprehensive logging for debugging session lifecycle.

---

81-84: Clear startup logging and proper type usage

The addition of startup logging and the updated map type properly align with the new session tracking approach.

src/sessionManager.ts (4)

10-14: Appropriate state management improvements

The nullable activeSessionId better represents session state, and the initialization logging aids debugging.

---

113-123: Clean session disconnect handling

The simplified disconnect handler properly manages the active session state without the complexity of default session handling.

---

256-305: Excellent session management refactoring

The refactored getSession function provides clear session validation, proper cleanup of stale sessions, and comprehensive logging. The removal of default session logic simplifies the code while maintaining robustness.

---

356-377: Well-implemented session closure function

The new closeSession function follows established patterns, handles edge cases properly, and maintains consistency with other session management functions.