claim assistant turns to avoid race conditions

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

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()
 

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...)
-}

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)
-}

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()

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)