Skip to content

Feat: Port TokenUsage and resolveErrorType from adk-python (Part 1/2) - #399

Open
AmaadMartin wants to merge 1 commit into
feat/opt-in-usage-metricsfrom
feat/telemetry-metrics-token-usage-part1
Open

Feat: Port TokenUsage and resolveErrorType from adk-python (Part 1/2)#399
AmaadMartin wants to merge 1 commit into
feat/opt-in-usage-metricsfrom
feat/telemetry-metrics-token-usage-part1

Conversation

@AmaadMartin

@AmaadMartin AmaadMartin commented Jul 31, 2026

Copy link
Copy Markdown
Owner

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

  1. Link to an existing issue (if applicable):
    Closes: N/A
    Related: N/A
  2. Or, if no issue exists, describe the change:
    Problem: adk-js has no equivalent of adk-python's telemetry/_token_usage.py or tracing.resolve_error_type. Token counts are read ad hoc wherever they are needed — tracing.ts sets gen_ai.usage.input_tokens from promptTokenCount alone, ignoring tool-use tokens, and gen_ai.usage.output_tokens from candidatesTokenCount alone, ignoring reasoning tokens — and there is no way to turn a caught error into an error.type label at all. The metrics port that follows (Part 2/2) needs both.

Solution: Port the two pieces, faithfully:

  • core/src/telemetry/token_usage.ts — a TokenUsage class over GenerateContentResponseUsageMetadata that aggregates prompt + tool-use tokens as input and candidate + reasoning tokens as output, plus toAttributes() and the four GenAI usage attribute-name constants. Ported field for field from adk-python/src/google/adk/telemetry/_token_usage.py.
  • core/src/telemetry/tracing.tsresolveErrorType(), the JS twin of tracing.resolve_error_type (same file placement as Python).

Stacked PR (Part 1/2). Part 2 rebuilds core/src/telemetry/metrics.ts on top of these.

Collision check. Before writing anything I checked the open PRs on this fork:

gh pr list --repo AmaadMartin/adk-js --state open --limit 300 --json number,title,headRefName
gh pr diff <n> --repo AmaadMartin/adk-js --name-only   # for every plausibly adjacent PR

299 open PRs; the adjacent ones are #22 (feat/opt-in-usage-metrics), #177, #348, #357, #387, #394. #22 owns core/src/telemetry/metrics.ts and core/test/telemetry/metrics_test.ts with a different, non-semantic-convention metric contract, and wires three call sites. It does not land this change (no TokenUsage, no resolveErrorType), but it overlaps the same module, so this work is stacked on feat/opt-in-usage-metrics rather than branched from main. This PR itself is purely additive and touches nothing #22 touches; Part 2 is where the two contracts are reconciled. Nothing else in the list overlaps: #357 and #308 touch tracing.ts/tracing_test.ts but only the ADK_CAPTURE_MESSAGE_CONTENT_IN_SPANS gate, which is a different region of both files.

