[kam-validmind]

Draft. Agent Authority (AA) needs this before it can merge validmind/agent-authority#152.

The short version

Atryum sits between AI agents and MCP servers. When an agent calls a tool, Atryum forwards the call to the MCP server (the "upstream") and attaches a credential, usually a bearer token.

Today Atryum attaches the same token for every caller of a server. This PR lets the program that embeds Atryum, in our case AA, choose the token for each call, based on who is calling.

Nothing changes for anyone who doesn't use the new option.

Why AA needs this

Many hosted MCP servers, such as Shortcut's, sign people in with OAuth. You don't paste a token; each person signs in with their own account. AA wants every agent's calls to run as the right account:

• A developer's own coding agent should act as that developer.
• A production bot should act as its own service account.

That way an agent can only do what its account is allowed to do, and the server's logs show who really did what.

Atryum couldn't support that, for four reasons:

1. One token per server. Atryum stores one OAuth token per server (oauth_credentials is keyed by server name), so everyone shares it.
2. No way in for the embedder. The code that adds the token is hard-wired, and it only knows which server is being called, not who is calling.
3. Shared sessions. Atryum keeps one MCP session open per server. If Sam's token opened the session, Priya's calls would reuse Sam's session.
4. Stale tokens on held calls. A call that waits for a person to approve it keeps the token it got before it was held. After an hour of waiting, that token may have expired.

What this PR adds

1. A new option: WithUpstreamCredentials

The embedder gives Atryum a function. Atryum calls it every time it is about to contact an upstream, and it tells the function:

• who is calling: the verified agent ID, issuer, sign-in method and key ID;
• which server: the server's name;
• why: one of four purposes, list (listing tools), call (a tool call, before the decision), forward (other MCP methods) or execute (right before the tool runs).

The function returns what to send:

Field	What it is
BearerToken	The token to send. Empty means "use the server's own token, if it has one".
Headers	Extra headers, merged over the server's own.
SessionKey	Which cached MCP session to use. Different credentials get different keys.
OnUnauthorized	Optional. Called if the upstream answers 401.

If the function returns an error, Atryum stops and passes the error back to the agent. It never quietly falls back to some other stored token.

sequenceDiagram
  accTitle: A tool call with the new option
  accDescr: Atryum asks the embedder for the caller's credential, the decider decides, and the call runs with the caller's own token.
  participant Agent
  participant Atryum
  participant AA as Embedder (AA)
  participant Server as MCP server
  Agent->>Atryum: tools/call
  Atryum->>AA: Credential for (agent, server, "call")?
  AA-->>Atryum: Sam's token, session key "shortcut/sam"
  Atryum->>AA: Decide this call
  AA-->>Atryum: Allow
  Atryum->>AA: Credential again (purpose "execute")
  AA-->>Atryum: Sam's current token
  Atryum->>Server: tools/call with Sam's token
  Server-->>Atryum: Result
  Atryum-->>Agent: Result

Loading

2. Asking again just before the tool runs

Right before Atryum runs a call (in finishExecution), it calls the function again with the purpose execute. This matters most for held calls:

sequenceDiagram
  accTitle: A held call released later
  accDescr: The token is fetched again when the held call is released, so it is current.
  participant Atryum
  participant AA as Embedder (AA)
  participant Person
  participant Server as MCP server
  Atryum->>AA: Credential (purpose "call")
  AA-->>Atryum: Token A
  Note over Atryum,Person: The call waits for approval.<br/>Token A may expire meanwhile.
  Person->>Atryum: Approve
  Atryum->>AA: Credential (purpose "execute")
  AA-->>Atryum: Token B (renewed)
  Atryum->>Server: tools/call with token B

Loading

If the function refuses at this point (for example, the person disconnected their account while the call was waiting), the call is marked failed with that message, and the server never sees the request.

3. One MCP session per credential

MCP over HTTP can use a session: the server hands back an Mcp-Session-Id, and later requests send it. Atryum caches it. Before, the cache was keyed by server name. Now it's keyed by SessionKey, which defaults to the server name, so nothing changes without the option.

