From 7736270a6ccfbc6ce1191a4bcd40be4291a1d630 Mon Sep 17 00:00:00 2001 From: shiv Date: Sun, 6 Sep 2026 09:48:01 +0530 Subject: [PATCH 1/3] feat(nubi): add query timeout flag, resilient polling, and recovery hints - Add --timeout / -t flag to nbctl nubi query (default: 0 / unlimited) - Add resilient polling in nubi query to retry transient errors and update spinner - Display Session ID, Conversation ID, nbctl nubi get command, and browser URL when query times out or is canceled - Emit structured JSON error payload when -o json / --format json is requested on trigger failure or timeout - Provide contextual diagnostic hints when account-level access is denied - Fix RunWithMockServer test isolation by ensuring NBCTL_TESTING=true is set - Add comprehensive test coverage for query timeouts, retries, and access errors --- cmd/nubi_query.go | 218 ++++++++++++++++++++++++++++++++--- cmd/nubi_test.go | 247 ++++++++++++++++++++++++++++++++++++++++ pkg/testutil/helpers.go | 3 + 3 files changed, 452 insertions(+), 16 deletions(-) diff --git a/cmd/nubi_query.go b/cmd/nubi_query.go index 148f0b8..732f030 100644 --- a/cmd/nubi_query.go +++ b/cmd/nubi_query.go @@ -21,7 +21,10 @@ import ( "github.com/spf13/viper" ) -var nubiQueryAsync bool +var ( + nubiQueryAsync bool + nubiQueryTimeout time.Duration +) var nubiQueryCmd = &cobra.Command{ Use: "query ", @@ -50,7 +53,13 @@ var nubiQueryCmd = &cobra.Command{ sessionID := uuid.New().String() nubiClient := nubi.New(client.NewClient(), accountID, username, sessionID, endpoint) - ctx, cancel := context.WithCancel(cmd.Context()) + var ctx context.Context + var cancel context.CancelFunc + if nubiQueryTimeout > 0 { + ctx, cancel = context.WithTimeout(cmd.Context(), nubiQueryTimeout) + } else { + ctx, cancel = context.WithCancel(cmd.Context()) + } defer cancel() sigChan := make(chan os.Signal, 1) @@ -67,6 +76,22 @@ var nubiQueryCmd = &cobra.Command{ async, _ := cmd.Flags().GetBool("async") if async { if err := nubiClient.TriggerInvestigation(ctx, query); err != nil { + if format.GetFormat().Get() == "json" { + jsonResp := map[string]interface{}{ + "error": fmt.Sprintf("failed to trigger investigation: %v", err), + "status": "ERROR", + "account_id": nubiClient.AccountID, + "query": query, + } + if hint := triggerErrorHint(err, nubiClient.AccountID); hint != "" { + jsonResp["hint"] = hint + } + format.GetFormat().Print(jsonResp) + return nil + } + if hint := triggerErrorHint(err, nubiClient.AccountID); hint != "" { + return fmt.Errorf("failed to trigger investigation: %w\nHint: %s", err, hint) + } return fmt.Errorf("failed to trigger investigation: %w", err) } if format.GetFormat().Get() == "json" { @@ -102,17 +127,137 @@ var nubiQueryCmd = &cobra.Command{ out := format.GetFormat().GetOutput() if err != nil { - if errors.Is(err, context.Canceled) { + isTimeout := errors.Is(err, context.DeadlineExceeded) || errors.Is(ctx.Err(), context.DeadlineExceeded) + isCanceled := errors.Is(err, context.Canceled) || errors.Is(ctx.Err(), context.Canceled) + + if isTimeout || isCanceled { + statusStr := "TIMED_OUT" + durStr := duration.Round(time.Second).String() + if duration < time.Second { + durStr = duration.Round(100 * time.Millisecond).String() + } + errMsg := fmt.Sprintf("Query timed out after %s.", durStr) + if isCanceled { + statusStr = "CANCELED" + errMsg = "Request canceled." + } + + endpointURL := strings.TrimSuffix(nubiClient.Endpoint, "/") + refID := nubiClient.ConversationID + if refID == "" { + refID = nubiClient.SessionID + } + conversationURL := fmt.Sprintf("%s/ask-nudgebee?accountId=%s&conversation_id=%s", endpointURL, nubiClient.AccountID, refID) + if format.GetFormat().Get() == "json" { - format.GetFormat().Print(map[string]interface{}{ - "error": "Request canceled.", - "status": "CANCELED", - }) + jsonResp := map[string]interface{}{ + "error": errMsg, + "status": statusStr, + "account_id": nubiClient.AccountID, + "session_id": nubiClient.SessionID, + "query": query, + "url": conversationURL, + "hint": "The investigation was triggered server-side. Retrieve results using 'nbctl nubi get' or increase timeout using '--timeout'.", + } + if nubiClient.ConversationID != "" { + jsonResp["conversation_id"] = nubiClient.ConversationID + } + format.GetFormat().Print(jsonResp) return nil } - _, _ = fmt.Fprintln(out, "Request canceled.") + + grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + boldStyle := lipgloss.NewStyle().Bold(true) + yellowStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("214")) + + _, _ = fmt.Fprintln(out, yellowStyle.Render(errMsg)) + _, _ = fmt.Fprintln(out, grayStyle.Render("\nThe investigation was triggered and may still be running or completed server-side.")) + if nubiClient.ConversationID != "" { + _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Conversation ID:"), nubiClient.ConversationID) + } + _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Session ID:"), nubiClient.SessionID) + + _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo retrieve the response once completed:")) + if nubiClient.ConversationID != "" { + _, _ = fmt.Fprintf(out, " nbctl nubi get %s\n", nubiClient.ConversationID) + } else { + _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s\n", nubiClient.SessionID) + } + + _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo view in browser:")) + _, _ = fmt.Fprintf(out, " %s\n", conversationURL) + + _, _ = fmt.Fprintln(out, grayStyle.Render("\nOptions to increase timeout or run in background:")) + _, _ = fmt.Fprintf(out, " nbctl nubi query %q --timeout 5m\n", query) + _, _ = fmt.Fprintf(out, " nbctl nubi query %q --async\n", query) + + return nil + } + + // If polling failed after triggering investigation, provide recovery information + if nubiClient.SessionID != "" && !strings.Contains(err.Error(), "triggering investigation") { + endpointURL := strings.TrimSuffix(nubiClient.Endpoint, "/") + refID := nubiClient.ConversationID + if refID == "" { + refID = nubiClient.SessionID + } + conversationURL := fmt.Sprintf("%s/ask-nudgebee?accountId=%s&conversation_id=%s", endpointURL, nubiClient.AccountID, refID) + + if format.GetFormat().Get() == "json" { + jsonResp := map[string]interface{}{ + "error": fmt.Sprintf("error executing query: %v", err), + "status": "ERROR", + "account_id": nubiClient.AccountID, + "session_id": nubiClient.SessionID, + "query": query, + "url": conversationURL, + "hint": "The investigation was triggered server-side. Retrieve results using 'nbctl nubi get'.", + } + if nubiClient.ConversationID != "" { + jsonResp["conversation_id"] = nubiClient.ConversationID + } + format.GetFormat().Print(jsonResp) + return nil + } + + grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + boldStyle := lipgloss.NewStyle().Bold(true) + redStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("196")) + + _, _ = fmt.Fprintln(out, redStyle.Render(fmt.Sprintf("Error executing query: %v", err))) + _, _ = fmt.Fprintln(out, grayStyle.Render("\nThe investigation was triggered and may still be running or completed server-side.")) + if nubiClient.ConversationID != "" { + _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Conversation ID:"), nubiClient.ConversationID) + } + _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Session ID:"), nubiClient.SessionID) + + _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo retrieve the response once completed:")) + if nubiClient.ConversationID != "" { + _, _ = fmt.Fprintf(out, " nbctl nubi get %s\n", nubiClient.ConversationID) + } else { + _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s\n", nubiClient.SessionID) + } + return nil + } + + if format.GetFormat().Get() == "json" { + jsonResp := map[string]interface{}{ + "error": fmt.Sprintf("error executing query: %v", err), + "status": "ERROR", + "account_id": nubiClient.AccountID, + "query": query, + } + if hint := triggerErrorHint(err, nubiClient.AccountID); hint != "" { + jsonResp["hint"] = hint + } + format.GetFormat().Print(jsonResp) return nil } + + if hint := triggerErrorHint(err, nubiClient.AccountID); hint != "" { + return fmt.Errorf("error executing query: %w\nHint: %s", err, hint) + } + return fmt.Errorf("error executing query: %w", err) } @@ -193,18 +338,48 @@ func (s *nubiQueryShell) triggerAndPoll(ctx context.Context, query string) (stri } func (s *nubiQueryShell) poll(ctx context.Context) (string, string, error) { + consecutiveErrors := 0 + const maxConsecutiveErrors = 5 + + check := func() (string, string, bool, error) { + resp, status, statusText, _, _, _, err := s.nubiClient.GetConversation(ctx) + if err != nil { + if errors.Is(ctx.Err(), context.Canceled) || errors.Is(ctx.Err(), context.DeadlineExceeded) { + return "", "", false, ctx.Err() + } + consecutiveErrors++ + if consecutiveErrors >= maxConsecutiveErrors { + return "", "", false, fmt.Errorf("getting conversation: %w", err) + } + return "", "", false, nil + } + consecutiveErrors = 0 + + if s.spinner != nil && statusText != "" { + s.spinner.Suffix = " " + statusText + } + + if status != "IN_PROGRESS" { + return resp, status, true, nil + } + return "", "", false, nil + } + + // Immediate check + if resp, status, done, err := check(); done || err != nil { + return resp, status, err + } + + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + for { select { case <-ctx.Done(): return "", "", ctx.Err() - case <-time.After(2 * time.Second): - resp, status, _, _, _, _, err := s.nubiClient.GetConversation(ctx) - if err != nil { - return "", "", fmt.Errorf("getting conversation: %w", err) - } - - if status != "IN_PROGRESS" { - return resp, status, nil + case <-ticker.C: + if resp, status, done, err := check(); done || err != nil { + return resp, status, err } } } @@ -212,6 +387,17 @@ func (s *nubiQueryShell) poll(ctx context.Context) (string, string, error) { func init() { nubiQueryCmd.Flags().BoolVar(&nubiQueryAsync, "async", false, "Trigger query asynchronously without waiting for response") + nubiQueryCmd.Flags().DurationVarP(&nubiQueryTimeout, "timeout", "t", 0, "Maximum time to wait for query completion (e.g. 2m, 5m). Default is 0 (no timeout)") nubiQueryCmd.Flags().String("account-id", "", "Account ID to execute the query against") nubiCmd.AddCommand(nubiQueryCmd) } + +func triggerErrorHint(err error, accountID string) string { + if err == nil { + return "" + } + if strings.Contains(strings.ToLower(err.Error()), "user does not have access") { + return fmt.Sprintf("User does not have access to account %s. Verify the account ID or assign an account role via 'nbctl auth assign-role'.", accountID) + } + return "" +} diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 670e326..425093d 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -17,6 +17,11 @@ func resetNubiFlags() { _ = f.Value.Set("false") f.Changed = false } + if f := nubiQueryCmd.Flags().Lookup("timeout"); f != nil { + _ = f.Value.Set("0s") + f.Changed = false + } + nubiQueryTimeout = 0 if f := nubiGetCmd.Flags().Lookup("session-id"); f != nil { _ = f.Value.Set("") f.Changed = false @@ -55,6 +60,81 @@ func TestNubiCmd_AsyncQuery(t *testing.T) { assert.Contains(t, output, "Session ID:") } +func TestNubiCmd_AsyncQuery_TriggerError_JSON(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": + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]any{ + "errors": []map[string]any{ + {"message": "api: user does not have access"}, + }, + }) + 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", "query", "hello", "--async", "-o", "json"}) + require.NoError(t, err) + + var result map[string]interface{} + err = json.Unmarshal([]byte(output), &result) + require.NoError(t, err) + + assert.Equal(t, "ERROR", result["status"]) + assert.Contains(t, result["error"], "user does not have access") + assert.Equal(t, "dummy-account", result["account_id"]) + assert.Contains(t, result["hint"], "User does not have access to account dummy-account") +} + +func TestNubiCmd_Query_AccessDenied_Text(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": + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(map[string]any{ + "errors": []map[string]any{ + {"message": "api: user does not have access"}, + }, + }) + 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", "query", "hello", "--async"}) + require.Error(t, err) + + assert.Contains(t, err.Error(), "user does not have access") + assert.Contains(t, err.Error(), "Hint: User does not have access to account dummy-account") + assert.Contains(t, output, "Hint: User does not have access to account dummy-account") +} + func TestNubiCmd_EmptyQuery(t *testing.T) { resetNubiFlags() viper.Set("username", "test-user") @@ -198,6 +278,173 @@ func TestNubiCmd_SyncQuery_JSON(t *testing.T) { assert.NotEmpty(t, result["url"]) } +func TestNubiCmd_Query_Timeout(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-timeout-1", + "status": "IN_PROGRESS", + }, + }, + }, + } + _ = 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", "query", "check performance", "--timeout", "100ms"}) + require.NoError(t, err) + + assert.Contains(t, output, "Query timed out after") + assert.Contains(t, output, "The investigation was triggered and may still be running or completed server-side.") + assert.Contains(t, output, "Conversation ID: conv-timeout-1") + assert.Contains(t, output, "Session ID:") + assert.Contains(t, output, "nbctl nubi get conv-timeout-1") + assert.Contains(t, output, "--timeout 5m") + assert.Contains(t, output, "--async") +} + +func TestNubiCmd_Query_Timeout_JSON(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-timeout-2", + "status": "IN_PROGRESS", + }, + }, + }, + } + _ = 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", "query", "check performance", "--timeout", "100ms", "--output", "json"}) + require.NoError(t, err) + + var result map[string]interface{} + err = json.Unmarshal([]byte(output), &result) + require.NoError(t, err) + + assert.Equal(t, "TIMED_OUT", result["status"]) + assert.Equal(t, "dummy-account", result["account_id"]) + assert.Equal(t, "conv-timeout-2", result["conversation_id"]) + assert.NotEmpty(t, result["session_id"]) + assert.Contains(t, result["url"], "conv-timeout-2") + assert.Contains(t, result["hint"], "nbctl nubi get") +} + +func TestNubiCmd_Query_TransientRetry(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + pollCount := 0 + 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": + pollCount++ + if pollCount == 2 { + // Simulate transient 500 error on first poll + w.WriteHeader(http.StatusInternalServerError) + _ = json.NewEncoder(w).Encode(map[string]any{"errors": []map[string]any{{"message": "transient error"}}}) + return + } + 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-retry-1", + "status": "COMPLETED", + }, + "messages": []map[string]interface{}{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "Recovered successfully", + "message_type": "generation", + }, + }, + }, + "ai_get_conversation_usage_metrics": map[string]interface{}{ + "data": map[string]interface{}{ + "conversation": map[string]interface{}{ + "total_cost": 0.001, + }, + }, + }, + }, + } + _ = 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", "query", "status check"}) + require.NoError(t, err) + assert.Contains(t, output, "Recovered") + assert.Contains(t, output, "successfully") +} + func TestNubiCmd_List(t *testing.T) { resetNubiFlags() viper.Set("username", "test-user") diff --git a/pkg/testutil/helpers.go b/pkg/testutil/helpers.go index 375edd0..8acb3fb 100644 --- a/pkg/testutil/helpers.go +++ b/pkg/testutil/helpers.go @@ -85,6 +85,9 @@ func RunCommandCaptureOutput(command *cobra.Command, args []string) (string, err // overrides, runs the provided cobra command, and tears down everything. // It returns the captured output and any error from running the command. func RunWithMockServer(handler http.HandlerFunc, viperOverrides map[string]any, cmd *cobra.Command, args []string) (string, error) { + _ = os.Setenv("NBCTL_TESTING", "true") + defer func() { _ = os.Unsetenv("NBCTL_TESTING") }() + // Reset the singleton client to ensure it picks up the new mock configuration client.ResetClient() From 546772fe908f6cefdd007b97a8d4543e80ce240c Mon Sep 17 00:00:00 2001 From: shiv Date: Sun, 6 Sep 2026 09:56:01 +0530 Subject: [PATCH 2/3] fix(nubi): round sub-second duration to ms and check err in poll --- cmd/nubi_query.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/cmd/nubi_query.go b/cmd/nubi_query.go index 732f030..31466a0 100644 --- a/cmd/nubi_query.go +++ b/cmd/nubi_query.go @@ -134,7 +134,10 @@ var nubiQueryCmd = &cobra.Command{ statusStr := "TIMED_OUT" durStr := duration.Round(time.Second).String() if duration < time.Second { - durStr = duration.Round(100 * time.Millisecond).String() + durStr = duration.Round(time.Millisecond).String() + if durStr == "0s" { + durStr = "<1ms" + } } errMsg := fmt.Sprintf("Query timed out after %s.", durStr) if isCanceled { @@ -344,8 +347,12 @@ func (s *nubiQueryShell) poll(ctx context.Context) (string, string, error) { check := func() (string, string, bool, error) { resp, status, statusText, _, _, _, err := s.nubiClient.GetConversation(ctx) if err != nil { - if errors.Is(ctx.Err(), context.Canceled) || errors.Is(ctx.Err(), context.DeadlineExceeded) { - return "", "", false, ctx.Err() + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || + errors.Is(ctx.Err(), context.Canceled) || errors.Is(ctx.Err(), context.DeadlineExceeded) { + if ctx.Err() != nil { + return "", "", false, ctx.Err() + } + return "", "", false, err } consecutiveErrors++ if consecutiveErrors >= maxConsecutiveErrors { From 9663519440ccdefcf7c06a4026b50e15a1c99034 Mon Sep 17 00:00:00 2001 From: shiv Date: Sun, 6 Sep 2026 20:37:10 +0530 Subject: [PATCH 3/3] feat(nubi): support --account-id on nubi get, isolate polling query parameters, and include account in recovery hints - cmd/nubi_get.go: Add --account-id flag and use resolveAccountID(cmd) to allow account overrides; return non-zero error in JSON mode when conversation is not found. - pkg/nubi/nubi.go: Send only conversationId when available, or only sessionId otherwise in GetConversation to prevent mutually exclusive AND filter failures in ai_get_conversation_v3. - cmd/nubi_query.go: Append --account-id to nubi get commands in text and JSON recovery hints. - cmd/nubi_test.go, pkg/nubi/nubi_test.go: Add unit tests for account scoping, parameter isolation, and JSON error handling. --- cmd/nubi_get.go | 14 ++- cmd/nubi_query.go | 30 ++++-- cmd/nubi_test.go | 235 +++++++++++++++++++++++++++++++++++++++++- pkg/nubi/nubi.go | 11 +- pkg/nubi/nubi_test.go | 50 +++++++++ 5 files changed, 320 insertions(+), 20 deletions(-) diff --git a/cmd/nubi_get.go b/cmd/nubi_get.go index 711e689..6b8a44e 100644 --- a/cmd/nubi_get.go +++ b/cmd/nubi_get.go @@ -28,9 +28,9 @@ var nubiGetCmd = &cobra.Command{ return fmt.Errorf("either conversation-id argument or --session-id flag must be provided") } - accountID := viper.GetString("account-id") - if accountID == "" { - return fmt.Errorf("account-id is required, please set it in your config file or pass via flag") + accountID, err := resolveAccountID(cmd) + if err != nil { + return err } username := viper.GetString("username") @@ -60,7 +60,7 @@ var nubiGetCmd = &cobra.Command{ if details != nil && details.Conversation.ID != "" { targetID = details.Conversation.ID } else { - return fmt.Errorf("conversation not found") + return fmt.Errorf("conversation not found for account %s", accountID) } } @@ -72,8 +72,11 @@ var nubiGetCmd = &cobra.Command{ return fmt.Errorf("failed to get conversation details: %w", err) } } + if details == nil || details.Conversation.ID == "" { + return fmt.Errorf("conversation not found for account %s", accountID) + } resolvedID := conversationID - if details != nil && details.Conversation.ID != "" { + if details.Conversation.ID != "" { resolvedID = details.Conversation.ID } stats, _ := nubiClient.GetConversationStats(ctx, resolvedID) @@ -173,5 +176,6 @@ var nubiGetCmd = &cobra.Command{ func init() { nubiGetCmd.Flags().String("session-id", "", "Optional session ID if conversation details lookup by session is explicitly needed") + nubiGetCmd.Flags().String("account-id", "", "Account ID to fetch the conversation from (overrides default profile)") nubiCmd.AddCommand(nubiGetCmd) } diff --git a/cmd/nubi_query.go b/cmd/nubi_query.go index 31466a0..49564ff 100644 --- a/cmd/nubi_query.go +++ b/cmd/nubi_query.go @@ -152,6 +152,11 @@ var nubiQueryCmd = &cobra.Command{ } conversationURL := fmt.Sprintf("%s/ask-nudgebee?accountId=%s&conversation_id=%s", endpointURL, nubiClient.AccountID, refID) + recoveryHint := fmt.Sprintf("The investigation was triggered server-side. Retrieve results using 'nbctl nubi get %s --account-id %s' or increase timeout using '--timeout'.", refID, nubiClient.AccountID) + if refID == nubiClient.SessionID { + recoveryHint = fmt.Sprintf("The investigation was triggered server-side. Retrieve results using 'nbctl nubi get --session-id %s --account-id %s' or increase timeout using '--timeout'.", refID, nubiClient.AccountID) + } + if format.GetFormat().Get() == "json" { jsonResp := map[string]interface{}{ "error": errMsg, @@ -160,7 +165,7 @@ var nubiQueryCmd = &cobra.Command{ "session_id": nubiClient.SessionID, "query": query, "url": conversationURL, - "hint": "The investigation was triggered server-side. Retrieve results using 'nbctl nubi get' or increase timeout using '--timeout'.", + "hint": recoveryHint, } if nubiClient.ConversationID != "" { jsonResp["conversation_id"] = nubiClient.ConversationID @@ -180,11 +185,15 @@ var nubiQueryCmd = &cobra.Command{ } _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Session ID:"), nubiClient.SessionID) + accountFlag := "" + if nubiClient.AccountID != "" { + accountFlag = fmt.Sprintf(" --account-id %s", nubiClient.AccountID) + } _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo retrieve the response once completed:")) if nubiClient.ConversationID != "" { - _, _ = fmt.Fprintf(out, " nbctl nubi get %s\n", nubiClient.ConversationID) + _, _ = fmt.Fprintf(out, " nbctl nubi get %s%s\n", nubiClient.ConversationID, accountFlag) } else { - _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s\n", nubiClient.SessionID) + _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s%s\n", nubiClient.SessionID, accountFlag) } _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo view in browser:")) @@ -206,6 +215,11 @@ var nubiQueryCmd = &cobra.Command{ } conversationURL := fmt.Sprintf("%s/ask-nudgebee?accountId=%s&conversation_id=%s", endpointURL, nubiClient.AccountID, refID) + recoveryHint := fmt.Sprintf("The investigation was triggered server-side. Retrieve results using 'nbctl nubi get %s --account-id %s'.", refID, nubiClient.AccountID) + if refID == nubiClient.SessionID { + recoveryHint = fmt.Sprintf("The investigation was triggered server-side. Retrieve results using 'nbctl nubi get --session-id %s --account-id %s'.", refID, nubiClient.AccountID) + } + if format.GetFormat().Get() == "json" { jsonResp := map[string]interface{}{ "error": fmt.Sprintf("error executing query: %v", err), @@ -214,7 +228,7 @@ var nubiQueryCmd = &cobra.Command{ "session_id": nubiClient.SessionID, "query": query, "url": conversationURL, - "hint": "The investigation was triggered server-side. Retrieve results using 'nbctl nubi get'.", + "hint": recoveryHint, } if nubiClient.ConversationID != "" { jsonResp["conversation_id"] = nubiClient.ConversationID @@ -234,11 +248,15 @@ var nubiQueryCmd = &cobra.Command{ } _, _ = fmt.Fprintf(out, " %s %s\n", boldStyle.Render("Session ID:"), nubiClient.SessionID) + accountFlag := "" + if nubiClient.AccountID != "" { + accountFlag = fmt.Sprintf(" --account-id %s", nubiClient.AccountID) + } _, _ = fmt.Fprintln(out, grayStyle.Render("\nTo retrieve the response once completed:")) if nubiClient.ConversationID != "" { - _, _ = fmt.Fprintf(out, " nbctl nubi get %s\n", nubiClient.ConversationID) + _, _ = fmt.Fprintf(out, " nbctl nubi get %s%s\n", nubiClient.ConversationID, accountFlag) } else { - _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s\n", nubiClient.SessionID) + _, _ = fmt.Fprintf(out, " nbctl nubi get --session-id %s%s\n", nubiClient.SessionID, accountFlag) } return nil } diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 425093d..fb376f8 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -21,11 +21,19 @@ func resetNubiFlags() { _ = f.Value.Set("0s") f.Changed = false } + if f := nubiQueryCmd.Flags().Lookup("account-id"); f != nil { + _ = f.Value.Set("") + f.Changed = false + } nubiQueryTimeout = 0 if f := nubiGetCmd.Flags().Lookup("session-id"); f != nil { _ = f.Value.Set("") f.Changed = false } + if f := nubiGetCmd.Flags().Lookup("account-id"); f != nil { + _ = f.Value.Set("") + f.Changed = false + } if f := rootCmd.PersistentFlags().Lookup("format"); f != nil { _ = f.Value.Set("text") f.Changed = false @@ -322,7 +330,7 @@ func TestNubiCmd_Query_Timeout(t *testing.T) { assert.Contains(t, output, "The investigation was triggered and may still be running or completed server-side.") assert.Contains(t, output, "Conversation ID: conv-timeout-1") assert.Contains(t, output, "Session ID:") - assert.Contains(t, output, "nbctl nubi get conv-timeout-1") + assert.Contains(t, output, "nbctl nubi get conv-timeout-1 --account-id dummy-account") assert.Contains(t, output, "--timeout 5m") assert.Contains(t, output, "--async") } @@ -376,7 +384,7 @@ func TestNubiCmd_Query_Timeout_JSON(t *testing.T) { assert.Equal(t, "conv-timeout-2", result["conversation_id"]) assert.NotEmpty(t, result["session_id"]) assert.Contains(t, result["url"], "conv-timeout-2") - assert.Contains(t, result["hint"], "nbctl nubi get") + assert.Contains(t, result["hint"], "nbctl nubi get conv-timeout-2 --account-id dummy-account") } func TestNubiCmd_Query_TransientRetry(t *testing.T) { @@ -794,3 +802,226 @@ func TestNubiCmd_Get_WithSessionIDFlag(t *testing.T) { assert.Contains(t, err.Error(), "conversation not found") }) } + +func TestNubiCmd_Get_WithAccountId(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + viper.Set("account-id", "default-account") + t.Cleanup(resetNubiFlags) + + var receivedAccountID string + var receivedConvID string + + 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": + var payload struct { + Query string `json:"query"` + Variables map[string]interface{} `json:"variables"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err == nil { + if acc, ok := payload.Variables["accountId"].(string); ok { + receivedAccountID = acc + } + if conv, ok := payload.Variables["conversationId"].(string); ok { + receivedConvID = conv + } + } + resp := map[string]interface{}{ + "data": map[string]interface{}{ + "ai_get_conversation_v3": map[string]interface{}{ + "conversation": map[string]interface{}{ + "id": "conv-explicit-1", + "status": "COMPLETED", + }, + "messages": []map[string]interface{}{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "Scoped to explicit account", + "message_type": "generation", + }, + }, + }, + }, + } + _ = json.NewEncoder(w).Encode(resp) + default: + http.NotFound(w, r) + } + }) + + defaults := map[string]any{ + "api-key": "dummy", + "username": "dummy-user", + "account-id": "default-account", + } + + output, err := testutil.RunWithMockServer(handler, defaults, nubiCmd, []string{ + "nubi", "get", "conv-explicit-1", "--account-id", "explicit-account-id", + }) + require.NoError(t, err) + assert.Equal(t, "explicit-account-id", receivedAccountID) + assert.Equal(t, "conv-explicit-1", receivedConvID) + assert.Contains(t, output, "Scoped") + assert.Contains(t, output, "explicit") +} + +func TestNubiCmd_Get_SessionId_WithAccountId(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + viper.Set("account-id", "default-account") + t.Cleanup(resetNubiFlags) + + var receivedAccountID string + var receivedSessionID string + + 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": + var payload struct { + Query string `json:"query"` + Variables map[string]interface{} `json:"variables"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err == nil { + if acc, ok := payload.Variables["accountId"].(string); ok { + receivedAccountID = acc + } + if sess, ok := payload.Variables["sessionId"].(string); ok { + receivedSessionID = sess + } + } + resp := map[string]interface{}{ + "data": map[string]interface{}{ + "ai_get_conversation_v3": map[string]interface{}{ + "conversation": map[string]interface{}{ + "id": "conv-session-1", + "status": "COMPLETED", + }, + "messages": []map[string]interface{}{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "Found by session ID", + "message_type": "generation", + }, + }, + }, + }, + } + _ = json.NewEncoder(w).Encode(resp) + default: + http.NotFound(w, r) + } + }) + + defaults := map[string]any{ + "api-key": "dummy", + "username": "dummy-user", + "account-id": "default-account", + } + + output, err := testutil.RunWithMockServer(handler, defaults, nubiCmd, []string{ + "nubi", "get", "--session-id", "sess-explicit-1", "--account-id", "explicit-account-id", + }) + require.NoError(t, err) + assert.Equal(t, "explicit-account-id", receivedAccountID) + assert.Equal(t, "sess-explicit-1", receivedSessionID) + assert.Contains(t, output, "Found") + assert.Contains(t, output, "session") +} + +func TestNubiCmd_Get_NotFound_JSON(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + emptyResponse := map[string]interface{}{ + "ai_get_conversation_v3": map[string]interface{}{ + "conversation": map[string]interface{}{"id": "", "status": ""}, + }, + } + _, err := testutil.RunWithSimpleGraphQL(emptyResponse, nubiCmd, []string{"nubi", "get", "conv-missing", "-o", "json"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "conversation not found for account dummy-account") +} + +func TestNubiCmd_SyncQuery_SessionIdDiffersFromConversationId(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + var recordedVars []map[string]interface{} + + 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": + var payload struct { + Query string `json:"query"` + Variables map[string]interface{} `json:"variables"` + } + _ = json.NewDecoder(r.Body).Decode(&payload) + recordedVars = append(recordedVars, payload.Variables) + + 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-different-uuid", + "status": "COMPLETED", + }, + "messages": []map[string]interface{}{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "Polled successfully with distinct session and conversation IDs", + "message_type": "generation", + }, + }, + }, + "ai_get_conversation_usage_metrics": map[string]interface{}{ + "data": map[string]interface{}{ + "conversation": map[string]interface{}{ + "total_cost": 0.001, + }, + }, + }, + }, + } + _ = 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", "query", "check distinct ids"}) + require.NoError(t, err) + assert.Contains(t, output, "Polled") + assert.Contains(t, output, "successfully") + + require.True(t, len(recordedVars) >= 2) + pollVars := recordedVars[1] + assert.NotEmpty(t, pollVars["sessionId"]) + assert.Nil(t, pollVars["conversationId"]) +} + diff --git a/pkg/nubi/nubi.go b/pkg/nubi/nubi.go index cd8055e..78afe7c 100644 --- a/pkg/nubi/nubi.go +++ b/pkg/nubi/nubi.go @@ -223,13 +223,10 @@ func (c *NubiClient) GetConversation(ctx context.Context) (string, string, strin `) req.Var("accountId", c.AccountID) - idToUse := c.ConversationID - if idToUse == "" { - idToUse = c.SessionID - } - if idToUse != "" { - req.Var("conversationId", idToUse) - req.Var("sessionId", idToUse) + if c.ConversationID != "" { + req.Var("conversationId", c.ConversationID) + } else if c.SessionID != "" { + req.Var("sessionId", c.SessionID) } var respData struct { diff --git a/pkg/nubi/nubi_test.go b/pkg/nubi/nubi_test.go index 5f8b22d..2374635 100644 --- a/pkg/nubi/nubi_test.go +++ b/pkg/nubi/nubi_test.go @@ -164,6 +164,56 @@ func TestNubiClient_GetConversation(t *testing.T) { assert.Equal(t, "COMPLETED", status) } +func TestNubiClient_GetConversation_ParameterIsolation(t *testing.T) { + var capturedVars map[string]any + + handler := func(w http.ResponseWriter, r *http.Request) { + var reqBody struct { + Variables map[string]any `json:"variables"` + } + _ = json.NewDecoder(r.Body).Decode(&reqBody) + capturedVars = reqBody.Variables + + resp := map[string]any{ + "data": map[string]any{ + "ai_get_conversation_v3": map[string]any{ + "conversation": map[string]any{ + "id": "conv-resolved-123", + "status": "COMPLETED", + }, + "messages": []map[string]any{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "Done", + "message_type": "generation", + }, + }, + }, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + + // 1. When ConversationID is empty, only sessionId should be sent + c.ConversationID = "" + c.SessionID = "sess-uuid-456" + _, _, _, _, _, _, err := c.GetConversation(context.Background()) + require.NoError(t, err) + assert.Equal(t, "sess-uuid-456", capturedVars["sessionId"]) + assert.Nil(t, capturedVars["conversationId"]) + assert.Equal(t, "conv-resolved-123", c.ConversationID) + + // 2. Subsequent call now has ConversationID populated, only conversationId should be sent + _, _, _, _, _, _, err = c.GetConversation(context.Background()) + require.NoError(t, err) + assert.Equal(t, "conv-resolved-123", capturedVars["conversationId"]) + assert.Nil(t, capturedVars["sessionId"]) +} + func TestNubiClient_SendFollowupResponse(t *testing.T) { handler := func(w http.ResponseWriter, r *http.Request) { resp := map[string]any{