feat: add outputSchema and infer execute output from it

README.md

@@ -67,6 +67,39 @@ you may need the following in `webpack.config.js`:
 "types": ["webmcp-types"]
 ```
 
+### Infer types from schemas
+
+When `inputSchema` and `outputSchema` are written inline, `registerTool()` infers the type of
+`execute`'s input and checks its return value against the output schema:
+
+```ts
+document.modelContext?.registerTool({
+  name: "search",
+  description: "Searches the catalog.",
+  inputSchema: {
+    type: "object",
+    properties: { query: { type: "string" }, limit: { type: "integer" } },
+    required: ["query"],
+  },
+  outputSchema: {
+    type: "object",
+    properties: { results: { type: "array", items: { type: "string" } } },
+    required: ["results"],
+  },
+  // input: { query: string; limit?: number }
+  // return type: { results: string[] } (or a Promise of it)
+  execute: async (input) => ({ results: [input.query] }),
+});
+```
+
+Tools without schemas, or whose schemas are typed as a plain `object`, fall back to a
+`Record<string, unknown>` input and an `unknown` return type.
+
+> [!NOTE]
+> If `execute` declares no parameters (e.g. `() => ({ status: "ok" })`), TypeScript checks it before
+> inferring the output schema, so returned `enum`/`const` values are widened (e.g. to `string`) and
+> rejected. Declaring the input parameter (e.g. `(input) => ...`) avoids this.
+
 ### Run the type tests
 
 - `npm install`

index.d.ts

@@ -39,6 +39,12 @@ type InferToolInput<TSchema> = [InferJsonSchema<TSchema>] extends [null | undefi
     : InferJsonSchema<TSchema> extends object ? InferJsonSchema<TSchema>
     : Record<string, unknown>;
 
+// Wrapping the callback in a deferred conditional type makes TypeScript instantiate the contextual
+// signature with the schemas inferred so far, so literals returned by execute are not widened.
+type InferredToolExecuteCallback<TInput extends object, TOutput> = [TOutput] extends [unknown]
+    ? WebMCP.ToolExecuteCallback<TInput, TOutput>
+    : never;
+
 declare global {
 /**
  * The WebMCP API enables web apps to provide JavaScript-based tools to AI agents.
@@ -63,9 +69,9 @@ namespace WebMCP {
      * Callback for tool execution.
      * @param inputObject The input parameters for the tool, as defined by its inputSchema.
      * @param options Options passed when executing the tool.
-     * @returns A promise that resolves with the tool's output.
+     * @returns A promise that resolves with the tool's output, as defined by its outputSchema.
      */
-    type ToolExecuteCallback<T extends object = Record<string, unknown>> = (inputObject: T, options: ToolExecuteCallbackOptions) => MaybePromise<unknown>;
+    type ToolExecuteCallback<TInput extends object = Record<string, unknown>, TOutput = unknown> = (inputObject: TInput, options: ToolExecuteCallbackOptions) => MaybePromise<TOutput>;
 
     /**
      * Metadata about a tool's behavior.
@@ -113,6 +119,11 @@ namespace WebMCP {
          * A JSON Schema describing the tool's input parameters.
          */
         inputSchema?: object;
+        /**
+         * A JSON Schema describing the expected output structure for the tool.
+         * Tool output is expected to conform to the provided schema, and agents may reject outputs that do not validate.
+         */
+        outputSchema?: object;
         /**
          * The function to execute when the tool is called.
          */
@@ -123,10 +134,11 @@ namespace WebMCP {
         annotations?: ToolAnnotations;
     }
 
-    /** A tool whose execute input is inferred from its input schema. */
-    type ModelContextToolFromSchema<TInputSchema extends object> = Omit<ModelContextTool, "inputSchema" | "execute"> & {
-        inputSchema: TInputSchema;
-        execute: ToolExecuteCallback<InferToolInput<TInputSchema>>;
+    /** A tool whose execute input and output are inferred from its input and output schemas. */
+    type ModelContextToolFromSchema<TInputSchema extends object = object, TOutputSchema extends object = object> = Omit<ModelContextTool, "inputSchema" | "outputSchema" | "execute"> & {
+        inputSchema?: TInputSchema;
+        outputSchema?: TOutputSchema;
+        execute: InferredToolExecuteCallback<InferToolInput<TInputSchema>, InferJsonSchema<TOutputSchema>>;
     };
 
     /**
@@ -184,6 +196,12 @@ namespace WebMCP {
          * It is a deep copy of the schema provided at tool registration, via `ModelContextTool.inputSchema`.
          */
         inputSchema?: object;
+        /**
+         * A JSON Schema object describing the expected output structure for the tool.
+         * Tool output is expected to conform to the provided schema, and agents may reject outputs that do not validate.
+         * It is a deep copy of the schema provided at tool registration, via `ModelContextTool.outputSchema`.
+         */
+        outputSchema?: object;
         /**
          * The Window of the document that registered the tool.
          */
@@ -229,17 +247,16 @@ namespace WebMCP {
      */
     interface ModelContext extends EventTarget {
         /**
-         * Registers a new tool and infers its execute input from its input schema.
-         * @param tool The tool definition.
-         * @param options Registration options.
-         */
-        registerTool<const TInputSchema extends object>(tool: ModelContextToolFromSchema<TInputSchema>, options?: ModelContextRegisterToolOptions): Promise<void>;
-        /**
-         * Registers a new tool.
+         * Registers a new tool and infers its execute input and output from its input and output schemas.
+         * Tools without schemas, or with widened `object` schemas, fall back to `Record<string, unknown>` input and `unknown` output.
+         *
+         * Note: when execute declares no parameters (e.g. `() => ({ status: "ok" })`), TypeScript checks it before
+         * inferring the output schema, so returned `enum`/`const` literals are widened (e.g. to `string`) and rejected.
+         * Declaring the input parameter (e.g. `(input) => ...`) avoids this.
          * @param tool The tool definition.
          * @param options Registration options.
          */
-        registerTool(tool: ModelContextTool, options?: ModelContextRegisterToolOptions): Promise<void>;
+        registerTool<const TInputSchema extends object = object, const TOutputSchema extends object = object>(tool: ModelContextToolFromSchema<TInputSchema, TOutputSchema>, options?: ModelContextRegisterToolOptions): Promise<void>;
         /**
          * Returns a list of registered tools exposed to this document.
          * @param options Filtering options.

index.test-d.ts

@@ -90,16 +90,32 @@ test('leaves fields optional when a required key is selected at runtime', () =>
 });
 
 declare const widenedSchema: object;
+declare const prebuiltTool: WebMCP.ModelContextTool;
 
-test('keeps the existing fallback for widened schemas and prebuilt tools', () => {
+test('keeps the fallback for widened schemas and prebuilt tools', () => {
     const tool = {
         name: 'prebuilt',
         description: 'Keeps the existing ModelContextTool contract.',
         inputSchema: widenedSchema,
-        execute: (input) => expectTypeOf(input).toEqualTypeOf<Record<string, unknown>>(),
+        outputSchema: widenedSchema,
+        execute: (input) => {
+            expectTypeOf(input).toEqualTypeOf<Record<string, unknown>>();
+            return 'anything';
+        },
     } satisfies WebMCP.ModelContextTool;
 
     void document.modelContext?.registerTool(tool);
+    void document.modelContext?.registerTool(prebuiltTool);
+    void document.modelContext?.registerTool({
+        name: 'widened',
+        description: 'Does not infer from widened schemas.',
+        inputSchema: widenedSchema,
+        outputSchema: widenedSchema,
+        execute: (input) => {
+            expectTypeOf(input).toEqualTypeOf<Record<string, unknown>>();
+            return 'anything';
+        },
+    });
 });
 
 type SearchSchema = {
@@ -108,16 +124,27 @@ type SearchSchema = {
     required: ['query'];
 };
 
+type SearchResultsSchema = {
+    type: 'array';
+    items: { type: 'string' };
+};
+
 test('exposes a named inferred tool type', () => {
     expectTypeOf<WebMCP.ModelContextToolFromSchema<SearchSchema>['execute']>()
         .parameter(0)
         .toEqualTypeOf<{ query: string }>();
+    expectTypeOf<WebMCP.ModelContextToolFromSchema<SearchSchema, SearchResultsSchema>['execute']>()
+        .returns
+        .toEqualTypeOf<WebMCP.MaybePromise<string[]>>();
 });
 
 test('uses an object fallback for a null schema', () => {
-    expectTypeOf<WebMCP.ModelContextToolFromSchema<{ type: 'null' }>['execute']>()
-        .parameter(0)
-        .toEqualTypeOf<Record<string, unknown>>();
+    void document.modelContext?.registerTool({
+        name: 'null-schema',
+        description: 'Falls back to an object input.',
+        inputSchema: { type: 'null' },
+        execute: (input) => expectTypeOf(input).toEqualTypeOf<Record<string, unknown>>(),
+    });
 });
 
 test('accepts a tool without an input schema', () => {
@@ -165,6 +192,91 @@ test('supports tool annotations', () => {
     expectTypeOf<WebMCP.RegisteredTool['annotations']>().toEqualTypeOf<WebMCP.ToolAnnotations | undefined>();
 });
 
+test('infers execute output from an output schema', () => {
+    void document.modelContext?.registerTool({
+        name: 'output-only',
+        description: 'Infers the return type from the output schema.',
+        outputSchema: {
+            type: 'object',
+            properties: {
+                total: { type: 'number' },
+                currency: { enum: ['EUR', 'USD'] },
+            },
+            required: ['total'],
+        },
+        execute: (input) => {
+            expectTypeOf(input).toEqualTypeOf<Record<string, unknown>>();
+            return { total: 42, currency: 'EUR' };
+        },
+    });
+
+    void document.modelContext?.registerTool({
+        name: 'input-and-output',
+        description: 'Infers both input and output.',
+        inputSchema: {
+            type: 'object',
+            properties: { query: { type: 'string' } },
+            required: ['query'],
+        },
+        outputSchema: {
+            type: 'object',
+            properties: { results: { type: 'array', items: { type: 'string' } } },
+        },
+        execute: (input) => {
+            expectTypeOf(input).toEqualTypeOf<{ query: string }>();
+            return { results: [input.query] };
+        },
+    });
+
+    void document.modelContext?.registerTool({
+        name: 'async-output',
+        description: 'Accepts an async execute callback.',
+        outputSchema: { type: 'array', items: { type: 'string' } },
+        execute: async () => ['a', 'b'],
+    });
+
+    expectTypeOf<WebMCP.ModelContextTool['outputSchema']>().toEqualTypeOf<object | undefined>();
+    expectTypeOf<WebMCP.RegisteredTool['outputSchema']>().toEqualTypeOf<object | undefined>();
+});
+
+test('rejects an invalid output schema or a return value that disagrees with it', () => {
+    void document.modelContext?.registerTool({
+        name: 'invalid-output-schema',
+        description: 'Output schema must be an object.',
+        // @ts-expect-error outputSchema must be an object.
+        outputSchema: '',
+        execute: () => {},
+    });
+
+    void document.modelContext?.registerTool({
+        name: 'output-mismatch',
+        description: 'Must return a value matching its output schema.',
+        outputSchema: {
+            type: 'object',
+            properties: { total: { type: 'number' } },
+            required: ['total'],
+        },
+        // @ts-expect-error total is a number in outputSchema.
+        execute: () => ({ total: 'many' }),
+    });
+
+    void document.modelContext?.registerTool({
+        name: 'output-missing',
+        description: 'Must return the required output properties.',
+        inputSchema: {
+            type: 'object',
+            properties: { query: { type: 'string' } },
+        },
+        outputSchema: {
+            type: 'object',
+            properties: { total: { type: 'number' } },
+            required: ['total'],
+        },
+        // @ts-expect-error total is required in outputSchema.
+        execute: async () => ({}),
+    });
+});
+
 declare const registeredTool: WebMCP.RegisteredTool;
 
 test('supports executeTool', async () => {