diff --git a/internal/cmd/functions.go b/internal/cmd/functions.go index 2079156..0514dee 100644 --- a/internal/cmd/functions.go +++ b/internal/cmd/functions.go @@ -24,6 +24,7 @@ var ( functionRunVariables []string // Variables as key=value pairs functionRunVariablesJSON string // Variables as JSON string functionRunNoStream bool // Opt out of log streaming + functionRunRuntime string // Which execution runtime to run on functionSecretValue string functionDownloadVersion string ) @@ -326,6 +327,10 @@ func init() { // a positive flag would be an opt-in to something already on. Same shape, // and the same reasoning, as --no-solve-captchas on sessions start. functionsRunCmd.Flags().BoolVar(&functionRunNoStream, "no-stream", false, "Return only the final response instead of streaming logs") + functionsRunCmd.Flags().StringVar(&functionRunRuntime, "runtime", "", fmt.Sprintf("Execution runtime: %s (Lambda) or %s (the configured AgentCore runtime). Server default when unset", api.Standard, api.Extended)) + _ = functionsRunCmd.RegisterFlagCompletionFunc("runtime", func(*cobra.Command, []string, string) ([]string, cobra.ShellCompDirective) { + return []string{string(api.Standard), string(api.Extended)}, cobra.ShellCompDirectiveNoFileComp + }) // Runs command flags functionsRunsCmd.Flags().StringVar(&functionID, "function-id", "", "Function ID (uses current function if not specified)") @@ -763,6 +768,19 @@ func runFunctionRun(cmd *cobra.Command, args []string) error { if functionRunNoStream { requestBody["stream"] = false } + // Same reasoning for omitting it when unset: the server picks the runtime, + // and sending its current choice back would pin it. The valid values are + // spelled out here rather than left to the API so a typo costs a message + // instead of a 422 - at the price of needing an edit here if the spec grows + // a third runtime. + if functionRunRuntime != "" { + switch api.RunFunctionRequestRuntime(functionRunRuntime) { + case api.Standard, api.Extended: + requestBody["runtime"] = functionRunRuntime + default: + return fmt.Errorf("invalid --runtime %q: expected %s or %s", functionRunRuntime, api.Standard, api.Extended) + } + } bodyJSON, err := json.Marshal(requestBody) if err != nil { diff --git a/internal/cmd/functionsextra_test.go b/internal/cmd/functionsextra_test.go index b70bffc..445fc00 100644 --- a/internal/cmd/functionsextra_test.go +++ b/internal/cmd/functionsextra_test.go @@ -382,3 +382,80 @@ func TestFunctionRun_SendsStreamOnlyWithNoStream(t *testing.T) { }) } } + +func TestFunctionRun_SendsRuntimeOnlyWhenSet(t *testing.T) { + for _, tc := range []struct { + name string + runtime string + want any + }{ + {name: "default omits runtime", runtime: "", want: nil}, + {name: "--runtime standard is sent", runtime: "standard", want: "standard"}, + {name: "--runtime extended is sent", runtime: "extended", want: "extended"}, + } { + t.Run(tc.name, func(t *testing.T) { + server := setupFunctionTest(t) + server.AddResponse("/functions/"+functionIDTest+"/runs/start", 200, functionRunJSON()) + + origFormat := outputFormat + origRuntime := functionRunRuntime + outputFormat = "json" + functionRunRuntime = tc.runtime + t.Cleanup(func() { + outputFormat = origFormat + functionRunRuntime = origRuntime + }) + + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + + testutil.CaptureOutput(func() { + if err := runFunctionRun(cmd, nil); err != nil { + t.Fatalf("unexpected error: %v", err) + } + }) + + requests := server.Requests("/functions/" + functionIDTest + "/runs/start") + if len(requests) != 1 { + t.Fatalf("got %d requests, want 1", len(requests)) + } + body := requestBody(t, requests[0]) + got, present := body["runtime"] + if tc.want == nil { + if present { + t.Errorf("runtime was sent as %v without --runtime", got) + } + return + } + if !present { + t.Fatal("runtime was not sent despite --runtime") + } + if got != tc.want { + t.Errorf("runtime = %v, want %v", got, tc.want) + } + }) + } +} + +func TestFunctionRun_RejectsUnknownRuntime(t *testing.T) { + server := setupFunctionTest(t) + server.AddResponse("/functions/"+functionIDTest+"/runs/start", 200, functionRunJSON()) + + origRuntime := functionRunRuntime + functionRunRuntime = "lambda" + t.Cleanup(func() { functionRunRuntime = origRuntime }) + + cmd := &cobra.Command{} + cmd.SetContext(context.Background()) + + err := runFunctionRun(cmd, nil) + if err == nil { + t.Fatal("expected an error for an unknown runtime") + } + if !strings.Contains(err.Error(), "invalid --runtime") { + t.Errorf("error = %q, want it to mention invalid --runtime", err) + } + if got := server.Requests("/functions/" + functionIDTest + "/runs/start"); len(got) != 0 { + t.Errorf("sent %d requests despite an invalid runtime, want 0", len(got)) + } +}