go-sip/internal/dispatcher/agent.go

package dispatcher

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"time"

	agentpb "git.ipao.vip/rogee/go-sip/gen/agent"
	"git.ipao.vip/rogee/go-sip/internal/configread"
	"git.ipao.vip/rogee/go-sip/internal/store"
	"google.golang.org/protobuf/proto"
)

var ErrRemoteResultUnknown = errors.New("remote Agent result is unknown; reconciliation required")

const (
	AppliedConfigKindSIP      = "sip"
	AppliedConfigStateApplied = "applied"
)

type AgentStatusProbe interface {
	Probe(ctx context.Context, agentID, cellID string) (*agentpb.AgentStatus, error)
}

type AgentSIPConfigVerifier struct {
	Probe   AgentStatusProbe
	AgentID string
	CellID  string
}

func (v *AgentSIPConfigVerifier) VerifyAppliedSIPConfig(ctx context.Context, snapshot configread.Snapshot) error {
	if v == nil || v.Probe == nil || v.AgentID == "" || v.CellID == "" {
		return errors.New("Agent status probe, agent ID, and configured cell ID are required")
	}
	status, err := v.Probe.Probe(ctx, v.AgentID, v.CellID)
	if err != nil {
		return fmt.Errorf("probe Agent applied SIP config: %w", err)
	}
	return verifyAppliedSIPConfigStatus(status, v.AgentID, v.CellID, snapshot)
}

func verifyAppliedSIPConfigStatus(status *agentpb.AgentStatus, agentID, cellID string, snapshot configread.Snapshot) error {
	if snapshot.SIPRevision <= 0 {
		return errors.New("SIP config revision is missing")
	}
	if status == nil || status.AgentId != agentID || status.CellId != cellID || status.BootId == "" {
		return errors.New("Agent status identity or boot ID is invalid")
	}
	var appliedSIP *agentpb.AppliedConfig
	for _, applied := range status.AppliedConfigs {
		if applied == nil || applied.Kind != AppliedConfigKindSIP {
			continue
		}
		if appliedSIP != nil {
			return errors.New("Agent reported multiple SIP applied-config entries")
		}
		appliedSIP = applied
	}
	if appliedSIP == nil || appliedSIP.State != AppliedConfigStateApplied || appliedSIP.ObservedAtUnixMs <= 0 {
		return errors.New("Agent has not reported one observed, applied SIP config")
	}
	revision := strconv.FormatInt(snapshot.SIPRevision, 10)
	if appliedSIP.Revision != revision {
		return fmt.Errorf("Agent applied SIP config does not match revision %s", revision)
	}
	return nil
}

type AgentSession struct {
	AgentID           string
	CellID            string
	BootID            string
	DispatcherEpoch   string
	SessionGeneration uint64
	ExpiresAtUnixMs   int64
}

type DispatchResult struct {
	Permit   *agentpb.ExecutionPermit
	Receipt  *agentpb.OperationReceipt
	State    agentpb.ExecutionState
	Unknown  bool
	Snapshot *agentpb.ExecutionSnapshot
}

// AgentCoordinator owns session binding and the no-retry-after-unknown rule.
// Quotas and reservations remain in Dispatcher SQLite; this type only turns a
// durable reservation into a fenced Agent permit and execution request.
type AgentCoordinator struct {
	mu       sync.Mutex
	now      func() time.Time
	clients  map[string]agentpb.AgentControlServiceClient
	sessions map[string]AgentSession
}

func NewAgentCoordinator(now func() time.Time) *AgentCoordinator {
	if now == nil {
		now = time.Now
	}
	return &AgentCoordinator{now: now, clients: make(map[string]agentpb.AgentControlServiceClient), sessions: make(map[string]AgentSession)}
}

