diff --git a/packages/stage-shared/src/perf/io-trace.ts b/packages/stage-shared/src/perf/io-trace.ts index ff5d50e44d..8acb01b193 100644 --- a/packages/stage-shared/src/perf/io-trace.ts +++ b/packages/stage-shared/src/perf/io-trace.ts @@ -114,6 +114,57 @@ export const IOAttributes = { Cancelled: `${customPrefix}.cancelled`, } as const +/** + * The attributes that the agent trace upload sends. The upload drops every other attribute before it leaves the client. + * + * Content is never on this list: ASR, TTS, and wake word text, streaming control parameters and raw tokens, + * tool arguments, and tool results. The API applies its own allowlist too. + */ +export const IOUploadAttributes: ReadonlySet = new Set([ + IOAttributes.GenAIRequestModel, + IOAttributes.GenAIProviderName, + IOAttributes.GenAIConversationId, + IOAttributes.GenAIToolName, + IOAttributes.GenAIToolCallId, + IOAttributes.ErrorType, + IOAttributes.TurnId, + IOAttributes.RequestId, + IOAttributes.Cancelled, + IOAttributes.LLM_TTFT, + IOAttributes.LLMStepIndex, + IOAttributes.LLMInputMessageCount, + IOAttributes.LLMInputUserMessageCount, + IOAttributes.LLMOutputChunkCount, + IOAttributes.LLMTextLength, + IOAttributes.WakeCharacterId, + IOAttributes.VoiceInputId, + IOAttributes.VoiceInputTrigger, + IOAttributes.VoiceInputOutcome, + IOAttributes.VoiceInputReason, + IOAttributes.VADAudioDurationMs, + IOAttributes.VADAborted, + IOAttributes.ASRAbort, + IOAttributes.EndJudgement, + IOAttributes.EndJudgeTranscriptLength, + IOAttributes.RephraseSegmentCount, + IOAttributes.RephraseChangedCount, + IOAttributes.TTSSegmentId, + IOAttributes.TTSChunkReason, + IOAttributes.TTSAudioDurationMs, + IOAttributes.TTSInterrupted, + IOAttributes.TTSInterruptReason, + IOAttributes.TTSCanceled, + IOAttributes.TTSNoAudio, + IOAttributes.StreamingControlParserName, + IOAttributes.StreamingControlParsed, + IOAttributes.StreamingControlMatched, + IOAttributes.StreamingControlHandlerCount, + IOAttributes.StreamingControlReason, + IOAttributes.StreamingControlCallName, + IOAttributes.StreamingControlTokenType, + IOAttributes.StreamingControlTokenLength, +]) + export const IOEvents = { // Non-standard LLMFirstToken: `${customPrefix}.llm.first_token`, diff --git a/packages/stage-ui/src/libs/agent-trace/exporter.test.ts b/packages/stage-ui/src/libs/agent-trace/exporter.test.ts new file mode 100644 index 0000000000..0a8bd9d84e --- /dev/null +++ b/packages/stage-ui/src/libs/agent-trace/exporter.test.ts @@ -0,0 +1,105 @@ +import type { AgentTraceUploadSpan } from './exporter' + +import { IOAttributes, IOSpanNames } from '@proj-airi/stage-shared' +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { startSpan, subscribeIOSpan } from '../../composables/use-io-tracer' +import { createAgentTraceExporter } from './exporter' + +const unsubscribers: (() => void)[] = [] + +afterEach(() => { + unsubscribers.splice(0).forEach(unsubscribe => unsubscribe()) +}) + +function setup(options: { enabled?: () => boolean, token?: string | null, fetch?: typeof globalThis.fetch } = {}) { + const uploads: { spans: AgentTraceUploadSpan[], resource: unknown }[] = [] + const fetch = options.fetch ?? vi.fn(async (_url: RequestInfo | URL, init?: RequestInit) => { + uploads.push(JSON.parse(init!.body as string)) + return new Response(null, { status: 202 }) + }) + const exporter = createAgentTraceExporter({ + enabled: options.enabled ?? (() => true), + resource: () => ({ appSurface: 'web', appVersion: '1.0.0' }), + token: () => options.token === undefined ? 'token' : options.token, + endpoint: 'https://api.example/api/v1/agent-traces', + fetch, + }) + unsubscribers.push(subscribeIOSpan(exporter.handleSpan)) + return { exporter, uploads, fetch } +} + +/** Runs one turn: a voice input root, the turn, one LLM inference, and a tool call. */ +function runTurn(steps: boolean[], mark: (traceId: string, official: boolean) => void) { + const voice = startSpan(IOSpanNames.VoiceInput, undefined, { [IOAttributes.VoiceInputTrigger]: 'speech' }) + const recognition = startSpan(IOSpanNames.SpeechRecognition, voice, { [IOAttributes.ASRText]: 'what the user said' }) + recognition.end() + voice.end() + const turn = startSpan(IOSpanNames.InteractionTurn, voice, { [IOAttributes.TurnId]: 'turn-1', [IOAttributes.GenAIConversationId]: 'session-1' }) + const traceId = turn.spanContext().traceId + for (const official of steps) + mark(traceId, official) + const llm = startSpan(IOSpanNames.LLMInference, turn, { [IOAttributes.GenAIRequestModel]: 'model', [IOAttributes.LLMInputMessageRoles]: 'user' }) + llm.end() + turn.end() + const playback = startSpan(IOSpanNames.AudioPlayback, turn, { [IOAttributes.TTSText]: 'what the model said', [IOAttributes.TTSSegmentId]: 'segment-1' }) + playback.end() + return traceId +} + +describe('agent trace exporter', () => { + it('uploads a turn that used only the official provider, without content attributes', async () => { + const { exporter, uploads } = setup() + const traceId = runTurn([true, true], exporter.markStep) + await exporter.flush() + + expect(uploads).toHaveLength(1) + expect(uploads[0].resource).toEqual({ appSurface: 'web', appVersion: '1.0.0' }) + const spans = uploads[0].spans + expect(spans.map(span => span.name)).toEqual([IOSpanNames.SpeechRecognition, IOSpanNames.VoiceInput, IOSpanNames.LLMInference, IOSpanNames.InteractionTurn, IOSpanNames.AudioPlayback]) + expect(spans.every(span => span.traceId === traceId)).toBe(true) + const turn = spans.find(span => span.name === IOSpanNames.InteractionTurn)! + expect(turn.attributes).toEqual({ [IOAttributes.TurnId]: 'turn-1', [IOAttributes.GenAIConversationId]: 'session-1' }) + const body = JSON.stringify(uploads) + expect(body).not.toContain('what the user said') + expect(body).not.toContain('what the model said') + expect(body).not.toContain(IOAttributes.Subsystem) + expect(body).not.toContain(IOAttributes.LLMInputMessageRoles) + }) + + it('drops a turn with a step that used another provider, and a turn with no step', async () => { + const { exporter, uploads, fetch } = setup() + runTurn([true, false], exporter.markStep) + runTurn([], exporter.markStep) + await exporter.flush() + expect(fetch).not.toHaveBeenCalled() + expect(uploads).toHaveLength(0) + }) + + it('uploads nothing while analytics is off or the user is signed out', async () => { + let enabled = false + const off = setup({ enabled: () => enabled }) + runTurn([true], off.exporter.markStep) + enabled = true + await off.exporter.flush() + expect(off.fetch).not.toHaveBeenCalled() + + const signedOut = setup({ token: null }) + runTurn([true], signedOut.exporter.markStep) + await signedOut.exporter.flush() + expect(signedOut.fetch).not.toHaveBeenCalled() + }) + + it('drops a batch whose upload fails, without throwing or retrying', async () => { + const fetch = vi.fn(async () => { + throw new TypeError('offline') + }) + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + const { exporter } = setup({ fetch }) + runTurn([true], exporter.markStep) + await expect(exporter.flush()).resolves.toBeUndefined() + await exporter.flush() + expect(fetch).toHaveBeenCalledTimes(1) + warn.mockRestore() + }) +}) diff --git a/packages/stage-ui/src/libs/agent-trace/exporter.ts b/packages/stage-ui/src/libs/agent-trace/exporter.ts new file mode 100644 index 0000000000..43dbbcd4c2 --- /dev/null +++ b/packages/stage-ui/src/libs/agent-trace/exporter.ts @@ -0,0 +1,170 @@ +import type { ReadableSpan } from '@opentelemetry/sdk-trace-base' + +import { SpanStatusCode } from '@opentelemetry/api' +import { hrTimeToMilliseconds } from '@opentelemetry/core' +import { IOSpanNames, IOUploadAttributes } from '@proj-airi/stage-shared' + +/** One span of `POST /api/v1/agent-traces`. */ +export interface AgentTraceUploadSpan { + traceId: string + spanId: string + parentSpanId: string | null + name: string + startTimeUnixMs: number + endTimeUnixMs: number + status: 'ok' | 'error' + attributes: Record +} + +export interface AgentTraceResource { + /** The names of the `x-airi-app-surface` chat header, which the request log stores. */ + appSurface: 'web' | 'electron' | 'mobile' + appVersion?: string +} + +export interface AgentTraceExporterOptions { + /** Reads the analytics setting. Spans that end while it is off are dropped. */ + enabled: () => boolean + resource: () => AgentTraceResource + /** The upload is skipped without a token. */ + token: () => string | null | undefined + endpoint: string + fetch?: typeof globalThis.fetch + now?: () => number +} + +/** The API accepts at most this many spans in one upload. */ +const MAX_BATCH_SPANS = 500 +/** A trace whose turn has not ended in this time is dropped, for example a voice input that was never sent. */ +const TRACE_TTL_MS = 10 * 60_000 +const MAX_TRACES = 64 +const MAX_PENDING_SPANS = 200 +const MAX_QUEUED_SPANS = 2_000 +/** Browsers refuse a `keepalive` request with a larger body. */ +const KEEPALIVE_BODY_LIMIT = 60_000 + +interface TraceState { + spans: AgentTraceUploadSpan[] + official: boolean + external: boolean + verdict?: 'upload' | 'drop' + touchedAt: number +} + +export function toUploadSpan(span: ReadableSpan): AgentTraceUploadSpan { + const attributes: AgentTraceUploadSpan['attributes'] = {} + for (const [key, value] of Object.entries(span.attributes)) { + if (IOUploadAttributes.has(key) && (typeof value === 'string' || typeof value === 'number' || typeof value === 'boolean')) + attributes[key] = value + } + return { + traceId: span.spanContext().traceId, + spanId: span.spanContext().spanId, + parentSpanId: span.parentSpanContext?.spanId ?? null, + name: span.name, + startTimeUnixMs: Math.round(hrTimeToMilliseconds(span.startTime)), + endTimeUnixMs: Math.round(hrTimeToMilliseconds(span.endTime)), + status: span.status.code === SpanStatusCode.ERROR ? 'error' : 'ok', + attributes, + } +} + +/** + * Collects ended IO tracer spans and uploads the traces of turns that used only the official provider. + * + * The chat store marks each step of a turn as official or external. The exporter holds the spans of a trace until + * its `Interaction turn` span ends, and then uploads them only if every step was official. A trace with an external + * step, or with no step, is dropped, so no data of a provider that AIRI does not run leaves the client. + * Spans that end after the turn, such as audio playback, follow the decision of their trace. + * + * A failed upload drops its batch. It does not retry, and it never blocks a turn. + */ +export function createAgentTraceExporter(options: AgentTraceExporterOptions) { + const now = options.now ?? Date.now + const traces = new Map() + const queue: AgentTraceUploadSpan[] = [] + + /** Returns the state of a trace. The map stays in the order of last use, so pruning removes the least recent first. */ + function stateOf(traceId: string) { + const state = traces.get(traceId) ?? { spans: [], official: false, external: false, touchedAt: 0 } + state.touchedAt = now() + traces.delete(traceId) + traces.set(traceId, state) + return state + } + + function prune() { + const expired = now() - TRACE_TTL_MS + for (const [traceId, state] of traces) { + if (state.touchedAt >= expired && traces.size <= MAX_TRACES) + break + traces.delete(traceId) + } + } + + function enqueue(spans: AgentTraceUploadSpan[]) { + queue.push(...spans) + if (queue.length > MAX_QUEUED_SPANS) + queue.splice(0, queue.length - MAX_QUEUED_SPANS) + } + + /** Records that a step of the turn in `traceId` sent its request to the official provider, or to another one. */ + function markStep(traceId: string, official: boolean) { + const state = stateOf(traceId) + if (official) + state.official = true + else + state.external = true + } + + function handleSpan(span: ReadableSpan) { + if (!options.enabled()) + return + const upload = toUploadSpan(span) + const state = stateOf(upload.traceId) + if (state.verdict === 'upload') { + enqueue([upload]) + } + else if (!state.verdict) { + if (state.spans.length < MAX_PENDING_SPANS) + state.spans.push(upload) + if (span.name === IOSpanNames.InteractionTurn) { + state.verdict = state.official && !state.external ? 'upload' : 'drop' + if (state.verdict === 'upload') + enqueue(state.spans) + state.spans = [] + } + } + prune() + } + + /** Uploads the queued spans. `keepalive` lets the request outlive a page that is being hidden or closed. */ + async function flush(flushOptions: { keepalive?: boolean } = {}) { + if (queue.length === 0) + return + const token = options.token() + if (!options.enabled() || !token) { + queue.length = 0 + return + } + const fetch = options.fetch ?? globalThis.fetch + while (queue.length > 0) { + const spans = queue.splice(0, MAX_BATCH_SPANS) + const body = JSON.stringify({ resource: options.resource(), spans }) + try { + await fetch(options.endpoint, { + method: 'POST', + headers: { 'Authorization': `Bearer ${token}`, 'Content-Type': 'application/json' }, + body, + credentials: 'omit', + keepalive: flushOptions.keepalive === true && body.length <= KEEPALIVE_BODY_LIMIT, + }) + } + catch (error) { + console.warn('[agent-trace] Failed to upload spans:', error) + } + } + } + + return { handleSpan, markStep, flush } +} diff --git a/packages/stage-ui/src/libs/agent-trace/index.ts b/packages/stage-ui/src/libs/agent-trace/index.ts new file mode 100644 index 0000000000..44638b7cb8 --- /dev/null +++ b/packages/stage-ui/src/libs/agent-trace/index.ts @@ -0,0 +1,37 @@ +import type { AgentTraceExporterOptions } from './exporter' + +import { subscribeIOSpan } from '../../composables/use-io-tracer' +import { createAgentTraceExporter } from './exporter' + +const FLUSH_INTERVAL_MS = 5_000 + +let exporter: ReturnType | undefined + +/** + * Starts the upload of IO tracer spans to the API. Later calls do nothing. + * + * Spans upload every few seconds, and when the page is hidden. + */ +export function startAgentTraceUpload(options: AgentTraceExporterOptions) { + if (exporter) + return + const started = createAgentTraceExporter(options) + exporter = started + subscribeIOSpan(started.handleSpan) + setInterval(() => void started.flush(), FLUSH_INTERVAL_MS) + if (typeof document !== 'undefined') { + document.addEventListener('visibilitychange', () => { + if (document.visibilityState === 'hidden') + void started.flush({ keepalive: true }) + }) + } +} + +/** + * Records the provider of one step of the turn in `traceId`. + * + * Only a trace whose steps all used the official provider is uploaded. Does nothing before the upload starts. + */ +export function markAgentTraceStep(traceId: string, official: boolean) { + exporter?.markStep(traceId, official) +} diff --git a/packages/stage-ui/src/libs/product-signals/index.ts b/packages/stage-ui/src/libs/product-signals/index.ts index cabc46059d..78cdac3cac 100644 --- a/packages/stage-ui/src/libs/product-signals/index.ts +++ b/packages/stage-ui/src/libs/product-signals/index.ts @@ -8,6 +8,9 @@ import { useBuildInfo } from '../../composables/use-build-info' import { useAuthStore } from '../../stores/auth' import { useAiriCardStore } from '../../stores/modules/airi-card' import { useSettingsAnalytics } from '../../stores/settings/analytics' +import { startAgentTraceUpload } from '../agent-trace' +import { getAuthToken } from '../auth' +import { SERVER_URL } from '../server' import { captureAnalyticsEvent, configureAnalyticsAdapter, @@ -82,6 +85,15 @@ class Analytics implements AnalyticsRecorder { if (authStore.isAuthenticated && authStore.user?.id) identifyAnalyticsUser(authStore.user.id) + // Agent traces follow the analytics setting. The exporter reads it again before each upload. + startAgentTraceUpload({ + enabled: () => isAnalyticsAvailableInBuild() && settingsAnalytics.analyticsEnabled, + // The API uses the surface names of the chat header, where the desktop app is `electron`. + resource: () => ({ appSurface: isStageTamagotchi() ? 'electron' : isStageCapacitor() ? 'mobile' : 'web', appVersion: buildInfo.version }), + token: getAuthToken, + endpoint: `${SERVER_URL}/api/v1/agent-traces`, + }) + authStore.onAuthenticated(() => { if (authStore.user?.id) identifyAnalyticsUser(authStore.user.id) diff --git a/packages/stage-ui/src/stores/chat.contract.test.ts b/packages/stage-ui/src/stores/chat.contract.test.ts index 00687a4887..6cf4158297 100644 --- a/packages/stage-ui/src/stores/chat.contract.test.ts +++ b/packages/stage-ui/src/stores/chat.contract.test.ts @@ -65,6 +65,8 @@ const ioTracerMocks = vi.hoisted(() => { } }) +const markAgentTraceStepMock = vi.hoisted(() => vi.fn()) + const llmStreamMock = vi.fn() const trackFirstMessageMock = vi.fn() const chatAnalyticsMocks = vi.hoisted(() => ({ @@ -122,6 +124,10 @@ vi.mock('../composables/use-io-tracer', () => ({ registerLlmInference: ioTracerMocks.registerLlmInference, })) +vi.mock('../libs/agent-trace', () => ({ + markAgentTraceStep: markAgentTraceStepMock, +})) + vi.mock('./chat/context-providers', () => ({ createMinecraftContext: () => createMinecraftContextMock(), createRuntimePromptContext: (prompt: string) => createRuntimePromptContextMock(prompt), @@ -385,6 +391,7 @@ describe('chat store contract', () => { ioTracerMocks.startSpanMock.mockClear() ioTracerMocks.registerLlmInference.mockClear() ioTracerMocks.unregisterLlmInference.mockClear() + markAgentTraceStepMock.mockClear() ioTracerMocks.takeVoiceTurnParent.mockReset() ioTracerMocks.takeComposerVoiceParent.mockReset() activeSessionIdRef.value = 'session-1' @@ -1773,6 +1780,22 @@ describe('chat store contract', () => { expect(ioTracerMocks.unregisterLlmInference).toHaveBeenCalledOnce() }) + it('marks the provider of each step on the trace of its turn', async () => { + cardSelections.set('alice', { provider: 'official-provider', model: 'chat-auto' }) + llmStreamMock.mockImplementationOnce(async (_model: string, _provider: GenerationProvider, _messages: Conversation, options: StreamOptions) => { + // A tool step changes the character to a provider that AIRI does not run. + cardSelections.set('alice', { provider: 'alice-provider', model: 'alice-model' }) + await options.resolveStep!() + await options.onStreamEvent?.({ type: 'finish' }) + }) + + const store = useChatStore() + await store.send({ sessionId: 'session-1', text: 'Use a tool' }) + + const traceId = ioTracerMocks.spans.find(span => span.name === IOSpanNames.InteractionTurn)!.spanContext().traceId + expect(markAgentTraceStepMock.mock.calls).toEqual([[traceId, true], [traceId, false]]) + }) + it('continues a composer dictation trace when the turn has no voice input of its own', async () => { const dictation = { name: IOSpanNames.VoiceInput } ioTracerMocks.takeComposerVoiceParent.mockReturnValueOnce(dictation) diff --git a/packages/stage-ui/src/stores/chat.ts b/packages/stage-ui/src/stores/chat.ts index 59de292263..7f08ef95f8 100644 --- a/packages/stage-ui/src/stores/chat.ts +++ b/packages/stage-ui/src/stores/chat.ts @@ -24,6 +24,7 @@ import { useAiriRuntimePrompt } from '../composables/use-airi-runtime-prompt' import { registerLlmInference, registerTurnSpan, startSpan, takeComposerVoiceParent, takeVoiceTurnParent } from '../composables/use-io-tracer' import { useChatVision } from '../composables/vision/use-chat-vision' import { useVisionInference } from '../composables/vision/use-vision-inference' +import { markAgentTraceStep } from '../libs/agent-trace' import { chatAssetIdFrom, inlineConversationAssets, storeChatAttachments } from '../libs/chat-assets' import { extractMessageText, isCloudSyncableMessage } from '../libs/chat-sync' import { createChatAnalyticsHooks, getProviderMode } from '../libs/product-signals/events/chat' @@ -660,7 +661,11 @@ export const useChatStore = defineStore('chat', () => { // Analytics correlation goes only to the official provider. A resolved step can change the provider, so each request rebuilds it. function requestHeaders(providerId: string | undefined) { const headers = { ...options?.headers } - if (getProviderMode(providerId ?? activeProvider.value) === 'official' && options?.requestCorrelation) { + const official = getProviderMode(providerId ?? activeProvider.value) === 'official' + // The trace of a turn uploads only when every step used the official provider. + if (correlation && turnSpan) + markAgentTraceStep(turnSpan.spanContext().traceId, official) + if (official && options?.requestCorrelation) { headers[AIRI_CHAT_SESSION_ID_HEADER] = options.requestCorrelation.conversationId headers[AIRI_CHAT_ROUND_ID_HEADER] = options.requestCorrelation.turnId headers[AIRI_CHAT_APP_SURFACE_HEADER] = getConversationAnalyticsSurface() @@ -861,9 +866,10 @@ export const useChatStore = defineStore('chat', () => { llmOutputChunkLengths.push(event.text.length) if (!llmFirstTokenEmitted) { llmFirstTokenEmitted = true - llmSpan.addEvent(IOEvents.LLMFirstToken, { - [IOAttributes.LLM_TTFT]: performance.now() - llmRequestTs, - }) + const ttft = performance.now() - llmRequestTs + llmSpan.addEvent(IOEvents.LLMFirstToken, { [IOAttributes.LLM_TTFT]: ttft }) + // The trace upload has no events, so the attribute carries the value too. + llmSpan.setAttribute(IOAttributes.LLM_TTFT, ttft) } llmTextLength += event.text.length }