diff --git a/pkg/environment/credential_helper_test.go b/pkg/environment/credential_helper_test.go index f7caccf28f..43272a822c 100644 --- a/pkg/environment/credential_helper_test.go +++ b/pkg/environment/credential_helper_test.go @@ -1,6 +1,7 @@ package environment import ( + "runtime" "testing" "github.com/stretchr/testify/assert" @@ -17,6 +18,19 @@ func TestNewCredentialHelperProvider(t *testing.T) { func TestCredentialHelperProvider_Get(t *testing.T) { t.Parallel() + echoCmd := "echo" + echoArgs := func(v string) []string { return []string{v} } + falseCmd := "false" + + if runtime.GOOS == "windows" { + echoCmd = "powershell" + echoArgs = func(v string) []string { + return []string{"-NoProfile", "-Command", "Write-Output '" + v + "'"} + } + falseCmd = "powershell" + // simulate 'false' by exiting with 1 + } + tests := []struct { name string command string @@ -25,11 +39,11 @@ func TestCredentialHelperProvider_Get(t *testing.T) { wantValue string wantFound bool }{ - {"ignores non-DOCKER_TOKEN vars", "echo", []string{"test-token"}, "OTHER_VAR", "", false}, - {"success", "echo", []string{"my-secret-token"}, DockerDesktopTokenEnv, "my-secret-token", true}, - {"trims whitespace", "echo", []string{" token-with-spaces "}, DockerDesktopTokenEnv, "token-with-spaces", true}, - {"empty output", "echo", []string{""}, DockerDesktopTokenEnv, "", false}, - {"command fails", "false", nil, DockerDesktopTokenEnv, "", false}, + {"ignores non-DOCKER_TOKEN vars", echoCmd, echoArgs("test-token"), "OTHER_VAR", "", false}, + {"success", echoCmd, echoArgs("my-secret-token"), DockerDesktopTokenEnv, "my-secret-token", true}, + {"trims whitespace", echoCmd, echoArgs(" token-with-spaces "), DockerDesktopTokenEnv, "token-with-spaces", true}, + {"empty output", echoCmd, echoArgs(""), DockerDesktopTokenEnv, "", false}, + {"command fails", falseCmd, []string{"-NoProfile", "-Command", "exit 1"}, DockerDesktopTokenEnv, "", false}, {"command not found", "nonexistent-command-12345", nil, DockerDesktopTokenEnv, "", false}, } diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 768689c12b..26c3cd0b78 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -1039,9 +1039,18 @@ func (sm *SessionManager) RunSession(ctx context.Context, sessionID, agentFilena defer cancel() defer runtimeSession.streaming.Unlock() - // Start title generation in parallel if needed + // Start title generation in parallel if needed, coordinating via WaitGroup + // so close(streamChan) does not fire while generateTitle is still sending. + var wg sync.WaitGroup if needsTitle { - go sm.generateTitle(ctx, sess, titleGen, userMessages, streamChan) + // Note: generateTitle runs on the request ctx, so wg.Wait() only unblocks + // quickly when that context is cancelled. In the DeleteSession path + // (where the stream context is cancelled but the request context stays alive), + // the wait blocks until the title LLM call completes, delaying stream teardown. + // This bounded delay (first turn only) prioritizes persisting the title. + wg.Go(func() { + sm.generateTitle(ctx, sess, titleGen, userMessages, streamChan) + }) } else if titleToEmit != "" { // Re-emit the existing title so late-joining SSE consumers // and boards can pick it up without an extra API call. @@ -1051,11 +1060,17 @@ func (sm *SessionManager) RunSession(ctx context.Context, sessionID, agentFilena stream := runtimeSession.runtime.RunStream(streamCtx, sess) for event := range stream { if streamCtx.Err() != nil { - return + break } streamChan <- event } + wg.Wait() + + if streamCtx.Err() != nil { + return + } + if err := sm.sessionStore.UpdateSession(ctx, sess); err != nil { return } diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 595a565cf1..1a5e01744b 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "errors" + "io" "net/http" "net/http/httptest" "path/filepath" @@ -25,6 +26,9 @@ import ( "github.com/docker/docker-agent/pkg/concurrent" "github.com/docker/docker-agent/pkg/config" "github.com/docker/docker-agent/pkg/config/types" + "github.com/docker/docker-agent/pkg/model/provider" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/runtime" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/session/sqlitestore" @@ -2301,3 +2305,80 @@ func TestExportSessionForRecovery_OmitsErrorsKeyWhenNone(t *testing.T) { _, present := export["errors"] assert.False(t, present) } + +type blockingStream struct { + chat.MessageStream + + blocker chan struct{} + returned bool +} + +func (s *blockingStream) Recv() (chat.MessageStreamResponse, error) { + if !s.returned { + <-s.blocker + s.returned = true + return chat.MessageStreamResponse{ + Choices: []chat.MessageStreamChoice{ + {Delta: chat.MessageDelta{Content: "A Generated Title"}}, + }, + }, nil + } + return chat.MessageStreamResponse{}, io.EOF +} + +func (s *blockingStream) Close() {} + +type blockingTitleProvider struct { + provider.Provider + + blocker chan struct{} +} + +func (p *blockingTitleProvider) ID() modelsdev.ID { return modelsdev.NewID("mock", "mock") } + +func (p *blockingTitleProvider) BaseConfig() base.Config { return base.Config{} } + +func (p *blockingTitleProvider) CreateChatCompletionStream(_ context.Context, _ []chat.Message, _ []tools.Tool) (chat.MessageStream, error) { + return &blockingStream{blocker: p.blocker}, nil +} + +func TestRunSession_GenerateTitleConcurrentClosePanic(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + + // Use fakeRuntime that ends the agent stream immediately by passing a nil release channel + fake := &fakeRuntime{} + sm := newTestSessionManager(t, sess, fake) + + blocker := make(chan struct{}) + mockProvider := &blockingTitleProvider{blocker: blocker} + titleGen := sessiontitle.New(mockProvider) + + // Inject the mock title generator into the active runtime + rt, _ := sm.runtimeSessions.Load(sess.ID) + rt.titleGen = titleGen + + // RunSession spawns generateTitle and fakeRuntime in parallel + streamChan, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "trigger title generation"}}, "") + require.NoError(t, err) + + go func() { + // Drain the stream. It will close once the title generator unblocks and wg.Wait() returns. + for range streamChan { + } + }() + + // Ensure the agent stream draining finishes, putting RunSession's defer stack + // in a position to execute if it weren't blocking on wg.Wait() + time.Sleep(50 * time.Millisecond) //nolint:forbidigo // giving the inner goroutine time to run past the fake stream and reach wg.Wait + + // Unblock the title generator. Without wg.Wait(), the streamChan would already + // be closed and this would panic on send in generateTitle. + close(blocker) + + // Test passes if it does not panic. Delete the session to clean up. + err = sm.DeleteSession(ctx, sess.ID) + require.NoError(t, err) +}