1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307

// Package compact keeps the conversation history inside the model's context
// window. When the history grows past a threshold, the OLD question turns are
// replaced by ONE summary written by the model itself, and the last N turns
// stay raw — the hybrid of CONTEXT_WINDOW.md § 2.e.
//
// Why it exists: nothing in the agent ever shortens `messages`. Measured with
// the fake engine of part 03, the history sent to the model grew 2 → 5 → 7 → 9
// messages over four requests, and every `bash` output (up to maxOutput
// characters) stays until /quit. On a local model the window is fixed at load
// time, so a long session first shows up as a false watchdog stall, then as a
// 400 from the engine.
//
// The package never calls the model directly: it builds the summary request
// and hands it to a Summarizer (engine.Summarize in production, a function
// returning a fixed text in tests). It knows nothing about the screen either —
// the report line is the agent's job.
package compact

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"

	"mm/internal/config"

	"github.com/firebase/genkit/go/ai"
)

// Summarizer answers a request built by this package — a system message, the
// old turns, then the prompt as a user message — with plain text.
type Summarizer func(ctx context.Context, request []*ai.Message) (string, error)

// Result describes one compression, for the one-line report.
type Result struct {
	Messages   []*ai.Message
	Compressed int // messages replaced by the summary pair
	Kept       int // raw messages kept after the summary (system excluded)
	Before     int // estimated tokens before
	After      int // estimated tokens after
	Elapsed    time.Duration
}

// Report is the one-line account of a compression, shared by the two front
// ends so the terminal and the editor say the same thing:
//
//	compressed 12 messages → 1 summary + 4 kept (18.2k → 6.1k tokens, 3.4s)
func (r Result) Report() string {
	return fmt.Sprintf("compressed %d messages → 1 summary + %d kept (%s → %s tokens, %s)",
		r.Compressed, r.Kept, kilo(r.Before), kilo(r.After), seconds(r.Elapsed))
}

// kilo prints a token count the way the banner does: 18.2k, not 18234.
func kilo(n int) string {
	if n < 1000 {
		return fmt.Sprintf("%d", n)
	}
	return fmt.Sprintf("%.1fk", float64(n)/1000)
}

// seconds prints a duration with one decimal under ten seconds, none above.
func seconds(d time.Duration) string {
	if d < 10*time.Second {
		return fmt.Sprintf("%.1fs", d.Seconds())
	}
	return fmt.Sprintf("%ds", int(d.Seconds()))
}

// ErrNothingToCompact is returned when the history holds no turn older than
// the ones to keep. It is not a failure: the threshold can be reached inside a
// single long turn, and there is then nothing to summarise yet.
var ErrNothingToCompact = errors.New("nothing to compact")

// metaKey marks the two messages a compression inserts, so that a later
// compression recognises an earlier summary (and merges it instead of
// summarising it a second time — detail decays exponentially otherwise).
const metaKey = "compact"

// Decision says whether to compress before the next question, and on what
// grounds. Tokens is the figure that was compared to the threshold; it is
// reported even when nothing triggers, for the /context-style displays.
type Decision struct {
	Compact bool
	Tokens  int
	Reason  string // "tokens", "messages", or "" when nothing triggered
}

// Decide applies the automatic trigger.
//
// `measured` is the input-token count the engine reported for the last call
// to the model (0 when it reported nothing, as the fake engine does). The
// larger of it and the local estimate is used: the measure sees what the
// estimate cannot (the tools' JSON, the chat template), but it lags one call
// behind and dies with the compression that made it stale — the caller
// forgets it after a successful compression.
//
// `window` is the served context size as far as the agent knows it: the
// yaml's `contextWindow`, or what the provider's probe learned from the server
// (llama-server's /props) — see engine.Engine.ContextWindow. 0 = unknown, and
// then only MaxMessages can trigger.
func Decide(msgs []*ai.Message, measured, window int, cfg config.ContextConfig) Decision {
	d := Decision{Tokens: max(Estimate(msgs), measured)}
	if !cfg.Enabled {
		return d
	}
	switch {
	// Cross-multiplied to stay in integers: tokens ≥ window × threshold / 100.
	case window > 0 && d.Tokens*100 >= window*cfg.Threshold:
		d.Compact, d.Reason = true, "tokens"
	case cfg.MaxMessages > 0 && len(msgs) >= cfg.MaxMessages:
		d.Compact, d.Reason = true, "messages"
	}
	return d
}

// Compact replaces every turn but the last cfg.KeepLastTurns by a summary pair
// (a user message carrying the notes, a short model acknowledgement), inserted
// right after the system prompt. The kept turns are reused as they are, so the
// tool_call/tool_result pairs inside them cannot be split; the cut itself falls
// on a turn boundary, never on a message count.
//
// On any error the caller's history is left untouched: the returned Result is
// empty and `msgs` was never modified.
func Compact(ctx context.Context, msgs []*ai.Message, cfg config.ContextConfig, summarize Summarizer) (Result, error) {
	start := time.Now()

	head, turns := Split(msgs)
	keep := max(cfg.KeepLastTurns, 1)
	if len(turns) <= keep {
		return Result{}, ErrNothingToCompact
	}
	old, recent := turns[:len(turns)-keep], turns[len(turns)-keep:]

	// An earlier summary that is the only old turn: re-summarising a summary
	// loses detail and gains nothing.
	merge := IsSummary(old[0][0])
	if merge && len(old) == 1 {
		return Result{}, ErrNothingToCompact
	}

	oldMsgs := flatten(old)
	text, err := summarize(ctx, Request(oldMsgs, cfg.Prompt, merge))
	if err != nil {
		return Result{}, err
	}
	if strings.TrimSpace(text) == "" {
		return Result{}, errors.New("empty summary")
	}

	questions := len(old)
	if merge {
		questions-- // the summary pair is not a question the user asked
	}
	pair := SummaryPair(text, len(oldMsgs), questions, commands(oldMsgs))

	out := make([]*ai.Message, 0, len(head)+len(pair)+len(msgs))
	out = append(out, head...)
	out = append(out, pair...)
	kept := flatten(recent)
	out = append(out, kept...)

	// The safety net: the cut above cannot split a pair, but the check is
	// cheap and the failure it guards against is a 400 on the next question.
	if err := Valid(out); err != nil {
		return Result{}, fmt.Errorf("compressed history is invalid: %w", err)
	}

	return Result{
		Messages:   out,
		Compressed: len(oldMsgs),
		Kept:       len(kept),
		Before:     Estimate(msgs),
		After:      Estimate(out),
		Elapsed:    time.Since(start),
	}, nil
}

