Fix Composio discovery and auth checks

Author: jahooma

sdk/src/composio.ts

@@ -18,6 +18,16 @@ type ComposioExecuteResponse = {
   output: ToolResultOutput[]
 }
 
+const COMPOSIO_DISCOVERY_TIMEOUT_MS = 750
+const COMPOSIO_DISCOVERY_SUCCESS_CACHE_MS = 5 * 60 * 1000
+const COMPOSIO_DISCOVERY_FAILURE_CACHE_MS = 60 * 1000
+const COMPOSIO_DISCOVERY_CACHE_MAX_ENTRIES = 64
+
+const composioToolDefinitionsCache = new Map<
+  string,
+  { expiresAt: number; tools: CustomToolDefinition[] }
+>()
+
 function toJsonValue(value: unknown): JSONValue {
   try {
     return JSON.parse(JSON.stringify(value ?? null)) as JSONValue
@@ -26,6 +36,70 @@ function toJsonValue(value: unknown): JSONValue {
   }
 }
 
+function createTimeoutSignal(
+  parentSignal: AbortSignal | undefined,
+  timeoutMs: number,
+) {
+  const controller = new AbortController()
+  const timeout = setTimeout(() => {
+    controller.abort(new Error('Composio discovery timed out'))
+  }, timeoutMs)
+  const onAbort = () => controller.abort(parentSignal?.reason)
+
+  if (parentSignal?.aborted) {
+    onAbort()
+  } else {
+    parentSignal?.addEventListener('abort', onAbort, { once: true })
+  }
+
+  return {
+    signal: controller.signal,
+    cleanup: () => {
+      clearTimeout(timeout)
+      parentSignal?.removeEventListener('abort', onAbort)
+    },
+  }
+}
+
+function getCachedComposioTools(apiKey: string): CustomToolDefinition[] | null {
+  const cached = composioToolDefinitionsCache.get(apiKey)
+  if (!cached) return null
+  if (Date.now() >= cached.expiresAt) {
+    composioToolDefinitionsCache.delete(apiKey)
+    return null
+  }
+  return cached.tools
+}
+
+function pruneComposioToolsCache(now: number) {
+  for (const [apiKey, cached] of composioToolDefinitionsCache) {
+    if (now >= cached.expiresAt) {
+      composioToolDefinitionsCache.delete(apiKey)
+    }
+  }
+
+  while (
+    composioToolDefinitionsCache.size >= COMPOSIO_DISCOVERY_CACHE_MAX_ENTRIES
+  ) {
+    const oldest = composioToolDefinitionsCache.keys().next()
+    if (oldest.done) break
+    composioToolDefinitionsCache.delete(oldest.value)
+  }
+}
+
+function cacheComposioTools(
+  apiKey: string,
+  tools: CustomToolDefinition[],
+  ttlMs: number,
+) {
+  const now = Date.now()
+  pruneComposioToolsCache(now)
+  composioToolDefinitionsCache.set(apiKey, {
+    tools,
+    expiresAt: now + ttlMs,
+  })
+}
+
 async function readErrorMessage(response: Response): Promise<string> {
   try {
     const body = (await response.json()) as {
@@ -90,24 +164,41 @@ async function executeComposioToolViaServer(params: {
 export async function getComposioCustomToolDefinitions(params: {
   apiKey: string
   logger?: Pick<Logger, 'warn'>
+  signal?: AbortSignal
 }): Promise<CustomToolDefinition[]> {
+  const cachedTools = getCachedComposioTools(params.apiKey)
+  if (cachedTools) return cachedTools
+  if (params.signal?.aborted) return []
+
+  const discoverySignal = createTimeoutSignal(
+    params.signal,
+    COMPOSIO_DISCOVERY_TIMEOUT_MS,
+  )
+
   let response: Response
   try {
     response = await fetch(new URL('/api/v1/composio/tools', WEBSITE_URL), {
       method: 'POST',
       headers: {
         Authorization: `Bearer ${params.apiKey}`,
       },
+      signal: discoverySignal.signal,
     })
   } catch (error) {
-    params.logger?.warn(
-      { error: error instanceof Error ? error.message : String(error) },
-      'Failed to fetch Composio tools',
-    )
+    if (!params.signal?.aborted) {
+      cacheComposioTools(params.apiKey, [], COMPOSIO_DISCOVERY_FAILURE_CACHE_MS)
+      params.logger?.warn(
+        { error: error instanceof Error ? error.message : String(error) },
+        'Failed to fetch Composio tools',
+      )
+    }
     return []
+  } finally {
+    discoverySignal.cleanup()
   }
 
   if (!response.ok) {
+    cacheComposioTools(params.apiKey, [], COMPOSIO_DISCOVERY_FAILURE_CACHE_MS)
     if (response.status !== 503) {
       params.logger?.warn(
         { status: response.status, error: await readErrorMessage(response) },
@@ -119,13 +210,13 @@ export async function getComposioCustomToolDefinitions(params: {
 
   try {
     const body = (await response.json()) as ComposioToolsResponse
-    return body.tools.map((tool) => ({
+    const tools = body.tools.map((tool) => ({
       toolName: tool.toolName,
       inputSchema: tool.inputSchema,
       description: tool.description,
       endsAgentStep: true,
       exampleInputs: [],
-      execute: async (input) => {
+      execute: async (input: unknown) => {
         return executeComposioToolViaServer({
           apiKey: params.apiKey,
           sessionId: body.sessionId,
@@ -137,7 +228,14 @@ export async function getComposioCustomToolDefinitions(params: {
         })
       },
     }))
+    cacheComposioTools(
+      params.apiKey,
+      tools,
+      COMPOSIO_DISCOVERY_SUCCESS_CACHE_MS,
+    )
+    return tools
   } catch (error) {
+    cacheComposioTools(params.apiKey, [], COMPOSIO_DISCOVERY_FAILURE_CACHE_MS)
     params.logger?.warn(
       { error: error instanceof Error ? error.message : String(error) },
       'Failed to parse Composio tools response',

sdk/src/run.ts

@@ -516,9 +516,14 @@ async function runOnce({
   }
   const userId = userInfo.id
 
+  if (signal?.aborted) {
+    return getCancelledRunState('Run cancelled by user.')
+  }
+
   const composioCustomToolDefinitions = await getComposioCustomToolDefinitions({
     apiKey,
     logger,
+    signal,
   })
 
   for (const toolName of COMPOSIO_META_TOOL_NAMES) {

web/src/app/api/v1/composio/__tests__/composio.test.ts

@@ -41,11 +41,20 @@ describe('/api/v1/composio', () => {
     }
     loggerWithContext = mock(() => logger)
     getUserInfoFromApiKey = mock(async ({ apiKey }) => {
+      if (apiKey === 'banned-key') {
+        return {
+          id: 'banned-user',
+          email: 'banned@example.com',
+          discord_id: null,
+          banned: true,
+        } as Awaited<ReturnType<GetUserInfoFromApiKeyFn>>
+      }
       if (apiKey !== 'valid-key') return null
       return {
         id: 'user-123',
         email: 'user@example.com',
         discord_id: null,
+        banned: false,
       } as Awaited<ReturnType<GetUserInfoFromApiKeyFn>>
     })
   })
@@ -291,4 +300,33 @@ describe('/api/v1/composio', () => {
       error: 'Missing or invalid Authorization header',
     })
   })
+
+  test('rejects suspended users before rate limiting or tool lookup', async () => {
+    const getToolsForUser = mock(async () => ({
+      sessionId: 'session-123',
+      tools: [],
+    }))
+    const checkRateLimit = mock(() => ({ limited: false as const }))
+    const req = new NextRequest('http://localhost/api/v1/composio/tools', {
+      method: 'POST',
+      headers: { Authorization: 'Bearer banned-key' },
+    })
+
+    const response = await postComposioTools({
+      req,
+      getUserInfoFromApiKey,
+      db: mockDb,
+      logger,
+      loggerWithContext,
+      getToolsForUser,
+      checkRateLimit,
+    })
+
+    expect(response.status).toBe(403)
+    const body = await response.json()
+    expect(body.error).toBe('account_suspended')
+    expect(body.message).toContain('Your account has been suspended')
+    expect(getToolsForUser).not.toHaveBeenCalled()
+    expect(checkRateLimit).not.toHaveBeenCalled()
+  })
 })

web/src/app/api/v1/composio/_auth.ts

@@ -1,4 +1,5 @@
 import { NextResponse } from 'next/server'
+import { env } from '@codebuff/internal/env'
 
 import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type {
@@ -13,6 +14,7 @@ type ComposioUser = {
   id: string
   email: string
   discord_id: string | null
+  banned: boolean
 }
 
 export async function requireComposioUser(params: {
@@ -39,7 +41,7 @@ export async function requireComposioUser(params: {
 
   const userInfo = await getUserInfoFromApiKey({
     apiKey,
-    fields: ['id', 'email', 'discord_id'],
+    fields: ['id', 'email', 'discord_id', 'banned'],
     logger,
   })
   if (!userInfo) {
@@ -52,6 +54,19 @@ export async function requireComposioUser(params: {
     }
   }
 
+  if (userInfo.banned) {
+    return {
+      ok: false,
+      response: NextResponse.json(
+        {
+          error: 'account_suspended',
+          message: `Your account has been suspended. Please contact ${env.NEXT_PUBLIC_SUPPORT_EMAIL} if you did not expect this.`,
+        },
+        { status: 403 },
+      ),
+    }
+  }
+
   return {
     ok: true,
     userInfo,