@decocms/runtime 1.3.0 → 1.4.0

package/README.md

@@ -152,37 +152,6 @@ const getUserDataTool = createPrivateTool({
 });
 ```
 
-### Streamable Tools
-
-For tools that return streaming responses:
-
-```typescript
-import { createStreamableTool } from "@decocms/runtime";
-
-const streamDataTool = createStreamableTool({
-  id: "streamData",
-  description: "Streams data as a response",
-  inputSchema: z.object({
-    query: z.string(),
-  }),
-  streamable: true,
-  execute: async ({ context }) => {
-    // Return a streaming Response
-    const stream = new ReadableStream({
-      async start(controller) {
-        controller.enqueue(new TextEncoder().encode("Chunk 1\n"));
-        controller.enqueue(new TextEncoder().encode("Chunk 2\n"));
-        controller.close();
-      },
-    });
-    
-    return new Response(stream, {
-      headers: { "Content-Type": "text/plain" },
-    });
-  },
-});
-```
-
 ### Registering Tools
 
 Tools can be registered in multiple ways:
@@ -701,7 +670,6 @@ export default withRuntime({
 ### Types
 
 - `Tool<TSchemaIn, TSchemaOut>` - Tool definition with typed input/output
-- `StreamableTool<TSchemaIn>` - Tool that returns streaming Response
 - `Prompt<TArgs>` - Prompt definition with typed arguments
 - `Resource` - Resource definition
 - `OAuthConfig` - OAuth configuration
@@ -714,7 +682,6 @@ export default withRuntime({
 - `withRuntime(options)` - Create an MCP server
 - `createTool(opts)` - Create a public tool
 - `createPrivateTool(opts)` - Create an authenticated tool
-- `createStreamableTool(opts)` - Create a streaming tool
 - `createPrompt(opts)` - Create an authenticated prompt
 - `createPublicPrompt(opts)` - Create a public prompt
 - `createResource(opts)` - Create an authenticated resource

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@decocms/runtime",
-  "version": "1.3.0",
+  "version": "1.4.0",
   "type": "module",
   "scripts": {
     "check": "tsc --noEmit",

package/src/bindings/binder.ts

@@ -7,7 +7,7 @@ import {
   type MCPClientFetchStub,
   type ToolBinder,
 } from "../mcp.ts";
-import { createPrivateTool, createStreamableTool } from "../tools.ts";
+import { createPrivateTool } from "../tools.ts";
 import { CHANNEL_BINDING } from "./channels.ts";
 
 // ToolLike is a simplified version of the Tool interface that matches what we need for bindings
@@ -77,13 +77,6 @@ export const bindingClient = <TDefinition extends readonly ToolBinder[]>(
     ): MCPClientFetchStub<TDefinition> => {
       return createMCPFetchStub<TDefinition>({
         connection: mcpConnection,
-        streamable: binder.reduce(
-          (acc, tool) => {
-            acc[tool.name] = tool.streamable === true;
-            return acc;
-          },
-          {} as Record<string, boolean>,
-        ),
       });
     },
   };
@@ -99,12 +92,8 @@ export const impl = <TBinder extends Binder>(
   schema: TBinder,
   implementation: BinderImplementation<TBinder>,
   createToolFn = createPrivateTool,
-  createStreamableToolFn = createStreamableTool,
 ) => {
-  const impl: (
-    | ReturnType<typeof createToolFn>
-    | ReturnType<typeof createStreamableToolFn>
-  )[] = [];
+  const impl: ReturnType<typeof createToolFn>[] = [];
   for (const key in schema) {
     const toolSchema = schema[key];
     const toolImplementation = implementation[key];
@@ -117,28 +106,17 @@ export const impl = <TBinder extends Binder>(
       throw new Error(`Implementation for ${key} is required`);
     }
 
-    const { name, handler, streamable, ...toolLike } = {
+    const { name, handler, ...toolLike } = {
       ...toolSchema,
       ...toolImplementation,
     };
-    if (streamable) {
-      impl.push(
-        createStreamableToolFn({
-          ...toolLike,
-          streamable,
-          id: name,
-          execute: ({ context }) => Promise.resolve(handler(context)),
-        }),
-      );
-    } else {
-      impl.push(
-        createToolFn({
-          ...toolLike,
-          id: name,
-          execute: ({ context }) => Promise.resolve(handler(context)),
-        }),
-      );
-    }
+    impl.push(
+      createToolFn({
+        ...toolLike,
+        id: name,
+        execute: ({ context }) => Promise.resolve(handler(context)),
+      }),
+    );
   }
   return impl;
 };

package/src/bindings.ts

@@ -161,7 +161,7 @@ export const AgentOf = () =>
     thinking: AgentModelInfoSchema.optional(),
     coding: AgentModelInfoSchema.optional(),
     fast: AgentModelInfoSchema.optional(),
-    toolApprovalLevel: z.enum(["auto", "readonly", "plan"]).default("readonly"),
+    toolApprovalLevel: z.enum(["auto", "readonly", "plan"]).default("auto"),
     temperature: z.number().default(0.5),
   });
 

package/src/index.ts

@@ -29,6 +29,7 @@ export {
   type ResourceContents,
   type CreatedResource,
   type WorkflowDefinition,
+  ensureAuthenticated,
 } from "./tools.ts";
 export { createWorkflow } from "./workflows.ts";
 import type { Binding } from "./wrangler.ts";

package/src/tools.ts

@@ -9,13 +9,14 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { WebStandardStreamableHTTPServerTransport as HttpServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
 import type {
   GetPromptResult,
+  Implementation,
   ToolAnnotations,
 } from "@modelcontextprotocol/sdk/types.js";
 import { z } from "zod";
 import type { ZodRawShape, ZodSchema, ZodTypeAny } from "zod";
 import { BindingRegistry, injectBindingSchemas } from "./bindings.ts";
 import { Event, type EventHandlers } from "./events.ts";
-import type { DefaultEnv } from "./index.ts";
+import type { DefaultEnv, User } from "./index.ts";
 import { State } from "./state.ts";
 import {
   type WorkflowDefinition,
@@ -66,25 +67,14 @@ export interface Tool<
   outputSchema?: TSchemaOut;
   execute(
     context: ToolExecutionContext<TSchemaIn>,
+    ctx?: AppContext,
   ): TSchemaOut extends ZodSchema
     ? Promise<z.infer<TSchemaOut>>
     : Promise<unknown>;
 }
 
 /**
- * Streamable tool interface for tools that return Response streams.
- */
-export interface StreamableTool<TSchemaIn extends ZodSchema = ZodSchema> {
-  _meta?: Record<string, unknown>;
-  id: string;
-  inputSchema: TSchemaIn;
-  streamable?: true;
-  description?: string;
-  execute(input: ToolExecutionContext<TSchemaIn>): Promise<Response>;
-}
-
-/**
- * CreatedTool is a permissive type that any Tool or StreamableTool can be assigned to.
+ * CreatedTool is a permissive type that any Tool can be assigned to.
  * Uses a structural type with relaxed execute signature to allow tools with any schema.
  */
 export type CreatedTool = {
@@ -94,12 +84,14 @@ export type CreatedTool = {
   annotations?: ToolAnnotations;
   inputSchema: ZodTypeAny;
   outputSchema?: ZodTypeAny;
-  streamable?: true;
   // Use a permissive execute signature - accepts any context shape
-  execute(context: {
-    context: unknown;
-    runtimeContext: AppContext;
-  }): Promise<unknown>;
+  execute(
+    context: {
+      context: unknown;
+      runtimeContext: AppContext;
+    },
+    ctx?: AppContext,
+  ): Promise<unknown>;
 };
 
 // Re-export types for external use
@@ -136,6 +128,7 @@ export interface Prompt<TArgs extends PromptArgsRawShape = PromptArgsRawShape> {
   argsSchema?: TArgs;
   execute(
     context: PromptExecutionContext<TArgs>,
+    ctx?: AppContext,
   ): Promise<GetPromptResult> | GetPromptResult;
 }
 
@@ -149,10 +142,13 @@ export type CreatedPrompt = {
   description?: string;
   argsSchema?: PromptArgsRawShape;
   // Use a permissive execute signature - accepts any args shape
-  execute(context: {
-    args: Record<string, string | undefined>;
-    runtimeContext: AppContext;
-  }): Promise<GetPromptResult> | GetPromptResult;
+  execute(
+    context: {
+      args: Record<string, string | undefined>;
+      runtimeContext: AppContext;
+    },
+    ctx?: AppContext,
+  ): Promise<GetPromptResult> | GetPromptResult;
 };
 
 // ============================================================================
@@ -198,6 +194,7 @@ export interface Resource {
   /** Handler function to read the resource content */
   read(
     context: ResourceExecutionContext,
+    ctx?: AppContext,
   ): Promise<ResourceContents> | ResourceContents;
 }
 
@@ -210,48 +207,58 @@ export type CreatedResource = {
   name: string;
   description?: string;
   mimeType?: string;
-  read(context: {
-    uri: URL;
-    runtimeContext: AppContext;
-  }): Promise<ResourceContents> | ResourceContents;
+  read(
+    context: {
+      uri: URL;
+      runtimeContext: AppContext;
+    },
+    ctx?: AppContext,
+  ): Promise<ResourceContents> | ResourceContents;
 };
 
 /**
- * creates a private tool that always ensure for athentication before being executed
+ * Ensure the current request is authenticated.
+ * Reads from the per-request AppContext (AsyncLocalStorage), not from a cached env.
+ *
+ * @param ctx - Per-request AppContext from the second arg of execute/read handlers
+ * @returns The authenticated User
+ * @throws Error if no request context or user is not authenticated
+ */
+export function ensureAuthenticated(ctx: AppContext): User {
+  const reqCtx = ctx?.env?.MESH_REQUEST_CONTEXT;
+  if (!reqCtx) {
+    throw new Error("Unauthorized: missing request context");
+  }
+  const user = reqCtx.ensureAuthenticated();
+  if (!user) {
+    throw new Error("Unauthorized");
+  }
+  return user;
+}
+
+let _warnedPrivateTool = false;
+
+/**
+ * @deprecated Use `createTool` with `ensureAuthenticated(ctx)` instead.
+ *
+ * Creates a private tool that ensures authentication before execution.
  */
 export function createPrivateTool<
   TSchemaIn extends ZodSchema = ZodSchema,
   TSchemaOut extends ZodSchema | undefined = undefined,
 >(opts: Tool<TSchemaIn, TSchemaOut>): Tool<TSchemaIn, TSchemaOut> {
-  const execute = opts.execute;
-  if (typeof execute === "function") {
-    opts.execute = (input: ToolExecutionContext<TSchemaIn>) => {
-      const env = input.runtimeContext.env;
-      if (env) {
-        env.MESH_REQUEST_CONTEXT?.ensureAuthenticated();
-      }
-      return execute(input);
-    };
+  if (!_warnedPrivateTool) {
+    console.warn(
+      "[runtime] createPrivateTool is deprecated. Use createTool with ensureAuthenticated(ctx) instead.",
+    );
+    _warnedPrivateTool = true;
   }
-  return createTool(opts);
-}
-
-export function createStreamableTool<TSchemaIn extends ZodSchema = ZodSchema>(
-  streamableTool: StreamableTool<TSchemaIn>,
-): StreamableTool<TSchemaIn> {
-  return {
-    ...streamableTool,
-    execute: (input: ToolExecutionContext<TSchemaIn>) => {
-      const env = input.runtimeContext.env;
-      if (env) {
-        env.MESH_REQUEST_CONTEXT?.ensureAuthenticated();
-      }
-      return streamableTool.execute({
-        ...input,
-        runtimeContext: createRuntimeContext(input.runtimeContext),
-      });
-    },
+  const execute = opts.execute;
+  opts.execute = (input: ToolExecutionContext<TSchemaIn>, ctx: AppContext) => {
+    ensureAuthenticated(ctx);
+    return execute(input, ctx);
   };
+  return createTool(opts);
 }
 
 export function createTool<
@@ -262,10 +269,8 @@ export function createTool<
   return {
     ...opts,
     execute: (input: ToolExecutionContext<TSchemaIn>) => {
-      return opts.execute({
-        ...input,
-        runtimeContext: createRuntimeContext(input.runtimeContext),
-      });
+      const ctx = createRuntimeContext(input.runtimeContext);
+      return opts.execute({ ...input, runtimeContext: ctx }, ctx);
     },
   };
 }
@@ -279,10 +284,8 @@ export function createPublicPrompt<TArgs extends PromptArgsRawShape>(
   return {
     ...opts,
     execute: (input: PromptExecutionContext<TArgs>) => {
-      return opts.execute({
-        ...input,
-        runtimeContext: createRuntimeContext(input.runtimeContext),
-      });
+      const ctx = createRuntimeContext(input.runtimeContext);
+      return opts.execute({ ...input, runtimeContext: ctx }, ctx);
     },
   };
 }
@@ -297,12 +300,9 @@ export function createPrompt<TArgs extends PromptArgsRawShape>(
   const execute = opts.execute;
   return createPublicPrompt({
     ...opts,
-    execute: (input: PromptExecutionContext<TArgs>) => {
-      const env = input.runtimeContext.env;
-      if (env) {
-        env.MESH_REQUEST_CONTEXT?.ensureAuthenticated();
-      }
-      return execute(input);
+    execute: (input: PromptExecutionContext<TArgs>, ctx: AppContext) => {
+      ensureAuthenticated(ctx);
+      return execute(input, ctx);
     },
   });
 }
@@ -314,10 +314,8 @@ export function createPublicResource(opts: Resource): Resource {
   return {
     ...opts,
     read: (input: ResourceExecutionContext) => {
-      return opts.read({
-        ...input,
-        runtimeContext: createRuntimeContext(input.runtimeContext),
-      });
+      const ctx = createRuntimeContext(input.runtimeContext);
+      return opts.read({ ...input, runtimeContext: ctx }, ctx);
     },
   };
 }
@@ -330,12 +328,9 @@ export function createResource(opts: Resource): Resource {
   const read = opts.read;
   return createPublicResource({
     ...opts,
-    read: (input: ResourceExecutionContext) => {
-      const env = input.runtimeContext.env;
-      if (env) {
-        env.MESH_REQUEST_CONTEXT?.ensureAuthenticated();
-      }
-      return read(input);
+    read: (input: ResourceExecutionContext, ctx: AppContext) => {
+      ensureAuthenticated(ctx);
+      return read(input, ctx);
     },
   });
 }
@@ -354,12 +349,6 @@ export interface Integration {
   appId: string;
 }
 
-export function isStreamableTool(
-  tool: CreatedTool,
-): tool is StreamableTool & CreatedTool {
-  return tool && "streamable" in tool && tool.streamable === true;
-}
-
 export interface OnChangeCallback<TState> {
   state: TState;
   scopes: string[];
@@ -485,6 +474,7 @@ export interface CreateMCPServerOptions<
   State extends
     TEnv["MESH_REQUEST_CONTEXT"]["state"] = TEnv["MESH_REQUEST_CONTEXT"]["state"],
 > {
+  serverInfo?: Partial<Implementation> & { instructions?: string };
   before?: (env: TEnv) => Promise<void> | void;
   oauth?: OAuthConfig;
   events?: {
@@ -498,35 +488,38 @@ export interface CreateMCPServerOptions<
   };
   tools?:
     | Array<
-        (
-          env: TEnv,
-        ) =>
-          | Promise<CreatedTool>
-          | CreatedTool
-          | CreatedTool[]
-          | Promise<CreatedTool[]>
+        | CreatedTool
+        | ((
+            env: TEnv,
+          ) =>
+            | Promise<CreatedTool>
+            | CreatedTool
+            | CreatedTool[]
+            | Promise<CreatedTool[]>)
       >
     | ((env: TEnv) => CreatedTool[] | Promise<CreatedTool[]>);
   prompts?:
     | Array<
-        (
-          env: TEnv,
-        ) =>
-          | Promise<CreatedPrompt>
-          | CreatedPrompt
-          | CreatedPrompt[]
-          | Promise<CreatedPrompt[]>
+        | CreatedPrompt
+        | ((
+            env: TEnv,
+          ) =>
+            | Promise<CreatedPrompt>
+            | CreatedPrompt
+            | CreatedPrompt[]
+            | Promise<CreatedPrompt[]>)
       >
     | ((env: TEnv) => CreatedPrompt[] | Promise<CreatedPrompt[]>);
   resources?:
     | Array<
-        (
-          env: TEnv,
-        ) =>
-          | Promise<CreatedResource>
-          | CreatedResource
-          | CreatedResource[]
-          | Promise<CreatedResource[]>
+        | CreatedResource
+        | ((
+            env: TEnv,
+          ) =>
+            | Promise<CreatedResource>
+            | CreatedResource
+            | CreatedResource[]
+            | Promise<CreatedResource[]>)
       >
     | ((env: TEnv) => CreatedResource[] | Promise<CreatedResource[]>);
   workflows?:
@@ -711,14 +704,21 @@ const toolsFor = <TSchema extends ZodTypeAny = never>({
     ...(workflows?.length
       ? workflows.map((wf) => {
           const id = wf.toolId ?? workflowToolId(wf.title);
+          const baseDescription = [
+            wf.description
+              ? `Run workflow: ${wf.description}`
+              : `Start the "${wf.title}" workflow.`,
+            "Returns an execution_id immediately. Use COLLECTION_WORKFLOW_EXECUTION_GET to track progress.",
+          ].join(" ");
           return createTool({
             id,
-            description: [
-              wf.description
-                ? `Run workflow: ${wf.description}`
-                : `Start the "${wf.title}" workflow.`,
-              "Returns an execution_id immediately. Use COLLECTION_WORKFLOW_EXECUTION_GET to track progress.",
-            ].join(" "),
+            description: (() => {
+              if (!wf.inputSchema) return baseDescription;
+              const schemaStr = JSON.stringify(wf.inputSchema, null, 2);
+              return schemaStr.length <= 2048
+                ? `${baseDescription}\n\nInput schema:\n${schemaStr}`
+                : `${baseDescription}\n\nThis workflow expects structured input. Use COLLECTION_WORKFLOW_GET to inspect the full input schema.`;
+            })(),
             inputSchema: z.object({
               input: z
                 .record(z.string(), z.unknown())
@@ -808,89 +808,155 @@ export const createMCPServer = <
 >(
   options: CreateMCPServerOptions<TEnv, TSchema, TBindings>,
 ): MCPServer<TEnv, TSchema, TBindings> => {
-  const createServer = async (bindings: TEnv) => {
-    await options.before?.(bindings);
+  // Tool/prompt/resource definitions are resolved once on first request and
+  // cached for the lifetime of the process. Tool *execution* reads per-request
+  // context from State (AsyncLocalStorage) via the second `ctx` argument, so
+  // reusing definitions is safe.
+  type Registrations = {
+    tools: CreatedTool[];
+    prompts: CreatedPrompt[];
+    resources: CreatedResource[];
+    workflows?: WorkflowDefinition[];
+  };
 
-    const server = new McpServer(
-      { name: "@deco/mcp-api", version: "1.0.0" },
-      { capabilities: { tools: {}, prompts: {}, resources: {} } },
-    );
+  let cached: Registrations | null = null;
+  let inflightResolve: Promise<Registrations> | null = null;
 
-    const toolsFn =
-      typeof options.tools === "function"
-        ? options.tools
-        : async (bindings: TEnv) => {
-            if (typeof options.tools === "function") {
-              return await options.tools(bindings);
-            }
-            return await Promise.all(
-              options.tools?.flatMap(async (tool) => {
-                const toolResult = tool(bindings);
-                const awaited = await toolResult;
-                if (Array.isArray(awaited)) {
-                  return awaited;
-                }
-                return [awaited];
-              }) ?? [],
-            ).then((t) => t.flat());
-          };
-    const tools = await toolsFn(bindings);
+  let _warnedFactoryDeprecation = false;
+  const warnFactoryDeprecation = () => {
+    if (!_warnedFactoryDeprecation) {
+      console.warn(
+        "[runtime] Passing factory functions to tools/prompts/resources is deprecated. " +
+          "Pass createTool()/createPrompt()/createResource() instances directly.",
+      );
+      _warnedFactoryDeprecation = true;
+    }
+  };
 
-    const resolvedWorkflows =
-      typeof options.workflows === "function"
-        ? await options.workflows(bindings)
-        : options.workflows;
+  /**
+   * Check whether a value is an already-created instance (has an `id` or `name` property)
+   * rather than a factory function.
+   */
+  const isInstance = (v: unknown): boolean =>
+    typeof v === "object" &&
+    v !== null &&
+    ("id" in v || "name" in v || "uri" in v);
 
-    tools.push(
-      ...toolsFor<TSchema>({ ...options, workflows: resolvedWorkflows }),
-    );
+  /**
+   * Resolve an array that may contain both direct instances and factory functions.
+   * Factories are called with `bindings` and trigger a deprecation warning.
+   */
+  async function resolveArray<T>(
+    items: Array<unknown> | undefined,
+    bindings: TEnv,
+  ): Promise<T[]> {
+    if (!items) return [];
+    return (
+      await Promise.all(
+        items.flatMap(async (item) => {
+          if (isInstance(item)) {
+            return [item as T];
+          }
+          // Factory function — deprecated path
+          warnFactoryDeprecation();
+          const result = await (item as (env: TEnv) => unknown)(bindings);
+          if (Array.isArray(result)) return result as T[];
+          return [result as T];
+        }),
+      )
+    ).flat();
+  }
+
+  const resolveRegistrations = async (
+    bindings: TEnv,
+  ): Promise<Registrations> => {
+    if (cached) return cached;
+    if (inflightResolve) return inflightResolve;
+
+    inflightResolve = (async (): Promise<Registrations> => {
+      try {
+        let tools: CreatedTool[];
+        if (typeof options.tools === "function") {
+          warnFactoryDeprecation();
+          tools = await options.tools(bindings);
+        } else {
+          tools = await resolveArray<CreatedTool>(options.tools, bindings);
+        }
+
+        const resolvedWorkflows =
+          typeof options.workflows === "function"
+            ? await options.workflows(bindings)
+            : options.workflows;
+
+        tools.push(
+          ...toolsFor<TSchema>({ ...options, workflows: resolvedWorkflows }),
+        );
+
+        let prompts: CreatedPrompt[];
+        if (typeof options.prompts === "function") {
+          warnFactoryDeprecation();
+          prompts = await options.prompts(bindings);
+        } else {
+          prompts = await resolveArray<CreatedPrompt>(
+            options.prompts,
+            bindings,
+          );
+        }
+
+        let resources: CreatedResource[];
+        if (typeof options.resources === "function") {
+          warnFactoryDeprecation();
+          resources = await options.resources(bindings);
+        } else {
+          resources = await resolveArray<CreatedResource>(
+            options.resources,
+            bindings,
+          );
+        }
+
+        const result = {
+          tools,
+          prompts,
+          resources,
+          workflows: resolvedWorkflows,
+        };
+        cached = result;
+        return result;
+      } catch (err) {
+        inflightResolve = null;
+        throw err;
+      }
+    })();
 
-    for (const tool of tools) {
+    return inflightResolve;
+  };
+
+  const registerAll = (server: McpServer, registrations: Registrations) => {
+    for (const tool of registrations.tools) {
       server.registerTool(
         tool.id,
         {
-          _meta: {
-            streamable: isStreamableTool(tool),
-            ...(tool._meta ?? {}),
-          },
+          _meta: tool._meta,
           description: tool.description,
           annotations: tool.annotations,
           inputSchema:
             tool.inputSchema && "shape" in tool.inputSchema
               ? (tool.inputSchema.shape as ZodRawShape)
               : z.object({}).shape,
-          outputSchema: isStreamableTool(tool)
-            ? z.object({ bytes: z.record(z.string(), z.number()) }).shape
-            : tool.outputSchema &&
-                typeof tool.outputSchema === "object" &&
-                "shape" in tool.outputSchema
+          outputSchema:
+            tool.outputSchema &&
+            typeof tool.outputSchema === "object" &&
+            "shape" in tool.outputSchema
               ? (tool.outputSchema.shape as ZodRawShape)
               : undefined,
         },
         async (args) => {
-          const result = await tool.execute({
-            context: args,
-            runtimeContext: createRuntimeContext(),
-          });
+          const ctx = createRuntimeContext();
+          const result = await tool.execute(
+            { context: args, runtimeContext: ctx },
+            ctx,
+          );
 
-          // For streamable tools, the Response is handled at the transport layer
-          // Do NOT call result.bytes() - it buffers the entire response in memory
-          // causing massive memory leaks (2GB+ Uint8Array accumulation)
-          if (isStreamableTool(tool) && result instanceof Response) {
-            return {
-              structuredContent: {
-                streamable: true,
-                status: result.status,
-                statusText: result.statusText,
-              },
-              content: [
-                {
-                  type: "text",
-                  text: `Streaming response: ${result.status} ${result.statusText}`,
-                },
-              ],
-            };
-          }
           return {
             structuredContent: result as Record<string, unknown>,
             content: [
@@ -904,28 +970,7 @@ export const createMCPServer = <
       );
     }
 
-    // Resolve and register prompts
-    const promptsFn =
-      typeof options.prompts === "function"
-        ? options.prompts
-        : async (bindings: TEnv) => {
-            if (typeof options.prompts === "function") {
-              return await options.prompts(bindings);
-            }
-            return await Promise.all(
-              options.prompts?.flatMap(async (prompt) => {
-                const promptResult = prompt(bindings);
-                const awaited = await promptResult;
-                if (Array.isArray(awaited)) {
-                  return awaited;
-                }
-                return [awaited];
-              }) ?? [],
-            ).then((p) => p.flat());
-          };
-    const prompts = await promptsFn(bindings);
-
-    for (const prompt of prompts) {
+    for (const prompt of registrations.prompts) {
       server.registerPrompt(
         prompt.name,
         {
@@ -936,36 +981,19 @@ export const createMCPServer = <
             : z.object({}).shape,
         },
         async (args) => {
-          return await prompt.execute({
-            args: args as Record<string, string | undefined>,
-            runtimeContext: createRuntimeContext(),
-          });
+          const ctx = createRuntimeContext();
+          return await prompt.execute(
+            {
+              args: args as Record<string, string | undefined>,
+              runtimeContext: ctx,
+            },
+            ctx,
+          );
         },
       );
     }
 
-    // Resolve and register resources
-    const resourcesFn =
-      typeof options.resources === "function"
-        ? options.resources
-        : async (bindings: TEnv) => {
-            if (typeof options.resources === "function") {
-              return await options.resources(bindings);
-            }
-            return await Promise.all(
-              options.resources?.flatMap(async (resource) => {
-                const resourceResult = resource(bindings);
-                const awaited = await resourceResult;
-                if (Array.isArray(awaited)) {
-                  return awaited;
-                }
-                return [awaited];
-              }) ?? [],
-            ).then((r) => r.flat());
-          };
-    const resources = await resourcesFn(bindings);
-
-    for (const resource of resources) {
+    for (const resource of registrations.resources) {
       server.resource(
         resource.name,
         resource.uri,
@@ -974,23 +1002,9 @@ export const createMCPServer = <
           mimeType: resource.mimeType,
         },
         async (uri) => {
-          const result = await resource.read({
-            uri,
-            runtimeContext: createRuntimeContext(),
-          });
-          // Build content object based on what's provided (text or blob, not both)
-          const content: {
-            uri: string;
-            mimeType?: string;
-            text?: string;
-            blob?: string;
-          } = { uri: result.uri };
-
-          if (result.mimeType) {
-            content.mimeType = result.mimeType;
-          }
+          const ctx = createRuntimeContext();
+          const result = await resource.read({ uri, runtimeContext: ctx }, ctx);
 
-          // MCP SDK expects either text or blob content, not both
           const meta =
             (result as { _meta?: Record<string, unknown> | null })._meta ??
             undefined;
@@ -1018,7 +1032,6 @@ export const createMCPServer = <
             };
           }
 
-          // Fallback to empty text if neither provided
           return {
             contents: [
               { uri: result.uri, mimeType: result.mimeType, text: "" },
@@ -1027,8 +1040,28 @@ export const createMCPServer = <
         },
       );
     }
+  };
+
+  const createServer = async (bindings: TEnv) => {
+    await options.before?.(bindings);
+
+    const { instructions, ...serverInfoOverrides } = options.serverInfo ?? {};
+    const server = new McpServer(
+      {
+        ...serverInfoOverrides,
+        name: serverInfoOverrides.name ?? "@deco/mcp-api",
+        version: serverInfoOverrides.version ?? "1.0.0",
+      },
+      {
+        capabilities: { tools: {}, prompts: {}, resources: {} },
+        ...(instructions && { instructions }),
+      },
+    );
+
+    const registrations = await resolveRegistrations(bindings);
+    registerAll(server, registrations);
 
-    return { server, tools, prompts, resources };
+    return { server, ...registrations };
   };
 
   const fetch = async (req: Request, env: TEnv) => {
@@ -1037,34 +1070,48 @@ export const createMCPServer = <
 
     await server.connect(transport);
 
+    const cleanup = () => {
+      try {
+        transport.close?.();
+      } catch {
+        /* ignore */
+      }
+      try {
+        server.close?.();
+      } catch {
+        /* ignore */
+      }
+    };
+
     try {
       const response = await transport.handleRequest(req);
 
-      // Check if this is a streaming response (SSE or streamable tool)
-      // SSE responses have text/event-stream content-type
-      // Note: response.body is always non-null for all HTTP responses, so we can't use it to detect streaming
       const contentType = response.headers.get("content-type");
       const isStreaming =
         contentType?.includes("text/event-stream") ||
         contentType?.includes("application/json-rpc");
 
-      // Only close transport for non-streaming responses
-      if (!isStreaming) {
-        try {
-          await transport.close?.();
-        } catch {
-          // Ignore close errors
-        }
+      if (!isStreaming || !response.body) {
+        cleanup();
+        return response;
       }
 
-      return response;
+      // Pipe the SSE body through a passthrough so that when the stream
+      // finishes (server sent the response) or the client disconnects
+      // (cancel), the server and transport are always cleaned up.
+      const { readable, writable } = new TransformStream();
+      response.body
+        .pipeTo(writable)
+        .catch(() => {})
+        .finally(cleanup);
+
+      return new Response(readable, {
+        status: response.status,
+        statusText: response.statusText,
+        headers: response.headers,
+      });
     } catch (error) {
-      // On error, always try to close transport to prevent leaks
-      try {
-        await transport.close?.();
-      } catch {
-        // Ignore close errors
-      }
+      cleanup();
       throw error;
     }
   };
@@ -1075,7 +1122,9 @@ export const createMCPServer = <
       throw new Error("Missing state, did you forget to call State.bind?");
     }
     const env = currentState?.env;
-    const { tools } = await createServer(env as TEnv & DefaultEnv<TSchema>);
+    const { tools } = await resolveRegistrations(
+      env as TEnv & DefaultEnv<TSchema>,
+    );
     const tool = tools.find((t) => t.id === toolCallId);
     const execute = tool?.execute;
     if (!execute) {
@@ -1084,10 +1133,8 @@ export const createMCPServer = <
       );
     }
 
-    return execute({
-      context: toolCallInput,
-      runtimeContext: createRuntimeContext(),
-    });
+    const ctx = createRuntimeContext();
+    return execute({ context: toolCallInput, runtimeContext: ctx }, ctx);
   };
 
   return {

package/src/workflows.ts

@@ -24,6 +24,12 @@ export interface WorkflowDefinition {
    * Defaults to START_WORKFLOW_<TITLE_SLUG> (e.g. START_WORKFLOW_FETCH_USERS).
    */
   toolId?: string;
+  /**
+   * JSON Schema describing the expected input for this workflow.
+   * When set, the mesh validates execution input against this schema
+   * before creating an execution.
+   */
+  inputSchema?: Record<string, unknown> | null;
 }
 
 interface WorkflowCollectionItem {
@@ -67,6 +73,7 @@ interface MeshWorkflowClient {
       description?: string;
       virtual_mcp_id?: string;
       steps: Step[];
+      input_schema?: Record<string, unknown> | null;
     };
   }) => Promise<{ item: WorkflowCollectionItem }>;
   COLLECTION_WORKFLOW_UPDATE: (input: {
@@ -76,6 +83,7 @@ interface MeshWorkflowClient {
       description?: string;
       virtual_mcp_id?: string;
       steps?: Step[];
+      input_schema?: Record<string, unknown> | null;
     };
   }) => Promise<{ success: boolean; error?: string }>;
   COLLECTION_WORKFLOW_DELETE: (input: {
@@ -254,10 +262,11 @@ function fingerprintWorkflows(declared: WorkflowDefinition[]): string {
   return JSON.stringify(
     declared.map((w) => ({
       title: w.title,
-      description: w.description ?? null,
-      virtual_mcp_id: w.virtual_mcp_id ?? null,
+      description: w.description ?? undefined,
+      virtual_mcp_id: w.virtual_mcp_id ?? undefined,
       steps: w.steps,
-      toolId: w.toolId ?? null,
+      toolId: w.toolId ?? undefined,
+      inputSchema: w.inputSchema ?? undefined,
     })),
   );
 }
@@ -383,6 +392,10 @@ async function doSyncWorkflows(
                 virtual_mcp_id: resolvedVmcpId,
               }),
               steps: wf.steps,
+              input_schema:
+                wf.inputSchema === undefined
+                  ? undefined
+                  : (wf.inputSchema ?? null),
             },
           });
           if (!result.success) {
@@ -402,6 +415,7 @@ async function doSyncWorkflows(
               description: wf.description,
               virtual_mcp_id: resolvedVmcpId,
               steps: wf.steps,
+              input_schema: wf.inputSchema ?? null,
             },
           });
           console.log(`${tag} CREATE "${wf.title}" OK`);
@@ -596,8 +610,30 @@ type InputForTool<
     : StepInput<TSteps>
   : StepInput<TSteps>;
 
-type BaseStepFields = Omit<Step, "name" | "input" | "action">;
-type BaseForEachFields = Omit<Step, "name" | "forEach" | "input" | "action">;
+/**
+ * Typed bail condition with @ref autocomplete.
+ * Self-references (e.g. `@thisStep.field` on step "thisStep") are valid for
+ * bail — the condition is evaluated after the step completes — but the current
+ * step name isn't in TSteps yet. Use the `(string & {})` escape hatch.
+ */
+type TypedBail<TSteps extends string> =
+  | true
+  | {
+      ref: KnownRefs<TSteps>;
+      eq?: unknown;
+      neq?: unknown;
+      gt?: number;
+      lt?: number;
+    };
+
+type BaseStepFields<TSteps extends string> = Omit<
+  Step,
+  "name" | "input" | "action" | "bail"
+> & { bail?: TypedBail<TSteps> };
+type BaseForEachFields<TSteps extends string> = Omit<
+  Step,
+  "name" | "forEach" | "input" | "action" | "bail"
+> & { bail?: TypedBail<TSteps> };
 
 /**
  * Tool-call variants of StepOpts — one discriminated member per tool ID so
@@ -608,12 +644,12 @@ type ToolCallStepOpts<
   TSteps extends string,
   TTools extends readonly ToolLike[],
 > = [TTools[number]] extends [never]
-  ? BaseStepFields & {
+  ? BaseStepFields<TSteps> & {
       action: { toolName: string & {}; transformCode?: string };
       input?: StepInput<TSteps>;
     }
   : {
-      [TId in TTools[number]["id"]]: BaseStepFields & {
+      [TId in TTools[number]["id"]]: BaseStepFields<TSteps> & {
         action: { toolName: TId; transformCode?: string };
         input?: InputForTool<TTools, TId, TSteps>;
       };
@@ -621,19 +657,22 @@ type ToolCallStepOpts<
 
 type StepOpts<TSteps extends string, TTools extends readonly ToolLike[]> =
   | ToolCallStepOpts<TSteps, TTools>
-  | (BaseStepFields & { action: { code: string }; input?: StepInput<TSteps> });
+  | (BaseStepFields<TSteps> & {
+      action: { code: string };
+      input?: StepInput<TSteps>;
+    });
 
 type ToolCallForEachOpts<
   TSteps extends string,
   TTools extends readonly ToolLike[],
 > = [TTools[number]] extends [never]
-  ? BaseForEachFields & {
+  ? BaseForEachFields<TSteps> & {
       action: { toolName: string & {}; transformCode?: string };
       input?: StepInput<TSteps>;
       concurrency?: number;
     }
   : {
-      [TId in TTools[number]["id"]]: BaseForEachFields & {
+      [TId in TTools[number]["id"]]: BaseForEachFields<TSteps> & {
         action: { toolName: TId; transformCode?: string };
         input?: InputForTool<TTools, TId, TSteps>;
         concurrency?: number;
@@ -645,7 +684,7 @@ type ForEachItemOpts<
   TTools extends readonly ToolLike[],
 > =
   | ToolCallForEachOpts<TSteps, TTools>
-  | (BaseForEachFields & {
+  | (BaseForEachFields<TSteps> & {
       action: { code: string };
       input?: StepInput<TSteps>;
       concurrency?: number;