feat: add custom model token weights in engine frontmatter#24134
Merged
feat: add custom model token weights in engine frontmatter#24134
Conversation
- Add `token-weights` field to engine schema (JSON Schema) - Add `EngineTokenWeights` and `EngineTokenClassWeights` structs to `EngineConfig` - Parse `token-weights` from frontmatter in `ExtractEngineConfig` - Embed custom weights in compiled YAML via `GH_AW_INFO_TOKEN_WEIGHTS` env var - Write token weights to `aw_info.json` in `generate_aw_info.cjs` - Add `CustomTokenWeights` type and `populateEffectiveTokensWithCustomWeights` in effective_tokens.go - Auto-read custom weights from `aw_info.json` in `analyzeTokenUsage` - Add comprehensive tests for all new functionality Agent-Logs-Url: https://github.com/github/gh-aw/sessions/e8a47829-09aa-429e-b0a5-7b1d0d5c6e7a Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
…weights Agent-Logs-Url: https://github.com/github/gh-aw/sessions/e8a47829-09aa-429e-b0a5-7b1d0d5c6e7a Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
Copilot created this pull request from a session on behalf of
pelikhan
April 2, 2026 14:47
View session
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support for specifying custom per-model multipliers and per-token-class weights in workflow frontmatter (engine.token-weights) so gh aw logs / gh aw audit can compute “effective tokens” using workflow-specific pricing data.
Changes:
- Extend workflow engine frontmatter parsing and JSON schema to accept
engine.token-weights(multipliers + token-class-weights). - Propagate parsed weights into the compiled workflow via
GH_AW_INFO_TOKEN_WEIGHTS, then persist them intoaw_info.jsonat runtime. - Update CLI effective-token computation and token usage analysis to merge custom weights with built-in defaults, with tests.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/workflow/engine.go | Adds structs + parser for engine.token-weights and stores result on EngineConfig. |
| pkg/workflow/engine_test.go | Adds unit tests for token-weights parsing and extraction into EngineConfig. |
| pkg/workflow/compiler_yaml.go | Embeds token weights JSON into generated workflow env (GH_AW_INFO_TOKEN_WEIGHTS). |
| pkg/parser/schemas/main_workflow_schema.json | Adds schema for engine.token-weights with multipliers + token-class-weights. |
| pkg/cli/token_usage.go | Loads token weights from aw_info.json and applies them when computing effective tokens. |
| pkg/cli/token_usage_test.go | Updates tests for new parseTokenUsageFile(..., customWeights) signature. |
| pkg/cli/logs_models.go | Extends AwInfo to deserialize token_weights from aw_info.json. |
| pkg/cli/effective_tokens.go | Adds custom-weight merge + effective-token computation using caller-provided weights. |
| pkg/cli/effective_tokens_test.go | Adds tests for custom-weight merging and effective token computation. |
| actions/setup/js/generate_aw_info.cjs | Reads GH_AW_INFO_TOKEN_WEIGHTS and writes it into aw_info.json as token_weights. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
pkg/workflow/compiler_yaml.go
Outdated
| // Embed custom token weights when specified in engine.token-weights | ||
| if data.EngineConfig != nil && data.EngineConfig.TokenWeights != nil { | ||
| if tokenWeightsJSON, err := json.Marshal(data.EngineConfig.TokenWeights); err == nil { | ||
| fmt.Fprintf(yaml, " GH_AW_INFO_TOKEN_WEIGHTS: '%s'\n", string(tokenWeightsJSON)) |
There was a problem hiding this comment.
The JSON string is emitted into YAML as a single-quoted scalar without escaping. If a user provides a model key containing a single quote (e.g. "bob's-model"), the generated workflow YAML becomes invalid. Escape single quotes for YAML (replace ' with '') or use a YAML-safe quoting/encoding helper before writing GH_AW_INFO_TOKEN_WEIGHTS.
Suggested change
| fmt.Fprintf(yaml, " GH_AW_INFO_TOKEN_WEIGHTS: '%s'\n", string(tokenWeightsJSON)) | |
| escapedTokenWeightsJSON := strings.ReplaceAll(string(tokenWeightsJSON), "'", "''") | |
| fmt.Fprintf(yaml, " GH_AW_INFO_TOKEN_WEIGHTS: '%s'\n", escapedTokenWeightsJSON) |
Comment on lines
+9159
to
+9179
| "minimum": 0, | ||
| "description": "Weight for input tokens (default: 1.0)" | ||
| }, | ||
| "cached-input": { | ||
| "type": "number", | ||
| "minimum": 0, | ||
| "description": "Weight for cached input tokens (default: 0.1)" | ||
| }, | ||
| "output": { | ||
| "type": "number", | ||
| "minimum": 0, | ||
| "description": "Weight for output tokens (default: 4.0)" | ||
| }, | ||
| "reasoning": { | ||
| "type": "number", | ||
| "minimum": 0, | ||
| "description": "Weight for reasoning tokens (default: 4.0)" | ||
| }, | ||
| "cache-write": { | ||
| "type": "number", | ||
| "minimum": 0, |
There was a problem hiding this comment.
Schema allows token-class-weights values of 0 (minimum: 0), but the implementation treats zero as “unset/use default” (see parseEngineTokenWeights + effective token merging). This makes a user-provided 0 silently ignored. Consider disallowing 0 in the schema (e.g. minimum > 0) or explicitly documenting that 0 is treated as not set / switching the Go structs to pointer fields so an explicit 0 can be distinguished from “missing”.
Suggested change
| "minimum": 0, | |
| "description": "Weight for input tokens (default: 1.0)" | |
| }, | |
| "cached-input": { | |
| "type": "number", | |
| "minimum": 0, | |
| "description": "Weight for cached input tokens (default: 0.1)" | |
| }, | |
| "output": { | |
| "type": "number", | |
| "minimum": 0, | |
| "description": "Weight for output tokens (default: 4.0)" | |
| }, | |
| "reasoning": { | |
| "type": "number", | |
| "minimum": 0, | |
| "description": "Weight for reasoning tokens (default: 4.0)" | |
| }, | |
| "cache-write": { | |
| "type": "number", | |
| "minimum": 0, | |
| "exclusiveMinimum": 0, | |
| "description": "Weight for input tokens (default: 1.0)" | |
| }, | |
| "cached-input": { | |
| "type": "number", | |
| "exclusiveMinimum": 0, | |
| "description": "Weight for cached input tokens (default: 0.1)" | |
| }, | |
| "output": { | |
| "type": "number", | |
| "exclusiveMinimum": 0, | |
| "description": "Weight for output tokens (default: 4.0)" | |
| }, | |
| "reasoning": { | |
| "type": "number", | |
| "exclusiveMinimum": 0, | |
| "description": "Weight for reasoning tokens (default: 4.0)" | |
| }, | |
| "cache-write": { | |
| "type": "number", | |
| "exclusiveMinimum": 0, |
Collaborator
|
@copilot review comments |
…eMinimum in schema
- Escape single quotes in GH_AW_INFO_TOKEN_WEIGHTS value using YAML single-quoted
scalar convention ('' for ') to prevent YAML parse failures with model keys
containing single quotes (e.g. "bob's-model")
- Change token-class-weights schema from minimum:0 to exclusiveMinimum:0 so that
zero weight values are rejected by schema validation, consistent with the
implementation treating zero as "use default"
- Add test TestTokenWeightsSingleQuoteEscapingInYAML to verify the escaping
Agent-Logs-Url: https://github.com/github/gh-aw/sessions/bc515783-422c-49bf-af85-8b62c099cb73
Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
Contributor
Author
Addressed both review comments in d927107:
|
Collaborator
|
@copilot reuse types to represent model usage data. Ensure data can be imported from shared agentic workflows. |
…across packages - Create pkg/types/token_weights.go with shared TokenWeights and TokenClassWeights types - Remove duplicate EngineTokenWeights/EngineTokenClassWeights from pkg/workflow/engine.go - Remove duplicate CustomTokenWeights/customTokenClassWeights from pkg/cli/effective_tokens.go - Update pkg/workflow/engine.go, compiler_yaml.go, engine_test.go to use types.TokenWeights - Update pkg/cli/effective_tokens.go, logs_models.go, token_usage.go, effective_tokens_test.go to use types.TokenWeights Follows the same pattern as pkg/types.BaseMCPServerConfig which is shared between pkg/parser and pkg/workflow. Agent-Logs-Url: https://github.com/github/gh-aw/sessions/615a693d-9408-4a21-9496-66073a045646 Co-authored-by: pelikhan <4175913+pelikhan@users.noreply.github.com>
Contributor
Author
Done in 56a2b4c. The duplicate Warning Firewall rules blocked me from connecting to one or more addresses (expand for details)I tried to connect to the following addresses, but was blocked by firewall rules:
If you need me to access, download, or install something from one of these locations, you can either:
|
github-actions bot
added a commit
that referenced
this pull request
Apr 2, 2026
- Add engine.token-weights reference section to engines.md (PR #24134) - Update frontmatter-full.md with token-weights schema and fix memory default (documents 6g default, closes #23950) - Add new audit.md reference page for gh aw audit/diff/report (closes #23957) - Add audit diff/report cross-references to debugging.md and cost-management.md - Quick Start: define frontmatter on first use, clarify COPILOT_GITHUB_TOKEN distinction from GITHUB_TOKEN (from issue #24115) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This was referenced Apr 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds support for users to provide custom model token weights in the workflow frontmatter under
engine.token-weights. This allows custom model cost information to be used when computing effective token data viagh aw logsandgh aw audit.Changes
Schema (
pkg/parser/schemas/main_workflow_schema.json)token-weightsproperty to the engine configuration object schema withmultipliersandtoken-class-weightssub-propertiesWorkflow parsing (
pkg/workflow/engine.go)EngineTokenWeightsandEngineTokenClassWeightsstructsTokenWeights *EngineTokenWeightsfield toEngineConfigparseEngineTokenWeightshelper to extracttoken-weightsfrom frontmatterCompilation (
pkg/workflow/compiler_yaml.go)GH_AW_INFO_TOKEN_WEIGHTSenv var when specifiedRuntime (
actions/setup/js/generate_aw_info.cjs)GH_AW_INFO_TOKEN_WEIGHTSand writes it toaw_info.jsonastoken_weightsCLI effective token computation (
pkg/cli/effective_tokens.go)CustomTokenWeightstype (used for deserialization from aw_info.json)populateEffectiveTokensWithCustomWeightsfunctionresolveEffectiveWeightshelper that merges custom weights with defaultscomputeModelEffectiveTokensWithWeightsfor direct computation with provided weightsCLI token analysis (
pkg/cli/token_usage.go,pkg/cli/logs_models.go)AwInfonow includesTokenWeights *CustomTokenWeightsto read from aw_info.jsonanalyzeTokenUsageautomatically reads custom weights fromaw_info.jsonin the run directory and applies themUsage Example
How it works
engine.token-weightsin their.mdworkflow filegh aw compileparses the weights and embeds them asGH_AW_INFO_TOKEN_WEIGHTSin the compiled YAMLgenerate_aw_info.cjsreads the env var and writes the weights toaw_info.jsongh aw logsorgh aw auditanalyzes a run, it reads the custom weights fromaw_info.jsonand uses them instead of (or merged with) the built-inmodel_multipliers.jsonvalues