diff --git a/cli/ask.go b/cli/ask.go new file mode 100644 index 000000000..897672636 --- /dev/null +++ b/cli/ask.go @@ -0,0 +1,222 @@ +package cli + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/url" + "strings" + + "github.com/spf13/cobra" + "k8s.io/cli-runtime/pkg/genericclioptions" +) + +const ( + ErrQueryEmpty = "no query provided" + ErrNoEndpoint = "endpoint not resolved" + ErrStreamIncomplete = "response may be incomplete (stream interrupted)" + ErrMalformedEnd = "failed to parse end event" + ErrMissingEnd = "stream ended without end event" +) + +// AskOptions holds the configuration for the ask command. +type AskOptions struct { + streams genericclioptions.IOStreams + query string + endpoint string + kubeConfig *KubeConfig + mode string + insecureAllowHTTP bool + + // conversationID is extracted from the start event during Run. + // Available after Run completes for conversation persistence (OLS-3636). + conversationID string + + // capturedEvents accumulates non-token, non-end events (start, reasoning, + // tool_call, tool_result) during Run. Not displayed in default mode but + // available for --output json (OLS-3639). + capturedEvents []SSEEvent +} + +// NewAskCmd creates the "ask" subcommand that sends a question to OLS +// and streams back the response. +func NewAskCmd(streams genericclioptions.IOStreams) *cobra.Command { + o := &AskOptions{ + streams: streams, + mode: "ask", + } + + cmd := &cobra.Command{ + Use: "ask [question]", + Short: "Ask OpenShift Lightspeed a question", + Long: "Send a question to OpenShift Lightspeed and stream the response.", + RunE: func(cmd *cobra.Command, args []string) error { + if err := o.Complete(cmd, args); err != nil { + return err + } + if err := o.Validate(); err != nil { + return err + } + return o.Run(cmd) + }, + Args: cobra.ArbitraryArgs, + SilenceUsage: true, + } + + return cmd +} + +// Complete resolves the query string, kubeconfig, and endpoint. +func (o *AskOptions) Complete(cmd *cobra.Command, args []string) error { + o.query = strings.Join(args, " ") + + kubeconfigPath, _ := cmd.Flags().GetString("kubeconfig") + contextName, _ := cmd.Flags().GetString("context") + insecureSkipTLS, _ := cmd.Flags().GetBool("insecure-skip-tls-verify") + caCertPath, _ := cmd.Flags().GetString("ca-cert") + insecureAllowHTTP, _ := cmd.Flags().GetBool("insecure-allow-http") + + kc, err := LoadKubeConfig(kubeconfigPath, contextName, insecureSkipTLS, caCertPath) + if err != nil { + return err + } + o.kubeConfig = kc + o.insecureAllowHTTP = insecureAllowHTTP + + endpoint, err := ResolveEndpoint(cmd, kc.ContextName) + if err != nil { + return err + } + o.endpoint = endpoint + + return nil +} + +// Validate checks that required fields are populated. +func (o *AskOptions) Validate() error { + if strings.TrimSpace(o.query) == "" { + return fmt.Errorf("%s: provide a question as arguments", ErrQueryEmpty) + } + if o.endpoint == "" { + return errors.New(ErrNoEndpoint) + } + // Reject cleartext HTTP to prevent sending bearer token unencrypted. + parsed, err := url.Parse(o.endpoint) + if err != nil { + return fmt.Errorf("invalid endpoint URL: %w", err) + } + if parsed.Scheme == "http" && !o.insecureAllowHTTP { + return fmt.Errorf("cleartext HTTP endpoint %q is not allowed: bearer token would be sent unencrypted. Use https:// or reconfigure with: oc ols config set-endpoint", o.endpoint) + } + return nil +} + +// Run executes the ask command: sends the query via SSE and streams +// token data to stdout. Referenced documents from the end event are +// printed to stdout along with the response. +func (o *AskOptions) Run(cmd *cobra.Command) error { + client := NewSSEClient(o.endpoint, o.kubeConfig.BearerToken, o.kubeConfig.TLSConfig) + + req := LLMRequest{ + Query: o.query, + Mode: o.mode, + MediaType: "application/json", + } + + ctx := cmd.Context() + if ctx == nil { + ctx = context.Background() + } + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + events, errc, err := client.StreamQuery(ctx, req) + if err != nil { + return err + } + + var endData *EndEventData + var hasTokens bool + var endParseErr error + o.capturedEvents = nil + + for ev := range events { + switch ev.Type { + case EventToken: + var td TokenEventData + if err := json.Unmarshal([]byte(ev.Data), &td); err != nil { + // If token data isn't JSON, use raw string as fallback + td.Token = ev.Data + } + hasTokens = true + if _, err := fmt.Fprint(o.streams.Out, td.Token); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + case EventStart: + // Extract conversation_id for persistence (OLS-3636). + var sd StartEventData + if err := json.Unmarshal([]byte(ev.Data), &sd); err == nil { + o.conversationID = sd.ConversationID + } + o.capturedEvents = append(o.capturedEvents, ev) + case EventEnd: + var ed EndEventData + if err := json.Unmarshal([]byte(ev.Data), &ed); err != nil { + endParseErr = err + } else { + endData = &ed + } + case EventReasoning, EventToolCall, EventToolResult: + // Captured but not displayed in default mode. + // Available via o.capturedEvents for --output json (OLS-3639). + o.capturedEvents = append(o.capturedEvents, ev) + default: + // Unknown event types are silently ignored. + } + } + + // Check for stream-level errors + if streamErr := <-errc; streamErr != nil { + if _, err := fmt.Fprintf(o.streams.ErrOut, "Warning: %s\n", ErrStreamIncomplete); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + return streamErr + } + + // Print trailing newline only if tokens were emitted + if hasTokens { + if _, err := fmt.Fprintln(o.streams.Out); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + } + + // Malformed or missing end event means the stream was not fully valid + if endParseErr != nil { + if _, err := fmt.Fprintf(o.streams.ErrOut, "Warning: %s: %v\n", ErrMalformedEnd, endParseErr); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + return fmt.Errorf("%s: %w", ErrMalformedEnd, endParseErr) + } + + if endData == nil { + if _, err := fmt.Fprintf(o.streams.ErrOut, "Warning: %s\n", ErrStreamIncomplete); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + return errors.New(ErrMissingEnd) + } + + // Display referenced documents on stdout + if endData != nil && len(endData.ReferencedDocuments) > 0 { + if _, err := fmt.Fprintf(o.streams.Out, "\nReferences:\n"); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + for _, doc := range endData.ReferencedDocuments { + if _, err := fmt.Fprintf(o.streams.Out, " - %s: %s\n", doc.DocTitle, doc.DocURL); err != nil { + return fmt.Errorf("%s: %w", ErrWriteOutput, err) + } + } + } + + return nil +} diff --git a/cli/ask_test.go b/cli/ask_test.go new file mode 100644 index 000000000..2b3b93ce6 --- /dev/null +++ b/cli/ask_test.go @@ -0,0 +1,331 @@ +package cli + +import ( + "fmt" + "net/http" + "net/http/httptest" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("AskCmd", func() { + // sseServer creates a test HTTP server that returns a canned SSE response. + sseServer := func(sseBody string) *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, sseBody) + })) + } + + // buildEndEvent builds an end event payload matching the real server format. + buildEndEvent := func(docs []ReferencedDocument) string { + payload := map[string]interface{}{ + "referenced_documents": docs, + "truncated": false, + "input_tokens": 100, + "output_tokens": 50, + "reasoning_tokens": 0, + } + return sseEndEvent(payload) + } + + Describe("Run", func() { + It("streams token events to stdout and extracts conversation_id", func() { + body := sseEvent(EventStart, map[string]interface{}{"conversation_id": "conv-123"}) + + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "Hello"}) + + sseEvent(EventToken, map[string]interface{}{"id": 1, "token": " world"}) + + buildEndEvent([]ReferencedDocument{}) + + server := sseServer(body) + defer server.Close() + + streams, out, _ := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test question", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + Expect(o.Run(cmd)).To(Succeed()) + Expect(out.String()).To(Equal("Hello world\n")) + Expect(o.conversationID).To(Equal("conv-123")) + }) + + It("displays referenced documents on stdout", func() { + docs := []ReferencedDocument{ + {DocTitle: "Pod Debugging", DocURL: "https://docs.example.com/pods"}, + {DocTitle: "Logs Guide", DocURL: "https://docs.example.com/logs"}, + } + body := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "answer"}) + + buildEndEvent(docs) + + server := sseServer(body) + defer server.Close() + + streams, out, _ := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + Expect(o.Run(cmd)).To(Succeed()) + Expect(out.String()).To(ContainSubstring("References:")) + Expect(out.String()).To(ContainSubstring("Pod Debugging")) + Expect(out.String()).To(ContainSubstring("https://docs.example.com/pods")) + }) + + It("handles stream with no referenced documents", func() { + body := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "simple answer"}) + + buildEndEvent([]ReferencedDocument{}) + + server := sseServer(body) + defer server.Close() + + streams, out, _ := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + Expect(o.Run(cmd)).To(Succeed()) + Expect(out.String()).To(ContainSubstring("simple answer")) + Expect(out.String()).NotTo(ContainSubstring("References:")) + }) + + It("propagates HTTP errors from SSEClient", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + streams, _, _ := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "bad-token", + }, + } + + cmd := NewAskCmd(streams) + err := o.Run(cmd) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrAuthFailed)) + }) + + It("captures reasoning, tool_call, and tool_result events without displaying them", func() { + body := sseEvent(EventStart, map[string]interface{}{"conversation_id": "conv-abc"}) + + sseEvent(EventReasoning, map[string]interface{}{"content": "thinking..."}) + + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "result"}) + + sseEvent(EventToolCall, map[string]interface{}{"name": "search", "args": map[string]string{"q": "test"}, "id": "call_1", "type": "tool_call"}) + + sseEvent(EventToolResult, map[string]interface{}{"id": "call_1", "name": "search", "status": "success", "content": "found it", "type": "tool_result", "round": 1}) + + buildEndEvent([]ReferencedDocument{}) + + server := sseServer(body) + defer server.Close() + + streams, out, errOut := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + Expect(o.Run(cmd)).To(Succeed()) + + // Not displayed to stdout or stderr + Expect(out.String()).To(Equal("result\n")) + Expect(errOut.String()).NotTo(ContainSubstring("thinking")) + Expect(errOut.String()).NotTo(ContainSubstring("search")) + Expect(errOut.String()).NotTo(ContainSubstring("found it")) + + // But captured internally for --output json (OLS-3639) + // start + reasoning + tool_call + tool_result = 4 captured events + Expect(o.capturedEvents).To(HaveLen(4)) + Expect(o.capturedEvents[0].Type).To(Equal(EventStart)) + Expect(o.capturedEvents[1].Type).To(Equal(EventReasoning)) + Expect(o.capturedEvents[2].Type).To(Equal(EventToolCall)) + Expect(o.capturedEvents[3].Type).To(Equal(EventToolResult)) + }) + + It("warns on malformed end event JSON", func() { + // Simulate an end event where the inner data is not valid EndEventData + // but the envelope is still valid JSON + body := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "answer"}) + + sseEvent(EventEnd, "not-a-json-object") + + server := sseServer(body) + defer server.Close() + + streams, _, errOut := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + err := o.Run(cmd) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrMalformedEnd)) + Expect(errOut.String()).To(ContainSubstring(ErrMalformedEnd)) + }) + + It("returns error when stream has no end event", func() { + body := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "partial"}) + + server := sseServer(body) + defer server.Close() + + streams, _, errOut := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + err := o.Run(cmd) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrMissingEnd)) + Expect(errOut.String()).To(ContainSubstring(ErrStreamIncomplete)) + }) + + It("does not print trailing newline when no tokens were emitted", func() { + body := buildEndEvent([]ReferencedDocument{}) + + server := sseServer(body) + defer server.Close() + + streams, out, _ := fakeStreams() + o := &AskOptions{ + streams: streams, + query: "test", + endpoint: server.URL, + insecureAllowHTTP: true, + mode: "ask", + kubeConfig: &KubeConfig{ + BearerToken: "test-token", + }, + } + + cmd := NewAskCmd(streams) + Expect(o.Run(cmd)).To(Succeed()) + Expect(out.String()).To(Equal("")) + }) + }) + + Describe("Validate", func() { + It("rejects empty query", func() { + o := &AskOptions{ + query: "", + endpoint: "https://example.com", + } + err := o.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrQueryEmpty)) + }) + + It("rejects whitespace-only query", func() { + o := &AskOptions{ + query: " ", + endpoint: "https://example.com", + } + err := o.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrQueryEmpty)) + }) + + It("rejects missing endpoint", func() { + o := &AskOptions{ + query: "test question", + endpoint: "", + } + err := o.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrNoEndpoint)) + }) + + It("accepts valid options", func() { + o := &AskOptions{ + query: "why is my pod crashing", + endpoint: "https://ols.example.com", + } + Expect(o.Validate()).To(Succeed()) + }) + + It("rejects cleartext HTTP endpoint by default", func() { + o := &AskOptions{ + query: "test", + endpoint: "http://ols.example.com", + } + err := o.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("cleartext HTTP")) + }) + + It("allows cleartext HTTP with --insecure-allow-http", func() { + o := &AskOptions{ + query: "test", + endpoint: "http://ols.example.com", + insecureAllowHTTP: true, + } + Expect(o.Validate()).To(Succeed()) + }) + }) + + Describe("NewAskCmd", func() { + It("creates command with correct use string", func() { + streams, _, _ := fakeStreams() + cmd := NewAskCmd(streams) + Expect(cmd.Use).To(Equal("ask [question]")) + }) + + It("accepts arbitrary args", func() { + streams, _, _ := fakeStreams() + cmd := NewAskCmd(streams) + Expect(cmd.Args).NotTo(BeNil()) + }) + }) +}) diff --git a/cli/integration_test.go b/cli/integration_test.go new file mode 100644 index 000000000..b95d09ffc --- /dev/null +++ b/cli/integration_test.go @@ -0,0 +1,195 @@ +package cli + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Integration", func() { + // createKubeconfig writes a minimal kubeconfig file with token auth + // pointing at the given server URL, returning the path. + createKubeconfig := func(serverURL, token string) string { + f, err := os.CreateTemp("", "kubeconfig-*.yaml") + Expect(err).NotTo(HaveOccurred()) + + content := fmt.Sprintf(`apiVersion: v1 +kind: Config +current-context: test-ctx +clusters: +- cluster: + server: %s + insecure-skip-tls-verify: true + name: test-cluster +contexts: +- context: + cluster: test-cluster + user: test-user + name: test-ctx +users: +- name: test-user + user: + token: %s +`, serverURL, token) + + _, err = f.WriteString(content) + Expect(err).NotTo(HaveOccurred()) + Expect(f.Close()).To(Succeed()) + return f.Name() + } + + // buildEndEvent builds an end event in the real server format. + buildEndEvent := func(docs []ReferencedDocument) string { + payload := map[string]interface{}{ + "referenced_documents": docs, + "truncated": false, + "input_tokens": 100, + "output_tokens": 50, + "reasoning_tokens": 0, + } + return sseEndEvent(payload) + } + + Describe("full command flow via default mode", func() { + It("sends query and streams response through root command dispatch", func() { + var capturedMethod string + var capturedPath string + var capturedAuth string + var capturedBody []byte + + docs := []ReferencedDocument{ + {DocTitle: "Pods Guide", DocURL: "https://docs.example.com/pods"}, + } + + sseBody := sseEvent(EventStart, map[string]interface{}{"conversation_id": "conv-integration-1"}) + + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "The pod is"}) + + sseEvent(EventToken, map[string]interface{}{"id": 1, "token": " crashing because"}) + + sseEvent(EventToken, map[string]interface{}{"id": 2, "token": " of OOM."}) + + buildEndEvent(docs) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + capturedMethod = r.Method + capturedPath = r.URL.Path + capturedAuth = r.Header.Get("Authorization") + capturedBody, _ = io.ReadAll(r.Body) + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, sseBody) + })) + defer server.Close() + + kubeconfigPath := createKubeconfig("https://k8s.example.com", "test-bearer-token") + defer os.Remove(kubeconfigPath) + + streams, out, _ := fakeStreams() + cmd := NewRootCmd(streams) + cmd.SetArgs([]string{ + "--kubeconfig", kubeconfigPath, + "--endpoint", server.URL, + "--insecure-skip-tls-verify", + "--insecure-allow-http", + "why is my pod crashing", + }) + + Expect(cmd.Execute()).To(Succeed()) + + // Verify correct HTTP request was sent + Expect(capturedMethod).To(Equal("POST")) + Expect(capturedPath).To(Equal("/v1/streaming_query")) + Expect(capturedAuth).To(Equal("Bearer test-bearer-token")) + + // Verify request body + var reqBody LLMRequest + Expect(json.Unmarshal(capturedBody, &reqBody)).To(Succeed()) + Expect(reqBody.Query).To(Equal("why is my pod crashing")) + Expect(reqBody.Mode).To(Equal("ask")) + Expect(reqBody.MediaType).To(Equal("application/json")) + + // Verify tokens appeared on stdout + Expect(out.String()).To(ContainSubstring("The pod is")) + Expect(out.String()).To(ContainSubstring("crashing because")) + Expect(out.String()).To(ContainSubstring("of OOM.")) + + // Verify referenced documents on stderr + Expect(out.String()).To(ContainSubstring("References:")) + Expect(out.String()).To(ContainSubstring("Pods Guide")) + Expect(out.String()).To(ContainSubstring("https://docs.example.com/pods")) + }) + + It("works with explicit ask subcommand", func() { + body := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "answer"}) + + buildEndEvent([]ReferencedDocument{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, body) + })) + defer server.Close() + + kubeconfigPath := createKubeconfig("https://k8s.example.com", "token-123") + defer os.Remove(kubeconfigPath) + + streams, out, _ := fakeStreams() + cmd := NewRootCmd(streams) + cmd.SetArgs([]string{ + "--kubeconfig", kubeconfigPath, + "--endpoint", server.URL, + "--insecure-skip-tls-verify", + "--insecure-allow-http", + "ask", "what is a deployment", + }) + + Expect(cmd.Execute()).To(Succeed()) + Expect(out.String()).To(ContainSubstring("answer")) + }) + + It("returns error on HTTP 401 through full dispatch", func() { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + kubeconfigPath := createKubeconfig("https://k8s.example.com", "expired-token") + defer os.Remove(kubeconfigPath) + + streams, _, _ := fakeStreams() + cmd := NewRootCmd(streams) + cmd.SetArgs([]string{ + "--kubeconfig", kubeconfigPath, + "--endpoint", server.URL, + "--insecure-skip-tls-verify", + "--insecure-allow-http", + "check my cluster", + }) + + err := cmd.Execute() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrAuthFailed)) + }) + + It("returns error when no query is provided via ask subcommand", func() { + kubeconfigPath := createKubeconfig("https://k8s.example.com", "token") + defer os.Remove(kubeconfigPath) + + streams, _, _ := fakeStreams() + cmd := NewRootCmd(streams) + cmd.SetArgs([]string{ + "--kubeconfig", kubeconfigPath, + "--endpoint", "https://ols.example.com", + "ask", + }) + + err := cmd.Execute() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrQueryEmpty)) + }) + }) +}) diff --git a/cli/kubeconfig.go b/cli/kubeconfig.go index 0c8225a38..54b74b275 100644 --- a/cli/kubeconfig.go +++ b/cli/kubeconfig.go @@ -55,6 +55,9 @@ func LoadKubeConfig(kubeconfigPath string, contextName string, insecureSkipTLS b restConfig, err := clientConfig.ClientConfig() if err != nil { + if clientcmd.IsEmptyConfig(err) || resolvedContext == "" { + return nil, fmt.Errorf("could not load kubeconfig: no valid configuration found. Try: oc login") + } return nil, fmt.Errorf("%s %q: %w", ErrResolveContext, resolvedContext, err) } diff --git a/cli/root.go b/cli/root.go index b7d5a49f7..f4e011b4c 100644 --- a/cli/root.go +++ b/cli/root.go @@ -18,6 +18,11 @@ const ( // NewRootCmd creates the root oc-ols command and registers subcommands. func NewRootCmd(streams genericclioptions.IOStreams) *cobra.Command { + askOpts := &AskOptions{ + streams: streams, + mode: "ask", + } + cmd := &cobra.Command{ Use: "oc-ols [command]", Short: "CLI for OpenShift Lightspeed", @@ -26,10 +31,14 @@ func NewRootCmd(streams genericclioptions.IOStreams) *cobra.Command { if len(args) == 0 { return cmd.Help() } - if _, err := fmt.Fprintf(streams.ErrOut, "ask command not yet implemented\n"); err != nil { - return fmt.Errorf("%s: %w", ErrWriteOutput, err) + // Default mode: dispatch unrecognized args to ask + if err := askOpts.Complete(cmd, args); err != nil { + return err + } + if err := askOpts.Validate(); err != nil { + return err } - return nil + return askOpts.Run(cmd) }, SilenceUsage: true, Args: cobra.ArbitraryArgs, @@ -43,6 +52,8 @@ func NewRootCmd(streams genericclioptions.IOStreams) *cobra.Command { "Path to kubeconfig file (default: $KUBECONFIG or ~/.kube/config)") cmd.PersistentFlags().Bool("insecure-skip-tls-verify", false, "Skip TLS certificate verification") + cmd.PersistentFlags().Bool("insecure-allow-http", false, + "Allow cleartext HTTP endpoints (development only)") cmd.PersistentFlags().String("context", "", "Kubeconfig context to use") cmd.PersistentFlags().String("ca-cert", "", @@ -53,6 +64,7 @@ func NewRootCmd(streams genericclioptions.IOStreams) *cobra.Command { cmd.AddCommand(NewVersionCmd(streams)) cmd.AddCommand(config.NewConfigCmd(streams)) + cmd.AddCommand(NewAskCmd(streams)) return cmd } diff --git a/cli/root_test.go b/cli/root_test.go index ff89f6983..85ba9c21e 100644 --- a/cli/root_test.go +++ b/cli/root_test.go @@ -29,7 +29,7 @@ var _ = Describe("RootCmd", func() { It("registers global flags", func() { streams, _, _ := fakeStreams() cmd := NewRootCmd(streams) - for _, name := range []string{"kubeconfig", "context", "insecure-skip-tls-verify", "ca-cert", "endpoint"} { + for _, name := range []string{"kubeconfig", "context", "insecure-skip-tls-verify", "insecure-allow-http", "ca-cert", "endpoint"} { Expect(cmd.PersistentFlags().Lookup(name)).NotTo(BeNil(), "expected persistent flag %q", name) } }) @@ -98,11 +98,17 @@ var _ = Describe("RootCmd", func() { }) }) - It("dispatches unrecognized args to the default mode stub", func() { - streams, _, errOut := fakeStreams() + It("dispatches unrecognized args to ask mode", func() { + // Point to a non-existent kubeconfig so the test doesn't pick up + // the real one from the environment. + streams, _, _ := fakeStreams() cmd := NewRootCmd(streams) - cmd.SetArgs([]string{"why is my pod crashing"}) - Expect(cmd.Execute()).To(Succeed()) - Expect(errOut.String()).To(ContainSubstring("not yet implemented")) + cmd.SetArgs([]string{"--kubeconfig", "/nonexistent/kubeconfig", "why is my pod crashing"}) + // Without a valid kubeconfig, this will fail at Complete() — + // but it proves dispatch happened (not "unknown command" error) + err := cmd.Execute() + Expect(err).To(HaveOccurred()) + // Should be a kubeconfig/auth error, not an "unknown command" error + Expect(err.Error()).NotTo(ContainSubstring("unknown command")) }) }) diff --git a/cli/streaming.go b/cli/streaming.go new file mode 100644 index 000000000..638c0639d --- /dev/null +++ b/cli/streaming.go @@ -0,0 +1,341 @@ +// Package cli implements the oc-ols kubectl plugin for querying OpenShift Lightspeed from the terminal. +package cli + +import ( + "bufio" + "bytes" + "context" + "crypto/tls" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "strings" + "time" +) + +// SSE event type constants matching lightspeed-service event types. +const ( + EventStart = "start" + EventToken = "token" + EventReasoning = "reasoning" + EventToolCall = "tool_call" + EventToolResult = "tool_result" + EventEnd = "end" + + ErrParseSSE = "failed to parse SSE stream" + ErrSendRequest = "failed to send request" //#nosec G101 -- error message, not a credential + ErrAuthFailed = "authentication failed" //#nosec G101 -- error message, not a credential + ErrAccessDenied = "access denied" + ErrPromptTooLong = "query exceeds maximum length" + ErrServiceError = "service error" + ErrStreamTimeout = "stream idle timeout" + ErrInsecureRedirect = "redirect to non-HTTPS URL blocked" //#nosec G101 -- error message, not a credential + + streamingQueryPath = "/v1/streaming_query" + + // HTTP transport timeouts. + connectTimeout = 10 * time.Second + tlsHandshakeTimeout = 10 * time.Second + responseHeaderTimeout = 30 * time.Second + + // IdleTimeout is the maximum duration to wait for the next SSE event + // before considering the stream stalled. + IdleTimeout = 120 * time.Second +) + +// SSEEvent represents a single parsed Server-Sent Event from the +// lightspeed-service streaming response. +type SSEEvent struct { + Type string + Data string +} + +// SSEClient sends queries to the lightspeed-service and streams back SSE events. +type SSEClient struct { + httpClient *http.Client + endpoint string + token string +} + +// NewSSEClient creates an SSE client configured with the given endpoint, +// bearer token, and TLS settings. +func NewSSEClient(endpoint, token string, tlsConfig *tls.Config) *SSEClient { + transport := &http.Transport{ + TLSClientConfig: tlsConfig, + DialContext: (&net.Dialer{ + Timeout: connectTimeout, + }).DialContext, + TLSHandshakeTimeout: tlsHandshakeTimeout, + ResponseHeaderTimeout: responseHeaderTimeout, + } + + return &SSEClient{ + httpClient: &http.Client{ + Transport: transport, + // Reject redirects to non-HTTPS targets to prevent leaking + // the bearer token over cleartext HTTP (CWE-319). + CheckRedirect: rejectHTTPRedirect, + }, + endpoint: endpoint, + token: token, + } +} + +// rejectHTTPRedirect prevents the HTTP client from following redirects to +// non-HTTPS URLs, which would expose the bearer token over cleartext. +func rejectHTTPRedirect(req *http.Request, _ []*http.Request) error { + if req.URL.Scheme != "https" { + return fmt.Errorf("%s: %s", ErrInsecureRedirect, req.URL) + } + return nil +} + +// StreamQuery sends a query to lightspeed-service and returns a channel of +// SSE events. The caller must drain the events channel. Immediate failures +// (connection errors, non-200 status) are returned as the error value. +// Stream-level errors (parse failures, idle timeout) are sent on the error +// channel. +func (c *SSEClient) StreamQuery(ctx context.Context, req LLMRequest) (<-chan SSEEvent, <-chan error, error) { + body, err := json.Marshal(req) + if err != nil { + return nil, nil, fmt.Errorf("%s: %w", ErrSendRequest, err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, + c.endpoint+streamingQueryPath, bytes.NewReader(body)) + if err != nil { + return nil, nil, fmt.Errorf("%s: %w", ErrSendRequest, err) + } + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("Accept", "text/event-stream") + httpReq.Header.Set("Authorization", "Bearer "+c.token) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, nil, fmt.Errorf("%s: %w", ErrSendRequest, err) + } + + if resp.StatusCode != http.StatusOK { + _ = resp.Body.Close() + return nil, nil, httpStatusError(resp.StatusCode) + } + + // Wrap the response body with an idle timeout so a stalled stream + // does not block forever. + reader := newIdleTimeoutReader(resp.Body, IdleTimeout) + + events, errc := parseSSEStream(ctx, reader) + return events, errc, nil +} + +// httpStatusError maps HTTP status codes to user-facing error messages +// matching the spec's error handling table. +func httpStatusError(status int) error { + switch status { + case http.StatusUnauthorized: + return fmt.Errorf("%s: is your login session active? Try: oc login", ErrAuthFailed) + case http.StatusForbidden: + return fmt.Errorf("%s: contact your cluster administrator to grant OLS access", ErrAccessDenied) + case http.StatusRequestEntityTooLarge: + return fmt.Errorf("%s: try a shorter question or fewer attachments", ErrPromptTooLong) + default: + return fmt.Errorf("%s: server returned %d", ErrServiceError, status) + } +} + +// idleTimeoutReader wraps an io.ReadCloser and enforces a maximum idle +// duration between reads. If no data arrives within the timeout window, +// subsequent reads return an error. +type idleTimeoutReader struct { + rc io.ReadCloser + timer *time.Timer + timeout time.Duration + failed error +} + +func newIdleTimeoutReader(rc io.ReadCloser, timeout time.Duration) *idleTimeoutReader { + return &idleTimeoutReader{ + rc: rc, + timer: time.NewTimer(timeout), + timeout: timeout, + } +} + +// Close stops the idle timer and closes the underlying reader. +func (r *idleTimeoutReader) Close() error { + r.timer.Stop() + return r.rc.Close() +} + +// readResult holds the outcome of a goroutine-based read. +type readResult struct { + n int + err error +} + +// Read implements io.Reader. Each successful read resets the idle timer. +// If the timer fires before data arrives, the underlying reader is closed +// and an error is returned. +func (r *idleTimeoutReader) Read(p []byte) (int, error) { + if r.failed != nil { + return 0, r.failed + } + + // Race the underlying read against the idle timer. + // A goroutine is needed because r.rc.Read may block indefinitely + // (e.g., waiting for the next SSE event from the network). + // The goroutine reads into a buffer it owns so it never touches p + // after this call returns. One goroutine per Read is acceptable + // because SSE streams have low read frequency (one per event frame). + buf := make([]byte, len(p)) + ch := make(chan readResult, 1) + go func() { + n, err := r.rc.Read(buf) + ch <- readResult{n, err} + }() + + select { + case res := <-ch: + copy(p, buf[:res.n]) + if res.n > 0 { + // Reset timer on successful read + if !r.timer.Stop() { + select { + case <-r.timer.C: + default: + } + } + r.timer.Reset(r.timeout) + } + if res.err != nil { + r.timer.Stop() + } + return res.n, res.err + case <-r.timer.C: + _ = r.rc.Close() + r.failed = fmt.Errorf("%s: no data received for %s", ErrStreamTimeout, r.timeout) + return 0, r.failed + } +} + +// parseSSEStream reads SSE frames from r and sends parsed events to the +// returned channel. Each frame is delimited by a blank line. The channel +// is closed when the reader is exhausted or an error occurs. +// +// The parser supports two framing styles: +// +// 1. Standard SSE with event: and data: lines: +// +// event: token +// data: {"id": 0, "token": "hello"} +// +// +// 2. JSON envelope (lightspeed-service current format) — event type is +// inside the JSON payload, no SSE event: line: +// +// data: {"event": "token", "data": {"id": 0, "token": "hello"}} +// +// +// When an SSE event: field is present, the parser uses it as the event +// type and passes data: content through as-is. When no event: field is +// present, the parser falls back to extracting the type from the JSON +// envelope. +func parseSSEStream(ctx context.Context, rc io.ReadCloser) (<-chan SSEEvent, <-chan error) { + events := make(chan SSEEvent) + errc := make(chan error, 1) + + go func() { + defer close(events) + defer close(errc) + defer func() { _ = rc.Close() }() + + scanner := bufio.NewScanner(rc) + // Increase scanner buffer for large tool_result payloads. + // Default is 64KB; tool_result events with resource dumps + // (e.g. pods_list) can reach ~35KB observed, ~500KB extreme. + scanner.Buffer(make([]byte, 0, bufio.MaxScanTokenSize), 1024*1024) + + var eventField string // SSE event: field for current frame + var dataLines []string + + for scanner.Scan() { + line := scanner.Text() + + // Blank line = end of frame, emit event if we have data + if line == "" { + if len(dataLines) > 0 { + raw := strings.Join(dataLines, "\n") + ev, err := buildSSEEvent(eventField, raw) + if err != nil { + errc <- fmt.Errorf("%s: %w", ErrParseSSE, err) + return + } + select { + case events <- ev: + case <-ctx.Done(): + return + } + dataLines = nil + } + eventField = "" + continue + } + + // SSE comment lines (starting with :) are ignored + if strings.HasPrefix(line, ":") { + continue + } + + // SSE event: field — sets the event type for this frame + if strings.HasPrefix(line, "event:") { + value := strings.TrimPrefix(line, "event:") + eventField = strings.TrimPrefix(value, " ") + continue + } + + if strings.HasPrefix(line, "data:") { + // SSE spec: strip exactly one leading space after "data:" + value := strings.TrimPrefix(line, "data:") + value = strings.TrimPrefix(value, " ") + dataLines = append(dataLines, value) + } + } + + if err := scanner.Err(); err != nil { + errc <- fmt.Errorf("%s: %w", ErrParseSSE, err) + } + }() + + return events, errc +} + +// buildSSEEvent constructs an SSEEvent from a parsed frame. When the SSE +// event: field is present, it is used as the event type and the raw data +// is passed through. Otherwise, the data is parsed as a JSON envelope +// to extract the event type (lightspeed-service format). +func buildSSEEvent(eventField, rawData string) (SSEEvent, error) { + if eventField != "" { + // Standard SSE framing: event type from event: line, data as-is. + return SSEEvent{Type: eventField, Data: rawData}, nil + } + // JSON envelope fallback: extract type from envelope. + return parseSSEEnvelope(rawData) +} + +// parseSSEEnvelope extracts the event type and inner data from a +// lightspeed-service JSON envelope: {"event": "", "data": }. +// The inner data is passed through as-is (json.RawMessage) to preserve +// numeric precision. +func parseSSEEnvelope(raw string) (SSEEvent, error) { + var envelope sseEnvelope + if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + return SSEEvent{}, fmt.Errorf("invalid JSON envelope: %w", err) + } + + return SSEEvent{ + Type: envelope.Event, + Data: string(envelope.Data), + }, nil +} diff --git a/cli/streaming_test.go b/cli/streaming_test.go new file mode 100644 index 000000000..e1209761d --- /dev/null +++ b/cli/streaming_test.go @@ -0,0 +1,478 @@ +package cli + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("parseSSEStream", func() { + // collectEvents drains the event channel and returns all events. + collectEvents := func(events <-chan SSEEvent, errc <-chan error) ([]SSEEvent, error) { + var result []SSEEvent + for ev := range events { + result = append(result, ev) + } + if err, ok := <-errc; ok && err != nil { + return result, err + } + return result, nil + } + + // nopCloser wraps a reader with a no-op Close for parseSSEStream. + nopCloser := io.NopCloser + + It("parses a single token event", func() { + input := sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "hello"}) + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result[0].Type).To(Equal(EventToken)) + + var td TokenEventData + Expect(json.Unmarshal([]byte(result[0].Data), &td)).To(Succeed()) + Expect(td.Token).To(Equal("hello")) + }) + + It("parses a complete SSE stream with multiple event types", func() { + input := sseEvent(EventStart, map[string]interface{}{"conversation_id": "abc-123"}) + + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "Hello"}) + + sseEvent(EventToken, map[string]interface{}{"id": 1, "token": " world"}) + + sseEvent(EventReasoning, map[string]interface{}{"content": "thinking about it"}) + + sseEvent(EventToolCall, map[string]interface{}{"name": "search", "args": map[string]string{"q": "test"}, "id": "call_1", "type": "tool_call"}) + + sseEvent(EventToolResult, map[string]interface{}{"id": "call_1", "name": "search", "status": "success", "content": "results", "type": "tool_result", "round": 1}) + + sseEndEvent(map[string]interface{}{ + "referenced_documents": []interface{}{}, + "truncated": false, + "input_tokens": 100, + "output_tokens": 50, + "reasoning_tokens": 0, + }) + + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(7)) + + Expect(result[0].Type).To(Equal(EventStart)) + Expect(result[1].Type).To(Equal(EventToken)) + Expect(result[2].Type).To(Equal(EventToken)) + Expect(result[3].Type).To(Equal(EventReasoning)) + Expect(result[4].Type).To(Equal(EventToolCall)) + Expect(result[5].Type).To(Equal(EventToolResult)) + Expect(result[6].Type).To(Equal(EventEnd)) + + // Verify start event has conversation_id + var start StartEventData + Expect(json.Unmarshal([]byte(result[0].Data), &start)).To(Succeed()) + Expect(start.ConversationID).To(Equal("abc-123")) + + // Verify token data + var tok TokenEventData + Expect(json.Unmarshal([]byte(result[1].Data), &tok)).To(Succeed()) + Expect(tok.Token).To(Equal("Hello")) + Expect(tok.ID).To(Equal(0)) + }) + + It("parses standard SSE framing with event: and data: lines", func() { + // Standard SSE format: event type in event: field, payload in data: field + input := "event: start\ndata: {\"conversation_id\": \"abc-123\"}\n\n" + + "event: token\ndata: {\"id\": 0, \"token\": \"Hello\"}\n\n" + + "event: token\ndata: {\"id\": 1, \"token\": \" world\"}\n\n" + + "event: end\ndata: {\"referenced_documents\": [], \"truncated\": false}\n\n" + + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(4)) + + Expect(result[0].Type).To(Equal(EventStart)) + Expect(result[1].Type).To(Equal(EventToken)) + Expect(result[2].Type).To(Equal(EventToken)) + Expect(result[3].Type).To(Equal(EventEnd)) + + // With standard framing, data is raw JSON (not envelope-wrapped) + var start StartEventData + Expect(json.Unmarshal([]byte(result[0].Data), &start)).To(Succeed()) + Expect(start.ConversationID).To(Equal("abc-123")) + + var tok TokenEventData + Expect(json.Unmarshal([]byte(result[1].Data), &tok)).To(Succeed()) + Expect(tok.Token).To(Equal("Hello")) + }) + + It("handles mixed framing (event: field present on some frames)", func() { + // Frame 1: standard SSE with event: field + // Frame 2: JSON envelope without event: field + input := "event: token\ndata: {\"id\": 0, \"token\": \"first\"}\n\n" + + sseEvent(EventToken, map[string]interface{}{"id": 1, "token": "second"}) + + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(2)) + + Expect(result[0].Type).To(Equal(EventToken)) + Expect(result[1].Type).To(Equal(EventToken)) + }) + + It("ignores SSE comment lines", func() { + input := ": keep-alive\n" + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "hi"}) + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result[0].Type).To(Equal(EventToken)) + }) + + It("returns no events for empty input", func() { + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(""))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(BeEmpty()) + }) + + It("propagates reader errors", func() { + events, errc := parseSSEStream(context.Background(), nopCloser(&failingReader{err: io.ErrUnexpectedEOF})) + result, err := collectEvents(events, errc) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrParseSSE)) + Expect(result).To(BeEmpty()) + }) + + It("returns error on invalid JSON in data line", func() { + input := "data: {not valid json}\n\n" + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrParseSSE)) + Expect(result).To(BeEmpty()) + }) + + It("handles tool_result events with large payloads", func() { + // Simulate a large tool_result like the real server sends + largeContent := strings.Repeat("pod info ", 1000) + input := sseEvent(EventToolResult, map[string]interface{}{ + "id": "call_abc", + "name": "pods_list", + "status": "success", + "content": largeContent, + "type": "tool_result", + "round": 1, + }) + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result[0].Type).To(Equal(EventToolResult)) + }) + + It("handles end event with available_quotas field", func() { + input := sseEndEvent(map[string]interface{}{ + "referenced_documents": []map[string]string{ + {"doc_url": "https://docs.example.com", "doc_title": "Test"}, + }, + "truncated": false, + "input_tokens": 2899, + "output_tokens": 191, + "reasoning_tokens": 0, + }) + events, errc := parseSSEStream(context.Background(), nopCloser(strings.NewReader(input))) + result, err := collectEvents(events, errc) + Expect(err).NotTo(HaveOccurred()) + Expect(result).To(HaveLen(1)) + Expect(result[0].Type).To(Equal(EventEnd)) + + var end EndEventData + Expect(json.Unmarshal([]byte(result[0].Data), &end)).To(Succeed()) + Expect(end.ReferencedDocuments).To(HaveLen(1)) + Expect(end.InputTokens).To(Equal(2899)) + }) +}) + +// failingReader always returns an error on Read. +type failingReader struct { + err error +} + +func (r *failingReader) Read(_ []byte) (int, error) { + return 0, r.err +} + +// sseHandler returns an http.HandlerFunc that writes SSE events to the response. +// It validates the request method, content type, and auth header. +func sseHandler(sseBody string) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + if r.Header.Get("Content-Type") != "application/json" { + http.Error(w, "bad content type", http.StatusBadRequest) + return + } + if r.Header.Get("Authorization") == "" { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(sseBody)) + } +} + +var _ = Describe("SSEClient", func() { + var ( + server *httptest.Server + client *SSEClient + ctx context.Context + cancel context.CancelFunc + ) + + AfterEach(func() { + if cancel != nil { + cancel() + } + if server != nil { + server.Close() + } + }) + + setupClient := func(handler http.HandlerFunc) { + server = httptest.NewServer(handler) + client = NewSSEClient(server.URL, "test-token", nil) + ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second) + } + + baseRequest := func() LLMRequest { + return LLMRequest{ + Query: "why is my pod crashing", + Mode: "ask", + MediaType: "application/json", + } + } + + Describe("StreamQuery", func() { + It("streams token events from a valid SSE response", func() { + body := sseEvent(EventStart, map[string]interface{}{"conversation_id": "id-1"}) + + sseEvent(EventToken, map[string]interface{}{"id": 0, "token": "Hello"}) + + sseEvent(EventToken, map[string]interface{}{"id": 1, "token": " world"}) + + sseEndEvent(map[string]interface{}{ + "referenced_documents": []interface{}{}, + "truncated": false, + "input_tokens": 50, + "output_tokens": 10, + "reasoning_tokens": 0, + }) + + setupClient(sseHandler(body)) + + events, errc, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).NotTo(HaveOccurred()) + + var received []SSEEvent + for ev := range events { + received = append(received, ev) + } + Expect(<-errc).NotTo(HaveOccurred()) + + Expect(received).To(HaveLen(4)) + Expect(received[0].Type).To(Equal(EventStart)) + Expect(received[1].Type).To(Equal(EventToken)) + Expect(received[2].Type).To(Equal(EventToken)) + Expect(received[3].Type).To(Equal(EventEnd)) + + var tok TokenEventData + Expect(json.Unmarshal([]byte(received[1].Data), &tok)).To(Succeed()) + Expect(tok.Token).To(Equal("Hello")) + }) + + It("sends correct request headers and body", func() { + var capturedReq *http.Request + var capturedBody []byte + + handler := func(w http.ResponseWriter, r *http.Request) { + capturedReq = r + capturedBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(sseEndEvent(map[string]interface{}{ + "referenced_documents": []interface{}{}, + "truncated": false, + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + }))) + } + + setupClient(handler) + events, errc, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).NotTo(HaveOccurred()) + + // Drain events + for range events { + } + Expect(<-errc).NotTo(HaveOccurred()) + + Expect(capturedReq.Method).To(Equal(http.MethodPost)) + Expect(capturedReq.URL.Path).To(Equal(streamingQueryPath)) + Expect(capturedReq.Header.Get("Authorization")).To(Equal("Bearer test-token")) + Expect(capturedReq.Header.Get("Content-Type")).To(Equal("application/json")) + Expect(capturedReq.Header.Get("Accept")).To(Equal("text/event-stream")) + + var reqBody LLMRequest + Expect(json.Unmarshal(capturedBody, &reqBody)).To(Succeed()) + Expect(reqBody.Query).To(Equal("why is my pod crashing")) + Expect(reqBody.Mode).To(Equal("ask")) + }) + + It("returns auth error on 401", func() { + handler := func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + } + setupClient(handler) + + _, _, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrAuthFailed)) + Expect(err.Error()).To(ContainSubstring("oc login")) + }) + + It("returns access denied on 403", func() { + handler := func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusForbidden) + } + setupClient(handler) + + _, _, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrAccessDenied)) + Expect(err.Error()).To(ContainSubstring("cluster administrator")) + }) + + It("returns prompt too long on 413", func() { + handler := func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusRequestEntityTooLarge) + } + setupClient(handler) + + _, _, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrPromptTooLong)) + }) + + It("returns generic error on other non-200 status", func() { + handler := func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + } + setupClient(handler) + + _, _, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrServiceError)) + Expect(err.Error()).To(ContainSubstring("500")) + }) + + It("rejects redirect to HTTP target", func() { + // An HTTPS-to-HTTP redirect would expose the bearer token + // over cleartext. The client must reject it. + httpTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Should never reach here + w.WriteHeader(http.StatusOK) + })) + defer httpTarget.Close() + + // Server that redirects to the HTTP target + handler := func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, httpTarget.URL+"/v1/streaming_query", http.StatusFound) + } + setupClient(handler) + + _, _, err := client.StreamQuery(ctx, baseRequest()) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrInsecureRedirect)) + }) + + It("respects context cancellation", func() { + // Server blocks until request context is done + handler := func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + // Block until the client cancels + <-r.Context().Done() + } + + setupClient(handler) + cancelCtx, cancelFn := context.WithCancel(ctx) + + events, errc, err := client.StreamQuery(cancelCtx, baseRequest()) + Expect(err).NotTo(HaveOccurred()) + + // Cancel the context — the stream should terminate + cancelFn() + + // Drain events — should close quickly + for range events { + } + + // Error channel may or may not have an error depending on timing, + // but it must close without blocking + Eventually(errc).Should(BeClosed()) + }) + }) +}) + +var _ = Describe("idleTimeoutReader", func() { + It("times out when no data arrives", func() { + // Create a reader that blocks forever + pr, _ := io.Pipe() + reader := newIdleTimeoutReader(pr, 50*time.Millisecond) + + buf := make([]byte, 64) + _, err := reader.Read(buf) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(ErrStreamTimeout)) + }) + + It("resets timeout on successful reads", func() { + // Use a pipe to control timing: send chunks within the idle window + // but make the total duration exceed the timeout. + pr, pw := io.Pipe() + timeout := 200 * time.Millisecond + reader := newIdleTimeoutReader(pr, timeout) + + // Write 3 chunks spaced 100ms apart (total 300ms > 200ms timeout). + // Each chunk resets the timer, so no timeout should occur. + go func() { + for i := 0; i < 3; i++ { + time.Sleep(100 * time.Millisecond) + _, _ = pw.Write([]byte("chunk")) + } + pw.Close() + }() + + var total int + buf := make([]byte, 64) + for { + n, err := reader.Read(buf) + total += n + if err != nil { + Expect(err).To(Equal(io.EOF)) + break + } + } + Expect(total).To(Equal(15)) // 3 x "chunk" + }) +}) diff --git a/cli/testutil_test.go b/cli/testutil_test.go index ef141d2a3..3a73caec4 100644 --- a/cli/testutil_test.go +++ b/cli/testutil_test.go @@ -2,6 +2,8 @@ package cli import ( "bytes" + "encoding/json" + "fmt" "k8s.io/cli-runtime/pkg/genericclioptions" ) @@ -16,3 +18,26 @@ func fakeStreams() (genericclioptions.IOStreams, *bytes.Buffer, *bytes.Buffer) { } return streams, out, errOut } + +// sseEvent builds a single SSE data line in the lightspeed-service format: +// +// data: {"event": "", "data": } +// +// payload is JSON-marshaled. Returns the line with trailing double newline. +func sseEvent(eventType string, payload interface{}) string { + payloadJSON, err := json.Marshal(payload) + if err != nil { + panic(fmt.Sprintf("sseEvent: failed to marshal payload: %v", err)) + } + // Build the envelope as raw JSON to avoid double-encoding the payload. + return fmt.Sprintf("data: {\"event\": %q, \"data\": %s}\n\n", eventType, string(payloadJSON)) +} + +// sseEndEvent builds an end event with optional available_quotas field. +func sseEndEvent(payload interface{}) string { + payloadJSON, err := json.Marshal(payload) + if err != nil { + panic(fmt.Sprintf("sseEndEvent: failed to marshal payload: %v", err)) + } + return fmt.Sprintf("data: {\"event\": \"end\", \"data\": %s, \"available_quotas\": {}}\n\n", string(payloadJSON)) +} diff --git a/cli/types.go b/cli/types.go new file mode 100644 index 000000000..8d64fbed1 --- /dev/null +++ b/cli/types.go @@ -0,0 +1,48 @@ +package cli + +import "encoding/json" + +// LLMRequest represents the request payload sent to the lightspeed-service +// /v1/streaming_query endpoint. +type LLMRequest struct { + Query string `json:"query"` + Mode string `json:"mode"` + ConversationID string `json:"conversation_id,omitempty"` + MediaType string `json:"media_type"` +} + +// ReferencedDocument represents a document cited in the LLM response. +type ReferencedDocument struct { + DocURL string `json:"doc_url"` + DocTitle string `json:"doc_title"` +} + +// StartEventData represents the parsed JSON payload of an SSE "start" event. +type StartEventData struct { + ConversationID string `json:"conversation_id"` +} + +// TokenEventData represents the parsed JSON payload of an SSE "token" event. +type TokenEventData struct { + ID int `json:"id"` + Token string `json:"token"` +} + +// EndEventData represents the parsed JSON payload of an SSE "end" event. +type EndEventData struct { + ReferencedDocuments []ReferencedDocument `json:"referenced_documents"` + Truncated bool `json:"truncated"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + ReasoningTokens int `json:"reasoning_tokens"` +} + +// sseEnvelope is the JSON structure sent by lightspeed-service inside +// each SSE data line: {"event": "", "data": } +// An optional top-level "available_quotas" field is ignored. +// Data uses json.RawMessage to preserve the original JSON without +// float64 conversion of numbers. +type sseEnvelope struct { + Event string `json:"event"` + Data json.RawMessage `json:"data"` +} diff --git a/cli/types_test.go b/cli/types_test.go new file mode 100644 index 000000000..a8141fc28 --- /dev/null +++ b/cli/types_test.go @@ -0,0 +1,119 @@ +package cli + +import ( + "encoding/json" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("LLMRequest", func() { + It("marshals to JSON with correct field names", func() { + req := LLMRequest{ + Query: "why is my pod crashing", + Mode: "ask", + MediaType: "application/json", + } + data, err := json.Marshal(req) + Expect(err).NotTo(HaveOccurred()) + + var m map[string]interface{} + Expect(json.Unmarshal(data, &m)).To(Succeed()) + Expect(m).To(HaveKeyWithValue("query", "why is my pod crashing")) + Expect(m).To(HaveKeyWithValue("mode", "ask")) + Expect(m).To(HaveKeyWithValue("media_type", "application/json")) + Expect(m).NotTo(HaveKey("conversation_id")) + }) + + It("includes conversation_id when set", func() { + req := LLMRequest{ + Query: "follow up", + Mode: "ask", + ConversationID: "abc-123", + MediaType: "application/json", + } + data, err := json.Marshal(req) + Expect(err).NotTo(HaveOccurred()) + + var m map[string]interface{} + Expect(json.Unmarshal(data, &m)).To(Succeed()) + Expect(m).To(HaveKeyWithValue("conversation_id", "abc-123")) + }) +}) + +var _ = Describe("StartEventData", func() { + It("unmarshals start event payload", func() { + payload := `{"conversation_id": "81dbe98a-7805-4090-af57-8410d2b08ee6"}` + + var start StartEventData + Expect(json.Unmarshal([]byte(payload), &start)).To(Succeed()) + Expect(start.ConversationID).To(Equal("81dbe98a-7805-4090-af57-8410d2b08ee6")) + }) +}) + +var _ = Describe("TokenEventData", func() { + It("unmarshals token event payload", func() { + payload := `{"id": 3, "token": " world"}` + + var tok TokenEventData + Expect(json.Unmarshal([]byte(payload), &tok)).To(Succeed()) + Expect(tok.ID).To(Equal(3)) + Expect(tok.Token).To(Equal(" world")) + }) +}) + +var _ = Describe("EndEventData", func() { + It("unmarshals end event payload matching real server format", func() { + payload := `{ + "referenced_documents": [ + {"doc_url": "https://docs.example.com/page", "doc_title": "Example Page"} + ], + "truncated": false, + "input_tokens": 2899, + "output_tokens": 191, + "reasoning_tokens": 0 + }` + + var end EndEventData + Expect(json.Unmarshal([]byte(payload), &end)).To(Succeed()) + Expect(end.ReferencedDocuments).To(HaveLen(1)) + Expect(end.ReferencedDocuments[0].DocURL).To(Equal("https://docs.example.com/page")) + Expect(end.ReferencedDocuments[0].DocTitle).To(Equal("Example Page")) + Expect(end.Truncated).To(BeFalse()) + Expect(end.InputTokens).To(Equal(2899)) + Expect(end.OutputTokens).To(Equal(191)) + Expect(end.ReasoningTokens).To(Equal(0)) + }) + + It("handles end event with empty referenced documents", func() { + payload := `{ + "referenced_documents": [], + "truncated": false, + "input_tokens": 100, + "output_tokens": 50, + "reasoning_tokens": 0 + }` + + var end EndEventData + Expect(json.Unmarshal([]byte(payload), &end)).To(Succeed()) + Expect(end.ReferencedDocuments).To(BeEmpty()) + }) +}) + +var _ = Describe("sseEnvelope", func() { + It("unmarshals a token envelope", func() { + raw := `{"event": "token", "data": {"id": 0, "token": "Hello"}}` + + var env sseEnvelope + Expect(json.Unmarshal([]byte(raw), &env)).To(Succeed()) + Expect(env.Event).To(Equal("token")) + }) + + It("unmarshals an end envelope with available_quotas", func() { + raw := `{"event": "end", "data": {"referenced_documents": [], "truncated": false, "input_tokens": 100, "output_tokens": 50, "reasoning_tokens": 0}, "available_quotas": {}}` + + var env sseEnvelope + Expect(json.Unmarshal([]byte(raw), &env)).To(Succeed()) + Expect(env.Event).To(Equal("end")) + }) +})