Notes for review

  1. Every emitted string is copied verbatim from adk-python, including the deliberately snake_case, non-semconv gen_ai.usage.experimental.system_instruction_tokens. Renaming any of them would split the dashboard series between the two runtimes.
  2. systemInstructionTokens is read through a local extension interface, not a cast. Python reaches it with a defensive getattr; @google/genai 2.9.0 (the version core/package.json declares, and the copy installed under core/node_modules) does not declare the field either, so UsageMetadataWithSystemInstructionTokens extends GenerateContentResponseUsageMetadata adds the optional field with no any involved.
  3. undefined vs 0 is load-bearing. addTokenCounts returns undefined only when both inputs are undefined; a reported 0 stays 0. A truthiness check (if (!a && !b)) would silently turn "the model used zero prompt tokens" into "no data", which is a different point on a dashboard. This is pinned by two tests and by mutation 1 below.
  4. resolveErrorType matches a genai ApiError structurally, not with instanceof. This is not hypothetical here: this workspace resolves two copies of @google/genainode_modules/@google/genai is 1.52.0 and core/node_modules/@google/genai is 2.9.0 — so an error thrown by one copy is not an instanceof the class exported by the other. The test constructs a real ApiError from the SDK to prove the structural read works against the actual class.
  5. error.name is not constructor.name-style type detection. error.name is a standard, serialisable Error property, which is the JS analogue of Python's type(error).__name__; the repo's "no constructor.name for type detection" rule is about cross-package class identity, which this is not. error.constructor.name is used only as a fallback when name has been blanked out, preserving behaviour that already has three tests on the base branch.
  6. Python reads error.code for its APIError; this reads status. The JS SDK's ApiError declares status: number and no code (core/node_modules/@google/genai/dist/genai.d.ts:444), so status is the same value under the SDK's own name. No speculative code fallback was added — there is nothing in the dependency tree that would set it.
  7. Both modules stay internal. Nothing is added to core/src/index.ts or core/src/common.ts; tests reach them by relative path, exactly like the existing tracing.ts tests.
  8. resolveErrorType lives in utils/error_utils.ts, not in tracing.ts. Python keeps it in tracing.py, but nothing in adk-js's tracing.ts calls it (traceToolCall takes no error parameter), and it is not tracing-specific — it is a generic error-classification helper, so it belongs with case_utils.ts / file_utils.ts and is named for what it does rather than for the feature that needed it first.
  9. It takes unknown, not Error. Anything can be thrown in JavaScript, and every caller gets its value from a catch. Typing the parameter honestly is what removes the e as Error cast from the three catch clauses in Part 2. errorType is only honoured when it is a string; the shape is matched with in narrowing, so there is no interface declaring an unknown field.

Testing Plan

Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes.
Unit Tests:
[x] I have added or updated unit tests for my change.
[x] All unit tests pass locally.

npx vitest run --project unit:core core/test/telemetry/token_usage_test.ts core/test/telemetry/tracing_test.ts
 ✓ core/test/telemetry/token_usage_test.ts (21 tests)
 ✓ core/test/telemetry/tracing_test.ts     (13 tests)
 Test Files  2 passed (2)
      Tests  34 passed (34)

core/test/telemetry/token_usage_test.ts ports all 18 cases from adk-python/tests/unittests/telemetry/test_token_usage.py, plus two JS-only ones for the systemInstructionTokens branch present and absent. Key absence is asserted with not.toHaveProperty, because the parity point is that the key is not emitted at all — not that it is emitted as undefined. The attribute-name constants are no longer exported (toAttributes() is their only consumer), so the tests spell the wire keys out as literals instead of importing them — which is a stronger guard, since an alias cannot drift with the module.

core/test/utils/error_utils_test.ts covers resolveErrorType, including the non-object and null thrown values the unknown parameter now admits. tracing_test.ts gains two new cases for the span attributes; no existing case was modified.

Coverage. core/src/telemetry/token_usage.ts and core/src/utils/error_utils.ts are both at 100% statements / branches / functions / lines.

Proof that the tests can fail. Each new test was run against mutated source and observed to fail:

Mutation Failing test Message
if (a === undefined && b === undefined) -> if (!a && !b) in addTokenCounts inputTokenCount > returns 0, not undefined, when both counts are zero expected undefined to be +0 // Object.is equality
delete the systemInstructionTokens branch of toAttributes toAttributes > emits the system instruction tokens the SDK does not declare expected { 'gen_ai.usage.input_tokens': 10 } to deeply equal { …(2) }
delete the numeric-status branch of resolveErrorType resolveErrorType > reports the HTTP status of a genai API error expected 'ApiError' to be '429' // Object.is equality
error.name || error.constructor.name -> error.name resolveErrorType > falls back to the class name when the name has been blanked out expected '' to be 'TypeError' // Object.is equality

Type suppressions: none. This PR adds zero @ts-expect-error, @ts-ignore, eslint-disable, any, as any, as never or coverage pragmas, in src/ or in tests.

npm run ts:check goes from 308 to 312 errors, and here is exactly why. No error is in src/. The four are in core/test/telemetry/tracing_test.ts, at the two new test cases, and they are four more instances of a defect that already produces eight errors in that same file: the shared fixtures are typed through @google/adk (the built core/dist declarations) while the module under test is imported from core/src, so TypeScript sees two nominally different InvocationContext/LlmRequest. Every existing traceCallLlm case in the file reports the same pair. Fixing it means re-pointing the test tree's program, which is what fork PRs #326 and #370 are for; dodging it here would mean either dropping the tests or duplicating the fixtures against deep src paths, and neither is worth it. Flagging it rather than claiming the count is unchanged.

