Skip to content

Commit 26959d0

Browse files
fy0fy0
committed
feat(web-session): estimate Codex budget from hidden usage
1 parent 1402f4c commit 26959d0

12 files changed

Lines changed: 562 additions & 89 deletions

File tree

‎model/tables/web_session.go‎

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -51,24 +51,31 @@ type WebSessionTable struct {
5151
LastMessageAt *time.Time `gorm:"type:datetime" json:"lastMessageAt"`
5252
LastEventSeq int64 `gorm:"type:integer;not null;default:0" json:"lastEventSeq"`
5353

54-
TotalInputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalInputTokens"`
55-
TotalCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalCachedInputTokens"`
56-
TotalOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalOutputTokens"`
57-
TotalCost float64 `gorm:"type:real;not null;default:0" json:"totalCost"`
58-
LastCompletedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
59-
LastCompletedCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
60-
LastCompletedOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
61-
LatestTurnInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
62-
LatestTurnCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
63-
LatestTurnOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
64-
LatestTurnUsageUpdatedAt *time.Time `gorm:"type:datetime" json:"-"`
65-
ContextBaselineInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
66-
ContextBaselineCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
67-
ContextBaselineOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
68-
LastContextCompactionAt *time.Time `gorm:"type:datetime" json:"-"`
69-
AutoRetryAttempt int `gorm:"type:integer;not null;default:0" json:"-"`
70-
AutoRetryNextAt *time.Time `gorm:"type:datetime" json:"-"`
71-
AutoRetryLastErrorCode *string `gorm:"type:text" json:"-"`
54+
TotalInputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalInputTokens"`
55+
TotalCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalCachedInputTokens"`
56+
TotalOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"totalOutputTokens"`
57+
TotalCost float64 `gorm:"type:real;not null;default:0" json:"totalCost"`
58+
LastCompletedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
59+
LastCompletedCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
60+
LastCompletedOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
61+
LatestTurnInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
62+
LatestTurnCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
63+
LatestTurnOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
64+
LatestTurnUsageUpdatedAt *time.Time `gorm:"type:datetime" json:"-"`
65+
LatestTokenCountInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
66+
LatestTokenCountCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
67+
LatestTokenCountOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
68+
LatestTokenCountTotalTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
69+
LatestTokenCountUpdatedAt *time.Time `gorm:"type:datetime" json:"-"`
70+
SessionContextWindowTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
71+
SessionContextWindowObservedAt *time.Time `gorm:"type:datetime" json:"-"`
72+
ContextBaselineInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
73+
ContextBaselineCachedInputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
74+
ContextBaselineOutputTokens int64 `gorm:"type:integer;not null;default:0" json:"-"`
75+
LastContextCompactionAt *time.Time `gorm:"type:datetime" json:"-"`
76+
AutoRetryAttempt int `gorm:"type:integer;not null;default:0" json:"-"`
77+
AutoRetryNextAt *time.Time `gorm:"type:datetime" json:"-"`
78+
AutoRetryLastErrorCode *string `gorm:"type:text" json:"-"`
7279

7380
LastError *string `gorm:"type:text" json:"lastError"`
7481
SyncError *string `gorm:"type:text" json:"syncError"`

‎service/websession/codex_context.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,11 @@ func (m *Manager) decorateSessionSummary(summary *SessionSummary) {
6666
summary.ContextWindowSource = ContextWindowSourceUnavailable
6767
return
6868
}
69+
if summary.ContextWindowTokens != nil &&
70+
*summary.ContextWindowTokens > 0 &&
71+
summary.ContextWindowSource == ContextWindowSourceSessionUsage {
72+
return
73+
}
6974
config := m.GetCodexRuntimeConfig()
7075
summary.ContextWindowTokens = ptr(config.ContextWindowTokens)
7176
summary.ContextWindowSource = config.Source

‎service/websession/manager.go‎

Lines changed: 62 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4320,6 +4320,18 @@ func mapSessionRecord(record tables.WebSessionTable) SessionSummary {
43204320
ContextEstimate: contextEstimate,
43214321
ContextEstimateMode: contextEstimateMode,
43224322
LastContextCompactionAt: record.LastContextCompactionAt,
4323+
ContextWindowTokens: func() *int64 {
4324+
if record.SessionContextWindowTokens <= 0 {
4325+
return nil
4326+
}
4327+
return ptr(record.SessionContextWindowTokens)
4328+
}(),
4329+
ContextWindowSource: func() ContextWindowSource {
4330+
if record.SessionContextWindowTokens > 0 {
4331+
return ContextWindowSourceSessionUsage
4332+
}
4333+
return ""
4334+
}(),
43234335
}
43244336
}
43254337

@@ -4376,7 +4388,26 @@ func buildLatestTurnUsage(record tables.WebSessionTable) (ContextEstimate, bool)
43764388
return buildRecordedLatestTurnUsage(record)
43774389
}
43784390

4391+
func buildLatestTokenCountUsage(record tables.WebSessionTable) (ContextEstimate, bool) {
4392+
if record.LatestTokenCountUpdatedAt == nil {
4393+
return ContextEstimate{}, false
4394+
}
4395+
estimate := ContextEstimate{
4396+
InputTokens: maxInt64(0, record.LatestTokenCountInputTokens),
4397+
CachedInputTokens: maxInt64(0, record.LatestTokenCountCachedInputTokens),
4398+
OutputTokens: maxInt64(0, record.LatestTokenCountOutputTokens),
4399+
}
4400+
estimate.UsedTokens = contextEstimateUsedTokens(estimate.InputTokens, estimate.OutputTokens)
4401+
if estimate.UsedTokens == 0 && record.LatestTokenCountTotalTokens > 0 {
4402+
estimate.UsedTokens = record.LatestTokenCountTotalTokens
4403+
}
4404+
return estimate, contextEstimateHasValue(estimate)
4405+
}
4406+
43794407
func buildContextEstimate(record tables.WebSessionTable) (ContextEstimate, ContextEstimateMode) {
4408+
if latestTokenCount, ok := buildLatestTokenCountUsage(record); ok {
4409+
return latestTokenCount, ContextEstimateModeLatestTokenCount
4410+
}
43804411
if latestTurnUsage, ok := buildLatestTurnUsage(record); ok {
43814412
return latestTurnUsage, ContextEstimateModeLatestTurnDelta
43824413
}
@@ -4408,19 +4439,29 @@ func maxInt64(left, right int64) int64 {
44084439

44094440
func contextEstimateTotalsUpdate(in, cin, out int64) map[string]any {
44104441
return map[string]any{
4411-
"total_input_tokens": in,
4412-
"total_cached_input_tokens": cin,
4413-
"total_output_tokens": out,
4414-
"updated_at": time.Now(),
4442+
"total_input_tokens": in,
4443+
"total_cached_input_tokens": cin,
4444+
"total_output_tokens": out,
4445+
"latest_token_count_input_tokens": 0,
4446+
"latest_token_count_cached_input_tokens": 0,
4447+
"latest_token_count_output_tokens": 0,
4448+
"latest_token_count_total_tokens": 0,
4449+
"latest_token_count_updated_at": nil,
4450+
"updated_at": time.Now(),
44154451
}
44164452
}
44174453

44184454
func contextEstimateIncrementUpdate(in, cin, out int64) map[string]any {
44194455
return map[string]any{
4420-
"total_input_tokens": gorm.Expr("total_input_tokens + ?", in),
4421-
"total_cached_input_tokens": gorm.Expr("total_cached_input_tokens + ?", cin),
4422-
"total_output_tokens": gorm.Expr("total_output_tokens + ?", out),
4423-
"updated_at": time.Now(),
4456+
"total_input_tokens": gorm.Expr("total_input_tokens + ?", in),
4457+
"total_cached_input_tokens": gorm.Expr("total_cached_input_tokens + ?", cin),
4458+
"total_output_tokens": gorm.Expr("total_output_tokens + ?", out),
4459+
"latest_token_count_input_tokens": 0,
4460+
"latest_token_count_cached_input_tokens": 0,
4461+
"latest_token_count_output_tokens": 0,
4462+
"latest_token_count_total_tokens": 0,
4463+
"latest_token_count_updated_at": nil,
4464+
"updated_at": time.Now(),
44244465
}
44254466
}
44264467

@@ -4429,11 +4470,16 @@ func contextEstimateBaselineResetUpdate(record tables.WebSessionTable, timestamp
44294470
timestamp = time.Now()
44304471
}
44314472
return map[string]any{
4432-
"context_baseline_input_tokens": record.TotalInputTokens,
4433-
"context_baseline_cached_input_tokens": record.TotalCachedInputTokens,
4434-
"context_baseline_output_tokens": record.TotalOutputTokens,
4435-
"last_context_compaction_at": timestamp,
4436-
"updated_at": time.Now(),
4473+
"context_baseline_input_tokens": record.TotalInputTokens,
4474+
"context_baseline_cached_input_tokens": record.TotalCachedInputTokens,
4475+
"context_baseline_output_tokens": record.TotalOutputTokens,
4476+
"last_context_compaction_at": timestamp,
4477+
"latest_token_count_input_tokens": 0,
4478+
"latest_token_count_cached_input_tokens": 0,
4479+
"latest_token_count_output_tokens": 0,
4480+
"latest_token_count_total_tokens": 0,
4481+
"latest_token_count_updated_at": nil,
4482+
"updated_at": time.Now(),
44374483
}
44384484
}
44394485

@@ -5194,6 +5240,9 @@ func extractContextCompactionText(item map[string]any) string {
51945240
if text := strings.TrimSpace(stringValue(item["text"])); text != "" {
51955241
sections = append(sections, text)
51965242
}
5243+
if message := strings.TrimSpace(stringValue(item["message"])); message != "" {
5244+
sections = append(sections, message)
5245+
}
51975246
if content := strings.TrimSpace(strings.Join(collectReasoningFragments(item["content"]), "")); content != "" {
51985247
sections = append(sections, content)
51995248
}

‎service/websession/manager_test.go‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -892,6 +892,7 @@ func TestManagerListSessionsIncludesConfiguredContextWindow(t *testing.T) {
892892

893893
homeDir := t.TempDir()
894894
t.Setenv("HOME", homeDir)
895+
t.Setenv("USERPROFILE", homeDir)
895896
configDir := filepath.Join(homeDir, ".codex")
896897
if err := os.MkdirAll(configDir, 0o755); err != nil {
897898
t.Fatalf("mkdir config dir failed: %v", err)
@@ -4073,6 +4074,32 @@ func TestHandleCodexAppServerUsageDefaultsContextEstimateToCumulativeTotal(t *te
40734074
}
40744075
}
40754076

4077+
func TestBuildContextEstimateUsesLatestTokenCountSnapshot(t *testing.T) {
4078+
now := time.Now()
4079+
record := tables.WebSessionTable{
4080+
TotalInputTokens: 999,
4081+
TotalCachedInputTokens: 111,
4082+
TotalOutputTokens: 88,
4083+
LatestTokenCountTotalTokens: 12656,
4084+
LatestTokenCountUpdatedAt: &now,
4085+
LastContextCompactionAt: &now,
4086+
ContextBaselineInputTokens: 900,
4087+
ContextBaselineCachedInputTokens: 100,
4088+
ContextBaselineOutputTokens: 80,
4089+
}
4090+
4091+
estimate, mode := buildContextEstimate(record)
4092+
if mode != ContextEstimateModeLatestTokenCount {
4093+
t.Fatalf("expected context estimate mode %q, got %q", ContextEstimateModeLatestTokenCount, mode)
4094+
}
4095+
if estimate.UsedTokens != 12656 {
4096+
t.Fatalf("expected usedTokens from latest token_count total, got %d", estimate.UsedTokens)
4097+
}
4098+
if estimate.InputTokens != 0 || estimate.CachedInputTokens != 0 || estimate.OutputTokens != 0 {
4099+
t.Fatalf("expected zero token_count breakdown, got %#v", estimate)
4100+
}
4101+
}
4102+
40764103
func TestFinalizeLatestTurnUsageUsesTurnDeltaEstimate(t *testing.T) {
40774104
cleanup := initTestDB(t)
40784105
defer cleanup()

0 commit comments

Comments
 (0)