diff --git a/pkg/model/provider/dmr/client.go b/pkg/model/provider/dmr/client.go index 6221431a3..77b9fbabc 100644 --- a/pkg/model/provider/dmr/client.go +++ b/pkg/model/provider/dmr/client.go @@ -9,7 +9,9 @@ import ( "log/slog" "maps" "net/http" + "net/url" "os" + "strings" "time" "github.com/openai/openai-go/v3" @@ -18,6 +20,7 @@ import ( "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/httpclient" "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/model/provider/dmr/dmrmodels" "github.com/docker/docker-agent/pkg/model/provider/oaistream" @@ -44,7 +47,8 @@ type Client struct { client openai.Client httpClient *http.Client - engine string + // engine is empty in gateway mode: engine-gated request shaping stays on defaults. + engine string // attachmentCaps records the document MIME types this DMR-hosted model is // declared to accept natively, parsed from provider_opts.supports_images / @@ -68,34 +72,48 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, opts ...options.Opt globalOptions := options.Apply(opts...) - // Skip docker model status query when BaseURL is explicitly provided. - // This avoids unnecessary exec calls and speeds up tests/CI scenarios. - var endpoint, engine string + var endpoint, engine, baseURL string + var httpClient *http.Client verifyViaAPI := false - if cfg.BaseURL == "" && os.Getenv("MODEL_RUNNER_HOST") == "" { - var err error - endpoint, engine, err = dmrmodels.DockerModelEndpointAndEngine(ctx) - switch { - case err == nil: - // Auto-pull the model if needed - if err := pullDockerModelIfNeeded(ctx, cfg.Model); err != nil { - slog.DebugContext(ctx, "docker model pull failed", "error", err) - return nil, err + gateway := globalOptions.Gateway() + if gateway != "" { + // Skip local discovery: the gateway may be the only reachable path + // (e.g. inside a sandbox). + u, err := url.Parse(gateway) + if err != nil { + return nil, fmt.Errorf("invalid models gateway URL: %w", err) + } + baseURL = fmt.Sprintf("%s://%s%s/v1/", u.Scheme, u.Host, strings.TrimSuffix(u.Path, "/")) + httpClient = httpclient.NewHTTPClient(ctx, base.GatewayHTTPOptions(u, dmrmodels.DefaultHostURL(), cfg, &globalOptions)...) + globalOptions.WrapTransport(ctx, httpClient) + } else { + // Skip docker model status query when BaseURL is explicitly provided. + // This avoids unnecessary exec calls and speeds up tests/CI scenarios. + if cfg.BaseURL == "" && os.Getenv("MODEL_RUNNER_HOST") == "" { + var err error + endpoint, engine, err = dmrmodels.DockerModelEndpointAndEngine(ctx) + switch { + case err == nil: + // Auto-pull the model if needed + if err := pullDockerModelIfNeeded(ctx, cfg.Model); err != nil { + slog.DebugContext(ctx, "docker model pull failed", "error", err) + return nil, err + } + case dmrmodels.IsNotInstalledError(err): + slog.DebugContext(ctx, "docker model status query failed", "error", err) + return nil, ErrNotInstalled + default: + // The `docker model` CLI is unusable (broken plugin, docker not on + // PATH, ...) but the DMR endpoint may still be up: check model + // availability through the HTTP API below so a missing model fails + // here instead of as a raw HTTP 404 at message time. + slog.ErrorContext(ctx, "docker model status query failed", "error", err) + verifyViaAPI = true } - case dmrmodels.IsNotInstalledError(err): - slog.DebugContext(ctx, "docker model status query failed", "error", err) - return nil, ErrNotInstalled - default: - // The `docker model` CLI is unusable (broken plugin, docker not on - // PATH, ...) but the DMR endpoint may still be up: check model - // availability through the HTTP API below so a missing model fails - // here instead of as a raw HTTP 404 at message time. - slog.ErrorContext(ctx, "docker model status query failed", "error", err) - verifyViaAPI = true } - } - baseURL, httpClient := dmrmodels.ResolveBaseURL(ctx, cfg, endpoint) + baseURL, httpClient = dmrmodels.ResolveBaseURL(ctx, cfg, endpoint) + } // A custom transport (e.g. the Docker Unix socket) must also be used by // the OpenAI adapter, not just the direct HTTP calls. @@ -137,8 +155,9 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, opts ...options.Opt ) // Skip model configuration for title-generation and compaction clones to // avoid reconfiguring the model with different settings (e.g., smaller - // max_tokens) that would affect the main agent. - if !globalOptions.GeneratingTitle() && !globalOptions.Compacting() { + // max_tokens) that would affect the main agent. It is local-only, so + // gateway mode skips it too. + if gateway == "" && !globalOptions.GeneratingTitle() && !globalOptions.Compacting() { if err := configureModel(ctx, httpClient, baseURL, cfg.Model, backendCfg, parsed.mode, parsed.rawRuntimeFlags); err != nil { slog.DebugContext(ctx, "model configure via API skipped or failed", "error", err) } diff --git a/pkg/model/provider/dmr/client_test.go b/pkg/model/provider/dmr/client_test.go index dc7c594cb..7fa735d5a 100644 --- a/pkg/model/provider/dmr/client_test.go +++ b/pkg/model/provider/dmr/client_test.go @@ -17,6 +17,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/dmr/dmrmodels" "github.com/docker/docker-agent/pkg/model/provider/options" ) @@ -34,6 +35,66 @@ func TestNewClientWithExplicitBaseURL(t *testing.T) { assert.Equal(t, "https://custom.example.com:8080/api/v1", client.BaseURL) } +func TestNewClientUsesModelsGateway(t *testing.T) { + t.Parallel() + + var received *http.Request + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + received = r.Clone(r.Context()) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"embedding":[0.1,0.2]}],"usage":{"prompt_tokens":1,"total_tokens":1}}`)) + })) + defer server.Close() + + cfg := &latest.ModelConfig{ + Provider: "dmr", + Model: "ai/qwen3", + } + + client, err := NewClient(t.Context(), cfg, options.WithGateway(server.URL+"/engines?trace=1")) + require.NoError(t, err) + assert.Equal(t, server.URL+"/engines/v1/", client.BaseURL) + + // A real SDK-backed operation must carry the gateway request contract: + // forward target, provider/model identity, and gateway query params. + _, err = client.CreateBatchEmbedding(t.Context(), []string{"hello"}) + require.NoError(t, err) + + require.NotNil(t, received) + assert.Equal(t, "/engines/v1/embeddings", received.URL.Path) + assert.Equal(t, "1", received.URL.Query().Get("trace")) + assert.Equal(t, dmrmodels.DefaultHostURL(), received.Header.Get("X-Cagent-Forward")) + assert.Equal(t, "dmr", received.Header.Get("X-Cagent-Provider")) + assert.Equal(t, "ai/qwen3", received.Header.Get("X-Cagent-Model")) +} + +func TestNewClientGatewaySkipsLocalDiscovery(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Skipping docker CLI shim test on Windows") + } + + // A docker CLI that predates Model Runner must not matter when a models + // gateway is configured: no local discovery should happen at all. + tempDir := t.TempDir() + dockerPath := filepath.Join(tempDir, "docker") + script := "#!/bin/sh\n" + + "printf 'unknown flag: --json\\n' >&2\n" + + "exit 1\n" + require.NoError(t, os.WriteFile(dockerPath, []byte(script), 0o755)) + + t.Setenv("PATH", tempDir+string(os.PathListSeparator)+os.Getenv("PATH")) + t.Setenv("MODEL_RUNNER_HOST", "") + + cfg := &latest.ModelConfig{ + Provider: "dmr", + Model: "ai/qwen3", + } + + client, err := NewClient(t.Context(), cfg, options.WithGateway("http://host.docker.internal:12434/engines/")) + require.NoError(t, err) + assert.Equal(t, "http://host.docker.internal:12434/engines/v1/", client.BaseURL) +} + func TestNewClientReturnsErrNotInstalledWhenDockerModelUnsupported(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("Skipping docker CLI shim test on Windows") diff --git a/pkg/model/provider/dmr/dmrmodels/resolve.go b/pkg/model/provider/dmr/dmrmodels/resolve.go index 821e1105a..f64aab6eb 100644 --- a/pkg/model/provider/dmr/dmrmodels/resolve.go +++ b/pkg/model/provider/dmr/dmrmodels/resolve.go @@ -61,9 +61,9 @@ func defaultContainerURL() string { return "http://model-runner.docker.internal" + dmrInferencePrefix + "/v1/" } -// defaultHostURL is the default DMR URL when running on the host with no +// DefaultHostURL is the default DMR URL when running on the host with no // explicit endpoint. It targets the standard local model-runner port. -func defaultHostURL() string { +func DefaultHostURL() string { return defaultURL("127.0.0.1", dmrDefaultPort) } @@ -73,7 +73,7 @@ func defaultForEnvironment() string { if inContainer() { return defaultContainerURL() } - return defaultHostURL() + return DefaultHostURL() } func inContainer() bool { @@ -121,7 +121,7 @@ func getDMRFallbackURLs(containerized bool) []string { defaultURL("172.17.0.1", dmrDefaultPort), } } - return []string{defaultHostURL()} + return []string{DefaultHostURL()} } // ResolveBaseURL determines the correct base URL to talk to Docker Model @@ -188,7 +188,7 @@ func resolvePrimaryDMRURL(endpoint string) (string, *http.Client) { // Legacy bug workaround: old DMR versions <= 0.1.44 could report http://:0/engines/v1/. if ep == "http://:0/engines/v1/" { - return defaultHostURL(), nil + return DefaultHostURL(), nil } if ep == "" {