Add managed local embedding helper

Author: rowantrollope

AGENTS.md

@@ -262,10 +262,16 @@ The most important implementation seams are:
 - Semantic query embeddings must use a real provider. QMD uses a local GGUF
   embedding model with explicit query/document formatting; deterministic hash
   embeddings are acceptable only as test doubles, never as product behavior.
+- Pure local embedding model work means downloading and caching the GGUF
+  artifact itself. Do not substitute Ollama or another local service shim when
+  the requirement is "pure GGUF".
+- Local GGUF embeddings are an explicit provider override using a managed Node
+  helper with `node-llama-cpp`, matching QMD's model lifecycle. Do not swap this
+  back to per-call `llama-embedding` process invocations.
 - Semantic query provider/model settings are global runtime settings, not
   workspace config. Embeddings should be treated as on by default; if provider
-  credentials are missing, return/report unavailable status without hard-failing
-  normal query flows.
+  runtime dependencies are missing, return/report unavailable status without
+  hard-failing normal query flows.
 - `AFS_EMBED_MODEL`, `AFS_EMBED_PROVIDER`, `AFS_EMBED_DIMENSIONS`, and
   `OPENAI_API_KEY` are read by the control-plane process. CLI help and
   troubleshooting copy must not imply that setting them only on an `afs query`

README.md

@@ -216,12 +216,16 @@ If you want to search workspace contents:
 currently falls back to keyword-ranked results until hybrid vector/rerank is
 complete. Keyword ranking uses RedisSearch BM25 query chunks when available,
 then falls back to direct content ranking. Use `query --semantic` for
-vector-only retrieval. Semantic embeddings are globally enabled and use OpenAI
-when `OPENAI_API_KEY` is available in the control-plane environment. Override
-the default `openai:text-embedding-3-small` model with `AFS_EMBED_MODEL` in the
-same environment, then restart the control plane. Semantic queries read
-existing embedding indexes; imports start embedding creation in the background,
-and existing workspaces can be prepared with
+vector-only retrieval. Semantic embeddings are globally enabled and default to
+OpenAI when `OPENAI_API_KEY` is set in the control-plane environment. For the
+explicit local provider, `afs query model download` asks the control plane to
+resolve, download, and load
+`hf:ggml-org/embeddinggemma-300M-GGUF/embeddinggemma-300M-Q8_0.gguf` through the
+same `node-llama-cpp` helper used for local indexing. On AFS Cloud, only an
+admin identity can trigger the server-side local model warm-up. Override model
+selection with `AFS_EMBED_MODEL` in the control-plane environment, then restart
+the control plane. Semantic queries read existing embedding indexes; imports start
+embedding creation in the background, and existing workspaces can be prepared with
 `afs fs <workspace> query index create --embeddings --wait`.
 
 If you want commands with an optional workspace argument to use `my-repo` by

cmd/afs/afs_commands_test.go

@@ -138,6 +138,14 @@ func (s stubAFSControlPlane) RebuildQueryIndex(context.Context, string, controlp
 	return controlplane.WorkspaceQueryIndexRebuildResponse{}, fmt.Errorf("unexpected RebuildQueryIndex call")
 }
 
+func (s stubAFSControlPlane) QueryModelStatus(context.Context, controlplane.QueryModelStatusRequest) (controlplane.QueryModelStatus, error) {
+	return controlplane.QueryModelStatus{}, fmt.Errorf("unexpected QueryModelStatus call")
+}
+
+func (s stubAFSControlPlane) DownloadQueryModel(context.Context, controlplane.QueryModelDownloadRequest) (controlplane.QueryModelDownloadResult, error) {
+	return controlplane.QueryModelDownloadResult{}, fmt.Errorf("unexpected DownloadQueryModel call")
+}
+
 func (s stubAFSControlPlane) DiffWorkspace(context.Context, string, string, string) (controlplane.WorkspaceDiffResponse, error) {
 	return controlplane.WorkspaceDiffResponse{}, fmt.Errorf("unexpected DiffWorkspace call")
 }

