Skip to content

Commit 04f026b

Browse files
chelojimenezchelojimenez
authored
Fix whoami UI gating by client capabilities (#1921)
1 parent a65ea7e commit 04f026b

2 files changed

Lines changed: 219 additions & 47 deletions

File tree

‎mcp/src/server.ts‎

Lines changed: 80 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,23 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
33
import { RESOURCE_MIME_TYPE } from "@modelcontextprotocol/ext-apps/server";
44
import type { ClientCapabilities } from "@modelcontextprotocol/sdk/types.js";
55
import type { JWTPayload } from "jose";
6-
import { createSessionToolRegistrar } from "./tools/sessionToolRegistrar.js";
6+
import {
7+
createSessionToolRegistrar,
8+
type SessionToolRegistrar,
9+
} from "./tools/sessionToolRegistrar.js";
710
import { registerDoctorTool } from "./tools/doctor.js";
811
import { registerGetOrgTool } from "./tools/getOrg.js";
912
import { registerGetWorkspacesTool } from "./tools/getWorkspaces.js";
1013
import { registerWhoamiTool } from "./tools/whoami.js";
1114

12-
const UI_EXTENSION_ID = "io.modelcontextprotocol/ui";
13-
1415
interface McpProps extends Record<string, unknown> {
1516
bearerToken: string;
1617
claims: JWTPayload;
1718
}
1819

1920
export class McpJamMcpServer extends McpAgent<Env, unknown, McpProps> {
21+
private sessionToolRegistrar?: SessionToolRegistrar;
22+
2023
server = new McpServer({
2124
name: "MCPJam MCP",
2225
version: "0.1.0",
@@ -31,39 +34,65 @@ export class McpJamMcpServer extends McpAgent<Env, unknown, McpProps> {
3134
}
3235

3336
async init(): Promise<void> {
34-
const registrar = createSessionToolRegistrar(this.server);
37+
const initializeRequest = await this.getInitializeRequest();
38+
const initializeClientCapabilities = (initializeRequest as
39+
| { params?: { capabilities?: ClientCapabilities } }
40+
| undefined)?.params?.capabilities;
41+
const registrar = createSessionToolRegistrar(
42+
this.server,
43+
uiSupportsResourceMime(initializeClientCapabilities)
44+
);
45+
this.sessionToolRegistrar = registrar;
3546

3647
registerWhoamiTool(registrar, this);
3748
registerDoctorTool(registrar, this);
3849
registerGetWorkspacesTool(registrar, this);
3950
registerGetOrgTool(registrar, this);
51+
}
4052

41-
const initializeRequest = await this.getInitializeRequest();
42-
const initializeClientCapabilities = (initializeRequest as
43-
| { params?: { capabilities?: ClientCapabilities } }
44-
| undefined)?.params?.capabilities;
53+
override async onConnect(conn: any, context: { request: Request }): Promise<void> {
54+
this.applyUiModeFromRawRequest(context.request);
55+
await super.onConnect(conn, context as any);
56+
}
4557

46-
registrar.setUiEnabled(isUiEnabled(initializeClientCapabilities));
58+
private applyUiModeFromRawRequest(request: Request): void {
59+
if (
60+
!this.sessionToolRegistrar ||
61+
this.getTransportType() !== "streamable-http" ||
62+
request.headers.get("cf-mcp-method") !== "POST"
63+
) {
64+
return;
65+
}
4766

48-
this.server.server.oninitialized = () => {
49-
registrar.setUiEnabled(
50-
isUiEnabled(this.server.server.getClientCapabilities())
51-
);
52-
};
53-
}
54-
}
67+
const payloadHeader = request.headers.get("cf-mcp-message");
68+
if (!payloadHeader) {
69+
return;
70+
}
5571

56-
function getUiCapability(
57-
clientCapabilities: ClientCapabilities | undefined
58-
): { mimeTypes?: string[] } | undefined {
59-
const extensions = (clientCapabilities as
60-
| (ClientCapabilities & { extensions?: Record<string, unknown> })
61-
| undefined)?.extensions;
72+
try {
73+
const rawPayload = Buffer.from(payloadHeader, "base64").toString("utf-8");
74+
const parsedBody = JSON.parse(rawPayload) as unknown;
75+
const messages = Array.isArray(parsedBody) ? parsedBody : [parsedBody];
76+
77+
for (const message of messages) {
78+
const clientCapabilities = getInitializeCapabilities(message);
79+
if (!clientCapabilities) {
80+
continue;
81+
}
6282

63-
return extensions?.[UI_EXTENSION_ID] as { mimeTypes?: string[] } | undefined;
83+
this.sessionToolRegistrar.setUiEnabled(
84+
uiSupportsResourceMime(clientCapabilities),
85+
{ notify: false }
86+
);
87+
return;
88+
}
89+
} catch {
90+
// Ignore malformed headers and let the transport surface the real error.
91+
}
92+
}
6493
}
6594

66-
function isUiEnabled(
95+
function uiSupportsResourceMime(
6796
clientCapabilities: ClientCapabilities | undefined
6897
): boolean {
6998
return (
@@ -72,3 +101,30 @@ function isUiEnabled(
72101
) ?? false
73102
);
74103
}
104+
105+
function getUiCapability(
106+
clientCapabilities:
107+
| (ClientCapabilities & { extensions?: Record<string, unknown> })
108+
| undefined
109+
): { mimeTypes?: string[] } | undefined {
110+
return clientCapabilities?.extensions?.["io.modelcontextprotocol/ui"] as
111+
| { mimeTypes?: string[] }
112+
| undefined;
113+
}
114+
115+
function getInitializeCapabilities(
116+
message: unknown
117+
): ClientCapabilities | undefined {
118+
if (!message || typeof message !== "object") {
119+
return undefined;
120+
}
121+
122+
const initializeMessage = message as {
123+
method?: unknown;
124+
params?: { capabilities?: ClientCapabilities };
125+
};
126+
127+
return initializeMessage.method === "initialize"
128+
? initializeMessage.params?.capabilities
129+
: undefined;
130+
}

‎mcp/src/tools/sessionToolRegistrar.ts‎

Lines changed: 139 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
import {
22
RESOURCE_MIME_TYPE,
3-
registerAppResource,
3+
RESOURCE_URI_META_KEY,
44
registerAppTool,
55
} from "@modelcontextprotocol/ext-apps/server";
66
import type {
77
McpServer,
8+
RegisteredResource,
89
RegisteredTool,
910
ToolCallback,
1011
} from "@modelcontextprotocol/sdk/server/mcp.js";
@@ -47,10 +48,39 @@ export interface SessionToolRegistrar {
4748
callback: ToolCallback<InputArgs>,
4849
ui?: ToolUiConfig<InputArgs>
4950
): RegisteredTool;
50-
setUiEnabled(enabled: boolean): void;
51+
setUiEnabled(enabled: boolean, options?: SetUiEnabledOptions): void;
5152
}
5253

53-
export function createSessionToolRegistrar(server: McpServer): SessionToolRegistrar {
54+
type SetUiEnabledOptions = {
55+
notify?: boolean;
56+
};
57+
58+
type MutableRegisteredTool = RegisteredTool & {
59+
_meta?: Record<string, unknown>;
60+
enabled: boolean;
61+
handler: ToolCallback<any>;
62+
};
63+
64+
type MutableRegisteredResource = RegisteredResource & {
65+
enabled: boolean;
66+
};
67+
68+
type UiAwareRegistration = {
69+
plainCallback: ToolCallback<any>;
70+
plainMeta: Record<string, unknown> | undefined;
71+
resource: MutableRegisteredResource;
72+
tool: MutableRegisteredTool;
73+
uiCallback: ToolCallback<any>;
74+
uiMeta: Record<string, unknown>;
75+
};
76+
77+
export function createSessionToolRegistrar(
78+
server: McpServer,
79+
uiEnabled: boolean
80+
): SessionToolRegistrar {
81+
const uiAwareRegistrations: UiAwareRegistration[] = [];
82+
let currentUiEnabled = uiEnabled;
83+
5484
return {
5585
registerTool<
5686
OutputArgs extends ZodRawShapeCompat | AnySchema,
@@ -65,26 +95,26 @@ export function createSessionToolRegistrar(server: McpServer): SessionToolRegist
6595
return server.registerTool(name, config, callback);
6696
}
6797

68-
const tool = registerAppTool(
69-
server,
70-
name,
71-
{
72-
...config,
73-
_meta: {
74-
...(config._meta ?? {}),
75-
ui: {
76-
resourceUri: ui.resourceUri,
77-
},
78-
},
79-
} as any,
80-
(ui.callback ?? callback) as any
81-
);
98+
const plainMeta = stripToolUiMeta(config._meta);
99+
const uiMeta = createToolUiMeta(config._meta, ui.resourceUri);
100+
const uiCallback = (ui.callback ?? callback) as ToolCallback<any>;
101+
const tool = currentUiEnabled
102+
? registerAppTool(
103+
server,
104+
name,
105+
{
106+
...config,
107+
_meta: uiMeta,
108+
} as any,
109+
uiCallback as any
110+
)
111+
: server.registerTool(name, { ...config, _meta: plainMeta }, callback);
82112

83-
registerAppResource(
84-
server,
113+
const resource = server.registerResource(
85114
ui.resourceName ?? `${config.title ?? name} UI`,
86115
ui.resourceUri,
87116
{
117+
mimeType: RESOURCE_MIME_TYPE,
88118
description:
89119
ui.resourceDescription ?? `${config.title ?? name} interactive UI`,
90120
_meta: ui.resourceMeta,
@@ -101,12 +131,98 @@ export function createSessionToolRegistrar(server: McpServer): SessionToolRegist
101131
})
102132
);
103133

134+
if (!currentUiEnabled) {
135+
resource.disable();
136+
}
137+
138+
uiAwareRegistrations.push({
139+
plainCallback: callback as ToolCallback<any>,
140+
plainMeta,
141+
resource: resource as MutableRegisteredResource,
142+
tool: tool as MutableRegisteredTool,
143+
uiCallback,
144+
uiMeta,
145+
});
146+
104147
return tool;
105148
},
106-
setUiEnabled() {
107-
// The hosted example server always advertises its app tool and relies on
108-
// the host to ignore UI metadata when unsupported. Keep the method for
109-
// compatibility with existing call sites.
149+
setUiEnabled(enabled, options) {
150+
if (currentUiEnabled === enabled) {
151+
return;
152+
}
153+
154+
currentUiEnabled = enabled;
155+
const notify = options?.notify ?? true;
156+
157+
for (const registration of uiAwareRegistrations) {
158+
applyUiEnabledState(registration, enabled, notify);
159+
}
160+
},
161+
};
162+
}
163+
164+
function applyUiEnabledState(
165+
registration: UiAwareRegistration,
166+
enabled: boolean,
167+
notify: boolean
168+
): void {
169+
if (notify) {
170+
registration.tool.update({
171+
_meta: enabled ? registration.uiMeta : (registration.plainMeta ?? {}),
172+
callback: enabled
173+
? registration.uiCallback
174+
: registration.plainCallback,
175+
} as any);
176+
177+
if (enabled) {
178+
registration.resource.enable();
179+
} else {
180+
registration.resource.disable();
181+
}
182+
183+
return;
184+
}
185+
186+
registration.tool._meta = enabled
187+
? registration.uiMeta
188+
: (registration.plainMeta ?? {});
189+
registration.tool.handler = enabled
190+
? registration.uiCallback
191+
: registration.plainCallback;
192+
registration.resource.enabled = enabled;
193+
}
194+
195+
function createToolUiMeta(
196+
meta: Record<string, unknown> | undefined,
197+
resourceUri: string
198+
): Record<string, unknown> {
199+
const uiMeta =
200+
meta?.ui && typeof meta.ui === "object" && !Array.isArray(meta.ui)
201+
? (meta.ui as Record<string, unknown>)
202+
: undefined;
203+
204+
return {
205+
...(stripToolUiMeta(meta) ?? {}),
206+
ui: {
207+
...(uiMeta ?? {}),
208+
resourceUri,
110209
},
210+
[RESOURCE_URI_META_KEY]: resourceUri,
111211
};
112212
}
213+
214+
function stripToolUiMeta(
215+
meta: Record<string, unknown> | undefined
216+
): Record<string, unknown> | undefined {
217+
if (!meta) {
218+
return undefined;
219+
}
220+
221+
const {
222+
[RESOURCE_URI_META_KEY]: _resourceUri,
223+
ui: _ui,
224+
...plainMeta
225+
} = meta;
226+
227+
return Object.keys(plainMeta).length > 0 ? plainMeta : undefined;
228+
}

0 commit comments

Comments
 (0)