diff --git a/apps/server/src/app.test.ts b/apps/server/src/app.test.ts index 3ea053a25..c16e1f641 100644 --- a/apps/server/src/app.test.ts +++ b/apps/server/src/app.test.ts @@ -55,6 +55,7 @@ function createTestDeps() { providerCatalogService: {} as any, productEventService: { track: vi.fn(async () => undefined), + trackGeneration: vi.fn(async () => undefined), countDistinctUsersByFeature: vi.fn(async () => []), }, configKV: { diff --git a/apps/server/src/routes/admin/voice-packs/route.test.ts b/apps/server/src/routes/admin/voice-packs/route.test.ts index e692b7e99..f7aa7560d 100644 --- a/apps/server/src/routes/admin/voice-packs/route.test.ts +++ b/apps/server/src/routes/admin/voice-packs/route.test.ts @@ -49,6 +49,7 @@ function createService() { function createProductEventService(): ProductEventService { return { track: vi.fn(async () => undefined), + trackGeneration: vi.fn(async () => undefined), countDistinctUsersByFeature: vi.fn(async () => []), } } diff --git a/apps/server/src/routes/audio-speech-ws/route.test.ts b/apps/server/src/routes/audio-speech-ws/route.test.ts index 9523fa6a2..16bfc0429 100644 --- a/apps/server/src/routes/audio-speech-ws/route.test.ts +++ b/apps/server/src/routes/audio-speech-ws/route.test.ts @@ -186,6 +186,7 @@ function makeFakeDeps(overrides: { } const productEventService = { track: vi.fn(async () => undefined), + trackGeneration: vi.fn(async () => undefined), countDistinctUsersByFeature: vi.fn(async () => []), } const configKV = { diff --git a/apps/server/src/routes/openai/v1/analytics.ts b/apps/server/src/routes/openai/v1/analytics.ts new file mode 100644 index 000000000..2799b1fc1 --- /dev/null +++ b/apps/server/src/routes/openai/v1/analytics.ts @@ -0,0 +1,14 @@ +import type { AiGenerationAppSurface } from '../../../services/domain/product-events' + +export const AIRI_CHAT_SESSION_ID_HEADER = 'x-airi-session-id' +export const AIRI_CHAT_ROUND_ID_HEADER = 'x-airi-round-id' +export const AIRI_CHAT_APP_SURFACE_HEADER = 'x-airi-app-surface' + +const CLIENT_CHAT_ANALYTICS_SURFACES = new Set(['web', 'mobile', 'electron']) + +export function resolveChatAnalyticsSurface(value: string | undefined): AiGenerationAppSurface { + if (CLIENT_CHAT_ANALYTICS_SURFACES.has(value as AiGenerationAppSurface)) + return value as AiGenerationAppSurface + + return 'server' +} diff --git a/apps/server/src/routes/openai/v1/index.ts b/apps/server/src/routes/openai/v1/index.ts index ddaeaadbc..e9bd2a46f 100644 --- a/apps/server/src/routes/openai/v1/index.ts +++ b/apps/server/src/routes/openai/v1/index.ts @@ -5,6 +5,12 @@ import type { LlmTracingDeps, V1RouteDeps } from './types' import { authGuard } from '../../../middlewares/auth' import { configGuard } from '../../../middlewares/config-guard' +import { + AIRI_CHAT_APP_SURFACE_HEADER, + AIRI_CHAT_ROUND_ID_HEADER, + AIRI_CHAT_SESSION_ID_HEADER, + resolveChatAnalyticsSurface, +} from './analytics' import { createV1Gateway } from './gateway' import { chatCompletionsRateLimit } from './middlewares' import { chatCompletions } from './operations/chat-completions' @@ -41,7 +47,9 @@ export function createV1Routes(input: CreateV1RoutesDeps) { return { userId: user.id, body, - sessionId: c.req.header('x-airi-session-id'), + sessionId: c.req.header(AIRI_CHAT_SESSION_ID_HEADER), + roundId: c.req.header(AIRI_CHAT_ROUND_ID_HEADER), + appSurface: resolveChatAnalyticsSurface(c.req.header(AIRI_CHAT_APP_SURFACE_HEADER)), abortSignal: c.req.raw.signal, } }, @@ -66,7 +74,7 @@ export function createV1Routes(input: CreateV1RoutesDeps) { return { userId: user.id, body, - sessionId: c.req.header('x-airi-session-id'), + sessionId: c.req.header(AIRI_CHAT_SESSION_ID_HEADER), abortSignal: c.req.raw.signal, } }, diff --git a/apps/server/src/routes/openai/v1/operations/chat-completions/index.ts b/apps/server/src/routes/openai/v1/operations/chat-completions/index.ts index a30642bd2..caaf4c3ba 100644 --- a/apps/server/src/routes/openai/v1/operations/chat-completions/index.ts +++ b/apps/server/src/routes/openai/v1/operations/chat-completions/index.ts @@ -1,5 +1,6 @@ import type { CapabilityAliasRoute } from '../../../../../schemas/provider-catalog' import type { UsageInfo } from '../../../../../services/domain/billing/billing' +import type { AiGenerationAppSurface } from '../../../../../services/domain/product-events' import type { GatewayCallback } from '../../gateway' import type { V1RouteDeps } from '../../types' @@ -20,9 +21,25 @@ export interface ChatCompletionsOperationRequest { userId: string body: Record sessionId?: string + roundId?: string + appSurface: AiGenerationAppSurface abortSignal?: AbortSignal } +interface GenerationCaptureInput { + deps: V1RouteDeps + userId: string + requestId: string + sessionId?: string + roundId?: string + appSurface: AiGenerationAppSurface + generationModel: string + routeCtxProvider: string + usage: UsageInfo + durationMs: number + stream: boolean +} + export function chatCompletions(deps: V1RouteDeps): GatewayCallback<'chat.completions'> { const logger = useLogger('v1-completions').useGlobalConfig() const telemetry = createRouteTelemetry({ @@ -182,7 +199,11 @@ export function chatCompletions(deps: V1RouteDeps): GatewayCallback<'chat.comple durationMs, requestId, userId: input.userId, + sessionId: input.sessionId, + roundId: input.roundId, + appSurface: input.appSurface, requestModel, + generationModel: langfuseModel, routeCtxProvider: routeCtx.provider, billing, billingPolicy, @@ -199,7 +220,11 @@ export function chatCompletions(deps: V1RouteDeps): GatewayCallback<'chat.comple durationMs, requestId, userId: input.userId, + sessionId: input.sessionId, + roundId: input.roundId, + appSurface: input.appSurface, requestModel, + generationModel: langfuseModel, routeCtxProvider: routeCtx.provider, billing, billingPolicy, @@ -213,6 +238,33 @@ interface ChatModelAliasPlan { modelIds: string[] } +function captureGeneration(input: GenerationCaptureInput): void { + const generationId = input.roundId ?? input.requestId + const totalTokens = input.usage.promptTokens != null && input.usage.completionTokens != null + ? input.usage.promptTokens + input.usage.completionTokens + : undefined + + input.deps.productEventService.trackGeneration({ + userId: input.userId, + traceId: input.sessionId ?? input.requestId, + generationId, + model: input.generationModel, + provider: input.routeCtxProvider || 'unknown', + providerType: 'official', + usageSource: input.usage.promptTokens != null || input.usage.completionTokens != null + ? 'reported' + : 'unavailable', + inputTokens: input.usage.promptTokens, + outputTokens: input.usage.completionTokens, + totalTokens, + conversationId: input.sessionId, + roundId: generationId, + appSurface: input.appSurface, + latencySeconds: input.durationMs / 1000, + stream: input.stream, + }) +} + async function resolveChatModelAliasPlan(deps: V1RouteDeps, aliasId: string): Promise { const alias = await deps.providerCatalogService.resolveEnabledAlias('llm', aliasId) const primaryRoutes = alias.routes.filter(route => route.pool === 'primary') @@ -302,7 +354,11 @@ function streamChatCompletion(input: { durationMs: number requestId: string userId: string + sessionId?: string + roundId?: string + appSurface: AiGenerationAppSurface requestModel: string + generationModel: string routeCtxProvider: string billing: ChatBilling billingPolicy: ChatBillingPolicy @@ -421,6 +477,20 @@ function streamChatCompletion(input: { }) input.telemetry.recordMetrics({ model: input.requestModel, status: input.response.status, type: 'chat', provider: input.routeCtxProvider, durationMs: input.durationMs, fluxConsumed, ...usage }) + captureGeneration({ + deps: input.deps, + userId: input.userId, + requestId: input.requestId, + sessionId: input.sessionId, + roundId: input.roundId, + appSurface: input.appSurface, + generationModel: input.generationModel, + routeCtxProvider: input.routeCtxProvider, + usage, + durationMs: input.durationMs, + stream: true, + }) + // Debit flux via DB transaction (source of truth) // NOTICE: streaming response is already sent, so we cannot reject on failure. // Log at error level so unpaid usage is visible in monitoring/alerts. @@ -507,7 +577,11 @@ async function completeNonStreamingChat(input: { durationMs: number requestId: string userId: string + sessionId?: string + roundId?: string + appSurface: AiGenerationAppSurface requestModel: string + generationModel: string routeCtxProvider: string billing: ChatBilling billingPolicy: ChatBillingPolicy @@ -556,6 +630,20 @@ async function completeNonStreamingChat(input: { }) input.telemetry.recordMetrics({ model: input.requestModel, status: input.response.status, type: 'chat', provider: input.routeCtxProvider, durationMs: input.durationMs, fluxConsumed, ...usage }) + captureGeneration({ + deps: input.deps, + userId: input.userId, + requestId: input.requestId, + sessionId: input.sessionId, + roundId: input.roundId, + appSurface: input.appSurface, + generationModel: input.generationModel, + routeCtxProvider: input.routeCtxProvider, + usage, + durationMs: input.durationMs, + stream: false, + }) + // Debit flux via DB transaction (source of truth). // The upstream call has already happened (cost incurred), so partial // debit + `fluxUnbilled` is the only sane recovery — same shape as the diff --git a/apps/server/src/routes/openai/v1/route.test.ts b/apps/server/src/routes/openai/v1/route.test.ts index 33b34876a..56111a157 100644 --- a/apps/server/src/routes/openai/v1/route.test.ts +++ b/apps/server/src/routes/openai/v1/route.test.ts @@ -14,6 +14,11 @@ import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' import { createV1Routes } from '.' import { ApiError } from '../../../utils/error' +import { + AIRI_CHAT_APP_SURFACE_HEADER, + AIRI_CHAT_ROUND_ID_HEADER, + AIRI_CHAT_SESSION_ID_HEADER, +} from './analytics' function createMockFluxService(flux = 100): FluxService { return { @@ -143,6 +148,7 @@ function createMockLlmRouter(impl?: Partial): LlmRouterService function createMockProductEventService(): ProductEventService { return { track: vi.fn(async () => undefined), + trackGeneration: vi.fn(async () => undefined), countDistinctUsersByFeature: vi.fn(async () => []), } } @@ -907,7 +913,7 @@ describe('v1CompletionsRoutes', () => { } }) - it('records Langfuse chat generation with the router-resolved upstream model', async () => { + it('records Langfuse and PostHog generations with authoritative usage and correlation', async () => { const llmRouter = createMockLlmRouter({ route: vi.fn(async (_req, ctx) => { if (ctx) { @@ -924,12 +930,27 @@ describe('v1CompletionsRoutes', () => { }) as any, }) const llmTracing = createMockLlmTracing() - const app = createTestApp(createMockFluxService(), createMockConfigKV(), undefined, undefined, undefined, llmRouter, llmTracing) + const productEventService = createMockProductEventService() + const app = createTestApp( + createMockFluxService(), + createMockConfigKV(), + undefined, + undefined, + undefined, + llmRouter, + llmTracing, + productEventService, + ) await app.fetch( new Request('http://localhost/api/v1/openai/chat/completions', { method: 'POST', - headers: { 'Content-Type': 'application/json' }, + headers: { + 'Content-Type': 'application/json', + [AIRI_CHAT_SESSION_ID_HEADER]: 'conversation-1', + [AIRI_CHAT_ROUND_ID_HEADER]: 'round-1', + [AIRI_CHAT_APP_SURFACE_HEADER]: 'electron', + }, body: JSON.stringify({ model: 'chat-auto', messages: [{ role: 'user', content: 'hi' }] }), }), { user: testUser } as any, @@ -942,6 +963,23 @@ describe('v1CompletionsRoutes', () => { userId: 'user-1', }), ) + expect(productEventService.trackGeneration).toHaveBeenCalledWith({ + userId: 'user-1', + traceId: 'conversation-1', + generationId: 'round-1', + model: 'openai/gpt-4o-mini', + provider: 'openrouter', + providerType: 'official', + usageSource: 'reported', + inputTokens: 1, + outputTokens: 2, + totalTokens: 3, + conversationId: 'conversation-1', + roundId: 'round-1', + appSurface: 'electron', + latencySeconds: expect.any(Number), + stream: false, + }) }) it('should not charge flux when upstream returns error', async () => { diff --git a/apps/server/src/routes/stripe/operations/checkout.ts b/apps/server/src/routes/stripe/operations/checkout.ts index dc8e29f94..3c5c3d9e8 100644 --- a/apps/server/src/routes/stripe/operations/checkout.ts +++ b/apps/server/src/routes/stripe/operations/checkout.ts @@ -33,6 +33,11 @@ export interface CheckoutOperationInput { request: Request } +interface PosthogIdentityHeaders { + distinctId?: string + sessionId?: string +} + /** * Creates Stripe checkout sessions for Flux packages. * @@ -75,6 +80,7 @@ export function createCheckoutOperation(deps: CheckoutOperationDeps) { const paymentMethods = await deps.configKV.getOptional('STRIPE_PAYMENT_METHODS') const paymentMethodOptions = await deps.configKV.getOptional('STRIPE_PAYMENT_METHOD_OPTIONS') ?? {} + const posthogIdentity = readPosthogIdentityHeaders(input.request) const sessionParams: CheckoutSessionCreateParams = { line_items: [{ price: stripePriceId, quantity: 1 }], @@ -87,6 +93,8 @@ export function createCheckoutOperation(deps: CheckoutOperationDeps) { metadata: { userId: input.user.id, fluxAmount: String(fluxAmount), + ...(posthogIdentity.distinctId && { posthogDistinctId: posthogIdentity.distinctId }), + ...(posthogIdentity.sessionId && { posthogSessionId: posthogIdentity.sessionId }), }, } @@ -133,9 +141,31 @@ export function createCheckoutOperation(deps: CheckoutOperationDeps) { flux_amount: fluxAmount, amount_total: session.amount_total, currency: session.currency, + ...(posthogIdentity.distinctId && { posthog_distinct_id: posthogIdentity.distinctId }), + ...(posthogIdentity.sessionId && { posthog_session_id: posthogIdentity.sessionId }), }, }) return { url: session.url } } } + +function readPosthogIdentityHeaders(request: Request): PosthogIdentityHeaders { + const distinctId = readStripeMetadataHeader(request, 'x-posthog-distinct-id') + const sessionId = readStripeMetadataHeader(request, 'x-posthog-session-id') + return { + ...(distinctId && { distinctId }), + ...(sessionId && { sessionId }), + } +} + +function readStripeMetadataHeader(request: Request, name: string): string | undefined { + const value = request.headers.get(name)?.trim() + if (!value) + return undefined + + // Stripe metadata values are capped and user-controlled headers can be + // oversized. Truncating keeps the checkout request valid without turning + // analytics identity into a payment blocker. + return value.slice(0, 200) +} diff --git a/apps/server/src/routes/stripe/operations/webhook.ts b/apps/server/src/routes/stripe/operations/webhook.ts index 8f593454b..7a1449b69 100644 --- a/apps/server/src/routes/stripe/operations/webhook.ts +++ b/apps/server/src/routes/stripe/operations/webhook.ts @@ -89,6 +89,8 @@ export function createWebhookOperation(deps: WebhookOperationDeps) { const userId = event.data.object.metadata?.userId if (userId) { const fluxAmount = Number(event.data.object.metadata?.fluxAmount) + const posthogDistinctId = event.data.object.metadata?.posthogDistinctId + const posthogSessionId = event.data.object.metadata?.posthogSessionId void deps.productEventService?.track({ userId, feature: 'billing', @@ -99,6 +101,10 @@ export function createWebhookOperation(deps: WebhookOperationDeps) { amount_total: event.data.object.amount_total, currency: event.data.object.currency, flux_amount: Number.isFinite(fluxAmount) ? fluxAmount : null, + stripe_checkout_session_id: event.data.object.id, + stripe_customer_id: typeof event.data.object.customer === 'string' ? event.data.object.customer : event.data.object.customer?.id ?? null, + ...(posthogDistinctId && { posthog_distinct_id: posthogDistinctId }), + ...(posthogSessionId && { posthog_session_id: posthogSessionId }), }, }) } diff --git a/apps/server/src/routes/stripe/route.test.ts b/apps/server/src/routes/stripe/route.test.ts index 1dbd76ee0..ebba3cb33 100644 --- a/apps/server/src/routes/stripe/route.test.ts +++ b/apps/server/src/routes/stripe/route.test.ts @@ -10,6 +10,7 @@ import { describe, expect, it, vi } from 'vitest' import { createStripeRoutes, formatPrice } from '.' import { ApiError } from '../../utils/error' +import { createCheckoutOperation } from './operations/checkout' import { createWebhookOperation } from './operations/webhook' // --- Mock helpers --- @@ -302,6 +303,77 @@ describe('stripeRoutes', () => { ) expect(res.status).toBe(503) }) + + it('stores browser PostHog identity in Stripe checkout metadata', async () => { + const createSession = vi.fn(async input => ({ + id: 'cs_1', + url: 'https://checkout.stripe.com/cs_1', + customer: null, + mode: 'payment', + status: 'open', + payment_status: 'unpaid', + amount_total: 500, + currency: 'usd', + success_url: 'http://localhost/settings/flux?success=true', + cancel_url: 'http://localhost/settings/flux?canceled=true', + payment_intent: null, + subscription: null, + metadata: input.metadata, + expires_at: null, + })) + const productEventService = { track: vi.fn() } + const operation = createCheckoutOperation({ + stripe: { + checkout: { + sessions: { + create: createSession, + }, + }, + } as any, + priceCatalog: { + findActivePrice: vi.fn(async () => ({ + id: 'price_test_500', + currency: 'usd', + unitAmount: 500, + currencyOptions: {}, + metadata: { fluxAmount: '500' }, + })), + getActivePrices: vi.fn(), + } as any, + stripeService: createMockStripeService(), + configKV: createMockConfigKV({ STRIPE_PAYMENT_METHODS: undefined }), + env: testEnv, + productEventService: productEventService as any, + }) + + await operation({ + user: testUser as any, + body: { stripePriceId: 'price_test_500' }, + request: new Request('http://localhost/api/v1/stripe/checkout', { + headers: { + 'x-posthog-distinct-id': 'anon-browser-1', + 'x-posthog-session-id': 'ph-session-1', + }, + }), + }) + + expect(createSession).toHaveBeenCalledWith(expect.objectContaining({ + metadata: { + userId: 'user-1', + fluxAmount: '500', + posthogDistinctId: 'anon-browser-1', + posthogSessionId: 'ph-session-1', + }, + })) + expect(productEventService.track).toHaveBeenCalledWith(expect.objectContaining({ + userId: 'user-1', + action: 'checkout_started', + metadata: expect.objectContaining({ + posthog_distinct_id: 'anon-browser-1', + posthog_session_id: 'ph-session-1', + }), + })) + }) }) describe('gET /api/v1/stripe/orders', () => { @@ -471,6 +543,75 @@ describe('stripeRoutes', () => { expect(res.status).toBe(503) }) + it('records payment completion with Stripe and PostHog identity from checkout metadata', async () => { + const checkoutEvent = { + id: 'evt_checkout_completed', + type: 'checkout.session.completed', + data: { + object: { + id: 'cs_1', + customer: 'cus_1', + customer_email: 'test@example.com', + mode: 'payment', + status: 'complete', + payment_status: 'paid', + amount_total: 500, + currency: 'usd', + success_url: 'http://localhost/settings/flux?success=true', + cancel_url: 'http://localhost/settings/flux?canceled=true', + payment_intent: 'pi_1', + subscription: null, + metadata: { + userId: 'user-1', + fluxAmount: '500', + posthogDistinctId: 'anon-browser-1', + posthogSessionId: 'ph-session-1', + }, + expires_at: null, + }, + }, + } + const productEventService = { track: vi.fn() } + const billingService = createMockBillingService() + const webhook = createWebhookOperation({ + stripe: { + webhooks: { + constructEvent: vi.fn(() => checkoutEvent), + }, + } as any, + webhookSecret: 'whsec_test', + fluxService: createMockFluxService(), + stripeService: createMockStripeService(), + billingService, + productEventService: productEventService as any, + }) + + await webhook({ signature: 'test_sig', body: '{}' }) + + expect(billingService.creditFluxFromStripeCheckout).toHaveBeenCalledWith(expect.objectContaining({ + stripeEventId: 'evt_checkout_completed', + userId: 'user-1', + stripeSessionId: 'cs_1', + fluxAmount: 500, + })) + expect(productEventService.track).toHaveBeenCalledWith({ + userId: 'user-1', + feature: 'billing', + action: 'payment_completed', + status: 'succeeded', + source: 'stripe.webhook', + metadata: { + amount_total: 500, + currency: 'usd', + flux_amount: 500, + stripe_checkout_session_id: 'cs_1', + stripe_customer_id: 'cus_1', + posthog_distinct_id: 'anon-browser-1', + posthog_session_id: 'ph-session-1', + }, + }) + }) + it('records subscription lifecycle product events from Stripe webhooks', async () => { const subscriptionEvent = { id: 'evt_sub_created', diff --git a/apps/server/src/services/adapters/posthog.ts b/apps/server/src/services/adapters/posthog.ts index 82c9ef4e9..bbbe2d8b9 100644 --- a/apps/server/src/services/adapters/posthog.ts +++ b/apps/server/src/services/adapters/posthog.ts @@ -18,6 +18,12 @@ export interface PosthogCaptureInput { * an interface so tests inject a fake instead of mocking the SDK. */ export interface PosthogSink { + /** + * Queue a high-volume analytics event without waiting for a network + * roundtrip. Use on request hot paths where occasional process-exit loss is + * preferable to user-visible latency. + */ + captureQueued?: (input: PosthogCaptureInput) => void capture: (input: PosthogCaptureInput) => Promise /** Flush and close the underlying client. Call on server shutdown. */ shutdown: () => Promise @@ -26,10 +32,10 @@ export interface PosthogSink { /** * PostHog sink for server-side product events. * - * Uses `captureImmediate` (one HTTP roundtrip per event, no background - * queue) on purpose: the forwarded events are low-frequency business facts - * (signup, payment, subscription lifecycle) fired from webhook/auth-hook - * paths where a queued batch could be lost on process exit. + * Low-frequency conversion facts use `captureImmediate` (one HTTP roundtrip + * per event) because they terminate money/auth funnels. High-frequency AI + * generation facts use `captureQueued`, which is buffered by the SDK and + * flushed on shutdown, so chat completion requests don't wait on PostHog. * * Capture failures are logged and swallowed — analytics forwarding must * never fail the Stripe webhook or auth flow that triggered it. The @@ -39,6 +45,19 @@ export function createPosthogSink(options: { projectKey: string, host: string }) const client = new PostHog(options.projectKey, { host: options.host }) return { + captureQueued(input: PosthogCaptureInput): void { + try { + client.capture({ + distinctId: input.distinctId, + event: input.event, + properties: input.properties, + }) + } + catch (err) { + logger.withError(err).withFields({ event: input.event }).warn('Failed to enqueue product event to PostHog') + } + }, + async capture(input: PosthogCaptureInput): Promise { try { await client.captureImmediate({ diff --git a/apps/server/src/services/domain/product-events.test.ts b/apps/server/src/services/domain/product-events.test.ts index c5109ca81..875697988 100644 --- a/apps/server/src/services/domain/product-events.test.ts +++ b/apps/server/src/services/domain/product-events.test.ts @@ -163,6 +163,7 @@ describe('productEventService', () => { event: 'payment_completed', properties: { app_surface: 'server', + airi_user_id: 'user-1', feature: 'billing', status: 'succeeded', source: 'stripe.webhook', @@ -175,6 +176,7 @@ describe('productEventService', () => { event: 'signup_completed', properties: { app_surface: 'server', + airi_user_id: 'user-2', feature: 'auth', status: 'succeeded', }, @@ -184,6 +186,50 @@ describe('productEventService', () => { expect(rows).toHaveLength(2) }) + it('merges Stripe webhook conversions with the browser PostHog person when a distinct id is present', async () => { + const capture = vi.fn(async () => {}) + const sink = { capture, shutdown: vi.fn(async () => {}) } + const service = createProductEventService(db, null, sink) + + await service.track({ + userId: 'user-1', + feature: 'billing', + action: 'payment_completed', + status: 'succeeded', + source: 'stripe.webhook', + metadata: { + posthog_distinct_id: 'anon-browser-1', + posthog_session_id: 'ph-session-1', + stripe_checkout_session_id: 'cs_1', + }, + }) + + expect(capture).toHaveBeenNthCalledWith(1, { + distinctId: 'user-1', + event: '$identify', + properties: { + $anon_distinct_id: 'anon-browser-1', + $session_id: 'ph-session-1', + airi_user_id: 'user-1', + }, + }) + expect(capture).toHaveBeenNthCalledWith(2, { + distinctId: 'user-1', + event: 'payment_completed', + properties: { + app_surface: 'server', + airi_user_id: 'user-1', + posthog_distinct_id: 'anon-browser-1', + $session_id: 'ph-session-1', + feature: 'billing', + status: 'succeeded', + source: 'stripe.webhook', + posthog_session_id: 'ph-session-1', + stripe_checkout_session_id: 'cs_1', + }, + }) + }) + it('does not forward high-volume per-request actions to PostHog', async () => { const capture = vi.fn(async () => {}) const sink = { capture, shutdown: vi.fn(async () => {}) } @@ -207,6 +253,54 @@ describe('productEventService', () => { expect(rows).toHaveLength(2) }) + it('captures an LLM generation as a PostHog AI fact without storing prompts or responses', async () => { + const capture = vi.fn(async () => {}) + const captureQueued = vi.fn() + const sink = { capture, captureQueued, shutdown: vi.fn(async () => {}) } + const service = createProductEventService(db, null, sink) + + service.trackGeneration({ + userId: 'user-1', + traceId: 'session-1', + generationId: 'round-1', + model: 'openai/gpt-5-mini', + provider: 'openai', + providerType: 'official', + usageSource: 'reported', + inputTokens: 12, + outputTokens: 8, + totalTokens: 20, + conversationId: 'session-1', + roundId: 'round-1', + appSurface: 'server', + }) + + expect(capture).not.toHaveBeenCalled() + expect(captureQueued).toHaveBeenCalledWith({ + distinctId: 'user-1', + event: '$ai_generation', + properties: { + $ai_trace_id: 'session-1', + $ai_session_id: 'session-1', + $ai_span_id: 'round-1', + $ai_model: 'openai/gpt-5-mini', + $ai_provider: 'openai', + $ai_input_tokens: 12, + $ai_output_tokens: 8, + $ai_total_tokens: 20, + $insert_id: 'ai-generation:round-1', + provider_type: 'official', + usage_source: 'reported', + conversation_id: 'session-1', + round_id: 'round-1', + app_surface: 'server', + }, + }) + + const rows = await db.select().from(schema.productEvents) + expect(rows).toHaveLength(0) + }) + it('does not forward to PostHog when the DB write fails', async () => { // ROOT CAUSE: // diff --git a/apps/server/src/services/domain/product-events.ts b/apps/server/src/services/domain/product-events.ts index 888acfee8..b39b40245 100644 --- a/apps/server/src/services/domain/product-events.ts +++ b/apps/server/src/services/domain/product-events.ts @@ -76,6 +76,27 @@ export interface ProductEventAggregateRow { distinctUsers: number } +export type AiGenerationAppSurface = 'server' | 'web' | 'mobile' | 'electron' + +/** Content-free PostHog AI generation fact keyed to the authenticated user. */ +export interface AiGenerationEventInput { + userId: string + traceId: string + generationId: string + model: string + provider: string + providerType: 'official' | 'custom' | 'unknown' + usageSource: 'reported' | 'estimated' | 'unavailable' + inputTokens?: number + outputTokens?: number + totalTokens?: number + conversationId?: string + roundId?: string + appSurface: AiGenerationAppSurface + latencySeconds?: number + stream?: boolean +} + /** * Server-side actions worth a PostHog copy, mapped to the event name the * client-side funnels expect. Only business facts that terminate or anchor @@ -115,6 +136,11 @@ function metricLabels(input: ProductEventInput): Record { return attrs } +function stringMetadata(input: ProductEventInput, key: string): string | undefined { + const value = input.metadata?.[key] + return typeof value === 'string' && value.length > 0 ? value : undefined +} + /** * Creates AIRI's first-party product analytics event writer. * @@ -134,6 +160,42 @@ function metricLabels(input: ProductEventInput): Record { */ export function createProductEventService(db: Database, metrics?: ProductMetrics | null, posthog?: PosthogSink | null) { return { + trackGeneration(input: AiGenerationEventInput): void { + if (!posthog) + return + + const event = { + distinctId: input.userId, + event: '$ai_generation', + properties: { + $ai_trace_id: input.traceId, + ...(input.conversationId && { $ai_session_id: input.conversationId }), + $ai_span_id: input.generationId, + $ai_model: input.model, + $ai_provider: input.provider, + ...(input.inputTokens != null && { $ai_input_tokens: input.inputTokens }), + ...(input.outputTokens != null && { $ai_output_tokens: input.outputTokens }), + ...(input.totalTokens != null && { $ai_total_tokens: input.totalTokens }), + ...(input.latencySeconds != null && { $ai_latency: input.latencySeconds }), + ...(input.stream != null && { $ai_stream: input.stream }), + $insert_id: `ai-generation:${input.generationId}`, + provider_type: input.providerType, + usage_source: input.usageSource, + ...(input.conversationId && { conversation_id: input.conversationId }), + ...(input.roundId && { round_id: input.roundId }), + app_surface: input.appSurface, + }, + } + + if (posthog.captureQueued) { + posthog.captureQueued(event) + return + } + + void posthog.capture(event) + .catch(err => logger.withError(err).withFields({ generationId: input.generationId }).warn('Failed to capture PostHog AI generation')) + }, + async track(input: ProductEventInput): Promise { // Postgres is the fact of record, so forwarding is gated both ways: // the DB write comes first (a PostHog outage can't lose the row) and @@ -171,11 +233,28 @@ export function createProductEventService(db: Database, metrics?: ProductMetrics const forwardedEvent = POSTHOG_FORWARDED_ACTIONS[input.action] if (persisted && posthog && forwardedEvent) { try { + const posthogDistinctId = stringMetadata(input, 'posthog_distinct_id') + const posthogSessionId = stringMetadata(input, 'posthog_session_id') + if (posthogDistinctId && posthogDistinctId !== input.userId) { + await posthog.capture({ + distinctId: input.userId, + event: '$identify', + properties: { + $anon_distinct_id: posthogDistinctId, + airi_user_id: input.userId, + ...(posthogSessionId && { $session_id: posthogSessionId }), + }, + }) + } + await posthog.capture({ distinctId: input.userId, event: forwardedEvent, properties: { app_surface: 'server', + airi_user_id: input.userId, + ...(posthogDistinctId && { posthog_distinct_id: posthogDistinctId }), + ...(posthogSessionId && { $session_id: posthogSessionId }), feature: input.feature, status: input.status, ...(input.source && { source: input.source }), diff --git a/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts b/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts index 20b880e6c..30f483c47 100644 --- a/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts +++ b/packages/core-agent/src/runtime/chat-orchestrator-runtime.test.ts @@ -2,7 +2,7 @@ import type { ChatProvider } from '@xsai-ext/providers/utils' import type { Message } from '@xsai/shared-chat' import type { ChatHistoryItem, ContextMessage, StreamingAssistantMessage } from '../types/chat' -import type { StreamEvent } from '../types/llm' +import type { StreamEvent, StreamOptions } from '../types/llm' import { ContextUpdateStrategy } from '@proj-airi/server-shared/types' import { describe, expect, it, vi } from 'vitest' @@ -42,12 +42,11 @@ function createHarness() { llmRequestStarted: [] as unknown[], llmFirstToken: [] as unknown[], assistantResponseRendered: [] as unknown[], + llmGeneration: [] as unknown[], messageRound: [] as unknown[], messageRoundFailed: [] as unknown[], } - const stream = vi.fn(async (_model: string, _chatProvider: ChatProvider, _messages: Message[], options?: { - onStreamEvent?: (event: StreamEvent) => Promise | void - }) => { + const stream = vi.fn(async (_model: string, _chatProvider: ChatProvider, _messages: Message[], options?: StreamOptions) => { await options?.onStreamEvent?.({ type: 'text-delta', text: 'assistant reply' }) await options?.onStreamEvent?.({ type: 'finish', finishReason: 'stop' }) }) @@ -100,6 +99,7 @@ function createHarness() { onLlmRequestStarted: event => telemetry.llmRequestStarted.push(event), onLlmFirstToken: event => telemetry.llmFirstToken.push(event), onAssistantResponseRendered: event => telemetry.assistantResponseRendered.push(event), + onLlmGeneration: event => telemetry.llmGeneration.push(event), onMessageRound: event => telemetry.messageRound.push(event), onMessageRoundFailed: event => telemetry.messageRoundFailed.push(event), }) @@ -297,6 +297,16 @@ describe('createChatOrchestratorRuntime', () => { it('emits telemetry milestones for a successful voice-backed message round', async () => { const harness = createHarness() harness.monotonicNow.set([100, 150, 250, 400, 460]) + harness.stream.mockImplementationOnce(async (_model, _chatProvider, _messages, options) => { + await options?.onStreamEvent?.({ type: 'text-delta', text: 'assistant reply' }) + await options?.onStreamEvent?.({ type: 'finish', finishReason: 'stop' }) + await options?.onUsage?.({ + inputTokens: 12, + outputTokens: 8, + totalTokens: 20, + source: 'reported', + }) + }) await harness.runtime.ingest('hello from voice', { model: 'gpt-test', @@ -338,13 +348,28 @@ describe('createChatOrchestratorRuntime', () => { latencyMs: 250, turnIndex: 1, }]) + expect(harness.telemetry.llmGeneration).toEqual([{ + conversationId: 'session-1', + roundId: 'user-id', + model: 'gpt-test', + provider: 'mock-provider', + inputTokens: 12, + outputTokens: 8, + totalTokens: 20, + usageSource: 'reported', + turnIndex: 1, + }]) expect(harness.telemetry.messageRound).toEqual([{ conversationId: 'session-1', roundId: 'user-id', durationMs: 360, hasVoice: true, + inputTokens: 12, model: 'gpt-test', + outputTokens: 8, + totalTokens: 20, turnIndex: 1, + usageSource: 'reported', }]) expect(harness.telemetry.chatActivationStarted).toEqual([{ conversationId: 'session-1', diff --git a/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts b/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts index 5ed767960..55ae4f3dd 100644 --- a/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts +++ b/packages/core-agent/src/runtime/chat-orchestrator-runtime.ts @@ -4,7 +4,7 @@ import type { CommonContentPart, Message, ToolMessage } from '@xsai/shared-chat' import type { AgentContextPort } from '../contracts/context-port' import type { AgentForegroundStreamPort } from '../contracts/stream-port' import type { ChatAssistantMessage, ChatHistoryItem, ChatSlices, ChatStreamEventContext, ContextMessage, StreamingAssistantMessage } from '../types/chat' -import type { StreamEvent, StreamOptions } from '../types/llm' +import type { LlmUsage, StreamEvent, StreamOptions } from '../types/llm' import { createQueue } from '@proj-airi/stream-kit' @@ -242,11 +242,24 @@ export interface ChatOrchestratorRuntimeDeps { model: string latencyMs: number }) => void + /** Called once per completed provider generation with content-free usage metadata. */ + onLlmGeneration?: (event: ChatRoundCorrelation & { + model: string + provider: string + inputTokens?: number + outputTokens?: number + totalTokens?: number + usageSource: LlmUsage['source'] + }) => void /** Called after one user-to-assistant message round completes successfully. */ onMessageRound?: (event: ChatRoundCorrelation & { durationMs: number hasVoice: boolean model: string + inputTokens?: number + outputTokens?: number + totalTokens?: number + usageSource: LlmUsage['source'] }) => void /** Called whenever a user-to-assistant round fails before completion. */ onMessageRoundFailed?: (event: ChatRoundCorrelation & { @@ -679,6 +692,7 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) const llmRequestStartedAt = monotonicNow() let llmFirstTokenEmitted = false + let generationUsage: LlmUsage = { source: 'unavailable' } deps.onLlmRequestStarted?.({ ...correlation, model: options.model, @@ -688,9 +702,25 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) await deps.llm.stream(options.model, options.chatProvider, newMessages as Message[], { headers, + requestCorrelation: { + conversationId: correlation.conversationId, + roundId: correlation.roundId, + }, tools: options.tools, waitForTools: true, captureToolErrors: true, + onUsage: (usage) => { + generationUsage = usage + deps.onLlmGeneration?.({ + ...correlation, + model: options.model, + provider: activeProvider, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + totalTokens: usage.totalTokens, + usageSource: usage.source, + }) + }, onStreamEvent: async (event: StreamEvent) => { switch (event.type) { case 'tool-call': @@ -794,6 +824,10 @@ export function createChatOrchestratorRuntime(deps: ChatOrchestratorRuntimeDeps) durationMs, hasVoice: !!options.input, model: options.model, + inputTokens: generationUsage.inputTokens, + outputTokens: generationUsage.outputTokens, + totalTokens: generationUsage.totalTokens, + usageSource: generationUsage.source, }) if (isActivationAttempt) { deps.onChatActivationSucceeded?.({ diff --git a/packages/core-agent/src/runtime/llm-service.test.ts b/packages/core-agent/src/runtime/llm-service.test.ts index 8fc13b44b..78f678441 100644 --- a/packages/core-agent/src/runtime/llm-service.test.ts +++ b/packages/core-agent/src/runtime/llm-service.test.ts @@ -27,16 +27,119 @@ const provider = { }), } as unknown as ChatProvider -function createMockStreamResult(steps: Promise = Promise.resolve([])) { +function createMockStreamResult( + steps: Promise = Promise.resolve([]), + totalUsage: Promise<{ prompt_tokens: number, completion_tokens: number, total_tokens: number } | undefined> = Promise.resolve(undefined), +) { return { steps, messages: Promise.resolve([]), usage: Promise.resolve(undefined), - totalUsage: Promise.resolve(undefined), + totalUsage, } } describe('streamFrom tool error capture', () => { + it('requests final streaming usage and emits the reported token totals once', async () => { + const onUsage = vi.fn() + streamTextMock.mockReturnValueOnce(createMockStreamResult( + Promise.resolve([]), + Promise.resolve({ prompt_tokens: 12, completion_tokens: 8, total_tokens: 20 }), + )) + + await streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + options: { onUsage }, + }) + + expect(streamTextMock).toHaveBeenCalledWith(expect.objectContaining({ + streamOptions: { includeUsage: true }, + })) + expect(onUsage).toHaveBeenCalledTimes(1) + expect(onUsage).toHaveBeenCalledWith({ + inputTokens: 12, + outputTokens: 8, + totalTokens: 20, + source: 'reported', + }) + }) + + it('marks usage unavailable when the provider omits the final usage chunk', async () => { + const onUsage = vi.fn() + streamTextMock.mockReturnValueOnce(createMockStreamResult()) + + await streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + options: { onUsage }, + }) + + expect(onUsage).toHaveBeenCalledWith({ source: 'unavailable' }) + }) + + it('marks usage unavailable when the final usage object has no token fields', async () => { + const onUsage = vi.fn() + streamTextMock.mockReturnValueOnce(createMockStreamResult( + Promise.resolve([]), + Promise.resolve({} as { prompt_tokens: number, completion_tokens: number, total_tokens: number }), + )) + + await streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + options: { onUsage }, + }) + + expect(onUsage).toHaveBeenCalledWith({ source: 'unavailable' }) + }) + + it('consumes totalUsage rejection when the stream fails before usage can be awaited', async () => { + const streamError = new Error('provider stream failed') + const totalUsageError = new Error('provider usage failed') + const unhandledRejections: unknown[] = [] + const onUnhandledRejection = (reason: unknown) => { + unhandledRejections.push(reason) + } + process.on('unhandledRejection', onUnhandledRejection) + + streamTextMock.mockReturnValueOnce(createMockStreamResult( + Promise.reject(streamError), + Promise.reject(totalUsageError), + )) + + try { + await expect(streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + })).rejects.toThrow('provider stream failed') + await new Promise(resolve => setImmediate(resolve)) + expect(unhandledRejections).toEqual([]) + } + finally { + process.off('unhandledRejection', onUnhandledRejection) + } + }) + + it('does not fail a completed generation when the usage observer throws', async () => { + streamTextMock.mockReturnValueOnce(createMockStreamResult()) + + await expect(streamFrom({ + model: 'model-a', + chatProvider: provider, + messages: [{ role: 'user', content: 'hello' }] as Message[], + options: { + onUsage: () => { + throw new Error('analytics unavailable') + }, + }, + })).resolves.toBeUndefined() + }) + /** * @example * await streamFrom({ model, chatProvider, messages, options: { captureToolErrors: true } }) diff --git a/packages/core-agent/src/runtime/llm-service.ts b/packages/core-agent/src/runtime/llm-service.ts index 23d47e1ed..84c5fe86c 100644 --- a/packages/core-agent/src/runtime/llm-service.ts +++ b/packages/core-agent/src/runtime/llm-service.ts @@ -1,5 +1,5 @@ import type { ChatProvider } from '@xsai-ext/providers/utils' -import type { Message, Tool } from '@xsai/shared-chat' +import type { Message, Tool, Usage } from '@xsai/shared-chat' import type { StreamFromOptions, StreamOptions } from '../types/llm' @@ -113,6 +113,19 @@ function createCapturedToolErrorResult(toolName: string, error: unknown): string return `Tool call error for "${toolName}": ${errorMessageFromValue(error)}` } +function normalizeUsage(usage: Usage | undefined) { + if (!usage || (usage.prompt_tokens == null && usage.completion_tokens == null && usage.total_tokens == null)) { + return { source: 'unavailable' as const } + } + + return { + inputTokens: usage.prompt_tokens, + outputTokens: usage.completion_tokens, + totalTokens: usage.total_tokens, + source: 'reported' as const, + } +} + function withCapturedToolErrors( tools: Tool[], capturedToolErrorByCallId: Map, @@ -187,6 +200,7 @@ export async function streamFrom({ return new Promise((resolve, reject) => { let settled = false + let stepsSettled = false const resolveOnce = () => { if (settled) return @@ -194,7 +208,7 @@ export async function streamFrom({ resolve() } const rejectOnce = (error: unknown) => { - if (settled) + if (settled || stepsSettled) return settled = true reject(error) @@ -204,13 +218,7 @@ export async function streamFrom({ try { const streamEvent = resolveCapturedToolErrorEvent(event, capturedToolErrorByCallId) await options?.onStreamEvent?.(streamEvent as any) - if (event && (event as any).type === 'finish') { - const finishReason = (event as any).finishReason - const waitingForToolRound = finishReason === 'tool_calls' || finishReason === 'tool-calls' - if (!waitingForToolRound || !options?.waitForTools) - resolveOnce() - } - else if (event && (event as any).type === 'error') { + if (event && (event as any).type === 'error') { rejectOnce((event as any).error ?? new Error('Stream error')) } } @@ -225,6 +233,7 @@ export async function streamFrom({ abortSignal: options?.abortSignal, messages: sanitized, headers: options?.headers, + streamOptions: { includeUsage: true }, stopWhen: stepCountAtLeast(10), // NOTICE: // Do not pass xsAI's `captureToolErrors` option here. In the installed @@ -249,12 +258,42 @@ export async function streamFrom({ // from starting. // Keep `steps.then(resolveOnce)` so evaluation runners observe the real end // of the stream lifecycle instead of an intermediate tool boundary. - void streamResult.steps.then(resolveOnce).catch((error) => { + void streamResult.steps.then(async () => { + // Ignore any late provider error event emitted after xsAI has already + // resolved the authoritative full-step lifecycle. + stepsSettled = true + let usage: Usage | undefined + try { + usage = await streamResult.totalUsage + } + catch (error) { + console.error('Stream totalUsage error:', error) + } + try { + await options?.onUsage?.(normalizeUsage(usage)) + } + catch (error) { + // Usage observers are telemetry-only and must not turn a completed + // provider response into a failed user message. + console.error('Stream usage callback error:', error) + } + resolveOnce() + }).catch((error) => { + // A failure after `steps` resolved belongs to optional usage + // observation and cannot invalidate the completed response. + if (stepsSettled) { + console.error('Stream usage observation error:', error) + resolveOnce() + return + } rejectOnce(error) console.error('Stream steps error:', error) }) void streamResult.messages.catch(error => console.error('Stream messages error:', error)) void streamResult.usage.catch(error => console.error('Stream usage error:', error)) + // `steps` and `totalUsage` reject independently when xsAI fails a + // stream. The success path awaits `totalUsage`, but if `steps` rejects + // first that await never runs, so keep this unconditional rejection sink. void streamResult.totalUsage.catch(error => console.error('Stream totalUsage error:', error)) } catch (error) { diff --git a/packages/core-agent/src/types/llm.ts b/packages/core-agent/src/types/llm.ts index cec744948..2f93f8432 100644 --- a/packages/core-agent/src/types/llm.ts +++ b/packages/core-agent/src/types/llm.ts @@ -1,6 +1,17 @@ import type { ChatProvider } from '@xsai-ext/providers/utils' import type { CommonContentPart, CompletionToolCall, CompletionToolResult, Message, Tool } from '@xsai/shared-chat' +/** Describes whether generation usage came from the provider or a local fallback. */ +export type LlmUsageSource = 'reported' | 'estimated' | 'unavailable' + +/** Provider-safe token usage emitted after one complete streamed generation. */ +export interface LlmUsage { + inputTokens?: number + outputTokens?: number + totalTokens?: number + source: LlmUsageSource +} + export type StreamEvent = | { type: 'text-delta', text: string } | { type: 'reasoning-delta', text: string } @@ -14,6 +25,13 @@ export interface StreamOptions { abortSignal?: AbortSignal headers?: Record onStreamEvent?: (event: StreamEvent) => void | Promise + /** Called once after the full stream, including tool rounds, has settled. */ + onUsage?: (usage: LlmUsage) => void | Promise + /** Internal correlation kept out of the provider request body. */ + requestCorrelation?: { + conversationId: string + roundId: string + } toolsCompatibility?: Map supportsTools?: boolean waitForTools?: boolean diff --git a/packages/stage-ui/src/composables/use-analytics.test.ts b/packages/stage-ui/src/composables/use-analytics.test.ts index 3bf53234b..282df1c48 100644 --- a/packages/stage-ui/src/composables/use-analytics.test.ts +++ b/packages/stage-ui/src/composables/use-analytics.test.ts @@ -83,6 +83,65 @@ describe('useAnalytics conversation product events', () => { }) }) + it('captures custom-provider token usage without prompt or response content', () => { + const analytics = useAnalytics() + + analytics.trackAiGeneration({ + conversation_id: 'session-1', + round_id: 'round-1', + provider_type: 'custom', + provider_id: 'openai-compatible', + model_id: 'custom-model', + usage_source: 'reported', + input_tokens: 12, + output_tokens: 8, + }) + + expect(analyticsMocks.posthogCaptureMock).toHaveBeenCalledWith('$ai_generation', { + $ai_trace_id: 'session-1', + $ai_session_id: 'session-1', + $ai_span_id: 'round-1', + $ai_model: 'custom-model', + $ai_provider: 'openai-compatible', + $ai_input_tokens: 12, + $ai_output_tokens: 8, + $ai_total_tokens: 20, + $insert_id: 'ai-generation:round-1', + app_surface: 'web', + conversation_id: 'session-1', + round_id: 'round-1', + provider_type: 'custom', + usage_source: 'reported', + }) + }) + + it('records unavailable custom-provider usage without inventing token or cost fields', () => { + const analytics = useAnalytics() + + analytics.trackAiGeneration({ + conversation_id: 'session-1', + round_id: 'round-2', + provider_type: 'custom', + provider_id: 'ollama', + model_id: 'local-model', + usage_source: 'unavailable', + }) + + expect(analyticsMocks.posthogCaptureMock).toHaveBeenCalledWith('$ai_generation', { + $ai_trace_id: 'session-1', + $ai_session_id: 'session-1', + $ai_span_id: 'round-2', + $ai_model: 'local-model', + $ai_provider: 'ollama', + $insert_id: 'ai-generation:round-2', + app_surface: 'web', + conversation_id: 'session-1', + round_id: 'round-2', + provider_type: 'custom', + usage_source: 'unavailable', + }) + }) + it('infers the mobile surface for capacitor conversation actions', () => { analyticsMocks.isStageCapacitorMock.mockReturnValue(true) const analytics = useAnalytics() diff --git a/packages/stage-ui/src/composables/use-analytics.ts b/packages/stage-ui/src/composables/use-analytics.ts index 9f212d8fd..4f4a78b5c 100644 --- a/packages/stage-ui/src/composables/use-analytics.ts +++ b/packages/stage-ui/src/composables/use-analytics.ts @@ -36,6 +36,7 @@ export type FeedbackDescriptionLengthBucket = 'empty' | 'short' | 'medium' | 'lo export type ProductAnalyticsEntry = 'app_start' | 'onboarding' | 'settings' | 'chat' | 'pricing' | 'quota_banner' | 'unknown' export type MessageInputMode = 'text' | 'voice' export type ConversationEventSource = 'new_session' | 'fork' | 'history' | 'share_button' | 'unknown' +export type AiUsageSource = 'reported' | 'estimated' | 'unavailable' /** * Full stage vocabulary of the cross-surface `oauth_callback_failed` event. @@ -108,7 +109,7 @@ interface ConversationBaseProperties { model: string } -function getConversationAnalyticsSurface(): ConversationAnalyticsSurface { +export function getConversationAnalyticsSurface(): ConversationAnalyticsSurface { if (isStageTamagotchi()) return 'electron' @@ -397,8 +398,54 @@ export function useAnalytics() { posthog.capture('assistant_response_rendered', properties) } + /** Cost-fact event for one custom-provider generation; content is intentionally excluded. */ + function trackAiGeneration(properties: { + conversation_id: string + round_id: string + provider_type: ProviderMode + provider_id: string + model_id: string + usage_source: AiUsageSource + input_tokens?: number + output_tokens?: number + total_tokens?: number + }) { + if (!canCapture()) + return + + const totalTokens = properties.total_tokens + ?? (properties.input_tokens != null && properties.output_tokens != null + ? properties.input_tokens + properties.output_tokens + : undefined) + + posthog.capture('$ai_generation', { + $ai_trace_id: properties.conversation_id, + $ai_session_id: properties.conversation_id, + $ai_span_id: properties.round_id, + $ai_model: properties.model_id, + $ai_provider: properties.provider_id, + ...(properties.input_tokens != null && { $ai_input_tokens: properties.input_tokens }), + ...(properties.output_tokens != null && { $ai_output_tokens: properties.output_tokens }), + ...(totalTokens != null && { $ai_total_tokens: totalTokens }), + $insert_id: `ai-generation:${properties.round_id}`, + app_surface: getConversationAnalyticsSurface(), + conversation_id: properties.conversation_id, + round_id: properties.round_id, + provider_type: properties.provider_type, + usage_source: properties.usage_source, + }) + } + /** Closing event for one full message round (user send → assistant render). */ - function trackMessageRound(properties: ChatRoundCorrelationProperties & { duration_ms: number, has_voice: boolean, model: string }) { + function trackMessageRound(properties: ChatRoundCorrelationProperties & { + duration_ms: number + has_voice: boolean + model: string + input_tokens?: number + output_tokens?: number + total_tokens?: number + usage_source?: AiUsageSource + }) { if (!canCapture()) return posthog.capture('message_round', properties) @@ -1234,6 +1281,7 @@ export function useAnalytics() { trackLlmRequestStarted, trackLlmFirstToken, trackAssistantResponseRendered, + trackAiGeneration, trackMessageRound, trackMessageRoundFailed, trackMessageSent, diff --git a/packages/stage-ui/src/libs/analytics-headers.ts b/packages/stage-ui/src/libs/analytics-headers.ts new file mode 100644 index 000000000..3a47b4b0a --- /dev/null +++ b/packages/stage-ui/src/libs/analytics-headers.ts @@ -0,0 +1,3 @@ +export const AIRI_CHAT_SESSION_ID_HEADER = 'x-airi-session-id' +export const AIRI_CHAT_ROUND_ID_HEADER = 'x-airi-round-id' +export const AIRI_CHAT_APP_SURFACE_HEADER = 'x-airi-app-surface' diff --git a/packages/stage-ui/src/libs/auth-fetch.test.ts b/packages/stage-ui/src/libs/auth-fetch.test.ts new file mode 100644 index 000000000..0c70675cf --- /dev/null +++ b/packages/stage-ui/src/libs/auth-fetch.test.ts @@ -0,0 +1,75 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { authedFetch } from './auth-fetch' + +const authMocks = vi.hoisted(() => ({ + getAuthToken: vi.fn(() => 'access-token'), +})) + +const posthogMocks = vi.hoisted(() => ({ + getPosthogIdentitySnapshot: vi.fn<() => { distinctId: string, sessionId: string } | null>(() => ({ + distinctId: 'distinct-1', + sessionId: 'session-1', + })), +})) + +vi.mock('./auth', () => ({ + getAuthToken: authMocks.getAuthToken, +})) + +vi.mock('../stores/analytics/posthog', () => ({ + getPosthogIdentitySnapshot: posthogMocks.getPosthogIdentitySnapshot, +})) + +describe('authedFetch', () => { + beforeEach(() => { + vi.restoreAllMocks() + authMocks.getAuthToken.mockReturnValue('access-token') + posthogMocks.getPosthogIdentitySnapshot.mockReturnValue({ + distinctId: 'distinct-1', + sessionId: 'session-1', + }) + }) + + it('sends PostHog identity headers with authenticated API requests', async () => { + const fetchMock = vi.fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>(async () => new Response('{}', { status: 200 })) + vi.stubGlobal('fetch', fetchMock) + + await authedFetch('https://api.airi.build/api/v1/stripe/checkout', { + method: 'POST', + }) + + const headers = fetchMock.mock.calls[0]?.[1]?.headers + expect(headers).toBeInstanceOf(Headers) + expect((headers as Headers).get('Authorization')).toBe('Bearer access-token') + expect((headers as Headers).get('x-posthog-distinct-id')).toBe('distinct-1') + expect((headers as Headers).get('x-posthog-session-id')).toBe('session-1') + }) + + it('omits PostHog identity headers when analytics has no active identity', async () => { + posthogMocks.getPosthogIdentitySnapshot.mockReturnValue(null) + const fetchMock = vi.fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>(async () => new Response('{}', { status: 200 })) + vi.stubGlobal('fetch', fetchMock) + + await authedFetch('https://api.example.test/api/v1/stripe/checkout') + + const headers = fetchMock.mock.calls[0]?.[1]?.headers + expect(headers).toBeInstanceOf(Headers) + expect((headers as Headers).get('Authorization')).toBe('Bearer access-token') + expect((headers as Headers).get('x-posthog-distinct-id')).toBeNull() + expect((headers as Headers).get('x-posthog-session-id')).toBeNull() + }) + + it('does not send PostHog identity headers to non-server origins', async () => { + const fetchMock = vi.fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>(async () => new Response('{}', { status: 200 })) + vi.stubGlobal('fetch', fetchMock) + + await authedFetch('https://third-party.example.test/resource') + + const headers = fetchMock.mock.calls[0]?.[1]?.headers + expect(headers).toBeInstanceOf(Headers) + expect((headers as Headers).get('Authorization')).toBe('Bearer access-token') + expect((headers as Headers).get('x-posthog-distinct-id')).toBeNull() + expect((headers as Headers).get('x-posthog-session-id')).toBeNull() + }) +}) diff --git a/packages/stage-ui/src/libs/auth-fetch.ts b/packages/stage-ui/src/libs/auth-fetch.ts index 5c5c0c7b4..6fc1c14a7 100644 --- a/packages/stage-ui/src/libs/auth-fetch.ts +++ b/packages/stage-ui/src/libs/auth-fetch.ts @@ -1,5 +1,7 @@ +import { getPosthogIdentitySnapshot } from '../stores/analytics/posthog' import { useAuthStore } from '../stores/auth' import { getAuthToken } from './auth' +import { SERVER_URL } from './server' /** * Fetch wrapper that transparently refreshes the OIDC access token on 401 @@ -25,6 +27,12 @@ export async function authedFetch( const headers = new Headers(init?.headers) if (token) headers.set('Authorization', `Bearer ${token}`) + const posthogIdentity = shouldAttachPosthogIdentity(input) ? getPosthogIdentitySnapshot() : null + if (posthogIdentity) { + headers.set('x-posthog-distinct-id', posthogIdentity.distinctId) + if (posthogIdentity.sessionId) + headers.set('x-posthog-session-id', posthogIdentity.sessionId) + } return fetch(input, { ...init, headers, credentials: 'omit' }) } @@ -52,6 +60,14 @@ export async function authedFetch( return retried } +function shouldAttachPosthogIdentity(input: RequestInfo | URL): boolean { + const url = typeof input === 'string' + ? input + : input instanceof URL ? input.toString() : input.url + + return new URL(url, SERVER_URL).origin === new URL(SERVER_URL).origin +} + function promptReLogin(authStore: ReturnType): void { authStore.clearAllAuthState() authStore.needsLogin = true diff --git a/packages/stage-ui/src/libs/providers/providers/official/shared.ts b/packages/stage-ui/src/libs/providers/providers/official/shared.ts index 2b8e7eb47..01379f805 100644 --- a/packages/stage-ui/src/libs/providers/providers/official/shared.ts +++ b/packages/stage-ui/src/libs/providers/providers/official/shared.ts @@ -1,6 +1,7 @@ import { createOpenAI } from '@xsai-ext/providers/create' import { getActivePinia } from 'pinia' +import { AIRI_CHAT_SESSION_ID_HEADER } from '../../../../libs/analytics-headers' import { getAuthToken } from '../../../../libs/auth' import { SERVER_URL } from '../../../../libs/server' import { useChatSessionStore } from '../../../../stores/chat/session-store' @@ -16,7 +17,7 @@ export function withCredentials() { } const chatSession = getActivePinia() ? useChatSessionStore() : null if (chatSession?.activeSessionId) - headers.set('x-airi-session-id', chatSession.activeSessionId) + headers.set(AIRI_CHAT_SESSION_ID_HEADER, chatSession.activeSessionId) const requestInit = { ...init, diff --git a/packages/stage-ui/src/stores/analytics/posthog.test.ts b/packages/stage-ui/src/stores/analytics/posthog.test.ts index 57bf31533..b58c467e3 100644 --- a/packages/stage-ui/src/stores/analytics/posthog.test.ts +++ b/packages/stage-ui/src/stores/analytics/posthog.test.ts @@ -1,8 +1,11 @@ import { describe, expect, it, vi } from 'vitest' -import { ensurePosthogInitialized } from './posthog' +import { ensurePosthogInitialized, getPosthogIdentitySnapshot } from './posthog' const posthogMocks = vi.hoisted(() => ({ + get_distinct_id: vi.fn(() => 'distinct-1'), + get_session_id: vi.fn(() => 'session-1'), + has_opted_out_capturing: vi.fn(() => false), init: vi.fn(), register: vi.fn(), })) @@ -32,4 +35,13 @@ describe('stage PostHog initialization', () => { expect(ensurePosthogInitialized(true)).toBe(true) expect(posthogMocks.register).toHaveBeenCalledWith({ app_surface: 'web' }) }) + + it('exposes the current PostHog identity for server-side conversion linking', () => { + expect(ensurePosthogInitialized(true)).toBe(true) + + expect(getPosthogIdentitySnapshot()).toEqual({ + distinctId: 'distinct-1', + sessionId: 'session-1', + }) + }) }) diff --git a/packages/stage-ui/src/stores/analytics/posthog.ts b/packages/stage-ui/src/stores/analytics/posthog.ts index f9eff92f3..4f1186e5a 100644 --- a/packages/stage-ui/src/stores/analytics/posthog.ts +++ b/packages/stage-ui/src/stores/analytics/posthog.ts @@ -12,6 +12,13 @@ import { let posthogInitialized = false +export interface PosthogIdentitySnapshot { + /** Current PostHog distinct id for the browser/device/user person. */ + distinctId: string + /** Current PostHog session id, when the SDK has established one. */ + sessionId?: string +} + // All AIRI surfaces (web, desktop, mobile) capture into a single PostHog // project. The platform is carried on every event via the `app_surface` super // property (registered at init), so cross-platform funnels live in one @@ -112,6 +119,25 @@ export function resetPosthog(): void { posthog.reset() } +/** + * Returns the current PostHog identity that server-side conversion events can + * use to merge Stripe webhook facts back into the same browser funnel. + */ +export function getPosthogIdentitySnapshot(): PosthogIdentitySnapshot | null { + if (!posthogInitialized || posthog.has_opted_out_capturing()) + return null + + const distinctId = posthog.get_distinct_id() + if (!distinctId) + return null + + const sessionId = posthog.get_session_id() + return { + distinctId, + ...(sessionId && { sessionId }), + } +} + interface PosthogCaptureOptions { send_instantly?: boolean transport?: 'XHR' | 'fetch' | 'sendBeacon' diff --git a/packages/stage-ui/src/stores/chat.contract.test.ts b/packages/stage-ui/src/stores/chat.contract.test.ts index 4004825f5..e10de4a43 100644 --- a/packages/stage-ui/src/stores/chat.contract.test.ts +++ b/packages/stage-ui/src/stores/chat.contract.test.ts @@ -6,6 +6,11 @@ import { createPinia, setActivePinia } from 'pinia' import { beforeEach, describe, expect, it, vi } from 'vitest' import { nextTick, ref } from 'vue' +import { + AIRI_CHAT_APP_SURFACE_HEADER, + AIRI_CHAT_ROUND_ID_HEADER, + AIRI_CHAT_SESSION_ID_HEADER, +} from '../libs/analytics-headers' import { useChatOrchestratorStore } from './chat' vi.hoisted(() => { @@ -40,6 +45,7 @@ const ioTracerMocks = vi.hoisted(() => { const llmStreamMock = vi.fn() const trackFirstMessageMock = vi.fn() const chatAnalyticsMocks = vi.hoisted(() => ({ + trackAiGeneration: vi.fn(), trackAssistantResponseRendered: vi.fn(), trackChatActivationFailed: vi.fn(), trackChatActivationStarted: vi.fn(), @@ -82,6 +88,7 @@ vi.mock('pinia', async () => { }) vi.mock('../composables', () => ({ + getConversationAnalyticsSurface: () => 'web', useAnalytics: () => ({ trackFirstMessage: trackFirstMessageMock, trackChatFailed: redundantChatAnalyticsMocks.trackChatFailed, @@ -90,6 +97,7 @@ vi.mock('../composables', () => ({ trackMessageSent: chatAnalyticsMocks.trackMessageSent, trackLlmRequestStarted: chatAnalyticsMocks.trackLlmRequestStarted, trackLlmFirstToken: chatAnalyticsMocks.trackLlmFirstToken, + trackAiGeneration: chatAnalyticsMocks.trackAiGeneration, trackAssistantResponseRendered: chatAnalyticsMocks.trackAssistantResponseRendered, trackAssistantResponseCompleted: redundantChatAnalyticsMocks.trackAssistantResponseCompleted, trackMessageRound: chatAnalyticsMocks.trackMessageRound, @@ -248,6 +256,51 @@ describe('chat orchestrator contract', () => { expect(chatAnalyticsMocks.trackChatActivationSucceeded).toHaveBeenCalledWith(expect.objectContaining(correlation)) }) + it('captures custom-provider usage once and leaves official generation capture to the server', async () => { + llmStreamMock.mockImplementation(async (_model: string, _chatProvider: ChatProvider, _messages: Message[], options: any) => { + await options.onStreamEvent({ type: 'text-delta', text: 'ok' }) + await options.onStreamEvent({ type: 'finish', finishReason: 'stop' }) + await options.onUsage({ + inputTokens: 12, + outputTokens: 8, + totalTokens: 20, + source: 'reported', + }) + }) + + const store = useChatOrchestratorStore() + await store.ingest('custom turn', { + model: 'gpt-test', + chatProvider: provider, + }) + + expect(chatAnalyticsMocks.trackAiGeneration).toHaveBeenCalledWith({ + conversation_id: 'session-1', + round_id: expect.any(String), + provider_type: 'custom', + provider_id: 'mock-provider', + model_id: 'gpt-test', + usage_source: 'reported', + input_tokens: 12, + output_tokens: 8, + total_tokens: 20, + }) + + chatAnalyticsMocks.trackAiGeneration.mockClear() + activeProviderRef.value = 'official-provider' + await store.ingest('official turn', { + model: 'chat-auto', + chatProvider: provider, + }) + + expect(chatAnalyticsMocks.trackAiGeneration).not.toHaveBeenCalled() + expect(llmStreamMock.mock.calls[1]?.[3]?.headers).toEqual({ + [AIRI_CHAT_APP_SURFACE_HEADER]: 'web', + [AIRI_CHAT_SESSION_ID_HEADER]: 'session-1', + [AIRI_CHAT_ROUND_ID_HEADER]: expect.any(String), + }) + }) + it('emits second turn analytics from chat sends', async () => { activeProviderRef.value = 'official-provider' llmStreamMock.mockImplementation(async (_model: string, _chatProvider: ChatProvider, _messages: Message[], options: any) => { diff --git a/packages/stage-ui/src/stores/chat.ts b/packages/stage-ui/src/stores/chat.ts index 665c57eeb..bba2cf3f3 100644 --- a/packages/stage-ui/src/stores/chat.ts +++ b/packages/stage-ui/src/stores/chat.ts @@ -10,8 +10,13 @@ import { nanoid } from 'nanoid' import { defineStore, storeToRefs } from 'pinia' import { ref, toRaw, watch } from 'vue' -import { useAnalytics } from '../composables' +import { getConversationAnalyticsSurface, useAnalytics } from '../composables' import { activeTurnSpan, startSpan } from '../composables/use-io-tracer' +import { + AIRI_CHAT_APP_SURFACE_HEADER, + AIRI_CHAT_ROUND_ID_HEADER, + AIRI_CHAT_SESSION_ID_HEADER, +} from '../libs/analytics-headers' import { extractMessageText, isCloudSyncableMessage } from '../libs/chat-sync' import { createMinecraftContext } from './chat/context-providers' import { useChatContextStore } from './chat/context-store' @@ -56,6 +61,7 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => { trackLlmRequestStarted, trackLlmFirstToken, trackAssistantResponseRendered, + trackAiGeneration, trackMessageRound, trackMessageRoundFailed, trackChatActivationStarted, @@ -83,6 +89,12 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => { options?: StreamOptions, ) { let llmTextLength = 0 + const headers = { ...options?.headers } + if (providerMode(activeProvider.value) === 'official' && options?.requestCorrelation) { + headers[AIRI_CHAT_SESSION_ID_HEADER] = options.requestCorrelation.conversationId + headers[AIRI_CHAT_ROUND_ID_HEADER] = options.requestCorrelation.roundId + headers[AIRI_CHAT_APP_SURFACE_HEADER] = getConversationAnalyticsSurface() + } const hadExistingTurn = !!activeTurnSpan.value if (!hadExistingTurn) { @@ -101,6 +113,7 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => { try { await llmStore.stream(model, chatProvider, messages, { ...options, + headers, onStreamEvent: async (event: StreamEvent) => { if (isTextDelta(event)) { if (!llmFirstTokenEmitted) { @@ -216,13 +229,35 @@ export const useChatOrchestratorStore = defineStore('chat-orchestrator', () => { latency_ms: latencyMs, }) }, - onMessageRound: ({ conversationId, roundId, turnIndex, durationMs, hasVoice, model }) => trackMessageRound({ + onLlmGeneration: ({ conversationId, roundId, model, provider, inputTokens, outputTokens, totalTokens, usageSource }) => { + const mode = providerMode(provider) + // The official path is captured server-side from authoritative upstream usage. + if (mode !== 'custom') + return + + trackAiGeneration({ + conversation_id: conversationId, + round_id: roundId, + provider_type: mode, + provider_id: provider, + model_id: model, + usage_source: usageSource, + input_tokens: inputTokens, + output_tokens: outputTokens, + total_tokens: totalTokens, + }) + }, + onMessageRound: ({ conversationId, roundId, turnIndex, durationMs, hasVoice, model, inputTokens, outputTokens, totalTokens, usageSource }) => trackMessageRound({ conversation_id: conversationId, round_id: roundId, turn_index: turnIndex, duration_ms: durationMs, has_voice: hasVoice, model, + input_tokens: inputTokens, + output_tokens: outputTokens, + total_tokens: totalTokens, + usage_source: usageSource, }), onMessageRoundFailed: ({ conversationId, roundId, turnIndex, model, provider, errorCode, failureStage, source }) => trackMessageRoundFailed({ conversation_id: conversationId, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 56ce14fb5..5fd5d9e5c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -6614,11 +6614,11 @@ packages: '@esbuild-kit/core-utils@3.3.2': resolution: {integrity: sha512-sPRAnw9CdSsRmEtnsl2WXWdyquogVpB3yZ3dgwJfe8zrOzTsV7cJvmwrKVa+0ma5BoiGJ+BoqkMvawbayKUsqQ==} - deprecated: 'Merged into tsx: https://tsx.is' + deprecated: 'Merged into tsx: https://tsx.hirok.io' '@esbuild-kit/esm-loader@2.6.5': resolution: {integrity: sha512-FxEMIkJKnodyA1OaCUoEvbYRkoZlLZ4d/eXFu9Fh8CbBBgP5EmZxrfTRyN0qpXZ4vOvqnE5YdRdcrmUUXuU+dA==} - deprecated: 'Merged into tsx: https://tsx.is' + deprecated: 'Merged into tsx: https://tsx.hirok.io' '@esbuild/aix-ppc64@0.25.12': resolution: {integrity: sha512-Hhmwd6CInZ3dwpuGTF8fJG6yoWmsToE+vYgD4nytZVxcu1ulHpUQRAB1UJ8+N1Am3Mz4+xOByoQoSZf4D+CpkA==}