From 3aa710a2335fc496a3362c19708ff2db6302d427 Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Mon, 5 Oct 2026 02:29:15 +0300 Subject: [PATCH 1/5] refactor(ai): drop run-option mapping between agent layers Agent calls spread the caller's options instead of remapping 18 fields, duplicate defaults and renamed fields are gone, the org fallback lives in one error helper, and casts in the component parser are removed. Built prompts are byte-identical. --- apps/api/src/routes/agent.ts | 80 +++++----- apps/api/src/routes/mcp.ts | 11 +- apps/slack/src/slack/blocks.test.ts | 2 +- .../ai/src/agent/conversation-history.test.ts | 11 +- packages/ai/src/agent/errors.ts | 2 +- packages/ai/src/agent/index.ts | 143 ++++++------------ packages/ai/src/agent/render.ts | 22 +-- packages/ai/src/ai/mcp/run-agent.ts | 124 ++++++--------- packages/ai/src/ai/prompts/shared.ts | 6 +- 9 files changed, 151 insertions(+), 250 deletions(-) diff --git a/apps/api/src/routes/agent.ts b/apps/api/src/routes/agent.ts index fd5ad192f..89f333e55 100644 --- a/apps/api/src/routes/agent.ts +++ b/apps/api/src/routes/agent.ts @@ -14,9 +14,9 @@ import { AGENT_THINKING_LEVELS, AGENT_TIERS } from "@databuddy/ai/agents/types"; import { type AgentModelKey, models } from "@databuddy/ai/config/models"; import { AgentError, - type AgentRequestInput, askDatabuddyAgent, type DatabuddyAgentActor, + type DatabuddyAgentOptions, prepareAgentRequest, streamDatabuddyAgent, toAgentErrorResponse, @@ -106,13 +106,14 @@ function agentActor( function agentFailure( error: unknown, - event: { - chatId: string; - organizationId: string | null; - source: "api" | "dashboard"; - userId: string | null; - websiteId?: string | null; - } + source: "api" | "dashboard", + chatId: string, + { + activeOrganizationId, + apiKey, + body, + user, + }: AgentAuth & { body: { organizationId?: string; websiteId?: string } } ): Response { if (error instanceof AgentError && error.status < 500) { mergeWideEvent({ agent_rejected: error.code }); @@ -121,21 +122,25 @@ function agentFailure( const errorType = getErrorName(error); trackAgentEvent("agent_activity", { action: "chat_error", - source: event.source, + source, agent_type: AGENT_TYPE, error_type: errorType, - organization_id: event.organizationId, - user_id: event.userId, - website_id: event.websiteId ?? null, + organization_id: + body.organizationId ?? + activeOrganizationId ?? + apiKey?.organizationId ?? + null, + user_id: user?.id ?? null, + website_id: body.websiteId ?? null, }); captureError(error, { agent_error: true, agent_type: AGENT_TYPE, - agent_chat_id: event.chatId, - ...(event.websiteId ? { agent_website_id: event.websiteId } : {}), - ...(event.userId ? { agent_user_id: event.userId } : {}), + agent_chat_id: chatId, + ...(body.websiteId ? { agent_website_id: body.websiteId } : {}), + ...(user?.id ? { agent_user_id: user.id } : {}), error_type: errorType, - source: event.source, + source, }); return toAgentErrorResponse(error); } @@ -499,28 +504,26 @@ export const agent = new Elysia({ prefix: "/v1/agent" }) mergeWideEvent({ agent_chat_id: conversationId, source: "api" }); try { - const agentRequest: AgentRequestInput = { + const principal = await prepareAgentRequest({ actor: agentActor( { activeOrganizationId, apiKey, user }, agentHeaders ), organizationId: body.organizationId, rateLimit: "agent:ask", - }; - const principal = await prepareAgentRequest(agentRequest); + }); mergeWideEvent({ organization_id: principal.organizationId, agent_user_id: user?.id ?? `apikey:${apiKey?.id}`, }); - const options = { - ...agentRequest, + const options: DatabuddyAgentOptions = { abortSignal: request.signal, conversationId, input: body.question, - mutationMode: "dry-run" as const, - output: "markdown" as const, + mutationMode: "dry-run", + output: "markdown", principal, - source: "api" as const, + source: "api", timezone: body.timezone, }; if (body.stream) { @@ -530,15 +533,11 @@ export const agent = new Elysia({ prefix: "/v1/agent" }) } return await askDatabuddyAgent(options); } catch (error) { - return agentFailure(error, { - chatId: conversationId, - organizationId: - body.organizationId ?? - activeOrganizationId ?? - apiKey?.organizationId ?? - null, - source: "api", - userId: user?.id ?? null, + return agentFailure(error, "api", conversationId, { + activeOrganizationId, + apiKey, + body, + user, }); } }, @@ -1074,16 +1073,11 @@ export const agent = new Elysia({ prefix: "/v1/agent" }) } return response; } catch (error) { - return agentFailure(error, { - chatId, - organizationId: - body.organizationId ?? - activeOrganizationId ?? - apiKey?.organizationId ?? - null, - source: "dashboard", - userId: user?.id ?? null, - websiteId: body.websiteId, + return agentFailure(error, "dashboard", chatId, { + activeOrganizationId, + apiKey, + body, + user, }); } })(); diff --git a/apps/api/src/routes/mcp.ts b/apps/api/src/routes/mcp.ts index ed2685f95..59c252474 100644 --- a/apps/api/src/routes/mcp.ts +++ b/apps/api/src/routes/mcp.ts @@ -1,5 +1,4 @@ import { isApiKeyPresent } from "@databuddy/api-keys/resolve"; -import { resolveAgentOrganizationId } from "@databuddy/ai/agent"; import { createMcpErrorResponse, createMcpUnauthorizedResponse, @@ -249,12 +248,10 @@ export const mcp = new Elysia({ name: "mcp" }) user, apiKey, oauthAccessToken: null, - organizationId: resolveAgentOrganizationId({ - activeOrganizationId: user - ? session?.session.activeOrganizationId - : null, - apiKey, - }), + organizationId: + apiKey?.organizationId ?? + (user ? session?.session.activeOrganizationId : null) ?? + null, }; }) .onBeforeHandle(({ user, apiKey, oauthAccessToken, set }) => { diff --git a/apps/slack/src/slack/blocks.test.ts b/apps/slack/src/slack/blocks.test.ts index b22e97093..ca06ec9c1 100644 --- a/apps/slack/src/slack/blocks.test.ts +++ b/apps/slack/src/slack/blocks.test.ts @@ -5,7 +5,7 @@ import { componentsToBlocks, componentToBlocks, } from "@/slack/blocks"; -import { buildAnalyticsInstructionsForMcp } from "../../../../packages/ai/src/ai/prompts/analytics"; +import { buildAnalyticsInstructionsForMcp } from "@databuddy/ai/prompts/analytics"; function firstBlock(spec: ComponentSpec): Block { const block = componentToBlocks(spec)[0]; diff --git a/packages/ai/src/agent/conversation-history.test.ts b/packages/ai/src/agent/conversation-history.test.ts index 084702fa5..4d157dcc4 100644 --- a/packages/ai/src/agent/conversation-history.test.ts +++ b/packages/ai/src/agent/conversation-history.test.ts @@ -23,7 +23,7 @@ mock.module("./slack-relevance", () => ({ function answer(options: RunMcpAgentOptions): string { runs.push(options); - return `Answer: ${options.question}`; + return `Answer: ${options.input}`; } mock.module("../ai/mcp/run-agent", () => ({ runMcpAgent: async (options: RunMcpAgentOptions) => answer(options), @@ -73,7 +73,6 @@ function principalFor(apiKey: ApiKeyRow | null): AgentPrincipal { }; } const options: DatabuddyAgentOptions = { - actor: { type: "api_key", apiKey: key }, principal: principalFor(key), conversationId: "slack-T_TEST-C_TEST-111_000", input: "First question", @@ -108,8 +107,8 @@ test.each([ }); await invoke({ ...options, input: "Third question" }); - expect(runs.map((run) => run.priorMessages?.length ?? 0)).toEqual([0, 2, 4]); - expect(runs[1]?.priorMessages).toEqual([ + expect(runs.map((run) => run.history?.length ?? 0)).toEqual([0, 2, 4]); + expect(runs[1]?.history).toEqual([ { role: "user", content: "First question" }, { role: "assistant", content: "Answer: First question" }, ]); @@ -130,7 +129,7 @@ test.each([ ] as const)("isolates Slack history by %s", async (_scope, isolated) => { await askDatabuddyAgent(options); await askDatabuddyAgent({ ...options, ...isolated }); - expect(runs.at(-1)?.priorMessages).toBeUndefined(); + expect(runs.at(-1)?.history).toBeUndefined(); }); test.each([ @@ -145,7 +144,7 @@ test.each([ await askDatabuddyAgent(input); await askDatabuddyAgent({ ...input, memoryUserId: "slack-T_TEST-U_B" }); await askDatabuddyAgent(input); - expect(runs.map((run) => run.priorMessages?.length ?? 0)).toEqual([0, 0, 2]); + expect(runs.map((run) => run.history?.length ?? 0)).toEqual([0, 0, 2]); }); test("renders component JSON as markdown for markdown output", async () => { diff --git a/packages/ai/src/agent/errors.ts b/packages/ai/src/agent/errors.ts index d318dc735..73e53cbf7 100644 --- a/packages/ai/src/agent/errors.ts +++ b/packages/ai/src/agent/errors.ts @@ -34,7 +34,7 @@ const AGENT_ERRORS = { }, } as const; -export type AgentErrorCode = keyof typeof AGENT_ERRORS; +type AgentErrorCode = keyof typeof AGENT_ERRORS; export class AgentError extends Error { readonly code: AgentErrorCode; diff --git a/packages/ai/src/agent/index.ts b/packages/ai/src/agent/index.ts index eb94f82d2..b6b40a3c5 100644 --- a/packages/ai/src/agent/index.ts +++ b/packages/ai/src/agent/index.ts @@ -30,18 +30,13 @@ import { mergeWideEvent } from "../lib/tracing"; import { matchesWebsiteDomain } from "../lib/website-domain"; import { AgentError } from "./errors"; import { - type AgentOutput, ComponentStreamSplitter, componentToPlainText, splitAgentText, } from "./render"; export type { ConversationMessage } from "../ai/mcp/conversation-store"; -export { - AgentError, - type AgentErrorCode, - toAgentErrorResponse, -} from "./errors"; +export { AgentError, toAgentErrorResponse } from "./errors"; export { classifySlackThreadReplyRelevance, type SlackThreadReplyRelevance, @@ -72,7 +67,7 @@ export type DatabuddyAgentActor = userId: string; }; -export interface AgentRequestInput { +interface AgentRequestInput { actor: DatabuddyAgentActor; billingMode?: DatabuddyAgentBillingMode; organizationId?: string | null; @@ -92,7 +87,10 @@ export interface AgentPrincipal { website: WebsiteSummary | null; } -export interface DatabuddyAgentOptions extends AgentRequestInput { +export type DatabuddyAgentOptions = ( + | AgentRequestInput + | { principal: AgentPrincipal } +) & { abortSignal?: AbortSignal; conversationId?: string; history?: ConversationMessage[]; @@ -104,14 +102,13 @@ export interface DatabuddyAgentOptions extends AgentRequestInput { onToolEvent?: (toolNames: string[]) => void; /** Streaming only: called once after completion and usage settlement. */ onToolTrace?: (trace: DatabuddyAgentToolTrace[]) => void; - output?: AgentOutput; + output?: "markdown"; persistConversation?: boolean; - principal?: AgentPrincipal; slackContext?: DatabuddyAgentSlackContext | null; source?: DatabuddyAgentSource; timeoutMs?: number; timezone?: string; -} +}; export interface DatabuddyAgentResult { answer: string; @@ -128,42 +125,26 @@ export interface DatabuddyAgentTraceResult extends DatabuddyAgentResult { const AGENT_RATE_LIMIT_PER_MINUTE = 30; -export function resolveAgentOrganizationId(input: { - activeOrganizationId?: string | null; - apiKey: ApiKeyRow | null; - requestedOrganizationId?: string | null; -}): string | null { - const keyOrganizationId = input.apiKey?.organizationId; +export async function prepareAgentRequest( + input: AgentRequestInput +): Promise { + const { actor } = input; + const apiKey = actor.type === "api_key" ? actor.apiKey : null; + const userId = actor.userId ?? apiKey?.userId ?? null; if ( - input.apiKey && - input.requestedOrganizationId && - input.requestedOrganizationId !== keyOrganizationId + apiKey && + input.organizationId && + input.organizationId !== apiKey.organizationId ) { throw new AgentError( "access_denied", "The API key does not belong to this organization." ); } - return ( - input.requestedOrganizationId ?? - keyOrganizationId ?? - input.activeOrganizationId ?? - null - ); -} - -export async function prepareAgentRequest( - input: AgentRequestInput -): Promise { - const { actor } = input; - const apiKey = actor.type === "api_key" ? actor.apiKey : null; - const userId = actor.userId ?? apiKey?.userId ?? null; - const organizationId = resolveAgentOrganizationId({ - activeOrganizationId: - actor.type === "session" ? actor.activeOrganizationId : null, - apiKey, - requestedOrganizationId: input.organizationId, - }); + const organizationId = + input.organizationId ?? + apiKey?.organizationId ?? + (actor.type === "session" ? actor.activeOrganizationId : null); if (!organizationId) { throw new AgentError("workspace_required"); } @@ -256,29 +237,26 @@ function selectRequestedWebsite( export async function askDatabuddyAgent( options: DatabuddyAgentOptions ): Promise { - const prepared = await prepareDatabuddyAgentCall(options); - const answer = renderAnswer( - await runMcpAgent(toRunOptions(options, prepared)), - options.output - ); + const run = await prepareDatabuddyAgentCall(options); + const answer = renderAnswer(await runMcpAgent(run), run.output); - await persistAgentConversation(options, prepared, answer); + await persistAgentConversation(run, answer); - return { answer, conversationId: prepared.conversationId }; + return { answer, conversationId: run.conversationId }; } export async function traceDatabuddyAgent( options: DatabuddyAgentOptions ): Promise { - const prepared = await prepareDatabuddyAgentCall(options); - const result = await runMcpAgentWithTrace(toRunOptions(options, prepared)); - const answer = renderAnswer(result.answer, options.output); + const run = await prepareDatabuddyAgentCall(options); + const result = await runMcpAgentWithTrace(run); + const answer = renderAnswer(result.answer, run.output); - await persistAgentConversation(options, prepared, answer); + await persistAgentConversation(run, answer); return { answer, - conversationId: prepared.conversationId, + conversationId: run.conversationId, steps: result.steps, toolCalls: result.toolCalls, usage: result.usage, @@ -288,16 +266,14 @@ export async function traceDatabuddyAgent( export async function* streamDatabuddyAgent( options: DatabuddyAgentOptions ): AsyncGenerator { - const prepared = await prepareDatabuddyAgentCall(options); + const run = await prepareDatabuddyAgentCall(options); const splitter = - options.output === "markdown" + run.output === "markdown" ? new ComponentStreamSplitter(componentToPlainText) : null; let answer = ""; - for await (const chunk of streamMcpAgentText( - toRunOptions(options, prepared) - )) { + for await (const chunk of streamMcpAgentText(run)) { const text = splitter ? splitter.push(chunk) : chunk; if (text) { answer += text; @@ -310,17 +286,22 @@ export async function* streamDatabuddyAgent( yield tail; } - await persistAgentConversation(options, prepared, answer); + await persistAgentConversation(run, answer); } -function renderAnswer(answer: string, output: AgentOutput | undefined): string { +function renderAnswer(answer: string, output: "markdown" | undefined): string { return output === "markdown" ? splitAgentText(answer, componentToPlainText).text : answer; } -async function prepareDatabuddyAgentCall(options: DatabuddyAgentOptions) { - const principal = options.principal ?? (await prepareAgentRequest(options)); +async function prepareDatabuddyAgentCall( + options: DatabuddyAgentOptions +): Promise { + const principal = + "principal" in options + ? options.principal + : await prepareAgentRequest(options); const conversationId = options.conversationId ?? crypto.randomUUID(); const memoryUserId = options.memoryUserId ?? principal.userId; // Slack threads belong to the integration; personal memory stays speaker-scoped. @@ -335,6 +316,7 @@ async function prepareDatabuddyAgentCall(options: DatabuddyAgentOptions) { )); return { + ...options, conversationId, conversationUserId, history: history.length > 0 ? history : undefined, @@ -344,45 +326,20 @@ async function prepareDatabuddyAgentCall(options: DatabuddyAgentOptions) { }; } -function toRunOptions( - options: DatabuddyAgentOptions, - prepared: Awaited> -): RunMcpAgentOptions { - return { - abortSignal: options.abortSignal, - conversationId: prepared.conversationId, - historyInput: options.historyInput, - memoryUserId: prepared.memoryUserId, - modelOverride: options.modelOverride, - mutationMode: options.mutationMode, - onToolEvent: options.onToolEvent, - onToolTrace: options.onToolTrace, - principal: prepared.principal, - priorMessages: prepared.history, - question: options.input, - slackContext: options.slackContext, - source: prepared.source, - storeMemory: options.persistConversation !== false, - timeoutMs: options.timeoutMs, - timezone: options.timezone, - }; -} - async function persistAgentConversation( - options: DatabuddyAgentOptions, - prepared: Awaited>, + run: RunMcpAgentOptions, answer: string ): Promise { - if (options.persistConversation === false) { + if (run.persistConversation === false) { return; } await appendToConversation( - prepared.conversationId, - prepared.conversationUserId, - prepared.principal.apiKey, - options.historyInput ?? options.input, + run.conversationId, + run.conversationUserId, + run.principal.apiKey, + run.historyInput ?? run.input, answer.trim(), - prepared.history + run.history ); } diff --git a/packages/ai/src/agent/render.ts b/packages/ai/src/agent/render.ts index ffadac418..e1ab783ab 100644 --- a/packages/ai/src/agent/render.ts +++ b/packages/ai/src/agent/render.ts @@ -29,8 +29,6 @@ export const AGENT_COMPONENT_TYPES = [ export type AgentComponentType = (typeof AGENT_COMPONENT_TYPES)[number]; -export type AgentOutput = "components" | "markdown"; - export interface ComponentSpec { type: string; [key: string]: unknown; @@ -85,19 +83,20 @@ function findCloseBrace(text: string, start: number): number { function parseComponent(json: string): ComponentSpec | null { try { - const parsed = JSON.parse(json) as unknown; + const parsed: unknown = JSON.parse(json); if ( parsed && typeof parsed === "object" && !Array.isArray(parsed) && - typeof (parsed as Record).type === "string" && - KNOWN_COMPONENT_TYPES.has( - (parsed as Record).type as string - ) + "type" in parsed && + typeof parsed.type === "string" && + KNOWN_COMPONENT_TYPES.has(parsed.type) ) { - return parsed as ComponentSpec; + return { ...parsed, type: parsed.type }; } - } catch {} + } catch { + // Malformed model text stays in the response. + } return null; } @@ -161,11 +160,6 @@ export class ComponentStreamSplitter { this.#buffer = rest.slice(1); } } - - if (final) { - return emitted; - } - this.#buffer = ""; return emitted; } diff --git a/packages/ai/src/ai/mcp/run-agent.ts b/packages/ai/src/ai/mcp/run-agent.ts index 6c3d04964..aec9a53a9 100644 --- a/packages/ai/src/ai/mcp/run-agent.ts +++ b/packages/ai/src/ai/mcp/run-agent.ts @@ -6,7 +6,7 @@ import { storeConversation, } from "../../lib/supermemory"; import type { LanguageModelUsage, StepResult, ToolSet } from "ai"; -import type { AgentPrincipal } from "../../agent"; +import type { AgentPrincipal, DatabuddyAgentOptions } from "../../agent"; import { createConversationAgent } from "../agents/conversation"; import { getAILogger } from "../../lib/ai-logger"; import { loadOrganizationBusinessContext } from "../../lib/organization-business-context"; @@ -14,33 +14,19 @@ import { captureError } from "../../lib/tracing"; import { trackAgentUsageAndBill } from "../agents/execution"; import { createMcpAgentConfig } from "../agents/mcp"; import { type AgentSource, getDefaultAgentModelId } from "../config/models"; -import type { AppMutationMode } from "../config/context"; import { prependBackgroundContext } from "../prompts/context"; -import type { DatabuddyAgentSlackContext } from "./slack-context"; const DEFAULT_MCP_AGENT_TIMEOUT_MS = 45_000; const EMPTY_ANSWER = "No answer was generated from the gathered evidence. Try a narrower question: one metric, one segment, or one time range."; -export interface RunMcpAgentOptions { - abortSignal?: AbortSignal; - conversationId?: string; - historyInput?: string; - memoryUserId?: string | null; - modelOverride?: string | null; - mutationMode?: AppMutationMode; - onToolEvent?: (toolNames: string[]) => void; - /** Called once after a stream completes and its usage has been settled. */ - onToolTrace?: (trace: McpAgentToolTrace[]) => void; +export type RunMcpAgentOptions = DatabuddyAgentOptions & { + conversationId: string; + conversationUserId: string | null; + memoryUserId: string | null; principal: AgentPrincipal; - priorMessages?: Array<{ role: "user" | "assistant"; content: string }>; - question: string; - slackContext?: DatabuddyAgentSlackContext | null; - source?: AgentSource; - storeMemory?: boolean; - timeoutMs?: number; - timezone?: string; -} + source: AgentSource; +}; export interface McpAgentToolTrace { index: number; @@ -72,9 +58,7 @@ export async function runMcpAgent( await trackPreparedUsage(prepared, result.totalUsage); const answer = result.text.trim() || EMPTY_ANSWER; - if (options.storeMemory !== false) { - storePreparedConversation(prepared, answer); - } + storePreparedConversation(prepared, answer); return answer; } finally { @@ -97,9 +81,7 @@ export async function runMcpAgentWithTrace( await trackPreparedUsage(prepared, result.totalUsage); const answer = result.text.trim() || EMPTY_ANSWER; - if (options.storeMemory !== false) { - storePreparedConversation(prepared, answer); - } + storePreparedConversation(prepared, answer); return { answer, @@ -217,9 +199,7 @@ export async function* streamMcpAgentText( } options.onToolTrace?.(collectToolTrace(prepared.capturedSteps)); - if (options.storeMemory !== false) { - storePreparedConversation(prepared, answer); - } + storePreparedConversation(prepared, answer); } finally { abort.cleanup(); await settleRemainingUsage(prepared); @@ -259,25 +239,19 @@ function createRunAbortController(options: RunMcpAgentOptions): { } async function prepareMcpAgentRun(options: RunMcpAgentOptions) { - const { principal } = options; - const { accessibleWebsites, apiKey, organizationId, website } = principal; - const sessionId = options.conversationId ?? crypto.randomUUID(); - const historyInput = options.historyInput ?? options.question; - const mcpUserId = principal.userId; - const memoryUserId = options.memoryUserId ?? mcpUserId; - const source = options.source ?? "mcp"; - const selectedModelId = - options.modelOverride ?? getDefaultAgentModelId(source); - const apiKeyId = apiKey?.id ?? null; + const { conversationId, memoryUserId, principal, source } = options; + const { accessibleWebsites, apiKey, organizationId, userId, website } = + principal; + const historyInput = options.historyInput ?? options.input; const [config, memoryCtx, businessContext] = await Promise.all([ createMcpAgentConfig({ billingCustomerId: principal.billingCustomerId, requestHeaders: principal.requestHeaders, apiKey, - userId: mcpUserId, + userId, timezone: options.timezone, - chatId: sessionId, + chatId: conversationId, latestUserMessage: historyInput, modelOverride: options.modelOverride, memoryUserId, @@ -290,7 +264,7 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) { websiteId: website?.id, }), isMemoryEnabled() - ? getMemoryContext(historyInput, memoryUserId, apiKeyId) + ? getMemoryContext(historyInput, memoryUserId, apiKey?.id ?? null) : Promise.resolve(null), loadOrganizationBusinessContext({ organizationId, @@ -320,40 +294,27 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) { authType: apiKey ? "api_key" : "session", timezone: options.timezone ?? "UTC", "tcc.conversational": "true", - ...(mcpUserId && { userId: mcpUserId }), + ...(userId && { userId }), organizationId, - "tcc.sessionId": sessionId, + "tcc.sessionId": conversationId, }, }, } ); const messages = prependBackgroundContext( - [ - ...(options.priorMessages ?? []), - { role: "user", content: options.question }, - ], + [...(options.history ?? []), { role: "user", content: options.input }], [businessContext, memoryCtx ? formatMemoryForPrompt(memoryCtx) : ""] ); return { agent, - apiKeyId, - billingCustomerId: principal.billingCustomerId, - billingAccess: principal.billingAccess, capturedSteps, historyInput, - usageSettlementAttempted: false, - memoryUserId, - mcpUserId, messages, - modelId: selectedModelId, - mutationMode: options.mutationMode, - organizationId, - sessionId, - source, - websiteDomain: website?.domain ?? undefined, - websiteId: website?.id, + modelId: options.modelOverride ?? getDefaultAgentModelId(source), + options, + usageSettlementAttempted: false, }; } @@ -366,6 +327,7 @@ async function trackPreparedUsage( } // An uncertain charge must not be replayed by cleanup. prepared.usageSettlementAttempted = true; + const { conversationId, principal, source } = prepared.options; await trackAgentUsageAndBill({ usage: { ...usage, @@ -374,12 +336,12 @@ async function trackPreparedUsage( : {}), }, modelId: prepared.modelId, - source: prepared.source, - organizationId: prepared.organizationId, - userId: prepared.mcpUserId, - chatId: prepared.sessionId, - billingCustomerId: prepared.billingCustomerId, - billingAccess: prepared.billingAccess, + source, + organizationId: principal.organizationId, + userId: principal.userId, + chatId: conversationId, + billingCustomerId: principal.billingCustomerId, + billingAccess: principal.billingAccess, }); } @@ -401,8 +363,8 @@ async function settleRemainingUsage( // Preserve the model error or consumer cancellation that entered cleanup. captureError(error, { agent_usage_billing_error: true, - agent_source: prepared.source, - agent_chat_id: prepared.sessionId, + agent_source: prepared.options.source, + agent_chat_id: prepared.options.conversationId, }); } } @@ -428,27 +390,29 @@ function collectToolTrace( } function storePreparedConversation( - prepared: Awaited>, + { historyInput, options }: Awaited>, answer: string ): void { if ( - prepared.mutationMode === "dry-run" || - !asksToRemember(prepared.historyInput) + options.persistConversation === false || + options.mutationMode === "dry-run" || + !asksToRemember(historyInput) ) { return; } + const { website } = options.principal; storeConversation( [ - { role: "user", content: prepared.historyInput }, + { role: "user", content: historyInput }, { role: "assistant", content: answer }, ], - prepared.memoryUserId, - prepared.apiKeyId, + options.memoryUserId, + options.principal.apiKey?.id ?? null, { - ...(prepared.websiteDomain ? { domain: prepared.websiteDomain } : {}), - metadata: { source: prepared.source }, - conversationId: prepared.sessionId, - websiteId: prepared.websiteId, + ...(website?.domain ? { domain: website.domain } : {}), + metadata: { source: options.source }, + conversationId: options.conversationId, + websiteId: website?.id, } ); } diff --git a/packages/ai/src/ai/prompts/shared.ts b/packages/ai/src/ai/prompts/shared.ts index 28514197d..d68b64dd2 100644 --- a/packages/ai/src/ai/prompts/shared.ts +++ b/packages/ai/src/ai/prompts/shared.ts @@ -75,14 +75,10 @@ const MUTATION_RULES = `**Mutations:** ${FEEDBACK_TOOL_RULES}`; -function mutationRules(mode: AppMutationMode = "allow"): string { - return mode === "dry-run" ? READ_ONLY_RULE : MUTATION_RULES; -} - export function agentRules(mode?: AppMutationMode): string { return `${COMMON_AGENT_RULES} -${[TOOL_ROUTING, mutationRules(mode), DATA_INTEGRITY_RULES, ATTRIBUTION_RULES, ANALYSIS_RULES].join("\n\n")} +${[TOOL_ROUTING, mode === "dry-run" ? READ_ONLY_RULE : MUTATION_RULES, DATA_INTEGRITY_RULES, ATTRIBUTION_RULES, ANALYSIS_RULES].join("\n\n")} `; } From 9a5bac179851b1c82654df86bf12871d06962203 Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:18:46 +0000 Subject: [PATCH 2/5] test(api): assert prepared agent identity --- .../src/routes/agent-business-context.test.ts | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/apps/api/src/routes/agent-business-context.test.ts b/apps/api/src/routes/agent-business-context.test.ts index 48ab23745..ea31e5939 100644 --- a/apps/api/src/routes/agent-business-context.test.ts +++ b/apps/api/src/routes/agent-business-context.test.ts @@ -875,17 +875,13 @@ describe("selected agent identity through native HTTP", () => { expect(response.status).toBe(200); expect(await response.text()).toContain("Synthetic answer."); const options = (stream ? state.stream : state.ask).mock.calls[0]?.[0]; - expect(options.actor).toMatchObject({ - type: "api_key", - apiKey: state.apiKey, - userId: null, - }); expect(options.principal).toMatchObject({ + apiKey: state.apiKey, organizationId: selectedKey.organizationId, userId: keyOwner, }); - expect(options.actor.requestHeaders.get("cookie")).toBeNull(); - expect(options.actor.requestHeaders.get("x-api-key")).toBe( + expect(options.principal.requestHeaders.get("cookie")).toBeNull(); + expect(options.principal.requestHeaders.get("x-api-key")).toBe( mixedHeaders["x-api-key"] ); expect(options.abortSignal).toBeInstanceOf(AbortSignal); @@ -1020,9 +1016,13 @@ describe("selected agent identity through native HTTP", () => { const response = await askResponse({}, { cookie: mixedHeaders.cookie }); expect(response.status).toBe(200); expect(await response.text()).toContain("Synthetic answer."); - const actor = state.ask.mock.calls[0]?.[0].actor; - expect(actor).toMatchObject({ type: "session", userId: "user-synthetic" }); - expect(actor.requestHeaders.get("cookie")).toBe(mixedHeaders.cookie); + const principal = state.ask.mock.calls[0]?.[0].principal; + expect(principal).toMatchObject({ + apiKey: null, + organizationId: "org-synthetic", + userId: "user-synthetic", + }); + expect(principal.requestHeaders.get("cookie")).toBe(mixedHeaders.cookie); }); it.each<{ From 5b988362b19ebd2208207522985344ad89720828 Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:25:29 +0000 Subject: [PATCH 3/5] fix(api): attribute agent failures to selected key --- apps/api/src/routes/agent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/api/src/routes/agent.ts b/apps/api/src/routes/agent.ts index 89f333e55..da32ac225 100644 --- a/apps/api/src/routes/agent.ts +++ b/apps/api/src/routes/agent.ts @@ -126,9 +126,9 @@ function agentFailure( agent_type: AGENT_TYPE, error_type: errorType, organization_id: + apiKey?.organizationId ?? body.organizationId ?? activeOrganizationId ?? - apiKey?.organizationId ?? null, user_id: user?.id ?? null, website_id: body.websiteId ?? null, From 7bb6502cde86531d797078fda16219bef3c63a19 Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Tue, 6 Oct 2026 13:37:09 +0000 Subject: [PATCH 4/5] fix(api): retain key identity in agent failure telemetry --- .../src/routes/agent-business-context.test.ts | 41 ++++++++++++++++++- apps/api/src/routes/agent.ts | 5 ++- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/apps/api/src/routes/agent-business-context.test.ts b/apps/api/src/routes/agent-business-context.test.ts index ea31e5939..687020bbf 100644 --- a/apps/api/src/routes/agent-business-context.test.ts +++ b/apps/api/src/routes/agent-business-context.test.ts @@ -17,6 +17,7 @@ const state = vi.hoisted(() => ({ contexts: [] as Record[], accessible: vi.fn(), errors: vi.fn(), + events: vi.fn(), sessionUserId: "user-synthetic" as string | null, sessionOrg: "org-synthetic" as string | null, apiKeyId: null as string | null, @@ -224,7 +225,9 @@ vi.mock("@databuddy/ai/agents/cache", () => ({ vi.mock("@databuddy/ai/lib/ai-logger", () => ({ getAILogger: () => ({ wrap: (model: MockLanguageModelV3) => model }), })); -vi.mock("@databuddy/ai/lib/databuddy", () => ({ trackAgentEvent: () => {} })); +vi.mock("@databuddy/ai/lib/databuddy", () => ({ + trackAgentEvent: state.events, +})); vi.mock("@databuddy/ai/lib/tracing", () => ({ captureError: state.errors, mergeWideEvent: () => {}, @@ -308,6 +311,7 @@ beforeEach(() => { state.contexts.length = 0; state.read.mockReset(); state.errors.mockReset(); + state.events.mockReset(); state.accessible.mockReset(); state.read.mockImplementation(async () => ({ profile: state.profile, @@ -858,6 +862,41 @@ const mixedHeaders = { "x-api-key": "dbdy_inert_selected_key", }; +describe("failure telemetry identity", () => { + it.each([ + { route: "ask", identity: "key", key: selectedKey }, + { route: "chat", identity: "key", key: selectedKey }, + { route: "ask", identity: "session", key: null }, + { route: "chat", identity: "session", key: null }, + ])("attributes failure telemetry to the selected $identity on /$route", async ({ + route, + key, + }) => { + state.apiKey = key; + const organizationId = key?.organizationId ?? "org-synthetic"; + const userId = key ? `apikey:${key.id}` : "user-synthetic"; + state.billing.mockRejectedValueOnce( + new BillingUnavailableError("Synthetic billing outage") + ); + const response = await (route === "ask" + ? ask({ organizationId }) + : chat({ organizationId })); + expect(response.status).toBe(503); + expect(state.events).toHaveBeenCalledWith( + "agent_activity", + expect.objectContaining({ + action: "chat_error", + organization_id: organizationId, + user_id: userId, + }) + ); + expect(state.errors).toHaveBeenCalledWith( + expect.any(Error), + expect.objectContaining({ agent_user_id: userId }) + ); + }); +}); + describe("selected agent identity through native HTTP", () => { it.each([ { stream: false, keyOwner: "key-owner-synthetic" }, diff --git a/apps/api/src/routes/agent.ts b/apps/api/src/routes/agent.ts index da32ac225..f37c5e6cb 100644 --- a/apps/api/src/routes/agent.ts +++ b/apps/api/src/routes/agent.ts @@ -120,6 +120,7 @@ function agentFailure( return toAgentErrorResponse(error); } const errorType = getErrorName(error); + const userId = apiKey ? `apikey:${apiKey.id}` : user?.id; trackAgentEvent("agent_activity", { action: "chat_error", source, @@ -130,7 +131,7 @@ function agentFailure( body.organizationId ?? activeOrganizationId ?? null, - user_id: user?.id ?? null, + user_id: userId ?? null, website_id: body.websiteId ?? null, }); captureError(error, { @@ -138,7 +139,7 @@ function agentFailure( agent_type: AGENT_TYPE, agent_chat_id: chatId, ...(body.websiteId ? { agent_website_id: body.websiteId } : {}), - ...(user?.id ? { agent_user_id: user.id } : {}), + ...(userId ? { agent_user_id: userId } : {}), error_type: errorType, source, }); From cc653b9a3c4181655dfbe7670c9db0ed61edee2f Mon Sep 17 00:00:00 2001 From: iza <59828082+izadoesdev@users.noreply.github.com> Date: Tue, 6 Oct 2026 14:14:29 +0000 Subject: [PATCH 5/5] refactor(ai): use the JSON value type for component parsing --- packages/ai/src/agent/render.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/agent/render.ts b/packages/ai/src/agent/render.ts index e1ab783ab..e373dc4de 100644 --- a/packages/ai/src/agent/render.ts +++ b/packages/ai/src/agent/render.ts @@ -1,3 +1,5 @@ +import type { JSONValue } from "@ai-sdk/provider"; + const COMPONENT_START = '{"type":"'; const CHART_COMPONENT_TYPES = [ @@ -83,7 +85,7 @@ function findCloseBrace(text: string, start: number): number { function parseComponent(json: string): ComponentSpec | null { try { - const parsed: unknown = JSON.parse(json); + const parsed: JSONValue = JSON.parse(json); if ( parsed && typeof parsed === "object" &&