Recover invalid Composio sessions

Author: jahooma

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

@@ -0,0 +1,137 @@
+import { beforeAll, beforeEach, describe, expect, mock, test } from 'bun:test'
+
+import type { Logger } from '@codebuff/common/types/contracts/logger'
+import type { getComposioToolsForUser as GetComposioToolsForUser } from '../composio'
+
+let getComposioToolsForUser: typeof GetComposioToolsForUser
+
+let createSession: ReturnType<typeof mock>
+let useSession: ReturnType<typeof mock>
+let getRawToolRouterSessionTools: ReturnType<typeof mock>
+
+beforeAll(async () => {
+  mock.module('server-only', () => ({}))
+  mock.module('@codebuff/internal/env', () => ({
+    env: { COMPOSIO_API_KEY: 'test-composio-api-key' },
+  }))
+  mock.module('@composio/core', () => ({
+    Composio: class {
+      tools = {
+        getRawToolRouterSessionTools,
+      }
+
+      create = createSession
+      use = useSession
+    },
+  }))
+  ;({ getComposioToolsForUser } = await import('../composio'))
+})
+
+describe('getComposioToolsForUser', () => {
+  let logger: Logger
+
+  beforeEach(() => {
+    logger = {
+      error: mock(() => {}),
+      warn: mock(() => {}),
+      info: mock(() => {}),
+      debug: mock(() => {}),
+    }
+    createSession = mock(async () => ({ sessionId: 'fresh-session' }))
+    useSession = mock(async () => ({ sessionId: 'stored-session' }))
+    getRawToolRouterSessionTools = mock(async () => [
+      {
+        slug: 'COMPOSIO_SEARCH_TOOLS',
+        inputParameters: { type: 'object', properties: {} },
+        description: 'Search tools',
+      },
+    ])
+  })
+
+  function makeDb(storedSessionId: string | null) {
+    const findFirst = mock(async () =>
+      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 }))
+    const whereDelete = mock(async () => undefined)
+
+    return {
+      db: {
+        query: {
+          composioSession: {
+            findFirst,
+          },
+        },
+        insert: mock(() => ({ values })),
+        delete: mock(() => ({ where: whereDelete })),
+      } as any,
+      findFirst,
+      onConflictDoUpdate,
+      values,
+      whereDelete,
+    }
+  }
+
+  test('replaces a stored session when Composio can no longer rehydrate it', async () => {
+    const notFound = Object.assign(new Error('Composio session not found'), {
+      status: 404,
+    })
+    useSession = mock(async () => {
+      throw notFound
+    })
+    const { db, whereDelete, values } = makeDb('stored-session')
+
+    const result = await getComposioToolsForUser({
+      db,
+      userId: 'user-123',
+      logger,
+    })
+
+    expect(result).toEqual({
+      sessionId: 'fresh-session',
+      tools: [
+        {
+          toolName: 'COMPOSIO_SEARCH_TOOLS',
+          inputSchema: { type: 'object', properties: {} },
+          description: 'Search tools',
+        },
+      ],
+    })
+    expect(useSession).toHaveBeenCalledWith('stored-session')
+    expect(whereDelete).toHaveBeenCalledTimes(1)
+    expect(createSession).toHaveBeenCalledWith('user-123')
+    expect(values).toHaveBeenCalledWith({
+      user_id: 'user-123',
+      session_id: 'fresh-session',
+    })
+  })
+
+  test('keeps the stored session row when rehydration fails transiently', async () => {
+    const transientError = Object.assign(new Error('Composio unavailable'), {
+      status: 502,
+    })
+    useSession = mock(async () => {
+      throw transientError
+    })
+    const { db, whereDelete } = makeDb('stored-session')
+
+    await expect(
+      getComposioToolsForUser({
+        db,
+        userId: 'user-123',
+        logger,
+      }),
+    ).rejects.toThrow('Composio unavailable')
+
+    expect(whereDelete).not.toHaveBeenCalled()
+    expect(createSession).not.toHaveBeenCalled()
+  })
+})

web/src/server/composio.ts

@@ -163,6 +163,21 @@ async function getStoredSessionById(params: {
   })
 }
 
+async function deleteStoredSession(params: {
+  db: CodebuffPgDatabase
+  userId: string
+  sessionId: string
+}) {
+  await params.db
+    .delete(schema.composioSession)
+    .where(
+      and(
+        eq(schema.composioSession.user_id, params.userId),
+        eq(schema.composioSession.session_id, params.sessionId),
+      ),
+    )
+}
+
 async function createSessionForUser(params: {
   db: CodebuffPgDatabase
   userId: string
@@ -188,6 +203,48 @@ async function createSessionForUser(params: {
   return cachedSession
 }
 
+function getErrorStatus(error: unknown): number | undefined {
+  if (!error || typeof error !== 'object') return undefined
+
+  const candidates = [
+    'status',
+    'statusCode',
+    'code',
+    'responseStatus',
+    'httpStatus',
+  ]
+  for (const key of candidates) {
+    const value = (error as Record<string, unknown>)[key]
+    if (typeof value === 'number') return value
+    if (typeof value === 'string' && /^\d+$/.test(value)) {
+      return Number(value)
+    }
+  }
+
+  const response = (error as Record<string, unknown>)['response']
+  return getErrorStatus(response)
+}
+
+function isInvalidStoredSessionError(error: unknown): boolean {
+  const status = getErrorStatus(error)
+  if (status && [400, 401, 403, 404, 410].includes(status)) {
+    return true
+  }
+
+  if (!(error instanceof Error)) return false
+
+  const message = error.message.toLowerCase()
+  return (
+    message.includes('session') &&
+    (message.includes('not found') ||
+      message.includes('not exist') ||
+      message.includes('invalid') ||
+      message.includes('expired') ||
+      message.includes('unauthorized') ||
+      message.includes('forbidden'))
+  )
+}
+
 async function rehydrateSession(params: {
   userId: string
   sessionId: string
@@ -228,12 +285,32 @@ async function getSessionForUser(params: {
         { userId: params.userId },
         'Rehydrating Composio session from database',
       )
-      return rehydrateSession({
-        userId: params.userId,
-        sessionId: storedSession.session_id,
-        apiKey,
-        includeTools: true,
-      })
+      try {
+        return await rehydrateSession({
+          userId: params.userId,
+          sessionId: storedSession.session_id,
+          apiKey,
+          includeTools: true,
+        })
+      } catch (error) {
+        if (!isInvalidStoredSessionError(error)) {
+          throw error
+        }
+
+        params.logger.warn(
+          {
+            error: getErrorObject(error),
+            userId: params.userId,
+            sessionId: storedSession.session_id,
+          },
+          'Stored Composio session is invalid; replacing it',
+        )
+        await deleteStoredSession({
+          db: params.db,
+          userId: params.userId,
+          sessionId: storedSession.session_id,
+        })
+      }
     }
 
     params.logger.info(