Manual End-to-End (E2E) Tests:
Please provide instructions on how to manually test your changes, including any necessary setup or configuration.

Both modules are internal, so the manual check is through the API they expose:

npm run build
node --input-type=module -e "
import {TokenUsage} from './core/dist/esm/telemetry/token_usage.js';
console.log(new TokenUsage({promptTokenCount: 10, toolUsePromptTokenCount: 5,
                            candidatesTokenCount: 20, thoughtsTokenCount: 8,
                            cachedContentTokenCount: 100}).toAttributes());
"
# { 'gen_ai.usage.input_tokens': 15, 'gen_ai.usage.output_tokens': 28,
#   'gen_ai.usage.cache_read.input_tokens': 100,
#   'gen_ai.usage.reasoning.output_tokens': 8 }

npm run build, npm run lint, npm run format:check and npm run docs:check all pass locally.

Checklist

[x] I have read the CONTRIBUTING.md document.
[x] I have performed a self-review of my own code.
[x] I have commented my code, particularly in hard-to-understand areas.
[x] I have added tests that prove my fix is effective or that my feature works.
[x] New and existing unit tests pass locally with my changes.

CI status: absent (validated locally)

This is a stacked PR whose base is feat/opt-in-usage-metrics, not main. The
repo's test workflow triggers on pull_request: branches: [main]
(.github/workflows/validation.yaml:5-6), so no run-tests job exists for this
PR — the only check that runs is auto-assign, which validates nothing. It was
therefore validated locally on the exact pushed commit 20271240:

npm run build                                                   # pass
npx vitest run --project unit:core \
  core/test/telemetry/token_usage_test.ts core/test/telemetry/tracing_test.ts
                                                                # 34 passed (2 files)
npm run lint                                                    # clean
npm run format:check                                            # All matched files use Prettier code style!
npm run docs:check                                              # clean (typedoc --treatWarningsAsErrors)
npm run ts:check                                                # 308 pre-existing errors, unchanged from the base branch

Revision — complexity review round 1

  • resolveErrorType moved from core/src/telemetry/tracing.ts to core/src/utils/error_utils.ts, and its tests from tracing_test.ts to core/test/utils/error_utils_test.ts. It had no caller inside tracing.ts and is not tracing-specific.
  • Its parameter is unknown, not Error, and the ClassifiedError interface with the errorType?: unknown field is gone — the shape is matched with in narrowing and errorType is honoured only when it is a string. The case that pinned String()-ing a numeric errorType was dropped with the behaviour it described; a new case asserts a non-string errorType is ignored, and two more cover the non-object thrown values unknown now admits.
  • toAttributes() and the attribute constants are no longer dead. The reviewer was right that nothing in production called them; the fix is the one adk-python already makes — traceCallLlm sets them — rather than deleting a faithful port. The four constants are now module-private, since toAttributes() is their only consumer.
  • The tautological attribute names test (asserting an imported constant equals its own literal) was dropped when the constants were un-exported; the same wire keys are now asserted as literals in every toAttributes case, which is what that test was reaching for.

adk-js reads token counts straight off usageMetadata in traceCallLlm,
with its own idea of what "input tokens" means: prompt tokens only,
dropping tool-use tokens, and candidate tokens only, dropping reasoning
tokens. Both are under-reported on every span, and there is no way at
all to turn a caught error into an error.type label.

TokenUsage centralises the aggregation the GenAI semantic conventions
ask for -- prompt + tool-use as input, candidates + reasoning as output,
keeping the distinction between "not reported" and "reported as zero"
that a truthiness check would lose -- and traceCallLlm now sets its
attributes, matching what adk-python's _set_usage_metadata_attributes
does. Spans consequently gain the cache-read and reasoning token
attributes, and their input/output counts are no longer short.

resolveErrorType is the JS twin of tracing.resolve_error_type: a
self-classified errorType wins, then the HTTP status of a @google/genai
ApiError (matched on shape, not instanceof, because two copies of the
SDK can coexist in one dependency tree), then the error name. It takes
unknown, because anything can be thrown in JavaScript, and lives in
utils/error_utils.ts rather than in tracing.ts -- it is not tracing
specific and tracing has no caller for it.

Neither module is added to the public entry points.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant