@ai-sdk/google/src/realtime/google-realtime-event-mapper.ts

import type {
  Experimental_RealtimeModelV4 as RealtimeModelV4,
  Experimental_RealtimeModelV4ClientEvent as RealtimeModelV4ClientEvent,
  Experimental_RealtimeModelV4FunctionCallOutput as RealtimeModelV4FunctionCallOutput,
  Experimental_RealtimeModelV4ServerEvent as RealtimeModelV4ServerEvent,
  Experimental_RealtimeModelV4SessionConfig as RealtimeModelV4SessionConfig,
  Experimental_RealtimeModelV4Usage as RealtimeModelV4Usage,
} from '@ai-sdk/provider';
import { isRecord, safeParseJSON } from '@ai-sdk/provider-utils';
import { getModelPath } from '../get-model-path';
import type { GoogleRealtimeModelOptions } from './google-realtime-model-options';

type GoogleRealtimeFunctionCall = {
  id: string;
  name: string;
  args?: Record<string, unknown>;
};

type GoogleRealtimeServerContent = {
  generationComplete?: boolean;
  interactionStatus?: string;
  interrupted?: boolean;
  modelTurn?: {
    parts?: Array<{
      inlineData?: { data?: string };
      text?: string;
    }>;
  };
  outputTranscription?: { text?: string; finished?: boolean };
  inputTranscription?: { text?: string; finished?: boolean };
  turnComplete?: boolean;
  waitingForInput?: boolean;
};

type GoogleRealtimeTokenDetail = {
  modality?: string;
  tokenCount?: number;
};

type GoogleRealtimeUsageMetadata = {
  promptTokenCount?: number;
  responseTokenCount?: number;
  totalTokenCount?: number;
  toolUsePromptTokenCount?: number;
  thoughtsTokenCount?: number;
  cachedContentTokenCount?: number;
  promptTokensDetails?: GoogleRealtimeTokenDetail[];
  responseTokensDetails?: GoogleRealtimeTokenDetail[];
};

type GoogleRealtimeWireEvent = {
  setupComplete?: unknown;
  toolCall?: {
    functionCalls?: GoogleRealtimeFunctionCall[];
  };
  toolCallCancellation?: unknown;
  serverContent?: GoogleRealtimeServerContent;
  usageMetadata?: GoogleRealtimeUsageMetadata;
  inputTranscription?: { text?: string; finished?: boolean };
  goAway?: { timeLeft?: string };
  sessionResumptionUpdate?: {
    newHandle?: string;
    resumable?: boolean;
    lastConsumedClientMessageIndex?: string;
  };
};

/**
 * Stateful event mapper for Google's Gemini Live API.
 *
 * Unlike OpenAI/xAI, Google's events don't have response/item IDs and
 * a single message can contain multiple pieces of data. This class
 * tracks turn state to generate consistent synthetic IDs.
 */
export class GoogleRealtimeEventMapper {
  private turnCounter = 0;
  private inputTranscriptionCounter = 0;
  private inputTranscriptionBuffer = '';
  private inputTranscriptionBoundary: 'none' | 'interrupted' | 'turn-complete' =
    'none';
  private preserveInputTranscriptionAcrossTurnComplete = false;
  private hasAudio = false;
  private hasText = false;
  private hasTranscript = false;
  private turnClosed = false;
  private inputAudioRate = 16000;

  private get responseId(): string {
    return `google-resp-${this.turnCounter}`;
  }

  private get itemId(): string {
    return `google-item-${this.turnCounter}`;
  }

  /**
   * Rolls over to the next turn lazily, once model content for the next turn
   * arrives. `turnComplete` merely marks the current turn closed.
   * This keeps an output transcript that arrives shortly after `turnComplete`
   * attached to the turn it belongs to, since Google delivers transcription
   * independently with no guaranteed ordering relative to `turnComplete`.
   */
  private beginTurnIfClosed(): void {
    if (!this.turnClosed) return;
    this.turnCounter++;
    this.hasAudio = false;
    this.hasText = false;
    this.hasTranscript = false;
    this.turnClosed = false;
    // Once the next response starts, its completion must delimit its own
    // input even if the interrupted response never sent a turnComplete.
    this.preserveInputTranscriptionAcrossTurnComplete = false;
  }

