From d34e18eed02b9996aa1c400311fc1add5588cb84 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 12:08:18 +0530 Subject: [PATCH 01/10] feat(nubi): add single query (-q/--query) and --async modes, cleanup unused subcommands --- README.md | 19 ++++++++++--- cmd/nubi.go | 69 +++++++++++++++++++++++++++++++++++++++++++++ cmd/nubi_create.go | 22 --------------- cmd/nubi_delete.go | 22 --------------- cmd/nubi_suggest.go | 22 --------------- cmd/nubi_test.go | 25 ++++++++++++++++ 6 files changed, 109 insertions(+), 70 deletions(-) delete mode 100644 cmd/nubi_create.go delete mode 100644 cmd/nubi_delete.go delete mode 100644 cmd/nubi_suggest.go create mode 100644 cmd/nubi_test.go diff --git a/README.md b/README.md index 171b8e4..406420b 100644 --- a/README.md +++ b/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 `: 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` diff --git a/cmd/nubi.go b/cmd/nubi.go index c67473c..7f0b4f4 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -105,6 +105,68 @@ var nubiCmd = &cobra.Command{ os.Exit(0) }() + // Single query mode (non-interactive) + if nubiQuery != "" { + ctx, cancel := context.WithCancel(cmd.Context()) + s.cancel = cancel + defer cancel() + + if nubiAsync { + if err := s.nubiClient.TriggerInvestigation(ctx, nubiQuery); 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, nubiQuery) + duration := time.Since(startTime) + s.spinner.Stop() + + out := format.GetFormat().GetOutput() + + if err != nil { + if 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) + } 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(context.Background()) + if err == nil && metrics != "" { + gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + _, _ = fmt.Fprintln(out, gray.Render(metrics)) + } + + gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + _, _ = fmt.Fprintln(out, gray.Render(fmt.Sprintf("Response time: %s", duration))) + + conversationURL := fmt.Sprintf("For more details: %s/ask-nudgebee?accountId=%s&conversation_id=%s", s.nubiClient.Endpoint, s.nubiClient.AccountID, s.nubiClient.ConversationID) + _, _ = fmt.Fprintln(out, gray.Render(conversationURL)) + + return nil + } + printNubiArt() // Welcome message styling @@ -599,6 +661,13 @@ func saveHistory(file string, history []string) error { return nil } +var ( + nubiQuery string + nubiAsync bool +) + func init() { + nubiCmd.Flags().StringVarP(&nubiQuery, "query", "q", "", "Execute a single query non-interactively and exit") + nubiCmd.Flags().BoolVar(&nubiAsync, "async", false, "Trigger query asynchronously without waiting for response (use with --query)") rootCmd.AddCommand(nubiCmd) } diff --git a/cmd/nubi_create.go b/cmd/nubi_create.go deleted file mode 100644 index ad0ec55..0000000 --- a/cmd/nubi_create.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var nubiCreateCmd = &cobra.Command{ - Use: "create", - Short: "Start a new Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi create called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(nubiCreateCmd) -} diff --git a/cmd/nubi_delete.go b/cmd/nubi_delete.go deleted file mode 100644 index e7e0ce4..0000000 --- a/cmd/nubi_delete.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var nubiDeleteCmd = &cobra.Command{ - Use: "delete", - Short: "Delete a Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi delete called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(nubiDeleteCmd) -} diff --git a/cmd/nubi_suggest.go b/cmd/nubi_suggest.go deleted file mode 100644 index 177677a..0000000 --- a/cmd/nubi_suggest.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var suggestCmd = &cobra.Command{ - Use: "suggest", - Short: "Get suggestions for a Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi suggest called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(suggestCmd) -} diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go new file mode 100644 index 0000000..7e90bac --- /dev/null +++ b/cmd/nubi_test.go @@ -0,0 +1,25 @@ +package cmd + +import ( + "testing" + + "github.com/nudgebee/nbctl/pkg/testutil" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNubiCmd_AsyncQuery(t *testing.T) { + 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:") +} From 779c21eeb897958311b6dc702fb54e13ce350551 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 12:15:35 +0530 Subject: [PATCH 02/10] fix(nubi): validate --async requires --query and pass active context to GetUsageMetrics --- cmd/nubi.go | 6 +++++- cmd/nubi_test.go | 20 ++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 7f0b4f4..ec9c949 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -105,6 +105,10 @@ var nubiCmd = &cobra.Command{ os.Exit(0) }() + if nubiAsync && nubiQuery == "" { + return fmt.Errorf("--async requires --query / -q") + } + // Single query mode (non-interactive) if nubiQuery != "" { ctx, cancel := context.WithCancel(cmd.Context()) @@ -152,7 +156,7 @@ var nubiCmd = &cobra.Command{ } } - metrics, err := s.nubiClient.GetUsageMetrics(context.Background()) + metrics, err := s.nubiClient.GetUsageMetrics(ctx) if err == nil && metrics != "" { gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) _, _ = fmt.Fprintln(out, gray.Render(metrics)) diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 7e90bac..de1cb31 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -9,6 +9,13 @@ import ( ) func TestNubiCmd_AsyncQuery(t *testing.T) { + nubiQuery = "" + nubiAsync = false + defer func() { + nubiQuery = "" + nubiAsync = false + }() + mockResponse := map[string]interface{}{ "ai_execute_investigation": map[string]interface{}{ "data": map[string]interface{}{ @@ -23,3 +30,16 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { assert.Contains(t, output, "Investigation triggered asynchronously.") assert.Contains(t, output, "Session ID:") } + +func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { + nubiQuery = "" + nubiAsync = false + defer func() { + nubiQuery = "" + nubiAsync = false + }() + + _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "--async"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "--async requires --query / -q") +} From 0b3598fd92cdc0b6a7c0564f9dadd6c2a12ff81e Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 13:21:40 +0530 Subject: [PATCH 03/10] refactor(nubi): retrieve flags dynamically from command context instead of global vars --- cmd/nubi.go | 28 ++++++++++++++++------------ cmd/nubi_test.go | 16 ++++------------ 2 files changed, 20 insertions(+), 24 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index ec9c949..e954328 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -105,18 +105,27 @@ var nubiCmd = &cobra.Command{ os.Exit(0) }() - if nubiAsync && nubiQuery == "" { + query, err := cmd.Flags().GetString("query") + if err != nil { + return err + } + async, err := cmd.Flags().GetBool("async") + if err != nil { + return err + } + + if async && query == "" { return fmt.Errorf("--async requires --query / -q") } // Single query mode (non-interactive) - if nubiQuery != "" { + if query != "" { ctx, cancel := context.WithCancel(cmd.Context()) s.cancel = cancel defer cancel() - if nubiAsync { - if err := s.nubiClient.TriggerInvestigation(ctx, nubiQuery); err != nil { + if async { + if err := s.nubiClient.TriggerInvestigation(ctx, query); err != nil { return fmt.Errorf("failed to trigger investigation: %w", err) } out := format.GetFormat().GetOutput() @@ -127,7 +136,7 @@ var nubiCmd = &cobra.Command{ s.spinner.Start() startTime := time.Now() - response, status, err := s.triggerAndPoll(ctx, nubiQuery) + response, status, err := s.triggerAndPoll(ctx, query) duration := time.Since(startTime) s.spinner.Stop() @@ -665,13 +674,8 @@ func saveHistory(file string, history []string) error { return nil } -var ( - nubiQuery string - nubiAsync bool -) - func init() { - nubiCmd.Flags().StringVarP(&nubiQuery, "query", "q", "", "Execute a single query non-interactively and exit") - nubiCmd.Flags().BoolVar(&nubiAsync, "async", false, "Trigger query asynchronously without waiting for response (use with --query)") + 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) } diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index de1cb31..383533b 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -9,12 +9,8 @@ import ( ) func TestNubiCmd_AsyncQuery(t *testing.T) { - nubiQuery = "" - nubiAsync = false - defer func() { - nubiQuery = "" - nubiAsync = false - }() + _ = nubiCmd.Flags().Set("query", "") + _ = nubiCmd.Flags().Set("async", "false") mockResponse := map[string]interface{}{ "ai_execute_investigation": map[string]interface{}{ @@ -32,12 +28,8 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { } func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { - nubiQuery = "" - nubiAsync = false - defer func() { - nubiQuery = "" - nubiAsync = false - }() + _ = nubiCmd.Flags().Set("query", "") + _ = nubiCmd.Flags().Set("async", "false") _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "--async"}) require.Error(t, err) From 01cdbbe278e4b7c944b28edf3358da0a8f383c55 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 14:09:57 +0530 Subject: [PATCH 04/10] fix(nubi): add WAITING status notice in single query mode and use t.Cleanup in tests --- cmd/nubi.go | 2 ++ cmd/nubi_test.go | 12 ++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index e954328..ba6e6f6 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -154,6 +154,8 @@ var nubiCmd = &cobra.Command{ if status == "WAITING" { _, _ = fmt.Fprintln(out, response) + gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + _, _ = fmt.Fprintln(out, gray.Render("\nNote: Nubi is waiting for a followup response. To continue, please use the interactive mode or visit the URL below.")) } else { rendered, err := renderMarkdown(response) if err != nil { diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 383533b..58d80ec 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -9,8 +9,10 @@ import ( ) func TestNubiCmd_AsyncQuery(t *testing.T) { - _ = nubiCmd.Flags().Set("query", "") - _ = nubiCmd.Flags().Set("async", "false") + t.Cleanup(func() { + _ = nubiCmd.Flags().Set("query", "") + _ = nubiCmd.Flags().Set("async", "false") + }) mockResponse := map[string]interface{}{ "ai_execute_investigation": map[string]interface{}{ @@ -28,8 +30,10 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { } func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { - _ = nubiCmd.Flags().Set("query", "") - _ = nubiCmd.Flags().Set("async", "false") + t.Cleanup(func() { + _ = nubiCmd.Flags().Set("query", "") + _ = nubiCmd.Flags().Set("async", "false") + }) _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "--async"}) require.Error(t, err) From 5b662a0c902fdc367bbc69eb68f2848a15fb8205 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 14:39:24 +0530 Subject: [PATCH 05/10] fix(nubi): use errors.Is for context check, set hermetic viper username in tests, and add sync query test --- cmd/nubi.go | 3 +- cmd/nubi_test.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 1 deletion(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index ba6e6f6..81e90e5 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "os" "os/signal" @@ -143,7 +144,7 @@ var nubiCmd = &cobra.Command{ out := format.GetFormat().GetOutput() if err != nil { - if err == context.Canceled { + if errors.Is(err, context.Canceled) { _, _ = fmt.Fprintln(out, "Request canceled.") return nil } diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 58d80ec..ea19aad 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -1,17 +1,22 @@ 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 TestNubiCmd_AsyncQuery(t *testing.T) { + viper.Set("username", "test-user") t.Cleanup(func() { _ = nubiCmd.Flags().Set("query", "") _ = nubiCmd.Flags().Set("async", "false") + viper.Set("username", "") }) mockResponse := map[string]interface{}{ @@ -30,12 +35,80 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { } func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { + viper.Set("username", "test-user") t.Cleanup(func() { _ = nubiCmd.Flags().Set("query", "") _ = nubiCmd.Flags().Set("async", "false") + viper.Set("username", "") }) _, 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_SyncQuery(t *testing.T) { + viper.Set("username", "test-user") + t.Cleanup(func() { + _ = nubiCmd.Flags().Set("query", "") + _ = nubiCmd.Flags().Set("async", "false") + viper.Set("username", "") + }) + + 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:") +} From b5f3bbb341a899a6a8dec89b74c93bf123cef128 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 14:57:13 +0530 Subject: [PATCH 06/10] fix(nubi): remove redundant cancel assignment and format resume command in WAITING notice --- cmd/nubi.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 81e90e5..5f75cce 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -122,7 +122,6 @@ var nubiCmd = &cobra.Command{ // Single query mode (non-interactive) if query != "" { ctx, cancel := context.WithCancel(cmd.Context()) - s.cancel = cancel defer cancel() if async { @@ -156,7 +155,7 @@ var nubiCmd = &cobra.Command{ if status == "WAITING" { _, _ = fmt.Fprintln(out, response) gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) - _, _ = fmt.Fprintln(out, gray.Render("\nNote: Nubi is waiting for a followup response. To continue, please use the interactive mode or visit the URL below.")) + _, _ = fmt.Fprintln(out, gray.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 { From d89bb74c202d1d0b86f5fdd9708cb2e5f6a48f27 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 15:12:46 +0530 Subject: [PATCH 07/10] fix(nubi): trim query whitespace, return context error on cancellation, and trim endpoint trailing slash --- cmd/nubi.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 5f75cce..53f0b44 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -115,6 +115,8 @@ var nubiCmd = &cobra.Command{ return err } + query = strings.TrimSpace(query) + if async && query == "" { return fmt.Errorf("--async requires --query / -q") } @@ -145,7 +147,7 @@ var nubiCmd = &cobra.Command{ if err != nil { if errors.Is(err, context.Canceled) { _, _ = fmt.Fprintln(out, "Request canceled.") - return nil + return err } return fmt.Errorf("error executing query: %w", err) } @@ -176,7 +178,8 @@ var nubiCmd = &cobra.Command{ gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) _, _ = fmt.Fprintln(out, gray.Render(fmt.Sprintf("Response time: %s", duration))) - conversationURL := fmt.Sprintf("For more details: %s/ask-nudgebee?accountId=%s&conversation_id=%s", s.nubiClient.Endpoint, s.nubiClient.AccountID, s.nubiClient.ConversationID) + 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, gray.Render(conversationURL)) return nil From 80146522bb22286f68c070d34be6eacf58642b87 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 15:17:40 +0530 Subject: [PATCH 08/10] fix(nubi): use ConversationID in GetUsageMetrics and deduplicate lipgloss gray style --- cmd/nubi.go | 12 +++++------- pkg/nubi/nubi.go | 6 +++++- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 53f0b44..1f1548a 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -143,6 +143,7 @@ var nubiCmd = &cobra.Command{ s.spinner.Stop() out := format.GetFormat().GetOutput() + grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) if err != nil { if errors.Is(err, context.Canceled) { @@ -156,8 +157,7 @@ var nubiCmd = &cobra.Command{ if status == "WAITING" { _, _ = fmt.Fprintln(out, response) - gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) - _, _ = fmt.Fprintln(out, gray.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))) + _, _ = 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 { @@ -171,16 +171,14 @@ var nubiCmd = &cobra.Command{ metrics, err := s.nubiClient.GetUsageMetrics(ctx) if err == nil && metrics != "" { - gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) - _, _ = fmt.Fprintln(out, gray.Render(metrics)) + _, _ = fmt.Fprintln(out, grayStyle.Render(metrics)) } - gray := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) - _, _ = fmt.Fprintln(out, gray.Render(fmt.Sprintf("Response time: %s", duration))) + _, _ = 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, gray.Render(conversationURL)) + _, _ = fmt.Fprintln(out, grayStyle.Render(conversationURL)) return nil } diff --git a/pkg/nubi/nubi.go b/pkg/nubi/nubi.go index f5debde..d57cae8 100644 --- a/pkg/nubi/nubi.go +++ b/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 { From 78abb12f5b8da3e99c2612fa355308d890a1e6b2 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 15:32:58 +0530 Subject: [PATCH 09/10] fix(nubi): prevent interactive fallback on empty query flag and restore cursor on signal interrupt --- cmd/nubi.go | 23 +++++++++++++++++++++-- cmd/nubi_test.go | 43 ++++++++++++++++++++++++++++--------------- 2 files changed, 49 insertions(+), 17 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 1f1548a..00714c9 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -117,15 +117,34 @@ var nubiCmd = &cobra.Command{ query = strings.TrimSpace(query) - if async && query == "" { + if async && !cmd.Flags().Changed("query") { return fmt.Errorf("--async requires --query / -q") } // Single query mode (non-interactive) - if query != "" { + 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: + if s.spinner.Active() { + s.spinner.Stop() + } + cancel() + os.Exit(1) + 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) diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index ea19aad..323a02e 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -11,13 +11,22 @@ import ( "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(func() { - _ = nubiCmd.Flags().Set("query", "") - _ = nubiCmd.Flags().Set("async", "false") - viper.Set("username", "") - }) + t.Cleanup(resetNubiFlags) mockResponse := map[string]interface{}{ "ai_execute_investigation": map[string]interface{}{ @@ -35,25 +44,29 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { } func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { + resetNubiFlags() viper.Set("username", "test-user") - t.Cleanup(func() { - _ = nubiCmd.Flags().Set("query", "") - _ = nubiCmd.Flags().Set("async", "false") - viper.Set("username", "") - }) + 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(func() { - _ = nubiCmd.Flags().Set("query", "") - _ = nubiCmd.Flags().Set("async", "false") - viper.Set("username", "") - }) + t.Cleanup(resetNubiFlags) handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") From a5c3d2daf9e1c19776364f15251716d4630bb049 Mon Sep 17 00:00:00 2001 From: shiv Date: Thu, 13 Aug 2026 16:12:34 +0530 Subject: [PATCH 10/10] fix(nubi): simplify signal handler goroutine to call cancel() and return nil on context cancellation --- cmd/nubi.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/cmd/nubi.go b/cmd/nubi.go index 00714c9..dba9de9 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -135,11 +135,7 @@ var nubiCmd = &cobra.Command{ go func() { select { case <-sigChan: - if s.spinner.Active() { - s.spinner.Stop() - } cancel() - os.Exit(1) case <-ctx.Done(): } }() @@ -167,7 +163,7 @@ var nubiCmd = &cobra.Command{ if err != nil { if errors.Is(err, context.Canceled) { _, _ = fmt.Fprintln(out, "Request canceled.") - return err + return nil } return fmt.Errorf("error executing query: %w", err) }