feat(go/ai): add typed-config NewTyped* constructors

go/ai/action_test.go

@@ -71,7 +71,7 @@ func defineProgrammableModel(r api.Registry) *programmableModel {
 		Tools:     true,
 		Multiturn: true,
 	}
-	DefineModel(r, "programmableModel", &ModelOptions{Supports: supports}, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
+	defineModel(r, "programmableModel", &ModelOptions{Supports: supports}, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
 		return pm.Generate(ctx, r, req, &ToolConfig{MaxTurns: 5}, cb)
 	})
 	return pm
@@ -97,7 +97,7 @@ func TestGenerateAction(t *testing.T) {
 
 			pm := defineProgrammableModel(r)
 
-			DefineTool(r, "testTool", "description",
+			defineTool(r, "testTool", "description",
 				func(ctx *ToolContext, input any) (any, error) {
 					return "tool called", nil
 				})

go/ai/background_model.go

@@ -19,6 +19,7 @@ package ai
 import (
 	"context"
 	"errors"
+	"maps"
 
 	"github.com/firebase/genkit/go/core"
 	"github.com/firebase/genkit/go/core/api"
@@ -41,55 +42,107 @@ type BackgroundModel interface {
 	SupportsCancel() bool
 }
 
-// backgroundModel is the concrete implementation of BackgroundModel interface.
-type backgroundModel struct {
-	core.BackgroundAction[*ModelRequest, *ModelResponse]
+// backgroundAction is an unexported alias of [core.BackgroundAction] used as
+// the embedded field in [BackgroundModelAction]; see the action alias in
+// generate.go for why.
+type backgroundAction[In, Out any] = core.BackgroundAction[In, Out]
+
+// BackgroundModelAction is a background model backed by registry actions. It
+// is the concrete type returned by [NewTypedBackgroundModel]; return it from
+// a plugin's Init for the framework to register.
+type BackgroundModelAction struct {
+	backgroundAction[*ModelRequest, *ModelResponse]
 }
 
+// BackgroundModelAction is a full registry action and can be passed anywhere
+// an [api.Action] is accepted as well as anywhere a [BackgroundModel] is
+// accepted. Its Register method registers the start, check, and cancel
+// actions together.
+var (
+	_ api.Action      = (*BackgroundModelAction)(nil)
+	_ BackgroundModel = (*BackgroundModelAction)(nil)
+)
+
 // ModelOperation is a background operation for a model.
 type ModelOperation = core.Operation[*ModelResponse]
 
 // StartModelOpFunc starts a background model operation.
 type StartModelOpFunc = func(ctx context.Context, req *ModelRequest) (*ModelOperation, error)
 
+// TypedStartModelOpFunc is a [StartModelOpFunc] that additionally
+// receives the request's typed Config: the framework deserializes the
+// request's raw config into it before calling the function (see
+// [NewTypedBackgroundModel]).
+type TypedStartModelOpFunc[Config any] = func(ctx context.Context, req *ModelRequest, config Config) (*ModelOperation, error)
+
 // CheckModelOpFunc checks the status of a background model operation.
 type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
 
 // CancelModelOpFunc cancels a background model operation.
 type CancelModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
 
-// BackgroundModelOptions holds configuration for defining a background model
+// TypedBackgroundModelOptions configures a background model created with
+// [NewTypedBackgroundModel]. It holds descriptor data plus optional lifecycle
+// hooks; the required start and check functions are constructor arguments.
+type TypedBackgroundModelOptions struct {
+	ConfigSchema map[string]any // JSON schema for the model's config.
+	Label        string         // User-friendly name for the model.
+	Stage        ModelStage     // Indicates the maturity stage of the model.
+	Supports     *ModelSupports // Capabilities of the model.
+	Versions     []string       // Available versions of the model.
+	Metadata     map[string]any // Arbitrary key-value data attached to the action descriptor.
+
+	// Cancel cancels a running operation. Optional: nil means the model does
+	// not support canceling operations.
+	Cancel CancelModelOpFunc
+}
+
+// BackgroundModelOptions holds configuration for defining a background model.
+//
+// Deprecated: Use [TypedBackgroundModelOptions] with
+// [NewTypedBackgroundModel].
 type BackgroundModelOptions struct {
 	ModelOptions
-	Cancel   CancelModelOpFunc // Function that cancels a background model operation.
-	Metadata map[string]any    // Additional metadata.
+	// Cancel is the function that cancels a background model operation.
+	Cancel   CancelModelOpFunc
+	Metadata map[string]any // Additional metadata.
 }
 
-// LookupBackgroundModel looks up a BackgroundAction registered by [DefineBackgroundModel].
+// LookupBackgroundModel looks up a registered [BackgroundModel] by name.
 // It returns nil if the background model was not found.
 func LookupBackgroundModel(r api.Registry, name string) BackgroundModel {
 	key := api.KeyFromName(api.ActionTypeBackgroundModel, name)
 	action := core.LookupBackgroundAction[*ModelRequest, *ModelResponse](r, key)
 	if action == nil {
 		return nil
 	}
-	return &backgroundModel{*action}
+	return &BackgroundModelAction{*action}
 }
 
-// NewBackgroundModel defines a new model that runs in the background.
-func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel {
+// NewTypedBackgroundModel creates a new [BackgroundModelAction]. Register it
+// with [BackgroundModelAction.Register] to make it resolvable by name.
+//
+// Config is the model's typed configuration; it is usually inferred from
+// startFn's signature. See [NewTypedModel] for how the request's config
+// is deserialized.
+func NewTypedBackgroundModel[Config any](
+	name string,
+	opts *TypedBackgroundModelOptions,
+	startFn TypedStartModelOpFunc[Config],
+	checkFn CheckModelOpFunc,
+) *BackgroundModelAction {
 	if name == "" {
-		panic("ai.NewBackgroundModel: name is required")
+		panic("ai.NewTypedBackgroundModel: name is required")
 	}
 	if startFn == nil {
-		panic("ai.NewBackgroundModel: startFn is required")
+		panic("ai.NewTypedBackgroundModel: startFn is required")
 	}
 	if checkFn == nil {
-		panic("ai.NewBackgroundModel: checkFn is required")
+		panic("ai.NewTypedBackgroundModel: checkFn is required")
 	}
 
 	if opts == nil {
-		opts = &BackgroundModelOptions{}
+		opts = &TypedBackgroundModelOptions{}
 	}
 	if opts.Label == "" {
 		opts.Label = name
@@ -98,41 +151,57 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start
 		opts.Supports = &ModelSupports{}
 	}
 
-	metadata := map[string]any{
-		"type": api.ActionTypeBackgroundModel,
-		"model": map[string]any{
-			"label": opts.Label,
-			"supports": map[string]any{
-				"media":       opts.Supports.Media,
-				"context":     opts.Supports.Context,
-				"multiturn":   opts.Supports.Multiturn,
-				"systemRole":  opts.Supports.SystemRole,
-				"tools":       opts.Supports.Tools,
-				"toolChoice":  opts.Supports.ToolChoice,
-				"constrained": opts.Supports.Constrained,
-				"output":      opts.Supports.Output,
-				"contentType": opts.Supports.ContentType,
-				"longRunning": opts.Supports.LongRunning,
-			},
-			"versions":      opts.Versions,
-			"stage":         opts.Stage,
-			"customOptions": opts.ConfigSchema,
+	configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config")
+
+	// The reserved type and model keys win over caller-provided metadata.
+	metadata := make(map[string]any, len(opts.Metadata)+2)
+	maps.Copy(metadata, opts.Metadata)
+	metadata["type"] = api.ActionTypeBackgroundModel
+	metadata["model"] = map[string]any{
+		"label": opts.Label,
+		"supports": map[string]any{
+			"media":       opts.Supports.Media,
+			"context":     opts.Supports.Context,
+			"multiturn":   opts.Supports.Multiturn,
+			"systemRole":  opts.Supports.SystemRole,
+			"tools":       opts.Supports.Tools,
+			"toolChoice":  opts.Supports.ToolChoice,
+			"constrained": opts.Supports.Constrained,
+			"output":      opts.Supports.Output,
+			"contentType": opts.Supports.ContentType,
+			"longRunning": opts.Supports.LongRunning,
 		},
+		"versions":      opts.Versions,
+		"stage":         opts.Stage,
+		"customOptions": configSchema,
 	}
 
-	inputSchema := core.InferSchemaMap(ModelRequest{})
-	if inputSchema != nil && opts.ConfigSchema != nil {
-		if props, ok := inputSchema["properties"].(map[string]any); ok {
-			props["config"] = opts.ConfigSchema
+	typedStartFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) {
+		// req.Config was normalized to the exact Config type by
+		// normalizeConfig below, so this hits the fast path.
+		cfg, err := resolveConfig[Config](req.Config)
+		if err != nil {
+			return nil, err
 		}
+		return startFn(ctx, req, cfg)
 	}
 
-	mws := []ModelMiddleware{
-		simulateSystemPrompt(&opts.ModelOptions, nil),
-		augmentWithContext(&opts.ModelOptions, nil),
-		validateSupport(name, &opts.ModelOptions),
+	mopts := &ModelOptions{
+		ConfigSchema: opts.ConfigSchema,
+		Label:        opts.Label,
+		Stage:        opts.Stage,
+		Supports:     opts.Supports,
+		Versions:     opts.Versions,
 	}
-	fn := core.ChainMiddleware(mws...)(backgroundModelToModelFn(startFn))
+
+	// normalizeConfig runs outermost so that the built-in wrappers and the
+	// start function all see the typed, converted config on the request.
+	fn := core.ChainMiddleware(
+		normalizeConfig[Config](name, opts.Versions),
+		simulateSystemPrompt(mopts, nil),
+		augmentWithContext(mopts, nil),
+		validateSupport(name, mopts),
+	)(backgroundModelToModelFn(typedStartFn))
 
 	wrappedFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) {
 		resp, err := fn(ctx, req, nil)
@@ -143,14 +212,35 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start
 		return modelOpFromResponse(resp)
 	}
 
-	return &backgroundModel{*core.NewBackgroundAction(name, api.ActionTypeBackgroundModel, metadata, wrappedFn, checkFn, opts.Cancel)}
+	return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{
+		Metadata:    metadata,
+		InputSchema: inputSchema,
+	}, wrappedFn, checkFn, opts.Cancel)}
 }
 
-// DefineBackgroundModel defines and registers a new model that runs in the background.
-func DefineBackgroundModel(r *registry.Registry, name string, opts *BackgroundModelOptions, fn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel {
-	m := NewBackgroundModel(name, opts, fn, checkFn)
-	m.Register(r)
-	return m
+// NewBackgroundModel defines a new model that runs in the background.
+//
+// Deprecated: Use [NewTypedBackgroundModel], which passes the request's
+// config to startFn as a typed value instead of leaving it type-erased on the
+// request.
+func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel {
+	if startFn == nil {
+		panic("ai.NewBackgroundModel: startFn is required")
+	}
+	if opts == nil {
+		opts = &BackgroundModelOptions{}
+	}
+	return NewTypedBackgroundModel(name, &TypedBackgroundModelOptions{
+		ConfigSchema: opts.ConfigSchema,
+		Label:        opts.Label,
+		Stage:        opts.Stage,
+		Supports:     opts.Supports,
+		Versions:     opts.Versions,
+		Metadata:     opts.Metadata,
+		Cancel:       opts.Cancel,
+	}, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) {
+		return startFn(ctx, req)
+	}, checkFn)
 }
 
 // GenerateOperation generates a model response as a long-running operation based on the provided options.

go/ai/config.go

@@ -0,0 +1,109 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package ai
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/firebase/genkit/go/core"
+	"github.com/firebase/genkit/go/internal/base"
+)
+
+// This file holds the typed-config plumbing shared by models, embedders, and
+// evaluators. Requests carry config as `any` on the wire; the
+// NewTyped* constructors wrap the user's typed function so that
+// the raw value is deserialized into the Config type parameter before the
+// function runs, and the request's type-erased config slot is normalized to
+// that same converted value so the two views never disagree.
+
+// nullableConfigSchema wraps a config schema for the request input-schema
+// slot so that an explicit JSON null is accepted on the wire: a typed-nil Go
+// config marshals to null and resolves to the zero Config value, so it must
+// not be rejected by input validation. The advertised config schema (the
+// customOptions metadata) stays unwrapped.
+func nullableConfigSchema(schema map[string]any) map[string]any {
+	if schema == nil {
+		return nil
+	}
+	return map[string]any{"anyOf": []any{schema, map[string]any{"type": "null"}}}
+}
+
+// normalizeConfig returns a model middleware that resolves the request's raw
+// config into Config and writes the converted value back into the request.
+// It runs as the outermost step of a model's built-in chain so that every
+// wrapper after it, and the model function itself, sees the typed value; by
+// then the config has already been validated against the model's config
+// schema at the action boundary.
+//
+// Version validation runs here, against the raw config, because conversion is
+// lossy: a "version" key sent by a JSON caller would be silently dropped when
+// deserializing into a Config type that has no such field.
+func normalizeConfig[Config any](model string, versions []string) ModelMiddleware {
+	return func(next ModelFunc) ModelFunc {
+		return func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
+			if err := validateVersion(model, versions, req.Config); err != nil {
+				return nil, err
+			}
+			cfg, err := resolveConfig[Config](req.Config)
+			if err != nil {
+				return nil, err
+			}
+			req.Config = cfg
+			return next(ctx, req, cb)
+		}
+	}
+}
+
+// actionConfigSchemas returns the action's effective config schema (the
+// explicit override when set, otherwise the schema inferred from Config) and
+// the request input schema with its config slot replaced by the null-tolerant
+// wrapping of that schema. reqZero is the zero request value to infer the
+// input schema from and key is the config slot's wire name ("config" for
+// models, "options" for embedders and evaluators).
+func actionConfigSchemas[Config any](override map[string]any, reqZero any, key string) (configSchema, inputSchema map[string]any) {
+	configSchema = override
+	if configSchema == nil {
+		configSchema = base.SchemaMapFor[Config]()
+	}
+
+	inputSchema = core.InferSchemaMap(reqZero)
+	if inputSchema != nil && configSchema != nil {
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props[key] = nullableConfigSchema(configSchema)
+		}
+	}
+	return configSchema, inputSchema
+}
+
+// resolveConfig converts the raw config value carried by a request into the
+// typed Config the action was defined with. It accepts the exact Config type
+// (or a pointer to it, which is dereferenced), a map[string]any (as sent by
+// the Dev UI and other JSON callers, deserialized via a JSON round-trip), or
+// nil (yielding the zero value). Any other type is rejected so one provider's
+// config cannot be silently passed to another provider's action.
+func resolveConfig[Config any](raw any) (Config, error) {
+	cfg, err := base.ConvertToExact[Config](raw)
+	if err != nil {
+		if errors.Is(err, base.ErrTypeMismatch) {
+			return cfg, core.NewPublicError(core.INVALID_ARGUMENT, fmt.Sprintf("invalid config type %T, want %T or map[string]any", raw, cfg), nil)
+		}
+		return cfg, core.NewPublicError(core.INVALID_ARGUMENT, fmt.Sprintf("invalid config for %T; check that field names and value types match: %v", cfg, err), nil)
+	}
+	return cfg, nil
+}