  // Google streams input transcription as a sequence of non-accumulating
  // fragments for a single user utterance (e.g. "The quick brown fox",
  // " jumps over the", " lazy dog."). Concatenate consecutive fragments into
  // one running transcript under a stable synthetic id so the realtime reducer
  // (which overwrites the message for a given id) surfaces them as a single
  // coherent user message rather than one message per fragment. Google's
  // `finished` signal is the authoritative utterance boundary. For older
  // payloads without that signal, response boundaries remain a fallback.
  //
  // Input transcriptions can arrive independently of response turn events, so
  // their synthetic IDs use a sequence independent of `turnCounter`.
  private accumulateInputTranscription({
    text,
    finished,
  }: {
    text?: string;
    finished?: boolean;
  }): { itemId: string; transcript: string } | undefined {
    // Google can send `finished` separately from text. Process it even when
    // text is absent or empty, without emitting an empty user message.
    if (!text) {
      if (finished === true) {
        this.finishInputTranscription();
      }
      return undefined;
    }

    const hasFinishedSignal = finished != null;
    if (
      this.inputTranscriptionBuffer !== '' &&
      (this.inputTranscriptionBoundary === 'interrupted' ||
        (this.inputTranscriptionBoundary === 'turn-complete' &&
          !hasFinishedSignal))
    ) {
      this.finishInputTranscription();
    }
    this.inputTranscriptionBoundary = 'none';
    this.inputTranscriptionBuffer += text;
    const result = {
      itemId: `google-input-${this.inputTranscriptionCounter}`,
      transcript: this.inputTranscriptionBuffer,
    };
    if (finished === true) {
      this.finishInputTranscription();
    }
    return result;
  }

  private finishInputTranscription(): void {
    if (this.inputTranscriptionBuffer === '') return;
    this.inputTranscriptionCounter++;
    this.inputTranscriptionBuffer = '';
  }

  parseServerEvent(
    raw: unknown,
  ): RealtimeModelV4ServerEvent | RealtimeModelV4ServerEvent[] {
    const events = this.parseEvent(raw);
    const usage = mapGoogleUsage(
      (raw as GoogleRealtimeWireEvent).usageMetadata,
    );
    if (usage == null) return events;
    const list = Array.isArray(events) ? events : [events];
    // Only a shared wire frame establishes an association with a completion.
    if (list.some(event => event.type === 'response-done')) {
      const mapped = list.map(event =>
        event.type === 'response-done' ? { ...event, usage } : event,
      );
      return Array.isArray(events) ? mapped : mapped[0];
    }
    const update: RealtimeModelV4ServerEvent = {
      type: 'custom',
      rawType: 'usageMetadata',
      usage,
      raw,
    };
    if (
      list.length === 1 &&
      list[0].type === 'custom' &&
      list[0].rawType === 'usageMetadata'
    )
      return update;
    return [...list, update];
  }

  private parseEvent(
    raw: unknown,
  ): RealtimeModelV4ServerEvent | RealtimeModelV4ServerEvent[] {
    const data = raw as GoogleRealtimeWireEvent;

    if (data.setupComplete != null) {
      return { type: 'session-created', raw };
    }

    if (data.toolCall != null) {
      this.beginTurnIfClosed();
      const functionCalls = data.toolCall.functionCalls ?? [];
      return functionCalls.flatMap(functionCall => {
        const args = JSON.stringify(functionCall.args ?? {});
        return [
          {
            type: 'function-call-arguments-delta' as const,
            responseId: this.responseId,
            itemId: this.itemId,
            callId: functionCall.id,
            delta: args,
            raw,
          },
          {
            type: 'function-call-arguments-done' as const,
            responseId: this.responseId,
            itemId: this.itemId,
            callId: functionCall.id,
            name: functionCall.name,
            arguments: args,
            raw,
          },
        ];
      });
    }

    if (data.toolCallCancellation != null) {
      return {
        type: 'custom',
        rawType: 'toolCallCancellation',
        raw,
      };
    }

    if (data.goAway != null) {
      return {
        type: 'custom',
        rawType: 'goAway',
        raw,
      };
    }

    if (data.sessionResumptionUpdate != null) {
      return {
        type: 'custom',
        rawType: 'sessionResumptionUpdate',
        raw,
      };
    }

    if (data.serverContent != null) {
      return this.parseServerContent(data.serverContent, raw);
    }

    if (data.inputTranscription != null) {
      const transcription = this.accumulateInputTranscription(
        data.inputTranscription,
      );
      if (transcription != null) {
        return {
          type: 'input-transcription-completed',
          ...transcription,
          raw,
        };
      }
    }

    return { type: 'custom', rawType: String(Object.keys(data)[0]), raw };
  }

