feat: add capability system for AI agent bubbles

Author: zhubzy

packages/bubble-core/src/bubbles/service-bubble/ai-agent.ts

@@ -38,6 +38,10 @@ import {
   zodSchemaToJsonString,
   buildJsonSchemaInstruction,
 } from '../../utils/zod-schema.js';
+import {
+  getCapability,
+  type CapabilityRuntimeContext,
+} from '../../capabilities/index.js';
 
 // Define tool hook context - provides access to messages and tool call details
 export type ToolHookContext = {
@@ -251,6 +255,23 @@ const ExpectedOutputSchema = z.union([
 /** Type for conversation history messages - enables KV cache optimization */
 export type ConversationMessage = z.infer<typeof ConversationMessageSchema>;
 
+// Schema for a single capability configuration on the AI agent
+const CapabilityConfigSchema = z.object({
+  id: z
+    .string()
+    .min(1)
+    .describe('Capability ID (e.g., "google-doc-knowledge-base")'),
+  inputs: z
+    .record(z.string(), z.union([z.string(), z.number(), z.boolean()]))
+    .default({})
+    .describe('Input parameter values for this capability'),
+  credentials: z
+    .record(z.nativeEnum(CredentialType), z.string())
+    .default({})
+    .optional()
+    .describe('Capability-specific credentials (injected at runtime)'),
+});
+
 // Define the parameters schema for the AI Agent bubble
 const AIAgentParamsSchema = z.object({
   message: z
@@ -322,6 +343,13 @@ const AIAgentParamsSchema = z.object({
     .describe(
       'Enable real-time streaming of tokens, tool calls, and iteration progress'
     ),
+  capabilities: z
+    .array(CapabilityConfigSchema)
+    .default([])
+    .optional()
+    .describe(
+      'Capabilities that extend the agent with bundled tools, prompts, and credentials. Example: [{ id: "google-doc-knowledge-base", inputs: { docId: "your-doc-id" } }]'
+    ),
   expectedOutputSchema: ExpectedOutputSchema.optional().describe(
     'Zod schema or JSON schema string that defines the expected structure of the AI response. When provided, automatically enables JSON mode and instructs the AI to output in the exact format. Example: z.object({ summary: z.string(), items: z.array(z.object({ name: z.string(), score: z.number() })) })'
   ),
@@ -522,6 +550,28 @@ export class AIAgentBubble extends ServiceBubble<
       );
       this.params.systemPrompt = `${this.params.systemPrompt}\n\n${buildJsonSchemaInstruction(schemaString)}`;
     }
+
+    // Inject capability system prompts
+    for (const capConfig of this.params.capabilities ?? []) {
+      const capDef = getCapability(capConfig.id);
+      if (!capDef) continue;
+
+      const ctx: CapabilityRuntimeContext = {
+        credentials:
+          (capConfig.credentials as Partial<Record<CredentialType, string>>) ??
+          {},
+        inputs: capConfig.inputs ?? {},
+        bubbleContext: this.context,
+      };
+
+      const addition =
+        capDef.createSystemPrompt?.(ctx) ??
+        capDef.metadata.systemPromptAddition;
+
+      if (addition) {
+        this.params.systemPrompt = `${this.params.systemPrompt}\n\n${addition}`;
+      }
+    }
   }
 
   protected async performAction(
@@ -1059,9 +1109,118 @@ export class AIAgentBubble extends ServiceBubble<
       }
     }
 
+    // 3. Capability tools
+    for (const capConfig of this.params.capabilities ?? []) {
+      const capDef = getCapability(capConfig.id);
+      if (!capDef) {
+        console.warn(
+          `[AIAgent] Capability '${capConfig.id}' not found in registry. Skipping.`
+        );
+        continue;
+      }
+
+      try {
+        const ctx: CapabilityRuntimeContext = {
+          credentials:
+            (capConfig.credentials as Partial<
+              Record<CredentialType, string>
+            >) ?? {},
+          inputs: capConfig.inputs ?? {},
+          bubbleContext: this.context,
+        };
+
+        const toolFuncs = capDef.createTools(ctx);
+
+        for (const toolMeta of capDef.metadata.tools) {
+          const func = toolFuncs[toolMeta.name];
+          if (!func) continue;
+
+          // Convert JSON schema back to Zod for DynamicStructuredTool
+          const toolSchema = this.jsonSchemaToZod(toolMeta.parameterSchema);
+
+          const dynamicTool = new DynamicStructuredTool({
+            name: toolMeta.name,
+            description: toolMeta.description,
+            schema: toolSchema,
+            func: func as (input: Record<string, unknown>) => Promise<unknown>,
+          } as any);
+
+          tools.push(dynamicTool);
+          console.log(
+            `🔧 [AIAgent] Registered capability tool: ${toolMeta.name} (from ${capConfig.id})`
+          );
+        }
+      } catch (error) {
+        console.error(
+          `Error initializing capability '${capConfig.id}':`,
+          error
+        );
+        continue;
+      }
+    }
+
     return tools;
   }
 
+  /**
+   * Converts a JSON Schema object to a Zod schema for DynamicStructuredTool.
+   * Handles common JSON Schema types used by capability tool definitions.
+   */
+  private jsonSchemaToZod(
+    jsonSchema: Record<string, unknown>
+  ): z.ZodObject<z.ZodRawShape> {
+    const properties = (
+      jsonSchema as { properties?: Record<string, Record<string, unknown>> }
+    ).properties;
+    const required = (jsonSchema as { required?: string[] }).required ?? [];
+
+    if (!properties || Object.keys(properties).length === 0) {
+      return z.object({});
+    }
+
+    const shape: z.ZodRawShape = {};
+    for (const [key, prop] of Object.entries(properties)) {
+      let fieldSchema: z.ZodTypeAny;
+
+      switch (prop.type) {
+        case 'string':
+          fieldSchema = z.string();
+          if (prop.description)
+            fieldSchema = fieldSchema.describe(prop.description as string);
+          break;
+        case 'number':
+        case 'integer':
+          fieldSchema = z.number();
+          if (prop.description)
+            fieldSchema = fieldSchema.describe(prop.description as string);
+          break;
+        case 'boolean':
+          fieldSchema = z.boolean();
+          if (prop.description)
+            fieldSchema = fieldSchema.describe(prop.description as string);
+          break;
+        case 'array':
+          fieldSchema = z.array(z.unknown());
+          if (prop.description)
+            fieldSchema = fieldSchema.describe(prop.description as string);
+          break;
+        default:
+          fieldSchema = z.unknown();
+          if (prop.description)
+            fieldSchema = fieldSchema.describe(prop.description as string);
+          break;
+      }
+
+      if (!required.includes(key)) {
+        fieldSchema = fieldSchema.optional();
+      }
+
+      shape[key] = fieldSchema;
+    }
+
+    return z.object(shape);
+  }
+
   /**
    * Custom tool execution node that supports hooks
    */

packages/bubble-core/src/capabilities/define-capability.ts

@@ -0,0 +1,123 @@
+import type {
+  CapabilityMetadata,
+  CapabilityInput,
+  CapabilityToolDef,
+} from '@bubblelab/shared-schemas';
+import type { CredentialType } from '@bubblelab/shared-schemas';
+import type {
+  ToolHookBefore,
+  ToolHookAfter,
+} from '../bubbles/service-bubble/ai-agent.js';
+import type { BubbleContext } from '../types/bubble.js';
+import { z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+
+/** Runtime context passed to capability tool factories and system prompt factories. */
+export interface CapabilityRuntimeContext {
+  credentials: Partial<Record<CredentialType, string>>;
+  inputs: Record<string, string | number | boolean>;
+  bubbleContext?: BubbleContext;
+}
+
+/** A single capability tool function that accepts parsed parameters and returns a result. */
+export type CapabilityToolFunc = (
+  params: Record<string, unknown>
+) => Promise<unknown>;
+
+/** Factory that creates tool functions given a runtime context. */
+export type CapabilityToolFactory = (
+  context: CapabilityRuntimeContext
+) => Record<string, CapabilityToolFunc>;
+
+/** Factory that creates a system prompt addition given a runtime context. */
+export type CapabilitySystemPromptFactory = (
+  context: CapabilityRuntimeContext
+) => string;
+
+/** Full runtime capability definition with metadata + factories. */
+export interface CapabilityDefinition {
+  metadata: CapabilityMetadata;
+  createTools: CapabilityToolFactory;
+  createSystemPrompt?: CapabilitySystemPromptFactory;
+  hooks?: {
+    beforeToolCall?: ToolHookBefore;
+    afterToolCall?: ToolHookAfter;
+  };
+}
+
+/** Options for the defineCapability() helper — ergonomic API for creating capabilities. */
+export interface DefineCapabilityOptions {
+  id: string;
+  name: string;
+  description: string;
+  icon?: string;
+  category?: string;
+  version?: string;
+  requiredCredentials: CredentialType[];
+  inputs: CapabilityInput[];
+  tools: Array<{
+    name: string;
+    description: string;
+    schema: z.ZodObject<z.ZodRawShape>;
+    func: (ctx: CapabilityRuntimeContext) => CapabilityToolFunc;
+  }>;
+  systemPrompt?: string | CapabilitySystemPromptFactory;
+  hooks?: CapabilityDefinition['hooks'];
+}
+
+/**
+ * Creates a CapabilityDefinition from a user-friendly options object.
+ * Converts Zod schemas to JSON Schema for serializable metadata,
+ * and wraps tool functions with context currying.
+ */
+export function defineCapability(
+  options: DefineCapabilityOptions
+): CapabilityDefinition {
+  // Build serializable tool definitions from Zod schemas
+  const toolDefs: CapabilityToolDef[] = options.tools.map((tool) => ({
+    name: tool.name,
+    description: tool.description,
+    parameterSchema: zodToJsonSchema(tool.schema, {
+      $refStrategy: 'none',
+    }) as Record<string, unknown>,
+  }));
+
+  // Build serializable metadata
+  const metadata: CapabilityMetadata = {
+    id: options.id,
+    name: options.name,
+    description: options.description,
+    icon: options.icon,
+    category: options.category,
+    version: options.version ?? '1.0.0',
+    requiredCredentials: options.requiredCredentials,
+    inputs: options.inputs,
+    tools: toolDefs,
+    systemPromptAddition:
+      typeof options.systemPrompt === 'string'
+        ? options.systemPrompt
+        : undefined,
+  };
+
+  // Build tool factory that curries context into each tool func
+  const createTools: CapabilityToolFactory = (ctx) => {
+    const toolFuncs: Record<string, CapabilityToolFunc> = {};
+    for (const tool of options.tools) {
+      toolFuncs[tool.name] = tool.func(ctx);
+    }
+    return toolFuncs;
+  };
+
+  // Build system prompt factory
+  const createSystemPrompt: CapabilitySystemPromptFactory | undefined =
+    typeof options.systemPrompt === 'function'
+      ? options.systemPrompt
+      : undefined;
+
+  return {
+    metadata,
+    createTools,
+    createSystemPrompt,
+    hooks: options.hooks,
+  };
+}

packages/bubble-core/src/capabilities/index.ts

@@ -0,0 +1,17 @@
+export {
+  defineCapability,
+  type CapabilityDefinition,
+  type CapabilityRuntimeContext,
+  type CapabilityToolFunc,
+  type CapabilityToolFactory,
+  type CapabilitySystemPromptFactory,
+  type DefineCapabilityOptions,
+} from './define-capability.js';
+
+export {
+  registerCapability,
+  getCapability,
+  getAllCapabilities,
+  getAllCapabilityMetadata,
+  getCapabilityMetadataById,
+} from './registry.js';

packages/bubble-core/src/capabilities/registry.ts

@@ -0,0 +1,37 @@
+import type { CapabilityMetadata } from '@bubblelab/shared-schemas';
+import type { CapabilityDefinition } from './define-capability.js';
+
+/** Global registry of capability definitions, keyed by capability ID. */
+const capabilityRegistry = new Map<string, CapabilityDefinition>();
+
+/** Registers a capability definition in the global registry. */
+export function registerCapability(cap: CapabilityDefinition): void {
+  if (capabilityRegistry.has(cap.metadata.id)) {
+    console.warn(
+      `[CapabilityRegistry] Overwriting existing capability: ${cap.metadata.id}`
+    );
+  }
+  capabilityRegistry.set(cap.metadata.id, cap);
+}
+
+/** Returns a registered capability by ID, or undefined if not found. */
+export function getCapability(id: string): CapabilityDefinition | undefined {
+  return capabilityRegistry.get(id);
+}
+
+/** Returns all registered capability definitions. */
+export function getAllCapabilities(): CapabilityDefinition[] {
+  return Array.from(capabilityRegistry.values());
+}
+
+/** Returns serializable metadata for all registered capabilities. */
+export function getAllCapabilityMetadata(): CapabilityMetadata[] {
+  return Array.from(capabilityRegistry.values()).map((cap) => cap.metadata);
+}
+
+/** Returns metadata for a single capability by ID, or undefined if not found. */
+export function getCapabilityMetadataById(
+  id: string
+): CapabilityMetadata | undefined {
+  return capabilityRegistry.get(id)?.metadata;
+}

packages/bubble-core/src/index.ts

@@ -4,6 +4,9 @@ export * from '@bubblelab/shared-schemas';
 export * from './types/credentials.js';
 export * from './types/available-tools.js';
 
+// Export capabilities framework
+export * from './capabilities/index.js';
+
 // Export error classes
 export {
   BubbleError,