fix checkpoint token metadata fallbacks

897386a·

peyton-alt·3w ago·3 files·+78 added/-9 removed

Sessions

7a1b83ebd234View transcript

Changes

3

• cmd/entire/cli

• Mcheckpoint_tokens.go+5/-4

• Mreview_context.go+2/-2

• Msessions_test.go+71/-3

136 unmodified lines

137
138
139
140
140
141
142
143
144
145
146
147
148
149
39 unmodified lines

189
190
191
189
190
191
192
193
194
195

136 unmodified lines

return buildCheckpointTokensReport(cpID, summary, metas, metadataWarnings), lookup, nil
}

func readCheckpointTokenSessionMetadata(ctx context.Context, store checkpoint.CommittedListReader, cpID id.CheckpointID, sessionCount int) ([]*checkpoint.CommittedMetadata, int, error) {
func readCheckpointTokenSessionMetadata(ctx context.Context, store checkpointSessionMetadataReader, cpID id.CheckpointID, sessionCount int) ([]*checkpoint.CommittedMetadata, int, error) {
    metas := make([]*checkpoint.CommittedMetadata, 0, sessionCount)
    var warnings int
    for i := range sessionCount {
        if ctxErr := ctx.Err(); ctxErr != nil {
            return nil, warnings, ctxErr //nolint:wrapcheck // Propagating context cancellation.
        }
        meta, err := store.ReadSessionMetadata(ctx, cpID, i)
        if err != nil {
            if ctxErr := ctx.Err(); ctxErr != nil {
39 unmodified lines

}

usage := aggregateCheckpointTokenUsage(metas)
    if metadataWarnings > 0 && summary != nil && summary.TokenUsage != nil {
        usage = summary.TokenUsage
    } else if usage == nil && summary != nil {
    if usage == nil && summary != nil {
        usage = summary.TokenUsage
    }
    if tokens := buildSessionTokensUsage(usage); tokens != nil {

Mcmd/entire/cli/checkpoint_tokens.go+5/-4

27 unmodified lines

28
29
30
31
31
32
33
34
276 unmodified lines

311
312
313
314
314
315
316
317

27 unmodified lines

reviewContextCommitSeparator = "\x1e"
)

type reviewContextSessionMetadataReader interface {
type checkpointSessionMetadataReader interface {
    ReadSessionMetadata(ctx context.Context, checkpointID checkpointid.CheckpointID, sessionIndex int) (*checkpoint.CommittedMetadata, error)
}

276 unmodified lines

cpID checkpointid.CheckpointID,
    sessionIndex int,
) (*checkpoint.CommittedMetadata, error) {
    if r, ok := reader.(reviewContextSessionMetadataReader); ok {
    if r, ok := reader.(checkpointSessionMetadataReader); ok {
        return r.ReadSessionMetadata(ctx, cpID, sessionIndex) //nolint:wrapcheck // Best-effort prompt context.
    }
    content, err := reader.ReadSessionContent(ctx, cpID, sessionIndex)

Mcmd/entire/cli/review_context.go+2/-2

1762 unmodified lines

1763
1764
1765
1766
1766
1767
1768
1769
1770
1771
23 unmodified lines

1795
1796
1797
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874

1762 unmodified lines

}
}

func TestCheckpointTokensReport_UsesRootSummaryWhenSessionMetadataIncomplete(t *testing.T) {
func TestCheckpointTokensReport_UsesReadableMetadataWhenSessionMetadataIncomplete(t *testing.T) {
    t.Parallel()

cpID := id.MustCheckpointID("abc123abc123")
    report := buildCheckpointTokensReport(
    cpID,
23 unmodified lines

if report.Tokens == nil {
        t.Fatalf("expected token data, got nil")
    }
    if report.Tokens.Total != 1500 {
        t.Fatalf("expected root summary total 1500, got %+v", report.Tokens)
    if report.Tokens.Total != 100 {
        t.Fatalf("expected readable session metadata total 100, got %+v", report.Tokens)
    }
    if len(report.Limitations) == 0 || !strings.Contains(report.Limitations[0], "1 checkpoint session metadata file could not be read") {
        t.Fatalf("expected incomplete metadata limitation, got %+v", report.Limitations)
    }
}

func TestCheckpointTokensReport_UsesRootSummaryWhenNoSessionMetadataReadable(t *testing.T) {
    t.Parallel()

cpID := id.MustCheckpointID("abc123def456")
    report := buildCheckpointTokensReport(
    cpID,
        &checkpoint.CheckpointSummary{
            CheckpointID: cpID,
            Sessions: []checkpoint.SessionFilePaths{
                {Metadata: "0/metadata.json"},
                {Metadata: "1/metadata.json"},
            },
            TokenUsage: &agent.TokenUsage{
                InputTokens:  1000,
                OutputTokens: 500,
                APICallCount: 7,
            },
            },
            nil,
            2,
        )

if report.Tokens == nil {
        t.Fatalf("expected token data, got nil")
    }
    if report.Tokens.Total != 1500 || report.Tokens.APICalls != 7 {
        t.Fatalf("expected root summary tokens, got %+v", report.Tokens)
    }
}

type cancelingCheckpointMetadataReader struct {
    cancel context.CancelFunc
    calls  int
}

func (r *cancelingCheckpointMetadataReader) ReadSessionMetadata(
    _ context.Context,
    _ id.CheckpointID,
    _ int,
) (*checkpoint.CommittedMetadata, error) {
    r.calls++
    if r.calls == 1 {
        r.cancel()
        return &checkpoint.CommittedMetadata{SessionID: "read-before-cancel"}, nil
    }
    return &checkpoint.CommittedMetadata{SessionID: "read-after-cancel"}, nil
}

func TestReadCheckpointTokenSessionMetadataStopsBetweenReadsWhenContextCanceled(t *testing.T) {
    t.Parallel()

ctx, cancel := context.WithCancel(context.Background())
    reader := &cancelingCheckpointMetadataReader{cancel: cancel}

metas, warnings, err := readCheckpointTokenSessionMetadata(ctx, reader, id.MustCheckpointID("abc123abc123"), 2)
    if !errors.Is(err, context.Canceled) {
        t.Fatalf("expected context.Canceled, got metas=%+v warnings=%d err=%v", metas, warnings, err)
    }
    if reader.calls != 1 {
        t.Fatalf("expected one metadata read before cancellation, got %d", reader.calls)
    }
    if metas != nil || warnings != 0 {
        t.Fatalf("expected canceled read to return no partial results, got metas=%+v warnings=%d", metas, warnings)
    }
}

func TestCheckpointTokensCmd_TextOutputWithComparison(t *testing.T) {
    repo, _ := runExplainAutoTestRepo(t)
    ctx := context.Background()