  private parseServerContent(
    serverContent: GoogleRealtimeServerContent,
    raw: unknown,
  ): RealtimeModelV4ServerEvent | RealtimeModelV4ServerEvent[] {
    const events: RealtimeModelV4ServerEvent[] = [];

    if (serverContent.interrupted) {
      // A barge-in ends the current user utterance; the next input
      // transcription fragment belongs to a fresh utterance. The interrupted
      // response's trailing turnComplete must not split that new utterance.
      this.inputTranscriptionBoundary = 'interrupted';
      this.preserveInputTranscriptionAcrossTurnComplete = true;
      this.turnClosed = true;
      events.push({
        type: 'speech-started',
        raw,
      });
    }

    if (serverContent.modelTurn?.parts) {
      // New model response content marks the start of the next turn.
      this.beginTurnIfClosed();
      for (const part of serverContent.modelTurn.parts) {
        if (part.inlineData?.data) {
          this.hasAudio = true;
          events.push({
            type: 'audio-delta',
            responseId: this.responseId,
            itemId: this.itemId,
            delta: part.inlineData.data,
            raw,
          });
        }
        if (part.text) {
          this.hasText = true;
          events.push({
            type: 'text-delta',
            responseId: this.responseId,
            itemId: this.itemId,
            delta: part.text,
            raw,
          });
        }
      }
    }

    if (serverContent.outputTranscription?.text) {
      this.hasTranscript = true;
      events.push({
        type: 'audio-transcript-delta',
        responseId: this.responseId,
        itemId: this.itemId,
        delta: serverContent.outputTranscription.text,
        raw,
      });
    }

    if (serverContent.inputTranscription != null) {
      const transcription = this.accumulateInputTranscription(
        serverContent.inputTranscription,
      );
      if (transcription != null) {
        events.push({
          type: 'input-transcription-completed',
          ...transcription,
          raw,
        });
      }
    }

    // `generationComplete` means generation has stopped, but playback and the
    // turn can remain open. Keep it distinct from `response-done`, which is
    // emitted only when Google sends `turnComplete`.
    if (serverContent.generationComplete) {
      events.push({
        type: 'custom',
        rawType: 'generationComplete',
        raw,
      });
    }

    // `interactionStatus` (IN_PROGRESS | IDLE | WAITING_FOR_INPUT) is the definitive
    // session-activity signal for background-reasoning models: `turnComplete`
    // no longer implies the model is idle, since asynchronous tool calls and
    // audio may still follow. Surface it as a custom event so clients can
    // coordinate state on it.
    if (serverContent.interactionStatus != null) {
      events.push({
        type: 'custom',
        rawType: 'interactionStatus',
        raw,
      });
    }

    // `waitingForInput` is the always-on Proactive Audio turn-taking signal:
    // the model has yielded the floor and is not generating because it
    // expects the user to continue.
    if (serverContent.waitingForInput) {
      events.push({
        type: 'custom',
        rawType: 'waitingForInput',
        raw,
      });
    }

    if (serverContent.turnComplete) {
      if (this.hasAudio) {
        events.push({
          type: 'audio-done',
          responseId: this.responseId,
          itemId: this.itemId,
          raw,
        });
      }
      if (this.hasText) {
        events.push({
          type: 'text-done',
          responseId: this.responseId,
          itemId: this.itemId,
          raw,
        });
      }
      if (this.hasTranscript) {
        events.push({
          type: 'audio-transcript-done',
          responseId: this.responseId,
          itemId: this.itemId,
          raw,
        });
      }
      events.push({
        type: 'response-done',
        responseId: this.responseId,
        status: 'completed',
        raw,
      });
      // Mark the turn closed but defer advancing the counter until the next
      // response actually begins (see `beginTurnIfClosed`).
      this.turnClosed = true;
      if (this.preserveInputTranscriptionAcrossTurnComplete) {
        this.preserveInputTranscriptionAcrossTurnComplete = false;
      } else {
        // For payloads without `finished`, the completed response remains the
        // fallback signal that the next transcription is a new utterance.
        // A transcription carrying `finished` can still arrive late and close
        // the preceding utterance after this response boundary.
        if (this.inputTranscriptionBoundary !== 'interrupted') {
          this.inputTranscriptionBoundary = 'turn-complete';
        }
      }
    }

    if (events.length === 0) {
      return { type: 'custom', rawType: 'serverContent', raw };
    }

    return events.length === 1 ? events[0] : events;
  }

