From 8ff3fb87802cb4c2e2c6d7b31c2d66123b09a269 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 3 Nov 2026 08:22:40 +0000 Subject: [PATCH] fix(studio): do not refuse a message that is already in the transcript When a turn was stopped between Send logging a message and queuing it, Send returned an error although the message was in the transcript, so a client retry logged it twice. The message now waits for the stopped turn to end and starts the next one. The fake harness gets an Inject hook and a step that ignores a stop, to reach that window in a test. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto --- core/services/studio/session/chat.go | 46 +++++----- core/services/studio/session/queue_test.go | 83 +++++++++++++++++++ .../studio/session/sessiontest/harness.go | 12 +++ 3 files changed, 117 insertions(+), 24 deletions(-) create mode 100644 core/services/studio/session/queue_test.go diff --git a/core/services/studio/session/chat.go b/core/services/studio/session/chat.go index 1726bf4e0c..163dc3ddae 100644 --- a/core/services/studio/session/chat.go +++ b/core/services/studio/session/chat.go @@ -701,7 +701,7 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { entry.Turn = turn c.appendEntry(ctx, entry) } - return c.queuePending(ctx, msg, withDirective(msg.Directive, c.modelText(entry))) + return c.queueLogged(ctx, msg, withDirective(msg.Directive, c.modelText(entry))) } // Taken under mu, so a swap that starts later waits for this inject. c.injectMu.RLock() @@ -721,18 +721,7 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { } // The turn is ending. Queue the message for the next one, or start // it here when the turn has ended meanwhile. - c.mu.Lock() - if c.running { - if c.stopped { - c.mu.Unlock() - return fmt.Errorf("%w: the turn is stopping", ErrTurnRunning) - } - c.pending = append(c.pending, withDirective(msg.Directive, c.modelText(entry))) - c.mu.Unlock() - return nil - } - c.mu.Unlock() - return c.send(ctx, msg, true) + return c.queueLogged(ctx, msg, withDirective(msg.Directive, c.modelText(entry))) } // Reserve the turn, then store the files without the lock. c.running, c.stopped = true, false @@ -763,21 +752,30 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { return nil } -// queuePending queues text for the next turn of the running loop, or starts -// it when the turn has ended meanwhile. The message is in the transcript. -func (c *Chat) queuePending(ctx context.Context, msg Message, text string) error { +// queueLogged queues text, which is in the transcript already, for the next +// turn of the running loop, or starts it when the turn has ended meanwhile. +// When the turn was stopped meanwhile, the message waits for the turn to end +// and starts the next one: it is in the transcript, so it must not be refused +// any more. +func (c *Chat) queueLogged(ctx context.Context, msg Message, text string) error { c.mu.Lock() - if c.running { - if c.stopped { - c.mu.Unlock() - return fmt.Errorf("%w: the turn is stopping", ErrTurnRunning) - } - c.pending = append(c.pending, text) + if !c.running { c.mu.Unlock() - return nil + return c.send(ctx, msg, true) } + if c.stopped { + done := c.done + c.mu.Unlock() + select { + case <-done: + case <-ctx.Done(): + return ctx.Err() + } + return c.send(ctx, msg, true) + } + c.pending = append(c.pending, text) c.mu.Unlock() - return c.send(ctx, msg, true) + return nil } // Attach stores files in the chat folder and returns them as attachments of diff --git a/core/services/studio/session/queue_test.go b/core/services/studio/session/queue_test.go new file mode 100644 index 0000000000..36dd6df29b --- /dev/null +++ b/core/services/studio/session/queue_test.go @@ -0,0 +1,83 @@ +package session_test + +import ( + "context" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/studio/session" + "github.com/mudler/LocalAI/core/services/studio/session/sessiontest" + "github.com/mudler/LocalAI/core/services/studio/workspace" +) + +// A harness whose turn is ending refuses Inject, so the message waits in the +// queue of the chat. These specs drive that path. +var _ = Describe("messages that wait for the next turn", func() { + var ( + h *sessiontest.FakeHarness + store session.Store + chat *session.Chat + block chan struct{} + once sync.Once + ) + + userTexts := func() []string { + entries, err := store.Entries(context.Background(), "c1", 0, 1000) + Expect(err).NotTo(HaveOccurred()) + var out []string + for _, e := range entries { + if e.Role == session.RoleUser { + out = append(out, e.Text) + } + } + return out + } + + release := func() { once.Do(func() { close(block) }) } + + BeforeEach(func() { + once = sync.Once{} + block = make(chan struct{}) + h = &sessiontest.FakeHarness{InjectErr: session.ErrNoTurn} + h.Turns = [][]sessiontest.Step{{{Hold: block}}} + store = sessiontest.NewMemoryStore() + m, err := session.New(session.Options{Store: store, Harness: h, StateDir: GinkgoT().TempDir()}) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = m.Close(context.Background()) }) + // The held turn ignores a stop, so a failing spec must let it go. + DeferCleanup(release) + chat, err = openAs(m, context.Background(), workspace.Actor{Principal: "alice", ChatID: "c1"}, "c1") + Expect(err).NotTo(HaveOccurred()) + Expect(chat.Send(context.Background(), session.Message{Text: "first"})).To(Succeed()) + Eventually(func() int { return len(h.Sent()) }).Should(Equal(1)) + }) + + It("does not refuse a message that reached the transcript when the turn is stopped meanwhile", func() { + stopped := make(chan struct{}) + var first atomic.Bool + h.InjectHook = func(string) error { + if first.CompareAndSwap(false, true) { + go func() { + _ = chat.Stop(context.Background()) + close(stopped) + }() + // Wait until Stop has marked the turn. The turn winds down slowly, + // so it still runs when Send looks at it again. + Eventually(func() error { + return chat.Send(context.Background(), session.Message{Text: "probe"}) + }).Should(MatchError(ContainSubstring("stopping"))) + time.AfterFunc(100*time.Millisecond, release) + } + return session.ErrNoTurn + } + err := chat.Send(context.Background(), session.Message{Text: "second"}) + Expect(err).NotTo(HaveOccurred(), "a message that is in the transcript was refused") + Eventually(stopped).Should(BeClosed()) + Expect(userTexts()).To(ContainElement("second"), "the message is not in the transcript") + Eventually(func() []string { return h.Sent() }).Should(HaveLen(2), "the message never reached a turn") + }) +}) diff --git a/core/services/studio/session/sessiontest/harness.go b/core/services/studio/session/sessiontest/harness.go index 4e7daaecd0..0afbffc25d 100644 --- a/core/services/studio/session/sessiontest/harness.go +++ b/core/services/studio/session/sessiontest/harness.go @@ -18,6 +18,7 @@ type Step struct { Tool string // call the session tool with this name ... Args string // ... with these JSON arguments Block <-chan struct{} // wait until the channel closes or the turn ends + Hold <-chan struct{} // wait until the channel closes, even when the turn is stopped Err error // end the turn with this error Panic string // panic with this value } @@ -33,6 +34,10 @@ type FakeHarness struct { // InjectErr, when set, is what every Inject returns, even in a turn. It // stands for a harness whose turn is ending. InjectErr error + // InjectHook, when set, runs at the start of every Inject and may return + // an error for it, as InjectErr does. It lets a test act in the window + // between a message reaching the transcript and its delivery. + InjectHook func(text string) error // StartDelay makes Start wait this long, as a harness that starts a // process does. StartDelay time.Duration @@ -179,6 +184,8 @@ func (s *fakeSession) Send(ctx context.Context, text string) error { case <-st.Block: case <-ctx.Done(): } + case st.Hold != nil: + <-st.Hold case st.Err != nil: return st.Err case st.Panic != "": @@ -189,6 +196,11 @@ func (s *fakeSession) Send(ctx context.Context, text string) error { } func (s *fakeSession) Inject(text string) error { + if s.h.InjectHook != nil { + if err := s.h.InjectHook(text); err != nil { + return err + } + } if s.h.InjectErr != nil { return s.h.InjectErr }