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

README.md

@@ -132,6 +132,7 @@ gitagent --repo https://github.com/org/repo "Add unit tests"
 | `--repo <url>` | `-r` | GitHub repo URL to clone and work on |
 | `--pat <token>` | | GitHub PAT (or set `GITHUB_TOKEN` / `GIT_TOKEN`) |
 | `--session <branch>` | | Resume an existing session branch |
+| `--session-id <id>` | | Session id sent on model requests, so a gateway groups the run (default: generated) |
 | `--model <provider:model>` | `-m` | Override model (e.g. `anthropic:claude-sonnet-4-5-20250929`) |
 | `--sandbox` | `-s` | Run in sandbox VM |
 | `--prompt <text>` | `-p` | Single-shot prompt (skip REPL) |

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(
@@ -304,6 +312,11 @@ async function ensureRepo(dir: string, model?: string): Promise<string> {
 	return absDir;
 }
 
+// The REPL outlives main(): main() resolves once the prompt loop is wired up, so
+// telemetry must be flushed by whichever exit path the user actually takes, not
+// when main()'s promise settles.
+let _replActive = false;
+
 async function main(): Promise<void> {
 	// Handle plugin subcommand: gitagent plugin <install|list|remove|...>
 	if (process.argv[2] === "plugin") {
@@ -321,7 +334,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 +478,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 +656,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(() => {});
@@ -702,6 +716,7 @@ async function main(): Promise<void> {
 				} catch {
 					/* ignore */
 				}
+				await shutdownTelemetry().catch(() => {});
 				process.exit(0);
 			}
 
@@ -804,6 +819,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}`));
@@ -843,10 +859,13 @@ async function main(): Promise<void> {
 			try {
 				_session.end({ "gitagent.cost_usd": _totalCostUsd });
 			} catch { /* ignore */ }
-			Promise.all([mcpSetup.cleanup(), stopSandbox()]).finally(() => process.exit(0));
+			Promise.all([mcpSetup.cleanup(), stopSandbox()])
+				.finally(() => shutdownTelemetry().catch(() => {}))
+				.finally(() => process.exit(0));
 		}
 	});
 
+	_replActive = true;
 	ask();
 }
 
@@ -856,7 +875,12 @@ process.on("SIGTERM", () => {
 });
 
 main()
-  .finally(() => shutdownTelemetry().catch(() => {}))
+  .finally(() => {
+    // Single-shot mode ends here; the REPL flushes from its own exit paths.
+    // finally runs before the catch below, so a prompt that throws still
+    // flushes before process.exit discards anything pending.
+    if (!_replActive) shutdownTelemetry().catch(() => {});
+  })
   .catch((err) => {
     console.error(red(`Fatal: ${err.message}`));
     process.exit(1);

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,40 @@ 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.
+ *
+ * Writes through to the model rather than returning a header map. The Agent is
+ * constructed with this exact object and pi-ai reads `headers` at request time,
+ * so a copy made here would never be seen. That is safe because the model is
+ * already this run's own: `loadAgent` clones it off the shared registry, and
+ * every `query()` loads its own, so concurrent runs never share one. Turns
+ * within a run are sequential, so the only writer per object is this function.
+ */
+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 {