polish v2: smoothness & delight

internal/agent/compaction.go

@@ -325,6 +325,43 @@ type compactionState struct {
 	// OnText is deliberately NOT forwarded (compaction stays invisible to the user),
 	// but its token COST must still be counted so usage reports and budgets include it.
 	onUsage func(Usage)
+
+	// calibrationRatio scales the raw byte/4 token estimate toward the provider's
+	// real prompt-token count. ApproxTextTokens over-counts code-heavy content by
+	// ~15-20%, which would trip compaction early (at ~60% of true capacity). It
+	// starts at 1.0 and converges via an EMA as each turn reports actual usage, so
+	// later turns compact nearer to real capacity. Zero is treated as 1.0.
+	calibrationRatio float64
+}
+
+// calibrate folds one turn's (rawEstimate, actualPromptTokens) sample into the
+// running calibration ratio. A single sample is clamped to a sane band so an
+// outlier (a huge cache-read turn, a provider-overhead spike) can't skew the
+// estimate enough to disable or thrash compaction.
+func (state *compactionState) calibrate(rawEstimate int, actualPromptTokens int) {
+	if !state.enabled || rawEstimate <= 0 || actualPromptTokens <= 0 {
+		return
+	}
+	sample := float64(actualPromptTokens) / float64(rawEstimate)
+	if sample < 0.5 {
+		sample = 0.5
+	} else if sample > 2.0 {
+		sample = 2.0
+	}
+	if state.calibrationRatio <= 0 {
+		state.calibrationRatio = 1.0
+	}
+	const alpha = 0.3 // weight on the newest sample; smooths jitter across turns
+	state.calibrationRatio = state.calibrationRatio*(1-alpha) + sample*alpha
+}
+
+// calibratedTokens applies the learned ratio to a raw estimate. Before any sample
+// arrives (ratio unset) it returns the raw estimate unchanged.
+func (state *compactionState) calibratedTokens(raw int) int {
+	if state.calibrationRatio <= 0 {
+		return raw
+	}
+	return int(float64(raw) * state.calibrationRatio)
 }
 
 func newCompactionState(options Options) *compactionState {
@@ -353,7 +390,7 @@ func (state *compactionState) maybeCompact(
 	// the messages; both the threshold check and the shrink check below use the
 	// same term so they stay consistent.
 	toolTokens := estimateToolDefTokens(tools)
-	size := estimateTokens(messages) + toolTokens
+	size := state.calibratedTokens(estimateTokens(messages) + toolTokens)
 	if size <= state.threshold {
 		return messages
 	}
@@ -370,7 +407,7 @@ func (state *compactionState) maybeCompact(
 	// entirely and preserve recent turns verbatim.
 	if pruned, reclaimed := pruneStaleToolOutput(messages, state.preserveLast); reclaimed > 0 {
 		messages = pruned
-		size = estimateTokens(messages) + toolTokens
+		size = state.calibratedTokens(estimateTokens(messages) + toolTokens)
 		if size <= state.threshold {
 			state.lowWaterMark = size
 			return messages
@@ -386,7 +423,7 @@ func (state *compactionState) maybeCompact(
 		// later turn) can try again; we never drop messages on failure here.
 		return messages
 	}
-	newSize := estimateTokens(compacted) + toolTokens
+	newSize := state.calibratedTokens(estimateTokens(compacted) + toolTokens)
 	if newSize >= size {
 		// Compaction did not actually shrink anything (e.g. nothing to
 		// summarize). Leave the history untouched and don't churn next turn.
@@ -444,7 +481,7 @@ func (state *compactionState) recover(
 	// the SAME combined (messages + tool-defs) domain maybeCompact uses, so the
 	// proactive shrink-guard compares like with like.
 	state.reactiveAttempted = true
-	state.lowWaterMark = estimateTokens(result) + estimateToolDefTokens(tools)
+	state.lowWaterMark = state.calibratedTokens(estimateTokens(result) + estimateToolDefTokens(tools))
 	return result, true, nil
 }
 

internal/agent/compaction_calibrate_test.go

@@ -0,0 +1,50 @@
+package agent
+
+import "testing"
+
+func TestCalibrationConvergesTowardActual(t *testing.T) {
+	state := &compactionState{enabled: true}
+
+	// Before any sample, the estimate is unchanged.
+	if got := state.calibratedTokens(1000); got != 1000 {
+		t.Fatalf("uncalibrated tokens = %d, want 1000", got)
+	}
+
+	// The estimator over-counts: our raw estimate is 1000 but the provider reports
+	// 850 (a 0.85 ratio). Feeding that sample repeatedly should pull the ratio
+	// toward 0.85, so a 1000-token raw estimate calibrates downward.
+	for range 20 {
+		state.calibrate(1000, 850)
+	}
+	got := state.calibratedTokens(1000)
+	if got >= 1000 || got < 820 || got > 880 {
+		t.Fatalf("calibrated tokens = %d, want ~850 after convergence", got)
+	}
+}
+
+func TestCalibrateIgnoresDegenerateSamples(t *testing.T) {
+	state := &compactionState{enabled: true}
+	state.calibrate(0, 500) // zero estimate
+	state.calibrate(500, 0) // zero actual
+	if state.calibrationRatio != 0 {
+		t.Fatalf("degenerate samples must not move the ratio, got %v", state.calibrationRatio)
+	}
+	// Disabled state never calibrates.
+	disabled := &compactionState{enabled: false}
+	disabled.calibrate(1000, 850)
+	if disabled.calibrationRatio != 0 {
+		t.Fatalf("disabled compaction must not calibrate, got %v", disabled.calibrationRatio)
+	}
+}
+
+func TestCalibrateClampsOutliers(t *testing.T) {
+	state := &compactionState{enabled: true}
+	// A wild outlier (10x) is clamped to 2.0, so one bad sample can't blow up the
+	// ratio and disable compaction.
+	for range 50 {
+		state.calibrate(1000, 10000)
+	}
+	if state.calibrationRatio > 2.0 {
+		t.Fatalf("ratio should be clamped at 2.0, got %v", state.calibrationRatio)
+	}
+}

internal/agent/compaction_preserve.go

@@ -20,6 +20,16 @@
 	toolNameUpdatePlan = "update_plan"
 	toolNameToolSearch = "tool_search"
 	toolNameSkill      = "skill"
+	toolNameWriteFile  = "write_file"
+	toolNameEditFile   = "edit_file"
+)
+
+// maxRecentEdits caps how many edited files are carried across a compaction, and
+// maxEditNoteBytes caps each edit's one-line note, so a long editing run can't
+// bloat the preserved-state block.
+const (
+	maxRecentEdits   = 20
+	maxEditNoteBytes = 160
 )
 
 const (
@@ -100,12 +110,123 @@
 }
 
 // skillEntry is a named preserved body. It began as loaded-skill state and is
-// reused for loaded tools and project instruction blocks.
+// reused for loaded tools, project instruction blocks, and recent file edits.
 type skillEntry struct {
 	name string
 	body string
 }
 
+// recentEdits returns the files mutated by write_file/edit_file calls in messages
+// — latest note per path, in last-seen order — as skillEntry{name: path, body:
+// note}. After compaction elides the editing turns, this tells the model WHAT it
+// changed in each file (from the tool's result) so it need not re-read to
+// rediscover its own footprint. Capped at maxRecentEdits paths.
+//
+// Ordering is by LAST edit, not first: re-editing a file moves it to the newest
+// position so the tail cap below keeps the file the model most recently touched
+// rather than an earlier, now-stale entry.
+func recentEdits(messages []zeroruntime.Message) []skillEntry {
+	pathByID := map[string]string{}
+	sequence := make([]string, 0)
+	for _, message := range messages {
+		for _, call := range message.ToolCalls {
+			if call.Name != toolNameWriteFile && call.Name != toolNameEditFile {
+				continue
+			}
+			path := editPathFromArguments(call.Arguments)
+			if path == "" {
+				continue
+			}
+			if call.ID != "" {
+				pathByID[call.ID] = path
+			}
+			sequence = append(sequence, path)
+		}
+	}
+	order := lastSeenOrder(sequence)
+	if len(order) == 0 {
+		return nil
+	}
+
+	noteByPath := map[string]string{}
+	for _, message := range messages {
+		if message.Role != zeroruntime.MessageRoleTool || message.ToolCallID == "" {
+			continue
+		}
+		if path, ok := pathByID[message.ToolCallID]; ok {
+			noteByPath[path] = editNote(message.Content)
+		}
+	}
+
+	// Keep the most recent maxRecentEdits paths (the tail of last-seen order).
+	if len(order) > maxRecentEdits {
+		order = order[len(order)-maxRecentEdits:]
+	}
+	entries := make([]skillEntry, 0, len(order))
+	for _, path := range order {
+		entries = append(entries, skillEntry{name: path, body: noteByPath[path]})
+	}
+	return entries
+}
+
+// lastSeenOrder dedupes paths keeping each at its LAST occurrence, preserving
+// chronological order otherwise. A re-edited path therefore lands at the newest
+// position rather than staying pinned to where it first appeared.
+func lastSeenOrder(paths []string) []string {
+	lastIdx := make(map[string]int, len(paths))
+	for i, p := range paths {
+		lastIdx[p] = i
+	}
+	order := make([]string, 0, len(lastIdx))
+	for i, p := range paths {
+		if lastIdx[p] == i {
+			order = append(order, p)
+		}
+	}
+	return order
+}
+
+// editPathFromArguments pulls the target file path from a write_file/edit_file
+// call's JSON arguments (path/file/file_path/filename aliases). Returns "" on
+// malformed arguments or no path.
+func editPathFromArguments(arguments string) string {
+	var parsed struct {
+		Path     string `json:"path"`
+		File     string `json:"file"`
+		FilePath string `json:"file_path"`
+		Filename string `json:"filename"`
+	}
+	if err := json.Unmarshal([]byte(strings.TrimSpace(arguments)), &parsed); err != nil {
+		return ""
+	}
+	for _, candidate := range []string{parsed.Path, parsed.File, parsed.FilePath, parsed.Filename} {
+		if s := strings.TrimSpace(candidate); s != "" {
+			return s
+		}
+	}
+	return ""
+}
+
+// editNote reduces an edit tool's result to a single short line (its first
+// non-empty line, byte-capped on a rune boundary) for the preserved summary.
+func editNote(content string) string {
+	for _, line := range strings.Split(content, "\n") {
+		s := strings.TrimSpace(line)
+		if s == "" {
+			continue
+		}
+		if len(s) > maxEditNoteBytes {
+			limit := maxEditNoteBytes
+			for limit > 0 && !utf8.RuneStart(s[limit]) {
+				limit--
+			}
+			s = s[:limit] + "…"
+		}
+		return s
+	}
+	return ""
+}
+
 // loadedSkills returns the skills loaded via the skill tool in messages — the
 // latest body per name, in first-seen order — matching each skill tool call to
 // its tool result by id.
@@ -310,11 +431,17 @@
 // preservedState is the JSON shape of the carried-across-compaction block.
 type preservedState struct {
 	Plan                string                 `json:"plan,omitempty"`
+	RecentEdits         []preservedEdit        `json:"recent_edits,omitempty"`
 	Tools               []preservedTool        `json:"tools,omitempty"`
 	Skills              []preservedSkill       `json:"skills,omitempty"`
 	ProjectInstructions []preservedInstruction `json:"project_instructions,omitempty"`
 }
 
+type preservedEdit struct {
+	Path string `json:"path"`
+	Note string `json:"note,omitempty"`
+}
+
 type preservedTool struct {
 	Name string `json:"name"`
 	Body string `json:"body"`
@@ -347,6 +474,12 @@
 		plan = priorState.Plan
 	}
 
+	// Recent edits: merge edits preserved earlier with fresh write_file/edit_file
+	// results in middle (newer note per path wins AND moves to the newest
+	// position), so the tail cap keeps the files the model most recently touched
+	// across repeated compactions rather than dropping a just-re-edited file.
+	edits := capRecentEdits(mergeRecentEdits(preservedEditsToEntries(priorState.RecentEdits), recentEdits(middle)))
+
 	// Tools: preserve deferred tool_search schemas from the transcript. Fresh
 	// loads override older carried copies by name.
 	tools := mergeSkillEntries(preservedToolsToEntries(priorState.Tools), loadedToolSchemas(middle))
@@ -360,19 +493,76 @@
 		projectInstructionEntries(middle),
 	)
 
-	if block := formatPreservedState(plan, tools, skills, instructions); block != "" {
+	if block := formatPreservedState(plan, edits, tools, skills, instructions); block != "" {
 		summary += "\n\n" + block
 	}
 	return summary
 }
 
+// mergeRecentEdits overlays fresh edits onto edits preserved by an earlier
+// compaction. Unlike mergeSkillEntries (which keeps refreshed entries in their
+// original slot), a path touched again by a fresh edit MOVES to the newest
+// position and takes the fresh note, so capRecentEdits — which keeps only the
+// most-recent tail — never drops a file the model just re-edited. Paths present
+// only in the older set keep their relative order, ahead of the fresh ones.
+func mergeRecentEdits(older, newer []skillEntry) []skillEntry {
+	if len(newer) == 0 {
+		return older
+	}
+	freshBody := make(map[string]string, len(newer))
+	freshOrder := make([]string, 0, len(newer))
+	for _, e := range newer {
+		if _, ok := freshBody[e.name]; !ok {
+			freshOrder = append(freshOrder, e.name)
+		}
+		freshBody[e.name] = e.body
+	}
+	merged := make([]skillEntry, 0, len(older)+len(freshOrder))
+	// Older entries that were NOT re-edited keep their position, oldest first.
+	for _, e := range older {
+		if _, refreshed := freshBody[e.name]; refreshed {
+			continue // re-appended at its newest position below
+		}
+		merged = append(merged, e)
+	}
+	// Fresh edits follow, in last-seen order, each carrying the fresh note — so
+	// the most recently touched files sit at the tail the cap preserves.
+	for _, name := range freshOrder {
+		merged = append(merged, skillEntry{name: name, body: freshBody[name]})
+	}
+	return merged
+}
+
+// capRecentEdits bounds the preserved edit list to the most recent maxRecentEdits
+// entries after a merge, so repeated compactions can't grow it without limit.
+func capRecentEdits(entries []skillEntry) []skillEntry {
+	if len(entries) > maxRecentEdits {
+		return entries[len(entries)-maxRecentEdits:]
+	}
+	return entries
+}
+
+func preservedEditsToEntries(edits []preservedEdit) []skillEntry {
+	entries := make([]skillEntry, 0, len(edits))
+	for _, e := range edits {
+		if e.Path == "" {
+			continue
+		}
+		entries = append(entries, skillEntry{name: e.Path, body: e.Note})
+	}
+	return entries
+}
+
 // formatPreservedState renders state as the labelled, single-line
 // JSON block. Returns "" when there is nothing to preserve.
-func formatPreservedState(plan string, tools, skills, instructions []skillEntry) string {
-	if plan == "" && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 {
+func formatPreservedState(plan string, edits, tools, skills, instructions []skillEntry) string {
+	if plan == "" && len(edits) == 0 && len(tools) == 0 && len(skills) == 0 && len(instructions) == 0 {
 		return ""
 	}
 	state := preservedState{Plan: plan}
+	for _, e := range edits {
+		state.RecentEdits = append(state.RecentEdits, preservedEdit{Path: e.name, Note: e.body})
+	}
 	for _, t := range tools {
 		state.Tools = append(state.Tools, preservedTool{Name: t.name, Body: t.body})
 	}
@@ -393,7 +583,7 @@
 // block. JSON escaping makes this lossless even when a skill body contains
 // markdown headings, code fences, or quotes. Returns ("", nil) when absent or
 // malformed.
 func parsePreservedState(summaryContent string) (string, []skillEntry) {
 	state := parsePreservedStateBlock(summaryContent)
 	return state.Plan, preservedSkillsToEntries(state.Skills)
 }