feat!: add context.Context to TokenStore methods #557
Conversation
In distributed environments a TokenStore will be implemented as a database, and often a context is required when running get/store operations to ensure cancellation propagation, instrumentation, etc works when running token store ops. This is a breaking change but I think it makes sense (and it's pretty easy to fix any breakages).
WalkthroughAdded context.Context parameters to TokenStore.GetToken and SaveToken, introduced ErrNoToken, updated MemoryTokenStore to check ctx and use locks, propagated ctx through getValidToken/refresh/ProcessAuthorizationResponse, and updated tests to exercise context propagation and cancellation/deadline scenarios. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (2)
🧰 Additional context used🧬 Code graph analysis (2)client/transport/oauth.go (1)
client/transport/oauth_test.go (2)
🔇 Additional comments (15)
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type 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: 1
🧹 Nitpick comments (8)
client/transport/oauth_test.go (2)
57-61: Context plumbed through tests — LGTM. Consider one cancellation-path test.The shift to ctx-aware store calls looks correct. As a nice-to-have, add a test that verifies behavior when context is already canceled before SaveToken/GetToken are invoked (useful for non-memory implementations).
I can add a small subtest that calls SaveToken/GetToken with a canceled context to assert timely return. Want me to draft it?
Also applies to: 75-82
219-221: Loosen brittle error-string assertion to reduce flakiness across platforms.Network error messages can vary (e.g., include “dial tcp” on some platforms). Include one more token to keep the test robust.
Apply this diff:
- if !strings.Contains(err.Error(), "connection refused") && - !strings.Contains(err.Error(), "failed to send protected resource request") { + if !strings.Contains(err.Error(), "connection refused") && + !strings.Contains(err.Error(), "failed to send protected resource request") && + !strings.Contains(err.Error(), "dial") {client/transport/oauth.go (4)
80-87: MemoryTokenStore: honor ctx cancellation (cheap and future-proof).Even though this is in-memory, early-returning on canceled/deadline contexts better matches the interface contract for drop-in replacements.
Apply this diff:
func (s *MemoryTokenStore) GetToken(ctx context.Context) (*Token, error) { - s.mu.RLock() + if err := ctx.Err(); err != nil { + return nil, err + } + s.mu.RLock() defer s.mu.RUnlock() if s.token == nil { return nil, errors.New("no token available") } return s.token, nil }
90-95: MemoryTokenStore.SaveToken: mirror cancellation check.Consistent with GetToken above.
Apply this diff:
func (s *MemoryTokenStore) SaveToken(ctx context.Context, token *Token) error { - s.mu.Lock() + if err := ctx.Err(); err != nil { + return err + } + s.mu.Lock() defer s.mu.Unlock() s.token = token return nil }
152-169: Preserve cancellation from TokenStore.GetToken.If the underlying store returns context.Canceled/DeadlineExceeded, surface it instead of collapsing to authorization required.
Apply this diff:
func (h *OAuthHandler) getValidToken(ctx context.Context) (*Token, error) { - token, err := h.config.TokenStore.GetToken(ctx) + token, err := h.config.TokenStore.GetToken(ctx) + if err != nil { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return nil, err + } + } if err == nil && !token.IsExpired() && token.AccessToken != "" { return token, nil }
221-229: Nit: handle canceled ctx while peeking old token.You ignore the error getting oldToken (fine), but if it’s due to ctx cancellation, better to propagate it. Optional tweak:
oldToken, oldErr := h.config.TokenStore.GetToken(ctx) if oldErr != nil && (errors.Is(oldErr, context.Canceled) || errors.Is(oldErr, context.DeadlineExceeded)) { return nil, oldErr }client/transport/streamable_http_oauth_test.go (1)
18-19: Avoid data races on requestCount in test server.The handler runs on the test HTTP server goroutine while assertions read requestCount in the main goroutine. Under -race this can flag a data race.
Apply these diffs:
- Import atomic:
import ( "context" "encoding/json" "errors" "net/http" "net/http/httptest" "testing" "time" + "sync/atomic" "github.com/mark3labs/mcp-go/mcp" )
- Use atomic counter:
- // Track request count to simulate 401 on first request, then success - requestCount := 0 + // Track request count to simulate 401 on first request, then success + var requestCount int32
- Increment atomically in handler:
- requestCount++ + atomic.AddInt32(&requestCount, 1)
- Read atomically in assertion:
- if requestCount != 1 { - t.Errorf("Expected server to receive 1 request, got %d", requestCount) + if got := atomic.LoadInt32(&requestCount); got != 1 { + t.Errorf("Expected server to receive 1 request, got %d", got) }Also applies to: 27-31, 115-117
client/transport/sse_oauth_test.go (1)
16-19: Prevent test data races on shared counters.Similar to the HTTP test, requestCount is written in the server goroutine and read in the test goroutine.
Apply these diffs:
- Import atomic:
import ( "context" "encoding/json" "errors" "net/http" "net/http/httptest" "testing" "time" + "sync/atomic" )
- Use atomic counter and boolean:
- requestCount := 0 + var requestCount int32 - sseEndpointSent := false + var sseEndpointSent atomic.Bool
- Increment and set atomically inside handler:
- requestCount++ + atomic.AddInt32(&requestCount, 1) ... - sseEndpointSent = true + sseEndpointSent.Store(true)
- Read atomically in assertions:
- if requestCount != 1 { - t.Errorf("Expected server to receive 1 request, got %d", requestCount) + if got := atomic.LoadInt32(&requestCount); got != 1 { + t.Errorf("Expected server to receive 1 request, got %d", got) } ... - if !sseEndpointSent { + if !sseEndpointSent.Load() { t.Errorf("Expected SSE endpoint to be sent") }Also applies to: 28-33, 135-137
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
client/oauth_test.go(2 hunks)client/transport/oauth.go(6 hunks)client/transport/oauth_test.go(4 hunks)client/transport/sse_oauth_test.go(2 hunks)client/transport/streamable_http_oauth_test.go(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
client/transport/oauth.go (1)
client/oauth.go (3)
Token(14-14)MemoryTokenStore(20-20)TokenStore(17-17)
🔇 Additional comments (5)
client/transport/oauth_test.go (1)
160-172: Good: ctx now passed into SaveToken in setup.This aligns with the new interface and keeps the test intent intact.
client/transport/oauth.go (1)
640-641: LGTM: SaveToken now ctx-aware.This keeps the token write cancellable and consistent with the new interface.
client/transport/streamable_http_oauth_test.go (1)
16-17: LGTM: tests updated to pass ctx into SaveToken.No logical changes beyond context propagation; setup remains clear and deterministic.
Also applies to: 63-65
client/oauth_test.go (1)
16-16: LGTM: ctx-aware SaveToken usage in client-side test.Matches the new public surface while keeping the test semantics unchanged.
Also applies to: 55-56
client/transport/sse_oauth_test.go (1)
14-15: LGTM: context passed into SaveToken in SSE tests.Change is minimal and correct.
Also applies to: 84-86
|
|
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
client/transport/oauth.go (1)
36-42: Document ctx/error semantics and introduce a sentinel ErrNoToken.Reiterating a prior suggestion: please document that implementations must honor ctx cancellation/deadlines and should return a sentinel error to distinguish “no token” from operational failures. Also, add that callers must propagate context errors unchanged.
Proposed doc tweak (example):
type TokenStore interface { - // GetToken returns the current token + // GetToken returns the current token. + // Implementations must honor ctx cancellation and deadlines (returning context.Canceled + // or context.DeadlineExceeded). If no token exists, return ErrNoToken. Any other + // operational/storage errors must be returned as-is. GetToken(ctx context.Context) (*Token, error) - // SaveToken saves a token + // SaveToken saves a token. + // Implementations must honor ctx cancellation/deadlines and return them unchanged. SaveToken(ctx context.Context, token *Token) error }And add a package-level sentinel:
// ErrNoToken indicates the store has no token for the client. var ErrNoToken = errors.New("no token available")
🧹 Nitpick comments (4)
client/transport/oauth.go (2)
79-90: Return ErrNoToken (not a new string error) and avoid external mutation by returning a copy.
- Use the sentinel ErrNoToken so call sites can distinguish “no token” from failures.
- Returning the stored pointer allows external mutation; returning a shallow copy is safer.
func (s *MemoryTokenStore) GetToken(ctx context.Context) (*Token, error) { if err := ctx.Err(); err != nil { return nil, err } s.mu.RLock() defer s.mu.RUnlock() if s.token == nil { - return nil, errors.New("no token available") + return nil, ErrNoToken } - return s.token, nil + t := *s.token + return &t, nil }
93-101: Define explicit behavior for nil token saves (clear vs. reject).Today a nil token will just set s.token = nil (allowed), but the behavior isn’t explicit. If “clear” is desired, make it explicit; otherwise reject nil to catch bugs.
Option A (explicitly clear):
func (s *MemoryTokenStore) SaveToken(ctx context.Context, token *Token) error { if err := ctx.Err(); err != nil { return err } - s.mu.Lock() - defer s.mu.Unlock() - s.token = token + s.mu.Lock() + defer s.mu.Unlock() + if token == nil { + s.token = nil + return nil + } + s.token = token return nil }Option B (reject):
- s.mu.Lock() + if token == nil { + return errors.New("SaveToken: token must not be nil") + } + s.mu.Lock() defer s.mu.Unlock()client/transport/streamable_http_oauth_test.go (1)
142-144: Avoid data race on authHeaderReceived (written in handler, read in test).The handler goroutine writes authHeaderReceived; the test goroutine reads it without synchronization. Use atomic.Value or a mutex, or remove this final read and rely on in-handler assertions.
One minimal approach:
// declare near other vars var authHeader atomic.Value // stores string // in handler, instead of direct assignment: authHeader.Store(r.Header.Get("Authorization")) // when asserting: got, _ := authHeader.Load().(string) if got != "Bearer test-token" { t.Errorf("Expected Authorization header 'Bearer test-token', got '%s'", got) }client/transport/oauth_test.go (1)
660-801: Production ‘stale context’ scenarios are well captured; consider reducing string matching.The scenarios are realistic and valuable. Minor: assertions that rely on substrings in error messages (e.g., ensuring errors don’t mention “database”) can be brittle. Prefer errors.Is/As on context errors, or sentinel errors, over string contains.
If you adopt ErrNoToken as suggested, I can generate complementary tests to assert that GetToken returns ErrNoToken for empty stores and that getValidToken treats ErrNoToken as “no token” while surfacing other errors.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
client/transport/oauth.go(6 hunks)client/transport/oauth_test.go(5 hunks)client/transport/streamable_http_oauth_test.go(4 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
client/transport/oauth_test.go (2)
client/transport/oauth.go (5)
NewMemoryTokenStore(75-77)Token(45-58)OAuthConfig(18-34)TokenStore(37-42)NewOAuthHandler(130-139)client/transport/streamable_http.go (1)
ErrOAuthAuthorizationRequired(231-231)
client/transport/oauth.go (1)
client/oauth.go (3)
Token(14-14)MemoryTokenStore(20-20)TokenStore(17-17)
🔇 Additional comments (8)
client/transport/oauth.go (2)
241-243: LGTM — SaveToken(ctx, …) correctly propagates context errors.This keeps refresh flows cancelable and observably propagates failures to callers.
654-656: LGTM — SaveToken(ctx, …) propagation on auth-code exchange.Consistent with the rest of the ctx-aware API surface.
client/transport/streamable_http_oauth_test.go (3)
9-9: Good call importing sync/atomic for request counting.Sets up for race-free increments in the handler.
64-66: LGTM — SaveToken(ctx, …) update matches the new API.
116-118: LGTM — atomic.LoadInt32 for verification.Race-free read for assertion.
client/transport/oauth_test.go (3)
57-61: LGTM — tests now pass context into GetToken/SaveToken.The updates reflect the breaking API and verify basic store behavior with ctx. Good coverage for the new signature.
Also applies to: 75-76, 81-82, 162-172, 171-173
399-467: Nice addition — MemoryTokenStore cancellation/deadline tests.These subtests pin the required ctx semantics and will catch regressions.
469-523: Solid ctx propagation tests through OAuth handler paths.These ensure GetAuthorizationHeader and getValidToken propagate ctx errors from the store and HTTP operations.
Also applies to: 525-600, 602-658
sd2k
force-pushed
the
tokenstore-takes-context
branch
from
d9986bd to
35199e0
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
client/transport/oauth.go (1)
80-89: Return ErrNoToken instead of a generic error.Standardizes the “no token” case and enables callers to distinguish absence from operational failures.
func (s *MemoryTokenStore) GetToken(ctx context.Context) (*Token, error) { if err := ctx.Err(); err != nil { return nil, err } s.mu.RLock() defer s.mu.RUnlock() if s.token == nil { - return nil, errors.New("no token available") + return nil, ErrNoToken } return s.token, nil }
♻️ Duplicate comments (3)
client/transport/oauth.go (3)
36-42: Clarify TokenStore ctx/error semantics and introduce ErrNoToken sentinel.Document cancellation/deadline behavior in the interface comments and standardize the “no token” case via a sentinel error for reliable handling downstream.
Apply this doc tweak:
type TokenStore interface { - // GetToken returns the current token - GetToken(ctx context.Context) (*Token, error) - // SaveToken saves a token - SaveToken(ctx context.Context, token *Token) error + // GetToken returns the current token. + // Implementations must honor ctx cancellation and deadlines, returning + // context.Canceled or context.DeadlineExceeded as appropriate. + // If no token is present, return ErrNoToken. + GetToken(ctx context.Context) (*Token, error) + // SaveToken saves a token. + // Implementations must honor ctx cancellation and deadlines and propagate ctx errors. + SaveToken(ctx context.Context, token *Token) error }And add this package-level sentinel (outside the shown range):
var ErrNoToken = errors.New("no token available")I can push a follow-up commit wiring this through the package if you’d like.
232-238: Remove unnecessary store read during refresh; reuse provided refresh token.Avoids extra IO and ctx-related failure risk; you already have refreshToken.
-// If no new refresh token is provided, keep the old one -oldToken, oldErr := h.config.TokenStore.GetToken(ctx) -if oldErr != nil && (errors.Is(oldErr, context.Canceled) || errors.Is(oldErr, context.DeadlineExceeded)) { - return nil, oldErr -} -if tokenResp.RefreshToken == "" && oldToken != nil { - tokenResp.RefreshToken = oldToken.RefreshToken -} +// If no new refresh token is provided, keep using the one we had +if tokenResp.RefreshToken == "" { + tokenResp.RefreshToken = refreshToken +}
159-165: Don’t swallow non-context errors from TokenStore; only treat ErrNoToken as “no token”.Propagate real failures; otherwise you’ll mask storage/IO issues and spuriously fall back to auth.
token, err := h.config.TokenStore.GetToken(ctx) if err != nil { - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - return nil, err - } + switch { + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + return nil, err + case errors.Is(err, ErrNoToken): + // proceed; authorization/refresh flow will decide next steps + token = nil + default: + return nil, fmt.Errorf("token store GetToken: %w", err) + } }
🧹 Nitpick comments (5)
client/transport/oauth.go (1)
327-415: Optional: metadataOnce + ctx can permanently poison discovery.If the first discovery call’s ctx is canceled/timed out, metadataOnce prevents retries. Consider retrying on context errors or caching only successful results.
I can draft a small refactor using singleflight or a retry-once path if desired.
client/transport/oauth_test.go (4)
57-61: LGTM: tests pass context through MemoryTokenStore.Post-sentinel follow-up: if ErrNoToken is introduced, assert it explicitly for the empty store path.
Proposed test tweak after adding ErrNoToken:
- _, err := store.GetToken(ctx) - if err == nil { - t.Errorf("Expected error when getting token from empty store") - } + _, err := store.GetToken(ctx) + if !errors.Is(err, ErrNoToken) { + t.Errorf("Expected ErrNoToken when getting token from empty store, got %v", err) + }Also applies to: 75-76, 81-82
217-223: Loosen brittle error assertion for metadata fetch.Kernel/transport differences can vary error strings. Just assert non-nil or a small set of known prefixes.
- // Verify the error message contains something about a connection error - // since we're now trying to connect to the well-known endpoint - if !strings.Contains(err.Error(), "connection refused") && - !strings.Contains(err.Error(), "failed to send protected resource request") { - t.Errorf("Expected error message to contain connection error, got %s", err.Error()) - } + // Verify we get an error (message text is platform-dependent) + if err == nil { + t.Fatalf("expected an error fetching server metadata") + }
629-641: Adjust comment to avoid coupling to current refresh implementation detail.If refresh stops reading from the store, the cancellation source may differ. Keep the assertion; relax the comment.
- // Should return context.Canceled error (from getting old token) + // Should return context.Canceled error (propagated via the refresh path)
11-52: Optional: flake guard for IsExpired tests.If clocks are tight, add a small skew (e.g., ExpiresAt = now + 100ms) or compare using time.Until > 0. Non-blocking.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
client/transport/oauth.go(6 hunks)client/transport/oauth_test.go(5 hunks)client/transport/streamable_http_oauth_test.go(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- client/transport/streamable_http_oauth_test.go
🧰 Additional context used
🧬 Code graph analysis (2)
client/transport/oauth.go (1)
client/oauth.go (3)
Token(14-14)MemoryTokenStore(20-20)TokenStore(17-17)
client/transport/oauth_test.go (3)
client/transport/oauth.go (5)
NewMemoryTokenStore(75-77)Token(45-58)OAuthConfig(18-34)TokenStore(37-42)NewOAuthHandler(130-139)client/oauth.go (4)
NewMemoryTokenStore(23-23)Token(14-14)OAuthConfig(11-11)TokenStore(17-17)client/transport/streamable_http.go (1)
ErrOAuthAuthorizationRequired(231-231)
🔇 Additional comments (8)
client/transport/oauth.go (3)
93-101: LGTM: SaveToken honors ctx and uses proper locking.
241-243: LGTM: SaveToken now correctly propagates ctx to the store.
654-656: LGTM: SaveToken(ctx, ...) used in auth-code exchange path as well.client/transport/oauth_test.go (5)
160-196: LGTM: empty access token path correctly yields authorization required.
399-467: LGTM: MemoryTokenStore cancellation/deadline behavior is well covered.
469-523: LGTM: Authorization header surfaces ctx errors from TokenStore.
525-600: LGTM: getValidToken ctx-cancellation cases are exercised.
660-800: LGTM: cached client scenarios for canceled/deadline contexts look solid.
|
@sd2k looks alright to me. Will get a second pair of eyes on this first though. Also could you address the coderabbit issues with either a justification or fix in the mean time? |
|
@ezynda3 yep all done, cheers. FWIW we're running this in a fork and it's working fine, the issues I was seeing before were because the clients were being closed in another codepath, so subsequent uses caused the client's context to be canceled (using |
Description
In distributed environments a TokenStore will be implemented as a database, and often a context is required when running get/store operations to ensure cancellation propagation, instrumentation, etc works when running token store ops. This is a breaking change but I think it makes sense (and it's pretty easy to fix any breakages).
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Summary by CodeRabbit
Refactor
Tests