From 17b2857dcd5229643a0c72168007663e38a735a5 Mon Sep 17 00:00:00 2001 From: Nivesh353 Date: Fri, 14 Aug 2026 14:04:53 +0530 Subject: [PATCH 1/4] 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 { From 6477cc9d0cf477bb06cf99a2e8684713d262ea06 Mon Sep 17 00:00:00 2001 From: Nivesh353 Date: Mon, 17 Aug 2026 12:19:47 +0530 Subject: [PATCH 2/4] fix(telemetry): flush spans on the exit path the user actually takes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit main() resolves as soon as the REPL is wired up, so `.finally(() => shutdownTelemetry())` fired at startup and tore the exporter down before the first prompt — interactive runs exported nothing. Flush from where the session really ends instead: /quit, Ctrl+C, or main() settling in single-shot mode. --- src/index.ts | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/index.ts b/src/index.ts index 7adedb0..a5be220 100644 --- a/src/index.ts +++ b/src/index.ts @@ -711,6 +711,7 @@ async function main(): Promise { } catch { /* ignore */ } + await shutdownTelemetry().catch(() => {}); process.exit(0); } @@ -853,20 +854,31 @@ async function main(): Promise { 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(); } +// 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; + // Flush OpenTelemetry exporters on SIGTERM. No-op when telemetry is disabled. process.on("SIGTERM", () => { shutdownTelemetry().catch(() => {}).finally(() => process.exit(0)); }); main() - .finally(() => shutdownTelemetry().catch(() => {})) + .finally(() => { + // Single-shot mode ends here; the REPL flushes from its own exit paths. + if (!_replActive) shutdownTelemetry().catch(() => {}); + }) .catch((err) => { console.error(red(`Fatal: ${err.message}`)); process.exit(1); From f960ae9d24f6a3037c67546cac7a41803ad408e0 Mon Sep 17 00:00:00 2001 From: Nivesh353 Date: Mon, 17 Aug 2026 15:50:12 +0530 Subject: [PATCH 3/4] docs: document --session-id in the CLI flag table --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 7d0257f..159bf06 100644 --- a/README.md +++ b/README.md @@ -132,6 +132,7 @@ gitagent --repo https://github.com/org/repo "Add unit tests" | `--repo ` | `-r` | GitHub repo URL to clone and work on | | `--pat ` | | GitHub PAT (or set `GITHUB_TOKEN` / `GIT_TOKEN`) | | `--session ` | | Resume an existing session branch | +| `--session-id ` | | Session id sent on model requests, so a gateway groups the run (default: generated) | | `--model ` | `-m` | Override model (e.g. `anthropic:claude-sonnet-4-5-20250929`) | | `--sandbox` | `-s` | Run in sandbox VM | | `--prompt ` | `-p` | Single-shot prompt (skip REPL) | From 5535ffc8121d87bdeb5224a236fe60e68cd5af31 Mon Sep 17 00:00:00 2001 From: Nivesh353 Date: Mon, 17 Aug 2026 16:01:36 +0530 Subject: [PATCH 4/4] docs(telemetry): explain the write-through, hoist the REPL flag Reviewer feedback: the model mutation in startTurnTrace and the flush on the single-shot error path both looked accidental. Say why they are deliberate, and declare _replActive above the function that reads it. --- src/index.ts | 12 +++++++----- src/telemetry.ts | 7 +++++++ 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/index.ts b/src/index.ts index a5be220..9e5e71e 100644 --- a/src/index.ts +++ b/src/index.ts @@ -312,6 +312,11 @@ async function ensureRepo(dir: string, model?: string): Promise { 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 { // Handle plugin subcommand: gitagent plugin if (process.argv[2] === "plugin") { @@ -864,11 +869,6 @@ async function main(): Promise { ask(); } -// 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; - // Flush OpenTelemetry exporters on SIGTERM. No-op when telemetry is disabled. process.on("SIGTERM", () => { shutdownTelemetry().catch(() => {}).finally(() => process.exit(0)); @@ -877,6 +877,8 @@ process.on("SIGTERM", () => { main() .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) => { diff --git a/src/telemetry.ts b/src/telemetry.ts index 511a889..a6a515e 100644 --- a/src/telemetry.ts +++ b/src/telemetry.ts @@ -200,6 +200,13 @@ export function isTelemetryEnabled(): boolean { * 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 {