feat(auth): add tenant_id to session and request context by andylim-duo · Pull Request #3 · andylim-duo/python-sdk

andylim-duo · 2026-03-11T19:09:55Z

Summary

This PR threads tenant_id from authentication tokens through the request lifecycle, enabling handlers to access tenant context for multi-tenant isolation.

Problem: After adding tenant_id to auth tokens (PR #2), there was no way for request handlers to access this information during request processing.

Solution: Propagate tenant_id through the session and context layers:

Add tenant_id field to base RequestContext (inherited by ServerRequestContext)
Add get_tenant_id() helper to extract tenant from auth context
Populate tenant_id in both ServerRequestContext instantiation sites
Add tenant_id property to ServerSession for session-level tenant tracking

Changes

src/mcp/shared/_context.py: Added tenant_id: str | None = None to base RequestContext
src/mcp/server/auth/middleware/auth_context.py: Added get_tenant_id() helper function
src/mcp/server/lowlevel/server.py: Updated both ServerRequestContext instantiation sites to populate tenant_id=get_tenant_id()
src/mcp/server/session.py: Added _tenant_id field and tenant_id property/setter to ServerSession

Part of Multi-Tenancy Implementation

This is Iteration 2 of a 6-part multi-tenancy implementation plan, building on Iteration 1 (PR #2).

After this PR, handlers can access ctx.tenant_id to determine which tenant is making a request:

@mcp.tool()
async def my_tool(ctx: Context) -> str:
    tenant_id = ctx.tenant_id  # Available from auth context
    # ... tenant-specific logic

Test Plan

New tests for get_tenant_id() in tests/server/auth/middleware/test_auth_context.py
New test file tests/server/test_multi_tenancy_session.py with 6 tests covering:
- RequestContext with/without tenant_id
- ServerRequestContext tenant_id inheritance
- ServerSession tenant_id property
- get_tenant_id() auth context extraction
All existing tests pass (1150 passed, 98 skipped, 1 xfailed)
Type checking passes: uv run --frozen pyright
Linting passes: uv run --frozen ruff check

Stacked PR

This PR is based on feature/multi-tenant-auth-tokens (PR #2). Merge PR #2 first, then this PR can be rebased onto main.

Thread tenant_id from authentication tokens through the request lifecycle to enable tenant-scoped operations in handlers. Changes: - Add tenant_id field to base RequestContext (inherited by ServerRequestContext) - Add get_tenant_id() helper in auth_context module to extract tenant from auth - Populate tenant_id in both ServerRequestContext instantiation sites in lowlevel/server.py - Add tenant_id property with getter/setter to ServerSession This is iteration 2 of the multi-tenancy implementation, building on the tenant_id field added to auth tokens in iteration 1.

Use proper async with pattern to ensure ServerSession's internal streams are cleaned up correctly, preventing resource warnings.

Add two tests to verify tenant_id doesn't leak between: - Concurrent async requests using the auth contextvar - Separate ServerSession instances These tests validate critical security properties for multi-tenant deployments where isolation between tenants must be guaranteed.

…tring Document the purpose and usage of the tenant_id field for multi-tenant server deployments.

andylim-duo

Review Summary

Overall the PR is well-structured and the get_tenant_id() helper + contextvar approach is clean. The test coverage is thorough, especially the concurrent isolation test. A few things to address before merging:

Key Issues

Disconnected tenant_id on ServerSession vs RequestContext — The PR adds tenant_id to both places, but only RequestContext.tenant_id is actually populated by the framework (via get_tenant_id()). ServerSession.tenant_id is never set by any code path, creating two disconnected sources of truth for the same concept. This should be reconciled — either remove the session property for now, or wire it up.
Mutable tenant_id setter on ServerSession — Allowing tenant reassignment mid-session is a potential security concern. Consider making it set-once or constructor-initialized.
Core server depends on auth module — lowlevel/server.py (used by all transports) now imports from the auth module. This couples the transport-agnostic server to HTTP-specific auth. Worth considering whether this should be made pluggable or at least documented as HTTP-only.

Minor

Dead code in test_get_tenant_id_with_tenant — unused MockApp and first middleware assignment.
Prefer anyio.lowlevel.checkpoint() over anyio.sleep(0.01) for context switching in tests.

src/mcp/server/session.py

src/mcp/server/lowlevel/server.py

tests/server/auth/middleware/test_auth_context.py

tests/server/test_multi_tenancy_session.py

…uest Wire up session.tenant_id so it is set automatically from the auth contextvar on the first authenticated request (set-once semantics). This connects RequestContext.tenant_id and ServerSession.tenant_id, ensuring the session is bound to a tenant for its lifetime.

Extract _simulate_tenant_binding helper to avoid pyright narrowing session.tenant_id to a literal type after assertion, which caused the subsequent `is None` check to be flagged as always-False.

…on handling Add two E2E tests using Client(server) that exercise the session.tenant_id set-once binding in lowlevel/server.py, covering the previously uncovered branches in _handle_request (line 456) and _handle_notification (line 504).

This line is now covered by the E2E tenant notification test.

Make the tenant_id setter raise ValueError if attempting to change to a different value once already set. This prevents accidental tenant reassignment which could be a security issue. Setting to the same value is allowed (idempotent).

…action Move tenant identification to a transport-agnostic contextvar (tenant_id_var) in the shared layer, removing the hard dependency from lowlevel/server.py on the auth middleware module. AuthContextMiddleware now sets tenant_id_var alongside auth_context_var, and the core server reads from the shared contextvar instead of calling get_tenant_id() from the auth module. This keeps the dependency direction correct (auth → shared, server → shared) and allows other transports to set tenant_id_var through their own mechanisms.

Remove unused MockApp() and AuthContextMiddleware(app) that were immediately overwritten by AuthContextMiddleware(TenantCheckApp()).

Use checkpoint() for deterministic context switching instead of a fixed-duration sleep in the tenant isolation test.

Import checkpoint directly from anyio.lowlevel to fix pyright reportAttributeAccessIssue on the lazy submodule.

…review concerns Add remove_resource() and remove_prompt() methods to MCPServer to match the existing remove_tool() API, enabling dynamic deprovisioning of all resource types for multi-tenant servers. PR review fixes: - Fix heading grammar: "Simple Registration of..." (concern #1) - Add private API warning for _lowlevel_server usage in example (concern #3) - Clarify example server needs TokenVerifier for production (concern #4) - Guard offboard_tenant against KeyError for unprovisioned tools (concern #5) - Add remove_resource/remove_prompt to MCPServer and docs (concern #6) Tests added for both single-tenant and multi-tenant remove operations, including cross-tenant isolation verification.

andylim-duo added 2 commits March 11, 2026 11:24

fix(test): use async context manager for ServerSession test

f597175

Use proper async with pattern to ensure ServerSession's internal streams are cleaned up correctly, preventing resource warnings.

Base automatically changed from feature/multi-tenant-auth-tokens to main March 11, 2026 20:08

andylim-duo added 3 commits March 11, 2026 16:10

Merge branch 'main' into feature/multi-tenant-session-context

f5351e0

Merge branch 'main' into feature/multi-tenant-session-context

2f6fa78

andylim-duo self-assigned this Mar 12, 2026

docs(context): add tenant_id field description to RequestContext docs…

9f4b679

…tring Document the purpose and usage of the tenant_id field for multi-tenant server deployments.

andylim-duo commented Mar 12, 2026

View reviewed changes

andylim-duo added 7 commits March 12, 2026 09:19

fix(test): resolve pyright reportUnnecessaryComparison error

161f123

Extract _simulate_tenant_binding helper to avoid pyright narrowing session.tenant_id to a literal type after assertion, which caused the subsequent `is None` check to be flagged as always-False.

style: apply ruff formatting to test file

2dcebd0

fix: remove stale pragma no cover from send_roots_list_changed

ec564e7

This line is now covered by the E2E tenant notification test.

Repository owner deleted a comment from alim Mar 13, 2026

andylim-duo added 3 commits March 16, 2026 13:59

fix(test): remove dead code in test_get_tenant_id_with_tenant

02725c8

Remove unused MockApp() and AuthContextMiddleware(app) that were immediately overwritten by AuthContextMiddleware(TenantCheckApp()).

fix(test): replace anyio.sleep(0.01) with anyio.lowlevel.checkpoint()

04c7535

Use checkpoint() for deterministic context switching instead of a fixed-duration sleep in the tenant isolation test.

fix(test): use explicit import for anyio.lowlevel.checkpoint

92dff29

Import checkpoint directly from anyio.lowlevel to fix pyright reportAttributeAccessIssue on the lazy submodule.

andylim-duo merged commit 0c36ac5 into main Mar 16, 2026
28 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(auth): add tenant_id to session and request context#3

feat(auth): add tenant_id to session and request context#3
andylim-duo merged 16 commits intomainfrom
feature/multi-tenant-session-context

andylim-duo commented Mar 11, 2026

Uh oh!

andylim-duo left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

andylim-duo commented Mar 11, 2026

Summary

Changes

Part of Multi-Tenancy Implementation

Test Plan

Stacked PR

Uh oh!

andylim-duo left a comment

Choose a reason for hiding this comment

Review Summary

Key Issues

Minor

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants