Merge pull request #485 from databuddy-analytics/staging

Land in-flight ai, db, and redis cleanup

Author: izadoesdev

package.json

@@ -54,6 +54,7 @@
     "db:deploy": "dotenv -- turbo run db:deploy",
     "db:seed": "dotenv -- sh -c 'cd packages/db && bun run db:seed \"$@\"' --",
     "email:dev": "dotenv -- sh -c 'cd packages/email && bun run dev'",
+    "email:segments": "dotenv -- sh -c 'cd packages/db && bun run email:segments \"$@\"' --",
     "sdk:build": "turbo run build --filter @databuddy/sdk --filter @databuddy/cache",
     "dev:dashboard": "dotenv -- turbo run dev --filter @databuddy/dashboard --filter @databuddy/api --filter @databuddy/insights",
     "dev:insights": "dotenv -- turbo run dev --filter @databuddy/insights",

packages/ai/src/ai/mcp/agent-tools.ts

@@ -147,7 +147,8 @@ export function createMcpAgentTools(
 				const data = await executeQuery(
 					queryRequest,
 					websiteDomain,
-					queryRequest.timezone
+					queryRequest.timezone,
+					options.abortSignal
 				);
 				return { data, rowCount: data.length, type: args.type };
 			},
@@ -182,6 +183,7 @@ Critical schema footguns: website id column is client_id (not website_id); times
 					sql: args.sql,
 					params: args.params,
 					toolName: "MCP Agent SQL",
+					abortSignal: options.abortSignal,
 				});
 			},
 		}),
@@ -233,6 +235,7 @@ Critical schema footguns: website id column is client_id (not website_id); times
 				const results = await executeBatch(requests, {
 					websiteDomain,
 					timezone: args.timezone ?? "UTC",
+					abortSignal: options.abortSignal,
 				});
 				return {
 					batch: true,

packages/ai/src/ai/mcp/investigate.ts

@@ -302,6 +302,12 @@ export async function runInvestigation(
 				"",
 				"Tool trace (what was actually queried and returned):",
 				compactTrace(trace.toolCalls),
+				...(trace.truncated
+					? [
+							"",
+							"NOTE: This investigation was stopped by a time limit before the agent finished gathering data. Base your memo only on the partial tool trace above. State in the verdict reason that the run was time-limited, and set confidence to 'low' unless the partial evidence is unambiguous.",
+						]
+					: []),
 			].join("\n"),
 		});
 

packages/ai/src/ai/mcp/run-agent.ts

@@ -5,7 +5,7 @@ import {
 	storeConversation,
 } from "../../lib/supermemory";
 import type { ApiKeyRow } from "@databuddy/api-keys/resolve";
-import type { LanguageModelUsage } from "ai";
+import type { LanguageModelUsage, StepResult, ToolSet } from "ai";
 import { ToolLoopAgent } from "ai";
 import { DatabuddyAgentUserError } from "../../agent/errors";
 import { getAILogger } from "../../lib/ai-logger";
@@ -55,6 +55,7 @@ export interface RunMcpAgentTraceResult {
 	answer: string;
 	steps: number;
 	toolCalls: McpAgentToolTrace[];
+	truncated?: boolean;
 	usage: LanguageModelUsage;
 }
 
@@ -114,11 +115,73 @@ export async function runMcpAgentWithTrace(
 			toolCalls: collectToolTrace(result.steps),
 			usage: result.totalUsage,
 		};
+	} catch (err) {
+		if (isInternalTimeoutAbort(err, options.abortSignal)) {
+			return await buildTruncatedTrace(prepared);
+		}
+		throw err;
 	} finally {
 		abort.cleanup();
 	}
 }
 
