diff --git a/apps/edge/internal/node/mapper.go b/apps/edge/internal/node/mapper.go new file mode 100644 index 0000000..9c6a558 --- /dev/null +++ b/apps/edge/internal/node/mapper.go @@ -0,0 +1,86 @@ +package node + +import ( + "fmt" + + "google.golang.org/protobuf/types/known/structpb" + + "iop/packages/config" + iop "iop/proto/gen/iop" +) + +// BuildConfigPayload converts a NodeRecord's adapter config to the proto payload +// sent to node during registration. +func BuildConfigPayload(rec *NodeRecord) (*iop.NodeConfigPayload, error) { + payload := &iop.NodeConfigPayload{ + Runtime: &iop.NodeRuntimeConfig{ + Concurrency: int32(rec.Runtime.Concurrency), + WorkspaceRoot: rec.Runtime.WorkspaceRoot, + }, + } + + mockSettings, err := structpb.NewStruct(nil) + if err != nil { + return nil, fmt.Errorf("buildConfigPayload: mock: %w", err) + } + payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ + Type: "mock", + Enabled: true, + Settings: mockSettings, + }) + + if rec.Adapters.Ollama.Enabled { + payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ + Type: "ollama", + Enabled: true, + Config: &iop.AdapterConfig_Ollama{ + Ollama: &iop.OllamaAdapterConfig{BaseUrl: rec.Adapters.Ollama.BaseURL}, + }, + }) + } + if rec.Adapters.Vllm.Enabled { + payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ + Type: "vllm", + Enabled: true, + Config: &iop.AdapterConfig_Vllm{ + Vllm: &iop.VllmAdapterConfig{Endpoint: rec.Adapters.Vllm.Endpoint}, + }, + }) + } + if rec.Adapters.CLI.Enabled { + profiles := make(map[string]*iop.CLIProfileConfig, len(rec.Adapters.CLI.Profiles)) + for name, p := range rec.Adapters.CLI.Profiles { + profiles[name] = cliProfileToProto(p) + } + payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ + Type: "cli", + Enabled: true, + Config: &iop.AdapterConfig_Cli{ + Cli: &iop.CLIAdapterConfig{Profiles: profiles}, + }, + }) + } + return payload, nil +} + +func cliProfileToProto(p config.CLIProfileConf) *iop.CLIProfileConfig { + out := &iop.CLIProfileConfig{ + Command: p.Command, + Args: append([]string(nil), p.Args...), + Env: append([]string(nil), p.Env...), + Persistent: p.Persistent, + Terminal: p.Terminal, + ResponseIdleTimeoutMs: int32(p.ResponseIdleTimeoutMS), + StartupIdleTimeoutMs: int32(p.StartupIdleTimeoutMS), + OutputFormat: p.OutputFormat, + Mode: p.Mode, + ResumeArgs: append([]string(nil), p.ResumeArgs...), + } + if !p.CompletionMarker.Empty() { + out.CompletionMarker = &iop.CLICompletionMarker{ + Line: p.CompletionMarker.Line, + Regex: p.CompletionMarker.Regex, + } + } + return out +} diff --git a/apps/edge/internal/node/mapper_test.go b/apps/edge/internal/node/mapper_test.go new file mode 100644 index 0000000..7e67aa0 --- /dev/null +++ b/apps/edge/internal/node/mapper_test.go @@ -0,0 +1,168 @@ +package node_test + +import ( + "testing" + + edgenode "iop/apps/edge/internal/node" + "iop/packages/config" +) + +func TestBuildConfigPayload_OllamaEnabled(t *testing.T) { + rec := &edgenode.NodeRecord{ + ID: "node-1", + Alias: "test", + Token: "token", + Adapters: config.AdaptersConf{ + Ollama: config.OllamaConf{ + Enabled: true, + BaseURL: "http://localhost:11434", + }, + }, + Runtime: config.RuntimeConf{ + Concurrency: 4, + WorkspaceRoot: "/tmp/ws", + }, + } + + payload, err := edgenode.BuildConfigPayload(rec) + if err != nil { + t.Fatalf("BuildConfigPayload failed: %v", err) + } + + if payload.Runtime == nil { + t.Fatal("expected runtime config, got nil") + } + if payload.Runtime.Concurrency != 4 { + t.Errorf("expected concurrency 4, got %d", payload.Runtime.Concurrency) + } + + if len(payload.Adapters) < 2 { + t.Fatalf("expected at least 2 adapters, got %d", len(payload.Adapters)) + } + + mockFound := false + ollamaFound := false + for _, a := range payload.Adapters { + if a.Type == "mock" { + mockFound = true + } + if a.Type == "ollama" { + ollamaFound = true + ollamaCfg := a.GetOllama() + if ollamaCfg == nil { + t.Fatal("expected ollama adapter config") + } + if ollamaCfg.BaseUrl != "http://localhost:11434" { + t.Errorf("expected ollama base url %q, got %q", "http://localhost:11434", ollamaCfg.BaseUrl) + } + } + } + if !mockFound { + t.Fatal("expected mock adapter") + } + if !ollamaFound { + t.Fatal("expected ollama adapter when enabled") + } +} + +func TestBuildConfigPayload_CLIProfiles(t *testing.T) { + rec := &edgenode.NodeRecord{ + ID: "node-1", + Alias: "test", + Token: "token", + Adapters: config.AdaptersConf{ + CLI: config.CLIConf{ + Enabled: true, + Profiles: map[string]config.CLIProfileConf{ + "default": { + Command: "echo", + Args: []string{"hello"}, + Env: []string{"FOO=bar"}, + Persistent: false, + Terminal: true, + ResponseIdleTimeoutMS: 5000, + StartupIdleTimeoutMS: 10000, + OutputFormat: "plain", + Mode: "implicit", + ResumeArgs: []string{"--resume"}, + CompletionMarker: config.CompletionMarkerConf{ + Line: "DONE", + Regex: "^DONE$", + }, + }, + }, + }, + }, + Runtime: config.RuntimeConf{ + WorkspaceRoot: "/tmp/ws", + }, + } + + payload, err := edgenode.BuildConfigPayload(rec) + if err != nil { + t.Fatalf("BuildConfigPayload failed: %v", err) + } + + cliFound := false + for _, a := range payload.Adapters { + if a.Type == "cli" { + cliFound = true + cliCfg := a.GetCli() + if cliCfg == nil { + t.Fatal("expected cli adapter config") + } + if len(cliCfg.Profiles) != 1 { + t.Fatalf("expected 1 cli profile, got %d", len(cliCfg.Profiles)) + } + p, ok := cliCfg.Profiles["default"] + if !ok { + t.Fatal("expected default profile") + } + if p.Command != "echo" { + t.Errorf("expected command %q, got %q", "echo", p.Command) + } + if len(p.Args) != 1 || p.Args[0] != "hello" { + t.Errorf("expected args [hello], got %v", p.Args) + } + if p.OutputFormat != "plain" { + t.Errorf("expected format %q, got %q", "plain", p.OutputFormat) + } + if cm := p.CompletionMarker; cm == nil || cm.Line != "DONE" { + t.Errorf("expected completion marker line %q, got %v", "DONE", cm) + } + } + } + if !cliFound { + t.Fatal("expected cli adapter when enabled") + } +} + +func TestBuildConfigPayload_MockAlwaysPresent(t *testing.T) { + rec := &edgenode.NodeRecord{ + ID: "node-1", + Alias: "test", + Token: "token", + Adapters: config.AdaptersConf{ + Ollama: config.OllamaConf{Enabled: false}, + Vllm: config.VllmConf{Enabled: false}, + CLI: config.CLIConf{Enabled: false}, + }, + Runtime: config.RuntimeConf{ + WorkspaceRoot: "/tmp/ws", + }, + } + + payload, err := edgenode.BuildConfigPayload(rec) + if err != nil { + t.Fatalf("BuildConfigPayload failed: %v", err) + } + + if len(payload.Adapters) < 1 { + t.Fatalf("expected at least 1 adapter, got %d", len(payload.Adapters)) + } + + first := payload.Adapters[0] + if first.Type != "mock" { + t.Errorf("expected first adapter type %q, got %q", "mock", first.Type) + } +} diff --git a/apps/edge/internal/node/store.go b/apps/edge/internal/node/store.go index e54128e..05b6298 100644 --- a/apps/edge/internal/node/store.go +++ b/apps/edge/internal/node/store.go @@ -4,6 +4,7 @@ import ( "fmt" "sync" + "github.com/google/uuid" "iop/packages/config" ) @@ -62,7 +63,6 @@ func (s *NodeStore) All() []*NodeRecord { } // LoadFromConfig seeds the store from EdgeConfig.Nodes. -// ID is derived as "node-{alias}" if not provided; use UUIDs in production. func LoadFromConfig(defs []config.NodeDefinition) (*NodeStore, error) { s := NewNodeStore() seenToken := make(map[string]bool) @@ -79,8 +79,12 @@ func LoadFromConfig(defs []config.NodeDefinition) (*NodeStore, error) { } seenToken[d.Token] = true seenAlias[d.Alias] = true + nodeID := d.ID + if nodeID == "" { + nodeID = uuid.NewString() + } s.Add(&NodeRecord{ - ID: "node-" + d.Alias, + ID: nodeID, Alias: d.Alias, Token: d.Token, Adapters: d.Adapters, diff --git a/apps/edge/internal/node/store_test.go b/apps/edge/internal/node/store_test.go index bfc8116..5a61c75 100644 --- a/apps/edge/internal/node/store_test.go +++ b/apps/edge/internal/node/store_test.go @@ -9,22 +9,21 @@ import ( func TestLoadFromConfig_Success(t *testing.T) { store, err := edgenode.LoadFromConfig([]config.NodeDefinition{ - {Alias: "alpha", Token: "token-alpha"}, {Alias: "beta", Token: "token-beta"}, }) if err != nil { t.Fatalf("load from config: %v", err) } - rec, ok := store.FindByToken("token-alpha") + rec, ok := store.FindByToken("token-beta") if !ok { - t.Fatal("expected record for token-alpha") + t.Fatal("expected record for token-beta") } - if rec.ID != "node-alpha" { - t.Fatalf("expected id %q, got %q", "node-alpha", rec.ID) + if rec.Alias != "beta" { + t.Fatalf("expected alias %q, got %q", "beta", rec.Alias) } - if rec.Alias != "alpha" { - t.Fatalf("expected alias %q, got %q", "alpha", rec.Alias) + if rec.ID == "" { + t.Fatal("expected non-empty ID") } if _, ok := store.FindByToken("missing"); ok { @@ -60,3 +59,62 @@ func TestLoadFromConfig_EmptyToken(t *testing.T) { t.Fatal("expected empty token error") } } + +func TestLoadFromConfig_ExplicitID(t *testing.T) { + store, err := edgenode.LoadFromConfig([]config.NodeDefinition{ + {ID: "my-custom-id", Alias: "alpha", Token: "token-alpha"}, + }) + if err != nil { + t.Fatalf("load from config: %v", err) + } + + rec, ok := store.FindByToken("token-alpha") + if !ok { + t.Fatal("expected record for token-alpha") + } + if rec.ID != "my-custom-id" { + t.Fatalf("expected id %q, got %q", "my-custom-id", rec.ID) + } +} + +func TestLoadFromConfig_AutoID(t *testing.T) { + store, err := edgenode.LoadFromConfig([]config.NodeDefinition{ + {Alias: "alpha", Token: "token-alpha"}, + }) + if err != nil { + t.Fatalf("load from config: %v", err) + } + + rec, ok := store.FindByToken("token-alpha") + if !ok { + t.Fatal("expected record for token-alpha") + } + if rec.ID == "" { + t.Fatal("expected auto-generated non-empty ID") + } + if rec.ID == "node-alpha" { + t.Fatal("expected UUID, not alias-derived ID") + } +} + +func TestLoadFromConfig_AutoIDUnique(t *testing.T) { + store, err := edgenode.LoadFromConfig([]config.NodeDefinition{ + {Alias: "alpha", Token: "token-alpha"}, + {Alias: "beta", Token: "token-beta"}, + }) + if err != nil { + t.Fatalf("load from config: %v", err) + } + + recA, ok := store.FindByToken("token-alpha") + if !ok { + t.Fatal("expected record for token-alpha") + } + recB, ok := store.FindByToken("token-beta") + if !ok { + t.Fatal("expected record for token-beta") + } + if recA.ID == recB.ID { + t.Fatalf("expected unique IDs, both are %q", recA.ID) + } +} diff --git a/apps/edge/internal/transport/integration_test.go b/apps/edge/internal/transport/integration_test.go index 9c10cc9..a43876c 100644 --- a/apps/edge/internal/transport/integration_test.go +++ b/apps/edge/internal/transport/integration_test.go @@ -111,8 +111,13 @@ func TestEdgeServerIntegration(t *testing.T) { if !resp.GetAccepted() { t.Fatalf("expected accepted register response, got reason %q", resp.GetReason()) } - if resp.GetNodeId() != "node-test-node" { - t.Fatalf("expected node id %q, got %q", "node-test-node", resp.GetNodeId()) + rec, ok := nodeStore.FindByToken("test-token") + if !ok { + t.Fatal("expected record for test-token in store") + } + wantNodeID := rec.ID + if resp.GetNodeId() != wantNodeID { + t.Fatalf("expected node id %q, got %q", wantNodeID, resp.GetNodeId()) } if resp.GetAlias() != "test-node" { t.Fatalf("expected alias %q, got %q", "test-node", resp.GetAlias()) @@ -125,7 +130,7 @@ func TestEdgeServerIntegration(t *testing.T) { } // 4. Registry 등록 여부 확인 - entry, ok := waitForRegistryEntry(ctx, registry, "node-test-node") + entry, ok := waitForRegistryEntry(ctx, registry, wantNodeID) if !ok { t.Fatal("node was not registered in edge registry within timeout") } @@ -136,15 +141,15 @@ func TestEdgeServerIntegration(t *testing.T) { if entry == nil { t.Fatal("expected registry entry, got nil") } - if got := entry.NodeID; got != "node-test-node" { - t.Fatalf("expected node id %q, got %q", "node-test-node", got) + if got := entry.NodeID; got != wantNodeID { + t.Fatalf("expected node id %q, got %q", wantNodeID, got) } if got := entry.Alias; got != "test-node" { t.Fatalf("expected alias %q, got %q", "test-node", got) } // 5. Alias로 Resolve 확인 - if e, err := registry.Resolve("test-node"); err != nil || e.NodeID != "node-test-node" { + if e, err := registry.Resolve("test-node"); err != nil || e.NodeID != wantNodeID { t.Fatalf("failed to resolve by alias %q: %v", "test-node", err) } } diff --git a/apps/edge/internal/transport/server.go b/apps/edge/internal/transport/server.go index 7815883..60431b1 100644 --- a/apps/edge/internal/transport/server.go +++ b/apps/edge/internal/transport/server.go @@ -2,18 +2,15 @@ package transport import ( "context" - "fmt" "net" "strconv" "sync" toki "git.toki-labs.com/toki/common-proto-socket/go" - "go.uber.org/zap" "google.golang.org/protobuf/proto" - "google.golang.org/protobuf/types/known/structpb" + "go.uber.org/zap" edgenode "iop/apps/edge/internal/node" - "iop/packages/config" iop "iop/proto/gen/iop" ) @@ -116,7 +113,7 @@ func (s *Server) onNodeConnected(client *toki.TcpClient) { return &iop.RegisterResponse{Accepted: false, Reason: "node already connected"}, nil } - cfg, err := buildConfigPayload(rec) + cfg, err := edgenode.BuildConfigPayload(rec) if err != nil { s.logger.Error("build config payload failed", zap.String("node_id", rec.ID), @@ -156,77 +153,3 @@ func safePrefix(s string) string { } return s } - -func buildConfigPayload(rec *edgenode.NodeRecord) (*iop.NodeConfigPayload, error) { - payload := &iop.NodeConfigPayload{ - Runtime: &iop.NodeRuntimeConfig{ - Concurrency: int32(rec.Runtime.Concurrency), - WorkspaceRoot: rec.Runtime.WorkspaceRoot, - }, - } - - mockSettings, err := structpb.NewStruct(nil) - if err != nil { - return nil, fmt.Errorf("buildConfigPayload: mock: %w", err) - } - payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ - Type: "mock", - Enabled: true, - Settings: mockSettings, - }) - - if rec.Adapters.Ollama.Enabled { - payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ - Type: "ollama", - Enabled: true, - Config: &iop.AdapterConfig_Ollama{ - Ollama: &iop.OllamaAdapterConfig{BaseUrl: rec.Adapters.Ollama.BaseURL}, - }, - }) - } - if rec.Adapters.Vllm.Enabled { - payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ - Type: "vllm", - Enabled: true, - Config: &iop.AdapterConfig_Vllm{ - Vllm: &iop.VllmAdapterConfig{Endpoint: rec.Adapters.Vllm.Endpoint}, - }, - }) - } - if rec.Adapters.CLI.Enabled { - profiles := make(map[string]*iop.CLIProfileConfig, len(rec.Adapters.CLI.Profiles)) - for name, p := range rec.Adapters.CLI.Profiles { - profiles[name] = cliProfileToProto(p) - } - payload.Adapters = append(payload.Adapters, &iop.AdapterConfig{ - Type: "cli", - Enabled: true, - Config: &iop.AdapterConfig_Cli{ - Cli: &iop.CLIAdapterConfig{Profiles: profiles}, - }, - }) - } - return payload, nil -} - -func cliProfileToProto(p config.CLIProfileConf) *iop.CLIProfileConfig { - out := &iop.CLIProfileConfig{ - Command: p.Command, - Args: append([]string(nil), p.Args...), - Env: append([]string(nil), p.Env...), - Persistent: p.Persistent, - Terminal: p.Terminal, - ResponseIdleTimeoutMs: int32(p.ResponseIdleTimeoutMS), - StartupIdleTimeoutMs: int32(p.StartupIdleTimeoutMS), - OutputFormat: p.OutputFormat, - Mode: p.Mode, - ResumeArgs: append([]string(nil), p.ResumeArgs...), - } - if !p.CompletionMarker.Empty() { - out.CompletionMarker = &iop.CLICompletionMarker{ - Line: p.CompletionMarker.Line, - Regex: p.CompletionMarker.Regex, - } - } - return out -} diff --git a/apps/edge/internal/transport/server_test.go b/apps/edge/internal/transport/server_test.go index 42ea411..adfe7c3 100644 --- a/apps/edge/internal/transport/server_test.go +++ b/apps/edge/internal/transport/server_test.go @@ -79,7 +79,7 @@ func TestBuildConfigPayload_CLIOneof(t *testing.T) { }, } - payload, err := buildConfigPayload(rec) + payload, err := edgenode.BuildConfigPayload(rec) if err != nil { t.Fatalf("build config payload: %v", err) } @@ -146,7 +146,7 @@ func TestBuildConfigPayload_CLICompletionMarker(t *testing.T) { }, } - payload, err := buildConfigPayload(rec) + payload, err := edgenode.BuildConfigPayload(rec) if err != nil { t.Fatalf("build config payload: %v", err) } @@ -174,7 +174,7 @@ func TestBuildConfigPayload_OllamaVllmOneof(t *testing.T) { Vllm: config.VllmConf{Enabled: true, Endpoint: "http://localhost:8000"}, }, } - payload, err := buildConfigPayload(rec) + payload, err := edgenode.BuildConfigPayload(rec) if err != nil { t.Fatalf("build: %v", err) } diff --git a/apps/node/internal/bootstrap/module.go b/apps/node/internal/bootstrap/module.go index 63722d6..421cf99 100644 --- a/apps/node/internal/bootstrap/module.go +++ b/apps/node/internal/bootstrap/module.go @@ -19,6 +19,7 @@ import ( "iop/packages/observability" ) + // Module returns the fx options that wire the node application. func Module(cfg *config.NodeConfig) fx.Option { return fx.Options( @@ -66,7 +67,7 @@ func Module(cfg *config.NodeConfig) fx.Option { } rtr := router.New(reg, logger) - n := node.New(result.NodeID, rtr, reg, st, logger) + n := node.New(result.NodeID, rtr, st, os.Stdout, logger) result.Session.SetHandler(n) sess = result.Session diff --git a/apps/node/internal/node/node.go b/apps/node/internal/node/node.go index 1133119..f86d62e 100644 --- a/apps/node/internal/node/node.go +++ b/apps/node/internal/node/node.go @@ -16,7 +16,6 @@ import ( "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/structpb" - "iop/apps/node/internal/adapters" "iop/apps/node/internal/runtime" "iop/apps/node/internal/store" "iop/apps/node/internal/transport" @@ -25,29 +24,33 @@ import ( // Node implements transport.Handler and coordinates the full execution pipeline. type Node struct { - nodeID string - router runtime.Router - registry *adapters.Registry - store *store.Store - runs *runManager - logger *zap.Logger + nodeID string + router runtime.Router + store *store.Store + runs *runManager + out io.Writer + logger *zap.Logger } // New creates a Node. It satisfies transport.Handler. +// out receives console debug output; pass os.Stdout for production, io.Discard in tests. func New( nodeID string, router runtime.Router, - registry *adapters.Registry, st *store.Store, + out io.Writer, logger *zap.Logger, ) *Node { + if out == nil { + out = os.Stdout + } return &Node{ - nodeID: nodeID, - router: router, - registry: registry, - store: st, - runs: newRunManager(), - logger: logger, + nodeID: nodeID, + router: router, + store: st, + runs: newRunManager(), + out: out, + logger: logger, } } @@ -72,18 +75,13 @@ func (n *Node) OnRunRequest(ctx context.Context, sess *transport.Session, req *i TimeoutSec: int(req.GetTimeoutSec()), Metadata: req.GetMetadata(), } - printEdgeMessage(os.Stdout, rr.Input) + printEdgeMessage(n.out, rr.Input) - spec, err := n.router.Resolve(ctx, rr) + spec, adapter, err := n.router.ResolveAdapter(ctx, rr) if err != nil { return fmt.Errorf("node: resolve: %w", err) } - adapter, ok := n.registry.Get(spec.Adapter) - if !ok { - return fmt.Errorf("node: adapter %q not found after routing", spec.Adapter) - } - if err := n.store.InsertRun(ctx, store.RunRecord{ RunID: spec.RunID, Adapter: spec.Adapter, @@ -113,7 +111,7 @@ func (n *Node) OnRunRequest(ctx context.Context, sess *transport.Session, req *i sink := &sessionSink{ sess: sess, - out: os.Stdout, + out: n.out, nodeID: n.nodeID, sessionID: normalizeSessionID(spec.SessionID), background: spec.Background, @@ -159,7 +157,7 @@ func (n *Node) OnCancel(_ context.Context, _ *transport.Session, req *iop.Cancel switch cancelActionFromProto(req.GetAction()) { case runtime.CancelActionTerminateSession: - adapter, ok := n.registry.Get(req.GetAdapter()) + adapter, ok := n.router.GetAdapter(req.GetAdapter()) if !ok { return fmt.Errorf("node: adapter %q not found", req.GetAdapter()) } @@ -183,7 +181,7 @@ func (n *Node) OnCommandRequest(ctx context.Context, sess *transport.Session, re zap.String("model", req.GetModel()), ) - adapter, ok := n.registry.Get(req.GetAdapter()) + adapter, ok := n.router.GetAdapter(req.GetAdapter()) if !ok { return &iop.NodeCommandResponse{ RequestId: req.GetRequestId(), diff --git a/apps/node/internal/node/node_test.go b/apps/node/internal/node/node_test.go index 7b92065..02672ab 100644 --- a/apps/node/internal/node/node_test.go +++ b/apps/node/internal/node/node_test.go @@ -3,6 +3,8 @@ package node_test import ( "context" "errors" + "fmt" + "io" "strings" "sync/atomic" "testing" @@ -10,7 +12,6 @@ import ( "go.uber.org/zap" - "iop/apps/node/internal/adapters" "iop/apps/node/internal/node" "iop/apps/node/internal/runtime" "iop/apps/node/internal/store" @@ -116,6 +117,8 @@ func (a *commandAdapter) HandleCommand(ctx context.Context, req runtime.CommandR type fixedRouter struct { adapterName string + adapter runtime.Adapter + adapters map[string]runtime.Adapter } func (r *fixedRouter) Resolve(_ context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, error) { @@ -134,27 +137,51 @@ func (r *fixedRouter) Resolve(_ context.Context, req runtime.RunRequest) (runtim }, nil } +func (r *fixedRouter) ResolveAdapter(ctx context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) { + spec, err := r.Resolve(ctx, req) + if err != nil { + return runtime.ExecutionSpec{}, nil, err + } + a, ok := r.adapters[spec.Adapter] + if !ok { + return runtime.ExecutionSpec{}, nil, fmt.Errorf("router: adapter %q not found", spec.Adapter) + } + return spec, a, nil +} + +func (r *fixedRouter) GetAdapter(adapterName string) (runtime.Adapter, bool) { + a, ok := r.adapters[adapterName] + return a, ok +} + type errorRouter struct{ err error } func (r *errorRouter) Resolve(_ context.Context, _ runtime.RunRequest) (runtime.ExecutionSpec, error) { return runtime.ExecutionSpec{}, r.err } -func makeNode(t *testing.T, rtr runtime.Router, reg *adapters.Registry) (*node.Node, *store.Store) { +func (r *errorRouter) ResolveAdapter(_ context.Context, _ runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) { + return runtime.ExecutionSpec{}, nil, r.err +} + +func (r *errorRouter) GetAdapter(_ string) (runtime.Adapter, bool) { + return nil, false +} + +func makeNode(t *testing.T, rtr runtime.Router) (*node.Node, *store.Store) { t.Helper() st, err := store.New(":memory:", zap.NewNop()) if err != nil { t.Fatalf("store: %v", err) } t.Cleanup(func() { _ = st.Close() }) - return node.New("test-node", rtr, reg, st, zap.NewNop()), st + return node.New("test-node", rtr, st, io.Discard, zap.NewNop()), st } // --- existing tests (updated) --- func TestOnRunRequest_RouterError(t *testing.T) { - reg := adapters.NewRegistry() - n, _ := makeNode(t, &errorRouter{err: errors.New("boom")}, reg) + n, _ := makeNode(t, &errorRouter{err: errors.New("boom")}) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1"}) if err == nil { @@ -166,23 +193,23 @@ func TestOnRunRequest_RouterError(t *testing.T) { } func TestOnRunRequest_AdapterNotFound(t *testing.T) { - reg := adapters.NewRegistry() - n, _ := makeNode(t, &fixedRouter{adapterName: "missing"}, reg) + router := &fixedRouter{adapterName: "missing", adapters: make(map[string]runtime.Adapter)} + n, _ := makeNode(t, router) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1"}) if err == nil { t.Fatal("expected error") } - if !strings.Contains(err.Error(), "not found after routing") { + if !strings.Contains(err.Error(), "not found") { t.Fatalf("expected adapter lookup error, got %v", err) } } func TestOnRunRequest_Success(t *testing.T) { - reg := adapters.NewRegistry() adapter := &countingAdapter{} - reg.Register(adapter) - n, st := makeNode(t, &fixedRouter{adapterName: "test"}, reg) + router := &fixedRouter{adapterName: "test", adapters: make(map[string]runtime.Adapter)} + router.adapters["test"] = adapter + n, st := makeNode(t, router) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{ RunId: "run-1", @@ -212,9 +239,9 @@ func TestOnRunRequest_Success(t *testing.T) { func TestOnRunRequest_ForegroundAdapterErrorReturned(t *testing.T) { adapter := &failingAdapter{err: errors.New("adapter boom")} - reg := adapters.NewRegistry() - reg.Register(adapter) - n, st := makeNode(t, &fixedRouter{adapterName: "failing"}, reg) + router := &fixedRouter{adapterName: "failing", adapters: make(map[string]runtime.Adapter)} + router.adapters["failing"] = adapter + n, st := makeNode(t, router) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{ RunId: "run-fail", @@ -245,9 +272,9 @@ func TestOnRunRequest_ForegroundAdapterErrorReturned(t *testing.T) { func TestOnRunRequest_ForegroundCancelReturnedAndStored(t *testing.T) { adapter := &failingAdapter{err: runtime.ErrRunCancelled} - reg := adapters.NewRegistry() - reg.Register(adapter) - n, st := makeNode(t, &fixedRouter{adapterName: "failing"}, reg) + router := &fixedRouter{adapterName: "failing", adapters: make(map[string]runtime.Adapter)} + router.adapters["failing"] = adapter + n, st := makeNode(t, router) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{ RunId: "run-cancel-fg", @@ -268,9 +295,9 @@ func TestOnRunRequest_ForegroundCancelReturnedAndStored(t *testing.T) { func TestOnRunRequest_BackgroundReturnsBeforeAdapterCompletes(t *testing.T) { ba := newBlockingAdapter() - reg := adapters.NewRegistry() - reg.Register(ba) - n, st := makeNode(t, &fixedRouter{adapterName: "blocking"}, reg) + router := &fixedRouter{adapterName: "blocking", adapters: make(map[string]runtime.Adapter)} + router.adapters["blocking"] = ba + n, st := makeNode(t, router) err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{ RunId: "run-bg", @@ -318,9 +345,9 @@ func TestOnRunRequest_BackgroundReturnsBeforeAdapterCompletes(t *testing.T) { func TestOnCancel_CancelsRunViaRunManager(t *testing.T) { ba := newBlockingAdapter() - reg := adapters.NewRegistry() - reg.Register(ba) - n, _ := makeNode(t, &fixedRouter{adapterName: "blocking"}, reg) + router := &fixedRouter{adapterName: "blocking", adapters: make(map[string]runtime.Adapter)} + router.adapters["blocking"] = ba + n, _ := makeNode(t, router) go func() { _ = n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{ @@ -345,9 +372,9 @@ func TestOnCancel_CancelsRunViaRunManager(t *testing.T) { func TestOnCancel_TerminatesAdapterSession(t *testing.T) { ta := &terminatingAdapter{} - reg := adapters.NewRegistry() - reg.Register(ta) - n, _ := makeNode(t, &fixedRouter{adapterName: "terminating"}, reg) + router := &fixedRouter{adapterName: "terminating", adapters: make(map[string]runtime.Adapter)} + router.adapters["terminating"] = ta + n, _ := makeNode(t, router) err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{ RunId: "", @@ -369,9 +396,9 @@ func TestOnCancel_TerminatesAdapterSession(t *testing.T) { func TestOnCommandRequest_Success(t *testing.T) { ca := &commandAdapter{} - reg := adapters.NewRegistry() - reg.Register(ca) - n, _ := makeNode(t, &fixedRouter{adapterName: "command"}, reg) + router := &fixedRouter{adapterName: "command", adapters: make(map[string]runtime.Adapter)} + router.adapters["command"] = ca + n, _ := makeNode(t, router) resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{ RequestId: "req-1", @@ -395,8 +422,8 @@ func TestOnCommandRequest_Success(t *testing.T) { } func TestOnCommandRequest_MissingAdapter(t *testing.T) { - reg := adapters.NewRegistry() - n, _ := makeNode(t, &fixedRouter{adapterName: "missing"}, reg) + router := &fixedRouter{adapterName: "missing", adapters: make(map[string]runtime.Adapter)} + n, _ := makeNode(t, router) resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{ RequestId: "req-1", @@ -413,9 +440,9 @@ func TestOnCommandRequest_MissingAdapter(t *testing.T) { func TestOnCommandRequest_NotSupported(t *testing.T) { ta := &terminatingAdapter{} - reg := adapters.NewRegistry() - reg.Register(ta) - n, _ := makeNode(t, &fixedRouter{adapterName: "terminating"}, reg) + router := &fixedRouter{adapterName: "terminating", adapters: make(map[string]runtime.Adapter)} + router.adapters["terminating"] = ta + n, _ := makeNode(t, router) resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{ RequestId: "req-1", @@ -432,9 +459,9 @@ func TestOnCommandRequest_NotSupported(t *testing.T) { func TestOnCommandRequest_AdapterError(t *testing.T) { ca := &commandAdapter{} - reg := adapters.NewRegistry() - reg.Register(ca) - n, _ := makeNode(t, &fixedRouter{adapterName: "command"}, reg) + router := &fixedRouter{adapterName: "command", adapters: make(map[string]runtime.Adapter)} + router.adapters["command"] = ca + n, _ := makeNode(t, router) resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{ RequestId: "req-1", diff --git a/apps/node/internal/router/router.go b/apps/node/internal/router/router.go index 998e155..ce0cc50 100644 --- a/apps/node/internal/router/router.go +++ b/apps/node/internal/router/router.go @@ -54,3 +54,19 @@ func (r *defaultRouter) Resolve(_ context.Context, req runtime.RunRequest) (runt ) return spec, nil } + +func (r *defaultRouter) ResolveAdapter(ctx context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) { + spec, err := r.Resolve(ctx, req) + if err != nil { + return runtime.ExecutionSpec{}, nil, err + } + adapter, ok := r.registry.Get(spec.Adapter) + if !ok { + return runtime.ExecutionSpec{}, nil, fmt.Errorf("router: adapter %q not found", spec.Adapter) + } + return spec, adapter, nil +} + +func (r *defaultRouter) GetAdapter(adapterName string) (runtime.Adapter, bool) { + return r.registry.Get(adapterName) +} diff --git a/apps/node/internal/router/router_test.go b/apps/node/internal/router/router_test.go index 226dfd4..36f518d 100644 --- a/apps/node/internal/router/router_test.go +++ b/apps/node/internal/router/router_test.go @@ -2,6 +2,7 @@ package router import ( "context" + "strings" "testing" "go.uber.org/zap" @@ -47,3 +48,46 @@ func TestResolve_PreservesSessionFields(t *testing.T) { t.Errorf("Background: got %v want %v", spec.Background, req.Background) } } + +func TestResolveAdapter_Found(t *testing.T) { + reg := adapters.NewRegistry() + expected := &stubAdapter{name: "cli"} + reg.Register(expected) + r := New(reg, zap.NewNop()) + + req := runtime.RunRequest{ + RunID: "run-1", + Adapter: "cli", + Model: "codex", + } + + spec, adapter, err := r.ResolveAdapter(context.Background(), req) + if err != nil { + t.Fatalf("ResolveAdapter: %v", err) + } + if spec.Adapter != "cli" { + t.Errorf("expected adapter name 'cli', got %q", spec.Adapter) + } + if adapter != expected { + t.Fatal("expected same adapter instance") + } +} + +func TestResolveAdapter_NotFound(t *testing.T) { + reg := adapters.NewRegistry() + r := New(reg, zap.NewNop()) + + req := runtime.RunRequest{ + RunID: "run-1", + Adapter: "nonexistent", + Model: "codex", + } + + _, _, err := r.ResolveAdapter(context.Background(), req) + if err == nil { + t.Fatal("expected error for unregistered adapter") + } + if !strings.Contains(err.Error(), "not found") { + t.Fatalf("expected 'not found' in error, got %v", err) + } +} diff --git a/apps/node/internal/runtime/types.go b/apps/node/internal/runtime/types.go index e02543e..abb1e7b 100644 --- a/apps/node/internal/runtime/types.go +++ b/apps/node/internal/runtime/types.go @@ -151,9 +151,11 @@ type Adapter interface { Execute(ctx context.Context, spec ExecutionSpec, sink EventSink) error } -// Router resolves a RunRequest into a concrete ExecutionSpec. +// Router resolves a RunRequest into a concrete ExecutionSpec and the Adapter to execute it. type Router interface { Resolve(ctx context.Context, req RunRequest) (ExecutionSpec, error) + ResolveAdapter(ctx context.Context, req RunRequest) (ExecutionSpec, Adapter, error) + GetAdapter(adapterName string) (Adapter, bool) } // PolicyEngine applies and validates policies on execution specs. diff --git a/configs/edge.yaml b/configs/edge.yaml index 4dfa242..e130e28 100644 --- a/configs/edge.yaml +++ b/configs/edge.yaml @@ -19,6 +19,7 @@ console: timeout_sec: 300 nodes: + # - id: "node-uuid" # optional; UUID v4 auto-assigned if empty - alias: "local-node" token: "changeme" adapters: diff --git a/go.mod b/go.mod index 85a5cd67..394e047 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.24 require ( git.toki-labs.com/toki/common-proto-socket/go v0.0.0-00010101000000-000000000000 github.com/creack/pty v1.1.24 + github.com/google/uuid v1.6.0 github.com/prometheus/client_golang v1.20.5 github.com/spf13/cobra v1.8.1 github.com/spf13/viper v1.19.0 @@ -20,7 +21,6 @@ require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/fsnotify/fsnotify v1.7.0 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect github.com/hashicorp/hcl v1.0.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect diff --git a/packages/config/config.go b/packages/config/config.go index ce0b3ae..1aa8724 100644 --- a/packages/config/config.go +++ b/packages/config/config.go @@ -21,6 +21,7 @@ type EdgeConfig struct { // NodeDefinition is the edge-side record for a pre-registered node. type NodeDefinition struct { + ID string `mapstructure:"id" yaml:"id"` // optional; UUID v4 auto-assigned if empty Alias string `mapstructure:"alias" yaml:"alias"` Token string `mapstructure:"token" yaml:"token"` Adapters AdaptersConf `mapstructure:"adapters" yaml:"adapters"`