@okrlinkhub/agent-bridge 0.1.0 → 2.0.0

package/src/client/index.ts

@@ -1,477 +1,316 @@
 import { httpActionGeneric } from "convex/server";
 import type {
-  GenericActionCtx,
+  FunctionHandle,
   GenericDataModel,
   GenericMutationCtx,
-  GenericQueryCtx,
-  FunctionHandle,
   HttpRouter,
 } from "convex/server";
 import type { ComponentApi } from "../component/_generated/component.js";
 
-// Convenient context types with minimal required capabilities.
-type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;
-type MutationCtx = Pick<
-  GenericMutationCtx<GenericDataModel>,
-  "runQuery" | "runMutation"
->;
-type ActionCtx = Pick<
-  GenericActionCtx<GenericDataModel>,
-  "runQuery" | "runMutation" | "runAction"
->;
-
-// --- Types ---
-
-export interface FunctionDef {
-  /** Alias name that agents use to call this function (e.g. "okr:getObjectives") */
-  name: string;
-  /** Function handle string from createFunctionHandle() */
-  handle: string;
-  /** Type of the Convex function */
-  type: "query" | "mutation" | "action";
-  /** Human-readable description */
-  description?: string;
+export type AgentBridgeFunctionType = "query" | "mutation" | "action";
+
+type UnknownFunctionReference = unknown;
+
+export interface AgentBridgeFunctionDefinition {
+  ref: UnknownFunctionReference;
+  type?: AgentBridgeFunctionType;
 }
 
-export interface DefaultPermission {
-  pattern: string;
-  permission: "allow" | "deny" | "rate_limited";
-  rateLimitConfig?: {
-    requestsPerHour: number;
-    tokenBudget: number;
-  };
+export interface AgentBridgeFunctionMetadata {
+  description?: string;
+  riskLevel?: "low" | "medium" | "high";
+  category?: string;
 }
 
 export interface AgentBridgeConfig {
-  /** Unique name identifying this app (e.g. "okr", "hr", "incentives") */
-  appName: string;
-  /** Default permissions applied to new agents during provisioning */
-  defaultPermissions?: DefaultPermission[];
+  functions: Record<
+    string,
+    UnknownFunctionReference | AgentBridgeFunctionDefinition
+  >;
+  metadata?: Record<string, AgentBridgeFunctionMetadata>;
 }
 
-// --- AgentBridge Client Class ---
-
-/**
- * Client class that wraps component calls for ergonomic use in the host app.
- *
- * Usage:
- * ```ts
- * import { AgentBridge } from "@okrlinkhub/agent-bridge";
- * import { components } from "./_generated/api";
- *
- * const bridge = new AgentBridge(components.agentBridge, { appName: "okr" });
- * ```
- */
-export class AgentBridge {
-  public component: ComponentApi;
-  public appName: string;
-  private defaultPermissions: DefaultPermission[];
-
-  constructor(component: ComponentApi, config: AgentBridgeConfig) {
-    this.component = component;
-    this.appName = config.appName;
-    this.defaultPermissions = config.defaultPermissions ?? [
-      { pattern: "*", permission: "deny" },
-    ];
-  }
+export function defineAgentBridgeConfig(
+  config: AgentBridgeConfig,
+): AgentBridgeConfig {
+  return config;
+}
 
-  /**
-   * Initialize the component configuration.
-   * Should be called once during app setup (e.g. in a seed/setup mutation).
-   */
-  async configure(ctx: MutationCtx): Promise<void> {
-    await ctx.runMutation(this.component.provisioning.configure, {
-      appName: this.appName,
-      defaultPermissions: this.defaultPermissions,
-    });
-  }
+export function generateAgentApiKey(prefix: string = "abk_live"): string {
+  const bytes = new Uint8Array(24);
+  crypto.getRandomValues(bytes);
+  const token = Array.from(bytes)
+    .map((byte) => byte.toString(16).padStart(2, "0"))
+    .join("");
+  return `${prefix}_${token}`;
+}
 
-  /**
-   * Register a function that agents can call.
-   * The host app must create the function handle via createFunctionHandle().
-   */
-  async registerFunction(ctx: MutationCtx, fn: FunctionDef): Promise<string> {
-    return await ctx.runMutation(this.component.registry.register, {
-      appName: this.appName,
-      functionName: fn.name,
-      functionHandle: fn.handle,
-      functionType: fn.type,
-      description: fn.description,
-    });
+type NormalizedFunctionDefinition = {
+  ref: UnknownFunctionReference;
+  type: AgentBridgeFunctionType;
+  metadata?: AgentBridgeFunctionMetadata;
+};
+
+type NormalizedAgentBridgeConfig = {
+  functions: Record<string, NormalizedFunctionDefinition>;
+};
+
+export function detectFunctionType(
+  fnRef: UnknownFunctionReference,
+): AgentBridgeFunctionType | null {
+  if (!fnRef || typeof fnRef !== "object") {
+    return null;
   }
 
-  /**
-   * Register multiple functions in bulk.
-   */
-  async registerFunctions(
-    ctx: MutationCtx,
-    functions: FunctionDef[],
-  ): Promise<string[]> {
-    const ids: string[] = [];
-    for (const fn of functions) {
-      const id = await this.registerFunction(ctx, fn);
-      ids.push(id);
+  if ("_type" in fnRef) {
+    const candidate = (fnRef as { _type?: unknown })._type;
+    if (
+      candidate === "query" ||
+      candidate === "mutation" ||
+      candidate === "action"
+    ) {
+      return candidate;
     }
-    return ids;
-  }
-
-  /**
-   * Unregister a function.
-   */
-  async unregisterFunction(
-    ctx: MutationCtx,
-    functionName: string,
-  ): Promise<boolean> {
-    return await ctx.runMutation(this.component.registry.unregister, {
-      appName: this.appName,
-      functionName,
-    });
-  }
-
-  /**
-   * List all registered functions for this app.
-   */
-  async listFunctions(ctx: QueryCtx) {
-    return await ctx.runQuery(this.component.registry.listFunctions, {
-      appName: this.appName,
-    });
-  }
-
-  /**
-   * Generate a provisioning token (admin operation).
-   * Returns the plaintext token -- store/communicate securely!
-   */
-  async generateProvisioningToken(
-    ctx: MutationCtx,
-    opts: {
-      employeeEmail: string;
-      department: string;
-      maxApps?: number;
-      expiresInDays?: number;
-      createdBy: string;
-    },
-  ) {
-    return await ctx.runMutation(
-      this.component.provisioning.generateProvisioningToken,
-      opts,
-    );
   }
 
-  /**
-   * List all registered agents.
-   */
-  async listAgents(ctx: QueryCtx, activeOnly?: boolean) {
-    return await ctx.runQuery(this.component.provisioning.listAgents, {
-      activeOnly,
-    });
-  }
-
-  /**
-   * Revoke an agent globally.
-   */
-  async revokeAgent(
-    ctx: MutationCtx,
-    opts: { agentId: string; revokedBy: string },
-  ): Promise<boolean> {
-    return await ctx.runMutation(this.component.provisioning.revokeAgent, opts);
-  }
+  return null;
+}
 
-  /**
-   * Revoke a specific app instance for an agent.
-   */
-  async revokeAppInstance(
-    ctx: MutationCtx,
-    opts: { agentId: string },
-  ): Promise<boolean> {
-    return await ctx.runMutation(
-      this.component.provisioning.revokeAppInstance,
-      {
-        agentId: opts.agentId,
-        appName: this.appName,
-      },
-    );
+export function normalizeAgentBridgeConfig(
+  config: AgentBridgeConfig,
+): NormalizedAgentBridgeConfig {
+  const functionEntries = Object.entries(config.functions);
+  if (functionEntries.length === 0) {
+    throw new Error("agent-bridge config requires at least one function");
   }
 
-  /**
-   * Set a permission for an agent.
-   */
-  async setPermission(
-    ctx: MutationCtx,
-    opts: {
-      agentId: string;
-      functionPattern: string;
-      permission: "allow" | "deny" | "rate_limited";
-      rateLimitConfig?: { requestsPerHour: number; tokenBudget: number };
-      createdBy: string;
-    },
-  ): Promise<string> {
-    return await ctx.runMutation(this.component.permissions.setPermission, {
-      agentId: opts.agentId,
-      appName: this.appName,
-      functionPattern: opts.functionPattern,
-      permission: opts.permission,
-      rateLimitConfig: opts.rateLimitConfig,
-      createdBy: opts.createdBy,
-    });
-  }
+  const normalizedFunctions: Record<string, NormalizedFunctionDefinition> = {};
+  for (const [functionKey, entry] of functionEntries) {
+    if (!functionKey.trim()) {
+      throw new Error("function keys cannot be empty");
+    }
 
-  /**
-   * Remove a permission for an agent.
-   */
-  async removePermission(
-    ctx: MutationCtx,
-    opts: { agentId: string; functionPattern: string },
-  ): Promise<boolean> {
-    return await ctx.runMutation(this.component.permissions.removePermission, {
-      agentId: opts.agentId,
-      appName: this.appName,
-      functionPattern: opts.functionPattern,
-    });
-  }
+    const metadata = config.metadata?.[functionKey];
+    const hasExplicitConfig =
+      entry && typeof entry === "object" && "ref" in (entry as object);
+    const ref = hasExplicitConfig
+      ? (entry as AgentBridgeFunctionDefinition).ref
+      : entry;
+    const explicitType = hasExplicitConfig
+      ? (entry as AgentBridgeFunctionDefinition).type
+      : undefined;
+
+    const detectedType = detectFunctionType(ref);
+    const functionType = explicitType ?? detectedType;
+    if (!functionType) {
+      throw new Error(
+        `Cannot detect function type for "${functionKey}". Set { ref, type } explicitly in agent-bridge config.`,
+      );
+    }
 
-  /**
-   * List permissions for an agent on this app.
-   */
-  async listPermissions(ctx: QueryCtx, agentId: string) {
-    return await ctx.runQuery(this.component.permissions.listPermissions, {
-      agentId,
-      appName: this.appName,
-    });
+    normalizedFunctions[functionKey] = {
+      ref,
+      type: functionType,
+      metadata,
+    };
   }
 
-  /**
-   * Query access logs.
-   */
-  async queryAccessLog(
-    ctx: QueryCtx,
-    opts?: { agentId?: string; limit?: number },
-  ) {
-    return await ctx.runQuery(this.component.gateway.queryAccessLog, {
-      agentId: opts?.agentId,
-      appName: this.appName,
-      limit: opts?.limit,
-    });
-  }
+  return { functions: normalizedFunctions };
 }
 
-// --- HTTP Route Registration ---
-
-/**
- * Register HTTP routes for the agent bridge component.
- * This exposes endpoints that OpenClaw agents call to execute functions,
- * provision themselves, and check health.
- *
- * Must be called in the host app's `convex/http.ts` file.
- *
- * @example
- * ```ts
- * import { httpRouter } from "convex/server";
- * import { registerRoutes } from "@okrlinkhub/agent-bridge";
- * import { components } from "./_generated/api";
- *
- * const http = httpRouter();
- * registerRoutes(http, components.agentBridge, { appName: "okr" });
- * export default http;
- * ```
- */
+type ExecuteRequestBody = {
+  functionKey?: string;
+  args?: Record<string, unknown>;
+  estimatedCost?: number;
+};
+
 export function registerRoutes(
   http: HttpRouter,
   component: ComponentApi,
-  config: { appName: string; pathPrefix?: string },
+  bridgeConfig: AgentBridgeConfig,
+  options?: { pathPrefix?: string },
 ) {
-  const prefix = config.pathPrefix ?? "/agent-bridge";
+  const prefix = options?.pathPrefix ?? "/agent";
+  const normalizedConfig = normalizeAgentBridgeConfig(bridgeConfig);
+  const availableFunctionKeys = Object.keys(normalizedConfig.functions);
 
-  // --- POST /agent-bridge/execute ---
-  // Gateway: execute a registered function on behalf of an agent
   http.route({
     path: `${prefix}/execute`,
     method: "POST",
     handler: httpActionGeneric(async (ctx, request) => {
-      let body: {
-        instanceToken?: string;
-        functionName?: string;
-        args?: Record<string, unknown>;
-      };
+      const apiKey = request.headers.get("X-Agent-API-Key");
+      if (!apiKey) {
+        return jsonResponse({ success: false, error: "Missing API key" }, 401);
+      }
+
+      let body: ExecuteRequestBody;
       try {
         body = await request.json();
       } catch {
-        return jsonResponse({ error: "Invalid JSON body" }, 400);
+        return jsonResponse(
+          { success: false, error: "Invalid JSON body" },
+          400,
+        );
       }
 
-      const { instanceToken, functionName, args } = body;
-
-      if (!instanceToken || !functionName) {
+      const functionKey = body.functionKey?.trim();
+      if (!functionKey) {
         return jsonResponse(
-          { error: "Missing required fields: instanceToken, functionName" },
+          { success: false, error: "Missing required field: functionKey" },
           400,
         );
       }
 
-      // Step 1: Authorize the request (mutation -- validates token, checks permissions,
-      // updates counters, returns function handle)
-      const authResult = await ctx.runMutation(
-        component.gateway.authorizeRequest,
-        {
-          instanceToken,
-          functionName,
-          appName: config.appName,
-        },
-      );
+      const functionDef = normalizedConfig.functions[functionKey];
+      if (!functionDef) {
+        return jsonResponse(
+          { success: false, error: `Function "${functionKey}" not found` },
+          404,
+        );
+      }
 
-      if (!authResult.authorized) {
-        const detailSuffix =
-          authResult.matchedPattern && authResult.matchedPermission
-            ? ` (matchedPattern="${authResult.matchedPattern}", permission="${authResult.matchedPermission}")`
-            : "";
-        // Log the denied access
-        await ctx.runMutation(component.gateway.logAccess, {
-          agentId: authResult.agentId ?? "unknown",
-          appName: config.appName,
-          functionCalled: functionName,
-          permission: "deny",
-          errorMessage: authResult.error + detailSuffix,
-        });
+      const authResult = await ctx.runMutation(component.gateway.authorizeRequest, {
+        apiKey,
+        functionKey,
+        estimatedCost: body.estimatedCost,
+      });
 
-        return jsonResponse(
-          { error: authResult.error },
+      if (!authResult.authorized) {
+        const response = jsonResponse(
+          { success: false, error: authResult.error },
           authResult.statusCode,
         );
+        if (authResult.statusCode === 429) {
+          response.headers.set(
+            "Retry-After",
+            String(authResult.retryAfterSeconds ?? 3600),
+          );
+        }
+        return response;
       }
 
-      // Step 2: Execute the function via the registered handle
       const startTime = Date.now();
       try {
+        const args = body.args ?? {};
         let result: unknown;
-
-        switch (authResult.functionType) {
-          case "query":
-            result = await ctx.runQuery(
-              authResult.functionHandle as FunctionHandle<"query">,
-              args ?? {},
-            );
-            break;
-          case "mutation":
-            result = await ctx.runMutation(
-              authResult.functionHandle as FunctionHandle<"mutation">,
-              args ?? {},
-            );
-            break;
-          case "action":
-            result = await ctx.runAction(
-              authResult.functionHandle as FunctionHandle<"action">,
-              args ?? {},
-            );
-            break;
+        if (functionDef.type === "query") {
+          result = await ctx.runQuery(
+            functionDef.ref as FunctionHandle<"query">,
+            args,
+          );
+        } else if (functionDef.type === "mutation") {
+          result = await ctx.runMutation(
+            functionDef.ref as FunctionHandle<"mutation">,
+            args,
+          );
+        } else {
+          result = await ctx.runAction(
+            functionDef.ref as FunctionHandle<"action">,
+            args,
+          );
         }
 
-        const durationMs = Date.now() - startTime;
-
-        // Step 3: Log successful access
         await ctx.runMutation(component.gateway.logAccess, {
           agentId: authResult.agentId,
-          appName: config.appName,
-          functionCalled: functionName,
-          permission: "allow",
-          durationMs,
+          functionKey,
+          args,
+          result,
+          duration: Date.now() - startTime,
+          timestamp: Date.now(),
         });
 
-        return jsonResponse({ result, durationMs }, 200);
+        return jsonResponse({ success: true, result }, 200);
       } catch (error: unknown) {
-        const durationMs = Date.now() - startTime;
         const errorMessage =
-          error instanceof Error ? error.message : "Unknown error";
+          error && typeof error === "object" && "message" in error
+            ? (error.message as string)
+            : "Unknown error";
 
-        // Log the failed execution
         await ctx.runMutation(component.gateway.logAccess, {
           agentId: authResult.agentId,
-          appName: config.appName,
-          functionCalled: functionName,
-          permission: "allow",
-          errorMessage,
-          durationMs,
+          functionKey,
+          args: body.args ?? {},
+          error: errorMessage,
+          duration: Date.now() - startTime,
+          timestamp: Date.now(),
         });
 
-        return jsonResponse({ error: errorMessage }, 500);
+        return jsonResponse({ success: false, error: errorMessage }, 500);
       }
     }),
   });
 
-  // --- POST /agent-bridge/provision ---
-  // Agent self-provisioning endpoint
   http.route({
-    path: `${prefix}/provision`,
-    method: "POST",
-    handler: httpActionGeneric(async (ctx, request) => {
-      let body: { provisioningToken?: string };
-      try {
-        body = await request.json();
-      } catch {
-        return jsonResponse({ error: "Invalid JSON body" }, 400);
-      }
-
-      const { provisioningToken } = body;
+    path: `${prefix}/functions`,
+    method: "GET",
+    handler: httpActionGeneric(async () => {
+      const functions = availableFunctionKeys.map((functionKey) => ({
+        functionKey,
+        type: normalizedConfig.functions[functionKey].type,
+        metadata: normalizedConfig.functions[functionKey].metadata,
+      }));
+      return jsonResponse({ functions }, 200);
+    }),
+  });
+}
 
-      if (!provisioningToken) {
-        return jsonResponse(
-          { error: "Missing required field: provisioningToken" },
-          400,
-        );
-      }
+type PermissionRule = {
+  pattern: string;
+  permission: "allow" | "deny" | "rate_limited";
+  rateLimitConfig?: {
+    requestsPerHour: number;
+    tokenBudget?: number;
+  };
+};
 
-      try {
-        const result = await ctx.runMutation(
-          component.provisioning.provisionAgent,
-          {
-            provisioningToken,
-            appName: config.appName,
-          },
-        );
+type MutationCtx = Pick<GenericMutationCtx<GenericDataModel>, "runMutation">;
 
-        return jsonResponse(result, 200);
-      } catch (error: unknown) {
-        const errorMessage =
-          error instanceof Error ? error.message : "Provisioning failed";
-        return jsonResponse({ error: errorMessage }, 400);
-      }
-    }),
+export async function setAgentPermissions(
+  ctx: MutationCtx,
+  component: ComponentApi,
+  args: {
+    agentId: string;
+    rules: PermissionRule[];
+    config: AgentBridgeConfig;
+  },
+) {
+  const availableFunctionKeys = Object.keys(args.config.functions);
+  return await ctx.runMutation(component.permissions.setAgentPermissions, {
+    agentId: args.agentId,
+    rules: args.rules,
+    availableFunctionKeys,
   });
+}
 
-  // --- GET /agent-bridge/health ---
-  // Health check endpoint
-  http.route({
-    path: `${prefix}/health`,
-    method: "GET",
-    handler: httpActionGeneric(async (ctx) => {
-      try {
-        const functions = await ctx.runQuery(
-          component.registry.listFunctions,
-          { appName: config.appName },
-        );
-
-        return jsonResponse(
-          {
-            status: "ok",
-            appName: config.appName,
-            registeredFunctions: functions.length,
-            timestamp: Date.now(),
-          },
-          200,
-        );
-      } catch {
-        return jsonResponse(
-          {
-            status: "error",
-            appName: config.appName,
-            timestamp: Date.now(),
-          },
-          500,
-        );
-      }
-    }),
+export async function setFunctionOverrides(
+  ctx: MutationCtx,
+  component: ComponentApi,
+  args: {
+    overrides: Array<{
+      key: string;
+      enabled: boolean;
+      globalRateLimit?: number;
+    }>;
+    config: AgentBridgeConfig;
+  },
+) {
+  const availableFunctionKeys = Object.keys(args.config.functions);
+  return await ctx.runMutation(component.permissions.setFunctionOverrides, {
+    overrides: args.overrides,
+    availableFunctionKeys,
   });
 }
 
-// --- Helper ---
+export function listConfiguredFunctions(config: AgentBridgeConfig) {
+  const normalizedConfig = normalizeAgentBridgeConfig(config);
+  return Object.entries(normalizedConfig.functions).map(
+    ([functionKey, functionDef]) => ({
+      functionKey,
+      type: functionDef.type,
+      metadata: functionDef.metadata,
+    }),
+  );
+}
 
 function jsonResponse(data: unknown, status: number): Response {
   return new Response(JSON.stringify(data), {