func (c *AgentCoordinator) Register(agentID string, client agentpb.AgentControlServiceClient) error {
	if agentID == "" || client == nil {
		return errors.New("agent ID and client are required")
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.clients[agentID] = client
	return nil
}

// Probe performs the pre-activation R01 status read. The response is limited
// to deployment status; Activate must still bind the returned boot ID before
// any session-authorized operation is attempted.
func (c *AgentCoordinator) Probe(ctx context.Context, agentID, cellID string) (*agentpb.AgentStatus, error) {
	if agentID == "" || cellID == "" {
		return nil, errors.New("agent ID and cell ID are required")
	}
	client, err := c.client(agentID)
	if err != nil {
		return nil, err
	}
	operationID := "status:" + agentID
	response, err := client.GetAgentStatus(ctx, &agentpb.GetAgentStatusRequest{
		Meta: &agentpb.RequestMeta{
			ProtocolVersion: "agent.v1",
			RequestId:       operationID + ":request",
			TraceId:         operationID,
			OperationId:     operationID,
			AgentId:         agentID,
			CellId:          cellID,
		},
		Target: &agentpb.AgentBinding{AgentId: agentID, CellId: cellID},
	})
	if err != nil {
		return nil, err
	}
	if response == nil || response.Status == nil {
		return nil, errors.New("Agent status response is empty")
	}
	if response.Status.AgentId != agentID || response.Status.CellId != cellID || response.Status.BootId == "" {
		return nil, errors.New("Agent status identity or boot ID is invalid")
	}
	return response.Status, nil
}

func (c *AgentCoordinator) Activate(ctx context.Context, agentID, cellID, bootID, epoch string, generation uint64) (AgentSession, error) {
	client, err := c.client(agentID)
	if err != nil {
		return AgentSession{}, err
	}
	if cellID == "" || bootID == "" || epoch == "" {
		return AgentSession{}, errors.New("cell, boot and dispatcher epoch are required")
	}
	operationID := fmt.Sprintf("activate:%s:%s", agentID, bootID)
	meta := &agentpb.RequestMeta{ProtocolVersion: "agent.v1", RequestId: operationID + ":request", TraceId: operationID, OperationId: operationID, DispatcherEpoch: epoch, AgentId: agentID, CellId: cellID, BootId: bootID}
	response, err := client.ActivateAgent(ctx, &agentpb.ActivateAgentRequest{Meta: meta, Binding: &agentpb.AgentBinding{AgentId: agentID, CellId: cellID, ExpectedBootId: bootID, DispatcherEpoch: epoch, SessionGeneration: generation}, ActivationOperationId: operationID})
	if err != nil {
		return AgentSession{}, err
	}
	if response.Session == nil || response.State != agentpb.ActivationState_ACTIVATION_STATE_ACTIVE {
		return AgentSession{}, errors.New("Agent activation was not active")
	}
	session := AgentSession{AgentID: agentID, CellID: cellID, BootID: bootID, DispatcherEpoch: epoch, SessionGeneration: response.Session.SessionGeneration, ExpiresAtUnixMs: response.Session.ExpiresAtUnixMs}
	c.mu.Lock()
	c.sessions[agentID] = session
	c.mu.Unlock()
	return session, nil
}

func (c *AgentCoordinator) ExecuteRaw(ctx context.Context, agentID string, binding *agentpb.ExecutionBinding, raw []byte, reservationID, configSHA256 string) (DispatchResult, error) {
	if binding == nil || binding.ExecutionId == "" || reservationID == "" {
		return DispatchResult{}, errors.New("execution binding and reservation are required")
	}
	client, session, err := c.clientAndSession(agentID)
	if err != nil {
		return DispatchResult{}, err
	}
	permitRequest := &agentpb.GetExecutionPermitRequest{Meta: c.meta(session, "permit:"+binding.ExecutionId, "permit:"+binding.ExecutionId), Binding: proto.Clone(binding).(*agentpb.ExecutionBinding), ResourceReservationId: reservationID, ExpectedTaskRevision: binding.TaskRevision, ConfigSha256: configSHA256}
	permitResponse, err := client.GetExecutionPermit(ctx, permitRequest)
	if err != nil {
		return c.reconcileUnknown(ctx, client, session, binding, err)
	}
	if permitResponse != nil && permitResponse.Permit == nil && permitResponse.Receipt != nil && permitResponse.Receipt.Result == agentpb.ResultCode_RESULT_CODE_APPLIED {
		// A durable receipt without the permit body is an incomplete read, not a
		// reason to originate. Re-read the same reservation under a new read key.
		permitRequest.Meta = c.meta(session, "permit-reconcile:"+binding.ExecutionId, "permit-reconcile:"+binding.ExecutionId)
		permitResponse, err = client.GetExecutionPermit(ctx, permitRequest)
		if err != nil {
			return c.reconcileUnknown(ctx, client, session, binding, err)
		}
	}
	if permitResponse == nil || permitResponse.Permit == nil || permitResponse.Receipt == nil || permitResponse.Receipt.Result == agentpb.ResultCode_RESULT_CODE_REJECTED || permitResponse.Receipt.Result == agentpb.ResultCode_RESULT_CODE_CONFLICT {
		if permitResponse == nil {
			return DispatchResult{}, fmt.Errorf("Agent did not grant execution permit: empty response")
		}
		return DispatchResult{}, fmt.Errorf("Agent did not grant execution permit: result=%s failure=%v permit=%v", permitResponse.Receipt.GetResult().String(), permitResponse.Receipt.GetFailure(), permitResponse.Permit)
	}
	executeMeta := c.meta(session, "execute:"+binding.ExecutionId, "execute:"+binding.ExecutionId)
	executeResponse, err := client.Execute(ctx, &agentpb.ExecuteRequest{Meta: executeMeta, Binding: proto.Clone(binding).(*agentpb.ExecutionBinding), CallExecuteJson: append([]byte(nil), raw...), ConfigSha256: configSHA256, PermitId: permitResponse.Permit.PermitId})
	if err != nil {
		return c.reconcileUnknown(ctx, client, session, binding, err)
	}
	if executeResponse == nil || executeResponse.Receipt == nil {
		return c.reconcileUnknown(ctx, client, session, binding, ErrRemoteResultUnknown)
	}
	return DispatchResult{Permit: permitResponse.Permit, Receipt: executeResponse.Receipt, State: executeResponse.State}, nil
}

func (c *AgentCoordinator) Control(ctx context.Context, agentID string, binding *agentpb.ExecutionBinding, action agentpb.ControlAction, policy agentpb.ActiveCallPolicy) (*agentpb.ApplyTaskControlResponse, error) {
	client, session, err := c.clientAndSession(agentID)
	if err != nil {
		return nil, err
	}
	meta := c.meta(session, fmt.Sprintf("control:%s:%d", binding.ExecutionId, binding.TaskRevision), fmt.Sprintf("control:%s:%d", binding.ExecutionId, binding.TaskRevision))
	return client.ApplyTaskControl(ctx, &agentpb.ApplyTaskControlRequest{Meta: meta, Binding: proto.Clone(binding).(*agentpb.ExecutionBinding), Action: action, ActiveCallPolicy: policy, ExpectedTaskRevision: binding.TaskRevision})
}

func (c *AgentCoordinator) client(agentID string) (agentpb.AgentControlServiceClient, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	client := c.clients[agentID]
	if client == nil {
		return nil, fmt.Errorf("Agent %q is not registered", agentID)
	}
	return client, nil
}

func (c *AgentCoordinator) clientAndSession(agentID string) (agentpb.AgentControlServiceClient, AgentSession, error) {
	client, err := c.client(agentID)
	if err != nil {
		return nil, AgentSession{}, err
	}
	c.mu.Lock()
	session, ok := c.sessions[agentID]
	c.mu.Unlock()
	if !ok {
		return nil, AgentSession{}, fmt.Errorf("Agent %q is not activated", agentID)
	}
	return client, session, nil
}

// AuthorizeInboundMeta accepts an Agent→Dispatcher fact only from the current
// activated Endpoint session. An old boot's durable fact can be replayed, but
// its request metadata must use the newly activated boot and generation.
func (c *AgentCoordinator) AuthorizeInboundMeta(meta *agentpb.RequestMeta) error {
	if c == nil || meta == nil || meta.ProtocolVersion != "agent.v1" || meta.AgentId == "" {
		return fmt.Errorf("active Agent session is required: %w", store.ErrCommandConflict)
	}
	c.mu.Lock()
	session, active := c.sessions[meta.AgentId]
	client, registered := c.clients[meta.AgentId]
	c.mu.Unlock()
	if !active || !registered || client == nil || session.ExpiresAtUnixMs <= c.now().UnixMilli() ||
		meta.CellId != session.CellID || meta.BootId != session.BootID ||
		meta.DispatcherEpoch != session.DispatcherEpoch || meta.SessionGeneration != session.SessionGeneration {
		return fmt.Errorf("Agent report does not match a current activated session: %w", store.ErrCommandConflict)
	}
	return nil
}

func (c *AgentCoordinator) meta(session AgentSession, operationID, idempotencyKey string) *agentpb.RequestMeta {
	return &agentpb.RequestMeta{ProtocolVersion: "agent.v1", RequestId: operationID + ":request", TraceId: operationID, OperationId: operationID, IdempotencyKey: idempotencyKey, DispatcherEpoch: session.DispatcherEpoch, AgentId: session.AgentID, CellId: session.CellID, BootId: session.BootID, SessionGeneration: session.SessionGeneration}
}

func (c *AgentCoordinator) reconcileUnknown(ctx context.Context, client agentpb.AgentControlServiceClient, session AgentSession, binding *agentpb.ExecutionBinding, cause error) (DispatchResult, error) {
	queryMeta := c.meta(session, "query:"+binding.ExecutionId, "query:"+binding.ExecutionId)
	response, err := client.QueryExecution(ctx, &agentpb.QueryExecutionRequest{Meta: queryMeta, Binding: proto.Clone(binding).(*agentpb.ExecutionBinding)})
	if err == nil && response != nil && response.Snapshot != nil {
		return DispatchResult{Unknown: true, Snapshot: response.Snapshot}, fmt.Errorf("%w: %v", ErrRemoteResultUnknown, cause)
	}
	return DispatchResult{Unknown: true}, fmt.Errorf("%w: %v", ErrRemoteResultUnknown, cause)
}