iop/apps/edge/internal/bootstrap/reconnect_readiness_integration_test.go

package bootstrap

import (
	"context"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	goruntime "runtime"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

// TestActualNodeReconnectReadyPumpsQueuedWaiterExactlyOnce covers the hole that
// fake queue and bootstrap tests cannot see independently. It starts the actual
// iop-node entrypoint once, closes it to make the configured resource offline,
// queues a waiter behind an alternate provider's occupied slot, then starts the
// same node entrypoint again. The queued run must reach the real Node handler and
// return exactly one terminal event after the second process completes the
// RegisterResponse → SetHandler → NodeReadyRequest/ack sequence.
func TestActualNodeReconnectReadyPumpsQueuedWaiterExactlyOnce(t *testing.T) {
	if goruntime.GOOS == "windows" {
		t.Skip("the local actual-node fixture uses the POSIX sh CLI profile")
	}
	if _, err := exec.LookPath("sh"); err != nil {
		t.Skipf("sh is required by the local actual-node fixture: %v", err)
	}

	dir := t.TempDir()
	edgeAddr := freeTCPAddr(t)
	edgeConfigPath := writeReconnectReadinessEdgeConfig(t, dir, edgeAddr)
	rt, err := NewRuntime(loadServeNormalizedConfig(t, edgeConfigPath))
	if err != nil {
		t.Fatalf("NewRuntime: %v", err)
	}
	if err := rt.Start(context.Background()); err != nil {
		t.Fatalf("start edge runtime: %v", err)
	}
	defer func() {
		if err := rt.Stop(); err != nil {
			t.Errorf("stop edge runtime: %v", err)
		}
	}()

	alt := dialReconnectReadinessFakeNode(t, edgeAddr)
	defer alt.Close()
	altRuns := make(chan *iop.RunRequest, 4)
	toki.AddListenerTyped[*iop.RunRequest](&alt.Communicator, func(req *iop.RunRequest) {
		altRuns <- req
	})
	altRegistration, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&alt.Communicator, &iop.RegisterRequest{Token: "alt-token"}, fakeNodeHandshakeTimeout)
	if err != nil {
		t.Fatalf("register alternate fake node: %v", err)
	}
	if !altRegistration.GetAccepted() {
		t.Fatalf("alternate registration rejected: %s", altRegistration.GetReason())
	}
	readyFakeNode(t, alt, altRegistration.GetNodeId())
	waitForReconnectReadiness(t, 3*time.Second, "alternate node dispatch-ready", func() bool {
		_, ok := rt.Registry.GetReady("node-alt")
		return ok
	})

	nodeBinary := buildReconnectReadinessNodeBinary(t, dir)
	nodeConfigPath := writeReconnectReadinessNodeConfig(t, dir, edgeAddr)

	// First process establishes the resource, then exits cleanly. The waiter below
	// is therefore testing an actual reconnect, rather than only the first Node
	// registration path.
	first := startReconnectReadinessNode(t, nodeBinary, nodeConfigPath, filepath.Join(dir, "node-first"))
	waitForReconnectReadiness(t, 6*time.Second, "initial actual node ready", func() bool {
		_, ok := rt.Registry.GetReady("node-recon")
		return ok
	})
	if err := first.Stop(); err != nil {
		t.Fatalf("stop initial actual node: %v", err)
	}
	waitForReconnectReadiness(t, 4*time.Second, "initial actual node removed", func() bool {
		_, ok := rt.Registry.Get("node-recon")
		return !ok
	})

	// The alternate provider receives one background run and deliberately keeps
	// its lease until the test sends its terminal event. With the reconnect node
	// offline, the next run must become a provider-pool waiter.
	blocking, err := rt.Service.SubmitRun(context.Background(), edgeservice.SubmitRunRequest{
		ModelGroupKey: "reconnect-ready-model",
		ProviderPool:  true,
		Background:    true,
		Input:         map[string]any{"prompt": "hold alternate capacity"},
	})
	if err != nil {
		t.Fatalf("submit alternate blocking run: %v", err)
	}
	defer blocking.Close()
	select {
	case req := <-altRuns:
		if req.GetRunId() != blocking.Dispatch().RunID {
			t.Fatalf("alternate run id=%q, want %q", req.GetRunId(), blocking.Dispatch().RunID)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("alternate provider did not receive the blocking RunRequest")
	}

	type submitOutcome struct {
		result edgeservice.RunResult
		err    error
	}
	waiterDone := make(chan submitOutcome, 1)
	waiterCtx, cancelWaiter := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancelWaiter()
	go func() {
		result, submitErr := rt.Service.SubmitRun(waiterCtx, edgeservice.SubmitRunRequest{
			ModelGroupKey: "reconnect-ready-model",
			ProviderPool:  true,
			Input:         map[string]any{"prompt": "reconnect-ready waiter"},
		})
		waiterDone <- submitOutcome{result: result, err: submitErr}
	}()

	waitForProviderSnapshot(t, rt.Service, "node-alt", "provider-alt", 1, 1, 0, 0)
	select {
	case outcome := <-waiterDone:
		if outcome.result != nil {
			outcome.result.Close()
		}
		t.Fatalf("waiter resolved before reconnect: %v", outcome.err)
	default:
	}

	second := startReconnectReadinessNode(t, nodeBinary, nodeConfigPath, filepath.Join(dir, "node-second"))
	waitForReconnectReadiness(t, 6*time.Second, "reconnected actual node ready", func() bool {
		_, ok := rt.Registry.GetReady("node-recon")
		return ok
	})

	var waiter submitOutcome
	select {
	case waiter = <-waiterDone:
	case <-time.After(6 * time.Second):
		t.Fatal("queued waiter was not dispatched after actual node ready ack")
	}
	if waiter.err != nil {
		t.Fatalf("queued waiter failed after reconnect: %v", waiter.err)
	}
	if waiter.result == nil {
		t.Fatal("queued waiter returned nil result")
	}
	defer waiter.result.Close()
	dispatch := waiter.result.Dispatch()
	if dispatch.NodeID != "node-recon" || dispatch.ProviderID != "provider-recon" {
		t.Fatalf("waiter dispatch=(node=%q provider=%q), want node-recon/provider-recon", dispatch.NodeID, dispatch.ProviderID)
	}
	if dispatch.Adapter != "cli" || dispatch.Target != "recon-target" {
		t.Fatalf("waiter dispatch=(adapter=%q target=%q), want cli/recon-target", dispatch.Adapter, dispatch.Target)
	}
	if dispatch.QueueReason != "capacity_full" {
		t.Fatalf("waiter queue reason=%q, want capacity_full", dispatch.QueueReason)
	}
	waitForProviderSnapshot(t, rt.Service, "node-recon", "provider-recon", 1, 0, 0, 0)

	// A rejected duplicate connection cannot turn itself ready or pump another
	// waiter while the actual reconnect owner executes its RunRequest.
	duplicate := dialReconnectReadinessFakeNode(t, edgeAddr)
	defer duplicate.Close()
	dupRegistration, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&duplicate.Communicator, &iop.RegisterRequest{Token: "recon-token"}, fakeNodeHandshakeTimeout)
	if err != nil {
		t.Fatalf("duplicate register request: %v", err)
	}
	if dupRegistration.GetAccepted() {
		t.Fatal("duplicate reconnect registration was accepted")
	}
	dupReady, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
		&duplicate.Communicator, &iop.NodeReadyRequest{NodeId: "node-recon"}, fakeNodeHandshakeTimeout)
	if err != nil {
		t.Fatalf("duplicate ready request: %v", err)
	}
	if dupReady.GetReady() {
		t.Fatal("duplicate connection received a successful ready acknowledgement")
	}

	terminalCount := waitForRunTerminal(t, waiter.result.Stream().Events, dispatch.RunID)
	if terminalCount != 1 {
		t.Fatalf("actual node terminal count=%d, want exactly 1", terminalCount)
	}
	waitForReconnectReadiness(t, 3*time.Second, "actual node payload", func() bool {
		return strings.Contains(second.Output(), "reconnect-ready-terminal")
	})

	// Late terminal delivery is idempotent and must neither re-release the lease
	// nor dispatch another request. The primary terminal above came from the
	// actual node process; this is only the competing callback regression probe.
	rt.Service.HandleRunLifecycleEvent(&iop.RunEvent{RunId: dispatch.RunID, Type: "complete", NodeId: "node-recon"})
	waitForProviderSnapshot(t, rt.Service, "node-recon", "provider-recon", 0, 0, 0, 0)
	select {
	case unexpected := <-altRuns:
		t.Fatalf("unexpected extra alternate dispatch after duplicate ready/terminal: %q", unexpected.GetRunId())
	default:
	}

	if err := alt.Send(&iop.RunEvent{RunId: blocking.Dispatch().RunID, Type: "complete", NodeId: "node-alt"}); err != nil {
		t.Fatalf("send alternate terminal: %v", err)
	}
	waitForProviderSnapshot(t, rt.Service, "node-alt", "provider-alt", 0, 0, 0, 0)
	assertReconnectReadinessCountersZero(t, rt.Service)

	if err := second.Stop(); err != nil {
		t.Fatalf("stop reconnected actual node: %v", err)
	}
}

func reconnectReadinessParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RegisterResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.NodeReadyResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.NodeReadyResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

func dialReconnectReadinessFakeNode(t *testing.T, edgeAddr string) *toki.TcpClient {
	t.Helper()
	host, portText, err := net.SplitHostPort(edgeAddr)
	if err != nil {
		t.Fatalf("split edge address %q: %v", edgeAddr, err)
	}
	port, err := strconv.Atoi(portText)
	if err != nil {
		t.Fatalf("parse edge port %q: %v", portText, err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	client, err := toki.DialTcp(ctx, host, port, 30, 10, reconnectReadinessParserMap())
	if err != nil {
		t.Fatalf("dial edge: %v", err)
	}
	return client
}

func writeReconnectReadinessEdgeConfig(t *testing.T, dir, edgeAddr string) string {
	t.Helper()
	path := filepath.Join(dir, "edge.yaml")
	yaml := fmt.Sprintf(`
server:
  listen: %q
bootstrap:
  listen: "127.0.0.1:0"
  artifact_dir: %q
logging:
  level: "error"
refresh:
  enabled: false
  listen: "127.0.0.1:0"
openai:
  enabled: false
  listen: "127.0.0.1:0"
a2a:
  enabled: false
  listen: "127.0.0.1:0"
metrics:
  port: 0
provider_pool:
  max_queue: 4
  queue_timeout_ms: 8000
models:
  - id: "reconnect-ready-model"
    providers:
      provider-alt: "alt-target"
      provider-recon: "recon-target"
nodes:
  - id: "node-alt"
    alias: "alternate"
    token: "alt-token"
    adapters:
      cli:
        enabled: true
        profiles:
          alt-target:
            command: "sh"
            args: ["-c", "sleep 30"]
    providers:
      - id: "provider-alt"
        type: "cli"
        category: "cli"
        adapter: "cli"
        models: ["alt-target"]
        health: "available"
        capacity: 1
  - id: "node-recon"
    alias: "reconnect"
    token: "recon-token"
    adapters:
      cli:
        enabled: true
        profiles:
          recon-target:
            command: "sh"
            args: ["-c", "sleep 0.4; printf 'reconnect-ready-terminal\\n'"]
    providers:
      - id: "provider-recon"
        type: "cli"
        category: "cli"
        adapter: "cli"
        models: ["recon-target"]
        health: "available"
        capacity: 1
`, edgeAddr, filepath.Join(dir, "artifacts"))
	if err := os.WriteFile(path, []byte(yaml), 0o600); err != nil {
		t.Fatalf("write edge config: %v", err)
	}
	return path
}

func writeReconnectReadinessNodeConfig(t *testing.T, dir, edgeAddr string) string {
	t.Helper()
	path := filepath.Join(dir, "node.yaml")
	yaml := fmt.Sprintf(`
transport:
  edge_addr: %q
  token: "recon-token"
reconnect:
  interval_sec: 1
  max_attempts: 1
logging:
  level: "error"
  path: %q
metrics:
  port: 0
`, edgeAddr, filepath.Join(dir, "node.log"))
	if err := os.WriteFile(path, []byte(yaml), 0o600); err != nil {
		t.Fatalf("write node config: %v", err)
	}
	return path
}

func buildReconnectReadinessNodeBinary(t *testing.T, dir string) string {
	t.Helper()
	root := reconnectReadinessRepoRoot(t)
	binary := filepath.Join(dir, "iop-node")
	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "go", "build", "-o", binary, "./apps/node/cmd/node")
	cmd.Dir = root
	output, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("build actual iop-node: %v\n%s", err, output)
	}
	return binary
}

