@okrlinkhub/agent-bridge 0.1.0 → 2.0.0

package/src/component/gateway.ts

@@ -1,59 +1,22 @@
 import { v } from "convex/values";
 import { mutation, query } from "./_generated/server.js";
-
-// --- Token hashing utility (same as provisioning.ts) ---
-
-async function hashToken(token: string): Promise<string> {
-  const encoder = new TextEncoder();
-  const data = encoder.encode(token);
-  const hash = await crypto.subtle.digest("SHA-256", data);
-  return Array.from(new Uint8Array(hash))
-    .map((b) => b.toString(16).padStart(2, "0"))
-    .join("");
-}
-
-// --- Pattern matching (same logic as permissions.ts) ---
-
-function patternSpecificity(pattern: string): number {
-  const wildcardIndex = pattern.indexOf("*");
-  if (wildcardIndex === -1) return pattern.length;
-  return wildcardIndex;
-}
-
-function matchesPattern(functionName: string, pattern: string): boolean {
-  if (pattern === "*") return true;
-  const escaped = pattern.replace(/[.+?^${}()|[\]\\]/g, "\\$&");
-  const regexStr = "^" + escaped.replace(/\*/g, ".*") + "$";
-  return new RegExp(regexStr).test(functionName);
-}
-
-// --- Authorize request result validator ---
+import {
+  findBestPermissionMatch,
+  hashApiKey,
+  type PermissionType,
+} from "./agentBridgeUtils.js";
 
 const authorizeResultValidator = v.union(
   v.object({
     authorized: v.literal(true),
-    agentId: v.string(),
-    appName: v.string(),
-    functionHandle: v.string(),
-    functionType: v.union(
-      v.literal("query"),
-      v.literal("mutation"),
-      v.literal("action"),
-    ),
+    agentId: v.id("agents"),
   }),
   v.object({
     authorized: v.literal(false),
     error: v.string(),
     statusCode: v.number(),
-    agentId: v.optional(v.string()),
-    matchedPattern: v.optional(v.string()),
-    matchedPermission: v.optional(
-      v.union(
-        v.literal("allow"),
-        v.literal("deny"),
-        v.literal("rate_limited"),
-      ),
-    ),
+    agentId: v.optional(v.id("agents")),
+    retryAfterSeconds: v.optional(v.number()),
   }),
 );
 
