fix checkpoint token metadata fallbacks

1735b02→main

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

1897 unmodified lines

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
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969

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