Agent-System/internal/app/intentflow/mux.go

package intentflow

import "context"

// OwnerMux routes generate keys to Task and everything else to Query.
type OwnerMux struct {
	Query ActionOwner
	Task  ActionOwner
}

func (m OwnerMux) pick(actionKey string) ActionOwner {
	switch actionKey {
	case "task.run", "task.run.adhoc":
		return m.Task
	default:
		return m.Query
	}
}

func (m OwnerMux) Authorize(ctx context.Context, exec StepExec) error {
	owner := m.pick(exec.ActionKey)
	if owner == nil {
		return ErrPlanActionForbidden
	}
	if a, ok := owner.(StepAuthorizer); ok {
		return a.Authorize(ctx, exec)
	}
	return nil
}

func (m OwnerMux) Execute(ctx context.Context, exec StepExec) (Receipt, error) {
	if err := m.Authorize(ctx, exec); err != nil {
		return Receipt{}, err
	}
	owner := m.pick(exec.ActionKey)
	if owner == nil {
		return Receipt{}, ErrPlanActionForbidden
	}
	return owner.Execute(ctx, exec)
}

func (m OwnerMux) LookupReceipt(ctx context.Context, idempotencyKey string) (Receipt, bool, error) {
	if m.Task != nil {
		if rec, ok, err := m.Task.LookupReceipt(ctx, idempotencyKey); err != nil || ok {
			return rec, ok, err
		}
	}
	if m.Query != nil {
		return m.Query.LookupReceipt(ctx, idempotencyKey)
	}
	return Receipt{}, false, nil
}

func (m OwnerMux) ObserveFence(stepID string, token int64) {
	if obs, ok := m.Query.(FenceObserver); ok {
		obs.ObserveFence(stepID, token)
	}
	if obs, ok := m.Task.(FenceObserver); ok {
		obs.ObserveFence(stepID, token)
	}
}