iop/apps/node/internal/node/node_test.go

3346 lines
102 KiB
Go

package node_test
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/structpb"
"iop/apps/node/internal/adapters"
"iop/apps/node/internal/node"
"iop/apps/node/internal/router"
"iop/apps/node/internal/runtime"
"iop/apps/node/internal/store"
"iop/apps/node/internal/transport"
iop "iop/proto/gen/iop"
)
// --- test doubles ---
type countingAdapter struct {
executeCalls int32
lastSpec runtime.ExecutionSpec
}
func (a *countingAdapter) Name() string { return "test" }
func (a *countingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "test", Targets: []string{"v1"}, MaxConcurrency: 1}, nil
}
func (a *countingAdapter) Execute(_ context.Context, spec runtime.ExecutionSpec, _ runtime.EventSink) error {
atomic.AddInt32(&a.executeCalls, 1)
a.lastSpec = spec
return nil
}
// failingAdapter returns a fixed error from Execute.
type failingAdapter struct{ err error }
func (a *failingAdapter) Name() string { return "failing" }
func (a *failingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "failing"}, nil
}
func (a *failingAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return a.err
}
// blockingAdapter blocks until its context is cancelled, then returns ErrRunCancelled.
type blockingAdapter struct {
started chan struct{}
done chan struct{}
}
func newBlockingAdapter() *blockingAdapter {
return &blockingAdapter{started: make(chan struct{}), done: make(chan struct{})}
}
func (a *blockingAdapter) Name() string { return "blocking" }
func (a *blockingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "blocking"}, nil
}
func (a *blockingAdapter) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
close(a.started)
<-ctx.Done()
close(a.done)
return runtime.ErrRunCancelled
}
// terminatingAdapter implements SessionTerminator.
type terminatingAdapter struct {
terminateCalls int32
lastTarget string
lastSessionID string
}
func (a *terminatingAdapter) Name() string { return "terminating" }
func (a *terminatingAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "terminating"}, nil
}
func (a *terminatingAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return nil
}
func (a *terminatingAdapter) TerminateSession(_ context.Context, target, sessionID string) error {
atomic.AddInt32(&a.terminateCalls, 1)
a.lastTarget = target
a.lastSessionID = sessionID
return nil
}
type instanceKeyAdapter struct {
instanceKey string
}
func (a *instanceKeyAdapter) Name() string { return "ollama" }
func (a *instanceKeyAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: "ollama",
InstanceKey: a.instanceKey,
Targets: []string{"llama3"},
MaxConcurrency: 4,
}, nil
}
func (a *instanceKeyAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return nil
}
type commandAdapter struct {
lastReq runtime.CommandRequest
providerStatus runtime.ProviderStatus
}
func (a *commandAdapter) Name() string { return "command" }
func (a *commandAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: "command",
Targets: []string{"v2", "v1"},
MaxConcurrency: 3,
ProviderStatus: a.providerStatus,
}, nil
}
func (a *commandAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return nil
}
func (a *commandAdapter) HandleCommand(ctx context.Context, req runtime.CommandRequest) (runtime.CommandResponse, error) {
a.lastReq = req
switch req.Type {
case runtime.CommandTypeUsageStatus:
return runtime.CommandResponse{
RequestID: req.RequestID,
Type: req.Type,
Adapter: req.Adapter,
Target: req.Target,
SessionID: req.SessionID,
UsageStatus: &runtime.AgentUsageStatus{
RawOutput: "success",
},
}, nil
default:
return runtime.CommandResponse{}, errors.New("command not supported by adapter")
}
}
type proberTestAdapter struct {
providerStatus runtime.ProviderStatus
probeTargets []string
probeErr error
}
func (a *proberTestAdapter) Name() string { return "prober" }
func (a *proberTestAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: "prober",
Targets: []string{"model-a"},
MaxConcurrency: 3,
ProviderStatus: a.providerStatus,
}, nil
}
func (a *proberTestAdapter) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return nil
}
func (a *proberTestAdapter) ProbeProvider(ctx context.Context, target string) (runtime.ProviderProbeResult, error) {
if a.probeErr != nil {
return runtime.ProviderProbeResult{}, a.probeErr
}
status := runtime.ProviderStatusAvailable
found := false
for _, t := range a.probeTargets {
if t == target {
found = true
break
}
}
if target != "" && !found {
status = runtime.ProviderStatusUnavailable
}
return runtime.ProviderProbeResult{
AdapterName: "prober",
Target: target,
Targets: a.probeTargets,
Status: status,
}, nil
}
type fixedRouter struct {
adapterName string
adapter runtime.Adapter
adapters map[string]runtime.Adapter
lookupErrors map[string]error // optional per-adapter errors for LookupAdapter
}
func (r *fixedRouter) Resolve(_ context.Context, req runtime.RunRequest) (runtime.ExecutionSpec, error) {
return runtime.ExecutionSpec{
RunID: req.RunID,
Adapter: r.adapterName,
Target: req.Target,
SessionID: req.SessionID,
SessionMode: req.SessionMode,
Background: req.Background,
Workspace: req.Workspace,
Policy: req.Policy,
Input: req.Input,
TimeoutSec: req.TimeoutSec,
Metadata: req.Metadata,
}, 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) LookupAdapter(adapterName string) (runtime.Adapter, error) {
if r.lookupErrors != nil {
if err, ok := r.lookupErrors[adapterName]; ok {
return nil, err
}
}
a, ok := r.adapters[adapterName]
if !ok {
return nil, fmt.Errorf("adapter %q not found", adapterName)
}
return 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 (r *errorRouter) ResolveAdapter(_ context.Context, _ runtime.RunRequest) (runtime.ExecutionSpec, runtime.Adapter, error) {
return runtime.ExecutionSpec{}, nil, r.err
}
func (r *errorRouter) LookupAdapter(_ string) (runtime.Adapter, error) {
return 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()
return makeNodeWithConcurrency(t, rtr, 0)
}
func makeNodeWithConcurrency(t *testing.T, rtr runtime.Router, concurrency int) (*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, st, concurrency, io.Discard, zap.NewNop(), nil), st
}
// --- existing tests (updated) ---
func TestOnRunRequest_RouterError(t *testing.T) {
n, _ := makeNode(t, &errorRouter{err: errors.New("boom")})
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1"})
if err == nil {
t.Fatal("expected error")
}
if !strings.Contains(err.Error(), "node: resolve:") {
t.Fatalf("expected resolve prefix, got %v", err)
}
}
func TestOnRunRequest_AdapterNotFound(t *testing.T) {
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") {
t.Fatalf("expected adapter lookup error, got %v", err)
}
}
func TestOnRunRequest_Success(t *testing.T) {
adapter := &countingAdapter{}
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",
Adapter: "test",
Target: "v1",
Workspace: "/config/workspace/iop",
})
if err != nil {
t.Fatalf("run request: %v", err)
}
if got := atomic.LoadInt32(&adapter.executeCalls); got != 1 {
t.Fatalf("expected 1 execute call, got %d", got)
}
if adapter.lastSpec.Workspace != "/config/workspace/iop" {
t.Fatalf("adapter workspace: got %q", adapter.lastSpec.Workspace)
}
run, err := st.GetRun(context.Background(), "run-1")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run == nil {
t.Fatal("expected persisted run record")
}
if run.Status != "completed" {
t.Fatalf("expected completed status, got %q", run.Status)
}
}
// --- new tests ---
func TestOnRunRequest_ForegroundAdapterErrorReturned(t *testing.T) {
adapter := &failingAdapter{err: errors.New("adapter boom")}
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",
Adapter: "failing",
Target: "v1",
})
if err == nil {
t.Fatal("expected error from OnRunRequest")
}
if !strings.Contains(err.Error(), "adapter boom") {
t.Fatalf("expected adapter boom in error, got %v", err)
}
run, err := st.GetRun(context.Background(), "run-fail")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run == nil {
t.Fatal("expected persisted run record")
}
if run.Status != "failed" {
t.Fatalf("expected failed status, got %q", run.Status)
}
if !strings.Contains(run.Error, "adapter boom") {
t.Fatalf("expected adapter boom in stored error, got %q", run.Error)
}
}
func TestOnRunRequest_ForegroundCancelReturnedAndStored(t *testing.T) {
adapter := &failingAdapter{err: runtime.ErrRunCancelled}
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",
Adapter: "failing",
Target: "v1",
})
if !errors.Is(err, runtime.ErrRunCancelled) {
t.Fatalf("expected ErrRunCancelled, got %v", err)
}
run, err := st.GetRun(context.Background(), "run-cancel-fg")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run == nil || run.Status != "cancelled" {
t.Fatalf("expected cancelled status, got %+v", run)
}
}
func TestOnRunRequest_BackgroundReturnsBeforeAdapterCompletes(t *testing.T) {
ba := newBlockingAdapter()
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",
Adapter: "blocking",
Target: "v1",
Background: true,
})
if err != nil {
t.Fatalf("OnRunRequest: %v", err)
}
// OnRunRequest should return before adapter finishes.
select {
case <-ba.started:
case <-time.After(2 * time.Second):
t.Fatal("adapter never started")
}
// Cancel via OnCancel so the blocking adapter finishes.
if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "run-bg"}); err != nil {
t.Fatalf("OnCancel: %v", err)
}
select {
case <-ba.done:
case <-time.After(2 * time.Second):
t.Fatal("adapter did not finish after cancel")
}
// Wait briefly for completeRun goroutine to update the store.
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
run, err := st.GetRun(context.Background(), "run-bg")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run != nil && run.Status == "cancelled" {
return
}
time.Sleep(10 * time.Millisecond)
}
run, _ := st.GetRun(context.Background(), "run-bg")
t.Fatalf("expected cancelled status, got %q", run.Status)
}
func TestOnCancel_CancelsRunViaRunManager(t *testing.T) {
ba := newBlockingAdapter()
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{
RunId: "run-cancel",
Adapter: "blocking",
Target: "v1",
})
}()
<-ba.started
if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "run-cancel"}); err != nil {
t.Fatalf("OnCancel: %v", err)
}
select {
case <-ba.done:
case <-time.After(2 * time.Second):
t.Fatal("adapter did not finish after cancel")
}
}
func TestOnCancel_TerminatesAdapterSession(t *testing.T) {
ta := &terminatingAdapter{}
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: "",
Adapter: "terminating",
Target: "codex",
SessionId: "session-a",
Action: iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
})
if err != nil {
t.Fatalf("OnCancel terminate: %v", err)
}
if atomic.LoadInt32(&ta.terminateCalls) != 1 {
t.Fatal("expected TerminateSession to be called once")
}
if ta.lastTarget != "codex" || ta.lastSessionID != "session-a" {
t.Fatalf("unexpected terminate args: target=%q session=%q", ta.lastTarget, ta.lastSessionID)
}
}
func TestOnCommandRequest_Success(t *testing.T) {
ca := &commandAdapter{}
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",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "command",
Target: "codex",
SessionId: "default",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if resp.GetUsageStatus().GetRawOutput() != "success" {
t.Fatalf("expected success raw output, got %q", resp.GetUsageStatus().GetRawOutput())
}
if ca.lastReq.Type != runtime.CommandTypeUsageStatus || ca.lastReq.Target != "codex" {
t.Fatalf("unexpected last req: %+v", ca.lastReq)
}
}
func TestOnCommandRequest_MissingAdapter(t *testing.T) {
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",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "missing",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "not found") {
t.Fatalf("expected not found error, got %q", resp.GetError())
}
}
func TestOnCommandRequest_NotSupported(t *testing.T) {
ta := &terminatingAdapter{}
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",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "terminating",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "does not support commands") {
t.Fatalf("expected not supported error, got %q", resp.GetError())
}
}
func TestOnCommandRequest_UnspecifiedRejected(t *testing.T) {
ca := &commandAdapter{}
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",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_UNSPECIFIED,
Adapter: "command",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "unsupported command type") {
t.Fatalf("expected unsupported command type error, got %q", resp.GetError())
}
// Adapter must not be invoked for unsupported command types.
if ca.lastReq.RequestID != "" {
t.Fatalf("expected adapter to be skipped, got lastReq %+v", ca.lastReq)
}
}
func TestOnCommandRequest_Capabilities(t *testing.T) {
ca := &commandAdapter{}
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-cap",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "command",
Target: "codex",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if resp.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
t.Fatalf("type: got %v want CAPABILITIES", resp.GetType())
}
if got := resp.GetResult()["targets"]; got != "v1,v2" {
t.Fatalf("result[targets]: got %q want %q", got, "v1,v2")
}
if got := resp.GetResult()["adapter"]; got != "command" {
t.Fatalf("result[adapter]: got %q want %q", got, "command")
}
if got := resp.GetResult()["max_concurrency"]; got != "3" {
t.Fatalf("result[max_concurrency]: got %q want %q", got, "3")
}
if got := resp.GetResult()["provider_status"]; got != "unknown" {
t.Fatalf("result[provider_status]: got %q want %q", got, "unknown")
}
if got := resp.GetResult()["capacity"]; got != "3" {
t.Fatalf("result[capacity]: got %q want %q", got, "3")
}
if got := resp.GetResult()["in_flight"]; got != "0" {
t.Fatalf("result[in_flight]: got %q want %q", got, "0")
}
if got := resp.GetResult()["queued"]; got != "0" {
t.Fatalf("result[queued]: got %q want %q", got, "0")
}
if len(resp.GetProviderSnapshots()) != 1 {
t.Fatalf("expected 1 provider snapshot, got %d", len(resp.GetProviderSnapshots()))
}
snap := resp.GetProviderSnapshots()[0]
if snap.GetAdapter() != "command" {
t.Fatalf("snap.Adapter: got %q want %q", snap.GetAdapter(), "command")
}
if snap.GetStatus() != "unknown" {
t.Fatalf("snap.Status: got %q want %q", snap.GetStatus(), "unknown")
}
if snap.GetCapacity() != 3 {
t.Fatalf("snap.Capacity: got %d want %d", snap.GetCapacity(), 3)
}
if snap.GetInFlight() != 0 {
t.Fatalf("snap.InFlight: got %d want %d", snap.GetInFlight(), 0)
}
if snap.GetQueued() != 0 {
t.Fatalf("snap.Queued: got %d want %d", snap.GetQueued(), 0)
}
// Capabilities must not go through the adapter CommandHandler.
if ca.lastReq.RequestID != "" {
t.Fatalf("expected CommandHandler to be skipped, got %+v", ca.lastReq)
}
}
func TestOnCommandRequest_Capabilities_InFlight(t *testing.T) {
ba := newBlockingAdapter()
router := &fixedRouter{adapterName: "blocking", adapters: make(map[string]runtime.Adapter)}
router.adapters["blocking"] = ba
n, _ := makeNode(t, router)
// Start a background run which will block.
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-in-flight",
Adapter: "blocking",
Target: "v1",
Background: true,
})
if err != nil {
t.Fatalf("OnRunRequest: %v", err)
}
select {
case <-ba.started:
case <-time.After(2 * time.Second):
t.Fatal("adapter never started")
}
// Now query Capabilities while the run is active.
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-cap",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "blocking",
Target: "v1",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
// Verify capacity, in_flight, queued values.
if got := resp.GetResult()["capacity"]; got != "0" {
t.Fatalf("result[capacity]: got %q want %q", got, "0")
}
if got := resp.GetResult()["in_flight"]; got != "1" {
t.Fatalf("result[in_flight]: got %q want %q", got, "1")
}
if got := resp.GetResult()["queued"]; got != "0" {
t.Fatalf("result[queued]: got %q want %q", got, "0")
}
if len(resp.GetProviderSnapshots()) != 1 {
t.Fatalf("expected 1 provider snapshot, got %d", len(resp.GetProviderSnapshots()))
}
snap := resp.GetProviderSnapshots()[0]
if snap.GetAdapter() != "blocking" {
t.Fatalf("snap.Adapter: got %q want %q", snap.GetAdapter(), "blocking")
}
if snap.GetStatus() != "unknown" {
t.Fatalf("snap.Status: got %q want %q", snap.GetStatus(), "unknown")
}
if snap.GetCapacity() != 0 {
t.Fatalf("snap.Capacity: got %d want %d", snap.GetCapacity(), 0)
}
if snap.GetInFlight() != 1 {
t.Fatalf("snap.InFlight: got %d want %d", snap.GetInFlight(), 1)
}
if snap.GetQueued() != 0 {
t.Fatalf("snap.Queued: got %d want %d", snap.GetQueued(), 0)
}
// Clean up.
if err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{RunId: "run-in-flight"}); err != nil {
t.Fatalf("OnCancel: %v", err)
}
<-ba.done
}
func TestOnCommandRequest_Capabilities_SafetyRejected(t *testing.T) {
sa := newQueuedSlowAdapter("slow", 1, 4, 0)
router := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
n, _ := makeNodeWithConcurrency(t, router, 0)
if err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "cap-hold", Adapter: "slow", Target: "v1", Background: true,
}); err != nil {
t.Fatalf("holding run: %v", err)
}
waitStarted(t, sa, "cap-hold")
// The second run request must fail immediately since there is no queue.
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "cap-rejected", Adapter: "slow", Target: "v1", Background: true,
})
if !errors.Is(err, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("expected ErrConcurrencyLimitExceeded, got %v", err)
}
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-cap-queued",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "slow",
Target: "v1",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["capacity"]; got != "1" {
t.Fatalf("result[capacity]: got %q want %q", got, "1")
}
if got := resp.GetResult()["in_flight"]; got != "1" {
t.Fatalf("result[in_flight]: got %q want %q", got, "1")
}
if got := resp.GetResult()["queued"]; got != "0" {
t.Fatalf("result[queued]: got %q want %q", got, "0")
}
if len(resp.GetProviderSnapshots()) != 1 {
t.Fatalf("expected 1 provider snapshot, got %d", len(resp.GetProviderSnapshots()))
}
snap := resp.GetProviderSnapshots()[0]
if snap.GetCapacity() != 1 {
t.Fatalf("snap.Capacity: got %d want %d", snap.GetCapacity(), 1)
}
if snap.GetInFlight() != 1 {
t.Fatalf("snap.InFlight: got %d want %d", snap.GetInFlight(), 1)
}
if snap.GetQueued() != 0 {
t.Fatalf("snap.Queued: got %d want %d", snap.GetQueued(), 0)
}
sa.releaseRun("cap-hold")
}
func TestOnCommandRequest_CapabilitiesProviderStatusModel(t *testing.T) {
cases := []struct {
name string
status runtime.ProviderStatus
want string
}{
{name: "available", status: runtime.ProviderStatusAvailable, want: "available"},
{name: "unavailable", status: runtime.ProviderStatusUnavailable, want: "unavailable"},
{name: "invalid status folds to unknown", status: runtime.ProviderStatus("degraded"), want: "unknown"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
ca := &commandAdapter{providerStatus: tc.status}
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-cap-status",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "command",
Target: "codex",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["provider_status"]; got != tc.want {
t.Fatalf("result[provider_status]: got %q want %q", got, tc.want)
}
})
}
}
func TestOnCommandRequest_CapabilitiesWithProber(t *testing.T) {
t.Run("probe_target_hit", func(t *testing.T) {
pa := &proberTestAdapter{
providerStatus: runtime.ProviderStatusAvailable,
probeTargets: []string{"model-a", "model-b"},
}
router := &fixedRouter{adapterName: "prober", adapters: make(map[string]runtime.Adapter)}
router.adapters["prober"] = pa
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-cap-prober",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "prober",
Target: "model-a",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["provider_status"]; got != "available" {
t.Fatalf("result[provider_status]: got %q want %q", got, "available")
}
if got := resp.GetResult()["targets"]; got != "model-a,model-b" {
t.Fatalf("result[targets]: got %q want %q", got, "model-a,model-b")
}
})
t.Run("probe_target_miss", func(t *testing.T) {
pa := &proberTestAdapter{
providerStatus: runtime.ProviderStatusAvailable,
probeTargets: []string{"model-a"},
}
router := &fixedRouter{adapterName: "prober", adapters: make(map[string]runtime.Adapter)}
router.adapters["prober"] = pa
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-cap-prober",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "prober",
Target: "model-c",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["provider_status"]; got != "unavailable" {
t.Fatalf("result[provider_status]: got %q want %q", got, "unavailable")
}
})
t.Run("probe_error_folds_to_unavailable", func(t *testing.T) {
pa := &proberTestAdapter{
providerStatus: runtime.ProviderStatusAvailable,
probeErr: fmt.Errorf("network error"),
}
router := &fixedRouter{adapterName: "prober", adapters: make(map[string]runtime.Adapter)}
router.adapters["prober"] = pa
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-cap-prober",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "prober",
Target: "model-a",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
// Probe error should NOT cause a command error, but set provider_status to unavailable
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["provider_status"]; got != "unavailable" {
t.Fatalf("result[provider_status]: got %q want %q", got, "unavailable")
}
if got := resp.GetResult()["provider_detail"]; got != "network error" {
t.Fatalf("result[provider_detail]: got %q want %q", got, "network error")
}
})
}
// TestOnCommandRequest_CapabilitiesWithoutCommandHandler verifies CAPABILITIES
// succeeds for an adapter that does not implement runtime.CommandHandler.
func TestOnCommandRequest_CapabilitiesWithoutCommandHandler(t *testing.T) {
ta := &terminatingAdapter{} // no CommandHandler
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-cap",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "terminating",
Target: "codex",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["adapter"]; got != "terminating" {
t.Fatalf("result[adapter]: got %q want %q", got, "terminating")
}
}
func TestOnCommandRequest_TransportStatus(t *testing.T) {
ta := &terminatingAdapter{} // adapter not required for transport_status
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-tx",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS,
Adapter: "terminating",
Target: "codex",
SessionId: "sess-a",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetResult()["node_id"]; got != "test-node" {
t.Fatalf("result[node_id]: got %q want %q", got, "test-node")
}
if got := resp.GetResult()["adapter"]; got != "terminating" {
t.Fatalf("result[adapter]: got %q want %q", got, "terminating")
}
if got := resp.GetResult()["session_id"]; got != "sess-a" {
t.Fatalf("result[session_id]: got %q want %q", got, "sess-a")
}
if _, ok := resp.GetResult()["connected"]; !ok {
t.Fatalf("result missing connected key: %+v", resp.GetResult())
}
if got, ok := resp.GetResult()["state"]; !ok || (got != "connected" && got != "disconnected") {
t.Fatalf("result[state]: got %q, want \"connected\" or \"disconnected\"", got)
}
}
func TestOnCommandRequest_AdapterError(t *testing.T) {
ca := &commandAdapter{}
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",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST,
Adapter: "command",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "command not supported by adapter") {
t.Fatalf("expected adapter unsupported error, got %q", resp.GetError())
}
}
func TestOnCommandRequest_TransportStatusDefaultsSessionID(t *testing.T) {
ta := &terminatingAdapter{}
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-ts-default",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS,
Adapter: "terminating",
Target: "codex",
SessionId: "",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error, got %q", resp.GetError())
}
if got := resp.GetSessionId(); got != "default" {
t.Fatalf("response SessionId: got %q want %q", got, "default")
}
if got := resp.GetResult()["session_id"]; got != "default" {
t.Fatalf("result[session_id]: got %q want %q", got, "default")
}
if got, ok := resp.GetResult()["state"]; !ok || (got != "connected" && got != "disconnected") {
t.Fatalf("result[state]: got %q, want \"connected\" or \"disconnected\"", got)
}
}
func TestOnCommandRequest_ErrorResponsesPreserveEnvelope(t *testing.T) {
cases := []struct {
name string
req *iop.NodeCommandRequest
adapter runtime.Adapter
wantErr string
}{
{
name: "missing adapter",
req: &iop.NodeCommandRequest{
RequestId: "req-missing",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "missing",
Target: "tgt",
SessionId: "sess-x",
},
adapter: nil,
wantErr: "not found",
},
{
name: "unsupported command type",
req: &iop.NodeCommandRequest{
RequestId: "req-unspec",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_UNSPECIFIED,
Adapter: "terminating",
Target: "tgt",
SessionId: "sess-y",
},
adapter: &terminatingAdapter{},
wantErr: "unsupported command type",
},
{
name: "adapter does not support commands",
req: &iop.NodeCommandRequest{
RequestId: "req-nosupport",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "terminating",
Target: "tgt",
SessionId: "sess-z",
},
adapter: &terminatingAdapter{},
wantErr: "does not support commands",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
adapters := make(map[string]runtime.Adapter)
if tc.adapter != nil {
adapters[tc.req.Adapter] = tc.adapter
}
router := &fixedRouter{adapterName: tc.req.Adapter, adapters: adapters}
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, tc.req)
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), tc.wantErr) {
t.Fatalf("error: got %q want contains %q", resp.GetError(), tc.wantErr)
}
if got := resp.GetRequestId(); got != tc.req.RequestId {
t.Fatalf("RequestId: got %q want %q", got, tc.req.RequestId)
}
if got := resp.GetType(); got != tc.req.Type {
t.Fatalf("Type: got %v want %v", got, tc.req.Type)
}
if got := resp.GetAdapter(); got != tc.req.Adapter {
t.Fatalf("Adapter: got %q want %q", got, tc.req.Adapter)
}
if got := resp.GetTarget(); got != tc.req.Target {
t.Fatalf("Target: got %q want %q", got, tc.req.Target)
}
if got := resp.GetSessionId(); got != tc.req.SessionId {
t.Fatalf("SessionId: got %q want %q", got, tc.req.SessionId)
}
})
}
}
// --- LookupAdapter / ambiguous error propagation tests ---
func TestOnCancel_TerminateSession_AmbiguousAdapterError(t *testing.T) {
ambigErr := fmt.Errorf("adapter \"ollama\" is ambiguous: matches instance keys [ollama@local ollama@dgx]; use an instance key")
router := &fixedRouter{
adapterName: "ollama",
adapters: make(map[string]runtime.Adapter),
lookupErrors: map[string]error{"ollama": ambigErr},
}
n, _ := makeNode(t, router)
err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{
Adapter: "ollama",
Target: "llama3",
SessionId: "sess-1",
Action: iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
})
if err == nil {
t.Fatal("expected error for ambiguous adapter lookup")
}
if !strings.Contains(err.Error(), "ambiguous") {
t.Errorf("expected 'ambiguous' in error, got %v", err)
}
if !strings.Contains(err.Error(), "instance key") {
t.Errorf("expected 'instance key' guidance in error, got %v", err)
}
}
func TestOnCancel_TerminateSession_ExactInstanceKey(t *testing.T) {
ta := &terminatingAdapter{}
router := &fixedRouter{
adapterName: "ollama@local",
adapters: map[string]runtime.Adapter{"ollama@local": ta},
}
n, _ := makeNode(t, router)
err := n.OnCancel(context.Background(), &transport.Session{}, &iop.CancelRequest{
Adapter: "ollama@local",
Target: "llama3",
SessionId: "sess-a",
Action: iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
})
if err != nil {
t.Fatalf("OnCancel with exact instance key: %v", err)
}
if atomic.LoadInt32(&ta.terminateCalls) != 1 {
t.Fatal("expected TerminateSession to be called once")
}
if ta.lastTarget != "llama3" || ta.lastSessionID != "sess-a" {
t.Errorf("terminate args: target=%q session=%q", ta.lastTarget, ta.lastSessionID)
}
}
func TestOnCommandRequest_Capabilities_AmbiguousAdapter(t *testing.T) {
ambigErr := fmt.Errorf("adapter \"cli\" is ambiguous: matches instance keys [cli@claude cli@codex]; use an instance key")
router := &fixedRouter{
adapterName: "cli",
adapters: make(map[string]runtime.Adapter),
lookupErrors: map[string]error{"cli": ambigErr},
}
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-ambig-cap",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "cli",
Target: "codex",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "ambiguous") {
t.Errorf("expected 'ambiguous' in response error, got %q", resp.GetError())
}
if !strings.Contains(resp.GetError(), "instance key") {
t.Errorf("expected 'instance key' guidance in response error, got %q", resp.GetError())
}
}
func TestOnCommandRequest_AdapterDispatch_AmbiguousAdapter(t *testing.T) {
ambigErr := fmt.Errorf("adapter \"cli\" is ambiguous: matches instance keys [cli@claude cli@codex]; use an instance key")
router := &fixedRouter{
adapterName: "cli",
adapters: make(map[string]runtime.Adapter),
lookupErrors: map[string]error{"cli": ambigErr},
}
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-ambig-cmd",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
Adapter: "cli",
Target: "codex",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if !strings.Contains(resp.GetError(), "ambiguous") {
t.Errorf("expected 'ambiguous' in response error, got %q", resp.GetError())
}
}
func TestOnCommandRequest_MultiAdapterCapabilities(t *testing.T) {
ika1 := &instanceKeyAdapter{instanceKey: "ollama@local"}
ika2 := &instanceKeyAdapter{instanceKey: "ollama@dgx"}
router := &fixedRouter{
adapters: map[string]runtime.Adapter{
"ollama@local": ika1,
"ollama@dgx": ika2,
},
}
n, _ := makeNode(t, router)
cases := []struct {
adapter string
wantInstKey string
}{
{"ollama@local", "ollama@local"},
{"ollama@dgx", "ollama@dgx"},
}
instKeys := make([]string, 0, len(cases))
for _, tc := range cases {
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-multi-" + tc.adapter,
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: tc.adapter,
Target: "llama3",
})
if err != nil {
t.Fatalf("CAPABILITIES %s: %v", tc.adapter, err)
}
if resp.GetError() != "" {
t.Fatalf("CAPABILITIES %s: unexpected error %q", tc.adapter, resp.GetError())
}
if got := resp.GetResult()["instance_key"]; got != tc.wantInstKey {
t.Errorf("CAPABILITIES %s: instance_key got %q want %q", tc.adapter, got, tc.wantInstKey)
}
if resp.GetAdapter() != tc.adapter {
t.Errorf("CAPABILITIES %s: response adapter got %q want %q", tc.adapter, resp.GetAdapter(), tc.adapter)
}
instKeys = append(instKeys, resp.GetResult()["instance_key"])
}
if instKeys[0] == instKeys[1] {
t.Errorf("expected distinct instance_key for distinct adapters, both got %q", instKeys[0])
}
}
func TestOnCommandRequest_Capabilities_ExactInstanceKey(t *testing.T) {
ika := &instanceKeyAdapter{instanceKey: "ollama@local"}
router := &fixedRouter{
adapterName: "ollama@local",
adapters: map[string]runtime.Adapter{"ollama@local": ika},
}
n, _ := makeNode(t, router)
resp, err := n.OnCommandRequest(context.Background(), &transport.Session{}, &iop.NodeCommandRequest{
RequestId: "req-exact-cap",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "ollama@local",
Target: "llama3",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetError() != "" {
t.Fatalf("expected no error for exact instance key, got %q", resp.GetError())
}
if got := resp.GetResult()["adapter"]; got != "ollama" {
t.Errorf("result[adapter]: got %q want ollama", got)
}
if got := resp.GetResult()["instance_key"]; got != "ollama@local" {
t.Errorf("result[instance_key]: got %q want ollama@local", got)
}
}
// --- concurrency gate tests ---
// slowAdapterUnlimited reports MaxConcurrency=0 (unlimited) and blocks until
// released, exposing a start channel so tests can synchronize.
type slowAdapterUnlimited struct {
started chan struct{}
release chan struct{}
}
func newSlowAdapterUnlimited() *slowAdapterUnlimited {
return &slowAdapterUnlimited{started: make(chan struct{}, 1), release: make(chan struct{})}
}
func (a *slowAdapterUnlimited) Name() string { return "slow" }
func (a *slowAdapterUnlimited) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "slow", MaxConcurrency: 0}, nil
}
func (a *slowAdapterUnlimited) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
select {
case a.started <- struct{}{}:
default:
}
select {
case <-a.release:
case <-ctx.Done():
return runtime.ErrRunCancelled
}
return nil
}
// queuedSlowAdapter is a configurable adapter used by concurrency safety tests.
// It advertises MaxConcurrency / MaxQueue / QueueTimeout via Capabilities, and
// gives per-run control: each Execute blocks until the test releases (or fails)
// that specific run id, so tests can observe execution ordering deterministically.
type queuedSlowAdapter struct {
name string
maxConcurrency int
maxQueue int
queueTimeout time.Duration
mu sync.Mutex
gates map[string]chan error // run_id → release signal carrying exec result
startSeq chan string // run_id pushed when Execute begins
}
func newQueuedSlowAdapter(name string, maxConcurrency, maxQueue int, queueTimeout time.Duration) *queuedSlowAdapter {
return &queuedSlowAdapter{
name: name,
maxConcurrency: maxConcurrency,
maxQueue: maxQueue,
queueTimeout: queueTimeout,
gates: make(map[string]chan error),
startSeq: make(chan string, 64),
}
}
func (a *queuedSlowAdapter) Name() string { return a.name }
func (a *queuedSlowAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: a.name,
MaxConcurrency: a.maxConcurrency,
MaxQueue: a.maxQueue,
QueueTimeoutMS: int(a.queueTimeout / time.Millisecond),
}, nil
}
func (a *queuedSlowAdapter) gateFor(runID string) chan error {
a.mu.Lock()
defer a.mu.Unlock()
ch, ok := a.gates[runID]
if !ok {
ch = make(chan error, 1)
a.gates[runID] = ch
}
return ch
}
func (a *queuedSlowAdapter) Execute(ctx context.Context, spec runtime.ExecutionSpec, _ runtime.EventSink) error {
gate := a.gateFor(spec.RunID)
a.startSeq <- spec.RunID
select {
case res := <-gate:
return res
case <-ctx.Done():
return runtime.ErrRunCancelled
}
}
// releaseRun lets the given run complete successfully.
func (a *queuedSlowAdapter) releaseRun(runID string) { a.gateFor(runID) <- nil }
// failRun lets the given run terminate with a non-cancel error.
func (a *queuedSlowAdapter) failRun(runID string) { a.gateFor(runID) <- errBoom }
// preRelease pre-arms a run so it completes the instant it starts executing.
func (a *queuedSlowAdapter) preRelease(runID string) { a.gateFor(runID) <- nil }
var errBoom = errors.New("boom")
// waitStarted blocks until the adapter reports the given run id started, or
// fails the test after a timeout.
func waitStarted(t *testing.T, a *queuedSlowAdapter, runID string) {
t.Helper()
deadline := time.After(2 * time.Second)
for {
select {
case got := <-a.startSeq:
if got == runID {
return
}
// Not the one we want; requeue for other waiters.
a.startSeq <- got
time.Sleep(2 * time.Millisecond)
case <-deadline:
t.Fatalf("run %s never started", runID)
}
}
}
// requireStatus asserts the run is stored with exactly the given status now.
func requireStatus(t *testing.T, st *store.Store, runID, status string) *store.RunRecord {
t.Helper()
return requireStatusEventually(t, st, runID, status)
}
// requireStatusEventually polls until the run reaches the expected status.
func requireStatusEventually(t *testing.T, st *store.Store, runID, status string) *store.RunRecord {
t.Helper()
deadline := time.After(2 * time.Second)
for {
run, err := st.GetRun(context.Background(), runID)
if err != nil {
t.Fatalf("GetRun(%s): %v", runID, err)
}
if run != nil && run.Status == status {
return run
}
select {
case <-deadline:
got := "<nil>"
if run != nil {
got = run.Status
}
t.Fatalf("run %s: expected status %q, got %q", runID, status, got)
case <-time.After(5 * time.Millisecond):
}
}
}
// TestOnRunRequest_OverDispatchSafetyRejectsWithoutQueue verifies that a run
// is rejected immediately with reason concurrency_unavailable when the concurrency limit is reached (no queue).
func TestOnRunRequest_OverDispatchSafetyRejectsWithoutQueue(t *testing.T) {
// capacity=1: one running is the ceiling.
sa := newQueuedSlowAdapter("slow", 1, 0, 0)
router := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
n, st := makeNodeWithConcurrency(t, router, 0) // global unlimited; adapter cap governs
hold := make(chan error, 1)
go func() {
hold <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "of-hold", Adapter: "slow", Target: "v1", Background: true,
})
}()
waitStarted(t, sa, "of-hold")
// Second run overflows the capacity and must be rejected immediately.
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "of-reject", Adapter: "slow", Target: "v1",
})
if !errors.Is(err, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("expected ErrConcurrencyLimitExceeded, got %v", err)
}
rec := requireStatus(t, st, "of-reject", "rejected")
if !strings.Contains(rec.Error, "concurrency unavailable") {
t.Fatalf("expected concurrency unavailable reason in error, got %q", rec.Error)
}
sa.releaseRun("of-hold")
if err := <-hold; err != nil {
t.Fatalf("of-hold: %v", err)
}
}
// TestOnRunRequest_PermitReleasedAfterCompletion verifies that after the first
// run completes, a subsequent run can acquire the slot immediately.
func TestOnRunRequest_PermitReleasedAfterCompletion(t *testing.T) {
sa := newQueuedSlowAdapter("slow", 1, 4, 0)
router := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
n, _ := makeNodeWithConcurrency(t, router, 1)
first := make(chan error, 1)
go func() {
first <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "rel-1", Adapter: "slow", Target: "v1",
})
}()
waitStarted(t, sa, "rel-1")
sa.releaseRun("rel-1")
if err := <-first; err != nil {
t.Fatalf("rel-1: %v", err)
}
// Second run must run immediately now that the slot is free.
sa.preRelease("rel-2")
if err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "rel-2", Adapter: "slow", Target: "v1",
}); err != nil {
t.Fatalf("rel-2 should succeed after slot released, got %v", err)
}
}
// TestConcurrencyLimit_Unlimited verifies that global=0 + adapter cap=0 imposes no cap.
func TestConcurrencyLimit_Unlimited(t *testing.T) {
const count = 5
shared := newSlowAdapterUnlimited() // MaxConcurrency=0 (unlimited)
router := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": shared}}
nd, _ := makeNodeWithConcurrency(t, router, 0) // global unlimited
errs := make(chan error, count)
for i := 0; i < count; i++ {
runID := fmt.Sprintf("run-unlimited-%d", i)
go func(id string) {
errs <- nd.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: id, Adapter: "slow", Target: "v1",
})
}(runID)
}
// Release all.
close(shared.release)
// All must complete without ErrConcurrencyLimitExceeded.
for i := 0; i < count; i++ {
if err := <-errs; errors.Is(err, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("unexpected concurrency rejection with limit=0: %v", err)
}
}
}
// TestConcurrencyLimit_RejectStoreAndEvent verifies that a rejected run
// is stored with terminal status "rejected" and an error message.
func TestConcurrencyLimit_RejectStoreAndEvent(t *testing.T) {
// capacity=1: any second concurrent run overflows immediately.
sa := newQueuedSlowAdapter("slow", 1, 0, 0)
router := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
nd, st := makeNodeWithConcurrency(t, router, 0)
errc := make(chan error, 1)
go func() {
errc <- nd.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-hold", Adapter: "slow", Target: "v1", Background: true,
})
}()
waitStarted(t, sa, "run-hold")
// Second run must be rejected.
err := nd.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-rejected", Adapter: "slow", Target: "v1",
})
if !errors.Is(err, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("expected ErrConcurrencyLimitExceeded, got %v", err)
}
// Rejected run must be stored with status "rejected".
run, err := st.GetRun(context.Background(), "run-rejected")
if err != nil {
t.Fatalf("GetRun: %v", err)
}
if run == nil {
t.Fatal("expected rejected run to be stored")
}
if run.Status != "rejected" {
t.Fatalf("expected status %q, got %q", "rejected", run.Status)
}
if run.Error == "" {
t.Fatal("expected non-empty error message for rejected run")
}
sa.releaseRun("run-hold")
if err := <-errc; err != nil {
t.Fatalf("run-hold: %v", err)
}
}
// TestNodeConfigRefreshWithoutApplyManagerReportsRestartRequired verifies that
// a Node without a live apply manager responds restart_required for any refresh
// request that carries changed paths, and applied for a no-op (empty) request.
func TestNodeConfigRefreshWithoutApplyManagerReportsRestartRequired(t *testing.T) {
router := &fixedRouter{adapterName: "test", adapters: map[string]runtime.Adapter{"test": &countingAdapter{}}}
n, _ := makeNode(t, router)
// Non-empty changed_paths → restart_required.
resp, err := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "req-restart",
ChangedPaths: []string{"nodes.0.providers.0.capacity"},
})
if err != nil {
t.Fatalf("OnConfigRefresh: %v", err)
}
if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_RESTART_REQUIRED {
t.Fatalf("expected restart_required, got %v", resp.GetStatus())
}
if resp.GetRequestId() != "req-restart" {
t.Fatalf("expected request_id=req-restart, got %q", resp.GetRequestId())
}
if len(resp.GetRestartRequiredPaths()) == 0 {
t.Fatal("expected restart_required_paths to be populated")
}
// No config and no changed_paths → applied (no-op).
resp2, err := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "req-noop",
})
if err != nil {
t.Fatalf("OnConfigRefresh noop: %v", err)
}
if resp2.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected applied for no-op, got %v", resp2.GetStatus())
}
}
func TestNodeConfigRefreshAppliesOpenAICompatCapacity(t *testing.T) {
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: "http://localhost:8080",
Capacity: 2,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 1, io.Discard, zap.NewNop(), set)
// Check initial capabilities
resp, err := n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-1",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp.GetResult()["capacity"] != "2" {
t.Fatalf("expected capacity 2, got %s", resp.GetResult()["capacity"])
}
// Refresh payload with new capacity
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: "http://localhost:8080",
Capacity: 5,
},
},
},
},
}
refreshResp, err := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if err != nil {
t.Fatalf("OnConfigRefresh: %v", err)
}
if refreshResp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected applied status, got %v", refreshResp.GetStatus())
}
// Check updated capabilities
resp2, err := n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-2",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if resp2.GetResult()["capacity"] != "5" {
t.Fatalf("expected capacity 5, got %s", resp2.GetResult()["capacity"])
}
}
func TestNodeConfigRefreshUsesUpdatedOpenAICompatEndpoint(t *testing.T) {
var oldCalls int32
var newCalls int32
serverOld := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&oldCalls, 1)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
}))
defer serverOld.Close()
serverNew := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
atomic.AddInt32(&newCalls, 1)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
}))
defer serverNew.Close()
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: serverOld.URL,
Capacity: 2,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 1, io.Discard, zap.NewNop(), set)
// Call capabilities command -> triggers probe (requests serverOld)
_, err = n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-1",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
Target: "gpt-4o",
})
if err != nil {
t.Fatalf("OnCommandRequest: %v", err)
}
if atomic.LoadInt32(&oldCalls) != 2 {
t.Fatalf("expected 2 calls to serverOld, got %d", oldCalls)
}
if atomic.LoadInt32(&newCalls) != 0 {
t.Fatalf("expected 0 calls to serverNew, got %d", newCalls)
}
// Refresh with new endpoint
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: serverNew.URL,
Capacity: 2,
},
},
},
},
}
refreshResp, err := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if err != nil {
t.Fatalf("OnConfigRefresh: %v", err)
}
if refreshResp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected applied status, got %v", refreshResp.GetStatus())
}
// Call capabilities command again -> triggers probe (requests serverNew)
_, err = n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-2",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
Target: "gpt-4o",
})
if err != nil {
t.Fatalf("OnCommandRequest 2: %v", err)
}
if atomic.LoadInt32(&oldCalls) != 2 {
t.Fatalf("expected still 2 calls to serverOld, got %d", oldCalls)
}
if atomic.LoadInt32(&newCalls) != 2 {
t.Fatalf("expected 2 calls to serverNew, got %d", newCalls)
}
}
func TestNodeConfigRefreshUsesUpdatedOpenAICompatHeaders(t *testing.T) {
var observedHeader string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
observedHeader = r.Header.Get("X-Test-Key")
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
}))
defer server.Close()
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Headers: map[string]string{"X-Test-Key": "initial-value"},
Capacity: 2,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 1, io.Discard, zap.NewNop(), set)
// Trigger request -> should see initial header
_, _ = n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-1",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
Target: "gpt-4o",
})
if observedHeader != "initial-value" {
t.Fatalf("expected header 'initial-value', got %q", observedHeader)
}
// Refresh with new header
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Headers: map[string]string{"X-Test-Key": "updated-value"},
Capacity: 2,
},
},
},
},
}
_, _ = n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
// Trigger request again -> should see updated header
_, _ = n.OnCommandRequest(context.Background(), nil, &iop.NodeCommandRequest{
RequestId: "cmd-2",
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
Adapter: "openai",
Target: "gpt-4o",
})
if observedHeader != "updated-value" {
t.Fatalf("expected header 'updated-value', got %q", observedHeader)
}
}
func TestNodeConfigRefreshUpdatesExistingAdapterGateCapacity(t *testing.T) {
// mock HTTP server that can be blocked or unblocked
blockChan := make(chan struct{})
defer func() {
select {
case <-blockChan:
default:
close(blockChan)
}
}()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if strings.Contains(r.URL.Path, "/v1/models") {
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
return
}
// Block completion requests until blockChan is closed or read
<-blockChan
// Write a dummy SSE chat completion response
_, _ = w.Write([]byte("data: {\"choices\": [{\"delta\": {\"content\": \"hello\"}}]}\n\ndata: [DONE]\n\n"))
}))
defer server.Close()
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 1,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), set)
inputMap := map[string]interface{}{
"prompt": "hello",
}
inputStruct, err := structpb.NewStruct(inputMap)
if err != nil {
t.Fatalf("failed to create input struct: %v", err)
}
// First request: should run and block inside the mock server
errChan1 := make(chan error, 1)
go func() {
errChan1 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-1",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
}()
// Wait a bit to ensure the first request is indeed running and blocked
time.Sleep(50 * time.Millisecond)
// Second request: should fail because capacity is 1
err2 := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-2",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
if err2 == nil {
t.Fatal("expected second request to fail due to concurrency limit")
}
if !errors.Is(err2, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("expected ErrConcurrencyLimitExceeded, got: %v", err2)
}
// Refresh config to increase capacity to 2
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 2,
},
},
},
},
}
refreshResp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if refreshResp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", refreshResp.Status)
}
// Third request: should now succeed (and block) instead of getting rejected
errChan3 := make(chan error, 1)
go func() {
errChan3 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-3",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
}()
// Wait a bit and check that third request hasn't errored immediately (it should be blocking)
time.Sleep(50 * time.Millisecond)
select {
case err3 := <-errChan3:
t.Fatalf("expected third request to block, but it failed immediately: %v", err3)
default:
// It is blocking as expected!
}
// Unblock mock server for all requests
close(blockChan)
// Wait for run-1 and run-3 to finish successfully
if err1 := <-errChan1; err1 != nil {
t.Fatalf("run-1 failed: %v", err1)
}
if err3 := <-errChan3; err3 != nil {
t.Fatalf("run-3 failed: %v", err3)
}
}
type lifecycleTestAdapter struct {
name string
stopCalls int32
started chan struct{}
blockChan chan struct{}
}
func (a *lifecycleTestAdapter) Name() string { return a.name }
func (a *lifecycleTestAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: a.name,
MaxConcurrency: 1,
}, nil
}
func (a *lifecycleTestAdapter) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
close(a.started)
select {
case <-a.blockChan:
case <-ctx.Done():
}
return nil
}
func (a *lifecycleTestAdapter) Start(_ context.Context) error {
return nil
}
func (a *lifecycleTestAdapter) Stop(_ context.Context) error {
atomic.AddInt32(&a.stopCalls, 1)
return nil
}
func TestNodeConfigRefreshDoesNotStopOldRegistryWithActiveRun(t *testing.T) {
reg1 := adapters.NewRegistry()
oldAdapter := &lifecycleTestAdapter{
name: "my-adapter",
started: make(chan struct{}),
blockChan: make(chan struct{}),
}
reg1.RegisterKeyed("my-adapter", "lifecycle", oldAdapter)
initialConfigSet := &adapters.ConfigSet{
Registry: reg1,
Items: map[string]adapters.ConfigItem{
"my-adapter": {
Key: "my-adapter",
Type: "lifecycle",
Fingerprint: "fingerprint-1",
},
},
}
rtr := router.New(reg1, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), initialConfigSet)
// Start active run on my-adapter
errChan := make(chan error, 1)
go func() {
errChan <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-1",
Adapter: "my-adapter",
Background: false,
})
}()
// Wait for run to start
<-oldAdapter.started
// Refresh to new config: "my-adapter" is removed (or updated)
// We'll use mock type in the new config payload
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "new-mock",
Type: "mock",
Enabled: true,
},
},
}
refreshResp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if refreshResp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", refreshResp.Status)
}
// Verify that the old adapter's Stop was NOT called because of the active run
if calls := atomic.LoadInt32(&oldAdapter.stopCalls); calls > 0 {
t.Fatalf("expected old adapter Stop calls to be 0, got %d", calls)
}
// Unblock the run
close(oldAdapter.blockChan)
if err := <-errChan; err != nil {
t.Fatalf("active run failed: %v", err)
}
// Verify that new requests now resolve to the new registry (new-mock should resolve, my-adapter should fail)
_, _, err = rtr.ResolveAdapter(context.Background(), runtime.RunRequest{
RunID: "run-2",
Adapter: "new-mock",
})
if err != nil {
t.Fatalf("expected new-mock to resolve on the new registry, got: %v", err)
}
_, _, err = rtr.ResolveAdapter(context.Background(), runtime.RunRequest{
RunID: "run-3",
Adapter: "my-adapter",
})
if err == nil {
t.Fatal("expected my-adapter lookup to fail on the new registry")
}
}
func TestNodeConfigRefreshDoesNotStopUnchangedActiveAdapterWhenSiblingChanges(t *testing.T) {
blockChan := make(chan struct{})
defer func() {
select {
case <-blockChan:
default:
close(blockChan)
}
}()
startedChan := make(chan struct{}, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if strings.Contains(r.URL.Path, "/v1/models") {
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
return
}
select {
case startedChan <- struct{}{}:
default:
}
<-blockChan
_, _ = w.Write([]byte("data: {\"choices\": [{\"delta\": {\"content\": \"hello\"}}]}\n\ndata: [DONE]\n\n"))
}))
defer server.Close()
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 2,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
sibling := &lifecycleTestAdapter{
name: "sibling-adapter",
started: make(chan struct{}),
blockChan: make(chan struct{}),
}
set.Registry.RegisterKeyed("sibling-adapter", "lifecycle", sibling)
set.Items["sibling-adapter"] = adapters.ConfigItem{
Key: "sibling-adapter",
Type: "lifecycle",
Fingerprint: "fingerprint-sibling-1",
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), set)
inputMap := map[string]interface{}{
"prompt": "hello",
}
inputStruct, err := structpb.NewStruct(inputMap)
if err != nil {
t.Fatalf("failed to create input struct: %v", err)
}
errChan := make(chan error, 1)
go func() {
errChan <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-openai-active",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
}()
<-startedChan
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 2,
},
},
},
},
}
refreshResp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if refreshResp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", refreshResp.Status)
}
if calls := atomic.LoadInt32(&sibling.stopCalls); calls > 0 {
t.Fatalf("expected sibling-adapter Stop calls to be 0, got %d", calls)
}
close(blockChan)
if err := <-errChan; err != nil {
t.Fatalf("active run failed: %v", err)
}
}
type blockingCapsAdapter struct {
name string
capsStarted chan struct{}
capsBlock chan struct{}
}
func (a *blockingCapsAdapter) Name() string { return a.name }
func (a *blockingCapsAdapter) Capabilities(ctx context.Context) (runtime.Capabilities, error) {
select {
case a.capsStarted <- struct{}{}:
default:
}
<-a.capsBlock
return runtime.Capabilities{
AdapterName: a.name,
MaxConcurrency: 1,
}, nil
}
func (a *blockingCapsAdapter) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return nil
}
func (a *blockingCapsAdapter) Start(_ context.Context) error { return nil }
func (a *blockingCapsAdapter) Stop(_ context.Context) error { return nil }
func TestNodeConfigRefreshWaitsForResolvedRunRegistrationBeforeStoppingOldRegistry(t *testing.T) {
reg1 := adapters.NewRegistry()
blockingAdapter := &blockingCapsAdapter{
name: "blocking-adapter",
capsStarted: make(chan struct{}, 1),
capsBlock: make(chan struct{}),
}
reg1.RegisterKeyed("blocking-adapter", "lifecycle", blockingAdapter)
initialConfigSet := &adapters.ConfigSet{
Registry: reg1,
Items: map[string]adapters.ConfigItem{
"blocking-adapter": {
Key: "blocking-adapter",
Type: "lifecycle",
Fingerprint: "fingerprint-1",
},
},
}
rtr := router.New(reg1, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), initialConfigSet)
errChan := make(chan error, 1)
go func() {
errChan <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-1",
Adapter: "blocking-adapter",
Background: false,
})
}()
<-blockingAdapter.capsStarted
refreshErrChan := make(chan error, 1)
refreshRespChan := make(chan *iop.NodeConfigRefreshResponse, 1)
go func() {
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "new-mock",
Type: "mock",
Enabled: true,
},
},
}
resp, err := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
refreshRespChan <- resp
refreshErrChan <- err
}()
time.Sleep(50 * time.Millisecond)
select {
case resp := <-refreshRespChan:
t.Fatalf("expected refresh to block, but it completed: %v", resp)
default:
// Refresh is successfully blocking!
}
close(blockingAdapter.capsBlock)
refreshErr := <-refreshErrChan
if refreshErr != nil {
t.Fatalf("refresh failed: %v", refreshErr)
}
resp := <-refreshRespChan
if resp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", resp.Status)
}
runErr := <-errChan
if runErr != nil {
t.Fatalf("run failed: %v", runErr)
}
}
func TestNodeConfigRefreshDecreasesExistingAdapterGateCapacity(t *testing.T) {
blockChan := make(chan struct{})
defer func() {
select {
case <-blockChan:
default:
close(blockChan)
}
}()
startedChan := make(chan struct{}, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
if strings.Contains(r.URL.Path, "/v1/models") {
_, _ = w.Write([]byte(`{"data": [{"id": "gpt-4o"}]}`))
return
}
select {
case startedChan <- struct{}{}:
default:
}
<-blockChan
_, _ = w.Write([]byte("data: {\"choices\": [{\"delta\": {\"content\": \"hello\"}}]}\n\ndata: [DONE]\n\n"))
}))
defer server.Close()
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 2,
},
},
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), set)
inputMap := map[string]interface{}{
"prompt": "hello",
}
inputStruct, err := structpb.NewStruct(inputMap)
if err != nil {
t.Fatalf("failed to create input struct: %v", err)
}
errChan1 := make(chan error, 1)
go func() {
errChan1 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-1",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
}()
<-startedChan
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "openai",
Type: "openai_compat",
Enabled: true,
Config: &iop.AdapterConfig_OpenaiCompat{
OpenaiCompat: &iop.OpenAICompatAdapterConfig{
Provider: "openai",
Endpoint: server.URL,
Capacity: 1,
},
},
},
},
}
refreshResp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: refreshPayload,
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if refreshResp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", refreshResp.Status)
}
ctx2, cancel2 := context.WithTimeout(context.Background(), 200*time.Millisecond)
defer cancel2()
err2 := n.OnRunRequest(ctx2, &transport.Session{}, &iop.RunRequest{
RunId: "run-2",
Adapter: "openai",
Target: "gpt-4o",
Background: false,
Input: inputStruct,
})
if err2 == nil {
t.Fatal("expected second request to fail due to concurrency limit")
}
if !errors.Is(err2, node.ErrConcurrencyLimitExceeded) {
t.Fatalf("expected ErrConcurrencyLimitExceeded, got: %v", err2)
}
close(blockChan)
if err1 := <-errChan1; err1 != nil {
t.Fatalf("run-1 failed: %v", err1)
}
}
// --- NCDRAIN-1: deferred old-registry stop drains after active runs finish ---
func TestNodeConfigRefreshDeferredStopRunsAfterActiveRunCompletes(t *testing.T) {
reg1 := adapters.NewRegistry()
oldAdapter := &lifecycleTestAdapter{
name: "my-adapter",
started: make(chan struct{}),
blockChan: make(chan struct{}),
}
reg1.RegisterKeyed("my-adapter", "lifecycle", oldAdapter)
initialConfigSet := &adapters.ConfigSet{
Registry: reg1,
Items: map[string]adapters.ConfigItem{
"my-adapter": {Key: "my-adapter", Type: "lifecycle", Fingerprint: "fingerprint-1"},
},
}
rtr := router.New(reg1, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), initialConfigSet)
errChan := make(chan error, 1)
go func() {
errChan <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-1",
Adapter: "my-adapter",
})
}()
<-oldAdapter.started
refreshResp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-1",
Config: &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{{Name: "new-mock", Type: "mock", Enabled: true}},
},
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if refreshResp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", refreshResp.Status)
}
// Old registry stop must be deferred while the run is still active.
if calls := atomic.LoadInt32(&oldAdapter.stopCalls); calls != 0 {
t.Fatalf("expected old adapter Stop to be deferred (0 calls), got %d", calls)
}
// Drain the active run; the deferred stop should fire afterwards.
close(oldAdapter.blockChan)
if err := <-errChan; err != nil {
t.Fatalf("active run failed: %v", err)
}
deadline := time.After(2 * time.Second)
for atomic.LoadInt32(&oldAdapter.stopCalls) == 0 {
select {
case <-deadline:
t.Fatal("deferred old registry stop never ran after the active run completed")
default:
time.Sleep(5 * time.Millisecond)
}
}
}
// --- NCDRAIN-2: runtime concurrency refresh remains admission no-op ---
func TestConfigRefreshRuntimeConcurrencyDoesNotAffectAdmission(t *testing.T) {
// Adapter is unlimited (MaxConcurrency=0), so the adapter gate imposes no limit.
sa := newQueuedSlowAdapter("slow", 0, 0, 0)
rtr := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
// Global concurrency is set to 1, but it must NOT affect admission.
// All concurrent runs should succeed because the adapter gate is unlimited.
n := node.New("test-node", rtr, st, 1, io.Discard, zap.NewNop(), nil)
errCh1 := make(chan error, 1)
go func() {
errCh1 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1", Adapter: "slow"})
}()
waitStarted(t, sa, "run-1")
// Second run should NOT be rejected: global concurrency is no longer used for admission.
errCh2 := make(chan error, 1)
go func() {
errCh2 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-2", Adapter: "slow"})
}()
waitStarted(t, sa, "run-2")
// Config refresh with runtime concurrency must succeed and store metadata.
resp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-conc",
Config: &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 2}},
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if resp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", resp.Status)
}
// Third run must also succeed: runtime concurrency is metadata only.
errCh3 := make(chan error, 1)
go func() {
errCh3 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-3", Adapter: "slow"})
}()
waitStarted(t, sa, "run-3")
sa.releaseRun("run-1")
sa.releaseRun("run-2")
sa.releaseRun("run-3")
if err := <-errCh1; err != nil {
t.Fatalf("run-1 failed: %v", err)
}
if err := <-errCh2; err != nil {
t.Fatalf("run-2 failed: %v", err)
}
if err := <-errCh3; err != nil {
t.Fatalf("run-3 failed: %v", err)
}
}
func TestConfigRefreshConcurrencyDecreaseDoesNotAffectAdmission(t *testing.T) {
sa := newQueuedSlowAdapter("slow", 0, 0, 0)
rtr := &fixedRouter{adapterName: "slow", adapters: map[string]runtime.Adapter{"slow": sa}}
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 2, io.Discard, zap.NewNop(), nil)
errCh1 := make(chan error, 1)
go func() {
errCh1 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-1", Adapter: "slow"})
}()
waitStarted(t, sa, "run-1")
errCh2 := make(chan error, 1)
go func() {
errCh2 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-2", Adapter: "slow"})
}()
waitStarted(t, sa, "run-2")
// Config refresh decreases runtime concurrency from 2 to 1.
// This must NOT affect admission: the adapter gate is unlimited.
resp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-dec",
Config: &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if resp.Status != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED {
t.Fatalf("expected refresh status APPLIED, got %v", resp.Status)
}
// Third run must be admitted: global concurrency is no longer used for admission.
errCh3 := make(chan error, 1)
go func() {
errCh3 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{RunId: "run-3", Adapter: "slow"})
}()
waitStarted(t, sa, "run-3")
// All three runs complete normally.
sa.releaseRun("run-1")
sa.releaseRun("run-2")
sa.releaseRun("run-3")
if err := <-errCh1; err != nil {
t.Fatalf("run-1 should complete normally, got %v", err)
}
if err := <-errCh2; err != nil {
t.Fatalf("run-2 should complete normally, got %v", err)
}
if err := <-errCh3; err != nil {
t.Fatalf("run-3 should complete normally, got %v", err)
}
}
// --- NCDRAIN-3: TestConfigRefreshFailedAdapterStartDoesNotRollbackConfig
// adapter gate preservation on failed adapter registry start ---
// TestConfigRefreshFailedAdapterStartDoesNotRollbackConfig verifies that when
// adapter registry Start fails during a config refresh, the existing adapter
// state and gates are preserved and the node returns FAILED status.
func TestConfigRefreshFailedAdapterStartDoesNotRollbackConfig(t *testing.T) {
initialPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "blocking",
Type: "mock",
Enabled: true,
},
},
}
set, err := adapters.BuildConfigSet(initialPayload, zap.NewNop())
if err != nil {
t.Fatalf("BuildConfigSet: %v", err)
}
rtr := router.New(set.Registry, zap.NewNop())
st, err := store.New(":memory:", zap.NewNop())
if err != nil {
t.Fatalf("store: %v", err)
}
t.Cleanup(func() { _ = st.Close() })
n := node.New("test-node", rtr, st, 0, io.Discard, zap.NewNop(), set)
// Add a blocking adapter that occupies the adapter slot.
blockingAdapt := newBlockingAdapterWithStart()
set.Registry.RegisterKeyed("blocking", "mock", blockingAdapt)
// Start an active run on the blocking adapter.
errCh1 := make(chan error, 1)
go func() {
errCh1 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-hold",
Adapter: "blocking",
Background: true,
})
}()
// Wait for the run to start.
select {
case <-blockingAdapt.started:
case <-time.After(2 * time.Second):
t.Fatal("blocking adapter never started")
}
// Second run on the same adapter must succeed: adapter gate is unlimited.
errCh2 := make(chan error, 1)
go func() {
errCh2 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-concurrent",
Adapter: "blocking",
Background: true,
})
}()
// Wait until the blocking adapter's Execute has been called (both runs started).
// The first run's started signal was consumed above, so we wait for the second run.
select {
case <-blockingAdapt.started:
case <-time.After(2 * time.Second):
t.Fatal("blocking adapter second Execute never started")
}
// Now send a refresh that includes a failing CLI adapter.
// The CLI adapter with a non-existent command path will cause Registry.Start to fail.
refreshPayload := &iop.NodeConfigPayload{
Adapters: []*iop.AdapterConfig{
{
Name: "blocking",
Type: "mock",
Enabled: true,
},
{
Name: "cli",
Type: "cli",
Enabled: true,
Config: &iop.AdapterConfig_Cli{
Cli: &iop.CLIAdapterConfig{
Profiles: map[string]*iop.CLIProfileConfig{
"failprofile": {
Command: "/nonexistent/path/to/binary",
Args: []string{"--fail"},
Persistent: true,
},
},
},
},
},
},
}
resp, refreshErr := n.OnConfigRefresh(context.Background(), nil, &iop.NodeConfigRefreshRequest{
RequestId: "refresh-fail",
Config: refreshPayload,
})
if refreshErr != nil {
t.Fatalf("OnConfigRefresh failed: %v", refreshErr)
}
if resp.GetStatus() != iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_FAILED {
t.Fatalf("expected FAILED status, got %v", resp.GetStatus())
}
// Clean up: unblock the held runs.
blockingAdapt.blockChan <- struct{}{}
blockingAdapt.blockChan <- struct{}{}
if err := <-errCh1; err != nil && !errors.Is(err, runtime.ErrRunCancelled) {
t.Fatalf("run-hold failed: %v", err)
}
if err := <-errCh2; err != nil && !errors.Is(err, runtime.ErrRunCancelled) {
t.Fatalf("run-concurrent failed: %v", err)
}
// 추가 검증: failed refresh 뒤에도 기존 blocking adapter가 계속 사용 가능함을 확인한다.
// old registry/router state가 보존되었음을 확인한다.
errCh3 := make(chan error, 1)
go func() {
errCh3 <- n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-post-verify",
Adapter: "blocking",
Background: true,
})
}()
select {
case <-blockingAdapt.started:
case <-time.After(2 * time.Second):
t.Fatal("blocking adapter never started for post-verify run")
}
blockingAdapt.blockChan <- struct{}{}
if err := <-errCh3; err != nil && !errors.Is(err, runtime.ErrRunCancelled) {
t.Fatalf("post-verify run failed: %v", err)
}
}
type blockingAdapterWithStart struct {
started chan struct{}
blockChan chan struct{}
}
func newBlockingAdapterWithStart() *blockingAdapterWithStart {
return &blockingAdapterWithStart{
started: make(chan struct{}, 3),
blockChan: make(chan struct{}, 3),
}
}
func (a *blockingAdapterWithStart) Name() string { return "blocking" }
func (a *blockingAdapterWithStart) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{
AdapterName: "blocking",
MaxConcurrency: 0, // unlimited per-adapter
}, nil
}
func (a *blockingAdapterWithStart) Execute(ctx context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
a.started <- struct{}{}
<-a.blockChan
return nil
}
func (a *blockingAdapterWithStart) Start(_ context.Context) error {
return nil
}
func (a *blockingAdapterWithStart) Stop(_ context.Context) error {
return nil
}
// --- terminal event guarantee tests (BUG-1) ---
// countingAdapterWithNoTerminal returns success without emitting any terminal event.
// It reuses countingAdapter's Execute but does not call sink at all.
type countingAdapterNoTerminal struct {
executeCalls int32
}
func (a *countingAdapterNoTerminal) Name() string { return "no-terminal" }
func (a *countingAdapterNoTerminal) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "no-terminal", MaxConcurrency: 1}, nil
}
func (a *countingAdapterNoTerminal) Execute(_ context.Context, spec runtime.ExecutionSpec, _ runtime.EventSink) error {
atomic.AddInt32(&a.executeCalls, 1)
return nil
}
// TestOnRunRequestEmitsCompleteWhenAdapterReturnsWithoutTerminal verifies that when an
// adapter returns nil (success) without emitting any terminal event, Node synthesizes
// a complete event so Edge can release the in_flight slot.
func TestOnRunRequestEmitsCompleteWhenAdapterReturnsWithoutTerminal(t *testing.T) {
adapter := &countingAdapterNoTerminal{}
router := &fixedRouter{adapterName: "no-terminal", adapters: make(map[string]runtime.Adapter)}
router.adapters["no-terminal"] = adapter
n, st := makeNode(t, router)
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-no-term",
Adapter: "no-terminal",
Target: "v1",
})
if err != nil {
t.Fatalf("run request: %v", err)
}
// Store should show completed (synthesized terminal event processed).
run, err := st.GetRun(context.Background(), "run-no-term")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run == nil || run.Status != "completed" {
t.Fatalf("expected completed status, got %q", run.Status)
}
if atomic.LoadInt32(&adapter.executeCalls) != 1 {
t.Fatalf("expected 1 execute call, got %d", adapter.executeCalls)
}
}
// TestOnRunRequestEmitsErrorWhenAdapterReturnsErrorWithoutTerminal verifies that when an
// adapter returns a non-cancel error without emitting a terminal event, Node synthesizes
// an error event.
func TestOnRunRequestEmitsErrorWhenAdapterReturnsErrorWithoutTerminal(t *testing.T) {
failing := &failingAdapterNoTerminal{err: fmt.Errorf("stream closed")}
router := &fixedRouter{adapterName: "no-term-err", adapters: make(map[string]runtime.Adapter)}
router.adapters["no-term-err"] = failing
n, st := makeNode(t, router)
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-no-term-err",
Adapter: "no-term-err",
Target: "v1",
})
if err == nil {
t.Fatal("expected error from OnRunRequest")
}
if !strings.Contains(err.Error(), "stream closed") {
t.Fatalf("expected 'stream closed' in error, got %v", err)
}
// Store should show failed.
run, err := st.GetRun(context.Background(), "run-no-term-err")
if err != nil {
t.Fatalf("get run: %v", err)
}
if run == nil || run.Status != "failed" {
t.Fatalf("expected failed status, got %q", run.Status)
}
}
// failingAdapterNoTerminal returns an error without emitting terminal events.
type failingAdapterNoTerminal struct {
err error
}
func (a *failingAdapterNoTerminal) Name() string { return "no-term-err" }
func (a *failingAdapterNoTerminal) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "no-term-err", MaxConcurrency: 1}, nil
}
func (a *failingAdapterNoTerminal) Execute(_ context.Context, _ runtime.ExecutionSpec, _ runtime.EventSink) error {
return a.err
}
// TestResolveAdapterErrorObservedByEdge verifies that when ResolveAdapter fails,
// Node returns an error. The Edge-observable RunEvent delivery is validated by
// the integration test TestIntegration_ResolveAdapterErrorObservedByEdge in
// node_concurrency_integration_test.go.
func TestResolveAdapterErrorObservedByEdge(t *testing.T) {
router := &errorRouter{err: fmt.Errorf("adapter not found")}
n, _ := makeNode(t, router)
err := n.OnRunRequest(context.Background(), &transport.Session{}, &iop.RunRequest{
RunId: "run-resolve-fail",
Adapter: "nonexistent",
Target: "v1",
})
if err == nil {
t.Fatal("expected error from OnRunRequest on ResolveAdapter failure")
}
if !strings.Contains(err.Error(), "node: resolve:") {
t.Fatalf("expected resolve prefix, got %v", err)
}
}
type mockTunnelAdapter struct {
countingAdapter
t *testing.T
expectedReq runtime.ProviderTunnelRequest
respondErr error
}
func (a *mockTunnelAdapter) Name() string { return "openai_compat" }
func (a *mockTunnelAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "openai_compat", Targets: []string{"qwen"}}, nil
}
func (a *mockTunnelAdapter) TunnelProvider(ctx context.Context, req runtime.ProviderTunnelRequest, sink runtime.ProviderTunnelSink) error {
if req.RunID != a.expectedReq.RunID || req.TunnelID != a.expectedReq.TunnelID {
a.t.Errorf("unexpected tunnel req: %+v", req)
}
if a.respondErr != nil {
return a.respondErr
}
err := sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: 0,
Kind: runtime.ProviderTunnelFrameKindResponseStart,
StatusCode: 200,
Headers: map[string]string{"Content-Type": "application/json"},
Timestamp: time.Now(),
})
if err != nil {
return err
}
err = sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: 1,
Kind: runtime.ProviderTunnelFrameKindBody,
Body: []byte(`{"choices":[{"delta":{"content":"ok"}}]}`),
Timestamp: time.Now(),
})
if err != nil {
return err
}
err = sink.EmitTunnelFrame(ctx, runtime.ProviderTunnelFrame{
RunID: req.RunID,
TunnelID: req.TunnelID,
Sequence: 2,
Kind: runtime.ProviderTunnelFrameKindEnd,
End: true,
Timestamp: time.Now(),
})
return err
}
type cancelAwareTunnelAdapter struct {
countingAdapter
started chan struct{}
observedCancel chan struct{}
}
func (a *cancelAwareTunnelAdapter) Name() string { return "openai_compat" }
func (a *cancelAwareTunnelAdapter) Capabilities(_ context.Context) (runtime.Capabilities, error) {
return runtime.Capabilities{AdapterName: "openai_compat", Targets: []string{"qwen"}}, nil
}
func (a *cancelAwareTunnelAdapter) TunnelProvider(ctx context.Context, _ runtime.ProviderTunnelRequest, _ runtime.ProviderTunnelSink) error {
close(a.started)
<-ctx.Done()
close(a.observedCancel)
return ctx.Err()
}
func TestNodeOnProviderTunnelRequest_Success(t *testing.T) {
mta := &mockTunnelAdapter{
t: t,
expectedReq: runtime.ProviderTunnelRequest{
RunID: "run-tunnel-1",
TunnelID: "tunnel-1",
},
}
router := &fixedRouter{adapterName: "openai_compat", adapters: make(map[string]runtime.Adapter)}
router.adapters["openai_compat"] = mta
n, _ := makeNode(t, router)
req := &iop.ProviderTunnelRequest{
RunId: "run-tunnel-1",
TunnelId: "tunnel-1",
Adapter: "openai_compat",
Target: "qwen",
Method: "POST",
Path: "/v1/chat/completions",
}
err := n.OnProviderTunnelRequest(context.Background(), nil, req)
if err != nil {
t.Fatalf("OnProviderTunnelRequest failed: %v", err)
}
}
func TestNodeOnProviderTunnelRequest_CancelRequestCancelsProviderContext(t *testing.T) {
adapter := &cancelAwareTunnelAdapter{
started: make(chan struct{}),
observedCancel: make(chan struct{}),
}
router := &fixedRouter{adapterName: "openai_compat", adapters: map[string]runtime.Adapter{"openai_compat": adapter}}
n, _ := makeNode(t, router)
req := &iop.ProviderTunnelRequest{
RunId: "run-tunnel-cancel",
TunnelId: "tunnel-cancel",
Adapter: "openai_compat",
Target: "qwen",
Method: "POST",
Path: "/v1/chat/completions",
TimeoutSec: 30,
}
errCh := make(chan error, 1)
go func() {
errCh <- n.OnProviderTunnelRequest(context.Background(), nil, req)
}()
select {
case <-adapter.started:
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for tunnel adapter to start")
}
err := n.OnCancel(context.Background(), nil, &iop.CancelRequest{
RunId: "run-tunnel-cancel",
Action: iop.CancelAction_CANCEL_ACTION_CANCEL_RUN,
})
if err != nil {
t.Fatalf("OnCancel failed: %v", err)
}
select {
case <-adapter.observedCancel:
case <-time.After(2 * time.Second):
t.Fatal("tunnel adapter did not observe cancel request")
}
select {
case err := <-errCh:
if !errors.Is(err, context.Canceled) {
t.Fatalf("OnProviderTunnelRequest error = %v, want context.Canceled", err)
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for tunnel request to finish after cancel")
}
}
func buildSessionTestPipeForNode(t *testing.T) (edgeSide *toki.TcpClient, sess *transport.Session) {
t.Helper()
edgeConn, nodeConn := net.Pipe()
edgeParserMap := toki.ParserMap{
toki.TypeNameOf(&iop.ProviderTunnelFrame{}): func(b []byte) (proto.Message, error) {
m := &iop.ProviderTunnelFrame{}
return m, proto.Unmarshal(b, m)
},
}
nodeParserMap := toki.ParserMap{}
edgeSide = toki.NewTcpClient(edgeConn, 0, 0, edgeParserMap)
nodeSide := toki.NewTcpClient(nodeConn, 0, 0, nodeParserMap)
t.Cleanup(func() { edgeSide.Close(); nodeSide.Close() })
sess = transport.ExportNewSession(nodeSide, zap.NewNop(), "node-id-1", "alias-1")
return edgeSide, sess
}
func TestNodeOnProviderTunnelRequest_LookupFailure(t *testing.T) {
router := &fixedRouter{
adapterName: "nonexistent",
adapters: make(map[string]runtime.Adapter),
lookupErrors: map[string]error{
"nonexistent": errors.New("adapter lookup error"),
},
}
n, _ := makeNode(t, router)
edgeSide, sess := buildSessionTestPipeForNode(t)
req := &iop.ProviderTunnelRequest{
RunId: "run-tunnel-1",
TunnelId: "tunnel-1",
Adapter: "nonexistent",
Target: "qwen",
}
frameCh := make(chan *iop.ProviderTunnelFrame, 10)
toki.AddListenerTyped[*iop.ProviderTunnelFrame](&edgeSide.Communicator, func(tf *iop.ProviderTunnelFrame) {
frameCh <- tf
})
err := n.OnProviderTunnelRequest(context.Background(), sess, req)
if err == nil {
t.Fatal("expected error, got nil")
}
select {
case frame := <-frameCh:
if frame.GetRunId() != "run-tunnel-1" || frame.GetTunnelId() != "tunnel-1" {
t.Errorf("unexpected IDs in frame: run_id=%q tunnel_id=%q", frame.GetRunId(), frame.GetTunnelId())
}
if frame.GetSequence() != 0 {
t.Errorf("expected sequence 0, got %d", frame.GetSequence())
}
if frame.GetKind() != iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR {
t.Errorf("expected ERROR frame, got %v", frame.GetKind())
}
if !strings.Contains(frame.GetError(), "adapter lookup error") {
t.Errorf("expected error message containing 'adapter lookup error', got %q", frame.GetError())
}
if frame.GetNodeId() != "node-id-1" || frame.GetNodeAlias() != "alias-1" {
t.Errorf("unexpected node ID or alias: node_id=%q alias=%q", frame.GetNodeId(), frame.GetNodeAlias())
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for ERROR frame")
}
// Verify no more frames
select {
case frame := <-frameCh:
t.Fatalf("unexpected duplicate frame received: %+v", frame)
default:
}
}
func TestNodeOnProviderTunnelRequest_UnsupportedAdapter(t *testing.T) {
// countingAdapter does not implement ProviderTunnelAdapter
mta := &countingAdapter{}
router := &fixedRouter{adapterName: "test", adapters: make(map[string]runtime.Adapter)}
router.adapters["test"] = mta
n, _ := makeNode(t, router)
edgeSide, sess := buildSessionTestPipeForNode(t)
req := &iop.ProviderTunnelRequest{
RunId: "run-tunnel-1",
TunnelId: "tunnel-1",
Adapter: "test",
Target: "qwen",
}
frameCh := make(chan *iop.ProviderTunnelFrame, 10)
toki.AddListenerTyped[*iop.ProviderTunnelFrame](&edgeSide.Communicator, func(tf *iop.ProviderTunnelFrame) {
frameCh <- tf
})
err := n.OnProviderTunnelRequest(context.Background(), sess, req)
if err == nil {
t.Fatal("expected error, got nil")
}
select {
case frame := <-frameCh:
if frame.GetRunId() != "run-tunnel-1" || frame.GetTunnelId() != "tunnel-1" {
t.Errorf("unexpected IDs in frame: run_id=%q tunnel_id=%q", frame.GetRunId(), frame.GetTunnelId())
}
if frame.GetSequence() != 0 {
t.Errorf("expected sequence 0, got %d", frame.GetSequence())
}
if frame.GetKind() != iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR {
t.Errorf("expected ERROR frame, got %v", frame.GetKind())
}
if !strings.Contains(frame.GetError(), "does not support tunneling") {
t.Errorf("expected error message containing 'does not support tunneling', got %q", frame.GetError())
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for ERROR frame")
}
// Verify no more frames
select {
case frame := <-frameCh:
t.Fatalf("unexpected duplicate frame received: %+v", frame)
default:
}
}
func TestNodeOnProviderTunnelRequest_AdapterErrorNoDuplicate(t *testing.T) {
// Adapter directly returns error
mta := &mockTunnelAdapter{
t: t,
expectedReq: runtime.ProviderTunnelRequest{
RunID: "run-tunnel-1",
TunnelID: "tunnel-1",
},
respondErr: errors.New("adapter runtime error"),
}
router := &fixedRouter{adapterName: "openai_compat", adapters: make(map[string]runtime.Adapter)}
router.adapters["openai_compat"] = mta
n, _ := makeNode(t, router)
edgeSide, sess := buildSessionTestPipeForNode(t)
req := &iop.ProviderTunnelRequest{
RunId: "run-tunnel-1",
TunnelId: "tunnel-1",
Adapter: "openai_compat",
Target: "qwen",
}
frameCh := make(chan *iop.ProviderTunnelFrame, 10)
toki.AddListenerTyped[*iop.ProviderTunnelFrame](&edgeSide.Communicator, func(tf *iop.ProviderTunnelFrame) {
frameCh <- tf
})
err := n.OnProviderTunnelRequest(context.Background(), sess, req)
if err == nil {
t.Fatal("expected error, got nil")
}
// The mockTunnelAdapter inside does not emit frames if respondErr is set, it just returns errors.
// Since we removed sendTunnelError for adapter failures from node.go, no frames should be received at all.
// Wait a brief moment to ensure no frames were sent.
select {
case frame := <-frameCh:
t.Fatalf("unexpected frame received: %+v", frame)
case <-time.After(100 * time.Millisecond):
}
}