func reconnectReadinessRepoRoot(t *testing.T) string {
	t.Helper()
	dir, err := os.Getwd()
	if err != nil {
		t.Fatalf("getwd: %v", err)
	}
	for {
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			return dir
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			t.Fatal("could not locate repository go.mod")
		}
		dir = parent
	}
}

type reconnectReadinessNodeProcess struct {
	cmd        *exec.Cmd
	stdoutPath string
	stderrPath string
	stdout     *os.File
	stderr     *os.File
	wait       chan error
	once       sync.Once
	err        error
}

func startReconnectReadinessNode(t *testing.T, binary, configPath, dir string) *reconnectReadinessNodeProcess {
	t.Helper()
	if err := os.MkdirAll(dir, 0o755); err != nil {
		t.Fatalf("create node work directory: %v", err)
	}
	stdoutPath := filepath.Join(dir, "stdout.log")
	stderrPath := filepath.Join(dir, "stderr.log")
	stdout, err := os.Create(stdoutPath)
	if err != nil {
		t.Fatalf("create node stdout log: %v", err)
	}
	stderr, err := os.Create(stderrPath)
	if err != nil {
		_ = stdout.Close()
		t.Fatalf("create node stderr log: %v", err)
	}

	cmd := exec.Command(binary, "--config", configPath, "serve")
	cmd.Dir = dir
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	if err := cmd.Start(); err != nil {
		_ = stdout.Close()
		_ = stderr.Close()
		t.Fatalf("start actual iop-node: %v", err)
	}
	p := &reconnectReadinessNodeProcess{
		cmd:        cmd,
		stdoutPath: stdoutPath,
		stderrPath: stderrPath,
		stdout:     stdout,
		stderr:     stderr,
		wait:       make(chan error, 1),
	}
	go func() { p.wait <- cmd.Wait() }()
	t.Cleanup(func() {
		if err := p.Stop(); err != nil && !t.Failed() {
			t.Errorf("stop actual iop-node during cleanup: %v", err)
		}
		if t.Failed() {
			t.Logf("actual iop-node stdout (%s):\n%s", p.stdoutPath, p.Output())
			t.Logf("actual iop-node stderr (%s):\n%s", p.stderrPath, p.ErrorOutput())
		}
	})
	return p
}

