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) +}