flowchart TD
  accTitle: Session cache before and after
  accDescr: Before, all callers shared one session per server. After, each credential has its own.
  subgraph Before
    A1["Sam's call"] --> S1["Session for 'shortcut'"]
    B1["Priya's call"] --> S1
  end
  subgraph After
    A2["Sam's call"] --> S2["Session 'shortcut/sam'"]
    B2["Priya's call"] --> S3["Session 'shortcut/priya'"]
  end

Loading

When the upstream answers 401 Unauthorized, Atryum now:

1. throws away that credential's session, so the next request starts a fresh one;
2. calls OnUnauthorized before returning the error, so the embedder can mark the token for renewal.

4. The old stored-token path is switched off when the option is set

With WithUpstreamCredentials set, Atryum's own oauth_credentials lookup is never used, on any path. A leftover token in that table can't leak into a call.

5. A way to say "not available to you yet"

A harness lists a server's tools, and reads its instructions, before it calls anything, and it swallows an error there. Claude Code showed a server whose listing failed as connected with no tools and never said why. So the credential function may now return an UpstreamUnavailable{Message, Tools, Instructions} as its error:

• a tools/list is answered with the Tools it carries (AA sends one placeholder whose description explains and links), instead of an error;
• an initialize adds Instructions to the session's instructions, through a new initialize purpose that lets the function look the caller up without sending or renewing anything;
• any other contact fails with Message;
• the upstream is never contacted.

Without the option, or with no Tools to show, nothing changes.

Where to look in the code

File	What changed
pkg/atryum/upstream_credentials.go	New. The public option and its types (UpstreamContact, UpstreamCredential, the purpose constants), and the adapter that turns them into an internal provider.
pkg/atryum/options.go, pkg/atryum/atryum.go	Store the option and wire it into the resolver.
internal/mcp/purpose.go	New. Puts the purpose and invocation ID on the context, so no function signatures had to change.
internal/mcp/client.go	CredentialProvider on the resolver; Upstream gains SessionKey, OnUnauthorized and ResolvedByProvider; session cache keyed by sessionKey(); 401 handling.
internal/invocation/service.go	Sets the purpose at each resolve, and asks again in finishExecution.
internal/api/handlers.go	Sets the forward purpose.
CHANGELOG.md	An Unreleased entry.

Two small behaviours to know about:

• If the function refuses a forwarded method (not tools/list or tools/call), the method isn't forwarded. That's what Atryum already did when resolving a server failed.
• ListAll skips a server the function refuses, instead of failing the whole list.

Tests

New tests:

• internal/mcp/credential_provider_test.go
  • Two callers on one server get two different tokens and two different sessions.
  • Returning an empty credential never picks up the stored oauth_credentials token.
  • An error from the function stops the resolve, and ListAll skips that server.
  • A 401 clears the session and calls OnUnauthorized exactly once, before the error comes back.
• internal/invocation/credential_provider_test.go
  • A call asks the function twice, first call then execute (with the invocation ID), and the upstream only ever sees the execute token.
  • If the function refuses at execute, the call fails with its message and the upstream receives nothing.

To run them:

go test ./internal/mcp/ ./internal/invocation/ -run 'CredentialProvider|Unauthorized|ProviderAgain|RefusesAtExecute|PurposeAnd'

The whole suite (go test ./...) passes on top of main.

How AA uses it

AA's PR, validmind/agent-authority#152, passes a function to WithUpstreamCredentials. For each call it:

1. finds which agent is calling;
2. picks whose account to use: a developer's for their own agent, the agent's own for a production agent;
3. returns that account's token, with session key <server>/<account id>.

If no account is connected, it returns nothing for call, so AA's decider refuses the call with a "connect your account" link, and it returns an error for the other purposes. AA's end-to-end test runs two agents through this proxy against a fake OAuth server and checks that each call ran as its own user.

AA pins Atryum to a released tag, so this needs to be merged and tagged before AA's PR can merge.

How this was made

