iop/apps/edge/internal/service/service.go

package service

import (
	"context"
	"fmt"
	"time"

	toki "git.toki-labs.com/toki/common-proto-socket/go"
	"google.golang.org/protobuf/types/known/structpb"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	eventpkg "iop/packages/events"
	iop "iop/proto/gen/iop"
)

const (
	DefaultSessionID  = "default"
	DefaultTimeoutSec = 30
)

type Service struct {
	registry *edgenode.Registry
	events   *edgeevents.Bus
}

func New(registry *edgenode.Registry, events *edgeevents.Bus) *Service {
	return &Service{registry: registry, events: events}
}

func (s *Service) ListNodes() []*edgenode.NodeEntry {
	return s.registry.All()
}

func (s *Service) ResolveNode(ref string) (*edgenode.NodeEntry, error) {
	return s.registry.Resolve(ref)
}

type SubmitRunRequest struct {
	NodeRef    string
	RunID      string
	Adapter    string
	Target     string
	SessionID  string
	Prompt     string
	Background bool
	TimeoutSec int
	Metadata   map[string]string
}

type RunHandle struct {
	RunID      string
	NodeID     string
	NodeLabel  string
	Adapter    string
	Target     string
	SessionID  string
	Background bool
	TimeoutSec int
	Events     <-chan *iop.RunEvent
	NodeEvents <-chan *iop.EdgeNodeEvent
	close      func()
}

func (h *RunHandle) Close() {
	if h != nil && h.close != nil {
		h.close()
	}
}

func (h *RunHandle) WaitTimeout() time.Duration {
	if h == nil {
		return time.Duration(DefaultTimeoutSec+5) * time.Second
	}
	return time.Duration(normalizeTimeoutSec(h.TimeoutSec)+5) * time.Second
}

func (s *Service) SubmitRun(_ context.Context, req SubmitRunRequest) (*RunHandle, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return nil, err
	}

	runReq, runID, err := BuildRunRequest(req)
	if err != nil {
		return nil, err
	}

	var runEvents <-chan *iop.RunEvent
	var unregisterRun func()
	var nodeEvents <-chan *iop.EdgeNodeEvent
	var unregisterNode func()
	if !runReq.GetBackground() {
		if s.events == nil {
			return nil, fmt.Errorf("event bus is not configured")
		}
		runEvents, unregisterRun = s.events.SubscribeRun(runID, 4096)
		nodeEvents, unregisterNode = s.events.SubscribeNode(entry.NodeID, 16)
	}

	if err := entry.Client.Send(runReq); err != nil {
		if unregisterRun != nil {
			unregisterRun()
		}
		if unregisterNode != nil {
			unregisterNode()
		}
		return nil, err
	}

	return &RunHandle{
		RunID:      runID,
		NodeID:     entry.NodeID,
		NodeLabel:  nodeLabel(entry),
		Adapter:    runReq.GetAdapter(),
		Target:     runReq.GetTarget(),
		SessionID:  runReq.GetSessionId(),
		Background: runReq.GetBackground(),
		TimeoutSec: int(runReq.GetTimeoutSec()),
		Events:     runEvents,
		NodeEvents: nodeEvents,
		close: func() {
			if unregisterRun != nil {
				unregisterRun()
			}
			if unregisterNode != nil {
				unregisterNode()
			}
		},
	}, nil
}

type TerminateSessionRequest struct {
	NodeRef   string
	Adapter   string
	Target    string
	SessionID string
}

type TerminateSessionResult struct {
	NodeID    string
	NodeLabel string
	SessionID string
}

func (s *Service) TerminateSession(_ context.Context, req TerminateSessionRequest) (TerminateSessionResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return TerminateSessionResult{}, err
	}

	sessionID := NormalizeSessionID(req.SessionID)
	cancelReq := &iop.CancelRequest{
		Adapter:   req.Adapter,
		Target:    req.Target,
		SessionId: sessionID,
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	}
	if err := entry.Client.Send(cancelReq); err != nil {
		return TerminateSessionResult{}, err
	}
	return TerminateSessionResult{
		NodeID:    entry.NodeID,
		NodeLabel: nodeLabel(entry),
		SessionID: sessionID,
	}, nil
}

