reject if run id does not exist in db

Author: charleslien

web/src/app/api/v1/chat/completions/route.ts

@@ -3,6 +3,7 @@ import { NextResponse } from 'next/server'
 
 import type { NextRequest } from 'next/server'
 
+import { getAgentRunFromId } from '@/db/agent-run'
 import { getUserInfoFromApiKey } from '@/db/user'
 import { handleOpenRouterStream } from '@/llm-api/openrouter'
 import { extractApiKeyFromHeader } from '@/util/auth'
@@ -45,37 +46,51 @@ export async function POST(req: NextRequest) {
       )
     }
 
-    if (body.stream) {
-      try {
-        const stream = await handleOpenRouterStream({
-          body,
-          userId,
-        })
+    if (!body.stream) {
+      return NextResponse.json(
+        { message: 'Not implemented. Use stream=true.' },
+        { status: 500 }
+      )
+    }
 
-        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 }
-        )
-      }
+    const runIdFromBody = body.codebuff_metadata?.agentRunId
+    if (!runIdFromBody) {
+      return NextResponse.json(
+        { message: 'No agent run ID found in request body' },
+        { status: 400 }
+      )
+    }
+    const runId = runIdFromBody
+      ? getAgentRunFromId({ agentRunId: runIdFromBody, userId, fields: ['id'] })
+      : null
+    if (!runId) {
+      return NextResponse.json(
+        { message: `Agent Run ID Not Found: ${runIdFromBody}` },
+        { status: 404 }
+      )
     }
 
-    return NextResponse.json(
-      { message: 'Not implemented. Use stream=true.' },
-      { status: 500 }
-    )
+    try {
+      const stream = await handleOpenRouterStream({
+        body,
+        userId,
+      })
+
+      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 }
+      )
+    }
   } catch (error) {
     logger.error(
       errorToObject(error),

web/src/db/agent-run.ts

@@ -0,0 +1,43 @@
+import db from '@codebuff/common/db'
+import * as schema from '@codebuff/common/db/schema'
+import { eq, and } from 'drizzle-orm'
+
+import type { InferSelectModel } from 'drizzle-orm'
+
+type AgentRunTable = typeof schema.agentRun
+type AgentRunColumn = AgentRunTable['_']['columns']
+type AgentRun = InferSelectModel<AgentRunTable>
+
+export async function getAgentRunFromId<
+  T extends readonly (keyof AgentRunColumn)[],
+>({
+  agentRunId,
+  userId,
+  fields,
+}: {
+  agentRunId: string
+  userId: string
+  fields: T
+}): Promise<
+  | {
+      [K in T[number]]: AgentRun[K]
+    }
+  | undefined
+> {
+  const selection = Object.fromEntries(
+    fields.map((field) => [field, schema.agentRun[field]])
+  ) as { [K in T[number]]: AgentRunColumn[K] }
+
+  const rows = await db
+    .select({ selection })
+    .from(schema.agentRun)
+    .where(
+      and(
+        eq(schema.agentRun.id, agentRunId),
+        eq(schema.agentRun.user_id, userId)
+      )
+    )
+    .limit(1)
+
+  return rows[0]?.selection
+}