Fix query index cleanup UX and docs (#18)

* Fix query index cleanup contract

* Keep query QA plan local

* Fix query QA regressions

* Handle post-QA edge cases

* Fix UI lint after rebase

* Align UI doc test with volume routes

---------

Co-authored-by: Rowan Trollope <rowantrollope@gmail.com>

Author: abrookins

cmd/afs/afs_commands_test.go

@@ -146,6 +146,10 @@ func (s stubAFSControlPlane) DownloadQueryModel(context.Context, controlplane.Qu
 	return controlplane.QueryModelDownloadResult{}, fmt.Errorf("unexpected DownloadQueryModel call")
 }
 
+func (s stubAFSControlPlane) CleanQueryIndex(context.Context, string, controlplane.WorkspaceQueryIndexCleanRequest) (controlplane.WorkspaceQueryIndexCleanResponse, error) {
+	return controlplane.WorkspaceQueryIndexCleanResponse{}, fmt.Errorf("unexpected CleanQueryIndex call")
+}
+
 func (s stubAFSControlPlane) DiffWorkspace(context.Context, string, string, string) (controlplane.WorkspaceDiffResponse, error) {
 	return controlplane.WorkspaceDiffResponse{}, fmt.Errorf("unexpected DiffWorkspace call")
 }

cmd/afs/afs_mcp_test.go

@@ -257,6 +257,34 @@ func TestAFSMCPCheckpointCreateAllowsUnchangedWorkspace(t *testing.T) {
 	}
 }
 
