Enrich Identity with upstream tokens and simplify upstreamswap by jhrozek · Pull Request #4355 · stacklok/toolhive

jhrozek · 2026-03-24T21:52:41Z

Summary

The upstreamswap middleware was coupled to the token storage layer via a ServiceGetter pattern, making it responsible for session lookup, token refresh, and error classification on every request. This moves all that complexity into the auth middleware where it belongs.
During JWT validation, when a tsid claim is present, the auth middleware now bulk-loads all upstream provider tokens via GetAllValidTokens and attaches them to Identity.UpstreamTokens. The upstreamswap middleware becomes a trivial map reader: identity.UpstreamTokens[cfg.ProviderName].
Infrastructure errors (storage outage) return 503 from the auth middleware — not 401 — preventing infinite re-authentication loops.
Expired tokens whose refresh fails are omitted from the result (fail-closed), matching the behavior on main.
The tsid session identifier is filtered from PrincipalInfo.Claims so it does not leak into webhook payloads or audit logs.
GetUpstreamTokenService is removed from the MiddlewareRunner interface (no remaining callers).

Closes #4139

Type of change

Refactoring (no behavior change)

Test plan

Unit tests (task test)
Linting (task lint-fix)

Changes

File	Change
`pkg/auth/identity.go`	Add `UpstreamTokens` field with redacted `MarshalJSON`
`pkg/auth/identity_test.go`	Tests for redaction (populated, empty, nil, mixed)
`pkg/auth/upstreamtoken/types.go`	Add `UpstreamTokenReader` interface
`pkg/auth/upstreamtoken/service.go`	Add `GetAllValidTokens` (bulk read with refresh)
`pkg/auth/upstreamtoken/service_test.go`	Tests for bulk path (fresh, expired, refresh, nil refresher)
`pkg/auth/token.go`	Add `WithUpstreamTokenReader` option, `loadUpstreamTokens`, 503 on storage failure
`pkg/auth/token_test.go`	Unit tests for `loadUpstreamTokens` + integration tests for full middleware pipeline
`pkg/auth/context.go`	Filter `tsid` from externalized claims
`pkg/auth/utils.go`	Accept variadic `TokenValidatorOption`
`pkg/auth/middleware.go`	Thread `UpstreamTokenReader` from runner into validator
`pkg/auth/middleware_test.go`	Updated mock expectations
`pkg/transport/types/transport.go`	Add `GetUpstreamTokenReader`, remove `GetUpstreamTokenService`
`pkg/transport/types/mocks/mock_transport.go`	Regenerated mock
`pkg/runner/runner.go`	Store `upstreamTokenReader`, remove dead `upstreamTokenService` field
`pkg/runner/runner_test.go`	Tests for `GetUpstreamTokenReader`
`pkg/auth/upstreamswap/middleware.go`	Simplified to read from `Identity.UpstreamTokens`
`pkg/auth/upstreamswap/middleware_test.go`	Simplified tests (no storage mocks)
`docs/middleware.md`	Updated upstream swap docs + `MiddlewareRunner` interface block

Does this introduce a user-facing change?

No. Non-AS deployments are completely unaffected (the enrichment guard is nil-checked). AS deployments get the same token injection behavior through a different internal path.

Special notes for reviewers

Commit structure: 6 commits, each self-contained and compilable. Read bottom-up: data model → service → plumbing → consumer → simplification → cleanup.
Non-AS safety: All callers of GetAuthenticationMiddleware outside the runner path pass zero options, so upstreamTokenReader is nil and the enrichment code is unreachable.
503 vs 401: The old upstreamswap distinguished infrastructure errors (503) from client errors (401). The new architecture preserves this — the auth middleware returns 503 when storage is down, and upstreamswap returns 401 when the provider token is missing.
Fail-closed on refresh failure: GetAllValidTokens omits providers whose tokens expired and could not be refreshed, matching main's behavior.
Performance: Single-upstream Redis: 2 calls vs 1 (marginal). Multi-upstream: 2 calls vs N (improvement). Memory storage: O(entries) scan vs O(1) (acceptable).

Generated with Claude Code

Large PR Justification

we moved the place where the upstream tokens are injected and had to redo the tests. The actual load-bearing code hasn't changed that much.

github-actions

Large PR Detected

This PR exceeds 1000 lines of changes and requires justification before it can be reviewed.

How to unblock this PR:

Add a section to your PR description with the following format:

## Large PR Justification

[Explain why this PR must be large, such as:]
- Generated code that cannot be split
- Large refactoring that must be atomic
- Multiple related changes that would break if separated
- Migration or data transformation

Alternative:

Consider splitting this PR into smaller, focused changes (< 1000 lines each) for easier review and reduced risk.

See our Contributing Guidelines for more details.

This review will be automatically dismissed once you add the justification section.

codecov · 2026-03-24T22:22:40Z

Codecov Report

