From 04d6447f7c492aafac97e2d2450b532650a85556 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 5 Aug 2026 22:03:14 +0200 Subject: [PATCH] feat: add typed telemetry contracts --- packages/agent/CHANGELOG.md | 1 + packages/agent/docs/harness-v2.md | 23 +- packages/agent/docs/telemetry-schema.md | 331 ++++++++++++ packages/agent/package.json | 2 + .../agent/scripts/generate-telemetry-docs.ts | 13 + packages/agent/src/harness/telemetry.ts | 506 ++++++++++++++++++ packages/agent/src/index.ts | 58 +- packages/agent/test/harness/telemetry.test.ts | 103 ++++ packages/ai/CHANGELOG.md | 1 + packages/ai/docs/telemetry-schema.md | 49 ++ packages/ai/package.json | 2 + .../ai/scripts/generate-telemetry-docs.ts | 120 +++++ packages/ai/src/api/simple-options.ts | 1 + packages/ai/src/index.ts | 1 + packages/ai/src/telemetry.ts | 375 +++++++++++++ packages/ai/src/types.ts | 3 + packages/ai/test/telemetry-options.test.ts | 171 ++++++ packages/ai/test/telemetry.test.ts | 154 ++++++ 18 files changed, 1900 insertions(+), 14 deletions(-) create mode 100644 packages/agent/docs/telemetry-schema.md create mode 100644 packages/agent/scripts/generate-telemetry-docs.ts create mode 100644 packages/agent/src/harness/telemetry.ts create mode 100644 packages/agent/test/harness/telemetry.test.ts create mode 100644 packages/ai/docs/telemetry-schema.md create mode 100644 packages/ai/scripts/generate-telemetry-docs.ts create mode 100644 packages/ai/src/telemetry.ts create mode 100644 packages/ai/test/telemetry-options.test.ts create mode 100644 packages/ai/test/telemetry.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 1a34b021b..b30ea248e 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -10,6 +10,7 @@ ### Added +- Added the typed harness telemetry schema, callback helpers, and generated schema reference. - Added bounded `Session.findEntriesOnBranch()` and `findEntryOnBranch()` queries with explicit traversal, filtering, ordering, and limit options. - Added a compile-complete `AgentHarness` v2 scaffold; unfinished operation paths reject with `HarnessNotImplemented` while durable execution is implemented. diff --git a/packages/agent/docs/harness-v2.md b/packages/agent/docs/harness-v2.md index c81d336cc..d2a1968df 100644 --- a/packages/agent/docs/harness-v2.md +++ b/packages/agent/docs/harness-v2.md @@ -1043,7 +1043,7 @@ Calls on a faulted harness reject with the same `HarnessFault` instance until th `finalMessage` is the run's newest entry that projects to an assistant message; `finalEntryId` is that entry's id. `leafId` is the lane's leaf when the operation finished — the race-free anchor for branch queries (`findEntriesOnBranch({ start: leafId })`). The two differ when a deferred write was applied after the final assistant message. Full transcripts are not duplicated into results; they are in the session and were delivered as events. -**Type provenance.** Core conversation and tool types (`AgentMessage`, `AgentTool`, `AgentToolResult`, `QueueMode`, `ThinkingLevel`) come from `packages/agent/src/types.ts`. Provider types (`Model`, `Models`, `Usage`, `RetryPolicy`, stream options, deferred handles) come from `packages/ai`. After I0, the generic telemetry contract, schema machinery, and pi-ai span schema come from `packages/ai/src/telemetry.ts`; the harness span schema comes from `packages/agent/src/harness/telemetry.ts`. Session, harness, hook, event, result, snapshot, navigation, and durable-record types are defined by the v2 implementation under `packages/agent/src/harness/`. Lowercase helpers in section 15 pseudocode without a definition (`preparation`, `runToolBatchForSingleCall`, request/option bags such as `AssistantRequest` and `FactWrite`) are constructive implementation detail, not contract. +**Type provenance.** Core conversation and tool types (`AgentMessage`, `AgentTool`, `AgentToolResult`, `QueueMode`, `ThinkingLevel`) come from `packages/agent/src/types.ts`. Provider types (`Model`, `Models`, `Usage`, `RetryPolicy`, stream options, deferred handles) come from `packages/ai`. The generic telemetry contract, schema machinery, and pi-ai span schema come from `packages/ai/src/telemetry.ts`; the harness span schema comes from `packages/agent/src/harness/telemetry.ts`. Session, harness, hook, event, result, snapshot, navigation, and durable-record types are defined by the v2 implementation under `packages/agent/src/harness/`. Lowercase helpers in section 15 pseudocode without a definition (`preparation`, `runToolBatchForSingleCall`, request/option bags such as `AssistantRequest` and `FactWrite`) are constructive implementation detail, not contract. ### Suspended operations @@ -2757,7 +2757,7 @@ Pi ships no exporter and requires no backend-specific telemetry implementation. ### Package ownership -I0 puts the generic contract, schema-definition machinery, shared no-op implementation, and pi-ai span schema in `packages/ai/src/telemetry.ts` and exports them from `@earendil-works/pi-ai`. `packages/agent/src/harness/telemetry.ts` imports and re-exports that generic surface, then adds the harness span schema and its inferred types; the agent root exports that module. There is one generic contract, never one copy per package. +The generic contract, schema-definition machinery, shared no-op implementation, and pi-ai span schema live in `packages/ai/src/telemetry.ts` and are exported from `@earendil-works/pi-ai`. `packages/agent/src/harness/telemetry.ts` imports and re-exports that generic surface, then adds the harness span schema and its inferred types; the agent root exports that module. There is one generic contract, never one copy per package. `ProviderRequestOptions.telemetryContext` carries the optional parent through every pi-ai operation, including stream/simple-stream, deferred fetch/cancel, and image generation. `AgentHarnessOptions.telemetryContext` supplies the harness root. Provider, `Models`, `ImagesModels`, and direct stream/image dispatch preserve the option unchanged. The agent-side request wrapper emits `pi.ai.request` through the pi-ai schema; pi-ai owns the vocabulary and helper, not a second ambient runtime. @@ -2966,7 +2966,7 @@ type AiSpanAttributes = AiSpanStartAttributes & AiSpanEndAttributes; ``` -The following tables are normative input to I0's schema objects. `!` means a required start attribute; `?` means an optional start attribute. Every end attribute is optional enrichment. Array element sets use `elementValues`; all other closed sets use `values`. The automatic throw/reject rule from the context contract applies to every span in addition to the explicit status rule shown. +The following tables are normative input to the schema objects. `!` means a required start attribute; `?` means an optional start attribute. Every end attribute is optional enrichment. Array element sets use `elementValues`; all other closed sets use `values`. The automatic throw/reject rule from the context contract applies to every span in addition to the explicit status rule shown. #### Pi-ai schema @@ -3204,7 +3204,7 @@ Gate invariants, asserted across Tier C: ## 21. Implementation status and work packages -Harness implementation lives directly in `packages/agent/src/harness/`, with v4 session tests in `packages/agent/test/harness/session/`. I0 will add the generic telemetry contract and pi-ai schema in `packages/ai/src/telemetry.ts`. The v4 session and SQLite backend are the default package interfaces; there is no experimental package surface. Retained compaction, message projection, resource, tool, environment, and utility modules stay under `src/harness/` and are adapted in place. +Harness implementation lives directly in `packages/agent/src/harness/`, with v4 session tests in `packages/agent/test/harness/session/`. The generic telemetry contract and pi-ai schema live in `packages/ai/src/telemetry.ts`. The v4 session and SQLite backend are the default package interfaces; there is no experimental package surface. Retained compaction, message projection, resource, tool, environment, and utility modules stay under `src/harness/` and are adapted in place. ### Landed foundation @@ -3215,10 +3215,11 @@ Harness implementation lives directly in `packages/agent/src/harness/`, with v4 - Compaction preparation, context projection, and branch summarization use v4 `Entry` queries and no longer depend on the legacy `SessionTreeEntry` model. Reusable compaction tests and dedicated context tests cover this interim behavior. - `Session.getBranch()` is not part of v4 and must not be reintroduced. Branch callers use `findEntriesOnBranch()` with an explicit start and deliberate ordering semantics; callers may use the documented `newestFirst` default. - The pi-ai deferred request types, `Models.fetchDeferred()` / `cancelDeferred()` dispatch, and faux-provider pending/ready/terminal/cancellation support are already implemented and tested. Only harness integration remains. +- The canonical telemetry contract, exact typed pi-ai and harness schemas, no-op context, request-option propagation, generated references, and focused tests are implemented; runtime span insertion remains O2. ### Scope boundary -`packages/coding-agent/**` remains untouched. This plan implements `packages/agent` and `packages/session-backends/sqlite-node`, plus I0's generic telemetry contract, pi-ai schema, exports, tests, and generated schema documentation under `packages/ai`. O2 consumes that pi-ai telemetry surface but otherwise keeps provider behavior unchanged. The v3 requirement is an input-format requirement for the new JSONL repository, not a coding-agent migration. No work package may modify coding-agent source, tests, RPC, UI, or package metadata. +`packages/coding-agent/**` remains untouched. This plan implements `packages/agent` and `packages/session-backends/sqlite-node`, plus the generic telemetry contract, pi-ai schema, exports, tests, and generated schema documentation under `packages/ai`. O2 consumes that pi-ai telemetry surface but otherwise keeps provider behavior unchanged. The v3 requirement is an input-format requirement for the new JSONL repository, not a coding-agent migration. No work package may modify coding-agent source, tests, RPC, UI, or package metadata. ### Package rules @@ -3245,7 +3246,7 @@ To avoid every parallel PR conflicting in this file, package authors do not edit 5. Run focused tests while iterating, then `npm run check` before handoff. 6. Report completion to the merge/status owner so the checkbox is updated immediately after merge. -For one serial worker, the recommended next package is **R0**, because it unblocks both reducer and JSONL work; **F0** follows. For a team, the dependency-ready roots are **F0, R0, QA1, I0, I1, and I2**. After I0 lands, L1 is also ready. Do not start J0 or R1 before R0. +The dependency-ready unreserved roots are **QA1, I1, I2, and L1**. Respect the active R2 and JSONL reservations below. ### Track F — scaffold truth and public ownership @@ -3356,9 +3357,7 @@ These packages own `packages/agent/src/harness/session/jsonl/**`, the concrete ` I0, I1, and I2 may proceed independently. I3 → I4 → I5 is serial and begins after R2 fixes the `LaneState` shape. These packages use separate modules with focused unit tests; I5 remains primitive-only and does not edit `agent-harness.ts`. -**In progress and reserved: I0 by @badlogic.** Other agents must not pick I0 while this ownership marker remains. - -- [ ] **I0 — telemetry contracts, typed schemas, and no-op context.** Dependencies: none. +- [x] **I0 — telemetry contracts, typed schemas, and no-op context.** Dependencies: none. - Primary files: `packages/ai/src/telemetry.ts`, `packages/ai/src/index.ts`, focused tests, package scripts, `packages/ai/scripts/generate-telemetry-docs.ts`, and generated `packages/ai/docs/telemetry-schema.md`; `packages/agent/src/harness/telemetry.ts`, `packages/agent/src/index.ts`, focused tests, package scripts, `packages/agent/scripts/generate-telemetry-docs.ts`, and generated `packages/agent/docs/telemetry-schema.md`. Do not edit `agent-harness.ts`; H0 replaces its temporary local type declarations after convergence. - In pi-ai, implement the one canonical section 18 callback-based `TelemetryContext` / `TelemetrySpan` contract, shared no-op context, serializable `defineTelemetrySchema()` machinery, the complete normative `AI_TELEMETRY_SCHEMA`, inferred span-name/attribute/union types, and typed `startAiSpan()` helper. Add optional `telemetryContext` to `ProviderRequestOptions` so every stream, deferred, and image option inherits it; provider, `Models`, `ImagesModels`, and direct dispatch preserve it. Export the telemetry surface from `@earendil-works/pi-ai`. - In agent, import and re-export the generic pi-ai surface; define the complete normative `HARNESS_TELEMETRY_SCHEMA`, inferred harness types, and typed `startHarnessSpan()` helper, and export the module from the agent package root. Do not duplicate the generic contract and do not adopt OTel or another external semantic convention. @@ -3389,8 +3388,6 @@ I0, I1, and I2 may proceed independently. I3 → I4 → I5 is serial and begins These packages all own `packages/agent/src/agent-loop.ts` and therefore merge strictly L1 → L2 → L3. Existing `agent-loop` and `agent` tests pass unchanged after each package. -**In progress and reserved: L1–L3 by @badlogic.** Other agents must not pick an L package while this ownership marker remains. - - [ ] **L1 — extract assistant streaming.** Dependencies: I0. - Add `streamAssistant()` and `StreamAssistantConfig`, including explicit telemetry context; route the compatibility loop's request path through it without changing events or results. - Acceptance: focused stream tests cover settled-result narrowing (a final `pending` value is a defect), plus unchanged existing loop tests. @@ -3485,8 +3482,8 @@ For a fresh implementation session, in this order. This document wins over older 5. `packages/session-backends/sqlite-node/src/sqlite/repo.ts` — v4 SQLite repository, leases, and forks. 6. `packages/session-backends/sqlite-node/src/sqlite/storage/branch-entries.ts` — branch cache queries. 7. `packages/agent/src/harness/agent-harness.ts` — v2 public API scaffold. -8. `packages/ai/src/telemetry.ts` — canonical telemetry contract, schema machinery, and pi-ai schema after I0. -9. `packages/agent/src/harness/telemetry.ts` — harness schema and typed helper after I0. +8. `packages/ai/src/telemetry.ts` — canonical telemetry contract, schema machinery, and pi-ai schema. +9. `packages/agent/src/harness/telemetry.ts` — harness schema and typed helper. 10. `packages/agent/src/agent-loop.ts` — the loop to split into the section 14 building blocks. 11. `packages/agent/src/agent.ts` — queues, continuation, abort, settlement to preserve in spirit. 12. `packages/agent/src/harness/messages.ts` — message conversion (`toProviderMessages` default). diff --git a/packages/agent/docs/telemetry-schema.md b/packages/agent/docs/telemetry-schema.md new file mode 100644 index 000000000..be32446bd --- /dev/null +++ b/packages/agent/docs/telemetry-schema.md @@ -0,0 +1,331 @@ +# Pi Agent Harness Telemetry Schema + + + +Schema version: 1 + +## `pi.harness.run` + +One admitted in-process run invocation + +- Parents: root or caller-owned external span +- Default status: `ok` +- Error when: The run fails or throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.session.id` | `string` | yes | | high cardinality | Session id | +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.operation.recovery` | `boolean` | yes | | | Whether this invocation resumes durable work | +| `pi.operation.kind` | `string` | yes | run | | Run operation kind | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.operation.outcome` | `string` | completed, aborted, failed, suspended | | Run invocation outcome | +| `pi.error.code` | `string` | | low cardinality | Stable operation error code | +| `pi.error.type` | `string` | | low cardinality | Low-cardinality operation error class | + +### Events + +No declared span events. + +## `pi.harness.compaction` + +One admitted in-process manual compaction invocation + +- Parents: root or caller-owned external span +- Default status: `ok` +- Error when: The compaction fails or throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.session.id` | `string` | yes | | high cardinality | Session id | +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.operation.recovery` | `boolean` | yes | | | Whether this invocation resumes durable work | +| `pi.operation.kind` | `string` | yes | compaction | | Compaction operation kind | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.operation.outcome` | `string` | completed, declined, aborted, failed | | Compaction invocation outcome | +| `pi.error.code` | `string` | | low cardinality | Stable operation error code | +| `pi.error.type` | `string` | | low cardinality | Low-cardinality operation error class | + +### Events + +No declared span events. + +## `pi.harness.navigation` + +One admitted in-process navigation invocation + +- Parents: root or caller-owned external span +- Default status: `ok` +- Error when: The navigation fails or throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.session.id` | `string` | yes | | high cardinality | Session id | +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.operation.recovery` | `boolean` | yes | | | Whether this invocation resumes durable work | +| `pi.operation.kind` | `string` | yes | navigation | | Navigation operation kind | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.operation.outcome` | `string` | completed, declined, aborted, failed | | Navigation invocation outcome | +| `pi.error.code` | `string` | | low cardinality | Stable operation error code | +| `pi.error.type` | `string` | | low cardinality | Low-cardinality operation error class | + +### Events + +No declared span events. + +## `pi.harness.checkpoint` + +One run checkpoint + +- Parents: `pi.harness.run` +- Default status: `ok` +- Error when: Checkpoint work throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.checkpoint.kind` | `string` | yes | normal, failure_drain, abort_reconcile | | Checkpoint purpose | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| _none_ | | | | | + +### Events + +No declared span events. + +## `pi.harness.turn` + +One assistant response and its tool batch + +- Parents: `pi.harness.run` +- Default status: `ok` +- Error when: Turn work throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.turn.id` | `string` | yes | | high cardinality | Invocation-local turn id | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| _none_ | | | | | + +### Events + +No declared span events. + +## `pi.harness.step` + +One durable retry attempt + +- Parents: `pi.harness.turn`, `pi.harness.checkpoint`, `pi.harness.compaction`, `pi.harness.navigation` +- Default status: `ok` +- Error when: The attempt retries, fails, or throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.step.kind` | `string` | yes | assistant, compaction, branch_summary | | Retryable step kind | +| `pi.step.attempt` | `number` | yes | | | One-based durable attempt number | +| `pi.compaction.reason` | `string` | no | manual, threshold, overflow | | Compaction trigger | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.step.outcome` | `string` | succeeded, retry, failed, aborted, deferred, overflow | | Attempt outcome | + +### Events + +No declared span events. + +## `pi.harness.tool` + +One raw phase-2 tool execution + +- Parents: `pi.harness.turn`, `pi.harness.run` +- Default status: `ok` +- Error when: Raw phase-2 execution returns an error + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.turn.id` | `string` | no | | high cardinality | Invocation-local live turn id | +| `pi.tool.name` | `string` | yes | | | Tool name | +| `pi.tool.call_id` | `string` | yes | | high cardinality | Tool call id | +| `pi.tool.replay` | `string` | yes | never, safe | | Declared replay policy | +| `pi.tool.recovery` | `boolean` | yes | | | Whether this is recovery execution | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.tool.is_error` | `boolean` | | | Whether raw phase-2 execution returned an error | + +### Events + +No declared span events. + +## `pi.harness.hook` + +One registered hook handler invocation + +- Parents: root or any caller span +- Default status: `ok` +- Error when: The handler throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | no | | high cardinality | Durable operation id when accepted | +| `pi.hook.name` | `string` | yes | before_run, before_resume, before_run_end, transform_context, before_request, before_payload, after_response, before_tool, after_tool, before_compaction, before_navigation | | Hook name | +| `pi.hook.registration_id` | `string` | no | | | Stable hook registration id | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.hook.outcome` | `string` | completed, skipped, blocked, failed | | Handler outcome | + +### Events + +No declared span events. + +## `pi.harness.sleep` + +One retry delay + +- Parents: `pi.harness.step`, `pi.harness.run` +- Default status: `ok` +- Error when: Sleep work throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.operation.id` | `string` | yes | | high cardinality | Durable operation id | +| `pi.sleep.delay_ms` | `number` | yes | | | Requested delay in milliseconds | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.sleep.outcome` | `string` | elapsed, aborted | | Delay outcome | + +### Events + +No declared span events. + +## `pi.harness.event_handler` + +One passive event listener invocation + +- Parents: root or any caller span +- Default status: `ok` +- Error when: The listener throws + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.event.type` | `string` | yes | run_start, run_resume, run_suspend, run_abort, run_end, fault, handler_error, turn_start, turn_end, retry_scheduled, retry_start, retry_end, message_start, message_update, message_end, tool_start, tool_update, tool_end, entry_added, write_pending, queue_update, fact_update, config_update, compaction_start, compaction_end, navigation_start, navigation_end, lane_created, usage | low cardinality | Delivered harness event type | +| `pi.lane.name` | `string` | no | | high cardinality | Lane name for lane-scoped events | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| _none_ | | | | | + +### Events + +No declared span events. + +## `pi.session.write` + +One committed session mutation + +- Parents: root or any caller span +- Default status: `ok` +- Error when: Storage rejects the mutation + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.lane.name` | `string` | yes | | high cardinality | Lane name | +| `pi.operation.id` | `string` | no | | high cardinality | Durable operation id when accepted | +| `pi.session.mutation` | `string` | yes | entry, record, lane, fact | | Session mutation kind | +| `pi.session.item_type` | `string` | no | | | Entry, record, lane, or fact subtype | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.session.seq` | `number` | | | Committed session sequence when exposed | + +### Events + +No declared span events. diff --git a/packages/agent/package.json b/packages/agent/package.json index d588fab12..f950da6a9 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -26,6 +26,8 @@ ], "scripts": { "clean": "shx rm -rf dist", + "generate-telemetry-docs": "node scripts/generate-telemetry-docs.ts", + "check:telemetry-docs": "node scripts/generate-telemetry-docs.ts --check", "build": "tsgo -p tsconfig.build.json", "test": "vitest --run", "test:harness": "vitest --run --config vitest.harness.config.ts", diff --git a/packages/agent/scripts/generate-telemetry-docs.ts b/packages/agent/scripts/generate-telemetry-docs.ts new file mode 100644 index 000000000..c048d653c --- /dev/null +++ b/packages/agent/scripts/generate-telemetry-docs.ts @@ -0,0 +1,13 @@ +import { resolve } from "node:path"; +import { pathToFileURL } from "node:url"; +import { generateTelemetryDocs } from "../../ai/scripts/generate-telemetry-docs.ts"; +import { HARNESS_TELEMETRY_SCHEMA } from "../src/harness/telemetry.ts"; + +if (process.argv[1] && pathToFileURL(resolve(process.argv[1])).href === import.meta.url) { + generateTelemetryDocs( + HARNESS_TELEMETRY_SCHEMA, + "Pi Agent Harness Telemetry Schema", + resolve(import.meta.dirname, "../docs/telemetry-schema.md"), + process.argv.includes("--check"), + ); +} diff --git a/packages/agent/src/harness/telemetry.ts b/packages/agent/src/harness/telemetry.ts new file mode 100644 index 000000000..2695c4aa2 --- /dev/null +++ b/packages/agent/src/harness/telemetry.ts @@ -0,0 +1,506 @@ +import type { + ExactTelemetryAttributes, + SchemaTelemetrySpan, + TelemetryContext, + TelemetrySchemaDefinition, + TelemetrySchemaSpanEndAttributes, + TelemetrySchemaSpanEventAttributes, + TelemetrySchemaSpanEventName, + TelemetrySchemaSpanName, + TelemetrySchemaSpanStartAttributes, + TelemetrySchemaSpanUnion, + TelemetrySpan, +} from "@earendil-works/pi-ai"; + +export type { + AttributeValue, + ExactTelemetryAttributes, + SchemaTelemetrySpan, + SpanAttributes, + SpanOptions, + SpanStatus, + TelemetryAttributeDefinition, + TelemetryAttributeMetadata, + TelemetryAttributeType, + TelemetryContext, + TelemetryEventAttributeDefinition, + TelemetryEventDefinition, + TelemetryParentDefinition, + TelemetrySchemaDefinition, + TelemetrySchemaSpanEndAttributes, + TelemetrySchemaSpanEventAttributes, + TelemetrySchemaSpanEventName, + TelemetrySchemaSpanName, + TelemetrySchemaSpanStartAttributes, + TelemetrySchemaSpanUnion, + TelemetrySpan, + TelemetrySpanDefinition, + TelemetryStartAttributeDefinition, +} from "@earendil-works/pi-ai"; + +const HOOK_NAMES = [ + "before_run", + "before_resume", + "before_run_end", + "transform_context", + "before_request", + "before_payload", + "after_response", + "before_tool", + "after_tool", + "before_compaction", + "before_navigation", +] as const; + +const EVENT_TYPES = [ + "run_start", + "run_resume", + "run_suspend", + "run_abort", + "run_end", + "fault", + "handler_error", + "turn_start", + "turn_end", + "retry_scheduled", + "retry_start", + "retry_end", + "message_start", + "message_update", + "message_end", + "tool_start", + "tool_update", + "tool_end", + "entry_added", + "write_pending", + "queue_update", + "fact_update", + "config_update", + "compaction_start", + "compaction_end", + "navigation_start", + "navigation_end", + "lane_created", + "usage", +] as const; + +const operationStartAttributes = { + "pi.session.id": { + type: "string", + required: true, + cardinality: "high", + description: "Session id", + }, + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.operation.recovery": { + type: "boolean", + required: true, + description: "Whether this invocation resumes durable work", + }, +} as const; + +const operationErrorAttributes = { + "pi.error.code": { + type: "string", + cardinality: "low", + description: "Stable operation error code", + }, + "pi.error.type": { + type: "string", + cardinality: "low", + description: "Low-cardinality operation error class", + }, +} as const; + +export const HARNESS_TELEMETRY_SCHEMA = { + version: 1, + spans: { + "pi.harness.run": { + description: "One admitted in-process run invocation", + parents: { kind: "root_or_external" }, + startAttributes: { + ...operationStartAttributes, + "pi.operation.kind": { + type: "string", + required: true, + values: ["run"], + description: "Run operation kind", + }, + }, + endAttributes: { + "pi.operation.outcome": { + type: "string", + values: ["completed", "aborted", "failed", "suspended"], + description: "Run invocation outcome", + }, + ...operationErrorAttributes, + }, + status: { default: "ok", errorWhen: "The run fails or throws" }, + }, + "pi.harness.compaction": { + description: "One admitted in-process manual compaction invocation", + parents: { kind: "root_or_external" }, + startAttributes: { + ...operationStartAttributes, + "pi.operation.kind": { + type: "string", + required: true, + values: ["compaction"], + description: "Compaction operation kind", + }, + }, + endAttributes: { + "pi.operation.outcome": { + type: "string", + values: ["completed", "declined", "aborted", "failed"], + description: "Compaction invocation outcome", + }, + ...operationErrorAttributes, + }, + status: { default: "ok", errorWhen: "The compaction fails or throws" }, + }, + "pi.harness.navigation": { + description: "One admitted in-process navigation invocation", + parents: { kind: "root_or_external" }, + startAttributes: { + ...operationStartAttributes, + "pi.operation.kind": { + type: "string", + required: true, + values: ["navigation"], + description: "Navigation operation kind", + }, + }, + endAttributes: { + "pi.operation.outcome": { + type: "string", + values: ["completed", "declined", "aborted", "failed"], + description: "Navigation invocation outcome", + }, + ...operationErrorAttributes, + }, + status: { default: "ok", errorWhen: "The navigation fails or throws" }, + }, + "pi.harness.checkpoint": { + description: "One run checkpoint", + parents: { kind: "spans", spans: ["pi.harness.run"] }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.checkpoint.kind": { + type: "string", + required: true, + values: ["normal", "failure_drain", "abort_reconcile"], + description: "Checkpoint purpose", + }, + }, + endAttributes: {}, + status: { default: "ok", errorWhen: "Checkpoint work throws" }, + }, + "pi.harness.turn": { + description: "One assistant response and its tool batch", + parents: { kind: "spans", spans: ["pi.harness.run"] }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.turn.id": { + type: "string", + required: true, + cardinality: "high", + description: "Invocation-local turn id", + }, + }, + endAttributes: {}, + status: { default: "ok", errorWhen: "Turn work throws" }, + }, + "pi.harness.step": { + description: "One durable retry attempt", + parents: { + kind: "spans", + spans: ["pi.harness.turn", "pi.harness.checkpoint", "pi.harness.compaction", "pi.harness.navigation"], + }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.step.kind": { + type: "string", + required: true, + values: ["assistant", "compaction", "branch_summary"], + description: "Retryable step kind", + }, + "pi.step.attempt": { + type: "number", + required: true, + description: "One-based durable attempt number", + }, + "pi.compaction.reason": { + type: "string", + required: false, + values: ["manual", "threshold", "overflow"], + description: "Compaction trigger", + }, + }, + endAttributes: { + "pi.step.outcome": { + type: "string", + values: ["succeeded", "retry", "failed", "aborted", "deferred", "overflow"], + description: "Attempt outcome", + }, + }, + status: { default: "ok", errorWhen: "The attempt retries, fails, or throws" }, + }, + "pi.harness.tool": { + description: "One raw phase-2 tool execution", + parents: { kind: "spans", spans: ["pi.harness.turn", "pi.harness.run"] }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.turn.id": { + type: "string", + required: false, + cardinality: "high", + description: "Invocation-local live turn id", + }, + "pi.tool.name": { + type: "string", + required: true, + description: "Tool name", + }, + "pi.tool.call_id": { + type: "string", + required: true, + cardinality: "high", + description: "Tool call id", + }, + "pi.tool.replay": { + type: "string", + required: true, + values: ["never", "safe"], + description: "Declared replay policy", + }, + "pi.tool.recovery": { + type: "boolean", + required: true, + description: "Whether this is recovery execution", + }, + }, + endAttributes: { + "pi.tool.is_error": { + type: "boolean", + description: "Whether raw phase-2 execution returned an error", + }, + }, + status: { default: "ok", errorWhen: "Raw phase-2 execution returns an error" }, + }, + "pi.harness.hook": { + description: "One registered hook handler invocation", + parents: { kind: "any" }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: false, + cardinality: "high", + description: "Durable operation id when accepted", + }, + "pi.hook.name": { + type: "string", + required: true, + values: HOOK_NAMES, + description: "Hook name", + }, + "pi.hook.registration_id": { + type: "string", + required: false, + description: "Stable hook registration id", + }, + }, + endAttributes: { + "pi.hook.outcome": { + type: "string", + values: ["completed", "skipped", "blocked", "failed"], + description: "Handler outcome", + }, + }, + status: { default: "ok", errorWhen: "The handler throws" }, + }, + "pi.harness.sleep": { + description: "One retry delay", + parents: { kind: "spans", spans: ["pi.harness.step", "pi.harness.run"] }, + startAttributes: { + "pi.operation.id": { + type: "string", + required: true, + cardinality: "high", + description: "Durable operation id", + }, + "pi.sleep.delay_ms": { + type: "number", + required: true, + description: "Requested delay in milliseconds", + }, + }, + endAttributes: { + "pi.sleep.outcome": { + type: "string", + values: ["elapsed", "aborted"], + description: "Delay outcome", + }, + }, + status: { default: "ok", errorWhen: "Sleep work throws" }, + }, + "pi.harness.event_handler": { + description: "One passive event listener invocation", + parents: { kind: "any" }, + startAttributes: { + "pi.event.type": { + type: "string", + required: true, + cardinality: "low", + values: EVENT_TYPES, + description: "Delivered harness event type", + }, + "pi.lane.name": { + type: "string", + required: false, + cardinality: "high", + description: "Lane name for lane-scoped events", + }, + }, + endAttributes: {}, + status: { default: "ok", errorWhen: "The listener throws" }, + }, + "pi.session.write": { + description: "One committed session mutation", + parents: { kind: "any" }, + startAttributes: { + "pi.lane.name": { + type: "string", + required: true, + cardinality: "high", + description: "Lane name", + }, + "pi.operation.id": { + type: "string", + required: false, + cardinality: "high", + description: "Durable operation id when accepted", + }, + "pi.session.mutation": { + type: "string", + required: true, + values: ["entry", "record", "lane", "fact"], + description: "Session mutation kind", + }, + "pi.session.item_type": { + type: "string", + required: false, + description: "Entry, record, lane, or fact subtype", + }, + }, + endAttributes: { + "pi.session.seq": { + type: "number", + description: "Committed session sequence when exposed", + }, + }, + status: { default: "ok", errorWhen: "Storage rejects the mutation" }, + }, + }, +} as const satisfies TelemetrySchemaDefinition; + +export type HarnessSpanName = TelemetrySchemaSpanName; +export type HarnessSpanStartAttributes = TelemetrySchemaSpanStartAttributes< + typeof HARNESS_TELEMETRY_SCHEMA, + Name +>; +export type HarnessSpanEndAttributes = TelemetrySchemaSpanEndAttributes< + typeof HARNESS_TELEMETRY_SCHEMA, + Name +>; +export type HarnessSpanAttributes = HarnessSpanStartAttributes & + HarnessSpanEndAttributes; +export type HarnessSpanEventName = TelemetrySchemaSpanEventName< + typeof HARNESS_TELEMETRY_SCHEMA, + Name +>; +export type HarnessSpanEventAttributes< + Name extends HarnessSpanName, + EventName extends HarnessSpanEventName, +> = TelemetrySchemaSpanEventAttributes; +export type HarnessTelemetrySpan = SchemaTelemetrySpan< + typeof HARNESS_TELEMETRY_SCHEMA, + Name +>; +export type HarnessSpan = TelemetrySchemaSpanUnion; + +export function startHarnessSpan< + Name extends HarnessSpanName, + const Attributes extends HarnessSpanStartAttributes, + Result, +>( + telemetryContext: TelemetryContext, + name: Name, + attributes: ExactTelemetryAttributes, Attributes>, + callback: (span: HarnessTelemetrySpan) => Result | Promise, +): Promise { + return telemetryContext.startSpan({ name, attributes }, (span: TelemetrySpan) => + callback(span as HarnessTelemetrySpan), + ); +} diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 216bda49b..4882547a4 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -1,5 +1,50 @@ // Core Agent -export { uuidv7 } from "@earendil-works/pi-ai"; + +export type { + AiSpan, + AiSpanAttributes, + AiSpanEndAttributes, + AiSpanEventAttributes, + AiSpanEventName, + AiSpanName, + AiSpanStartAttributes, + AiTelemetrySpan, + AttributeValue, + ExactTelemetryAttributes, + InferEventAttributes, + InferOptionalAttributes, + InferRequiredAndOptionalAttributes, + InferStartAttributes, + SchemaTelemetrySpan, + SpanAttributes, + SpanAttributes as TelemetrySpanAttributes, + SpanOptions, + SpanStatus, + TelemetryAttributeDefinition, + TelemetryAttributeMetadata, + TelemetryAttributeType, + TelemetryContext, + TelemetryEventAttributeDefinition, + TelemetryEventDefinition, + TelemetryParentDefinition, + TelemetrySchemaDefinition, + TelemetrySchemaSpanEndAttributes, + TelemetrySchemaSpanEventAttributes, + TelemetrySchemaSpanEventName, + TelemetrySchemaSpanName, + TelemetrySchemaSpanStartAttributes, + TelemetrySchemaSpanUnion, + TelemetrySpan, + TelemetrySpanDefinition, + TelemetryStartAttributeDefinition, +} from "@earendil-works/pi-ai"; +export { + AI_TELEMETRY_SCHEMA, + defineTelemetrySchema, + NOOP_TELEMETRY_CONTEXT, + startAiSpan, + uuidv7, +} from "@earendil-works/pi-ai"; export * from "./agent.ts"; // Loop functions export * from "./agent-loop.ts"; @@ -41,6 +86,17 @@ export * from "./harness/session/index.ts"; export * from "./harness/session/search.ts"; export * from "./harness/skills.ts"; export * from "./harness/system-prompt.ts"; +export type { + HarnessSpan, + HarnessSpanAttributes, + HarnessSpanEndAttributes, + HarnessSpanEventAttributes, + HarnessSpanEventName, + HarnessSpanName, + HarnessSpanStartAttributes, + HarnessTelemetrySpan, +} from "./harness/telemetry.ts"; +export { HARNESS_TELEMETRY_SCHEMA, startHarnessSpan } from "./harness/telemetry.ts"; export * from "./harness/tools/index.ts"; export { type AgentHarnessResources, diff --git a/packages/agent/test/harness/telemetry.test.ts b/packages/agent/test/harness/telemetry.test.ts new file mode 100644 index 000000000..23c9e2e9f --- /dev/null +++ b/packages/agent/test/harness/telemetry.test.ts @@ -0,0 +1,103 @@ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { NOOP_TELEMETRY_CONTEXT, type TelemetryContext } from "@earendil-works/pi-ai"; +import { describe, expect, expectTypeOf, it } from "vitest"; +import { renderTelemetrySchemaMarkdown } from "../../../ai/scripts/generate-telemetry-docs.ts"; +import { + HARNESS_TELEMETRY_SCHEMA, + type HarnessSpanEndAttributes, + type HarnessSpanStartAttributes, + startHarnessSpan, +} from "../../src/harness/telemetry.ts"; + +describe("HARNESS_TELEMETRY_SCHEMA", () => { + it("is serializable, exhaustive, and generates the checked-in reference", () => { + expect(Object.keys(HARNESS_TELEMETRY_SCHEMA.spans)).toEqual([ + "pi.harness.run", + "pi.harness.compaction", + "pi.harness.navigation", + "pi.harness.checkpoint", + "pi.harness.turn", + "pi.harness.step", + "pi.harness.tool", + "pi.harness.hook", + "pi.harness.sleep", + "pi.harness.event_handler", + "pi.session.write", + ]); + expect(() => JSON.stringify(HARNESS_TELEMETRY_SCHEMA)).not.toThrow(); + const expected = renderTelemetrySchemaMarkdown(HARNESS_TELEMETRY_SCHEMA, "Pi Agent Harness Telemetry Schema"); + const actual = readFileSync(resolve(import.meta.dirname, "../../docs/telemetry-schema.md"), "utf8"); + expect(actual).toBe(expected); + }); + + it("infers per-span literals and optional completion enrichment", async () => { + type RunStart = HarnessSpanStartAttributes<"pi.harness.run">; + type RunEnd = HarnessSpanEndAttributes<"pi.harness.run">; + expectTypeOf().toEqualTypeOf<"run">(); + expectTypeOf().toEqualTypeOf< + "completed" | "aborted" | "failed" | "suspended" | undefined + >(); + + const telemetryContext: TelemetryContext = NOOP_TELEMETRY_CONTEXT; + await startHarnessSpan( + telemetryContext, + "pi.harness.run", + { + "pi.session.id": "session", + "pi.lane.name": "main", + "pi.operation.id": "operation", + "pi.operation.kind": "run", + "pi.operation.recovery": false, + }, + (span) => { + span.setAttributes({ "pi.operation.outcome": "completed" }); + span.setAttributes({}); + // @ts-expect-error the harness schema declares no span events + span.addEvent("result"); + }, + ); + + const compileTimeFailures = () => { + const extraRunAttributes = { + "pi.session.id": "session", + "pi.lane.name": "main", + "pi.operation.id": "operation", + "pi.operation.kind": "run", + "pi.operation.recovery": false, + "pi.unknown": true, + } as const; + // @ts-expect-error variables with unknown attributes are rejected + void startHarnessSpan(telemetryContext, "pi.harness.run", extraRunAttributes, () => {}); + void startHarnessSpan( + telemetryContext, + "pi.harness.checkpoint", + { + "pi.lane.name": "main", + "pi.operation.id": "operation", + "pi.checkpoint.kind": "normal", + }, + (span) => { + // @ts-expect-error empty end schemas reject every attribute + span.setAttributes({ "pi.unknown": true }); + }, + ); + void startHarnessSpan( + telemetryContext, + "pi.harness.run", + { + "pi.session.id": "session", + "pi.lane.name": "main", + "pi.operation.id": "operation", + // @ts-expect-error run spans accept only the run operation kind + "pi.operation.kind": "navigation", + "pi.operation.recovery": false, + }, + () => {}, + ); + // @ts-expect-error missing required run start attributes + void startHarnessSpan(telemetryContext, "pi.harness.run", {}, () => {}); + }; + expectTypeOf(compileTimeFailures).toBeFunction(); + }); +}); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 5ce5acbc0..064e5190a 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -68,6 +68,7 @@ ### Added +- Added the callback-based telemetry context contract, typed pi-ai telemetry schema, generated schema reference, and explicit telemetry context propagation across stream, deferred, and image request options. - Added deferred provider request contracts, durable response handles, authenticated fetch/cancel dispatch, and faux-provider support for pending, ready, failed, and cancelled responses ([#7339](https://github.com/earendil-works/pi/pull/7339) by [@davidbrai](https://github.com/davidbrai)). - Added Baseten as a built-in OpenAI-compatible provider with models.dev catalog generation and native `chat_template_args` reasoning controls. diff --git a/packages/ai/docs/telemetry-schema.md b/packages/ai/docs/telemetry-schema.md new file mode 100644 index 000000000..9d582e991 --- /dev/null +++ b/packages/ai/docs/telemetry-schema.md @@ -0,0 +1,49 @@ +# Pi AI Telemetry Schema + + + +Schema version: 1 + +## `pi.ai.request` + +One logical request to an AI provider + +- Parents: root or any caller span +- Default status: `ok` +- Error when: The operation throws or returns an error result + +### Start attributes + +| Name | Type | Required | Values | Notes | Description | +|---|---|---:|---|---|---| +| `pi.ai.operation` | `string` | yes | stream, fetch_deferred, cancel_deferred, generate_images | | Logical provider operation | +| `pi.ai.provider` | `string` | yes | | | Selected provider id | +| `pi.ai.model` | `string` | yes | | | Requested model id | +| `pi.ai.api` | `string` | yes | | | Provider API id | +| `pi.ai.streaming` | `boolean` | yes | | | Whether this operation returns a stream | +| `pi.ai.deferred` | `boolean` | no | | | Whether the operation requests or participates in deferred execution | + +### End attributes + +All end attributes are optional completion enrichment. + +| Name | Type | Values | Notes | Description | +|---|---|---|---|---| +| `pi.ai.response.model` | `string` | | | Concrete response model | +| `pi.ai.response.id` | `string` | | high cardinality | Provider response id | +| `pi.ai.response.stop_reason` | `string` | stop, length, tool_use, error, aborted, deferred | | Normalized terminal response reason | +| `pi.ai.http.status_code` | `number` | | | Final HTTP status | +| `pi.ai.usage.input_tokens` | `number` | | | Reported input tokens | +| `pi.ai.usage.output_tokens` | `number` | | | Reported output tokens | +| `pi.ai.usage.cache_read_tokens` | `number` | | | Reported cache-read tokens | +| `pi.ai.usage.cache_write_tokens` | `number` | | | Reported cache-write tokens | +| `pi.ai.usage.reasoning_tokens` | `number` | | | Reported reasoning tokens | +| `pi.ai.usage.total_tokens` | `number` | | | Reported total tokens | +| `pi.ai.usage.cost` | `number` | | | Reported total cost | +| `pi.ai.stream.chunk_count` | `number` | | | Streamed update chunk count | +| `pi.ai.stream.time_to_first_chunk_ms` | `number` | | | Elapsed milliseconds to first update chunk | +| `pi.ai.error.type` | `string` | | low cardinality | Provider or transport error class | + +### Events + +No declared span events. diff --git a/packages/ai/package.json b/packages/ai/package.json index a6773c7b1..cb38d2051 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -53,6 +53,8 @@ "hydrate-model-data": "node scripts/generate-models.ts --strict --data-only", "generate-model-catalog": "node scripts/generate-models.ts --strict --json-only --json-output ../../.artifacts/model-catalog", "generate-image-models": "node scripts/generate-image-models.ts --strict", + "generate-telemetry-docs": "node scripts/generate-telemetry-docs.ts", + "check:telemetry-docs": "node scripts/generate-telemetry-docs.ts --check", "check:model-data": "node scripts/check-model-data.ts", "build": "npm run generate-models && npm run build:offline", "build:offline": "npm run check:model-data && tsgo -p tsconfig.build.json && shx rm -rf dist/providers/data && shx cp -r src/providers/data dist/providers/data", diff --git a/packages/ai/scripts/generate-telemetry-docs.ts b/packages/ai/scripts/generate-telemetry-docs.ts new file mode 100644 index 000000000..decd96250 --- /dev/null +++ b/packages/ai/scripts/generate-telemetry-docs.ts @@ -0,0 +1,120 @@ +import { mkdirSync, readFileSync, writeFileSync } from "node:fs"; +import { dirname, resolve } from "node:path"; +import { pathToFileURL } from "node:url"; +import { AI_TELEMETRY_SCHEMA, type TelemetryAttributeDefinition, type TelemetrySchemaDefinition } from "../src/telemetry.ts"; + +function escapeCell(value: string): string { + return value.replaceAll("|", "\\|").replaceAll("\n", " "); +} + +function allowedValues(definition: TelemetryAttributeDefinition): string { + if ("values" in definition && definition.values) return definition.values.map(String).join(", "); + if ("elementValues" in definition && definition.elementValues) { + return `elements: ${definition.elementValues.map(String).join(", ")}`; + } + return ""; +} + +function attributeNotes(definition: TelemetryAttributeDefinition): string { + return [definition.cardinality ? `${definition.cardinality} cardinality` : "", definition.sensitive ? "sensitive" : ""] + .filter(Boolean) + .join(", "); +} + +function parentDescription(parent: TelemetrySchemaDefinition["spans"][string]["parents"]): string { + switch (parent.kind) { + case "any": + return "root or any caller span"; + case "root_or_external": + return "root or caller-owned external span"; + case "spans": + return parent.spans.map((span) => `\`${span}\``).join(", "); + } +} + +export function renderTelemetrySchemaMarkdown(schema: TelemetrySchemaDefinition, title: string): string { + const lines = [ + `# ${title}`, + "", + "", + "", + `Schema version: ${schema.version}`, + "", + ]; + + for (const [spanName, span] of Object.entries(schema.spans)) { + lines.push(`## \`${spanName}\``, "", span.description, ""); + lines.push(`- Parents: ${parentDescription(span.parents)}`); + lines.push(`- Default status: \`${span.status.default}\``); + lines.push(`- Error when: ${span.status.errorWhen}`, ""); + lines.push("### Start attributes", ""); + lines.push("| Name | Type | Required | Values | Notes | Description |"); + lines.push("|---|---|---:|---|---|---|"); + for (const [name, definition] of Object.entries(span.startAttributes)) { + lines.push( + `| \`${name}\` | \`${definition.type}\` | ${definition.required ? "yes" : "no"} | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`, + ); + } + if (Object.keys(span.startAttributes).length === 0) lines.push("| _none_ | | | | | |"); + lines.push("", "### End attributes", ""); + lines.push("All end attributes are optional completion enrichment.", ""); + lines.push("| Name | Type | Values | Notes | Description |"); + lines.push("|---|---|---|---|---|"); + for (const [name, definition] of Object.entries(span.endAttributes)) { + lines.push( + `| \`${name}\` | \`${definition.type}\` | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`, + ); + } + if (Object.keys(span.endAttributes).length === 0) lines.push("| _none_ | | | | |"); + lines.push("", "### Events", ""); + const events = Object.entries(span.events ?? {}); + if (events.length === 0) { + lines.push("No declared span events.", ""); + continue; + } + for (const [eventName, event] of events) { + lines.push(`#### \`${eventName}\``, "", event.description, ""); + lines.push("| Name | Type | Required | Values | Notes | Description |"); + lines.push("|---|---|---:|---|---|---|"); + for (const [name, definition] of Object.entries(event.attributes)) { + lines.push( + `| \`${name}\` | \`${definition.type}\` | ${definition.required ? "yes" : "no"} | ${escapeCell(allowedValues(definition))} | ${escapeCell(attributeNotes(definition))} | ${escapeCell(definition.description)} |`, + ); + } + if (Object.keys(event.attributes).length === 0) lines.push("| _none_ | | | | | |"); + lines.push(""); + } + } + + return `${lines.join("\n").trimEnd()}\n`; +} + +export function generateTelemetryDocs( + schema: TelemetrySchemaDefinition, + title: string, + outputPath: string, + check: boolean, +): void { + const expected = renderTelemetrySchemaMarkdown(schema, title); + if (check) { + let actual = ""; + try { + actual = readFileSync(outputPath, "utf8"); + } catch { + throw new Error(`${outputPath} is missing; run the telemetry documentation generator`); + } + if (actual !== expected) throw new Error(`${outputPath} is stale; run the telemetry documentation generator`); + return; + } + mkdirSync(dirname(outputPath), { recursive: true }); + writeFileSync(outputPath, expected); +} + +if (process.argv[1] && pathToFileURL(resolve(process.argv[1])).href === import.meta.url) { + generateTelemetryDocs( + AI_TELEMETRY_SCHEMA, + "Pi AI Telemetry Schema", + resolve(import.meta.dirname, "../docs/telemetry-schema.md"), + process.argv.includes("--check"), + ); +} diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index abbb25428..42499aa68 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -33,6 +33,7 @@ export function buildBaseOptions( samplingParams, maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens), signal: options?.signal, + telemetryContext: options?.telemetryContext, apiKey: apiKey || options?.apiKey, fetch: options?.fetch, transport: options?.transport, diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 4ff678102..63d507400 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -35,6 +35,7 @@ export * from "./models.ts"; export * from "./models-store.ts"; export * from "./providers/faux.ts"; export * from "./session-resources.ts"; +export * from "./telemetry.ts"; export * from "./types.ts"; export * from "./utils/diagnostics.ts"; export * from "./utils/event-stream.ts"; diff --git a/packages/ai/src/telemetry.ts b/packages/ai/src/telemetry.ts new file mode 100644 index 000000000..6bd4852e9 --- /dev/null +++ b/packages/ai/src/telemetry.ts @@ -0,0 +1,375 @@ +export type AttributeValue = string | number | boolean | readonly string[] | readonly number[] | readonly boolean[]; + +export interface SpanAttributes { + [name: string]: AttributeValue | undefined; +} + +export interface SpanOptions { + name: string; + attributes?: SpanAttributes; +} + +export type SpanStatus = { status: "ok" } | { status: "error"; error?: { name: string; message: string } }; + +export interface TelemetryContext { + startSpan(options: SpanOptions, callback: (span: TelemetrySpan) => T | Promise): Promise; +} + +export interface TelemetrySpan extends TelemetryContext { + addEvent(name: string, attributes?: SpanAttributes): void; + setAttributes(attributes: SpanAttributes): void; + setStatus(status: SpanStatus): void; +} + +function startNoopSpan(_options: SpanOptions, callback: (span: TelemetrySpan) => T | Promise): Promise { + try { + return Promise.resolve(callback(noopTelemetrySpan)); + } catch (error) { + return Promise.reject(error); + } +} + +const noopTelemetrySpan: TelemetrySpan = { + startSpan: startNoopSpan, + addEvent: () => {}, + setAttributes: () => {}, + setStatus: () => {}, +}; +Object.freeze(noopTelemetrySpan); + +/** Shared telemetry context used when an application does not provide one. */ +export const NOOP_TELEMETRY_CONTEXT: TelemetryContext = noopTelemetrySpan; + +export type TelemetryAttributeType = "string" | "number" | "boolean" | "string[]" | "number[]" | "boolean[]"; + +export interface TelemetryAttributeMetadata { + description: string; + sensitive?: boolean; + cardinality?: "low" | "high"; +} + +export type TelemetryAttributeDefinition = TelemetryAttributeMetadata & + ( + | { type: "string"; values?: readonly string[]; examples?: readonly string[] } + | { type: "number"; values?: readonly number[]; examples?: readonly number[] } + | { type: "boolean"; values?: readonly boolean[]; examples?: readonly boolean[] } + | { type: "string[]"; elementValues?: readonly string[]; examples?: readonly (readonly string[])[] } + | { type: "number[]"; elementValues?: readonly number[]; examples?: readonly (readonly number[])[] } + | { type: "boolean[]"; elementValues?: readonly boolean[]; examples?: readonly (readonly boolean[])[] } + ); + +export type TelemetryStartAttributeDefinition = TelemetryAttributeDefinition & { required: boolean }; +export type TelemetryEventAttributeDefinition = TelemetryAttributeDefinition & { required: boolean }; + +export interface TelemetryEventDefinition { + description: string; + attributes: Record; +} + +export type TelemetryParentDefinition = + | { kind: "any" } + | { kind: "root_or_external" } + | { kind: "spans"; spans: readonly string[] }; + +export interface TelemetrySpanDefinition { + description: string; + parents: TelemetryParentDefinition; + startAttributes: Record; + endAttributes: Record; + events?: Record; + status: { default: "ok"; errorWhen: string }; +} + +export interface TelemetrySchemaDefinition { + version: number; + spans: Record; +} + +/** Typed identity helper for serializable telemetry schema data. */ +export function defineTelemetrySchema(schema: T): T { + return schema; +} + +type AttributeDefinitionValue = Definition extends { + type: "string"; + values: readonly (infer Value extends string)[]; +} + ? Value + : Definition extends { type: "string" } + ? string + : Definition extends { type: "number"; values: readonly (infer Value extends number)[] } + ? Value + : Definition extends { type: "number" } + ? number + : Definition extends { type: "boolean"; values: readonly (infer Value extends boolean)[] } + ? Value + : Definition extends { type: "boolean" } + ? boolean + : Definition extends { + type: "string[]"; + elementValues: readonly (infer Value extends string)[]; + } + ? readonly Value[] + : Definition extends { type: "string[]" } + ? readonly string[] + : Definition extends { + type: "number[]"; + elementValues: readonly (infer Value extends number)[]; + } + ? readonly Value[] + : Definition extends { type: "number[]" } + ? readonly number[] + : Definition extends { + type: "boolean[]"; + elementValues: readonly (infer Value extends boolean)[]; + } + ? readonly Value[] + : readonly boolean[]; + +type RequiredAttributeNames< + Definitions extends Record, +> = { + [Name in keyof Definitions]-?: Definitions[Name]["required"] extends true ? Name : never; +}[keyof Definitions]; + +type OptionalAttributeNames< + Definitions extends Record, +> = Exclude>; + +export type InferRequiredAndOptionalAttributes< + Definitions extends Record, +> = keyof Definitions extends never + ? Record + : { + [Name in RequiredAttributeNames]: AttributeDefinitionValue; + } & { + [Name in OptionalAttributeNames]?: AttributeDefinitionValue; + }; + +export type InferStartAttributes> = + InferRequiredAndOptionalAttributes; + +export type InferOptionalAttributes> = + keyof Definitions extends never + ? Record + : { [Name in keyof Definitions]?: AttributeDefinitionValue }; + +export type ExactTelemetryAttributes = Actual & + Record, never>; + +export type InferEventAttributes> = + InferRequiredAndOptionalAttributes; + +export type TelemetrySchemaSpanName = keyof Schema["spans"] & string; + +type SchemaSpan< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = Schema["spans"][Name]; + +export type TelemetrySchemaSpanStartAttributes< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = SchemaSpan["startAttributes"] extends infer Definitions extends Record< + string, + TelemetryStartAttributeDefinition +> + ? InferStartAttributes + : never; + +export type TelemetrySchemaSpanEndAttributes< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = SchemaSpan["endAttributes"] extends infer Definitions extends Record< + string, + TelemetryAttributeDefinition +> + ? InferOptionalAttributes + : never; + +type SchemaSpanEvents< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = SchemaSpan extends { events: infer Events extends Record } + ? Events + : Record; + +export type TelemetrySchemaSpanEventName< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = keyof SchemaSpanEvents & string; + +type SchemaSpanEvent< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, + EventName extends TelemetrySchemaSpanEventName, +> = SchemaSpanEvents extends infer Events + ? EventName extends keyof Events + ? Events[EventName] + : never + : never; + +export type TelemetrySchemaSpanEventAttributes< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, + EventName extends TelemetrySchemaSpanEventName, +> = SchemaSpanEvent extends { + attributes: infer Definitions extends Record; +} + ? InferEventAttributes + : never; + +type SchemaSpanEventAttributeDefinitions< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, + EventName extends TelemetrySchemaSpanEventName, +> = SchemaSpanEvent extends { + attributes: infer Definitions extends Record; +} + ? Definitions + : Record; + +type EventArguments< + Definitions extends Record, + Attributes extends InferEventAttributes, +> = [RequiredAttributeNames] extends [never] + ? [attributes?: ExactTelemetryAttributes, Attributes>] + : [attributes: ExactTelemetryAttributes, Attributes>]; + +export type SchemaTelemetrySpan< + Schema extends TelemetrySchemaDefinition, + Name extends TelemetrySchemaSpanName, +> = Omit & { + addEvent< + EventName extends TelemetrySchemaSpanEventName, + const Attributes extends InferEventAttributes< + SchemaSpanEventAttributeDefinitions + > = InferEventAttributes>, + >( + name: EventName, + ...args: EventArguments, Attributes> + ): void; + setAttributes>( + attributes: ExactTelemetryAttributes, Attributes>, + ): void; +}; + +export type TelemetrySchemaSpanUnion = { + [Name in TelemetrySchemaSpanName]: { + name: Name; + startAttributes: TelemetrySchemaSpanStartAttributes; + endAttributes: TelemetrySchemaSpanEndAttributes; + events: { + [EventName in TelemetrySchemaSpanEventName]: TelemetrySchemaSpanEventAttributes< + Schema, + Name, + EventName + >; + }; + }; +}[TelemetrySchemaSpanName]; + +export const AI_TELEMETRY_SCHEMA = defineTelemetrySchema({ + version: 1, + spans: { + "pi.ai.request": { + description: "One logical request to an AI provider", + parents: { kind: "any" }, + startAttributes: { + "pi.ai.operation": { + type: "string", + required: true, + values: ["stream", "fetch_deferred", "cancel_deferred", "generate_images"], + description: "Logical provider operation", + }, + "pi.ai.provider": { + type: "string", + required: true, + description: "Selected provider id", + }, + "pi.ai.model": { + type: "string", + required: true, + description: "Requested model id", + }, + "pi.ai.api": { + type: "string", + required: true, + description: "Provider API id", + }, + "pi.ai.streaming": { + type: "boolean", + required: true, + description: "Whether this operation returns a stream", + }, + "pi.ai.deferred": { + type: "boolean", + required: false, + description: "Whether the operation requests or participates in deferred execution", + }, + }, + endAttributes: { + "pi.ai.response.model": { type: "string", description: "Concrete response model" }, + "pi.ai.response.id": { + type: "string", + cardinality: "high", + description: "Provider response id", + }, + "pi.ai.response.stop_reason": { + type: "string", + values: ["stop", "length", "tool_use", "error", "aborted", "deferred"], + description: "Normalized terminal response reason", + }, + "pi.ai.http.status_code": { type: "number", description: "Final HTTP status" }, + "pi.ai.usage.input_tokens": { type: "number", description: "Reported input tokens" }, + "pi.ai.usage.output_tokens": { type: "number", description: "Reported output tokens" }, + "pi.ai.usage.cache_read_tokens": { type: "number", description: "Reported cache-read tokens" }, + "pi.ai.usage.cache_write_tokens": { + type: "number", + description: "Reported cache-write tokens", + }, + "pi.ai.usage.reasoning_tokens": { type: "number", description: "Reported reasoning tokens" }, + "pi.ai.usage.total_tokens": { type: "number", description: "Reported total tokens" }, + "pi.ai.usage.cost": { type: "number", description: "Reported total cost" }, + "pi.ai.stream.chunk_count": { type: "number", description: "Streamed update chunk count" }, + "pi.ai.stream.time_to_first_chunk_ms": { + type: "number", + description: "Elapsed milliseconds to first update chunk", + }, + "pi.ai.error.type": { + type: "string", + cardinality: "low", + description: "Provider or transport error class", + }, + }, + status: { default: "ok", errorWhen: "The operation throws or returns an error result" }, + }, + }, +} as const); + +export type AiSpanName = TelemetrySchemaSpanName; +export type AiSpanStartAttributes = TelemetrySchemaSpanStartAttributes< + typeof AI_TELEMETRY_SCHEMA, + Name +>; +export type AiSpanEndAttributes = TelemetrySchemaSpanEndAttributes< + typeof AI_TELEMETRY_SCHEMA, + Name +>; +export type AiSpanAttributes = AiSpanStartAttributes & AiSpanEndAttributes; +export type AiSpanEventName = TelemetrySchemaSpanEventName; +export type AiSpanEventAttributes< + Name extends AiSpanName, + EventName extends AiSpanEventName, +> = TelemetrySchemaSpanEventAttributes; +export type AiTelemetrySpan = SchemaTelemetrySpan; +export type AiSpan = TelemetrySchemaSpanUnion; + +export function startAiSpan, Result>( + telemetryContext: TelemetryContext, + name: Name, + attributes: ExactTelemetryAttributes, Attributes>, + callback: (span: AiTelemetrySpan) => Result | Promise, +): Promise { + return telemetryContext.startSpan({ name, attributes }, (span) => callback(span as AiTelemetrySpan)); +} diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 8e1ca5124..e2777ed3e 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -8,6 +8,7 @@ import type { OpenAICodexResponsesOptions } from "./api/openai-codex-responses.t import type { OpenAICompletionsOptions } from "./api/openai-completions.ts"; import type { OpenAIResponsesOptions } from "./api/openai-responses.ts"; import type { PiMessagesOptions } from "./api/pi-messages.ts"; +import type { TelemetryContext } from "./telemetry.ts"; import type { AssistantMessageDiagnostic } from "./utils/diagnostics.ts"; import type { AssistantMessageEventStream } from "./utils/event-stream.ts"; @@ -117,6 +118,8 @@ export interface ProviderResponse { /** Authentication, HTTP transport, and lifecycle callbacks shared by provider requests. */ export interface ProviderRequestOptions> { signal?: AbortSignal; + /** Explicit parent context for telemetry produced by this logical request. */ + telemetryContext?: TelemetryContext; apiKey?: string; /** * Optional fetch implementation for provider HTTP requests. diff --git a/packages/ai/test/telemetry-options.test.ts b/packages/ai/test/telemetry-options.test.ts new file mode 100644 index 000000000..732332db8 --- /dev/null +++ b/packages/ai/test/telemetry-options.test.ts @@ -0,0 +1,171 @@ +import { describe, expect, it } from "vitest"; +import { buildBaseOptions } from "../src/api/simple-options.ts"; +import { generateImages } from "../src/images.ts"; +import { registerImagesApiProvider } from "../src/images-api-registry.ts"; +import { createImagesModels, createImagesProvider } from "../src/images-models.ts"; +import { createModels, createProvider } from "../src/models.ts"; +import { NOOP_TELEMETRY_CONTEXT, type TelemetryContext } from "../src/telemetry.ts"; +import type { + Context, + DeferredHandle, + ImagesContext, + ImagesModel, + Model, + ProviderRequestOptions, +} from "../src/types.ts"; +import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; + +const telemetryContext: TelemetryContext = NOOP_TELEMETRY_CONTEXT; +const context: Context = { messages: [] }; +const imagesContext: ImagesContext = { input: [{ type: "text", text: "circle" }] }; + +const model: Model<"telemetry-test"> = { + id: "model", + name: "Model", + api: "telemetry-test", + provider: "telemetry-provider", + baseUrl: "https://example.test", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000, + maxTokens: 100, +}; + +const imageModel: ImagesModel<"telemetry-test-images"> = { + id: "image-model", + name: "Image Model", + api: "telemetry-test-images", + provider: "telemetry-image-provider", + baseUrl: "https://example.test", + input: ["text"], + output: ["image"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, +}; + +function completedStream(requestModel: Model): AssistantMessageEventStream { + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "stop", + message: { + role: "assistant", + content: [], + api: requestModel.api, + provider: requestModel.provider, + model: requestModel.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: 0, + }, + }); + }); + return stream; +} + +describe("ProviderRequestOptions.telemetryContext", () => { + it("is inherited by every request option surface and simple-stream conversion", () => { + const options = { telemetryContext } satisfies ProviderRequestOptions; + expect(options.telemetryContext).toBe(telemetryContext); + expect(buildBaseOptions(model, context, { telemetryContext }).telemetryContext).toBe(telemetryContext); + }); + + it("survives provider and Models stream/deferred dispatch", async () => { + const observed: Array = []; + const handle: DeferredHandle = { + provider: model.provider, + modelId: model.id, + api: model.api, + id: "response", + }; + const provider = createProvider({ + id: model.provider, + auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } }, + models: [model], + api: { + stream: (requestModel, _context, options) => { + observed.push(options?.telemetryContext); + return completedStream(requestModel); + }, + streamSimple: (requestModel, _context, options) => { + observed.push(options?.telemetryContext); + return completedStream(requestModel); + }, + fetchDeferred: (requestModel, _handle, options) => { + observed.push(options?.telemetryContext); + return completedStream(requestModel); + }, + cancelDeferred: async (_requestModel, _handle, options) => { + observed.push(options?.telemetryContext); + }, + }, + }); + + await provider.stream(model, context, { telemetryContext }).result(); + await provider.streamSimple(model, context, { telemetryContext }).result(); + await provider.fetchDeferred!(model, handle, { telemetryContext }).result(); + await provider.cancelDeferred!(model, handle, { telemetryContext }); + + const models = createModels(); + models.setProvider(provider); + await models.stream(model, context, { telemetryContext }).result(); + await models.streamSimple(model, context, { telemetryContext }).result(); + await models.fetchDeferred(model, handle, { telemetryContext }); + await models.cancelDeferred(model, handle, { telemetryContext }); + + expect(observed).toHaveLength(8); + expect(observed.every((value) => value === telemetryContext)).toBe(true); + }); + + it("survives direct and ImagesModels image dispatch", async () => { + const observed: Array = []; + registerImagesApiProvider({ + api: imageModel.api, + generateImages: async (requestModel, _context, options) => { + observed.push(options?.telemetryContext); + return { + api: requestModel.api, + provider: requestModel.provider, + model: requestModel.id, + output: [], + stopReason: "stop", + timestamp: 0, + }; + }, + }); + await generateImages(imageModel, imagesContext, { telemetryContext }); + + const models = createImagesModels(); + models.setProvider( + createImagesProvider({ + id: imageModel.provider, + auth: { apiKey: { name: "Test", resolve: async () => ({ auth: {} }) } }, + models: [imageModel], + api: { + generateImages: async (requestModel, _context, options) => { + observed.push(options?.telemetryContext); + return { + api: requestModel.api, + provider: requestModel.provider, + model: requestModel.id, + output: [], + stopReason: "stop", + timestamp: 0, + }; + }, + }, + }), + ); + await models.generateImages(imageModel, imagesContext, { telemetryContext }); + + expect(observed).toEqual([telemetryContext, telemetryContext]); + }); +}); diff --git a/packages/ai/test/telemetry.test.ts b/packages/ai/test/telemetry.test.ts new file mode 100644 index 000000000..8e037a663 --- /dev/null +++ b/packages/ai/test/telemetry.test.ts @@ -0,0 +1,154 @@ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { describe, expect, expectTypeOf, it } from "vitest"; +import { renderTelemetrySchemaMarkdown } from "../scripts/generate-telemetry-docs.ts"; +import { + AI_TELEMETRY_SCHEMA, + type AiSpanEndAttributes, + type AiSpanStartAttributes, + NOOP_TELEMETRY_CONTEXT, + startAiSpan, + type TelemetryContext, + type TelemetrySpan, +} from "../src/telemetry.ts"; + +function unreadable(value: T): T { + return new Proxy(value, { + get: () => { + throw new Error("read"); + }, + getOwnPropertyDescriptor: () => { + throw new Error("inspect"); + }, + ownKeys: () => { + throw new Error("enumerate"); + }, + }); +} + +describe("NOOP_TELEMETRY_CONTEXT", () => { + it("admits callbacks synchronously and reuses one inert span", async () => { + let admitted = false; + let firstSpan: TelemetrySpan | undefined; + const result = NOOP_TELEMETRY_CONTEXT.startSpan({ name: "first" }, async (span) => { + admitted = true; + firstSpan = span; + const child = await span.startSpan({ name: "child" }, (childSpan) => childSpan); + expect(child).toBe(span); + return 42; + }); + + expect(admitted).toBe(true); + expect(await result).toBe(42); + expect(Object.isFrozen(firstSpan)).toBe(true); + }); + + it("preserves synchronous and asynchronous rejection values", async () => { + const syncError = { kind: "sync" }; + const sync = NOOP_TELEMETRY_CONTEXT.startSpan({ name: "sync" }, () => { + throw syncError; + }); + await expect(sync).rejects.toBe(syncError); + + const asyncError = { kind: "async" }; + const asyncResult = NOOP_TELEMETRY_CONTEXT.startSpan({ name: "async" }, async () => { + throw asyncError; + }); + await expect(asyncResult).rejects.toBe(asyncError); + }); + + it("does not inspect or retain telemetry payloads", async () => { + const options = unreadable({ name: "operation", attributes: { secret: "prompt content" } }); + await NOOP_TELEMETRY_CONTEXT.startSpan(options, (span) => { + expect(() => span.addEvent("event", unreadable({ secret: "content" }))).not.toThrow(); + expect(() => span.setAttributes(unreadable({ secret: "content" }))).not.toThrow(); + expect(() => span.setStatus(unreadable({ status: "ok" as const }))).not.toThrow(); + }); + }); +}); + +describe("AI_TELEMETRY_SCHEMA", () => { + it("is serializable and generates the checked-in reference", () => { + expect(() => JSON.stringify(AI_TELEMETRY_SCHEMA)).not.toThrow(); + const expected = renderTelemetrySchemaMarkdown(AI_TELEMETRY_SCHEMA, "Pi AI Telemetry Schema"); + const actual = readFileSync(resolve(import.meta.dirname, "../docs/telemetry-schema.md"), "utf8"); + expect(actual).toBe(expected); + }); + + it("infers exact start and optional end attributes", async () => { + type Start = AiSpanStartAttributes<"pi.ai.request">; + type End = AiSpanEndAttributes<"pi.ai.request">; + expectTypeOf().toMatchTypeOf<{ + "pi.ai.operation": "stream" | "fetch_deferred" | "cancel_deferred" | "generate_images"; + "pi.ai.provider": string; + "pi.ai.model": string; + "pi.ai.api": string; + "pi.ai.streaming": boolean; + "pi.ai.deferred"?: boolean; + }>(); + expectTypeOf().toEqualTypeOf< + "stop" | "length" | "tool_use" | "error" | "aborted" | "deferred" | undefined + >(); + + const context: TelemetryContext = NOOP_TELEMETRY_CONTEXT; + await startAiSpan( + context, + "pi.ai.request", + { + "pi.ai.operation": "stream", + "pi.ai.provider": "provider", + "pi.ai.model": "model", + "pi.ai.api": "api", + "pi.ai.streaming": true, + }, + (span) => { + span.setAttributes({ "pi.ai.response.stop_reason": "tool_use" }); + // @ts-expect-error pi.ai.request declares no span events + span.addEvent("chunk"); + }, + ); + + const compileTimeFailures = () => { + const extraAttributes = { + "pi.ai.operation": "stream", + "pi.ai.provider": "provider", + "pi.ai.model": "model", + "pi.ai.api": "api", + "pi.ai.streaming": true, + "pi.ai.unknown": true, + } as const; + // @ts-expect-error variables with unknown attributes are rejected + void startAiSpan(context, "pi.ai.request", extraAttributes, () => {}); + // @ts-expect-error missing required start attributes + void startAiSpan(context, "pi.ai.request", { "pi.ai.operation": "stream" }, () => {}); + void startAiSpan( + context, + "pi.ai.request", + { + "pi.ai.operation": "stream", + "pi.ai.provider": "provider", + "pi.ai.model": "model", + "pi.ai.api": "api", + "pi.ai.streaming": true, + // @ts-expect-error unknown start attribute + "pi.ai.unknown": true, + }, + () => {}, + ); + void startAiSpan( + context, + "pi.ai.request", + { + // @ts-expect-error invalid closed-set value + "pi.ai.operation": "invalid", + "pi.ai.provider": "provider", + "pi.ai.model": "model", + "pi.ai.api": "api", + "pi.ai.streaming": true, + }, + () => {}, + ); + }; + expectTypeOf(compileTimeFailures).toBeFunction(); + }); +});