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
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
|
// Package engine wraps the connection to the LLM engine — Docker Model Runner,
// llama-server, or any other Provider of the registry — and the generation of
// answers. Everything here is engine-agnostic: the only lines that know WHICH
// server is behind the URL live in the Provider implementations (see
// provider.go and openai_compat.go). This package used to be called `dmr`;
// measured before the split, exactly one function (Init) and two string
// constants were specific to Docker Model Runner. The rest — the whitespace
// gap, the watchdog, the retry, the command counter — already behaved the same
// against the fake engine, which is not DMR either.
package engine
import (
"context"
"encoding/json"
"fmt"
"maps"
"slices"
"strings"
"sync"
"sync/atomic"
"time"
"mm/internal/config"
"mm/internal/spinner"
"mm/internal/ui"
"github.com/firebase/genkit/go/ai"
"github.com/firebase/genkit/go/genkit"
)
// Engine is what main.go hands to the agent and to the tools: an initialised
// Genkit, the FULL model reference Genkit expects ("dmr/<id>",
// "llamacpp/<id>"), and the provider that built it — kept around for the one
// thing that stays provider-specific after start-up: explaining an error in one
// line (see Explain).
type Engine struct {
G *genkit.Genkit
Model string
Provider Provider
Backend Backend
// ContextWindow is the served context size as far as the agent knows it:
// the config's hint, or what Probe learned from the server (llama-server's
// /props). 0 = unknown. The context compression (internal/compact) reads
// it; the banner shows it with its origin, because a number without an
// origin never gets corrected.
ContextWindow int
// lastInput is the input-token count the server reported on the most
// recent call to the model, captured in Generate's WrapModel hook. It is
// the exact figure the compression trigger prefers over its estimate.
lastInput atomic.Int64
}
// New reads the provider named in the config, resolves its backend (URL, key,
// model) and opens it. Unknown providers are refused HERE, with the list of
// known names: a typo in `provider:` must not turn into a "model not found"
// twenty seconds later, once the model has tried to load.
func New(ctx context.Context) (*Engine, error) {
p, err := Lookup(config.Cfg.Provider)
if err != nil {
return nil, err
}
b, err := p.Resolve(config.Cfg)
if err != nil {
return nil, err
}
g, model, err := p.Open(ctx, b)
if err != nil {
return nil, err
}
return &Engine{G: g, Model: model, Provider: p, Backend: b, ContextWindow: b.ContextWindow}, nil
}
// Probe asks the provider what it can learn about the server BEFORE the first
// question: reachability, the context window it serves, a model that is not
// pulled. Never fatal — the user may well start the server after the agent.
//
// The context window it learns is remembered on the Engine. The provider has
// already applied the precedence — the config's value wins over the server's,
// because the operator knows what they started the server with — so what is
// stored here is the one number the whole agent reasons with.
func (e *Engine) Probe(ctx context.Context) Info {
info := e.Provider.Probe(ctx, e.Backend)
if info.ContextWindow > 0 {
e.ContextWindow = info.ContextWindow
}
return info
}
// EnsureContextWindow re-runs the provider's probe when the window is still
// unknown, and reports what it learned: the window and its origin ("/props"),
// or 0 and "" when the server still tells nothing.
//
// Why: the start-up probe runs once (main.go), and on a demo machine the
// server is often started AFTER the agent — observed: bob started before
// llama-server left `ctx: unknown` for the whole session, so the compression
// could only ever trigger on maxMessages. Asking again, lazily, fixes that
// without a background poller. Cheap by construction: the caller only asks
// while the window is unknown, and a known window is returned as-is, without
// a request — a `/props` hit per question would be noise the server does not
// deserve. Warnings are deliberately dropped here: the start-up probe already
// printed them, and repeating "nothing answers" at every question is nagging.
func (e *Engine) EnsureContextWindow(ctx context.Context) (window int, source string) {
if e.ContextWindow > 0 {
return e.ContextWindow, ""
}
info := e.Provider.Probe(ctx, e.Backend)
if info.ContextWindow > 0 {
e.ContextWindow = info.ContextWindow
return info.ContextWindow, info.ContextSource
}
return 0, ""
}
// LastInputTokens returns the server's own count of the context it last read,
// or 0 when it reports no usage (the fake engine of part 03, some proxies).
func (e *Engine) LastInputTokens() int { return int(e.lastInput.Load()) }
// ForgetInputTokens clears that count. Called after a compression: the measure
// described the history that was just replaced, and keeping it would trigger a
// second compression on the very next question.
func (e *Engine) ForgetInputTokens() { e.lastInput.Store(0) }
// Explain turns a transport error into the one-line, actionable message the
// REPL prints between brackets. See Provider.Explain for the rationale.
func (e *Engine) Explain(err error) string {
return e.Provider.Explain(err, e.Backend)
}
// blanks is the set of characters the model adds around its text without any of
// it showing on screen.
const blanks = " \t\r\n"
// gap holds back the TRAILING whitespace of the streamed text: it is only
// printed if visible text follows. Without that, the line breaks the model puts
// after its sentence are printed as-is, and the screen hollows out.
//
// The model produces MORE AND MORE of it as the turns go by — it imitates its
// own messages, which it re-reads in the history every turn. Measured: 2 blank
// lines before the first 🛠️, 4 before the second, 7, then 11. Holding the
// whitespace back is enough to remove all of it, since that whitespace is never
// followed by text: it is a tool that writes next.
type gap struct {
// The `bash` tool may call drop() from SEVERAL goroutines — Genkit runs the
// calls of a single turn in parallel (see printMu in tools).
mu sync.Mutex
held string // whitespace waiting to be followed by text
started bool // visible text has already been printed
}
// next returns what should be printed for this chunk, "" when there is nothing
// to print right now.
func (g *gap) next(text string) string {
g.mu.Lock()
defer g.mu.Unlock()
body := strings.TrimRight(text, blanks)
if body == "" { // an all-whitespace chunk: set aside, not printed
if g.started {
g.held += text
}
return ""
}
// The chunk splits in three: leading whitespace, text, trailing whitespace.
tail := text[len(body):]
visible := strings.TrimLeft(body, blanks)
lead := body[:len(body)-len(visible)]
out := ""
if g.started { // before the first word, all whitespace is dropped
out = squeeze(g.held + lead)
}
g.held, g.started = tail, true
return out + visible
}
// drop forgets the pending whitespace. Called when someone ELSE is about to
// write — a tool announcing its command: that whitespace belonged to the text
// before it, and printing it now would reopen the hole we just closed.
func (g *gap) drop() {
g.mu.Lock()
g.held = ""
g.mu.Unlock()
}
// squeeze brings any run of line breaks down to two, that is ONE blank line at
// most. Spaces are left alone: they carry the indentation of a code block, and
// cutting it would shift the displayed code.
func squeeze(s string) string {
var b strings.Builder
run := 0
for _, r := range s {
if r == '\n' {
run++
if run > 2 {
continue
}
} else {
run = 0
}
b.WriteRune(r)
}
return b.String()
}
// Generate queries the model, printing the answer as it comes (streaming). Some
// OpenAI-compatible servers sometimes cut the SSE stream abruptly ("unexpected
// end of JSON input"): if nothing has been printed AND no command has run yet,
// we retry once WITHOUT streaming — which is more robust. The "no command"
// condition matters: see the comment on the retry below.
//
// It returns the FULL conversation alongside the response — every intermediate
// turn included — because Genkit only hands back the last message. The caller
// needs the rest, or the agent forgets the commands it just ran.
//
// A spinner covers every moment where the program is waiting rather than
// printing. On a local engine the first of those is long — the model may still
// be loading, and the whole prompt has to be processed before a single token
// comes back — and without it the agent looks hung.
//
// Watchdog: if no tokens arrive for a certain duration, we assume the connection
// is hung and we cancel the context. This is the "silent loop" of
// STALL_DETECTION_PROBLEM.md — a liveness failure, not a behavioural one.
func (e *Engine) Generate(ctx context.Context, messages []*ai.Message, tools []ai.ToolRef) (*ai.ModelResponse, []*ai.Message, error) {
// Genkit only hands us back the LAST message. The intermediate turns — "I
// am calling bash", "here is the output" — exist only in the request it
// builds for the next turn. This middleware captures that request on its way
// through: the last turn seen therefore holds the whole conversation.
//
// The `WrapModel` hook wraps every call to the model — one per turn — so the
// last request seen is the most complete one. `MiddlewareFunc` adapts a
// closure to the `ai.Middleware` contract without going through a plugin:
// `New` is called once per `Generate`, which keeps `full` captured at the
// right level.
var full []*ai.Message
var gaps gap
capture := ai.MiddlewareFunc(func(context.Context) (*ai.Hooks, error) {
return &ai.Hooks{
WrapModel: func(ctx context.Context, params *ai.ModelParams, next ai.ModelNext) (*ai.ModelResponse, error) {
full = slices.Clone(params.Request.Messages)
resp, err := next(ctx, params)
// The server's own count of what it just read: the one exact
// measure of the context size, kept for the compression
// trigger. Zero means the server reported nothing.
if resp != nil && resp.Usage != nil && resp.Usage.InputTokens > 0 {
e.lastInput.Store(int64(resp.Usage.InputTokens))
}
return resp, err
},
// The tool is about to write to the screen: we forget the
// whitespace the previous text left pending. The hook lands at
// exactly the right moment — just before the tool prints its 🛠️
// line.
WrapTool: func(ctx context.Context, params *ai.ToolParams, next ai.ToolNext) (*ai.MultipartToolResponse, error) {
gaps.drop()
return next(ctx, params)
},
}, nil
})
opts := []ai.GenerateOption{
ai.WithModelName(e.Model),
ai.WithMessages(messages...),
ai.WithTools(tools...),
ai.WithMaxTurns(config.Cfg.MaxTurns), // max model ↔ tools round trips
ai.WithConfig(config.Cfg.Sampling),
ai.WithUse(capture),
}
// The prefix is printed by the first token rather than up front, so the
// spinner has a line to itself and vanishes without taking anything with it.
printed := false
emit := func(text string) {
// An event-consuming front end (ACP) takes the chunk raw: the editor
// renders Markdown and owns the layout, so the gap/squeeze machinery —
// which fights the terminal screen, not the text — stays out of the
// way. `printed` is still maintained: the no-streaming retry below must
// know whether the client already saw part of an answer.
if s := ui.ActiveSink(); s != nil {
if text == "" {
return
}
printed = true
s.Text(text)
return
}
// Trailing whitespace is held back, not printed: see the gap type.
out := gaps.next(text)
if out == "" {
return
}
// Stop BEFORE every print, not only the first one.
//
// A question can take several turns, and the tool restarts the spinner
// on its way back to the model. The text of the next turn then landed on
// the spinner's line, which erased it on the following frame
// (`\r\033[2K`): "Hello 👋 to Sam" became "⠦ Thinking… 0s 👋 to Sam".
// Worse, the last sentence of an answer vanished entirely, erased by the
// final Stop().
//
// Stop() is idempotent and returns immediately when nothing is running,
// so calling it on every chunk costs nothing.
spinner.Stop()
if !printed {
fmt.Print("🤖 ")
printed = true
}
fmt.Print(out)
}
// --- Watchdog ---
// A sub-context the watchdog cancels when no token has arrived for
// WatchdogTimeout; the streaming callback beats it on every chunk.
genCtx, beat, cancelGen := watchdog(ctx, config.Cfg.WatchdogTimeout)
defer cancelGen()
spinner.Start("Thinking")
defer spinner.Stop()
resp, err := genkit.Generate(genCtx, e.G, append(opts,
ai.WithStreaming(func(_ context.Context, chunk *ai.ModelResponseChunk) error {
beat()
emit(chunk.Text())
return nil
}),
)...)
// `full` grows as soon as a tool turn has happened: it is the witness that
// says whether commands ran, and it is used twice below.
toolsRan := len(full) > len(messages)
// The watchdog cut `genCtx` while the PARENT context is intact: this is a
// silence of the connection, not a cancellation by the user.
//
// We question the contexts rather than the error. Testing
// `err == context.Canceled` does not work: Genkit wraps the error, the
// comparison fails, and the turn then went off into the non-streamed retry —
// that is, it called the MUTE engine again. Measured: 5 s of watchdog then
// 20 s of "Retrying without streaming…" against a silent server, and the
// "stalled" message never shown.
stalled := genCtx.Err() != nil && ctx.Err() == nil
if err != nil && stalled {
// We return `full` despite the failure: if commands ran before the
// silence, their work must not be lost.
return nil, full, fmt.Errorf("generation timed out (stalled)")
}
// The retry is only worth it for an SSE stream cut before the first byte.
// Any other error (maxTurns exceeded, a failing tool, invalid arguments)
// happens AFTER commands have run, and replaying the turn would replay them:
// measured, an `echo >>` wrote itself twice. `full` grows as soon as a tool
// turn has happened, which is enough to tell the two apart.
if err != nil && !printed && !toolsRan {
spinner.Start("Retrying without streaming")
if resp, err = genkit.Generate(ctx, e.G, opts...); err == nil {
emit(resp.Text())
}
}
// `full` stops at the last call to the model; its answer is appended after.
if resp != nil && resp.Message != nil {
full = append(full, resp.Message)
}
return resp, full, err
}
// watchdog returns a child of ctx that is cancelled when beat() has not been
// called for `timeout`. This is the "silent loop" of STALL_DETECTION_PROBLEM.md
// — a liveness failure, not a behavioural one — and it is shared by Generate
// and Summarize so that both fail the same way.
//
// The timestamp is atomic because beat() runs on the streaming goroutine while
// the ticker reads it from another.
func watchdog(ctx context.Context, timeout time.Duration) (genCtx context.Context, beat func(), cancel context.CancelFunc) {
var last atomic.Int64 // UnixNano
last.Store(time.Now().UnixNano())
genCtx, cancel = context.WithCancel(ctx)
go func() {
ticker := time.NewTicker(1 * time.Second)
defer ticker.Stop()
for {
select {
case <-genCtx.Done():
return
case <-ticker.C:
if time.Since(time.Unix(0, last.Load())) > timeout {
// Stop() before printing, like everything else: without it
// the watchdog message lands on the spinner's line and the
// next frame erases it — the diagnosis would vanish at the
// very moment it is needed.
spinner.Stop()
// ui.Out, not stdout: in ACP mode this diagnosis belongs to
// the logs (stderr), never to the JSON-RPC stream.
fmt.Fprintf(ui.Out, "\n[watchdog: connection hang detected after %v]\n", timeout)
cancel()
return
}
}
}
}()
return genCtx, func() { last.Store(time.Now().UnixNano()) }, cancel
}
// Summarize asks the model for a plain-text answer to `messages` — no tools,
// no screen output. It exists for the context compression (internal/compact):
// the summary must not stream onto the screen as if it were an answer, and a
// request that offers tools invites the model to call them instead of writing.
// It goes through the same Genkit and the same model reference as Generate,
// so it lands on the same provider and the same server.
//
// The watchdog is kept, with twice the usual patience: this is the longest
// prefill of the session (the whole old history at once) and nothing streams
// until it is done. Streaming stays on so that the summary's own tokens beat
// the watchdog while it is being written.
func (e *Engine) Summarize(ctx context.Context, messages []*ai.Message, maxTokens int) (string, error) {
// The agent's sampling, minus what does not apply here: `parallel_tool_calls`
// without tools is a 400 on the OpenAI API, and `max_tokens` is the summary's
// budget, not the answer's.
sampling := maps.Clone(config.Cfg.Sampling)
if sampling == nil {
sampling = map[string]any{}
}
delete(sampling, "parallel_tool_calls")
sampling["max_tokens"] = maxTokens
genCtx, beat, cancel := watchdog(ctx, 2*config.Cfg.WatchdogTimeout)
defer cancel()
resp, err := genkit.Generate(genCtx, e.G,
ai.WithModelName(e.Model),
ai.WithMessages(messages...),
ai.WithConfig(sampling),
ai.WithStreaming(func(_ context.Context, _ *ai.ModelResponseChunk) error {
beat()
return nil
}),
)
if err != nil {
// Same test as in Generate: the contexts, not the wrapped error.
if genCtx.Err() != nil && ctx.Err() == nil {
return "", fmt.Errorf("summary timed out (stalled)")
}
return "", err
}
text := strings.TrimSpace(resp.Text())
if text == "" {
return "", fmt.Errorf("the model returned an empty summary")
}
return text, nil
}
// Names of the tools whose responses we count. They live here rather than
// scattered through the display code: the counter is a measurement, and a
// measurement needs a single definition.
const (
toolBash = "bash"
toolSkill = "read_skill"
toolRead = "read_file"
toolWrite = "write_file"
toolEdit = "edit_file"
)
// executed returns, in order, the text of the calls to a given tool that
// ACTUALLY ran.
//
// We start from the RESPONSES: a call the model asked for but which never ran
// (interrupted turn, maxTurns exceeded) produces none. The text of the call, on
// the other hand, is ONLY in the request — hence the pairing by Ref, the
// identifier the server gives to each call.
func executed(history []*ai.Message, name string) []string {
// Calls with no Ref (some servers do not set one) are paired in order, BUT
// only against the same tool name: a turn mixes `bash` and `read_skill`, and
// the first pending call is not necessarily for the right tool.
type call struct{ ref, name, text string }
var pending []call
var out []string
for _, m := range history {
for _, p := range m.Content {
switch {
case p.IsToolRequest():
pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, commandText(p.ToolRequest)})
case p.IsToolResponse() && p.ToolResponse != nil && p.ToolResponse.Name == name:
ref := p.ToolResponse.Ref
i := slices.IndexFunc(pending, func(c call) bool {
if ref != "" {
return c.ref == ref
}
return c.name == name
})
if i < 0 { // lost call: the name beats nothing
out = append(out, name)
continue
}
out = append(out, pending[i].text)
pending = slices.Delete(pending, i, i+1)
}
}
}
return out
}
// commandText describes a tool call in one line. For `bash` it is the command
// itself — the rest is only wrapping; for any other tool, its name followed by
// its arguments.
func commandText(r *ai.ToolRequest) string {
if in, ok := r.Input.(map[string]any); ok {
if c, ok := in["command"].(string); ok && strings.TrimSpace(c) != "" {
return strings.TrimSpace(c)
}
}
args, err := json.Marshal(r.Input)
if err != nil || string(args) == "null" {
return r.Name
}
return r.Name + " " + string(args)
}
// responses counts the responses of a given tool in a history.
func responses(history []*ai.Message, name string) int {
return len(executed(history, name))
}
// Commands counts the commands actually run in a history — one per response of
// the `bash` tool. It is the only reliable measure of the work done: a model
// that tells stories without acting produces a fine answer and a count of zero.
//
// The other tools do NOT count. Loading a skill runs nothing, and counting those
// reads would make the number say the opposite of what it is for: measured
// before the fix, a turn with two reads and a single command displayed
// "⚙ 3 command(s)".
func Commands(history []*ai.Message) int {
return responses(history, toolBash)
}
// CommandList returns, in order, the text of the commands Commands counts — the
// one shown in the recap under the count when `displayCommands` is on. The count
// says HOW MANY, the list says WHAT.
func CommandList(history []*ai.Message) []string {
return executed(history, toolBash)
}
// FileOps returns, in order, the file-tool calls that ran — read_file,
// write_file, edit_file — as one line each ("edit_file hello/main.go (2 edits)").
//
// They are counted APART from the commands, not folded into them: this part
// compares editing through the `edit` CLI (one bash command per edit) with
// editing through built-in tools, and a recap that counted both as "commands"
// would hide the very difference being measured. `⚙ 3 command(s) · 📝 2 file
// op(s)` keeps the two columns readable side by side.
func FileOps(history []*ai.Message) []string {
return executedAny(history, toolRead, toolWrite, toolEdit)
}
// executedAny is executed() across several tool names, keeping the order in
// which the responses appear — a per-name pass would lose the interleaving of
// a read followed by an edit.
func executedAny(history []*ai.Message, names ...string) []string {
type call struct{ ref, name, text string }
var pending []call
var out []string
isOurs := func(n string) bool { return slices.Contains(names, n) }
for _, m := range history {
for _, p := range m.Content {
switch {
case p.IsToolRequest():
pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, fileOpText(p.ToolRequest)})
case p.IsToolResponse() && p.ToolResponse != nil && isOurs(p.ToolResponse.Name):
ref, name := p.ToolResponse.Ref, p.ToolResponse.Name
i := slices.IndexFunc(pending, func(c call) bool {
if ref != "" {
return c.ref == ref
}
return c.name == name
})
if i < 0 {
out = append(out, name)
continue
}
out = append(out, pending[i].text)
pending = slices.Delete(pending, i, i+1)
}
}
}
return out
}
// fileOpText describes a file-tool call in one line: the tool, the path, and
// for edit_file the number of edits — never the texts themselves, which can
// run to dozens of lines and belong in the diff, not in the recap.
func fileOpText(r *ai.ToolRequest) string {
in, _ := r.Input.(map[string]any)
path, _ := in["path"].(string)
switch r.Name {
case toolEdit:
if edits, ok := in["edits"].([]any); ok {
return fmt.Sprintf("%s %s (%d edit(s))", r.Name, path, len(edits))
}
case toolRead:
if s, ok := in["start"].(float64); ok && s > 0 {
e, _ := in["end"].(float64)
return fmt.Sprintf("%s %s %d-%d", r.Name, path, int(s), int(e))
}
}
if path == "" {
return commandText(r)
}
return r.Name + " " + path
}
// Skills counts the skills loaded — one per response of `read_skill`. Shown
// next to the command count: seeing which procedure the agent chose to follow is
// half the point of this part.
func Skills(history []*ai.Message) int {
return responses(history, toolSkill)
}
|