From ad5df6dee0b400b6780f66b3ba648490e8a090cc Mon Sep 17 00:00:00 2001 From: david <58231312+davseby@users.noreply.github.com> Date: Tue, 25 Aug 2026 18:32:06 +0300 Subject: [PATCH] claim assistant turns to avoid race conditions --- server/core/internal/assistant/manager.go | 41 ++++- .../core/internal/assistant/manager_test.go | 33 +++- server/core/internal/assistant/session.go | 74 ++++---- .../core/internal/assistant/session_test.go | 169 ++++++++++++++---- server/core/internal/assistant/tools/input.go | 24 ++- .../internal/assistant/tools/input_test.go | 8 +- 6 files changed, 264 insertions(+), 85 deletions(-) diff --git a/server/core/internal/assistant/manager.go b/server/core/internal/assistant/manager.go index 264f9b2c..42a79607 100644 --- a/server/core/internal/assistant/manager.go +++ b/server/core/internal/assistant/manager.go @@ -16,6 +16,7 @@ import ( "io" "log/slog" "net/http" + "sync" "time" "github.com/cloudwego/eino/components/model" @@ -56,6 +57,16 @@ type Manager struct { pendings *persist.Pendings offload *persist.Offload + // turns tracks the conversation keys with a turn in flight. The + // claim lives here rather than on the session because every + // connection of the same (org, user) shares one history, pending + // record and checkpoint: two turns running at once — even from + // different connections — would overwrite each other's state. + turns struct { + mu sync.Mutex + m map[string]struct{} + } + metrics *metrics } @@ -93,7 +104,7 @@ func NewManager( // belonging to the same conversation, expiring with it. blobs := redkit.NewBytesStore(pool, _sessionExpiration) - return &Manager{ + m := &Manager{ log: log, db: db, search: searcher, @@ -116,6 +127,10 @@ func NewManager( metrics: newMetrics(fc, providerName), } + + m.turns.m = make(map[string]struct{}) + + return m } // Configured reports whether the assistant is configured on this @@ -152,6 +167,30 @@ func (m *Manager) ToolSet(orgID, userID string) *tools.Set { )) } +// claimTurn claims the conversation key for one turn, reporting whether +// the claim succeeded. It fails while any session of the same key — the +// same (org, user) on any connection — is already running one. +func (m *Manager) claimTurn(key string) bool { + m.turns.mu.Lock() + defer m.turns.mu.Unlock() + + if _, ok := m.turns.m[key]; ok { + return false + } + + m.turns.m[key] = struct{}{} + + return true +} + +// releaseTurn releases the conversation key a turn claimed. +func (m *Manager) releaseTurn(key string) { + m.turns.mu.Lock() + defer m.turns.mu.Unlock() + + delete(m.turns.m, key) +} + // Chat runs an assistant chat over the given connection for the given // organisation and user. It owns the read loop: inbound messages are // dispatched to a fresh session until the connection reports io.EOF (a diff --git a/server/core/internal/assistant/manager_test.go b/server/core/internal/assistant/manager_test.go index ded7f4fb..d25a1446 100644 --- a/server/core/internal/assistant/manager_test.go +++ b/server/core/internal/assistant/manager_test.go @@ -180,7 +180,7 @@ func testManagerStores() (*Manager, *stores) { st := stubStores() log := discardLog() - return &Manager{ + m := &Manager{ log: log, db: &toolsMock.DB{}, search: &toolsMock.Searcher{ @@ -195,7 +195,11 @@ func testManagerStores() (*Manager, *stores) { metrics: newMetrics(metricutil.NewFactory("test", prometheus.NewRegistry()), "claude"), model: stubChatModel(), summary: stubChatModel(), - }, st + } + + m.turns.m = make(map[string]struct{}) + + return m, st } func Test_Manager_Configured(t *testing.T) { @@ -218,6 +222,7 @@ func Test_NewManager(t *testing.T) { assert.NotNil(t, m.pendings) assert.NotNil(t, m.offload) assert.NotNil(t, m.metrics) + assert.NotNil(t, m.turns.m) assert.Nil(t, m.tree) // an unset summarization model falls back to the chat model, so @@ -250,6 +255,30 @@ func Test_Manager_ToolSet(t *testing.T) { assert.Len(t, s.Tools(), 28) } +func Test_Manager_claimTurn(t *testing.T) { + t.Parallel() + + m := testManager() + + require.True(t, m.claimTurn("k1")) + + // held until released, and per key: another conversation's turns + // are not serialized with this one's. + assert.False(t, m.claimTurn("k1")) + assert.True(t, m.claimTurn("k2")) +} + +func Test_Manager_releaseTurn(t *testing.T) { + t.Parallel() + + m := testManager() + + require.True(t, m.claimTurn("k1")) + m.releaseTurn("k1") + + assert.True(t, m.claimTurn("k1")) +} + func Test_Manager_Chat(t *testing.T) { t.Parallel() diff --git a/server/core/internal/assistant/session.go b/server/core/internal/assistant/session.go index aac7e269..f7cb5064 100644 --- a/server/core/internal/assistant/session.go +++ b/server/core/internal/assistant/session.go @@ -30,8 +30,7 @@ type session struct { supv *xync.Supervisor - mu sync.Mutex - processing bool + mu sync.Mutex // activeDocumentID is the document the user is currently viewing, // surfaced into the system prompt so the model can resolve "this @@ -230,8 +229,13 @@ func (s *session) handleMessage(ctx context.Context, content string) { tn := s.newTurn() s.goTurn(ctx, func(turnCtx context.Context) *persist.PendingConfirm { + // the turn builds on the stored conversation rather than this + // session's copy: another connection of the same user may have + // completed turns since this session loaded it, and a turn built + // on the stale copy would save a history discarding theirs. + msgs := append(s.man.history.Load(turnCtx, s.key), schema.UserMessage(content)) + s.mu.Lock() - msgs := append(cloneMessages(s.messages), schema.UserMessage(content)) s.messages = msgs s.mu.Unlock() @@ -242,15 +246,18 @@ func (s *session) handleMessage(ctx context.Context, content string) { } // handleReset clears the conversation, provided a turn is not running. +// It claims the conversation the way a turn does, so it cannot clear +// the stores while another connection's turn is mid-run. func (s *session) handleReset(ctx context.Context) { - s.mu.Lock() - if s.processing { - s.mu.Unlock() + if !s.man.claimTurn(s.key) { s.writer.WriteJSON(ctx, protocol.NewErrorMessage("cannot reset while processing")) return } + defer s.man.releaseTurn(s.key) + + s.mu.Lock() s.messages = nil s.mu.Unlock() @@ -263,6 +270,9 @@ func (s *session) handleReset(ctx context.Context) { // mismatched turn id is dropped: only the outstanding confirmation can // be answered, and a stale reply from a previous turn must not fire. func (s *session) handleConfirmResponse(ctx context.Context, turnID string, approved, all bool) { + // looked up before claiming so a duplicate answer to an already + // consumed confirmation is dropped silently instead of being told + // the session is busy. pending := s.man.pendings.Load(ctx, s.key) if pending == nil || pending.TurnID != turnID { return @@ -272,6 +282,16 @@ func (s *session) handleConfirmResponse(ctx context.Context, turnID string, appr return } + // re-read under the claim: between the first look and the claim, + // a new message on another connection of the same user may have + // abandoned this confirmation and deleted its checkpoint. + pending = s.man.pendings.Load(ctx, s.key) + if pending == nil || pending.TurnID != turnID { + s.man.releaseTurn(s.key) + + return + } + s.man.pendings.Clear(ctx, s.key) s.man.log.Info("assistant confirm answered", @@ -319,30 +339,28 @@ func (s *session) handleConfirmResponse(ctx context.Context, turnID string, appr }) } -// beginTurn claims the session for a turn, reporting to the client when -// one is already running. +// beginTurn claims the conversation for a turn, reporting to the client +// when one is already running. The claim is held in the manager, keyed +// like the conversation itself: every connection of the same (org, +// user) shares one history and checkpoint, so two turns must not run at +// once even when they arrive on different connections. func (s *session) beginTurn(ctx context.Context) bool { - s.mu.Lock() - defer s.mu.Unlock() - - if s.processing { + if !s.man.claimTurn(s.key) { s.writer.WriteJSON(ctx, protocol.NewErrorMessage("already processing a message")) return false } - s.processing = true - return true } -// goTurn runs one turn on the supervisor, guaranteeing the session is +// goTurn runs one turn on the supervisor, guaranteeing the claim is // released and the client hears about it even if the turn panics. // // A turn that parks on a confirmation returns it rather than sending it -// itself, so the session is already free by the time the client is asked. +// itself, so the claim is already free by the time the client is asked. // Otherwise a fast answer would arrive while the finished turn still held -// the session and be rejected as "already processing". +// the claim and be rejected as "already processing". func (s *session) goTurn(ctx context.Context, fn func(context.Context) *persist.PendingConfirm) { s.supv.Go(func(sctx context.Context) { defer logutil.Recover(s.man.log, nil) @@ -354,19 +372,17 @@ func (s *session) goTurn(ctx context.Context, fn func(context.Context) *persist. // checkpoint is already in Redis, and losing the record // would make the parked turn unreachable on reconnect. // - // It is written before the session is released, because the - // read loop is a separate goroutine: a message arriving on a - // free session looks up the pending record to decide whether - // to abandon a parked turn, and would find nothing while this - // write was still in flight — then delete the checkpoint this - // record points at. + // It is written before the claim is released, because a + // message arriving on any connection of the same user the + // moment the claim frees looks up the pending record to + // decide whether to abandon a parked turn, and would find + // nothing while this write was still in flight — then delete + // the checkpoint this record points at. if pending != nil { s.man.pendings.Save(context.WithoutCancel(ctx), s.key, *pending) } - s.mu.Lock() - s.processing = false - s.mu.Unlock() + s.man.releaseTurn(s.key) if pending != nil { s.writer.WriteJSON(ctx, protocol.NewConfirmRequest(pending.TurnID, pending.Actions)) @@ -415,9 +431,3 @@ func (s *session) runOptions(extra map[string]any) []adk.AgentRunOption { adk.WithSessionValues(values), } } - -// cloneMessages copies the conversation so an in-flight turn cannot be -// mutated by the session underneath it. -func cloneMessages(msgs []*schema.Message) []*schema.Message { - return append(make([]*schema.Message, 0, len(msgs)+1), msgs...) -} diff --git a/server/core/internal/assistant/session_test.go b/server/core/internal/assistant/session_test.go index 2aa0b35d..166a2985 100644 --- a/server/core/internal/assistant/session_test.go +++ b/server/core/internal/assistant/session_test.go @@ -17,6 +17,7 @@ import ( "github.com/oxynote/oxynote/server/core/internal/assistant/tools" toolsMock "github.com/oxynote/oxynote/server/core/internal/assistant/tools/_mock" "github.com/oxynote/oxynote/server/core/internal/document" + "github.com/oxynote/oxynote/server/core/pkg/errutil" "github.com/rs/xid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -207,7 +208,6 @@ func Test_Manager_newSession(t *testing.T) { assert.Same(t, writer, s.writer) assert.NotNil(t, s.runner) assert.NotNil(t, s.supv) - assert.False(t, s.processing) assert.Empty(t, s.activeDocumentID) require.Len(t, s.messages, 1) @@ -353,10 +353,37 @@ func Test_session_handleMessage(t *testing.T) { cc := map[string]struct { Responses []*schema.Message Content string + Stored []*schema.Message Pending bool Settles bool Expect func(t *testing.T, rec *recorder, cm *mock.ChatModel, st *stores) }{ + "Turn builds on the stored conversation": { + Content: "and what about this", + Stored: []*schema.Message{ + schema.UserMessage("from another tab"), + schema.AssistantMessage("noted", nil), + }, + Settles: true, + Expect: func(t *testing.T, _ *recorder, cm *mock.ChatModel, _ *stores) { + // the turn starts from what the store holds, not from this + // session's copy: another connection of the same user may + // have finished turns since this one loaded it, and a turn + // built on the stale copy would save a history discarding + // theirs. + require.NotEmpty(t, cm.StreamCalls()) + + var all strings.Builder + + for _, msg := range cm.StreamCalls()[0].Input { + all.WriteString(msg.Content) + all.WriteString("\n") + } + + assert.Contains(t, all.String(), "from another tab") + assert.NotContains(t, all.String(), "stale local copy") + }, + }, "Outstanding confirmation is abandoned by a new message": { Responses: []*schema.Message{ schema.AssistantMessage("fresh answer", nil), @@ -445,6 +472,16 @@ func Test_session_handleMessage(t *testing.T) { }) } + if c.Stored != nil { + st.history.GetFunc = func(_ context.Context, _ string) (*[]*schema.Message, error) { + return &c.Stored, nil + } + + s.mu.Lock() + s.messages = []*schema.Message{schema.UserMessage("stale local copy")} + s.mu.Unlock() + } + s.handleMessage(context.Background(), c.Content) if c.Settles { @@ -494,9 +531,15 @@ func Test_session_handleReset(t *testing.T) { s.mu.Lock() s.messages = []*schema.Message{schema.UserMessage("hi")} - s.processing = c.Processing s.mu.Unlock() + // a turn in flight holds the conversation's claim; it may be + // running on any connection of the same user, so it is held + // in the manager rather than on this session. + if c.Processing { + require.True(t, m.claimTurn(s.key)) + } + s.handleReset(context.Background()) s.mu.Lock() @@ -521,6 +564,10 @@ func Test_session_handleReset(t *testing.T) { assert.Len(t, st.history.DeleteCalls(), 1) assert.Len(t, st.pendings.DeleteCalls(), 1) assert.Len(t, st.blobs.DeleteCalls(), 1) + + // the reset's own claim is released, so the next turn can + // start. + assert.True(t, m.claimTurn(s.key)) }) } } @@ -538,13 +585,25 @@ func Test_session_handleConfirmResponse(t *testing.T) { } cc := map[string]struct { - Responses []*schema.Message - StaleTurn bool - CorruptCheckpoint bool - Approved bool - All bool - Applied int + Responses []*schema.Message + StaleTurn bool + AbandonedUnderClaim bool + CorruptCheckpoint bool + Approved bool + All bool + Applied int }{ + "Answer abandoned between the look and the claim is dropped": { + // between the answer's first look at the pending record and + // its claim of the conversation, a new message on another + // connection of the same user can abandon the confirmation + // and delete its checkpoint; the re-read under the claim has + // to catch that and drop the answer instead of resuming into + // a checkpoint that no longer exists. + Responses: oneWrite, + AbandonedUnderClaim: true, + Approved: true, + }, "Approving resumes the parked write": { Responses: oneWrite, Approved: true, @@ -591,7 +650,7 @@ func Test_session_handleConfirmResponse(t *testing.T) { applier := stubEditApplier() cm := stubChatModel(c.Responses...) - m := testManager() + m, st := testManagerStores() m.model = cm m.db = stubDocumentDB() m.applier = applier @@ -619,8 +678,34 @@ func Test_session_handleConfirmResponse(t *testing.T) { )) } + if c.AbandonedUnderClaim { + // the first look still sees the record; the re-read under + // the claim finds it gone, as it would after another + // connection's message abandoned the confirmation. + var loads atomic.Int32 + + orig := st.pendings.GetFunc + st.pendings.GetFunc = func(ctx context.Context, key string) (*persist.PendingConfirm, error) { + if loads.Add(1) > 1 { + return nil, errutil.ErrNotFound + } + + return orig(ctx, key) + } + } + s.handleConfirmResponse(context.Background(), turnID, c.Approved, c.All) + if c.AbandonedUnderClaim { + // nothing resumed and nothing was applied — and the claim + // was handed back, so the abandoning turn's successor can + // start. + assert.Empty(t, applier.ApplyCalls()) + assert.True(t, m.claimTurn(s.key)) + + return + } + if c.StaleTurn { assert.Empty(t, applier.ApplyCalls()) @@ -661,7 +746,8 @@ func Test_session_handleConfirmResponse(t *testing.T) { func Test_session_beginTurn(t *testing.T) { t.Parallel() - s, rec := prepSession(t, testManager()) + m := testManager() + s, rec := prepSession(t, m) require.True(t, s.beginTurn(context.Background())) @@ -672,6 +758,21 @@ func Test_session_beginTurn(t *testing.T) { msg, ok := find[protocol.ErrorMessage](rec) require.True(t, ok) assert.Equal(t, "already processing a message", msg.Message) + + // the claim covers every connection of the same (org, user): a + // second session over the same conversation shares the one history + // and checkpoint, so its turn is refused just the same. + other, otherRec := prepSession(t, m) + + assert.False(t, other.beginTurn(context.Background())) + + msg, ok = find[protocol.ErrorMessage](otherRec) + require.True(t, ok) + assert.Equal(t, "already processing a message", msg.Message) + + // once released, either connection may begin the next turn. + m.releaseTurn(s.key) + assert.True(t, other.beginTurn(context.Background())) } func Test_session_goTurn(t *testing.T) { @@ -703,16 +804,18 @@ func Test_session_goTurn(t *testing.T) { s, rec := prepSession(t, m) // a parked turn's record has to be durable before anything - // else can claim the session: the read loop runs on its own - // goroutine, and a message arriving in that window would see - // no pending confirmation and orphan the checkpoint the - // parked turn is waiting on. - var processingAtSave atomic.Bool + // else can claim the conversation: a message arriving in that + // window — on this connection or another of the same user — + // would see no pending confirmation and orphan the checkpoint + // the parked turn is waiting on. + var claimedAtSave atomic.Bool st.pendings.SetFunc = func(context.Context, string, persist.PendingConfirm) error { - s.mu.Lock() - processingAtSave.Store(s.processing) - s.mu.Unlock() + m.turns.mu.Lock() + _, held := m.turns.m[s.key] + m.turns.mu.Unlock() + + claimedAtSave.Store(held) return nil } @@ -753,11 +856,13 @@ func Test_session_goTurn(t *testing.T) { s.supv.Wait() } - // the session is released before the question is asked, so - // an immediate answer is not refused as a second turn. - s.mu.Lock() - assert.False(t, s.processing) - s.mu.Unlock() + // the claim is released before the question is asked, so an + // immediate answer is not refused as a second turn. + m.turns.mu.Lock() + _, held := m.turns.m[s.key] + m.turns.mu.Unlock() + + assert.False(t, held) if c.Pending == nil { assert.Empty(t, rec.types()) @@ -776,8 +881,8 @@ func Test_session_goTurn(t *testing.T) { require.Len(t, ff, 1) assert.NoError(t, ff[0].Ctx.Err()) - assert.True(t, processingAtSave.Load(), - "the confirmation must be stored before the session is released") + assert.True(t, claimedAtSave.Load(), + "the confirmation must be stored before the claim is released") }) } } @@ -822,17 +927,3 @@ func Test_session_resendPendingConfirmation(t *testing.T) { assert.Equal(t, first.TurnID, again.TurnID) assert.Equal(t, first.Actions, again.Actions) } - -func Test_cloneMessages(t *testing.T) { - t.Parallel() - - orig := []*schema.Message{schema.UserMessage("a")} - clone := cloneMessages(orig) - - clone = append(clone, schema.UserMessage("b")) - - // appending to the clone must not reach the original, or an - // in-flight turn could be mutated underneath itself. - assert.Len(t, orig, 1) - assert.Len(t, clone, 2) -} diff --git a/server/core/internal/assistant/tools/input.go b/server/core/internal/assistant/tools/input.go index da39cdb2..4fbf1b7f 100644 --- a/server/core/internal/assistant/tools/input.go +++ b/server/core/internal/assistant/tools/input.go @@ -301,11 +301,19 @@ func (i *input) CreateDocument(doc document.Document) error { return nil } +// errUnknownDocument is what DeleteDocument reports for an id that +// names nothing in the session's organisation. Another organisation's +// id lands here too, which is the point: the tools cannot be used to +// discover that a document exists elsewhere. +var errUnknownDocument = errors.New("no document with that id in this organisation; call list_documents for the ids that exist") + // DeleteDocument removes the document. It records nothing as touched: // the document is gone, so there is nothing left to point a caller at. // The delete reports the ids of the destroyed subtree, and their // search-index removal is queued in the same transaction: after the -// commit nothing else knows what went away. +// commit nothing else knows what went away. An empty subtree means the +// delete matched nothing, which is reported as an error rather than a +// silent success. func (i *input) DeleteDocument(id xid.ID) error { var tx Tx @@ -320,12 +328,14 @@ func (i *input) DeleteDocument(id xid.ID) error { return err } - if len(ids) != 0 { - if err := i.jobs.Enqueue(i.ctx, tx, search.BlocksDifference{ - RemovedDocuments: ids, - }); err != nil { - return fmt.Errorf("insert search job: %w", err) - } + if len(ids) == 0 { + return errUnknownDocument + } + + if err := i.jobs.Enqueue(i.ctx, tx, search.BlocksDifference{ + RemovedDocuments: ids, + }); err != nil { + return fmt.Errorf("insert search job: %w", err) } return tx.Commit() diff --git a/server/core/internal/assistant/tools/input_test.go b/server/core/internal/assistant/tools/input_test.go index 60936529..3f1b444e 100644 --- a/server/core/internal/assistant/tools/input_test.go +++ b/server/core/internal/assistant/tools/input_test.go @@ -467,12 +467,12 @@ func Test_input_DeleteDocument(t *testing.T) { return tcase{DB: stubDB(tx, nil), Tx: tx, Commits: 1, Jobs: 1, Err: assert.AnError} }(), - // a document unknown to the delete reports no subtree, so there - // is no removal to queue. - "Empty subtree queues nothing": func() tcase { + // a delete that matched nothing means the id names no document + // in this organisation — a miss, not a success. + "Empty subtree reports an unknown document": func() tcase { tx := stubTx(nil, nil, nil, nil) - return tcase{DB: stubDB(tx, nil), Tx: tx, Commits: 1} + return tcase{DB: stubDB(tx, nil), Tx: tx, Err: errUnknownDocument} }(), "Delete and search removal land together": func() tcase { tx := stubTx([]xid.ID{docID, xid.New()}, nil, nil, nil)