diff --git a/cmd/auth_assign_role.go b/cmd/auth_assign_role.go index 4c8d2b9..fd882f7 100644 --- a/cmd/auth_assign_role.go +++ b/cmd/auth_assign_role.go @@ -27,17 +27,21 @@ var authAssignRoleCmd = &cobra.Command{ if accountID != "" { // Account-scoped role assignment req := client.NewRequest(` - mutation AssignAccountGroupRole($request: UserRolesUpsertAccountGroupInput!) { - userroles_upsert_account_group(request: $request) { + mutation AssignAccountGroupRole($role: auth_account_group_roles_upsert_one_input!) { + userroles_upsert_account_group(role: $role) { status message } } `) - req.Var("request", map[string]any{ - "group_id": groupID, - "role": role, - "account_id": accountID, + req.Var("role", map[string]any{ + "group_id": groupID, + "account_roles": []map[string]string{ + { + "account_id": accountID, + "role": role, + }, + }, }) var respData struct { @@ -59,14 +63,14 @@ var authAssignRoleCmd = &cobra.Command{ } else { // Tenant-level role assignment req := client.NewRequest(` - mutation AssignGroupRole($request: UserRolesUpsertGroupInput!) { - userroles_upsert_group(request: $request) { + mutation AssignGroupRole($role: auth_tenant_group_roles_upsert_one_input!) { + userroles_upsert_group(role: $role) { status message } } `) - req.Var("request", map[string]any{ + req.Var("role", map[string]any{ "group_id": groupID, "role": role, }) diff --git a/cmd/auth_groups.go b/cmd/auth_groups.go index 3a8e0e0..21d455f 100644 --- a/cmd/auth_groups.go +++ b/cmd/auth_groups.go @@ -1,6 +1,8 @@ package cmd import ( + "bytes" + "encoding/json" "strings" "github.com/nudgebee/nbctl/pkg/client" @@ -13,6 +15,47 @@ var authGroupsCmd = &cobra.Command{ Short: "Manage tenant user groups", } +type groupRoleItem struct { + Role string `json:"role"` + EntityType string `json:"entity_type"` + EntityID string `json:"entity_id"` +} + +type groupRolesField []groupRoleItem + +func (g *groupRolesField) UnmarshalJSON(data []byte) error { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 || string(trimmed) == "null" || string(trimmed) == `""` { + *g = nil + return nil + } + + var items []groupRoleItem + err := json.Unmarshal(trimmed, &items) + if err == nil { + *g = items + return nil + } + + var str string + if errStr := json.Unmarshal(trimmed, &str); errStr != nil { + return err + } + + str = strings.TrimSpace(str) + if str == "" || str == "null" { + *g = nil + return nil + } + + if errArr := json.Unmarshal([]byte(str), &items); errArr != nil { + return errArr + } + + *g = items + return nil +} + var authGroupsListCmd = &cobra.Command{ Use: "list", Short: "List user groups with assigned roles and member count", @@ -26,8 +69,8 @@ var authGroupsListCmd = &cobra.Command{ id name description - roles - user_count + group_roles + member_count created_at } } @@ -37,12 +80,12 @@ var authGroupsListCmd = &cobra.Command{ var respData struct { UsergroupsList struct { Rows []struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description"` - Roles []string `json:"roles"` - UserCount int `json:"user_count"` - CreatedAt string `json:"created_at"` + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + GroupRoles groupRolesField `json:"group_roles"` + MemberCount int `json:"member_count"` + CreatedAt string `json:"created_at"` } `json:"rows"` } `json:"usergroups_list"` } @@ -61,16 +104,23 @@ var authGroupsListCmd = &cobra.Command{ } var rows []groupRow for _, r := range respData.UsergroupsList.Rows { + var roles []string + for _, item := range r.GroupRoles { + if item.Role != "" { + roles = append(roles, item.Role) + } + } + rolesStr := "-" - if len(r.Roles) > 0 { - rolesStr = strings.Join(r.Roles, ", ") + if len(roles) > 0 { + rolesStr = strings.Join(roles, ", ") } rows = append(rows, groupRow{ ID: r.ID, Name: r.Name, Description: r.Description, Roles: rolesStr, - UserCount: r.UserCount, + UserCount: r.MemberCount, CreatedAt: r.CreatedAt, }) } @@ -103,24 +153,22 @@ var authGroupsCreateCmd = &cobra.Command{ graphqlClient := client.NewClient() req := client.NewRequest(` - mutation CreateUserGroup($request: UserGroupCreateInput!) { - usergroup_create(request: $request) { + mutation CreateUserGroup($name: String!, $description: String) { + usergroup_create(name: $name, description: $description) { id - name - status } } `) - req.Var("request", map[string]any{ - "name": groupName, - "description": desc, - }) + req.Var("name", groupName) + if desc != "" { + req.Var("description", desc) + } else { + req.Var("description", nil) + } var respData struct { UsergroupCreate struct { - ID string `json:"id"` - Name string `json:"name"` - Status string `json:"status"` + ID string `json:"id"` } `json:"usergroup_create"` } @@ -128,7 +176,15 @@ var authGroupsCreateCmd = &cobra.Command{ return err } - format.GetFormat().Print(respData.UsergroupCreate) + format.GetFormat().Print(struct { + ID string `json:"id"` + Name string `json:"name"` + Status string `json:"status"` + }{ + ID: respData.UsergroupCreate.ID, + Name: groupName, + Status: "created", + }) return nil }, } diff --git a/cmd/auth_roles.go b/cmd/auth_roles.go index aa46e10..c2e3825 100644 --- a/cmd/auth_roles.go +++ b/cmd/auth_roles.go @@ -2,6 +2,7 @@ package cmd import ( "fmt" + "strings" "github.com/nudgebee/nbctl/pkg/client" "github.com/nudgebee/nbctl/pkg/format" @@ -26,9 +27,11 @@ var authRolesListCmd = &cobra.Command{ value } customroles_list { - id - name - description + roles { + id + name + description + } } } `) @@ -38,10 +41,12 @@ var authRolesListCmd = &cobra.Command{ DisplayName string `json:"display_name"` Value string `json:"value"` } `json:"roles_list"` - CustomrolesList []struct { - ID string `json:"id"` - Name string `json:"name"` - Description string `json:"description"` + CustomrolesList struct { + Roles []struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + } `json:"roles"` } `json:"customroles_list"` } @@ -76,7 +81,7 @@ var authRolesListCmd = &cobra.Command{ Description: "-", }) } - for _, cr := range respData.CustomrolesList { + for _, cr := range respData.CustomrolesList.Roles { rows = append(rows, roleRow{ Type: "Custom", Name: cr.Name, @@ -111,28 +116,53 @@ var authRolesCreateCmd = &cobra.Command{ graphqlClient := client.NewClient() + type customRolePermissionInput struct { + Module string `json:"module"` + Class string `json:"class,omitempty"` + } + + var permInputs []customRolePermissionInput + for _, p := range permissions { + p = strings.TrimSpace(p) + if p == "" { + continue + } + parts := strings.SplitN(p, ":", 2) + if len(parts) == 2 { + module := strings.TrimSpace(parts[0]) + class := strings.TrimSpace(parts[1]) + if module == "" { + continue + } + permInputs = append(permInputs, customRolePermissionInput{ + Module: module, + Class: class, + }) + } else { + permInputs = append(permInputs, customRolePermissionInput{ + Module: p, + }) + } + } + req := client.NewRequest(` - mutation CreateCustomRole($request: CustomRoleCreateInput!) { - customroles_create(request: $request) { + mutation CreateCustomRole($name: String!, $description: String, $permissions: [CustomRolePermissionInput!]) { + customroles_create(name: $name, description: $description, permissions: $permissions) { id - name - status } } `) - - input := map[string]any{ - "name": roleName, - "description": desc, - "permissions": permissions, + req.Var("name", roleName) + if desc != "" { + req.Var("description", desc) + } else { + req.Var("description", nil) } - req.Var("request", input) + req.Var("permissions", permInputs) var respData struct { CustomrolesCreate struct { - ID string `json:"id"` - Name string `json:"name"` - Status string `json:"status"` + ID string `json:"id"` } `json:"customroles_create"` } @@ -140,7 +170,15 @@ var authRolesCreateCmd = &cobra.Command{ return err } - format.GetFormat().Print(respData.CustomrolesCreate) + format.GetFormat().Print(struct { + ID string `json:"id"` + Name string `json:"name"` + Status string `json:"status"` + }{ + ID: respData.CustomrolesCreate.ID, + Name: roleName, + Status: "created", + }) return nil }, } diff --git a/cmd/auth_test.go b/cmd/auth_test.go new file mode 100644 index 0000000..78e8b84 --- /dev/null +++ b/cmd/auth_test.go @@ -0,0 +1,248 @@ +package cmd + +import ( + "encoding/json" + "strings" + "testing" + + "github.com/nudgebee/nbctl/pkg/testutil" +) + +func TestAuthGroupsList_Unit(t *testing.T) { + mockData := map[string]any{ + "usergroups_list": map[string]any{ + "rows": []map[string]any{ + { + "id": "ug-1", + "name": "DevOps", + "description": "DevOps Team", + "group_roles": `[{"role":"tenant_admin"}]`, + "member_count": 5, + "created_at": "2026-01-01T00:00:00Z", + }, + { + "id": "ug-2", + "name": "SRE", + "description": "SRE Team", + "group_roles": []map[string]any{{"role": "account_admin"}}, + "member_count": 3, + "created_at": "2026-01-02T00:00:00Z", + }, + }, + }, + } + + got, err := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{"auth", "groups", "list"}) + if err != nil { + t.Fatalf("auth groups list failed: %v", err) + } + if got == "" { + t.Fatalf("expected output, got empty string") + } + + expected := []string{"Group Name", "Assigned Roles", "DevOps", "tenant_admin", "SRE", "account_admin"} + for _, exp := range expected { + if !strings.Contains(got, exp) { + t.Errorf("expected output to contain %q, got %q", exp, got) + } + } +} + +func TestAuthRolesList_Unit(t *testing.T) { + mockData := map[string]any{ + "roles_list": []map[string]any{ + { + "display_name": "Admin", + "value": "tenant_admin", + }, + }, + "customroles_list": map[string]any{ + "roles": []map[string]any{ + { + "id": "cr-1", + "name": "Security Auditor", + "description": "Read-only access to audit logs", + }, + }, + }, + } + + got, err := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{"auth", "roles", "list"}) + if err != nil { + t.Fatalf("auth roles list failed: %v", err) + } + if got == "" { + t.Fatalf("expected output, got empty string") + } + + expected := []string{"Built-in", "Admin", "tenant_admin", "Custom", "Security Auditor"} + for _, exp := range expected { + if !strings.Contains(got, exp) { + t.Errorf("expected output to contain %q, got %q", exp, got) + } + } +} + +func TestAuthUsersList_Unit(t *testing.T) { + mockData := map[string]any{ + "users_list_by_tenant": map[string]any{ + "rows": []map[string]any{ + { + "id": "user-1", + "username": "alice@example.com", + "display_name": "Alice Smith", + "status": "active", + "created_at": "2026-01-01T00:00:00Z", + }, + }, + }, + } + + got, err := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{"auth", "users", "list"}) + if err != nil { + t.Fatalf("auth users list failed: %v", err) + } + if got == "" { + t.Fatalf("expected output, got empty string") + } + + expected := []string{"Username", "Display Name", "alice@example.com", "Alice Smith"} + for _, exp := range expected { + if !strings.Contains(got, exp) { + t.Errorf("expected output to contain %q, got %q", exp, got) + } + } +} + +func TestAuthUsersGet_Unit(t *testing.T) { + mockData := map[string]any{ + "users_list_by_tenant": map[string]any{ + "rows": []map[string]any{ + { + "id": "user-1", + "username": "alice@example.com", + "display_name": "Alice Smith", + "status": "active", + "created_at": "2026-01-01T00:00:00Z", + }, + }, + }, + } + + got, err := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{"auth", "users", "get", "alice@example.com"}) + if err != nil { + t.Fatalf("auth users get failed: %v", err) + } + if got == "" { + t.Fatalf("expected output, got empty string") + } + + if !strings.Contains(got, "alice@example.com") { + t.Errorf("expected output to contain alice@example.com, got %q", got) + } +} + +func TestAuthRolesCreate_Unit(t *testing.T) { + mockData := map[string]any{ + "customroles_create": map[string]any{ + "id": "cr-100", + }, + } + + got, err := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{ + "auth", "roles", "create", "test-role", + "--description", "test role description", + "--permission", "events:read", + "--permission", "logs", + }) + if err != nil { + t.Fatalf("auth roles create failed: %v", err) + } + if got == "" { + t.Fatalf("expected output, got empty string") + } + + expected := []string{"cr-100", "test-role", "created"} + for _, exp := range expected { + if !strings.Contains(got, exp) { + t.Errorf("expected output to contain %q, got %q", exp, got) + } + } + + // Test permissions with empty strings and empty modules (should be skipped) + got2, err2 := testutil.RunWithSimpleGraphQL(mockData, rootCmd, []string{ + "auth", "roles", "create", "test-role-2", + "--permission", "", + "--permission", ":read", + "--permission", "audit:view", + }) + if err2 != nil { + t.Fatalf("expected empty module and empty string to be skipped without error, got: %v", err2) + } + if !strings.Contains(got2, "test-role-2") { + t.Errorf("expected output to contain test-role-2, got %q", got2) + } +} + +func TestGroupRolesField_UnmarshalJSON(t *testing.T) { + // Test direct array + var g1 groupRolesField + if err := json.Unmarshal([]byte(`[{"role":"admin","entity_type":"tenant","entity_id":"t-1"}]`), &g1); err != nil { + t.Fatalf("unexpected error for direct array: %v", err) + } + if len(g1) != 1 || g1[0].Role != "admin" { + t.Errorf("unexpected content for direct array: %+v", g1) + } + + // Test stringified array + var g2 groupRolesField + if err := json.Unmarshal([]byte(`"[{\"role\":\"viewer\"}]"`), &g2); err != nil { + t.Fatalf("unexpected error for stringified array: %v", err) + } + if len(g2) != 1 || g2[0].Role != "viewer" { + t.Errorf("unexpected content for stringified array: %+v", g2) + } + + // Test empty string + var g3 groupRolesField + if err := json.Unmarshal([]byte(`""`), &g3); err != nil { + t.Fatalf("unexpected error for empty string: %v", err) + } + if len(g3) != 0 { + t.Errorf("expected 0 items for empty string, got %d", len(g3)) + } + + // Test invalid structure returns meaningful array unmarshal error + var g4 groupRolesField + if err := json.Unmarshal([]byte(`[123]`), &g4); err == nil { + t.Fatalf("expected error for invalid array element, got nil") + } else if !strings.Contains(err.Error(), "cannot unmarshal number") { + t.Errorf("expected error about number unmarshaling, got: %v", err) + } + + // Test stringified invalid JSON returns string unmarshal error + var g5 groupRolesField + if err := json.Unmarshal([]byte(`"not a valid json array"`), &g5); err == nil { + t.Fatalf("expected error for stringified invalid JSON, got nil") + } else if !strings.Contains(err.Error(), "invalid character") { + t.Errorf("expected invalid character error, got: %v", err) + } + + // Test null + var g6 groupRolesField + if err := json.Unmarshal([]byte(` null `), &g6); err != nil { + t.Fatalf("unexpected error for null with whitespace: %v", err) + } + if len(g6) != 0 { + t.Errorf("expected 0 items for null, got %d", len(g6)) + } + + // Test string with whitespace + var g7 groupRolesField + if err := json.Unmarshal([]byte(` " " `), &g7); err != nil { + t.Fatalf("unexpected error for empty string with whitespace: %v", err) + } + if len(g7) != 0 { + t.Errorf("expected 0 items for empty string with whitespace, got %d", len(g7)) + } +} diff --git a/cmd/nubi.go b/cmd/nubi.go index 0a18549..0ee239b 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -375,7 +375,7 @@ func (s *nubiShell) handleSlashCommand(in string) { case "/bookmarks": s.handleBookmarkCommand(parts) case "/agents": - agents, err := s.nubiClient.ListAgents() + agents, err := s.nubiClient.ListAgents(context.Background()) if err != nil { fmt.Printf("Error listing agents: %v\n", err) return @@ -390,7 +390,7 @@ func (s *nubiShell) handleSlashCommand(in string) { }, }) case "/tools": - tools, err := s.nubiClient.ListTools() + tools, err := s.nubiClient.ListTools(context.Background()) if err != nil { fmt.Printf("Error listing tools: %v\n", err) return @@ -600,6 +600,14 @@ func saveHistory(file string, history []string) error { } func initNubiClient(args []string) (*nubi.NubiClient, error) { + return initNubiClientWithAccountRequirement(args, true) +} + +func initNubiClientOptionalAccount(args []string) (*nubi.NubiClient, error) { + return initNubiClientWithAccountRequirement(args, false) +} + +func initNubiClientWithAccountRequirement(args []string, requireAccount bool) (*nubi.NubiClient, error) { var accountID string if len(args) > 0 { accountID = args[0] @@ -607,7 +615,7 @@ func initNubiClient(args []string) (*nubi.NubiClient, error) { accountID = viper.GetString("account-id") } - if accountID == "" { + if requireAccount && accountID == "" { return nil, fmt.Errorf("account-id is required, please provide it as an argument or set it in your config file") } diff --git a/cmd/nubi_agents.go b/cmd/nubi_agents.go index ab40b64..778b58e 100644 --- a/cmd/nubi_agents.go +++ b/cmd/nubi_agents.go @@ -20,12 +20,12 @@ var nubiAgentsCmd = &cobra.Command{ Long: `Display registered AI agents, descriptions, status, and assigned toolsets.`, Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - nubiClient, err := initNubiClient(args) + nubiClient, err := initNubiClientOptionalAccount(args) if err != nil { return err } - agents, err := nubiClient.ListAgents() + agents, err := nubiClient.ListAgents(cmd.Context()) if err != nil { return err } diff --git a/cmd/nubi_query.go b/cmd/nubi_query.go index 0b035b6..148f0b8 100644 --- a/cmd/nubi_query.go +++ b/cmd/nubi_query.go @@ -36,9 +36,9 @@ var nubiQueryCmd = &cobra.Command{ return fmt.Errorf("query cannot be empty") } - 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") @@ -212,5 +212,6 @@ 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().String("account-id", "", "Account ID to execute the query against") nubiCmd.AddCommand(nubiQueryCmd) } diff --git a/cmd/nubi_tools.go b/cmd/nubi_tools.go index 967ae05..ae3fe0b 100644 --- a/cmd/nubi_tools.go +++ b/cmd/nubi_tools.go @@ -11,12 +11,12 @@ var nubiToolsCmd = &cobra.Command{ Long: `Display registered tools, descriptions, status, and tool types.`, Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { - nubiClient, err := initNubiClient(args) + nubiClient, err := initNubiClientOptionalAccount(args) if err != nil { return err } - tools, err := nubiClient.ListTools() + tools, err := nubiClient.ListTools(cmd.Context()) if err != nil { return err } diff --git a/pkg/nubi/nubi.go b/pkg/nubi/nubi.go index ec1b786..cd8055e 100644 --- a/pkg/nubi/nubi.go +++ b/pkg/nubi/nubi.go @@ -3,6 +3,7 @@ package nubi import ( "context" "encoding/json" + "errors" "fmt" "os" "strings" @@ -832,27 +833,105 @@ type AgentItem struct { Tools []string `json:"tools"` } -func (c *NubiClient) ListAgents() ([]AgentItem, error) { - req := client.NewRequest(` - query ListAgents($accountId: String!) { - ai_list_agents(request: {account_id: $accountId}) { +func isAccessDeniedError(err error) bool { + if err == nil { + return false + } + var gqlErrors client.GraphQLErrors + if errors.As(err, &gqlErrors) { + for _, ge := range gqlErrors { + if strings.Contains(strings.ToLower(ge.Message), "user does not have access") || + strings.Contains(strings.ToLower(ge.Message), "access-denied") || + strings.Contains(strings.ToLower(ge.Message), "unauthorized") { + return true + } + if len(ge.Extensions) > 0 { + var ext struct { + Code string `json:"code"` + } + if json.Unmarshal(ge.Extensions, &ext) == nil { + if strings.EqualFold(ext.Code, "access-denied") || strings.EqualFold(ext.Code, "unauthorized") { + return true + } + } + } + } + } + errStr := strings.ToLower(err.Error()) + return strings.Contains(errStr, "user does not have access") || + strings.Contains(errStr, "access-denied") || + strings.Contains(errStr, "unauthorized") +} + +func (c *NubiClient) listWithAccountFallback(ctx context.Context, queryName string) (json.RawMessage, error) { + switch queryName { + case "ai_list_agents", "ai_list_tools": + // Allowed query names + default: + return nil, fmt.Errorf("invalid or disallowed query name: %q", queryName) + } + + if ctx == nil { + ctx = context.Background() + } + req := client.NewRequest(fmt.Sprintf(` + query ($accountId: String!) { + %s(request: {account_id: $accountId}) { data } } - `) + `, queryName)) req.Var("accountId", c.AccountID) - var respData struct { - AiListAgents struct { - Data []AgentItem `json:"data"` - } `json:"ai_list_agents"` + var respData map[string]struct { + Data json.RawMessage `json:"data"` } - if err := c.Client.Run(context.Background(), req, &respData); err != nil { + if err := c.Client.Run(ctx, req, &respData); err != nil { + if c.AccountID != "" && isAccessDeniedError(err) { + reqFallback := client.NewRequest(fmt.Sprintf(` + query { + %s(request: {account_id: ""}) { + data + } + } + `, queryName)) + var respDataFallback map[string]struct { + Data json.RawMessage `json:"data"` + } + if errFallback := c.Client.Run(ctx, reqFallback, &respDataFallback); errFallback == nil { + res, ok := respDataFallback[queryName] + if !ok || len(res.Data) == 0 { + return nil, fmt.Errorf("fallback query %s returned no data", queryName) + } + return res.Data, nil + } else { + return nil, fmt.Errorf("%w (fallback error: %v)", err, errFallback) + } + } return nil, err } - return respData.AiListAgents.Data, nil + res, ok := respData[queryName] + if !ok || len(res.Data) == 0 { + return nil, fmt.Errorf("query %s returned no data", queryName) + } + return res.Data, nil +} + +func (c *NubiClient) ListAgents(ctx context.Context) ([]AgentItem, error) { + rawData, err := c.listWithAccountFallback(ctx, "ai_list_agents") + if err != nil { + return nil, err + } + if len(rawData) == 0 || string(rawData) == "null" { + return nil, nil + } + var agents []AgentItem + if err := json.Unmarshal(rawData, &agents); err != nil { + return nil, fmt.Errorf("failed to unmarshal agents data: %w", err) + } + return agents, nil } type ToolItem struct { @@ -862,27 +941,19 @@ type ToolItem struct { NBToolType string `json:"nb_tool_type"` } -func (c *NubiClient) ListTools() ([]ToolItem, error) { - req := client.NewRequest(` - query ListTools($accountId: String!) { - ai_list_tools(request: {account_id: $accountId}) { - data - } - } - `) - req.Var("accountId", c.AccountID) - - var respData struct { - AiListTools struct { - Data []ToolItem `json:"data"` - } `json:"ai_list_tools"` - } - - if err := c.Client.Run(context.Background(), req, &respData); err != nil { +func (c *NubiClient) ListTools(ctx context.Context) ([]ToolItem, error) { + rawData, err := c.listWithAccountFallback(ctx, "ai_list_tools") + if err != nil { return nil, err } - - return respData.AiListTools.Data, nil + if len(rawData) == 0 || string(rawData) == "null" { + return nil, nil + } + var tools []ToolItem + if err := json.Unmarshal(rawData, &tools); err != nil { + return nil, fmt.Errorf("failed to unmarshal tools data: %w", err) + } + return tools, nil } type FunctionItem struct { diff --git a/pkg/nubi/nubi_test.go b/pkg/nubi/nubi_test.go index 060a40a..5f8b22d 100644 --- a/pkg/nubi/nubi_test.go +++ b/pkg/nubi/nubi_test.go @@ -310,13 +310,78 @@ func TestNubiClient_ListAgents(t *testing.T) { c, teardown := newTestNubiClient(handler) defer teardown() - agents, err := c.ListAgents() + agents, err := c.ListAgents(context.Background()) assert.NoError(t, err) assert.Len(t, agents, 1) assert.Equal(t, "agent1", agents[0].Name) assert.Equal(t, "desc1", agents[0].Description) } +func TestNubiClient_ListAgents_Fallback(t *testing.T) { + callCount := 0 + handler := func(w http.ResponseWriter, r *http.Request) { + callCount++ + if callCount == 1 { + resp := map[string]any{ + "errors": []map[string]any{ + {"message": "User does not have access"}, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + return + } + resp := map[string]any{ + "data": map[string]any{ + "ai_list_agents": map[string]any{ + "data": json.RawMessage(`[{"name":"fallback-agent","description":"fallback-desc"}]`), + }, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + c.AccountID = "acc-restricted" + + agents, err := c.ListAgents(context.Background()) + assert.NoError(t, err) + assert.Len(t, agents, 1) + assert.Equal(t, "fallback-agent", agents[0].Name) + assert.Equal(t, 2, callCount) +} + +func TestNubiClient_ListAgents_Fallback_BothFail(t *testing.T) { + callCount := 0 + handler := func(w http.ResponseWriter, r *http.Request) { + callCount++ + if callCount == 1 { + resp := map[string]any{ + "errors": []map[string]any{ + {"message": "User does not have access"}, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + return + } + resp := map[string]any{ + "errors": []map[string]any{ + {"message": "server down"}, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + c.AccountID = "acc-restricted" + + _, err := c.ListAgents(context.Background()) + assert.Error(t, err) + assert.Contains(t, err.Error(), "fallback error") + assert.Equal(t, 2, callCount) +} + func TestNubiClient_ListTools(t *testing.T) { handler := func(w http.ResponseWriter, r *http.Request) { resp := map[string]any{ @@ -332,13 +397,93 @@ func TestNubiClient_ListTools(t *testing.T) { c, teardown := newTestNubiClient(handler) defer teardown() - tools, err := c.ListTools() + tools, err := c.ListTools(context.Background()) assert.NoError(t, err) assert.Len(t, tools, 1) assert.Equal(t, "tool1", tools[0].Name) assert.Equal(t, "desc1", tools[0].Description) } +func TestNubiClient_ListTools_Fallback(t *testing.T) { + callCount := 0 + handler := func(w http.ResponseWriter, r *http.Request) { + callCount++ + if callCount == 1 { + resp := map[string]any{ + "errors": []map[string]any{ + {"message": "access-denied to tools"}, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + return + } + resp := map[string]any{ + "data": map[string]any{ + "ai_list_tools": map[string]any{ + "data": json.RawMessage(`[{"name":"fallback-tool","description":"fallback-tool-desc"}]`), + }, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + c.AccountID = "acc-restricted" + + tools, err := c.ListTools(context.Background()) + assert.NoError(t, err) + assert.Len(t, tools, 1) + assert.Equal(t, "fallback-tool", tools[0].Name) + assert.Equal(t, 2, callCount) +} + +func TestNubiClient_ListAgents_EmptyData(t *testing.T) { + handler := func(w http.ResponseWriter, r *http.Request) { + resp := map[string]any{ + "data": map[string]any{ + "ai_list_agents": map[string]any{}, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + + agents, err := c.ListAgents(context.Background()) + assert.Error(t, err) + assert.Contains(t, err.Error(), "returned no data") + assert.Nil(t, agents) +} + +func TestNubiClient_ListTools_NullData(t *testing.T) { + handler := func(w http.ResponseWriter, r *http.Request) { + resp := map[string]any{ + "data": map[string]any{ + "ai_list_tools": map[string]any{ + "data": nil, + }, + }, + } + require.NoError(t, json.NewEncoder(w).Encode(resp)) + } + + c, teardown := newTestNubiClient(handler) + defer teardown() + + tools, err := c.ListTools(context.Background()) + assert.NoError(t, err) + assert.Nil(t, tools) +} + +func TestNubiClient_ListWithAccountFallback_DisallowedQuery(t *testing.T) { + c := &NubiClient{} + _, err := c.listWithAccountFallback(context.Background(), "malicious_query") + assert.Error(t, err) + assert.Contains(t, err.Error(), "invalid or disallowed query name") +} + func TestNubiClient_ListFunctions(t *testing.T) { handler := func(w http.ResponseWriter, r *http.Request) { resp := map[string]any{