fix(session): re-baseline subagent tokens on session adopt · Entire

fix(session): re-baseline subagent tokens on session adopt

d39ecfd→main·

suhaanthayyil·3d ago·4 files·+85 added/-3 removed

Cross-repo session adoption opens a fresh target-local checkpoint window
(StepCount=0, CheckpointTokenUsage=nil) but the cloned TokenUsage carries the
source session's full cumulative subagent total. Without re-baselining
SubagentTokensBaseline, the first post-adopt checkpoint subtracted the source's
stale-or-nil baseline and over-reported — potentially the source session's entire
subagent usage.

Add a shared session.State.RebaselineSubagentTokens helper so the invariant
"every site that starts a fresh checkpoint window must re-baseline" is enforced
in one place. Call it from the adopt reset block and from resetCheckpointWindow.

Changes

4

454 unmodified lines

455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475

454 unmodified lines

s.TranscriptLinesAtStart = 0
}

// RebaselineSubagentTokens snapshots the current cumulative subagent total
// (TokenUsage.SubagentTokens) into SubagentTokensBaseline so the next checkpoint
// window's CheckpointTokenUsage.SubagentTokens is rescoped to "since this
// re-baseline" rather than re-reporting the full cumulative subagent total.
//
// The invariant is: every site that starts a fresh checkpoint window by clearing
// CheckpointTokenUsage MUST also re-baseline. Callers: the condensation reset
// helper (resetCheckpointWindow) and cross-repo session adoption, which likewise
// opens a fresh target-local window. Sharing this here keeps the two in step.
func (s *State) RebaselineSubagentTokens() {
    if s.TokenUsage != nil {
        s.SubagentTokensBaseline = s.TokenUsage.SubagentTokens
    }
}

// RealignAttributionBase sets AttributionBaseCommit to newBase and clears any
// bookkeeping whose meaning depends on attribution being diverged from the
// shadow-branch base. Call this every time a code path intentionally brings

Mcmd/entire/cli/session/state.go+15

479 unmodified lines

480
481
482
483
484
485
486
487
488
489
490
491
492
493

479 unmodified lines

adopted.LastCheckpointID = id.EmptyCheckpointID
    adopted.LastCheckpointCommitHash = ""
    adopted.CheckpointTokenUsage = nil
    // Re-baseline the subagent cumulative for the fresh target-local window. The
    // cloned TokenUsage carries the SOURCE session's full cumulative subagent
    // total; without re-baselining here, the first post-adopt checkpoint would
    // subtract the source's (stale or nil) baseline and over-report — potentially
    // the source session's entire subagent usage. Mirrors resetCheckpointWindow's
    // baseline capture so the first adopted checkpoint only counts target-side
    // subagent growth, consistent with the PromptWindowBase reset below.
    adopted.RebaselineSubagentTokens()

adopted.FullyCondensed = false
    adopted.UntrackedFilesAtStart = untrackedFiles

Mcmd/entire/cli/session_adopt.go+8

12 unmodified lines

13
14
15
16
17
18
19
1082 unmodified lines

1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167

12 unmodified lines

"time"

"github.com/entireio/cli/cmd/entire/cli/agent"
    "github.com/entireio/cli/cmd/entire/cli/agent/types"
    "github.com/entireio/cli/cmd/entire/cli/checkpoint/id"
    "github.com/entireio/cli/cmd/entire/cli/internal/flock"
    "github.com/entireio/cli/cmd/entire/cli/paths"
1082 unmodified lines

}

// TestSessionAdopt_RebaselinesSubagentTokens pins finding 019f5ebf-dc42: cross-repo
// adoption opens a fresh target-local checkpoint window (StepCount=0,
// CheckpointTokenUsage=nil), but the cloned TokenUsage carries the SOURCE
// session's full cumulative subagent total. If SubagentTokensBaseline is not
// re-baselined to that cumulative, the first post-adopt checkpoint subtracts a
// stale/nil baseline and over-reports the source session's subagent usage.
func TestSessionAdopt_RebaselinesSubagentTokens(t *testing.T) {
    for _, tc := range []struct {
        name           string
        sourceBaseline *agent.TokenUsage
    } {
        // Source never condensed: baseline is nil, so the first adopted
        // checkpoint would report the entire cumulative subagent total.
        {name: "never-condensed-source", sourceBaseline: nil},
        // Source condensed at an earlier window: its baseline is stale relative
        // to the current cumulative and must not carry into the target window.
        {name: "previously-condensed-source", sourceBaseline: &agent.TokenUsage{InputTokens: 200, OutputTokens: 100, APICallCount: 2}},
    } {
        t.Run(tc.name, func(t *testing.T) {
            targetRepo := setupAdoptRepo(t)
            testutil.WriteFile(t, targetRepo, "feature.txt", "agent change\n")
            t.Chdir(targetRepo)

adopted, _, err := buildAdoptedSessionState(context.Background(), &session.State{
                SessionID:    "test-adopt-subagent-baseline-" + tc.name,
                AgentType:    agent.AgentTypeClaudeCode,
                StartedAt:    time.Now().Add(-5 * time.Minute),
                Phase:        session.PhaseActive,
                BaseCommit:   "source-head",
                WorktreePath: "/source/repo",
                TokenUsage: &agent.TokenUsage{
                    InputTokens:    1000,
                    OutputTokens:   500,
                    APICallCount:   10,
                    SubagentTokens: &agent.TokenUsage{InputTokens: 500, OutputTokens: 250, APICallCount: 5},
                },
                SubagentTokensBaseline: tc.sourceBaseline,
            })
            if err != nil {
                t.Fatalf("buildAdoptedSessionState failed: %v", err)
            }

if adopted.SubagentTokensBaseline == nil {
                t.Fatal("adopted SubagentTokensBaseline = nil, want re-baselined to the cumulative subagent total")
            }
            if adopted.SubagentTokensBaseline.InputTokens != 500 || adopted.SubagentTokensBaseline.OutputTokens != 250 {
                t.Fatalf("adopted SubagentTokensBaseline = %#v, want cumulative subagent total 500/250",
                    adopted.SubagentTokensBaseline)
            }

// The first post-adopt checkpoint delta (cumulative - baseline) must be
            // zero: adoption should count only target-side subagent growth.
            delta := types.SubtractTokenUsage(adopted.TokenUsage.SubagentTokens, adopted.SubagentTokensBaseline)
            if delta.InputTokens != 0 || delta.OutputTokens != 0 || delta.APICallCount != 0 {
                t.Fatalf("first post-adopt subagent delta = %#v, want zero", delta)
            }
        })
    }
}

func TestSessionAdopt_PreservesReviewAndInvestigateMetadata(t *testing.T) {
    for _, tc := range []struct {
        name string