+function isInternalTimeoutAbort(
+	err: unknown,
+	externalSignal: AbortSignal | undefined
+): boolean {
+	return (
+		err instanceof Error &&
+		err.name === "AbortError" &&
+		!externalSignal?.aborted
+	);
+}
+
+async function buildTruncatedTrace(
+	prepared: Awaited<ReturnType<typeof prepareMcpAgentRun>>
+): Promise<RunMcpAgentTraceResult> {
+	const steps = prepared.capturedSteps;
+	const usage = aggregateStepUsage(steps);
+	await trackPreparedUsage(prepared, usage);
+	const stepCount = steps.length;
+	return {
+		answer: `The investigation reached its time budget after ${stepCount} step${stepCount === 1 ? "" : "s"} and was stopped before composing a final summary. The partial tool trace is the only evidence gathered.`,
+		steps: stepCount,
+		toolCalls: collectToolTrace(steps),
+		truncated: true,
+		usage,
+	};
+}
+
+function aggregateStepUsage(
+	steps: ReadonlyArray<{ usage: LanguageModelUsage }>
+): LanguageModelUsage {
+	let inputTokens = 0;
+	let outputTokens = 0;
+	let totalTokens = 0;
+	let noCacheTokens = 0;
+	let cacheReadTokens = 0;
+	let cacheWriteTokens = 0;
+	let textTokens = 0;
+	let reasoningTokens = 0;
+	for (const { usage } of steps) {
+		inputTokens += usage.inputTokens ?? 0;
+		outputTokens += usage.outputTokens ?? 0;
+		totalTokens += usage.totalTokens ?? 0;
+		noCacheTokens += usage.inputTokenDetails?.noCacheTokens ?? 0;
+		cacheReadTokens += usage.inputTokenDetails?.cacheReadTokens ?? 0;
+		cacheWriteTokens += usage.inputTokenDetails?.cacheWriteTokens ?? 0;
+		textTokens += usage.outputTokenDetails?.textTokens ?? 0;
+		reasoningTokens += usage.outputTokenDetails?.reasoningTokens ?? 0;
+	}
+	return {
+		inputTokens,
+		outputTokens,
+		totalTokens,
+		inputTokenDetails: { noCacheTokens, cacheReadTokens, cacheWriteTokens },
+		outputTokenDetails: { textTokens, reasoningTokens },
+	};
+}
+
 export async function* streamMcpAgentText(
 	options: RunMcpAgentOptions
 ): AsyncGenerator<string> {
@@ -189,12 +252,7 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) {
 	const selectedModelId =
 		options.modelOverride ?? getDefaultAgentModelId(source);
 
-	const apiKeyId =
-		options.apiKey &&
-		typeof options.apiKey === "object" &&
-		"id" in options.apiKey
-			? (options.apiKey as { id: string }).id
-			: null;
+	const apiKeyId = options.apiKey?.id ?? null;
 
 	const billingCustomerId =
 		options.billingMode === "skip"
@@ -265,6 +323,7 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) {
 	mcpTelemetryMetadata["tcc.sessionId"] = sessionId;
 
 	const ai = getAILogger();
+	const capturedSteps: StepResult<ToolSet>[] = [];
 	const agent = new ToolLoopAgent({
 		model: ai.wrap(config.model),
 		instructions,
@@ -273,6 +332,9 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) {
 		stopWhen: config.stopWhen,
 		temperature: config.temperature,
 		experimental_context: config.experimental_context,
+		onStepFinish: (step) => {
+			capturedSteps.push(step);
+		},
 		experimental_telemetry: {
 			isEnabled: true,
 			functionId: `databuddy.${source}.ask`,
@@ -296,6 +358,7 @@ async function prepareMcpAgentRun(options: RunMcpAgentOptions) {
 		agent,
 		apiKeyId,
 		billingCustomerId,
+		capturedSteps,
 		memoryUserId,
 		mcpUserId,
 		messages,

packages/ai/src/ai/tools/execute-sql-query.ts

@@ -32,12 +32,14 @@ export async function executeAgentSqlForWebsite({
 	sql,
 	params,
 	toolName = "Execute SQL Tool",
+	abortSignal,
 }: {
 	websiteId: string;
 	websiteDomain?: string;
 	sql: string;
 	params?: Record<string, unknown>;
 	toolName?: string;
+	abortSignal?: AbortSignal;
 }): Promise<QueryResult> {
 	const validation = validateAgentSQL(sql);
 	if (!validation.valid) {
@@ -66,7 +68,8 @@ export async function executeAgentSqlForWebsite({
 			max_result_bytes: 50_000_000,
 			result_overflow_mode: "break",
 			use_query_cache: 0,
-		}
+		},
+		abortSignal
 	);
 
 	return result.data.length > MAX_MODEL_ROWS
@@ -104,6 +107,7 @@ export const executeSqlQueryTool = tool({
 			websiteDomain: resolved.domain,
 			sql,
 			params,
+			abortSignal: options.abortSignal,
 		});
 	},
 });

packages/ai/src/ai/tools/utils/query.ts

@@ -73,13 +73,15 @@ export async function executeTimedQuery<T extends Record<string, unknown>>(
 	sql: string,
 	params: Record<string, unknown> = {},
 	logContext?: Record<string, unknown>,
-	clickhouseSettings?: Record<string, string | number>
+	clickhouseSettings?: Record<string, string | number>,
+	abortSignal?: AbortSignal
 ): Promise<QueryResult<T>> {
 	const logger = createToolLogger(toolName);
 	const queryStart = Date.now();
 
 	try {
 		const raw = await chQuery<T>(sql, params, {
+			abort_signal: abortSignal,
 			readonly: true,
 			...(clickhouseSettings && { clickhouse_settings: clickhouseSettings }),
 		});

packages/ai/src/query/batch-executor.ts

@@ -15,6 +15,7 @@ interface BatchResult {
 	type: string;
 }
 interface BatchOptions {
+	abortSignal?: AbortSignal;
 	timezone?: string;
 	websiteDomain?: string | null;
 }
@@ -252,7 +253,7 @@ async function runSingle(
 			{ ...req, timezone: opts?.timezone ?? req.timezone },
 			opts?.websiteDomain
 		);
-		const data = await builder.execute();
+		const data = await builder.execute(opts?.abortSignal);
 
 		mergeWideEvent({
 			query_type: req.type,
@@ -422,6 +423,7 @@ export async function executeBatch(
 				({ req }) => QueryBuilders[req.type]?.noCache
 			);
 			const rawRows = await chQuery(sql, params, {
+				abort_signal: opts?.abortSignal,
 				clickhouse_settings: getClickHouseQuerySettings(groupNoCache),
 			});
 

packages/ai/src/query/index.ts

@@ -83,8 +83,9 @@ function createBuilder(
 export const executeQuery = async (
 	request: QueryRequest,
 	websiteDomain?: string | null,
-	timezone?: string
-) => createBuilder(request, websiteDomain, timezone).execute();
+	timezone?: string,
+	abortSignal?: AbortSignal
+) => createBuilder(request, websiteDomain, timezone).execute(abortSignal);
 
 export const compileQuery = (
 	request: QueryRequest,

packages/ai/src/query/simple-builder.ts

@@ -1104,9 +1104,10 @@ export class SimpleQueryBuilder {
 		return this.request.offset ? ` OFFSET ${this.request.offset}` : "";
 	}
 
-	async execute(): Promise<Record<string, unknown>[]> {
+	async execute(abortSignal?: AbortSignal): Promise<Record<string, unknown>[]> {
 		const { sql, params } = this.compile();
 		const rawData = await chQuery<Record<string, unknown>>(sql, params, {
+			abort_signal: abortSignal,
 			clickhouse_settings: getClickHouseQuerySettings(this.config.noCache),
 		});
 		return applyPlugins(rawData, this.config, this.websiteDomain);

packages/db/package.json

@@ -18,6 +18,7 @@
     "db:push": "dotenv -- drizzle-kit push",
     "db:studio": "dotenv -- drizzle-kit studio",
     "db:seed": "tsx src/seed.ts",
+    "email:segments": "dotenv -- bun src/scripts/export-user-email-segments.ts",
     "clickhouse:init": "tsx src/clickhouse/setup.ts",
     "test": "bun test src",
     "e2e:db:create": "bun ./scripts/e2e-db-lifecycle.ts create",