// Split separates the leading system message(s) from the question turns. A
// turn starts at a user message and runs until the next one, so it holds the
// whole model ↔ tools exchange the question caused; an earlier summary pair is
// a turn of its own, since it starts with a user message too.
//
// Every strategy of CONTEXT_WINDOW.md cuts on these boundaries: cutting on a
// message count can leave a `tool` message without the `assistant` that asked
// for it, which the OpenAI API refuses outright.
func Split(msgs []*ai.Message) (head []*ai.Message, turns [][]*ai.Message) {
	i := 0
	for i < len(msgs) && msgs[i].Role == ai.RoleSystem {
		i++
	}
	head = msgs[:i]
	for _, m := range msgs[i:] {
		if m.Role == ai.RoleUser || len(turns) == 0 {
			turns = append(turns, nil)
		}
		turns[len(turns)-1] = append(turns[len(turns)-1], m)
	}
	return head, turns
}

// Valid checks what an OpenAI-compatible server checks before answering: the
// history starts with the system prompt, every tool request sits in a model
// message, and every tool response answers a request seen earlier — paired by
// Ref when the server set one, by tool name otherwise (same rule as
// engine.executed).
func Valid(msgs []*ai.Message) error {
	if len(msgs) == 0 || msgs[0].Role != ai.RoleSystem {
		return errors.New("history does not start with the system prompt")
	}
	type call struct{ ref, name string }
	var pending []call
	for i, m := range msgs {
		for _, p := range m.Content {
			switch {
			case p.IsToolRequest() && p.ToolRequest != nil:
				if m.Role != ai.RoleModel {
					return fmt.Errorf("message %d: tool request in a %s message", i, m.Role)
				}
				pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name})
			case p.IsToolResponse() && p.ToolResponse != nil:
				r := p.ToolResponse
				j := slices.IndexFunc(pending, func(c call) bool {
					if r.Ref != "" {
						return c.ref == r.Ref
					}
					return c.name == r.Name
				})
				if j < 0 {
					return fmt.Errorf("message %d: tool response %q answers no pending request", i, r.Name)
				}
				pending = slices.Delete(pending, j, j+1)
			}
		}
	}
	return nil
}

// Estimate approximates the token count of a history without asking the
// engine: characters divided by 3.5, plus a few tokens of framing per message.
//
// 4 characters per token holds for English prose; code, JSON, paths and French
// tokenise closer to 3. Taking 3.5 over-estimates prose a little, which is the
// safe side for a threshold. The engine's own count, when it reports one, is
// preferred by Decide — this is the fallback, and the number the report shows
// for a history that was never sent.
func Estimate(msgs []*ai.Message) int {
	chars := 0
	for _, m := range msgs {
		for _, p := range m.Content {
			switch {
			case p.IsText():
				chars += len(p.Text)
			case p.IsToolRequest() && p.ToolRequest != nil:
				chars += len(p.ToolRequest.Name) + len(jsonOf(p.ToolRequest.Input))
			case p.IsToolResponse() && p.ToolResponse != nil:
				chars += len(p.ToolResponse.Name) + len(jsonOf(p.ToolResponse.Output))
			}
		}
	}
	return (chars*2+6)/7 + 4*len(msgs)
}

// jsonOf renders a tool payload the way the request will carry it. A string
// output is taken as-is: marshalling it would count the escaping twice.
func jsonOf(v any) string {
	if s, ok := v.(string); ok {
		return s
	}
	b, err := json.Marshal(v)
	if err != nil {
		return fmt.Sprint(v)
	}
	return string(b)
}

// IsSummary reports whether a message is the notes half of a summary pair.
func IsSummary(m *ai.Message) bool {
	return m != nil && m.Role == ai.RoleUser && m.Metadata[metaKey] == "summary"
}

// flatten joins turns back into one message list.
func flatten(turns [][]*ai.Message) []*ai.Message {
	var out []*ai.Message
	for _, t := range turns {
		out = append(out, t...)
	}
	return out
}

// commands counts the bash commands that actually ran in a history — one per
// response of the `bash` tool, the same measure as engine.Commands. It is not
// imported from engine so that this package stays free of Genkit's wiring and
// testable with a plain function.
func commands(msgs []*ai.Message) int {
	n := 0
	for _, m := range msgs {
		for _, p := range m.Content {
			if p.IsToolResponse() && p.ToolResponse != nil && p.ToolResponse.Name == "bash" {
				n++
			}
		}
	}
	return n
}