diff --git a/apps/api/src/routes/agent-business-context.test.ts b/apps/api/src/routes/agent-business-context.test.ts index 48ab23745..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" }, @@ -875,17 +914,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 +1055,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<{ diff --git a/apps/api/src/routes/agent.ts b/apps/api/src/routes/agent.ts index fd5ad192f..f37c5e6cb 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,36 +106,42 @@ 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 }); return toAgentErrorResponse(error); } const errorType = getErrorName(error); + const userId = apiKey ? `apikey:${apiKey.id}` : user?.id; 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: + apiKey?.organizationId ?? + body.organizationId ?? + activeOrganizationId ?? + null, + user_id: userId ?? 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 } : {}), + ...(userId ? { agent_user_id: userId } : {}), error_type: errorType, - source: event.source, + source, }); return toAgentErrorResponse(error); } @@ -499,28 +505,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 +534,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 +1074,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..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 = [ @@ -29,8 +31,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 +85,20 @@ function findCloseBrace(text: string, start: number): number { function parseComponent(json: string): ComponentSpec | null { try { - const parsed = JSON.parse(json) as unknown; + const parsed: JSONValue = 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 +162,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")} `; }