diff --git a/packages/ai/src/ai/mcp/business-context-delivery.test.ts b/packages/ai/src/ai/mcp/business-context-delivery.test.ts index 42a49cecc..da074101a 100644 --- a/packages/ai/src/ai/mcp/business-context-delivery.test.ts +++ b/packages/ai/src/ai/mcp/business-context-delivery.test.ts @@ -52,11 +52,12 @@ mock.module("@databuddy/auth", () => ({ auth: { api: { getSession: async () => session } }, })); let allowed = true; +let sites = [site]; const accessible = mock(async (auth: AccessibleWebsitesAuth) => allowed && auth.organizationId === "org-synthetic" && (auth.apiKey || auth.user) - ? [site] + ? sites : [] ); mock.module("../../lib/accessible-websites", () => ({ @@ -173,9 +174,74 @@ beforeEach(() => { model.doStreamCalls.length = 0; session = null; allowed = true; + sites = [site]; }); describe("canonical business context at the native shared-agent model boundary", () => { + it.each([ + ["Reports.Example.com", "reports.example.com"], + ["reports.example.com", "REPORTS.EXAMPLE.COM"], + ["Reports.Example.com", "http://reports.example.com"], + ["reports.example.com", "HTTPS://REPORTS.EXAMPLE.COM"], + ["WWW.Example.com", "https://www.example.com"], + ["Reports.Example.com:8443", "HtTpS://reports.example.com:8443"], + ])("accepts stored %s selected as %s through ask, stream and trace", async ( + domain, + websiteDomain + ) => { + sites = [{ ...site, domain }]; + for (const websiteId of [undefined, site.id]) { + const input = { ...options, websiteId, websiteDomain }; + await askDatabuddyAgent(input); + await traceDatabuddyAgent(input); + for await (const _chunk of streamDatabuddyAgent(input)) { + /* consume native stream */ + } + } + const calls = [...model.doGenerateCalls, ...model.doStreamCalls]; + expect(calls).toHaveLength(6); + for (const call of calls) { + expect(JSON.stringify(call.prompt)).toContain(meaning); + } + expect(read).toHaveBeenCalledTimes(6); + expect(read.mock.calls.every(([id]) => id === "org-synthetic")).toBe(true); + }); + it("retains the selected ID when accessible sites share a normalized domain", async () => { + sites = [ + { ...site, id: "earlier-site", domain: "REPORTS.EXAMPLE.COM" }, + { ...site, domain: "Reports.Example.com" }, + ]; + await askDatabuddyAgent({ + ...options, + websiteId: site.id, + websiteDomain: "https://reports.example.com", + }); + expect(JSON.stringify(model.doGenerateCalls[0].prompt)).toContain(meaning); + expect(read).toHaveBeenCalledTimes(1); + }); + it("rejects unsupported domain rewrites and mismatched accessible IDs before profile or model access", async () => { + sites = [ + site, + { ...site, id: "other-allowed-site", domain: "other.example.com" }, + ]; + for (const websiteDomain of [ + "www.reports.example.com", + "reports.example.com/", + " https://reports.example.com", + "https://reports.example.com/path", + "https://reports.example.com.evil.example", + "//reports.example.com", + "reports.example.com:443", + "ftp://reports.example.com", + "https://other.example.com", + ]) { + await expect( + askDatabuddyAgent({ ...options, websiteId: site.id, websiteDomain }) + ).rejects.toThrow("not accessible"); + } + expect(read).not.toHaveBeenCalled(); + expect(model.doGenerateCalls).toHaveLength(0); + }); for (const source of ["slack", "mcp", "dashboard"] as const) { it(`${source}: pairs absent/saved context through ask, stream and trace`, async () => { for (const present of [false, true]) { diff --git a/packages/ai/src/ai/mcp/run-agent.ts b/packages/ai/src/ai/mcp/run-agent.ts index a4e934923..ddfb0c8ad 100644 --- a/packages/ai/src/ai/mcp/run-agent.ts +++ b/packages/ai/src/ai/mcp/run-agent.ts @@ -12,6 +12,7 @@ import { DatabuddyAgentUserError } from "../../agent/errors"; import { getAILogger } from "../../lib/ai-logger"; import { getAccessibleWebsites } from "../../lib/accessible-websites"; import { loadOrganizationBusinessContext } from "../../lib/organization-business-context"; +import { matchesWebsiteDomain } from "../../lib/website-domain"; import { mergeWideEvent } from "../../lib/tracing"; import { ensureAgentCreditsAvailable, @@ -262,14 +263,15 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) { organizationId, user: session?.user.id === mcpUserId ? session.user : null, }); + const websiteDomain = options.websiteDomain; // A caller-supplied site must not bind another organization's brief or tools. if ( (options.websiteId && !accessibleWebsites.some((site) => site.id === options.websiteId)) || - (options.websiteDomain && + (websiteDomain && !accessibleWebsites.some( (site) => - site.domain === options.websiteDomain && + matchesWebsiteDomain(site.domain, websiteDomain) && (!options.websiteId || site.id === options.websiteId) )) ) { diff --git a/packages/ai/src/ai/mcp/tool-context.test.ts b/packages/ai/src/ai/mcp/tool-context.test.ts index a47562830..bc41b77e9 100644 --- a/packages/ai/src/ai/mcp/tool-context.test.ts +++ b/packages/ai/src/ai/mcp/tool-context.test.ts @@ -1,6 +1,30 @@ import { describe, expect, it, mock } from "bun:test"; +import type { WebsiteSummary } from "../../lib/accessible-websites"; const permission = mock(async () => ({ success: true })); +const sites: WebsiteSummary[] = [ + { + id: "site", + domain: "Reports.Example.com", + name: "Reports", + createdAt: null, + isPublic: false, + }, + { + id: "www-site", + domain: "WWW.Example.com", + name: "WWW", + createdAt: null, + isPublic: false, + }, + { + id: "port-site", + domain: "Reports.Example.com:8443", + name: "Port", + createdAt: null, + isPublic: false, + }, +]; mock.module("@databuddy/auth", () => ({ websitesApi: { hasPermission: permission }, })); @@ -12,7 +36,7 @@ mock.module("../../lib/website-utils", () => ({ getCachedWebsite: async () => null, })); mock.module("../../lib/accessible-websites", () => ({ - getAccessibleWebsites: async () => [], + getAccessibleWebsites: async () => sites, })); mock.module("@databuddy/api-keys/resolve", () => ({ hasKeyScope: () => true, @@ -20,7 +44,36 @@ mock.module("@databuddy/api-keys/resolve", () => ({ })); mock.module("@databuddy/redis", () => ({ getRedisCache: () => null })); -const { ensureWebsiteAccess } = await import("./tool-context"); +const { ensureWebsiteAccess, resolveWebsiteId } = await import("./tool-context"); + +describe("MCP domain selector compatibility", () => { + it.each([ + ["reports.example.com", "site"], + ["REPORTS.EXAMPLE.COM", "site"], + ["http://reports.example.com", "site"], + ["HTTPS://REPORTS.EXAMPLE.COM", "site"], + ["https://www.example.com", "www-site"], + ["HtTpS://reports.example.com:8443", "port-site"], + ])("resolves %s to %s", async (websiteDomain, expected) => { + expect( + await resolveWebsiteId({ websiteDomain }, { apiKey: null, userId: null }) + ).toBe(expected); + }); + it.each([ + "www.reports.example.com", + "reports.example.com/", + " https://reports.example.com", + "https://reports.example.com/path", + "https://reports.example.com.evil.example", + "//reports.example.com", + "reports.example.com:443", + "ftp://reports.example.com", + ])("does not rewrite unsupported selector %s", async (websiteDomain) => { + expect( + await resolveWebsiteId({ websiteDomain }, { apiKey: null, userId: null }) + ).toBeInstanceOf(Error); + }); +}); describe("shared agent's business-context organization boundary", () => { it("rejects a site in another organization even if the session could read both", async () => { diff --git a/packages/ai/src/ai/mcp/tool-context.ts b/packages/ai/src/ai/mcp/tool-context.ts index e1140e60c..4936a8619 100644 --- a/packages/ai/src/ai/mcp/tool-context.ts +++ b/packages/ai/src/ai/mcp/tool-context.ts @@ -11,8 +11,8 @@ import { websitesApi } from "@databuddy/auth"; import { getRedisCache } from "@databuddy/redis"; import type { AppContext } from "../config/context"; import { getCachedWebsite, validateWebsite } from "../../lib/website-utils"; +import { matchesWebsiteDomain } from "../../lib/website-domain"; -const PROTOCOL_RE = /^https?:\/\//; const ACCESSIBLE_WEBSITES_TTL_SEC = 30; const ACCESSIBLE_WEBSITES_KEY_PREFIX = "mcp:accessible_websites:v2:"; @@ -135,9 +135,11 @@ export async function resolveWebsiteId( const list = await getCachedAccessibleWebsites(principal); - if (input.websiteDomain) { - const domain = input.websiteDomain.toLowerCase().replace(PROTOCOL_RE, ""); - const match = list.find((w) => w.domain?.toLowerCase() === domain); + const domain = input.websiteDomain; + if (domain) { + const match = list.find((website) => + matchesWebsiteDomain(website.domain, domain) + ); if (match) { return match.id; } diff --git a/packages/ai/src/lib/website-domain.ts b/packages/ai/src/lib/website-domain.ts new file mode 100644 index 000000000..85f6ca5e7 --- /dev/null +++ b/packages/ai/src/lib/website-domain.ts @@ -0,0 +1,9 @@ +const PROTOCOL_RE = /^https?:\/\//; + +/** Match MCP domain selectors without changing hostname, port or path semantics. */ +export function matchesWebsiteDomain( + domain: string | null, + input: string +): boolean { + return domain?.toLowerCase() === input.toLowerCase().replace(PROTOCOL_RE, ""); +}