1010package ai
1111
1212import (
13- "context"
14- "fmt"
15- "os"
1613 "slices"
17- "strings"
1814
1915 "github.com/redpanda-data/redpanda/src/go/rpk/pkg/cobraext"
2016 "github.com/redpanda-data/redpanda/src/go/rpk/pkg/config"
21- "github.com/redpanda-data/redpanda/src/go/rpk/pkg/publicapi"
22- "github.com/spf13/afero"
2317 "github.com/spf13/cobra"
2418 "go.uber.org/zap"
2519)
2620
27- // The rpk ai plugin reads the following environment variables and flag:
28- //
29- // - envAuthToken: bearer token forwarded as Authorization to the AI
30- // Gateway. rpk reads the cached token off the active cloud profile
31- // and exports it; a missing or stale token is reported by the
32- // gateway as a 401, at which point the user runs `rpk cloud login`.
33- // - envEndpoint: AI Gateway v2 base URL. rpk resolves it from the
34- // active cloud profile's cached AIGatewayURL (or falls back to a
35- // publicapi lookup) and exports it before exec.
36- // - flagEndpoint: same intent as envEndpoint, but on the command
37- // line. rpk only watches for it so we can skip the cluster lookup;
38- // the flag is parsed and consumed by the plugin itself.
39- //
40- // Any explicit value the user supplies (env or flag) wins — rpk only
41- // fills in missing pieces.
42- const (
43- envAuthToken = "RPAI_TOKEN"
44- envEndpoint = "RPAI_ENDPOINT"
45-
46- flagEndpoint = "rpai-endpoint"
47- )
48-
49- // Exports RPAI_TOKEN and RPAI_ENDPOINT for the child plugin using cached
50- // values from rpk.yaml and the profile's AIGatewayURL (no OAuth refresh).
51- // Skips env var writes if the vars are already set or --rpai-endpoint is
52- // present.
53- func resolveAndInjectEnv (ctx context.Context , fs afero.Fs , p * config.Params , pluginArgs []string ) error {
54- cfg , err := p .Load (fs )
55- if err != nil {
56- return fmt .Errorf ("unable to load rpk config: %w" , err )
57- }
58-
59- if os .Getenv (envAuthToken ) == "" {
60- // Mirrors the lookup chain in oauth.LoadCloudToken (without the
61- // refresh) so --config, -X cloud_auth.*, and RPK_PROFILE all work.
62- auth := cfg .VirtualProfile ().VirtualAuth ()
63- if auth == nil {
64- auth = cfg .VirtualRpkYaml ().CurrentAuth ()
65- }
66- if auth != nil && auth .AuthToken != "" {
67- if err := os .Setenv (envAuthToken , auth .AuthToken ); err != nil {
68- return fmt .Errorf ("unable to set %s: %w" , envAuthToken , err )
69- }
70- }
71- }
72-
73- if os .Getenv (envEndpoint ) == "" && ! hasEndpointFlag (pluginArgs ) {
74- endpoint , err := resolveAigwEndpoint (ctx , cfg )
75- if err != nil {
76- return err
77- }
78- if err := os .Setenv (envEndpoint , endpoint ); err != nil {
79- return fmt .Errorf ("unable to set %s: %w" , envEndpoint , err )
80- }
81- }
82-
83- return nil
84- }
85-
86- // skipCloudForHelp reports whether a --help / -h / --version flag is present,
87- // in which case we must not reach out to the cloud API or trigger OAuth. The
88- // rpk ai plugin child process renders its own help/version output locally.
89- func skipCloudForHelp (args []string ) bool {
90- for _ , a := range args {
91- if a == "--help" || a == "-h" || a == "--version" {
92- return true
93- }
94- }
95- return false
96- }
97-
9821// parseFlags splits args into plugin args + rpk-global-flags consumed by rpk,
9922// and parses the rpk-global-flags so the logger and config loader pick them up.
23+ //
24+ // rpk does not inject any cloud context (token or endpoint) into the plugin:
25+ // the rpk ai plugin owns its own login (`rpk ai auth login`) and environment
26+ // selection (`rpk ai env use`), so it runs without a selected rpk cloud
27+ // cluster. Everything except rpk's own globals is forwarded untouched.
10028func parseFlags (p * config.Params , cmd * cobra.Command , args []string ) ([]string , error ) {
10129 f := cmd .Flags ()
10230
@@ -114,44 +42,3 @@ func parseFlags(p *config.Params, cmd *cobra.Command, args []string) ([]string,
11442 }
11543 return keepForPlugin , nil
11644}
117-
118- // hasEndpointFlag reports whether the plugin args carry an explicit endpoint
119- // flag (in any supported form: --flag=value or --flag followed by its
120- // value).
121- func hasEndpointFlag (args []string ) bool {
122- prefix := "--" + flagEndpoint
123- for _ , a := range args {
124- if a == prefix || strings .HasPrefix (a , prefix + "=" ) {
125- return true
126- }
127- }
128- return false
129- }
130-
131- // resolveAigwEndpoint returns the AI Gateway v2 URL for the active rpk cloud
132- // profile's cluster. It prefers the value cached on the profile at creation
133- // time and only falls back to a live publicapi lookup when the cache is empty
134- // (older profiles created before AIGatewayURL existed, or profiles whose
135- // cluster had no AI Gateway attached at creation but does now).
136- func resolveAigwEndpoint (ctx context.Context , cfg * config.Config ) (string , error ) {
137- prof := cfg .VirtualProfile ()
138- if prof == nil || ! prof .FromCloud || prof .CloudCluster .ClusterID == "" {
139- return "" , fmt .Errorf ("no cluster selected for this rpk profile; run 'rpk cloud cluster use <id>' or pass --%s" , flagEndpoint )
140- }
141- if prof .CloudCluster .AIGatewayURL != "" {
142- return prof .CloudCluster .AIGatewayURL , nil
143- }
144- clusterID := prof .CloudCluster .ClusterID
145-
146- token := os .Getenv (envAuthToken )
147- cl := publicapi .NewCloudClientSet (cfg .DevOverrides ().PublicAPIURL , token )
148- cluster , err := cl .ClusterForID (ctx , clusterID )
149- if err != nil {
150- return "" , fmt .Errorf ("unable to resolve aigw endpoint for cluster %s: %w" , clusterID , err )
151- }
152- endpoint := cluster .GetAiGateway ().GetV2Url ()
153- if endpoint == "" {
154- return "" , fmt .Errorf ("cluster %s does not have an AI Gateway v2 endpoint; pick a cluster that does, or pass --%s" , clusterID , flagEndpoint )
155- }
156- return endpoint , nil
157- }
0 commit comments