From 984eb01eacf64d0e9054321669ca09ad2daa5969 Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 22:57:38 +0530 Subject: [PATCH 1/8] feat(nubi): add kb and memory command groups for Knowledge Base and Cortex context --- cmd/nubi_kb.go | 334 +++++++++++++++++++++++++++++++++++++++++++++ cmd/nubi_memory.go | 119 ++++++++++++++++ 2 files changed, 453 insertions(+) create mode 100644 cmd/nubi_kb.go create mode 100644 cmd/nubi_memory.go diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go new file mode 100644 index 0000000..5688f02 --- /dev/null +++ b/cmd/nubi_kb.go @@ -0,0 +1,334 @@ +package cmd + +import ( + "fmt" + "strings" + + "github.com/nudgebee/nbctl/pkg/client" + "github.com/nudgebee/nbctl/pkg/format" + "github.com/spf13/cobra" +) + +var nubiKbCmd = &cobra.Command{ + Use: "kb", + Short: "Manage Knowledge Base vector sources for AI retrieval", +} + +type kbItem struct { + ID string `json:"id"` + Name string `json:"name"` + KBType string `json:"kb_type"` + KBSource string `json:"kb_source"` + Status string `json:"status"` + DocumentCount int `json:"document_count"` + DataSizeBytes float64 `json:"data_size_bytes"` + LastLoadedAt string `json:"last_loaded_at"` + UpdatedAt string `json:"updated_at"` +} + +var nubiKbListCmd = &cobra.Command{ + Use: "list [account-id]", + Short: "List Knowledge Base sources", + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + accountID, err := resolveAccountID(cmd) + if err != nil && len(args) > 0 { + accountID = strings.TrimSpace(args[0]) + } + if accountID == "" { + return fmt.Errorf("account-id is required, please provide it via --account-id flag or argument") + } + + req := client.NewRequest(` + query ListKB($request: ListKBRequest!) { + ai_list_kb(request: $request) { + data { + id + name + kb_type + kb_source + status + document_count + data_size_bytes + last_loaded_at + updated_at + } + errors { + message + } + } + } + `) + req.Var("request", map[string]any{ + "account_id": accountID, + }) + + var respData struct { + AiListKb struct { + Data []kbItem `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } `json:"ai_list_kb"` + } + + if err := client.Run(cmd.Context(), req, &respData); err != nil { + return err + } + + if len(respData.AiListKb.Errors) > 0 { + return fmt.Errorf("backend error: %s", respData.AiListKb.Errors[0].Message) + } + + table := format.TabularData{ + Data: respData.AiListKb.Data, + Fields: []format.TableField{ + {Header: "KB ID", Field: "ID"}, + {Header: "Name", Field: "Name"}, + {Header: "KB Type", Field: "KBType"}, + {Header: "Status", Field: "Status"}, + {Header: "Documents", Field: "DocumentCount"}, + {Header: "Last Loaded", Field: "LastLoadedAt"}, + }, + } + format.GetFormat().Print(table) + + return nil + }, +} + +var nubiKbGetCmd = &cobra.Command{ + Use: "get ", + Short: "Get details for a Knowledge Base source", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + kbID := strings.TrimSpace(args[0]) + if kbID == "" { + return fmt.Errorf("kb-id cannot be empty") + } + + accountID, err := resolveAccountID(cmd) + if err != nil { + return fmt.Errorf("account-id is required, please provide it via --account-id flag") + } + + req := client.NewRequest(` + query GetKB($request: GetKBRequest!) { + ai_get_kb(request: $request) { + data { + id + tenant_id + account_id + name + description + data_format + data_filename + data_size_bytes + status + kb_type + kb_source + integration_id + document_count + last_loaded_at + created_at + updated_at + } + errors { + message + } + } + } + `) + req.Var("request", map[string]any{ + "account_id": accountID, + "id": kbID, + }) + + var respData struct { + AiGetKb struct { + Data *struct { + ID string `json:"id"` + TenantID string `json:"tenant_id"` + AccountID string `json:"account_id"` + Name string `json:"name"` + Description string `json:"description"` + DataFormat string `json:"data_format"` + DataFilename string `json:"data_filename"` + DataSizeBytes float64 `json:"data_size_bytes"` + Status string `json:"status"` + KBType string `json:"kb_type"` + KBSource string `json:"kb_source"` + IntegrationID string `json:"integration_id"` + DocumentCount int `json:"document_count"` + LastLoadedAt string `json:"last_loaded_at"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + } `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } `json:"ai_get_kb"` + } + + if err := client.Run(cmd.Context(), req, &respData); err != nil { + return err + } + + if len(respData.AiGetKb.Errors) > 0 { + return fmt.Errorf("backend error: %s", respData.AiGetKb.Errors[0].Message) + } + + if respData.AiGetKb.Data == nil { + return fmt.Errorf("knowledge base '%s' not found", kbID) + } + + format.GetFormat().Print(*respData.AiGetKb.Data) + return nil + }, +} + +var nubiKbSyncCmd = &cobra.Command{ + Use: "sync ", + Short: "Trigger manual re-indexing / vector embedding sync for a Knowledge Base", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + kbID := strings.TrimSpace(args[0]) + if kbID == "" { + return fmt.Errorf("kb-id cannot be empty") + } + + accountID, err := resolveAccountID(cmd) + if err != nil { + return fmt.Errorf("account-id is required, please provide it via --account-id flag") + } + + req := client.NewRequest(` + mutation SyncKB($request: RetriggerKBRequest!) { + ai_sync_kb(request: $request) { + data + errors { + message + } + } + } + `) + req.Var("request", map[string]any{ + "account_id": accountID, + "id": kbID, + }) + + var respData struct { + AiSyncKb struct { + Data any `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } `json:"ai_sync_kb"` + } + + if err := client.Run(cmd.Context(), req, &respData); err != nil { + return err + } + + if len(respData.AiSyncKb.Errors) > 0 { + return fmt.Errorf("backend error: %s", respData.AiSyncKb.Errors[0].Message) + } + + format.GetFormat().Print(map[string]any{ + "status": "triggered", + "kb_id": kbID, + "message": "Knowledge Base vector re-indexing triggered successfully", + }) + return nil + }, +} + +var nubiKbEnableCmd = &cobra.Command{ + Use: "enable ", + Short: "Enable a Knowledge Base source for AI retrieval", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + return toggleKBEnabled(cmd, args[0], true) + }, +} + +var nubiKbDisableCmd = &cobra.Command{ + Use: "disable ", + Short: "Disable a Knowledge Base source from AI retrieval", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + return toggleKBEnabled(cmd, args[0], false) + }, +} + +func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { + kbID := strings.TrimSpace(rawKBID) + if kbID == "" { + return fmt.Errorf("kb-id cannot be empty") + } + + accountID, err := resolveAccountID(cmd) + if err != nil { + return fmt.Errorf("account-id is required, please provide it via --account-id flag") + } + + req := client.NewRequest(` + mutation UpdateKBEnabled($request: UpdateKBEnabledRequest!) { + ai_update_kb_enabled(request: $request) { + data + errors { + message + } + } + } + `) + req.Var("request", map[string]any{ + "account_id": accountID, + "kb_id": kbID, + "enabled": enabled, + }) + + var respData struct { + AiUpdateKbEnabled struct { + Data any `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } `json:"ai_update_kb_enabled"` + } + + if err := client.Run(cmd.Context(), req, &respData); err != nil { + return err + } + + if len(respData.AiUpdateKbEnabled.Errors) > 0 { + return fmt.Errorf("backend error: %s", respData.AiUpdateKbEnabled.Errors[0].Message) + } + + statusStr := "enabled" + if !enabled { + statusStr = "disabled" + } + + format.GetFormat().Print(map[string]any{ + "kb_id": kbID, + "status": statusStr, + "message": fmt.Sprintf("Knowledge Base successfully %s", statusStr), + }) + return nil +} + +func init() { + nubiCmd.AddCommand(nubiKbCmd) + nubiKbCmd.AddCommand(nubiKbListCmd) + nubiKbCmd.AddCommand(nubiKbGetCmd) + nubiKbCmd.AddCommand(nubiKbSyncCmd) + nubiKbCmd.AddCommand(nubiKbEnableCmd) + nubiKbCmd.AddCommand(nubiKbDisableCmd) + + nubiKbListCmd.Flags().String("account-id", "", "Account ID (overrides profile)") + nubiKbGetCmd.Flags().String("account-id", "", "Account ID (overrides profile)") + nubiKbSyncCmd.Flags().String("account-id", "", "Account ID (overrides profile)") + nubiKbEnableCmd.Flags().String("account-id", "", "Account ID (overrides profile)") + nubiKbDisableCmd.Flags().String("account-id", "", "Account ID (overrides profile)") +} diff --git a/cmd/nubi_memory.go b/cmd/nubi_memory.go new file mode 100644 index 0000000..94983f6 --- /dev/null +++ b/cmd/nubi_memory.go @@ -0,0 +1,119 @@ +package cmd + +import ( + "fmt" + "strings" + + "github.com/nudgebee/nbctl/pkg/client" + "github.com/nudgebee/nbctl/pkg/format" + "github.com/spf13/cobra" +) + +var nubiMemoryCmd = &cobra.Command{ + Use: "memory", + Aliases: []string{"bcortex", "app-context"}, + Short: "Manage AI-learned operational memory, architecture decisions, and application context", +} + +type memoryItem struct { + ID string `json:"id"` + AccountID string `json:"account_id"` + ConversationID string `json:"conversation_id"` + MessageID string `json:"message_id"` + Content string `json:"content"` + MemoryType string `json:"memory_type"` + CreatedAt string `json:"created_at"` +} + +var nubiMemoryListCmd = &cobra.Command{ + Use: "list [account-id]", + Short: "List AI operational memory items, architecture decisions, and learned patterns", + Args: cobra.MaximumNArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + accountID, err := resolveAccountID(cmd) + if err != nil && len(args) > 0 { + accountID = strings.TrimSpace(args[0]) + } + if accountID == "" { + return fmt.Errorf("account-id is required, please provide it via --account-id flag or argument") + } + + memoryType, _ := cmd.Flags().GetString("type") + queryFilter, _ := cmd.Flags().GetString("query") + limit, _ := cmd.Flags().GetInt("limit") + + memoryType = strings.TrimSpace(memoryType) + queryFilter = strings.TrimSpace(queryFilter) + + req := client.NewRequest(` + query ListAIMemory($request: ListAIMemoryRequest!) { + ai_list_memory(request: $request) { + data { + id + account_id + conversation_id + message_id + content + memory_type + created_at + } + errors { + message + } + } + } + `) + + input := map[string]any{ + "account_id": accountID, + "limit": limit, + } + if memoryType != "" { + input["memory_type"] = memoryType + } + if queryFilter != "" { + input["query"] = queryFilter + } + req.Var("request", input) + + var respData struct { + AiListMemory struct { + Data []memoryItem `json:"data"` + Errors []struct { + Message string `json:"message"` + } `json:"errors"` + } `json:"ai_list_memory"` + } + + if err := client.Run(cmd.Context(), req, &respData); err != nil { + return err + } + + if len(respData.AiListMemory.Errors) > 0 { + return fmt.Errorf("backend error: %s", respData.AiListMemory.Errors[0].Message) + } + + table := format.TabularData{ + Data: respData.AiListMemory.Data, + Fields: []format.TableField{ + {Header: "Memory ID", Field: "ID"}, + {Header: "Memory Type", Field: "MemoryType"}, + {Header: "Content", Field: "Content"}, + {Header: "Created At", Field: "CreatedAt"}, + }, + } + format.GetFormat().Print(table) + + return nil + }, +} + +func init() { + nubiCmd.AddCommand(nubiMemoryCmd) + nubiMemoryCmd.AddCommand(nubiMemoryListCmd) + + nubiMemoryListCmd.Flags().String("account-id", "", "Account ID (overrides profile)") + nubiMemoryListCmd.Flags().String("type", "", "Filter memory by type (e.g. pattern, decision, architecture)") + nubiMemoryListCmd.Flags().String("query", "", "Filter memory content by search term") + nubiMemoryListCmd.Flags().Int("limit", 50, "Maximum number of memory items to return") +} From 10469de495e666276769fd7722013d41ebc26f8a Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:02:42 +0530 Subject: [PATCH 2/8] test(nubi): add unit tests for nubi kb and nubi memory subcommands --- cmd/nubi_kb_test.go | 142 ++++++++++++++++++++++++++++++++++++++++ cmd/nubi_memory_test.go | 67 +++++++++++++++++++ 2 files changed, 209 insertions(+) create mode 100644 cmd/nubi_kb_test.go create mode 100644 cmd/nubi_memory_test.go diff --git a/cmd/nubi_kb_test.go b/cmd/nubi_kb_test.go new file mode 100644 index 0000000..2c3dea3 --- /dev/null +++ b/cmd/nubi_kb_test.go @@ -0,0 +1,142 @@ +package cmd + +import ( + "encoding/json" + "testing" + + "github.com/nudgebee/nbctl/pkg/testutil" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNubiKbCmd_List(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_list_kb": map[string]any{ + "data": []map[string]any{ + { + "id": "kb-101", + "name": "Payment Runbook", + "kb_type": "manual", + "status": "active", + "document_count": 3, + "last_loaded_at": "2026-09-07T12:00:00Z", + "data_size_bytes": 1024, + }, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "list"}) + require.NoError(t, err) + + assert.Contains(t, output, "kb-101") + assert.Contains(t, output, "Payment Runbook") + assert.Contains(t, output, "manual") + assert.Contains(t, output, "active") +} + +func TestNubiKbCmd_List_JSON(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_list_kb": map[string]any{ + "data": []map[string]any{ + { + "id": "kb-102", + "name": "K8s Troubleshooting", + "kb_type": "integration", + "status": "active", + "document_count": 15, + }, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "list", "-o", "json"}) + require.NoError(t, err) + + var result []map[string]any + err = json.Unmarshal([]byte(output), &result) + require.NoError(t, err) + + require.Len(t, result, 1) + assert.Equal(t, "kb-102", result[0]["id"]) + assert.Equal(t, "K8s Troubleshooting", result[0]["name"]) +} + +func TestNubiKbCmd_Get(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_get_kb": map[string]any{ + "data": map[string]any{ + "id": "kb-101", + "name": "Payment Runbook", + "description": "SOP for payment gateway latency", + "status": "active", + "kb_type": "manual", + "document_count": 3, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "get", "kb-101"}) + require.NoError(t, err) + + assert.Contains(t, output, "kb-101") + assert.Contains(t, output, "Payment Runbook") + assert.Contains(t, output, "SOP for payment gateway latency") +} + +func TestNubiKbCmd_Sync(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_sync_kb": map[string]any{ + "data": "ok", + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "sync", "kb-101"}) + require.NoError(t, err) + + assert.Contains(t, output, "triggered") + assert.Contains(t, output, "kb-101") +} + +func TestNubiKbCmd_Enable_Disable(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_update_kb_enabled": map[string]any{ + "data": "ok", + }, + } + + t.Run("enable", func(t *testing.T) { + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "enable", "kb-101"}) + require.NoError(t, err) + assert.Contains(t, output, "enabled") + assert.Contains(t, output, "kb-101") + }) + + t.Run("disable", func(t *testing.T) { + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "disable", "kb-101"}) + require.NoError(t, err) + assert.Contains(t, output, "disabled") + assert.Contains(t, output, "kb-101") + }) +} diff --git a/cmd/nubi_memory_test.go b/cmd/nubi_memory_test.go new file mode 100644 index 0000000..6f9e7c3 --- /dev/null +++ b/cmd/nubi_memory_test.go @@ -0,0 +1,67 @@ +package cmd + +import ( + "encoding/json" + "testing" + + "github.com/nudgebee/nbctl/pkg/testutil" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNubiMemoryCmd_List(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_list_memory": map[string]any{ + "data": []map[string]any{ + { + "id": "mem-201", + "memory_type": "investigation_result", + "content": "RabbitMQ queue backlog pattern", + "created_at": "2026-09-07T10:00:00Z", + }, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "memory", "list"}) + require.NoError(t, err) + + assert.Contains(t, output, "mem-201") + assert.Contains(t, output, "investigation_result") + assert.Contains(t, output, "RabbitMQ queue backlog pattern") +} + +func TestNubiMemoryCmd_List_JSON(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_list_memory": map[string]any{ + "data": []map[string]any{ + { + "id": "mem-202", + "memory_type": "architecture_decision", + "content": "Use PostgreSQL for transaction log persistence", + "created_at": "2026-09-07T11:00:00Z", + }, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "memory", "list", "-o", "json", "--type", "architecture_decision"}) + require.NoError(t, err) + + var result []map[string]any + err = json.Unmarshal([]byte(output), &result) + require.NoError(t, err) + + require.Len(t, result, 1) + assert.Equal(t, "mem-202", result[0]["id"]) + assert.Equal(t, "architecture_decision", result[0]["memory_type"]) +} From 2fcce2196c2c613f8c82c5d467c46b48cf30e78f Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:04:08 +0530 Subject: [PATCH 3/8] fix(nubi): prioritize positional account-id argument over profile default and wrap error context with %w --- cmd/nubi_kb.go | 19 +++++++++++-------- cmd/nubi_memory.go | 13 ++++++++----- cmd/nubi_test.go | 17 ++++++++--------- 3 files changed, 27 insertions(+), 22 deletions(-) diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go index 5688f02..c1967a3 100644 --- a/cmd/nubi_kb.go +++ b/cmd/nubi_kb.go @@ -31,12 +31,15 @@ var nubiKbListCmd = &cobra.Command{ Short: "List Knowledge Base sources", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - accountID, err := resolveAccountID(cmd) - if err != nil && len(args) > 0 { + var accountID string + if len(args) > 0 && strings.TrimSpace(args[0]) != "" { accountID = strings.TrimSpace(args[0]) - } - if accountID == "" { - return fmt.Errorf("account-id is required, please provide it via --account-id flag or argument") + } else { + var err error + accountID, err = resolveAccountID(cmd) + if err != nil { + return fmt.Errorf("resolving account ID: %w", err) + } } req := client.NewRequest(` @@ -109,7 +112,7 @@ var nubiKbGetCmd = &cobra.Command{ accountID, err := resolveAccountID(cmd) if err != nil { - return fmt.Errorf("account-id is required, please provide it via --account-id flag") + return fmt.Errorf("resolving account ID: %w", err) } req := client.NewRequest(` @@ -199,7 +202,7 @@ var nubiKbSyncCmd = &cobra.Command{ accountID, err := resolveAccountID(cmd) if err != nil { - return fmt.Errorf("account-id is required, please provide it via --account-id flag") + return fmt.Errorf("resolving account ID: %w", err) } req := client.NewRequest(` @@ -269,7 +272,7 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { accountID, err := resolveAccountID(cmd) if err != nil { - return fmt.Errorf("account-id is required, please provide it via --account-id flag") + return fmt.Errorf("resolving account ID: %w", err) } req := client.NewRequest(` diff --git a/cmd/nubi_memory.go b/cmd/nubi_memory.go index 94983f6..f5e551e 100644 --- a/cmd/nubi_memory.go +++ b/cmd/nubi_memory.go @@ -30,12 +30,15 @@ var nubiMemoryListCmd = &cobra.Command{ Short: "List AI operational memory items, architecture decisions, and learned patterns", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - accountID, err := resolveAccountID(cmd) - if err != nil && len(args) > 0 { + var accountID string + if len(args) > 0 && strings.TrimSpace(args[0]) != "" { accountID = strings.TrimSpace(args[0]) - } - if accountID == "" { - return fmt.Errorf("account-id is required, please provide it via --account-id flag or argument") + } else { + var err error + accountID, err = resolveAccountID(cmd) + if err != nil { + return fmt.Errorf("resolving account ID: %w", err) + } } memoryType, _ := cmd.Flags().GetString("type") diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index fb376f8..b1827e9 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -189,8 +189,8 @@ func TestNubiCmd_SyncQuery(t *testing.T) { "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_cost": 0.001, + "total_input_tokens": 50, "total_output_tokens": 100, }, }, @@ -252,8 +252,8 @@ func TestNubiCmd_SyncQuery_JSON(t *testing.T) { "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_cost": 0.001, + "total_input_tokens": 50, "total_output_tokens": 100, }, }, @@ -690,8 +690,8 @@ func TestNubiCmd_Stats(t *testing.T) { "ai_get_conversation_usage_metrics": map[string]interface{}{ "data": map[string]interface{}{ "conversation": map[string]interface{}{ - "total_cost": 0.005, - "total_input_tokens": 1200, + "total_cost": 0.005, + "total_input_tokens": 1200, "total_output_tokens": 350, }, }, @@ -717,8 +717,8 @@ func TestNubiCmd_Stats_JSON(t *testing.T) { "data": map[string]interface{}{ "conversation": map[string]interface{}{ "total_cost_usd": 0.0275597, - "total_input_tokens": 57270, - "total_output_tokens": 996, + "total_input_tokens": 57270, + "total_output_tokens": 996, "total_cached_input_tokens": 39857, "total_cache_hit_rate_percentage": 69.59, "model_usage": []map[string]interface{}{ @@ -1024,4 +1024,3 @@ func TestNubiCmd_SyncQuery_SessionIdDiffersFromConversationId(t *testing.T) { assert.NotEmpty(t, pollVars["sessionId"]) assert.Nil(t, pollVars["conversationId"]) } - From b9c19ee563710547e1df1113b37fbc831528a2e6 Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:05:41 +0530 Subject: [PATCH 4/8] fix(nubi): add resolveAccountIDWithPositional helper for 3-tier account ID precedence and join GraphQL backend errors --- cmd/account_id.go | 20 +++++++++++++ cmd/nubi_kb.go | 70 ++++++++++++++++++++++++--------------------- cmd/nubi_kb_test.go | 21 ++++++++++++++ cmd/nubi_memory.go | 23 +++++---------- 4 files changed, 86 insertions(+), 48 deletions(-) diff --git a/cmd/account_id.go b/cmd/account_id.go index 239d5bc..c7962eb 100644 --- a/cmd/account_id.go +++ b/cmd/account_id.go @@ -20,3 +20,23 @@ func resolveAccountID(cmd *cobra.Command) (string, error) { } return accountID, nil } + +// resolveAccountIDWithPositional returns the account-id prioritizing an explicit +// --account-id flag (if set), followed by a positional argument, and falling +// back to profile config via viper. +func resolveAccountIDWithPositional(cmd *cobra.Command, args []string) (string, error) { + if cmd.Flags().Changed("account-id") { + flagVal, _ := cmd.Flags().GetString("account-id") + if flagVal != "" { + return flagVal, nil + } + } + if len(args) > 0 && args[0] != "" { + return args[0], nil + } + accountID, err := resolveAccountID(cmd) + if err != nil { + return "", fmt.Errorf("resolving account ID: %w", err) + } + return accountID, nil +} diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go index c1967a3..3a4c633 100644 --- a/cmd/nubi_kb.go +++ b/cmd/nubi_kb.go @@ -26,20 +26,34 @@ type kbItem struct { UpdatedAt string `json:"updated_at"` } +type graphqlErrorItem struct { + Message string `json:"message"` +} + +func joinGraphQLErrors(errs []graphqlErrorItem) error { + if len(errs) == 0 { + return nil + } + var msgs []string + for _, e := range errs { + if strings.TrimSpace(e.Message) != "" { + msgs = append(msgs, strings.TrimSpace(e.Message)) + } + } + if len(msgs) == 0 { + return nil + } + return fmt.Errorf("backend error: %s", strings.Join(msgs, "; ")) +} + var nubiKbListCmd = &cobra.Command{ Use: "list [account-id]", Short: "List Knowledge Base sources", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - var accountID string - if len(args) > 0 && strings.TrimSpace(args[0]) != "" { - accountID = strings.TrimSpace(args[0]) - } else { - var err error - accountID, err = resolveAccountID(cmd) - if err != nil { - return fmt.Errorf("resolving account ID: %w", err) - } + accountID, err := resolveAccountIDWithPositional(cmd, args) + if err != nil { + return err } req := client.NewRequest(` @@ -68,10 +82,8 @@ var nubiKbListCmd = &cobra.Command{ var respData struct { AiListKb struct { - Data []kbItem `json:"data"` - Errors []struct { - Message string `json:"message"` - } `json:"errors"` + Data []kbItem `json:"data"` + Errors []graphqlErrorItem `json:"errors"` } `json:"ai_list_kb"` } @@ -79,8 +91,8 @@ var nubiKbListCmd = &cobra.Command{ return err } - if len(respData.AiListKb.Errors) > 0 { - return fmt.Errorf("backend error: %s", respData.AiListKb.Errors[0].Message) + if err := joinGraphQLErrors(respData.AiListKb.Errors); err != nil { + return err } table := format.TabularData{ @@ -167,9 +179,7 @@ var nubiKbGetCmd = &cobra.Command{ CreatedAt string `json:"created_at"` UpdatedAt string `json:"updated_at"` } `json:"data"` - Errors []struct { - Message string `json:"message"` - } `json:"errors"` + Errors []graphqlErrorItem `json:"errors"` } `json:"ai_get_kb"` } @@ -177,8 +187,8 @@ var nubiKbGetCmd = &cobra.Command{ return err } - if len(respData.AiGetKb.Errors) > 0 { - return fmt.Errorf("backend error: %s", respData.AiGetKb.Errors[0].Message) + if err := joinGraphQLErrors(respData.AiGetKb.Errors); err != nil { + return err } if respData.AiGetKb.Data == nil { @@ -222,10 +232,8 @@ var nubiKbSyncCmd = &cobra.Command{ var respData struct { AiSyncKb struct { - Data any `json:"data"` - Errors []struct { - Message string `json:"message"` - } `json:"errors"` + Data any `json:"data"` + Errors []graphqlErrorItem `json:"errors"` } `json:"ai_sync_kb"` } @@ -233,8 +241,8 @@ var nubiKbSyncCmd = &cobra.Command{ return err } - if len(respData.AiSyncKb.Errors) > 0 { - return fmt.Errorf("backend error: %s", respData.AiSyncKb.Errors[0].Message) + if err := joinGraphQLErrors(respData.AiSyncKb.Errors); err != nil { + return err } format.GetFormat().Print(map[string]any{ @@ -293,10 +301,8 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { var respData struct { AiUpdateKbEnabled struct { - Data any `json:"data"` - Errors []struct { - Message string `json:"message"` - } `json:"errors"` + Data any `json:"data"` + Errors []graphqlErrorItem `json:"errors"` } `json:"ai_update_kb_enabled"` } @@ -304,8 +310,8 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { return err } - if len(respData.AiUpdateKbEnabled.Errors) > 0 { - return fmt.Errorf("backend error: %s", respData.AiUpdateKbEnabled.Errors[0].Message) + if err := joinGraphQLErrors(respData.AiUpdateKbEnabled.Errors); err != nil { + return err } statusStr := "enabled" diff --git a/cmd/nubi_kb_test.go b/cmd/nubi_kb_test.go index 2c3dea3..a7b60be 100644 --- a/cmd/nubi_kb_test.go +++ b/cmd/nubi_kb_test.go @@ -40,6 +40,27 @@ func TestNubiKbCmd_List(t *testing.T) { assert.Contains(t, output, "active") } +func TestNubiKbCmd_List_PositionalAccountPrecedence(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "profile-account-id") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]any{ + "ai_list_kb": map[string]any{ + "data": []map[string]any{ + { + "id": "kb-override", + "name": "Overridden KB", + }, + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "kb", "list", "override-account-id"}) + require.NoError(t, err) + assert.Contains(t, output, "kb-override") +} + func TestNubiKbCmd_List_JSON(t *testing.T) { resetNubiFlags() viper.Set("account-id", "test-account-id") diff --git a/cmd/nubi_memory.go b/cmd/nubi_memory.go index f5e551e..2c2aa38 100644 --- a/cmd/nubi_memory.go +++ b/cmd/nubi_memory.go @@ -1,7 +1,6 @@ package cmd import ( - "fmt" "strings" "github.com/nudgebee/nbctl/pkg/client" @@ -30,15 +29,9 @@ var nubiMemoryListCmd = &cobra.Command{ Short: "List AI operational memory items, architecture decisions, and learned patterns", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - var accountID string - if len(args) > 0 && strings.TrimSpace(args[0]) != "" { - accountID = strings.TrimSpace(args[0]) - } else { - var err error - accountID, err = resolveAccountID(cmd) - if err != nil { - return fmt.Errorf("resolving account ID: %w", err) - } + accountID, err := resolveAccountIDWithPositional(cmd, args) + if err != nil { + return err } memoryType, _ := cmd.Flags().GetString("type") @@ -81,10 +74,8 @@ var nubiMemoryListCmd = &cobra.Command{ var respData struct { AiListMemory struct { - Data []memoryItem `json:"data"` - Errors []struct { - Message string `json:"message"` - } `json:"errors"` + Data []memoryItem `json:"data"` + Errors []graphqlErrorItem `json:"errors"` } `json:"ai_list_memory"` } @@ -92,8 +83,8 @@ var nubiMemoryListCmd = &cobra.Command{ return err } - if len(respData.AiListMemory.Errors) > 0 { - return fmt.Errorf("backend error: %s", respData.AiListMemory.Errors[0].Message) + if err := joinGraphQLErrors(respData.AiListMemory.Errors); err != nil { + return err } table := format.TabularData{ From 3a175d654310e2f74ca870dd5c186a35ee07a24d Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:07:27 +0530 Subject: [PATCH 5/8] fix(nubi): use defensive pointer structs for top-level GraphQL keys to prevent nil panics --- cmd/nubi_kb.go | 24 ++++++++++++++++++++---- cmd/nubi_memory.go | 7 ++++++- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go index 3a4c633..9384173 100644 --- a/cmd/nubi_kb.go +++ b/cmd/nubi_kb.go @@ -81,7 +81,7 @@ var nubiKbListCmd = &cobra.Command{ }) var respData struct { - AiListKb struct { + AiListKb *struct { Data []kbItem `json:"data"` Errors []graphqlErrorItem `json:"errors"` } `json:"ai_list_kb"` @@ -91,6 +91,10 @@ var nubiKbListCmd = &cobra.Command{ return err } + if respData.AiListKb == nil { + return fmt.Errorf("empty response from backend") + } + if err := joinGraphQLErrors(respData.AiListKb.Errors); err != nil { return err } @@ -160,7 +164,7 @@ var nubiKbGetCmd = &cobra.Command{ }) var respData struct { - AiGetKb struct { + AiGetKb *struct { Data *struct { ID string `json:"id"` TenantID string `json:"tenant_id"` @@ -187,6 +191,10 @@ var nubiKbGetCmd = &cobra.Command{ return err } + if respData.AiGetKb == nil { + return fmt.Errorf("empty response from backend") + } + if err := joinGraphQLErrors(respData.AiGetKb.Errors); err != nil { return err } @@ -231,7 +239,7 @@ var nubiKbSyncCmd = &cobra.Command{ }) var respData struct { - AiSyncKb struct { + AiSyncKb *struct { Data any `json:"data"` Errors []graphqlErrorItem `json:"errors"` } `json:"ai_sync_kb"` @@ -241,6 +249,10 @@ var nubiKbSyncCmd = &cobra.Command{ return err } + if respData.AiSyncKb == nil { + return fmt.Errorf("empty response from backend") + } + if err := joinGraphQLErrors(respData.AiSyncKb.Errors); err != nil { return err } @@ -300,7 +312,7 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { }) var respData struct { - AiUpdateKbEnabled struct { + AiUpdateKbEnabled *struct { Data any `json:"data"` Errors []graphqlErrorItem `json:"errors"` } `json:"ai_update_kb_enabled"` @@ -310,6 +322,10 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { return err } + if respData.AiUpdateKbEnabled == nil { + return fmt.Errorf("empty response from backend") + } + if err := joinGraphQLErrors(respData.AiUpdateKbEnabled.Errors); err != nil { return err } diff --git a/cmd/nubi_memory.go b/cmd/nubi_memory.go index 2c2aa38..771b1ac 100644 --- a/cmd/nubi_memory.go +++ b/cmd/nubi_memory.go @@ -1,6 +1,7 @@ package cmd import ( + "fmt" "strings" "github.com/nudgebee/nbctl/pkg/client" @@ -73,7 +74,7 @@ var nubiMemoryListCmd = &cobra.Command{ req.Var("request", input) var respData struct { - AiListMemory struct { + AiListMemory *struct { Data []memoryItem `json:"data"` Errors []graphqlErrorItem `json:"errors"` } `json:"ai_list_memory"` @@ -83,6 +84,10 @@ var nubiMemoryListCmd = &cobra.Command{ return err } + if respData.AiListMemory == nil { + return fmt.Errorf("empty response from backend") + } + if err := joinGraphQLErrors(respData.AiListMemory.Errors); err != nil { return err } From c00434160407d16d92da8b92cc93a4095f3afc77 Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:17:26 +0530 Subject: [PATCH 6/8] refactor(nubi): use dedicated kbActionResponse struct for formatted tabular output in text mode --- cmd/nubi_kb.go | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go index 9384173..a83af93 100644 --- a/cmd/nubi_kb.go +++ b/cmd/nubi_kb.go @@ -208,6 +208,12 @@ var nubiKbGetCmd = &cobra.Command{ }, } +type kbActionResponse struct { + KBID string `json:"kb_id"` + Status string `json:"status"` + Message string `json:"message"` +} + var nubiKbSyncCmd = &cobra.Command{ Use: "sync ", Short: "Trigger manual re-indexing / vector embedding sync for a Knowledge Base", @@ -257,10 +263,10 @@ var nubiKbSyncCmd = &cobra.Command{ return err } - format.GetFormat().Print(map[string]any{ - "status": "triggered", - "kb_id": kbID, - "message": "Knowledge Base vector re-indexing triggered successfully", + format.GetFormat().Print(kbActionResponse{ + KBID: kbID, + Status: "triggered", + Message: "Knowledge Base vector re-indexing triggered successfully", }) return nil }, @@ -335,10 +341,10 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { statusStr = "disabled" } - format.GetFormat().Print(map[string]any{ - "kb_id": kbID, - "status": statusStr, - "message": fmt.Sprintf("Knowledge Base successfully %s", statusStr), + format.GetFormat().Print(kbActionResponse{ + KBID: kbID, + Status: statusStr, + Message: fmt.Sprintf("Knowledge Base successfully %s", statusStr), }) return nil } From bd93a1d8084810448709d8349bd286d47adf7b90 Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:19:45 +0530 Subject: [PATCH 7/8] fix(nubi): validate limit flag > 0 in nubi memory list command --- cmd/nubi_memory.go | 3 +++ cmd/nubi_memory_test.go | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/cmd/nubi_memory.go b/cmd/nubi_memory.go index 771b1ac..4ad97fd 100644 --- a/cmd/nubi_memory.go +++ b/cmd/nubi_memory.go @@ -38,6 +38,9 @@ var nubiMemoryListCmd = &cobra.Command{ memoryType, _ := cmd.Flags().GetString("type") queryFilter, _ := cmd.Flags().GetString("query") limit, _ := cmd.Flags().GetInt("limit") + if limit <= 0 { + return fmt.Errorf("limit must be greater than 0") + } memoryType = strings.TrimSpace(memoryType) queryFilter = strings.TrimSpace(queryFilter) diff --git a/cmd/nubi_memory_test.go b/cmd/nubi_memory_test.go index 6f9e7c3..a0fe874 100644 --- a/cmd/nubi_memory_test.go +++ b/cmd/nubi_memory_test.go @@ -65,3 +65,13 @@ func TestNubiMemoryCmd_List_JSON(t *testing.T) { assert.Equal(t, "mem-202", result[0]["id"]) assert.Equal(t, "architecture_decision", result[0]["memory_type"]) } + +func TestNubiMemoryCmd_List_InvalidLimit(t *testing.T) { + resetNubiFlags() + viper.Set("account-id", "test-account-id") + t.Cleanup(resetNubiFlags) + + _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "memory", "list", "--limit", "0"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "limit must be greater than 0") +} From 2aa8d6ee8cd18b9f537c5e3a67b54ac34176d770 Mon Sep 17 00:00:00 2001 From: shiv Date: Mon, 7 Sep 2026 23:21:53 +0530 Subject: [PATCH 8/8] refactor(nubi): promote account-id to PersistentFlags on nubiKbCmd and handle flag errors in account_id.go --- cmd/account_id.go | 16 ++++++++++------ cmd/nubi_kb.go | 8 ++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/cmd/account_id.go b/cmd/account_id.go index c7962eb..eca27b5 100644 --- a/cmd/account_id.go +++ b/cmd/account_id.go @@ -2,6 +2,7 @@ package cmd import ( "fmt" + "strings" "github.com/spf13/cobra" "github.com/spf13/viper" @@ -11,7 +12,10 @@ import ( // prefers an explicit --account-id flag, falls back to the configured // profile via viper, and returns an error if neither is set. func resolveAccountID(cmd *cobra.Command) (string, error) { - accountID, _ := cmd.Flags().GetString("account-id") + accountID, err := cmd.Flags().GetString("account-id") + if err != nil { + accountID = "" + } if accountID == "" { accountID = viper.GetString("account-id") } @@ -26,13 +30,13 @@ func resolveAccountID(cmd *cobra.Command) (string, error) { // back to profile config via viper. func resolveAccountIDWithPositional(cmd *cobra.Command, args []string) (string, error) { if cmd.Flags().Changed("account-id") { - flagVal, _ := cmd.Flags().GetString("account-id") - if flagVal != "" { - return flagVal, nil + flagVal, err := cmd.Flags().GetString("account-id") + if err == nil && strings.TrimSpace(flagVal) != "" { + return strings.TrimSpace(flagVal), nil } } - if len(args) > 0 && args[0] != "" { - return args[0], nil + if len(args) > 0 && strings.TrimSpace(args[0]) != "" { + return strings.TrimSpace(args[0]), nil } accountID, err := resolveAccountID(cmd) if err != nil { diff --git a/cmd/nubi_kb.go b/cmd/nubi_kb.go index a83af93..a1f742b 100644 --- a/cmd/nubi_kb.go +++ b/cmd/nubi_kb.go @@ -351,15 +351,11 @@ func toggleKBEnabled(cmd *cobra.Command, rawKBID string, enabled bool) error { func init() { nubiCmd.AddCommand(nubiKbCmd) + nubiKbCmd.PersistentFlags().String("account-id", "", "Account ID (overrides profile)") + nubiKbCmd.AddCommand(nubiKbListCmd) nubiKbCmd.AddCommand(nubiKbGetCmd) nubiKbCmd.AddCommand(nubiKbSyncCmd) nubiKbCmd.AddCommand(nubiKbEnableCmd) nubiKbCmd.AddCommand(nubiKbDisableCmd) - - nubiKbListCmd.Flags().String("account-id", "", "Account ID (overrides profile)") - nubiKbGetCmd.Flags().String("account-id", "", "Account ID (overrides profile)") - nubiKbSyncCmd.Flags().String("account-id", "", "Account ID (overrides profile)") - nubiKbEnableCmd.Flags().String("account-id", "", "Account ID (overrides profile)") - nubiKbDisableCmd.Flags().String("account-id", "", "Account ID (overrides profile)") }