toki/iop
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: edge config mapper, node writer injection, router registry implementation

Browse source
This commit is contained in:
toki 2026-05-05 08:57:21 +09:00
parent cb96fb58eb
commit d89af2e726
16 changed files with 495 additions and 161 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

86
apps/edge/internal/node/mapper.go Normal file
View file

@ -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
}

168
apps/edge/internal/node/mapper_test.go Normal file
View file

@ -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)
}
}

8
apps/edge/internal/node/store.go
View file

@ -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,

72
apps/edge/internal/node/store_test.go
View file

@ -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)
}
}

17
apps/edge/internal/transport/integration_test.go
View file

@ -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)
}
}

81
apps/edge/internal/transport/server.go
View file

@ -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
}

6
apps/edge/internal/transport/server_test.go
View file

@ -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)
}

3
apps/node/internal/bootstrap/module.go
View file

@ -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

46
apps/node/internal/node/node.go
View file

@ -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(),

101
apps/node/internal/node/node_test.go
View file

@ -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",

16
apps/node/internal/router/router.go
View file

@ -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)
}

44
apps/node/internal/router/router_test.go
View file

@ -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)
}
}

4
apps/node/internal/runtime/types.go
View file

@ -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.

1
configs/edge.yaml
View file

@ -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:

2
go.mod
View file

@ -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

1
packages/config/config.go
View file

@ -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"`