  serializeClientEvent(
    event: RealtimeModelV4ClientEvent,
    modelId: string,
  ): ReturnType<RealtimeModelV4['serializeClientEvent']> {
    switch (event.type) {
      case 'session-update':
        // Capture the configured capture rate so input audio blobs advertise
        // the real rate. Google accepts any rate as long as the blob's mimeType
        // matches; a mismatched label corrupts custom-rate audio.
        if (event.config.inputAudioFormat?.rate != null) {
          this.inputAudioRate = event.config.inputAudioFormat.rate;
        }
        return {
          setup: buildGoogleSessionConfig(event.config, modelId),
        };

      case 'input-audio-append':
        return {
          realtimeInput: {
            audio: {
              data: event.audio,
              mimeType: `audio/pcm;rate=${this.inputAudioRate}`,
            },
          },
        };

      case 'input-audio-commit':
        return {
          realtimeInput: {
            audioStreamEnd: true,
          },
        };

      case 'input-audio-clear':
      case 'response-create':
      case 'response-cancel':
      case 'conversation-item-truncate':
        return null;

      case 'conversation-item-create': {
        const item = event.item;
        switch (item.type) {
          case 'text-message':
            return {
              realtimeInput: {
                text: item.text,
              },
            };
          case 'function-call-output':
            return serializeFunctionCallOutput(item);
          case 'audio-message':
            return null;
        }
        break;
      }
    }

    return null;
  }
}

/**
 * Maps Gemini Live `usageMetadata` to normalized usage. Google reports usage
 * per modality in `promptTokensDetails` / `responseTokensDetails`.
 *
 * Field names are Live-API-specific: the Live `UsageMetadata` uses
 * `responseTokenCount` / `responseTokensDetails` for output, NOT the
 * `candidatesTokenCount` / `candidatesTokensDetails` of the regular
 * `generateContent` response (see google-language-model.ts). Do not "align"
 * these to `candidates*` — they are different wire shapes.
 *
 * Gross totals preserve provider counts even when modality details are absent
 * or contain modalities without a dedicated normalized bucket. Never guess a
 * text/audio split from those totals.
 */
function mapGoogleUsage(
  usageMetadata: GoogleRealtimeUsageMetadata | undefined,
): RealtimeModelV4Usage | undefined {
  if (usageMetadata == null) return undefined;

  const input = sumByModality(usageMetadata.promptTokensDetails);
  const output = sumByModality(usageMetadata.responseTokensDetails);

  return compactUsage({
    inputTokens:
      usageMetadata.promptTokenCount != null
        ? usageMetadata.promptTokenCount +
          (usageMetadata.toolUsePromptTokenCount ?? 0)
        : undefined,
    outputTokens:
      usageMetadata.responseTokenCount != null
        ? usageMetadata.responseTokenCount +
          (usageMetadata.thoughtsTokenCount ?? 0)
        : undefined,
    cachedInputTokens: usageMetadata.cachedContentTokenCount,
    inputTextTokens: input.text,
    inputAudioTokens: input.audio,
    outputTextTokens: output.text,
    outputAudioTokens: output.audio,
  });
}

/** Aggregate Gemini per-modality token details into text/audio buckets. */
function sumByModality(details: GoogleRealtimeTokenDetail[] | undefined): {
  text?: number;
  audio?: number;
} {
  if (details == null) return {};
  const totals: { text?: number; audio?: number } = {};
  for (const detail of details) {
    const count = detail.tokenCount;
    if (count == null) continue;
    const modality = detail.modality?.toUpperCase();
    if (modality === 'AUDIO') {
      totals.audio = (totals.audio ?? 0) + count;
    } else if (modality === 'TEXT') {
      totals.text = (totals.text ?? 0) + count;
    }
  }
  return totals;
}

/** Drop `undefined` fields; return `undefined` when nothing is observable. */
function compactUsage(
  usage: RealtimeModelV4Usage,
): RealtimeModelV4Usage | undefined {
  const result: RealtimeModelV4Usage = {};
  for (const [key, value] of Object.entries(usage)) {
    if (value != null) {
      result[key as keyof RealtimeModelV4Usage] = value;
    }
  }
  return Object.keys(result).length > 0 ? result : undefined;
}

/**
 * The value Gemini accepts for `functionResponse.response`.
 *
 * The field is typed `google.protobuf.Struct`, which is an object and nothing else, so
 * a string, number, array or `null` on the wire is a protocol violation and the socket
 * closes with 1007. `onToolCall` returns `unknown` and `addToolOutput` takes `unknown`,
 * so those shapes reach here as perfectly valid JSON.
 *
 * The field's own docstring says what to do with one: "Use `output` key to specify
 * function output ... If `output` and `error` keys are not specified, then whole
 * `response` is treated as function output." An object is passed through unchanged and
 * everything else is wrapped.
 */
function toFunctionResponseStruct(value: unknown): Record<string, unknown> {
  const isStruct =
    typeof value === 'object' && value !== null && !Array.isArray(value);
  return isStruct ? (value as Record<string, unknown>) : { output: value };
}

