fix: properly distinguish auth and internal userId

Author: sinedied

packages/agent-api/src/auth.ts

@@ -1,5 +1,6 @@
 import { HttpRequest } from '@azure/functions';
 import { DefaultAzureCredential, getBearerTokenProvider } from '@azure/identity';
+import { UserDbService } from './user-db-service';
 
 const azureOpenAiScope = 'https://cognitiveservices.azure.com/.default';
 
@@ -17,18 +18,31 @@ export function getAzureOpenAiTokenProvider() {
   return getBearerTokenProvider(getCredentials(), azureOpenAiScope);
 }
 
-export function getUserId(request: HttpRequest, body?: any): string | undefined {
+export function getAuthenticationUserId(request: HttpRequest, body?: any): string | undefined {
   let userId: string | undefined;
 
-  // Get the user ID from Azure easy auth if it's available
+  // Get the user ID from Azure easy auth
   try {
     const token = Buffer.from(request.headers.get('x-ms-client-principal') ?? '', 'base64').toString('ascii');
     const infos = token && JSON.parse(token);
     userId = infos?.userId;
   } catch {}
 
-  // Get the user ID from the request as a fallback
-  userId ??= body?.context?.userId ?? request.query.get('userId') ?? undefined;
-
   return userId;
 }
+
+export async function getInternalUserId(request: HttpRequest, body?: any): Promise<string | undefined> {
+  // Get the user ID from Azure easy auth if it's available,
+  let authUserId = getAuthenticationUserId(request, body);
+  if (authUserId) {
+    // Exchange the auth user ID to the internal user ID
+    const db = await UserDbService.getInstance();
+    let user = await db.getUserById(authUserId);
+    if (user) {
+      return user.id;
+    }
+  }
+
+  // Get the user ID from the request as a fallback
+  return body?.context?.userId ?? request.query.get('userId') ?? undefined;
+}

packages/agent-api/src/functions/chats-delete.ts

@@ -1,12 +1,12 @@
 import process from 'node:process';
 import { HttpRequest, HttpResponseInit, InvocationContext, app } from '@azure/functions';
 import { AzureCosmsosDBNoSQLChatMessageHistory } from '@langchain/azure-cosmosdb';
-import { getCredentials, getUserId } from '../auth.js';
+import { getCredentials, getInternalUserId } from '../auth.js';
 
 async function deleteChats(request: HttpRequest, context: InvocationContext): Promise<HttpResponseInit> {
   const azureCosmosDbEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT;
   const { sessionId } = request.params;
-  const userId = getUserId(request);
+  const userId = await getInternalUserId(request);
 
   if (!userId) {
     return {

packages/agent-api/src/functions/chats-get.ts

@@ -1,12 +1,12 @@
 import process from 'node:process';
 import { HttpRequest, HttpResponseInit, InvocationContext, app } from '@azure/functions';
 import { AzureCosmsosDBNoSQLChatMessageHistory } from '@langchain/azure-cosmosdb';
-import { getCredentials, getUserId } from '../auth.js';
+import { getCredentials, getInternalUserId } from '../auth.js';
 
 async function getChats(request: HttpRequest, context: InvocationContext): Promise<HttpResponseInit> {
   const azureCosmosDbEndpoint = process.env.AZURE_COSMOSDB_NOSQL_ENDPOINT;
   const { sessionId } = request.params;
-  const userId = getUserId(request);
+  const userId = await getInternalUserId(request);
 
   if (!userId) {
     return {

packages/agent-api/src/functions/chats-post.ts

@@ -11,7 +11,7 @@ import { createToolCallingAgent } from 'langchain/agents';
 import { AgentExecutor } from 'langchain/agents';
 import { loadMcpTools } from '@langchain/mcp-adapters';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
-import { getAzureOpenAiTokenProvider, getCredentials, getUserId } from '../auth.js';
+import { getAzureOpenAiTokenProvider, getCredentials, getInternalUserId } from '../auth.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { ChainValues } from '@langchain/core/utils/types.js';
 
@@ -50,16 +50,23 @@ export async function postChats(request: HttpRequest, context: InvocationContext
     const requestBody = (await request.json()) as AIChatCompletionRequest;
     const { messages, context: chatContext } = requestBody;
 
-    // Wrong userID!!! need to exchange the one from the token to the me-get
-    const userId = getUserId(request, requestBody);
+    const userId = await getInternalUserId(request, requestBody);
+    if (!userId) {
+      return {
+        status: 400,
+        jsonBody: {
+          error: 'Invalid or missing userId in the request',
+        },
+      }
+    }
 
     if (!messages || messages.length === 0 || !messages.at(-1)?.content) {
       return {
         status: 400,
         jsonBody: {
           error: 'Invalid or missing messages in the request body',
         },
-      }
+      };
     }
 
     let model: BaseChatModel;

packages/agent-api/src/functions/me-get.ts

@@ -1,23 +1,23 @@
 import { createHash } from 'node:crypto';
 import { app, type HttpRequest, type InvocationContext } from '@azure/functions';
 import { UserDbService } from '../user-db-service.js';
-import { getUserId } from '../auth.js';
+import { getAuthenticationUserId } from '../auth.js';
 
 app.http('me-get', {
   methods: ['GET'],
   authLevel: 'anonymous',
   route: 'me',
   async handler(request: HttpRequest, context: InvocationContext) {
     try {
-      const rawUserId = getUserId(request);
-      if (!rawUserId) {
+      const authenticationUserId = getAuthenticationUserId(request);
+      if (!authenticationUserId) {
         return {
           status: 401,
           jsonBody: { error: 'Unauthorized' },
         };
       }
 
-      const id = createHash('sha256').update(rawUserId).digest('hex').substring(0, 32);
+      const id = createHash('sha256').update(authenticationUserId).digest('hex').substring(0, 32);
       context.log(`User ID ${id}`);
 
       const db = await UserDbService.getInstance();