type UsageStatusRequest struct {
	NodeRef    string
	Adapter    string
	Target     string
	SessionID  string
	TimeoutSec int
}

type UsageStatusResult struct {
	NodeID      string
	NodeLabel   string
	Target      string
	SessionID   string
	UsageStatus *iop.AgentUsageStatus
}

func (s *Service) UsageStatus(_ context.Context, req UsageStatusRequest) (UsageStatusResult, error) {
	entry, err := s.ResolveNode(req.NodeRef)
	if err != nil {
		return UsageStatusResult{}, err
	}

	commandReq := BuildUsageStatusRequest(req.Adapter, req.Target, req.SessionID, req.TimeoutSec)
	resp, err := toki.SendRequestTyped[*iop.NodeCommandRequest, *iop.NodeCommandResponse](
		&entry.Client.Communicator,
		commandReq,
		StatusWaitTimeout(commandReq),
	)
	if err != nil {
		return UsageStatusResult{}, fmt.Errorf("transport error: %w", err)
	}
	if resp.GetError() != "" {
		return UsageStatusResult{}, fmt.Errorf("node reported error: %s", resp.GetError())
	}

	return UsageStatusResult{
		NodeID:      entry.NodeID,
		NodeLabel:   nodeLabel(entry),
		Target:      commandReq.GetTarget(),
		SessionID:   commandReq.GetSessionId(),
		UsageStatus: resp.GetUsageStatus(),
	}, nil
}

func BuildUsageStatusRequest(adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
	reqID := fmt.Sprintf("status-%d", time.Now().UnixNano())
	return &iop.NodeCommandRequest{
		RequestId:  reqID,
		Type:       iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:    adapter,
		Target:     targetName,
		SessionId:  NormalizeSessionID(sessionID),
		TimeoutSec: int32(normalizeTimeoutSec(timeoutSec)),
	}
}

func StatusWaitTimeout(req *iop.NodeCommandRequest) time.Duration {
	return time.Duration(normalizeTimeoutSec(int(req.GetTimeoutSec()))+5) * time.Second
}

func NormalizeSessionID(id string) string {
	if id == "" {
		return DefaultSessionID
	}
	return id
}

func NewRunID() string {
	return fmt.Sprintf("manual-%d", time.Now().UnixNano())
}

func IsNodeDisconnected(event *iop.EdgeNodeEvent) bool {
	return event.GetType() == eventpkg.TypeNodeDisconnected
}

func BuildRunRequest(req SubmitRunRequest) (*iop.RunRequest, string, error) {
	input, err := structpb.NewStruct(map[string]any{"prompt": req.Prompt})
	if err != nil {
		return nil, "", err
	}
	runID := req.RunID
	if runID == "" {
		runID = NewRunID()
	}
	metadata := map[string]string{"source": "edge-console"}
	for k, v := range req.Metadata {
		metadata[k] = v
	}
	return &iop.RunRequest{
		RunId:       runID,
		Adapter:     req.Adapter,
		Target:      req.Target,
		SessionId:   NormalizeSessionID(req.SessionID),
		SessionMode: iop.RunSessionMode_RUN_SESSION_MODE_CREATE_IF_MISSING,
		Background:  req.Background,
		Input:       input,
		TimeoutSec:  int32(normalizeTimeoutSec(req.TimeoutSec)),
		Metadata:    metadata,
	}, runID, nil
}

func normalizeTimeoutSec(timeoutSec int) int {
	if timeoutSec <= 0 {
		return DefaultTimeoutSec
	}
	return timeoutSec
}

func nodeLabel(entry *edgenode.NodeEntry) string {
	if entry.Alias != "" {
		return entry.Alias
	}
	return entry.NodeID
}