fix: order command controls after dispatch

This commit is contained in:
2026-09-11 08:56:00 +00:00
parent f6f900e597
commit 684981c235
13 changed files with 360 additions and 38 deletions
+40 -8
View File
@@ -180,7 +180,7 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
return
}
go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), reconciled, &capacityMu, capacity, reservations, stdinSent, signalSent, scriptTransfers, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
go server.dispatchLoop(sessionContext, queue, handle.DispatchWake(), handle.SignalDispatch, reconciled, &capacityMu, capacity, reservations, stdinSent, signalSent, scriptTransfers, hello.GetClientId(), hello.GetPlatform(), encodeSessionID(sessionID), registration.Generation, cancel)
encodedSessionID := encodeSessionID(sessionID)
welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
SessionId: encodedSessionID, SessionGeneration: registration.Generation,
@@ -413,7 +413,7 @@ 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, beforeEnqueue func(*store.DispatchCandidate)) (domain.UUID, bool, error) {
func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *WriterQueue, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, beforeEnqueue func(*store.DispatchCandidate), onWritten func(domain.UUID)) (domain.UUID, bool, error) {
candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
if err != nil || candidate == nil {
return domain.UUID{}, false, err
@@ -451,7 +451,11 @@ func (server *AgentServer) enqueueNextDispatch(ctx context.Context, queue *Write
if beforeEnqueue != nil {
beforeEnqueue(candidate)
}
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded, OnWritten: func() {
if onWritten != nil {
onWritten(candidate.IssueUUID)
}
}}) {
sent, requeueErr := requeue(ErrDispatchDataFull)
return candidate.IssueUUID, sent, requeueErr
}
@@ -531,12 +535,18 @@ func (server *AgentServer) enqueueNextScript(queue *WriterQueue, sessionID strin
// lane. Session-local sent tracking suppresses duplicate frames while a live
// connection remains usable; reconnecting naturally replays unacknowledged
// writes from storage.
func (server *AgentServer) enqueueNextStdin(ctx context.Context, queue *WriterQueue, clientID, sessionID string, generation uint64, sent map[string]struct{}, sentMu *sync.Mutex) (bool, error) {
func (server *AgentServer) enqueueNextStdin(ctx context.Context, queue *WriterQueue, clientID, sessionID string, generation uint64, sent map[string]struct{}, dispatchWritten map[string]bool, sentMu *sync.Mutex) (bool, error) {
intents, err := server.Store.PendingStdin(ctx, clientID)
if err != nil {
return false, err
}
for _, intent := range intents {
sentMu.Lock()
written, waitingForDispatch := dispatchWritten[intent.IssueUUID.String()]
sentMu.Unlock()
if waitingForDispatch && !written {
continue
}
key := stdinIntentKey(intent.IssueUUID, intent.WriteSeq)
sentMu.Lock()
_, alreadySent := sent[key]
@@ -579,12 +589,18 @@ func signalIntentKey(issue domain.UUID, revision uint64, signal rvboxv1.SignalKi
return issue.String() + ":" + fmt.Sprint(revision) + ":" + fmt.Sprint(signal)
}
func (server *AgentServer) enqueueNextSignal(ctx context.Context, queue *WriterQueue, clientID, sessionID string, generation uint64, sent map[string]struct{}, sentMu *sync.Mutex) (bool, error) {
func (server *AgentServer) enqueueNextSignal(ctx context.Context, queue *WriterQueue, clientID, sessionID string, generation uint64, sent map[string]struct{}, dispatchWritten map[string]bool, sentMu *sync.Mutex) (bool, error) {
intents, err := server.Store.PendingSignals(ctx, clientID, generation)
if err != nil {
return false, err
}
for _, intent := range intents {
sentMu.Lock()
written, waitingForDispatch := dispatchWritten[intent.IssueUUID.String()]
sentMu.Unlock()
if waitingForDispatch && !written {
continue
}
key := signalIntentKey(intent.IssueUUID, intent.CommandRevision, intent.Signal)
sentMu.Lock()
_, alreadySent := sent[key]
@@ -617,7 +633,11 @@ func (server *AgentServer) enqueueNextSignal(ctx context.Context, queue *WriterQ
// 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, stdinSent, signalSent map[string]struct{}, scriptTransfers map[string]*scriptTransfer, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue, wake <-chan struct{}, signalWake func(), reconciled <-chan struct{}, capacityMu *sync.Mutex, capacity *CapacityShadow, reservations map[string]DispatchLane, stdinSent, signalSent map[string]struct{}, scriptTransfers map[string]*scriptTransfer, clientID string, platform rvboxv1.Platform, sessionID string, generation uint64, cancel context.CancelFunc) {
// A control frame may use the essential queue and therefore overtake data.
// Keep its issue blocked only until the dispatch frame has actually crossed
// the socket writer; after that, WebSocket ordering preserves the dependency.
dispatchWritten := make(map[string]bool)
select {
case <-reconciled:
case <-ctx.Done():
@@ -625,7 +645,7 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
}
for {
for {
stdinQueued, stdinErr := server.enqueueNextStdin(ctx, queue, clientID, sessionID, generation, stdinSent, capacityMu)
stdinQueued, stdinErr := server.enqueueNextStdin(ctx, queue, clientID, sessionID, generation, stdinSent, dispatchWritten, capacityMu)
if stdinErr != nil {
server.closeForDispatchFailure(cancel)
return
@@ -633,7 +653,7 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
if stdinQueued {
continue
}
signalQueued, signalErr := server.enqueueNextSignal(ctx, queue, clientID, sessionID, generation, signalSent, capacityMu)
signalQueued, signalErr := server.enqueueNextSignal(ctx, queue, clientID, sessionID, generation, signalSent, dispatchWritten, capacityMu)
if signalErr != nil {
server.closeForDispatchFailure(cancel)
return
@@ -659,16 +679,25 @@ func (server *AgentServer) dispatchLoop(ctx context.Context, queue *WriterQueue,
issue, sent, err := server.enqueueNextDispatch(ctx, queue, clientID, platform, sessionID, generation, func(candidate *store.DispatchCandidate) {
capacityMu.Lock()
reservations[candidate.IssueUUID.String()] = lane
dispatchWritten[candidate.IssueUUID.String()] = false
if candidate.ScriptPresent {
scriptTransfers[candidate.IssueUUID.String()] = &scriptTransfer{Body: append([]byte(nil), candidate.ScriptContent...), Digest: sha256.Sum256(candidate.ScriptContent)}
}
inserted = true
capacityMu.Unlock()
}, func(issue domain.UUID) {
capacityMu.Lock()
dispatchWritten[issue.String()] = true
capacityMu.Unlock()
if signalWake != nil {
signalWake()
}
})
if !sent {
capacityMu.Lock()
if inserted {
delete(reservations, issue.String())
delete(dispatchWritten, issue.String())
delete(scriptTransfers, issue.String())
}
capacity.Release(lane)
@@ -710,6 +739,9 @@ func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.
if frame.Written != nil {
close(frame.Written)
}
if frame.OnWritten != nil {
frame.OnWritten()
}
continue
}
if !errors.Is(err, context.DeadlineExceeded) {