cmd/afs/afs_query_commands.go

@@ -68,6 +68,9 @@ func runWorkspaceQuery(mode, workspace string, args []string) error {
 	if mode == mcptools.FileQueryModeHybrid && isWorkspaceQueryIndexInvocation(args) {
 		return runWorkspaceQueryIndex(workspace, args[1:])
 	}
+	if mode == mcptools.FileQueryModeHybrid && workspace == "" && isWorkspaceQueryModelInvocation(args) {
+		return runWorkspaceQueryModel(args[1:])
+	}
 	opts, err := parseWorkspaceQueryArgs(mode, args)
 	if err != nil {
 		return err
@@ -84,6 +87,121 @@ func runWorkspaceQuery(mode, workspace string, args []string) error {
 	return runWorkspaceQueryRequest(ctx, remote, opts, request)
 }
 
+func isWorkspaceQueryModelInvocation(args []string) bool {
+	if len(args) == 0 || args[0] != "model" {
+		return false
+	}
+	if len(args) == 1 || isHelpArg(args[1]) {
+		return true
+	}
+	switch strings.TrimSpace(args[1]) {
+	case "status", "download":
+		return true
+	default:
+		return false
+	}
+}
+
+func runWorkspaceQueryModel(args []string) error {
+	if len(args) == 0 || isHelpArg(args[0]) {
+		fmt.Fprint(os.Stderr, workspaceQueryModelUsageText(filepath.Base(os.Args[0])))
+		return nil
+	}
+	switch strings.TrimSpace(args[0]) {
+	case "status":
+		return runWorkspaceQueryModelStatus(args[1:])
+	case "download":
+		return runWorkspaceQueryModelDownload(args[1:])
+	default:
+		return fmt.Errorf("unknown query model subcommand %q\n\n%s", args[0], workspaceQueryModelUsageText(filepath.Base(os.Args[0])))
+	}
+}
+
+func runWorkspaceQueryModelStatus(args []string) error {
+	fs := flag.NewFlagSet("query model status", flag.ContinueOnError)
+	fs.SetOutput(io.Discard)
+	var jsonOut bool
+	var model string
+	fs.BoolVar(&jsonOut, "json", false, "write JSON output")
+	fs.StringVar(&model, "model", "", "local model id")
+	if err := fs.Parse(args); err != nil || fs.NArg() != 0 {
+		return fmt.Errorf("%s", workspaceQueryModelUsageText(filepath.Base(os.Args[0])))
+	}
+	ctx := context.Background()
+	_, service, closeFn, err := openAFSControlPlane(ctx)
+	if err != nil {
+		return err
+	}
+	defer closeFn()
+	status, err := service.QueryModelStatus(ctx, controlplane.QueryModelStatusRequest{Model: model})
+	if err != nil {
+		return workspaceQueryModelControlPlaneError(err)
+	}
+	if jsonOut {
+		enc := json.NewEncoder(os.Stdout)
+		enc.SetIndent("", "  ")
+		return enc.Encode(status)
+	}
+	fmt.Fprintln(os.Stdout, "Query model")
+	fmt.Fprintln(os.Stdout)
+	fmt.Fprintf(os.Stdout, "model      %s\n", status.Spec.ID)
+	fmt.Fprintf(os.Stdout, "cache_dir  %s\n", status.CacheDir)
+	fmt.Fprintf(os.Stdout, "path       %s\n", status.Path)
+	fmt.Fprintf(os.Stdout, "downloaded %t\n", status.Exists)
+	if status.SizeBytes > 0 {
+		fmt.Fprintf(os.Stdout, "size       %s\n", formatBytes(status.SizeBytes))
+	}
+	return nil
+}
+
+func runWorkspaceQueryModelDownload(args []string) error {
+	fs := flag.NewFlagSet("query model download", flag.ContinueOnError)
+	fs.SetOutput(io.Discard)
+	var jsonOut bool
+	var model string
+	fs.BoolVar(&jsonOut, "json", false, "write JSON output")
+	fs.StringVar(&model, "model", "", "local model id")
+	if err := fs.Parse(args); err != nil || fs.NArg() != 0 {
+		return fmt.Errorf("%s", workspaceQueryModelUsageText(filepath.Base(os.Args[0])))
+	}
+	ctx := context.Background()
+	_, service, closeFn, err := openAFSControlPlane(ctx)
+	if err != nil {
+		return err
+	}
+	defer closeFn()
+	if !jsonOut {
+		fmt.Fprintln(os.Stderr, "Resolving local embedding model on the control plane...")
+	}
+	result, err := service.DownloadQueryModel(ctx, controlplane.QueryModelDownloadRequest{Model: model})
+	if err != nil {
+		return workspaceQueryModelControlPlaneError(err)
+	}
+	if jsonOut {
+		enc := json.NewEncoder(os.Stdout)
+		enc.SetIndent("", "  ")
+		return enc.Encode(result)
+	}
+	fmt.Fprintln(os.Stdout, "Query model")
+	fmt.Fprintln(os.Stdout)
+	fmt.Fprintf(os.Stdout, "model      %s\n", result.Spec.ID)
+	fmt.Fprintf(os.Stdout, "cache_dir  %s\n", result.CacheDir)
+	fmt.Fprintf(os.Stdout, "path       %s\n", result.Path)
+	fmt.Fprintf(os.Stdout, "cached     %t\n", result.Exists)
+	fmt.Fprintf(os.Stdout, "resolved   %t\n", result.Downloaded || result.Exists)
+	if result.SizeBytes > 0 {
+		fmt.Fprintf(os.Stdout, "size       %s\n", formatBytes(result.SizeBytes))
+	}
+	return nil
+}
+
+func workspaceQueryModelControlPlaneError(err error) error {
+	if errors.Is(err, os.ErrNotExist) {
+		return fmt.Errorf("query model routes are not available on this control plane; rebuild and restart afs-control-plane")
+	}
+	return err
+}
+
 func isWorkspaceQueryIndexInvocation(args []string) bool {
 	if len(args) == 0 || args[0] != "index" {
 		return false
@@ -726,16 +844,18 @@ func workspaceQueryUsageText(bin, mode string) string {
 	return brandHeaderString() + fmt.Sprintf(`Usage:
   %s query [flags] <query>
   %s fs [workspace] query [flags] <query>
-  %s query index <status|rebuild|clean> [flags]
+  %s query index <status|create|rebuild|clean> [flags]
+  %s query model <status|download> [flags]
 
 QMD-style hybrid + rerank workspace query.
 Plain text runs hybrid retrieval by default. Use --keyword for keyword-ranked
 retrieval only, or --semantic for vector-only semantic search.
 
 Default query currently falls back to keyword ranked results until hybrid
 vector/rerank is complete. Use --semantic for vector-only retrieval. Semantic
-embeddings are globally enabled and use OpenAI when OPENAI_API_KEY is set in
-the control-plane environment.
+embeddings are globally enabled and default to OpenAI when OPENAI_API_KEY is set
+in the control-plane environment. Set AFS_EMBED_PROVIDER=local to use the local
+GGUF helper.
 Use grep when you know the exact text.
 
 Typed query documents:
@@ -770,6 +890,33 @@ Examples:
   %s query --keyword "checkpoint savepoint"
   %s query --semantic "how do I save a snapshot?"
   %s query index status
+  %s query model download
   %s fs repo query $'lex: checkpoint\nvec: how do I save a snapshot?'
-`, bin, bin, bin, bin, bin, bin, bin, bin)
+`, bin, bin, bin, bin, bin, bin, bin, bin, bin, bin)
+}
+
+func workspaceQueryModelUsageText(bin string) string {
+	return brandHeaderString() + fmt.Sprintf(`Usage:
+  %[1]s query model <status|download> [flags]
+
+Manage the control-plane global local GGUF embedding model cache.
+
+Subcommands:
+  status             Show the configured local model and expected cache path
+  download           Ask the Node helper to resolve/download/load the model now
+
+Flags:
+  --model <model>    Model id, default hf:ggml-org/embeddinggemma-300M-GGUF/embeddinggemma-300M-Q8_0.gguf
+  --json             Write JSON output
+
+Environment:
+  AFS_EMBED_MODEL_DIR   Control-plane cache directory override
+  AFS_EMBED_HELPER_CMD  Control-plane Node.js command override
+  AFS_NODE_LLAMA_CPP_MODULE
+                        node-llama-cpp module specifier override
+
+Examples:
+  %[1]s query model status
+  %[1]s query model download
+`, bin)
 }

cmd/afs/afs_query_commands_test.go

@@ -6,6 +6,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"path/filepath"
 	"strings"
 	"testing"
 
@@ -124,6 +125,40 @@ func TestWorkspaceQueryIndexInvocationIncludesCreate(t *testing.T) {
 	}
 }
 
+func TestWorkspaceQueryModelInvocationAndStatus(t *testing.T) {
+	for _, args := range [][]string{
+		{"model"},
+		{"model", "status"},
+		{"model", "download"},
+	} {
+		if !isWorkspaceQueryModelInvocation(args) {
+			t.Fatalf("isWorkspaceQueryModelInvocation(%#v) = false, want true", args)
+		}
+	}
+	if isWorkspaceQueryModelInvocation([]string{"model", "files"}) {
+		t.Fatal("isWorkspaceQueryModelInvocation(model files) = true, want natural query")
+	}
+
+	_, _, closeStore := setupAFSGrepTest(t)
+	defer closeStore()
+	t.Setenv("AFS_EMBED_MODEL_DIR", t.TempDir())
+	output, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "model", "status"})
+	})
+	if err != nil {
+		t.Fatalf("cmdQuery(model status) returned error: %v", err)
+	}
+	for _, want := range []string{
+		"Query model",
+		"hf:ggml-org/embeddinggemma-300M-GGUF/embeddinggemma-300M-Q8_0.gguf",
+		"downloaded false",
+	} {
+		if !strings.Contains(output, want) {
+			t.Fatalf("output = %q, missing %q", output, want)
+		}
+	}
+}
+
 func TestParseWorkspaceQueryArgsRejectsModeFlagsWithTypedDocuments(t *testing.T) {
 	_, err := parseWorkspaceQueryArgs(mcptools.FileQueryModeHybrid, []string{
 		"--semantic",
@@ -151,8 +186,9 @@ func TestParseWorkspaceQueryArgsRejectsKeywordAndSemanticTogether(t *testing.T)
 	}
 }
 
-func TestCmdQuerySemanticMissingKeyFailsClearly(t *testing.T) {
-	t.Setenv("OPENAI_API_KEY", "")
+func TestCmdQuerySemanticMissingLocalHelperFailsClearly(t *testing.T) {
+	t.Setenv("AFS_EMBED_PROVIDER", "local")
+	t.Setenv("AFS_EMBED_HELPER_CMD", filepath.Join(t.TempDir(), "missing-node"))
 	_, _, closeStore := setupAFSGrepTest(t)
 	defer closeStore()
 
@@ -162,8 +198,8 @@ func TestCmdQuerySemanticMissingKeyFailsClearly(t *testing.T) {
 	if err == nil {
 		t.Fatal("cmdQuery(--semantic) returned nil error, want unavailable")
 	}
-	if !strings.Contains(err.Error(), "OPENAI_API_KEY") {
-		t.Fatalf("error = %q, want OPENAI_API_KEY guidance", err)
+	if !strings.Contains(err.Error(), "local embedding helper runtime") {
+		t.Fatalf("error = %q, want local helper guidance", err)
 	}
 }
 
@@ -643,7 +679,7 @@ func TestCmdQueryIndexStatusReportsGlobalEmbeddingStatus(t *testing.T) {
 	if status.Workspace != "repo" || status.Keyword.Files != 0 {
 		t.Fatalf("status = %+v, want empty repo keyword status", status)
 	}
-	if !status.Embeddings.Enabled || status.Embeddings.Model != "openai:text-embedding-3-small" || status.Embeddings.Available {
+	if !status.Embeddings.Enabled || status.Embeddings.Provider != "openai" || status.Embeddings.Model != "openai:text-embedding-3-small" || status.Embeddings.Available {
 		t.Fatalf("embedding status = %+v, want global OpenAI unavailable without key", status)
 	}
 }
@@ -844,6 +880,7 @@ func TestCmdQueryContractCoversHybridFallbacksAndIndexDisambiguation(t *testing.
 	})
 
 	t.Run("semantic JSON reports unavailable without failing command", func(t *testing.T) {
+		t.Setenv("AFS_EMBED_PROVIDER", "")
 		t.Setenv("OPENAI_API_KEY", "")
 		output, err := captureStdout(t, func() error {
 			return cmdQuery([]string{"query", "--semantic", "--json", "workspace recovery"})

cmd/afs/backend.go

@@ -144,6 +144,8 @@ type afsControlPlane interface {
 	QueryWorkspace(ctx context.Context, workspace string, request mcptools.FileQueryRequest) (mcptools.FileQueryResponse, error)
 	QueryIndexStatus(ctx context.Context, workspace string, request controlplane.WorkspaceQueryIndexStatusRequest) (controlplane.WorkspaceQueryIndexStatus, error)
 	RebuildQueryIndex(ctx context.Context, workspace string, request controlplane.WorkspaceQueryIndexRebuildRequest) (controlplane.WorkspaceQueryIndexRebuildResponse, error)
+	QueryModelStatus(ctx context.Context, request controlplane.QueryModelStatusRequest) (controlplane.QueryModelStatus, error)
+	DownloadQueryModel(ctx context.Context, request controlplane.QueryModelDownloadRequest) (controlplane.QueryModelDownloadResult, error)
 	DiffWorkspace(ctx context.Context, workspace, baseView, headView string) (controlplane.WorkspaceDiffResponse, error)
 	RestoreCheckpoint(ctx context.Context, workspace, checkpointID string) error
 	SaveCheckpoint(ctx context.Context, input controlplane.SaveCheckpointRequest) (bool, error)

cmd/afs/controlplane_http_client.go

@@ -416,6 +416,26 @@ func (c *httpControlPlaneClient) RebuildQueryIndex(ctx context.Context, workspac
 	return out, err
 }
 
+func (c *httpControlPlaneClient) QueryModelStatus(ctx context.Context, request controlplane.QueryModelStatusRequest) (controlplane.QueryModelStatus, error) {
+	params := url.Values{}
+	if strings.TrimSpace(request.Model) != "" {
+		params.Set("model", strings.TrimSpace(request.Model))
+	}
+	rel := "/v1/query/model/status"
+	if encoded := params.Encode(); encoded != "" {
+		rel += "?" + encoded
+	}
+	var out controlplane.QueryModelStatus
+	err := c.doJSON(ctx, http.MethodGet, rel, nil, &out, http.StatusOK)
+	return out, err
+}
+
+func (c *httpControlPlaneClient) DownloadQueryModel(ctx context.Context, request controlplane.QueryModelDownloadRequest) (controlplane.QueryModelDownloadResult, error) {
+	var out controlplane.QueryModelDownloadResult
+	err := c.doJSONWithClient(ctx, c.queryer, http.MethodPost, "/v1/query/model/download", request, &out, http.StatusOK)
+	return out, err
+}
+
 func (c *httpControlPlaneClient) DiffWorkspace(ctx context.Context, workspace, baseView, headView string) (controlplane.WorkspaceDiffResponse, error) {
 	params := url.Values{}
 	if strings.TrimSpace(baseView) != "" {