mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 04:23:15 -05:00
fix(studio): cap the messages that wait for the next turn
The queue of messages that wait while a turn ends or the harness swaps had no limit. Cap it at 32 messages and 256 KiB, and refuse a message over the cap with ErrTurnRunning before it reaches the transcript. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
3 files changed
+54
-4
No files matched your search
@@ -689,6 +689,12 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error {
|
||||
c.mu.Unlock()
|
||||
return fmt.Errorf("%w: the turn is stopping", ErrTurnRunning)
|
||||
}
|
||||
if c.pendingFullLocked(len(msg.Text) + len(msg.Note)) {
|
||||
// Refused before the message reaches the transcript, so a retry
|
||||
// does not log it twice.
|
||||
c.mu.Unlock()
|
||||
return fmt.Errorf("%w: too many messages are waiting for the turn", ErrTurnRunning)
|
||||
}
|
||||
hs, turn, swapping := c.hs, c.turn, c.swapping
|
||||
if hs == nil || swapping {
|
||||
// The harness is being swapped: queue the message for the next
|
||||
@@ -752,18 +758,39 @@ func (c *Chat) send(ctx context.Context, msg Message, logged bool) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Limits of the queue of messages that wait for the next turn. The queue only
|
||||
// fills while a turn ends or the harness swaps, so a full one means a client
|
||||
// that sends faster than the chat can take.
|
||||
const (
|
||||
maxPendingMessages = 32
|
||||
maxPendingBytes = 256 << 10
|
||||
)
|
||||
|
||||
// pendingFullLocked says whether a message of n more bytes does not fit in the
|
||||
// queue. The caller holds c.mu.
|
||||
func (c *Chat) pendingFullLocked(n int) bool {
|
||||
if len(c.pending) >= maxPendingMessages {
|
||||
return true
|
||||
}
|
||||
total := n
|
||||
for _, p := range c.pending {
|
||||
total += len(p)
|
||||
}
|
||||
return total > maxPendingBytes
|
||||
}
|
||||
|
||||
// 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.
|
||||
// When the turn was stopped, or the queue filled up since Send checked it, the
|
||||
// message waits for the turn to end and starts the next one: it was accepted
|
||||
// into 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 {
|
||||
c.mu.Unlock()
|
||||
return c.send(ctx, msg, true)
|
||||
}
|
||||
if c.stopped {
|
||||
if c.stopped || c.pendingFullLocked(len(text)) {
|
||||
done := c.done
|
||||
c.mu.Unlock()
|
||||
select {
|
||||
|
||||
@@ -7,3 +7,9 @@ func (c *Chat) HarnessStateLen() int {
|
||||
defer c.mu.Unlock()
|
||||
return len(c.rec.HarnessState)
|
||||
}
|
||||
|
||||
// The limits of the queue of waiting messages, for the specs.
|
||||
const (
|
||||
MaxPendingMessages = maxPendingMessages
|
||||
MaxPendingBytes = maxPendingBytes
|
||||
)
|
||||
@@ -2,6 +2,8 @@ package session_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -80,4 +82,19 @@ var _ = Describe("messages that wait for the next turn", func() {
|
||||
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")
|
||||
})
|
||||
|
||||
It("refuses a message when too many wait, before it reaches the transcript", func() {
|
||||
var refused error
|
||||
for i := 0; i < 2*session.MaxPendingMessages && refused == nil; i++ {
|
||||
refused = chat.Send(context.Background(), session.Message{Text: fmt.Sprintf("m%d", i)})
|
||||
}
|
||||
Expect(refused).To(MatchError(session.ErrTurnRunning))
|
||||
Expect(userTexts()).To(HaveLen(1+session.MaxPendingMessages), "the refused message was logged")
|
||||
})
|
||||
|
||||
It("refuses a message when the waiting ones are too big", func() {
|
||||
big := strings.Repeat("x", session.MaxPendingBytes/2+1)
|
||||
Expect(chat.Send(context.Background(), session.Message{Text: big})).To(Succeed())
|
||||
Expect(chat.Send(context.Background(), session.Message{Text: big})).To(MatchError(session.ErrTurnRunning))
|
||||
})
|
||||
})
|
||||
Reference in new issue
Block a user