From 5f997570701f278fb06c8d44812f859b84c40a02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sun, 16 Aug 2026 21:25:13 +0200 Subject: [PATCH] fix(#3996): keep image-output models eligible for session titles Title generation is an explicitly text-only call, so image-output-capable models remain valid title candidates and may fall through to the next model on failure. Preserve title_model -> agent model -> fallbacks order, skip nil providers, and return nil only when no configured candidate remains. Apply the same construction behavior to LocalRuntime, the server session path, and the debug title command. Tests pin image-output candidate ordering, fallback, and isolation from the main generation stream. --- cmd/root/debug.go | 5 +- docs/configuration/models/index.md | 7 ++ pkg/runtime/runtime.go | 13 ++- pkg/runtime/title_generator_test.go | 131 ++++++++++++++++++++++++++++ pkg/server/session_manager.go | 8 +- pkg/server/title_generator_test.go | 73 ++++++++++++++++ pkg/sessiontitle/generator.go | 24 ++--- pkg/sessiontitle/generator_test.go | 92 +++++++++++++++++++ 8 files changed, 330 insertions(+), 23 deletions(-) create mode 100644 pkg/runtime/title_generator_test.go create mode 100644 pkg/server/title_generator_test.go diff --git a/cmd/root/debug.go b/cmd/root/debug.go index 0c282a65c9..0882d1873c 100644 --- a/cmd/root/debug.go +++ b/cmd/root/debug.go @@ -222,7 +222,10 @@ func (f *debugFlags) runDebugTitleCommand(cmd *cobra.Command, args []string) (co if len(models) == 0 { return fmt.Errorf("agent %q has no model configured", agent.Name()) } - gen := sessiontitle.New(models[0], models[1:]...) + gen := sessiontitle.New(models...) + if gen == nil { + return fmt.Errorf("agent %q has no usable title model", agent.Name()) + } title, err := gen.Generate(ctx, "debug", []string{args[1]}) if err != nil { diff --git a/docs/configuration/models/index.md b/docs/configuration/models/index.md index f0cf4fd481..e9be86cb4c 100644 --- a/docs/configuration/models/index.md +++ b/docs/configuration/models/index.md @@ -179,6 +179,13 @@ behavior that specifically keys off a declared image-output capability (for example, a provider-specific request-shape guard); it does not by itself change what Docker Agent sends to or renders from the model. +One side effect of declaring `image: true`: the model is skipped as a +session-title candidate, because titles are generated by a plain text-only +completion that image-output routes can reject. Title generation uses the +first non-image-output candidate (dedicated `title_model`, then the agent's +model, then its fallbacks); when every candidate declares image output, the +automatic title is skipped and the session keeps its default title. + > [!WARNING] > **Constraint** > diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 073ef41759..8c2c498378 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -123,7 +123,8 @@ type Runtime interface { UpdateSessionTitle(ctx context.Context, sess *session.Session, title string) error // TitleGenerator returns a generator for automatic session titles, or nil - // if the runtime does not support local title generation (e.g. remote runtimes). + // if the runtime does not support local title generation (e.g. remote + // runtimes) or no configured model is a usable title candidate. TitleGenerator(ctx context.Context) *sessiontitle.Generator // Steer enqueues a user message for urgent mid-turn injection into the @@ -1359,17 +1360,15 @@ func (r *LocalRuntime) ExecuteMCPPrompt(ctx context.Context, promptName string, return "", fmt.Errorf("MCP prompt '%s' not found in any active toolset", promptName) } -// TitleGenerator returns a title generator for automatic session title generation. +// TitleGenerator returns a title generator for automatic session title +// generation, or nil when no configured model is a usable title candidate +// (see [sessiontitle.New]). func (r *LocalRuntime) TitleGenerator(ctx context.Context) *sessiontitle.Generator { a := r.CurrentAgent() if a == nil { return nil } - models := a.TitleModels(ctx) - if len(models) == 0 { - return nil - } - return sessiontitle.New(models[0], models[1:]...) + return sessiontitle.New(a.TitleModels(ctx)...) } // getAgentModelID returns the model ID for an agent. The zero ID is diff --git a/pkg/runtime/title_generator_test.go b/pkg/runtime/title_generator_test.go new file mode 100644 index 0000000000..1c49138f4f --- /dev/null +++ b/pkg/runtime/title_generator_test.go @@ -0,0 +1,131 @@ +package runtime + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + + "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/base" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" +) + +type imageOutputProvider struct { + id string + stream chat.MessageStream + calls atomic.Int64 +} + +func (p *imageOutputProvider) ID() modelsdev.ID { return modelsdev.ParseIDOrZero(p.id) } + +func (p *imageOutputProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + p.calls.Add(1) + if p.stream == nil { + return nil, errors.New("no stream configured") + } + return p.stream, nil +} + +func (p *imageOutputProvider) BaseConfig() base.Config { + return base.Config{ModelConfig: latest.ModelConfig{ + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}, + }} +} + +func (p *imageOutputProvider) MaxTokens() int { return 0 } + +func newTitleTestRuntime(t *testing.T, root *agent.Agent) *LocalRuntime { + t.Helper() + rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)), + WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + return rt +} + +// Image-output models are safe title candidates because Gemini title requests +// deliberately omit response modalities and the media marker instruction. +func TestLocalRuntime_TitleGenerator_AllImageOutputCandidatesRemainEligible(t *testing.T) { + t.Parallel() + + primary := &imageOutputProvider{id: "google/image-primary"} + title := &imageOutputProvider{id: "google/image-title"} + fallback := &imageOutputProvider{id: "google/image-fallback"} + root := agent.New("root", "test", + agent.WithModel(primary), + agent.WithTitleModel(title), + agent.WithFallbackModel(fallback), + ) + rt := newTitleTestRuntime(t, root) + + assert.NotNil(t, rt.TitleGenerator(t.Context()), + "image-output models remain eligible for text-only title generation") + assert.Zero(t, primary.calls.Load()) + assert.Zero(t, title.calls.Load()) + assert.Zero(t, fallback.calls.Load()) +} + +// A dedicated image-output title model remains first in the candidate order; +// title requests are text-only, and failure falls through to the agent model. +func TestLocalRuntime_TitleGenerator_ImageOutputTitleModelFallsThrough(t *testing.T) { + t.Parallel() + + imageTitle := &imageOutputProvider{id: "google/image-title"} + safe := &countingProvider{ + id: "safe/model", + stream: newStreamBuilder().AddContent("A Title").AddStopWithUsage(5, 3).Build(), + } + root := agent.New("root", "test", + agent.WithModel(safe), + agent.WithTitleModel(imageTitle), + ) + rt := newTitleTestRuntime(t, root) + + gen := rt.TitleGenerator(t.Context()) + require.NotNil(t, gen) + + generated, err := gen.Generate(t.Context(), "sess-1", []string{"hello"}) + require.NoError(t, err) + assert.Equal(t, "A Title", generated) + assert.Equal(t, int64(1), imageTitle.calls.Load(), "the image-output title model is attempted first") + assert.Equal(t, 1, safe.callCount) +} + +// The title candidate list must not affect the main generation stream: an +// image-output model remains usable by both normal and text-only title calls. +func TestRunStream_ImageOutputOnlyModel_MainStreamUnaffected(t *testing.T) { + t.Parallel() + + provider := &imageOutputProvider{ + id: "google/image-model", + stream: newStreamBuilder().AddContent("here is your cat").AddStopWithUsage(10, 5).Build(), + } + root := agent.New("root", "test", agent.WithModel(provider)) + rt := newTitleTestRuntime(t, root) + + require.NotNil(t, rt.TitleGenerator(t.Context())) + + sess := session.New(session.WithUserMessage("draw a cat")) + var content strings.Builder + for ev := range rt.RunStream(t.Context(), sess) { + switch e := ev.(type) { + case *ErrorEvent: + t.Fatalf("main stream must stay successful, got ErrorEvent %q", e.Error) + case *AgentChoiceEvent: + content.WriteString(e.Content) + } + } + + assert.Equal(t, "here is your cat", content.String()) + assert.Equal(t, int64(1), provider.calls.Load(), "exactly the main generation request reaches the provider") + assert.Empty(t, sess.Title, "the session keeps its default title when title generation is skipped") +} diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 768689c12b..605deba102 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -1540,11 +1540,9 @@ func (sm *SessionManager) runtimeForSession(ctx context.Context, sess *session.S // the requested models instead of the agent's defaults. applyStoredOverrides(ctx, sess.ID, run, sess.AgentModelOverrides) - titleModels := agt.TitleModels(ctx) - var titleGen *sessiontitle.Generator - if len(titleModels) > 0 { - titleGen = sessiontitle.New(titleModels[0], titleModels[1:]...) - } + // May be nil when the agent has no configured title candidate; a nil + // generator skips title generation. + titleGen := sessiontitle.New(agt.TitleModels(ctx)...) // Construction succeeded: the selected author default may now be // committed and the pending marker consumed, exactly once. diff --git a/pkg/server/title_generator_test.go b/pkg/server/title_generator_test.go new file mode 100644 index 0000000000..a5cf0a8031 --- /dev/null +++ b/pkg/server/title_generator_test.go @@ -0,0 +1,73 @@ +package server + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/session" +) + +const imageOnlyTitleConfig = `models: + imagegen: + provider: openai + model: fake-image-model + output_capabilities: + image: true +agents: + root: + model: imagegen + instruction: Be helpful. +` + +const mixedTitleConfig = `models: + imagegen: + provider: openai + model: fake-image-model + output_capabilities: + image: true + title_model: safe + safe: + provider: openai + model: gpt-4o-mini +agents: + root: + model: imagegen + instruction: Be helpful. +` + +func TestRuntimeForSession_TitleGeneratorKeepsImageOutputModels(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "dummy") + ctx := t.Context() + + sources := config.Sources{ + "image-only.yaml": config.NewBytesSource("image-only.yaml", []byte(imageOnlyTitleConfig)), + "mixed.yaml": config.NewBytesSource("mixed.yaml", []byte(mixedTitleConfig)), + } + store := session.NewInMemorySessionStore() + sm := NewSessionManager(ctx, sources, store, 0, &config.RuntimeConfig{}) + + t.Run("image-output candidates retain title generation", func(t *testing.T) { + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + + run, titleGen, err := sm.runtimeForSession(ctx, sess, "image-only.yaml", "", &config.RuntimeConfig{}) + require.NoError(t, err) + t.Cleanup(func() { _ = run.Close() }) + + assert.NotNil(t, titleGen, "an image-output-only agent can generate text-only titles") + }) + + t.Run("safe dedicated title model keeps titles enabled", func(t *testing.T) { + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + + run, titleGen, err := sm.runtimeForSession(ctx, sess, "mixed.yaml", "", &config.RuntimeConfig{}) + require.NoError(t, err) + t.Cleanup(func() { _ = run.Close() }) + + assert.NotNil(t, titleGen, "a safe title_model must keep title generation available") + }) +} diff --git a/pkg/sessiontitle/generator.go b/pkg/sessiontitle/generator.go index 6c6df138fb..a7b39e1f50 100644 --- a/pkg/sessiontitle/generator.go +++ b/pkg/sessiontitle/generator.go @@ -9,7 +9,6 @@ import ( "fmt" "io" "log/slog" - "slices" "strings" "time" @@ -46,15 +45,20 @@ type Generator struct { models []provider.Provider } -// New creates a new title Generator. The first model is the primary; any -// additional ones are fallbacks tried in order if earlier attempts fail. -// Nil providers are silently ignored. -func New(model provider.Provider, fallbackModels ...provider.Provider) *Generator { - models := slices.DeleteFunc( - append([]provider.Provider{model}, fallbackModels...), - func(p provider.Provider) bool { return p == nil }, - ) - return &Generator{models: models} +// New creates a title Generator from the ordered candidate models. Nil +// providers are skipped. Image-output-capable providers remain eligible +// because title-generation calls are explicitly text-only. +func New(models ...provider.Provider) *Generator { + usable := make([]provider.Provider, 0, len(models)) + for _, model := range models { + if model != nil { + usable = append(usable, model) + } + } + if len(usable) == 0 { + return nil + } + return &Generator{models: usable} } // Generate produces a title for a session based on the provided user messages. diff --git a/pkg/sessiontitle/generator_test.go b/pkg/sessiontitle/generator_test.go index 7041021abb..62b7a25e5e 100644 --- a/pkg/sessiontitle/generator_test.go +++ b/pkg/sessiontitle/generator_test.go @@ -10,6 +10,7 @@ import ( "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/model/provider/base" "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/tools" @@ -189,3 +190,94 @@ func TestGenerator_Generate_FallsBackOnEmptyOutput(t *testing.T) { assert.Equal(t, 1, primary.calls) assert.Equal(t, 1, fallback.calls) } + +// imageOutputProvider is an image-capable title candidate. +func imageOutputProvider(id modelsdev.ID) *mockProvider { + return &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)}, + }} + }, + } +} + +func TestNew_ImageOutputCandidateFallsThroughOnFailure(t *testing.T) { + t.Parallel() + + img := imageOutputProvider(modelsdev.NewID("google", "image-primary")) + safe := &mockProvider{ + id: modelsdev.NewID("safe", "fallback"), + createFn: func() (chat.MessageStream, error) { + return streamWithContent("Safe Title"), nil + }, + } + + title, err := New(img, safe).Generate(t.Context(), "sess-1", []string{"hello"}) + require.NoError(t, err) + assert.Equal(t, "Safe Title", title) + assert.Equal(t, 1, img.calls) + assert.Equal(t, 1, safe.calls) +} + +func TestNew_AllImageOutputCandidatesRemainEligible(t *testing.T) { + t.Parallel() + + img1 := imageOutputProvider(modelsdev.NewID("google", "image-one")) + img2 := &mockProvider{ + id: modelsdev.NewID("google", "image-two"), + createFn: func() (chat.MessageStream, error) { + return streamWithContent("Image Model Title"), nil + }, + } + + gen := New(img1, img2) + require.NotNil(t, gen) + title, err := gen.Generate(t.Context(), "sess-1", []string{"hello"}) + require.NoError(t, err) + assert.Equal(t, "Image Model Title", title) + assert.Equal(t, 1, img1.calls) + assert.Equal(t, 1, img2.calls) +} + +func TestNew_NoUsableCandidates_ReturnsNil(t *testing.T) { + t.Parallel() + + assert.Nil(t, New()) + assert.Nil(t, New(nil, nil)) +} + +func TestNew_KeepsUndeclaredAndExplicitFalseCandidates(t *testing.T) { + t.Parallel() + + undeclared := &mockProvider{ + id: modelsdev.NewID("safe", "undeclared"), + createFn: func() (chat.MessageStream, error) { + return streamWithContent("Ordinary Title"), nil + }, + } + declaredFalse := &mockProvider{ + id: modelsdev.NewID("safe", "declared-false"), + createFn: func() (chat.MessageStream, error) { + return streamWithContent("Fallback Title"), nil + }, + baseCfgFn: func() base.Config { + return base.Config{ModelConfig: latest.ModelConfig{ + OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)}, + }} + }, + } + + gen := New(undeclared, declaredFalse) + require.NotNil(t, gen) + + title, err := gen.Generate(t.Context(), "sess-1", []string{"hello"}) + require.NoError(t, err) + assert.Equal(t, "Ordinary Title", title) + assert.Equal(t, 1, undeclared.calls) + assert.Zero(t, declaredFalse.calls) +}