pyric 0.1.0-alpha.21 → 0.1.0-alpha.24

package/src/ai/activity.ts

@@ -0,0 +1,43 @@
+import { sdkActivity } from '../sandbox/internal/sdk-activity.js';
+import { getAiEvidence, type AiEvidence } from '../sandbox/internal/ai-evidence.js';
+import { aiCompletionUsage } from '../sandbox/internal/usage-evidence.js';
+import type { AITarget } from './types.js';
+
+/** One SDK request, regardless of streamed chunk count or chat entry point. */
+export function beginAiActivity(target: AITarget, model: string, method: string) {
+  const start = performance.now();
+  let detail: AiEvidence = target.kind === 'sandbox' ? target.broker.observationIdentity(model)
+    : { requestedModel: model, engine: 'unknown', usageSource: 'unknown' };
+  const activity = sdkActivity.begin({ app: target, kind: 'operation', method,
+    source: { service: 'ai', target: model, key: model } });
+  activity.ai(detail);
+  const observe = (value: object) => {
+    detail = { ...detail, ...getAiEvidence(value) };
+    activity.ai(detail);
+  };
+  let ended = false;
+  return {
+    chunk(value: object) {
+      observe(value);
+      detail = { ...detail, firstChunkMs: detail.firstChunkMs ?? performance.now() - start };
+    },
+    progress: () => activity.progress(),
+    finish(value: object, response: object = value) {
+      if (ended) return;
+      ended = true;
+      observe(value);
+      activity.response(response);
+      detail = { ...detail, durationMs: performance.now() - start };
+      activity.ai(detail);
+      activity.delivered(undefined, aiCompletionUsage(detail, method));
+      activity.complete();
+    },
+    fail(error?: unknown) {
+      if (ended) return;
+      ended = true;
+      if (error && typeof error === 'object') observe(error);
+      activity.ai({ ...detail, durationMs: performance.now() - start });
+      activity.fail({ aiFailures: 1, ...(method === 'countTokens' ? {} : { aiUnknownUsage: 1 }) });
+    },
+  };
+}

package/src/ai/broker/broker.ts

@@ -1,3 +1,4 @@
+import { aiEndpoint, getAiEvidence, setAiEvidence, type AiEvidence } from '../../sandbox/internal/ai-evidence.js';
 /**
  * The sandbox AI broker — the ONE in-process Gemini-wire model every plane
  * consumes (house pattern: the messaging broker):
@@ -37,13 +38,14 @@
  */
 
 import type { Sandbox } from 'pyric/sandbox';
-import { emitSandboxEvent, makeServiceMutationEvent } from 'pyric/sandbox/internal';
+import { emitSandboxEvent, getClock, makeServiceMutationEvent } from 'pyric/sandbox/internal';
+import type { AiEventOperation } from '../events.js';
 
 import { describeResponseBlock } from '../blocked.js';
 import { AiBrokerError, Synthesizer, badRole, emptyContents, missingThoughtSignature } from './synthesizer.js';
 import { ScriptedEngine } from './scripted-engine.js';
 import { OpenAiEngine } from './openai-engine.js';
