feat(telemetry): one session id and one trace per turn (RPT-057)

src/index.ts

@@ -31,6 +31,7 @@ import {
 	initTelemetry,
 	wrapToolWithOtel,
 	startSessionSpan,
+	startTurnTrace,
 	recordGenAiCall,
 	shutdownTelemetry,
 } from "./telemetry.js";
@@ -52,6 +53,7 @@ interface ParsedArgs {
 	repo?: string;
 	pat?: string;
 	session?: string;
+	sessionId?: string;
 	voice?: string;
 }
 
@@ -67,6 +69,7 @@ function parseArgs(argv: string[]): ParsedArgs {
 	let repo: string | undefined;
 	let pat: string | undefined;
 	let session: string | undefined;
+	let sessionId: string | undefined;
 	let voice: string | undefined;
 
 	for (let i = 0; i < args.length; i++) {
@@ -107,6 +110,11 @@ function parseArgs(argv: string[]): ParsedArgs {
 			case "--session":
 				session = args[++i];
 				break;
+			// Distinct from --session (a git branch for repo/sandbox mode): this is
+			// the id carried on model requests so a gateway can group the run.
+			case "--session-id":
+				sessionId = args[++i];
+				break;
 			case "--voice":
 			case "-v":
 				// Accept optional backend name: --voice, --voice openai, --voice gemini
@@ -124,7 +132,7 @@ function parseArgs(argv: string[]): ParsedArgs {
 		}
 	}
 
-	return { model, dir, prompt, env, sandbox, sandboxRepo, sandboxToken, repo, pat, session, voice };
+	return { model, dir, prompt, env, sandbox, sandboxRepo, sandboxToken, repo, pat, session, sessionId, voice };
 }
 
 function handleEvent(
@@ -321,7 +329,7 @@ async function main(): Promise<void> {
 		return;
 	}
 
-	const { model, dir: rawDir, prompt, env, sandbox: useSandbox, sandboxRepo, sandboxToken, repo, pat, session: sessionBranch, voice } = parseArgs(process.argv);
+	const { model, dir: rawDir, prompt, env, sandbox: useSandbox, sandboxRepo, sandboxToken, repo, pat, session: sessionBranch, sessionId: sessionIdFlag, voice } = parseArgs(process.argv);
 
 	// If --repo is given, derive a default dir from the repo URL (skip interactive prompt)
 	let dir = rawDir;
@@ -465,7 +473,7 @@ async function main(): Promise<void> {
 
 	let loaded;
 	try {
-		loaded = await loadAgent(dir, model, env);
+		loaded = await loadAgent(dir, model, env, sessionIdFlag);
 	} catch (err: any) {
 		console.error(red(`Error: ${err.message}`));
 		process.exit(1);
@@ -643,6 +651,7 @@ async function main(): Promise<void> {
 	// Single-shot mode
 	if (prompt) {
 		try {
+			startTurnTrace(loaded.model);
 			await otelContext.with(_session.ctx, () => agent.prompt(prompt));
 		} catch (err: any) {
 			auditLogger?.logError(err.message).catch(() => {});
@@ -804,6 +813,7 @@ async function main(): Promise<void> {
 			}
 
 			try {
+				startTurnTrace(loaded.model);
 				await otelContext.with(_session.ctx, () => agent.prompt(promptText));
 			} catch (err: any) {
 				console.error(red(`Error: ${err.message}`));

src/loader.ts

@@ -117,8 +117,10 @@ async function ensureGitagentDir(agentDir: string): Promise<string> {
 	return gitagentDir;
 }
 
-async function writeSessionState(gitagentDir: string): Promise<string> {
-	const sessionId = randomUUID();
+async function writeSessionState(gitagentDir: string, override?: string): Promise<string> {
+	// A caller-supplied id wins so an embedding host (Studio, a web UI, a test)
+	// can tie this run to a session it already knows about.
+	const sessionId = override || randomUUID();
 	const state = {
 		session_id: sessionId,
 		started_at: new Date().toISOString(),
@@ -239,6 +241,7 @@ export async function loadAgent(
 	agentDir: string,
 	modelFlag?: string,
 	envFlag?: string,
+	sessionIdOverride?: string,
 ): Promise<LoadedAgent> {
 	// Parse agent.yaml
 	const manifestRaw = await readFile(join(agentDir, "agent.yaml"), "utf-8");
@@ -249,7 +252,7 @@ export async function loadAgent(
 
 	// Ensure .gitagent/ directory and write session state
 	const gitagentDir = await ensureGitagentDir(agentDir);
-	const sessionId = await writeSessionState(gitagentDir);
+	const sessionId = await writeSessionState(gitagentDir, sessionIdOverride);
 
 	// Resolve inheritance (Phase 2.4)
 	let parentRules = "";
@@ -405,6 +408,19 @@ Do NOT track trivial single-command tasks (e.g. "what time is it"). But DO check
 		model = getModel(provider as any, modelId as any);
 	}
 
+	// One run is many model requests: every turn of the agent loop, plus the
+	// off-loop reflection, repair and compaction calls. A gateway that groups
+	// telemetry per request sees each of those as a separate session unless the
+	// client says otherwise, so carry this run's id on every request.
+	//
+	// Cloned rather than mutated — getModel() returns a shared registry object,
+	// and writing to it would leak this run's id into every other model built in
+	// the same process.
+	model = {
+		...model,
+		headers: { ...(model as any).headers, "X-Session-Id": sessionId },
+	};
+
 	// For custom providers not in pi-ai's env key map, ensure an API key is available.
 	// pi-ai calls getEnvApiKey(model.provider) which only knows built-in providers.
 	// For unknown providers using openai-completions API, set provider to "openai" so

src/sdk.ts

@@ -29,6 +29,7 @@ import { context as otelContext } from "@opentelemetry/api";
 import {
 	wrapToolWithOtel,
 	startSessionSpan,
+	startTurnTrace,
 	recordGenAiCall,
 } from "./telemetry.js";
 
@@ -151,7 +152,10 @@ export function query(options: QueryOptions): Query {
 		}
 
 		// 1. Load agent
-		const loaded = await loadAgent(dir, options.model, options.env);
+		// options.sessionId, when given, becomes the agent's session id — so a host
+		// that already tracks a conversation sees its own id on the model requests
+		// rather than a fresh one per run.
+		const loaded = await loadAgent(dir, options.model, options.env, options.sessionId);
 		_manifest = loaded.manifest;
 		_sessionId = _sessionId || loaded.sessionId;
 
@@ -515,6 +519,7 @@ export function query(options: QueryOptions): Query {
 					return;
 				}
 			}
+			startTurnTrace(loaded.model);
 			await otelContext.with(_session.ctx, () =>
 				agent.prompt(options.prompt as string),
 			);
@@ -539,6 +544,7 @@ export function query(options: QueryOptions): Query {
 						return;
 					}
 				}
+				startTurnTrace(loaded.model);
 				await otelContext.with(_session.ctx, () =>
 					agent.prompt(userMsg.content),
 				);

src/telemetry.ts

@@ -24,6 +24,7 @@ import type {
 	Counter,
 } from "@opentelemetry/api";
 import type { AgentTool } from "@mariozechner/pi-agent-core";
+import { randomBytes } from "crypto";
 
 // ── Public types ───────────────────────────────────────────────────────
 
@@ -184,6 +185,33 @@ export function isTelemetryEnabled(): boolean {
 	return _initialized;
 }
 
+// ── Turn-scoped trace propagation ──────────────────────────────────────
+
+/**
+ * Start a new W3C trace for one user turn.
+ *
+ * A single user message costs several HTTP calls to the model gateway — one
+ * that comes back with a tool call, another with the answer, and so on. Each
+ * call is a separate request, so a gateway that traces per request records one
+ * trace per call and the turn arrives split across several of them. Sending the
+ * same `traceparent` on every call of the turn lets the gateway stitch them
+ * into one trace.
+ *
+ * No-op once telemetry is initialised: the undici instrumentation already
+ * injects `traceparent` from the active span, and a header written here would
+ * fight it.
+ */
+export function startTurnTrace(model: unknown): void {
+	try {
+		if (_initialized || !model) return;
+		const m = model as { headers?: Record<string, string> };
+		const traceparent = `00-${randomBytes(16).toString("hex")}-${randomBytes(8).toString("hex")}-01`;
+		m.headers = { ...(m.headers ?? {}), traceparent };
+	} catch {
+		// Telemetry must never break a run.
+	}
+}
+
 // ── Tracer / meter accessors ───────────────────────────────────────────
 
 export function getTracer(): Tracer {