diff --git a/lib/chat/__tests__/handleChatGenerate.test.ts b/lib/chat/__tests__/handleChatGenerate.test.ts index 05e219f5..c1d3ede8 100644 --- a/lib/chat/__tests__/handleChatGenerate.test.ts +++ b/lib/chat/__tests__/handleChatGenerate.test.ts @@ -1,32 +1,15 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { NextResponse } from "next/server"; -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; -import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; +import { validateAuthContext } from "@/lib/auth/validateAuthContext"; import { setupChatRequest } from "@/lib/chat/setupChatRequest"; import { saveChatCompletion } from "@/lib/chat/saveChatCompletion"; import { setupConversation } from "@/lib/chat/setupConversation"; import { handleChatGenerate } from "../handleChatGenerate"; // Mock all dependencies before importing the module under test -vi.mock("@/lib/auth/getApiKeyAccountId", () => ({ - getApiKeyAccountId: vi.fn(), -})); - -vi.mock("@/lib/auth/getAuthenticatedAccountId", () => ({ - getAuthenticatedAccountId: vi.fn(), -})); - -vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ - validateOverrideAccountId: vi.fn(), -})); - -vi.mock("@/lib/keys/getApiKeyDetails", () => ({ - getApiKeyDetails: vi.fn(), -})); - -vi.mock("@/lib/organizations/validateOrganizationAccess", () => ({ - validateOrganizationAccess: vi.fn(), +vi.mock("@/lib/auth/validateAuthContext", () => ({ + validateAuthContext: vi.fn(), })); vi.mock("@/lib/chat/setupChatRequest", () => ({ @@ -61,8 +44,7 @@ vi.mock("@/lib/chat/setupConversation", () => ({ setupConversation: vi.fn(), })); -const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); -const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId); +const mockValidateAuthContext = vi.mocked(validateAuthContext); const mockSetupChatRequest = vi.mocked(setupChatRequest); const mockSaveChatCompletion = vi.mocked(saveChatCompletion); const mockSetupConversation = vi.mocked(setupConversation); @@ -112,7 +94,11 @@ describe("handleChatGenerate", () => { describe("validation", () => { it("returns 400 error when neither messages nor prompt is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ roomId: "room-123" }, { "x-api-key": "test-key" }); @@ -125,20 +111,25 @@ describe("handleChatGenerate", () => { }); it("returns 401 error when no auth header is provided", async () => { + mockValidateAuthContext.mockResolvedValue( + NextResponse.json({ status: "error", error: "Unauthorized" }, { status: 401 }), + ); const request = createMockRequest({ prompt: "Hello" }, {}); const result = await handleChatGenerate(request as any); expect(result).toBeInstanceOf(NextResponse); expect(result.status).toBe(401); - const json = await result.json(); - expect(json.message).toBe("Exactly one of x-api-key or Authorization must be provided"); }); }); describe("text generation", () => { it("returns generated text using agent.generate() for valid requests", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = createMockAgent({ text: "Hello! How can I help you?", @@ -171,7 +162,11 @@ describe("handleChatGenerate", () => { }); it("uses messages array when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = createMockAgent({ text: "Response", @@ -199,7 +194,11 @@ describe("handleChatGenerate", () => { }); it("passes through optional parameters", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "room-xyz", memoryId: "memory-id", @@ -241,7 +240,11 @@ describe("handleChatGenerate", () => { }); it("includes reasoningText when present", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = createMockAgent({ text: "Response", @@ -270,7 +273,11 @@ describe("handleChatGenerate", () => { describe("error handling", () => { it("returns 500 error when setupChatRequest fails", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupChatRequest.mockRejectedValue(new Error("Setup failed")); const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); @@ -284,7 +291,11 @@ describe("handleChatGenerate", () => { }); it("returns 500 error when agent.generate() fails", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = { generate: vi.fn().mockRejectedValue(new Error("Generation failed")), @@ -309,10 +320,11 @@ describe("handleChatGenerate", () => { }); describe("accountId override", () => { - it("allows org API key to override accountId", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockValidateOverrideAccountId.mockResolvedValue({ + it("allows accountId override", async () => { + mockValidateAuthContext.mockResolvedValue({ accountId: "target-account-456", + orgId: null, + authToken: "token", }); const mockAgent = createMockAgent({ @@ -344,7 +356,11 @@ describe("handleChatGenerate", () => { describe("message persistence", () => { it("saves assistant message to database when roomId is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "room-abc-123", memoryId: "memory-id", @@ -378,7 +394,11 @@ describe("handleChatGenerate", () => { }); it("saves message with auto-generated roomId when roomId is not provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "auto-generated-room-id", memoryId: "memory-id", @@ -410,7 +430,11 @@ describe("handleChatGenerate", () => { }); it("includes roomId in HTTP response when provided by client", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "client-provided-room-id", memoryId: "memory-id", @@ -443,7 +467,11 @@ describe("handleChatGenerate", () => { }); it("includes auto-generated roomId in HTTP response when not provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "auto-generated-room-456", memoryId: "memory-id", @@ -473,7 +501,11 @@ describe("handleChatGenerate", () => { }); it("passes correct text to saveChatCompletion", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "room-xyz", memoryId: "memory-id", @@ -507,7 +539,11 @@ describe("handleChatGenerate", () => { }); it("still returns success response even if saveChatCompletion fails", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "room-abc", memoryId: "memory-id", diff --git a/lib/chat/__tests__/handleChatStream.test.ts b/lib/chat/__tests__/handleChatStream.test.ts index ab9b9e79..c670b1ef 100644 --- a/lib/chat/__tests__/handleChatStream.test.ts +++ b/lib/chat/__tests__/handleChatStream.test.ts @@ -1,32 +1,15 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { NextResponse } from "next/server"; -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; -import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; +import { validateAuthContext } from "@/lib/auth/validateAuthContext"; import { setupChatRequest } from "@/lib/chat/setupChatRequest"; import { setupConversation } from "@/lib/chat/setupConversation"; import { createUIMessageStream, createUIMessageStreamResponse } from "ai"; import { handleChatStream } from "../handleChatStream"; // Mock all dependencies before importing the module under test -vi.mock("@/lib/auth/getApiKeyAccountId", () => ({ - getApiKeyAccountId: vi.fn(), -})); - -vi.mock("@/lib/auth/getAuthenticatedAccountId", () => ({ - getAuthenticatedAccountId: vi.fn(), -})); - -vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ - validateOverrideAccountId: vi.fn(), -})); - -vi.mock("@/lib/keys/getApiKeyDetails", () => ({ - getApiKeyDetails: vi.fn(), -})); - -vi.mock("@/lib/organizations/validateOrganizationAccess", () => ({ - validateOrganizationAccess: vi.fn(), +vi.mock("@/lib/auth/validateAuthContext", () => ({ + validateAuthContext: vi.fn(), })); vi.mock("@/lib/chat/setupConversation", () => ({ @@ -67,8 +50,7 @@ vi.mock("ai", () => ({ createUIMessageStreamResponse: vi.fn(), })); -const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); -const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId); +const mockValidateAuthContext = vi.mocked(validateAuthContext); const mockSetupConversation = vi.mocked(setupConversation); const mockSetupChatRequest = vi.mocked(setupChatRequest); const mockCreateUIMessageStream = vi.mocked(createUIMessageStream); @@ -107,7 +89,11 @@ describe("handleChatStream", () => { describe("validation", () => { it("returns 400 error when neither messages nor prompt is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ roomId: "room-123" }, { "x-api-key": "test-key" }); @@ -120,20 +106,25 @@ describe("handleChatStream", () => { }); it("returns 401 error when no auth header is provided", async () => { + mockValidateAuthContext.mockResolvedValue( + NextResponse.json({ status: "error", error: "Unauthorized" }, { status: 401 }), + ); const request = createMockRequest({ prompt: "Hello" }, {}); const result = await handleChatStream(request as any); expect(result).toBeInstanceOf(NextResponse); expect(result.status).toBe(401); - const json = await result.json(); - expect(json.message).toBe("Exactly one of x-api-key or Authorization must be provided"); }); }); describe("streaming", () => { it("creates a streaming response for valid requests", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = { stream: vi.fn().mockResolvedValue({ @@ -173,7 +164,11 @@ describe("handleChatStream", () => { }); it("uses messages array when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = { stream: vi.fn().mockResolvedValue({ @@ -206,7 +201,11 @@ describe("handleChatStream", () => { }); it("passes through optional parameters", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const mockAgent = { stream: vi.fn().mockResolvedValue({ @@ -251,7 +250,11 @@ describe("handleChatStream", () => { describe("error handling", () => { it("returns 500 error when setupChatRequest fails", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupChatRequest.mockRejectedValue(new Error("Setup failed")); const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); @@ -266,10 +269,11 @@ describe("handleChatStream", () => { }); describe("accountId override", () => { - it("allows org API key to override accountId", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockValidateOverrideAccountId.mockResolvedValue({ + it("allows accountId override", async () => { + mockValidateAuthContext.mockResolvedValue({ accountId: "target-account-456", + orgId: null, + authToken: "token", }); const mockAgent = { diff --git a/lib/chat/__tests__/integration/chatEndToEnd.test.ts b/lib/chat/__tests__/integration/chatEndToEnd.test.ts index 25841a5e..b54e51f5 100644 --- a/lib/chat/__tests__/integration/chatEndToEnd.test.ts +++ b/lib/chat/__tests__/integration/chatEndToEnd.test.ts @@ -1,7 +1,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { NextResponse } from "next/server"; -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; +import { validateAuthContext } from "@/lib/auth/validateAuthContext"; import selectAccountEmails from "@/lib/supabase/account_emails/selectAccountEmails"; import { selectAccountInfo } from "@/lib/supabase/account_info/selectAccountInfo"; import { getAccountWithDetails } from "@/lib/supabase/accounts/getAccountWithDetails"; @@ -33,24 +33,8 @@ import { setupChatRequest } from "../../setupChatRequest"; */ // Mock auth dependencies -vi.mock("@/lib/auth/getApiKeyAccountId", () => ({ - getApiKeyAccountId: vi.fn(), -})); - -vi.mock("@/lib/auth/getAuthenticatedAccountId", () => ({ - getAuthenticatedAccountId: vi.fn(), -})); - -vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ - validateOverrideAccountId: vi.fn(), -})); - -vi.mock("@/lib/keys/getApiKeyDetails", () => ({ - getApiKeyDetails: vi.fn(), -})); - -vi.mock("@/lib/organizations/validateOrganizationAccess", () => ({ - validateOrganizationAccess: vi.fn(), +vi.mock("@/lib/auth/validateAuthContext", () => ({ + validateAuthContext: vi.fn(), })); // Mock Supabase dependencies @@ -155,7 +139,7 @@ vi.mock("ai", () => ({ })), })); -const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); +const mockValidateAuthContext = vi.mocked(validateAuthContext); const mockSelectAccountEmails = vi.mocked(selectAccountEmails); const mockSelectAccountInfo = vi.mocked(selectAccountInfo); const mockGetAccountWithDetails = vi.mocked(getAccountWithDetails); @@ -205,7 +189,11 @@ describe("Chat Integration Tests", () => { describe("validateChatRequest integration", () => { it("validates and returns body for valid request with prompt", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); @@ -218,7 +206,11 @@ describe("Chat Integration Tests", () => { }); it("validates and returns body for valid request with messages", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { @@ -233,22 +225,11 @@ describe("Chat Integration Tests", () => { expect((result as any).messages).toHaveLength(1); }); - it("returns 401 when no auth header is provided", async () => { - const request = createMockRequest({ prompt: "Hello" }, {}); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - expect((result as NextResponse).status).toBe(401); - }); - - it("returns 401 when API key lookup fails", async () => { - // getApiKeyAccountId returns a NextResponse when authentication fails - mockGetApiKeyAccountId.mockResolvedValue( - NextResponse.json({ status: "error", message: "Unauthorized" }, { status: 401 }), + it("returns 401 when auth fails", async () => { + mockValidateAuthContext.mockResolvedValue( + NextResponse.json({ status: "error", error: "Unauthorized" }, { status: 401 }), ); - - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "invalid-key" }); + const request = createMockRequest({ prompt: "Hello" }, {}); const result = await validateChatRequest(request as any); @@ -257,7 +238,11 @@ describe("Chat Integration Tests", () => { }); it("returns 400 when neither messages nor prompt is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ roomId: "room-123" }, { "x-api-key": "valid-key" }); @@ -268,7 +253,11 @@ describe("Chat Integration Tests", () => { }); it("returns 400 when both prompt and messages are provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { @@ -285,7 +274,11 @@ describe("Chat Integration Tests", () => { }); it("passes through optional parameters", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { @@ -604,7 +597,11 @@ describe("Chat Integration Tests", () => { describe("end-to-end validation flow", () => { it("validates prompt-based requests through full pipeline", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ prompt: "What is 2+2?" }, { "x-api-key": "valid-key" }); @@ -617,7 +614,11 @@ describe("Chat Integration Tests", () => { }); it("validates messages-based requests through full pipeline", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { @@ -639,7 +640,11 @@ describe("Chat Integration Tests", () => { }); it("handles complete chat flow with post-completion", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSelectRoom.mockResolvedValue(null); mockGenerateChatTitle.mockResolvedValue("Math Question"); diff --git a/lib/chat/__tests__/validateChatRequest.test.ts b/lib/chat/__tests__/validateChatRequest.test.ts index b420c598..0fd081fa 100644 --- a/lib/chat/__tests__/validateChatRequest.test.ts +++ b/lib/chat/__tests__/validateChatRequest.test.ts @@ -2,11 +2,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { NextResponse } from "next/server"; import { validateChatRequest, chatRequestSchema } from "../validateChatRequest"; -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; -import { getAuthenticatedAccountId } from "@/lib/auth/getAuthenticatedAccountId"; -import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; -import { getApiKeyDetails } from "@/lib/keys/getApiKeyDetails"; -import { validateOrganizationAccess } from "@/lib/organizations/validateOrganizationAccess"; +import { validateAuthContext } from "@/lib/auth/validateAuthContext"; import { generateUUID } from "@/lib/uuid/generateUUID"; import { createNewRoom } from "@/lib/chat/createNewRoom"; import insertMemories from "@/lib/supabase/memories/insertMemories"; @@ -14,24 +10,8 @@ import filterMessageContentForMemories from "@/lib/messages/filterMessageContent import { setupConversation } from "@/lib/chat/setupConversation"; // Mock dependencies -vi.mock("@/lib/auth/getApiKeyAccountId", () => ({ - getApiKeyAccountId: vi.fn(), -})); - -vi.mock("@/lib/auth/getAuthenticatedAccountId", () => ({ - getAuthenticatedAccountId: vi.fn(), -})); - -vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ - validateOverrideAccountId: vi.fn(), -})); - -vi.mock("@/lib/keys/getApiKeyDetails", () => ({ - getApiKeyDetails: vi.fn(), -})); - -vi.mock("@/lib/organizations/validateOrganizationAccess", () => ({ - validateOrganizationAccess: vi.fn(), +vi.mock("@/lib/auth/validateAuthContext", () => ({ + validateAuthContext: vi.fn(), })); vi.mock("@/lib/uuid/generateUUID", () => { @@ -58,11 +38,7 @@ vi.mock("@/lib/chat/setupConversation", () => ({ setupConversation: vi.fn(), })); -const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); -const mockGetAuthenticatedAccountId = vi.mocked(getAuthenticatedAccountId); -const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId); -const mockGetApiKeyDetails = vi.mocked(getApiKeyDetails); -const mockValidateOrganizationAccess = vi.mocked(validateOrganizationAccess); +const mockValidateAuthContext = vi.mocked(validateAuthContext); const mockGenerateUUID = vi.mocked(generateUUID); const mockCreateNewRoom = vi.mocked(createNewRoom); const mockInsertMemories = vi.mocked(insertMemories); @@ -97,7 +73,11 @@ describe("validateChatRequest", () => { describe("schema validation", () => { it("rejects when neither messages nor prompt is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ roomId: "room-123" }, { "x-api-key": "test-key" }); @@ -110,7 +90,11 @@ describe("validateChatRequest", () => { }); it("rejects when both messages and prompt are provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { @@ -129,7 +113,11 @@ describe("validateChatRequest", () => { }); it("accepts valid messages array", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { messages: [{ role: "user", content: "Hello" }] }, @@ -143,7 +131,11 @@ describe("validateChatRequest", () => { }); it("accepts valid prompt string", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ prompt: "Hello, world!" }, { "x-api-key": "test-key" }); @@ -156,178 +148,75 @@ describe("validateChatRequest", () => { describe("authentication", () => { it("rejects request without any auth header", async () => { - const request = createMockRequest({ prompt: "Hello" }, {}); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - const json = await (result as NextResponse).json(); - expect(json.status).toBe("error"); - expect(json.message).toBe("Exactly one of x-api-key or Authorization must be provided"); - }); - - it("rejects request with both x-api-key and Authorization headers", async () => { - const request = createMockRequest( - { prompt: "Hello" }, - { "x-api-key": "test-key", authorization: "Bearer test-token" }, - ); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - const json = await (result as NextResponse).json(); - expect(json.status).toBe("error"); - expect(json.message).toBe("Exactly one of x-api-key or Authorization must be provided"); - }); - - it("rejects request with invalid API key", async () => { - mockGetApiKeyAccountId.mockResolvedValue( - NextResponse.json({ status: "error", message: "Invalid API key" }, { status: 401 }), - ); - - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "invalid-key" }); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - }); - - it("uses accountId from valid API key", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-abc-123"); - - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); - - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("account-abc-123"); - }); - - it("accepts valid Authorization Bearer token", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("account-from-jwt-456"); - - const request = createMockRequest( - { prompt: "Hello" }, - { authorization: "Bearer valid-jwt-token" }, - ); - - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("account-from-jwt-456"); - }); - - it("rejects request with invalid Authorization token", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue( + mockValidateAuthContext.mockResolvedValue( NextResponse.json( - { status: "error", message: "Failed to verify authentication token" }, + { status: "error", error: "Exactly one of x-api-key or Authorization must be provided" }, { status: 401 }, ), ); - - const request = createMockRequest( - { prompt: "Hello" }, - { authorization: "Bearer invalid-token" }, - ); + const request = createMockRequest({ prompt: "Hello" }, {}); const result = await validateChatRequest(request as any); expect(result).toBeInstanceOf(NextResponse); const json = await (result as NextResponse).json(); expect(json.status).toBe("error"); + expect(json.error).toBe("Exactly one of x-api-key or Authorization must be provided"); }); - it("returns accountId for org API key", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockGetApiKeyDetails.mockResolvedValue({ - accountId: "org-account-123", - }); - - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "org-api-key" }); + it("returns auth error when validateAuthContext fails", async () => { + mockValidateAuthContext.mockResolvedValue( + NextResponse.json({ status: "error", message: "Unauthorized" }, { status: 401 }), + ); + const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "invalid-key" }); const result = await validateChatRequest(request as any); - expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("org-account-123"); + expect(result).toBeInstanceOf(NextResponse); + expect((result as NextResponse).status).toBe(401); }); - it("returns accountId for personal API key", async () => { - mockGetApiKeyAccountId.mockResolvedValue("personal-account-123"); - mockGetApiKeyDetails.mockResolvedValue({ - accountId: "personal-account-123", + it("uses accountId from validateAuthContext", async () => { + mockValidateAuthContext.mockResolvedValue({ + accountId: "resolved-account-123", + orgId: null, + authToken: "token", }); - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "personal-api-key" }); - - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("personal-account-123"); - }); - - it("returns accountId for bearer token auth", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("jwt-account-456"); - - const request = createMockRequest( - { prompt: "Hello" }, - { authorization: "Bearer valid-jwt-token" }, - ); - + const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "valid-key" }); const result = await validateChatRequest(request as any); expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("jwt-account-456"); + expect((result as any).accountId).toBe("resolved-account-123"); }); - }); - describe("accountId override", () => { - it("allows org API key to override accountId", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockValidateOverrideAccountId.mockResolvedValue({ - accountId: "target-account-456", + it("passes accountId and organizationId to validateAuthContext", async () => { + mockValidateAuthContext.mockResolvedValue({ + accountId: "target-456", + orgId: "org-789", + authToken: "token", }); const request = createMockRequest( - { prompt: "Hello", accountId: "target-account-456" }, - { "x-api-key": "org-api-key" }, + { prompt: "Hello", accountId: "target-456", organizationId: "org-789" }, + { "x-api-key": "valid-key" }, ); + await validateChatRequest(request as any); - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - expect((result as any).accountId).toBe("target-account-456"); - expect(mockValidateOverrideAccountId).toHaveBeenCalledWith({ - apiKey: "org-api-key", - targetAccountId: "target-account-456", + expect(mockValidateAuthContext).toHaveBeenCalledWith(expect.anything(), { + accountId: "target-456", + organizationId: "org-789", }); }); - - it("rejects unauthorized accountId override", async () => { - mockGetApiKeyAccountId.mockResolvedValue("personal-account-123"); - mockValidateOverrideAccountId.mockResolvedValue( - NextResponse.json( - { status: "error", message: "Access denied to specified accountId" }, - { status: 403 }, - ), - ); - - const request = createMockRequest( - { prompt: "Hello", accountId: "target-account-456" }, - { "x-api-key": "personal-api-key" }, - ); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - const json = await (result as NextResponse).json(); - expect(json.status).toBe("error"); - expect(json.message).toBe("Access denied to specified accountId"); - }); }); describe("message normalization", () => { it("converts prompt to messages array", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest({ prompt: "Hello, world!" }, { "x-api-key": "test-key" }); @@ -340,7 +229,11 @@ describe("validateChatRequest", () => { }); it("preserves original messages when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const originalMessages = [ { role: "user", content: "Hi" }, @@ -360,7 +253,11 @@ describe("validateChatRequest", () => { describe("optional fields", () => { it("passes through roomId", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "room-xyz", memoryId: "memory-id", @@ -378,7 +275,11 @@ describe("validateChatRequest", () => { }); it("passes through artistId", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { prompt: "Hello", artistId: "artist-abc" }, @@ -392,7 +293,11 @@ describe("validateChatRequest", () => { }); it("passes through model selection", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { prompt: "Hello", model: "gpt-4" }, @@ -406,7 +311,11 @@ describe("validateChatRequest", () => { }); it("passes through excludeTools array", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { prompt: "Hello", excludeTools: ["tool1", "tool2"] }, @@ -420,7 +329,11 @@ describe("validateChatRequest", () => { }); it("passes through topic", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); const request = createMockRequest( { prompt: "Hello", topic: "Pulse Feb 2" }, @@ -474,110 +387,47 @@ describe("validateChatRequest", () => { expect(result.success).toBe(true); }); - it("uses provided organizationId when user is member of org (bearer token)", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("user-account-123"); - mockValidateOrganizationAccess.mockResolvedValue(true); - - const request = createMockRequest( - { prompt: "Hello", organizationId: "org-456" }, - { authorization: "Bearer valid-jwt-token" }, - ); - - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - expect(mockValidateOrganizationAccess).toHaveBeenCalledWith({ - accountId: "user-account-123", - organizationId: "org-456", - }); - }); - - it("uses provided organizationId when user is member of org (API key)", async () => { - mockGetApiKeyAccountId.mockResolvedValue("api-key-account-123"); - mockGetApiKeyDetails.mockResolvedValue({ - accountId: "api-key-account-123", - }); - mockValidateOrganizationAccess.mockResolvedValue(true); - - const request = createMockRequest( - { prompt: "Hello", organizationId: "org-789" }, - { "x-api-key": "personal-api-key" }, - ); - - const result = await validateChatRequest(request as any); - expect(mockValidateOrganizationAccess).toHaveBeenCalledWith({ - accountId: "api-key-account-123", - organizationId: "org-789", - }); - }); - - it("overwrites API key orgId with provided organizationId when user is member", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockGetApiKeyDetails.mockResolvedValue({ - accountId: "org-account-123", + it("uses orgId from validateAuthContext when organizationId is provided", async () => { + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: "org-456", + authToken: "token", }); - mockValidateOrganizationAccess.mockResolvedValue(true); const request = createMockRequest( - { prompt: "Hello", organizationId: "different-org-456" }, - { "x-api-key": "org-api-key" }, + { prompt: "Hello", organizationId: "org-456" }, + { "x-api-key": "test-key" }, ); const result = await validateChatRequest(request as any); expect(result).not.toBeInstanceOf(NextResponse); + expect((result as any).orgId).toBe("org-456"); }); - it("rejects organizationId when user is NOT a member of org", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("user-account-123"); - mockValidateOrganizationAccess.mockResolvedValue(false); - - const request = createMockRequest( - { prompt: "Hello", organizationId: "org-not-member" }, - { authorization: "Bearer valid-jwt-token" }, - ); - - const result = await validateChatRequest(request as any); - - expect(result).toBeInstanceOf(NextResponse); - const json = await (result as NextResponse).json(); - expect(json.status).toBe("error"); - expect(json.message).toBe("Access denied to specified organizationId"); - }); - - it("uses API key orgId when no organizationId is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); - mockGetApiKeyDetails.mockResolvedValue({ - accountId: "org-account-123", + it("returns null orgId when no organizationId provided", async () => { + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", }); - const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "org-api-key" }); - - const result = await validateChatRequest(request as any); - - expect(result).not.toBeInstanceOf(NextResponse); - // Should not validate org access when no organizationId is provided - expect(mockValidateOrganizationAccess).not.toHaveBeenCalled(); - }); - - it("returns null orgId when no organizationId provided and bearer token auth", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("user-account-123"); - - const request = createMockRequest( - { prompt: "Hello" }, - { authorization: "Bearer valid-jwt-token" }, - ); + const request = createMockRequest({ prompt: "Hello" }, { "x-api-key": "test-key" }); const result = await validateChatRequest(request as any); expect(result).not.toBeInstanceOf(NextResponse); - expect(mockValidateOrganizationAccess).not.toHaveBeenCalled(); + expect((result as any).orgId).toBeNull(); }); }); describe("auto room creation", () => { it("returns roomId from setupConversation when roomId is not provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "generated-uuid-456", memoryId: "memory-id", @@ -592,7 +442,11 @@ describe("validateChatRequest", () => { }); it("calls setupConversation with correct params when roomId is not provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "generated-uuid-789", memoryId: "memory-id", @@ -618,7 +472,11 @@ describe("validateChatRequest", () => { }); it("passes artistId to setupConversation when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "generated-uuid-abc", memoryId: "memory-id", @@ -639,7 +497,11 @@ describe("validateChatRequest", () => { }); it("passes topic to setupConversation when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "generated-uuid-topic", memoryId: "memory-id", @@ -660,7 +522,11 @@ describe("validateChatRequest", () => { }); it("returns provided roomId when roomId is provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "existing-room-123", memoryId: "memory-id", @@ -678,7 +544,11 @@ describe("validateChatRequest", () => { }); it("passes roomId to setupConversation when provided", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "existing-room-456", memoryId: "memory-id", @@ -699,7 +569,11 @@ describe("validateChatRequest", () => { }); it("works with bearer token auth for auto room creation", async () => { - mockGetAuthenticatedAccountId.mockResolvedValue("jwt-account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "jwt-account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "jwt-generated-uuid", memoryId: "memory-id", @@ -722,7 +596,11 @@ describe("validateChatRequest", () => { }); it("calls setupConversation when roomId is auto-created", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "new-room-uuid", memoryId: "new-room-uuid", @@ -750,7 +628,11 @@ describe("validateChatRequest", () => { }); it("calls setupConversation for existing rooms", async () => { - mockGetApiKeyAccountId.mockResolvedValue("account-123"); + mockValidateAuthContext.mockResolvedValue({ + accountId: "account-123", + orgId: null, + authToken: "token", + }); mockSetupConversation.mockResolvedValue({ roomId: "existing-room-id", memoryId: "memory-uuid", diff --git a/lib/chat/validateChatRequest.ts b/lib/chat/validateChatRequest.ts index e0d80adb..ff7eb3a7 100644 --- a/lib/chat/validateChatRequest.ts +++ b/lib/chat/validateChatRequest.ts @@ -2,12 +2,9 @@ import type { NextRequest } from "next/server"; import { NextResponse } from "next/server"; import { z } from "zod"; import { getCorsHeaders } from "@/lib/networking/getCorsHeaders"; -import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; -import { getAuthenticatedAccountId } from "@/lib/auth/getAuthenticatedAccountId"; -import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; +import { validateAuthContext } from "@/lib/auth/validateAuthContext"; import { getMessages } from "@/lib/messages/getMessages"; import convertToUiMessages from "@/lib/messages/convertToUiMessages"; -import { validateOrganizationAccess } from "@/lib/organizations/validateOrganizationAccess"; import { setupConversation } from "@/lib/chat/setupConversation"; import { validateMessages } from "@/lib/chat/validateMessages"; @@ -86,81 +83,15 @@ export async function validateChatRequest( const validatedBody: BaseChatRequestBody = validationResult.data; - // Check which auth mechanism is provided - const apiKey = request.headers.get("x-api-key"); - const authHeader = request.headers.get("authorization"); - const hasApiKey = !!apiKey; - const hasAuth = !!authHeader; - - // Enforce that exactly one auth mechanism is provided - if ((hasApiKey && hasAuth) || (!hasApiKey && !hasAuth)) { - return NextResponse.json( - { - status: "error", - message: "Exactly one of x-api-key or Authorization must be provided", - }, - { - status: 401, - headers: getCorsHeaders(), - }, - ); - } - - // Authenticate and get accountId and orgId - let accountId: string; - let orgId: string | null = null; - - if (hasApiKey) { - // Validate API key authentication - const accountIdOrError = await getApiKeyAccountId(request); - if (accountIdOrError instanceof NextResponse) { - return accountIdOrError; - } - accountId = accountIdOrError; - - // Handle accountId override - if (validatedBody.accountId) { - const overrideResult = await validateOverrideAccountId({ - apiKey, - targetAccountId: validatedBody.accountId, - }); - if (overrideResult instanceof NextResponse) { - return overrideResult; - } - accountId = overrideResult.accountId; - } - } else { - // Validate bearer token authentication (no org context for JWT auth) - const accountIdOrError = await getAuthenticatedAccountId(request); - if (accountIdOrError instanceof NextResponse) { - return accountIdOrError; - } - accountId = accountIdOrError; - } - - // Handle organizationId override from request body - if (validatedBody.organizationId) { - const hasOrgAccess = await validateOrganizationAccess({ - accountId, - organizationId: validatedBody.organizationId, - }); - - if (!hasOrgAccess) { - return NextResponse.json( - { - status: "error", - message: "Access denied to specified organizationId", - }, - { - status: 403, - headers: getCorsHeaders(), - }, - ); - } - - // Use the provided organizationId as orgId - orgId = validatedBody.organizationId; + // Authenticate, handle accountId/organizationId overrides + const authResult = await validateAuthContext(request, { + accountId: validatedBody.accountId, + organizationId: validatedBody.organizationId, + }); + if (authResult instanceof NextResponse) { + return authResult; } + const { accountId, orgId } = authResult; // Normalize chat content: // - If only prompt is provided, convert it into a single user UIMessage @@ -188,14 +119,11 @@ export async function validateChatRequest( memoryId: lastMessage.id, }); - // Extract the auth token to forward to MCP server - const authToken = hasApiKey ? apiKey! : authHeader!.replace(/^Bearer\s+/i, ""); - return { ...validatedBody, accountId, orgId, roomId: finalRoomId, - authToken, + authToken: authResult.authToken, } as ChatRequestBody; }