-import { GeminiEngine } from './gemini-engine.js';
+import { GeminiEngine, GEMINI_DEFAULT_BASE_URL } from './gemini-engine.js';
 import type {
   AnswerEngine,
   CountTokensRequest,
@@ -91,6 +93,19 @@ function isAnswerEngine(value: EngineConfig | AnswerEngine): value is AnswerEngi
 /** What Studio's stream (and the construction log line) name the engine. */
 type EngineKind = 'scripted' | 'openai' | 'gemini' | 'custom';
 
+/**
+ * What a status read reports about the resolved engine: the mode it is
+ * running in, the model and upstream it answers through when it has one, and
+ * whether a key is configured. Never the key itself, its prefix, or its
+ * length.
+ */
+export interface AiEngineStatus {
+  engine: EngineKind;
+  model?: string;
+  upstream?: string;
+  keyPresent: boolean;
+}
+
 export class AiBroker {
   readonly engine: AnswerEngine;
   private readonly sandbox: Sandbox | undefined;
@@ -120,6 +135,24 @@ export class AiBroker {
     console.info(this.describeEngine());
   }
 
+  /** Snapshot routing before execution; never treat a scripted modelVersion as identity. */
+  observationIdentity(model: string): AiEvidence {
+    let routed: ModelResolution | undefined;
+    try { routed = this.engine.resolveEffectiveModel?.(model); } catch { /* Diagnostic only. */ }
+    return { requestedModel: model, engine: this.engineKind,
+      ...(routed ? { routedModel: routed.model, mappingReason: routed.reason } : {}),
+      endpoint: aiEndpoint(this.engineBaseUrl ?? (this.engineKind === 'gemini' ? GEMINI_DEFAULT_BASE_URL : undefined)),
+      usageSource: this.engineKind === 'scripted' ? 'scripted' : 'unknown' };
+  }
+  private annotate<T extends WireResponse>(response: T, identity: AiEvidence): T {
+    const usage = response.usageMetadata;
+    const original = getAiEvidence(response);
+    const valid = (n: unknown) => typeof n === 'number' && Number.isFinite(n) && n >= 0 ? n : undefined;
+    return setAiEvidence(response, { ...identity,
+      ...(identity.engine === 'gemini' ? { reportedModel: response.modelVersion, usageSource: usage ? 'backend' as const : 'unknown' as const } : {}),
+      ...original, inputTokens: valid(usage?.promptTokenCount), outputTokens: valid(usage?.candidatesTokenCount), totalTokens: valid(usage?.totalTokenCount) });
+  }
+
   /** One-line construction-time summary of what this broker resolved to. */
   private describeEngine(): string {
     const isOpenAiEngine = this.engineKind === 'openai';
@@ -137,6 +170,21 @@ export class AiBroker {
     return '[pyric/ai] engine resolved: scripted (zero-config unless a script is queued)';
   }
 
+  /**
+   * The engine, model, upstream, and whether a key is configured, for a
+   * status read. Never the key itself: only {@link GeminiEngine.hasKeyConfigured}'s
+   * boolean crosses this seam.
+   */
+  describeStatus(): AiEngineStatus {
+    const keyPresent = this.engine instanceof GeminiEngine ? this.engine.hasKeyConfigured() : false;
+    const status: AiEngineStatus = { engine: this.engineKind, keyPresent };
+    if (this.engineModel !== undefined) status.model = this.engineModel;
+    const upstream =
+      this.engineKind === 'gemini' ? (this.engineBaseUrl ?? GEMINI_DEFAULT_BASE_URL) : this.engineBaseUrl;
+    if (upstream !== undefined) status.upstream = upstream;
+    return status;
+  }
+
   /** Additive detail fields every emitted event carries so Studio can show the engine. */
   private engineDetail(): Record<string, unknown> {
     const detail: Record<string, unknown> = {
@@ -226,6 +274,7 @@ export class AiBroker {
   }
 
   async generateContent(req: GenerateContentRequest, model: string): Promise<WireResponse> {
+    const identity = this.observationIdentity(model);
     this.validate(req, model);
     this.emitModelSubstitution(model);
     try {
@@ -237,7 +286,7 @@ export class AiBroker {
         responseId: response.responseId,
       });
       this.emitBlockIfBlocked(model, response);
-      return response;
+      return this.annotate(response, identity);
     } catch (err) {
       this.emitRejectionIfBrokerError(model, err);
       throw err;
@@ -251,9 +300,11 @@ export class AiBroker {
    * an HTTP error instead of a stream.
    */
   streamGenerateContent(req: GenerateContentRequest, model: string): AsyncIterable<WireChunk> {
+    const identity = this.observationIdentity(model);
     this.validate(req, model);
     this.emitModelSubstitution(model);
     const inner = this.engine.streamGenerateContent(req, model);
+    const annotate = (chunk: WireChunk) => this.annotate(chunk, identity);
     const emit = (chunkCount: number) =>
       this.emit('stream_generate_content', model, {
         contentCount: req.contents.length,
@@ -275,7 +326,7 @@ export class AiBroker {
         for await (const chunk of inner) {
           chunkCount++;
           emitBlock(chunk);
-          yield chunk;
+          yield annotate(chunk);
         }
         emit(chunkCount);
       } catch (err) {
@@ -286,11 +337,12 @@ export class AiBroker {
   }
 
   async countTokens(req: CountTokensRequest, model: string): Promise<CountTokensResponse> {
+    const identity = this.observationIdentity(model);
     this.validateContents(req, model);
     try {
       const response = await this.engine.countTokens(req, model);
       this.emit('count_tokens', model, { totalTokens: response.totalTokens });
-      return response;
+      return setAiEvidence(response, { ...identity, totalTokens: response.totalTokens, usageSource: this.engineKind === 'gemini' ? 'backend' : this.engineKind === 'openai' ? 'estimated' : this.engineKind === 'scripted' ? 'scripted' : 'unknown' });
     } catch (err) {
       this.emitRejectionIfBrokerError(model, err);
       throw err;
@@ -345,12 +397,13 @@ export class AiBroker {
    * Best-effort, storage-precedent: a throw from the emit path must never
    * fail the AI operation the caller just completed.
    */
-  private emit(op: string, model: string, detail: Record<string, unknown>): void {
+  private emit(op: AiEventOperation, model: string, detail: Record<string, unknown>): void {
     if (this.sandbox === undefined) return;
     try {
       emitSandboxEvent(
         this.sandbox,
         makeServiceMutationEvent({
+          at: getClock(this.sandbox).now(),
           service: 'ai',
           op,
           path: model,

package/src/ai/broker/gemini-engine.ts

@@ -67,6 +67,22 @@ export class GeminiEngine implements AnswerEngine {
     this.explicitKey = options?.apiKey;
   }
 
+  /**
+   * Whether a static key is configured, checking the explicit option and the
+   * same environment variables {@link resolveAuthToken} does, without
+   * resolving Application Default Credentials: a status read must never shell
+   * out, and ADC's own presence answers only at call time.
+   */
+  hasKeyConfigured(): boolean {
+    const candidates = [
+      this.explicitKey,
+      process.env.GEMINI_API_KEY,
+      process.env.GOOGLE_GENAI_API_KEY,
+      process.env.VITE_GEMINI_API_KEY,
+    ];
+    return candidates.some((key) => key !== undefined && key.trim() !== '');
+  }
+
   /**
    * Resolves authentication credentials sequentially, checking static options,
    * environment variables, and finally Google Cloud Application Default

package/src/ai/broker/index.ts

@@ -5,7 +5,7 @@
  * stage) owns the public surface.
  */
 
-export { AiBroker, loadObservationEnvelope, type AiBrokerOptions } from './broker.js';
+export { AiBroker, loadObservationEnvelope, type AiBrokerOptions, type AiEngineStatus } from './broker.js';
 export { ScriptedEngine, lastUserText, promptTextOf } from './scripted-engine.js';
 export { GeminiEngine, GEMINI_DEFAULT_BASE_URL, type GeminiEngineOptions } from './gemini-engine.js';
 export {

package/src/ai/broker/openai-engine.ts

@@ -1,3 +1,4 @@
+import { setAiEvidence } from '../../sandbox/internal/ai-evidence.js';
 /**
  * The OpenAI AnswerEngine: Gemini wire in, OpenAI-compatible upstream
  * (Ollama, llama.cpp, any /v1/chat/completions) out, Gemini wire back.
@@ -518,7 +519,7 @@ export class OpenAiEngine implements AnswerEngine {
     const res = await this.post(body);
     const upstream = (await res.json()) as OpenAIResponse;
     const translated = openAIToGeminiResponse(upstream);
-    return this.synth.decorate(translated, { model, promptText: promptTextOf(req) });
+    return setAiEvidence(this.synth.decorate(translated, { model, promptText: promptTextOf(req) }), { reportedModel: upstream.model, usageSource: upstream.usage ? 'backend' : 'estimated' });
   }
 
   streamGenerateContent(req: GenerateContentRequest, model: string): AsyncIterable<WireChunk> {
@@ -540,6 +541,7 @@ export class OpenAiEngine implements AnswerEngine {
       const decoder = new TextDecoder();
       let emittedText = '';
       let upstreamUsage: OpenAIStreamChunk['usage'];
+      let reportedModel: string | undefined;
       let pendingFinish: string | null = null;
 
       /**
@@ -563,7 +565,7 @@ export class OpenAiEngine implements AnswerEngine {
               candidatesTokenCount,
               totalTokenCount: promptTokenCount + candidatesTokenCount,
             };
-        return {
+        return setAiEvidence({
           candidates: [
             {
               content: { parts: parts.length ? parts : [{ text: '' }], role: 'model' },
@@ -578,7 +580,7 @@ export class OpenAiEngine implements AnswerEngine {
           },
           modelVersion,
           responseId,
-        };
+        }, { reportedModel, usageSource: upstreamUsage ? 'backend' : 'estimated' });
       }
 
       const reader = res.body.getReader();
@@ -599,6 +601,7 @@ export class OpenAiEngine implements AnswerEngine {
               continue; // non-JSON keepalive/comment frames
             }
             if (chunk.usage) upstreamUsage = chunk.usage;
+            if (chunk.model) reportedModel = chunk.model;
             const choice = chunk.choices?.[0];
             if (!choice) continue; // usage-only final frame folds into the finish chunk
 

package/src/ai/broker/scripted-engine.ts

@@ -11,7 +11,8 @@
  * directly) or a shorthand the synthesizer expands. Matcher semantics:
  *   - entries are consumed at most once, scanned in queue order;
  *   - a matcher (substring / regex / predicate) is checked against the last
- *     user turn's text (predicates get the whole request);
+ *     user turn's text (predicates get the whole request and the requested
+ *     model);
  *   - the FIRST unconsumed matching entry wins;
  *   - an entry without `match` is unconditional — next-in-queue.
  */
@@ -105,7 +106,7 @@ export class ScriptedEngine implements AnswerEngine {
 
   async generateContent(req: GenerateContentRequest, model: string): Promise<WireResponse> {
     const opts = this.opts(req, model);
-    const respond = this.take(req);
+    const respond = this.take(req, model);
     if (respond === undefined) {
       this.warnSyntheticDefaultOnce();
       return specialDefault(this.synth, req, opts) ?? this.synth.text(defaultText(req), opts);
@@ -116,7 +117,7 @@ export class ScriptedEngine implements AnswerEngine {
 
   streamGenerateContent(req: GenerateContentRequest, model: string): AsyncIterable<WireChunk> {
     const opts = this.opts(req, model);
-    const respond = this.take(req);
+    const respond = this.take(req, model);
     const synth = this.synth;
     if (respond === undefined) this.warnSyntheticDefaultOnce();
 
@@ -161,7 +162,7 @@ export class ScriptedEngine implements AnswerEngine {
   }
 
   /** First unconsumed matching entry wins; marks it consumed. */
-  private take(req: GenerateContentRequest): ScriptRespond | undefined {
+  private take(req: GenerateContentRequest, model: string): ScriptRespond | undefined {
     const text = lastUserText(req);
     for (const slot of this.queue) {
       if (slot.consumed) continue;
@@ -173,7 +174,7 @@ export class ScriptedEngine implements AnswerEngine {
             ? text.includes(m)
             : m instanceof RegExp
               ? m.test(text)
-              : m(req);
+              : m(req, model);
       if (hit) {
         slot.consumed = true;
         return slot.entry.respond;
@@ -182,6 +183,16 @@ export class ScriptedEngine implements AnswerEngine {
     return undefined;
   }
 
+  /** Empty the queue. Nothing already answered is affected. */
+  clear(): void {
+    this.queue.length = 0;
+  }
+
+  /** The queue as it stands, each entry beside whether it has been consumed. */
+  list(): ReadonlyArray<{ entry: ScriptEntry; consumed: boolean }> {
+    return this.queue.map((slot) => ({ entry: slot.entry, consumed: slot.consumed }));
+  }
+
   private expandUnary(respond: ScriptShorthand, opts: SynthesizeOptions): WireResponse {
     if ('error' in respond) throw new AiBrokerError({ error: { ...respond.error } });
     return expandNonError(this.synth, respond, opts);

package/src/ai/broker/types.ts

@@ -193,8 +193,12 @@ export type RawEnvelope = WireResponse &
 
 export type ScriptRespond = ScriptShorthand | RawEnvelope;
 
-/** Matcher: substring / regex against the last user turn text, or a predicate on the request. */
-export type ScriptMatcher = string | RegExp | ((req: GenerateContentRequest) => boolean);
+/**
+ * Matcher: substring / regex against the last user turn text, or a predicate
+ * on the request. The predicate also receives the requested model, so a
+ * matcher can select on the model a call names as well as its text.
+ */
+export type ScriptMatcher = string | RegExp | ((req: GenerateContentRequest, model?: string) => boolean);
 
 export interface ScriptEntry {
   /** Absent ⇒ unconditional next-in-queue. */

package/src/ai/events.ts

@@ -0,0 +1,29 @@
+/**
+ * What the AI surface puts on the sandbox event stream.
+ *
+ * The broker emits one event per answered call, per rejection, and per
+ * substitution. The scripted engine's pending entry queue is not here: a
+ * queued entry is a registration waiting to be used, not something that
+ * happened, and its matcher and responder are functions no event can carry.
+ */
+import type { ServiceEventRecord } from '../sandbox/types/service-event-record.js';
+
+export const AI_EVENT_RECORD = {
+  service: 'ai',
+  operations: [
+    'generate_content',
+    'stream_generate_content',
+    'count_tokens',
+    'request_rejected',
+    'response_blocked',
+    'model_substituted',
+  ],
+  target: {
+    name: 'model',
+    description: 'The model name the call named, e.g. gemini-2.5-flash.',
+    always: false,
+  },
+} as const satisfies ServiceEventRecord;
+
+/** Every operation the AI surface emits. */
+export type AiEventOperation = (typeof AI_EVENT_RECORD)['operations'][number];

package/src/ai/generation-enums.ts

@@ -0,0 +1,48 @@
+/** Runtime values from the installed Firebase SDK. */
+
+export const ImagenAspectRatio = {
+  SQUARE: "1:1",
+  LANDSCAPE_3x4: "3:4",
+  PORTRAIT_4x3: "4:3",
+  LANDSCAPE_16x9: "16:9",
+  PORTRAIT_9x16: "9:16",
+} as const;
+export type ImagenAspectRatio = (typeof ImagenAspectRatio)[keyof typeof ImagenAspectRatio];
+
+export const ImagenPersonFilterLevel = {
+  BLOCK_ALL: "dont_allow",
+  ALLOW_ADULT: "allow_adult",
+  ALLOW_ALL: "allow_all",
+} as const;
+export type ImagenPersonFilterLevel = (typeof ImagenPersonFilterLevel)[keyof typeof ImagenPersonFilterLevel];
+
+export const ImagenSafetyFilterLevel = {
+  BLOCK_LOW_AND_ABOVE: "block_low_and_above",
+  BLOCK_MEDIUM_AND_ABOVE: "block_medium_and_above",
+  BLOCK_ONLY_HIGH: "block_only_high",
+  BLOCK_NONE: "block_none",
+} as const;
+export type ImagenSafetyFilterLevel = (typeof ImagenSafetyFilterLevel)[keyof typeof ImagenSafetyFilterLevel];
+
+export const InferenceMode = {
+  PREFER_ON_DEVICE: "prefer_on_device",
+  ONLY_ON_DEVICE: "only_on_device",
+  ONLY_IN_CLOUD: "only_in_cloud",
+  PREFER_IN_CLOUD: "prefer_in_cloud",
+} as const;
+export type InferenceMode = (typeof InferenceMode)[keyof typeof InferenceMode];
+
+export const InferenceSource = {
+  ON_DEVICE: "on_device",
+  IN_CLOUD: "in_cloud",
+} as const;
+export type InferenceSource = (typeof InferenceSource)[keyof typeof InferenceSource];
+
+export const LiveResponseType = {
+  SERVER_CONTENT: "serverContent",
+  TOOL_CALL: "toolCall",
+  TOOL_CALL_CANCELLATION: "toolCallCancellation",
+  GOING_AWAY_NOTICE: "goingAwayNotice",
+  SESSION_RESUMPTION_UPDATE: "sessionResumptionUpdate",
+} as const;
+export type LiveResponseType = (typeof LiveResponseType)[keyof typeof LiveResponseType];

package/src/ai/index.ts

@@ -137,3 +137,12 @@ export type {
   AI,
   AIOptions,
 } from './types.js';
+
+export {
+  ImagenAspectRatio,
+  ImagenPersonFilterLevel,
+  ImagenSafetyFilterLevel,
+  InferenceMode,
+  InferenceSource,
+  LiveResponseType,
+} from './generation-enums.js';

package/src/ai/instances.ts

@@ -97,7 +97,7 @@ export function getAI(target?: Sandbox | FirebaseApp, options?: AIOptions): AI {
 
 function sandboxAI(sandbox: Sandbox, options?: AIOptions, app?: FirebaseApp): AI {
   const backend = options?.backend ?? new GoogleAIBackend();
-  const { key, location } = describeBackend(backend);
+  const { key } = describeBackend(backend);
   const handles = cachedHandles(app ?? sandbox);
   const existing = handles.get(key);
   if (existing) {
@@ -113,6 +113,27 @@ function sandboxAI(sandbox: Sandbox, options?: AIOptions, app?: FirebaseApp): AI
     brokersBySandbox.set(sandbox, broker);
     brokerEngineBySandbox.set(sandbox, options?.engine);
   }
+  const handle = sandboxHandle(sandbox, broker, options, app);
+  handles.set(key, handle);
+  return handle;
+}
+
+/**
+ * Host-only configuration lifecycle: a new handle binds a new broker to the
+ * SAME sandbox. Existing handles and in-flight requests keep their broker.
+ * Does not alter getAI's per-owner cache or first-configuration-wins contract.
+ */
+export function createConfiguredSandboxAI(sandbox: Sandbox, options: AIOptions): AI {
+  const supplied = options.engine;
+  const engine = supplied && !isCustomEngine(supplied) && supplied.kind === 'openai'
+    ? { ...supplied, modelMap: { ...supplied.modelMap } }
+    : supplied;
+  return sandboxHandle(sandbox, new AiBroker({ sandbox, engine }), { ...options, engine });
+}
+
+function sandboxHandle(sandbox: Sandbox, broker: AiBroker, options?: AIOptions, app?: FirebaseApp): AI {
+  const backend = options?.backend ?? new GoogleAIBackend();
+  const { location } = describeBackend(backend);
   const handle: AI = {
     ...(app !== undefined ? { app } : {}),
     backend,
@@ -131,6 +152,5 @@ function sandboxAI(sandbox: Sandbox, options?: AIOptions, app?: FirebaseApp): AI
       : {}),
   };
   Object.defineProperty(handle, TARGET_SYMBOL, { value: target, enumerable: false });
-  handles.set(key, handle);
   return handle;
 }

package/src/ai/internal.ts

@@ -60,3 +60,8 @@ export function createTransportAI(
   handles.set(key, handle);
   return handle;
 }
+
+export { packAiEvidence, unpackAiEvidence, getAiEvidence, setAiEvidence, type AiEvidence } from '../sandbox/internal/ai-evidence.js';
+
+/** Explicit host configuration revisions sharing one sandbox event/state owner. */
+export { createConfiguredSandboxAI } from './instances.js';

package/src/ai/sandbox-plane.ts

@@ -1,3 +1,4 @@
+import { beginAiActivity } from './activity.js';
 /**
  * The sandbox model plane for `pyric/ai` — everything between the public
  * `GenerativeModel`/`ChatSession` classes and the {@link AiBroker}:
@@ -110,15 +111,18 @@ export async function planeGenerateContent(
   request: Record<string, unknown>,
   singleRequestOptions?: SingleRequestOptions,
 ): Promise<GenerateContentResult> {
-  target.assertAlive?.();
-  checkSignal(singleRequestOptions);
+  const activity = beginAiActivity(target, modelResource, 'generateContent');
   try {
+    target.assertAlive?.();
+    checkSignal(singleRequestOptions);
     validateModel(modelResource);
     const normalized = normalizeRequest(request) as unknown as GenerateContentRequest;
     const response = await (target.kind === 'transport' ? target.transport : target.broker)
       .generateContent(normalized, modelResource);
+    activity.finish(response);
     return { response: createEnhancedContentResponse(response) };
   } catch (err) {
+    activity.fail(err);
     throw toAIError(err, modelResource, 'generateContent');
   }
 }
@@ -129,10 +133,11 @@ export async function planeGenerateContentStream(
   request: Record<string, unknown>,
   singleRequestOptions?: SingleRequestOptions,
 ): Promise<GenerateContentStreamResult> {
-  target.assertAlive?.();
-  checkSignal(singleRequestOptions);
+  const activity = beginAiActivity(target, modelResource, 'generateContentStream');
   let inner: AsyncIterable<WireResponse>;
   try {
+    target.assertAlive?.();
+    checkSignal(singleRequestOptions);
     validateModel(modelResource);
     const normalized = normalizeRequest(request) as unknown as GenerateContentRequest;
     // Broker validation runs eagerly here — a bad request throws before
@@ -140,6 +145,7 @@ export async function planeGenerateContentStream(
     inner = (target.kind === 'transport' ? target.transport : target.broker)
       .streamGenerateContent(normalized, modelResource);
   } catch (err) {
+    activity.fail(err);
     throw toAIError(err, modelResource, 'streamGenerateContent');
   }
 
@@ -157,10 +163,12 @@ export async function planeGenerateContentStream(
   const pumped = (async () => {
     try {
       for await (const chunk of inner) {
+        activity.chunk(chunk);
         buffered.push(chunk);
         wake();
       }
     } catch (err) {
+      activity.fail(err);
       failure = toAIError(err, modelResource, 'streamGenerateContent');
       wake();
       throw failure;
@@ -168,7 +176,9 @@ export async function planeGenerateContentStream(
       done = true;
       wake();
     }
-    return aggregateResponses(buffered);
+    const aggregate = aggregateResponses(buffered);
+    activity.finish(buffered.at(-1) ?? {}, aggregate);
+    return aggregate;
   })();
 
   const response: Promise<EnhancedResponse> = pumped.then((aggregate) =>
@@ -184,6 +194,7 @@ export async function planeGenerateContentStream(
     for (;;) {
       if (next < buffered.length) {
         // Chunks are broker-owned fresh objects; enhance in place.
+        activity.progress();
         yield createEnhancedContentResponse(structuredClone(buffered[next++]!));
         continue;
       }
@@ -204,16 +215,20 @@ export async function planeCountTokens(
   request: Record<string, unknown>,
   singleRequestOptions?: SingleRequestOptions,
 ): Promise<CountTokensResponse> {
-  target.assertAlive?.();
-  checkSignal(singleRequestOptions);
+  const activity = beginAiActivity(target, modelResource, 'countTokens');
   try {
+    target.assertAlive?.();
+    checkSignal(singleRequestOptions);
     validateModel(modelResource);
     const dispatch = target.kind === 'transport' ? target.transport : target.broker;
     const normalized = normalizeRequest(request) as unknown as Parameters<
       typeof dispatch.countTokens
     >[0];
-    return await dispatch.countTokens(normalized, modelResource);
+    const result = await dispatch.countTokens(normalized, modelResource);
+    activity.finish(result);
+    return result;
   } catch (err) {
+    activity.fail(err);
     throw toAIError(err, modelResource, 'countTokens');
   }
 }