❌ Patch coverage is 96.36364% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.65%. Comparing base (a9f92a0) to head (bd92306).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
pkg/auth/middleware.go	50.00%	1 Missing and 1 partial ⚠️
pkg/auth/utils.go	0.00%	1 Missing ⚠️
pkg/runner/runner.go	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4355      +/-   ##
==========================================
+ Coverage   69.55%   69.65%   +0.09%     
==========================================
  Files         480      480              
  Lines       48837    48948     +111     
==========================================
+ Hits        33971    34095     +124     
+ Misses      12249    12229      -20     
- Partials     2617     2624       +7

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Large PR justification has been provided. Thank you!

github-actions · 2026-03-25T11:19:22Z

✅ Large PR justification has been provided. The size review has been dismissed and this PR can now proceed with normal review.

pkg/auth/upstreamtoken/service.go

pkg/auth/token_test.go

The auth middleware will populate upstream provider access tokens on the Identity struct, allowing downstream middleware (upstreamswap) to read tokens without coupling to the storage layer. MarshalJSON redacts token values while preserving provider keys, and both nil and empty maps are omitted via omitempty. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Introduce a narrow UpstreamTokenReader interface that decouples the auth middleware from storage internals. InProcessService.GetAllValidTokens performs a bulk read of all upstream providers for a session, refreshing expired tokens transparently and falling back to expired tokens when refresh fails. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Add GetUpstreamTokenReader to MiddlewareRunner so the auth middleware can access the bulk token reader for identity enrichment. The runner stores the reader as a separate field (set alongside the token service) avoiding type assertions. GetAuthenticationMiddleware now accepts variadic TokenValidatorOption for forward-compatible option passing. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Add WithUpstreamTokenReader option to TokenValidator that loads all upstream provider tokens from storage when a tsid claim is present in the JWT. The enrichment happens between claimsToIdentity and context injection, populating Identity.UpstreamTokens for downstream middleware consumption. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Replace the ServiceGetter/GetValidTokens pattern with a direct read from the pre-enriched Identity.UpstreamTokens map. The auth middleware now handles storage reads and token refresh, so upstreamswap becomes a pure reader that looks up the provider token and injects it into the request header. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

The upstreamswap middleware no longer calls the token service directly, so GetUpstreamTokenService is dead code. Remove the interface method, its Runner implementation, the mock, tests, and docs reference. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Rename the type to avoid stutter (upstreamtoken.TokenReader instead of upstreamtoken.UpstreamTokenReader). Method and function names that include "UpstreamTokenReader" are unchanged since they don't stutter. Also remove the unused providerName parameter from the test helper requestWithIdentity in upstreamswap tests. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Replace hand-written mockUpstreamTokenReader with go.uber.org/mock generated MockTokenReader, per project convention of never hand-writing mocks. Add //go:generate directive to types.go. Add providerName to log lines in refreshOrFail so operators can identify which upstream provider failed during concurrent refreshes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

jhrozek requested review from ChrisJBurns, JAORMX, amirejaz, blkt, lujunsan and yrobla as code owners March 24, 2026 21:52

github-actions bot added the size/XL Extra large PR: 1000+ lines changed label Mar 24, 2026

github-actions bot previously requested changes Mar 24, 2026

View reviewed changes

jhrozek force-pushed the authserver-multi-upstream/4-tokens-in-identity branch from c54cd11 to 0b369e5 Compare March 24, 2026 22:18

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Mar 24, 2026

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Mar 25, 2026

jhrozek force-pushed the authserver-multi-upstream/4-tokens-in-identity branch from 70131b5 to 0d2fbeb Compare March 25, 2026 11:39

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Mar 25, 2026

tgrunnagle previously approved these changes Mar 25, 2026

View reviewed changes

pkg/auth/upstreamtoken/service.go Show resolved Hide resolved

pkg/auth/token_test.go Outdated Show resolved Hide resolved

jhrozek and others added 8 commits March 25, 2026 15:44

jhrozek dismissed tgrunnagle’s stale review via bd92306 March 25, 2026 15:45

jhrozek force-pushed the authserver-multi-upstream/4-tokens-in-identity branch from 0d2fbeb to bd92306 Compare March 25, 2026 15:45

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Mar 25, 2026

tgrunnagle approved these changes Mar 25, 2026

View reviewed changes

jhrozek merged commit 7c66773 into main Mar 25, 2026
80 checks passed

jhrozek deleted the authserver-multi-upstream/4-tokens-in-identity branch March 25, 2026 16:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enrich Identity with upstream tokens and simplify upstreamswap#4355

Enrich Identity with upstream tokens and simplify upstreamswap#4355
jhrozek merged 8 commits intomainfrom
authserver-multi-upstream/4-tokens-in-identity

jhrozek commented Mar 24, 2026 •

edited

Loading

Uh oh!

github-actions bot left a comment

Uh oh!

codecov bot commented Mar 24, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Mar 25, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jhrozek commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Type of change

Test plan

Changes

Does this introduce a user-facing change?

Special notes for reviewers

Large PR Justification

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Large PR Detected

How to unblock this PR:

Alternative:

Uh oh!

codecov bot commented Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions bot commented Mar 25, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jhrozek commented Mar 24, 2026 •

edited

Loading

codecov bot commented Mar 24, 2026 •

edited

Loading