Fix Composio session race and timeout cache

Author: jahooma

sdk/src/__tests__/composio.test.ts

@@ -0,0 +1,66 @@
+import { afterEach, describe, expect, mock, test } from 'bun:test'
+
+import { getComposioCustomToolDefinitions } from '../composio'
+
+describe('getComposioCustomToolDefinitions', () => {
+  const originalFetch = globalThis.fetch
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch
+  })
+
+  test('does not cache an empty tool list after discovery timeout', async () => {
+    const apiKey = `timeout-key-${Date.now()}`
+    const timeoutFetch = mock(
+      async (_url: string | URL | Request, init?: RequestInit) => {
+        const signal = init?.signal
+        return new Promise<Response>((_resolve, reject) => {
+          if (signal?.aborted) {
+            reject(new Error('aborted'))
+            return
+          }
+
+          signal?.addEventListener(
+            'abort',
+            () => reject(new Error('aborted')),
+            { once: true },
+          )
+        })
+      },
+    )
+    globalThis.fetch = timeoutFetch as unknown as typeof fetch
+
+    const timedOutTools = await getComposioCustomToolDefinitions({
+      apiKey,
+      logger: { warn: mock(() => {}) },
+    })
+    expect(timedOutTools).toEqual([])
+    expect(timeoutFetch).toHaveBeenCalledTimes(1)
+
+    const successFetch = mock(async () => {
+      return new Response(
+        JSON.stringify({
+          sessionId: 'session-123',
+          tools: [
+            {
+              toolName: 'COMPOSIO_SEARCH_TOOLS',
+              inputSchema: { type: 'object', properties: {} },
+              description: 'Search tools',
+            },
+          ],
+        }),
+        { status: 200 },
+      )
+    })
+    globalThis.fetch = successFetch as unknown as typeof fetch
+
+    const tools = await getComposioCustomToolDefinitions({
+      apiKey,
+      logger: { warn: mock(() => {}) },
+    })
+
+    expect(successFetch).toHaveBeenCalledTimes(1)
+    expect(tools).toHaveLength(1)
+    expect(tools[0]?.toolName).toBe('COMPOSIO_SEARCH_TOOLS')
+  })
+})

sdk/src/composio.ts

@@ -185,13 +185,23 @@ export async function getComposioCustomToolDefinitions(params: {
       signal: discoverySignal.signal,
     })
   } catch (error) {
-    if (!params.signal?.aborted) {
-      cacheComposioTools(params.apiKey, [], COMPOSIO_DISCOVERY_FAILURE_CACHE_MS)
+    if (params.signal?.aborted) {
+      return []
+    }
+
+    if (discoverySignal.signal.aborted) {
       params.logger?.warn(
         { error: error instanceof Error ? error.message : String(error) },
-        'Failed to fetch Composio tools',
+        'Timed out fetching Composio tools',
       )
+      return []
     }
+
+    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()

web/src/server/__tests__/composio.test.ts

@@ -45,19 +45,27 @@ describe('getComposioToolsForUser', () => {
     ])
   })
 
