From 1e1d8ef97a8e021140dae7c461934dffe7c27c62 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 4 Nov 2026 11:16:30 +0100 Subject: [PATCH 1/3] feat(apps): add model_service app resource and picker Maps the model_service manifest resource type to uc_securable with securable_type MODEL_SERVICE and default permission EXECUTE, and adds a catalog -> schema -> model service picker backed by AiGateway.ListModelServices. Co-authored-by: Isaac --- libs/apps/generator/generator.go | 6 +++++ libs/apps/generator/generator_test.go | 12 ++++++++++ libs/apps/prompt/listers.go | 21 ++++++++++++++++++ libs/apps/prompt/listers_test.go | 32 +++++++++++++++++++++++++++ libs/apps/prompt/prefetch.go | 9 ++++---- libs/apps/prompt/prompt.go | 5 +++++ libs/apps/prompt/resource_registry.go | 3 +++ 7 files changed, 84 insertions(+), 4 deletions(-) create mode 100644 libs/apps/prompt/listers_test.go diff --git a/libs/apps/generator/generator.go b/libs/apps/generator/generator.go index 3f99d80f900..3d38bed46b7 100644 --- a/libs/apps/generator/generator.go +++ b/libs/apps/generator/generator.go @@ -410,6 +410,12 @@ var appResourceSpecs = map[string]appResourceSpec{ staticFields: [][2]string{{"securable_type", "FUNCTION"}}, permission: "EXECUTE", }, + "model_service": { + yamlKey: "uc_securable", + varFields: [][2]string{{"name", "securable_full_name"}}, + staticFields: [][2]string{{"securable_type", "MODEL_SERVICE"}}, + permission: "EXECUTE", + }, "uc_connection": { yamlKey: "uc_securable", varFields: [][2]string{{"id", "securable_full_name"}}, diff --git a/libs/apps/generator/generator_test.go b/libs/apps/generator/generator_test.go index d7a1de85cc8..a9c63e3e273 100644 --- a/libs/apps/generator/generator_test.go +++ b/libs/apps/generator/generator_test.go @@ -246,6 +246,7 @@ func TestGenerateBundleResourcesDefaultPermissions(t *testing.T) { {"database", "CAN_CONNECT_AND_CREATE"}, {"volume", "READ_VOLUME"}, {"uc_function", "EXECUTE"}, + {"model_service", "EXECUTE"}, {"uc_connection", "USE_CONNECTION"}, {"genie_space", "CAN_VIEW"}, {"vector_search_index", "SELECT"}, @@ -466,6 +467,17 @@ func TestGenerateResourceYAMLAllTypes(t *testing.T) { "permission: EXECUTE", }, }, + { + name: "model_service maps to uc_securable MODEL_SERVICE", + resource: manifest.Resource{Type: "model_service", Alias: "Model Service", ResourceKey: "model", Permission: "EXECUTE"}, + expectContains: []string{ + "- name: model", + "uc_securable:", + "securable_full_name: ${var.model_name}", + "securable_type: MODEL_SERVICE", + "permission: EXECUTE", + }, + }, { name: "uc_connection maps to uc_securable CONNECTION", resource: manifest.Resource{Type: "uc_connection", Alias: "UC Connection", ResourceKey: "conn", Permission: "USE_CONNECTION"}, diff --git a/libs/apps/prompt/listers.go b/libs/apps/prompt/listers.go index 4bc2b1e301e..834a0fa7abc 100644 --- a/libs/apps/prompt/listers.go +++ b/libs/apps/prompt/listers.go @@ -155,6 +155,27 @@ func ListFunctionsInSchema(ctx context.Context, catalogName, schemaName string) return capResults(out), nil } +// ListModelServicesInSchema returns AI Gateway model services within a catalog.schema as selectable items. +// IDs are the full catalog.schema.name, without the API's model-services/ prefix. +func ListModelServicesInSchema(ctx context.Context, catalogName, schemaName string) ([]ListItem, error) { + w, err := workspaceClient(ctx) + if err != nil { + return nil, err + } + svcs, err := w.AiGateway.ListModelServicesAll(ctx, catalog.ListModelServicesRequest{ + Parent: "schemas/" + catalogName + "." + schemaName, + }) + if err != nil { + return nil, err + } + out := make([]ListItem, 0, min(len(svcs), maxListResults)) + for _, s := range svcs { + fullName := strings.TrimPrefix(s.Name, "model-services/") + out = append(out, ListItem{ID: fullName, Label: fullName}) + } + return capResults(out), nil +} + // listDatabasesResponse is the response from the /databases endpoint. type listDatabasesResponse struct { Databases []struct { diff --git a/libs/apps/prompt/listers_test.go b/libs/apps/prompt/listers_test.go new file mode 100644 index 00000000000..71049848dcc --- /dev/null +++ b/libs/apps/prompt/listers_test.go @@ -0,0 +1,32 @@ +package prompt_test + +import ( + "testing" + + "github.com/databricks/cli/libs/apps/prompt" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/catalog" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestListModelServicesInSchema(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/main.default"}). + Return([]catalog.ModelService{ + {Name: "model-services/main.default.chat"}, + {Name: "model-services/main.default.embed"}, + }, nil) + + items, err := prompt.ListModelServicesInSchema(ctx, "main", "default") + require.NoError(t, err) + assert.Equal(t, []prompt.ListItem{ + {ID: "main.default.chat", Label: "main.default.chat"}, + {ID: "main.default.embed", Label: "main.default.embed"}, + }, items) +} diff --git a/libs/apps/prompt/prefetch.go b/libs/apps/prompt/prefetch.go index d3f38bade08..59373ba107a 100644 --- a/libs/apps/prompt/prefetch.go +++ b/libs/apps/prompt/prefetch.go @@ -34,10 +34,11 @@ var firstStepPrefetch = map[string]struct { cacheKey string ctor pagedConstructor }{ - ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, - ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, - ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, - ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, + ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeModelService: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, + ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, } // PrefetchResources kicks off a background goroutine for every resource type diff --git a/libs/apps/prompt/prompt.go b/libs/apps/prompt/prompt.go index 6bbbaaa2d28..5f5e744eaf1 100644 --- a/libs/apps/prompt/prompt.go +++ b/libs/apps/prompt/prompt.go @@ -852,6 +852,11 @@ func PromptForUCFunction(ctx context.Context, r manifest.Resource, required bool return promptUCResource(ctx, r, required, "UC Function", "Fetching functions...", ListFunctionsInSchema) } +// PromptForModelService shows a three-step picker for AI Gateway model services: catalog -> schema -> model service. +func PromptForModelService(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { + return promptUCResource(ctx, r, required, "Model Service", "Fetching model services...", ListModelServicesInSchema) +} + // PromptForUCConnection shows a picker for UC connections. func PromptForUCConnection(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { title := resourceTitle("Select UC Connection", r) diff --git a/libs/apps/prompt/resource_registry.go b/libs/apps/prompt/resource_registry.go index e6846ad2f2b..8744aa5594c 100644 --- a/libs/apps/prompt/resource_registry.go +++ b/libs/apps/prompt/resource_registry.go @@ -16,6 +16,7 @@ const ( ResourceTypeVectorSearchIndex = "vector_search_index" ResourceTypeUCFunction = "uc_function" ResourceTypeUCConnection = "uc_connection" + ResourceTypeModelService = "model_service" ResourceTypeDatabase = "database" // Lakebase Provisioned (V1) ResourceTypePostgres = "postgres" // Lakebase Autoscaling (V2) ResourceTypeGenieSpace = "genie_space" @@ -62,6 +63,8 @@ func GetPromptFunc(resourceType string) (PromptResourceFunc, bool) { return PromptForVectorSearchIndex, true case ResourceTypeUCFunction: return PromptForUCFunction, true + case ResourceTypeModelService: + return PromptForModelService, true case ResourceTypeUCConnection: return PromptForUCConnection, true case ResourceTypeDatabase: From 877bd4522ac07b380dc3d5d62ce9c6cdbfeee3ee Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 4 Nov 2026 14:37:12 +0100 Subject: [PATCH 2/3] feat(apps): open model service picker on system.ai List the system.ai model services first with a "Browse other catalogs..." entry that falls back to the catalog -> schema -> model service picker. Finding the system catalog otherwise means paging through every catalog. Co-authored-by: Isaac --- libs/apps/prompt/prompt.go | 41 +++++++++++++++++++++++++++++++-- libs/apps/prompt/prompt_test.go | 31 +++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/libs/apps/prompt/prompt.go b/libs/apps/prompt/prompt.go index 5f5e744eaf1..90b61e88ffd 100644 --- a/libs/apps/prompt/prompt.go +++ b/libs/apps/prompt/prompt.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/apps" "github.com/databricks/databricks-sdk-go/service/postgres" @@ -852,9 +853,45 @@ func PromptForUCFunction(ctx context.Context, r manifest.Resource, required bool return promptUCResource(ctx, r, required, "UC Function", "Fetching functions...", ListFunctionsInSchema) } -// PromptForModelService shows a three-step picker for AI Gateway model services: catalog -> schema -> model service. +const browseCatalogsID = "__browse_catalogs__" + +// systemAIModelServiceItems lists the system.ai model services followed by a "Browse other catalogs..." +// entry. Without access to system.ai only the browse entry is returned. +func systemAIModelServiceItems(ctx context.Context) ([]ListItem, error) { + items, err := ListModelServicesInSchema(ctx, "system", "ai") + if errors.Is(err, apierr.ErrPermissionDenied) || errors.Is(err, apierr.ErrNotFound) { + items, err = nil, nil + } + if err != nil { + return nil, err + } + return append(items, ListItem{ID: browseCatalogsID, Label: "Browse other catalogs..."}), nil +} + +// PromptForModelService shows the system.ai model services first, because finding the system catalog +// means paging through every catalog in large workspaces. "Browse other catalogs..." falls back to the +// catalog -> schema -> model service picker. func PromptForModelService(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - return promptUCResource(ctx, r, required, "Model Service", "Fetching model services...", ListModelServicesInSchema) + var items []ListItem + err := RunWithSpinnerCtx(ctx, "Fetching model services...", func() error { + var fetchErr error + items, fetchErr = systemAIModelServiceItems(ctx) + return fetchErr + }) + if err != nil { + return nil, err + } + value := browseCatalogsID + if len(items) > 1 { + value, err = PromptFromList(ctx, resourceTitle("Select Model Service", r), "", items, required) + if err != nil { + return nil, err + } + } + if value == browseCatalogsID { + return promptUCResource(ctx, r, required, "Model Service", "Fetching model services...", ListModelServicesInSchema) + } + return singleValueResult(r, value), nil } // PromptForUCConnection shows a picker for UC connections. diff --git a/libs/apps/prompt/prompt_test.go b/libs/apps/prompt/prompt_test.go index 239fb852cfa..38e21b7fad1 100644 --- a/libs/apps/prompt/prompt_test.go +++ b/libs/apps/prompt/prompt_test.go @@ -9,7 +9,11 @@ import ( "time" "github.com/databricks/cli/libs/apps/manifest" + "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/catalog" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -524,3 +528,30 @@ func TestPrintSuccessNotInPlace(t *testing.T) { assert.Contains(t, got, "cd my-app") assert.Contains(t, got, "npm run dev") } + +func TestSystemAIModelServiceItems(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). + Return([]catalog.ModelService{{Name: "model-services/system.ai.chat"}}, nil) + + items, err := systemAIModelServiceItems(ctx) + require.NoError(t, err) + assert.Equal(t, []ListItem{ + {ID: "system.ai.chat", Label: "system.ai.chat"}, + {ID: browseCatalogsID, Label: "Browse other catalogs..."}, + }, items) +} + +func TestSystemAIModelServiceItemsWithoutAccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) + m.GetMockAiGatewayAPI().EXPECT(). + ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). + Return(nil, apierr.ErrPermissionDenied) + + items, err := systemAIModelServiceItems(ctx) + require.NoError(t, err) + assert.Equal(t, []ListItem{{ID: browseCatalogsID, Label: "Browse other catalogs..."}}, items) +} From c01e56dbac608735cfddc7579b5380cb38a255a0 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 4 Nov 2026 15:33:18 +0100 Subject: [PATCH 3/3] refactor(apps): prefetch system.ai model services and trim picker code Prefetch the system.ai model service list for model_service resources instead of the catalog list, so the first picker no longer blocks on a fetch; the catalog list is fetched on demand when browsing. Merge the single-caller helper into listSystemAIModelServices and shorten the new doc comments. Co-authored-by: Isaac --- libs/apps/prompt/listers.go | 3 +-- libs/apps/prompt/prefetch.go | 3 ++- libs/apps/prompt/prompt.go | 22 +++++++--------------- libs/apps/prompt/prompt_test.go | 12 ++++++------ 4 files changed, 16 insertions(+), 24 deletions(-) diff --git a/libs/apps/prompt/listers.go b/libs/apps/prompt/listers.go index 834a0fa7abc..de8c3f1a404 100644 --- a/libs/apps/prompt/listers.go +++ b/libs/apps/prompt/listers.go @@ -155,8 +155,7 @@ func ListFunctionsInSchema(ctx context.Context, catalogName, schemaName string) return capResults(out), nil } -// ListModelServicesInSchema returns AI Gateway model services within a catalog.schema as selectable items. -// IDs are the full catalog.schema.name, without the API's model-services/ prefix. +// ListModelServicesInSchema returns model services in catalog.schema, with IDs as catalog.schema.name. func ListModelServicesInSchema(ctx context.Context, catalogName, schemaName string) ([]ListItem, error) { w, err := workspaceClient(ctx) if err != nil { diff --git a/libs/apps/prompt/prefetch.go b/libs/apps/prompt/prefetch.go index 59373ba107a..e2a2b017f19 100644 --- a/libs/apps/prompt/prefetch.go +++ b/libs/apps/prompt/prefetch.go @@ -26,6 +26,7 @@ const ( cacheKeyCatalogs = "_catalogs" cacheKeyDatabaseInstances = "_database_instances" cacheKeyPostgresProjects = "_postgres_projects" + cacheKeySystemAIModels = "_system_ai_model_services" ) // firstStepPrefetch maps multi-step resource types to the cache key and @@ -36,7 +37,7 @@ var firstStepPrefetch = map[string]struct { }{ ResourceTypeVolume: {cacheKeyCatalogs, ListCatalogs}, ResourceTypeUCFunction: {cacheKeyCatalogs, ListCatalogs}, - ResourceTypeModelService: {cacheKeyCatalogs, ListCatalogs}, + ResourceTypeModelService: {cacheKeySystemAIModels, listSystemAIModelServices}, ResourceTypeDatabase: {cacheKeyDatabaseInstances, ListDatabaseInstances}, ResourceTypePostgres: {cacheKeyPostgresProjects, ListPostgresProjects}, } diff --git a/libs/apps/prompt/prompt.go b/libs/apps/prompt/prompt.go index 90b61e88ffd..915bc6b4ec3 100644 --- a/libs/apps/prompt/prompt.go +++ b/libs/apps/prompt/prompt.go @@ -855,9 +855,8 @@ func PromptForUCFunction(ctx context.Context, r manifest.Resource, required bool const browseCatalogsID = "__browse_catalogs__" -// systemAIModelServiceItems lists the system.ai model services followed by a "Browse other catalogs..." -// entry. Without access to system.ai only the browse entry is returned. -func systemAIModelServiceItems(ctx context.Context) ([]ListItem, error) { +// listSystemAIModelServices lists system.ai model services plus the browse entry; no access yields only the latter. +func listSystemAIModelServices(ctx context.Context) (*PagedFetcher, error) { items, err := ListModelServicesInSchema(ctx, "system", "ai") if errors.Is(err, apierr.ErrPermissionDenied) || errors.Is(err, apierr.ErrNotFound) { items, err = nil, nil @@ -865,25 +864,18 @@ func systemAIModelServiceItems(ctx context.Context) ([]ListItem, error) { if err != nil { return nil, err } - return append(items, ListItem{ID: browseCatalogsID, Label: "Browse other catalogs..."}), nil + return &PagedFetcher{Items: append(items, ListItem{ID: browseCatalogsID, Label: "Browse other catalogs..."})}, nil } -// PromptForModelService shows the system.ai model services first, because finding the system catalog -// means paging through every catalog in large workspaces. "Browse other catalogs..." falls back to the -// catalog -> schema -> model service picker. +// PromptForModelService opens on system.ai, since paging to the system catalog is slow in large workspaces. func PromptForModelService(ctx context.Context, r manifest.Resource, required bool) (map[string]string, error) { - var items []ListItem - err := RunWithSpinnerCtx(ctx, "Fetching model services...", func() error { - var fetchErr error - items, fetchErr = systemAIModelServiceItems(ctx) - return fetchErr - }) + f, err := getFetcherByKey(ctx, cacheKeySystemAIModels, "Fetching model services...", listSystemAIModelServices) if err != nil { return nil, err } value := browseCatalogsID - if len(items) > 1 { - value, err = PromptFromList(ctx, resourceTitle("Select Model Service", r), "", items, required) + if len(f.Items) > 1 { + value, err = PromptFromList(ctx, resourceTitle("Select Model Service", r), "", f.Items, required) if err != nil { return nil, err } diff --git a/libs/apps/prompt/prompt_test.go b/libs/apps/prompt/prompt_test.go index 38e21b7fad1..9e77eb76aca 100644 --- a/libs/apps/prompt/prompt_test.go +++ b/libs/apps/prompt/prompt_test.go @@ -529,29 +529,29 @@ func TestPrintSuccessNotInPlace(t *testing.T) { assert.Contains(t, got, "npm run dev") } -func TestSystemAIModelServiceItems(t *testing.T) { +func TestListSystemAIModelServices(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) m.GetMockAiGatewayAPI().EXPECT(). ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). Return([]catalog.ModelService{{Name: "model-services/system.ai.chat"}}, nil) - items, err := systemAIModelServiceItems(ctx) + f, err := listSystemAIModelServices(ctx) require.NoError(t, err) assert.Equal(t, []ListItem{ {ID: "system.ai.chat", Label: "system.ai.chat"}, {ID: browseCatalogsID, Label: "Browse other catalogs..."}, - }, items) + }, f.Items) } -func TestSystemAIModelServiceItemsWithoutAccess(t *testing.T) { +func TestListSystemAIModelServicesWithoutAccess(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) m.GetMockAiGatewayAPI().EXPECT(). ListModelServicesAll(ctx, catalog.ListModelServicesRequest{Parent: "schemas/system.ai"}). Return(nil, apierr.ErrPermissionDenied) - items, err := systemAIModelServiceItems(ctx) + f, err := listSystemAIModelServices(ctx) require.NoError(t, err) - assert.Equal(t, []ListItem{{ID: browseCatalogsID, Label: "Browse other catalogs..."}}, items) + assert.Equal(t, []ListItem{{ID: browseCatalogsID, Label: "Browse other catalogs..."}}, f.Items) }