diff --git a/lib/agents/generalAgent/__tests__/getGeneralAgent.test.ts b/lib/agents/generalAgent/__tests__/getGeneralAgent.test.ts index 832e56fc..0b0ce5c3 100644 --- a/lib/agents/generalAgent/__tests__/getGeneralAgent.test.ts +++ b/lib/agents/generalAgent/__tests__/getGeneralAgent.test.ts @@ -84,6 +84,7 @@ describe("getGeneralAgent", () => { it("returns a RoutingDecision object with required properties", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -98,6 +99,7 @@ describe("getGeneralAgent", () => { it("returns a ToolLoopAgent instance", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -109,6 +111,7 @@ describe("getGeneralAgent", () => { it("uses DEFAULT_MODEL when no model is specified", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -120,6 +123,7 @@ describe("getGeneralAgent", () => { it("uses custom model when specified in body", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], model: "anthropic/claude-3-opus", }; @@ -132,6 +136,7 @@ describe("getGeneralAgent", () => { it("sets stopWhen to stepCountIs(111)", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -147,6 +152,7 @@ describe("getGeneralAgent", () => { it("fetches account emails using accountId", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -165,6 +171,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -180,6 +187,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -195,6 +203,7 @@ describe("getGeneralAgent", () => { it("fetches artist info when artistId is provided", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, artistId: "artist-456", messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -207,6 +216,7 @@ describe("getGeneralAgent", () => { it("does not fetch artist info when artistId is not provided", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -223,6 +233,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, artistId: "artist-456", messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -248,6 +259,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, artistId: "artist-456", messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -265,6 +277,7 @@ describe("getGeneralAgent", () => { it("calls getSystemPrompt with all required parameters", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, artistId: "artist-456", roomId: "room-789", messages: [{ id: "1", role: "user", content: "Hello" }], @@ -276,6 +289,7 @@ describe("getGeneralAgent", () => { roomId: "room-789", artistId: "artist-456", accountId: "account-123", + orgId: null, email: "user@example.com", artistInstruction: undefined, knowledgeBaseText: undefined, @@ -294,6 +308,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -333,6 +348,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Edit these images" }], }; @@ -349,6 +365,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -362,6 +379,7 @@ describe("getGeneralAgent", () => { it("calls setupToolsForRequest with body", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], excludeTools: ["dangerous_tool"], }; @@ -377,6 +395,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -390,6 +409,7 @@ describe("getGeneralAgent", () => { it("creates ToolLoopAgent with model matching result model", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], model: "anthropic/claude-sonnet-4", }; @@ -406,6 +426,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -422,6 +443,7 @@ describe("getGeneralAgent", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -433,6 +455,7 @@ describe("getGeneralAgent", () => { it("returns stopWhen in the routing decision", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; diff --git a/lib/agents/generalAgent/getGeneralAgent.ts b/lib/agents/generalAgent/getGeneralAgent.ts index 1e9064eb..e5794c20 100644 --- a/lib/agents/generalAgent/getGeneralAgent.ts +++ b/lib/agents/generalAgent/getGeneralAgent.ts @@ -18,7 +18,7 @@ import { getAccountWithDetails } from "@/lib/supabase/accounts/getAccountWithDet * @returns The general agent */ export default async function getGeneralAgent(body: ChatRequestBody): Promise { - const { accountId, messages, artistId, model: bodyModel } = body; + const { accountId, orgId, messages, artistId, model: bodyModel } = body; const accountEmails = await selectAccountEmails({ accountIds: accountId }); const email = accountEmails[0]?.email || undefined; @@ -37,6 +37,7 @@ export default async function getGeneralAgent(body: ChatRequestBody): Promise { it("calls getConnectedAccount with accountId", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -72,6 +73,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Create a spreadsheet for me" }], }; @@ -90,6 +92,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Test" }], }; @@ -114,6 +117,7 @@ describe("getGoogleSheetsTools", () => { it("returns Google Sheets tools from Composio", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -131,6 +135,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -152,6 +157,7 @@ describe("getGoogleSheetsTools", () => { it("returns google_sheets_login tool", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -168,6 +174,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -185,6 +192,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -200,6 +208,7 @@ describe("getGoogleSheetsTools", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; diff --git a/lib/chat/__tests__/handleChatGenerate.test.ts b/lib/chat/__tests__/handleChatGenerate.test.ts index 1ba8dbea..4274e96b 100644 --- a/lib/chat/__tests__/handleChatGenerate.test.ts +++ b/lib/chat/__tests__/handleChatGenerate.test.ts @@ -14,6 +14,10 @@ vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ validateOverrideAccountId: vi.fn(), })); +vi.mock("@/lib/keys/getApiKeyDetails", () => ({ + getApiKeyDetails: vi.fn(), +})); + vi.mock("@/lib/chat/setupChatRequest", () => ({ setupChatRequest: vi.fn(), })); diff --git a/lib/chat/__tests__/handleChatStream.test.ts b/lib/chat/__tests__/handleChatStream.test.ts index 652befaa..734964ed 100644 --- a/lib/chat/__tests__/handleChatStream.test.ts +++ b/lib/chat/__tests__/handleChatStream.test.ts @@ -14,6 +14,10 @@ vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ validateOverrideAccountId: vi.fn(), })); +vi.mock("@/lib/keys/getApiKeyDetails", () => ({ + getApiKeyDetails: vi.fn(), +})); + vi.mock("@/lib/chat/setupChatRequest", () => ({ setupChatRequest: vi.fn(), })); diff --git a/lib/chat/__tests__/integration/chatEndToEnd.test.ts b/lib/chat/__tests__/integration/chatEndToEnd.test.ts index 0019cb46..86ef311e 100644 --- a/lib/chat/__tests__/integration/chatEndToEnd.test.ts +++ b/lib/chat/__tests__/integration/chatEndToEnd.test.ts @@ -27,6 +27,10 @@ vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ validateOverrideAccountId: vi.fn(), })); +vi.mock("@/lib/keys/getApiKeyDetails", () => ({ + getApiKeyDetails: vi.fn(), +})); + // Mock Supabase dependencies vi.mock("@/lib/supabase/account_emails/selectAccountEmails", () => ({ default: vi.fn(), diff --git a/lib/chat/__tests__/setupChatRequest.test.ts b/lib/chat/__tests__/setupChatRequest.test.ts index b715d768..96ca7195 100644 --- a/lib/chat/__tests__/setupChatRequest.test.ts +++ b/lib/chat/__tests__/setupChatRequest.test.ts @@ -46,6 +46,7 @@ describe("setupChatRequest", () => { it("returns a ChatConfig object with all required properties", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -61,6 +62,7 @@ describe("setupChatRequest", () => { it("calls getGeneralAgent with the body", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -72,6 +74,7 @@ describe("setupChatRequest", () => { it("uses instructions from routing decision as system prompt", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -100,6 +103,7 @@ describe("setupChatRequest", () => { it("passes tools to convertToModelMessages", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -123,6 +127,7 @@ describe("setupChatRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: manyMessages, }; @@ -137,6 +142,7 @@ describe("setupChatRequest", () => { it("provides a function that generates unique UUIDs", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -157,6 +163,7 @@ describe("setupChatRequest", () => { it("includes anthropic thinking configuration", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -169,6 +176,7 @@ describe("setupChatRequest", () => { it("includes google thinkingConfig", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -181,6 +189,7 @@ describe("setupChatRequest", () => { it("includes openai reasoning configuration", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -195,6 +204,7 @@ describe("setupChatRequest", () => { it("spreads routing decision properties into result", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -217,6 +227,7 @@ describe("setupChatRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -230,6 +241,7 @@ describe("setupChatRequest", () => { it("includes a prepareStep function", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -241,6 +253,7 @@ describe("setupChatRequest", () => { it("prepareStep returns options when no tool chain matches", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -256,6 +269,7 @@ describe("setupChatRequest", () => { it("prepareStep returns next tool when tool chain is triggered", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -284,6 +298,7 @@ describe("setupChatRequest", () => { it("prepareStep merges tool chain result with original options", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; diff --git a/lib/chat/__tests__/setupToolsForRequest.test.ts b/lib/chat/__tests__/setupToolsForRequest.test.ts index 5734a194..d4107b30 100644 --- a/lib/chat/__tests__/setupToolsForRequest.test.ts +++ b/lib/chat/__tests__/setupToolsForRequest.test.ts @@ -10,6 +10,22 @@ vi.mock("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({ StreamableHTTPClientTransport: vi.fn().mockImplementation(() => ({})), })); +vi.mock("@modelcontextprotocol/sdk/server/mcp.js", () => ({ + McpServer: vi.fn().mockImplementation(() => ({ + connect: vi.fn(), + })), +})); + +vi.mock("@modelcontextprotocol/sdk/inMemory.js", () => ({ + InMemoryTransport: { + createLinkedPair: vi.fn().mockReturnValue([{}, {}]), + }, +})); + +vi.mock("@/lib/mcp/tools", () => ({ + registerAllTools: vi.fn(), +})); + vi.mock("@/lib/agents/googleSheetsAgent", () => ({ getGoogleSheetsTools: vi.fn(), })); @@ -53,6 +69,7 @@ describe("setupToolsForRequest", () => { it("creates MCP client with correct URL", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -64,6 +81,7 @@ describe("setupToolsForRequest", () => { it("fetches tools from MCP client", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -72,12 +90,43 @@ describe("setupToolsForRequest", () => { expect(result).toHaveProperty("tool1"); expect(result).toHaveProperty("tool2"); }); + + it("passes accountId to MCP client via authenticated transport", async () => { + const body: ChatRequestBody = { + accountId: "account-123", + orgId: null, + messages: [{ id: "1", role: "user", content: "Hello" }], + }; + + await setupToolsForRequest(body); + + // Verify MCP client was created with a transport that includes auth info + expect(mockCreateMCPClient).toHaveBeenCalledWith( + expect.objectContaining({ + transport: expect.any(Object), + }), + ); + }); + + it("passes orgId to MCP client via authenticated transport", async () => { + const body: ChatRequestBody = { + accountId: "account-123", + orgId: "org-456", + messages: [{ id: "1", role: "user", content: "Hello" }], + }; + + await setupToolsForRequest(body); + + // Verify MCP client was created + expect(mockCreateMCPClient).toHaveBeenCalled(); + }); }); describe("Google Sheets tools integration", () => { it("calls getGoogleSheetsTools with request body", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Create a spreadsheet" }], }; @@ -91,6 +140,7 @@ describe("setupToolsForRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Create a spreadsheet" }], }; @@ -105,6 +155,7 @@ describe("setupToolsForRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Create a spreadsheet" }], }; @@ -120,6 +171,7 @@ describe("setupToolsForRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -145,6 +197,7 @@ describe("setupToolsForRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -161,6 +214,7 @@ describe("setupToolsForRequest", () => { it("excludes tools specified in excludeTools array", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], excludeTools: ["tool1"], }; @@ -176,6 +230,7 @@ describe("setupToolsForRequest", () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], excludeTools: ["tool1", "googlesheets_create"], }; @@ -191,6 +246,7 @@ describe("setupToolsForRequest", () => { it("returns all tools when excludeTools is undefined", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], }; @@ -203,6 +259,7 @@ describe("setupToolsForRequest", () => { it("returns all tools when excludeTools is empty array", async () => { const body: ChatRequestBody = { accountId: "account-123", + orgId: null, messages: [{ id: "1", role: "user", content: "Hello" }], excludeTools: [], }; diff --git a/lib/chat/__tests__/validateChatRequest.test.ts b/lib/chat/__tests__/validateChatRequest.test.ts index 50efe204..6057b55e 100644 --- a/lib/chat/__tests__/validateChatRequest.test.ts +++ b/lib/chat/__tests__/validateChatRequest.test.ts @@ -15,13 +15,19 @@ vi.mock("@/lib/accounts/validateOverrideAccountId", () => ({ validateOverrideAccountId: vi.fn(), })); +vi.mock("@/lib/keys/getApiKeyDetails", () => ({ + getApiKeyDetails: vi.fn(), +})); + import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; import { getAuthenticatedAccountId } from "@/lib/auth/getAuthenticatedAccountId"; import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; +import { getApiKeyDetails } from "@/lib/keys/getApiKeyDetails"; const mockGetApiKeyAccountId = vi.mocked(getApiKeyAccountId); const mockGetAuthenticatedAccountId = vi.mocked(getAuthenticatedAccountId); const mockValidateOverrideAccountId = vi.mocked(validateOverrideAccountId); +const mockGetApiKeyDetails = vi.mocked(getApiKeyDetails); // Helper to create mock NextRequest function createMockRequest(body: unknown, headers: Record = {}): Request { @@ -195,6 +201,59 @@ describe("validateChatRequest", () => { const json = await (result as NextResponse).json(); expect(json.status).toBe("error"); }); + + it("returns orgId for org API key", async () => { + mockGetApiKeyAccountId.mockResolvedValue("org-account-123"); + mockGetApiKeyDetails.mockResolvedValue({ + accountId: "org-account-123", + orgId: "org-account-123", + }); + + const request = createMockRequest( + { prompt: "Hello" }, + { "x-api-key": "org-api-key" }, + ); + + const result = await validateChatRequest(request as any); + + expect(result).not.toBeInstanceOf(NextResponse); + expect((result as any).accountId).toBe("org-account-123"); + expect((result as any).orgId).toBe("org-account-123"); + }); + + it("returns null orgId for personal API key", async () => { + mockGetApiKeyAccountId.mockResolvedValue("personal-account-123"); + mockGetApiKeyDetails.mockResolvedValue({ + accountId: "personal-account-123", + orgId: null, + }); + + 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"); + expect((result as any).orgId).toBeNull(); + }); + + it("returns null orgId for bearer token auth", async () => { + mockGetAuthenticatedAccountId.mockResolvedValue("jwt-account-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("jwt-account-456"); + expect((result as any).orgId).toBeNull(); + }); }); describe("accountId override", () => { diff --git a/lib/chat/validateChatRequest.ts b/lib/chat/validateChatRequest.ts index 164356a5..1808f771 100644 --- a/lib/chat/validateChatRequest.ts +++ b/lib/chat/validateChatRequest.ts @@ -6,6 +6,7 @@ import { getApiKeyAccountId } from "@/lib/auth/getApiKeyAccountId"; import { getAuthenticatedAccountId } from "@/lib/auth/getAuthenticatedAccountId"; import { validateOverrideAccountId } from "@/lib/accounts/validateOverrideAccountId"; import { getMessages } from "@/lib/messages/getMessages"; +import { getApiKeyDetails } from "@/lib/keys/getApiKeyDetails"; export const chatRequestSchema = z .object({ @@ -41,6 +42,7 @@ type BaseChatRequestBody = z.infer; export type ChatRequestBody = BaseChatRequestBody & { accountId: string; + orgId: string | null; }; /** @@ -64,7 +66,7 @@ export async function validateChatRequest( { status: "error", message: "Invalid input", - errors: validationResult.error.issues.map((err) => ({ + errors: validationResult.error.issues.map(err => ({ field: err.path.join("."), message: err.message, })), @@ -98,8 +100,9 @@ export async function validateChatRequest( ); } - // Authenticate and get accountId + // Authenticate and get accountId and orgId let accountId: string; + let orgId: string | null = null; if (hasApiKey) { // Validate API key authentication @@ -109,6 +112,12 @@ export async function validateChatRequest( } accountId = accountIdOrError; + // Get org context from API key details + const keyDetails = await getApiKeyDetails(apiKey!); + if (keyDetails) { + orgId = keyDetails.orgId; + } + // Handle accountId override for org API keys if (validatedBody.accountId) { const overrideResult = await validateOverrideAccountId({ @@ -121,7 +130,7 @@ export async function validateChatRequest( accountId = overrideResult.accountId; } } else { - // Validate bearer token authentication + // Validate bearer token authentication (no org context for JWT auth) const accountIdOrError = await getAuthenticatedAccountId(request); if (accountIdOrError instanceof NextResponse) { return accountIdOrError; @@ -132,11 +141,9 @@ export async function validateChatRequest( // Normalize chat content: // - If messages are provided, keep them as-is // - If only prompt is provided, convert it into a single user UIMessage - const hasMessages = - Array.isArray(validatedBody.messages) && validatedBody.messages.length > 0; + const hasMessages = Array.isArray(validatedBody.messages) && validatedBody.messages.length > 0; const hasPrompt = - typeof validatedBody.prompt === "string" && - validatedBody.prompt.trim().length > 0; + typeof validatedBody.prompt === "string" && validatedBody.prompt.trim().length > 0; if (!hasMessages && hasPrompt) { validatedBody.messages = getMessages(validatedBody.prompt); @@ -145,5 +152,6 @@ export async function validateChatRequest( return { ...validatedBody, accountId, + orgId, } as ChatRequestBody; } diff --git a/lib/emails/inbound/validateNewEmailMemory.ts b/lib/emails/inbound/validateNewEmailMemory.ts index 805f8c39..ed166814 100644 --- a/lib/emails/inbound/validateNewEmailMemory.ts +++ b/lib/emails/inbound/validateNewEmailMemory.ts @@ -76,6 +76,7 @@ export async function validateNewEmailMemory( const chatRequestBody: ChatRequestBody = { accountId, + orgId: null, messages: getMessages(emailText), roomId: finalRoomId, }; diff --git a/lib/prompts/getSystemPrompt.ts b/lib/prompts/getSystemPrompt.ts index 3d29afef..54964670 100644 --- a/lib/prompts/getSystemPrompt.ts +++ b/lib/prompts/getSystemPrompt.ts @@ -19,6 +19,7 @@ export function getSystemPrompt({ roomId, artistId, accountId, + orgId, email, knowledgeBaseText, artistInstruction, @@ -28,16 +29,18 @@ export function getSystemPrompt({ roomId?: string; artistId?: string; accountId: string; + orgId?: string | null; email?: string; knowledgeBaseText?: string; artistInstruction?: string; conversationName?: string; accountWithDetails?: AccountWithDetails; }): string { - let systemPrompt = `${SYSTEM_PROMPT} + let systemPrompt = `${SYSTEM_PROMPT} **IMPORTANT CONTEXT VALUES (use these exact values in tools):** - account_id: ${accountId || "Unknown"} (use this for ALL tools that require account_id parameter) + - organization_id: ${orgId || "None. Personal Account."} - artist_account_id: ${artistId} - active_account_email: ${email || "Unknown"} - active_conversation_id: ${roomId || "No ID"}