@narumitw/pi-langfuse 0.18.0 → 0.26.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/src/tracing.ts CHANGED
@@ -1,11 +1,30 @@
1
- export const MAX_CAPTURE_BYTES = 64 * 1024;
2
- const MAX_STRING_LENGTH = 50_000;
3
- const MAX_COLLECTION_LENGTH = 200;
4
- const MAX_DEPTH = 12;
1
+ import { sanitizeTraceValue } from "./sanitizer.js";
2
+
3
+ export { MAX_CAPTURE_BYTES, sanitizeTraceValue } from "./sanitizer.js";
4
+
5
+ export const TRACE_SCHEMA_VERSION = "2";
6
+ const MAX_HEADER_VALUE_LENGTH = 1_024;
5
7
  const CONTENT_DISABLED = "[content capture disabled]";
6
- const TRUNCATED = "[truncated: content budget exceeded]";
7
- const BASE64_DATA_URI = /data:[^,\s"'`]*;base64,[^\s"'`]*/gi;
8
- const BASE64_DATA_URI_OMITTED = "[base64 data URI omitted]";
8
+ const ALLOWED_RESPONSE_HEADERS = new Set([
9
+ "anthropic-ratelimit-requests-limit",
10
+ "anthropic-ratelimit-requests-remaining",
11
+ "anthropic-ratelimit-requests-reset",
12
+ "anthropic-ratelimit-tokens-limit",
13
+ "anthropic-ratelimit-tokens-remaining",
14
+ "anthropic-ratelimit-tokens-reset",
15
+ "cf-ray",
16
+ "openai-request-id",
17
+ "request-id",
18
+ "retry-after",
19
+ "x-amzn-requestid",
20
+ "x-request-id",
21
+ "x-ratelimit-limit-requests",
22
+ "x-ratelimit-limit-tokens",
23
+ "x-ratelimit-remaining-requests",
24
+ "x-ratelimit-remaining-tokens",
25
+ "x-ratelimit-reset-requests",
26
+ "x-ratelimit-reset-tokens",
27
+ ]);
9
28
 
10
29
  export interface ObservationAttributes {
11
30
  input?: unknown;
@@ -13,9 +32,12 @@ export interface ObservationAttributes {
13
32
  metadata?: Record<string, unknown>;
14
33
  level?: "DEBUG" | "DEFAULT" | "WARNING" | "ERROR";
15
34
  statusMessage?: string;
35
+ completionStartTime?: Date;
16
36
  model?: string;
37
+ modelParameters?: Record<string, string | number>;
17
38
  usageDetails?: Record<string, number>;
18
39
  costDetails?: Record<string, number>;
40
+ version?: string;
19
41
  name?: string;
20
42
  sessionId?: string;
21
43
  tags?: string[];
@@ -49,27 +71,66 @@ interface RecorderContext {
49
71
  interface ModelDescriptor {
50
72
  provider?: string;
51
73
  id?: string;
74
+ api?: string;
75
+ }
76
+
77
+ export interface GitMetadata {
78
+ branch?: string;
79
+ commit?: string;
80
+ detached: boolean;
81
+ }
82
+
83
+ export interface ContextSnapshot {
84
+ leafId?: string | null;
85
+ contextUsage?: {
86
+ tokens: number | null;
87
+ contextWindow: number;
88
+ percent: number | null;
89
+ };
52
90
  }
53
91
 
54
92
  interface BeginAgentInput {
55
93
  prompt: unknown;
56
94
  images?: unknown;
57
95
  model?: ModelDescriptor;
96
+ git?: GitMetadata;
97
+ snapshot?: ContextSnapshot;
98
+ }
99
+
100
+ interface AttemptInput {
101
+ reason?: string;
102
+ }
103
+
104
+ interface GenerationInput {
105
+ payload?: unknown;
106
+ payloadStage?: "before_provider_request";
107
+ model?: ModelDescriptor;
108
+ thinkingLevel?: string;
58
109
  }
59
110
 
60
111
  interface AssistantMessage {
61
112
  role: string;
62
113
  content?: unknown;
63
114
  provider?: string;
115
+ api?: string;
64
116
  model?: string;
65
117
  responseModel?: string;
118
+ responseId?: string;
66
119
  usage?: {
67
120
  input?: number;
68
121
  output?: number;
69
122
  cacheRead?: number;
70
123
  cacheWrite?: number;
124
+ cacheWrite1h?: number;
125
+ reasoning?: number;
71
126
  totalTokens?: number;
72
- cost?: { total?: number };
127
+ cost?: {
128
+ input?: number;
129
+ output?: number;
130
+ cacheRead?: number;
131
+ cacheWrite?: number;
132
+ total?: number;
133
+ };
73
134
  };
74
135
  stopReason?: string;
75
136
  errorMessage?: string;
@@ -86,14 +147,87 @@ interface TurnResult {
86
147
  toolResultCount: number;
87
148
  }
88
149
 
150
+ interface CompactionStart {
151
+ reason: string;
152
+ willRetry: boolean;
153
+ tokensBefore?: number;
154
+ messagesToSummarize: number;
155
+ turnPrefixMessages: number;
156
+ branchEntries: number;
157
+ isSplitTurn: boolean;
158
+ }
159
+
160
+ interface CompactionFinish {
161
+ reason: string;
162
+ willRetry: boolean;
163
+ fromExtension: boolean;
164
+ tokensBefore?: number;
165
+ details?: unknown;
166
+ usage?: {
167
+ input?: number;
168
+ output?: number;
169
+ cacheRead?: number;
170
+ cacheWrite?: number;
171
+ totalTokens?: number;
172
+ cost?: {
173
+ input?: number;
174
+ output?: number;
175
+ cacheRead?: number;
176
+ cacheWrite?: number;
177
+ total?: number;
178
+ };
179
+ };
180
+ }
181
+
182
+ type Outcome = "success" | "error" | "aborted" | "length" | "interrupted";
183
+
184
+ interface GenerationState {
185
+ observation: Observation;
186
+ endTime?: number;
187
+ completionStartTime?: number;
188
+ requestMetadata: Record<string, unknown>;
189
+ statuses: number[];
190
+ headers: Record<string, string>;
191
+ }
192
+
193
+ interface ToolState {
194
+ observation: Observation;
195
+ startedAt: number;
196
+ progressUpdates: number;
197
+ firstProgressAt?: number;
198
+ }
199
+
200
+ interface CompactionState {
201
+ observation: Observation;
202
+ }
203
+
204
+ interface Counters {
205
+ attempts: number;
206
+ turns: number;
207
+ generations: number;
208
+ tools: number;
209
+ toolErrors: number;
210
+ compactions: number;
211
+ recoveredErrors: number;
212
+ failedAttempts: number;
213
+ }
214
+
89
215
  export class TraceRecorder {
90
216
  private root: Observation | undefined;
217
+ private attempt: Observation | undefined;
218
+ private attemptIndex: number | undefined;
91
219
  private turn: Observation | undefined;
92
220
  private turnIndex: number | undefined;
93
- private generation: Observation | undefined;
94
- private generationEndTime: number | undefined;
95
- private readonly tools = new Map<string, Observation>();
221
+ private generation: GenerationState | undefined;
222
+ private readonly tools = new Map<string, ToolState>();
223
+ private readonly duplicateToolIds = new Set<string>();
224
+ private compaction: CompactionState | undefined;
96
225
  private lastOutput: unknown;
226
+ private lastAssistant: AssistantMessage | undefined;
227
+ private lastAttemptOutcome: Outcome | undefined;
228
+ private unresolvedToolErrors = 0;
229
+ private startSnapshot: ContextSnapshot | undefined;
230
+ private counters: Counters = emptyCounters();
97
231
 
98
232
  constructor(
99
233
  private readonly backend: TraceBackend,
@@ -104,40 +238,98 @@ export class TraceRecorder {
104
238
  return this.root !== undefined;
105
239
  }
106
240
 
241
+ hasActiveAttempt(): boolean {
242
+ return this.attempt !== undefined;
243
+ }
244
+
107
245
  beginAgent(input: BeginAgentInput): void {
108
- if (this.root) this.closeActiveTrace("Interrupted by a new Pi agent run.");
246
+ if (this.root)
247
+ this.closeActiveTrace("Interrupted by a new Pi run.", undefined, "interrupted");
109
248
 
110
249
  this.lastOutput = undefined;
250
+ this.lastAssistant = undefined;
251
+ this.lastAttemptOutcome = undefined;
252
+ this.unresolvedToolErrors = 0;
253
+ this.duplicateToolIds.clear();
254
+ this.startSnapshot = input.snapshot;
255
+ this.counters = emptyCounters();
111
256
  const traceInput = this.capture({
112
257
  prompt: input.prompt,
113
258
  ...(hasItems(input.images) ? { images: input.images } : {}),
114
259
  });
115
260
  const metadata: Record<string, unknown> = {
116
261
  "pi.cwd": this.context.cwd,
262
+ ...(input.git?.branch ? { "pi.git.branch": input.git.branch } : {}),
263
+ ...(input.git?.commit ? { "pi.git.commit": input.git.commit } : {}),
264
+ ...(input.git ? { "pi.git.detached": input.git.detached } : {}),
117
265
  "pi.mode": this.context.mode,
118
266
  ...(input.model?.id ? { "pi.model": input.model.id } : {}),
119
267
  ...(input.model?.provider ? { "pi.provider": input.model.provider } : {}),
120
268
  "pi.session.id": this.context.sessionId,
269
+ "pi.trace.schema_version": Number(TRACE_SCHEMA_VERSION),
270
+ ...snapshotMetadata("start", input.snapshot),
121
271
  };
122
-
123
- this.root = this.backend.start("pi.agent", { input: traceInput, metadata }, { asType: "agent" });
272
+ const attributes = { input: traceInput, metadata, version: TRACE_SCHEMA_VERSION };
273
+ const gitTag = input.git?.branch
274
+ ? `branch:${input.git.branch}`
275
+ : input.git?.detached
276
+ ? "git:detached"
277
+ : undefined;
278
+
279
+ this.root = this.backend.start("pi.agent", attributes, { asType: "agent" });
124
280
  this.root.updateTrace?.({
125
281
  name: "pi.trace",
126
282
  sessionId: this.context.sessionId,
283
+ version: TRACE_SCHEMA_VERSION,
127
284
  input: traceInput,
128
285
  metadata,
129
- tags: ["pi"],
286
+ tags: ["pi", ...(gitTag ? [gitTag] : [])],
130
287
  });
131
288
  }
132
289
 
290
+ beginAttempt(input: AttemptInput = {}): void {
291
+ if (!this.root) return;
292
+ if (this.attempt) {
293
+ this.closeAttempt(undefined, "Interrupted by the next Pi attempt.", "interrupted");
294
+ }
295
+ const index = this.counters.attempts;
296
+ this.counters.attempts += 1;
297
+ this.lastAssistant = undefined;
298
+ this.unresolvedToolErrors = 0;
299
+ this.attemptIndex = index;
300
+ this.attempt = this.backend.start(
301
+ "pi.attempt",
302
+ {
303
+ metadata: {
304
+ "pi.attempt.index": index,
305
+ ...(input.reason ? { "pi.attempt.reason": input.reason } : {}),
306
+ },
307
+ version: TRACE_SCHEMA_VERSION,
308
+ },
309
+ { asType: "span", parent: this.root },
310
+ );
311
+ }
312
+
313
+ finishAttempt(message?: AssistantMessage): void {
314
+ if (message?.role === "assistant") this.finishAssistant(message);
315
+ if (!this.attempt) return;
316
+ const outcome = this.unresolvedToolErrors > 0 ? "error" : classifyOutcome(message);
317
+ this.closeAttempt(
318
+ message,
319
+ this.unresolvedToolErrors > 0 ? "Pi attempt ended after a tool failure." : undefined,
320
+ outcome,
321
+ );
322
+ }
323
+
133
324
  beginTurn(turnIndex: number): void {
134
325
  if (!this.root) return;
135
- if (this.turn) this.closeTurn("Interrupted by the next Pi turn.");
326
+ if (this.turn) this.closeTurn("Interrupted by the next Pi turn.", "ERROR");
327
+ this.counters.turns += 1;
136
328
  this.turnIndex = turnIndex;
137
329
  this.turn = this.backend.start(
138
330
  "pi.turn",
139
- { metadata: { "pi.turn.index": turnIndex } },
140
- { asType: "span", parent: this.root },
331
+ { metadata: { "pi.turn.index": turnIndex }, version: TRACE_SCHEMA_VERSION },
332
+ { asType: "span", parent: this.attempt ?? this.root },
141
333
  );
142
334
  }
143
335
 
@@ -147,66 +339,86 @@ export class TraceRecorder {
147
339
  this.closeTools("Tool span ended when the Pi turn finished.");
148
340
 
149
341
  const mismatched = this.turnIndex !== turnIndex;
150
- const failed =
151
- result.message.role === "assistant" &&
152
- (result.message.stopReason === "error" || Boolean(result.message.errorMessage));
342
+ const outcome = classifyOutcome(result.message);
153
343
  this.turn.update({
154
344
  metadata: {
155
345
  "pi.turn.index": this.turnIndex ?? turnIndex,
156
346
  "pi.turn.tool_result_count": result.toolResultCount,
157
- ...(result.message.stopReason
158
- ? { "pi.turn.stop_reason": result.message.stopReason }
159
- : {}),
347
+ ...(result.message.stopReason ? { "pi.turn.stop_reason": result.message.stopReason } : {}),
160
348
  },
161
- ...(failed
349
+ ...(mismatched && outcome === "success"
162
350
  ? {
163
- level: "ERROR" as const,
164
- statusMessage: result.message.errorMessage ?? "The Pi turn failed.",
351
+ level: "WARNING" as const,
352
+ statusMessage: `Turn ${turnIndex} ended while turn ${this.turnIndex} was active.`,
165
353
  }
166
- : mismatched
167
- ? {
168
- level: "WARNING" as const,
169
- statusMessage: `Turn ${turnIndex} ended while turn ${this.turnIndex} was active.`,
170
- }
171
- : {}),
354
+ : severityForOutcome(outcome, result.message.errorMessage)),
172
355
  });
173
356
  this.turn.end();
174
357
  this.turn = undefined;
175
358
  this.turnIndex = undefined;
176
359
  }
177
360
 
178
- beginGeneration(): void {
361
+ beginGeneration(input: GenerationInput = {}): void {
179
362
  if (!this.root) return;
180
363
  this.closeGeneration("Interrupted by the next provider request.");
181
- this.generationEndTime = undefined;
182
- this.generation = this.backend.start(
364
+ if (this.unresolvedToolErrors > 0) {
365
+ this.counters.recoveredErrors += this.unresolvedToolErrors;
366
+ this.unresolvedToolErrors = 0;
367
+ }
368
+ this.counters.generations += 1;
369
+ const requestMetadata: Record<string, unknown> = {
370
+ ...(input.payloadStage ? { "pi.request.payload_stage": input.payloadStage } : {}),
371
+ ...(input.model?.provider ? { "pi.request.provider": input.model.provider } : {}),
372
+ ...(input.model?.id ? { "pi.request.model": input.model.id } : {}),
373
+ ...(input.model?.api ? { "pi.request.api": input.model.api } : {}),
374
+ ...(input.thinkingLevel ? { "pi.request.thinking_level": input.thinkingLevel } : {}),
375
+ };
376
+ const observation = this.backend.start(
183
377
  "pi.llm",
184
- {},
185
- { asType: "generation", parent: this.turn ?? this.root },
378
+ {
379
+ ...(input.payload !== undefined ? { input: this.capture(input.payload) } : {}),
380
+ ...(input.model?.id ? { model: input.model.id } : {}),
381
+ ...(input.thinkingLevel
382
+ ? { modelParameters: { thinking_level: input.thinkingLevel } }
383
+ : {}),
384
+ ...(Object.keys(requestMetadata).length > 0 ? { metadata: requestMetadata } : {}),
385
+ version: TRACE_SCHEMA_VERSION,
386
+ },
387
+ { asType: "generation", parent: this.turn ?? this.attempt ?? this.root },
186
388
  );
389
+ this.generation = {
390
+ observation,
391
+ requestMetadata,
392
+ statuses: [],
393
+ headers: {},
394
+ };
187
395
  }
188
396
 
189
- recordProviderResponse(status: number): void {
397
+ recordProviderResponse(status: number, headers: Record<string, string> = {}): void {
190
398
  if (!this.generation) return;
191
- this.generation.update({
192
- metadata: { "http.response.status_code": status },
193
- ...(status >= 400
194
- ? { level: "ERROR" as const, statusMessage: `Provider returned HTTP ${status}.` }
195
- : {}),
196
- });
399
+ this.generation.statuses.push(status);
400
+ Object.assign(this.generation.headers, allowlistedResponseHeaders(headers));
401
+ }
402
+
403
+ markGenerationFirstOutput(timestamp = Date.now()): void {
404
+ if (this.generation && this.generation.completionStartTime === undefined) {
405
+ this.generation.completionStartTime = timestamp;
406
+ }
197
407
  }
198
408
 
199
409
  markGenerationEnd(endTime = Date.now()): void {
200
- if (this.generation && this.generationEndTime === undefined) {
201
- this.generationEndTime = endTime;
410
+ if (this.generation && this.generation.endTime === undefined) {
411
+ this.generation.endTime = endTime;
202
412
  }
203
413
  }
204
414
 
205
415
  finishAssistant(message: AssistantMessage): void {
206
416
  if (message.role !== "assistant") return;
417
+ this.lastAssistant = message;
207
418
  this.lastOutput = this.capture(message.content);
208
419
  if (!this.generation) return;
209
420
 
421
+ const state = this.generation;
210
422
  const usageDetails = numericRecord({
211
423
  input: message.usage?.input,
212
424
  output: message.usage?.output,
@@ -214,104 +426,305 @@ export class TraceRecorder {
214
426
  cache_creation_input_tokens: message.usage?.cacheWrite,
215
427
  total: message.usage?.totalTokens,
216
428
  });
217
- const totalCost = message.usage?.cost?.total;
218
- const failed = message.stopReason === "error" || Boolean(message.errorMessage);
429
+ const costDetails = positiveNumericRecord({
430
+ input: message.usage?.cost?.input,
431
+ output: message.usage?.cost?.output,
432
+ cache_read: message.usage?.cost?.cacheRead,
433
+ cache_write: message.usage?.cost?.cacheWrite,
434
+ total: message.usage?.cost?.total,
435
+ });
436
+ const outcome = classifyOutcome(message);
219
437
  const responseModel = message.responseModel ?? message.model;
220
- this.generation.update({
438
+ const httpMetadata = generationHttpMetadata(state);
439
+ const recoveredStatuses = state.statuses.filter((status) => status >= 400).length;
440
+ if (outcome === "success" && recoveredStatuses > 0) {
441
+ this.counters.recoveredErrors += recoveredStatuses;
442
+ }
443
+ state.observation.update({
221
444
  output: this.lastOutput,
445
+ ...(state.completionStartTime !== undefined
446
+ ? { completionStartTime: new Date(state.completionStartTime) }
447
+ : {}),
222
448
  ...(responseModel ? { model: responseModel } : {}),
223
449
  ...(Object.keys(usageDetails).length > 0 ? { usageDetails } : {}),
224
- ...(typeof totalCost === "number" && Number.isFinite(totalCost) && totalCost > 0
225
- ? { costDetails: { total: totalCost } }
226
- : {}),
450
+ ...(Object.keys(costDetails).length > 0 ? { costDetails } : {}),
227
451
  metadata: {
228
- ...(message.provider ? { "pi.provider": message.provider } : {}),
452
+ ...state.requestMetadata,
453
+ ...httpMetadata,
454
+ ...(message.provider
455
+ ? { "pi.provider": message.provider, "pi.response.provider": message.provider }
456
+ : {}),
457
+ ...(message.api ? { "pi.response.api": message.api } : {}),
458
+ ...(responseModel ? { "pi.response.model": responseModel } : {}),
459
+ ...(message.responseId ? { "pi.response.id": message.responseId } : {}),
229
460
  ...(message.responseModel && message.model && message.responseModel !== message.model
230
461
  ? { "pi.requested_model": message.model }
231
462
  : {}),
232
463
  ...(message.stopReason ? { "pi.stop_reason": message.stopReason } : {}),
464
+ ...(finiteNumber(message.usage?.reasoning)
465
+ ? { "pi.usage.reasoning_tokens": message.usage.reasoning }
466
+ : {}),
467
+ ...(finiteNumber(message.usage?.cacheWrite1h)
468
+ ? { "pi.usage.cache_write_1h_tokens": message.usage.cacheWrite1h }
469
+ : {}),
233
470
  },
234
- ...(failed
235
- ? {
236
- level: "ERROR" as const,
237
- statusMessage: message.errorMessage ?? "The model returned an error.",
238
- }
239
- : {}),
471
+ ...severityForOutcome(outcome, message.errorMessage),
240
472
  });
241
- this.generation.end(this.generationEndTime);
473
+ state.observation.end(state.endTime);
242
474
  this.generation = undefined;
243
- this.generationEndTime = undefined;
244
475
  }
245
476
 
246
- beginTool(toolCallId: string, toolName: string, args?: unknown): void {
477
+ beginTool(toolCallId: string, toolName: string, args?: unknown, startedAt = Date.now()): void {
247
478
  if (!this.root) return;
479
+ this.counters.tools += 1;
480
+ if (this.duplicateToolIds.has(toolCallId)) {
481
+ this.counters.toolErrors += 1;
482
+ const duplicate = this.startToolObservation(toolCallId, toolName, args);
483
+ duplicate.update({
484
+ level: "ERROR",
485
+ statusMessage: "Duplicate tool call ID cannot be correlated safely.",
486
+ });
487
+ duplicate.end();
488
+ return;
489
+ }
490
+
248
491
  const existing = this.tools.get(toolCallId);
249
492
  if (existing) {
250
- existing.update({ level: "ERROR", statusMessage: "Duplicate tool execution started." });
251
- existing.end();
493
+ existing.observation.update({
494
+ level: "ERROR",
495
+ statusMessage: "Duplicate tool execution started.",
496
+ });
497
+ existing.observation.end();
498
+ this.tools.delete(toolCallId);
499
+ this.counters.toolErrors += 2;
500
+ const duplicate = this.startToolObservation(toolCallId, toolName, args);
501
+ duplicate.update({
502
+ level: "ERROR",
503
+ statusMessage: "Duplicate tool call ID cannot be correlated safely.",
504
+ });
505
+ duplicate.end();
506
+ this.duplicateToolIds.add(toolCallId);
507
+ return;
252
508
  }
253
- this.tools.set(
254
- toolCallId,
255
- this.backend.start(
256
- `pi.tool.${toolName}`,
257
- {
258
- ...(args !== undefined ? { input: this.capture(args) } : {}),
259
- metadata: { "pi.tool.call_id": toolCallId, "pi.tool.name": toolName },
260
- },
261
- { asType: "tool", parent: this.turn ?? this.root },
262
- ),
263
- );
509
+
510
+ this.tools.set(toolCallId, {
511
+ observation: this.startToolObservation(toolCallId, toolName, args),
512
+ startedAt,
513
+ progressUpdates: 0,
514
+ });
264
515
  }
265
516
 
266
- updateToolInput(toolCallId: string, input: unknown): void {
267
- this.tools.get(toolCallId)?.update({ input: this.capture(input) });
517
+ recordToolInput(toolCallId: string, input: unknown): void {
518
+ if (input === undefined || this.duplicateToolIds.has(toolCallId)) return;
519
+ this.tools.get(toolCallId)?.observation.update({ input: this.capture(input) });
520
+ }
521
+
522
+ recordToolProgress(toolCallId: string, timestamp = Date.now()): void {
523
+ const tool = this.tools.get(toolCallId);
524
+ if (!tool) return;
525
+ tool.progressUpdates += 1;
526
+ tool.firstProgressAt ??= timestamp;
268
527
  }
269
528
 
270
529
  finishTool(toolCallId: string, result: ToolResult): void {
530
+ if (this.duplicateToolIds.has(toolCallId)) return;
271
531
  const tool = this.tools.get(toolCallId);
272
532
  if (!tool) return;
273
- tool.update({
533
+ if (result.isError) {
534
+ this.counters.toolErrors += 1;
535
+ this.unresolvedToolErrors += 1;
536
+ }
537
+ tool.observation.update({
274
538
  output: this.capture({ content: result.content, details: result.details }),
539
+ metadata: {
540
+ "pi.tool.progress_update_count": tool.progressUpdates,
541
+ ...(tool.firstProgressAt !== undefined
542
+ ? {
543
+ "pi.tool.time_to_first_progress_ms": Math.max(
544
+ 0,
545
+ tool.firstProgressAt - tool.startedAt,
546
+ ),
547
+ }
548
+ : {}),
549
+ },
275
550
  ...(result.isError
276
551
  ? { level: "ERROR" as const, statusMessage: "Pi tool execution failed." }
277
552
  : {}),
278
553
  });
279
- tool.end();
554
+ tool.observation.end();
280
555
  this.tools.delete(toolCallId);
281
556
  }
282
557
 
283
- settle(): void {
284
- this.closeActiveTrace();
558
+ beginCompaction(input: CompactionStart): void {
559
+ if (!this.root) return;
560
+ this.closeCompaction("Interrupted by another Pi compaction.");
561
+ this.counters.compactions += 1;
562
+ this.compaction = {
563
+ observation: this.backend.start(
564
+ "pi.compaction",
565
+ {
566
+ metadata: {
567
+ "pi.compaction.reason": input.reason,
568
+ "pi.compaction.will_retry": input.willRetry,
569
+ ...(finiteNumber(input.tokensBefore)
570
+ ? { "pi.compaction.tokens_before": input.tokensBefore }
571
+ : {}),
572
+ "pi.compaction.messages_to_summarize": input.messagesToSummarize,
573
+ "pi.compaction.turn_prefix_messages": input.turnPrefixMessages,
574
+ "pi.compaction.branch_entries": input.branchEntries,
575
+ "pi.compaction.is_split_turn": input.isSplitTurn,
576
+ },
577
+ version: TRACE_SCHEMA_VERSION,
578
+ },
579
+ { asType: "span", parent: this.root },
580
+ ),
581
+ };
582
+ }
583
+
584
+ finishCompaction(input: CompactionFinish): void {
585
+ if (!this.compaction) return;
586
+ const details = structuralCompactionDetails(input.details);
587
+ const usage = numericRecord({
588
+ input: input.usage?.input,
589
+ output: input.usage?.output,
590
+ cache_read: input.usage?.cacheRead,
591
+ cache_write: input.usage?.cacheWrite,
592
+ total: input.usage?.totalTokens,
593
+ });
594
+ const cost = positiveNumericRecord({
595
+ input: input.usage?.cost?.input,
596
+ output: input.usage?.cost?.output,
597
+ cache_read: input.usage?.cost?.cacheRead,
598
+ cache_write: input.usage?.cost?.cacheWrite,
599
+ total: input.usage?.cost?.total,
600
+ });
601
+ this.compaction.observation.update({
602
+ metadata: {
603
+ "pi.compaction.reason": input.reason,
604
+ "pi.compaction.will_retry": input.willRetry,
605
+ "pi.compaction.from_extension": input.fromExtension,
606
+ ...(finiteNumber(input.tokensBefore)
607
+ ? { "pi.compaction.tokens_before": input.tokensBefore }
608
+ : {}),
609
+ ...details,
610
+ ...prefixRecord("pi.compaction.usage", usage),
611
+ ...prefixRecord("pi.compaction.cost", cost),
612
+ },
613
+ });
614
+ this.compaction.observation.end();
615
+ this.compaction = undefined;
616
+ }
617
+
618
+ settle(snapshot?: ContextSnapshot): void {
619
+ this.closeActiveTrace(undefined, snapshot);
620
+ }
621
+
622
+ interrupt(statusMessage: string, snapshot?: ContextSnapshot): void {
623
+ this.closeActiveTrace(statusMessage, snapshot, "interrupted");
285
624
  }
286
625
 
287
626
  async flush(): Promise<void> {
288
627
  await this.backend.forceFlush();
289
628
  }
290
629
 
291
- async shutdown(): Promise<void> {
292
- this.closeActiveTrace("Pi shut down before the active trace settled.");
630
+ async shutdown(snapshot?: ContextSnapshot): Promise<void> {
631
+ this.closeActiveTrace("Pi shut down before the active trace settled.", snapshot, "interrupted");
293
632
  await this.backend.shutdown();
294
633
  }
295
634
 
296
- private closeActiveTrace(statusMessage?: string): void {
297
- this.closeTurn(statusMessage ?? "Pi turn ended when the agent settled.");
635
+ private closeActiveTrace(
636
+ statusMessage?: string,
637
+ snapshot?: ContextSnapshot,
638
+ forcedOutcome?: Outcome,
639
+ ): void {
640
+ this.closeTurn(statusMessage ?? "Pi turn ended when the agent run settled.");
641
+ this.closeAttempt(undefined, statusMessage, forcedOutcome ?? "interrupted");
642
+ this.closeCompaction(statusMessage ?? "Pi compaction ended when the agent run settled.");
298
643
 
299
644
  if (!this.root) return;
300
- this.root.update({
645
+ const baseOutcome =
646
+ forcedOutcome ?? this.lastAttemptOutcome ?? classifyOutcome(this.lastAssistant);
647
+ const recoveredErrorCount =
648
+ this.counters.recoveredErrors +
649
+ (baseOutcome === "success" ? this.counters.failedAttempts : 0);
650
+ const traceOutcome =
651
+ baseOutcome === "success" && recoveredErrorCount > 0 ? "recovered_success" : baseOutcome;
652
+ const finalMetadata: Record<string, unknown> = {
653
+ "pi.trace.schema_version": Number(TRACE_SCHEMA_VERSION),
654
+ "pi.trace.outcome": traceOutcome,
655
+ ...(this.lastAssistant?.stopReason
656
+ ? { "pi.trace.stop_reason": this.lastAssistant.stopReason }
657
+ : {}),
658
+ "pi.trace.attempt_count": this.counters.attempts,
659
+ "pi.trace.turn_count": this.counters.turns,
660
+ "pi.trace.generation_count": this.counters.generations,
661
+ "pi.trace.tool_count": this.counters.tools,
662
+ "pi.trace.tool_error_count": this.counters.toolErrors,
663
+ "pi.trace.compaction_count": this.counters.compactions,
664
+ "pi.trace.recovered_error_count": recoveredErrorCount,
665
+ ...snapshotMetadata("start", this.startSnapshot),
666
+ ...snapshotMetadata("end", snapshot),
667
+ };
668
+ const finalAttributes: ObservationAttributes = {
301
669
  ...(this.lastOutput !== undefined ? { output: this.lastOutput } : {}),
302
- ...(statusMessage ? { level: "WARNING" as const, statusMessage } : {}),
303
- });
670
+ metadata: finalMetadata,
671
+ ...severityForOutcome(baseOutcome, statusMessage ?? this.lastAssistant?.errorMessage),
672
+ };
673
+ this.root.update(finalAttributes);
304
674
  this.root.updateTrace?.({
305
675
  ...(this.lastOutput !== undefined ? { output: this.lastOutput } : {}),
306
- ...(statusMessage ? { metadata: { "pi.status": statusMessage } } : {}),
676
+ metadata: finalMetadata,
677
+ version: TRACE_SCHEMA_VERSION,
307
678
  });
308
679
  this.root.end();
309
680
  this.root = undefined;
310
681
  this.lastOutput = undefined;
682
+ this.lastAssistant = undefined;
683
+ this.lastAttemptOutcome = undefined;
684
+ this.unresolvedToolErrors = 0;
685
+ this.duplicateToolIds.clear();
686
+ this.startSnapshot = undefined;
311
687
  }
312
688
 
313
- private closeTurn(statusMessage: string): void {
314
- this.closeGeneration(statusMessage);
689
+ private closeAttempt(
690
+ message?: AssistantMessage,
691
+ statusMessage?: string,
692
+ forcedOutcome?: Outcome,
693
+ ): void {
694
+ if (!this.attempt) return;
695
+ this.closeTurn(statusMessage ?? "Pi turn ended when the attempt finished.");
696
+ const outcome = forcedOutcome ?? classifyOutcome(message);
697
+ if (outcome === "error") this.counters.failedAttempts += 1;
698
+ this.lastAttemptOutcome = outcome;
699
+ this.attempt.update({
700
+ ...(message?.content !== undefined ? { output: this.capture(message.content) } : {}),
701
+ metadata: {
702
+ "pi.attempt.index": this.attemptIndex ?? Math.max(0, this.counters.attempts - 1),
703
+ "pi.attempt.outcome": outcome,
704
+ ...(message?.stopReason ? { "pi.attempt.stop_reason": message.stopReason } : {}),
705
+ },
706
+ ...severityForOutcome(outcome, statusMessage ?? message?.errorMessage),
707
+ });
708
+ this.attempt.end();
709
+ this.attempt = undefined;
710
+ this.attemptIndex = undefined;
711
+ this.unresolvedToolErrors = 0;
712
+ }
713
+
714
+ private startToolObservation(toolCallId: string, toolName: string, args?: unknown): Observation {
715
+ return this.backend.start(
716
+ `pi.tool.${toolName}`,
717
+ {
718
+ ...(args !== undefined ? { input: this.capture(args) } : {}),
719
+ metadata: { "pi.tool.call_id": toolCallId, "pi.tool.name": toolName },
720
+ version: TRACE_SCHEMA_VERSION,
721
+ },
722
+ { asType: "tool", parent: this.turn ?? this.attempt ?? this.root },
723
+ );
724
+ }
725
+
726
+ private closeTurn(statusMessage: string, generationLevel: "ERROR" | "WARNING" = "WARNING"): void {
727
+ this.closeGeneration(statusMessage, generationLevel);
315
728
  this.closeTools(statusMessage);
316
729
  if (!this.turn) return;
317
730
  this.turn.update({ level: "WARNING", statusMessage });
@@ -322,18 +735,45 @@ export class TraceRecorder {
322
735
 
323
736
  private closeTools(statusMessage: string): void {
324
737
  for (const tool of this.tools.values()) {
325
- tool.update({ level: "WARNING", statusMessage });
326
- tool.end();
738
+ tool.observation.update({
739
+ level: "WARNING",
740
+ statusMessage,
741
+ metadata: {
742
+ "pi.tool.progress_update_count": tool.progressUpdates,
743
+ ...(tool.firstProgressAt !== undefined
744
+ ? {
745
+ "pi.tool.time_to_first_progress_ms": Math.max(
746
+ 0,
747
+ tool.firstProgressAt - tool.startedAt,
748
+ ),
749
+ }
750
+ : {}),
751
+ },
752
+ });
753
+ tool.observation.end();
327
754
  }
328
755
  this.tools.clear();
329
756
  }
330
757
 
331
- private closeGeneration(statusMessage: string): void {
758
+ private closeGeneration(statusMessage: string, level: "ERROR" | "WARNING" = "ERROR"): void {
332
759
  if (!this.generation) return;
333
- this.generation.update({ level: "ERROR", statusMessage });
334
- this.generation.end(this.generationEndTime);
760
+ this.generation.observation.update({
761
+ level,
762
+ statusMessage,
763
+ metadata: {
764
+ ...this.generation.requestMetadata,
765
+ ...generationHttpMetadata(this.generation),
766
+ },
767
+ });
768
+ this.generation.observation.end(this.generation.endTime);
335
769
  this.generation = undefined;
336
- this.generationEndTime = undefined;
770
+ }
771
+
772
+ private closeCompaction(statusMessage: string): void {
773
+ if (!this.compaction) return;
774
+ this.compaction.observation.update({ level: "WARNING", statusMessage });
775
+ this.compaction.observation.end();
776
+ this.compaction = undefined;
337
777
  }
338
778
 
339
779
  private capture(value: unknown): unknown {
@@ -341,117 +781,103 @@ export class TraceRecorder {
341
781
  }
342
782
  }
343
783
 
344
- export function sanitizeTraceValue(value: unknown): unknown {
345
- const budget = { remaining: MAX_CAPTURE_BYTES };
346
- const sanitized = sanitize(value, new WeakSet<object>(), 0, budget);
347
- return serializedBytes(sanitized) <= MAX_CAPTURE_BYTES ? sanitized : TRUNCATED;
784
+ function generationHttpMetadata(state: GenerationState): Record<string, unknown> {
785
+ const finalStatus = state.statuses.at(-1);
786
+ return {
787
+ ...(state.statuses.length > 0
788
+ ? {
789
+ "http.response.status_codes": [...state.statuses],
790
+ "http.response.status_code": finalStatus,
791
+ "http.response.attempt_count": state.statuses.length,
792
+ "http.response.retry_count": Math.max(0, state.statuses.length - 1),
793
+ }
794
+ : {}),
795
+ ...(Object.keys(state.headers).length > 0
796
+ ? { "http.response.headers": { ...state.headers } }
797
+ : {}),
798
+ };
348
799
  }
349
800
 
350
- function sanitize(
351
- value: unknown,
352
- active: WeakSet<object>,
353
- depth: number,
354
- budget: { remaining: number },
355
- ): unknown {
356
- if (budget.remaining <= byteLength(TRUNCATED)) return TRUNCATED;
357
- if (value === null || typeof value === "number" || typeof value === "boolean") {
358
- return consume(value, budget);
359
- }
360
- if (typeof value === "string") {
361
- const bounded = truncateString(value, Math.min(MAX_STRING_LENGTH, budget.remaining));
362
- const redacted = bounded.replace(BASE64_DATA_URI, BASE64_DATA_URI_OMITTED);
363
- return consume(redacted, budget);
364
- }
365
- if (typeof value === "bigint") return consume(value.toString(), budget);
366
- if (typeof value === "undefined" || typeof value === "function" || typeof value === "symbol") {
367
- return undefined;
368
- }
369
- if (value instanceof Date) {
370
- return consume(Number.isNaN(value.getTime()) ? "[invalid date]" : value.toISOString(), budget);
371
- }
372
- if (value instanceof Error) {
373
- return sanitize({ name: value.name, message: value.message }, active, depth, budget);
801
+ function allowlistedResponseHeaders(headers: Record<string, string>): Record<string, string> {
802
+ const output: Record<string, string> = {};
803
+ for (const [rawName, rawValue] of Object.entries(headers)) {
804
+ const name = rawName.toLowerCase();
805
+ if (!ALLOWED_RESPONSE_HEADERS.has(name) || typeof rawValue !== "string") continue;
806
+ output[name] = truncateString(rawValue, MAX_HEADER_VALUE_LENGTH);
374
807
  }
375
- if (depth >= MAX_DEPTH) return consume("[maximum depth reached]", budget);
376
- if (active.has(value)) return consume("[circular]", budget);
377
- active.add(value);
808
+ return output;
809
+ }
378
810
 
379
- let result: unknown;
380
- if (Array.isArray(value)) {
381
- const items: unknown[] = [];
382
- for (const item of value.slice(0, MAX_COLLECTION_LENGTH)) {
383
- if (budget.remaining <= byteLength(TRUNCATED)) break;
384
- items.push(sanitize(item, active, depth + 1, budget));
385
- }
386
- if (items.length < value.length) items.push(`[${value.length - items.length} items omitted]`);
387
- result = items;
388
- } else {
389
- const record = value as Record<string, unknown>;
390
- const output: Record<string, unknown> = {};
391
- const objectType = readProperty(record, "type").value;
392
- const redactData = objectType === "image" || objectType === "base64";
393
- let visited = 0;
394
- let processed = 0;
395
- let omitted = false;
396
- try {
397
- for (const key in record) {
398
- if (visited >= MAX_COLLECTION_LENGTH || budget.remaining <= byteLength(TRUNCATED)) {
399
- omitted = true;
400
- break;
401
- }
402
- visited += 1;
403
- const keyBudget = Math.min(MAX_STRING_LENGTH, Math.max(0, budget.remaining - 4));
404
- if (key.length > keyBudget) {
405
- omitted = true;
406
- break;
407
- }
408
- const keyBytes = byteLength(key);
409
- if (keyBytes > keyBudget) {
410
- omitted = true;
411
- break;
811
+ function structuralCompactionDetails(details: unknown): Record<string, unknown> {
812
+ if (!details || typeof details !== "object" || Array.isArray(details)) return {};
813
+ const record = details as Record<string, unknown>;
814
+ return {
815
+ ...(Array.isArray(record.readFiles)
816
+ ? { "pi.compaction.read_file_count": record.readFiles.length }
817
+ : {}),
818
+ ...(Array.isArray(record.modifiedFiles)
819
+ ? { "pi.compaction.modified_file_count": record.modifiedFiles.length }
820
+ : {}),
821
+ };
822
+ }
823
+
824
+ function snapshotMetadata(prefix: string, snapshot?: ContextSnapshot): Record<string, unknown> {
825
+ if (!snapshot) return {};
826
+ return {
827
+ ...(snapshot.leafId ? { [`pi.trace.${prefix}_leaf_id`]: snapshot.leafId } : {}),
828
+ ...(snapshot.contextUsage
829
+ ? {
830
+ ...(finiteNumber(snapshot.contextUsage.tokens)
831
+ ? { [`pi.trace.${prefix}_context_tokens`]: snapshot.contextUsage.tokens }
832
+ : {}),
833
+ [`pi.trace.${prefix}_context_window`]: snapshot.contextUsage.contextWindow,
834
+ ...(finiteNumber(snapshot.contextUsage.percent)
835
+ ? { [`pi.trace.${prefix}_context_percent`]: snapshot.contextUsage.percent }
836
+ : {}),
412
837
  }
413
- budget.remaining -= keyBytes + 4;
414
- if (!Object.hasOwn(record, key)) continue;
415
- const property = readProperty(record, key);
416
- if (redactData && key === "data") {
417
- output[key] = consume("[base64 omitted]", budget);
418
- } else if (!property.ok) output[key] = consume("[unreadable property]", budget);
419
- else output[key] = sanitize(property.value, active, depth + 1, budget);
420
- processed += 1;
421
- }
422
- } catch {
423
- if (processed === 0) {
424
- active.delete(value);
425
- return consume("[unreadable object]", budget);
426
- }
427
- omitted = true;
428
- }
429
- if (omitted) output["$truncated"] = "additional object entries omitted";
430
- result = output;
431
- }
432
- active.delete(value);
433
- return result;
838
+ : {}),
839
+ };
434
840
  }
435
841
 
436
- function readProperty(
437
- record: Record<string, unknown>,
438
- key: string,
439
- ): { ok: true; value: unknown } | { ok: false; value?: undefined } {
440
- try {
441
- return { ok: true, value: record[key] };
442
- } catch {
443
- return { ok: false };
444
- }
842
+ function classifyOutcome(message?: {
843
+ role?: string;
844
+ stopReason?: string;
845
+ errorMessage?: string;
846
+ }): Outcome {
847
+ if (message?.role !== "assistant") return "interrupted";
848
+ if (message.stopReason === "error" || message.errorMessage) return "error";
849
+ if (message.stopReason === "aborted") return "aborted";
850
+ if (message.stopReason === "length") return "length";
851
+ return "success";
445
852
  }
446
853
 
447
- function consume<T>(value: T, budget: { remaining: number }): T | string {
448
- const size = serializedBytes(value);
449
- if (size > budget.remaining) {
450
- budget.remaining -= byteLength(TRUNCATED);
451
- return TRUNCATED;
854
+ function severityForOutcome(
855
+ outcome: Outcome,
856
+ statusMessage?: string,
857
+ ): Pick<ObservationAttributes, "level" | "statusMessage"> {
858
+ if (outcome === "success") return {};
859
+ if (outcome === "error") {
860
+ return {
861
+ level: "ERROR",
862
+ statusMessage: statusMessage ?? "The Pi operation failed.",
863
+ };
452
864
  }
453
- budget.remaining -= size;
454
- return value;
865
+ return {
866
+ level: "WARNING",
867
+ statusMessage:
868
+ statusMessage ??
869
+ (outcome === "aborted"
870
+ ? "The Pi operation was aborted."
871
+ : outcome === "length"
872
+ ? "The model reached its output limit."
873
+ : "The Pi operation was interrupted."),
874
+ };
875
+ }
876
+
877
+ function prefixRecord(prefix: string, values: Record<string, number>): Record<string, number> {
878
+ return Object.fromEntries(
879
+ Object.entries(values).map(([key, value]) => [`${prefix}.${key}`, value]),
880
+ );
455
881
  }
456
882
 
457
883
  function truncateString(value: string, maxBytes: number): string {
@@ -472,20 +898,41 @@ function truncateString(value: string, maxBytes: number): string {
472
898
  return `${output}${suffix}`;
473
899
  }
474
900
 
475
- function serializedBytes(value: unknown): number {
476
- return byteLength(JSON.stringify(value) ?? "null");
477
- }
478
-
479
901
  function byteLength(value: string): number {
480
902
  return Buffer.byteLength(value, "utf8");
481
903
  }
482
904
 
483
905
  function numericRecord(values: Record<string, number | undefined>): Record<string, number> {
484
906
  return Object.fromEntries(
485
- Object.entries(values).filter((entry): entry is [string, number] => Number.isFinite(entry[1])),
907
+ Object.entries(values).filter((entry): entry is [string, number] => finiteNumber(entry[1])),
908
+ );
909
+ }
910
+
911
+ function positiveNumericRecord(values: Record<string, number | undefined>): Record<string, number> {
912
+ return Object.fromEntries(
913
+ Object.entries(values).filter(
914
+ (entry): entry is [string, number] => finiteNumber(entry[1]) && entry[1] > 0,
915
+ ),
486
916
  );
487
917
  }
488
918
 
919
+ function finiteNumber(value: unknown): value is number {
920
+ return typeof value === "number" && Number.isFinite(value);
921
+ }
922
+
489
923
  function hasItems(value: unknown): boolean {
490
924
  return Array.isArray(value) ? value.length > 0 : value !== undefined;
491
925
  }
926
+
927
+ function emptyCounters(): Counters {
928
+ return {
929
+ attempts: 0,
930
+ turns: 0,
931
+ generations: 0,
932
+ tools: 0,
933
+ toolErrors: 0,
934
+ compactions: 0,
935
+ recoveredErrors: 0,
936
+ failedAttempts: 0,
937
+ };
938
+ }