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

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"mm/internal/ask"
	"mm/internal/config"
	"mm/internal/engine"
	"mm/internal/spinner"
)

// autoPull is `autoPull: true` (ticket 61): before anything else starts, it
// downloads the models this mm will use and the server does not have yet —
// its own, then its mmConfigs agents' and its routing deciders' — with the
// progress on w (stderr: never the answer of -p, never the ACP stream).
//
// Only the agent's own model can stop mm: without it there is nothing to
// run. Another model that cannot be pulled is a warning — that agent or that
// decider fails later, where it is used. An asked agent (a child of an /ask)
// pulls nothing: its asker already did, and its start-up is bounded.
func autoPull(ctx context.Context, e *engine.Engine, w io.Writer, tty bool) error {
	if !config.Cfg.AutoPull || ask.Depth() > 0 {
		return nil
	}
	report := newPullReport(w, tty)
	pulled, err := e.EnsureModel(ctx, report.step)
	report.done(e.Backend.Model, pulled, err)
	switch {
	case errors.Is(err, engine.ErrCannotPull), errors.Is(err, engine.ErrServerDown):
		fmt.Fprintf(w, "[warning: autoPull: %s (%s): %v]\n", e.Backend.Model, e.Backend.Provider, err)
	case err != nil:
		return fmt.Errorf("autoPull: %s: %w", e.Backend.Model, err)
	}
	for _, path := range otherModelFiles() {
		model, pulled, err := engine.EnsureAgentModel(ctx, path, report.step)
		report.done(model, pulled, err)
		// Every failure is worth a line here: the model is NOT on its
		// server (a listed one returns no error), and for a provider that
		// cannot pull, the line says what to do instead.
		if err != nil {
			fmt.Fprintf(w, "[warning: autoPull: %s (%s): %v]\n", model, path, err)
		}
	}
	return nil
}

// otherModelFiles lists the config files of the other models this mm will
// use: its mmConfigs agents (in name order), then its routing deciders.
func otherModelFiles() []string {
	var paths []string
	for _, name := range config.AgentNames() {
		paths = append(paths, config.Cfg.MMConfigs[name].Path)
	}
	for _, decider := range []string{config.Cfg.SkillRouting.Decider, config.Cfg.ToolRouting.Decider} {
		if decider != "" {
			paths = append(paths, decider)
		}
	}
	return paths
}

// pullReport shows a download. On a terminal, one line rewritten in place;
// elsewhere (a log, an editor's stderr) a line per tenth, so a long download
// does not write a thousand lines.
type pullReport struct {
	w       io.Writer
	tty     bool
	started time.Time
	shown   int // the last tenth printed, off a terminal
	live    bool
}

func newPullReport(w io.Writer, tty bool) *pullReport {
	return &pullReport{w: w, tty: tty, shown: -1}
}

// step is engine.PullProgress's receiver.
func (r *pullReport) step(p engine.PullProgress) {
	if !r.live {
		r.live, r.started, r.shown = true, time.Now(), -1
	}
	line := "⬇ pulling " + p.Model
	tenth := 0
	if p.Total > 0 {
		line += fmt.Sprintf(" — %s / %s (%d %%)", size(p.Completed, p.Total), size(p.Total, p.Total), p.Completed*100/p.Total)
		tenth = int(p.Completed * 10 / p.Total)
	}
	if r.tty {
		fmt.Fprintf(r.w, "\r%s\033[K", line)
		return
	}
	if tenth > r.shown {
		r.shown = tenth
		fmt.Fprintln(r.w, line)
	}
}

// done closes a download's line: what was pulled and how long it took, or
// nothing when the model was already there.
func (r *pullReport) done(model string, pulled bool, err error) {
	if r.live && r.tty {
		fmt.Fprint(r.w, "\r\033[K")
	}
	if pulled && err == nil {
		fmt.Fprintf(r.w, "⬇ pulled %s in %s\n", model, elapsed(time.Since(r.started)))
	}
	r.live = false
}

// size prints n in the unit that suits the whole download: GB from 1 GB
// (8.1 GB, as the model cards say), MB below (145 MB rather than 0.1 GB).
func size(n, total int64) string {
	if total >= 1e9 {
		return fmt.Sprintf("%.1f GB", float64(n)/1e9)
	}
	return fmt.Sprintf("%d MB", n/1e6)
}

// elapsed gives tenths under ten seconds (a small model pulls in 0.8s, not
// "0s"), whole seconds beyond.
func elapsed(d time.Duration) string {
	if d < 10*time.Second {
		return fmt.Sprintf("%.1fs", d.Seconds())
	}
	return d.Round(time.Second).String()
}

// stderrIsTerminal reports whether the progress can be rewritten in place.
func stderrIsTerminal() bool {
	return spinner.Styled() && isCharDevice(os.Stderr)
}

func isCharDevice(f *os.File) bool {
	info, err := f.Stat()
	return err == nil && info.Mode()&os.ModeCharDevice != 0
}