initial openrouter wrapper

charleslien · charleslien · commit e4e66a0667ac · 2025-10-01T18:18:08.000-07:00
diff --git a/backend/src/llm-apis/openrouter.ts b/backend/src/llm-apis/openrouter.ts
@@ -2,8 +2,11 @@ import { models } from '@codebuff/common/old-constants'
 import { isExplicitlyDefinedModel } from '@codebuff/common/util/model-utils'
 import { env } from '@codebuff/internal/env'
 import { createOpenRouter } from '@codebuff/internal/openrouter-ai-sdk'
+import { cloneDeep } from 'lodash'
+import z from 'zod/v4'
 
 import type { Model } from '@codebuff/common/old-constants'
+import type { Request, Response } from 'express'
 
 // Provider routing documentation: https://openrouter.ai/docs/features/provider-routing
 const providerOrder = {
@@ -45,3 +48,178 @@ export function openRouterLanguageModel(model: Model) {
     logprobs: true,
   })
 }
+
+const openrouterUsageSchema = z
+  .object({
+    prompt_tokens: z.number(),
+    prompt_tokens_details: z
+      .object({
+        cached_tokens: z.number(),
+      })
+      .nullish(),
+    completion_tokens: z.number(),
+    completion_tokens_details: z
+      .object({
+        reasoning_tokens: z.number(),
+      })
+      .nullish(),
+    total_tokens: z.number(),
+    cost: z.number().optional(),
+    cost_details: z
+      .object({
+        upstream_inference_cost: z.number().nullish(),
+      })
+      .nullish(),
+  })
+  .nullish()
+
+export async function handleOpenrouterStream({
+  req,
+  res,
+  userId,
+}: {
+  req: Request
+  res: Response
+  userId: string
+}) {
+  res.writeHead(200, {
+    // Mandatory SSE headers
+    'Content-Type': 'text/event-stream',
+    'Cache-Control': 'no-cache',
+    Connection: 'keep-alive',
+    // (optional) allow local browser demos
+    'Access-Control-Allow-Origin': '*',
+  })
+
+  res.write(`: connected ${new Date().toISOString()}\n`)
+  const heartbeat = setInterval(() => {
+    res.write(`: heartbeat ${new Date().toISOString()}\n\n`)
+  }, 30000)
+  res.on('close', () => {
+    clearInterval(heartbeat)
+  })
+
+  const body = cloneDeep(req.body)
+  if (body.usage === undefined) {
+    body.usage = {}
+  }
+  body.usage.include = true
+  const response = await fetch(
+    'https://openrouter.ai/api/v1/chat/completions',
+    {
+      method: 'POST',
+      headers: {
+        Authorization: `Bearer ${process.env.OPEN_ROUTER_API_KEY}`,
+        'HTTP-Referer': 'https://codebuff.com',
+        'X-Title': 'Codebuff',
+        'Content-Type': 'application/json',
+      },
+      body: JSON.stringify(body),
+    },
+  )
+
+  const reader = response.body?.getReader()
+  if (!reader) {
+    res.status(500).json({ message: 'Failed to get response reader' })
+    return
+  }
+
+  const decoder = new TextDecoder()
+  let buffer = ''
+  try {
+    while (true) {
+      const { done, value } = await reader.read()
+      console.log('asdf', {
+        done,
+        value: decoder.decode(value, { stream: true }),
+      })
+      if (done) {
+        break
+      }
+
+      buffer += decoder.decode(value, { stream: true })
+      let lineEnd = buffer.indexOf('\n')
+      while (lineEnd !== -1) {
+        const line = buffer.slice(0, lineEnd + 1)
+        buffer = buffer.slice(lineEnd + 1)
+        // if (line.startsWith('data: ')) {
+        //   const data = line.trim().slice('data: '.length)
+        //   await processData(data, userId)
+        // }
+        res.write(line)
+        lineEnd = buffer.indexOf('\n')
+      }
+    }
+  } finally {
+    reader.cancel()
+  }
+  res.end()
+}
+
+/*
+
+async function processData(data: string, userId: string) {
+  if (data === '[DONE]') {
+    return
+  }
+
+  let obj
+  try {
+    obj = JSON.parse(data)
+  } catch (error) {
+    trackEvent(
+      AnalyticsEvent.OPENROUTER_MALFORMED_JSON_RESPONSE_CHUNK,
+      userId,
+      {
+        data,
+      },
+    )
+    return
+  }
+
+  if (typeof obj !== 'object') {
+    return
+  }
+  if (typeof obj.usage !== 'object') {
+    return
+  }
+  const parseResult = openrouterUsageSchema.safeParse(obj.usage)
+  if (!parseResult.success) {
+    trackEvent(
+      AnalyticsEvent.OPENROUTER_MALFORMED_JSON_RESPONSE_CHUNK,
+      userId,
+      {
+        message: `Usage does not match schema:\n${parseResult.error.message}`,
+        data,
+      },
+    )
+    return
+  }
+
+  const directCost = parseResult?.data?.cost ?? 0
+  const upstreamCost = parseResult?.data?.cost_details?.upstream_inference_cost
+
+  saveMessage({
+    messageId: obj.id,
+    userId,
+    clientSessionId: generateCompactId('direct-'),
+    fingerprintId: generateCompactId('direct-'),
+    userInputId: generateCompactId('direct-'),
+    model,
+    request,
+  request: Message[]
+  response: string
+  inputTokens: number
+  outputTokens: number
+  cacheCreationInputTokens?: number
+  cacheReadInputTokens?: number
+  finishedAt: Date
+  latencyMs: number
+  usesUserApiKey?: boolean
+  chargeUser?: boolean
+  costOverrideDollars?: number
+  agentId?: string
+  })
+}
+
+*/
diff --git a/backend/src/util/auth-helpers.ts b/backend/src/util/auth-helpers.ts
@@ -1,10 +1,20 @@
 import type { Request } from 'express'
 
 /**
- * Extract auth token from x-codebuff-api-key header
+ * Extract auth token from x-codebuff-api-key header or authorization header
  */
 export function extractAuthTokenFromHeader(req: Request): string | undefined {
-  const token = req.headers['x-codebuff-api-key'] as string | undefined
-  // Trim any whitespace that might be present
-  return token?.trim()
-}
+  const token = req.headers['x-codebuff-api-key']
+  if (typeof token === 'string' && token) {
+    return token
+  }
+
+  const authorization = req.headers['authorization']
+  if (!authorization) {
+    return undefined
+  }
+  if (!authorization.startsWith('Bearer ')) {
+    return undefined
+  }
+  return authorization.slice('Bearer '.length)
+}
diff --git a/common/src/constants/analytics-events.ts b/common/src/constants/analytics-events.ts
@@ -33,6 +33,7 @@ export enum AnalyticsEvent {
   TOOL_USE = 'backend.tool_use',
   UNKNOWN_TOOL_CALL = 'backend.unknown_tool_call',
   USER_INPUT = 'backend.user_input',
+  OPENROUTER_MALFORMED_JSON_RESPONSE_CHUNK = 'backend.openrouter_malformed_json_response_chunk',
 
   // Web
   SIGNUP = 'web.signup',
diff --git a/npm-app/package.json b/npm-app/package.json
@@ -34,6 +34,7 @@
     "bun": ">=1.2.11"
   },
   "dependencies": {
+    "@ai-sdk/openai-compatible": "^1.0.19",
     "@codebuff/code-map": "workspace:*",
     "@codebuff/common": "workspace:*",
     "@types/diff": "8.0.0",
diff --git a/npm-app/src/asdf.ts b/npm-app/src/asdf.ts
@@ -0,0 +1,35 @@
+import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
+import { streamText } from 'ai'
+
+import { websiteUrl } from './config'
+
+const codebuffBackendProvider = createOpenAICompatible({
+  name: 'codebuff',
+  apiKey: '12345',
+  baseURL: websiteUrl + '/api/v1',
+})
+
+const response = streamText({
+  model: codebuffBackendProvider('anthropic/claude-sonnet-4.5'),
+  messages: [
+    {
+      role: 'user',
+      content:
+        'This is a bunch of text just to fill out some space. Ignore this.'.repeat(
+          1000,
+        ),
+    },
+    {
+      role: 'user',
+      content: 'Hello',
+      providerOptions: {
+        codebuff: {
+          cacheControl: { type: 'ephemeral' },
+        },
+      },
+    },
+  ],
+})
+for await (const chunk of response.fullStream) {
+  console.log('asdf', { chunk })
+}
diff --git a/web/src/app/api/v1/chat/completions/route.ts b/web/src/app/api/v1/chat/completions/route.ts
@@ -0,0 +1,73 @@
+import { NextResponse } from 'next/server'
+
+import type { NextRequest } from 'next/server'
+
+import { getUserInfoFromApiKey } from '@/db/user'
+import { handleOpenrouterStream } from '@/llm-api/openrouter'
+import { extractApiKeyFromHeader } from '@/util/auth'
+import { errorToObject } from '@/util/error'
+import { logger } from '@/util/logger'
+
+export async function POST(req: NextRequest) {
+  try {
+    const body = await req.json()
+
+    console.log('asdf', {
+      req: { headers: Object.fromEntries(req.headers), body },
+    })
+
+    const apiKey = extractApiKeyFromHeader(req)
+
+    if (!apiKey) {
+      return NextResponse.json({ message: 'Unauthorized' }, { status: 401 })
+    }
+
+    const userInfo = await getUserInfoFromApiKey({ apiKey, fields: ['id'] })
+    if (!userInfo) {
+      return NextResponse.json(
+        { message: 'Invalid Codebuff API key' },
+        { status: 401 }
+      )
+    }
+
+    if (body.stream) {
+      try {
+        const stream = await handleOpenrouterStream({
+          body,
+        })
+
+        return new NextResponse(stream, {
+          headers: {
+            'Content-Type': 'text/event-stream',
+            'Cache-Control': 'no-cache',
+            Connection: 'keep-alive',
+            'Access-Control-Allow-Origin': '*',
+          },
+        })
+      } catch (error) {
+        logger.error(
+          errorToObject(error),
+          'Error setting up OpenRouter stream:'
+        )
+        return NextResponse.json(
+          { error: 'Failed to initialize stream' },
+          { status: 500 }
+        )
+      }
+    }
+
+    return NextResponse.json(
+      { message: 'Not implemented. Use stream=true.' },
+      { status: 500 }
+    )
+  } catch (error) {
+    logger.error(
+      errorToObject(error),
+      'Error processing chat completions request:'
+    )
+    return NextResponse.json(
+      { error: 'Internal server error' },
+      { status: 500 }
+    )
+  }
+}
diff --git a/web/src/db/user.ts b/web/src/db/user.ts
@@ -0,0 +1,39 @@
+import db from '@codebuff/common/db'
+import * as schema from '@codebuff/common/db/schema'
+import { eq } from 'drizzle-orm'
+
+import type { InferSelectModel } from 'drizzle-orm'
+
+type UserTable = typeof schema.user
+type UserColumn = UserTable['_']['columns']
+type User = InferSelectModel<UserTable>
+
+export async function getUserInfoFromApiKey<
+  T extends readonly (keyof UserColumn)[],
+>({
+  apiKey,
+  fields,
+}: {
+  apiKey: string
+  fields: T
+}): Promise<
+  | {
+      [K in T[number]]: User[K]
+    }
+  | undefined
+> {
+  // Build a typed selection object for user columns
+  const userSelection = Object.fromEntries(
+    fields.map((field) => [field, schema.user[field]])
+  ) as { [K in T[number]]: UserColumn[K] }
+
+  const rows = await db
+    .select({ user: userSelection }) // <-- important: nest under 'user'
+    .from(schema.user)
+    .leftJoin(schema.session, eq(schema.user.id, schema.session.userId))
+    .where(eq(schema.session.sessionToken, apiKey))
+    .limit(1)
+
+  // Drizzle returns { user: ..., session: ... }, we return only the user part
+  return rows[0]?.user
+}
diff --git a/web/src/llm-api/openrouter.ts b/web/src/llm-api/openrouter.ts
diff --git a/web/src/util/auth.ts b/web/src/util/auth.ts
diff --git a/web/src/util/error.ts b/web/src/util/error.ts