func (p *reconnectReadinessNodeProcess) Stop() error {
	p.once.Do(func() {
		if p.cmd.ProcessState == nil {
			if err := p.cmd.Process.Signal(os.Interrupt); err != nil && !strings.Contains(err.Error(), "process already finished") {
				p.err = err
			}
		}
		select {
		case waitErr := <-p.wait:
			if p.err == nil {
				p.err = waitErr
			}
		case <-time.After(5 * time.Second):
			_ = p.cmd.Process.Kill()
			waitErr := <-p.wait
			if p.err == nil {
				p.err = fmt.Errorf("iop-node did not stop after interrupt: %w", waitErr)
			}
		}
		_ = p.stdout.Close()
		_ = p.stderr.Close()
	})
	return p.err
}

func (p *reconnectReadinessNodeProcess) Output() string {
	b, err := os.ReadFile(p.stdoutPath)
	if err != nil {
		return fmt.Sprintf("read stdout: %v", err)
	}
	return string(b)
}

func (p *reconnectReadinessNodeProcess) ErrorOutput() string {
	b, err := os.ReadFile(p.stderrPath)
	if err != nil {
		return fmt.Sprintf("read stderr: %v", err)
	}
	return string(b)
}

func waitForReconnectReadiness(t *testing.T, timeout time.Duration, description string, condition func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if condition() {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	if condition() {
		return
	}
	t.Fatalf("timed out waiting for %s", description)
}

func waitForProviderSnapshot(t *testing.T, service *edgeservice.Service, nodeID, providerID string, inFlight, queued, longInFlight, longQueued int32) {
	t.Helper()
	waitForReconnectReadiness(t, 4*time.Second, fmt.Sprintf("provider snapshot %s/%s", nodeID, providerID), func() bool {
		for _, node := range service.ListNodeSnapshots() {
			if node.NodeID != nodeID {
				continue
			}
			for _, provider := range node.ProviderSnapshots {
				if provider.GetId() == providerID {
					return provider.GetInFlight() == inFlight && provider.GetQueued() == queued &&
						provider.GetLongInFlight() == longInFlight && provider.GetLongQueued() == longQueued
				}
			}
		}
		return false
	})
}

func waitForRunTerminal(t *testing.T, events <-chan *iop.RunEvent, runID string) int {
	t.Helper()
	terminalCount := 0
	deadline := time.NewTimer(6 * time.Second)
	defer deadline.Stop()
	for terminalCount == 0 {
		select {
		case event := <-events:
			if event != nil && event.GetRunId() == runID && isReconnectTerminal(event.GetType()) {
				terminalCount++
			}
		case <-deadline.C:
			t.Fatalf("timed out waiting for actual node terminal run_id=%q", runID)
		}
	}

	quiet := time.NewTimer(300 * time.Millisecond)
	defer quiet.Stop()
	for {
		select {
		case event := <-events:
			if event != nil && event.GetRunId() == runID && isReconnectTerminal(event.GetType()) {
				terminalCount++
			}
		case <-quiet.C:
			return terminalCount
		}
	}
}

func isReconnectTerminal(eventType string) bool {
	switch eventType {
	case "complete", "error", "cancelled":
		return true
	default:
		return false
	}
}

func assertReconnectReadinessCountersZero(t *testing.T, service *edgeservice.Service) {
	t.Helper()
	for _, node := range service.ListNodeSnapshots() {
		for _, provider := range node.ProviderSnapshots {
			if provider.GetInFlight() != 0 || provider.GetQueued() != 0 || provider.GetLongInFlight() != 0 || provider.GetLongQueued() != 0 {
				t.Errorf("provider counters not settled node=%q provider=%q normal=%d queued=%d long=%d long_queued=%d",
					node.NodeID, provider.GetId(), provider.GetInFlight(), provider.GetQueued(), provider.GetLongInFlight(), provider.GetLongQueued())
			}
		}
	}
}