From afaa5819abde6a134874efab4ae87fec9477a0d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 15:52:11 +0000 Subject: [PATCH 1/5] fix(provider): preserve transport and surface clone failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/guides/go-sdk/index.md | 59 ++++- docs/providers/google/index.md | 12 + pkg/acp/agent_test.go | 9 +- pkg/model/provider/anthropic/vertex/vertex.go | 19 +- pkg/model/provider/clone.go | 127 ++++++++--- pkg/model/provider/clone_checked_test.go | 114 ++++++++++ pkg/model/provider/clone_matrix_test.go | 156 +++++++++++++ pkg/model/provider/clone_merge_test.go | 26 ++- pkg/model/provider/clone_test.go | 3 +- pkg/model/provider/clone_vertex_test.go | 207 ++++++++++++++++++ pkg/model/provider/contracts/contracts.go | 19 +- pkg/model/provider/options/options.go | 64 ++++-- pkg/model/provider/options/options_test.go | 21 ++ pkg/model/provider/rulebased/client.go | 5 +- pkg/model/provider/vertexai/modelgarden.go | 36 ++- .../provider/vertexai/modelgarden_test.go | 11 + pkg/runtime/clone_failure_test.go | 138 ++++++++++++ pkg/runtime/compaction_context_limit_test.go | 6 +- pkg/runtime/compactor/compactor.go | 6 +- pkg/runtime/compactor/compactor_test.go | 41 +++- pkg/runtime/fallback_test.go | 11 +- pkg/runtime/live_sessions_test.go | 10 +- pkg/runtime/model_context_test.go | 6 +- pkg/runtime/runtime_test.go | 14 +- pkg/runtime/sampling.go | 10 +- pkg/runtime/session_compaction.go | 26 ++- pkg/runtime/session_compaction_test.go | 3 +- pkg/runtime/title_generator_test.go | 8 +- pkg/runtime/transforms_test.go | 8 +- pkg/server/session_manager_test.go | 11 +- pkg/sessiontitle/generator.go | 5 +- pkg/sessiontitle/generator_test.go | 58 +++-- 32 files changed, 1131 insertions(+), 118 deletions(-) create mode 100644 pkg/model/provider/clone_checked_test.go create mode 100644 pkg/model/provider/clone_matrix_test.go create mode 100644 pkg/model/provider/clone_vertex_test.go create mode 100644 pkg/runtime/clone_failure_test.go diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index 88c747d9ea..23ff25cbde 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -847,7 +847,64 @@ client, err := openai.NewClient(ctx, &latest.ModelConfig{ }, env, options.WithTokenSource(tokenSource)) ``` -The OpenAI client checks for a configured `TokenSource` before falling back to `token_key`, and uses it to set the `Authorization` header on both HTTP and WebSocket requests. Static API-key, ChatGPT, and gateway auth paths are unaffected. `FromModelOptions` round-trips the token source, so a cloned provider config keeps it. Vertex AI Model Garden's default constructor refreshes GCP access tokens through ADC. Its `NewClientWithTokenSource` constructor instead accepts a host-managed request-time token source; neither host credential files nor instance metadata are consulted on that path. +The OpenAI client checks for a configured `TokenSource` before falling back to `token_key`, and uses it to set the `Authorization` header on both HTTP and WebSocket requests. Static API-key, ChatGPT, and gateway auth paths are unaffected. `FromModelOptions` round-trips the token source, so a cloned provider config keeps it. Vertex AI Model Garden's default constructor refreshes GCP access tokens through ADC unless `options.WithTokenSource` is supplied. Its `NewClientWithTokenSource` constructor instead accepts a host-managed request-time token source; neither host credential files nor instance metadata are consulted on that path. + +## Cloning Providers + +Use the checked `provider.Clone` API when adjusting a provider for an auxiliary +request. Options are applied in order; later overrides win. + +```go +summaryModel, err := provider.Clone(ctx, model, + options.WithStructuredOutput(nil), + options.WithMaxTokens(4096), + options.WithCompacting(), +) +if err != nil { + return fmt.Errorf("creating summary model: %w", err) +} +``` + +Providers constructed through a `provider.Registry` retain their reconstruction +function, environment and routing model references. Custom implementations must +supply `contracts.Config.RebuildProvider` (or `SetProviderRebuilder`), including +in test doubles. Missing reconstruction support returns `provider.ErrNoRebuilder`; +reconstruction errors are wrapped in `*provider.CloneError` and support +`errors.Is` / `errors.As`. + +Clones isolate mutable model configuration, custom-provider definitions and +structured-output schemas. Transport wrappers, token-source callbacks and the +models.dev store remain shared services. Public option setters retain reference +semantics; isolation happens when cloning. Programmatic configurations must be +JSON-serializable: unsupported values in `provider_opts`, schemas or custom +providers fail closed with `provider.ErrUncopyableConfig`, rather than panic or +share mutable state. Cumulative clones preserve explicit structured-output +removal, gateway clearing and no-thinking overrides; existing provider token +floors and routed-child token limits still apply. + +`CloneWithOptions` retains its public signature but is deprecated. On failure it +returns an error-only provider whose stream creation returns the clone error. It +never performs inference on the original provider or exposes its optional +embedding, context-discovery or compaction capabilities. + +| Caller | Clone failure behavior | +| --- | --- | +| Session title | Tries the next configured candidate; exhaustion does not interrupt chat. | +| MCP sampling, with or without tools | Returns the actionable clone error before inference. | +| LLM compaction | Reports failure and leaves the session history intact. | +| Compaction window resolution | The checked resolver returns zero plus the error; the LLM compaction path reports it. UI gauges remain best-effort. | + +Hook-supplied summaries and provider-native compaction remain usable without a +clonable summary model. For Model Garden, catalogue identity and dispatch identity +are separate: `ModelConfig` supplies capability lookups, while +`DispatchConfig` / `CloneSource()` preserves the original Google configuration for +reconstruction. Both ADC and request-context token sources retain native Vertex +transport, project, location and model-in-URL semantics. Model Garden token-source +precedence is an explicit `NewClientWithTokenSource` argument, then a stored +`options.WithTokenSource`, then ADC. Register that source only for the intended +Google provider, not as a process-wide option for unrelated routed models. + +This API change adds no YAML configuration surface or schema version. ## Using Different Providers diff --git a/docs/providers/google/index.md b/docs/providers/google/index.md index 00ef7a93ce..2a29c4acc6 100644 --- a/docs/providers/google/index.md +++ b/docs/providers/google/index.md @@ -301,6 +301,12 @@ Gemini SDK: - **Other publishers** (e.g. `meta`, `mistral`) use Vertex AI's OpenAI-compatible `/chat/completions` endpoint. +Title, sampling and compaction clones keep the selected Model Garden transport. +Claude requests remain on the Anthropic-native Vertex endpoints, with the model +in the URL and `anthropic_version: vertex-2023-10-16` in the body. Project, +location, regional/global host selection and Google OAuth authentication survive +cloning; clones never switch to the direct Claude API or send its API key. + ### Authentication Vertex AI uses Google Cloud Application Default Credentials (ADC). Make sure you @@ -310,6 +316,12 @@ are authenticated: gcloud auth application-default login ``` +Go embedders can instead supply a request-context token source through +`vertexai.NewClientWithTokenSource` or `options.WithTokenSource`. That source is +preserved across clones and replaces ADC, including credential-file and metadata +lookups. HTTP transport wrappers are preserved on both authentication paths; +ADC still applies OAuth authorization around the wrapped transport. + ### Configuration ```yaml diff --git a/pkg/acp/agent_test.go b/pkg/acp/agent_test.go index a44fb5ab9a..1c0a157dae 100644 --- a/pkg/acp/agent_test.go +++ b/pkg/acp/agent_test.go @@ -14,7 +14,10 @@ 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/config/latest" + "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/session" "github.com/docker/docker-agent/pkg/session/sqlitestore" @@ -52,7 +55,11 @@ func (m *mockProvider) CreateChatCompletionStream(context.Context, []chat.Messag return m.stream, nil } -func (m *mockProvider) BaseConfig() base.Config { return base.Config{} } +func (m *mockProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return m, nil } + return cfg +} func (m *mockProvider) MaxTokens() int { return 0 } diff --git a/pkg/model/provider/anthropic/vertex/vertex.go b/pkg/model/provider/anthropic/vertex/vertex.go index 86b0d2a0b8..acfa95cdcb 100644 --- a/pkg/model/provider/anthropic/vertex/vertex.go +++ b/pkg/model/provider/anthropic/vertex/vertex.go @@ -99,10 +99,15 @@ func newClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro // OAuth2 (via the google transport in vertex.WithCredentials), so we // must clear the stray X-Api-Key header that would otherwise leak a // direct-API credential into Google's infrastructure. - sdkOptions := []option.RequestOption{ + var sdkOptions []option.RequestOption + if source == nil { + // Before WithCredentials, OAuth wraps our transport instead of being replaced. + sdkOptions = append(sdkOptions, option.WithHTTPClient(wrappedHTTPClient(ctx, globalOptions))) + } + sdkOptions = append(sdkOptions, sdkvertex.WithCredentials(ctx, location, project, creds), option.WithAPIKey(""), - } + ) if source != nil { sdkOptions = append(sdkOptions, tokenSourceOptions(ctx, globalOptions, source)...) } @@ -141,11 +146,9 @@ func credentials(ctx context.Context, source options.TokenSource) (*google.Crede // as-is, so the SDK never overrides the Authorization header that the // middleware sets from the request context. func tokenSourceOptions(ctx context.Context, globalOptions options.ModelOptions, source options.TokenSource) []option.RequestOption { - httpClient := httpclient.NewHTTPClient(ctx) - globalOptions.WrapTransport(ctx, httpClient) return []option.RequestOption{ option.WithMiddleware(bearerTokenMiddleware(source)), - option.WithHTTPClient(httpClient), + option.WithHTTPClient(wrappedHTTPClient(ctx, globalOptions)), } } @@ -168,3 +171,9 @@ func bearerTokenMiddleware(source options.TokenSource) option.Middleware { return next(req) } } + +func wrappedHTTPClient(ctx context.Context, globalOptions options.ModelOptions) *http.Client { + client := httpclient.NewHTTPClient(ctx) + globalOptions.WrapTransport(ctx, client) + return client +} diff --git a/pkg/model/provider/clone.go b/pkg/model/provider/clone.go index d0b2c977bc..132854f67a 100644 --- a/pkg/model/provider/clone.go +++ b/pkg/model/provider/clone.go @@ -2,56 +2,121 @@ package provider import ( "context" + "encoding/json" + "errors" + "fmt" "log/slog" + "slices" + "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "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/tools" ) -// CloneWithOptions returns a new Provider instance using the same provider/model -// as the base provider, applying the provided options. If cloning fails, the -// original base provider is returned. -func CloneWithOptions(ctx context.Context, baseProvider Provider, opts ...options.Opt) Provider { - cfg := baseProvider.BaseConfig() - modelConfig, mergedOpts := mergeCloneOptions(cfg, opts) +// ErrNoRebuilder means a provider has no construction function for cloning. +var ErrNoRebuilder = errors.New("provider cannot be cloned without its rebuild function") +// ErrUncopyableConfig means mutable configuration contains non-JSON values. +var ErrUncopyableConfig = errors.New("provider clone requires JSON-serializable configuration") + +// CloneError identifies the provider whose reconstruction failed. +type CloneError struct { + ID modelsdev.ID + Err error +} + +func (e *CloneError) Error() string { return fmt.Sprintf("cloning provider %q: %v", e.ID, e.Err) } +func (e *CloneError) Unwrap() error { return e.Err } + +// Clone reconstructs a provider with ordered option overrides, or returns an error. +func Clone(ctx context.Context, p Provider, opts ...options.Opt) (Provider, error) { + clone, _, err := cloneProvider(ctx, p, opts) + return clone, err +} + +// CloneWithOptions reconstructs a provider, returning an error-only provider on failure. +// +// Deprecated: use Clone to handle reconstruction errors before making a request. +func CloneWithOptions(ctx context.Context, p Provider, opts ...options.Opt) Provider { + clone, requested, err := cloneProvider(ctx, p, opts) + if err == nil { + return clone + } + slog.WarnContext(ctx, "Failed to clone provider", "error", err) + var id modelsdev.ID + if p != nil { + id = p.ID() + } + return &failedCloneProvider{config: requested, id: id, err: err} +} + +func cloneProvider(ctx context.Context, p Provider, opts []options.Opt) (Provider, base.Config, error) { + if p == nil { + return nil, base.Config{}, errors.New("cannot clone a nil provider") + } + cfg := p.BaseConfig() + modelConfig, mergedOpts, err := mergeCloneOptions(cfg, opts) + if err != nil { + return nil, base.Config{}, &CloneError{ID: p.ID(), Err: err} + } + requested := base.Config{ModelConfig: modelConfig, ModelOptions: options.Apply(mergedOpts...), Env: cfg.Env, BaseURL: cfg.BaseURL} if cfg.RebuildProvider == nil { - slog.WarnContext(ctx, "Provider cannot be cloned without its rebuild function; using base provider", "id", baseProvider.ID()) - return baseProvider + return nil, requested, &CloneError{ID: p.ID(), Err: ErrNoRebuilder} } clone, err := cfg.RebuildProvider(ctx, &modelConfig, mergedOpts...) + if err == nil && clone == nil { + err = errors.New("rebuild returned a nil provider") + } if err != nil { - slog.WarnContext(ctx, "Failed to clone provider; using base provider", "error", err, "id", baseProvider.ID()) - return baseProvider + return nil, requested, &CloneError{ID: p.ID(), Err: err} } + return clone, requested, nil +} - return clone +type failedCloneProvider struct { + config base.Config + id modelsdev.ID + err error } -// mergeCloneOptions is the pure half of CloneWithOptions. Given the base -// provider's configuration and the user-supplied overrides, it returns: -// -// - a copy of the base ModelConfig with explicit overrides applied (currently -// MaxTokens and the no-thinking flag), and -// - the full ordered slice of options that should be passed to NewWithModels -// (existing options first, then user overrides; later opts win). -// -// Splitting this out from the impure NewWithModels call lets us table-test the -// option-merging logic without spinning up an HTTP server. -func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, []options.Opt) { +func (p *failedCloneProvider) ID() modelsdev.ID { return p.id } +func (p *failedCloneProvider) BaseConfig() base.Config { return p.config } +func (p *failedCloneProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + return nil, p.err +} + +// mergeCloneOptions snapshots mutable state before reconstructing a provider. +func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, []options.Opt, error) { // Preserve existing options, then apply overrides. Later opts take precedence. - baseOpts := options.FromModelOptions(cfg.ModelOptions) - mergedOpts := append(baseOpts, opts...) + baseOptions, err := cfg.ModelOptions.Clone() + if err != nil { + return latest.ModelConfig{}, nil, fmt.Errorf("%w: %w", ErrUncopyableConfig, err) + } + baseOpts := options.FromModelOptions(baseOptions) + mergedOpts := slices.Concat(baseOpts, opts) + + // Snapshot overrides too: option setters can carry caller-owned maps. + merged := options.Apply(mergedOpts...) + merged, err = merged.Clone() + if err != nil { + return latest.ModelConfig{}, nil, fmt.Errorf("%w: %w", ErrUncopyableConfig, err) + } + mergedOpts = options.FromModelOptions(merged) - // Apply every option to a single accumulator so we can read the final - // effective values directly. "Later opt wins" semantics fall out naturally. - var merged options.ModelOptions - for _, opt := range mergedOpts { - opt(&merged) + source := cfg.CloneSource() + var modelConfig latest.ModelConfig + data, err := json.Marshal(source) + if err != nil { + return latest.ModelConfig{}, nil, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) + } + if err := json.Unmarshal(data, &modelConfig); err != nil { + return latest.ModelConfig{}, nil, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) } + modelConfig.Name, modelConfig.DisplayModel = source.Name, source.DisplayModel - modelConfig := cfg.ModelConfig if mt := merged.MaxTokens(); mt != 0 { modelConfig.MaxTokens = &mt } @@ -61,5 +126,5 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, // thinking_budget via its setIfNil merge. modelConfig.ThinkingBudget = &latest.ThinkingBudget{Effort: "none"} } - return modelConfig, mergedOpts + return modelConfig, mergedOpts, nil } diff --git a/pkg/model/provider/clone_checked_test.go b/pkg/model/provider/clone_checked_test.go new file mode 100644 index 0000000000..5c2db58378 --- /dev/null +++ b/pkg/model/provider/clone_checked_test.go @@ -0,0 +1,114 @@ +package provider + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/contracts" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/modelsdev" +) + +func TestCloneErrors(t *testing.T) { + t.Parallel() + failure := errors.New("synthetic rebuild failure") + for _, tt := range []struct { + name string + rebuild base.RebuildProviderFunc + want error + }{ + {name: "missing", want: ErrNoRebuilder}, + {name: "reconstruction", rebuild: func(context.Context, *latest.ModelConfig, ...options.Opt) (Provider, error) { return nil, failure }, want: failure}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + original := &fakeProvider{id: modelsdev.NewID("test", "original"), config: base.Config{RebuildProvider: tt.rebuild}} + got, err := Clone(t.Context(), original) + require.ErrorIs(t, err, tt.want) + assert.Nil(t, got) + var cloneErr *CloneError + require.ErrorAs(t, err, &cloneErr) + assert.Equal(t, original.ID(), cloneErr.ID) + legacy := CloneWithOptions(t.Context(), original, options.WithNoThinking()) + stream, err := legacy.CreateChatCompletionStream(t.Context(), nil, nil) + require.ErrorIs(t, err, tt.want) + assert.Nil(t, stream) + _, embeds := legacy.(EmbeddingProvider) + assert.False(t, embeds) + _, compacts := legacy.(contracts.ConversationCompactor) + assert.False(t, compacts) + assert.Nil(t, ContextWindowResolver(legacy)) + assert.True(t, legacy.BaseConfig().ModelConfig.ThinkingBudget.IsDisabled()) + }) + } + _, err := Clone(t.Context(), nil) + require.ErrorContains(t, err, "nil provider") +} + +func TestCloneCumulativeOptionsAreIsolated(t *testing.T) { + t.Parallel() + cfg := base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "original", ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"nested": map[string]any{"value": "original"}}}, ModelOptions: options.Apply( + options.WithGateway("https://gateway.example"), options.WithStructuredOutput(&latest.StructuredOutput{Schema: map[string]any{"nested": map[string]any{"type": "string"}}}), + options.WithProviders(map[string]latest.ProviderConfig{"custom": {ThinkingBudget: &latest.ThinkingBudget{Tokens: 4096}}}), + )} + var rebuild base.RebuildProviderFunc + rebuild = func(_ context.Context, c *latest.ModelConfig, opts ...options.Opt) (Provider, error) { + return &fakeProvider{id: modelsdev.NewID(c.Provider, c.Model), config: base.Config{ModelConfig: *c, ModelOptions: options.Apply(opts...), RebuildProvider: rebuild}}, nil + } + cfg.RebuildProvider = rebuild + original := &fakeProvider{id: modelsdev.NewID("test", "original"), config: cfg} + preserved, err := Clone(t.Context(), original, nil) + require.NoError(t, err) + c := preserved.BaseConfig() + c.ModelConfig.ThinkingBudget.Tokens = 1 + c.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"] = "changed" + c.ModelOptions.StructuredOutput().Schema["nested"].(map[string]any)["type"] = "number" + c.ModelOptions.Providers()["custom"].ThinkingBudget.Tokens = 1 + b := original.BaseConfig() + assert.Equal(t, 2048, b.ModelConfig.ThinkingBudget.Tokens) + assert.Equal(t, "original", b.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"]) + assert.Equal(t, "string", b.ModelOptions.StructuredOutput().Schema["nested"].(map[string]any)["type"]) + assert.Equal(t, 4096, b.ModelOptions.Providers()["custom"].ThinkingBudget.Tokens) + first, err := Clone(t.Context(), original, options.WithStructuredOutput(nil), options.WithGateway(""), options.WithMaxTokens(128), options.WithNoThinking()) + require.NoError(t, err) + second, err := Clone(t.Context(), first, options.WithGeneratingTitle(), options.WithMaxTokens(256)) + require.NoError(t, err) + final := second.BaseConfig() + assert.Nil(t, final.ModelOptions.StructuredOutput()) + assert.Empty(t, final.ModelOptions.Gateway()) + assert.True(t, final.ModelOptions.NoThinking()) + assert.Equal(t, int64(256), *final.ModelConfig.MaxTokens) + assert.True(t, final.ModelConfig.ThinkingBudget.IsDisabled()) +} + +func TestCloneRestrictedRegistryReturnsActionableError(t *testing.T) { + t.Parallel() + registry := EmptyRegistry() + cfg := base.Config{ModelConfig: latest.ModelConfig{Provider: "google", Model: "claude-sonnet-4-6"}, RebuildProvider: func(ctx context.Context, cfg *latest.ModelConfig, opts ...options.Opt) (Provider, error) { + return registry.New(ctx, cfg, nil, opts...) + }} + _, err := Clone(t.Context(), &fakeProvider{config: cfg}) + require.ErrorContains(t, err, `unknown provider type "google"`) + var cloneErr *CloneError + require.ErrorAs(t, err, &cloneErr) +} + +func TestClonePreservesEmbeddingCapability(t *testing.T) { + t.Parallel() + p := &instrumentTestProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "embedding"}}} + p.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (Provider, error) { + return instrumentProvider(&instrumentTestProvider{}), nil + } + clone, err := Clone(t.Context(), instrumentProvider(p)) + require.NoError(t, err) + _, ok := clone.(EmbeddingProvider) + assert.True(t, ok) + _, ok = clone.(BatchEmbeddingProvider) + assert.True(t, ok) +} diff --git a/pkg/model/provider/clone_matrix_test.go b/pkg/model/provider/clone_matrix_test.go new file mode 100644 index 0000000000..9a1e6fba04 --- /dev/null +++ b/pkg/model/provider/clone_matrix_test.go @@ -0,0 +1,156 @@ +//go:build !js && !docker_agent_no_openai + +package provider + +import ( + "context" + "encoding/json" + "math" + "reflect" + "testing" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/options" +) + +func TestCloneOptionMatrixOnWire(t *testing.T) { + t.Parallel() + for _, shape := range []string{"literal", "environment", "inherited", "custom", "routed"} { + t.Run(shape, func(t *testing.T) { + t.Parallel() + server, body := captureNamedCustomProviderRequestBody(t, "openai_chatcompletions") + env := environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "secret", "CUSTOM_TOKEN": "secret", "ENDPOINT": server.URL, "MODEL": "gpt-5.6"}) + cfg := &latest.ModelConfig{Provider: "openai", Model: "gpt-5.6", BaseURL: server.URL, ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}, ProviderOpts: map[string]any{"api_type": "openai_chatcompletions"}} + providers := map[string]latest.ProviderConfig{} + models := map[string]latest.ModelConfig{} + switch shape { + case "environment": + cfg.BaseURL = "${ENDPOINT}" + cfg.Model = "${MODEL}" + case "inherited", "custom", "routed": + cfg.Provider = "custom" + cfg.BaseURL = "" + cfg.ProviderOpts = nil + cfg.ThinkingBudget = nil + providers["custom"] = latest.ProviderConfig{Provider: "openai", BaseURL: server.URL, TokenKey: "CUSTOM_TOKEN", APIType: "openai_chatcompletions", ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}} + if shape == "routed" { + models["target"] = latest.ModelConfig{Provider: "custom", Model: "gpt-5.6"} + cfg.Routing = []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}} + } + } + original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, env, options.WithProviders(providers), options.WithStructuredOutput(&latest.StructuredOutput{Name: "test", Schema: map[string]any{"type": "object", "properties": map[string]any{"count": map[string]any{"type": "integer", "minimum": 1}}}}), options.WithGateway("https://must-bypass.example")) + require.NoError(t, err) + clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128), options.WithStructuredOutput(nil)) + require.NoError(t, err) + clone, err = Clone(t.Context(), clone, options.WithGeneratingTitle()) + require.NoError(t, err) + stream, err := clone.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + var payload map[string]any + require.NoError(t, json.Unmarshal(body(), &payload)) + assert.Equal(t, "none", payload["reasoning_effort"]) + assert.NotContains(t, payload, "response_format") + if shape != "routed" { + assert.InDelta(t, float64(256), payload["max_completion_tokens"], 0) + } + cfgClone := clone.BaseConfig() + assert.True(t, cfgClone.ModelOptions.NoThinking()) + assert.Nil(t, cfgClone.ModelOptions.StructuredOutput()) + if shape == "inherited" || shape == "custom" || shape == "routed" { + assert.NotEmpty(t, cfgClone.ModelOptions.Providers()) + } + }) + } +} + +func TestCloneRejectsUncopyableConfigWithoutPanic(t *testing.T) { + t.Parallel() + for _, value := range []any{math.Inf(1), func() {}, make(chan int)} { + for _, field := range []string{"provider_opts", "structured_output", "providers"} { + cfg := base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "base"}, RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (Provider, error) { + t.Error("must reject before rebuilding") + return nil, nil + }} + switch field { + case "provider_opts": + cfg.ModelConfig.ProviderOpts = map[string]any{"bad": value} + case "structured_output": + cfg.ModelOptions = options.Apply(options.WithStructuredOutput(&latest.StructuredOutput{Schema: map[string]any{"bad": value}})) + case "providers": + cfg.ModelOptions = options.Apply(options.WithProviders(map[string]latest.ProviderConfig{"bad": {ProviderOpts: map[string]any{"bad": value}}})) + } + p := &fakeProvider{config: cfg} + clone, err := Clone(t.Context(), p) + assert.Nil(t, clone) + require.ErrorIs(t, err, ErrUncopyableConfig) + legacy := CloneWithOptions(t.Context(), p) + _, err = legacy.CreateChatCompletionStream(t.Context(), nil, nil) + require.ErrorIs(t, err, ErrUncopyableConfig) + } + } +} + +func TestCloneCopiesYAMLAndHiddenFields(t *testing.T) { + t.Parallel() + var cfg latest.ModelConfig + require.NoError(t, yaml.Unmarshal([]byte(`provider: test +model: base +max_tokens: 1024 +track_usage: false +parallel_tool_calls: false +provider_opts: + extra_body: + chat_template_kwargs: + enable_thinking: false + strict_tools: [one, two] +`), &cfg)) + cfg.Name = "named" + cfg.DisplayModel = "display" + copied, _, err := mergeCloneOptions(base.Config{ModelConfig: cfg}, nil) + require.NoError(t, err) + assert.Equal(t, "named", copied.Name) + assert.Equal(t, "display", copied.DisplayModel) + require.NotNil(t, copied.TrackUsage) + assert.False(t, *copied.TrackUsage) + require.NotNil(t, copied.ParallelToolCalls) + assert.False(t, *copied.ParallelToolCalls) + assert.Equal(t, int64(1024), *copied.MaxTokens) + assert.Equal(t, []any{"one", "two"}, copied.ProviderOpts["strict_tools"]) + assert.Equal(t, false, copied.ProviderOpts["extra_body"].(map[string]any)["chat_template_kwargs"].(map[string]any)["enable_thinking"]) + var hidden []string + typ := reflect.TypeFor[latest.ModelConfig]() + for field := range typ.Fields() { + if field.Tag.Get("json") == "-" { + hidden = append(hidden, field.Name) + } + } + assert.ElementsMatch(t, []string{"Name", "DisplayModel"}, hidden, "checked copy must preserve every JSON-excluded field") +} + +func TestClonePreservesStructuredOutputOnWire(t *testing.T) { + t.Parallel() + server, body := captureNamedCustomProviderRequestBody(t, "openai_chatcompletions") + original, err := fullTestRegistry().New(t.Context(), &latest.ModelConfig{Provider: "openai", Model: "gpt-4o", BaseURL: server.URL, ProviderOpts: map[string]any{"api_type": "openai_chatcompletions"}}, environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "secret"}), options.WithStructuredOutput(&latest.StructuredOutput{Name: "counter", Schema: map[string]any{"type": "object", "properties": map[string]any{"count": map[string]any{"type": "integer", "minimum": 1}}}})) + require.NoError(t, err) + clone, err := Clone(t.Context(), original, options.WithMaxTokens(512)) + require.NoError(t, err) + stream, err := clone.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + var payload map[string]any + require.NoError(t, json.Unmarshal(body(), &payload)) + require.Contains(t, payload, "response_format") + format := payload["response_format"].(map[string]any)["json_schema"].(map[string]any) + assert.Equal(t, "counter", format["name"]) + assert.InDelta(t, float64(1), format["schema"].(map[string]any)["properties"].(map[string]any)["count"].(map[string]any)["minimum"], 0) +} diff --git a/pkg/model/provider/clone_merge_test.go b/pkg/model/provider/clone_merge_test.go index fee91833fd..955a572457 100644 --- a/pkg/model/provider/clone_merge_test.go +++ b/pkg/model/provider/clone_merge_test.go @@ -29,7 +29,8 @@ func TestMergeCloneOptions_NoOverrides(t *testing.T) { }, } - got, mergedOpts := mergeCloneOptions(cfg, nil) + got, mergedOpts, err := mergeCloneOptions(cfg, nil) + require.NoError(t, err) require.NotNil(t, got.MaxTokens) assert.Equal(t, originalMax, *got.MaxTokens, "MaxTokens should be untouched when no overrides") @@ -47,7 +48,8 @@ func TestMergeCloneOptions_MaxTokensOverride(t *testing.T) { } newMax := int64(2048) - got, mergedOpts := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(newMax)}) + got, mergedOpts, err := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(newMax)}) + require.NoError(t, err) require.NotNil(t, got.MaxTokens) assert.Equal(t, newMax, *got.MaxTokens, "MaxTokens should follow explicit override") @@ -69,7 +71,8 @@ func TestMergeCloneOptions_NoThinking(t *testing.T) { }, } - got, _ := mergeCloneOptions(cfg, []options.Opt{options.WithNoThinking()}) + got, _, err := mergeCloneOptions(cfg, []options.Opt{options.WithNoThinking()}) + require.NoError(t, err) require.NotNil(t, got.ThinkingBudget, "WithNoThinking must write the disabled sentinel, not nil") assert.True(t, got.ThinkingBudget.IsDisabled(), @@ -90,7 +93,8 @@ func TestMergeCloneOptions_PreservesBaseOptions(t *testing.T) { ModelConfig: latest.ModelConfig{Provider: "openai", Model: "gpt-4o"}, } - _, mergedOpts := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(1024)}) + _, mergedOpts, err := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(1024)}) + require.NoError(t, err) // Replay the merged opts and check all three flags survived. var probe options.ModelOptions @@ -115,16 +119,15 @@ func TestMergeCloneOptions_LaterOverridesWin(t *testing.T) { ModelConfig: latest.ModelConfig{Provider: "openai", Model: "gpt-4o"}, } - got, _ := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(int64(4096))}) + got, _, err := mergeCloneOptions(cfg, []options.Opt{options.WithMaxTokens(int64(4096))}) + require.NoError(t, err) require.NotNil(t, got.MaxTokens) assert.Equal(t, int64(4096), *got.MaxTokens, "later opt must override earlier opt") } -// TestCloneWithOptions_FallbackOnError verifies the previously-uncovered -// fallback path: when NewWithModels fails, CloneWithOptions returns the -// original base provider unchanged. -func TestCloneWithOptions_FallbackOnError(t *testing.T) { +// A legacy clone must expose reconstruction failures without base inference. +func TestCloneWithOptions_ErrorProvider(t *testing.T) { t.Parallel() original := &fakeProvider{ @@ -136,5 +139,8 @@ func TestCloneWithOptions_FallbackOnError(t *testing.T) { got := CloneWithOptions(t.Context(), original, options.WithMaxTokens(int64(2048))) - assert.Same(t, original, got, "should fall back to base provider when cloning fails") + assert.NotSame(t, original, got) + stream, err := got.CreateChatCompletionStream(t.Context(), nil, nil) + assert.Nil(t, stream) + require.ErrorContains(t, err, "rebuild failed") } diff --git a/pkg/model/provider/clone_test.go b/pkg/model/provider/clone_test.go index 9b388a73c7..cf156b1362 100644 --- a/pkg/model/provider/clone_test.go +++ b/pkg/model/provider/clone_test.go @@ -184,7 +184,8 @@ func TestCloneWithOptions_PreservesOutputCapabilities(t *testing.T) { OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: &image}, }} - cloned, _ := mergeCloneOptions(cfg, []options.Opt{options.WithGeneratingTitle()}) + cloned, _, err := mergeCloneOptions(cfg, []options.Opt{options.WithGeneratingTitle()}) + require.NoError(t, err) require.NotNil(t, cloned.OutputCapabilities) require.NotNil(t, cloned.OutputCapabilities.Image) diff --git a/pkg/model/provider/clone_vertex_test.go b/pkg/model/provider/clone_vertex_test.go new file mode 100644 index 0000000000..d40cc7d6bb --- /dev/null +++ b/pkg/model/provider/clone_vertex_test.go @@ -0,0 +1,207 @@ +//go:build !js && !docker_agent_no_openai + +package provider + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/httpclient" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/model/provider/vertexai" +) + +type cloneRoundTripper func(*http.Request) (*http.Response, error) + +func (f cloneRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestCloneVertexNativeDispatch(t *testing.T) { + t.Run("ADC", func(t *testing.T) { + tokens := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "/", r.URL.Path) + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"access_token":"fake-adc","token_type":"Bearer","expires_in":3600}`) + })) + defer tokens.Close() + wireCalls := 0 + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return cloneRoundTripper(func(req *http.Request) (*http.Response, error) { + wireCalls++ + assert.Equal(t, "us-east5-aiplatform.googleapis.com", req.URL.Host) + assert.Equal(t, "/v1/projects/test-project/locations/us-east5/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", req.URL.Path) + assert.Equal(t, "Bearer fake-adc", req.Header.Get("Authorization")) + assert.Empty(t, req.Header.Get("X-Api-Key")) + body, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"anthropic_version":"vertex-2023-10-16"`) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")), Request: req}, nil + }) + }) + file := filepath.Join(t.TempDir(), "adc.json") + require.NoError(t, os.WriteFile(file, fmt.Appendf(nil, `{"type":"authorized_user","client_id":"offline","client_secret":"offline","refresh_token":"offline","token_uri":%q}`, tokens.URL), 0o600)) + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", file) + googleCalls := 0 + registry := NewRegistry(map[string]Factory{"google": func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) { + googleCalls++ + return vertexai.NewClient(ctx, cfg, env, opts...) + }, "anthropic": func(context.Context, *latest.ModelConfig, environment.Provider, ...options.Opt) (Provider, error) { + t.Fatal("direct Anthropic dispatch") + return nil, nil + }}) + original, err := registry.New(t.Context(), vertexCloneConfig("us-east5"), environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "must-not-leak"}), wrapper) + require.NoError(t, err) + clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128)) + require.NoError(t, err) + cfg := clone.BaseConfig() + assert.Equal(t, "anthropic", cfg.ModelConfig.Provider) + assert.Equal(t, "google", cfg.CloneSource().Provider) + assert.Equal(t, "test-project", cfg.CloneSource().ProviderOpts["project"]) + assert.Equal(t, "us-east5", cfg.CloneSource().ProviderOpts["location"]) + assert.True(t, cfg.ModelOptions.NoThinking()) + assert.Equal(t, 2, googleCalls) + for _, p := range []Provider{original, clone} { + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + for { + _, err = stream.Recv() + if err != nil { + break + } + } + stream.Close() + } + assert.Equal(t, 2, wireCalls) + }) + t.Run("contextual token source", func(t *testing.T) { + t.Parallel() + for _, location := range []string{"us-east5", "global", "eu"} { + t.Run(location, func(t *testing.T) { + t.Parallel() + var sessions []string + source := func(ctx context.Context) (string, error) { + sessions = append(sessions, httpclient.SessionIDFromContext(ctx)) + return "oauth-token", ctx.Err() + } + calls := 0 + registry := NewRegistry(map[string]Factory{"google": func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) { + calls++ + if calls == 1 { + return vertexai.NewClientWithTokenSource(ctx, cfg, env, source, opts...) + } + return vertexai.NewClient(ctx, cfg, env, opts...) + }, "anthropic": func(context.Context, *latest.ModelConfig, environment.Provider, ...options.Opt) (Provider, error) { + t.Fatal("direct Anthropic dispatch") + return nil, nil + }}) + wireCalls := 0 + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return cloneRoundTripper(func(req *http.Request) (*http.Response, error) { + wireCalls++ + host := location + "-aiplatform.googleapis.com" + if location == "global" { + host = "aiplatform.googleapis.com" + } + if location == "eu" { + host = "aiplatform.eu.rep.googleapis.com" + } + assert.Equal(t, host, req.URL.Host) + assert.Equal(t, "/v1/projects/test-project/locations/"+location+"/publishers/anthropic/models/claude-sonnet-4-6:streamRawPredict", req.URL.Path) + assert.Equal(t, "Bearer oauth-token", req.Header.Get("Authorization")) + assert.Empty(t, req.Header.Get("X-Api-Key")) + body, err := io.ReadAll(req.Body) + require.NoError(t, err) + var payload map[string]any + require.NoError(t, json.Unmarshal(body, &payload)) + assert.Equal(t, "vertex-2023-10-16", payload["anthropic_version"]) + assert.NotContains(t, payload, "model") + assert.NotContains(t, payload, "thinking") + assert.InDelta(t, float64(256), payload["max_tokens"], 0) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")), Request: req}, nil + }) + }) + original, err := registry.New(t.Context(), vertexCloneConfig(location), environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "must-not-leak"}), options.WithGateway("https://ignored-gateway.example"), wrapper) + require.NoError(t, err) + clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128), options.WithStructuredOutput(nil)) + require.NoError(t, err) + clone, err = Clone(t.Context(), clone, options.WithGeneratingTitle()) + require.NoError(t, err) + assert.Equal(t, 3, calls) + stream, err := clone.CreateChatCompletionStream(httpclient.ContextWithSessionID(t.Context(), "clone-session"), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + for { + _, err = stream.Recv() + if err != nil { + break + } + } + stream.Close() + assert.Equal(t, 1, wireCalls) + assert.Equal(t, []string{"", "", "", "clone-session"}, sessions) + }) + } + }) +} + +func vertexCloneConfig(location string) *latest.ModelConfig { + return &latest.ModelConfig{Provider: "google", Model: "claude-sonnet-4-6", BypassModelsGateway: true, ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"publisher": "anthropic", "project": "test-project", "location": location}} +} + +func TestCloneModelGardenCompatibleTransportAndSiblingAuth(t *testing.T) { + t.Parallel() + cfg := &latest.ModelConfig{Provider: "google", Model: "meta/llama-test", ProviderOpts: map[string]any{"publisher": "meta", "project": "test-project", "location": "us-central1"}} + var sources, wire int + source := func(ctx context.Context) (string, error) { sources++; return "vertex-only", ctx.Err() } + registry := NewRegistry(map[string]Factory{"google": vertexai.NewClient, "openai": openAITestFactory}) + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return cloneRoundTripper(func(req *http.Request) (*http.Response, error) { + wire++ + assert.Equal(t, "us-central1-aiplatform.googleapis.com", req.URL.Host) + assert.Equal(t, "/v1beta1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions", req.URL.Path) + assert.Equal(t, "Bearer vertex-only", req.Header.Get("Authorization")) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")), Request: req}, nil + }) + }) + original, err := registry.New(t.Context(), cfg, environment.NewNoEnvProvider(), options.WithTokenSource(source), wrapper) + require.NoError(t, err) + clone, err := Clone(t.Context(), original, options.WithNoThinking()) + require.NoError(t, err) + stream, err := clone.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + assert.Equal(t, 1, wire) + assert.Positive(t, sources) + // A leaf's stored OAuth source must not enter a router's sibling options. + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + assert.Equal(t, "Bearer sibling-key", req.Header.Get("Authorization")) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer server.Close() + factories := NewRegistry(map[string]Factory{"google": func(context.Context, *latest.ModelConfig, environment.Provider, ...options.Opt) (Provider, error) { + return original, nil + }, "openai": openAITestFactory}) + router, err := factories.NewWithModels(t.Context(), &latest.ModelConfig{Provider: "openai", Model: "gpt-4o", Routing: []latest.RoutingRule{{Model: "vertex", Examples: []string{"vertex"}}}}, map[string]latest.ModelConfig{"vertex": *cfg}, environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "sibling-key"}), options.WithProviders(map[string]latest.ProviderConfig{"openai": {BaseURL: server.URL, TokenKey: "OPENAI_API_KEY"}})) + require.NoError(t, err) + router, err = Clone(t.Context(), router, options.WithNoThinking()) + require.NoError(t, err) + stream, err = router.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "unmatched fallback request"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + assert.Equal(t, 1, wire, "fallback must not call the Vertex transport") +} diff --git a/pkg/model/provider/contracts/contracts.go b/pkg/model/provider/contracts/contracts.go index 40b26a5662..77e740f1ea 100644 --- a/pkg/model/provider/contracts/contracts.go +++ b/pkg/model/provider/contracts/contracts.go @@ -26,9 +26,11 @@ type RebuildProviderFunc func( // Config is the common configuration embedded by provider clients. type Config struct { - ModelConfig latest.ModelConfig - ModelOptions options.ModelOptions - Env environment.Provider + // DispatchConfig keeps construction identity separate from catalogue identity. + DispatchConfig *latest.ModelConfig + ModelConfig latest.ModelConfig + ModelOptions options.ModelOptions + Env environment.Provider // RebuildProvider preserves the construction context needed to clone a provider. RebuildProvider RebuildProviderFunc // BaseURL is the resolved endpoint actually used by the provider, unlike @@ -36,6 +38,17 @@ type Config struct { BaseURL string } +// SetDispatchConfig preserves the original configuration for reconstruction. +func (c *Config) SetDispatchConfig(cfg *latest.ModelConfig) { c.DispatchConfig = cfg.Clone() } + +// CloneSource returns the construction configuration, not the catalogue identity. +func (c *Config) CloneSource() latest.ModelConfig { + if c.DispatchConfig != nil { + return *c.DispatchConfig + } + return c.ModelConfig +} + // SetProviderRebuilder sets the function used to reconstruct the provider. func (c *Config) SetProviderRebuilder(rebuild RebuildProviderFunc) { c.RebuildProvider = rebuild diff --git a/pkg/model/provider/options/options.go b/pkg/model/provider/options/options.go index 12ca03c0da..897fcedc47 100644 --- a/pkg/model/provider/options/options.go +++ b/pkg/model/provider/options/options.go @@ -2,6 +2,8 @@ package options import ( "context" + "encoding/json" + "fmt" "log/slog" "net/http" @@ -13,18 +15,20 @@ import ( type TokenSource func(context.Context) (string, error) type ModelOptions struct { - gateway string - encryptedConfig string - structuredOutput *latest.StructuredOutput - generatingTitle bool - compacting bool - noThinking bool - maxTokens int64 - providers map[string]latest.ProviderConfig - modelsDevStore *modelsdev.Store - transportWrapper func(http.RoundTripper) http.RoundTripper - tokenSource TokenSource - openAIVendor bool + gateway string + gatewaySet bool + structuredOutputSet bool + encryptedConfig string + structuredOutput *latest.StructuredOutput + generatingTitle bool + compacting bool + noThinking bool + maxTokens int64 + providers map[string]latest.ProviderConfig + modelsDevStore *modelsdev.Store + transportWrapper func(http.RoundTripper) http.RoundTripper + tokenSource TokenSource + openAIVendor bool } func (c *ModelOptions) Gateway() string { @@ -131,6 +135,7 @@ func ForEvaluator(cfg latest.EvaluatorConfig, opts ...Opt) ModelOptions { func WithGateway(gateway string) Opt { return func(cfg *ModelOptions) { cfg.gateway = gateway + cfg.gatewaySet = true } } @@ -149,7 +154,8 @@ func WithEncryptedConfig(encryptedConfig string) Opt { // constrain the response natively in that mode. func WithStructuredOutput(structuredOutput *latest.StructuredOutput) Opt { return func(cfg *ModelOptions) { - if structuredOutput.ToolMode() { + cfg.structuredOutputSet = true + if structuredOutput == nil || structuredOutput.ToolMode() { cfg.structuredOutput = nil return } @@ -251,13 +257,13 @@ func WithHTTPTransportWrapper(fn func(base http.RoundTripper) http.RoundTripper) // Opt configuration functions. Later Opts override earlier ones when applied. func FromModelOptions(m ModelOptions) []Opt { var out []Opt - if g := m.Gateway(); g != "" { - out = append(out, WithGateway(g)) + if m.gatewaySet { + out = append(out, WithGateway(m.gateway)) } if m.encryptedConfig != "" { out = append(out, WithEncryptedConfig(m.encryptedConfig)) } - if m.structuredOutput != nil { + if m.structuredOutputSet { out = append(out, WithStructuredOutput(m.structuredOutput)) } if m.generatingTitle { @@ -289,3 +295,29 @@ func FromModelOptions(m ModelOptions) []Opt { } return out } + +// Clone copies mutable option state while preserving shared services and callbacks. +func (c *ModelOptions) Clone() (ModelOptions, error) { + cloned := *c + if c.structuredOutput != nil { + data, err := json.Marshal(c.structuredOutput) + if err != nil { + return ModelOptions{}, fmt.Errorf("copying structured_output: %w", err) + } + cloned.structuredOutput = new(latest.StructuredOutput) + if err := json.Unmarshal(data, cloned.structuredOutput); err != nil { + return ModelOptions{}, err + } + } + if c.providers != nil { + data, err := json.Marshal(c.providers) + if err != nil { + return ModelOptions{}, fmt.Errorf("copying providers: %w", err) + } + cloned.providers = nil + if err := json.Unmarshal(data, &cloned.providers); err != nil { + return ModelOptions{}, err + } + } + return cloned, nil +} diff --git a/pkg/model/provider/options/options_test.go b/pkg/model/provider/options/options_test.go index ef710b9743..7293a729f4 100644 --- a/pkg/model/provider/options/options_test.go +++ b/pkg/model/provider/options/options_test.go @@ -9,6 +9,7 @@ import ( "github.com/stretchr/testify/require" "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/modelsdev" ) // sentinelTransport is a minimal http.RoundTripper used only for identity checks. @@ -159,3 +160,23 @@ func TestFromModelOptions_NilWrapperNotIncluded(t *testing.T) { assert.Nil(t, dst.TransportWrapper()) } + +func TestModelOptionsClonePreservesServicesAndExplicitClears(t *testing.T) { + t.Parallel() + store := modelsdev.NewDatabaseStore(&modelsdev.Database{}) + wrapper := func(rt http.RoundTripper) http.RoundTripper { return rt } + source := func(context.Context) (string, error) { return "token", nil } + m := Apply(WithModelsDevStore(store), WithHTTPTransportWrapper(wrapper), WithTokenSource(source), WithGateway("old"), WithGateway(""), WithStructuredOutput(nil)) + clone, err := m.Clone() + require.NoError(t, err) + probe := Apply(FromModelOptions(clone)...) + assert.Same(t, store, probe.ModelsDevStore()) + assert.NotNil(t, probe.TransportWrapper()) + assert.NotNil(t, probe.TokenSource()) + assert.Empty(t, probe.Gateway()) + assert.True(t, probe.gatewaySet) + assert.True(t, probe.structuredOutputSet) + token, err := probe.TokenSource()(t.Context()) + require.NoError(t, err) + assert.Equal(t, "token", token) +} diff --git a/pkg/model/provider/rulebased/client.go b/pkg/model/provider/rulebased/client.go index ee152ad685..5dafa3dc9c 100644 --- a/pkg/model/provider/rulebased/client.go +++ b/pkg/model/provider/rulebased/client.go @@ -66,8 +66,9 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, models map[string]l client := &Client{ Config: base.Config{ - ModelConfig: *cfg, - Env: env, + ModelConfig: *cfg, + ModelOptions: options.Apply(opts...), + Env: env, }, matcher: newMatcher(), fallback: fallback, diff --git a/pkg/model/provider/vertexai/modelgarden.go b/pkg/model/provider/vertexai/modelgarden.go index a294afdc50..358015f478 100644 --- a/pkg/model/provider/vertexai/modelgarden.go +++ b/pkg/model/provider/vertexai/modelgarden.go @@ -75,22 +75,32 @@ func IsModelGardenConfig(cfg *latest.ModelConfig) bool { // Without this rewrite, attachments and other capability-gated features are // silently dropped because `/` IDs like // `google/claude-sonnet-4-20250514` do not exist in the models.dev database. -// See modelsDevProvider for the exact mapping. +// See modelsDevProvider for the exact mapping. A stored token-source option +// replaces ADC and survives reconstruction. func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Client, error) { + o := options.Apply(opts...) + opts = options.FromModelOptions(o) + if o.TokenSource() != nil { + return NewClientWithTokenSource(ctx, cfg, env, o.TokenSource(), opts...) + } project, location, err := resolveProjectLocation(ctx, cfg, env) if err != nil { return nil, err } rewritten := withModelsDevProvider(cfg) + var client Client if strings.EqualFold(publisher(cfg), "anthropic") { - return vertex.NewClient(ctx, rewritten, env, project, location, opts...) + client, err = vertex.NewClient(ctx, rewritten, env, project, location, opts...) + } else { + client, err = newOpenAIClient(ctx, rewritten, env, project, location, opts...) } - return newOpenAIClient(ctx, rewritten, env, project, location, opts...) + return withDispatchConfig(client, cfg, err) } // NewClientWithTokenSource is NewClient authenticated with source instead of // Application Default Credentials, which are never looked up. The token is -// resolved with each request's context. +// resolved with each request's context. The explicit source overrides any +// token-source option and is retained for subsequent clones. func NewClientWithTokenSource(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, source options.TokenSource, opts ...options.Opt) (Client, error) { if source == nil { return nil, errors.New("token source is required") @@ -99,11 +109,25 @@ func NewClientWithTokenSource(ctx context.Context, cfg *latest.ModelConfig, env if err != nil { return nil, err } + opts = append(opts, options.WithTokenSource(source)) rewritten := withModelsDevProvider(cfg) + var client Client if strings.EqualFold(publisher(cfg), "anthropic") { - return vertex.NewClientWithTokenSource(ctx, rewritten, env, project, location, source, opts...) + client, err = vertex.NewClientWithTokenSource(ctx, rewritten, env, project, location, source, opts...) + } else { + client, err = newOpenAIClientWithTokenSource(ctx, rewritten, env, project, location, source, opts...) + } + return withDispatchConfig(client, cfg, err) +} + +func withDispatchConfig(client Client, cfg *latest.ModelConfig, err error) (Client, error) { + if err != nil { + return nil, err + } + if setter, ok := client.(interface{ SetDispatchConfig(cfg *latest.ModelConfig) }); ok { + setter.SetDispatchConfig(cfg) } - return newOpenAIClientWithTokenSource(ctx, rewritten, env, project, location, source, opts...) + return client, nil } // withModelsDevProvider returns a deep copy of cfg with its Provider field diff --git a/pkg/model/provider/vertexai/modelgarden_test.go b/pkg/model/provider/vertexai/modelgarden_test.go index a5ed4772c3..5b4add26cb 100644 --- a/pkg/model/provider/vertexai/modelgarden_test.go +++ b/pkg/model/provider/vertexai/modelgarden_test.go @@ -542,3 +542,14 @@ func TestResolveProjectLocation(t *testing.T) { }) } } + +func TestNewClientAppliesOptionsOnce(t *testing.T) { + t.Parallel() + calls := 0 + _, err := NewClient(t.Context(), &latest.ModelConfig{Provider: "google", Model: "claude-sonnet-4-6", ProviderOpts: map[string]any{"publisher": "anthropic", "project": "test-project", "location": "us-east5"}}, environment.NewNoEnvProvider(), func(o *options.ModelOptions) { + calls++ + options.WithTokenSource(func(ctx context.Context) (string, error) { return "token", ctx.Err() })(o) + }) + require.NoError(t, err) + assert.Equal(t, 1, calls) +} diff --git a/pkg/runtime/clone_failure_test.go b/pkg/runtime/clone_failure_test.go new file mode 100644 index 0000000000..5f2025cf4f --- /dev/null +++ b/pkg/runtime/clone_failure_test.go @@ -0,0 +1,138 @@ +//nolint:staticcheck // Sampling is supported during the MCP deprecation window. +package runtime + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "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/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" +) + +type cloneFailureProvider struct { + base.Config + + endpoint string +} + +func (p *cloneFailureProvider) CreateChatCompletionStream(ctx context.Context, _ []chat.Message, _ []tools.Tool) (chat.MessageStream, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.endpoint, http.NoBody) + if err != nil { + return nil, err + } + res, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + res.Body.Close() + return &mockStream{}, nil +} + +func (p *cloneFailureProvider) ContextWindow(ctx context.Context) (int64, error) { + _, err := p.CreateChatCompletionStream(ctx, nil, nil) + return 16_000, err +} + +func TestCloneFailuresReachRuntimeCallers(t *testing.T) { + t.Parallel() + synthetic := errors.New("synthetic reconstruction failure") + for _, tt := range []struct { + name string + rebuild base.RebuildProviderFunc + want error + }{ + {name: "missing rebuilder", want: provider.ErrNoRebuilder}, + {name: "reconstruction failure", want: synthetic, rebuild: func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { + return nil, synthetic + }}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + var requests atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { requests.Add(1); w.WriteHeader(http.StatusOK) })) + defer server.Close() + p := &cloneFailureProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "base", ProviderOpts: map[string]any{"context_size": 100_000}}, RebuildProvider: tt.rebuild}, endpoint: server.URL} + a := agent.New("root", "test", agent.WithModel(p), agent.WithCompactionModel(p)) + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(a)), WithCurrentAgent("root"), WithModelStore(mockModelStoreWithLimit{limit: 100_000})) + require.NoError(t, err) + _, err = rt.samplingHandler(t.Context(), &mcp.CreateMessageParams{Messages: []*mcp.SamplingMessage{{Role: "user", Content: &mcp.TextContent{Text: "hello"}}}, MaxTokens: 128}) + require.ErrorIs(t, err, tt.want) + require.ErrorContains(t, err, "cloning model for sampling") + _, err = rt.samplingWithToolsHandler(t.Context(), &mcp.CreateMessageWithToolsParams{Messages: []*mcp.SamplingMessageV2{{Role: "user", Content: []mcp.Content{&mcp.TextContent{Text: "hello"}}}}, MaxTokens: 128}) + require.ErrorIs(t, err, tt.want) + require.ErrorContains(t, err, "cloning model for sampling") + n, err := rt.compactionContextLimitChecked(t.Context(), a) + assert.Zero(t, n) + require.ErrorIs(t, err, tt.want) + assert.Zero(t, rt.compactionContextLimit(t.Context(), a)) + assert.Equal(t, int64(100_000), rt.effectiveContextLimit(t.Context(), a, 100_000)) + sess := session.New(session.WithMessages([]session.Item{session.NewMessageItem(session.UserMessage("keep this conversation"))})) + before := sess.MessagesSnapshot() + events := make(chan Event, 32) + rt.Summarize(t.Context(), sess, "", NewChannelSink(events)) + close(events) + var failures int + for event := range events { + switch e := event.(type) { + case *ErrorEvent: + failures++ + assert.Contains(t, e.Error, "resolving compaction context window") + assert.Contains(t, e.Error, tt.want.Error()) + case *SessionCompactionEvent: + if e.Status == "completed" { + assert.Equal(t, CompactionOutcomeFailed, e.Outcome) + } + case *SessionSummaryEvent: + t.Error("unexpected summary") + } + } + assert.Equal(t, 1, failures) + assert.Equal(t, before, sess.MessagesSnapshot()) + assert.Zero(t, requests.Load(), "neither inference nor base discovery is allowed") + }) + } +} + +func TestSamplingCloneOptionsReachBothHandlers(t *testing.T) { + t.Parallel() + for _, withTools := range []bool{false, true} { + t.Run(map[bool]string{false: "without tools", true: "with tools"}[withTools], func(t *testing.T) { + t.Parallel() + clone := &mockProvider{id: "test/clone", stream: newStreamBuilder().AddContent("sampled").Build()} + p := &cloneFailureProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "base"}, RebuildProvider: func(_ context.Context, cfg *latest.ModelConfig, opts ...options.Opt) (provider.Provider, error) { + m := options.Apply(opts...) + assert.Nil(t, m.StructuredOutput()) + assert.True(t, m.NoThinking()) + assert.Equal(t, int64(512), m.MaxTokens()) + assert.Equal(t, int64(512), *cfg.MaxTokens) + return clone, nil + }}} + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(agent.New("root", "", agent.WithModel(p)))), WithCurrentAgent("root"), WithModelStore(mockModelStore{})) + require.NoError(t, err) + if withTools { + res, err := rt.samplingWithToolsHandler(t.Context(), &mcp.CreateMessageWithToolsParams{MaxTokens: 512, Messages: []*mcp.SamplingMessageV2{{Role: "user", Content: []mcp.Content{&mcp.TextContent{Text: "hello"}}}}}) + require.NoError(t, err) + assert.Equal(t, "test/clone", res.Model) + } else { + res, err := rt.samplingHandler(t.Context(), &mcp.CreateMessageParams{MaxTokens: 512, Messages: []*mcp.SamplingMessage{{Role: "user", Content: &mcp.TextContent{Text: "hello"}}}}) + require.NoError(t, err) + assert.Equal(t, "test/clone", res.Model) + } + }) + } +} diff --git a/pkg/runtime/compaction_context_limit_test.go b/pkg/runtime/compaction_context_limit_test.go index c3e03fe099..da8d2a7849 100644 --- a/pkg/runtime/compaction_context_limit_test.go +++ b/pkg/runtime/compaction_context_limit_test.go @@ -11,7 +11,9 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" + "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" @@ -32,9 +34,11 @@ func (p *providerOptsProvider) CreateChatCompletionStream(context.Context, []cha } func (p *providerOptsProvider) BaseConfig() base.Config { - return base.Config{ + cfg := base.Config{ ModelConfig: latest.ModelConfig{ProviderOpts: p.opts}, } + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil } + return cfg } func (p *providerOptsProvider) MaxTokens() int { return 0 } diff --git a/pkg/runtime/compactor/compactor.go b/pkg/runtime/compactor/compactor.go index 4efe79a89e..a0d20075f8 100644 --- a/pkg/runtime/compactor/compactor.go +++ b/pkg/runtime/compactor/compactor.go @@ -204,11 +204,15 @@ func RunLLM(ctx context.Context, args LLMArgs) (result *Result, err error) { return nil, errors.New("compactor: agent has no model") } - summaryModel := provider.CloneWithOptions(ctx, baseModel, + summaryModel, err := provider.Clone(ctx, baseModel, options.WithStructuredOutput(nil), options.WithMaxTokens(summaryTokenBudget(args.ContextLimit)), options.WithCompacting(), ) + if err != nil { + return nil, fmt.Errorf("compactor: cloning summary model: %w", err) + } + compactionAgent := agent.New(args.Agent.Name(), "", agent.WithModel(summaryModel)) messages, firstKeptEntry := extractMessages(args.Session, compactionAgent, args.ContextLimit, args.AdditionalPrompt) diff --git a/pkg/runtime/compactor/compactor_test.go b/pkg/runtime/compactor/compactor_test.go index 9547dceb89..63dcd53c9b 100644 --- a/pkg/runtime/compactor/compactor_test.go +++ b/pkg/runtime/compactor/compactor_test.go @@ -14,7 +14,10 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/compaction" + "github.com/docker/docker-agent/pkg/config/latest" + "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/session" "github.com/docker/docker-agent/pkg/tools" @@ -24,7 +27,9 @@ type fakeProvider struct{ id modelsdev.ID } func (p fakeProvider) ID() modelsdev.ID { return p.id } -func (p fakeProvider) BaseConfig() base.Config { return base.Config{} } +func (p fakeProvider) BaseConfig() base.Config { + return base.Config{RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil }} +} func (p fakeProvider) CreateChatCompletionStream( context.Context, @@ -755,3 +760,37 @@ func TestRunLLM_RequiresContextLimit(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "ContextLimit") } + +type failedCloneModel struct { + base.Config +} + +func (p *failedCloneModel) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + panic("base inference must not run") +} + +func TestRunLLMCloneFailureLeavesSessionIntact(t *testing.T) { + t.Parallel() + synthetic := errors.New("synthetic reconstruction failure") + for _, failure := range []error{nil, synthetic} { + p := &failedCloneModel{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "broken"}}} + if failure != nil { + p.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { + return nil, failure + } + } + sess := session.New(session.WithMessages([]session.Item{session.NewMessageItem(session.UserMessage("keep history"))})) + before := sess.MessagesSnapshot() + result, err := RunLLM(t.Context(), LLMArgs{Session: sess, Agent: agent.New("root", "", agent.WithModel(p)), ContextLimit: 100_000, RunAgent: func(context.Context, *agent.Agent, *session.Session) error { + t.Error("summary runtime must not run") + return nil + }}) + want := failure + if want == nil { + want = provider.ErrNoRebuilder + } + require.ErrorIs(t, err, want) + assert.Nil(t, result) + assert.Equal(t, before, sess.MessagesSnapshot()) + } +} diff --git a/pkg/runtime/fallback_test.go b/pkg/runtime/fallback_test.go index caea8936c3..01613f15c4 100644 --- a/pkg/runtime/fallback_test.go +++ b/pkg/runtime/fallback_test.go @@ -14,8 +14,10 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" "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/modelerrors" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/session" @@ -53,8 +55,13 @@ func (p *countingProvider) CreateChatCompletionStream(context.Context, []chat.Me } return p.stream, nil } -func (p *countingProvider) BaseConfig() base.Config { return base.Config{} } -func (p *countingProvider) MaxTokens() int { return 0 } + +func (p *countingProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil } + return cfg +} +func (p *countingProvider) MaxTokens() int { return 0 } // Verify interface compliance var ( diff --git a/pkg/runtime/live_sessions_test.go b/pkg/runtime/live_sessions_test.go index 771a40e8b1..be79a57661 100644 --- a/pkg/runtime/live_sessions_test.go +++ b/pkg/runtime/live_sessions_test.go @@ -14,7 +14,9 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/hooks" + "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/session" "github.com/docker/docker-agent/pkg/team" @@ -67,8 +69,12 @@ func (p *stepProvider) CreateChatCompletionStream(ctx context.Context, _ []chat. return step.stream, nil } -func (p *stepProvider) BaseConfig() base.Config { return base.Config{} } -func (p *stepProvider) MaxTokens() int { return 0 } +func (p *stepProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil } + return cfg +} +func (p *stepProvider) MaxTokens() int { return 0 } // newLiveSessionsRuntime builds a runtime with a "root" agent and a "worker" // agent whose model is the supplied provider. diff --git a/pkg/runtime/model_context_test.go b/pkg/runtime/model_context_test.go index d4342f447f..8959a95ab8 100644 --- a/pkg/runtime/model_context_test.go +++ b/pkg/runtime/model_context_test.go @@ -13,7 +13,9 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/config/latest" + "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" ) @@ -26,10 +28,12 @@ type contextProvider struct { func (p *contextProvider) ContextWindow(ctx context.Context) (int64, error) { return p.lookup(ctx) } func newContextProvider(endpoint string, lookup func(context.Context) (int64, error)) *contextProvider { - return &contextProvider{ + p := &contextProvider{ configProvider: configProvider{Config: base.Config{BaseURL: endpoint, ModelConfig: latest.ModelConfig{Provider: "dmr", Model: "ai/qwen3"}}}, lookup: lookup, } + p.SetProviderRebuilder(func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil }) + return p } func TestResolveContextLimitDiscovery(t *testing.T) { diff --git a/pkg/runtime/runtime_test.go b/pkg/runtime/runtime_test.go index b940bbe467..1e5a386f88 100644 --- a/pkg/runtime/runtime_test.go +++ b/pkg/runtime/runtime_test.go @@ -22,7 +22,9 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/hooks" + "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/modelerrors" "github.com/docker/docker-agent/pkg/modelinfo" "github.com/docker/docker-agent/pkg/modelsdev" @@ -231,7 +233,11 @@ func (m *mockProvider) CreateChatCompletionStream(context.Context, []chat.Messag return m.stream, nil } -func (m *mockProvider) BaseConfig() base.Config { return base.Config{} } +func (m *mockProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return m, nil } + return cfg +} func (m *mockProvider) MaxTokens() int { return 0 } @@ -906,7 +912,11 @@ func (p *queueProvider) CreateChatCompletionStream(context.Context, []chat.Messa return s, nil } -func (p *queueProvider) BaseConfig() base.Config { return base.Config{} } +func (p *queueProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil } + return cfg +} func (p *queueProvider) MaxTokens() int { return 0 } diff --git a/pkg/runtime/sampling.go b/pkg/runtime/sampling.go index 4b07498d3e..a694d4710c 100644 --- a/pkg/runtime/sampling.go +++ b/pkg/runtime/sampling.go @@ -77,7 +77,10 @@ func (r *LocalRuntime) samplingHandler(ctx context.Context, req *mcp.CreateMessa return nil, errors.New("current agent has no model configured") } - model := provider.CloneWithOptions(ctx, baseModel, samplingModelOptions(req)...) + model, err := provider.Clone(ctx, baseModel, samplingModelOptions(req)...) + if err != nil { + return nil, fmt.Errorf("cloning model for sampling: %w", err) + } stream, err := model.CreateChatCompletionStream(ctx, messages, nil) if err != nil { @@ -328,7 +331,10 @@ func (r *LocalRuntime) samplingWithToolsHandler(ctx context.Context, req *mcp.Cr return nil, errors.New("current agent has no model configured") } - model := provider.CloneWithOptions(ctx, baseModel, samplingModelOptionsFor(req.MaxTokens)...) + model, err := provider.Clone(ctx, baseModel, samplingModelOptionsFor(req.MaxTokens)...) + if err != nil { + return nil, fmt.Errorf("cloning model for sampling: %w", err) + } stream, err := model.CreateChatCompletionStream(ctx, messages, chatTools) if err != nil { diff --git a/pkg/runtime/session_compaction.go b/pkg/runtime/session_compaction.go index 2c5d36f383..da08d95f6b 100644 --- a/pkg/runtime/session_compaction.go +++ b/pkg/runtime/session_compaction.go @@ -69,7 +69,7 @@ func (r *LocalRuntime) sessionCompactionEnabled(a *agent.Agent) bool { // env / cwd / OS info is automatically present after a compaction without // any extra dispatch. func (r *LocalRuntime) doCompact(ctx context.Context, sess *session.Session, a *agent.Agent, additionalPrompt, reason string, events EventSink) { - contextLimit := r.compactionContextLimit(ctx, a) + contextLimit, limitErr := r.compactionContextLimitChecked(ctx, a) // before_compaction: hooks can veto or supply a custom summary. pre := r.executeBeforeCompactionHooks(ctx, sess, a, reason, contextLimit, events) @@ -108,6 +108,11 @@ func (r *LocalRuntime) doCompact(ctx context.Context, sess *session.Session, a * if native != nil { result, err = r.compactNatively(ctx, sess, a, native, additionalPrompt, events) } else { + if limitErr != nil { + events.Emit(ErrorForSession(sess.ID, limitErr.Error())) + outcome = CompactionOutcomeFailed + return + } if contextLimit <= 0 { slog.ErrorContext(ctx, "Failed to generate session summary", "error", "model definition unavailable") @@ -239,21 +244,32 @@ func summaryFromHook(sess *session.Session, a *agent.Agent, pre *hooks.Result, c // pass the cloned summary-call provider so its provider_opts (which // match the underlying model) are considered. func (r *LocalRuntime) compactionContextLimit(ctx context.Context, a *agent.Agent) int64 { + limit, err := r.compactionContextLimitChecked(ctx, a) + if err != nil { + slog.DebugContext(ctx, "Compaction context window unresolvable", "error", err) + } + return limit +} + +func (r *LocalRuntime) compactionContextLimitChecked(ctx context.Context, a *agent.Agent) (int64, error) { if a == nil { - return 0 + return 0, nil } model := a.CompactionModel() if model == nil { model = a.Model(ctx) } if model == nil { - return 0 + return 0, nil } - summaryModel := provider.CloneWithOptions(ctx, model, + summaryModel, err := provider.Clone(ctx, model, options.WithStructuredOutput(nil), options.WithMaxTokens(compactor.MaxSummaryTokens), ) - return r.resolveContextLimit(ctx, summaryModel, summaryModel.ID()) + if err != nil { + return 0, fmt.Errorf("resolving compaction context window: %w", err) + } + return r.resolveContextLimit(ctx, summaryModel, summaryModel.ID()), nil } // effectiveContextLimit returns the context budget the running session diff --git a/pkg/runtime/session_compaction_test.go b/pkg/runtime/session_compaction_test.go index 14421cbeb5..5558785087 100644 --- a/pkg/runtime/session_compaction_test.go +++ b/pkg/runtime/session_compaction_test.go @@ -15,6 +15,7 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" + providerbase "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/session/sqlitestore" @@ -188,7 +189,7 @@ func TestDoCompactBeforeHookSuppliesSummary(t *testing.T) { // The provider must NOT be called — if it is, we'll consume from // the (empty) mockStream and the test will catch it. - prov := &mockProvider{id: "test/mock-model", stream: &mockStream{}} + prov := &cloneFailureProvider{Config: providerbase.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "unclonable"}}} root := agent.New("root", "test", agent.WithModel(prov), agent.WithHooks(hookCfg), diff --git a/pkg/runtime/title_generator_test.go b/pkg/runtime/title_generator_test.go index 2cc2f32349..60afa6a8c1 100644 --- a/pkg/runtime/title_generator_test.go +++ b/pkg/runtime/title_generator_test.go @@ -13,7 +13,9 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" + modelprovider "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/session" "github.com/docker/docker-agent/pkg/team" @@ -37,9 +39,13 @@ func (p *imageOutputProvider) CreateChatCompletionStream(context.Context, []chat } func (p *imageOutputProvider) BaseConfig() base.Config { - return base.Config{ModelConfig: latest.ModelConfig{ + cfg := base.Config{ModelConfig: latest.ModelConfig{ OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, }} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (modelprovider.Provider, error) { + return p, nil + } + return cfg } func (p *imageOutputProvider) MaxTokens() int { return 0 } diff --git a/pkg/runtime/transforms_test.go b/pkg/runtime/transforms_test.go index dc3beeb059..f497086b5c 100644 --- a/pkg/runtime/transforms_test.go +++ b/pkg/runtime/transforms_test.go @@ -18,7 +18,9 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/hooks" + "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/modelinfo" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/paths" @@ -53,7 +55,11 @@ type recordingMsgProvider struct { baseConfig base.Config } -func (p *recordingMsgProvider) BaseConfig() base.Config { return p.baseConfig } +func (p *recordingMsgProvider) BaseConfig() base.Config { + cfg := p.baseConfig + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil } + return cfg +} func (p *recordingMsgProvider) CreateChatCompletionStream(_ context.Context, msgs []chat.Message, _ []tools.Tool) (chat.MessageStream, error) { snapshot := slices.Clone(msgs) diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 090d0751bb..847989d370 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -27,8 +27,11 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/concurrent" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/config/types" + modelprovider "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/runtime" "github.com/docker/docker-agent/pkg/session" @@ -103,7 +106,13 @@ func (p *delayedTitleProvider) CreateChatCompletionStream(context.Context, []cha return &delayedTitleStream{release: p.release}, nil } -func (p *delayedTitleProvider) BaseConfig() base.Config { return base.Config{} } +func (p *delayedTitleProvider) BaseConfig() base.Config { + cfg := base.Config{} + cfg.RebuildProvider = func(context.Context, *latest.ModelConfig, ...options.Opt) (modelprovider.Provider, error) { + return p, nil + } + return cfg +} type delayedTitleStream struct { release chan struct{} diff --git a/pkg/sessiontitle/generator.go b/pkg/sessiontitle/generator.go index 8665949748..169de25929 100644 --- a/pkg/sessiontitle/generator.go +++ b/pkg/sessiontitle/generator.go @@ -140,7 +140,7 @@ func (g *Generator) Generate(ctx context.Context, sessionID string, userMessages func generateOnce(ctx context.Context, baseModel provider.Provider, messages []chat.Message) (string, error) { // Clone the model with title-generation-specific options so each attempt // gets a consistent, low-token one-shot call. - titleModel := provider.CloneWithOptions( + titleModel, err := provider.Clone( ctx, baseModel, options.WithStructuredOutput(nil), @@ -148,6 +148,9 @@ func generateOnce(ctx context.Context, baseModel provider.Provider, messages []c options.WithNoThinking(), options.WithGeneratingTitle(), ) + if err != nil { + return "", fmt.Errorf("model %q: cloning title model: %w", baseModel.ID(), err) + } stream, err := titleModel.CreateChatCompletionStream(ctx, messages, nil) if err != nil { diff --git a/pkg/sessiontitle/generator_test.go b/pkg/sessiontitle/generator_test.go index c525832117..d537b72e9c 100644 --- a/pkg/sessiontitle/generator_test.go +++ b/pkg/sessiontitle/generator_test.go @@ -11,6 +11,7 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" + "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" @@ -39,7 +40,7 @@ func (p *mockProvider) BaseConfig() base.Config { if p.baseCfgFn != nil { return p.baseCfgFn() } - return base.Config{} + return base.Config{RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil }} } type mockStream struct { @@ -80,22 +81,22 @@ func streamWithContent(content string) chat.MessageStream { func TestGenerateOnceUsesTitleHeadroomAndClearsStructuredOutput(t *testing.T) { t.Parallel() + clone := &mockProvider{id: modelsdev.NewID("google", "gemini-3-flash"), createFn: func() (chat.MessageStream, error) { return streamWithContent("Title"), nil }} structured := &latest.StructuredOutput{Schema: map[string]any{"type": "object"}} - baseProvider := &mockProvider{ - id: modelsdev.NewID("google", "gemini-3-flash"), - baseCfgFn: func() base.Config { - maxTokens := int64(7) - return base.Config{ - ModelConfig: latest.ModelConfig{MaxTokens: &maxTokens}, - ModelOptions: options.Apply(options.WithStructuredOutput(structured)), - } - }, - createFn: func() (chat.MessageStream, error) { return streamWithContent("Title"), nil }, - } - + baseProvider := &mockProvider{id: clone.id, baseCfgFn: func() base.Config { + return base.Config{ModelOptions: options.Apply(options.WithStructuredOutput(structured)), RebuildProvider: func(_ context.Context, cfg *latest.ModelConfig, opts ...options.Opt) (provider.Provider, error) { + assert.Equal(t, int64(titleMaxTokens), *cfg.MaxTokens) + m := options.Apply(opts...) + assert.Nil(t, m.StructuredOutput()) + assert.True(t, m.NoThinking()) + assert.True(t, m.GeneratingTitle()) + return clone, nil + }} + }} _, err := generateOnce(t.Context(), baseProvider, buildPrompt([]string{"hello"})) require.NoError(t, err) - assert.Equal(t, 1, baseProvider.calls, "a failed clone would call the base provider and lose title-specific options") + assert.Zero(t, baseProvider.calls) + assert.Equal(t, 1, clone.calls) assert.Equal(t, 128, titleMaxTokens) } @@ -216,17 +217,16 @@ func TestGenerator_Generate_FallsBackOnEmptyOutput(t *testing.T) { // imageOutputProvider is an image-capable title candidate. func imageOutputProvider(id modelsdev.ID) *mockProvider { - return &mockProvider{ + p := &mockProvider{ id: id, createFn: func() (chat.MessageStream, error) { return nil, errors.New("image-output title attempt failed") }, - baseCfgFn: func() base.Config { - return base.Config{ModelConfig: latest.ModelConfig{ - OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, - }} - }, } + p.baseCfgFn = func() base.Config { + return base.Config{ModelConfig: latest.ModelConfig{OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}}, RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { return p, nil }} + } + return p } func TestNew_ImageOutputCandidateFallsThroughOnFailure(t *testing.T) { @@ -304,3 +304,21 @@ func TestNew_KeepsUndeclaredAndExplicitFalseCandidates(t *testing.T) { assert.Equal(t, 1, undeclared.calls) assert.Zero(t, declaredFalse.calls) } + +func TestGeneratorCloneFailuresExhaustCandidates(t *testing.T) { + t.Parallel() + synthetic := errors.New("synthetic rebuild failure") + missing := &mockProvider{id: modelsdev.NewID("test", "missing"), baseCfgFn: func() base.Config { return base.Config{} }} + broken := &mockProvider{id: modelsdev.NewID("test", "broken"), baseCfgFn: func() base.Config { + return base.Config{RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (provider.Provider, error) { + return nil, synthetic + }} + }} + title, err := New(missing, broken).Generate(t.Context(), "session", []string{"hello"}) + assert.Empty(t, title) + require.ErrorIs(t, err, provider.ErrNoRebuilder) + require.ErrorIs(t, err, synthetic) + require.ErrorContains(t, err, "all 2 title model(s) failed") + assert.Zero(t, missing.calls) + assert.Zero(t, broken.calls) +} From 979a07fc9041ba2da36b7300ecb542b1195cf06f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 16:18:11 +0000 Subject: [PATCH 2/5] fix(provider): isolate routed clones and retain thinking-off intent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/guides/go-sdk/index.md | 9 +- pkg/model/provider/clone.go | 22 +++- pkg/model/provider/clone_matrix_test.go | 43 +++++++ pkg/model/provider/clone_routed_test.go | 89 +++++++++++++ pkg/model/provider/clone_thinking_test.go | 150 ++++++++++++++++++++++ pkg/model/provider/factory.go | 53 +++++++- 6 files changed, 356 insertions(+), 10 deletions(-) create mode 100644 pkg/model/provider/clone_routed_test.go create mode 100644 pkg/model/provider/clone_thinking_test.go diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index 23ff25cbde..00bb908f7c 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -880,7 +880,14 @@ JSON-serializable: unsupported values in `provider_opts`, schemas or custom providers fail closed with `provider.ErrUncopyableConfig`, rather than panic or share mutable state. Cumulative clones preserve explicit structured-output removal, gateway clearing and no-thinking overrides; existing provider token -floors and routed-child token limits still apply. +floors and routed-child token limits still apply. Named routing definitions are +captured at construction and copied for each child, so later input-map mutations +do not change clones. + +Configuration-declared thinking off (`none` or `0`, including provider defaults) +is retained as an option before normalization. Reconstruction applies that option +to each leaf's config too, so config-driven and option-driven providers agree. +Routed children inherit strict-off overrides but keep their own token limits. `CloneWithOptions` retains its public signature but is deprecated. On failure it returns an error-only provider whose stream creation returns the clone error. It diff --git a/pkg/model/provider/clone.go b/pkg/model/provider/clone.go index 132854f67a..914496c70c 100644 --- a/pkg/model/provider/clone.go +++ b/pkg/model/provider/clone.go @@ -107,15 +107,10 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, mergedOpts = options.FromModelOptions(merged) source := cfg.CloneSource() - var modelConfig latest.ModelConfig - data, err := json.Marshal(source) + modelConfig, err := copyModelConfig(source) if err != nil { - return latest.ModelConfig{}, nil, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) - } - if err := json.Unmarshal(data, &modelConfig); err != nil { - return latest.ModelConfig{}, nil, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) + return latest.ModelConfig{}, nil, err } - modelConfig.Name, modelConfig.DisplayModel = source.Name, source.DisplayModel if mt := merged.MaxTokens(); mt != 0 { modelConfig.MaxTokens = &mt @@ -128,3 +123,16 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, } return modelConfig, mergedOpts, nil } + +func copyModelConfig(source latest.ModelConfig) (latest.ModelConfig, error) { + var copied latest.ModelConfig + data, err := json.Marshal(source) + if err != nil { + return copied, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) + } + if err := json.Unmarshal(data, &copied); err != nil { + return copied, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) + } + copied.Name, copied.DisplayModel = source.Name, source.DisplayModel + return copied, nil +} diff --git a/pkg/model/provider/clone_matrix_test.go b/pkg/model/provider/clone_matrix_test.go index 9a1e6fba04..f45334f72f 100644 --- a/pkg/model/provider/clone_matrix_test.go +++ b/pkg/model/provider/clone_matrix_test.go @@ -154,3 +154,46 @@ func TestClonePreservesStructuredOutputOnWire(t *testing.T) { assert.Equal(t, "counter", format["name"]) assert.InDelta(t, float64(1), format["schema"].(map[string]any)["properties"].(map[string]any)["count"].(map[string]any)["minimum"], 0) } + +func TestCloneConfiguredOpenAIThinkingOffOnWire(t *testing.T) { + t.Parallel() + for _, model := range []string{"o3", "gpt-4o"} { + for _, shape := range []string{"literal", "custom", "routed"} { + for _, budget := range []*latest.ThinkingBudget{{Effort: "none"}, {}} { + t.Run(model+"/"+shape+"/"+budget.Effort, func(t *testing.T) { + t.Parallel() + server, body := captureNamedCustomProviderRequestBody(t, "openai_chatcompletions") + cfg := &latest.ModelConfig{Provider: "openai", Model: model, BaseURL: server.URL, ThinkingBudget: budget, ProviderOpts: map[string]any{"api_type": "openai_chatcompletions"}} + providers := map[string]latest.ProviderConfig{"openai": {ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}}} + models := map[string]latest.ModelConfig{} + if shape != "literal" { + cfg.Provider = "custom" + cfg.BaseURL = "" + providers["custom"] = latest.ProviderConfig{BaseURL: server.URL, TokenKey: "OPENAI_API_KEY", APIType: "openai_chatcompletions", ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}} + } + if shape == "routed" { + models["target"] = *cfg + cfg.Routing = []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}} + } + p, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "fake"}), options.WithProviders(providers)) + require.NoError(t, err) + for range 3 { + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + var payload map[string]any + require.NoError(t, json.Unmarshal(body(), &payload)) + if model == "o3" { + assert.Equal(t, "low", payload["reasoning_effort"]) + } else { + assert.NotContains(t, payload, "reasoning_effort") + } + p, err = Clone(t.Context(), p, options.WithCompacting(), options.WithStructuredOutput(nil), options.WithMaxTokens(8192)) + require.NoError(t, err) + } + }) + } + } + } +} diff --git a/pkg/model/provider/clone_routed_test.go b/pkg/model/provider/clone_routed_test.go new file mode 100644 index 0000000000..244b512f01 --- /dev/null +++ b/pkg/model/provider/clone_routed_test.go @@ -0,0 +1,89 @@ +package provider + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/options" +) + +func TestCloneRoutedConfigurationsAreIsolated(t *testing.T) { + t.Parallel() + var children []*instrumentTestProvider + registry := NewRegistry(map[string]Factory{"test": func(_ context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) { + p := &instrumentTestProvider{Config: base.Config{ModelConfig: *cfg, ModelOptions: options.Apply(opts...), Env: env}} + if cfg.Model == "child" { + children = append(children, p) + } + return p, nil + }}) + models := map[string]latest.ModelConfig{ + "target": {Provider: "test", Model: "child", Name: "named", DisplayModel: "display", ThinkingBudget: &latest.ThinkingBudget{Tokens: 8192}, MaxTokens: new(int64(16384)), ProviderOpts: map[string]any{"nested": map[string]any{"value": "initial"}}}, + "unused": {ProviderOpts: map[string]any{"bad": make(chan int)}}, + } + cfg := &latest.ModelConfig{Provider: "test", Model: "fallback", Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + original, err := registry.NewWithModels(t.Context(), cfg, models, environment.NewNoEnvProvider()) + require.NoError(t, err) + models["target"].ProviderOpts["nested"].(map[string]any)["value"] = "input changed" + models["target"].ThinkingBudget.Tokens = 1 + first, err := Clone(t.Context(), original, options.WithMaxTokens(128)) + require.NoError(t, err) + second, err := Clone(t.Context(), first) + require.NoError(t, err) + require.NotNil(t, second) + require.Len(t, children, 3) + for _, child := range children { + assert.Equal(t, "initial", child.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"]) + assert.Equal(t, 8192, child.ModelConfig.ThinkingBudget.Tokens) + assert.Equal(t, int64(16384), *child.ModelConfig.MaxTokens) + assert.Equal(t, "named", child.ModelConfig.Name) + assert.Equal(t, "display", child.ModelConfig.DisplayModel) + } + for i, child := range children { + child.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"] = i + child.ModelConfig.ThinkingBudget.Tokens = i + *child.ModelConfig.MaxTokens = int64(i) + } + for i, child := range children { + assert.Equal(t, i, child.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"]) + assert.Equal(t, i, child.ModelConfig.ThinkingBudget.Tokens) + assert.Equal(t, int64(i), *child.ModelConfig.MaxTokens) + } + assert.Equal(t, "input changed", models["target"].ProviderOpts["nested"].(map[string]any)["value"]) + assert.Equal(t, 1, models["target"].ThinkingBudget.Tokens) + models["target"] = latest.ModelConfig{Provider: "test", Model: "replacement"} + _, err = Clone(t.Context(), second) + require.NoError(t, err) + require.Len(t, children, 4) + assert.Equal(t, "initial", children[3].ModelConfig.ProviderOpts["nested"].(map[string]any)["value"]) + assert.Equal(t, 8192, children[3].ModelConfig.ThinkingBudget.Tokens) +} + +func TestCloneRejectsUncopyableRoutedChild(t *testing.T) { + t.Parallel() + calls := 0 + registry := NewRegistry(map[string]Factory{"test": func(context.Context, *latest.ModelConfig, environment.Provider, ...options.Opt) (Provider, error) { + calls++ + return &instrumentTestProvider{}, nil + }}) + cfg := latest.ModelConfig{Provider: "test", Model: "fallback", Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + models := map[string]latest.ModelConfig{"target": {Provider: "test", Model: "child", ProviderOpts: map[string]any{"bad": make(chan int)}}} + p, err := registry.NewWithModels(t.Context(), &cfg, models, environment.NewNoEnvProvider()) + require.ErrorIs(t, err, ErrUncopyableConfig) + require.ErrorContains(t, err, `routed model "target"`) + assert.Nil(t, p) + assert.Zero(t, calls) + original := &instrumentTestProvider{Config: base.Config{ModelConfig: cfg}} + registry.attachRebuilder(original, models, environment.NewNoEnvProvider()) + clone, err := Clone(t.Context(), original) + require.ErrorIs(t, err, ErrUncopyableConfig) + require.ErrorContains(t, err, `routed model "target"`) + assert.Nil(t, clone) + assert.Zero(t, calls) +} diff --git a/pkg/model/provider/clone_thinking_test.go b/pkg/model/provider/clone_thinking_test.go new file mode 100644 index 0000000000..9d3cfdfb55 --- /dev/null +++ b/pkg/model/provider/clone_thinking_test.go @@ -0,0 +1,150 @@ +//go:build !js + +package provider + +import ( + "encoding/json" + "io" + "net/http" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/options" +) + +func TestCloneRoutedClaudeStrictOffOnWire(t *testing.T) { + t.Parallel() + for _, provider := range []string{"anthropic", "amazon-bedrock"} { + t.Run(provider, func(t *testing.T) { + t.Parallel() + var body map[string]any + calls := 0 + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return thinkingRoundTripper(func(req *http.Request) (*http.Response, error) { + calls++ + body = nil + require.NoError(t, json.NewDecoder(req.Body).Decode(&body)) + return &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"offline capture"}}`)), Request: req}, nil + }) + }) + model := "claude-haiku-4-5" + if provider == "amazon-bedrock" { + model = "anthropic.claude-haiku-4-5-20251001-v1:0" + } + models := map[string]latest.ModelConfig{"target": {Provider: provider, Model: model, MaxTokens: new(int64(16384)), ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"region": "us-east-1"}}} + cfg := &latest.ModelConfig{Provider: provider, Model: model, Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + env := environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "fake", "AWS_BEARER_TOKEN_BEDROCK": "fake"}) + original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, env, wrapper) + require.NoError(t, err) + request := func(p Provider, thinking bool) { + t.Helper() + before := calls + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + if err == nil { + _, err = stream.Recv() + stream.Close() + } + require.Error(t, err) + require.Equal(t, before+1, calls) + fields := body + if provider == "amazon-bedrock" { + if additional, exists := body["additionalModelRequestFields"]; exists { + fields = additional.(map[string]any) + } + assert.InDelta(t, float64(16384), body["inferenceConfig"].(map[string]any)["maxTokens"], 0) + } else { + assert.InDelta(t, float64(8192), body["max_tokens"], 0) + } + if thinking { + require.Contains(t, fields, "thinking") + assert.InDelta(t, float64(2048), fields["thinking"].(map[string]any)["budget_tokens"], 0) + } else { + assert.NotContains(t, fields, "thinking") + assert.NotContains(t, fields, "reasoning_config") + } + } + request(original, true) + for _, purpose := range []options.Opt{options.WithGeneratingTitle(), options.WithStructuredOutput(nil), options.WithCompacting()} { + clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128), purpose) + require.NoError(t, err) + clone, err = Clone(t.Context(), clone) + require.NoError(t, err) + request(clone, false) + } + request(original, true) + }) + } +} + +type thinkingRoundTripper func(*http.Request) (*http.Response, error) + +func (f thinkingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } + +func TestCloneConfiguredThinkingOffOnWire(t *testing.T) { + t.Parallel() + for _, shape := range []string{"literal-none", "literal-zero", "environment", "provider-none", "provider-zero", "routed-none", "routed-zero", "model-wins"} { + t.Run(shape, func(t *testing.T) { + t.Parallel() + var body map[string]any + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return thinkingRoundTripper(func(req *http.Request) (*http.Response, error) { + body = nil + require.NoError(t, json.NewDecoder(req.Body).Decode(&body)) + return &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"offline capture"}}`)), Request: req}, nil + }) + }) + cfg := &latest.ModelConfig{Provider: "custom", Model: "claude-haiku-4-5", ThinkingBudget: &latest.ThinkingBudget{Effort: "none"}} + pc := latest.ProviderConfig{Provider: "anthropic", BaseURL: "https://offline.invalid", TokenKey: "KEY", ThinkingBudget: &latest.ThinkingBudget{Tokens: 8192}} + switch shape { + case "literal-zero", "routed-zero": + cfg.ThinkingBudget = &latest.ThinkingBudget{} + case "environment": + cfg.Model = "${MODEL}" + pc.BaseURL = "${ENDPOINT}" + case "provider-none", "provider-zero": + cfg.ThinkingBudget = nil + pc.ThinkingBudget = &latest.ThinkingBudget{} + if shape == "provider-none" { + pc.ThinkingBudget.Effort = "none" + } + case "model-wins": + cfg.ThinkingBudget = &latest.ThinkingBudget{Tokens: 8192} + pc.ThinkingBudget = &latest.ThinkingBudget{Effort: "none"} + } + models := map[string]latest.ModelConfig{} + if strings.HasPrefix(shape, "routed-") { + models["target"] = *cfg + cfg.Routing = []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}} + } + original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"KEY": "fake", "MODEL": "claude-haiku-4-5", "ENDPOINT": "https://offline.invalid"}), wrapper, options.WithProviders(map[string]latest.ProviderConfig{"custom": pc})) + require.NoError(t, err) + p := original + for range 3 { + baseCfg := p.BaseConfig() + if len(baseCfg.ModelConfig.Routing) == 0 { + assert.Equal(t, shape != "model-wins", baseCfg.ModelOptions.NoThinking()) + } + stream, requestErr := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + if requestErr == nil { + _, requestErr = stream.Recv() + stream.Close() + } + require.Error(t, requestErr) + if shape == "model-wins" { + require.Contains(t, body, "thinking") + assert.InDelta(t, float64(8192), body["thinking"].(map[string]any)["budget_tokens"], 0) + } else { + assert.NotContains(t, body, "thinking") + } + p, err = Clone(t.Context(), p, options.WithStructuredOutput(nil), options.WithMaxTokens(16384), options.WithCompacting()) + require.NoError(t, err) + } + }) + } +} diff --git a/pkg/model/provider/factory.go b/pkg/model/provider/factory.go index 51fbf55b13..615a0108e6 100644 --- a/pkg/model/provider/factory.go +++ b/pkg/model/provider/factory.go @@ -33,6 +33,10 @@ func (r *Registry) New(ctx context.Context, cfg *latest.ModelConfig, env environ 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 { + models, err := copyRoutedModels(cfg, models) + if err != nil { + return nil, err + } // A router makes no HTTP calls itself; its fallback and routed targets // do. When the router model opts out of the gateway, propagate the // bypass to every child so the whole routing subtree dials directly. @@ -54,6 +58,28 @@ func (r *Registry) NewWithModels(ctx context.Context, cfg *latest.ModelConfig, m return r.createDirectProvider(ctx, cfg, env, opts...) } +// Snapshot only referenced definitions; the rebuilder must not capture caller-owned state. +func copyRoutedModels(cfg *latest.ModelConfig, models map[string]latest.ModelConfig) (map[string]latest.ModelConfig, error) { + names := []string{cfg.Provider + "/" + cfg.Model} + for _, rule := range cfg.Routing { + names = append(names, rule.Model) + } + copied := make(map[string]latest.ModelConfig) + for _, name := range names { + if _, exists := copied[name]; exists { + continue + } + if model, exists := models[name]; exists { + clone, err := copyModelConfig(model) + if err != nil { + return nil, fmt.Errorf("copying routed model %q: %w", name, err) + } + copied[name] = clone + } + } + return copied, nil +} + 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...) } @@ -63,7 +89,11 @@ func (r *Registry) resolveRoutedModel(ctx context.Context, modelSpec string, mod if len(modelCfg.Routing) > 0 { return nil, fmt.Errorf("model %q has routing rules and cannot be used as a routing target", modelSpec) } - return r.createDirectProvider(ctx, &modelCfg, env, factoryOpts...) + copied, err := copyModelConfig(modelCfg) + if err != nil { + return nil, fmt.Errorf("copying routed model %q: %w", modelSpec, err) + } + return r.createDirectProvider(ctx, &copied, env, factoryOpts...) } inlineCfg, parseErr := latest.ParseModelRef(modelSpec) if parseErr != nil { @@ -77,7 +107,26 @@ func (r *Registry) createDirectProvider(ctx context.Context, cfg *latest.ModelCo return nil, errors.New("provider registry is required") } globalOptions := options.Apply(opts...) - enhancedCfg := applyProviderDefaults(cfg, globalOptions.Providers()) + // Capture explicit off before defaults normalize it to nil. Model settings win. + budget := cfg.ThinkingBudget + if budget == nil { + budget = globalOptions.Providers()[cfg.Provider].ThinkingBudget + } + noThinking := globalOptions.NoThinking() || budget.IsDisabled() + constructionCfg := cfg + // Config-driven leaves need the same override as option-driven leaves. + if globalOptions.NoThinking() { + copied, err := copyModelConfig(*cfg) + if err != nil { + return nil, err + } + copied.ThinkingBudget = &latest.ThinkingBudget{Effort: "none"} + constructionCfg = &copied + } + if noThinking && !globalOptions.NoThinking() { + opts = append(opts, options.WithNoThinking()) + } + enhancedCfg := applyProviderDefaults(constructionCfg, globalOptions.Providers()) if err := expandModelConfigEnv(ctx, enhancedCfg, env); err != nil { return nil, err } From 755e1a4c2ba4d56de0a636f6a666ff2a6c33f8af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Tue, 13 Oct 2026 17:05:15 +0000 Subject: [PATCH 3/5] fix(provider): preserve numeric precision in checked clones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/guides/go-sdk/index.md | 10 +- pkg/model/provider/clone.go | 15 +- pkg/model/provider/clone_matrix_test.go | 29 +++- pkg/model/provider/clone_numeric_test.go | 122 ++++++++++++++ pkg/model/provider/clone_vertex_test.go | 83 ++++++++-- pkg/model/provider/contracts/contracts.go | 4 +- .../provider/internal/deepcopy/deepcopy.go | 150 ++++++++++++++++++ .../internal/deepcopy/deepcopy_test.go | 102 ++++++++++++ pkg/model/provider/options/options.go | 27 +--- pkg/model/provider/options/options_test.go | 25 ++- pkg/model/provider/vertexai/modelgarden.go | 31 +++- .../provider/vertexai/modelgarden_test.go | 6 +- 12 files changed, 543 insertions(+), 61 deletions(-) create mode 100644 pkg/model/provider/clone_numeric_test.go create mode 100644 pkg/model/provider/internal/deepcopy/deepcopy.go create mode 100644 pkg/model/provider/internal/deepcopy/deepcopy_test.go diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index 00bb908f7c..d47a47265e 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -875,10 +875,12 @@ reconstruction errors are wrapped in `*provider.CloneError` and support Clones isolate mutable model configuration, custom-provider definitions and structured-output schemas. Transport wrappers, token-source callbacks and the models.dev store remain shared services. Public option setters retain reference -semantics; isolation happens when cloning. Programmatic configurations must be -JSON-serializable: unsupported values in `provider_opts`, schemas or custom -providers fail closed with `provider.ErrUncopyableConfig`, rather than panic or -share mutable state. Cumulative clones preserve explicit structured-output +semantics; isolation happens when cloning. Scalar types and exact numeric values are preserved. Values in `provider_opts`, +schemas and custom providers must be copyable scalars, maps, slices, arrays or +plain structs. Functions, channels, unsafe pointers, cyclic values and structs +with unexported reference-bearing fields fail closed with +`provider.ErrUncopyableConfig`, rather than panic or share mutable state. Copies +are bounded to 128 levels and 100,000 visited values. Cumulative clones preserve explicit structured-output removal, gateway clearing and no-thinking overrides; existing provider token floors and routed-child token limits still apply. Named routing definitions are captured at construction and copied for each child, so later input-map mutations diff --git a/pkg/model/provider/clone.go b/pkg/model/provider/clone.go index 914496c70c..ab438156f2 100644 --- a/pkg/model/provider/clone.go +++ b/pkg/model/provider/clone.go @@ -2,7 +2,6 @@ package provider import ( "context" - "encoding/json" "errors" "fmt" "log/slog" @@ -11,6 +10,7 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/internal/deepcopy" "github.com/docker/docker-agent/pkg/model/provider/options" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/tools" @@ -19,8 +19,8 @@ import ( // ErrNoRebuilder means a provider has no construction function for cloning. var ErrNoRebuilder = errors.New("provider cannot be cloned without its rebuild function") -// ErrUncopyableConfig means mutable configuration contains non-JSON values. -var ErrUncopyableConfig = errors.New("provider clone requires JSON-serializable configuration") +// ErrUncopyableConfig means mutable configuration cannot be safely copied. +var ErrUncopyableConfig = errors.New("provider clone requires copyable configuration") // CloneError identifies the provider whose reconstruction failed. type CloneError struct { @@ -125,14 +125,9 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, } func copyModelConfig(source latest.ModelConfig) (latest.ModelConfig, error) { - var copied latest.ModelConfig - data, err := json.Marshal(source) + copied, err := deepcopy.Value(source) if err != nil { - return copied, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) + return latest.ModelConfig{}, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) } - if err := json.Unmarshal(data, &copied); err != nil { - return copied, fmt.Errorf("%w: copying model configuration: %w", ErrUncopyableConfig, err) - } - copied.Name, copied.DisplayModel = source.Name, source.DisplayModel return copied, nil } diff --git a/pkg/model/provider/clone_matrix_test.go b/pkg/model/provider/clone_matrix_test.go index f45334f72f..2c7d365f23 100644 --- a/pkg/model/provider/clone_matrix_test.go +++ b/pkg/model/provider/clone_matrix_test.go @@ -3,9 +3,9 @@ package provider import ( + "bytes" "context" "encoding/json" - "math" "reflect" "testing" @@ -74,7 +74,9 @@ func TestCloneOptionMatrixOnWire(t *testing.T) { func TestCloneRejectsUncopyableConfigWithoutPanic(t *testing.T) { t.Parallel() - for _, value := range []any{math.Inf(1), func() {}, make(chan int)} { + cycle := map[string]any{} + cycle["self"] = cycle + for _, value := range []any{func() {}, make(chan int), cycle} { for _, field := range []string{"provider_opts", "structured_output", "providers"} { cfg := base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "base"}, RebuildProvider: func(context.Context, *latest.ModelConfig, ...options.Opt) (Provider, error) { t.Error("must reject before rebuilding") @@ -197,3 +199,26 @@ func TestCloneConfiguredOpenAIThinkingOffOnWire(t *testing.T) { } } } + +func TestCloneNumericPrecisionOnWire(t *testing.T) { + t.Parallel() + server, body := captureNamedCustomProviderRequestBody(t, "openai_chatcompletions") + schema := map[string]any{"type": "object", "properties": map[string]any{"count": map[string]any{"type": "integer", "minimum": int64(9007199254740993)}}} + original, err := fullTestRegistry().New(t.Context(), &latest.ModelConfig{Provider: "openai", Model: "gpt-4o", BaseURL: server.URL, ProviderOpts: map[string]any{"api_type": "openai_chatcompletions", "extra_body": map[string]any{"exact_limit": int64(9007199254740993)}}}, environment.NewMapEnvProvider(map[string]string{"OPENAI_API_KEY": "secret"}), options.WithStructuredOutput(&latest.StructuredOutput{Name: "counter", Schema: schema})) + require.NoError(t, err) + for range 3 { + original, err = Clone(t.Context(), original) + require.NoError(t, err) + stream, err := original.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + var payload map[string]any + decoder := json.NewDecoder(bytes.NewReader(body())) + decoder.UseNumber() + require.NoError(t, decoder.Decode(&payload)) + assert.Equal(t, json.Number("9007199254740993"), payload["exact_limit"]) + bound := payload["response_format"].(map[string]any)["json_schema"].(map[string]any)["schema"].(map[string]any)["properties"].(map[string]any)["count"].(map[string]any)["minimum"] + assert.Equal(t, json.Number("9007199254740993"), bound) + } +} diff --git a/pkg/model/provider/clone_numeric_test.go b/pkg/model/provider/clone_numeric_test.go new file mode 100644 index 0000000000..66e2594d98 --- /dev/null +++ b/pkg/model/provider/clone_numeric_test.go @@ -0,0 +1,122 @@ +package provider + +import ( + "context" + "math" + "testing" + + "github.com/goccy/go-yaml" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/model/provider/providerutil" +) + +func cloneNumericOptions() map[string]any { + return map[string]any{ + "int64": int64(9007199254740993), "uint64": uint64(math.MaxUint64), + "int8": int8(-7), "float32": float32(0.1), "top_k": uint64(40), + "extra_body": map[string]any{"limit": int64(9007199254740993), "items": []any{int8(-7), float32(0.1)}}, + } +} + +func TestClonePreservesNumericTypesCumulatively(t *testing.T) { + t.Parallel() + source := cloneNumericOptions() + var p Provider = &instrumentTestProvider{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "test", Model: "base", ProviderOpts: source}}} + var rebuild base.RebuildProviderFunc + rebuild = func(_ context.Context, cfg *latest.ModelConfig, opts ...options.Opt) (Provider, error) { + return &instrumentTestProvider{Config: base.Config{ModelConfig: *cfg, ModelOptions: options.Apply(opts...), RebuildProvider: rebuild}}, nil + } + p.(*instrumentTestProvider).RebuildProvider = rebuild + for range 3 { + clone, err := Clone(t.Context(), p) + require.NoError(t, err) + copied := clone.BaseConfig().ModelConfig.ProviderOpts + require.Equal(t, source, copied) + for key, value := range source { + assert.IsType(t, value, copied[key]) + } + baseTopK, baseOK := providerutil.GetProviderOptInt64(source, "top_k") + cloneTopK, cloneOK := providerutil.GetProviderOptInt64(copied, "top_k") + assert.Equal(t, baseTopK, cloneTopK) + assert.Equal(t, baseOK, cloneOK) + p = clone + } + copied := p.BaseConfig().ModelConfig.ProviderOpts + copied["extra_body"].(map[string]any)["items"].([]any)[0] = int8(0) + assert.Equal(t, int8(-7), source["extra_body"].(map[string]any)["items"].([]any)[0]) +} + +func TestClonePreservesYAMLNumericTypes(t *testing.T) { + t.Parallel() + var cfg latest.ModelConfig + require.NoError(t, yaml.Unmarshal([]byte(`provider: test +model: base +provider_opts: + big: 9007199254740993 + negative: -5 + fractional: 0.5 + date: 2026-01-01 + quoted_date: "2026-01-01" + timestamp: 2026-01-01T10:00:00Z + nested: {items: [9007199254740993, -5, 0.5]} +`), &cfg)) + require.IsType(t, uint64(0), cfg.ProviderOpts["big"]) + require.Equal(t, uint64(9007199254740993), cfg.ProviderOpts["big"]) + require.IsType(t, int64(0), cfg.ProviderOpts["negative"]) + require.IsType(t, float64(0), cfg.ProviderOpts["fractional"]) + for _, key := range []string{"date", "quoted_date", "timestamp"} { + require.IsType(t, "", cfg.ProviderOpts[key]) + } + registry := NewRegistry(map[string]Factory{"test": func(_ context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) { + return &instrumentTestProvider{Config: base.Config{ModelConfig: *cfg, ModelOptions: options.Apply(opts...), Env: env}}, nil + }}) + p, err := registry.New(t.Context(), &cfg, environment.NewNoEnvProvider()) + require.NoError(t, err) + for range 3 { + p, err = Clone(t.Context(), p) + require.NoError(t, err) + assert.Equal(t, cfg.ProviderOpts, p.BaseConfig().ModelConfig.ProviderOpts) + } +} + +func TestCloneRoutedPreservesNumericTypes(t *testing.T) { + t.Parallel() + var children []*instrumentTestProvider + registry := NewRegistry(map[string]Factory{"test": func(_ context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) { + p := &instrumentTestProvider{Config: base.Config{ModelConfig: *cfg, ModelOptions: options.Apply(opts...), Env: env}} + if cfg.Model == "child" { + children = append(children, p) + } + return p, nil + }}) + source := cloneNumericOptions() + p, err := registry.NewWithModels(t.Context(), &latest.ModelConfig{Provider: "test", Model: "fallback", Routing: []latest.RoutingRule{{Model: "target"}}}, map[string]latest.ModelConfig{"target": {Provider: "test", Model: "child", ProviderOpts: source}}, environment.NewNoEnvProvider()) + require.NoError(t, err) + for range 3 { + p, err = Clone(t.Context(), p) + require.NoError(t, err) + } + require.Len(t, children, 4) + for _, child := range children { + assert.Equal(t, source, child.ModelConfig.ProviderOpts) + } + children[3].ModelConfig.ProviderOpts["extra_body"].(map[string]any)["limit"] = int64(0) + assert.Equal(t, int64(9007199254740993), children[0].ModelConfig.ProviderOpts["extra_body"].(map[string]any)["limit"]) + assert.Equal(t, int64(9007199254740993), source["extra_body"].(map[string]any)["limit"]) +} + +func TestClonePreservesNonFiniteConfigurationScalars(t *testing.T) { + t.Parallel() + source := map[string]any{"inf": math.Inf(1), "nan": math.Float64frombits(0x7ff8000000000042)} + copied, _, err := mergeCloneOptions(base.Config{ModelConfig: latest.ModelConfig{ProviderOpts: source}, ModelOptions: options.Apply(options.WithStructuredOutput(&latest.StructuredOutput{Schema: source}))}, nil) + require.NoError(t, err) + for key, value := range source { + assert.Equal(t, math.Float64bits(value.(float64)), math.Float64bits(copied.ProviderOpts[key].(float64))) + } +} diff --git a/pkg/model/provider/clone_vertex_test.go b/pkg/model/provider/clone_vertex_test.go index d40cc7d6bb..c8df9b36b4 100644 --- a/pkg/model/provider/clone_vertex_test.go +++ b/pkg/model/provider/clone_vertex_test.go @@ -21,6 +21,7 @@ import ( "github.com/docker/docker-agent/pkg/config/latest" "github.com/docker/docker-agent/pkg/environment" "github.com/docker/docker-agent/pkg/httpclient" + "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/model/provider/options" "github.com/docker/docker-agent/pkg/model/provider/vertexai" ) @@ -48,6 +49,7 @@ func TestCloneVertexNativeDispatch(t *testing.T) { body, err := io.ReadAll(req.Body) require.NoError(t, err) assert.Contains(t, string(body), `"anthropic_version":"vertex-2023-10-16"`) + assert.Contains(t, string(body), `"top_k":9007199254740993`) return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")), Request: req}, nil }) }) @@ -74,6 +76,7 @@ func TestCloneVertexNativeDispatch(t *testing.T) { assert.True(t, cfg.ModelOptions.NoThinking()) assert.Equal(t, 2, googleCalls) for _, p := range []Provider{original, clone} { + assertVertexNumericConfig(t, p) stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) require.NoError(t, err) for { @@ -124,12 +127,17 @@ func TestCloneVertexNativeDispatch(t *testing.T) { assert.Empty(t, req.Header.Get("X-Api-Key")) body, err := io.ReadAll(req.Body) require.NoError(t, err) + assert.Contains(t, string(body), `"top_k":9007199254740993`) var payload map[string]any require.NoError(t, json.Unmarshal(body, &payload)) assert.Equal(t, "vertex-2023-10-16", payload["anthropic_version"]) assert.NotContains(t, payload, "model") - assert.NotContains(t, payload, "thinking") - assert.InDelta(t, float64(256), payload["max_tokens"], 0) + if httpclient.SessionIDFromContext(req.Context()) == "clone-session" { + assert.NotContains(t, payload, "thinking") + assert.InDelta(t, float64(256), payload["max_tokens"], 0) + } else { + assert.Contains(t, payload, "thinking") + } return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")), Request: req}, nil }) }) @@ -140,6 +148,16 @@ func TestCloneVertexNativeDispatch(t *testing.T) { clone, err = Clone(t.Context(), clone, options.WithGeneratingTitle()) require.NoError(t, err) assert.Equal(t, 3, calls) + assertVertexNumericConfig(t, original) + assertVertexNumericConfig(t, clone) + baseStream, err := original.CreateChatCompletionStream(httpclient.ContextWithSessionID(t.Context(), "base-session"), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + for { + if _, err := baseStream.Recv(); err != nil { + break + } + } + baseStream.Close() stream, err := clone.CreateChatCompletionStream(httpclient.ContextWithSessionID(t.Context(), "clone-session"), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) require.NoError(t, err) for { @@ -149,20 +167,20 @@ func TestCloneVertexNativeDispatch(t *testing.T) { } } stream.Close() - assert.Equal(t, 1, wireCalls) - assert.Equal(t, []string{"", "", "", "clone-session"}, sessions) + assert.Equal(t, 2, wireCalls) + assert.Equal(t, []string{"", "", "", "base-session", "clone-session"}, sessions) }) } }) } func vertexCloneConfig(location string) *latest.ModelConfig { - return &latest.ModelConfig{Provider: "google", Model: "claude-sonnet-4-6", BypassModelsGateway: true, ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"publisher": "anthropic", "project": "test-project", "location": location}} + return &latest.ModelConfig{Provider: "google", Model: "claude-sonnet-4-6", BypassModelsGateway: true, ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"publisher": "anthropic", "project": "test-project", "location": location, "top_k": int64(9007199254740993), "numeric": map[string]any{"unsigned": uint64(18446744073709551615)}}} } func TestCloneModelGardenCompatibleTransportAndSiblingAuth(t *testing.T) { t.Parallel() - cfg := &latest.ModelConfig{Provider: "google", Model: "meta/llama-test", ProviderOpts: map[string]any{"publisher": "meta", "project": "test-project", "location": "us-central1"}} + cfg := &latest.ModelConfig{Provider: "google", Model: "meta/llama-test", ProviderOpts: map[string]any{"publisher": "meta", "project": "test-project", "location": "us-central1", "extra_body": map[string]any{"exact_limit": int64(9007199254740993)}}} var sources, wire int source := func(ctx context.Context) (string, error) { sources++; return "vertex-only", ctx.Err() } registry := NewRegistry(map[string]Factory{"google": vertexai.NewClient, "openai": openAITestFactory}) @@ -172,6 +190,9 @@ func TestCloneModelGardenCompatibleTransportAndSiblingAuth(t *testing.T) { assert.Equal(t, "us-central1-aiplatform.googleapis.com", req.URL.Host) assert.Equal(t, "/v1beta1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions", req.URL.Path) assert.Equal(t, "Bearer vertex-only", req.Header.Get("Authorization")) + body, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Contains(t, string(body), `"exact_limit":9007199254740993`) return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")), Request: req}, nil }) }) @@ -179,11 +200,16 @@ func TestCloneModelGardenCompatibleTransportAndSiblingAuth(t *testing.T) { require.NoError(t, err) clone, err := Clone(t.Context(), original, options.WithNoThinking()) require.NoError(t, err) - stream, err := clone.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) - require.NoError(t, err) - drainStream(t, stream) - stream.Close() - assert.Equal(t, 1, wire) + for _, p := range []Provider{original, clone} { + cfg := p.BaseConfig() + assert.Equal(t, int64(9007199254740993), cfg.CloneSource().ProviderOpts["extra_body"].(map[string]any)["exact_limit"]) + assert.Equal(t, int64(9007199254740993), cfg.ModelConfig.ProviderOpts["extra_body"].(map[string]any)["exact_limit"]) + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + } + assert.Equal(t, 2, wire) assert.Positive(t, sources) // A leaf's stored OAuth source must not enter a router's sibling options. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { @@ -199,9 +225,40 @@ func TestCloneModelGardenCompatibleTransportAndSiblingAuth(t *testing.T) { require.NoError(t, err) router, err = Clone(t.Context(), router, options.WithNoThinking()) require.NoError(t, err) - stream, err = router.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "unmatched fallback request"}}, nil) + stream, err := router.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "unmatched fallback request"}}, nil) require.NoError(t, err) drainStream(t, stream) stream.Close() - assert.Equal(t, 1, wire, "fallback must not call the Vertex transport") + assert.Equal(t, 2, wire, "fallback must not call the Vertex transport") +} + +func assertVertexNumericConfig(t *testing.T, p Provider) { + t.Helper() + cfg := p.BaseConfig() + for _, model := range []latest.ModelConfig{cfg.ModelConfig, cfg.CloneSource()} { + assert.Equal(t, int64(9007199254740993), model.ProviderOpts["top_k"]) + assert.Equal(t, uint64(18446744073709551615), model.ProviderOpts["numeric"].(map[string]any)["unsigned"]) + } +} + +func TestCloneVertexRejectsUncopyableConfiguration(t *testing.T) { + t.Parallel() + cfg := vertexCloneConfig("us-east5") + cfg.ProviderOpts["bad"] = func() {} + source := options.WithTokenSource(func(context.Context) (string, error) { + t.Error("must reject before token lookup") + return "", nil + }) + p, err := vertexai.NewClient(t.Context(), cfg, environment.NewNoEnvProvider(), source) + require.ErrorContains(t, err, "provider_opts[bad]") + assert.Nil(t, p) + original := &instrumentTestProvider{Config: base.Config{ModelConfig: *cfg, RebuildProvider: func(ctx context.Context, cfg *latest.ModelConfig, opts ...options.Opt) (Provider, error) { + return vertexai.NewClient(ctx, cfg, environment.NewNoEnvProvider(), append(opts, source)...) + }}} + p, err = Clone(t.Context(), original) + require.ErrorIs(t, err, ErrUncopyableConfig) + require.ErrorContains(t, err, "provider_opts[bad]") + var cloneErr *CloneError + require.ErrorAs(t, err, &cloneErr) + assert.Nil(t, p) } diff --git a/pkg/model/provider/contracts/contracts.go b/pkg/model/provider/contracts/contracts.go index 77e740f1ea..6cbcef2a38 100644 --- a/pkg/model/provider/contracts/contracts.go +++ b/pkg/model/provider/contracts/contracts.go @@ -38,8 +38,8 @@ type Config struct { BaseURL string } -// SetDispatchConfig preserves the original configuration for reconstruction. -func (c *Config) SetDispatchConfig(cfg *latest.ModelConfig) { c.DispatchConfig = cfg.Clone() } +// SetDispatchConfig takes ownership of a caller-isolated configuration for reconstruction. +func (c *Config) SetDispatchConfig(cfg *latest.ModelConfig) { c.DispatchConfig = cfg } // CloneSource returns the construction configuration, not the catalogue identity. func (c *Config) CloneSource() latest.ModelConfig { diff --git a/pkg/model/provider/internal/deepcopy/deepcopy.go b/pkg/model/provider/internal/deepcopy/deepcopy.go new file mode 100644 index 0000000000..ca1963ae41 --- /dev/null +++ b/pkg/model/provider/internal/deepcopy/deepcopy.go @@ -0,0 +1,150 @@ +package deepcopy + +import ( + "errors" + "fmt" + "reflect" + "strings" +) + +var ( + ErrUnsupported = errors.New("unsupported configuration value") + ErrCycle = errors.New("cyclic configuration value") + ErrLimit = errors.New("configuration copy limit exceeded") +) + +// Value isolates mutable values, preserving scalar types and duplicating shared references. +func Value[T any](source T) (T, error) { + var zero T + w := walker{active: make(map[reference]bool)} + copied, err := w.copy(reflect.ValueOf(&source).Elem(), "$", 0) + if err != nil { + return zero, err + } + reflect.ValueOf(&zero).Elem().Set(copied) + return zero, nil +} + +type reference struct { + typ reflect.Type + ptr uintptr + len int +} + +type walker struct { + active map[reference]bool + nodes int +} + +func (w *walker) copy(v reflect.Value, path string, depth int) (reflect.Value, error) { + w.nodes++ + if depth > 128 || w.nodes > 100_000 { + return reflect.Value{}, fmt.Errorf("%w at %s", ErrLimit, path) + } + kind := v.Kind() + if kind == reflect.Func || kind == reflect.Chan || kind == reflect.UnsafePointer { + return reflect.Value{}, fmt.Errorf("%w at %s: %s", ErrUnsupported, path, v.Type()) + } + if kind == reflect.Pointer || kind == reflect.Map || kind == reflect.Slice { + if v.IsNil() { + return reflect.Zero(v.Type()), nil + } + ref := reference{typ: v.Type(), ptr: v.Pointer()} + if kind == reflect.Slice { + ref.len = v.Len() + } + if w.active[ref] { + return reflect.Value{}, fmt.Errorf("%w at %s", ErrCycle, path) + } + w.active[ref] = true + defer delete(w.active, ref) + } + out := reflect.New(v.Type()).Elem() + switch kind { + case reflect.Interface, reflect.Pointer: + if v.IsNil() { + return out, nil + } + child, err := w.copy(v.Elem(), path, depth+1) + if err != nil { + return reflect.Value{}, err + } + if kind == reflect.Pointer { + out.Set(reflect.New(v.Type().Elem())) + out.Elem().Set(child) + } else { + out.Set(child) + } + case reflect.Map: + if v.Len() > (100_000-w.nodes)/2 { + return reflect.Value{}, fmt.Errorf("%w at %s", ErrLimit, path) + } + out.Set(reflect.MakeMapWithSize(v.Type(), v.Len())) + iter := v.MapRange() + for iter.Next() { + key, err := w.copy(iter.Key(), path+"[key]", depth+1) + if err != nil { + return reflect.Value{}, err + } + child, err := w.copy(iter.Value(), fmt.Sprintf("%s[%v]", path, iter.Key()), depth+1) + if err != nil { + return reflect.Value{}, err + } + out.SetMapIndex(key, child) + } + case reflect.Slice, reflect.Array: + if v.Len() > 100_000-w.nodes { + return reflect.Value{}, fmt.Errorf("%w at %s", ErrLimit, path) + } + if kind == reflect.Slice { + out.Set(reflect.MakeSlice(v.Type(), v.Len(), v.Len())) + } + for i := range v.Len() { + child, err := w.copy(v.Index(i), fmt.Sprintf("%s[%d]", path, i), depth+1) + if err != nil { + return reflect.Value{}, err + } + out.Index(i).Set(child) + } + case reflect.Struct: + out.Set(v) + for i := range v.NumField() { + field := v.Type().Field(i) + name := strings.Split(field.Tag.Get("json"), ",")[0] + if name == "" || name == "-" { + name = field.Name + } + fieldPath := path + "." + name + if !field.IsExported() { + if containsReferences(field.Type) { + return reflect.Value{}, fmt.Errorf("%w at %s: unexported reference", ErrUnsupported, fieldPath) + } + continue + } + child, err := w.copy(v.Field(i), fieldPath, depth+1) + if err != nil { + return reflect.Value{}, err + } + out.Field(i).Set(child) + } + default: + out.Set(v) + } + return out, nil +} + +func containsReferences(typ reflect.Type) bool { + switch typ.Kind() { + case reflect.Pointer, reflect.Map, reflect.Slice, reflect.Interface, reflect.Func, reflect.Chan, reflect.UnsafePointer: + return true + case reflect.Array: + return containsReferences(typ.Elem()) + case reflect.Struct: + for field := range typ.Fields() { + if containsReferences(field.Type) { + return true + } + } + } + return false +} diff --git a/pkg/model/provider/internal/deepcopy/deepcopy_test.go b/pkg/model/provider/internal/deepcopy/deepcopy_test.go new file mode 100644 index 0000000000..7665ccd947 --- /dev/null +++ b/pkg/model/provider/internal/deepcopy/deepcopy_test.go @@ -0,0 +1,102 @@ +package deepcopy + +import ( + "encoding/json" + "math" + "reflect" + "testing" + "unsafe" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValuePreservesTypesAndIsolatesReferences(t *testing.T) { + t.Parallel() + type number int64 + type plain struct { + Value *int64 + hidden int + } + shared := map[string]any{"items": []any{int64(9007199254740993)}} + source := map[string]any{ + "int64": int64(9007199254740993), "uint64": uint64(math.MaxUint64), + "int8": int8(-7), "float32": float32(0.1), "number": json.Number("9007199254740993"), + "named": number(9007199254740993), "array": [1]*int64{new(int64(7))}, + "struct": plain{Value: new(int64(9)), hidden: 4}, + "first": shared, "second": shared, "nil": nil, + "typedNil": (*int)(nil), "nilMap": map[string]any(nil), "empty": []any{}, + } + copied, err := Value(source) + require.NoError(t, err) + require.Equal(t, source, copied) + for key, value := range source { + assert.Equal(t, reflect.TypeOf(value), reflect.TypeOf(copied[key])) + } + copied["first"].(map[string]any)["items"].([]any)[0] = int64(0) + assert.Equal(t, int64(9007199254740993), shared["items"].([]any)[0]) + assert.Equal(t, int64(9007199254740993), copied["second"].(map[string]any)["items"].([]any)[0]) + *copied["array"].([1]*int64)[0] = 0 + *copied["struct"].(plain).Value = 0 + assert.Equal(t, int64(7), *source["array"].([1]*int64)[0]) + assert.Equal(t, int64(9), *source["struct"].(plain).Value) + var absent any + got, err := Value(absent) + require.NoError(t, err) + assert.Nil(t, got) +} + +func TestValuePreservesNonFiniteBits(t *testing.T) { + t.Parallel() + for _, bits := range []uint64{0x7ff8000000000042, math.Float64bits(math.Inf(1)), math.Float64bits(math.Copysign(0, -1))} { + source := map[string]any{"value": math.Float64frombits(bits)} + copied, err := Value(source) + require.NoError(t, err) + assert.Equal(t, bits, math.Float64bits(copied["value"].(float64))) + } +} + +func TestValueRejectsUnsupportedValues(t *testing.T) { + t.Parallel() + type private struct{ hidden map[string]any } + hidden := private{hidden: map[string]any{"value": 1}} + require.NotNil(t, hidden.hidden) + for _, value := range []any{func() {}, make(chan int), unsafe.Pointer(new(int)), hidden} { + copied, err := Value(map[string]any{"nested": []any{value}}) + require.ErrorIs(t, err, ErrUnsupported) + require.ErrorContains(t, err, "$[nested][0]") + assert.Nil(t, copied) + } +} + +func TestValueRejectsCycles(t *testing.T) { + t.Parallel() + cycle := map[string]any{} + cycle["self"] = cycle + slice := []any{nil} + slice[0] = slice + type node struct{ Next *node } + pointer := &node{} + pointer.Next = pointer + for _, value := range []any{cycle, slice, pointer} { + copied, err := Value(value) + require.ErrorIs(t, err, ErrCycle) + assert.Nil(t, copied) + } +} + +func TestValueBoundsDepthAndBreadth(t *testing.T) { + t.Parallel() + var deep any = "end" + for range 130 { + deep = []any{deep} + } + _, err := Value(deep) + require.ErrorIs(t, err, ErrLimit) + var diamond any = "end" + for range 18 { + diamond = []any{diamond, diamond} + } + _, err = Value(diamond) + require.ErrorIs(t, err, ErrLimit) +} diff --git a/pkg/model/provider/options/options.go b/pkg/model/provider/options/options.go index 897fcedc47..f561e8b71b 100644 --- a/pkg/model/provider/options/options.go +++ b/pkg/model/provider/options/options.go @@ -2,12 +2,12 @@ package options import ( "context" - "encoding/json" "fmt" "log/slog" "net/http" "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/model/provider/internal/deepcopy" "github.com/docker/docker-agent/pkg/modelsdev" ) @@ -299,25 +299,14 @@ func FromModelOptions(m ModelOptions) []Opt { // Clone copies mutable option state while preserving shared services and callbacks. func (c *ModelOptions) Clone() (ModelOptions, error) { cloned := *c - if c.structuredOutput != nil { - data, err := json.Marshal(c.structuredOutput) - if err != nil { - return ModelOptions{}, fmt.Errorf("copying structured_output: %w", err) - } - cloned.structuredOutput = new(latest.StructuredOutput) - if err := json.Unmarshal(data, cloned.structuredOutput); err != nil { - return ModelOptions{}, err - } + var err error + cloned.structuredOutput, err = deepcopy.Value(c.structuredOutput) + if err != nil { + return ModelOptions{}, fmt.Errorf("copying structured_output: %w", err) } - if c.providers != nil { - data, err := json.Marshal(c.providers) - if err != nil { - return ModelOptions{}, fmt.Errorf("copying providers: %w", err) - } - cloned.providers = nil - if err := json.Unmarshal(data, &cloned.providers); err != nil { - return ModelOptions{}, err - } + cloned.providers, err = deepcopy.Value(c.providers) + if err != nil { + return ModelOptions{}, fmt.Errorf("copying providers: %w", err) } return cloned, nil } diff --git a/pkg/model/provider/options/options_test.go b/pkg/model/provider/options/options_test.go index 7293a729f4..a3a50f7b84 100644 --- a/pkg/model/provider/options/options_test.go +++ b/pkg/model/provider/options/options_test.go @@ -3,6 +3,7 @@ package options import ( "context" "net/http" + "reflect" "testing" "github.com/stretchr/testify/assert" @@ -171,8 +172,8 @@ func TestModelOptionsClonePreservesServicesAndExplicitClears(t *testing.T) { require.NoError(t, err) probe := Apply(FromModelOptions(clone)...) assert.Same(t, store, probe.ModelsDevStore()) - assert.NotNil(t, probe.TransportWrapper()) - assert.NotNil(t, probe.TokenSource()) + assert.Equal(t, reflect.ValueOf(wrapper).Pointer(), reflect.ValueOf(probe.TransportWrapper()).Pointer()) + assert.Equal(t, reflect.ValueOf(source).Pointer(), reflect.ValueOf(probe.TokenSource()).Pointer()) assert.Empty(t, probe.Gateway()) assert.True(t, probe.gatewaySet) assert.True(t, probe.structuredOutputSet) @@ -180,3 +181,23 @@ func TestModelOptionsClonePreservesServicesAndExplicitClears(t *testing.T) { require.NoError(t, err) assert.Equal(t, "token", token) } + +func TestModelOptionsClonePreservesNumericTypesAndIsolation(t *testing.T) { + t.Parallel() + values := map[string]any{"minimum": int64(9007199254740993), "maximum": uint64(18446744073709551615), "small": int8(-7), "fraction": float32(0.1), "nested": []any{map[string]any{"limit": int64(9007199254740993)}}} + m := Apply(WithStructuredOutput(&latest.StructuredOutput{Schema: values}), WithProviders(map[string]latest.ProviderConfig{"custom": {ProviderOpts: values}})) + for range 3 { + clone, err := m.Clone() + require.NoError(t, err) + require.Equal(t, values, clone.StructuredOutput().Schema) + require.Equal(t, values, clone.Providers()["custom"].ProviderOpts) + for key, value := range values { + assert.IsType(t, value, clone.StructuredOutput().Schema[key]) + assert.IsType(t, value, clone.Providers()["custom"].ProviderOpts[key]) + } + m = Apply(FromModelOptions(clone)...) + } + m.StructuredOutput().Schema["nested"].([]any)[0].(map[string]any)["limit"] = int64(0) + assert.Equal(t, int64(9007199254740993), values["nested"].([]any)[0].(map[string]any)["limit"]) + assert.Equal(t, int64(9007199254740993), m.Providers()["custom"].ProviderOpts["nested"].([]any)[0].(map[string]any)["limit"]) +} diff --git a/pkg/model/provider/vertexai/modelgarden.go b/pkg/model/provider/vertexai/modelgarden.go index 358015f478..dd60cf3c31 100644 --- a/pkg/model/provider/vertexai/modelgarden.go +++ b/pkg/model/provider/vertexai/modelgarden.go @@ -40,6 +40,7 @@ import ( "github.com/docker/docker-agent/pkg/environment" "github.com/docker/docker-agent/pkg/model/provider/anthropic/vertex" "github.com/docker/docker-agent/pkg/model/provider/contracts" + "github.com/docker/docker-agent/pkg/model/provider/internal/deepcopy" "github.com/docker/docker-agent/pkg/model/provider/openai" "github.com/docker/docker-agent/pkg/model/provider/options" ) @@ -87,7 +88,10 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro if err != nil { return nil, err } - rewritten := withModelsDevProvider(cfg) + rewritten, err := withModelsDevProvider(cfg) + if err != nil { + return nil, err + } var client Client if strings.EqualFold(publisher(cfg), "anthropic") { client, err = vertex.NewClient(ctx, rewritten, env, project, location, opts...) @@ -110,7 +114,10 @@ func NewClientWithTokenSource(ctx context.Context, cfg *latest.ModelConfig, env return nil, err } opts = append(opts, options.WithTokenSource(source)) - rewritten := withModelsDevProvider(cfg) + rewritten, err := withModelsDevProvider(cfg) + if err != nil { + return nil, err + } var client Client if strings.EqualFold(publisher(cfg), "anthropic") { client, err = vertex.NewClientWithTokenSource(ctx, rewritten, env, project, location, source, opts...) @@ -125,7 +132,11 @@ func withDispatchConfig(client Client, cfg *latest.ModelConfig, err error) (Clie return nil, err } if setter, ok := client.(interface{ SetDispatchConfig(cfg *latest.ModelConfig) }); ok { - setter.SetDispatchConfig(cfg) + copied, err := deepcopy.Value(cfg) + if err != nil { + return nil, fmt.Errorf("vertex ai: copying dispatch configuration: %w", err) + } + setter.SetDispatchConfig(copied) } return client, nil } @@ -134,12 +145,15 @@ func withDispatchConfig(client Client, cfg *latest.ModelConfig, err error) (Clie // rewritten so that `Provider + "/" + Model` becomes a valid models.dev key. // The original cfg is never mutated so callers — and the wider runtime — // keep their "google" provider view. -func withModelsDevProvider(cfg *latest.ModelConfig) *latest.ModelConfig { - out := cfg.Clone() +func withModelsDevProvider(cfg *latest.ModelConfig) (*latest.ModelConfig, error) { + out, err := deepcopy.Value(cfg) + if err != nil { + return nil, fmt.Errorf("vertex ai: copying model configuration: %w", err) + } if p := modelsDevProvider(cfg); p != "" { out.Provider = p } - return out + return out, nil } // modelsDevProvider returns the models.dev provider key that hosts the model @@ -254,7 +268,10 @@ func openAIBaseURL(project, location string) string { } func newOpenAIClientWithTokenSource(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, project, location string, tokenSource options.TokenSource, opts ...options.Opt) (*openai.Client, error) { - oaiCfg := cfg.Clone() + oaiCfg, err := deepcopy.Value(cfg) + if err != nil { + return nil, fmt.Errorf("vertex ai: copying OpenAI configuration: %w", err) + } oaiCfg.BaseURL = openAIBaseURL(project, location) oaiCfg.TokenKey = "" diff --git a/pkg/model/provider/vertexai/modelgarden_test.go b/pkg/model/provider/vertexai/modelgarden_test.go index 5b4add26cb..4a040f8e29 100644 --- a/pkg/model/provider/vertexai/modelgarden_test.go +++ b/pkg/model/provider/vertexai/modelgarden_test.go @@ -268,7 +268,8 @@ func TestWithModelsDevProvider(t *testing.T) { originalProvider := tt.cfg.Provider originalPublisher, _ := tt.cfg.ProviderOpts["publisher"].(string) - got := withModelsDevProvider(tt.cfg) + got, err := withModelsDevProvider(tt.cfg) + require.NoError(t, err) require.NotNil(t, got) require.NotSame(t, tt.cfg, got, "withModelsDevProvider must return a copy") assert.Equal(t, tt.wantProvider, got.Provider) @@ -344,7 +345,8 @@ func TestWithModelsDevProvider_CapabilityLookupIDs(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - rewritten := withModelsDevProvider(tt.cfg) + rewritten, err := withModelsDevProvider(tt.cfg) + require.NoError(t, err) gotID := rewritten.Provider + "/" + rewritten.DisplayOrModel() assert.Equal(t, tt.wantID, gotID) }) From 1a4cd20391e8f55db43a9d5a27ad8a2710e7287d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Wed, 14 Oct 2026 09:21:33 +0000 Subject: [PATCH 4/5] fix: honor explicit token caps in routed provider clones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/guides/go-sdk/index.md | 9 ++- pkg/model/provider/clone.go | 8 +++ pkg/model/provider/clone_matrix_test.go | 78 +++++++++++++++++++++- pkg/model/provider/clone_routed_test.go | 8 ++- pkg/model/provider/clone_thinking_test.go | 77 ++++++++++++++++++--- pkg/model/provider/factory.go | 9 ++- pkg/model/provider/options/options.go | 49 +++++++++----- pkg/model/provider/options/options_test.go | 33 +++++++++ pkg/model/provider/rulebased/client.go | 5 ++ pkg/runtime/sampling_routed_cap_test.go | 60 +++++++++++++++++ pkg/sessiontitle/routed_cap_test.go | 50 ++++++++++++++ 11 files changed, 351 insertions(+), 35 deletions(-) create mode 100644 pkg/runtime/sampling_routed_cap_test.go create mode 100644 pkg/sessiontitle/routed_cap_test.go diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index d47a47265e..780a9aa228 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -881,15 +881,18 @@ plain structs. Functions, channels, unsafe pointers, cyclic values and structs with unexported reference-bearing fields fail closed with `provider.ErrUncopyableConfig`, rather than panic or share mutable state. Copies are bounded to 128 levels and 100,000 visited values. Cumulative clones preserve explicit structured-output -removal, gateway clearing and no-thinking overrides; existing provider token -floors and routed-child token limits still apply. Named routing definitions are +removal, gateway clearing and no-thinking overrides. A non-zero `WithMaxTokens` +cap supplied to `Clone` overrides direct, matched-route and fallback limits, +subject to existing provider token floors. Cumulative clones retain that cap +when it is omitted or zero; a later non-zero cap replaces it. Named routing definitions are captured at construction and copied for each child, so later input-map mutations do not change clones. Configuration-declared thinking off (`none` or `0`, including provider defaults) is retained as an option before normalization. Reconstruction applies that option to each leaf's config too, so config-driven and option-driven providers agree. -Routed children inherit strict-off overrides but keep their own token limits. +Routed children inherit strict-off overrides and explicit clone caps. Ordinary +construction-time router limits still leave each child's token defaults unchanged. `CloneWithOptions` retains its public signature but is deprecated. On failure it returns an error-only provider whose stream creation returns the clone error. It diff --git a/pkg/model/provider/clone.go b/pkg/model/provider/clone.go index ab438156f2..a0318d1df5 100644 --- a/pkg/model/provider/clone.go +++ b/pkg/model/provider/clone.go @@ -32,6 +32,8 @@ func (e *CloneError) Error() string { return fmt.Sprintf("cloning provider %q: % func (e *CloneError) Unwrap() error { return e.Err } // Clone reconstructs a provider with ordered option overrides, or returns an error. +// Non-zero token caps also constrain routed children, subject to provider floors. +// An omitted or zero cap retains a previous explicit clone cap. func Clone(ctx context.Context, p Provider, opts ...options.Opt) (Provider, error) { clone, _, err := cloneProvider(ctx, p, opts) return clone, err @@ -97,6 +99,12 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, } baseOpts := options.FromModelOptions(baseOptions) mergedOpts := slices.Concat(baseOpts, opts) + requested := options.Apply(opts...) + if mt := requested.MaxTokens(); mt != 0 { + mergedOpts = append(mergedOpts, options.WithAuxiliaryMaxTokens(mt)) + } else if baseOptions.MaxTokensInheritable() && baseOptions.MaxTokens() != 0 { + mergedOpts = append(mergedOpts, options.WithAuxiliaryMaxTokens(baseOptions.MaxTokens())) + } // Snapshot overrides too: option setters can carry caller-owned maps. merged := options.Apply(mergedOpts...) diff --git a/pkg/model/provider/clone_matrix_test.go b/pkg/model/provider/clone_matrix_test.go index 2c7d365f23..ac6f496d90 100644 --- a/pkg/model/provider/clone_matrix_test.go +++ b/pkg/model/provider/clone_matrix_test.go @@ -59,9 +59,7 @@ func TestCloneOptionMatrixOnWire(t *testing.T) { require.NoError(t, json.Unmarshal(body(), &payload)) assert.Equal(t, "none", payload["reasoning_effort"]) assert.NotContains(t, payload, "response_format") - if shape != "routed" { - assert.InDelta(t, float64(256), payload["max_completion_tokens"], 0) - } + assert.InDelta(t, float64(256), payload["max_completion_tokens"], 0) cfgClone := clone.BaseConfig() assert.True(t, cfgClone.ModelOptions.NoThinking()) assert.Nil(t, cfgClone.ModelOptions.StructuredOutput()) @@ -222,3 +220,77 @@ func TestCloneNumericPrecisionOnWire(t *testing.T) { assert.Equal(t, json.Number("9007199254740993"), bound) } } + +func TestCloneRoutedTokenCapsPreserveConstructionDefaults(t *testing.T) { + t.Parallel() + server, body := captureNamedCustomProviderRequestBody(t, "openai_chatcompletions") + cfg := &latest.ModelConfig{Provider: "custom", Model: "gpt-4o", MaxTokens: new(int64(4096)), Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + models := map[string]latest.ModelConfig{"target": {Provider: "custom", Model: "gpt-5.6", MaxTokens: new(int64(8192))}} + original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"KEY": "fake"}), options.WithMaxTokens(4096), options.WithProviders(map[string]latest.ProviderConfig{"custom": {BaseURL: server.URL, TokenKey: "KEY", APIType: "openai_chatcompletions", MaxTokens: new(int64(16384))}})) + require.NoError(t, err) + request := func(p Provider, tokenCap int64) { + t.Helper() + for _, message := range []string{"hello", "unrelated"} { + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: message}}, nil) + require.NoError(t, err) + drainStream(t, stream) + stream.Close() + var payload map[string]any + require.NoError(t, json.Unmarshal(body(), &payload)) + wantModel, wantCap := "gpt-5.6", tokenCap + if message == "unrelated" { + wantModel = "gpt-4o" + } + if tokenCap == 0 { + wantCap = 8192 + if message == "unrelated" { + wantCap = 16384 + } + } else if message == "hello" && tokenCap < 256 { + wantCap = 256 + } + assert.Equal(t, wantModel, payload["model"]) + field := "max_completion_tokens" + if message == "unrelated" { + field = "max_tokens" + } + assert.InDelta(t, float64(wantCap), payload[field], 0) + } + } + request(original, 0) + for _, tc := range []struct { + opts []options.Opt + tokenCap int64 + inheritable bool + }{ + {[]options.Opt{options.WithAuxiliaryMaxTokens(128), options.WithMaxTokens(64)}, 64, false}, + {[]options.Opt{options.WithMaxTokens(64), options.WithAuxiliaryMaxTokens(128)}, 128, true}, + } { + p, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"KEY": "fake"}), append(options.FromModelOptions(original.BaseConfig().ModelOptions), append(tc.opts, options.WithNoThinking())...)...) + require.NoError(t, err) + baseCfg := p.BaseConfig() + assert.Equal(t, tc.tokenCap, baseCfg.ModelOptions.MaxTokens()) + assert.Equal(t, tc.inheritable, baseCfg.ModelOptions.MaxTokensInheritable()) + wireCap := tc.tokenCap + if !tc.inheritable { + wireCap = 0 + } + request(p, wireCap) + } + plain, err := Clone(t.Context(), original) + require.NoError(t, err) + request(plain, 0) + for _, opts := range [][]options.Opt{ + {options.WithMaxTokens(128)}, + {options.WithAuxiliaryMaxTokens(1024), options.WithMaxTokens(128)}, + {options.WithMaxTokens(1024), options.WithAuxiliaryMaxTokens(128)}, + } { + clone, err := Clone(t.Context(), original, append(opts, options.WithNoThinking())...) + require.NoError(t, err) + request(clone, 128) + clone, err = Clone(t.Context(), clone, options.WithMaxTokens(0)) + require.NoError(t, err) + request(clone, 128) + } + request(original, 0) +} diff --git a/pkg/model/provider/clone_routed_test.go b/pkg/model/provider/clone_routed_test.go index 244b512f01..d469067627 100644 --- a/pkg/model/provider/clone_routed_test.go +++ b/pkg/model/provider/clone_routed_test.go @@ -38,10 +38,14 @@ func TestCloneRoutedConfigurationsAreIsolated(t *testing.T) { require.NoError(t, err) require.NotNil(t, second) require.Len(t, children, 3) - for _, child := range children { + for i, child := range children { assert.Equal(t, "initial", child.ModelConfig.ProviderOpts["nested"].(map[string]any)["value"]) assert.Equal(t, 8192, child.ModelConfig.ThinkingBudget.Tokens) - assert.Equal(t, int64(16384), *child.ModelConfig.MaxTokens) + if i == 0 { + assert.Equal(t, int64(16384), *child.ModelConfig.MaxTokens) + } else { + assert.Equal(t, int64(128), *child.ModelConfig.MaxTokens) + } assert.Equal(t, "named", child.ModelConfig.Name) assert.Equal(t, "display", child.ModelConfig.DisplayModel) } diff --git a/pkg/model/provider/clone_thinking_test.go b/pkg/model/provider/clone_thinking_test.go index 9d3cfdfb55..3c421e25de 100644 --- a/pkg/model/provider/clone_thinking_test.go +++ b/pkg/model/provider/clone_thinking_test.go @@ -40,12 +40,12 @@ func TestCloneRoutedClaudeStrictOffOnWire(t *testing.T) { models := map[string]latest.ModelConfig{"target": {Provider: provider, Model: model, MaxTokens: new(int64(16384)), ThinkingBudget: &latest.ThinkingBudget{Tokens: 2048}, ProviderOpts: map[string]any{"region": "us-east-1"}}} cfg := &latest.ModelConfig{Provider: provider, Model: model, Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} env := environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "fake", "AWS_BEARER_TOKEN_BEDROCK": "fake"}) - original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, env, wrapper) + original, err := fullTestRegistry().NewWithModels(t.Context(), cfg, models, env, wrapper, options.WithMaxTokens(4096)) require.NoError(t, err) - request := func(p Provider, thinking bool) { + request := func(p Provider, message string, thinking bool, tokenCap int64) { t.Helper() before := calls - stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: message}}, nil) if err == nil { _, err = stream.Recv() stream.Close() @@ -57,9 +57,9 @@ func TestCloneRoutedClaudeStrictOffOnWire(t *testing.T) { if additional, exists := body["additionalModelRequestFields"]; exists { fields = additional.(map[string]any) } - assert.InDelta(t, float64(16384), body["inferenceConfig"].(map[string]any)["maxTokens"], 0) + assert.InDelta(t, float64(tokenCap), body["inferenceConfig"].(map[string]any)["maxTokens"], 0) } else { - assert.InDelta(t, float64(8192), body["max_tokens"], 0) + assert.InDelta(t, float64(tokenCap), body["max_tokens"], 0) } if thinking { require.Contains(t, fields, "thinking") @@ -69,15 +69,37 @@ func TestCloneRoutedClaudeStrictOffOnWire(t *testing.T) { assert.NotContains(t, fields, "reasoning_config") } } - request(original, true) + defaultCap := int64(8192) + if provider == "amazon-bedrock" { + defaultCap = 16384 + } + request(original, "hello", true, defaultCap) for _, purpose := range []options.Opt{options.WithGeneratingTitle(), options.WithStructuredOutput(nil), options.WithCompacting()} { clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128), purpose) require.NoError(t, err) - clone, err = Clone(t.Context(), clone) - require.NoError(t, err) - request(clone, false) + for _, tokenCap := range []int64{128, 0, 64, 1024} { + clone, err = Clone(t.Context(), clone, options.WithMaxTokens(tokenCap)) + require.NoError(t, err) + requested := tokenCap + if tokenCap == 0 { + requested = 128 + } + assert.Equal(t, requested, *clone.BaseConfig().ModelConfig.MaxTokens) + baseCfg := clone.BaseConfig() + assert.Equal(t, requested, baseCfg.ModelOptions.MaxTokens()) + wireCap := requested + if provider == "anthropic" && wireCap < 256 { + wireCap = 256 + } + request(clone, "hello", false, wireCap) + request(clone, "unrelated", false, wireCap) + clone, err = Clone(t.Context(), clone) + require.NoError(t, err) + request(clone, "hello", false, wireCap) + request(clone, "unrelated", false, wireCap) + } } - request(original, true) + request(original, "hello", true, defaultCap) }) } } @@ -148,3 +170,38 @@ func TestCloneConfiguredThinkingOffOnWire(t *testing.T) { }) } } + +func TestCloneDirectClaudeRetainsExplicitCapOnZero(t *testing.T) { + t.Parallel() + var body map[string]any + wrapper := options.WithHTTPTransportWrapper(func(http.RoundTripper) http.RoundTripper { + return thinkingRoundTripper(func(req *http.Request) (*http.Response, error) { + require.NoError(t, json.NewDecoder(req.Body).Decode(&body)) + return &http.Response{StatusCode: http.StatusBadRequest, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"offline capture"}}`)), Request: req}, nil + }) + }) + original, err := fullTestRegistry().New(t.Context(), &latest.ModelConfig{Provider: "anthropic", Model: "claude-haiku-4-5"}, environment.NewMapEnvProvider(map[string]string{"ANTHROPIC_API_KEY": "fake"}), wrapper) + require.NoError(t, err) + request := func(p Provider, tokenCap int64) { + t.Helper() + stream, err := p.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "hello"}}, nil) + if err == nil { + _, err = stream.Recv() + stream.Close() + } + require.Error(t, err) + assert.InDelta(t, float64(tokenCap), body["max_tokens"], 0) + } + request(original, 8192) + p, err := Clone(t.Context(), original, options.WithMaxTokens(128), options.WithNoThinking()) + require.NoError(t, err) + for _, opts := range [][]options.Opt{nil, {options.WithMaxTokens(0)}, nil} { + p, err = Clone(t.Context(), p, opts...) + require.NoError(t, err) + cfg := p.BaseConfig() + assert.Equal(t, int64(128), cfg.ModelOptions.MaxTokens()) + assert.Equal(t, int64(128), *cfg.ModelConfig.MaxTokens) + request(p, 256) + } + request(original, 8192) +} diff --git a/pkg/model/provider/factory.go b/pkg/model/provider/factory.go index 615a0108e6..bca56ac176 100644 --- a/pkg/model/provider/factory.go +++ b/pkg/model/provider/factory.go @@ -115,12 +115,17 @@ func (r *Registry) createDirectProvider(ctx context.Context, cfg *latest.ModelCo noThinking := globalOptions.NoThinking() || budget.IsDisabled() constructionCfg := cfg // Config-driven leaves need the same override as option-driven leaves. - if globalOptions.NoThinking() { + if globalOptions.NoThinking() || globalOptions.MaxTokensInheritable() { copied, err := copyModelConfig(*cfg) if err != nil { return nil, err } - copied.ThinkingBudget = &latest.ThinkingBudget{Effort: "none"} + if globalOptions.NoThinking() { + copied.ThinkingBudget = &latest.ThinkingBudget{Effort: "none"} + } + if mt := globalOptions.MaxTokens(); globalOptions.MaxTokensInheritable() && mt != 0 { + copied.MaxTokens = &mt + } constructionCfg = &copied } if noThinking && !globalOptions.NoThinking() { diff --git a/pkg/model/provider/options/options.go b/pkg/model/provider/options/options.go index f561e8b71b..57442b27bc 100644 --- a/pkg/model/provider/options/options.go +++ b/pkg/model/provider/options/options.go @@ -15,20 +15,21 @@ import ( type TokenSource func(context.Context) (string, error) type ModelOptions struct { - gateway string - gatewaySet bool - structuredOutputSet bool - encryptedConfig string - structuredOutput *latest.StructuredOutput - generatingTitle bool - compacting bool - noThinking bool - maxTokens int64 - providers map[string]latest.ProviderConfig - modelsDevStore *modelsdev.Store - transportWrapper func(http.RoundTripper) http.RoundTripper - tokenSource TokenSource - openAIVendor bool + gateway string + gatewaySet bool + structuredOutputSet bool + encryptedConfig string + structuredOutput *latest.StructuredOutput + generatingTitle bool + compacting bool + noThinking bool + maxTokens int64 + maxTokensInheritable bool + providers map[string]latest.ProviderConfig + modelsDevStore *modelsdev.Store + transportWrapper func(http.RoundTripper) http.RoundTripper + tokenSource TokenSource + openAIVendor bool } func (c *ModelOptions) Gateway() string { @@ -57,6 +58,11 @@ func (c *ModelOptions) MaxTokens() int64 { return c.maxTokens } +// MaxTokensInheritable reports whether an auxiliary cap applies to routed children. +func (c *ModelOptions) MaxTokensInheritable() bool { + return c.maxTokensInheritable +} + func (c *ModelOptions) NoThinking() bool { return c.noThinking } @@ -178,6 +184,15 @@ func WithCompacting() Opt { func WithMaxTokens(maxTokens int64) Opt { return func(cfg *ModelOptions) { cfg.maxTokens = maxTokens + cfg.maxTokensInheritable = false + } +} + +// WithAuxiliaryMaxTokens imposes a cap on both direct and routed providers. +func WithAuxiliaryMaxTokens(maxTokens int64) Opt { + return func(cfg *ModelOptions) { + cfg.maxTokens = maxTokens + cfg.maxTokensInheritable = true } } @@ -276,7 +291,11 @@ func FromModelOptions(m ModelOptions) []Opt { out = append(out, WithNoThinking()) } if m.maxTokens != 0 { - out = append(out, WithMaxTokens(m.maxTokens)) + if m.maxTokensInheritable { + out = append(out, WithAuxiliaryMaxTokens(m.maxTokens)) + } else { + out = append(out, WithMaxTokens(m.maxTokens)) + } } if len(m.providers) > 0 { out = append(out, WithProviders(m.providers)) diff --git a/pkg/model/provider/options/options_test.go b/pkg/model/provider/options/options_test.go index a3a50f7b84..7f45801243 100644 --- a/pkg/model/provider/options/options_test.go +++ b/pkg/model/provider/options/options_test.go @@ -201,3 +201,36 @@ func TestModelOptionsClonePreservesNumericTypesAndIsolation(t *testing.T) { assert.Equal(t, int64(9007199254740993), values["nested"].([]any)[0].(map[string]any)["limit"]) assert.Equal(t, int64(9007199254740993), m.Providers()["custom"].ProviderOpts["nested"].([]any)[0].(map[string]any)["limit"]) } + +func TestMaxTokensPropagationAndRoundTrip(t *testing.T) { + t.Parallel() + for _, tc := range []struct { + name string + opts []Opt + want int64 + inheritable bool + }{ + {"ordinary", []Opt{WithMaxTokens(64)}, 64, false}, + {"auxiliary", []Opt{WithAuxiliaryMaxTokens(128)}, 128, true}, + {"ordinary last", []Opt{WithAuxiliaryMaxTokens(128), WithMaxTokens(64)}, 64, false}, + {"auxiliary last", []Opt{WithMaxTokens(64), WithAuxiliaryMaxTokens(128)}, 128, true}, + {"ordinary zero unchanged", []Opt{WithMaxTokens(64), WithMaxTokens(0)}, 0, false}, + } { + t.Run(tc.name, func(t *testing.T) { + original := Apply(tc.opts...) + cloned, err := original.Clone() + require.NoError(t, err) + replayed := FromModelOptions(cloned) + if tc.want != 0 { + require.Len(t, replayed, 1) + } else { + require.Empty(t, replayed) + } + roundTrip := Apply(replayed...) + for _, m := range []ModelOptions{original, cloned, roundTrip} { + assert.Equal(t, tc.want, m.MaxTokens()) + assert.Equal(t, tc.inheritable, m.MaxTokensInheritable()) + } + }) + } +} diff --git a/pkg/model/provider/rulebased/client.go b/pkg/model/provider/rulebased/client.go index 5dafa3dc9c..b0e782fc87 100644 --- a/pkg/model/provider/rulebased/client.go +++ b/pkg/model/provider/rulebased/client.go @@ -55,7 +55,12 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, models map[string]l return nil, errors.New("no routing rules configured") } + effective := options.Apply(opts...) routeOpts := filterOutMaxTokens(opts) + // Ordinary router limits stay local; only the effective auxiliary cap propagates. + if effective.MaxTokensInheritable() && effective.MaxTokens() != 0 { + routeOpts = append(routeOpts, options.WithAuxiliaryMaxTokens(effective.MaxTokens())) + } // Create fallback provider from the model's provider/model fields. fallbackSpec := cfg.Provider + "/" + cfg.Model diff --git a/pkg/runtime/sampling_routed_cap_test.go b/pkg/runtime/sampling_routed_cap_test.go new file mode 100644 index 0000000000..542ea47a18 --- /dev/null +++ b/pkg/runtime/sampling_routed_cap_test.go @@ -0,0 +1,60 @@ +//go:build !js && !docker_agent_no_openai + +//nolint:staticcheck // Sampling is supported during the MCP deprecation window. +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "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/openai" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/team" +) + +func TestSamplingHandlersHonorRoutedTokenCap(t *testing.T) { + t.Parallel() + for _, withTools := range []bool{false, true} { + for _, message := range []string{"hello", "unrelated"} { + t.Run(fmt.Sprintf("tools=%t/%s", withTools, message), func(t *testing.T) { + t.Parallel() + var payload map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.NoError(t, json.NewDecoder(r.Body).Decode(&payload)) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"sampled\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + })) + defer server.Close() + registry := provider.NewRegistry(map[string]provider.Factory{"openai_chatcompletions": func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) { + return openai.NewClient(ctx, cfg, env, opts...) + }}) + cfg := &latest.ModelConfig{Provider: "custom", Model: "gpt-5.6", MaxTokens: new(int64(4096)), Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + models := map[string]latest.ModelConfig{"target": {Provider: "custom", Model: "gpt-5.6", MaxTokens: new(int64(8192))}} + p, err := registry.NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"KEY": "fake"}), options.WithMaxTokens(4096), options.WithProviders(map[string]latest.ProviderConfig{"custom": {BaseURL: server.URL, TokenKey: "KEY", APIType: "openai_chatcompletions", MaxTokens: new(int64(16384)), ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}}})) + require.NoError(t, err) + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(agent.New("root", "", agent.WithModel(p)))), WithCurrentAgent("root"), WithModelStore(mockModelStore{})) + require.NoError(t, err) + if withTools { + _, err = rt.samplingWithToolsHandler(t.Context(), &mcp.CreateMessageWithToolsParams{MaxTokens: 512, Messages: []*mcp.SamplingMessageV2{{Role: "user", Content: []mcp.Content{&mcp.TextContent{Text: message}}}}}) + } else { + _, err = rt.samplingHandler(t.Context(), &mcp.CreateMessageParams{MaxTokens: 512, Messages: []*mcp.SamplingMessage{{Role: "user", Content: &mcp.TextContent{Text: message}}}}) + } + require.NoError(t, err) + assert.InDelta(t, float64(512), payload["max_completion_tokens"], 0) + assert.Equal(t, "none", payload["reasoning_effort"]) + }) + } + } +} diff --git a/pkg/sessiontitle/routed_cap_test.go b/pkg/sessiontitle/routed_cap_test.go new file mode 100644 index 0000000000..bdfa52886a --- /dev/null +++ b/pkg/sessiontitle/routed_cap_test.go @@ -0,0 +1,50 @@ +//go:build !js && !docker_agent_no_openai + +package sessiontitle + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "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/openai" + "github.com/docker/docker-agent/pkg/model/provider/options" +) + +func TestGenerateOnceHonorsRoutedTokenCap(t *testing.T) { + t.Parallel() + for _, message := range []string{"hello", "unrelated"} { + t.Run(message, func(t *testing.T) { + t.Parallel() + var payload map[string]any + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.NoError(t, json.NewDecoder(r.Body).Decode(&payload)) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"Title\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n") + })) + defer server.Close() + registry := provider.NewRegistry(map[string]provider.Factory{"openai_chatcompletions": func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (provider.Provider, error) { + return openai.NewClient(ctx, cfg, env, opts...) + }}) + cfg := &latest.ModelConfig{Provider: "custom", Model: "gpt-5.6", MaxTokens: new(int64(4096)), Routing: []latest.RoutingRule{{Model: "target", Examples: []string{"hello"}}}} + models := map[string]latest.ModelConfig{"target": {Provider: "custom", Model: "gpt-5.6", MaxTokens: new(int64(8192))}} + p, err := registry.NewWithModels(t.Context(), cfg, models, environment.NewMapEnvProvider(map[string]string{"KEY": "fake"}), options.WithMaxTokens(4096), options.WithProviders(map[string]latest.ProviderConfig{"custom": {BaseURL: server.URL, TokenKey: "KEY", APIType: "openai_chatcompletions", MaxTokens: new(int64(16384)), ThinkingBudget: &latest.ThinkingBudget{Effort: "high"}}})) + require.NoError(t, err) + title, err := generateOnce(t.Context(), p, []chat.Message{{Role: chat.MessageRoleUser, Content: message}}) + require.NoError(t, err) + assert.Equal(t, "Title", title) + assert.InDelta(t, float64(256), payload["max_completion_tokens"], 0) + assert.Equal(t, "none", payload["reasoning_effort"]) + }) + } +} From ac3407391a4ca953c29202e98e4263886f6c9fda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Wed, 14 Oct 2026 09:35:38 +0000 Subject: [PATCH 5/5] fix: retain the last nonzero clone token override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- docs/guides/go-sdk/index.md | 4 ++- pkg/model/provider/clone.go | 18 +++++++++--- pkg/model/provider/clone_matrix_test.go | 3 ++ pkg/model/provider/clone_thinking_test.go | 35 +++++++++++++++++++++++ 4 files changed, 55 insertions(+), 5 deletions(-) diff --git a/docs/guides/go-sdk/index.md b/docs/guides/go-sdk/index.md index 780a9aa228..2c6a687199 100644 --- a/docs/guides/go-sdk/index.md +++ b/docs/guides/go-sdk/index.md @@ -884,7 +884,9 @@ are bounded to 128 levels and 100,000 visited values. Cumulative clones preserve removal, gateway clearing and no-thinking overrides. A non-zero `WithMaxTokens` cap supplied to `Clone` overrides direct, matched-route and fallback limits, subject to existing provider token floors. Cumulative clones retain that cap -when it is omitted or zero; a later non-zero cap replaces it. Named routing definitions are +when it is omitted or zero; the last non-zero cap wins, including within one +clone call. Ordinary construction-time `WithMaxTokens(0)` reset semantics are +unchanged. Named routing definitions are captured at construction and copied for each child, so later input-map mutations do not change clones. diff --git a/pkg/model/provider/clone.go b/pkg/model/provider/clone.go index a0318d1df5..825bb9f98a 100644 --- a/pkg/model/provider/clone.go +++ b/pkg/model/provider/clone.go @@ -33,7 +33,7 @@ func (e *CloneError) Unwrap() error { return e.Err } // Clone reconstructs a provider with ordered option overrides, or returns an error. // Non-zero token caps also constrain routed children, subject to provider floors. -// An omitted or zero cap retains a previous explicit clone cap. +// The last non-zero cap wins; zero retains the latest explicit cap. func Clone(ctx context.Context, p Provider, opts ...options.Opt) (Provider, error) { clone, _, err := cloneProvider(ctx, p, opts) return clone, err @@ -99,9 +99,19 @@ func mergeCloneOptions(cfg base.Config, opts []options.Opt) (latest.ModelConfig, } baseOpts := options.FromModelOptions(baseOptions) mergedOpts := slices.Concat(baseOpts, opts) - requested := options.Apply(opts...) - if mt := requested.MaxTokens(); mt != 0 { - mergedOpts = append(mergedOpts, options.WithAuxiliaryMaxTokens(mt)) + var requested options.ModelOptions + var explicitCap int64 + for _, opt := range opts { + if opt == nil { + continue + } + opt(&requested) + if mt := requested.MaxTokens(); mt != 0 { + explicitCap = mt + } + } + if explicitCap != 0 { + mergedOpts = append(mergedOpts, options.WithAuxiliaryMaxTokens(explicitCap)) } else if baseOptions.MaxTokensInheritable() && baseOptions.MaxTokens() != 0 { mergedOpts = append(mergedOpts, options.WithAuxiliaryMaxTokens(baseOptions.MaxTokens())) } diff --git a/pkg/model/provider/clone_matrix_test.go b/pkg/model/provider/clone_matrix_test.go index ac6f496d90..38a78cab45 100644 --- a/pkg/model/provider/clone_matrix_test.go +++ b/pkg/model/provider/clone_matrix_test.go @@ -280,6 +280,9 @@ func TestCloneRoutedTokenCapsPreserveConstructionDefaults(t *testing.T) { plain, err := Clone(t.Context(), original) require.NoError(t, err) request(plain, 0) + plain, err = Clone(t.Context(), original, options.WithNoThinking(), options.WithGeneratingTitle()) + require.NoError(t, err) + request(plain, 0) for _, opts := range [][]options.Opt{ {options.WithMaxTokens(128)}, {options.WithAuxiliaryMaxTokens(1024), options.WithMaxTokens(128)}, diff --git a/pkg/model/provider/clone_thinking_test.go b/pkg/model/provider/clone_thinking_test.go index 3c421e25de..4c37b57133 100644 --- a/pkg/model/provider/clone_thinking_test.go +++ b/pkg/model/provider/clone_thinking_test.go @@ -99,6 +99,25 @@ func TestCloneRoutedClaudeStrictOffOnWire(t *testing.T) { request(clone, "unrelated", false, wireCap) } } + for _, opts := range [][]options.Opt{ + {options.WithMaxTokens(1024), nil, options.WithMaxTokens(0)}, + {options.WithAuxiliaryMaxTokens(1024), options.WithMaxTokens(0)}, + {options.WithMaxTokens(1024), options.WithAuxiliaryMaxTokens(0)}, + } { + clone, err := Clone(t.Context(), original, options.WithNoThinking(), options.WithMaxTokens(128)) + require.NoError(t, err) + clone, err = Clone(t.Context(), clone, opts...) + require.NoError(t, err) + for range 2 { + cfg := clone.BaseConfig() + assert.Equal(t, int64(1024), *cfg.ModelConfig.MaxTokens) + assert.Equal(t, int64(1024), cfg.ModelOptions.MaxTokens()) + request(clone, "hello", false, 1024) + request(clone, "unrelated", false, 1024) + clone, err = Clone(t.Context(), clone) + require.NoError(t, err) + } + } request(original, "hello", true, defaultCap) }) } @@ -203,5 +222,21 @@ func TestCloneDirectClaudeRetainsExplicitCapOnZero(t *testing.T) { assert.Equal(t, int64(128), *cfg.ModelConfig.MaxTokens) request(p, 256) } + for _, opts := range [][]options.Opt{ + {options.WithMaxTokens(1024), options.WithMaxTokens(0)}, + {options.WithAuxiliaryMaxTokens(1024), options.WithMaxTokens(0)}, + {options.WithMaxTokens(1024), options.WithAuxiliaryMaxTokens(0)}, + } { + clone, err := Clone(t.Context(), p, opts...) + require.NoError(t, err) + for range 2 { + cfg := clone.BaseConfig() + assert.Equal(t, int64(1024), cfg.ModelOptions.MaxTokens()) + assert.Equal(t, int64(1024), *cfg.ModelConfig.MaxTokens) + request(clone, 1024) + clone, err = Clone(t.Context(), clone) + require.NoError(t, err) + } + } request(original, 8192) }