async function serializeFunctionCallOutput(
  item: RealtimeModelV4FunctionCallOutput,
): Promise<unknown> {
  const parseResult = await safeParseJSON({ text: item.output });
  const response = parseResult.success
    ? toFunctionResponseStruct(parseResult.value)
    : // Preserve non-JSON output in the required object wrapper.
      { output: item.output };

  return {
    toolResponse: {
      functionResponses: [
        {
          id: item.callId,
          name: item.name,
          response,
        },
      ],
    },
  };
}

/**
 * Live models that reason in the background (e.g. `gemini-3.8-live-extended-thinking`).
 * Google requires exactly one of `thinkingLevel` / `thinkingBudget` in their setup and
 * rejects `thinkingConfig` on every other Live model.
 */
function isThinkingLiveModel(modelId: string): boolean {
  const modelName = modelId.split('/').at(-1)?.toLowerCase() ?? '';
  return /^gemini-\d+\.\d+-live\b.*thinking/.test(modelName);
}

/**
 * Builds a Google-specific session configuration from a normalized config.
 * Used to construct the `bidiGenerateContentSetup` payload for auth token creation.
 */
export function buildGoogleSessionConfig(
  config: RealtimeModelV4SessionConfig | undefined,
  modelId: string,
): Record<string, unknown> {
  const setup: Record<string, unknown> = {
    model: getModelPath(modelId),
  };

  const { google, ...restProviderOptions } = config?.providerOptions ?? {};
  const googleOptions = isRecord(google)
    ? (google as GoogleRealtimeModelOptions)
    : undefined;

  const generationConfig: Record<string, unknown> = {};

  if (config?.outputModalities != null) {
    generationConfig.responseModalities = config.outputModalities.map(m =>
      m.toUpperCase(),
    );
  } else {
    generationConfig.responseModalities = ['AUDIO'];
  }

  if (config?.voice != null) {
    generationConfig.speechConfig = {
      voiceConfig: {
        prebuiltVoiceConfig: {
          voiceName: config.voice,
        },
      },
    };
  }

  setup.generationConfig = generationConfig;

  if (config?.instructions != null) {
    setup.systemInstruction = {
      parts: [{ text: config.instructions }],
    };
  }

  if (config?.tools != null && config.tools.length > 0) {
    setup.tools = [
      {
        functionDeclarations: config.tools.map(tool => ({
          name: tool.name,
          description: tool.description,
          parametersJsonSchema: tool.parameters,
          ...(googleOptions?.defaultToolBehavior != null
            ? { behavior: googleOptions.defaultToolBehavior }
            : {}),
        })),
      },
    ];
  }

  if (config?.inputAudioTranscription != null) {
    setup.inputAudioTranscription = {};
  }

  if (config?.outputAudioTranscription != null) {
    setup.outputAudioTranscription = {};
  }

  // Background-reasoning models require a `thinkingLevel` or `thinkingBudget`.
  // Default to the lowest-latency level when the effective thinking config (the
  // typed option, else a raw `providerOptions.generationConfig.thinkingConfig`)
  // sets neither. Merged last so it survives a raw
  // `providerOptions.generationConfig`.
  const rawGenerationConfig = restProviderOptions.generationConfig;
  const explicitThinkingConfig =
    googleOptions?.thinkingConfig ??
    (isRecord(rawGenerationConfig) &&
    isRecord(rawGenerationConfig.thinkingConfig)
      ? rawGenerationConfig.thinkingConfig
      : undefined);
  const thinkingConfig =
    isThinkingLiveModel(modelId) &&
    explicitThinkingConfig?.thinkingLevel == null &&
    explicitThinkingConfig?.thinkingBudget == null
      ? { ...explicitThinkingConfig, thinkingLevel: 'low' as const }
      : googleOptions?.thinkingConfig;
  const applyThinkingConfig = () => {
    if (thinkingConfig == null) return;
    const target = isRecord(setup.generationConfig)
      ? setup.generationConfig
      : generationConfig;
    setup.generationConfig = { ...target, thinkingConfig };
  };

  if (config?.providerOptions == null) {
    applyThinkingConfig();
    return setup;
  }

  Object.assign(setup, restProviderOptions);

  if (googleOptions?.translationConfig != null) {
    const target = isRecord(setup.generationConfig)
      ? setup.generationConfig
      : generationConfig;
    setup.generationConfig = {
      ...target,
      translationConfig: googleOptions.translationConfig,
    };
  }

  applyThinkingConfig();
  return setup;
}