diff --git a/package-lock.json b/package-lock.json index 54f09a9..1c3de38 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "@open-gitagent/gitagent", - "version": "2.0.2", + "version": "2.2.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@open-gitagent/gitagent", - "version": "2.0.2", + "version": "2.2.0", "license": "MIT", "dependencies": { "@mariozechner/pi-agent-core": "^0.70.2", diff --git a/package.json b/package.json index 77f3060..73468c0 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@open-gitagent/gitagent", - "version": "2.1.0", + "version": "2.2.0", "description": "A universal git-native multimodal always learning AI Agent (TinyHuman)", "author": "shreyaskapale", "license": "MIT", diff --git a/src/index.ts b/src/index.ts index ba4eeb2..1119334 100644 --- a/src/index.ts +++ b/src/index.ts @@ -31,6 +31,7 @@ import { initTelemetry, wrapToolWithOtel, startSessionSpan, + startTurnTrace, recordGenAiCall, shutdownTelemetry, } from "./telemetry.js"; @@ -304,6 +305,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") { @@ -643,6 +649,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(() => {}); @@ -702,6 +709,7 @@ async function main(): Promise { } catch { /* ignore */ } + await shutdownTelemetry().catch(() => {}); process.exit(0); } @@ -804,6 +812,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}`)); @@ -843,10 +852,13 @@ 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(); } @@ -856,7 +868,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); 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..a6a515e 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,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 }; + 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 {