Assistant/modules/work/application/graph_runtime.go

package application

import (
	"context"
	"errors"
	"time"

	"git.cin.su/m/Assistant/contracts/work"
	"git.cin.su/m/Assistant/modules/work/domain"
	"git.cin.su/m/Assistant/modules/work/ports"
)

var ErrGraphDispatchBinding = errors.New("graph node dispatch binding is missing")

type GraphScheduleRequest struct {
	Graph   work.WorkGraph             `json:"graph"`
	Intents []work.AgentDispatchIntent `json:"intents"`
	Now     int64                      `json:"now_unix_millis"`
}

type GraphCompleteRequest struct {
	GraphID   string           `json:"graph_id"`
	Revision  int64            `json:"revision"`
	NodeID    work.GraphNodeID `json:"node_id"`
	Succeeded bool             `json:"succeeded"`
}

func AdmitAndSchedule(ctx context.Context, store ports.GraphStore, request GraphScheduleRequest) (work.WorkGraph, []work.GraphNodeID, error) {
	if request.Now < 1 {
		request.Now = time.Now().UnixMilli()
	}
	admitted, err := domain.AdmitGraph(request.Graph)
	if err != nil {
		return work.WorkGraph{}, nil, err
	}
	if err := store.SaveWorkGraph(ctx, admitted); err != nil {
		return work.WorkGraph{}, nil, err
	}
	ready, err := domain.ScheduleReadyNodes(admitted, 0)
	if err != nil {
		return work.WorkGraph{}, nil, err
	}
	intentByDispatch := make(map[string]work.AgentDispatchIntent, len(request.Intents))
	for _, intent := range request.Intents {
		intentByDispatch[intent.DispatchID] = intent
	}
	for _, node := range admitted.Nodes {
		if node.DispatchID == "" {
			return work.WorkGraph{}, nil, ErrGraphDispatchBinding
		}
		if _, ok := intentByDispatch[node.DispatchID]; !ok {
			return work.WorkGraph{}, nil, ErrGraphDispatchBinding
		}
	}
	for _, intent := range request.Intents {
		if err := store.SaveDispatchIntent(ctx, intent); err != nil {
			return work.WorkGraph{}, nil, err
		}
	}
	next := admitted
	next.Revision++
	next.State = work.GraphRunning
	selected := make(map[work.GraphNodeID]bool, len(ready))
	for _, nodeID := range ready {
		selected[nodeID] = true
	}
	for i := range next.Nodes {
		if !selected[next.Nodes[i].ID] {
			continue
		}
		if next.Nodes[i].DispatchID == "" {
			return work.WorkGraph{}, nil, ErrGraphDispatchBinding
		}
		intent := intentByDispatch[next.Nodes[i].DispatchID]
		if err := store.EnqueueDispatch(ctx, work.AgentDispatchQueueItem{DispatchID: intent.DispatchID, State: work.AgentDispatchQueued, AvailableAtUnixMillis: request.Now, UpdatedAtUnixMillis: request.Now}); err != nil {
			return work.WorkGraph{}, nil, err
		}
		next.Nodes[i].State = work.GraphNodeLeased
	}
	if err := store.SaveWorkGraph(ctx, next); err != nil {
		return work.WorkGraph{}, nil, err
	}
	return next, ready, nil
}

func CompleteGraphNode(ctx context.Context, store ports.GraphStore, graphID string, revision int64, nodeID work.GraphNodeID, success bool) (work.WorkGraph, error) {
	graph, err := store.LoadWorkGraph(ctx, graphID, revision)
	if err != nil {
		return work.WorkGraph{}, err
	}
	found := false
	for i := range graph.Nodes {
		if graph.Nodes[i].ID == nodeID {
			found = true
			if success {
				graph.Nodes[i].State = work.GraphNodeSuccess
			} else {
				graph.Nodes[i].State = work.GraphNodeFailed
			}
		}
	}
	if !found {
		return work.WorkGraph{}, domain.ErrInvalidGraph
	}
	if !success {
		graph.State = work.GraphFailed
	} else {
		ready, err := domain.ReadyGraphNodes(graph)
		if err != nil {
			return work.WorkGraph{}, err
		}
		for i := range graph.Nodes {
			for _, candidate := range ready {
				if graph.Nodes[i].ID == candidate && graph.Nodes[i].State == work.GraphNodeBlocked {
					graph.Nodes[i].State = work.GraphNodeReady
					if graph.Nodes[i].DispatchID == "" {
						return work.WorkGraph{}, ErrGraphDispatchBinding
					}
					if err := store.EnqueueDispatch(ctx, work.AgentDispatchQueueItem{DispatchID: graph.Nodes[i].DispatchID, State: work.AgentDispatchQueued, AvailableAtUnixMillis: time.Now().UnixMilli(), UpdatedAtUnixMillis: time.Now().UnixMilli()}); err != nil {
						return work.WorkGraph{}, err
					}
				}
			}
		}
		allSuccess := true
		for _, node := range graph.Nodes {
			if node.State != work.GraphNodeSuccess {
				allSuccess = false
				break
			}
		}
		if allSuccess {
			graph.State = work.GraphCompleted
		} else {
			graph.State = work.GraphPartial
		}
	}
	graph.Revision++
	if err := store.SaveWorkGraph(ctx, graph); err != nil {
		return work.WorkGraph{}, err
	}
	return graph, nil
}