toki
b4124f0bd6
- Add CLI core setup for edge and node services - Refactor edge transport layer (server, integration tests) - Refactor node transport layer (parser, session, heartbeat, client) - Add main_test.go files for edge and node commands - Add input package for edge service - Add go.work and go.work.sum for workspace support - Update configs, docs, and project rules
461 lines
13 KiB
Go
461 lines
13 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
toki "git.toki-labs.com/toki/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()
|
|
}
|
|
|
|
// NodeSnapshot is a surface-neutral view of a registered node, suitable for
|
|
// CLI, HTTP, or other transports. Label is the short display label (node0, node1).
|
|
type NodeSnapshot struct {
|
|
NodeID string
|
|
Alias string
|
|
Label string
|
|
}
|
|
|
|
// ListNodeSnapshots returns surface-neutral node descriptors derived from the
|
|
// current registry contents.
|
|
func (s *Service) ListNodeSnapshots() []NodeSnapshot {
|
|
entries := s.registry.All()
|
|
out := make([]NodeSnapshot, 0, len(entries))
|
|
for _, entry := range entries {
|
|
out = append(out, NodeEntrySnapshot(entry))
|
|
}
|
|
return out
|
|
}
|
|
|
|
// NodeEntrySnapshot converts a registry entry into the surface-neutral DTO.
|
|
func NodeEntrySnapshot(entry *edgenode.NodeEntry) NodeSnapshot {
|
|
if entry == nil {
|
|
return NodeSnapshot{}
|
|
}
|
|
return NodeSnapshot{
|
|
NodeID: entry.NodeID,
|
|
Alias: entry.Alias,
|
|
Label: nodeLabel(entry),
|
|
}
|
|
}
|
|
|
|
func (s *Service) ResolveNode(ref string) (*edgenode.NodeEntry, error) {
|
|
return s.registry.Resolve(ref)
|
|
}
|
|
|
|
// ResolveNodeSnapshot returns a surface-neutral DTO for the resolved node.
|
|
func (s *Service) ResolveNodeSnapshot(ref string) (NodeSnapshot, error) {
|
|
entry, err := s.ResolveNode(ref)
|
|
if err != nil {
|
|
return NodeSnapshot{}, err
|
|
}
|
|
return NodeEntrySnapshot(entry), nil
|
|
}
|
|
|
|
type SubmitRunRequest struct {
|
|
NodeRef string
|
|
RunID string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
Prompt string
|
|
Background bool
|
|
TimeoutSec int
|
|
Metadata map[string]string
|
|
}
|
|
|
|
// RunDispatch describes a dispatched run in surface-neutral terms. It is the
|
|
// metadata any caller (console, HTTP, future RPC) needs after submission.
|
|
type RunDispatch struct {
|
|
RunID string
|
|
NodeID string
|
|
NodeLabel string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
Background bool
|
|
TimeoutSec int
|
|
}
|
|
|
|
// RunStream carries asynchronous events for a dispatched foreground run.
|
|
// Background runs leave both channels nil.
|
|
type RunStream struct {
|
|
Events <-chan *iop.RunEvent
|
|
NodeEvents <-chan *iop.EdgeNodeEvent
|
|
}
|
|
|
|
// RunHandle is the legacy combined surface kept for existing callers; it
|
|
// embeds the surface-neutral DTOs so HTTP/API code can consume RunDispatch
|
|
// or RunStream directly without depending on this struct.
|
|
type RunHandle struct {
|
|
RunDispatch
|
|
RunStream
|
|
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{
|
|
RunDispatch: RunDispatch{
|
|
RunID: runID,
|
|
NodeID: entry.NodeID,
|
|
NodeLabel: nodeLabel(entry),
|
|
Adapter: runReq.GetAdapter(),
|
|
Target: runReq.GetTarget(),
|
|
SessionID: runReq.GetSessionId(),
|
|
Background: runReq.GetBackground(),
|
|
TimeoutSec: int(runReq.GetTimeoutSec()),
|
|
},
|
|
RunStream: RunStream{
|
|
Events: runEvents,
|
|
NodeEvents: nodeEvents,
|
|
},
|
|
close: func() {
|
|
if unregisterRun != nil {
|
|
unregisterRun()
|
|
}
|
|
if unregisterNode != nil {
|
|
unregisterNode()
|
|
}
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type CancelRunRequest struct {
|
|
NodeRef string
|
|
RunID string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
}
|
|
|
|
func BuildCancelRunRequest(req CancelRunRequest) *iop.CancelRequest {
|
|
return &iop.CancelRequest{
|
|
RunId: req.RunID,
|
|
Adapter: req.Adapter,
|
|
Target: req.Target,
|
|
SessionId: NormalizeSessionID(req.SessionID),
|
|
Action: iop.CancelAction_CANCEL_ACTION_CANCEL_RUN,
|
|
}
|
|
}
|
|
|
|
func (s *Service) CancelRun(_ context.Context, req CancelRunRequest) (CommandResult, error) {
|
|
entry, err := s.ResolveNode(req.NodeRef)
|
|
if err != nil {
|
|
return CommandResult{}, err
|
|
}
|
|
cancelReq := BuildCancelRunRequest(req)
|
|
if err := entry.Client.Send(cancelReq); err != nil {
|
|
return CommandResult{}, err
|
|
}
|
|
return CommandResult{
|
|
NodeID: entry.NodeID,
|
|
NodeLabel: nodeLabel(entry),
|
|
SessionID: cancelReq.GetSessionId(),
|
|
}, nil
|
|
}
|
|
|
|
type TerminateSessionRequest struct {
|
|
NodeRef string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
}
|
|
|
|
// CommandResult is the surface-neutral acknowledgement for one-shot node
|
|
// commands (terminate-session, future control RPCs).
|
|
type CommandResult struct {
|
|
NodeID string
|
|
NodeLabel string
|
|
SessionID string
|
|
}
|
|
|
|
// TerminateSessionResult is retained as an alias for backward compatibility.
|
|
type TerminateSessionResult = CommandResult
|
|
|
|
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
|
|
}
|
|
|
|
// UsageStatusView is the surface-neutral DTO for the usage-status command.
|
|
// proto types are retained for the embedded AgentUsageStatus payload.
|
|
type UsageStatusView struct {
|
|
NodeID string
|
|
NodeLabel string
|
|
Target string
|
|
SessionID string
|
|
UsageStatus *iop.AgentUsageStatus
|
|
}
|
|
|
|
// UsageStatusResult is retained as an alias for backward compatibility.
|
|
type UsageStatusResult = UsageStatusView
|
|
|
|
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 {
|
|
return buildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS, "status", adapter, targetName, sessionID, timeoutSec)
|
|
}
|
|
|
|
// BuildNodeCommandRequest builds a NodeCommandRequest for any supported type.
|
|
// idPrefix is used to namespace request_id (e.g. "status", "caps", "sessions").
|
|
func BuildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
|
|
return buildNodeCommandRequest(cmdType, idPrefix, adapter, targetName, sessionID, timeoutSec)
|
|
}
|
|
|
|
func buildNodeCommandRequest(cmdType iop.NodeCommandType, idPrefix, adapter, targetName, sessionID string, timeoutSec int) *iop.NodeCommandRequest {
|
|
if idPrefix == "" {
|
|
idPrefix = "cmd"
|
|
}
|
|
return &iop.NodeCommandRequest{
|
|
RequestId: fmt.Sprintf("%s-%d", idPrefix, time.Now().UnixNano()),
|
|
Type: cmdType,
|
|
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
|
|
}
|
|
|
|
// NodeCommandRequestSpec is the surface-neutral input for ops console node
|
|
// commands (capabilities, session_list, transport_status). UsageStatus keeps
|
|
// its own typed request shape.
|
|
type NodeCommandRequestSpec struct {
|
|
NodeRef string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
TimeoutSec int
|
|
}
|
|
|
|
// NodeCommandView is the surface-neutral result for non-usage-status node
|
|
// commands. Result mirrors the proto map and is empty when the node returned
|
|
// no payload.
|
|
type NodeCommandView struct {
|
|
NodeID string
|
|
NodeLabel string
|
|
Adapter string
|
|
Target string
|
|
SessionID string
|
|
Type iop.NodeCommandType
|
|
Result map[string]string
|
|
}
|
|
|
|
// Capabilities dispatches a CAPABILITIES node command and returns the result map.
|
|
func (s *Service) Capabilities(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
|
|
return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps")
|
|
}
|
|
|
|
// SessionList dispatches a SESSION_LIST node command and returns the result map.
|
|
func (s *Service) SessionList(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
|
|
return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, "sessions")
|
|
}
|
|
|
|
// TransportStatus dispatches a TRANSPORT_STATUS node command and returns the result map.
|
|
func (s *Service) TransportStatus(_ context.Context, req NodeCommandRequestSpec) (NodeCommandView, error) {
|
|
return s.sendNodeCommand(req, iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "transport")
|
|
}
|
|
|
|
func (s *Service) sendNodeCommand(req NodeCommandRequestSpec, cmdType iop.NodeCommandType, idPrefix string) (NodeCommandView, error) {
|
|
entry, err := s.ResolveNode(req.NodeRef)
|
|
if err != nil {
|
|
return NodeCommandView{}, err
|
|
}
|
|
commandReq := buildNodeCommandRequest(cmdType, idPrefix, 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 NodeCommandView{}, fmt.Errorf("transport error: %w", err)
|
|
}
|
|
if resp.GetError() != "" {
|
|
return NodeCommandView{}, fmt.Errorf("node reported error: %s", resp.GetError())
|
|
}
|
|
return NodeCommandView{
|
|
NodeID: entry.NodeID,
|
|
NodeLabel: nodeLabel(entry),
|
|
Adapter: commandReq.GetAdapter(),
|
|
Target: commandReq.GetTarget(),
|
|
SessionID: commandReq.GetSessionId(),
|
|
Type: resp.GetType(),
|
|
Result: resp.GetResult(),
|
|
}, nil
|
|
}
|
|
|
|
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 := make(map[string]string, len(req.Metadata))
|
|
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 {
|
|
return entry.DisplayLabel()
|
|
}
|