Fix MCP recovery for invalid agent tokens#137
Merged
Conversation
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThis PR implements structured error handling for expired or invalid pre-registered agent tokens. It adds the ChangesAgent token invalidation and recovery
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![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 (2)
packages/mcp/src/server.ts (1)
477-480: ⚡ Quick winPreserve original tool error details when adding recovery guidance.
This currently replaces
result.contententirely, which drops upstream error context. Append the recovery text instead of overwriting.♻️ Proposed patch
- return { - ...result, - content: [{ type: 'text' as const, text: agentTokenRecoveryMessage() }], - }; + return { + ...result, + content: [ + ...(Array.isArray(result.content) ? result.content : []), + { type: 'text' as const, text: agentTokenRecoveryMessage() }, + ], + };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/mcp/src/server.ts` around lines 477 - 480, The current return replaces result.content with only the recovery message, losing upstream error details; instead append the recovery guidance to the existing result.content array (preserving its items and types) by concatenating or pushing the object produced by agentTokenRecoveryMessage() as { type: 'text', text: agentTokenRecoveryMessage() } so the returned object spreads ...result and returns content: [...result.content, { type: 'text' as const, text: agentTokenRecoveryMessage() }].packages/mcp/src/__tests__/integration.test.ts (1)
196-231: ⚡ Quick winAdd a routed-workspace invalid-token regression case.
Please add a variant where the failed call is routed via
workspace_id/workspace_aliasand assert the active workspace token remains intact. That will lock in the intended multi-workspace recovery behavior.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/mcp/src/__tests__/integration.test.ts` around lines 196 - 231, Add a new test variant alongside "clears an invalid pre-registered token so strict registration can recover" that repeats the failure/recovery flow but routes the failed call via a workspace identifier (use callTool arguments like workspace_id or workspace_alias when calling message.post) and then performs the same agent.register recovery; use the same setup symbols (invalidAgentTokens, createRelayMcpServer, staleClient, InMemoryTransport.createLinkedPair, captured) and assert the register request still uses the primary apiKey (expect registerReq!.headers.authorization toBe 'Bearer rk_live_int1') and additionally assert the active workspace token for that workspace remains unchanged after recovery (check whatever in-memory workspace token store or token map your test suite exposes to confirm the workspace-specific token is intact).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@packages/mcp/src/server.ts`:
- Around line 292-326: The invalidateActiveAgentToken function currently reads
and mutates session.workspaceKey and the workspace entry for the active
workspace, so calls that are meant to invalidate a token for a different
workspace can accidentally clear the active workspace; update
invalidateActiveAgentToken to accept an explicit workspace identifier (e.g.,
workspaceKey or workspaceAlias) in addition to asAgent, use that identifier to
look up and update the specific session.workspaces entry (and avoid touching
session.workspaceKey) and only stop session.wsBridge / clear
session.subscriptions / setSession global state when the invalidated workspace
matches the current session.workspaceKey; update callers that pass only asAgent
(places that route by workspace_id/workspace_alias) to supply the correct
workspace identifier so the correct workspace entry is updated without mutating
the active workspace state.
---
Nitpick comments:
In `@packages/mcp/src/__tests__/integration.test.ts`:
- Around line 196-231: Add a new test variant alongside "clears an invalid
pre-registered token so strict registration can recover" that repeats the
failure/recovery flow but routes the failed call via a workspace identifier (use
callTool arguments like workspace_id or workspace_alias when calling
message.post) and then performs the same agent.register recovery; use the same
setup symbols (invalidAgentTokens, createRelayMcpServer, staleClient,
InMemoryTransport.createLinkedPair, captured) and assert the register request
still uses the primary apiKey (expect registerReq!.headers.authorization toBe
'Bearer rk_live_int1') and additionally assert the active workspace token for
that workspace remains unchanged after recovery (check whatever in-memory
workspace token store or token map your test suite exposes to confirm the
workspace-specific token is intact).
In `@packages/mcp/src/server.ts`:
- Around line 477-480: The current return replaces result.content with only the
recovery message, losing upstream error details; instead append the recovery
guidance to the existing result.content array (preserving its items and types)
by concatenating or pushing the object produced by agentTokenRecoveryMessage()
as { type: 'text', text: agentTokenRecoveryMessage() } so the returned object
spreads ...result and returns content: [...result.content, { type: 'text' as
const, text: agentTokenRecoveryMessage() }].
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 96559800-c2c2-4a5c-8cbe-cc21a411eebd
📒 Files selected for processing (10)
README.mdopenapi.yamlpackages/mcp/src/__tests__/integration.test.tspackages/mcp/src/server.tspackages/sdk-rust/src/types.rspackages/sdk-rust/tests/parity.rspackages/sdk-typescript/src/__tests__/errors.test.tspackages/sdk-typescript/src/errors.tspackages/server/src/middleware/__tests__/auth.test.tspackages/server/src/middleware/auth.ts
…-agent-token # Conflicts: # packages/sdk-rust/tests/parity.rs
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
1 issue found across 10 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="packages/mcp/src/server.ts">
<violation number="1" location="packages/mcp/src/server.ts:476">
P2: Use the routed workspace identity when invalidating tokens; this path can clear the active session instead of the workspace that actually failed auth.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| && typeof args.as === 'string' | ||
| ? args.as | ||
| : undefined; | ||
| invalidateActiveAgentToken(asAgent); |
Contributor
There was a problem hiding this comment.
P2: Use the routed workspace identity when invalidating tokens; this path can clear the active session instead of the workspace that actually failed auth.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/mcp/src/server.ts, line 476:
<comment>Use the routed workspace identity when invalidating tokens; this path can clear the active session instead of the workspace that actually failed auth.</comment>
<file context>
@@ -368,7 +463,37 @@ export function createRelayMcpServer(options: McpServerOptions): McpServer {
+ && typeof args.as === 'string'
+ ? args.as
+ : undefined;
+ invalidateActiveAgentToken(asAgent);
+ return {
+ ...result,
</file context>
|
Preview deployed!
This preview shares the staging database and will be cleaned up when the PR is merged or closed. Run E2E testsnpm run e2e -- https://pr137-api.relaycast.dev --ciOpen observer dashboard |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
agent_token_invalidAPI error for invalid agent tokensagent.registercan rotate/re-register under strict identityCloses #136.
Tests
npm --workspace @relaycast/sdk testnpm --workspace @relaycast/server test -- --run src/middleware/__tests__/auth.test.tsnpm --workspace @relaycast/mcp testnpm --workspace @relaycast/sdk run buildnpm --workspace @relaycast/server run buildnpm --workspace @relaycast/mcp run buildnpm --workspace @relaycast/sdk run lintnpm --workspace @relaycast/server run lint(existing warnings only)npm --workspace @relaycast/mcp run lint