diff --git a/src/go/rpk/pkg/cli/ai/BUILD b/src/go/rpk/pkg/cli/ai/BUILD index 3072c549f0324..67b52e48e7839 100644 --- a/src/go/rpk/pkg/cli/ai/BUILD +++ b/src/go/rpk/pkg/cli/ai/BUILD @@ -19,7 +19,6 @@ go_library( "//src/go/rpk/pkg/osutil", "//src/go/rpk/pkg/out", "//src/go/rpk/pkg/plugin", - "//src/go/rpk/pkg/publicapi", "//src/go/rpk/pkg/redpanda", "@com_github_spf13_afero//:afero", "@com_github_spf13_cobra//:cobra", @@ -37,10 +36,8 @@ go_test( embed = [":ai"], deps = [ "//src/go/rpk/pkg/config", - "@build_buf_gen_go_redpandadata_cloud_protocolbuffers_go//redpanda/api/controlplane/v1:controlplane", "@com_github_spf13_afero//:afero", "@com_github_spf13_cobra//:cobra", "@com_github_stretchr_testify//require", - "@org_golang_google_protobuf//proto", ], ) diff --git a/src/go/rpk/pkg/cli/ai/ai.go b/src/go/rpk/pkg/cli/ai/ai.go index a6763a07f9b8a..5accbde84e485 100644 --- a/src/go/rpk/pkg/cli/ai/ai.go +++ b/src/go/rpk/pkg/cli/ai/ai.go @@ -27,21 +27,15 @@ import ( func init() { // Whenever a `rpk ai ` managed-plugin leaf is dispatched, - // inject the plugin's token/endpoint env vars and strip rpk global - // flags before the child process is exec'd. Reaching this wrapper - // means cobra already routed to a real plugin leaf, so unless the user - // asked for --help / --version on the leaf itself (in which case the - // plugin renders its own local help), we always need cloud context — - // regardless of whether the leaf takes positional args. - plugin.RegisterManaged(rpaiPluginSlug, []string{"ai"}, func(cmd *cobra.Command, fs afero.Fs, p *config.Params) *cobra.Command { + // strip rpk's global flags before the child process is exec'd. rpk does + // not inject any cloud context: the rpk ai plugin owns its own login + // (`rpk ai auth login`) and environment selection (`rpk ai env use`), so + // it runs without a selected rpk cloud cluster. + plugin.RegisterManaged(rpaiPluginSlug, []string{"ai"}, func(cmd *cobra.Command, _ afero.Fs, p *config.Params) *cobra.Command { run := cmd.Run cmd.Run = func(cmd *cobra.Command, args []string) { pluginArgs, err := parseFlags(p, cmd, args) out.MaybeDie(err, "unable to prepare rpk ai invocation: %v", err) - if !skipCloudForHelp(pluginArgs) { - err = resolveAndInjectEnv(cmd.Context(), fs, p, pluginArgs) - out.MaybeDie(err, "unable to prepare rpk ai invocation: %v", err) - } run(cmd, pluginArgs) } return cmd @@ -83,11 +77,6 @@ func NewCommand(fs afero.Fs, p *config.Params, execFn func(string, []string) err return } - if !skipCloudForHelp(pluginArgs) { - err = resolveAndInjectEnv(cmd.Context(), fs, p, pluginArgs) - out.MaybeDie(err, "unable to prepare rpk ai invocation: %v", err) - } - var pluginPath string if !pluginExists { // FIPS is gated here, after the help/version short-circuits, diff --git a/src/go/rpk/pkg/cli/ai/hook.go b/src/go/rpk/pkg/cli/ai/hook.go index b6732952bcb20..57ef334c7ded7 100644 --- a/src/go/rpk/pkg/cli/ai/hook.go +++ b/src/go/rpk/pkg/cli/ai/hook.go @@ -10,93 +10,21 @@ package ai import ( - "context" - "fmt" - "os" "slices" - "strings" "github.com/redpanda-data/redpanda/src/go/rpk/pkg/cobraext" "github.com/redpanda-data/redpanda/src/go/rpk/pkg/config" - "github.com/redpanda-data/redpanda/src/go/rpk/pkg/publicapi" - "github.com/spf13/afero" "github.com/spf13/cobra" "go.uber.org/zap" ) -// The rpk ai plugin reads the following environment variables and flag: -// -// - envAuthToken: bearer token forwarded as Authorization to the AI -// Gateway. rpk reads the cached token off the active cloud profile -// and exports it; a missing or stale token is reported by the -// gateway as a 401, at which point the user runs `rpk cloud login`. -// - envEndpoint: AI Gateway v2 base URL. rpk resolves it from the -// active cloud profile's cached AIGatewayURL (or falls back to a -// publicapi lookup) and exports it before exec. -// - flagEndpoint: same intent as envEndpoint, but on the command -// line. rpk only watches for it so we can skip the cluster lookup; -// the flag is parsed and consumed by the plugin itself. -// -// Any explicit value the user supplies (env or flag) wins — rpk only -// fills in missing pieces. -const ( - envAuthToken = "RPAI_TOKEN" - envEndpoint = "RPAI_ENDPOINT" - - flagEndpoint = "rpai-endpoint" -) - -// Exports RPAI_TOKEN and RPAI_ENDPOINT for the child plugin using cached -// values from rpk.yaml and the profile's AIGatewayURL (no OAuth refresh). -// Skips env var writes if the vars are already set or --rpai-endpoint is -// present. -func resolveAndInjectEnv(ctx context.Context, fs afero.Fs, p *config.Params, pluginArgs []string) error { - cfg, err := p.Load(fs) - if err != nil { - return fmt.Errorf("unable to load rpk config: %w", err) - } - - if os.Getenv(envAuthToken) == "" { - // Mirrors the lookup chain in oauth.LoadCloudToken (without the - // refresh) so --config, -X cloud_auth.*, and RPK_PROFILE all work. - auth := cfg.VirtualProfile().VirtualAuth() - if auth == nil { - auth = cfg.VirtualRpkYaml().CurrentAuth() - } - if auth != nil && auth.AuthToken != "" { - if err := os.Setenv(envAuthToken, auth.AuthToken); err != nil { - return fmt.Errorf("unable to set %s: %w", envAuthToken, err) - } - } - } - - if os.Getenv(envEndpoint) == "" && !hasEndpointFlag(pluginArgs) { - endpoint, err := resolveAigwEndpoint(ctx, cfg) - if err != nil { - return err - } - if err := os.Setenv(envEndpoint, endpoint); err != nil { - return fmt.Errorf("unable to set %s: %w", envEndpoint, err) - } - } - - return nil -} - -// skipCloudForHelp reports whether a --help / -h / --version flag is present, -// in which case we must not reach out to the cloud API or trigger OAuth. The -// rpk ai plugin child process renders its own help/version output locally. -func skipCloudForHelp(args []string) bool { - for _, a := range args { - if a == "--help" || a == "-h" || a == "--version" { - return true - } - } - return false -} - // parseFlags splits args into plugin args + rpk-global-flags consumed by rpk, // and parses the rpk-global-flags so the logger and config loader pick them up. +// +// rpk does not inject any cloud context (token or endpoint) into the plugin: +// the rpk ai plugin owns its own login (`rpk ai auth login`) and environment +// selection (`rpk ai env use`), so it runs without a selected rpk cloud +// cluster. Everything except rpk's own globals is forwarded untouched. func parseFlags(p *config.Params, cmd *cobra.Command, args []string) ([]string, error) { f := cmd.Flags() @@ -114,44 +42,3 @@ func parseFlags(p *config.Params, cmd *cobra.Command, args []string) ([]string, } return keepForPlugin, nil } - -// hasEndpointFlag reports whether the plugin args carry an explicit endpoint -// flag (in any supported form: --flag=value or --flag followed by its -// value). -func hasEndpointFlag(args []string) bool { - prefix := "--" + flagEndpoint - for _, a := range args { - if a == prefix || strings.HasPrefix(a, prefix+"=") { - return true - } - } - return false -} - -// resolveAigwEndpoint returns the AI Gateway v2 URL for the active rpk cloud -// profile's cluster. It prefers the value cached on the profile at creation -// time and only falls back to a live publicapi lookup when the cache is empty -// (older profiles created before AIGatewayURL existed, or profiles whose -// cluster had no AI Gateway attached at creation but does now). -func resolveAigwEndpoint(ctx context.Context, cfg *config.Config) (string, error) { - prof := cfg.VirtualProfile() - if prof == nil || !prof.FromCloud || prof.CloudCluster.ClusterID == "" { - return "", fmt.Errorf("no cluster selected for this rpk profile; run 'rpk cloud cluster use ' or pass --%s", flagEndpoint) - } - if prof.CloudCluster.AIGatewayURL != "" { - return prof.CloudCluster.AIGatewayURL, nil - } - clusterID := prof.CloudCluster.ClusterID - - token := os.Getenv(envAuthToken) - cl := publicapi.NewCloudClientSet(cfg.DevOverrides().PublicAPIURL, token) - cluster, err := cl.ClusterForID(ctx, clusterID) - if err != nil { - return "", fmt.Errorf("unable to resolve aigw endpoint for cluster %s: %w", clusterID, err) - } - endpoint := cluster.GetAiGateway().GetV2Url() - if endpoint == "" { - return "", fmt.Errorf("cluster %s does not have an AI Gateway v2 endpoint; pick a cluster that does, or pass --%s", clusterID, flagEndpoint) - } - return endpoint, nil -} diff --git a/src/go/rpk/pkg/cli/ai/hook_test.go b/src/go/rpk/pkg/cli/ai/hook_test.go index 14eb3ba698d4e..5cf9fdccce01b 100644 --- a/src/go/rpk/pkg/cli/ai/hook_test.go +++ b/src/go/rpk/pkg/cli/ai/hook_test.go @@ -10,19 +10,11 @@ package ai import ( - "fmt" - "net/http" - "net/http/httptest" - "os" - "strings" "testing" - controlplanev1 "buf.build/gen/go/redpandadata/cloud/protocolbuffers/go/redpanda/api/controlplane/v1" "github.com/redpanda-data/redpanda/src/go/rpk/pkg/config" - "github.com/spf13/afero" "github.com/spf13/cobra" "github.com/stretchr/testify/require" - "google.golang.org/protobuf/proto" ) func TestParseFlags_StripsRpkGlobals(t *testing.T) { @@ -55,263 +47,3 @@ func TestParseFlags_StripsRpkGlobals(t *testing.T) { require.NoError(t, gotErr) require.Equal(t, []string{"llm", "list", "--foo=bar"}, got) } - -func TestSkipCloudForHelp(t *testing.T) { - cases := []struct { - name string - args []string - want bool - }{ - {"empty", nil, false}, - {"help long", []string{"--help"}, true}, - {"help short", []string{"-h"}, true}, - {"version", []string{"--version"}, true}, - {"subcommand no help", []string{"llm", "list"}, false}, - {"subcommand then help", []string{"llm", "list", "--help"}, true}, - {"format flag only", []string{"--format", "json"}, false}, - } - for _, c := range cases { - t.Run(c.name, func(t *testing.T) { - require.Equal(t, c.want, skipCloudForHelp(c.args)) - }) - } -} - -// TestResolveAndInjectEnv_LeafDispatchHappy covers leaf dispatch: cobra -// hands `rpk ai llm list` straight to the `list` leaf with args=nil. -// The hook still has to inject both env vars even though args is empty. -// The token comes from the cached cloud_auth entry on rpk.yaml — no -// OAuth refresh, just a passthrough — so an expired token just means -// the gateway returns 401 and the plugin tells the user to re-login. -func TestResolveAndInjectEnv_LeafDispatchHappy(t *testing.T) { - t.Setenv(envAuthToken, "") - t.Setenv(envEndpoint, "") - - cluster := &controlplanev1.Cluster{ - Id: "clu-1", - AiGateway: &controlplanev1.Cluster_AIGateway{ - V2Url: "https://aigw.example.com", - }, - } - ts := httptest.NewServer(clusterHandler(t, cluster)) - defer ts.Close() - - loadCloudProfile(t, ts.URL, "clu-1") - - fs := afero.NewMemMapFs() - path, err := config.DefaultRpkYamlPath() - require.NoError(t, err) - yaml := `version: 6 -current_profile: dev -current_cloud_auth_org_id: org-1 -current_cloud_auth_kind: cloud-sso -cloud_auth: - - name: cloud - organization: cloud - org_id: org-1 - kind: cloud-sso - auth_token: cached-cloud-token -profiles: - - name: dev - from_cloud: true - cloud_cluster: - cluster_id: "clu-1" -` - require.NoError(t, afero.WriteFile(fs, path, []byte(yaml), 0o600)) - - // args=nil simulates leaf dispatch: cobra consumed the path tokens. - require.NoError(t, resolveAndInjectEnv(t.Context(), fs, new(config.Params), nil)) - require.Equal(t, "cached-cloud-token", os.Getenv(envAuthToken), "RPAI_TOKEN must be set from the cached cloud_auth entry") - require.Equal(t, "https://aigw.example.com", os.Getenv(envEndpoint), "RPAI_ENDPOINT must be set from aigw v2 url") -} - -// TestResolveAndInjectEnv_NoCachedToken covers the "fresh install" -// case: rpk.yaml has no cloud_auth entry. resolveAndInjectEnv must not -// fail — it just leaves RPAI_TOKEN unset and lets the gateway return -// 401 (or the plugin's own auth chain pick up an alternative source). -func TestResolveAndInjectEnv_NoCachedToken(t *testing.T) { - t.Setenv(envAuthToken, "") - t.Setenv(envEndpoint, "") - - cluster := &controlplanev1.Cluster{ - Id: "clu-1", - AiGateway: &controlplanev1.Cluster_AIGateway{ - V2Url: "https://aigw.example.com", - }, - } - ts := httptest.NewServer(clusterHandler(t, cluster)) - defer ts.Close() - - loadCloudProfile(t, ts.URL, "clu-1") - - fs := afero.NewMemMapFs() - path, err := config.DefaultRpkYamlPath() - require.NoError(t, err) - // No cloud_auth section. - yaml := `version: 6 -current_profile: dev -profiles: - - name: dev - from_cloud: true - cloud_cluster: - cluster_id: "clu-1" -` - require.NoError(t, afero.WriteFile(fs, path, []byte(yaml), 0o600)) - - require.NoError(t, resolveAndInjectEnv(t.Context(), fs, new(config.Params), nil)) - require.Empty(t, os.Getenv(envAuthToken), "RPAI_TOKEN must remain unset when rpk.yaml has no cached auth") - require.Equal(t, "https://aigw.example.com", os.Getenv(envEndpoint)) -} - -// TestResolveAndInjectEnv_SkipsWhenEndpointFlagPresent confirms that passing -// --rpai-endpoint on the command line suppresses the cluster lookup (and -// therefore works even with no aigw-attached cluster). -func TestResolveAndInjectEnv_SkipsWhenEndpointFlagPresent(t *testing.T) { - t.Setenv(envAuthToken, "already-set") - t.Setenv(envEndpoint, "") - - fs := afero.NewMemMapFs() - pluginArgs := []string{"--rpai-endpoint=https://custom.example.com", "llm", "list"} - require.NoError(t, resolveAndInjectEnv(t.Context(), fs, new(config.Params), pluginArgs)) - require.Empty(t, os.Getenv(envEndpoint), "RPAI_ENDPOINT must stay unset when flag is present") -} - -func TestHasEndpointFlag(t *testing.T) { - cases := []struct { - args []string - want bool - }{ - {nil, false}, - {[]string{"llm", "list"}, false}, - {[]string{"--rpai-endpoint", "https://foo"}, true}, - {[]string{"--rpai-endpoint=https://foo"}, true}, - {[]string{"llm", "list", "--rpai-endpoint=https://foo"}, true}, - // Not a match: prefix only, same-flag-family shouldn't false-positive. - {[]string{"--rpai-endpoint-other"}, false}, - } - for _, c := range cases { - t.Run(strings.Join(c.args, " "), func(t *testing.T) { - require.Equal(t, c.want, hasEndpointFlag(c.args)) - }) - } -} - -// clusterHandler stubs the public-API GetCluster endpoint to return the given -// cluster. If cluster is nil, it returns a cluster with no AI Gateway set. -func clusterHandler(t *testing.T, cluster *controlplanev1.Cluster) http.HandlerFunc { - t.Helper() - return func(w http.ResponseWriter, r *http.Request) { - if !strings.Contains(r.URL.Path, "ClusterService/GetCluster") { - http.NotFound(w, r) - return - } - resp := &controlplanev1.GetClusterResponse{Cluster: cluster} - marshal, err := proto.Marshal(resp) - require.NoError(t, err) - w.Header().Set("Content-Type", "application/proto") - w.Write(marshal) - } -} - -func TestResolveAigwEndpoint_NoClusterSelected(t *testing.T) { - ts := httptest.NewServer(clusterHandler(t, nil)) - defer ts.Close() - - // No profile written. - t.Setenv("RPK_PUBLIC_API_URL", ts.URL) - fs := afero.NewMemMapFs() - p := new(config.Params) - cfg, err := p.Load(fs) - require.NoError(t, err) - - _, err = resolveAigwEndpoint(t.Context(), cfg) - require.Error(t, err) - require.Contains(t, err.Error(), "no cluster selected") -} - -func TestResolveAigwEndpoint_NoAiGatewayV2(t *testing.T) { - // Cluster exists but has no AiGateway attached. - cluster := &controlplanev1.Cluster{Id: "clu-1"} - ts := httptest.NewServer(clusterHandler(t, cluster)) - defer ts.Close() - - cfg := loadCloudProfile(t, ts.URL, "clu-1") - _, err := resolveAigwEndpoint(t.Context(), cfg) - require.Error(t, err) - require.Contains(t, err.Error(), "does not have an AI Gateway v2 endpoint") -} - -func TestResolveAigwEndpoint_Happy(t *testing.T) { - cluster := &controlplanev1.Cluster{ - Id: "clu-1", - AiGateway: &controlplanev1.Cluster_AIGateway{ - V2Url: "https://aigw.example.com", - }, - } - ts := httptest.NewServer(clusterHandler(t, cluster)) - defer ts.Close() - - cfg := loadCloudProfile(t, ts.URL, "clu-1") - endpoint, err := resolveAigwEndpoint(t.Context(), cfg) - require.NoError(t, err) - require.Equal(t, "https://aigw.example.com", endpoint) -} - -// TestResolveAigwEndpoint_UsesCachedURL verifies that when the profile already -// carries an AIGatewayURL (populated at profile creation), we skip the -// publicapi lookup entirely. The httptest server below has no handler — if -// resolveAigwEndpoint touched the network, the test would fail. -func TestResolveAigwEndpoint_UsesCachedURL(t *testing.T) { - ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - t.Fatalf("unexpected publicapi call: cached AIGatewayURL should short-circuit") - w.WriteHeader(http.StatusInternalServerError) - })) - defer ts.Close() - - cfg := loadCloudProfile(t, ts.URL, "clu-1") - prof := cfg.VirtualProfile() - require.NotNil(t, prof) - prof.CloudCluster.AIGatewayURL = "https://cached-aigw.example.com" - - endpoint, err := resolveAigwEndpoint(t.Context(), cfg) - require.NoError(t, err) - require.Equal(t, "https://cached-aigw.example.com", endpoint) -} - -// loadCloudProfile builds a *config.Config with a cloud profile selecting the -// given cluster ID, pointing at publicAPIURL. -func loadCloudProfile(t *testing.T, publicAPIURL, clusterID string) *config.Config { - t.Helper() - t.Setenv("RPK_PUBLIC_API_URL", publicAPIURL) - t.Setenv("RPK_CLOUD_TOKEN", "test-token") - // Bazel's test sandbox clears HOME, which breaks - // config.DefaultRpkYamlPath -> os.UserConfigDir. Pin HOME (and - // XDG_CONFIG_HOME for Linux) to a stable tmpdir so the path we write - // to is the same path the loader reads from. - home := t.TempDir() - t.Setenv("HOME", home) - t.Setenv("XDG_CONFIG_HOME", home+"/.config") - - fs := afero.NewMemMapFs() - // DefaultRpkYamlPath resolves via os.UserConfigDir, which is - // platform-dependent (~/.config/rpk on Linux, ~/Library/Application - // Support/rpk on darwin). Writing to that exact path means the config - // loader will actually find our yaml. - path, err := config.DefaultRpkYamlPath() - require.NoError(t, err) - - yaml := fmt.Sprintf(`version: 6 -current_profile: dev -profiles: - - name: dev - from_cloud: true - cloud_cluster: - cluster_id: %q -`, clusterID) - require.NoError(t, afero.WriteFile(fs, path, []byte(yaml), 0o600)) - - p := new(config.Params) - cfg, err := p.Load(fs) - require.NoError(t, err) - return cfg -}