refactor: make toolsets and providers explicit

cmd/root/debug.go

@@ -13,6 +13,7 @@ import (
 	"github.com/docker/docker-agent/pkg/sessiontitle"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/telemetry"
 	"github.com/docker/docker-agent/pkg/tools"
 	skillstool "github.com/docker/docker-agent/pkg/tools/builtin/skills"
@@ -76,6 +77,7 @@ func (f *debugFlags) loadTeam(ctx context.Context, agentFilename string, opts ..
 		return nil, err
 	}
 
+	opts = append(loaderdefaults.Opts(), opts...)
 	t, err := teamloader.Load(ctx, agentSource, &f.runConfig, opts...)
 	if err != nil {
 		return nil, err

cmd/root/new.go

@@ -57,13 +57,14 @@ func (f *newFlags) runNewCommand(cmd *cobra.Command, args []string) (commandErr
 		telemetry.TrackCommandError(ctx, "new", args, commandErr)
 	}()
 
-	t, err := creator.Agent(ctx, &f.runConfig, f.modelParam)
+	loadResult, err := creator.Load(ctx, &f.runConfig, f.modelParam)
 	if err != nil {
 		return err
 	}
+	t := loadResult.Team
 	defer stopToolSets(t)
 
-	rt, err := runtime.New(t)
+	rt, err := runtime.New(t, runtime.WithProviderRegistry(loadResult.ProviderRegistry))
 	if err != nil {
 		return err
 	}

cmd/root/run.go

@@ -31,6 +31,7 @@ import (
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/telemetry"
 	"github.com/docker/docker-agent/pkg/tui"
 	"github.com/docker/docker-agent/pkg/tui/styles"
@@ -693,9 +694,7 @@ func (f *runExecFlags) dispatchWorktreeCreate(ctx context.Context, out *cli.Prin
 }
 
 func (f *runExecFlags) loadAgentFrom(ctx context.Context, req runtime.LoadTeamRequest) (*teamloader.LoadResult, error) {
-	opts := []teamloader.Opt{
-		teamloader.WithModelOverrides(req.ModelOverrides),
-	}
+	opts := append(loaderdefaults.Opts(), teamloader.WithModelOverrides(req.ModelOverrides))
 	if len(req.PromptFiles) > 0 {
 		opts = append(opts, teamloader.WithPromptFiles(req.PromptFiles))
 	}
@@ -712,6 +711,7 @@ func (f *runExecFlags) runtimeOpts(loadResult *teamloader.LoadResult, runConfig
 		Providers:          loadResult.Providers,
 		ModelsGateway:      runConfig.ModelsGateway,
 		EnvProvider:        runConfig.EnvProvider(),
+		ProviderRegistry:   loadResult.ProviderRegistry,
 		AgentDefaultModels: loadResult.AgentDefaultModels,
 	}
 	opts := []runtime.Opt{

cmd/root/toolsets.go

@@ -1,13 +0,0 @@
-package root
-
-// Blank-import optional toolsets so they register themselves with teamloader
-// (via their init()) and are available in the default toolset registry.
-//
-// These toolsets are deliberately kept out of teamloader's static import graph
-// so embedders that don't need them aren't forced to link their dependencies.
-// "rag" in particular pulls in a cgo tree-sitter dependency; importing it here
-// keeps the docker-agent CLI's behavior unchanged while letting lighter
-// embedders opt out.
-import (
-	_ "github.com/docker/docker-agent/pkg/tools/builtin/rag"
-)

cmd/root/toolsets_test.go

@@ -7,22 +7,13 @@ import (
 
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/teamloader"
+	loadertoolsets "github.com/docker/docker-agent/pkg/teamloader/toolsets"
 )
 
-// TestRAGToolsetRegistered guards the blank import in toolsets.go: the "rag"
-// toolset registers itself with teamloader via init(), so it must resolve in
-// the default registry of the CLI binary. If someone drops the blank import,
-// "rag" would silently become an unknown toolset type at runtime; this test
-// turns that into a build-time failure instead.
 func TestRAGToolsetRegistered(t *testing.T) {
-	registry := teamloader.NewDefaultToolsetRegistry()
+	registry := loadertoolsets.NewDefaultToolsetRegistry()
 
-	// A rag toolset with no config fails inside the creator with a config
-	// error — not with "unknown toolset type". Reaching that creator-level
-	// error proves the "rag" creator is registered.
 	_, err := registry.CreateTool(t.Context(), latest.Toolset{Type: "rag"}, "", &config.RuntimeConfig{}, "")
 	require.Error(t, err)
-	require.NotContains(t, err.Error(), "unknown toolset type",
-		"the rag toolset is not registered; check the blank import in cmd/root/toolsets.go")
+	require.NotContains(t, err.Error(), "unknown toolset type")
 }

e2e/mcp_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/mcp"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 )
 
 func TestMCP_SingleAgent(t *testing.T) {
@@ -19,7 +20,7 @@ func TestMCP_SingleAgent(t *testing.T) {
 	require.NoError(t, err)
 
 	_, runConfig := startRecordingAIProxy(t)
-	team, err := teamloader.Load(ctx, agentSource, runConfig)
+	team, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	require.NoError(t, err, "Failed to load agent")
 	t.Cleanup(func() {
 		require.NoError(t, team.StopToolSets(ctx))
@@ -42,7 +43,7 @@ func TestMCP_MultiAgent(t *testing.T) {
 	require.NoError(t, err)
 
 	_, runConfig := startRecordingAIProxy(t)
-	team, err := teamloader.Load(ctx, agentSource, runConfig)
+	team, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	require.NoError(t, err, "Failed to load team")
 	t.Cleanup(func() {
 		require.NoError(t, team.StopToolSets(ctx))

e2e/runtime_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/docker/docker-agent/pkg/runtime"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 )
 
 func TestRuntime_OpenAI_Basic(t *testing.T) {
@@ -21,7 +22,7 @@ func TestRuntime_OpenAI_Basic(t *testing.T) {
 	require.NoError(t, err)
 
 	_, runConfig := startRecordingAIProxy(t)
-	team, err := teamloader.Load(ctx, agentSource, runConfig)
+	team, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	require.NoError(t, err)
 
 	rt, err := runtime.New(team)
@@ -58,7 +59,7 @@ func TestRuntime_MultiAgent_SessionReload(t *testing.T) {
 	require.NoError(t, err)
 
 	_, runConfig := startRecordingAIProxy(t)
-	team, err := teamloader.Load(ctx, agentSource, runConfig)
+	team, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	require.NoError(t, err)
 
 	// Use a SQLite store so we test real persistence and reload.
@@ -107,7 +108,7 @@ func TestRuntime_Mistral_Basic(t *testing.T) {
 	require.NoError(t, err)
 
 	_, runConfig := startRecordingAIProxy(t)
-	team, err := teamloader.Load(ctx, agentSource, runConfig, teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))
+	team, err := teamloader.Load(ctx, agentSource, runConfig, append(loaderdefaults.Opts(), teamloader.WithModelOverrides([]string{"mistral/mistral-small"}))...)
 	require.NoError(t, err)
 
 	rt, err := runtime.New(team)

pkg/a2a/adapter_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 )
 
 func TestNewDockerAgentAdapter(t *testing.T) {
@@ -17,7 +18,7 @@ func TestNewDockerAgentAdapter(t *testing.T) {
 	agentSource, err := config.Resolve("testdata/basic.yaml", nil)
 	require.NoError(t, err)
 
-	team, err := teamloader.Load(t.Context(), agentSource, &config.RuntimeConfig{})
+	team, err := teamloader.Load(t.Context(), agentSource, &config.RuntimeConfig{}, loaderdefaults.Opts()...)
 	require.NoError(t, err)
 	defer func() {
 		require.NoError(t, team.StopToolSets(t.Context()))
@@ -36,7 +37,7 @@ func TestNewCAgentAdapter_NonExistent(t *testing.T) {
 	agentSource, err := config.Resolve("testdata/basic.yaml", nil)
 	require.NoError(t, err)
 
-	team, err := teamloader.Load(t.Context(), agentSource, &config.RuntimeConfig{})
+	team, err := teamloader.Load(t.Context(), agentSource, &config.RuntimeConfig{}, loaderdefaults.Opts()...)
 	require.NoError(t, err)
 	defer func() {
 		require.NoError(t, team.StopToolSets(t.Context()))

pkg/a2a/server.go

@@ -22,6 +22,7 @@ import (
 	pathx "github.com/docker/docker-agent/pkg/path"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/version"
 )
 
@@ -46,7 +47,7 @@ func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runCon
 		return err
 	}
 
-	t, err := teamloader.Load(ctx, agentSource, runConfig)
+	t, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	if err != nil {
 		return fmt.Errorf("failed to load agents: %w", err)
 	}

pkg/acp/agent.go

@@ -17,10 +17,12 @@ import (
 	"github.com/docker/docker-agent/pkg/agent"
 	"github.com/docker/docker-agent/pkg/chat"
 	"github.com/docker/docker-agent/pkg/config"
+	"github.com/docker/docker-agent/pkg/model/provider"
 	"github.com/docker/docker-agent/pkg/runtime"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/version"
 )
 
@@ -31,10 +33,11 @@ type Agent struct {
 	sessionStore session.Store
 	sessions     map[string]*Session
 
-	conn     *acp.AgentSideConnection
-	clientFS acp.FileSystemCapabilities
-	team     *team.Team
-	mu       sync.Mutex
+	conn             *acp.AgentSideConnection
+	clientFS         acp.FileSystemCapabilities
+	team             *team.Team
+	providerRegistry *provider.Registry
+	mu               sync.Mutex
 }
 
 var _ acp.Agent = (*Agent)(nil)
@@ -82,11 +85,14 @@ func (a *Agent) Initialize(ctx context.Context, params acp.InitializeRequest) (a
 	a.mu.Lock()
 	a.clientFS = params.ClientCapabilities.Fs
 	defer a.mu.Unlock()
-	t, err := teamloader.Load(ctx, a.agentSource, a.runConfig, teamloader.WithToolsetRegistry(createToolsetRegistry(a)))
+	loadOpts := append(loaderdefaults.Opts(), teamloader.WithToolsetRegistry(createToolsetRegistry(a)))
+	loadResult, err := teamloader.LoadWithConfig(ctx, a.agentSource, a.runConfig, loadOpts...)
 	if err != nil {
 		return acp.InitializeResponse{}, fmt.Errorf("failed to load teams: %w", err)
 	}
+	t := loadResult.Team
 	a.team = t
+	a.providerRegistry = loadResult.ProviderRegistry
 	slog.DebugContext(ctx, "Teams loaded successfully", "source", a.agentSource.Name(), "agent_count", t.Size())
 
 	agentTitle := "docker agent"
@@ -132,6 +138,7 @@ func (a *Agent) newRuntime(workingDir string) (runtime.Runtime, *agent.Agent, er
 	opts := []runtime.Opt{
 		runtime.WithCurrentAgent(defaultAgent.Name()),
 		runtime.WithSessionStore(a.sessionStore),
+		runtime.WithProviderRegistry(a.providerRegistry),
 	}
 	if workingDir != "" {
 		opts = append(opts, runtime.WithWorkingDir(workingDir))

pkg/acp/registry.go

@@ -7,6 +7,7 @@ import (
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loadertoolsets "github.com/docker/docker-agent/pkg/teamloader/toolsets"
 	"github.com/docker/docker-agent/pkg/tools"
 	"github.com/docker/docker-agent/pkg/tools/builtin/filesystem"
 )
@@ -15,7 +16,7 @@ import (
 func createToolsetRegistry(agent *Agent) teamloader.ToolsetRegistry {
 	return &acpToolsetRegistry{
 		agent:    agent,
-		registry: teamloader.NewDefaultToolsetRegistry(),
+		registry: loadertoolsets.NewDefaultToolsetRegistry(),
 	}
 }
 

pkg/chatserver/server.go

@@ -43,6 +43,7 @@ import (
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 )
 
 // Options configures the chat completions server. Future improvements
@@ -152,7 +153,7 @@ func loadTeam(ctx context.Context, agentFilename string, runConfig *config.Runti
 	if err != nil {
 		return nil, err
 	}
-	t, err := teamloader.Load(ctx, src, runConfig)
+	t, err := teamloader.Load(ctx, src, runConfig, loaderdefaults.Opts()...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to load agents: %w", err)
 	}

pkg/creator/agent.go

@@ -13,6 +13,7 @@ import (
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 )
 
 //go:embed instructions.txt
@@ -25,7 +26,7 @@ const (
 	creatorWelcomeMessage = "Hello! I'm here to create agents for you.\n\nCan you explain to me what the agent will be used for?"
 )
 
-// Agent creates and returns a team configured for the agent builder functionality.
+// Load creates and returns the agent-builder team load result.
 // The agent builder helps users create their own agent configurations interactively.
 //
 // Parameters:
@@ -34,7 +35,7 @@ const (
 //   - modelNameOverride: Optional model override (empty string uses auto-selection)
 //
 // Returns the configured team or an error if configuration fails.
-func Agent(ctx context.Context, runConfig *config.RuntimeConfig, modelNameOverride string) (*team.Team, error) {
+func Load(ctx context.Context, runConfig *config.RuntimeConfig, modelNameOverride string) (*teamloader.LoadResult, error) {
 	instructions := buildInstructions(ctx, runConfig)
 
 	configYAML, err := buildCreatorConfigYAML(instructions)
@@ -42,14 +43,23 @@ func Agent(ctx context.Context, runConfig *config.RuntimeConfig, modelNameOverri
 		return nil, fmt.Errorf("building creator config: %w", err)
 	}
 
-	return teamloader.Load(
+	return teamloader.LoadWithConfig(
 		ctx,
 		config.NewBytesSource("creator", configYAML),
 		runConfig,
-		teamloader.WithModelOverrides([]string{modelNameOverride}),
+		append(loaderdefaults.Opts(), teamloader.WithModelOverrides([]string{modelNameOverride}))...,
 	)
 }
 
+// Agent creates and returns a team configured for the agent builder functionality.
+func Agent(ctx context.Context, runConfig *config.RuntimeConfig, modelNameOverride string) (*team.Team, error) {
+	result, err := Load(ctx, runConfig, modelNameOverride)
+	if err != nil {
+		return nil, err
+	}
+	return result.Team, nil
+}
+
 // buildInstructions creates the full instruction set for the creator agent,
 // including provider-specific model configuration examples.
 func buildInstructions(ctx context.Context, runConfig *config.RuntimeConfig) string {

pkg/embeddedchat/embeddedchat.go

@@ -13,6 +13,7 @@ import (
 	dagentruntime "github.com/docker/docker-agent/pkg/runtime"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
@@ -38,7 +39,7 @@ type Config struct {
 	// config is used.
 	RuntimeConfig *dagentcfg.RuntimeConfig
 	// ToolsetRegistry resolves toolsets declared by AgentSource. When nil,
-	// docker-agent's default registry is used.
+	// docker-agent's full YAML-loading registry is used.
 	ToolsetRegistry teamloader.ToolsetRegistry
 	// RuntimeOptions are appended when constructing the runtime.
 	RuntimeOptions []dagentruntime.Opt
@@ -108,7 +109,7 @@ func New(ctx context.Context, cfg Config) (*Session, error) {
 		runConfig = &dagentcfg.RuntimeConfig{}
 	}
 
-	var loadOpts []teamloader.Opt
+	loadOpts := loaderdefaults.Opts()
 	if cfg.ToolsetRegistry != nil {
 		loadOpts = append(loadOpts, teamloader.WithToolsetRegistry(cfg.ToolsetRegistry))
 	}
@@ -122,6 +123,7 @@ func New(ctx context.Context, cfg Config) (*Session, error) {
 		Providers:          loaded.Providers,
 		ModelsGateway:      runConfig.ModelsGateway,
 		EnvProvider:        runConfig.EnvProvider(),
+		ProviderRegistry:   loaded.ProviderRegistry,
 		AgentDefaultModels: loaded.AgentDefaultModels,
 	}
 

pkg/mcp/server.go

@@ -20,6 +20,7 @@ import (
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/tools"
 	"github.com/docker/docker-agent/pkg/version"
 )
@@ -91,7 +92,7 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo
 		return nil, nil, err
 	}
 
-	t, err := teamloader.Load(ctx, agentSource, runConfig)
+	t, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to load agents: %w", err)
 	}

pkg/model/provider/base/base.go

@@ -18,6 +18,10 @@ type Config struct {
 	// Models stores the full models map for providers that need it (e.g., routers).
 	// This enables proper cloning of providers that reference other models.
 	Models map[string]latest.ModelConfig
+	// ProviderRegistry stores the registry that created this provider. It is kept
+	// as any to avoid a package cycle; pkg/model/provider type-asserts it when
+	// cloning providers with adjusted options.
+	ProviderRegistry any
 	// BaseURL is the resolved HTTP base URL the client talks to, when
 	// the provider is reachable over a configurable HTTP endpoint.
 	// Distinct from [latest.ModelConfig.BaseURL] (the user-typed input):
@@ -31,6 +35,10 @@ type Config struct {
 	BaseURL string
 }
 
+func (c *Config) SetProviderRegistry(registry any) {
+	c.ProviderRegistry = registry
+}
+
 // ID returns the provider and model identity as a [modelsdev.ID] so
 // callers cannot accidentally pass a bare model string where a
 // provider-qualified identity is required. The model component uses

pkg/model/provider/clone.go

@@ -18,7 +18,11 @@ func CloneWithOptions(ctx context.Context, baseProvider Provider, opts ...option
 
 	// Use NewWithModels to support cloning routers that reference other models.
 	// cfg.Models is populated by routers; for other providers it's nil (which is fine).
-	clone, err := NewWithModels(ctx, &modelConfig, cfg.Models, cfg.Env, mergedOpts...)
+	registry, _ := cfg.ProviderRegistry.(*Registry)
+	if registry == nil {
+		registry = DefaultRegistry()
+	}
+	clone, err := registry.NewWithModels(ctx, &modelConfig, cfg.Models, cfg.Env, mergedOpts...)
 	if err != nil {
 		slog.DebugContext(ctx, "Failed to clone provider; using base provider", "error", err, "id", baseProvider.ID())
 		return baseProvider

pkg/model/provider/clone_test.go

@@ -79,7 +79,7 @@ func TestCloneWithOptions_RouterWithModelReferences(t *testing.T) {
 	})
 
 	// Create the router with the models map
-	router, err := NewWithModels(t.Context(), routerCfg, models, env)
+	router, err := fullTestRegistry().NewWithModels(t.Context(), routerCfg, models, env)
 	require.NoError(t, err)
 
 	// Verify the original router has the models map stored
@@ -122,7 +122,7 @@ func TestCloneWithOptions_DirectProvider(t *testing.T) {
 		"OPENAI_API_KEY": "test-key",
 	})
 
-	provider, err := New(t.Context(), cfg, env)
+	provider, err := fullTestRegistry().New(t.Context(), cfg, env)
 	require.NoError(t, err)
 
 	// Clone with max tokens
@@ -159,7 +159,7 @@ func TestCloneWithOptions_PreservesMaxTokens(t *testing.T) {
 		"OPENAI_API_KEY": "test-key",
 	})
 
-	provider, err := New(t.Context(), cfg, env, options.WithMaxTokens(maxTokens))
+	provider, err := fullTestRegistry().New(t.Context(), cfg, env, options.WithMaxTokens(maxTokens))
 	require.NoError(t, err)
 
 	// Clone with an option that doesn't affect max_tokens (e.g., WithGeneratingTitle)
@@ -200,7 +200,7 @@ func TestCloneWithOptions_OverridesMaxTokens(t *testing.T) {
 		"OPENAI_API_KEY": "test-key",
 	})
 
-	provider, err := New(t.Context(), cfg, env, options.WithMaxTokens(originalMaxTokens))
+	provider, err := fullTestRegistry().New(t.Context(), cfg, env, options.WithMaxTokens(originalMaxTokens))
 	require.NoError(t, err)
 
 	// Clone with an explicit max_tokens override

pkg/model/provider/custom_provider_openai_test.go

@@ -70,7 +70,7 @@ func TestCustomProvider_WithProvidersOption(t *testing.T) {
 	})
 
 	// Create provider with WithProviders option (as teamloader does)
-	provider, err := New(t.Context(), modelCfg, env, options.WithProviders(customProviders))
+	provider, err := fullTestRegistry().New(t.Context(), modelCfg, env, options.WithProviders(customProviders))
 	require.NoError(t, err)
 
 	stream, err := provider.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "Hi"}}, []tools.Tool{})
@@ -152,7 +152,7 @@ func TestCustomProvider_RequestReachesServer(t *testing.T) {
 		customTokenKey: expectedToken,
 	})
 
-	provider, err := New(t.Context(), modelCfg, env)
+	provider, err := fullTestRegistry().New(t.Context(), modelCfg, env)
 	require.NoError(t, err)
 
 	stream, err := provider.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "Hello"}}, []tools.Tool{})
@@ -213,7 +213,7 @@ func TestCustomProvider_ResponsesAPIType(t *testing.T) {
 
 	env := environment.NewMapEnvProvider(map[string]string{"API_KEY": "test"})
 
-	provider, err := New(t.Context(), modelCfg, env)
+	provider, err := fullTestRegistry().New(t.Context(), modelCfg, env)
 	require.NoError(t, err)
 
 	stream, err := provider.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "Hello"}}, []tools.Tool{})
@@ -273,7 +273,7 @@ func TestCustomProvider_ChatCompletionsAPIType(t *testing.T) {
 
 	env := environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "test"})
 
-	provider, err := New(t.Context(), modelCfg, env)
+	provider, err := fullTestRegistry().New(t.Context(), modelCfg, env)
 	require.NoError(t, err)
 
 	stream, err := provider.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "Test"}}, []tools.Tool{})
@@ -308,7 +308,7 @@ func TestCustomProvider_MissingAPIKey(t *testing.T) {
 
 	env := environment.NewNoEnvProvider() // key not set
 
-	_, err := New(t.Context(), modelCfg, env)
+	_, err := fullTestRegistry().New(t.Context(), modelCfg, env)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "MISSING_API_KEY", "Error should mention the missing env var")
 }

pkg/model/provider/factory.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"maps"
 
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/environment"
@@ -14,101 +15,89 @@ import (
 	"github.com/docker/docker-agent/pkg/model/provider/rulebased"
 )
 
-// createRuleBasedRouter creates a rule-based routing provider.
-func createRuleBasedRouter(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return rulebased.NewClient(ctx, cfg, models, env, resolveRoutedModel, opts...)
+type Factory func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error)
+
+type Registry struct {
+	factories map[string]Factory
 }
 
-// resolveRoutedModel is the rulebased.ProviderFactory used by
-// createRuleBasedRouter. It resolves a routing target — which is either a name
-// from the models map or an inline "provider/model" spec — and returns the
-// provider for it. Routing targets cannot themselves have routing rules.
-//
-// Defined as a package-level function (rather than an inline closure) so the
-// recursion-prevention, parse-error and factory-error paths can be unit-tested
-// directly without going through rulebased.NewClient.
-func resolveRoutedModel(
-	ctx context.Context,
-	modelSpec string,
-	models map[string]latest.ModelConfig,
-	env environment.Provider,
-	factoryOpts ...options.Opt,
-) (rulebased.Provider, error) {
-	// Check if modelSpec is a reference to a model in the models map.
+func NewRegistry(factories map[string]Factory) *Registry {
+	copied := make(map[string]Factory, len(factories)+1)
+	copied["dmr"] = dmrFactory
+	maps.Copy(copied, factories)
+	return &Registry{factories: copied}
+}
+
+func (r *Registry) New(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return r.NewWithModels(ctx, cfg, nil, env, opts...)
+}
+
+func (r *Registry) NewWithModels(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	slog.DebugContext(ctx, "Creating model provider", "type", cfg.Provider, "model", cfg.Model)
+	if len(cfg.Routing) > 0 {
+		p, err := r.createRuleBasedRouter(ctx, cfg, models, env, opts...)
+		if err != nil {
+			return nil, err
+		}
+		if setter, ok := p.(interface{ SetProviderRegistry(registry any) }); ok {
+			setter.SetProviderRegistry(r)
+		}
+		return p, nil
+	}
+	return r.createDirectProvider(ctx, cfg, env, opts...)
+}
+
+func (r *Registry) createRuleBasedRouter(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return rulebased.NewClient(ctx, cfg, models, env, r.resolveRoutedModel, opts...)
+}
+
+func (r *Registry) resolveRoutedModel(ctx context.Context, modelSpec string, models map[string]latest.ModelConfig, env environment.Provider, factoryOpts ...options.Opt) (rulebased.Provider, error) {
 	if modelCfg, exists := models[modelSpec]; exists {
-		// Prevent infinite recursion - referenced models cannot have routing rules.
 		if len(modelCfg.Routing) > 0 {
 			return nil, fmt.Errorf("model %q has routing rules and cannot be used as a routing target", modelSpec)
 		}
-		return createDirectProvider(ctx, &modelCfg, env, factoryOpts...)
+		return r.createDirectProvider(ctx, &modelCfg, env, factoryOpts...)
 	}
-
-	// Otherwise, treat as an inline model spec (e.g., "openai/gpt-4o").
 	inlineCfg, parseErr := latest.ParseModelRef(modelSpec)
 	if parseErr != nil {
 		return nil, fmt.Errorf("invalid model spec %q: expected 'provider/model' format or a model reference", modelSpec)
 	}
-	return createDirectProvider(ctx, &inlineCfg, env, factoryOpts...)
+	return r.createDirectProvider(ctx, &inlineCfg, env, factoryOpts...)
 }
 
-// createDirectProvider creates a provider without routing (direct model access).
-func createDirectProvider(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+func (r *Registry) createDirectProvider(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	if r == nil {
+		r = DefaultRegistry()
+	}
 	var globalOptions options.ModelOptions
 	for _, opt := range opts {
 		opt(&globalOptions)
 	}
-
-	// Apply defaults from custom providers (from config) or built-in aliases
 	enhancedCfg := applyProviderDefaults(cfg, globalOptions.Providers())
-
 	providerType := resolveProviderType(enhancedCfg)
-
-	factory, ok := providerFactories[providerType]
+	factory, ok := r.factories[providerType]
 	if !ok {
 		slog.ErrorContext(ctx, "Unknown provider type", "type", providerType)
 		return nil, unknownProviderError(providerType)
 	}
-	return factory(ctx, enhancedCfg, env, opts...)
-}
-
-// unknownProviderError builds the error returned when no factory is registered
-// for providerType. The provider may be genuinely unknown, or it may have been
-// compiled out via a `no_<provider>` build tag (openai/anthropic/google/bedrock
-// are optional). The message hints at the latter so trimmed builds are easy to
-// diagnose.
-func unknownProviderError(providerType string) error {
-	return fmt.Errorf("unknown provider type %q (it may not be compiled into this build)", providerType)
-}
-
-// providerFactory builds a Provider from a fully-resolved ModelConfig.
-// Tests may swap entries in providerFactories to exercise dispatch logic
-// without spinning up real provider clients.
-type providerFactory func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error)
-
-// providerFactories maps a resolved provider type (the value returned by
-// resolveProviderType) to its constructor.
-//
-// Entries are contributed by per-provider files via registerProviderFactory in
-// their init(). The optional providers (openai, anthropic, google,
-// amazon-bedrock) each live behind a `no_<provider>` build tag so embedders can
-// drop them — and their transitive SDK dependencies — from a build. The map is
-// package-private but modifiable; tests must restore the original entries with
-// t.Cleanup.
-var providerFactories = map[string]providerFactory{}
-
-// registerProviderFactory records f under name in providerFactories. It is
-// called from the init() of each per-provider file; duplicate registration of
-// the same name panics to surface wiring mistakes at startup.
-func registerProviderFactory(name string, f providerFactory) {
-	if _, exists := providerFactories[name]; exists {
-		panic("provider: duplicate factory registration for " + name)
+	p, err := factory(ctx, enhancedCfg, env, opts...)
+	if err != nil {
+		return nil, err
 	}
-	providerFactories[name] = f
+	if setter, ok := p.(interface{ SetProviderRegistry(registry any) }); ok {
+		setter.SetProviderRegistry(r)
+	}
+	return p, nil
 }
 
-//nolint:gochecknoinits // Intentional: providers self-register so optional ones can be dropped via build tags.
-func init() {
-	registerProviderFactory("dmr", dmrFactory)
+var defaultFactories map[string]Factory
+
+func DefaultRegistry() *Registry {
+	return NewRegistry(defaultFactories)
+}
+
+func unknownProviderError(providerType string) error {
+	return fmt.Errorf("unknown provider type %q (register it with provider.NewRegistry or use providers.NewDefaultRegistry)", providerType)
 }
 
 func dmrFactory(ctx context.Context, cfg *latest.ModelConfig, _ environment.Provider, opts ...options.Opt) (Provider, error) {

pkg/model/provider/factory_anthropic.go

@@ -1,28 +0,0 @@
-//go:build !js && !docker_agent_no_anthropic
-
-package provider
-
-import (
-	"context"
-
-	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/environment"
-	"github.com/docker/docker-agent/pkg/model/provider/anthropic"
-	"github.com/docker/docker-agent/pkg/model/provider/options"
-)
-
-// Anthropic is an optional provider. Build with `-tags docker_agent_no_anthropic` to drop
-// it and its transitive dependencies (github.com/anthropics/anthropic-sdk-go).
-//
-// Note: the google provider's Vertex Model Garden support imports the
-// anthropic package directly, so `docker_agent_no_anthropic` only removes the dependency
-// when the google provider is also dropped (`-tags docker_agent_no_google`).
-//
-//nolint:gochecknoinits // Intentional: self-registration enables dropping this provider via build tags.
-func init() {
-	registerProviderFactory("anthropic", anthropicFactory)
-}
-
-func anthropicFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return anthropic.NewClient(ctx, cfg, env, opts...)
-}

pkg/model/provider/factory_bedrock.go

@@ -1,25 +0,0 @@
-//go:build !js && !docker_agent_no_bedrock
-
-package provider
-
-import (
-	"context"
-
-	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/environment"
-	"github.com/docker/docker-agent/pkg/model/provider/bedrock"
-	"github.com/docker/docker-agent/pkg/model/provider/options"
-)
-
-// Amazon Bedrock is an optional provider. Build with `-tags docker_agent_no_bedrock` to drop
-// it and its transitive dependencies (the github.com/aws/aws-sdk-go-v2 stack),
-// which is the largest provider-specific dependency tree.
-//
-//nolint:gochecknoinits // Intentional: self-registration enables dropping this provider via build tags.
-func init() {
-	registerProviderFactory("amazon-bedrock", bedrockFactory)
-}
-
-func bedrockFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return bedrock.NewClient(ctx, cfg, env, opts...)
-}

pkg/model/provider/factory_compat_test.go

@@ -0,0 +1,37 @@
+package provider
+
+import (
+	"context"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/model/provider/rulebased"
+	"github.com/docker/docker-agent/pkg/model/provider/vertexai"
+)
+
+type providerFactory = Factory
+
+var providerFactories = DefaultRegistry().factories
+
+func testRegistry() *Registry { return NewRegistry(providerFactories) }
+
+func createDirectProvider(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return testRegistry().createDirectProvider(ctx, cfg, env, opts...)
+}
+
+func resolveRoutedModel(ctx context.Context, modelSpec string, models map[string]latest.ModelConfig, env environment.Provider, factoryOpts ...options.Opt) (rulebased.Provider, error) {
+	return testRegistry().resolveRoutedModel(ctx, modelSpec, models, env, factoryOpts...)
+}
+
+var (
+	geminiClientFactory providerFactory
+	vertexClientFactory providerFactory
+)
+
+func googleFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider) (Provider, error) {
+	if vertexai.IsModelGardenConfig(cfg) {
+		return vertexClientFactory(ctx, cfg, env)
+	}
+	return geminiClientFactory(ctx, cfg, env)
+}

pkg/model/provider/factory_google.go

@@ -1,45 +0,0 @@
-//go:build !js && !docker_agent_no_google
-
-package provider
-
-import (
-	"context"
-
-	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/environment"
-	"github.com/docker/docker-agent/pkg/model/provider/gemini"
-	"github.com/docker/docker-agent/pkg/model/provider/options"
-	"github.com/docker/docker-agent/pkg/model/provider/vertexai"
-)
-
-// Google is an optional provider. Build with `-tags docker_agent_no_google` to drop it and
-// its transitive dependencies (google.golang.org/genai, the Vertex AI / cloud
-// auth stack, and — via Vertex Model Garden — the anthropic and openai SDKs).
-//
-//nolint:gochecknoinits // Intentional: self-registration enables dropping this provider via build tags.
-func init() {
-	registerProviderFactory("google", googleFactory)
-}
-
-func googleFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	// Route non-Gemini models on Vertex AI (Model Garden) through the
-	// vertexai package, which picks the right endpoint per publisher.
-	if vertexai.IsModelGardenConfig(cfg) {
-		return vertexClientFactory(ctx, cfg, env, opts...)
-	}
-	return geminiClientFactory(ctx, cfg, env, opts...)
-}
-
-// geminiClientFactory and vertexClientFactory are the inner constructors used
-// by googleFactory. They are package-level variables (rather than direct
-// references to gemini.NewClient / vertexai.NewClient) so that tests can swap
-// them with fakes via t.Cleanup and assert that googleFactory routes correctly
-// based on vertexai.IsModelGardenConfig — without spinning up real clients.
-var (
-	geminiClientFactory providerFactory = func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-		return gemini.NewClient(ctx, cfg, env, opts...)
-	}
-	vertexClientFactory providerFactory = func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-		return vertexai.NewClient(ctx, cfg, env, opts...)
-	}
-)

pkg/model/provider/factory_js.go

@@ -1,105 +1,93 @@
 //go:build js && wasm
 
-// Package provider's js/wasm factory.
-//
-// The non-js factory.go pulls in every provider — including dmr (os/exec),
-// bedrock and vertexai (cloud SDKs). None of those can be cross-compiled to
-// js/wasm, so this file replaces factory.go under js/wasm with a slim variant
-// that only knows about the providers that work over plain net/http (which
-// the Go runtime maps to fetch in the browser):
-//
-//   - openai / openai_chatcompletions / openai_responses
-//   - anthropic
-//   - google (Gemini API; Vertex AI is unsupported under wasm)
-//
-// Docker Model Runner is unsupported and returns an error. Rule-based routing
-// works under wasm because the rulebased provider no longer depends on bleve.
 package provider
 
 import (
 	"context"
 	"fmt"
 	"log/slog"
+	"maps"
 
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/environment"
-	"github.com/docker/docker-agent/pkg/model/provider/anthropic"
-	"github.com/docker/docker-agent/pkg/model/provider/gemini"
-	"github.com/docker/docker-agent/pkg/model/provider/openai"
 	"github.com/docker/docker-agent/pkg/model/provider/options"
 	"github.com/docker/docker-agent/pkg/model/provider/rulebased"
 )
 
-// createRuleBasedRouter creates a rule-based routing provider.
-func createRuleBasedRouter(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return rulebased.NewClient(ctx, cfg, models, env, resolveRoutedModel, opts...)
+type Factory func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error)
+
+type Registry struct {
+	factories map[string]Factory
 }
 
-// resolveRoutedModel is the rulebased.ProviderFactory used by
-// createRuleBasedRouter. It resolves a routing target — which is either a name
-// from the models map or an inline "provider/model" spec — and returns the
-// provider for it. Routing targets cannot themselves have routing rules.
-func resolveRoutedModel(
-	ctx context.Context,
-	modelSpec string,
-	models map[string]latest.ModelConfig,
-	env environment.Provider,
-	factoryOpts ...options.Opt,
-) (rulebased.Provider, error) {
+func NewRegistry(factories map[string]Factory) *Registry {
+	copied := make(map[string]Factory, len(factories))
+	maps.Copy(copied, factories)
+	return &Registry{factories: copied}
+}
+
+var defaultFactories map[string]Factory
+
+func DefaultRegistry() *Registry { return NewRegistry(defaultFactories) }
+
+func (r *Registry) New(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return r.NewWithModels(ctx, cfg, nil, env, opts...)
+}
+
+func (r *Registry) NewWithModels(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	slog.DebugContext(ctx, "Creating model provider", "type", cfg.Provider, "model", cfg.Model)
+	if len(cfg.Routing) > 0 {
+		p, err := r.createRuleBasedRouter(ctx, cfg, models, env, opts...)
+		if err != nil {
+			return nil, err
+		}
+		if setter, ok := p.(interface{ SetProviderRegistry(registry any) }); ok {
+			setter.SetProviderRegistry(r)
+		}
+		return p, nil
+	}
+	return r.createDirectProvider(ctx, cfg, env, opts...)
+}
+
+func (r *Registry) createRuleBasedRouter(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return rulebased.NewClient(ctx, cfg, models, env, r.resolveRoutedModel, opts...)
+}
+
+func (r *Registry) resolveRoutedModel(ctx context.Context, modelSpec string, models map[string]latest.ModelConfig, env environment.Provider, factoryOpts ...options.Opt) (rulebased.Provider, error) {
 	if modelCfg, exists := models[modelSpec]; exists {
 		if len(modelCfg.Routing) > 0 {
 			return nil, fmt.Errorf("model %q has routing rules and cannot be used as a routing target", modelSpec)
 		}
-		return createDirectProvider(ctx, &modelCfg, env, factoryOpts...)
+		return r.createDirectProvider(ctx, &modelCfg, env, factoryOpts...)
 	}
-
 	inlineCfg, parseErr := latest.ParseModelRef(modelSpec)
 	if parseErr != nil {
 		return nil, fmt.Errorf("invalid model spec %q: expected 'provider/model' format or a model reference", modelSpec)
 	}
-	return createDirectProvider(ctx, &inlineCfg, env, factoryOpts...)
+	return r.createDirectProvider(ctx, &inlineCfg, env, factoryOpts...)
 }
 
-// createDirectProvider mirrors the non-wasm version but only dispatches to
-// providers that are reachable from a browser.
-func createDirectProvider(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+func (r *Registry) createDirectProvider(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	if r == nil {
+		r = DefaultRegistry()
+	}
 	var globalOptions options.ModelOptions
 	for _, opt := range opts {
 		opt(&globalOptions)
 	}
-
 	enhancedCfg := applyProviderDefaults(cfg, globalOptions.Providers())
-
 	providerType := resolveProviderType(enhancedCfg)
-
-	factory, ok := providerFactories[providerType]
+	factory, ok := r.factories[providerType]
 	if !ok {
 		slog.ErrorContext(ctx, "Unknown or unsupported provider type under js/wasm", "type", providerType)
-		return nil, fmt.Errorf("provider type %q is not supported under js/wasm (only openai/anthropic/google work in the browser)", providerType)
+		return nil, fmt.Errorf("provider type %q is not registered", providerType)
 	}
-	return factory(ctx, enhancedCfg, env, opts...)
-}
-
-type providerFactory func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error)
-
-// providerFactories: js/wasm-only subset. dmr (os/exec), amazon-bedrock and
-// vertex AI (cloud SDKs that don't compile to wasm) are deliberately absent.
-var providerFactories = map[string]providerFactory{
-	"openai":                 openaiFactory,
-	"openai_chatcompletions": openaiFactory,
-	"openai_responses":       openaiFactory,
-	"anthropic":              anthropicFactory,
-	"google":                 googleFactory,
-}
-
-func openaiFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return openai.NewClient(ctx, cfg, env, opts...)
-}
-
-func anthropicFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return anthropic.NewClient(ctx, cfg, env, opts...)
-}
-
-func googleFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return gemini.NewClient(ctx, cfg, env, opts...)
+	p, err := factory(ctx, enhancedCfg, env, opts...)
+	if err != nil {
+		return nil, err
+	}
+	if setter, ok := p.(interface{ SetProviderRegistry(registry any) }); ok {
+		setter.SetProviderRegistry(r)
+	}
+	return p, nil
 }

pkg/model/provider/factory_openai.go

@@ -1,26 +0,0 @@
-//go:build !js && !docker_agent_no_openai
-
-package provider
-
-import (
-	"context"
-
-	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/environment"
-	"github.com/docker/docker-agent/pkg/model/provider/openai"
-	"github.com/docker/docker-agent/pkg/model/provider/options"
-)
-
-// OpenAI is an optional provider. Build with `-tags docker_agent_no_openai` to drop it and
-// its transitive dependencies (github.com/openai/openai-go).
-//
-//nolint:gochecknoinits // Intentional: self-registration enables dropping this provider via build tags.
-func init() {
-	registerProviderFactory("openai", openaiFactory)
-	registerProviderFactory("openai_chatcompletions", openaiFactory)
-	registerProviderFactory("openai_responses", openaiFactory)
-}
-
-func openaiFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return openai.NewClient(ctx, cfg, env, opts...)
-}

pkg/model/provider/provider.go

@@ -14,52 +14,16 @@
 //   - factory.go: shared dispatch from a resolved provider type to the
 //     concrete client constructor, plus the always-available dmr provider
 //     and the rule-based router.
-//   - factory_<name>.go: one file per optional provider (openai, anthropic,
-//     google, amazon-bedrock); each registers itself with the dispatch table
-//     and is gated by a build tag (see "Build tags" below).
 //
-// # Build tags
-//
-// The openai, anthropic, google and amazon-bedrock providers are optional.
-// Each lives behind a negative build tag so a project embedding docker-agent
-// can compile a provider out — together with its transitive SDK dependencies —
-// to shrink the binary and dependency graph. All providers are included by
-// default; pass the relevant tag(s) to opt out.
-//
-// Build tags are global to a build, not scoped per module: a tag set by the
-// top-level project applies to every dependency too. The tags are therefore
-// prefixed with "docker_agent_" so an embedding project can use its own build
-// tags (even a plain "no_openai") without accidentally toggling these
-// providers.
-//
-// The available tags are:
-//
-//   - docker_agent_no_openai: drop the OpenAI provider (github.com/openai/openai-go).
-//   - docker_agent_no_anthropic: drop the Anthropic provider
-//     (github.com/anthropics/anthropic-sdk-go). The google provider's Vertex
-//     Model Garden support also imports the anthropic package, so the
-//     dependency is only fully removed when combined with docker_agent_no_google.
-//   - docker_agent_no_google: drop the Google provider (google.golang.org/genai, the
-//     Vertex AI / cloud auth stack, and — via Vertex Model Garden — the
-//     anthropic and openai SDKs). Vertex AI is unsupported either way under
-//     js/wasm.
-//   - docker_agent_no_bedrock: drop the Amazon Bedrock provider
-//     (the github.com/aws/aws-sdk-go-v2 stack), the largest provider-specific
-//     dependency tree.
-//
-// For example, to build without Bedrock and OpenAI:
-//
-//	go build -tags 'docker_agent_no_bedrock docker_agent_no_openai' ...
-//
-// Requesting a model whose provider was compiled out fails at construction
-// time with a clear "not compiled into this build" error rather than at
-// compile time. The dmr provider and the rule-based router are always
-// compiled in (except under js/wasm, which has its own slim factory).
+// Optional SDK-backed providers live outside this package's default import
+// graph. YAML-loading applications that need Docker Agent's full provider set
+// should import pkg/model/provider/providers and pass its explicit registry;
+// embedders that build agents manually can import only the concrete provider
+// packages they use and pass their factories to [NewRegistry].
 package provider
 
 import (
 	"context"
-	"log/slog"
 
 	"github.com/docker/docker-agent/pkg/chat"
 	"github.com/docker/docker-agent/pkg/config/latest"
@@ -117,21 +81,16 @@ type RerankingProvider interface {
 	Rerank(ctx context.Context, query string, documents []types.Document, criteria string) ([]float64, error)
 }
 
-// New creates a new provider from a model config.
-// This is a convenience wrapper for NewWithModels with no models map.
+// New creates a new provider from a model config using the default registry.
+// The default registry only contains providers that the core package can expose
+// without optional SDK dependencies. YAML-loading applications that need all
+// docker-agent providers should use pkg/model/provider/providers.NewDefaultRegistry.
 func New(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	return NewWithModels(ctx, cfg, nil, env, opts...)
+	return DefaultRegistry().New(ctx, cfg, env, opts...)
 }
 
 // NewWithModels creates a new provider from a model config with access to the full models map.
 // The models map is used to resolve model references in routing rules.
 func NewWithModels(ctx context.Context, cfg *latest.ModelConfig, models map[string]latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
-	slog.DebugContext(ctx, "Creating model provider", "type", cfg.Provider, "model", cfg.Model)
-
-	// Check if this model has routing rules - if so, create a rule-based router
-	if len(cfg.Routing) > 0 {
-		return createRuleBasedRouter(ctx, cfg, models, env, opts...)
-	}
-
-	return createDirectProvider(ctx, cfg, env, opts...)
+	return DefaultRegistry().NewWithModels(ctx, cfg, models, env, opts...)
 }

pkg/model/provider/providers/providers.go

@@ -0,0 +1,49 @@
+package providers
+
+import (
+	"context"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider"
+	"github.com/docker/docker-agent/pkg/model/provider/anthropic"
+	"github.com/docker/docker-agent/pkg/model/provider/bedrock"
+	"github.com/docker/docker-agent/pkg/model/provider/gemini"
+	"github.com/docker/docker-agent/pkg/model/provider/openai"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/model/provider/vertexai"
+)
+
+func NewDefaultRegistry() *provider.Registry {
+	return provider.NewRegistry(DefaultFactories())
+}
+
+func DefaultFactories() map[string]provider.Factory {
+	return map[string]provider.Factory{
+		"openai":                 openaiFactory,
+		"openai_chatcompletions": openaiFactory,
+		"openai_responses":       openaiFactory,
+		"anthropic":              anthropicFactory,
+		"google":                 googleFactory,
+		"amazon-bedrock":         bedrockFactory,
+	}
+}
+
+func openaiFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) {
+	return openai.NewClient(ctx, cfg, env, opts...)
+}
+
+func anthropicFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) {
+	return anthropic.NewClient(ctx, cfg, env, opts...)
+}
+
+func googleFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) {
+	if vertexai.IsModelGardenConfig(cfg) {
+		return vertexai.NewClient(ctx, cfg, env, opts...)
+	}
+	return gemini.NewClient(ctx, cfg, env, opts...)
+}
+
+func bedrockFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) {
+	return bedrock.NewClient(ctx, cfg, env, opts...)
+}

pkg/model/provider/test_registry_test.go

@@ -0,0 +1,44 @@
+package provider
+
+import (
+	"context"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/anthropic"
+	"github.com/docker/docker-agent/pkg/model/provider/bedrock"
+	"github.com/docker/docker-agent/pkg/model/provider/gemini"
+	"github.com/docker/docker-agent/pkg/model/provider/openai"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+	"github.com/docker/docker-agent/pkg/model/provider/vertexai"
+)
+
+func fullTestRegistry() *Registry {
+	return NewRegistry(map[string]Factory{
+		"openai":                 openAITestFactory,
+		"openai_chatcompletions": openAITestFactory,
+		"openai_responses":       openAITestFactory,
+		"anthropic":              anthropicTestFactory,
+		"google":                 googleTestFactory,
+		"amazon-bedrock":         bedrockTestFactory,
+	})
+}
+
+func openAITestFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return openai.NewClient(ctx, cfg, env, opts...)
+}
+
+func anthropicTestFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return anthropic.NewClient(ctx, cfg, env, opts...)
+}
+
+func googleTestFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	if vertexai.IsModelGardenConfig(cfg) {
+		return vertexai.NewClient(ctx, cfg, env, opts...)
+	}
+	return gemini.NewClient(ctx, cfg, env, opts...)
+}
+
+func bedrockTestFactory(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+	return bedrock.NewClient(ctx, cfg, env, opts...)
+}

