Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions packages/ai/src/ai/tools/annotations.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
import { tool } from "ai";
import dayjs from "dayjs";
import { z } from "zod";
import { callRPCProcedure, createToolLogger, getAppContext } from "./utils";
import {
callRPCProcedure,
createToolLogger,
getAppContext,
resolveToolWebsite,
} from "./utils";

const logger = createToolLogger("Annotations Tools");

Expand Down Expand Up @@ -95,10 +100,11 @@ export function createAnnotationTools() {
execute: async ({ websiteId, chartType, chartContext }, options) => {
const context = getAppContext(options);
try {
const resolved = resolveToolWebsite(context, websiteId);
const result = await callRPCProcedure(
"annotations",
"list",
{ websiteId, chartType, chartContext },
{ websiteId: resolved.websiteId, chartType, chartContext },
context
);
return {
Expand Down Expand Up @@ -141,6 +147,8 @@ export function createAnnotationTools() {
) => {
const context = getAppContext(options);
try {
const resolved = resolveToolWebsite(context, websiteId);

if (!confirmed) {
const dateRangePreview = `${chartContext.dateRange.start_date} to ${chartContext.dateRange.end_date} (${chartContext.dateRange.granularity})`;

Expand All @@ -149,7 +157,7 @@ export function createAnnotationTools() {
message:
"Please review the annotation details below and confirm if you want to create it:",
annotation: {
websiteId,
websiteId: resolved.websiteId,
chartType,
dateRange: dateRangePreview,
annotationType,
Expand All @@ -170,7 +178,7 @@ export function createAnnotationTools() {
"annotations",
"create",
{
websiteId,
websiteId: resolved.websiteId,
chartType,
chartContext,
annotationType,
Expand Down
39 changes: 39 additions & 0 deletions packages/ai/src/ai/tools/utils/context.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -86,4 +86,43 @@ describe("resolveToolWebsite", () => {

expect(() => resolveToolWebsite(ctx)).toThrow(/multiple websites/);
});

it("resolves a domain name to the matching website id in the accessible set", () => {
const ctx = makeCtx({
accessibleWebsites: [
{ id: "web_a", domain: "app.example.com", name: null, isPublic: null, createdAt: null },
],
});

expect(resolveToolWebsite(ctx, "app.example.com")).toEqual({
websiteId: "web_a",
domain: "app.example.com",
});
});

it("resolves a domain name to the context websiteId when accessibleWebsites is absent", () => {
const ctx = makeCtx({
websiteId: "web_ctx",
websiteDomain: "ctx.com",
});

expect(resolveToolWebsite(ctx, "ctx.com")).toEqual({
websiteId: "web_ctx",
domain: "ctx.com",
});
});

it("still rejects a domain name that does not match any accessible website or the context domain", () => {
const ctx = makeCtx({
accessibleWebsites: [
{ id: "web_a", domain: "a.com", name: null, isPublic: null, createdAt: null },
],
websiteId: "web_ctx",
websiteDomain: "ctx.com",
});

expect(() => resolveToolWebsite(ctx, "unknown.com")).toThrow(
/not in this workspace/
);
});
});
26 changes: 20 additions & 6 deletions packages/ai/src/ai/tools/utils/context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,29 @@ export function resolveToolWebsite(
(id === ctx.websiteId ? ctx.websiteDomain : undefined);

if (inputWebsiteId) {
const isAccessible =
// Fast path: exact UUID match in the accessible set or the context default
const isAccessibleById =
accessible.some((w) => w.id === inputWebsiteId) ||
inputWebsiteId === ctx.websiteId;
if (!isAccessible) {
throw new Error(
`Website "${inputWebsiteId}" is not in this workspace. Call list_websites to see available websites.`
);
if (isAccessibleById) {
return { websiteId: inputWebsiteId, domain: domainFor(inputWebsiteId) };
}
return { websiteId: inputWebsiteId, domain: domainFor(inputWebsiteId) };

// Slow path: the model may have passed a domain name instead of an id.
// Check the accessible list first, then the single-site context default.
const byDomain = accessible.find(
(w) => w.domain != null && w.domain === inputWebsiteId
);
if (byDomain) {
return { websiteId: byDomain.id, domain: byDomain.domain ?? undefined };

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 The ?? undefined is dead code here: the find predicate already guards w.domain != null, so byDomain.domain is always a non-null string at this point. The nullish coalescing just adds noise.

Suggested change
return { websiteId: byDomain.id, domain: byDomain.domain ?? undefined };
return { websiteId: byDomain.id, domain: byDomain.domain };

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

}
if (inputWebsiteId === ctx.websiteDomain && ctx.websiteId) {
Comment on lines +40 to +46

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Domain matching is case-sensitive via ===, but RFC 1035 treats hostnames as case-insensitive. If the AI model passes a domain with different casing than what is stored (e.g. App.Quiver.AI vs app.quiver.ai), the slow-path lookup silently falls through to the error throw, leaving the original bug intact for that casing variant. Normalising both sides to lowercase costs nothing and makes the fix robust.

Suggested change
const byDomain = accessible.find(
(w) => w.domain != null && w.domain === inputWebsiteId
);
if (byDomain) {
return { websiteId: byDomain.id, domain: byDomain.domain ?? undefined };
}
if (inputWebsiteId === ctx.websiteDomain && ctx.websiteId) {
const normalised = inputWebsiteId.toLowerCase();
const byDomain = accessible.find(
(w) => w.domain != null && w.domain.toLowerCase() === normalised
);
if (byDomain) {
return { websiteId: byDomain.id, domain: byDomain.domain ?? undefined };
}
if (ctx.websiteDomain?.toLowerCase() === normalised && ctx.websiteId) {

return { websiteId: ctx.websiteId, domain: ctx.websiteDomain };
}

throw new Error(
`Website "${inputWebsiteId}" is not in this workspace. Call list_websites to see available websites.`
);
}

const fallbackId = ctx.defaultWebsiteId ?? ctx.websiteId;
Expand Down
Loading