diff --git a/src/Capacitor.Cli.Core/Kiro/KiroUsage.cs b/src/Capacitor.Cli.Core/Kiro/KiroUsage.cs index a3e60b82..1a37c30b 100644 --- a/src/Capacitor.Cli.Core/Kiro/KiroUsage.cs +++ b/src/Capacitor.Cli.Core/Kiro/KiroUsage.cs @@ -4,19 +4,29 @@ namespace Capacitor.Cli.Core.Kiro; /// /// Reads Kiro's per-turn usage from a session's {id}.json metadata and -/// injects it onto the matching JSONL transcript line. Kiro records NO token -/// counts, but each `user_turn_metadatas[]` entry carries billing -/// metering_usage[] (credits) and an end-of-turn -/// context_usage_percentage — and that lives in the .json, not the -/// .jsonl the server normalizer sees. So the CLI maps each turn's usage to -/// its final message id (the turn's last entry in message_ids[] — -/// the final assistant response) and stamps a _kcap_usage object on that -/// transcript line; KiroTranscriptNormalizer lifts it onto -/// extensions.kiro. Import-only today (the live path would key off Kiro's -/// per-turn stop hook). +/// injects it onto the matching JSONL transcript line. Each +/// user_turn_metadatas[] entry carries billing metering_usage[] +/// (credits) and an end-of-turn context_usage_percentage, plus +/// input_token_count/output_token_count that Kiro persists as 0 +/// today (the CLI discards the Amazon-Q API's TokenUsage — upstream aws#2397). +/// This lives in the .json, not the .jsonl the server normalizer +/// sees. So the CLI maps each turn's usage to its final message id (the +/// turn's last entry in message_ids[] — the final assistant response) and +/// stamps a _kcap_usage object on that transcript line; +/// KiroTranscriptNormalizer lifts credits/context% onto +/// extensions.kiro and, when the token counts are non-zero, stamps a +/// canonical $usage (the AI-1196 hedge: dormant now, auto-lights if a +/// future kiro-cli release populates them). The session model_id rides +/// alongside the token counts so that $usage has a model. Import-only today +/// (the live path would key off Kiro's per-turn stop hook). /// 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); /// /// Builds anchor message_id → usage from the metadata JSON. The anchor @@ -28,10 +38,20 @@ public static IReadOnlyDictionary 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() 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(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() is not { Length: > 0 } anchor) continue; @@ -50,14 +70,30 @@ public static IReadOnlyDictionary AnchorMap(string? metadataJ double? ctx = t["context_usage_percentage"] is JsonValue c && c.TryGetValue(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(out var n) && n > 0 ? n : null; } /// @@ -74,10 +110,14 @@ public static string EnrichLine(string line, IReadOnlyDictionary() 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; data["_kcap_usage"] = usage; return root.ToJsonString(); diff --git a/test/Capacitor.Cli.Tests.Unit/KiroUsageTests.cs b/test/Capacitor.Cli.Tests.Unit/KiroUsageTests.cs index a1a9a4cf..df87c3a1 100644 --- a/test/Capacitor.Cli.Tests.Unit/KiroUsageTests.cs +++ b/test/Capacitor.Cli.Tests.Unit/KiroUsageTests.cs @@ -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"); + } }