-  function makeDb(storedSessionId: string | null) {
-    const findFirst = mock(async () =>
-      storedSessionId
+  function makeDb(storedSessionIds: string | null | Array<string | null>) {
+    const storedSessionIdSequence = Array.isArray(storedSessionIds)
+      ? [...storedSessionIds]
+      : [storedSessionIds]
+    const findFirst = mock(async () => {
+      const storedSessionId =
+        storedSessionIdSequence.length > 1
+          ? storedSessionIdSequence.shift()
+          : storedSessionIdSequence[0]
+
+      return storedSessionId
         ? {
             user_id: 'user-123',
             session_id: storedSessionId,
             created_at: new Date(),
             updated_at: new Date(),
           }
-        : null,
-    )
-    const onConflictDoUpdate = mock(async () => undefined)
-    const values = mock(() => ({ onConflictDoUpdate }))
+        : null
+    })
+    const onConflictDoNothing = mock(async () => undefined)
+    const values = mock(() => ({ onConflictDoNothing }))
     const whereDelete = mock(async () => undefined)
 
     return {
@@ -71,7 +79,7 @@ describe('getComposioToolsForUser', () => {
         delete: mock(() => ({ where: whereDelete })),
       } as any,
       findFirst,
-      onConflictDoUpdate,
+      onConflictDoNothing,
       values,
       whereDelete,
     }
@@ -84,7 +92,10 @@ describe('getComposioToolsForUser', () => {
     useSession = mock(async () => {
       throw notFound
     })
-    const { db, whereDelete, values } = makeDb('stored-session')
+    const { db, whereDelete, values } = makeDb([
+      'stored-session',
+      'fresh-session',
+    ])
 
     const result = await getComposioToolsForUser({
       db,
@@ -112,6 +123,32 @@ describe('getComposioToolsForUser', () => {
     })
   })
 
+  test('returns the persisted session when concurrent creation stores a different session', async () => {
+    createSession = mock(async () => ({ sessionId: 'losing-session' }))
+    useSession = mock(async () => ({ sessionId: 'winning-session' }))
+    const { db, values, onConflictDoNothing } = makeDb([
+      null,
+      'winning-session',
+    ])
+
+    const result = await getComposioToolsForUser({
+      db,
+      userId: 'user-123',
+      logger,
+      apiKey: 'test-composio-api-key',
+    })
+
+    expect(result?.sessionId).toBe('winning-session')
+    expect(createSession).toHaveBeenCalledWith('user-123')
+    expect(values).toHaveBeenCalledWith({
+      user_id: 'user-123',
+      session_id: 'losing-session',
+    })
+    expect(onConflictDoNothing).toHaveBeenCalledTimes(1)
+    expect(useSession).toHaveBeenCalledWith('winning-session')
+    expect(getRawToolRouterSessionTools).toHaveBeenCalledWith('winning-session')
+  })
+
   test('keeps the stored session row when rehydration fails transiently', async () => {
     const transientError = Object.assign(new Error('Composio unavailable'), {
       status: 502,

web/src/server/composio.ts

@@ -121,7 +121,7 @@ async function getToolDefinitionsForSession(params: {
     }))
 }
 
-async function saveSession(params: {
+async function insertSessionIfAbsent(params: {
   db: CodebuffPgDatabase
   userId: string
   sessionId: string
@@ -132,12 +132,8 @@ async function saveSession(params: {
       user_id: params.userId,
       session_id: params.sessionId,
     })
-    .onConflictDoUpdate({
+    .onConflictDoNothing({
       target: schema.composioSession.user_id,
-      set: {
-        session_id: params.sessionId,
-        updated_at: new Date(),
-      },
     })
 }
 
@@ -182,15 +178,41 @@ async function createSessionForUser(params: {
   db: CodebuffPgDatabase
   userId: string
   apiKey: string
+  logger: Logger
 }): Promise<CachedComposioSession> {
   const composio = getComposioClient(params.apiKey)
   const session = await composio.create(params.userId)
-  await saveSession({
+  await insertSessionIfAbsent({
     db: params.db,
     userId: params.userId,
     sessionId: session.sessionId,
   })
 
+  const storedSession = await getStoredSessionByUser({
+    db: params.db,
+    userId: params.userId,
+  })
+  if (!storedSession) {
+    throw new Error('Failed to persist Composio session')
+  }
+
+  if (storedSession.session_id !== session.sessionId) {
+    params.logger.info(
+      {
+        userId: params.userId,
+        createdSessionId: session.sessionId,
+        storedSessionId: storedSession.session_id,
+      },
+      'Using existing persisted Composio session after concurrent creation',
+    )
+    return rehydrateSession({
+      userId: params.userId,
+      sessionId: storedSession.session_id,
+      apiKey: params.apiKey,
+      includeTools: true,
+    })
+  }
+
   const cachedSession: CachedComposioSession = {
     userId: params.userId,
     sessionId: session.sessionId,
@@ -322,6 +344,7 @@ async function getSessionForUser(params: {
       db: params.db,
       userId: params.userId,
       apiKey,
+      logger: params.logger,
     })
   } catch (error) {
     params.logger.error(