Written with Claude Code, from an AA feature request: "AA MCP doesn't seem to have an OAuth flow; in Atryum you can put in a server URL and it talks you through doing an OAuth flow." It went through a spec, a plan and a task list. The person's decisions:

• each person signs in with their own account, with no shared account;
• production agents get their own accounts;
• encrypting stored tokens is left for later.

This option was the smallest Atryum change that let AA do the rest itself. The full contract is in AA's specs/001-mcp-upstream-oauth/contracts/atryum-upstream-credentials.md.

[cachafla]

Nice work on this 👏 The provider design reads well, and the execute re-resolve does the right thing for held calls. I left three things I'd fix before we tag it:

1. ReportStuck fails for a caller with no connected account (internal/invocation/hooks.go:278)

ReportStuck calls ResolveContext(ctx, server) with no purpose, so it defaults to PurposeCall. It only needs the server's name and never contacts the upstream. With WithUpstreamCredentials set, AA refuses a caller with no connected account (an error or UpstreamUnavailable), and ReportStuck turns that into "unknown server". So the person who's stuck can't report it, and we ask the embedder for a token (maybe renewing it) for nothing.

I'd suggest looking the server up without going through credential() here, or resolving with a purpose that skips the credential (PurposeInitialize would work).

2. The provider also runs for stdio servers (internal/mcp/client.go:349)

Resolver.credential() calls the provider for every upstream, including ones with a Command and no BaseURL. The changelog and the option's doc say "every path that resolves an HTTP upstream". An embedder that refuses callers with no connected account would also block local servers that don't need a token at all.

I'd suggest skipping the provider when BaseURL is empty. Alternatively, document that it covers every server and add the transport to UpstreamContact so the embedder can tell them apart.

3. Notice adds a credential lookup on every initialize, even without the option (internal/api/handlers.go:1789, internal/invocation/service.go:2095)

The handler calls Notice on every initialize, and Notice resolves the server whether or not a provider is set. Without the option, that goes through the RefreshingOAuthCredentialStore overlay: an extra DB read per initialize, plus an OAuth refresh against the provider when the stored token is close to expiry. That breaks the "without it, nothing changes" promise in the changelog.

I'd suggest returning "" early from Notice when no provider is set. The resolver could expose that, or the service could keep a flag.

Note: I didn't run the tests locally, this is from reading the diff and the surrounding code at 05a4e50.

[kam-validmind]

Thanks, all three fixed in 3e07756, and main merged in (the session key now combines the tenant with the credential's key).

1. ReportStuck with no connected account. The resolver has a Lookup now: the server as configured, no credential, no provider call. The report uses it, so a caller the embedder refuses can still say they're stuck, and the embedder isn't asked for a token it never sends. Test: TestStuckReportAsksNoCredential.

2. stdio servers. The provider is asked only about upstreams with a BaseURL; a stdio server resolves as configured. Doc and changelog say so. Test: TestProviderSkipsStdioUpstreams.

3. Notice without the option. Returns "" when no provider is set, before any read. Test: TestNoticeWithoutProviderReadsNothing counts store reads with and without a provider.

go test ./... and gofmt pass locally.

[cachafla]

Thanks @kam-validmind, all three are addressed in 3e07756 👏

1. ReportStuck now uses the new Resolver.Lookup (internal/invocation/hooks.go:280), so it never asks the provider for a credential. Covered by TestStuckReportAsksNoCredential.
2. Resolver.credential returns stdio upstreams (no BaseURL) as configured before calling the provider, and the changelog says so. Covered by TestProviderSkipsStdioUpstreams.
3. Notice returns "" when HasProvider() is false, before any read. Covered by TestNoticeWithoutProviderReadsNothing.

CI is green (Go, Baseline). Ready for Andres's approval 🚀

[cachafla]

Head is 3e07756 and CI is green: Go, Baseline (fake-agent), and CLA pass. Live harnesses skipped since that job is opt-in (workflow_dispatch or RUN_LIVE_HARNESS_INTEGRATION_TESTS).