feat(pi): trace model requests

docs/content/docs/observability.mdx

@@ -25,11 +25,13 @@ agent/prompt                                 the caller's trace
 pi.durable.submit                            the run's trace, linked to agent/prompt
 ├── pi.durable.phase  pi.generation prepare
 ├── pi.durable.phase  pi.generation request
+│   ├── chat claude-opus-5-5                 the model request
 │   └── pi.durable.phase  pi.tool call
 │       └── app.lookup                       your span
 └── pi.durable.phase  pi.generation tools
     ├── pi.durable.phase  pi.generation prepare
     └── pi.durable.phase  pi.generation request
+        └── chat claude-opus-5-5
 ```
 
 - **Each run is its own trace**, linked to its caller. RivetKit treats scheduled actions and workflow runs the same way. A `conversation.submit` run has the same trace as a `prompt` run, without `invoke_agent pi`.
@@ -47,6 +49,7 @@ Every span Pi records has `rivet.actor.id`, the Actor's id.
 | `invoke_agent pi` | A `prompt`, from saving its input until its run settles. A child of the `prompt` action span. | `gen_ai.operation.name`, `gen_ai.agent.name`, `gen_ai.conversation.id`, `gen_ai.request.model` |
 | `pi.durable.submit` | Pi admits a submission. | `submission.type`, `submission.id`, `conversation.id`, `result` |
 | `pi.durable.phase` | One step of a task, such as a model request or a tool call. | `task.id`, `task.name`, `phase`, `kind`, `attempt`, `resumed`, `outcome`, `tool.name`, `tool.call_id` |
+| `chat <model>` | One model call. Each retry and each poll of a deferred answer is its own call. | `gen_ai.operation.name`, `gen_ai.provider.name`, `gen_ai.request.model`, `gen_ai.usage.input_tokens` (with cached tokens), `gen_ai.usage.output_tokens`, `gen_ai.usage.cost_usd` |
 
 ## Failures
 
@@ -57,6 +60,7 @@ A failed span has an error status and `error.type`:
 | `unanswered` | `invoke_agent pi` | The run ended without an answer, such as after a model error. |
 | `run_error` | `invoke_agent pi` | The `prompt` threw. |
 | `task_faulted`, `task_failed`, `task_orphaned` | `pi.durable.phase` | A task stopped in this phase, such as a tool that threw. |
+| `provider_error` | `chat <model>` | The model call failed. Pi may retry it. |
 
 Tracing can't break a run. A collector that is down or a tracer that throws loses spans, and the run still answers.
 

packages/pi/src/models.ts

@@ -5,6 +5,7 @@ import {
 import { UserError } from "rivetkit";
 import type { PiCredentialSource } from "./credentials.js";
 import type { PiContext } from "./runtime.js";
+import { traceModelRequests } from "./tracing.js";
 
 /** Pi's credential store interface. `read`, `list`, `modify`, and `delete`, keyed by provider id. */
 export type PiCredentialStore = NonNullable<
@@ -62,7 +63,7 @@ export async function createActorModelRuntime(
 		await runtime.setRuntimeApiKey(providerId, apiKey);
 	}
 	await runtime.refresh({ allowNetwork: false });
-	return runtime;
+	return traceModelRequests(runtime);
 }
 
 /** A credential store that holds nothing and refuses writes. */

packages/pi/src/tracing.ts

@@ -1,5 +1,12 @@
 import type { Context } from "@earendil-works/chord";
 import { withTelemetry } from "@earendil-works/chord/context";
+import type {
+	Api,
+	AssistantMessage,
+	AssistantMessageEventStream,
+	Model,
+} from "@earendil-works/pi-ai";
+import type { ModelRuntime } from "@earendil-works/pi-coding-agent";
 import type {
 	SpanAttributes,
 	SpanLink,
@@ -51,6 +58,37 @@ function runActive<T>(active: OtelContext, body: () => T): T {
 	}
 }
 
+/** The spans Pi's model layer emits. Users see each as `chat <model>`. */
+const MODEL_SPANS: ReadonlySet<string> = new Set([
+	"pi.ai.request",
+	"pi.ai.poll",
+]);
+
+/**
+ * Pi's model span attributes under the names users build dashboards on, from
+ * OpenTelemetry's GenAI conventions. Input tokens include cached ones, as those
+ * conventions ask. A deferred or pending response has no usage to report yet.
+ */
+function modelAttributes(attributes: SpanAttributes = {}): SpanAttributes {
+	const hasUsage =
+		attributes.stop !== "deferred" && attributes.stop !== "pending";
+	const input = attributes["usage.input"];
+	return {
+		"gen_ai.provider.name": attributes.provider,
+		"gen_ai.request.model": attributes.model,
+		...(hasUsage && {
+			"gen_ai.usage.input_tokens":
+				typeof input === "number"
+					? input +
+						Number(attributes["usage.cache_read"] ?? 0) +
+						Number(attributes["usage.cache_write"] ?? 0)
+					: undefined,
+			"gen_ai.usage.output_tokens": attributes["usage.output"],
+			"gen_ai.usage.cost_usd": attributes["cost.total"],
+		}),
+	};
+}
+
 function attributesOf(attributes: SpanAttributes | undefined): Attributes {
 	const defined: Attributes = {};
 	for (const [name, value] of Object.entries(attributes ?? {})) {
@@ -91,23 +129,31 @@ class OtelTelemetry implements TelemetryContext {
 		callback: (span: TelemetrySpan) => T | Promise<T>,
 	): Promise<T> {
 		const { adapter, parent } = this;
+		const isModelSpan = MODEL_SPANS.has(options.name);
 		const native = this.settled
 			? undefined
 			: guarded(() =>
 					adapter.tracer().startSpan(
-						options.name,
+						isModelSpan ? `chat ${options.attributes?.model}` : options.name,
 						{
 							attributes: {
 								...adapter.attributes,
-								...attributesOf(options.attributes),
+								...attributesOf(
+									isModelSpan
+										? {
+												"gen_ai.operation.name": "chat",
+												...modelAttributes(options.attributes),
+											}
+										: options.attributes,
+								),
 							},
 							links: options.links?.flatMap(linksOf),
 						},
 						// Never the active context, so a background phase does not nest under an unrelated action.
 						parent,
 					),
 				);
-		const span = new OtelSpan(adapter, native);
+		const span = new OtelSpan(adapter, native, isModelSpan);
 		try {
 			const result = await runActive(
 				native ? trace.setSpan(parent, native) : parent,
@@ -135,6 +181,7 @@ class OtelSpan extends OtelTelemetry implements TelemetrySpan {
 	constructor(
 		adapter: OtelAdapter,
 		readonly native: Span | undefined,
+		readonly isModelSpan = false,
 	) {
 		super(adapter, native ? trace.setSpan(ROOT_CONTEXT, native) : ROOT_CONTEXT);
 	}
@@ -159,7 +206,10 @@ class OtelSpan extends OtelTelemetry implements TelemetrySpan {
 	}
 
 	setAttributes(attributes: SpanAttributes): void {
-		this.#record((native) => native.setAttributes(attributesOf(attributes)));
+		const recorded = this.isModelSpan
+			? modelAttributes(attributes)
+			: attributes;
+		this.#record((native) => native.setAttributes(attributesOf(recorded)));
 	}
 
 	setStatus(status: SpanStatus): void {
@@ -323,3 +373,86 @@ export async function traceRun<T extends { status: string; reason?: string }>(
 		guarded(() => span?.end());
 	}
 }
+
+const MODEL_CALL_FAILED: SpanStatus = {
+	status: "error",
+	error: { name: "provider_error", message: "" },
+};
+
+/** Records what a finished model call reports, never its content or its error text. */
+function recordModelResult(span: TelemetrySpan, message: AssistantMessage) {
+	const { input, output, cacheRead, cacheWrite, cost } = message.usage;
+	span.setAttributes({
+		stop: message.stopReason,
+		"usage.input": input,
+		"usage.output": output,
+		"usage.cache_read": cacheRead,
+		"usage.cache_write": cacheWrite,
+		"cost.total": cost.total,
+	});
+	if (message.stopReason === "error") span.setStatus(MODEL_CALL_FAILED);
+}
+
+/** Runs one model call in a span of Pi's model vocabulary, ending when its stream does. */
+function traceModelCall(
+	spanName: "pi.ai.request" | "pi.ai.poll",
+	model: Model<Api>,
+	telemetry: TelemetrySpan | undefined,
+	call: () => AssistantMessageEventStream,
+): AssistantMessageEventStream {
+	if (!telemetry) return call();
+	let stream: AssistantMessageEventStream | undefined;
+	let thrown: unknown;
+	void telemetry.startSpan(
+		{
+			name: spanName,
+			attributes: { provider: model.provider, model: model.id },
+		},
+		async (span) => {
+			try {
+				stream = call();
+				recordModelResult(span, await stream.result());
+			} catch (error) {
+				thrown = error;
+				span.setStatus(MODEL_CALL_FAILED);
+			}
+		},
+	);
+	if (!stream) throw thrown;
+	return stream;
+}
+
+/**
+ * The default model runtime with a span per model call. `ModelRuntime` skips
+ * the `observeModelRequest` that pi-ai's own `Models` runs, and pi-ai does not
+ * export it, so this emits the same spans. An app's own `models` stays as it is,
+ * so a call never has two spans. tests/tracing-models.test.ts fails when Pi
+ * fixes `ModelRuntime`. Remove this wrapper then: users see the same spans.
+ *
+ * A Proxy and not patched methods: `completeSimple` calls `streamSimple`
+ * inside `ModelRuntime`, so patching both would give that call two spans.
+ */
+export function traceModelRequests(runtime: ModelRuntime): ModelRuntime {
+	const traced: Partial<Record<string | symbol, unknown>> = {
+		streamSimple: ((model, context, options) =>
+			traceModelCall("pi.ai.request", model, options?.telemetry, () =>
+				runtime.streamSimple(model, context, options),
+			)) satisfies ModelRuntime["streamSimple"],
+		completeSimple: ((model, context, options) =>
+			traceModelCall("pi.ai.request", model, options?.telemetry, () =>
+				runtime.streamSimple(model, context, options),
+			).result()) satisfies ModelRuntime["completeSimple"],
+		fetchDeferred: ((model, handle, options) =>
+			traceModelCall("pi.ai.poll", model, options?.telemetry, () =>
+				runtime.streamDeferred(model, handle, options),
+			).result()) satisfies ModelRuntime["fetchDeferred"],
+	};
+	return new Proxy(runtime, {
+		get: (target, name) => {
+			const method = traced[name];
+			if (method) return method;
+			const value: unknown = Reflect.get(target, name, target);
+			return typeof value === "function" ? value.bind(target) : value;
+		},
+	});
+}