From 547567d5168679c0b59aad1450dbfbf4ec3e8d76 Mon Sep 17 00:00:00 2001 From: Bryan Font Date: Mon, 13 Jul 2026 12:52:24 -0400 Subject: [PATCH] Adapt rollout snapshots for lineage accounting --- .../Generated/CodexParserHash.generated.swift | 2 +- .../Vendored/CostUsage/CostUsageScanner.swift | 56 +++++++++++++++-- .../CodexLineageLedgerTests.swift | 61 +++++++++++++++++++ 3 files changed, 114 insertions(+), 5 deletions(-) diff --git a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift index 085fda188b..2a1047a65a 100644 --- a/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift +++ b/Sources/CodexBarCore/Generated/CodexParserHash.generated.swift @@ -1,5 +1,5 @@ // Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand. enum CodexParserHash { - static let value = "cdef6eb9658a43e2" + static let value = "865e101428b49ab7" } diff --git a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift index fef88de04f..9bda419158 100644 --- a/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift +++ b/Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift @@ -98,6 +98,13 @@ enum CostUsageScanner { let totals: CostUsageCodexTotals } + private struct CodexParsedTokenEvidence { + let sessionId: String? + let forkedFromId: String? + let snapshots: [CodexTimestampedTotals] + let observations: [CodexLineageLedger.Observation] + } + enum CodexForkBaseline { case resolved(CostUsageCodexTotals?) case unresolved @@ -1680,13 +1687,13 @@ enum CostUsageScanner { private static func parseCodexTokenSnapshots( fileURL: URL, - checkCancellation: CancellationCheck? = nil) throws -> ( - sessionId: String?, - snapshots: [CodexTimestampedTotals]) + checkCancellation: CancellationCheck? = nil) throws -> CodexParsedTokenEvidence { var sessionId: String? + var forkedFromId: String? var accumulator = CodexSnapshotAccumulator() var snapshots: [CodexTimestampedTotals] = [] + var observations: [CodexLineageLedger.Observation] = [] var warnedAboutUnparsedTimestamp = false func parsedSnapshotDate(timestamp: String) -> Date? { @@ -1708,6 +1715,12 @@ enum CostUsageScanner { timestamp: timestamp, date: parsedSnapshotDate(timestamp: timestamp), totals: counted)) + if let last, let total { + observations.append(CodexLineageLedger.Observation( + timestamp: timestamp, + last: Self.lineageTotals(last), + total: Self.lineageTotals(total))) + } } do { @@ -1724,6 +1737,9 @@ enum CostUsageScanner { if sessionId == nil { sessionId = metadata.sessionId } + if forkedFromId == nil { + forkedFromId = metadata.forkedFromId + } case let .tokenCount(record): appendSnapshot(timestamp: record.timestamp, last: record.last, total: record.total) case .turnContext, .taskStarted: @@ -1746,6 +1762,9 @@ enum CostUsageScanner { ?? obj["sessionId"] as? String ?? obj["id"] as? String } + if forkedFromId == nil { + forkedFromId = Self.codexForkParentId(from: payload) + } return } @@ -1785,7 +1804,36 @@ enum CostUsageScanner { metadata: ["path": fileURL.path, "error": error.localizedDescription]) } - return (sessionId, snapshots) + return CodexParsedTokenEvidence( + sessionId: sessionId, + forkedFromId: forkedFromId, + snapshots: snapshots, + observations: observations) + } + + static func parseCodexLineageDocument( + fileURL: URL, + checkCancellation: CancellationCheck? = nil) throws -> CodexLineageLedger.Document + { + let parsed = try Self.parseCodexTokenSnapshots( + fileURL: fileURL, + checkCancellation: checkCancellation) + return CodexLineageLedger.Document( + ownerID: Self.codexRolloutOwnerID(fileURL: fileURL) ?? parsed.sessionId ?? fileURL.standardizedFileURL.path, + metadataSessionID: parsed.sessionId, + parentSessionID: parsed.forkedFromId, + observations: parsed.observations) + } + + private static func lineageTotals(_ totals: CostUsageCodexTotals) -> CodexLineageLedger.Totals { + CodexLineageLedger.Totals(input: totals.input, cached: totals.cached, output: totals.output) + } + + private static func codexRolloutOwnerID(fileURL: URL) -> String? { + let stem = fileURL.deletingPathExtension().lastPathComponent + guard stem.count >= 36 else { return nil } + let candidate = String(stem.suffix(36)) + return UUID(uuidString: candidate) == nil ? nil : candidate.lowercased() } static func parseCodexFile( diff --git a/Tests/CodexBarTests/CodexLineageLedgerTests.swift b/Tests/CodexBarTests/CodexLineageLedgerTests.swift index 434e1e7572..9a6c46c808 100644 --- a/Tests/CodexBarTests/CodexLineageLedgerTests.swift +++ b/Tests/CodexBarTests/CodexLineageLedgerTests.swift @@ -3,6 +3,56 @@ import Testing @testable import CodexBarCore struct CodexLineageLedgerTests { + @Test + func `fast rollout parser adapts complete token states into a ledger document`() throws { + let environment = try CostUsageTestEnvironment() + defer { environment.cleanup() } + let ownerID = "019f55a1-7f6e-70c0-8e4f-f5bbefa9b7ac" + let fileURL = environment.root.appendingPathComponent( + "rollout-2026-07-09T12-00-00-\(ownerID).jsonl") + let contents = [ + #"{"type":"session_meta","payload":{"id":"metadata-id","forked_from_id":"parent-id"}}"#, + Self.tokenCountLine( + timestamp: "2026-07-09T12:00:00Z", + last: (input: 100, cached: 40, output: 10), + total: (input: 100, cached: 40, output: 10)), + Self.tokenCountLine( + timestamp: "2026-07-09T12:01:00Z", + last: (input: 50, cached: 20, output: 5), + total: (input: 150, cached: 60, output: 15)), + #"{"type":"event_msg","timestamp":"2026-07-09T12:02:00Z","payload":{"type":"token_count","info":{"# + + #""total_token_usage":{"input_tokens":200,"cached_input_tokens":80,"output_tokens":20}}}}"#, + ].joined(separator: "\n") + try contents.write(to: fileURL, atomically: true, encoding: .utf8) + + let document = try CostUsageScanner.parseCodexLineageDocument(fileURL: fileURL) + + #expect(document.ownerID == ownerID) + #expect(document.metadataSessionID == "metadata-id") + #expect(document.parentSessionID == "parent-id") + #expect(document.observations.count == 2) + #expect(document.observations[0].last == .init(input: 100, cached: 40, output: 10)) + #expect(document.observations[1].total == .init(input: 150, cached: 60, output: 15)) + } + + @Test + func `ledger document parsing propagates cancellation`() throws { + let environment = try CostUsageTestEnvironment() + defer { environment.cleanup() } + let fileURL = environment.root.appendingPathComponent("rollout-without-owner.jsonl") + try Self.tokenCountLine( + timestamp: "2026-07-09T12:00:00Z", + last: (input: 100, cached: 0, output: 10), + total: (input: 100, cached: 0, output: 10)) + .write(to: fileURL, atomically: true, encoding: .utf8) + + #expect(throws: CancellationError.self) { + _ = try CostUsageScanner.parseCodexLineageDocument( + fileURL: fileURL, + checkCancellation: { throw CancellationError() }) + } + } + @Test func `transitive lineage counts copied observations once`() throws { let first = Self.observation(timestamp: "2026-07-09T12:00:00Z", input: 100, totalInput: 100) @@ -192,4 +242,15 @@ struct CodexLineageLedgerTests { cached: event.total.cachedInputTokens, output: event.total.outputTokens)) } + + private static func tokenCountLine( + timestamp: String, + last: (input: Int, cached: Int, output: Int), + total: (input: Int, cached: Int, output: Int)) -> String + { + #"{"type":"event_msg","timestamp":"\#(timestamp)","payload":{"type":"token_count","info":{"# + + #""last_token_usage":{"input_tokens":\#(last.input),"cached_input_tokens":\#(last.cached),"# + + #""output_tokens":\#(last.output)},"total_token_usage":{"input_tokens":\#(total.input),"# + + #""cached_input_tokens":\#(total.cached),"output_tokens":\#(total.output)}}}}"# + } }