feat: dispatch queued commands after reconciliation

internal/server/session/agent_server.go

@@ -28,6 +28,7 @@ var (
 	ErrUnexpectedOrigin  = errors.New("agent connections must not send an Origin header")
 	ErrUnexpectedMessage = errors.New("agent message is not a binary protobuf envelope")
 	ErrStaleSession      = errors.New("agent message is for an unknown or fenced session")
+	ErrDispatchDataFull  = errors.New("session dispatch data lane is full")
 )
 
 // AgentServer is the transport edge for agent WebSocket sessions. Command
@@ -137,6 +138,11 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
 
 	queue := NewWriterQueue(16, 64, 8)
 	defer queue.Close()
+	capacity := &CapacityShadow{}
+	if !capacity.UpdateAdvertised(0, 0, hello.GetMaxRunningCommands(), hello.GetMaxQueuedCommands()) {
+		server.close(connection, websocket.StatusPolicyViolation, "invalid initial client capacity")
+		return
+	}
 	encodedSessionID := encodeSessionID(sessionID)
 	welcome, err := proto.Marshal(&rvboxv1.AgentEnvelope{
 		SessionId: encodedSessionID, SessionGeneration: registration.Generation,
@@ -194,14 +200,80 @@ func (server *AgentServer) serveConnection(parent context.Context, connection *w
 				SessionId: encodedSessionID, SessionGeneration: registration.Generation,
 				Payload: &rvboxv1.AgentEnvelope_ReconcileResult{ReconcileResult: result},
 			})
-			if err != nil || queue.EnqueueControl(Frame{Kind: FrameControl, Payload: encoded}) != nil {
+			resultWritten := make(chan struct{})
+			if err != nil || queue.EnqueueControl(Frame{Kind: FrameControl, Payload: encoded, Written: resultWritten}) != nil {
 				server.close(connection, websocket.StatusInternalError, "could not queue reconciliation result")
 				return
 			}
+			select {
+			case <-resultWritten:
+			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
+			}
+			continue
+		}
+		if advertised := envelope.GetClientCapacity(); advertised != nil {
+			if !capacity.UpdateAdvertised(advertised.GetRunningCommands(), advertised.GetQueuedCommands(), advertised.GetMaxRunningCommands(), advertised.GetMaxQueuedCommands()) {
+				server.close(connection, websocket.StatusPolicyViolation, "invalid client capacity")
+				return
+			}
 		}
 	}
 }
 
+// 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) {
+	candidate, err := server.Store.ClaimNextDispatch(ctx, clientID, generation, server.now())
+	if err != nil || candidate == nil {
+		return false, err
+	}
+	requeue := func(cause error) (bool, error) {
+		_, rollbackErr := server.Store.RequeueDispatch(ctx, candidate.IssueUUID, clientID, generation)
+		if rollbackErr != nil {
+			return false, rollbackErr
+		}
+		return false, cause
+	}
+	spec := &rvboxv1.ExecutionSpec{}
+	if err := proto.Unmarshal(candidate.ExecutionSpec, spec); err != nil {
+		return requeue(fmt.Errorf("decode persisted execution spec: %w", err))
+	}
+	if err := agentproto.ValidateExecutionSpec(spec, server.limits(), platform); err != nil {
+		return requeue(fmt.Errorf("validate persisted execution spec: %w", err))
+	}
+	dispatch := &rvboxv1.CommandDispatch{
+		IssueUuid: candidate.IssueUUID.String(), CommandRevision: candidate.Revision, TargetSessionGeneration: generation,
+		IssueTime: timestamppb.New(candidate.IssueTime), Spec: spec,
+	}
+	if candidate.QueueExpiryTime != nil {
+		dispatch.QueueExpiryTime = timestamppb.New(*candidate.QueueExpiryTime)
+	}
+	encoded, err := proto.Marshal(&rvboxv1.AgentEnvelope{
+		SessionId: sessionID, SessionGeneration: generation, Payload: &rvboxv1.AgentEnvelope_CommandDispatch{CommandDispatch: dispatch},
+	})
+	if err != nil {
+		return requeue(err)
+	}
+	if !queue.EnqueueData(Frame{Kind: FrameData, Payload: encoded}) {
+		return requeue(ErrDispatchDataFull)
+	}
+	return true, nil
+}
+
 func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.Conn, queue *WriterQueue, heartbeat *synchronizedHeartbeat, started time.Time) {
 	for {
 		frameContext, cancel := context.WithTimeout(ctx, server.heartbeatPollInterval())
@@ -214,6 +286,9 @@ func (server *AgentServer) writeLoop(ctx context.Context, connection *websocket.
 			if err != nil {
 				return
 			}
+			if frame.Written != nil {
+				close(frame.Written)
+			}
 			continue
 		}
 		if !errors.Is(err, context.DeadlineExceeded) {