/*
* Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT.
*/
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import {
ShapeOutput,
ZodRawShapeCompat,
} from "@modelcontextprotocol/sdk/server/zod-compat.js";
import { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
import {
CallToolResult,
ServerNotification,
ServerRequest,
} from "@modelcontextprotocol/sdk/types.js";
import { FireHydrantCore } from "../core.js";
import { ConsoleLogger } from "./console-logger.js";
import { MCPScope } from "./scopes.js";
import { isAsyncIterable, isBinaryData, valueToBase64 } from "./shared.js";
export type ToolDefinition<
Args extends undefined | ZodRawShapeCompat = undefined,
> = Args extends ZodRawShapeCompat ? {
name: string;
description: string;
scopes?: MCPScope[];
args: Args;
annotations: {
title: string;
destructiveHint: boolean;
idempotentHint: boolean;
openWorldHint: boolean;
readOnlyHint: boolean;
};
tool: (
client: FireHydrantCore,
args: ShapeOutput<Args>,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => CallToolResult | Promise<CallToolResult>;
}
: {
name: string;
description: string;
scopes?: MCPScope[];
args?: undefined;
annotations: {
title: string;
destructiveHint: boolean;
idempotentHint: boolean;
openWorldHint: boolean;
readOnlyHint: boolean;
};
tool: (
client: FireHydrantCore,
extra: RequestHandlerExtra<ServerRequest, ServerNotification>,
) => CallToolResult | Promise<CallToolResult>;
};
// Optional function to assist with formatting tool results
export async function formatResult(
value: unknown,
init: { response?: Response | undefined },
): Promise<CallToolResult> {
if (typeof value === "undefined") {
return { content: [] };
}
const { response } = init;
const contentType = response?.headers.get("content-type") ?? "";
let content: CallToolResult["content"] = [];
if (contentType.search(/\bjson\b/g)) {
content = [{ type: "text", text: JSON.stringify(value) }];
} else if (
contentType.startsWith("text/event-stream")
&& isAsyncIterable(value)
) {
content = await consumeSSE(value);
} else if (contentType.startsWith("text/") && typeof value === "string") {
content = [{ type: "text", text: value }];
} else if (isBinaryData(value) && contentType.startsWith("image/")) {
const data = await valueToBase64(value);
content = data == null
? []
: [{ type: "image", data, mimeType: contentType }];
} else {
return {
content: [{
type: "text",
text: `Unsupported content type: "${contentType}"`,
}],
isError: true,
};
}
return { content };
}
async function consumeSSE(
value: AsyncIterable<unknown>,
): Promise<CallToolResult["content"]> {
const content: CallToolResult["content"] = [];
for await (const chunk of value) {
if (typeof chunk === "string") {
content.push({ type: "text", text: chunk });
} else {
content.push({ type: "text", text: JSON.stringify(chunk) });
}
}
return content;
}
export function createRegisterTool(
logger: ConsoleLogger,
server: McpServer,
getSDK: () => FireHydrantCore,
allowedScopes: Set<MCPScope>,
allowedTools?: Set<string>,
): [
<A extends ZodRawShapeCompat | undefined>(tool: ToolDefinition<A>) => void,
Array<{ name: string; description: string }>,
] {
const tools: Array<{ name: string; description: string }> = [];
const registerTool = <A extends ZodRawShapeCompat | undefined>(
tool: ToolDefinition<A>,
): void => {
if (allowedTools && !allowedTools.has(tool.name)) {
return;
}
const scopes = tool.scopes ?? [];
if (allowedScopes.size > 0 && scopes.length === 0) {
return;
}
if (
allowedScopes.size > 0
&& !scopes.every((s: MCPScope) => allowedScopes.has(s))
) {
return;
}
if (tool.args) {
server.tool(
tool.name,
tool.description,
tool.args,
tool.annotations,
async (args, ctx) => {
return tool.tool(getSDK(), args, ctx);
},
);
} else {
server.tool(
tool.name,
tool.description,
tool.annotations,
async (_, ctx) => {
return tool.tool(getSDK(), ctx);
},
);
}
logger.debug("Registered tool", { name: tool.name });
tools.push({ name: tool.name, description: tool.description });
};
return [registerTool, tools];
}