+func TestAFSMCPCheckpointCreateGeneratesCheckpointNameWhenOmitted(t *testing.T) {
+	t.Helper()
+
+	server, closeFn := setupAFSMCPTestServer(t)
+	defer closeFn()
+	server.profile = controlplane.MCPProfileWorkspaceRWCheckpoint
+
+	checkpointResult := server.callTool(context.Background(), "checkpoint_create", map[string]any{})
+	if checkpointResult.IsError {
+		t.Fatalf("checkpoint_create without name returned error result: %+v", checkpointResult)
+	}
+
+	var checkpointPayload map[string]any
+	if err := decodeStructuredContent(checkpointResult.StructuredContent, &checkpointPayload); err != nil {
+		t.Fatalf("decodeStructuredContent(checkpoint generated) returned error: %v", err)
+	}
+	if created, _ := checkpointPayload["created"].(bool); !created {
+		t.Fatalf("checkpoint_create created = %#v, want true", checkpointPayload["created"])
+	}
+	checkpoint, _ := checkpointPayload["checkpoint"].(string)
+	if !strings.HasPrefix(checkpoint, "save-") {
+		t.Fatalf("checkpoint_create checkpoint = %#v, want save-*", checkpointPayload["checkpoint"])
+	}
+	if _, err := server.store.getSavepointMeta(context.Background(), "repo", checkpoint); err != nil {
+		t.Fatalf("getSavepointMeta(%s) returned error: %v", checkpoint, err)
+	}
+}
+
 func TestAFSMCPFileWriteDoesNotRematerializeLocalWorkspaceCache(t *testing.T) {
 	t.Helper()
 

cmd/afs/afs_query_commands_test.go

@@ -738,6 +738,92 @@ func findLineWithPrefix(output, prefix string) string {
 	return ""
 }
 
+func TestCmdQueryIndexCleanClearsGeneratedData(t *testing.T) {
+	_, store, closeStore := setupAFSGrepTest(t)
+	defer closeStore()
+
+	writeLiveAFSFile(t, store, "repo", "/docs/checkpoints.md", "checkpoint recovery guide\n")
+	if _, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "index", "status", "--json"})
+	}); err != nil {
+		t.Fatalf("cmdQuery(index status) returned error: %v", err)
+	}
+
+	output, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "index", "clean", "--yes", "--json"})
+	})
+	if err != nil {
+		t.Fatalf("cmdQuery(index clean) returned error: %v", err)
+	}
+
+	var response controlplane.WorkspaceQueryIndexCleanResponse
+	if err := json.Unmarshal([]byte(output), &response); err != nil {
+		t.Fatalf("Unmarshal(clean response) returned error: %v\n%s", err, output)
+	}
+	if !response.Cleared || response.RemovedFiles != 1 || response.RemovedChunks == 0 {
+		t.Fatalf("clean response = %+v, want cleared response with removed files/chunks", response)
+	}
+	if response.Status.Keyword.Chunks != 0 {
+		t.Fatalf("post-clean chunks = %d, want 0", response.Status.Keyword.Chunks)
+	}
+}
+
+func TestCmdQueryIndexCleanPromptsBeforeClearing(t *testing.T) {
+	_, store, closeStore := setupAFSGrepTest(t)
+	defer closeStore()
+
+	writeLiveAFSFile(t, store, "repo", "/docs/checkpoints.md", "checkpoint recovery guide\n")
+	if _, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "index", "status", "--json"})
+	}); err != nil {
+		t.Fatalf("cmdQuery(index status) returned error: %v", err)
+	}
+
+	input, err := os.CreateTemp(t.TempDir(), "query-index-clean-stdin")
+	if err != nil {
+		t.Fatalf("CreateTemp() returned error: %v", err)
+	}
+	if _, err := input.WriteString("n\n"); err != nil {
+		t.Fatalf("WriteString() returned error: %v", err)
+	}
+	if _, err := input.Seek(0, 0); err != nil {
+		t.Fatalf("Seek() returned error: %v", err)
+	}
+	origStdin := os.Stdin
+	os.Stdin = input
+	t.Cleanup(func() {
+		os.Stdin = origStdin
+		_ = input.Close()
+	})
+
+	output, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "index", "clean"})
+	})
+	if err != nil {
+		t.Fatalf("cmdQuery(index clean) returned error: %v", err)
+	}
+	if !strings.Contains(output, "Are you sure you want to clear generated query index data for repo?") {
+		t.Fatalf("cmdQuery(index clean) output = %q, want confirmation prompt", output)
+	}
+	if !strings.Contains(output, "Query index clean cancelled.") {
+		t.Fatalf("cmdQuery(index clean) output = %q, want cancellation message", output)
+	}
+
+	statusOutput, err := captureStdout(t, func() error {
+		return cmdQuery([]string{"query", "index", "status", "--json"})
+	})
+	if err != nil {
+		t.Fatalf("cmdQuery(index status after cancel) returned error: %v", err)
+	}
+	var status controlplane.WorkspaceQueryIndexStatus
+	if err := json.Unmarshal([]byte(statusOutput), &status); err != nil {
+		t.Fatalf("Unmarshal(status) returned error: %v\n%s", err, statusOutput)
+	}
+	if status.Keyword.Chunks == 0 {
+		t.Fatalf("status after cancelled clean = %+v, want indexed chunks to remain", status)
+	}
+}
+
 func TestWorkspaceQueryConfigFallsBackWhenConfigRouteIsMissing(t *testing.T) {
 	cfg, err := workspaceQueryConfig(context.Background(), stubAFSControlPlane{
 		workspaceConfigErr: os.ErrNotExist,

cmd/afs/afs_query_index_commands.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bufio"
 	"context"
 	"encoding/json"
 	"flag"
@@ -142,20 +143,42 @@ func runWorkspaceQueryIndexClean(workspace string, args []string) error {
 	fs := flag.NewFlagSet("query index clean", flag.ContinueOnError)
 	fs.SetOutput(io.Discard)
 	var jsonOut bool
+	var yes bool
 	fs.BoolVar(&jsonOut, "json", false, "write JSON output")
+	fs.BoolVar(&yes, "yes", false, "confirm removal of generated query index data")
+	fs.BoolVar(&yes, "y", false, "confirm removal of generated query index data")
 	if err := fs.Parse(args); err != nil {
 		return fmt.Errorf("%s", workspaceQueryIndexUsageText(filepath.Base(os.Args[0])))
 	}
 	if fs.NArg() != 0 {
 		return fmt.Errorf("%s", workspaceQueryIndexUsageText(filepath.Base(os.Args[0])))
 	}
-	status, err := workspaceQueryIndexStatusForWorkspace(workspace, "/")
+	ctx := context.Background()
+	remote, err := openFSRemoteWorkspace(ctx, workspace)
 	if err != nil {
 		return err
 	}
-	status.State = "clean"
-	status.Message = "No query index data was removed."
-	return writeWorkspaceQueryIndexStatus(status, jsonOut)
+	defer remote.close()
+	if !yes {
+		ok, err := confirmWorkspaceQueryIndexClean(remote.selection.Name)
+		if err != nil {
+			return err
+		}
+		if !ok {
+			fmt.Println()
+			fmt.Println("Query index clean cancelled.")
+			fmt.Println()
+			return nil
+		}
+	}
+	response, err := remote.controlPlane.CleanQueryIndex(ctx, remote.selection.ID, controlplane.WorkspaceQueryIndexCleanRequest{
+		Workspace: remote.selection.Name,
+		Confirm:   true,
+	})
+	if err != nil {
+		return err
+	}
+	return writeWorkspaceQueryIndexClean(response, jsonOut)
 }
 
 func workspaceQueryIndexStatusForWorkspace(workspace, path string) (controlplane.WorkspaceQueryIndexStatus, error) {
@@ -271,6 +294,47 @@ func embeddingBackfillLabel(result controlplane.QueryEmbeddingBackfillResult) st
 	return "off"
 }
 
+func writeWorkspaceQueryIndexClean(response controlplane.WorkspaceQueryIndexCleanResponse, jsonOut bool) error {
+	if jsonOut {
+		enc := json.NewEncoder(os.Stdout)
+		enc.SetIndent("", "  ")
+		return enc.Encode(response)
+	}
+	fmt.Fprintln(os.Stdout, "Query index clean")
+	fmt.Fprintln(os.Stdout)
+	fmt.Fprintf(os.Stdout, "workspace  %s\n", response.Workspace)
+	fmt.Fprintf(os.Stdout, "cleared    %t\n", response.Cleared)
+	fmt.Fprintf(os.Stdout, "files      %d\n", response.RemovedFiles)
+	fmt.Fprintf(os.Stdout, "chunks     %d\n", response.RemovedChunks)
+	fmt.Fprintf(os.Stdout, "state      %s\n", response.Status.State)
+	if response.Message != "" {
+		fmt.Fprintln(os.Stdout)
+		fmt.Fprintln(os.Stdout, response.Message)
+	}
+	return nil
+}
+
+func confirmWorkspaceQueryIndexClean(workspace string) (bool, error) {
+	return confirmWorkspaceQueryIndexCleanWithReader(workspace, bufio.NewReader(os.Stdin))
+}
+
+func confirmWorkspaceQueryIndexCleanWithReader(workspace string, reader *bufio.Reader) (bool, error) {
+	workspace = strings.TrimSpace(workspace)
+	if workspace == "" {
+		return false, nil
+	}
+	fmt.Println()
+	fmt.Printf("Are you sure you want to clear generated query index data for %s? Workspace files will not change. [y/N] ", workspace)
+	raw, err := reader.ReadString('\n')
+	if err != nil && strings.TrimSpace(raw) == "" {
+		fmt.Println()
+		return false, nil
+	}
+	fmt.Println()
+	answer := strings.ToLower(strings.TrimSpace(raw))
+	return answer == "y" || answer == "yes", nil
+}
+
 func workspaceQueryIndexUsageText(bin string) string {
 	return brandHeaderString() + fmt.Sprintf(`Usage:
   %[1]s query index <status|create|rebuild|clean> [flags]
@@ -282,18 +346,20 @@ Subcommands:
   status             Show keyword query projection and embedding state
   create             Build keyword chunks and semantic embeddings
   rebuild            Enqueue existing files for keyword query indexing
-  clean              Remove stale query index data
+  clean              Clear generated query index data for the workspace
 
 Flags:
   --json             Write JSON output
   --path <path>      Scope status or rebuild to a workspace path
   --wait             Wait for rebuild completion
   --force            Rebuild existing chunks
   --embeddings       Build semantic embeddings
+  --yes, -y          Confirm full-workspace query index cleanup
 
 Examples:
   %[1]s query index status
   %[1]s fs repo query index create --embeddings --wait
   %[1]s fs repo query index rebuild --path /cmd/afs --wait
+  %[1]s query index clean --yes
 `, bin)
 }

cmd/afs/backend.go

@@ -146,6 +146,7 @@ type afsControlPlane interface {
 	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)
+	CleanQueryIndex(ctx context.Context, workspace string, request controlplane.WorkspaceQueryIndexCleanRequest) (controlplane.WorkspaceQueryIndexCleanResponse, 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/config_commands_test.go

@@ -304,6 +304,110 @@ func TestCmdConfigResetRemovesConfigAndState(t *testing.T) {
 	}
 }
 
+func TestCmdConfigResetWithAlternateConfigPreservesDefaultState(t *testing.T) {
+	t.Helper()
+
+	withTempHome(t)
+	configFile := filepath.Join(t.TempDir(), "afs.config.json")
+	origConfigPath := cfgPathOverride
+	cfgPathOverride = configFile
+	t.Cleanup(func() {
+		cfgPathOverride = origConfigPath
+	})
+
+	if err := os.WriteFile(configFile, []byte(`{"mode":"sync"}`), 0o600); err != nil {
+		t.Fatalf("WriteFile(config) returned error: %v", err)
+	}
+
+	rawDefaultState, err := json.MarshalIndent(state{
+		Mode:             modeSync,
+		CurrentWorkspace: "default-workspace",
+		LocalPath:        t.TempDir(),
+	}, "", "  ")
+	if err != nil {
+		t.Fatalf("json.MarshalIndent(default state) returned error: %v", err)
+	}
+	if err := os.MkdirAll(filepath.Dir(defaultStatePath()), 0o700); err != nil {
+		t.Fatalf("MkdirAll(default state dir) returned error: %v", err)
+	}
+	if err := os.WriteFile(defaultStatePath(), rawDefaultState, 0o600); err != nil {
+		t.Fatalf("WriteFile(defaultStatePath) returned error: %v", err)
+	}
+
+	out, err := captureStdout(t, func() error {
+		return cmdConfig([]string{"config", "reset"})
+	})
+	if err != nil {
+		t.Fatalf("cmdConfig(reset) returned error: %v", err)
+	}
+	if strings.Contains(out, "Unmounted workspace default-workspace") {
+		t.Fatalf("cmdConfig(reset) output = %q, should not unmount default runtime for an alternate config", out)
+	}
+	if _, err := os.Stat(defaultStatePath()); err != nil {
+		t.Fatalf("defaultStatePath removed by alternate-config reset: %v", err)
+	}
+}
+
+func TestCmdConfigResetWithAlternateConfigRemovesOnlyScopedState(t *testing.T) {
+	t.Helper()
+
+	withTempHome(t)
+	configFile := filepath.Join(t.TempDir(), "afs.config.json")
+	origConfigPath := cfgPathOverride
+	cfgPathOverride = configFile
+	t.Cleanup(func() {
+		cfgPathOverride = origConfigPath
+	})
+
+	if err := os.WriteFile(configFile, []byte(`{"mode":"sync"}`), 0o600); err != nil {
+		t.Fatalf("WriteFile(config) returned error: %v", err)
+	}
+
+	scopedStatePath := statePath()
+	rawScopedState, err := json.MarshalIndent(state{
+		Mode:             modeSync,
+		CurrentWorkspace: "scoped-workspace",
+		LocalPath:        t.TempDir(),
+	}, "", "  ")
+	if err != nil {
+		t.Fatalf("json.MarshalIndent(scoped state) returned error: %v", err)
+	}
+	if err := os.MkdirAll(filepath.Dir(scopedStatePath), 0o700); err != nil {
+		t.Fatalf("MkdirAll(scoped state dir) returned error: %v", err)
+	}
+	if err := os.WriteFile(scopedStatePath, rawScopedState, 0o600); err != nil {
+		t.Fatalf("WriteFile(scopedStatePath) returned error: %v", err)
+	}
+
+	rawDefaultState, err := json.MarshalIndent(state{
+		Mode:             modeSync,
+		CurrentWorkspace: "default-workspace",
+		LocalPath:        t.TempDir(),
+	}, "", "  ")
+	if err != nil {
+		t.Fatalf("json.MarshalIndent(default state) returned error: %v", err)
+	}
+	if err := os.MkdirAll(filepath.Dir(defaultStatePath()), 0o700); err != nil {
+		t.Fatalf("MkdirAll(default state dir) returned error: %v", err)
+	}
+	if err := os.WriteFile(defaultStatePath(), rawDefaultState, 0o600); err != nil {
+		t.Fatalf("WriteFile(defaultStatePath) returned error: %v", err)
+	}
+
+	if _, err := captureStdout(t, func() error {
+		return cmdConfig([]string{"config", "reset"})
+	}); err != nil {
+		t.Fatalf("cmdConfig(reset) returned error: %v", err)
+	}
+
+	if _, err := os.Stat(scopedStatePath); !errors.Is(err, os.ErrNotExist) {
+		t.Fatalf("scoped state still exists after reset: %v", err)
+	}
+	if _, err := os.Stat(defaultStatePath()); err != nil {
+		t.Fatalf("default state removed by scoped reset: %v", err)
+	}
+}
+
 func TestCmdConfigSetAgentNamePersistsFriendlyAgentName(t *testing.T) {
 	t.Helper()