Skip to content

Commit b5d845c

Browse files
zmofeizmofei
committed
Add authentication token support for Mapbox API tools
- Add support for passing authentication tokens via MCP request context - Modified BaseTool to accept and pass through authentication tokens - Update all Mapbox API tools to use provided token or fall back to env variable - Fix failing tests related to token validation - Simplify execute method signature by removing unused extra parameter
1 parent ef10091 commit b5d845c

13 files changed

Lines changed: 115 additions & 46 deletions

File tree

‎plop-templates/mapbox-api-tool.hbs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,14 +36,15 @@ export class {{pascalCase name}}Tool extends MapboxApiBasedTool<
3636
}
3737

3838
protected async execute(
39-
input: {{pascalCase name}}Input
39+
input: {{pascalCase name}}Input,
40+
accessToken?: string
4041
): Promise<{ type: 'text'; text: string }> {
4142
try {
4243
// TODO: Implement your Mapbox API call here
4344

4445
// Example implementation:
4546
// const username = MapboxApiBasedTool.getUserNameFromToken();
46-
// const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}/${input.styleId}?access_token=${MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN}`;
47+
// const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}/${input.styleId}?access_token=${accessToken}`;
4748
//
4849
// const response = await fetch(url);
4950
//

‎src/tools/BaseTool.ts‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import {
22
McpServer,
33
RegisteredTool
44
} from '@modelcontextprotocol/sdk/server/mcp';
5+
import { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';
56
import { z, ZodTypeAny } from 'zod';
67

78
const ContentItemSchema = z.union([
@@ -38,10 +39,15 @@ export abstract class BaseTool<InputSchema extends ZodTypeAny> {
3839
/**
3940
* Validates and runs the tool logic.
4041
*/
41-
async run(rawInput: unknown): Promise<z.infer<typeof OutputSchema>> {
42+
async run(
43+
rawInput: unknown,
44+
extra?: RequestHandlerExtra<any, any>
45+
): Promise<z.infer<typeof OutputSchema>> {
4246
try {
4347
const input = this.inputSchema.parse(rawInput);
44-
const result = await this.execute(input);
48+
const accessToken =
49+
extra?.authInfo?.token || process.env.MAPBOX_ACCESS_TOKEN;
50+
const result = await this.execute(input, accessToken);
4551

4652
// Check if result is already a content object (image or text)
4753
if (
@@ -86,7 +92,8 @@ export abstract class BaseTool<InputSchema extends ZodTypeAny> {
8692
* Tool logic to be implemented by subclasses.
8793
*/
8894
protected abstract execute(
89-
_input: z.infer<InputSchema>
95+
_input: z.infer<InputSchema>,
96+
accessToken?: string
9097
): Promise<ContentItem | unknown>;
9198

9299
/**
@@ -99,7 +106,7 @@ export abstract class BaseTool<InputSchema extends ZodTypeAny> {
99106
this.description,
100107
(this.inputSchema as unknown as z.ZodObject<Record<string, z.ZodTypeAny>>)
101108
.shape,
102-
this.run.bind(this)
109+
(args, extra) => this.run(args, extra)
103110
);
104111
}
105112

‎src/tools/MapboxApiBasedTool.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ describe('MapboxApiBasedTool', () => {
7979
});
8080

8181
expect(() => MapboxApiBasedTool.getUserNameFromToken()).toThrow(
82-
'MAPBOX_ACCESS_TOKEN is not set'
82+
'No access token provided. Please set MAPBOX_ACCESS_TOKEN environment variable or pass it as an argument.'
8383
);
8484
} finally {
8585
Object.defineProperty(MapboxApiBasedTool, 'MAPBOX_ACCESS_TOKEN', {

‎src/tools/MapboxApiBasedTool.ts‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { RequestHandlerExtra } from '@modelcontextprotocol/sdk/shared/protocol.js';
12
import { z, ZodTypeAny } from 'zod';
23
import { BaseTool, OutputSchema } from './BaseTool.js';
34

@@ -17,14 +18,19 @@ export abstract class MapboxApiBasedTool<
1718
* Mapbox tokens are JWT tokens where the payload contains the username.
1819
* @throws Error if the token is not set, invalid, or doesn't contain username
1920
*/
20-
static getUserNameFromToken(): string {
21-
if (!MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN) {
22-
throw new Error('MAPBOX_ACCESS_TOKEN is not set');
21+
static getUserNameFromToken(access_token?: string): string {
22+
if (!access_token) {
23+
if (!MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN) {
24+
throw new Error(
25+
'No access token provided. Please set MAPBOX_ACCESS_TOKEN environment variable or pass it as an argument.'
26+
);
27+
}
28+
access_token = MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN;
2329
}
2430

2531
try {
2632
// JWT format: header.payload.signature
27-
const parts = MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN.split('.');
33+
const parts = access_token.split('.');
2834
if (parts.length !== 3) {
2935
throw new Error('MAPBOX_ACCESS_TOKEN is not in valid JWT format');
3036
}
@@ -68,19 +74,30 @@ export abstract class MapboxApiBasedTool<
6874
/**
6975
* Validates Mapbox token and runs the tool logic.
7076
*/
71-
async run(rawInput: unknown): Promise<z.infer<typeof OutputSchema>> {
77+
async run(
78+
rawInput: unknown,
79+
extra?: RequestHandlerExtra<any, any>
80+
): Promise<z.infer<typeof OutputSchema>> {
7281
try {
73-
if (!MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN) {
74-
throw new Error('MAPBOX_ACCESS_TOKEN is not set');
82+
// First check if token is provided via authentication context
83+
// Check both standard token field and accessToken in extra for compatibility
84+
// In the streamableHttp, the authInfo is injected into extra from `req.auth`
85+
// https://github.com/modelcontextprotocol/typescript-sdk/blob/main/src/server/streamableHttp.ts#L405
86+
const authToken = extra?.authInfo?.token;
87+
const accessToken = authToken || MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN;
88+
if (!accessToken) {
89+
throw new Error(
90+
'No access token available. Please provide via Bearer auth or MAPBOX_ACCESS_TOKEN env var'
91+
);
7592
}
7693

7794
// Validate that the token has the correct JWT format
78-
if (!this.isValidJwtFormat(MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN)) {
95+
if (!this.isValidJwtFormat(accessToken)) {
7996
throw new Error('MAPBOX_ACCESS_TOKEN is not in valid JWT format');
8097
}
8198

8299
// Call parent run method which handles the rest
83-
return await super.run(rawInput);
100+
return await super.run(rawInput, extra);
84101
} catch (error) {
85102
const errorMessage =
86103
error instanceof Error ? error.message : String(error);

‎src/tools/create-style-tool/CreateStyleTool.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,12 @@ export class CreateStyleTool extends MapboxApiBasedTool<
1414
super({ inputSchema: CreateStyleSchema });
1515
}
1616

17-
protected async execute(input: CreateStyleInput): Promise<any> {
18-
const username = MapboxApiBasedTool.getUserNameFromToken();
19-
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}?access_token=${MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN}`;
17+
protected async execute(
18+
input: CreateStyleInput,
19+
accessToken?: string
20+
): Promise<any> {
21+
const username = MapboxApiBasedTool.getUserNameFromToken(accessToken);
22+
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}?access_token=${accessToken}`;
2023

2124
const payload = {
2225
name: input.name,

‎src/tools/create-token-tool/CreateTokenTool.test.ts‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,17 +79,31 @@ describe('CreateTokenTool', () => {
7979

8080
it('throws error when unable to extract username from token', async () => {
8181
const originalToken = MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN;
82+
const originalEnvToken = process.env.MAPBOX_ACCESS_TOKEN;
8283

8384
try {
8485
// Set a token without username in payload
8586
const invalidPayload = Buffer.from(
8687
JSON.stringify({ sub: 'test' })
8788
).toString('base64');
89+
const invalidToken = `eyJhbGciOiJIUzI1NiJ9.${invalidPayload}.signature`;
90+
8891
Object.defineProperty(MapboxApiBasedTool, 'MAPBOX_ACCESS_TOKEN', {
89-
value: `eyJhbGciOiJIUzI1NiJ9.${invalidPayload}.signature`,
92+
value: invalidToken,
9093
writable: true,
9194
configurable: true
9295
});
96+
process.env.MAPBOX_ACCESS_TOKEN = invalidToken;
97+
98+
// Setup fetch mock to prevent actual API calls
99+
const fetchMock = setupFetch();
100+
fetchMock.mockResolvedValueOnce({
101+
ok: true,
102+
status: 200,
103+
statusText: 'OK',
104+
headers: new Headers(),
105+
json: async () => ({ token: 'test-token' })
106+
} as Response);
93107

94108
const toolWithInvalidToken = new CreateTokenTool();
95109
toolWithInvalidToken['log'] = jest.fn();
@@ -112,6 +126,7 @@ describe('CreateTokenTool', () => {
112126
writable: true,
113127
configurable: true
114128
});
129+
process.env.MAPBOX_ACCESS_TOKEN = originalEnvToken;
115130
}
116131
});
117132
});

‎src/tools/create-token-tool/CreateTokenTool.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,14 @@ export class CreateTokenTool extends MapboxApiBasedTool<
1717
}
1818

1919
protected async execute(
20-
input: CreateTokenInput
20+
input: CreateTokenInput,
21+
accessToken?: string
2122
): Promise<{ type: 'text'; text: string }> {
22-
if (!MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN) {
23+
if (!accessToken) {
2324
throw new Error('MAPBOX_ACCESS_TOKEN is not set');
2425
}
2526

26-
const username = MapboxApiBasedTool.getUserNameFromToken();
27+
const username = MapboxApiBasedTool.getUserNameFromToken(accessToken);
2728

2829
this.log(
2930
'info',
@@ -47,7 +48,7 @@ export class CreateTokenTool extends MapboxApiBasedTool<
4748
);
4849
}
4950

50-
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}tokens/v2/${username}?access_token=${MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN}`;
51+
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}tokens/v2/${username}?access_token=${accessToken}`;
5152

5253
const body: {
5354
note: string;

‎src/tools/delete-style-tool/DeleteStyleTool.ts‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,12 @@ export class DeleteStyleTool extends MapboxApiBasedTool<
1414
super({ inputSchema: DeleteStyleSchema });
1515
}
1616

17-
protected async execute(input: DeleteStyleInput): Promise<any> {
18-
const username = MapboxApiBasedTool.getUserNameFromToken();
19-
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}/${input.styleId}?access_token=${MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN}`;
17+
protected async execute(
18+
input: DeleteStyleInput,
19+
accessToken?: string
20+
): Promise<any> {
21+
const username = MapboxApiBasedTool.getUserNameFromToken(accessToken);
22+
const url = `${MapboxApiBasedTool.MAPBOX_API_ENDPOINT}styles/v1/${username}/${input.styleId}?access_token=${accessToken}`;
2023

2124
const response = await fetch(url, {
2225
method: 'DELETE'

‎src/tools/list-styles-tool/ListStylesTool.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,18 @@ export class ListStylesTool extends MapboxApiBasedTool<
1111
super({ inputSchema: ListStylesSchema });
1212
}
1313

14-
protected async execute(input: ListStylesInput): Promise<any> {
15-
const username = MapboxApiBasedTool.getUserNameFromToken();
14+
protected async execute(
15+
input: ListStylesInput,
16+
accessToken?: string
17+
): Promise<any> {
18+
const username = MapboxApiBasedTool.getUserNameFromToken(accessToken);
1619

1720
// Build query parameters
1821
const params = new URLSearchParams();
19-
if (!MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN) {
22+
if (!accessToken) {
2023
throw new Error('MAPBOX_ACCESS_TOKEN is not set');
2124
}
22-
params.append('access_token', MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN);
25+
params.append('access_token', accessToken);
2326

2427
if (input.limit) {
2528
params.append('limit', input.limit.toString());

‎src/tools/list-tokens-tool/ListTokensTool.test.ts‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,17 +70,31 @@ describe('ListTokensTool', () => {
7070

7171
it('throws error when unable to extract username from token', async () => {
7272
const originalToken = MapboxApiBasedTool.MAPBOX_ACCESS_TOKEN;
73+
const originalEnvToken = process.env.MAPBOX_ACCESS_TOKEN;
7374

7475
try {
7576
// Set a token without username in payload
7677
const invalidPayload = Buffer.from(
7778
JSON.stringify({ sub: 'test' })
7879
).toString('base64');
80+
const invalidToken = `eyJhbGciOiJIUzI1NiJ9.${invalidPayload}.signature`;
81+
7982
Object.defineProperty(MapboxApiBasedTool, 'MAPBOX_ACCESS_TOKEN', {
80-
value: `eyJhbGciOiJIUzI1NiJ9.${invalidPayload}.signature`,
83+
value: invalidToken,
8184
writable: true,
8285
configurable: true
8386
});
87+
process.env.MAPBOX_ACCESS_TOKEN = invalidToken;
88+
89+
// Setup fetch mock to prevent actual API calls
90+
const fetchMock = setupFetch();
91+
fetchMock.mockResolvedValueOnce({
92+
ok: true,
93+
status: 200,
94+
statusText: 'OK',
95+
headers: new Headers(),
96+
json: async () => []
97+
} as Response);
8498

8599
const toolWithInvalidToken = new ListTokensTool();
86100
toolWithInvalidToken['log'] = jest.fn();
@@ -100,6 +114,7 @@ describe('ListTokensTool', () => {
100114
writable: true,
101115
configurable: true
102116
});
117+
process.env.MAPBOX_ACCESS_TOKEN = originalEnvToken;
103118
}
104119
});
105120
});

0 commit comments

Comments
 (0)