From 17b2857dcd5229643a0c72168007663e38a735a5 Mon Sep 17 00:00:00 2001 From: Nivesh353 Date: Fri, 14 Aug 2026 14:04:53 +0530 Subject: [PATCH] feat(telemetry): group a run's model requests into one session and trace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One user turn costs several requests to the model gateway — one returns a tool call, the next the answer. The gateway traced each separately, so a single "hi" showed up as two unrelated traces with no way to total them. Send X-Session-Id (this run's id, overridable via --session-id or the SDK's sessionId) on every request, and a per-turn traceparent so the gateway can stitch a turn's requests into one trace. Both no-op safely: the header is set on a cloned model, and traceparent yields to the undici instrumentation when OTel is initialised. --- src/index.ts | 16 +++++++++++++--- src/loader.ts | 22 +++++++++++++++++++--- src/sdk.ts | 8 +++++++- src/telemetry.ts | 28 ++++++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 7 deletions(-) diff --git a/src/index.ts b/src/index.ts index ba4eeb2..7adedb0 100644 --- a/src/index.ts +++ b/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 { 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 { 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 { // 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 { } try { + startTurnTrace(loaded.model); await otelContext.with(_session.ctx, () => agent.prompt(promptText)); } catch (err: any) { console.error(red(`Error: ${err.message}`)); diff --git a/src/loader.ts b/src/loader.ts index 3fe06ba..6909c76 100644 --- a/src/loader.ts +++ b/src/loader.ts @@ -117,8 +117,10 @@ async function ensureGitagentDir(agentDir: string): Promise { return gitagentDir; } -async function writeSessionState(gitagentDir: string): Promise { - const sessionId = randomUUID(); +async function writeSessionState(gitagentDir: string, override?: string): Promise { + // 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 { // 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 diff --git a/src/sdk.ts b/src/sdk.ts index 55b941c..d14a171 100644 --- a/src/sdk.ts +++ b/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), ); diff --git a/src/telemetry.ts b/src/telemetry.ts index 10cf562..511a889 100644 --- a/src/telemetry.ts +++ b/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 }; + 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 {