@@ -62,132 +25,121 @@ const authorizeResultValidator = v.union(
  * This is a mutation (not a query) because it updates counters and last activity.
  *
  * Steps:
- * 1. Validate instance token
+ * 1. Validate API key
  * 2. Check agent is active
  * 3. Check function permissions
- * 4. Look up function handle in registry
- * 5. Update activity counters
+ * 4. Check function global override
+ * 5. Check rate limits
  *
- * Returns the function handle info if authorized, or an error.
+ * Returns the agent id if authorized, or an error.
  */
 export const authorizeRequest = mutation({
   args: {
-    instanceToken: v.string(),
-    functionName: v.string(),
-    appName: v.string(),
+    apiKey: v.string(),
+    functionKey: v.string(),
+    estimatedCost: v.optional(v.number()),
   },
   returns: authorizeResultValidator,
   handler: async (ctx, args) => {
-    // 1. Validate instance token
-    const tokenHash = await hashToken(args.instanceToken);
-    const instance = await ctx.db
-      .query("agentAppInstances")
-      .withIndex("by_instance_token_hash", (q) =>
-        q.eq("instanceTokenHash", tokenHash),
-      )
+    const apiKeyHash = await hashApiKey(args.apiKey);
+    const agent = await ctx.db
+      .query("agents")
+      .withIndex("by_apiKeyHash", (q) => q.eq("apiKeyHash", apiKeyHash))
       .unique();
-
-    if (!instance) {
+    if (!agent) {
       return {
         authorized: false as const,
-        error: "Invalid instance token",
+        error: "Invalid API key",
         statusCode: 401,
       };
     }
 
-    if (instance.expiresAt < Date.now()) {
+    if (!agent.enabled) {
       return {
         authorized: false as const,
-        error: "Instance token has expired",
-        statusCode: 401,
+        error: "Agent disabled",
+        statusCode: 403,
+        agentId: agent._id,
       };
     }
 
-    if (instance.appName !== args.appName) {
+    const permissions = await ctx.db
+      .query("agentPermissions")
+      .withIndex("by_agentId", (q) => q.eq("agentId", agent._id))
+      .collect();
+    const matchedRule = findBestPermissionMatch(args.functionKey, permissions);
+    if (!matchedRule || matchedRule.permission === "deny") {
       return {
         authorized: false as const,
-        error: "Token does not match this app",
+        error: `Function ${args.functionKey} not allowed`,
         statusCode: 403,
+        agentId: agent._id,
       };
     }
 
-    const agentId = instance.agentId;
-
-    // 2. Check agent is active
-    const agent = await ctx.db
-      .query("registeredAgents")
-      .withIndex("by_agent_id", (q) => q.eq("agentId", agentId))
+    const functionOverride = await ctx.db
+      .query("agentFunctions")
+      .withIndex("by_key", (q) => q.eq("key", args.functionKey))
       .unique();
-
-    if (!agent || !agent.isActive) {
+    if (functionOverride && !functionOverride.enabled) {
       return {
         authorized: false as const,
-        error: "Agent has been revoked",
+        error: `Function ${args.functionKey} disabled`,
         statusCode: 403,
-        agentId,
+        agentId: agent._id,
       };
     }
 
-    // 3. Check permissions
-    const permissions = await ctx.db
-      .query("functionPermissions")
-      .withIndex("by_agent_and_app", (q) =>
-        q.eq("agentId", agentId).eq("appName", args.appName),
-      )
+    const effectiveHourlyLimit = resolveEffectiveHourlyLimit(
+      agent.rateLimit,
+      matchedRule.permission,
+      matchedRule.rateLimitConfig?.requestsPerHour,
+      functionOverride?.globalRateLimit,
+    );
+    const oneHourAgo = Date.now() - 60 * 60 * 1000;
+    const recentLogs = await ctx.db
+      .query("agentLogs")
+      .withIndex("by_agentId_and_timestamp", (q) => q.eq("agentId", agent._id))
       .collect();
-
-    const matches = permissions
-      .filter((p) => matchesPattern(args.functionName, p.functionPattern))
-      .sort(
-        (a, b) =>
-          patternSpecificity(b.functionPattern) -
-          patternSpecificity(a.functionPattern),
-      );
-
-    if (matches.length === 0 || matches[0].permission === "deny") {
-      const bestMatch = matches[0];
+    const recentCallCount = recentLogs.filter(
+      (log) => log.timestamp >= oneHourAgo,
+    ).length;
+    if (recentCallCount >= effectiveHourlyLimit) {
       return {
         authorized: false as const,
-        error: "Function not authorized for this agent",
-        statusCode: 403,
-        agentId,
-        matchedPattern: bestMatch?.functionPattern,
-        matchedPermission: bestMatch?.permission,
+        error: "Rate limit exceeded",
+        statusCode: 429,
+        retryAfterSeconds: 3600,
+        agentId: agent._id,
       };
     }
 
-    // TODO: For rate_limited permissions, check rate counters here
-    // (circuit breaker - deferred to post-MVP)
-
-    // 4. Look up function handle
-    const fnEntry = await ctx.db
-      .query("functionRegistry")
-      .withIndex("by_app_and_function", (q) =>
-        q.eq("appName", args.appName).eq("functionName", args.functionName),
-      )
-      .unique();
-
-    if (!fnEntry) {
-      return {
-        authorized: false as const,
-        error: `Function "${args.functionName}" is not registered`,
-        statusCode: 404,
-        agentId,
-      };
+    if (
+      matchedRule.permission === "rate_limited" &&
+      matchedRule.rateLimitConfig?.tokenBudget !== undefined
+    ) {
+      const estimatedCost = args.estimatedCost ?? 0;
+      const tokenEstimate = recentLogs
+        .filter((log) => log.timestamp >= oneHourAgo)
+        .reduce((sum, log) => sum + estimateCostFromLog(log.args), 0);
+      if (tokenEstimate + estimatedCost > matchedRule.rateLimitConfig.tokenBudget) {
+        return {
+          authorized: false as const,
+          error: "Token budget exceeded",
+          statusCode: 429,
+          retryAfterSeconds: 3600,
+          agentId: agent._id,
+        };
+      }
     }
 
-    // 5. Update activity counters
-    await ctx.db.patch(instance._id, {
-      lastActivityAt: Date.now(),
-      monthlyRequests: instance.monthlyRequests + 1,
+    await ctx.db.patch(agent._id, {
+      lastUsed: Date.now(),
     });
 
     return {
       authorized: true as const,
-      agentId,
-      appName: args.appName,
-      functionHandle: fnEntry.functionHandle,
-      functionType: fnEntry.functionType,
+      agentId: agent._id,
     };
   },
 });
@@ -198,23 +150,24 @@ export const authorizeRequest = mutation({
  */
 export const logAccess = mutation({
   args: {
-    agentId: v.string(),
-    appName: v.string(),
-    functionCalled: v.string(),
-    permission: v.string(),
-    errorMessage: v.optional(v.string()),
-    durationMs: v.optional(v.number()),
+    agentId: v.id("agents"),
+    functionKey: v.string(),
+    args: v.any(),
+    result: v.optional(v.any()),
+    error: v.optional(v.string()),
+    duration: v.number(),
+    timestamp: v.number(),
   },
   returns: v.null(),
   handler: async (ctx, args) => {
-    await ctx.db.insert("accessLog", {
-      timestamp: Date.now(),
+    await ctx.db.insert("agentLogs", {
+      timestamp: args.timestamp,
       agentId: args.agentId,
-      appName: args.appName,
-      functionCalled: args.functionCalled,
-      permission: args.permission,
-      errorMessage: args.errorMessage,
-      durationMs: args.durationMs,
+      functionKey: args.functionKey,
+      args: args.args,
+      result: args.result,
+      error: args.error,
+      duration: args.duration,
     });
     return null;
   },
@@ -225,58 +178,91 @@ export const logAccess = mutation({
  */
 export const queryAccessLog = query({
   args: {
-    agentId: v.optional(v.string()),
-    appName: v.optional(v.string()),
+    agentId: v.optional(v.id("agents")),
+    functionKey: v.optional(v.string()),
     limit: v.optional(v.number()),
   },
   returns: v.array(
     v.object({
+      _id: v.id("agentLogs"),
       timestamp: v.number(),
-      agentId: v.string(),
-      appName: v.string(),
-      functionCalled: v.string(),
-      permission: v.string(),
-      errorMessage: v.optional(v.string()),
-      durationMs: v.optional(v.number()),
+      agentId: v.id("agents"),
+      functionKey: v.string(),
+      args: v.any(),
+      result: v.optional(v.any()),
+      error: v.optional(v.string()),
+      duration: v.number(),
     }),
   ),
   handler: async (ctx, args) => {
     const limit = args.limit ?? 50;
 
-    if (args.agentId) {
+    const agentId = args.agentId;
+    if (agentId !== undefined) {
       const logs = await ctx.db
-        .query("accessLog")
-        .withIndex("by_agent_and_timestamp", (q) =>
-          q.eq("agentId", args.agentId!),
+        .query("agentLogs")
+        .withIndex("by_agentId_and_timestamp", (q) =>
+          q.eq("agentId", agentId),
         )
         .order("desc")
         .take(limit);
 
       return logs.map((l) => ({
+        _id: l._id,
         timestamp: l.timestamp,
         agentId: l.agentId,
-        appName: l.appName,
-        functionCalled: l.functionCalled,
-        permission: l.permission,
-        errorMessage: l.errorMessage,
-        durationMs: l.durationMs,
+        functionKey: l.functionKey,
+        args: l.args,
+        result: l.result,
+        error: l.error,
+        duration: l.duration,
       }));
     }
 
-    // No filter: get recent logs
     const logs = await ctx.db
-      .query("accessLog")
+      .query("agentLogs")
       .order("desc")
       .take(limit);
+    const filteredLogs =
+      args.functionKey !== undefined
+        ? logs.filter((log) => log.functionKey === args.functionKey)
+        : logs;
 
-    return logs.map((l) => ({
+    return filteredLogs.map((l) => ({
+      _id: l._id,
       timestamp: l.timestamp,
       agentId: l.agentId,
-      appName: l.appName,
-      functionCalled: l.functionCalled,
-      permission: l.permission,
-      errorMessage: l.errorMessage,
-      durationMs: l.durationMs,
+      functionKey: l.functionKey,
+      args: l.args,
+      result: l.result,
+      error: l.error,
+      duration: l.duration,
     }));
   },
 });
+
+function resolveEffectiveHourlyLimit(
+  baseAgentLimit: number,
+  permissionType: PermissionType,
+  permissionLimit?: number,
+  globalLimit?: number,
+) {
+  let effective = baseAgentLimit;
+  if (permissionType === "rate_limited" && permissionLimit !== undefined) {
+    effective = Math.min(effective, permissionLimit);
+  }
+  if (globalLimit !== undefined) {
+    effective = Math.min(effective, globalLimit);
+  }
+  return effective;
+}
+
+function estimateCostFromLog(args: unknown): number {
+  if (!args || typeof args !== "object") {
+    return 0;
+  }
+  if ("estimatedCost" in args && typeof args.estimatedCost === "number") {
+    return args.estimatedCost;
+  }
+  return 0;
+}