Maven/cmd/mavend/modelswap.go

package main

import (
	"context"
	"fmt"
	"log"
	"path/filepath"
	"time"

	"github.com/kami/maven/internal/config"
	"github.com/kami/maven/internal/ipc"
	"github.com/kami/maven/internal/llm"
	"github.com/kami/maven/internal/phraser"
)

// Swapping the resident model while the daemon runs (Vikunja #250).
//
// Off unless configured: with no phraser.swap_models allowlist the two IPC
// methods are never wired, so they answer ErrUnknownMethod. When it is wired the
// swap method is AuthStepUp (internal/auth), which means an authed human surface
// only — there is no act, no intent and no timer that reaches it. The daemon
// never decides to change its own brain.
//
// The allowlist is exact-match against paths a human wrote in mavend.json. The
// request carries a path and llama-server is started with it as `-m`, so
// anything looser would turn "swap the model" into "load any file on my disk".
func wireModelSwap(srv *ipc.Server, phr phraser.Phraser, cfg *config.Config) {
	if cfg.Phraser == nil || len(cfg.Phraser.SwapModels) == 0 {
		return
	}
	lp, ok := phr.(*phraser.LLMPhraser)
	if !ok {
		log.Printf("model swap: phraser.swap_models is set but there is no llama-server phraser — swap disabled")
		return
	}
	allowed := map[string]bool{}
	for _, m := range cfg.Phraser.SwapModels {
		allowed[filepath.Clean(m)] = true
	}
	// The configured model is always swappable back to, listed or not: the way
	// out of a bad swap must not depend on remembering to allowlist the model
	// you are already running.
	allowed[filepath.Clean(cfg.Phraser.ModelPath)] = true

	srv.SwapModelFn = func(ctx context.Context, req ipc.SwapModelReq) (ipc.SwapModelResp, error) {
		path := filepath.Clean(req.ModelPath)
		if !allowed[path] {
			log.Printf("model swap: REFUSED %q — not in phraser.swap_models", req.ModelPath)
			return ipc.SwapModelResp{}, fmt.Errorf("%w: %q is not in phraser.swap_models", ipc.ErrForbidden, req.ModelPath)
		}
		res, err := lp.Swap(ctx, phraser.SwapSpec{
			ModelPath:  path,
			NGpuLayers: req.NGpuLayers,
			NCtx:       req.NCtx,
		})
		resp := ipc.SwapModelResp{
			Model:      res.Model,
			ModelPath:  res.ModelPath,
			BaseURL:    res.BaseURL,
			RolledBack: res.RolledBack,
			TookMs:     res.Took.Milliseconds(),
		}
		if err != nil {
			// A rolled-back swap is a failure that left a working daemon behind.
			// Both halves matter to the caller, so the response is filled in even
			// though the error is returned.
			log.Printf("model swap: %v", err)
			return resp, err
		}
		return resp, nil
	}

	srv.ModelStatusFn = func(ctx context.Context) (ipc.ModelStatusResp, error) {
		path, ngl, nctx := lp.LiveModel()
		base := lp.BaseURL()
		resp := ipc.ModelStatusResp{
			ModelPath:  path,
			BaseURL:    base,
			NGpuLayers: ngl,
			NCtx:       nctx,
			Swappable:  cfg.Phraser.SwapModels,
		}
		if base == "" {
			resp.Model = llm.UnknownModel
			return resp, nil
		}
		id, err := llm.ModelID(ctx, base)
		if err != nil {
			// Report the honest "I could not confirm it" rather than echoing the
			// configured filename as if the server had said it.
			resp.Model = llm.UnknownModel
			return resp, nil
		}
		resp.Model = id
		return resp, nil
	}

	log.Printf("model swap: enabled, %d allowlisted model(s) — step-up required", len(cfg.Phraser.SwapModels))
}

// llmClientFor builds a completion client on the phraser's llama-server and
// keeps it pointed at the right one across a model swap.
//
// Without the OnSwap registration every holder of a base URL — the LLM router,
// the replier, the mail extractor, the memory evaluator — would keep talking to
// the port of a server that no longer exists, and the daemon would degrade to
// the classifier permanently after the first swap. The client is re-pointed, not
// rebuilt, so nothing that holds it has to know a swap happened.
func llmClientFor(lp *phraser.LLMPhraser, timeout time.Duration) *llm.Client {
	c := llm.New(lp.BaseURL(), timeout)
	lp.OnSwap(func(base string) { c.SetBaseURL(base) })
	return c
}