Merge pull request #1420 from rumpl/title-gen

Fix races

Author: rumpl

pkg/runtime/runtime.go

@@ -703,7 +703,8 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 		r.registerDefaultTools()
 
 		if sess.Title == "" {
-			r.titleGen.Generate(ctx, sess, events)
+			userMessage := sess.GetLastUserMessageContent()
+			r.titleGen.Generate(ctx, sess, userMessage, events)
 		}
 
 		iteration := 0

pkg/runtime/runtime_test.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"reflect"
+	"sync"
 	"testing"
 	"time"
 
@@ -568,12 +569,15 @@ func TestToolCallVariations(t *testing.T) {
 // queueProvider returns a different stream on each CreateChatCompletionStream call.
 type queueProvider struct {
 	id      string
+	mu      sync.Mutex
 	streams []chat.MessageStream
 }
 
 func (p *queueProvider) ID() string { return p.id }
 
 func (p *queueProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
 	if len(p.streams) == 0 {
 		return &mockStream{}, nil
 	}

pkg/runtime/title_generator.go

@@ -29,24 +29,22 @@ func newTitleGenerator(model provider.Provider) *titleGenerator {
 	}
 }
 
-func (t *titleGenerator) Generate(ctx context.Context, sess *session.Session, events chan<- Event) {
+func (t *titleGenerator) Generate(ctx context.Context, sess *session.Session, userMessage string, events chan<- Event) {
+	if userMessage == "" {
+		return
+	}
 	t.wg.Go(func() {
-		t.generate(ctx, sess, events)
+		t.generate(ctx, sess, userMessage, events)
 	})
 }
 
 func (t *titleGenerator) Wait() {
 	t.wg.Wait()
 }
 
-func (t *titleGenerator) generate(ctx context.Context, sess *session.Session, events chan<- Event) {
+func (t *titleGenerator) generate(ctx context.Context, sess *session.Session, firstUserMessage string, events chan<- Event) {
 	slog.Debug("Generating title for session", "session_id", sess.ID)
 
-	firstUserMessage := sess.GetLastUserMessageContent()
-	if firstUserMessage == "" {
-		return
-	}
-
 	userPrompt := fmt.Sprintf(titleUserPromptFormat, firstUserMessage)
 
 	titleModel := provider.CloneWithOptions(