diff --git a/core/services/studio/session/chat.go b/core/services/studio/session/chat.go index 26bc9f3277..760e044f78 100644 --- a/core/services/studio/session/chat.go +++ b/core/services/studio/session/chat.go @@ -729,7 +729,10 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { // it here when the turn has ended meanwhile. return c.queueLogged(ctx, msg, withDirective(msg.Directive, c.modelText(entry))) } - // Reserve the turn, then store the files without the lock. + // Reserve the turn, then store the files without the lock. Messages that a + // stopped turn left queued go first. + carried := c.pending + c.pending = nil c.running, c.stopped = true, false c.done = make(chan struct{}) done := c.done @@ -739,6 +742,7 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { if err != nil { c.mu.Lock() c.endLoopLocked(done) + c.pending = append(carried, c.pending...) c.mu.Unlock() return err } @@ -748,7 +752,7 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error { // The transcript keeps the words and the files apart; the model gets // them together, as it always has. entry.Attachments = append(slices.Clip(entry.Attachments), saved...) - full := withDirective(msg.Directive, c.modelText(entry)) + full := strings.Join(append(carried, withDirective(msg.Directive, c.modelText(entry))), "\n\n") if !logged { entry.Turn = c.turnNow() + 1 c.appendEntry(ctx, entry) @@ -828,7 +832,7 @@ func (c *Chat) turnNow() int { // endLoopLocked ends the running state. The caller holds c.mu. func (c *Chat) endLoopLocked(done chan struct{}) { - c.running, c.cancel, c.pending = false, nil, nil + c.running, c.cancel = false, nil close(done) } @@ -992,14 +996,15 @@ func (c *Chat) checkpoint(n int, text string) ChangeEvent { // Stop interrupts the running turn and waits until it has ended. The turn // still checkpoints, so partial work is kept and can be undone. Messages that -// are queued for the next turn are dropped. +// are queued for the next turn are in the transcript, so they stay queued and +// go to the model with the next message. func (c *Chat) Stop(ctx context.Context) error { c.mu.Lock() if !c.running { c.mu.Unlock() return nil } - c.stopped, c.pending = true, nil + c.stopped = true cancel, done, hs := c.cancel, c.done, c.hs c.mu.Unlock() if cancel != nil { diff --git a/core/services/studio/session/queue_test.go b/core/services/studio/session/queue_test.go index c9a98fa2e8..665d86f6d1 100644 --- a/core/services/studio/session/queue_test.go +++ b/core/services/studio/session/queue_test.go @@ -20,11 +20,11 @@ import ( // 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 + h *sessiontest.FakeHarness + store session.Store + chat *session.Chat + block chan struct{} + release func() ) userTexts := func() []string { @@ -39,11 +39,13 @@ var _ = Describe("messages that wait for the next turn", func() { return out } - release := func() { once.Do(func() { close(block) }) } - BeforeEach(func() { - once = sync.Once{} - block = make(chan struct{}) + // Each spec has its own channel and release: a timer of the previous + // spec may still fire. + blocked := make(chan struct{}) + var once sync.Once + block = blocked + release = func() { once.Do(func() { close(blocked) }) } h = &sessiontest.FakeHarness{InjectErr: session.ErrNoTurn} h.Turns = [][]sessiontest.Step{{{Hold: block}}} store = sessiontest.NewMemoryStore() @@ -97,4 +99,13 @@ var _ = Describe("messages that wait for the next turn", func() { Expect(chat.Send(context.Background(), session.Message{Text: big})).To(Succeed()) Expect(chat.Send(context.Background(), session.Message{Text: big})).To(MatchError(session.ErrTurnRunning)) }) + + It("keeps the waiting messages when the turn is stopped and sends them with the next one", func() { + Expect(chat.Send(context.Background(), session.Message{Text: "queued"})).To(Succeed()) + time.AfterFunc(50*time.Millisecond, release) + Expect(chat.Stop(context.Background())).To(Succeed()) + Expect(chat.Send(context.Background(), session.Message{Text: "next"})).To(Succeed()) + Eventually(h.Sent).Should(HaveLen(2)) + Expect(h.Sent()[1]).To(Equal("queued\n\nnext"), "a message in the transcript never reached the model") + }) })