From 27d7e1de0d315e3f1313738fe99fde0b704e582c Mon Sep 17 00:00:00 2001 From: John Howard Date: Thu, 16 Jul 2026 15:58:06 -0700 Subject: [PATCH] cli: add auto-completion for atespace and actors This mirrors how kubectl works and makes it much easier to select an actor. --- cmd/kubectl-ate/internal/cmd/create_actor.go | 1 + cmd/kubectl-ate/internal/cmd/delete_actor.go | 1 + .../internal/cmd/delete_atespace.go | 1 + cmd/kubectl-ate/internal/cmd/get_actors.go | 1 + cmd/kubectl-ate/internal/cmd/get_atespaces.go | 1 + cmd/kubectl-ate/internal/cmd/logs_actors.go | 1 + cmd/kubectl-ate/internal/cmd/pause_actor.go | 1 + .../internal/cmd/resource_completions.go | 143 ++++++++++++++++++ .../internal/cmd/resource_completions_test.go | 73 +++++++++ cmd/kubectl-ate/internal/cmd/resume_actor.go | 1 + cmd/kubectl-ate/internal/cmd/suspend_actor.go | 1 + 11 files changed, 225 insertions(+) create mode 100644 cmd/kubectl-ate/internal/cmd/resource_completions.go create mode 100644 cmd/kubectl-ate/internal/cmd/resource_completions_test.go diff --git a/cmd/kubectl-ate/internal/cmd/create_actor.go b/cmd/kubectl-ate/internal/cmd/create_actor.go index 6a8c750cd..41b72b234 100644 --- a/cmd/kubectl-ate/internal/cmd/create_actor.go +++ b/cmd/kubectl-ate/internal/cmd/create_actor.go @@ -68,5 +68,6 @@ func init() { _ = createActorCmd.MarkFlagRequired("template") createActorCmd.Flags().StringVarP(&atespaceFlag, "atespace", "a", "", "Atespace to create the actor in (required)") _ = createActorCmd.MarkFlagRequired("atespace") + registerAtespaceFlagCompletion(createActorCmd) createCmd.AddCommand(createActorCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/delete_actor.go b/cmd/kubectl-ate/internal/cmd/delete_actor.go index 578b8928e..a04cb98bb 100644 --- a/cmd/kubectl-ate/internal/cmd/delete_actor.go +++ b/cmd/kubectl-ate/internal/cmd/delete_actor.go @@ -52,5 +52,6 @@ var deleteActorCmd = &cobra.Command{ func init() { deleteActorCmd.Flags().StringVarP(&deleteAtespaceFlag, "atespace", "a", "", "Atespace the actor lives in") _ = deleteActorCmd.MarkFlagRequired("atespace") + registerActorCompletions(deleteActorCmd) deleteCmd.AddCommand(deleteActorCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/delete_atespace.go b/cmd/kubectl-ate/internal/cmd/delete_atespace.go index aecdde71b..778966fe3 100644 --- a/cmd/kubectl-ate/internal/cmd/delete_atespace.go +++ b/cmd/kubectl-ate/internal/cmd/delete_atespace.go @@ -45,5 +45,6 @@ var deleteAtespaceCmd = &cobra.Command{ } func init() { + registerAtespaceArgCompletion(deleteAtespaceCmd) deleteCmd.AddCommand(deleteAtespaceCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/get_actors.go b/cmd/kubectl-ate/internal/cmd/get_actors.go index 05dd77b31..9cc6ebeae 100644 --- a/cmd/kubectl-ate/internal/cmd/get_actors.go +++ b/cmd/kubectl-ate/internal/cmd/get_actors.go @@ -96,5 +96,6 @@ var getActorsCmd = &cobra.Command{ func init() { getActorsCmd.Flags().StringVarP(&getActorsAtespaceFlag, "atespace", "a", "", "Atespace to list/get actors in. Required for a single actor; for listing, use this or -A.") getActorsCmd.Flags().BoolVarP(&getActorsAllAtespaces, "all-atespaces", "A", false, "List actors across all atespaces (listing only; mutually exclusive with --atespace)") + registerActorCompletions(getActorsCmd) getCmd.AddCommand(getActorsCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/get_atespaces.go b/cmd/kubectl-ate/internal/cmd/get_atespaces.go index 9c456c253..c657510a9 100644 --- a/cmd/kubectl-ate/internal/cmd/get_atespaces.go +++ b/cmd/kubectl-ate/internal/cmd/get_atespaces.go @@ -52,5 +52,6 @@ var getAtespacesCmd = &cobra.Command{ } func init() { + registerAtespaceArgCompletion(getAtespacesCmd) getCmd.AddCommand(getAtespacesCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/logs_actors.go b/cmd/kubectl-ate/internal/cmd/logs_actors.go index 5c78f422d..5f9389746 100644 --- a/cmd/kubectl-ate/internal/cmd/logs_actors.go +++ b/cmd/kubectl-ate/internal/cmd/logs_actors.go @@ -53,6 +53,7 @@ func init() { logsActorsCmd.Flags().BoolVarP(&followLogs, "follow", "f", false, "Specify if the logs should be streamed.") logsActorsCmd.Flags().StringVarP(&logsAtespaceFlag, "atespace", "a", "", "Atespace the actor lives in") _ = logsActorsCmd.MarkFlagRequired("atespace") + registerActorCompletions(logsActorsCmd) logsCmd.AddCommand(logsActorsCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/pause_actor.go b/cmd/kubectl-ate/internal/cmd/pause_actor.go index 6b22749fd..275fdf074 100644 --- a/cmd/kubectl-ate/internal/cmd/pause_actor.go +++ b/cmd/kubectl-ate/internal/cmd/pause_actor.go @@ -51,5 +51,6 @@ var pauseActorCmd = &cobra.Command{ func init() { pauseActorCmd.Flags().StringVarP(&pauseAtespaceFlag, "atespace", "a", "", "Atespace the actor lives in") _ = pauseActorCmd.MarkFlagRequired("atespace") + registerActorCompletions(pauseActorCmd) pauseCmd.AddCommand(pauseActorCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/resource_completions.go b/cmd/kubectl-ate/internal/cmd/resource_completions.go new file mode 100644 index 000000000..9a841856d --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/resource_completions.go @@ -0,0 +1,143 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "context" + "sort" + "strings" + "time" + + "github.com/agent-substrate/substrate/internal/ateclient" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "github.com/spf13/cobra" + "google.golang.org/grpc" +) + +const resourceCompletionTimeout = 5 * time.Second + +type resourceCompletionClient interface { + ListAtespaces(context.Context, *ateapipb.ListAtespacesRequest, ...grpc.CallOption) (*ateapipb.ListAtespacesResponse, error) + ListActors(context.Context, *ateapipb.ListActorsRequest, ...grpc.CallOption) (*ateapipb.ListActorsResponse, error) + Close() +} + +var newResourceCompletionClient = func(ctx context.Context) (resourceCompletionClient, error) { + return ateclient.NewClient(ctx, kubeconfig, k8sContext, endpoint, traceEnabled) +} + +func registerAtespaceFlagCompletion(cmd *cobra.Command) { + _ = cmd.RegisterFlagCompletionFunc("atespace", completeAtespaces) +} + +func registerActorCompletions(cmd *cobra.Command) { + registerAtespaceFlagCompletion(cmd) + cmd.ValidArgsFunction = completeActors +} + +func registerAtespaceArgCompletion(cmd *cobra.Command) { + cmd.ValidArgsFunction = completeAtespaceArg +} + +func completeAtespaces(cmd *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) { + ctx, cancel := resourceCompletionContext(cmd) + defer cancel() + + client, err := newResourceCompletionClient(ctx) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + defer client.Close() + + resp, err := client.ListAtespaces(ctx, &ateapipb.ListAtespacesRequest{}) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + + names := make([]string, 0, len(resp.GetAtespaces())) + for _, atespace := range resp.GetAtespaces() { + names = append(names, atespace.GetMetadata().GetName()) + } + return matchingResourceNames(names, toComplete), cobra.ShellCompDirectiveNoFileComp +} + +func completeAtespaceArg(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { + if len(args) != 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + return completeAtespaces(cmd, args, toComplete) +} + +func completeActors(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { + if len(args) != 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + + atespace, err := cmd.Flags().GetString("atespace") + if err != nil || atespace == "" { + return nil, cobra.ShellCompDirectiveNoFileComp + } + + ctx, cancel := resourceCompletionContext(cmd) + defer cancel() + + client, err := newResourceCompletionClient(ctx) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + defer client.Close() + + var names []string + pageToken := "" + for { + resp, err := client.ListActors(ctx, &ateapipb.ListActorsRequest{ + Atespace: atespace, + PageSize: 1000, + PageToken: pageToken, + }) + if err != nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + for _, actor := range resp.GetActors() { + names = append(names, actor.GetMetadata().GetName()) + } + + pageToken = resp.GetNextPageToken() + if pageToken == "" { + break + } + } + + return matchingResourceNames(names, toComplete), cobra.ShellCompDirectiveNoFileComp +} + +func resourceCompletionContext(cmd *cobra.Command) (context.Context, context.CancelFunc) { + parent := cmd.Context() + if parent == nil { + parent = context.Background() + } + return context.WithTimeout(parent, resourceCompletionTimeout) +} + +func matchingResourceNames(names []string, prefix string) []string { + matches := names[:0] + for _, name := range names { + if name != "" && strings.HasPrefix(name, prefix) { + matches = append(matches, name) + } + } + sort.Strings(matches) + return matches +} diff --git a/cmd/kubectl-ate/internal/cmd/resource_completions_test.go b/cmd/kubectl-ate/internal/cmd/resource_completions_test.go new file mode 100644 index 000000000..e545fed29 --- /dev/null +++ b/cmd/kubectl-ate/internal/cmd/resource_completions_test.go @@ -0,0 +1,73 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package cmd + +import ( + "context" + "strings" + "testing" + + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc" +) + +type fakeCompletionClient struct { + actorRequest *ateapipb.ListActorsRequest +} + +func (*fakeCompletionClient) ListAtespaces(context.Context, *ateapipb.ListAtespacesRequest, ...grpc.CallOption) (*ateapipb.ListAtespacesResponse, error) { + return &ateapipb.ListAtespacesResponse{Atespaces: []*ateapipb.Atespace{ + {Metadata: &ateapipb.ResourceMetadata{Name: "team-z"}}, + {Metadata: &ateapipb.ResourceMetadata{Name: "team-a"}}, + }}, nil +} + +func (f *fakeCompletionClient) ListActors(_ context.Context, req *ateapipb.ListActorsRequest, _ ...grpc.CallOption) (*ateapipb.ListActorsResponse, error) { + f.actorRequest = req + return &ateapipb.ListActorsResponse{Actors: []*ateapipb.Actor{ + {Metadata: &ateapipb.ResourceMetadata{Name: "alpha"}}, + {Metadata: &ateapipb.ResourceMetadata{Name: "beta"}}, + }}, nil +} + +func (*fakeCompletionClient) Close() {} + +func TestLogsActorCompletions(t *testing.T) { + fake := &fakeCompletionClient{} + original := newResourceCompletionClient + newResourceCompletionClient = func(context.Context) (resourceCompletionClient, error) { return fake, nil } + t.Cleanup(func() { + newResourceCompletionClient = original + _ = logsActorsCmd.Flags().Set("atespace", "") + }) + + completeAtespace, ok := logsActorsCmd.GetFlagCompletionFunc("atespace") + if !ok { + t.Fatal("--atespace has no completion function") + } + atespaces, _ := completeAtespace(logsActorsCmd, nil, "team-") + if got := strings.Join(atespaces, ","); got != "team-a,team-z" { + t.Errorf("atespace completions = %q, want %q", got, "team-a,team-z") + } + + _ = logsActorsCmd.Flags().Set("atespace", "team-a") + actors, _ := logsActorsCmd.ValidArgsFunction(logsActorsCmd, nil, "a") + if got := strings.Join(actors, ","); got != "alpha" { + t.Errorf("actor completions = %q, want %q", got, "alpha") + } + if got := fake.actorRequest.GetAtespace(); got != "team-a" { + t.Errorf("ListActors atespace = %q, want %q", got, "team-a") + } +} diff --git a/cmd/kubectl-ate/internal/cmd/resume_actor.go b/cmd/kubectl-ate/internal/cmd/resume_actor.go index d744688f8..1e1a630de 100644 --- a/cmd/kubectl-ate/internal/cmd/resume_actor.go +++ b/cmd/kubectl-ate/internal/cmd/resume_actor.go @@ -54,5 +54,6 @@ func init() { resumeActorCmd.Flags().BoolVarP(&bootFlag, "boot", "", false, "Skip golden snapshot and boot from scratch.") resumeActorCmd.Flags().StringVarP(&resumeAtespaceFlag, "atespace", "a", "", "Atespace the actor lives in") _ = resumeActorCmd.MarkFlagRequired("atespace") + registerActorCompletions(resumeActorCmd) resumeCmd.AddCommand(resumeActorCmd) } diff --git a/cmd/kubectl-ate/internal/cmd/suspend_actor.go b/cmd/kubectl-ate/internal/cmd/suspend_actor.go index 3bca211cd..7ac7d624c 100644 --- a/cmd/kubectl-ate/internal/cmd/suspend_actor.go +++ b/cmd/kubectl-ate/internal/cmd/suspend_actor.go @@ -51,5 +51,6 @@ var suspendActorCmd = &cobra.Command{ func init() { suspendActorCmd.Flags().StringVarP(&suspendAtespaceFlag, "atespace", "a", "", "Atespace the actor lives in") _ = suspendActorCmd.MarkFlagRequired("atespace") + registerActorCompletions(suspendActorCmd) suspendCmd.AddCommand(suspendActorCmd) }