Skip to content

Kiro: forward per-turn token counts + session model for the pricing hedge #276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alexeyzimarev merged 2 commits into from
Jul 7, 2026
Merged

Kiro: forward per-turn token counts + session model for the pricing hedge #276

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2d4f6b3
[AI-1196] Kiro: forward per-turn token counts + session model for the…
realtonyyoung Jul 6, 2026
482b91d
[AI-1196] Address review: defensive model_id read in KiroUsage.AnchorMap
realtonyyoung Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 57 additions & 17 deletions src/Capacitor.Cli.Core/Kiro/KiroUsage.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,29 @@ namespace Capacitor.Cli.Core.Kiro;

/// <summary>
/// Reads Kiro's per-turn usage from a session's <c>{id}.json</c> metadata and
/// injects it onto the matching JSONL transcript line. Kiro records NO token
/// counts, but each `user_turn_metadatas[]` entry carries billing
/// <c>metering_usage[]</c> (credits) and an end-of-turn
/// <c>context_usage_percentage</c> — and that lives in the <c>.json</c>, not the
/// <c>.jsonl</c> the server normalizer sees. So the CLI maps each turn's usage to
/// its <b>final</b> message id (the turn's last entry in <c>message_ids[]</c> —
/// the final assistant response) and stamps a <c>_kcap_usage</c> object on that
/// transcript line; <c>KiroTranscriptNormalizer</c> lifts it onto
/// <c>extensions.kiro</c>. Import-only today (the live path would key off Kiro's
/// per-turn <c>stop</c> hook).
/// injects it onto the matching JSONL transcript line. Each
/// <c>user_turn_metadatas[]</c> entry carries billing <c>metering_usage[]</c>
/// (credits) and an end-of-turn <c>context_usage_percentage</c>, plus
/// <c>input_token_count</c>/<c>output_token_count</c> that Kiro persists as 0
/// today (the CLI discards the Amazon-Q API's TokenUsage — upstream aws#2397).
/// This lives in the <c>.json</c>, not the <c>.jsonl</c> the server normalizer
/// sees. So the CLI maps each turn's usage to its <b>final</b> message id (the
/// turn's last entry in <c>message_ids[]</c> — the final assistant response) and
/// stamps a <c>_kcap_usage</c> object on that transcript line;
/// <c>KiroTranscriptNormalizer</c> lifts credits/context% onto
/// <c>extensions.kiro</c> and, when the token counts are non-zero, stamps a
/// canonical <c>$usage</c> (the AI-1196 hedge: dormant now, auto-lights if a
/// future kiro-cli release populates them). The session <c>model_id</c> rides
/// alongside the token counts so that <c>$usage</c> has a model. Import-only today
/// (the live path would key off Kiro's per-turn <c>stop</c> hook).
/// </summary>
public static class KiroUsage {
public readonly record struct TurnUsage(double Credits, double? ContextPct);
public readonly record struct TurnUsage(
double Credits,
double? ContextPct,
long? InputTokens = null,
long? OutputTokens = null,
string? Model = null);

/// <summary>
/// Builds <c>anchor message_id → usage</c> from the metadata JSON. The anchor
Expand All @@ -28,10 +38,20 @@ public static IReadOnlyDictionary<string, TurnUsage> AnchorMap(string? metadataJ
if (string.IsNullOrWhiteSpace(metadataJson)) return map;

try {
var turns = JsonNode.Parse(metadataJson)
var root = JsonNode.Parse(metadataJson);
var turns = root
?["session_state"]?["conversation_metadata"]?["user_turn_metadatas"] as JsonArray;
if (turns is null) return map;

// Session-level model id (e.g. "minimax-m2.5", "auto"). Rides alongside the
// token counts so the server can stamp a canonical $usage with a model —
// AccumulateTokens drops entries with no model (HasModel guard).
// Defensive read: GetValue<string>() throws if model_id is present but not a
// JSON string, and this whole method is one best-effort try/catch — a bad
// optional field must not abort the map and drop credits/context% enrichment.
var sessionModel = root?["session_state"]?["rts_model_state"]?["model_info"]?["model_id"]
is JsonValue mv && mv.TryGetValue<string>(out var sm) ? sm : null;

foreach (var t in turns) {
if (t?["message_ids"] is not JsonArray mids || mids.Count == 0) continue;
if (mids[^1]?.GetValue<string>() is not { Length: > 0 } anchor) continue;
Expand All @@ -50,14 +70,30 @@ public static IReadOnlyDictionary<string, TurnUsage> AnchorMap(string? metadataJ

double? ctx = t["context_usage_percentage"] is JsonValue c && c.TryGetValue<double>(out var cp) ? cp : null;

if (hasCredits || ctx is not null)
map[anchor] = new TurnUsage(hasCredits ? credits : 0, ctx);
// Token-count hedge (AI-1196): Kiro persists these as 0 today (upstream
// aws#2397 — the CLI discards the API's TokenUsage). Carry them ONLY when
// non-zero, so the hedge stays dormant now and lights up automatically if
// a future kiro-cli release starts populating them.
var inTok = Tokens(t, "input_token_count");
var outTok = Tokens(t, "output_token_count");
var hasTokens = inTok is not null || outTok is not null;

if (hasCredits || ctx is not null || hasTokens)
map[anchor] = new TurnUsage(
hasCredits ? credits : 0, ctx,
inTok, outTok,
hasTokens ? sessionModel : null);
}
} catch {
// Malformed metadata — usage is best-effort enrichment, never fatal.
}

return map;

// Reads a per-turn token counter, returning null for absent/non-positive
// values so a 0 (Kiro's value today) never flows through as a real count.
static long? Tokens(JsonNode? turn, string prop) =>
turn?[prop] is JsonValue v && v.TryGetValue<long>(out var n) && n > 0 ? n : null;
}

/// <summary>
Expand All @@ -74,10 +110,14 @@ public static string EnrichLine(string line, IReadOnlyDictionary<string, TurnUsa
if (root["data"] is not JsonObject data) return line;
if (data["message_id"]?.GetValue<string>() is not { } mid || !anchors.TryGetValue(mid, out var u)) return line;

// bool/int/double assign to JsonObject is AOT-reflection-free; credits
// and context% are doubles, so no JsonNode.Parse-for-strings dance.
// bool/int/double/long/string assign to JsonObject is AOT-reflection-free,
// so no JsonNode.Parse dance. Credits/context% are doubles; token counts are
// longs and only present when non-zero (AI-1196 hedge); model rides with them.
var usage = new JsonObject { ["credits"] = u.Credits };
if (u.ContextPct is { } pct) usage["context_usage_percentage"] = pct;
if (u.ContextPct is { } pct) usage["context_usage_percentage"] = pct;
if (u.InputTokens is { } inTok) usage["input_token_count"] = inTok;
if (u.OutputTokens is { } outTok) usage["output_token_count"] = outTok;
if (u.Model is { Length: > 0 } model) usage["model"] = model;
Comment thread
realtonyyoung marked this conversation as resolved.
data["_kcap_usage"] = usage;

return root.ToJsonString();
Expand Down
80 changes: 80 additions & 0 deletions test/Capacitor.Cli.Tests.Unit/KiroUsageTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,4 +56,84 @@ public async Task enrich_line_injects_usage_on_anchor_assistant_line_only() {
public async Task anchor_map_empty_on_missing_or_malformed(string? json) {
await Assert.That(KiroUsage.AnchorMap(json).Count).IsEqualTo(0);
}

// ── token-count hedge (AI-1196) ──────────────────────────────────────────
// Kiro's schema carries input/output_token_count per turn and a session model
// id, but the CLI persists them as 0 today (upstream aws#2397). When a future
// release populates them, the CLI must forward the (non-zero) counts + model so
// the server can stamp a canonical $usage — no further code change needed.

// One turn WITH real token counts and a session-level model id.
const string MetaWithTokens = """
{"session_state":{
"rts_model_state":{"model_info":{"model_id":"minimax-m2.5","context_window_tokens":196000}},
"conversation_metadata":{"user_turn_metadatas":[
{"message_ids":["u1","a1"],
"input_token_count":1200,"output_token_count":340,
"context_usage_percentage":7.5,
"metering_usage":[{"value":0.1,"unit":"credit"}]}
]}}}
""";

[Test]
public async Task anchor_map_captures_nonzero_token_counts_and_session_model() {
var map = KiroUsage.AnchorMap(MetaWithTokens);

await Assert.That(map.ContainsKey("a1")).IsTrue();
await Assert.That(map["a1"].InputTokens).IsEqualTo(1200L);
await Assert.That(map["a1"].OutputTokens).IsEqualTo(340L);
await Assert.That(map["a1"].Model).IsEqualTo("minimax-m2.5");
}

[Test]
public async Task enrich_line_injects_token_counts_and_model_when_present() {
var map = KiroUsage.AnchorMap(MetaWithTokens);
var anchor = """{"version":"v1","kind":"AssistantMessage","data":{"message_id":"a1","content":[{"kind":"text","data":"done"}]}}""";

var enriched = KiroUsage.EnrichLine(anchor, map);

await Assert.That(enriched).Contains("input_token_count");
await Assert.That(enriched).Contains("1200");
await Assert.That(enriched).Contains("output_token_count");
await Assert.That(enriched).Contains("340");
await Assert.That(enriched).Contains("minimax-m2.5");
}

[Test]
public async Task anchor_map_tolerates_non_string_model_id() {
// Review finding: a malformed/variant model_id (present but not a JSON string)
// must NOT abort the whole anchor map — usage enrichment is best-effort. Here the
// turn's credits/context% must still be captured.
const string metaBadModel = """
{"session_state":{
"rts_model_state":{"model_info":{"model_id":12345}},
"conversation_metadata":{"user_turn_metadatas":[
{"message_ids":["u1","a1"],
"context_usage_percentage":7.5,
"metering_usage":[{"value":0.1,"unit":"credit"}]}
]}}}
""";

var map = KiroUsage.AnchorMap(metaBadModel);

await Assert.That(map.ContainsKey("a1")).IsTrue();
await Assert.That(map["a1"].Credits).IsEqualTo(0.1);
await Assert.That(map["a1"].ContextPct).IsEqualTo(7.5);
}

[Test]
public async Task zero_token_counts_stay_dormant() {
// The default fixture has input/output_token_count = 0 — the hedge must NOT
// fabricate token/model fields (today's behaviour: credits/context% only).
var map = KiroUsage.AnchorMap(Meta);

await Assert.That(map["a2"].InputTokens).IsNull();
await Assert.That(map["a2"].OutputTokens).IsNull();

var anchor = """{"version":"v1","kind":"AssistantMessage","data":{"message_id":"a2","content":[{"kind":"text","data":"done"}]}}""";
var enriched = KiroUsage.EnrichLine(anchor, map);

await Assert.That(enriched).DoesNotContain("input_token_count");
await Assert.That(enriched).DoesNotContain("output_token_count");
}
}
Loading