@vitest-agent/mcp 3.0.4 → 4.0.1

package/register-toolkit.js

@@ -0,0 +1,311 @@
+import { RenderText } from "./annotations.js";
+import { buildUnexpectedToolErrorEnvelope } from "./utils/tool-error-envelope.js";
+import { Cause, Context, Effect, Layer, Option, Result, Schema, Sink, Stream } from "effect";
+import { AiError, McpSchema, McpServer, Tool } from "effect/unstable/ai";
+
+//#region src/register-toolkit.ts
+const isPlainObject = (value) => value !== null && typeof value === "object" && !Array.isArray(value);
+/** Top-level discriminants a union of parameter shapes may dispatch on. */
+const DISCRIMINANT_KEYS = ["action", "kind"];
+/** Length cap for every echoed unknown key, so a hostile payload cannot balloon the error. */
+const ECHO_LIMIT = 200;
+const truncate = (value) => value.length > ECHO_LIMIT ? `${value.slice(0, ECHO_LIMIT)}…` : value;
+const literalValues = (node) => {
+	if (!isPlainObject(node)) return void 0;
+	if (Array.isArray(node.enum)) return node.enum;
+	if ("const" in node) return [node.const];
+};
+/**
+* The discriminant every member of a union shares as a literal-valued
+* property, or `undefined` when the members do not form such a union.
+*/
+const findDiscriminant = (members) => {
+	if (members.length === 0) return void 0;
+	for (const key of DISCRIMINANT_KEYS) if (members.every((member) => {
+		if (!isPlainObject(member) || !isPlainObject(member.properties)) return false;
+		return literalValues(member.properties[key]) !== void 0;
+	})) return key;
+};
+const REF_PREFIX = "#/$defs/";
+const lookupRef = (ref, root) => {
+	if (typeof ref !== "string" || !ref.startsWith(REF_PREFIX)) return void 0;
+	const defs = root.$defs;
+	if (!isPlainObject(defs)) return void 0;
+	const target = defs[ref.slice(8)];
+	return isPlainObject(target) ? target : void 0;
+};
+/**
+* Replace a `$ref` node with its `$defs` target merged under the node's
+* own keywords (the node's keys win). Follows chained refs; gives up on a
+* cycle or a dangling ref and returns the node as-is.
+*/
+const inlineRef = (node, root, seen = /* @__PURE__ */ new Set()) => {
+	const target = lookupRef(node.$ref, root);
+	if (target === void 0 || seen.has(node.$ref)) return node;
+	const { $ref, ...rest } = node;
+	return inlineRef({
+		...target,
+		...rest
+	}, root, /* @__PURE__ */ new Set([...seen, $ref]));
+};
+/**
+* A schema whose root (or whose top-level union members) is a `$ref` —
+* what Effect emits for any schema carrying an `identifier` annotation —
+* is inlined so the object / discriminant checks below see the real
+* shape. `$defs` is kept for any nested references.
+*
+* @internal
+*/
+const inlineRootRefs = (schema) => {
+	const root = inlineRef(schema, schema);
+	const combinator = Array.isArray(root.anyOf) ? "anyOf" : Array.isArray(root.oneOf) ? "oneOf" : void 0;
+	if (combinator === void 0) return root;
+	const members = root[combinator].map((member) => isPlainObject(member) ? inlineRef(member, schema) : member);
+	return {
+		...root,
+		[combinator]: members
+	};
+};
+/**
+* Deep-copy `schema`, setting `additionalProperties: false` on every
+* object node that declares `properties` (or is a bare object with no
+* combinator), and rewriting a top-level union of discriminated object
+* shapes to `{ type: "object", oneOf, "x-discriminator" }` so the served
+* schema still satisfies MCP's object-root requirement.
+*
+* @internal
+*/
+const strictifyJsonSchema = (schema) => {
+	const visit = (node) => {
+		if (Array.isArray(node)) return node.map(visit);
+		if (!isPlainObject(node)) return node;
+		const out = {};
+		for (const [key, value] of Object.entries(node)) switch (key) {
+			case "properties":
+			case "$defs":
+			case "definitions":
+				out[key] = isPlainObject(value) ? Object.fromEntries(Object.entries(value).map(([k, v]) => [k, visit(v)])) : value;
+				break;
+			case "items":
+			case "prefixItems":
+			case "anyOf":
+			case "oneOf":
+			case "allOf":
+			case "not":
+				out[key] = visit(value);
+				break;
+			case "additionalProperties":
+				out[key] = isPlainObject(value) ? visit(value) : value;
+				break;
+			default: out[key] = value;
+		}
+		const hasCombinator = Array.isArray(out.anyOf) || Array.isArray(out.oneOf) || Array.isArray(out.allOf);
+		if ((out.type === "object" || isPlainObject(out.properties)) && !hasCombinator && !isPlainObject(out.additionalProperties)) out.additionalProperties = false;
+		return out;
+	};
+	const strict = visit(inlineRootRefs(schema));
+	const members = strict.anyOf ?? strict.oneOf;
+	if (Array.isArray(members) && strict.type === void 0) {
+		const discriminant = findDiscriminant(members);
+		if (discriminant !== void 0) {
+			const { anyOf: _anyOf, oneOf: _oneOf, ...rest } = strict;
+			return {
+				type: "object",
+				...rest,
+				oneOf: members,
+				"x-discriminator": discriminant
+			};
+		}
+	}
+	return strict;
+};
+const resolveRef = (node, root) => lookupRef(node.$ref, root) ?? node;
+/** Pick the union member a payload object matches, or `undefined` when none does unambiguously. */
+const selectMember = (members, value, root) => {
+	const resolved = members.map((m) => isPlainObject(m) ? resolveRef(m, root) : m);
+	const discriminant = findDiscriminant(resolved);
+	if (discriminant !== void 0) {
+		const actual = value[discriminant];
+		return resolved.find((member) => {
+			if (!isPlainObject(member) || !isPlainObject(member.properties)) return false;
+			return literalValues(member.properties[discriminant])?.includes(actual) ?? false;
+		});
+	}
+	const objectMembers = resolved.filter((m) => isPlainObject(m) && isPlainObject(m.properties));
+	return objectMembers.length === 1 ? objectMembers[0] : void 0;
+};
+/**
+* Walk `value` against `schema`, collecting — at every object level — the
+* keys the payload carries that the schema does not declare.
+*
+* @internal
+*/
+const collectUnknownKeys = (value, schema) => {
+	const out = [];
+	const walk = (current, rawNode, path) => {
+		if (!isPlainObject(rawNode)) return;
+		const node = resolveRef(rawNode, schema);
+		if (Array.isArray(node.allOf)) {
+			const merged = {};
+			for (const member of node.allOf) {
+				const resolvedMember = isPlainObject(member) ? resolveRef(member, schema) : member;
+				if (isPlainObject(resolvedMember) && isPlainObject(resolvedMember.properties)) Object.assign(merged, resolvedMember.properties);
+			}
+			if (Object.keys(merged).length > 0) {
+				walk(current, {
+					...node,
+					allOf: void 0,
+					properties: {
+						...merged,
+						...isPlainObject(node.properties) ? node.properties : {}
+					}
+				}, path);
+				return;
+			}
+		}
+		if (isPlainObject(node.properties)) {
+			if (!isPlainObject(current)) return;
+			const accepted = Object.keys(node.properties);
+			const unknown = Object.keys(current).filter((key) => !accepted.includes(key));
+			if (unknown.length > 0 && node.additionalProperties !== true && !isPlainObject(node.additionalProperties)) out.push({
+				path,
+				unknown,
+				accepted
+			});
+			for (const key of accepted) if (key in current) walk(current[key], node.properties[key], [...path, key]);
+			if (isPlainObject(node.additionalProperties)) for (const key of unknown) walk(current[key], node.additionalProperties, [...path, key]);
+			return;
+		}
+		const members = node.oneOf ?? node.anyOf;
+		if (Array.isArray(members)) {
+			if (!isPlainObject(current)) return;
+			const member = selectMember(members, current, schema);
+			if (member !== void 0) walk(current, member, path);
+			return;
+		}
+		if (Array.isArray(current)) {
+			if (Array.isArray(node.prefixItems)) for (const [index, item] of node.prefixItems.entries()) walk(current[index], item, [...path, String(index)]);
+			if (isPlainObject(node.items)) for (const [index, item] of current.entries()) walk(item, node.items, [...path, String(index)]);
+			return;
+		}
+		if (isPlainObject(current) && node.type === "object" && node.additionalProperties === false) {
+			const unknown = Object.keys(current);
+			if (unknown.length > 0) out.push({
+				path,
+				unknown,
+				accepted: []
+			});
+		}
+	};
+	walk(value, schema, []);
+	return out;
+};
+/**
+* The exact wording the pre-Effect `strict()` helper produced, so agents
+* that learned to self-correct on it keep working.
+*/
+const formatUnknownKeys = (levels) => levels.map((level) => {
+	const qualify = (key) => truncate([...level.path, key].join("."));
+	const accepted = level.accepted.length === 0 ? "(none)" : level.accepted.join(", ");
+	return `Unrecognized parameter(s): ${level.unknown.map(qualify).join(", ")}. Accepted params: ${accepted}`;
+}).join(" ");
+/**
+* MCP models `structuredContent` as a JSON object, so a `null`, array or
+* primitive encoded result is omitted rather than sent through as-is.
+*/
+const toStructuredContent = (value) => isPlainObject(value) ? value : void 0;
+const declaredFailureResult = (message) => new McpSchema.CallToolResult({
+	isError: true,
+	content: [{
+		type: "text",
+		text: message
+	}]
+});
+const envelopeResult = (toolName, err) => {
+	const envelope = buildUnexpectedToolErrorEnvelope(toolName, err);
+	return new McpSchema.CallToolResult({
+		isError: true,
+		structuredContent: envelope,
+		content: [{
+			type: "text",
+			text: JSON.stringify(envelope)
+		}]
+	});
+};
+/**
+* Register every tool of `toolkit` with the ambient `McpServer` under the
+* strict contract described in the module docs.
+*
+* @internal
+*/
+const registerStrictToolkitEffect = Effect.fnUntraced(function* (toolkit) {
+	const registry = yield* McpServer.McpServer;
+	const built = yield* toolkit;
+	const services = yield* Effect.context();
+	for (const tool of Object.values(built.tools)) {
+		const annotations = tool.annotations;
+		const render = Context.get(annotations, RenderText);
+		const isDeclaredFailure = Schema.is(tool.failureSchema);
+		const toolMeta = Context.getOrUndefined(annotations, Tool.Meta);
+		const outputJsonSchema = inlineRootRefs(Tool.getJsonSchemaFromSchema(tool.successSchema));
+		const outputSchema = outputJsonSchema.type === "object" ? yield* Schema.decodeUnknownEffect(McpSchema.ToolJsonSchema)(outputJsonSchema).pipe(Effect.orDie) : void 0;
+		const servedInput = strictifyJsonSchema(Tool.getJsonSchema(tool));
+		const inputSchema = yield* Schema.decodeUnknownEffect(McpSchema.ToolJsonSchema)(servedInput).pipe(Effect.orDie);
+		const mcpTool = new McpSchema.Tool({
+			name: tool.name,
+			description: Tool.getDescription(tool),
+			inputSchema,
+			...outputSchema === void 0 ? {} : { outputSchema },
+			annotations: {
+				...Context.getOption(annotations, Tool.Title).pipe(Option.map((title) => ({ title })), Option.getOrUndefined),
+				readOnlyHint: Context.get(annotations, Tool.Readonly),
+				destructiveHint: Context.get(annotations, Tool.Destructive),
+				idempotentHint: Context.get(annotations, Tool.Idempotent),
+				openWorldHint: Context.get(annotations, Tool.OpenWorld)
+			},
+			_meta: toolMeta
+		});
+		yield* registry.addTool({
+			tool: mcpTool,
+			annotations,
+			handle(payload) {
+				const raw = payload ?? {};
+				const unknownKeys = collectUnknownKeys(raw, servedInput);
+				if (unknownKeys.length > 0) return Effect.fail(new McpSchema.InvalidParams({ message: formatUnknownKeys(unknownKeys) }));
+				return built.handle(tool.name, raw).pipe(Stream.unwrap, Stream.run(Sink.last()), Effect.flatMap(Effect.fromOption), Effect.map((result) => {
+					const encoded = result.encodedResult;
+					const text = render?.(encoded) ?? JSON.stringify(encoded);
+					return new McpSchema.CallToolResult({
+						isError: false,
+						structuredContent: toStructuredContent(encoded),
+						content: encoded === void 0 ? [] : [{
+							type: "text",
+							text
+						}]
+					});
+				}), Effect.provideContext(services), Effect.catchCause((cause) => {
+					const failure = Cause.findError(cause);
+					if (Result.isFailure(failure) && !Cause.hasDies(cause)) return Effect.failCause(failure.failure);
+					const err = Result.isSuccess(failure) ? failure.success : Cause.squash(cause);
+					if (AiError.isAiError(err) && err.reason._tag === "ToolParameterValidationError") return Effect.fail(new McpSchema.InvalidParams({ message: err.reason.message }));
+					const logged = Effect.logError(`tool ${tool.name} failed`, cause);
+					if (isDeclaredFailure(err) && err instanceof Error) return Effect.as(logged, declaredFailureResult(err.message));
+					return Effect.as(logged, envelopeResult(tool.name, err));
+				}));
+			}
+		});
+	}
+});
+/**
+* A layer that registers `toolkit` with the strict contract. Same shape
+* as `McpServer.toolkit`: it requires the toolkit's handlers (from
+* `kit.toLayer(handlers)`) and every service the tools declare in
+* `dependencies`, and shares the memoized `McpServer` with the transport
+* layer it is provided into.
+*
+* @public
+*/
+const registerStrictToolkit = (toolkit) => Layer.effectDiscard(registerStrictToolkitEffect(toolkit)).pipe(Layer.provide(McpServer.McpServer.layer));
+
+//#endregion
+export { collectUnknownKeys, inlineRootRefs, registerStrictToolkit, registerStrictToolkitEffect, strictifyJsonSchema };