feat(nubi): add single query (-q/--query) and --async modes

README.md

@@ -609,11 +609,14 @@ nbctl metrics query --account-id 123e4567-e89b-12d3-a456-426614174000 --query "n
 
 #### `nbctl nubi`
 
-Starts an interactive shell session with Nudgebee AI to ask questions and get insights.
+Starts an interactive shell session or executes a single query with Nudgebee AI to ask questions and get insights.
 
-*   **Usage**: `nbctl nubi [account-id]`
+*   **Usage**: `nbctl nubi [account-id] [flags]`
 *   **Arguments**:
     *   `[account-id]` (optional): The account ID to use for the Nubi session. If not provided, `nbctl` will attempt to use the `default-account-id` from your configuration.
+*   **Flags**:
+    *   `-q, --query <string>`: Execute a single query non-interactively and exit.
+    *   `--async`: Trigger the query asynchronously without waiting for the response (used with `--query`).
 *   **Prerequisites**: Your `account-id` and `username` must be configured (see `nbctl configure`).
 
 **Interactive Features (within the Nubi shell):**
@@ -633,14 +636,22 @@ The Nubi shell supports various slash commands to manage your session and intera
 *   `/functions`: Lists all available functions that Nubi can execute.
 *   `/exit`: Exits the Nubi interactive shell.
 
-**Example:**
+**Examples:**
 
+*Start interactive session:*
 ```bash
 nbctl nubi
+nbctl nubi my-dev-account-id
 ```
 
+*Single query mode (synchronous):*
 ```bash
-nbctl nubi my-dev-account-id
+nbctl nubi -q "What is the health status of my services?"
+```
+
+*Single query mode (asynchronous):*
+```bash
+nbctl nubi -q "Analyze latest deployment logs" --async
 ```
 
 #### `nbctl optimizations`

cmd/nubi.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"os"
 	"os/signal"
@@ -105,6 +106,98 @@ var nubiCmd = &cobra.Command{
 			os.Exit(0)
 		}()
 
+		query, err := cmd.Flags().GetString("query")
+		if err != nil {
+			return err
+		}
+		async, err := cmd.Flags().GetBool("async")
+		if err != nil {
+			return err
+		}
+
+		query = strings.TrimSpace(query)
+
+		if async && !cmd.Flags().Changed("query") {
+			return fmt.Errorf("--async requires --query / -q")
+		}
+
+		// Single query mode (non-interactive)
+		if cmd.Flags().Changed("query") {
+			if query == "" {
+				return fmt.Errorf("query cannot be empty")
+			}
+
+			ctx, cancel := context.WithCancel(cmd.Context())
+			defer cancel()
+
+			sigChan := make(chan os.Signal, 1)
+			signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+			go func() {
+				select {
+				case <-sigChan:
+					cancel()
+				case <-ctx.Done():
+				}
+			}()
+			defer signal.Stop(sigChan)
+
+			if async {
+				if err := s.nubiClient.TriggerInvestigation(ctx, query); err != nil {
+					return fmt.Errorf("failed to trigger investigation: %w", err)
+				}
+				out := format.GetFormat().GetOutput()
+				_, _ = fmt.Fprintln(out, "Investigation triggered asynchronously.")
+				_, _ = fmt.Fprintf(out, "Session ID: %s\n", s.nubiClient.SessionID)
+				return nil
+			}
+
+			s.spinner.Start()
+			startTime := time.Now()
+			response, status, err := s.triggerAndPoll(ctx, query)
+			duration := time.Since(startTime)
+			s.spinner.Stop()
+
+			out := format.GetFormat().GetOutput()
+			grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240"))
+
+			if err != nil {
+				if errors.Is(err, context.Canceled) {
+					_, _ = fmt.Fprintln(out, "Request canceled.")
+					return nil
+				}
+				return fmt.Errorf("error executing query: %w", err)
+			}
+
+			s.lastResponse = response
+
+			if status == "WAITING" {
+				_, _ = fmt.Fprintln(out, response)
+				_, _ = fmt.Fprintln(out, grayStyle.Render(fmt.Sprintf("\nNote: Nubi is waiting for a followup response. To continue interactively, run 'nbctl nubi' and switch to this conversation using:\n  /conversation %s\nOr visit the URL below.", s.nubiClient.ConversationID)))
+			} else {
+				rendered, err := renderMarkdown(response)
+				if err != nil {
+					_, _ = fmt.Fprintf(out, "Error rendering markdown: %v\n", err)
+					_, _ = fmt.Fprintln(out, response)
+				} else {
+					borderStyle := lipgloss.NewStyle().BorderStyle(lipgloss.RoundedBorder()).Padding(0, 1)
+					_, _ = fmt.Fprintln(out, borderStyle.Render(rendered))
+				}
+			}
+
+			metrics, err := s.nubiClient.GetUsageMetrics(ctx)
+			if err == nil && metrics != "" {
+				_, _ = fmt.Fprintln(out, grayStyle.Render(metrics))
+			}
+
+			_, _ = fmt.Fprintln(out, grayStyle.Render(fmt.Sprintf("Response time: %s", duration)))
+
+			endpoint := strings.TrimSuffix(s.nubiClient.Endpoint, "/")
+			conversationURL := fmt.Sprintf("For more details: %s/ask-nudgebee?accountId=%s&conversation_id=%s", endpoint, s.nubiClient.AccountID, s.nubiClient.ConversationID)
+			_, _ = fmt.Fprintln(out, grayStyle.Render(conversationURL))
+
+			return nil
+		}
+
 		printNubiArt()
 
 		// Welcome message styling
