fix checkpoint token metadata fallbacks

bbc16e8·

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

Sessions

7a1b83ebd234View transcript

Changes

3

• cmd/entire/cli

• Mcheckpoint_tokens.go+4/-1

• Mreview_context.go+2/-2

• Msessions_test.go+66

136 unmodified lines

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

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 {

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

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

1868 unmodified lines

1869
1870
1871
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940

1868 unmodified lines

}
}

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()