feat: wake reconciled sessions for queued dispatch

internal/server/session/agent_server.go

@@ -139,10 +139,15 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
 	queue := NewWriterQueue(16, 64, 8)
 	defer queue.Close()
 	capacity := &CapacityShadow{}
+	var capacityMu sync.Mutex
+	reservations := make(map[string]DispatchLane)
+	reconciled := make(chan struct{})
+	var reconcileOnce sync.Once
 	if !capacity.UpdateAdvertised(0, 0, hello.GetMaxRunningCommands(), hello.GetMaxQueuedCommands()) {
 		server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
 		return
 	}
+	go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), reconciled, &capacityMu, capacity, reservations, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
 	encodedSessionID := encodeSessionID(sessionID)
 	welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
 		SessionId: encodedSessionID, SessionGeneration: registration.Generation,
@@ -210,25 +215,19 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
 			case <-sessionContext.Done():
 				return
 			}
-			lane := capacity.Reserve()
-			dispatched, dispatchErr := false, error(nil)
-			if lane != DispatchNone {
-				dispatched, dispatchErr = server.enqueueNextDispatch(sessionContext, queue, hello.GetClientId(), hello.GetPlatform(), encodedSessionID, registration.Generation)
-				if !dispatched {
-					capacity.Release(lane)
-				}
-			}
-			if dispatchErr != nil && !errors.Is(dispatchErr, ErrDispatchDataFull) {
-				server.close(connection, websocket.StatusInternalError, "could not queue command dispatch")
-				return
-			}
+			reconcileOnce.Do(func() { close(reconciled) })
+			handle.SignalDispatch()
 			continue
 		}
 		if advertised := envelope.GetClientCapacity(); advertised != nil {
-			if !capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands()) {
+			capacityMu.Lock()
+			validCapacity := capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands())
+			capacityMu.Unlock()
+			if !validCapacity {
 				server.close(connection, websocket.StatusPolicyViolation, "invalid client capacity")
 				return
 			}
+			handle.SignalDispatch()
 			continue
 		}
 		if acknowledgement := envelope.GetCommandAccepted(); acknowledgement != nil {
@@ -241,6 +240,13 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
 				server.close(connection, websocket.StatusPolicyViolation, "invalid command acceptance")
 				return
 			}
+			capacityMu.Lock()
+			if lane, reserved := reservations[issue.String()]; reserved {
+				capacity.Release(lane)
+				delete(reservations, issue.String())
+			}
+			capacityMu.Unlock()
+			handle.SignalDispatch()
 			continue
 		}
 		if event := envelope.GetCommandEvent(); event != nil {
@@ -315,10 +321,10 @@ func eventType(event *rvboxv1.CommandEvent) uint16 {
 // enqueueNextDispatch records the queued-to-dispatched transition before
 // exposing work to the network. A full data lane is a pre-write failure, so
 // only the owning generation can put the command back into the queue.
-func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64) (bool, error) {
+func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, beforeEnqueue func(domain.UUID)) (domain.UUID, bool, error) {
 	candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
 	if err != nil || candidate == nil {
-		return false, err
+		return domain.UUID{}, false, err
 	}
 	requeue := func(cause error) (bool, error) {
 		_, rollbackErr := server.Store.RequeueDispatch(ctx, candidate.IssueUUID, clientID, generation)
@@ -329,10 +335,12 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
 	}
 	spec := &rvboxv1.ExecutionSpec{}
 	if err := proto.Unmarshal(candidate.ExecutionSpec, spec); err != nil {
-		return requeue(fmt.Errorf("decode persisted execution spec: %w", err))
+		sent, requeueErr := requeue(fmt.Errorf("decode persisted execution spec: %w", err))
+		return candidate.IssueUUID, sent, requeueErr
 	}
 	if err := agentproto.ValidateExecutionSpec(spec, server.limits(), platform); err != nil {
-		return requeue(fmt.Errorf("validate persisted execution spec: %w", err))
+		sent, requeueErr := requeue(fmt.Errorf("validate persisted execution spec: %w", err))
+		return candidate.IssueUUID, sent, requeueErr
 	}
 	dispatch := &rvboxv1.CommandDispatch{
 		IssueUuid: candidate.IssueUUID.String(), CommandRevision: candidate.Revision, TargetSessionGeneration: generation,
@@ -345,12 +353,70 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
 		SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_CommandDispatch{CommandDispatch: dispatch},
 	})
 	if err != nil {
-		return requeue(err)
+		sent, requeueErr := requeue(err)
+		return candidate.IssueUUID, sent, requeueErr
+	}
+	if beforeEnqueue != nil {
+		beforeEnqueue(candidate.IssueUUID)
 	}
 	if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
-		return requeue(ErrDispatchDataFull)
+		sent, requeueErr := requeue(ErrDispatchDataFull)
+		return candidate.IssueUUID, sent, requeueErr
+	}
+	return candidate.IssueUUID, true, nil
+}
+
+// dispatchLoop is the per-session serialized dispatcher. It waits for a
+// complete reconciliation result before consuming queued work, then coalesces
+// wakeups from local control RPCs, capacity advertisements, and acceptances.
+func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue, wake <-chan struct{}, reconciled <-chan struct{}, capacityMu *sync.Mutex, capacity *CapacityShadow, reservations map[string]DispatchLane, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
+	select {
+	case <-reconciled:
+	case <-ctx.Done():
+		return
+	}
+	for {
+		for {
+			capacityMu.Lock()
+			lane := capacity.Reserve()
+			capacityMu.Unlock()
+			if lane == DispatchNone {
+				break
+			}
+			inserted := false
+			issue, sent, err := server.enqueueNextDispatch(ctx, queue, clientID, platform, sessionID, generation, func(issue domain.UUID) {
+				capacityMu.Lock()
+				reservations[issue.String()] = lane
+				inserted = true
+				capacityMu.Unlock()
+			})
+			if !sent {
+				capacityMu.Lock()
+				if inserted {
+					delete(reservations, issue.String())
+				}
+				capacity.Release(lane)
+				capacityMu.Unlock()
+				if err != nil && !errors.Is(err, ErrDispatchDataFull) {
+					server.closeForDispatchFailure(cancel)
+					return
+				}
+				break
+			}
+		}
+		select {
+		case <-wake:
+			continue
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+func (server *AgentServer) closeForDispatchFailure(cancel context.CancelFunc) {
+	if cancel != nil {
+		cancel()
 	}
-	return true, nil
 }
 
 func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.Conn, queue *WriterQueue, heartbeat *synchronizedHeartbeat, started time.Time) {