pkg/runtime/model_hook.go

@@ -24,7 +24,9 @@ import (
 // times per turn and the marginal cost of construction is negligible
 // compared to the LLM round-trip itself. If profiling later shows
 // otherwise, a sync.Map keyed on modelSpec would be a drop-in.
-type providerModelClient struct{}
+type providerModelClient struct {
+	registry *provider.Registry
+}
 
 // Ask implements [hooks.ModelClient].
 //
@@ -33,7 +35,7 @@ type providerModelClient struct{}
 // fail-closed semantics deny PreToolUse calls cleanly. The caller
 // (the model handler) is responsible for surfacing the error to the
 // hook's [HandlerResult].
-func (providerModelClient) Ask(
+func (c providerModelClient) Ask(
 	ctx context.Context,
 	modelSpec, system, user string,
 	schema *latest.StructuredOutput,
@@ -47,7 +49,11 @@ func (providerModelClient) Ask(
 	if schema != nil {
 		opts = append(opts, options.WithStructuredOutput(schema))
 	}
-	p, err := provider.New(ctx, &cfg, environment.NewDefaultProvider(), opts...)
+	registry := c.registry
+	if registry == nil {
+		registry = provider.DefaultRegistry()
+	}
+	p, err := registry.New(ctx, &cfg, environment.NewDefaultProvider(), opts...)
 	if err != nil {
 		return "", fmt.Errorf("create provider: %w", err)
 	}
@@ -80,6 +86,6 @@ func (providerModelClient) Ask(
 // registerModelHook installs the [hooks.HookTypeModel] factory on r
 // using the runtime's default [hooks.ModelClient]. It is called once
 // from [NewLocalRuntime] alongside the builtins.
-func registerModelHook(r *hooks.Registry) {
-	r.Register(hooks.HookTypeModel, hooks.NewModelFactory(providerModelClient{}))
+func registerModelHook(r *hooks.Registry, registry *provider.Registry) {
+	r.Register(hooks.HookTypeModel, hooks.NewModelFactory(providerModelClient{registry: registry}))
 }

pkg/runtime/model_switcher.go

@@ -161,6 +161,8 @@ type ModelSwitcherConfig struct {
 	ModelsGateway string
 	// EnvProvider provides access to environment variables
 	EnvProvider environment.Provider
+	// ProviderRegistry instantiates providers for runtime model switching.
+	ProviderRegistry *provider.Registry
 	// AgentDefaultModels maps agent names to their configured default model references
 	AgentDefaultModels map[string]string
 }
@@ -737,7 +739,14 @@ func (r *LocalRuntime) createProviderFromConfig(ctx context.Context, cfg *latest
 		}
 	}
 
-	return provider.NewWithModels(ctx,
+	registry := r.modelSwitcherCfg.ProviderRegistry
+	if registry == nil {
+		registry = r.providerRegistry
+	}
+	if registry == nil {
+		registry = provider.DefaultRegistry()
+	}
+	return registry.NewWithModels(ctx,
 		cfg,
 		r.modelSwitcherCfg.Models,
 		r.modelSwitcherCfg.EnvProvider,
@@ -748,6 +757,14 @@ func (r *LocalRuntime) createProviderFromConfig(ctx context.Context, cfg *latest
 // WithModelSwitcherConfig sets the model switcher configuration for the runtime.
 func WithModelSwitcherConfig(cfg *ModelSwitcherConfig) Opt {
 	return func(r *LocalRuntime) {
+		if cfg != nil && cfg.ProviderRegistry == nil {
+			cfgCopy := *cfg
+			cfgCopy.ProviderRegistry = r.providerRegistry
+			cfg = &cfgCopy
+		}
+		if cfg != nil && cfg.ProviderRegistry != nil {
+			r.providerRegistry = cfg.ProviderRegistry
+		}
 		r.modelSwitcherCfg = cfg
 	}
 }

pkg/runtime/model_switcher_test.go

@@ -13,7 +13,9 @@ import (
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/effort"
 	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider"
 	"github.com/docker/docker-agent/pkg/model/provider/base"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
 	"github.com/docker/docker-agent/pkg/modelsdev"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/tools"
@@ -44,6 +46,20 @@ func newConfigProvider(cfg latest.ModelConfig) *configProvider {
 	return &configProvider{Config: base.Config{ModelConfig: cfg}}
 }
 
+func testProviderRegistry() *provider.Registry {
+	factory := func(_ context.Context, cfg *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (provider.Provider, error) {
+		return newConfigProvider(*cfg), nil
+	}
+	return provider.NewRegistry(map[string]provider.Factory{
+		"openai":                 factory,
+		"openai_chatcompletions": factory,
+		"openai_responses":       factory,
+		"anthropic":              factory,
+		"google":                 factory,
+		"amazon-bedrock":         factory,
+	})
+}
+
 func TestCurrentThinkingLevel(t *testing.T) {
 	t.Parallel()
 
@@ -177,7 +193,8 @@ func TestCycleAgentThinkingLevel_AdvancesAndOverrides(t *testing.T) {
 	r := &LocalRuntime{
 		team: team.New(team.WithAgents(root)),
 		modelSwitcherCfg: &ModelSwitcherConfig{
-			EnvProvider: environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "sk-test"}),
+			ProviderRegistry: testProviderRegistry(),
+			EnvProvider:      environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "sk-test"}),
 		},
 	}
 
@@ -266,7 +283,8 @@ func TestCycleAgentThinkingLevel_PerModelTopTier(t *testing.T) {
 			r := &LocalRuntime{
 				team: team.New(team.WithAgents(root)),
 				modelSwitcherCfg: &ModelSwitcherConfig{
-					EnvProvider: environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "sk-test"}),
+					ProviderRegistry: testProviderRegistry(),
+					EnvProvider:      environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "sk-test"}),
 				},
 			}
 
@@ -483,8 +501,9 @@ func TestGetAvailableProviders(t *testing.T) {
 
 			r := &LocalRuntime{
 				modelSwitcherCfg: &ModelSwitcherConfig{
-					EnvProvider:   environment.NewMapEnvProvider(tt.envVars),
-					ModelsGateway: tt.modelsGateway,
+					ProviderRegistry: testProviderRegistry(),
+					EnvProvider:      environment.NewMapEnvProvider(tt.envVars),
+					ModelsGateway:    tt.modelsGateway,
 				},
 			}
 
@@ -516,7 +535,8 @@ func TestGetAvailableProviders_AnthropicWIF(t *testing.T) {
 		t.Parallel()
 		r := &LocalRuntime{
 			modelSwitcherCfg: &ModelSwitcherConfig{
-				EnvProvider: environment.NewMapEnvProvider(nil),
+				ProviderRegistry: testProviderRegistry(),
+				EnvProvider:      environment.NewMapEnvProvider(nil),
 				Models: map[string]latest.ModelConfig{
 					"claude": {Provider: "anthropic", Model: "claude-x", Auth: wifAuth},
 				},
@@ -530,7 +550,8 @@ func TestGetAvailableProviders_AnthropicWIF(t *testing.T) {
 		t.Parallel()
 		r := &LocalRuntime{
 			modelSwitcherCfg: &ModelSwitcherConfig{
-				EnvProvider: environment.NewMapEnvProvider(nil),
+				ProviderRegistry: testProviderRegistry(),
+				EnvProvider:      environment.NewMapEnvProvider(nil),
 				Providers: map[string]latest.ProviderConfig{
 					"claude": {Provider: "anthropic", Auth: wifAuth},
 				},
@@ -547,7 +568,8 @@ func TestGetAvailableProviders_AnthropicWIF(t *testing.T) {
 		t.Parallel()
 		r := &LocalRuntime{
 			modelSwitcherCfg: &ModelSwitcherConfig{
-				EnvProvider: environment.NewMapEnvProvider(nil),
+				ProviderRegistry: testProviderRegistry(),
+				EnvProvider:      environment.NewMapEnvProvider(nil),
 				Models: map[string]latest.ModelConfig{
 					"claude": {Provider: "anthropic", Model: "claude-x"},
 				},
@@ -613,6 +635,7 @@ func TestBuildCatalogChoices(t *testing.T) {
 	r := &LocalRuntime{
 		modelsStore: &mockCatalogStore{db: db},
 		modelSwitcherCfg: &ModelSwitcherConfig{
+			ProviderRegistry: testProviderRegistry(),
 			EnvProvider: environment.NewMapEnvProvider(map[string]string{
 				"OPENAI_API_KEY":    "sk-test",
 				"ANTHROPIC_API_KEY": "sk-ant-test",
@@ -673,6 +696,7 @@ func TestBuildCatalogChoicesWithDuplicates(t *testing.T) {
 	r := &LocalRuntime{
 		modelsStore: &mockCatalogStore{db: db},
 		modelSwitcherCfg: &ModelSwitcherConfig{
+			ProviderRegistry: testProviderRegistry(),
 			EnvProvider: environment.NewMapEnvProvider(map[string]string{
 				"OPENAI_API_KEY": "sk-test",
 			}),
@@ -696,6 +720,7 @@ func TestResolveModelRef_RejectsAlloyConfig(t *testing.T) {
 
 	r := &LocalRuntime{
 		modelSwitcherCfg: &ModelSwitcherConfig{
+			ProviderRegistry: testProviderRegistry(),
 			Models: map[string]latest.ModelConfig{
 				// Alloy config: no provider, comma-separated models
 				"alloy_model": {Model: "openai/gpt-4o,anthropic/claude-sonnet-4-0"},
@@ -723,7 +748,8 @@ func TestResolveModelRef_InvalidFormat(t *testing.T) {
 
 	r := &LocalRuntime{
 		modelSwitcherCfg: &ModelSwitcherConfig{
-			Models: map[string]latest.ModelConfig{},
+			ProviderRegistry: testProviderRegistry(),
+			Models:           map[string]latest.ModelConfig{},
 		},
 	}
 

pkg/runtime/runtime.go

@@ -22,6 +22,7 @@ import (
 	"github.com/docker/docker-agent/pkg/hooks"
 	"github.com/docker/docker-agent/pkg/hooks/builtins"
 	"github.com/docker/docker-agent/pkg/httpclient"
+	"github.com/docker/docker-agent/pkg/model/provider"
 	"github.com/docker/docker-agent/pkg/modelsdev"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/sessiontitle"
@@ -209,6 +210,7 @@ type LocalRuntime struct {
 	workingDir                string   // Working directory for hooks execution
 	env                       []string // Environment variables for hooks execution
 	modelSwitcherCfg          *ModelSwitcherConfig
+	providerRegistry          *provider.Registry
 	gatewayModels             gatewayModelsCache
 
 	// hooksRegistry is the runtime-private hooks.Registry used to build
@@ -372,6 +374,14 @@ func WithSessionCompaction(sessionCompaction bool) Opt {
 	}
 }
 
+func WithProviderRegistry(registry *provider.Registry) Opt {
+	return func(r *LocalRuntime) {
+		if registry != nil {
+			r.providerRegistry = registry
+		}
+	}
+}
+
 func WithModelStore(store ModelStore) Opt {
 	return func(r *LocalRuntime) {
 		r.modelsStore = store
@@ -536,6 +546,7 @@ func NewLocalRuntime(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
 		fallback:               newFallbackExecutor(),
 		now:                    time.Now,
 		telemetry:              defaultTelemetry{},
+		providerRegistry:       provider.DefaultRegistry(),
 		maxOverflowCompactions: defaultMaxOverflowCompactions,
 		toolListTimeout:        defaultToolListTimeout,
 	}
@@ -575,7 +586,7 @@ func NewLocalRuntime(agents *team.Team, opts ...Opt) (*LocalRuntime, error) {
 	if err := builtins.Register(r.hooksRegistry); err != nil {
 		return nil, fmt.Errorf("register builtin hooks: %w", err)
 	}
-	registerModelHook(r.hooksRegistry)
+	registerModelHook(r.hooksRegistry, r.providerRegistry)
 
 	// cache_response is registered here (not in pkg/hooks/builtins)
 	// because it needs to capture the runtime to resolve the agent

pkg/server/session_manager.go

@@ -23,6 +23,7 @@ import (
 	"github.com/docker/docker-agent/pkg/sessiontitle"
 	"github.com/docker/docker-agent/pkg/team"
 	"github.com/docker/docker-agent/pkg/teamloader"
+	loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
@@ -829,6 +830,7 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S
 		Providers:          loadResult.Providers,
 		ModelsGateway:      rc.ModelsGateway,
 		EnvProvider:        rc.EnvProvider(),
+		ProviderRegistry:   loadResult.ProviderRegistry,
 		AgentDefaultModels: loadResult.AgentDefaultModels,
 	}
 
@@ -867,7 +869,7 @@ func (sm *SessionManager) loadTeam(ctx context.Context, agentFilename string, ru
 		return nil, fmt.Errorf("agent not found: %s", agentFilename)
 	}
 
-	return teamloader.Load(ctx, agentSource, runConfig)
+	return teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 }
 
 // loadTeamWithConfig is like loadTeam but also returns the loaded model and
@@ -878,7 +880,7 @@ func (sm *SessionManager) loadTeamWithConfig(ctx context.Context, agentFilename
 		return nil, fmt.Errorf("agent not found: %s", agentFilename)
 	}
 
-	return teamloader.LoadWithConfig(ctx, agentSource, runConfig)
+	return teamloader.LoadWithConfig(ctx, agentSource, runConfig, loaderdefaults.Opts()...)
 }
 
 // applyRunModelOverride applies modelRef as the per-agent model override

pkg/teamloader/defaults/defaults.go

@@ -0,0 +1,14 @@
+package defaults
+
+import (
+	"github.com/docker/docker-agent/pkg/model/provider/providers"
+	"github.com/docker/docker-agent/pkg/teamloader"
+	"github.com/docker/docker-agent/pkg/teamloader/toolsets"
+)
+
+func Opts() []teamloader.Opt {
+	return []teamloader.Opt{
+		teamloader.WithToolsetRegistry(toolsets.NewDefaultToolsetRegistry()),
+		teamloader.WithProviderRegistry(providers.NewDefaultRegistry()),
+	}
+}

pkg/teamloader/registry.go

@@ -9,118 +9,30 @@ import (
 	"github.com/docker/docker-agent/pkg/config"
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/tools"
-	"github.com/docker/docker-agent/pkg/tools/a2a"
-	agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent"
-	"github.com/docker/docker-agent/pkg/tools/builtin/api"
-	"github.com/docker/docker-agent/pkg/tools/builtin/fetch"
-	"github.com/docker/docker-agent/pkg/tools/builtin/filesystem"
-	"github.com/docker/docker-agent/pkg/tools/builtin/lsp"
-	"github.com/docker/docker-agent/pkg/tools/builtin/mcpcatalog"
-	"github.com/docker/docker-agent/pkg/tools/builtin/memory"
-	"github.com/docker/docker-agent/pkg/tools/builtin/modelpicker"
-	"github.com/docker/docker-agent/pkg/tools/builtin/openapi"
-	"github.com/docker/docker-agent/pkg/tools/builtin/shell"
-	"github.com/docker/docker-agent/pkg/tools/builtin/tasks"
-	"github.com/docker/docker-agent/pkg/tools/builtin/think"
-	"github.com/docker/docker-agent/pkg/tools/builtin/todo"
-	"github.com/docker/docker-agent/pkg/tools/builtin/userprompt"
-	"github.com/docker/docker-agent/pkg/tools/mcp"
 )
 
 // ToolsetCreator is a function that creates a toolset based on the provided configuration.
 // configName identifies the agent config file (e.g. "memory_agent" from "memory_agent.yaml").
 type ToolsetCreator func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, configName string) (tools.ToolSet, error)
 
-// extraCreators holds toolset creators contributed by packages that opt in via
-// RegisterToolsetCreator. This keeps optional toolsets (e.g. "rag", which pulls
-// in a cgo tree-sitter dependency) out of teamloader's static import graph:
-// they are linked only when the owning package is blank-imported by the binary.
-var extraCreators = map[string]ToolsetCreator{}
-
-// RegisterToolsetCreator registers a creator for the given toolset type, to be
-// included in every registry returned by NewDefaultToolsetRegistry. Packages
-// call it from an init() so that a blank import is enough to make their toolset
-// available; binaries that don't import the package don't pay for its
-// dependencies. It panics on a duplicate registration to surface wiring bugs at
-// startup. Not safe for concurrent use; call only from init().
-func RegisterToolsetCreator(toolsetType string, creator ToolsetCreator) {
-	if _, exists := extraCreators[toolsetType]; exists {
-		panic(fmt.Sprintf("teamloader: toolset creator %q already registered", toolsetType))
-	}
-	extraCreators[toolsetType] = creator
-}
-
 // ToolsetRegistry manages the registration of toolset creators by type.
 type ToolsetRegistry interface {
 	CreateTool(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, agentName string) (tools.ToolSet, error)
 }
 
+func NewToolsetRegistry(creators map[string]ToolsetCreator) ToolsetRegistry {
+	copied := make(map[string]ToolsetCreator, len(creators))
+	maps.Copy(copied, creators)
+	return &toolsetRegistry{creators: copied}
+}
+
+var defaultToolsetCreators map[string]ToolsetCreator
+
+// NewDefaultToolsetRegistry returns the package-level default registry. It is
+// empty unless an application explicitly wires creators into this package; YAML
+// applications should use pkg/teamloader/toolsets.NewDefaultToolsetRegistry.
 func NewDefaultToolsetRegistry() ToolsetRegistry {
-	r := &toolsetRegistry{
-		creators: map[string]ToolsetCreator{
-			"todo": func(_ context.Context, toolset latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return todo.CreateToolSet(toolset)
-			},
-			"tasks": func(_ context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return tasks.CreateToolSet(toolset, parentDir, runConfig)
-			},
-			"memory": func(_ context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, configName string) (tools.ToolSet, error) {
-				return memory.CreateToolSet(toolset, parentDir, runConfig, configName)
-			},
-			"think": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return think.CreateToolSet()
-			},
-			"shell": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return shell.CreateToolSet(ctx, toolset, runConfig)
-			},
-			"script": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return shell.CreateScriptToolSet(ctx, toolset, runConfig)
-			},
-			"filesystem": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return filesystem.CreateToolSet(toolset, runConfig)
-			},
-			"fetch": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return fetch.CreateToolSet(toolset, runConfig)
-			},
-			"mcp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return mcp.CreateToolSet(ctx, toolset, runConfig)
-			},
-			"mcp_catalog": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				var opts []mcpcatalog.Option
-				if len(toolset.AllowedServers) > 0 {
-					opts = append(opts, mcpcatalog.WithAllowedServers(toolset.AllowedServers))
-				}
-				if len(toolset.BlockedServers) > 0 {
-					opts = append(opts, mcpcatalog.WithBlockedServers(toolset.BlockedServers))
-				}
-				return mcpcatalog.New(opts...), nil
-			},
-			"api": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return api.CreateToolSet(toolset, runConfig)
-			},
-			"a2a": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return a2a.CreateToolSet(ctx, toolset, runConfig)
-			},
-			"lsp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return lsp.CreateToolSet(ctx, toolset, runConfig)
-			},
-			"user_prompt": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return userprompt.CreateToolSet()
-			},
-			"openapi": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return openapi.CreateToolSet(ctx, toolset, runConfig)
-			},
-			"model_picker": func(_ context.Context, toolset latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return modelpicker.CreateToolSet(toolset)
-			},
-			"background_agents": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-				return agenttool.CreateToolSet()
-			},
-		},
-	}
-	// Merge in creators contributed via RegisterToolsetCreator (e.g. "rag").
-	maps.Copy(r.creators, extraCreators)
-	return r
+	return NewToolsetRegistry(defaultToolsetCreators)
 }
 
 // toolsetRegistry manages the registration of toolset creators by type.
@@ -129,13 +41,6 @@ type toolsetRegistry struct {
 }
 
 // CreateTool creates a toolset using the registered creator for the given type.
-//
-// Every successful toolset is decorated with tools.WithName so status
-// surfaces (the /tools dialog, error messages, …) always have a stable
-// user-facing label. The decoration is a no-op for toolsets that
-// already advertise a non-empty Name(): it only fills the gap left by
-// built-in toolsets that don't take a `name:` field in YAML, replacing
-// the previous fallback to fmt.Sprintf("%T", ts).
 func (r *toolsetRegistry) CreateTool(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, agentName string) (tools.ToolSet, error) {
 	creator, ok := r.creators[toolset.Type]
 	if !ok {

pkg/teamloader/registry_test.go

@@ -1,6 +1,7 @@
 package teamloader
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 	"testing"
@@ -12,16 +13,35 @@ import (
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/environment"
 	"github.com/docker/docker-agent/pkg/tools"
+	"github.com/docker/docker-agent/pkg/tools/builtin/fetch"
 	"github.com/docker/docker-agent/pkg/tools/builtin/lsp"
+	"github.com/docker/docker-agent/pkg/tools/builtin/shell"
 	mcptool "github.com/docker/docker-agent/pkg/tools/mcp"
 )
 
+func testToolsetRegistry() ToolsetRegistry {
+	return NewToolsetRegistry(map[string]ToolsetCreator{
+		"shell": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return shell.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"mcp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return mcptool.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"fetch": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return fetch.CreateToolSet(toolset, runConfig)
+		},
+		"lsp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return lsp.CreateToolSet(ctx, toolset, runConfig)
+		},
+	})
+}
+
 func TestCreateShellTool(t *testing.T) {
 	toolset := latest.Toolset{
 		Type: "shell",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
@@ -41,7 +61,7 @@ func TestCreateMCPTool_CommandNotFound_CreatesToolsetAnyway(t *testing.T) {
 		Command: "./bin/nonexistent-mcp-server",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
@@ -62,7 +82,7 @@ func TestCreateMCPTool_BareCommandNotFound_CreatesToolsetAnyway(t *testing.T) {
 		Command: "some-nonexistent-mcp-binary",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
@@ -88,7 +108,7 @@ func TestCreateMCPTool_WorkingDir_ReachesSubprocess(t *testing.T) {
 		WorkingDir: customDir,
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: agentDir},
 		EnvProviderForTests: environment.NewOsEnvProvider(),
@@ -120,7 +140,7 @@ func TestCreateMCPTool_RelativeWorkingDir_ResolvedAgainstAgentDir(t *testing.T)
 		WorkingDir: "tools/mcp",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: agentDir},
 		EnvProviderForTests: environment.NewOsEnvProvider(),
@@ -151,7 +171,7 @@ func TestCreateMCPTool_NonexistentWorkingDir_ReturnsError(t *testing.T) {
 		WorkingDir: nonExistent,
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
 		EnvProviderForTests: environment.NewOsEnvProvider(),
@@ -176,7 +196,7 @@ func TestCreateLSPTool_WorkingDir_ReachesHandler(t *testing.T) {
 		WorkingDir: customDir,
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: agentDir},
 		EnvProviderForTests: environment.NewOsEnvProvider(),
@@ -202,7 +222,7 @@ func TestCreateMCPTool_RefRemote_WorkingDir_ReturnsError(t *testing.T) {
 		WorkingDir: "./workspace",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
 		EnvProviderForTests: environment.NewOsEnvProvider(),
@@ -224,7 +244,7 @@ func TestCreateMCPTool_RefRemote_NoWorkingDir_Succeeds(t *testing.T) {
 		Ref:  "docker:remote-server",
 	}
 
-	registry := NewDefaultToolsetRegistry()
+	registry := testToolsetRegistry()
 	runConfig := &config.RuntimeConfig{
 		Config:              config.Config{WorkingDir: t.TempDir()},
 		EnvProviderForTests: environment.NewOsEnvProvider(),

pkg/teamloader/teamloader.go

@@ -37,9 +37,10 @@ import (
 var defaultMaxTokens int64 = 32000
 
 type loadOptions struct {
-	modelOverrides  []string
-	promptFiles     []string
-	toolsetRegistry ToolsetRegistry
+	modelOverrides   []string
+	promptFiles      []string
+	toolsetRegistry  ToolsetRegistry
+	providerRegistry *provider.Registry
 }
 
 type Opt func(*loadOptions) error
@@ -60,7 +61,7 @@ func WithPromptFiles(files []string) Opt {
 	}
 }
 
-// WithToolsetRegistry allows using a custom toolset registry instead of the default
+// WithToolsetRegistry allows using a custom toolset registry instead of the default.
 func WithToolsetRegistry(registry ToolsetRegistry) Opt {
 	return func(opts *loadOptions) error {
 		opts.toolsetRegistry = registry
@@ -68,12 +69,24 @@ func WithToolsetRegistry(registry ToolsetRegistry) Opt {
 	}
 }
 
+// WithProviderRegistry allows using a custom model provider registry instead of the default.
+func WithProviderRegistry(registry *provider.Registry) Opt {
+	return func(opts *loadOptions) error {
+		if registry != nil {
+			opts.providerRegistry = registry
+		}
+		return nil
+	}
+}
+
 // LoadResult contains the result of loading an agent team, including
 // the team and configuration needed for runtime model switching.
 type LoadResult struct {
 	Team      *team.Team
 	Models    map[string]latest.ModelConfig
 	Providers map[string]latest.ProviderConfig
+	// ProviderRegistry is the registry used to instantiate model providers for this load.
+	ProviderRegistry *provider.Registry
 	// AgentDefaultModels maps agent names to their configured default model references
 	AgentDefaultModels map[string]string
 }
@@ -92,6 +105,7 @@ func Load(ctx context.Context, agentSource config.Source, runConfig *config.Runt
 func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *config.RuntimeConfig, opts ...Opt) (*LoadResult, error) {
 	var loadOpts loadOptions
 	loadOpts.toolsetRegistry = NewDefaultToolsetRegistry()
+	loadOpts.providerRegistry = provider.DefaultRegistry()
 
 	for _, o := range opts {
 		if err := o(&loadOpts); err != nil {
@@ -200,7 +214,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 			}
 			opts = append(opts, agent.WithHarness(&harnessCfg))
 		} else {
-			models, err := getModelsForAgent(ctx, cfg, &agentConfig, autoModel, runConfig)
+			models, err := getModelsForAgent(ctx, cfg, &agentConfig, autoModel, runConfig, loadOpts.providerRegistry)
 			if err != nil {
 				// Return auto model fallback errors and DMR not installed errors directly
 				// without wrapping to provide cleaner messages
@@ -216,7 +230,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 			// Load fallback models if configured
 			fallbackModelRefs := agentConfig.GetFallbackModels()
 			if len(fallbackModelRefs) > 0 {
-				fallbackModels, err := getFallbackModelsForAgent(ctx, cfg, &agentConfig, runConfig)
+				fallbackModels, err := getFallbackModelsForAgent(ctx, cfg, &agentConfig, runConfig, loadOpts.providerRegistry)
 				if err != nil {
 					return nil, fmt.Errorf("failed to get fallback models: %w", err)
 				}
@@ -230,7 +244,7 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 			}
 
 			// A model may delegate session-title generation to another model.
-			titleModel, err := getTitleModelForAgent(ctx, cfg, &agentConfig, runConfig)
+			titleModel, err := getTitleModelForAgent(ctx, cfg, &agentConfig, runConfig, loadOpts.providerRegistry)
 			if err != nil {
 				return nil, fmt.Errorf("failed to get title model: %w", err)
 			}
@@ -321,11 +335,12 @@ func LoadWithConfig(ctx context.Context, agentSource config.Source, runConfig *c
 		),
 		Models:             cfg.Models,
 		Providers:          cfg.Providers,
+		ProviderRegistry:   loadOpts.providerRegistry,
 		AgentDefaultModels: agentDefaultModels,
 	}, nil
 }
 
-func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, autoModelFn func() latest.ModelConfig, runConfig *config.RuntimeConfig) ([]provider.Provider, error) {
+func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, autoModelFn func() latest.ModelConfig, runConfig *config.RuntimeConfig, providerRegistry *provider.Registry) ([]provider.Provider, error) {
 	var models []provider.Provider
 
 	// Obtain the singleton store once, outside the loop.
@@ -368,7 +383,7 @@ func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentC
 		}
 
 		// Pass the full models map for routing rules to resolve model references
-		model, err := provider.NewWithModels(ctx,
+		model, err := providerRegistry.NewWithModels(ctx,
 			&modelCfg,
 			cfg.Models,
 			runConfig.EnvProvider(),
@@ -389,7 +404,7 @@ func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentC
 
 // getFallbackModelsForAgent returns fallback providers for an agent based on its fallback configuration.
 // It uses the same resolution logic as primary models (named model, inline provider/model format).
-func getFallbackModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, runConfig *config.RuntimeConfig) ([]provider.Provider, error) {
+func getFallbackModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, runConfig *config.RuntimeConfig, providerRegistry *provider.Registry) ([]provider.Provider, error) {
 	var fallbackModels []provider.Provider
 
 	// Obtain the singleton store once, outside the loop.
@@ -431,7 +446,7 @@ func getFallbackModelsForAgent(ctx context.Context, cfg *latest.Config, a *lates
 		}
 
 		// Pass the full models map for routing rules to resolve model references
-		model, err := provider.NewWithModels(ctx,
+		model, err := providerRegistry.NewWithModels(ctx,
 			&modelCfg,
 			cfg.Models,
 			runConfig.EnvProvider(),
@@ -449,7 +464,7 @@ func getFallbackModelsForAgent(ctx context.Context, cfg *latest.Config, a *lates
 // getTitleModelForAgent resolves the dedicated title-generation model for an
 // agent, if any. It returns the model named by the `title_model` field of the
 // first of the agent's configured models that sets it, or nil when none do.
-func getTitleModelForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, runConfig *config.RuntimeConfig) (provider.Provider, error) {
+func getTitleModelForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentConfig, runConfig *config.RuntimeConfig, providerRegistry *provider.Registry) (provider.Provider, error) {
 	var titleRef string
 	for name := range strings.SplitSeq(a.Model, ",") {
 		if modelCfg, ok := cfg.Models[name]; ok && modelCfg.TitleModel != "" {
@@ -495,7 +510,7 @@ func getTitleModelForAgent(ctx context.Context, cfg *latest.Config, a *latest.Ag
 		opts = append(opts, options.WithModelsDevStore(modelsStore))
 	}
 
-	model, err := provider.NewWithModels(ctx, &modelCfg, cfg.Models, runConfig.EnvProvider(), opts...)
+	model, err := providerRegistry.NewWithModels(ctx, &modelCfg, cfg.Models, runConfig.EnvProvider(), opts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create title model '%s': %w", titleRef, err)
 	}
@@ -743,6 +758,10 @@ func loadExternalAgent(ctx context.Context, ref string, runConfig *config.Runtim
 		opts = append(opts, WithToolsetRegistry(loadOpts.toolsetRegistry))
 	}
 
+	if loadOpts.providerRegistry != nil {
+		opts = append(opts, WithProviderRegistry(loadOpts.providerRegistry))
+	}
+
 	result, err := Load(contextWithExternalDepth(ctx, depth+1), source, runConfig, opts...)
 	if err != nil {
 		return nil, err

pkg/teamloader/teamloader_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/docker/docker-agent/pkg/environment"
 	"github.com/docker/docker-agent/pkg/js"
 	"github.com/docker/docker-agent/pkg/model/provider/dmr"
+	providerdefaults "github.com/docker/docker-agent/pkg/model/provider/providers"
 	"github.com/docker/docker-agent/pkg/tools"
 )
 
@@ -27,6 +28,13 @@ var skipExamples = map[string]string{
 	"pr-reviewer-bedrock.yaml": "requires AWS profile configuration",
 }
 
+func withTestProviderRegistry(opts ...Opt) []Opt {
+	return append([]Opt{
+		WithProviderRegistry(providerdefaults.NewDefaultRegistry()),
+		WithToolsetRegistry(testToolsetRegistry()),
+	}, opts...)
+}
+
 func collectExamples(t *testing.T) []string {
 	t.Helper()
 
@@ -107,7 +115,7 @@ func TestLoadExamples(t *testing.T) {
 			runConfig := &config.RuntimeConfig{}
 			runConfig.WorkingDir = t.TempDir()
 
-			teams, err := Load(t.Context(), agentSource, runConfig)
+			teams, err := Load(t.Context(), agentSource, runConfig, withTestProviderRegistry()...)
 			if errors.Is(err, dmr.ErrNotInstalled) {
 				t.Skipf("Skipping %s: Docker Model Runner not installed", agentFilename)
 			}
@@ -154,7 +162,7 @@ func TestLoadDefaultAgent(t *testing.T) {
 		}),
 	}
 
-	teams, err := Load(t.Context(), agentSource, runConfig)
+	teams, err := Load(t.Context(), agentSource, runConfig, withTestProviderRegistry()...)
 	require.NoError(t, err)
 	require.NotEmpty(t, teams)
 }
@@ -189,7 +197,7 @@ func TestOverrideModel(t *testing.T) {
 			agentSource, err := config.Resolve("testdata/basic.yaml", nil)
 			require.NoError(t, err)
 
-			team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, WithModelOverrides(test.overrides))
+			team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, withTestProviderRegistry(WithModelOverrides(test.overrides))...)
 			if test.expectedErr != "" {
 				require.Contains(t, err.Error(), test.expectedErr)
 			} else {
@@ -221,7 +229,7 @@ agents:
     instruction: test
 `)
 
-		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{})
+		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...)
 		require.NoError(t, err)
 
 		root, err := team.Agent("root")
@@ -250,7 +258,7 @@ agents:
     instruction: test
 `)
 
-		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{})
+		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...)
 		require.NoError(t, err)
 
 		root, err := team.Agent("root")
@@ -267,7 +275,7 @@ agents:
     instruction: test
 `)
 
-		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{})
+		team, err := Load(t.Context(), config.NewBytesSource("title.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...)
 		require.NoError(t, err)
 
 		root, err := team.Agent("root")
@@ -297,7 +305,7 @@ func TestLoadHarnessAgentWithoutModel(t *testing.T) {
       type: codex
 `)
 
-	team, err := Load(t.Context(), config.NewBytesSource("harness.yaml", data), &config.RuntimeConfig{})
+	team, err := Load(t.Context(), config.NewBytesSource("harness.yaml", data), &config.RuntimeConfig{}, withTestProviderRegistry()...)
 	require.NoError(t, err)
 
 	coder, err := team.Agent("coder")
@@ -313,7 +321,7 @@ func TestToolsetInstructions(t *testing.T) {
 	agentSource, err := config.Resolve("testdata/tool-instruction.yaml", nil)
 	require.NoError(t, err)
 
-	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{})
+	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, withTestProviderRegistry()...)
 	require.NoError(t, err)
 
 	agent, err := team.Agent("root")
@@ -337,7 +345,7 @@ func TestInstructionExpansion(t *testing.T) {
 	agentSource, err := config.Resolve("testdata/instruction-expansion.yaml", nil)
 	require.NoError(t, err)
 
-	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{})
+	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, withTestProviderRegistry()...)
 	require.NoError(t, err)
 
 	rootAgent, err := team.Agent("root")
@@ -376,7 +384,7 @@ func TestAutoModelFallbackError(t *testing.T) {
 		EnvProviderForTests: &noEnvProvider{},
 	}
 
-	_, err = Load(t.Context(), agentSource, runConfig)
+	_, err = Load(t.Context(), agentSource, runConfig, withTestProviderRegistry()...)
 	require.Error(t, err)
 
 	var autoErr *config.AutoModelFallbackError
@@ -446,7 +454,7 @@ func TestWithPromptFiles(t *testing.T) {
 				opts = append(opts, WithPromptFiles(tt.cliPromptFiles))
 			}
 
-			team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, opts...)
+			team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{}, withTestProviderRegistry(opts...)...)
 			require.NoError(t, err)
 
 			rootAgent, err := team.Agent("root")
@@ -478,7 +486,7 @@ agents:
 
 	// Load with CLI prompt files - should merge with config
 	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{},
-		WithPromptFiles([]string{"cli-file.md"}))
+		withTestProviderRegistry(WithPromptFiles([]string{"cli-file.md"}))...)
 	require.NoError(t, err)
 
 	rootAgent, err := team.Agent("root")
@@ -511,7 +519,7 @@ agents:
 
 	// CLI specifies a file that's already in config - should deduplicate
 	team, err := Load(t.Context(), agentSource, &config.RuntimeConfig{},
-		WithPromptFiles([]string{"AGENTS.md", "extra.md"}))
+		withTestProviderRegistry(WithPromptFiles([]string{"AGENTS.md", "extra.md"}))...)
 	require.NoError(t, err)
 
 	rootAgent, err := team.Agent("root")
@@ -549,7 +557,7 @@ func TestGetToolsForAgent_MultipleLSPToolsetsAreCombined(t *testing.T) {
 
 	expander := js.NewJsExpander(runConfig.EnvProvider())
 
-	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, NewDefaultToolsetRegistry(), "test-config", expander)
+	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, testToolsetRegistry(), "test-config", expander)
 	require.Empty(t, warnings)
 
 	// Should have exactly one toolset (the multiplexer)
@@ -591,7 +599,7 @@ func TestGetToolsForAgent_SingleLSPToolsetNotWrapped(t *testing.T) {
 
 	expander := js.NewJsExpander(runConfig.EnvProvider())
 
-	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, NewDefaultToolsetRegistry(), "test-config", expander)
+	got, warnings := getToolsForAgent(t.Context(), a, ".", &runConfig, testToolsetRegistry(), "test-config", expander)
 	require.Empty(t, warnings)
 
 	// Should have exactly one toolset that provides LSP tools.
@@ -647,7 +655,7 @@ agents:
 	runConfig := &config.RuntimeConfig{}
 	runConfig.WorkingDir = tmpDir
 
-	_, err := LoadWithConfig(t.Context(), source, runConfig)
+	_, err := LoadWithConfig(t.Context(), source, runConfig, withTestProviderRegistry()...)
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "escapes parent directory")
 }

pkg/teamloader/toolsets/toolsets.go

@@ -0,0 +1,97 @@
+package toolsets
+
+import (
+	"context"
+
+	"github.com/docker/docker-agent/pkg/config"
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/teamloader"
+	"github.com/docker/docker-agent/pkg/tools"
+	"github.com/docker/docker-agent/pkg/tools/a2a"
+	agenttool "github.com/docker/docker-agent/pkg/tools/builtin/agent"
+	"github.com/docker/docker-agent/pkg/tools/builtin/api"
+	"github.com/docker/docker-agent/pkg/tools/builtin/fetch"
+	"github.com/docker/docker-agent/pkg/tools/builtin/filesystem"
+	"github.com/docker/docker-agent/pkg/tools/builtin/lsp"
+	"github.com/docker/docker-agent/pkg/tools/builtin/mcpcatalog"
+	"github.com/docker/docker-agent/pkg/tools/builtin/memory"
+	"github.com/docker/docker-agent/pkg/tools/builtin/modelpicker"
+	"github.com/docker/docker-agent/pkg/tools/builtin/openapi"
+	"github.com/docker/docker-agent/pkg/tools/builtin/rag"
+	"github.com/docker/docker-agent/pkg/tools/builtin/shell"
+	"github.com/docker/docker-agent/pkg/tools/builtin/tasks"
+	"github.com/docker/docker-agent/pkg/tools/builtin/think"
+	"github.com/docker/docker-agent/pkg/tools/builtin/todo"
+	"github.com/docker/docker-agent/pkg/tools/builtin/userprompt"
+	"github.com/docker/docker-agent/pkg/tools/mcp"
+)
+
+func NewDefaultToolsetRegistry() teamloader.ToolsetRegistry {
+	return teamloader.NewToolsetRegistry(DefaultToolsetCreators())
+}
+
+func DefaultToolsetCreators() map[string]teamloader.ToolsetCreator {
+	return map[string]teamloader.ToolsetCreator{
+		"todo": func(_ context.Context, toolset latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return todo.CreateToolSet(toolset)
+		},
+		"tasks": func(_ context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return tasks.CreateToolSet(toolset, parentDir, runConfig)
+		},
+		"memory": func(_ context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, configName string) (tools.ToolSet, error) {
+			return memory.CreateToolSet(toolset, parentDir, runConfig, configName)
+		},
+		"think": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return think.CreateToolSet()
+		},
+		"shell": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return shell.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"script": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return shell.CreateScriptToolSet(ctx, toolset, runConfig)
+		},
+		"filesystem": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return filesystem.CreateToolSet(toolset, runConfig)
+		},
+		"fetch": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return fetch.CreateToolSet(toolset, runConfig)
+		},
+		"mcp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return mcp.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"mcp_catalog": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			var opts []mcpcatalog.Option
+			if len(toolset.AllowedServers) > 0 {
+				opts = append(opts, mcpcatalog.WithAllowedServers(toolset.AllowedServers))
+			}
+			if len(toolset.BlockedServers) > 0 {
+				opts = append(opts, mcpcatalog.WithBlockedServers(toolset.BlockedServers))
+			}
+			return mcpcatalog.New(opts...), nil
+		},
+		"api": func(_ context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return api.CreateToolSet(toolset, runConfig)
+		},
+		"a2a": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return a2a.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"lsp": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return lsp.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"user_prompt": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return userprompt.CreateToolSet()
+		},
+		"openapi": func(ctx context.Context, toolset latest.Toolset, _ string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return openapi.CreateToolSet(ctx, toolset, runConfig)
+		},
+		"model_picker": func(_ context.Context, toolset latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return modelpicker.CreateToolSet(toolset)
+		},
+		"background_agents": func(_ context.Context, _ latest.Toolset, _ string, _ *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return agenttool.CreateToolSet()
+		},
+		"rag": func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
+			return rag.CreateToolSet(ctx, toolset, parentDir, runConfig)
+		},
+	}
+}

pkg/tools/builtin/rag/register.go

@@ -1,23 +0,0 @@
-package rag
-
-import (
-	"context"
-
-	"github.com/docker/docker-agent/pkg/config"
-	"github.com/docker/docker-agent/pkg/config/latest"
-	"github.com/docker/docker-agent/pkg/teamloader"
-	"github.com/docker/docker-agent/pkg/tools"
-)
-
-// init registers the "rag" toolset with teamloader. This package is kept out of
-// teamloader's static import graph (it pulls in a cgo tree-sitter dependency via
-// pkg/rag), so the "rag" toolset is only linked into binaries that blank-import
-// this package. The docker-agent CLI does so; embedders that don't need RAG can
-// omit it and build without the cgo dependency.
-//
-//nolint:gochecknoinits // Intentional: self-registers the optional rag toolset so a blank import enables it without pulling cgo into teamloader's import graph.
-func init() {
-	teamloader.RegisterToolsetCreator("rag", func(ctx context.Context, toolset latest.Toolset, parentDir string, runConfig *config.RuntimeConfig, _ string) (tools.ToolSet, error) {
-		return CreateToolSet(ctx, toolset, parentDir, runConfig)
-	})
-}