feat(autumn): track workspace AI message usage

Author: urjitc

.dev.vars.example

@@ -27,6 +27,9 @@ BETTER_AUTH_URL="http://localhost:3000"
 # AI chat and title generation (Vercel AI Gateway).
 # AI_GATEWAY_API_KEY=""
 
+# Autumn usage analytics. When unset, usage tracking is skipped.
+# AUTUMN_SECRET_KEY=""
+
 # Web search and URL/PDF extraction (Firecrawl).
 # FIRECRAWL_API_KEY=""
 

docs/ENVIRONMENT.md

@@ -37,6 +37,7 @@ The dev server boots and lets you sign in with just the two variables below. Eve
 | --- | --- | --- |
 | `GOOGLE_CLIENT_ID` / `GOOGLE_CLIENT_SECRET` | Google sign-in | Google sign-in fails; use **Continue as guest** (shown in local dev) |
 | `AI_GATEWAY_API_KEY` | AI chat / title generation (Vercel AI Gateway) | AI calls error when invoked |
+| `AUTUMN_SECRET_KEY` | Autumn usage analytics for completed AI chat messages | Usage tracking is skipped |
 | `FIRECRAWL_API_KEY` | Web search + URL/PDF extraction (Firecrawl) | Those calls error when invoked |
 | `TCC_API_KEY` | Agent observability export (The Context Company) | AI chat still works; TCC export is skipped |
 | `LLAMA_CLOUD_API_KEY` | Document extraction (LlamaParse) | Those calls error when invoked |

package.json

@@ -104,6 +104,7 @@
 		"@tiptap/y-tiptap": "^3.0.5",
 		"agents": "^0.17.0",
 		"ai": "^6.0.208",
+		"autumn-js": "^1.2.34",
 		"better-auth": "^1.6.19",
 		"class-variance-authority": "^0.7.1",
 		"clsx": "^2.1.1",

pnpm-lock.yaml

@@ -20,6 +20,7 @@ import { createCompactFunction } from "agents/experimental/memory/utils";
 import { generateText, type LanguageModel, type ToolSet } from "ai";
 
 import type { AIInspectorSnapshot } from "#/features/workspaces/ai/ai-inspector";
+import type { AIThreadContext } from "#/features/workspaces/ai/ai-thread-metadata";
 import { AIThreadTelemetryRecorder } from "#/features/workspaces/ai/ai-thread-telemetry-recorder";
 import {
 	createAIThreadTools,
@@ -34,8 +35,13 @@ import {
 	DEFAULT_WORKSPACE_AI_CHAT_MODEL_ID,
 	getWorkspaceAiChatModel,
 	resolveWorkspaceAiChatModelId,
+	type WorkspaceAiChatModelId,
 } from "#/features/workspaces/ai/models";
 import type { UserAIStore } from "#/features/workspaces/ai/user-ai-agents";
+import {
+	checkWorkspaceAiMessageAccess,
+	trackWorkspaceAiMessageUsage,
+} from "#/integrations/autumn/workspace-ai-usage";
 
 const AI_THREAD_CHAT_RECOVERY_NO_PROGRESS_TIMEOUT_MS = 90_000;
 const AI_THREAD_CHAT_RECOVERY_TERMINAL_MESSAGE =
@@ -53,6 +59,11 @@ type AIThreadRunSettlement =
 			kind: "failed";
 	  };
 
+interface AIThreadUsageContext {
+	modelId: WorkspaceAiChatModelId;
+	thread: AIThreadContext;
+}
+
 export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
 	return class AIThread extends Think<Cloudflare.Env> {
 		override chatRecovery = {
@@ -75,6 +86,7 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
 		override sendReasoning = false;
 		private shouldRefreshSessionPrompt = false;
 		private activeRunStartedAt: number | undefined;
+		private activeUsageContext: AIThreadUsageContext | undefined;
 		private readonly telemetry = new AIThreadTelemetryRecorder({
 			env: this.env,
 			host: this,
@@ -153,6 +165,24 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
 			}
 
 			const modelId = resolveWorkspaceAiChatModelId(ctx.body?.modelId);
+
+			if (!ctx.continuation) {
+				const access = await checkWorkspaceAiMessageAccess({
+					env: this.env,
+					modelId,
+					userId: thread.userId,
+				});
+
+				if (!access.allowed) {
+					throw new Error("Usage limit reached");
+				}
+
+				this.activeUsageContext = {
+					modelId,
+					thread,
+				};
+			}
+
 			const system = getAIThreadSystemPromptForWorkspace(ctx.system, thread.promptScope, {
 				timeZone: getBodyString(ctx.body, "timeZone"),
 				workspaceAiContext: ctx.body?.workspaceAiContext,
@@ -204,6 +234,7 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
 
 		override async onChatResponse(result: ChatResponseResult) {
 			this.telemetry.recordTurnFinished(result);
+			this._trackCompletedMessageUsage(result);
 			if (!this._shouldSettleRunAfterResponse(result)) {
 				await this._refreshSessionPromptIfNeeded();
 				return;
@@ -307,13 +338,35 @@ export function createAIThreadClass(getUserAIStore: () => typeof UserAIStore) {
 				}
 
 				this.activeRunStartedAt = undefined;
+				this.activeUsageContext = undefined;
 			} catch (error) {
 				onError(error);
 			} finally {
 				await this._refreshSessionPromptIfNeeded();
 			}
 		}
 
+		private _trackCompletedMessageUsage(result: ChatResponseResult) {
+			if (result.continuation || result.status !== "completed") {
+				return;
+			}
+
+			const usageContext = this.activeUsageContext;
+			if (!usageContext) {
+				return;
+			}
+
+			void this.keepAliveWhile(() =>
+				trackWorkspaceAiMessageUsage({
+					env: this.env,
+					modelId: usageContext.modelId,
+					threadId: usageContext.thread.id,
+					userId: usageContext.thread.userId,
+					workspaceId: usageContext.thread.workspaceId,
+				}),
+			);
+		}
+
 		private async _handleChatRecoveryExhausted(ctx: ChatRecoveryExhaustedContext) {
 			await this._recordAuxiliaryError({
 				error: new Error(AI_THREAD_CHAT_RECOVERY_TERMINAL_MESSAGE),