@@ -600,5 +693,7 @@ func saveHistory(file string, history []string) error {
 }
 
 func init() {
+	nubiCmd.Flags().StringP("query", "q", "", "Execute a single query non-interactively and exit")
+	nubiCmd.Flags().Bool("async", false, "Trigger query asynchronously without waiting for response (use with --query)")
 	rootCmd.AddCommand(nubiCmd)
 }

cmd/nubi_test.go

@@ -0,0 +1,127 @@
+package cmd
+
+import (
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/nudgebee/nbctl/pkg/testutil"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func resetNubiFlags() {
+	if f := nubiCmd.Flags().Lookup("query"); f != nil {
+		_ = f.Value.Set("")
+		f.Changed = false
+	}
+	if f := nubiCmd.Flags().Lookup("async"); f != nil {
+		_ = f.Value.Set("false")
+		f.Changed = false
+	}
+	viper.Set("username", "")
+}
+
+func TestNubiCmd_AsyncQuery(t *testing.T) {
+	resetNubiFlags()
+	viper.Set("username", "test-user")
+	t.Cleanup(resetNubiFlags)
+
+	mockResponse := map[string]interface{}{
+		"ai_execute_investigation": map[string]interface{}{
+			"data": map[string]interface{}{
+				"response": "ok",
+			},
+		},
+	}
+
+	output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "test-account-id", "-q", "hello", "--async"})
+	require.NoError(t, err)
+
+	assert.Contains(t, output, "Investigation triggered asynchronously.")
+	assert.Contains(t, output, "Session ID:")
+}
+
+func TestNubiCmd_AsyncWithoutQuery(t *testing.T) {
+	resetNubiFlags()
+	viper.Set("username", "test-user")
+	t.Cleanup(resetNubiFlags)
+
+	_, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "--async"})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "--async requires --query / -q")
+}
+
+func TestNubiCmd_EmptyQuery(t *testing.T) {
+	resetNubiFlags()
+	viper.Set("username", "test-user")
+	t.Cleanup(resetNubiFlags)
+
+	_, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "-q", "   "})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "query cannot be empty")
+}
+
+func TestNubiCmd_SyncQuery(t *testing.T) {
+	resetNubiFlags()
+	viper.Set("username", "test-user")
+	t.Cleanup(resetNubiFlags)
+
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch r.URL.Path {
+		case "/api/auth/token":
+			_ = json.NewEncoder(w).Encode(map[string]any{"token": "fake-token", "expiry": 3600})
+		case "/api/graphql":
+			resp := map[string]interface{}{
+				"data": map[string]interface{}{
+					"ai_execute_investigation": map[string]interface{}{
+						"data": map[string]interface{}{
+							"response": "started",
+						},
+					},
+					"ai_get_conversation_v3": map[string]interface{}{
+						"conversation": map[string]interface{}{
+							"id":     "conv-123",
+							"status": "COMPLETED",
+						},
+						"messages": []map[string]interface{}{
+							{
+								"id":           "msg-1",
+								"status":       "COMPLETED",
+								"response":     "System status is healthy",
+								"message_type": "generation",
+							},
+						},
+					},
+					"ai_get_conversation_usage_metrics": map[string]interface{}{
+						"data": map[string]interface{}{
+							"conversation": map[string]interface{}{
+								"total_cost":         0.001,
+								"total_input_tokens": 50,
+								"total_output_tokens": 100,
+							},
+						},
+					},
+				},
+			}
+			_ = json.NewEncoder(w).Encode(resp)
+		default:
+			http.NotFound(w, r)
+		}
+	})
+
+	defaults := map[string]any{
+		"api-key":    "dummy",
+		"username":   "dummy-user",
+		"account-id": "dummy-account",
+	}
+	output, err := testutil.RunWithMockServer(handler, defaults, nubiCmd, []string{"nubi", "test-account-id", "-q", "system status"})
+	require.NoError(t, err)
+
+	assert.Contains(t, output, "System status")
+	assert.Contains(t, output, "healthy")
+	assert.Contains(t, output, "Cost: $0.001000")
+	assert.Contains(t, output, "Response time:")
+}

pkg/nubi/nubi.go

@@ -416,8 +416,12 @@ func (c *NubiClient) GetUsageMetrics(ctx context.Context) (string, error) {
 		}
 	`)
 
+	convID := c.ConversationID
+	if convID == "" {
+		convID = c.SessionID
+	}
 	req.Var("accountId", c.AccountID)
-	req.Var("conversationId", c.SessionID)
+	req.Var("conversationId", convID)
 
 	var respData struct {
 		AiGetConversationUsageMetrics struct {