Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions cmd/afs/afs_commands_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}
Expand Down
28 changes: 28 additions & 0 deletions cmd/afs/afs_mcp_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
86 changes: 86 additions & 0 deletions cmd/afs/afs_query_commands_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
76 changes: 71 additions & 5 deletions cmd/afs/afs_query_index_commands.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package main

import (
"bufio"
"context"
"encoding/json"
"flag"
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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]
Expand All @@ -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)
}
1 change: 1 addition & 0 deletions cmd/afs/backend.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
104 changes: 104 additions & 0 deletions cmd/afs/config_commands_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Loading
Loading