feat(node): 연결 supervisor 및 CLI adaptors 개선

- runtimeSupervisor 단일 goroutine으로 연결·재연결 수명주기 관리
- 초기 연결 실패 시 fx OnStart 후크 중단 방지 (SDD S16)
- finite 재연결 시도 고갈 시 노드 종료
- CLI adaptors: oneshot, opencode_sse, persistent 보강 및 테스트 확장
This commit is contained in:
toki 2026-07-22 18:11:06 +09:00
parent 4dcf6f0cf9
commit 4a76851254
18 changed files with 1677 additions and 119 deletions

View file

@ -39,9 +39,36 @@ func (c *CLI) executeCommand(ctx context.Context, spec runtime.ExecutionSpec, pr
}
if err := cmd.Start(); err != nil {
// The context can be cancelled or time out before the process starts,
// which surfaces as a start error. Honor the cancellation contract with a
// single cancelled event and ErrRunCancelled instead of a start error.
if ctx.Err() != nil {
_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeCancelled,
Message: cancelEventForContext(ctx.Err()),
Timestamp: time.Now(),
})
return "", runtime.ErrRunCancelled
}
return "", emitReturnedError(ctx, sink, spec.RunID, fmt.Errorf("cli adapter: start %q: %w", profile.Command, err))
}
// Closing the stdout/stderr pipes on cancellation unblocks pending reads
// promptly. CommandContext only signals the direct child, so a forked
// grandchild that keeps the pipes open would otherwise stall the read until
// it exits on its own; closing the reader ends the read-vs-wait race.
readsDone := make(chan struct{})
defer close(readsDone)
go func() {
select {
case <-ctx.Done():
_ = stdout.Close()
_ = stderr.Close()
case <-readsDone:
}
}()
_ = sink.Emit(ctx, runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeStart,
@ -71,6 +98,19 @@ func (c *CLI) executeCommand(ctx context.Context, spec runtime.ExecutionSpec, pr
if readErr != nil {
_ = <-stderrDone
_ = cmd.Wait()
// A context cancellation can close the stdout pipe before cmd.Wait()
// observes the process exit, surfacing as a read error. Treat that as a
// cancellation so the user-cancel/timeout contract wins over the generic
// read-stdout error, emitting exactly one cancelled event.
if ctx.Err() != nil {
_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeCancelled,
Message: cancelEventForContext(ctx.Err()),
Timestamp: time.Now(),
})
return combinedOutput(), runtime.ErrRunCancelled
}
_ = sink.Emit(ctx, runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeError,

View file

@ -513,9 +513,11 @@ func TestCLIExecuteOneShot_UserCancelEmitsUserCancelMessage(t *testing.T) {
cfg := config.CLIConf{
Enabled: true,
Profiles: map[string]config.CLIProfileConf{
// Emit a marker before blocking so the test can cancel deterministically
// once the process is running and the stdout reader is active.
"slow-cancel": {
Command: "sh",
Args: []string{"-c", `sleep 10`},
Args: []string{"-c", `echo running; sleep 10`},
},
},
}
@ -523,33 +525,48 @@ func TestCLIExecuteOneShot_UserCancelEmitsUserCancelMessage(t *testing.T) {
sink := &testutil.FakeSink{}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
done := make(chan error, 1)
go func() {
time.Sleep(100 * time.Millisecond)
cancel()
done <- c.Execute(ctx, noderuntime.ExecutionSpec{
RunID: "run-usercancel",
Target: "slow-cancel",
Input: map[string]any{"prompt": "hello"},
}, sink)
}()
err := c.Execute(ctx, noderuntime.ExecutionSpec{
RunID: "run-usercancel",
Target: "slow-cancel",
Input: map[string]any{"prompt": "hello"},
}, sink)
// Cancel only after the first stdout chunk proves the process started and the
// read loop is active, so this exercises the stdout-read/command-wait
// cancellation race rather than a pre-start cancellation.
testutil.RequireEventually(t, 5*time.Second, 5*time.Millisecond, func() bool {
return testutil.CollectDeltas(sink.Events()) != ""
})
cancel()
var err error
select {
case err = <-done:
case <-time.After(5 * time.Second):
t.Fatal("execute did not return promptly after cancellation")
}
if err != noderuntime.ErrRunCancelled {
t.Fatalf("expected ErrRunCancelled, got %v", err)
}
var cancelCount int
var lastMsg string
var cancelEvent bool
events := sink.Events()
if len(events) > 0 {
for i := len(events) - 1; i >= 0; i-- {
if events[i].Type == noderuntime.EventTypeCancelled {
cancelEvent = true
lastMsg = events[i].Message
}
for _, ev := range sink.Events() {
switch ev.Type {
case noderuntime.EventTypeCancelled:
cancelCount++
lastMsg = ev.Message
case noderuntime.EventTypeError:
t.Fatalf("expected no error event on cancellation, got %q", ev.Error)
}
}
if !cancelEvent {
t.Fatal("expected cancelled event")
if cancelCount != 1 {
t.Fatalf("expected exactly one cancelled event, got %d", cancelCount)
}
if lastMsg != "user-cancel" {
t.Fatalf("expected cancel Message 'user-cancel', got %q", lastMsg)

View file

@ -166,6 +166,24 @@ func (e *opencodeExecutor) Execute(ctx context.Context, spec runtime.ExecutionSp
prompt := extractPrompt(spec.Input)
if err := opencodePromptAsync(ctx, sess.serverURL, sess.sessionID, prompt, opts); err != nil {
// A cancellation while the async prompt POST is in flight surfaces as a
// request error; honor the cancellation contract with a bounded background
// abort on the remote session, a single cancelled event, and
// ErrRunCancelled instead of a generic prompt-async error.
if ctx.Err() != nil {
bg, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
if sess.sessionID != "" {
_ = opencodeAbort(bg, sess.serverURL, sess.sessionID)
}
_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeCancelled,
Message: cancelEventForContext(ctx.Err()),
Timestamp: time.Now(),
})
return runtime.ErrRunCancelled
}
_ = sink.Emit(ctx, runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeError,

View file

@ -31,13 +31,37 @@ type opencodeFakeServer struct {
createCalls atomic.Int32
lastPromptBody []byte
createdID string
// promptReady is signalled once the prompt_async handler has fully read the
// request body. Tests wait on this before issuing ordering decisions.
promptReady chan struct{}
// promptBlock, when non-nil, is drained once to unblock the prompt_async
// response write. nil means the handler replies immediately.
promptBlock *promptGate
}
// promptGate is a one-shot response gate for the prompt_async handler.
type promptGate struct {
mu sync.Mutex
closed bool
release chan struct{} // closed to unblock the handler
blocked chan struct{} // closed when the handler enters the blocked wait
}
func (g *promptGate) signalBlocked() {
select {
case <-g.blocked:
default:
close(g.blocked)
}
}
func newOpencodeFakeServer(t *testing.T, sessionID string) (*opencodeFakeServer, *httptest.Server) {
t.Helper()
s := &opencodeFakeServer{
events: make(chan map[string]any, 32),
createdID: sessionID,
events: make(chan map[string]any, 32),
createdID: sessionID,
promptReady: make(chan struct{}, 1),
}
mux := http.NewServeMux()
eventHandler := func(w http.ResponseWriter, r *http.Request) {
@ -94,6 +118,22 @@ func newOpencodeFakeServer(t *testing.T, sessionID string) (*opencodeFakeServer,
s.mu.Lock()
s.lastPromptBody = append([]byte(nil), buf[:n]...)
s.mu.Unlock()
// Signal that the prompt body has been received so callers can
// observe the acceptance event before the response is written.
select {
case s.promptReady <- struct{}{}:
default:
}
// If a response gate is armed, wait for release or request cancellation.
gate := s.currentPromptGate()
if gate != nil {
gate.signalBlocked()
select {
case <-gate.release:
case <-r.Context().Done():
return
}
}
w.WriteHeader(http.StatusAccepted)
case "abort":
s.abortCalls.Add(1)
@ -127,6 +167,68 @@ func (s *opencodeFakeServer) push(ev map[string]any) {
}
}
// BlockPromptResponse arms the prompt response gate. The prompt_async handler
// will not reply with 202 until UnblockPromptResponse is called or the server
// closes. Safe to call multiple times; only the most recent gate is in effect.
// Returns the gate reference for handler-block ordering checks.
func (s *opencodeFakeServer) BlockPromptResponse() *promptGate {
s.mu.Lock()
g := &promptGate{release: make(chan struct{}), blocked: make(chan struct{})}
s.promptBlock = g
s.mu.Unlock()
return g
}
// UnblockPromptResponse releases the most recent prompt response gate. If no
// gate is armed, this is a no-op. Safe to call multiple times; only the
// first call closes the release channel.
func (s *opencodeFakeServer) UnblockPromptResponse() {
s.mu.Lock()
block := s.promptBlock
s.promptBlock = nil
s.mu.Unlock()
if block == nil {
return
}
block.mu.Lock()
if block.closed {
block.mu.Unlock()
return
}
block.closed = true
block.mu.Unlock()
close(block.release)
}
// currentPromptGate returns the currently armed gate, or nil.
func (s *opencodeFakeServer) currentPromptGate() *promptGate {
s.mu.Lock()
defer s.mu.Unlock()
return s.promptBlock
}
// waitPromptBodyRead waits until the handler has fully read the prompt body,
// or times out.
func waitPromptBodyRead(t *testing.T, fake *opencodeFakeServer) {
t.Helper()
select {
case <-fake.promptReady:
case <-time.After(3 * time.Second):
t.Fatal("prompt body was not read by handler")
}
}
// waitPromptResponseBlocked waits until the handler has entered the blocked
// wait on the response gate, or times out.
func waitPromptResponseBlocked(t *testing.T, g *promptGate) {
t.Helper()
select {
case <-g.blocked:
case <-time.After(3 * time.Second):
t.Fatal("handler did not enter blocked wait on response gate")
}
}
func opencodeGlobalEvent(eventType string, props map[string]any) map[string]any {
return map[string]any{
"directory": "/tmp/iop-test",
@ -779,3 +881,76 @@ func TestCLIExecuteOpencodeSSE_ContextCancelAbortsSession(t *testing.T) {
t.Fatal("expected /session/{id}/abort to be called")
}
}
// TestCLIExecuteOpencodeSSE_ContextCancelAfterPromptAcceptedAbortsSession
// verifies the regression: when prompt_async POST reaches the server (202) but
// the response is blocked, cancelling the run context must abort the remote
// session exactly once via /abort, emit exactly one cancelled event, and
// return ErrRunCancelled with no error event.
func TestCLIExecuteOpencodeSSE_ContextCancelAfterPromptAcceptedAbortsSession(t *testing.T) {
fake, srv := newOpencodeFakeServer(t, "ses_abort")
cfg := config.CLIConf{
Enabled: true,
Profiles: map[string]config.CLIProfileConf{
"opencode": opencodeSSEProfile(srv.URL),
},
}
c := clipkg.New(cfg, zap.NewNop())
sink := &testutil.FakeSink{}
ctx, cancel := context.WithCancel(context.Background())
// Arm the response gate before Execute so the handler blocks on it.
blocked := fake.BlockPromptResponse()
defer fake.UnblockPromptResponse()
done := make(chan error, 1)
go func() {
done <- c.Execute(ctx, noderuntime.ExecutionSpec{
RunID: "run-abort-after-accept",
Target: "opencode",
Input: map[string]any{"prompt": "abort-test"},
}, sink)
}()
// Wait for the handler to confirm it has the body and is blocked on the gate.
waitPromptBodyRead(t, fake)
waitPromptResponseBlocked(t, blocked)
// Cancel the run context. The adapter must handle this by sending a
// bounded /abort, emitting one cancelled event, and returning ErrRunCancelled.
cancel()
// The execute call must return ErrRunCancelled within a bounded window.
select {
case err := <-done:
if err != noderuntime.ErrRunCancelled {
t.Fatalf("expected ErrRunCancelled, got %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("execute did not return after prompt-accepted cancellation")
}
// Execute returns only after the in-flight cancellation path completes /abort.
if got := fake.abortCalls.Load(); got != 1 {
t.Fatalf("expected /abort call count == 1 when Execute returned, got %d", got)
}
// Exactly one cancelled event, no error event.
var cancelCount, errCount int
for _, e := range sink.Events() {
switch e.Type {
case noderuntime.EventTypeCancelled:
cancelCount++
case noderuntime.EventTypeError:
errCount++
}
}
if cancelCount != 1 {
t.Fatalf("expected exactly 1 cancelled event, got %d", cancelCount)
}
if errCount != 0 {
t.Fatalf("expected 0 error events, got %d", errCount)
}
}

View file

@ -264,6 +264,18 @@ func (e *persistentExecutor) Execute(ctx context.Context, spec runtime.Execution
}
if err := run.setup(); err != nil {
// Writing the prompt aborts with the context error when the run is
// cancelled or times out mid-setup. Emit a single cancelled event and
// return ErrRunCancelled so cancellation is not reported as a setup error.
if ctx.Err() != nil {
_ = sink.Emit(context.Background(), runtime.RuntimeEvent{
RunID: spec.RunID,
Type: runtime.EventTypeCancelled,
Message: cancelEventForContext(ctx.Err()),
Timestamp: time.Now(),
})
return runtime.ErrRunCancelled
}
return emitRuntimeError(ctx, sink, spec.RunID, err.Error())
}

View file

@ -210,7 +210,7 @@ func TestCLIExecutePersistent_CompletionMarkerLineEndsRun(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}
@ -261,7 +261,7 @@ func TestCLIExecutePersistent_CompletionMarkerRegexEndsRun(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}
@ -311,7 +311,7 @@ func TestCLIExecutePersistent_NoMarkerFallsBackToIdleTimeout(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}
@ -361,7 +361,7 @@ func TestCLIExecutePersistent_IdleTimeoutMessageReason(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}

View file

@ -36,7 +36,7 @@ func TestCLIExecutePersistentWaitsForSlowFirstOutput(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}

View file

@ -86,7 +86,12 @@ func TestCLIExecutePersistentTerminalSendsCarriageReturn(t *testing.T) {
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 100,
StartupIdleTimeoutMS: 50,
// This scenario asserts the raw "\r" terminator, so the prompt must be
// written only after the helper has entered raw mode. The readiness
// marker anchors the drain to that point; the timeout must exceed the
// re-exec'd helper's worst-case cold start (~700ms observed) so the
// drain never settles before the marker arrives.
StartupIdleTimeoutMS: 800,
},
},
}
@ -98,7 +103,10 @@ func TestCLIExecutePersistentTerminalSendsCarriageReturn(t *testing.T) {
}
defer func() { _ = c.Stop(ctx) }()
execCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
// Generous bound: the startup drain alone is 800ms (see profile) and process
// spawn is slow and variable in constrained CI, so a 1s cap would race the run
// to completion. This still asserts completion, not a timeout.
execCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
sink := &testutil.FakeSink{}
@ -136,7 +144,10 @@ func TestCLIExecutePersistentTerminalAcceptsClaudeBypassWarning(t *testing.T) {
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 100,
StartupIdleTimeoutMS: 250,
// The startup drain must observe the bypass-warning screen (and
// auto-accept it) before the prompt is sent; cover the re-exec'd
// helper's worst-case cold start, which is larger under -race.
StartupIdleTimeoutMS: 800,
},
},
}
@ -182,7 +193,10 @@ func TestCLIExecutePersistentTerminalAcceptsClaudeWorkspaceTrustAndBypassWarning
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 100,
StartupIdleTimeoutMS: 250,
// The startup drain must observe the trust + bypass-warning screens
// (and auto-accept them) before the prompt is sent; cover the
// re-exec'd helper's worst-case cold start, which is larger under -race.
StartupIdleTimeoutMS: 800,
},
},
}
@ -228,7 +242,9 @@ func TestCLIExecutePersistentClaudeTUICancelsSessionLimitUpgradePrompt(t *testin
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 500,
StartupIdleTimeoutMS: 50,
// Cover the re-exec'd helper's variable cold start so the startup
// drain waits for its readiness marker before the prompt is sent.
StartupIdleTimeoutMS: 300,
},
},
}
@ -282,7 +298,10 @@ func TestCLIExecutePersistentClaudeTUIReplaysPromptAfterStartupReadyRace(t *test
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 100,
StartupIdleTimeoutMS: 40,
// Cover the re-exec'd helper's variable cold start so the startup
// drain observes the trust screen and auto-accepts it before the
// prompt is sent, keeping the replay-after-ready path deterministic.
StartupIdleTimeoutMS: 300,
},
},
}
@ -406,8 +425,12 @@ func TestCLIExecutePersistentClaudeTUIFiltersTerminalChrome(t *testing.T) {
Enabled: true,
Profiles: map[string]config.CLIProfileConf{
"claude-tui": {
Command: "sh",
Args: []string{"-c", `stty -echo; while IFS= read -r line; do printf '\033[2C\r\033[7A\342\227\217\033[1Cmessage:%s\r\n\302\267 Cogitating... (1s, 1 tokens)\r\n\342\235\257 ' "$line"; done`},
Command: "sh",
// The frame begins with cursor-home + erase-display, mirroring a real
// claude TUI redraw, so the rendered assistant screen is independent
// of whether the prompt raced `stty -echo` and got echoed: any echoed
// characters are erased before the assistant line is drawn.
Args: []string{"-c", `stty -echo; while IFS= read -r line; do printf '\033[H\033[2J\033[2C\r\033[7A\342\227\217\033[1Cmessage:%s\r\n\302\267 Cogitating... (1s, 1 tokens)\r\n\342\235\257 ' "$line"; done`},
Persistent: true,
Terminal: true,
ResponseIdleTimeoutMS: 100,

View file

@ -5,7 +5,6 @@ import (
"os"
osexec "os/exec"
"testing"
"time"
)
func TestRawTUIHelperProcess(t *testing.T) {
@ -36,8 +35,20 @@ func setRawTerminal() {
}
}
// emitStartupReady announces that the helper has finished configuring its raw
// terminal and is now reading input. Re-exec'd test binaries pay a variable
// cold-start cost, so the adapter's startup drain must observe this marker (and
// idle after it) before it writes the prompt. Otherwise the prompt would race
// terminal setup and be echoed or line-translated. The marker is consumed by
// the startup drain and never surfaces as a delta.
func emitStartupReady() {
fmt.Fprint(os.Stdout, "iop-tui-ready\r\n")
_ = os.Stdout.Sync()
}
func runRawTUIHelper() {
setRawTerminal()
emitStartupReady()
readPromptAndReply()
}
@ -114,10 +125,15 @@ func runClaudeTrustAndWarningHelper() {
func runClaudeRateLimitHelper() {
setRawTerminal()
emitStartupReady()
for {
buf := make([]byte, 1)
n, err := os.Stdin.Read(buf)
if n > 0 && buf[0] == '\r' {
// Trigger on either terminator: in raw mode the adapter's terminator is
// "\r", but if the prompt raced terminal setup the line discipline may have
// translated it to "\n". Accepting both keeps session-limit detection
// independent of the re-exec'd helper's cold-start timing.
if n > 0 && (buf[0] == '\r' || buf[0] == '\n') {
fmt.Fprint(os.Stdout, "\u25cf \u23bf You've hit your session limit \u00b7 resets 12:50pm (Asia/Seoul)\r\n")
fmt.Fprint(os.Stdout, "What do you want to do?\r\n 1. Stop and wait for limit to reset\r\n2. Upgrade your plan\r\n")
_ = os.Stdout.Sync()
@ -130,17 +146,17 @@ func runClaudeRateLimitHelper() {
func runClaudeReadyReplayHelper() {
setRawTerminal()
fmt.Fprint(os.Stdout, "Quick safety check: Is this a project you created or one you trust?\r\n 1. Yes, I trust this folder\r\n2. No, exit\r\n")
_ = os.Stdout.Sync()
emitStartupReady()
_ = readRawLine()
time.Sleep(150 * time.Millisecond)
_ = readRawLine()
// Handshake via blocking reads instead of a fixed sleep: swallow the first
// prompt delivery to model a CLI that is not ready for input yet, then
// advertise the ready screen so the adapter replays the prompt exactly once.
_ = readRawLine() // first prompt, intentionally dropped before the ready screen
fmt.Fprint(os.Stdout, "\r\n ")
_ = os.Stdout.Sync()
prompt := readRawLine()
prompt := readRawLine() // replayed prompt
fmt.Fprintf(os.Stdout, "\r\n\u25cf replay:%s\r\n\u276f ", prompt)
_ = os.Stdout.Sync()

View file

@ -23,6 +23,11 @@ import (
iop "iop/proto/gen/iop"
)
// readySignalTimeout bounds the post-handler NodeReadyRequest/ack round trip. A
// ready failure is retryable (not a *transport.ConnectError), so the supervisor
// reconnects rather than shutting the node down.
const readySignalTimeout = 10 * time.Second
// DialFunc is the signature used to connect to Edge. Can be replaced in tests.
type DialFunc func(ctx context.Context, addr, token string, logger *zap.Logger) (*transport.RegisterResult, error)
@ -107,6 +112,17 @@ func connectRuntime(ctx context.Context, cfg *config.NodeConfig, logger *zap.Log
})
result.Session.SetHandler(n)
// The handler is installed; only now tell edge we are ready to receive
// dispatch. Edge opens run/tunnel/command eligibility and pumps any waiter
// stranded while this node was offline strictly in response to this signal, so
// the first request can never race ahead of the handler. A failed ready
// handshake is retryable: close the partial owner and let the supervisor
// reconnect rather than sit connected-but-unreachable.
if err := result.Session.SignalReady(readySignalTimeout); err != nil {
owner.close()
return nil, fmt.Errorf("signal ready: %w", err)
}
logger.Info("connected to edge",
zap.String("node_id", result.NodeID),
zap.String("alias", result.Alias),

View file

@ -12,6 +12,7 @@ import (
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
"git.toki-labs.com/toki/proto-socket/go/packets"
"go.uber.org/fx"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
@ -28,9 +29,48 @@ func bootstrapEdgeParserMap() toki.ParserMap {
m := &iop.RegisterRequest{}
return m, proto.Unmarshal(b, m)
},
toki.TypeNameOf(&iop.NodeReadyRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.NodeReadyRequest{}
return m, proto.Unmarshal(b, m)
},
}
}
// bindEdgeReadyResponder answers the node's post-handler NodeReadyRequest with a
// positive ack, so connectRuntime's SignalReady completes and the session is
// treated as fully established. Mock edges that accept registration must install
// it or the node's ready handshake times out and reconnects. When readyCh is
// non-nil it fires only after the ready response has been written to the socket.
// This lets a test safely close the edge side after a fully established session,
// instead of racing the in-flight ready handshake.
func bindEdgeReadyResponder(client *toki.TcpClient, readyCh chan<- struct{}) {
client.Communicator.AddRequestListener(
toki.TypeNameOf(&iop.NodeReadyRequest{}),
func(message proto.Message, requestNonce int32) {
if _, ok := message.(*iop.NodeReadyRequest); !ok {
return
}
data, err := proto.Marshal(&iop.NodeReadyResponse{Ready: true})
if err != nil {
return
}
if err := client.QueuePacket(&packets.PacketBase{
TypeName: toki.TypeNameOf(&iop.NodeReadyResponse{}),
ResponseNonce: requestNonce,
Data: data,
}); err != nil {
return
}
if readyCh != nil {
select {
case readyCh <- struct{}{}:
default:
}
}
},
)
}
func freeAddrForBootstrap(t *testing.T) (host string, port int, addr string) {
t.Helper()
l, err := net.Listen("tcp", "127.0.0.1:0")
@ -60,6 +100,453 @@ func useTempCwd(t *testing.T) {
})
}
// startAcceptingEdge starts a mock edge bound to host:port that accepts every
// RegisterRequest. The returned channels fire once per registration and
// disconnect, respectively.
func startAcceptingEdge(t *testing.T, ctx context.Context, host string, port int, nodeID string) (chan struct{}, chan struct{}) {
t.Helper()
regCh := make(chan struct{}, 16)
disconnectCh := make(chan struct{}, 16)
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
client.AddDisconnectListener(func(_ *toki.TcpClient) {
disconnectCh <- struct{}{}
})
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
&client.Communicator,
func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
regCh <- struct{}{}
return &iop.RegisterResponse{
Accepted: true,
NodeId: nodeID,
Config: &iop.NodeConfigPayload{Runtime: &iop.NodeRuntimeConfig{Concurrency: 1}},
}, nil
},
)
bindEdgeReadyResponder(client, nil)
return client
})
if err := server.Start(ctx); err != nil {
t.Fatalf("start accepting edge: %v", err)
}
t.Cleanup(func() { server.Stop() })
return regCh, disconnectCh
}
// startRejectingEdge starts a mock edge bound to host:port that rejects every
// RegisterRequest with the given reason.
func startRejectingEdge(t *testing.T, ctx context.Context, host string, port int, reason string) {
t.Helper()
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
&client.Communicator,
func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
return &iop.RegisterResponse{Accepted: false, Reason: reason}, nil
},
)
return client
})
if err := server.Start(ctx); err != nil {
t.Fatalf("start rejecting edge: %v", err)
}
t.Cleanup(func() { server.Stop() })
}
// TestSupervisorConnectsWhenEdgeStartsAfterNode verifies the actual DialEdge
// connection-refused path: Node starts first against a closed port, keeps
// retrying in the same process, and establishes exactly one session after Edge
// begins listening on that same port. SDD S16.
func TestSupervisorConnectsWhenEdgeStartsAfterNode(t *testing.T) {
useTempCwd(t)
host, port, addr := freeAddrForBootstrap(t)
const refusedAttempts = 3
var dialCount int32
retryReady := make(chan struct{}, refusedAttempts)
releaseRetry := make(chan struct{})
observedDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
return transport.DialEdge(ctx, edgeAddr, token, logger)
})
fakeSleeper := func(ctx context.Context, _ time.Duration) {
retryReady <- struct{}{}
if atomic.LoadInt32(&dialCount) >= refusedAttempts {
select {
case <-releaseRetry:
case <-ctx.Done():
}
}
}
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(observedDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
defer func() { _ = app.Stop(context.Background()) }()
// Edge is not listening yet. Each completed sleep entry follows a real
// DialEdge connection-refused result. Hold the third retry until the listener
// is live so the next actual dial succeeds deterministically.
for attempt := 1; attempt <= refusedAttempts; attempt++ {
select {
case <-retryReady:
case <-time.After(5 * time.Second):
t.Fatalf("timeout waiting for refused DialEdge attempt %d", attempt)
}
}
if count := atomic.LoadInt32(&dialCount); count != refusedAttempts {
t.Fatalf("expected %d real refused dials before Edge startup, got %d", refusedAttempts, count)
}
serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(serverCancel)
regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "delayed-node")
close(releaseRetry)
select {
case <-regCh:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for delayed initial registration")
}
if count := atomic.LoadInt32(&dialCount); count != refusedAttempts+1 {
t.Fatalf("expected %d dials (3 real refusals + 1 success), got %d", refusedAttempts+1, count)
}
// Exactly one session must be established.
select {
case <-regCh:
t.Fatal("unexpected second registration; supervisor should hold one session")
case <-time.After(200 * time.Millisecond):
}
}
// TestSupervisorUnlimitedAttemptsPassesTen verifies that an explicit
// max_attempts=0 (unlimited) keeps retrying past the former hard-coded 10-attempt
// cap and eventually connects. SDD S16/S17.
func TestSupervisorUnlimitedAttemptsPassesTen(t *testing.T) {
useTempCwd(t)
host, port, addr := freeAddrForBootstrap(t)
serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(serverCancel)
regCh, _ := startAcceptingEdge(t, serverCtx, host, port, "unlimited-node")
const failFirst = 15 // exceeds the old hard-coded 10-attempt cap
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
if atomic.AddInt32(&dialCount, 1) <= failFirst {
return nil, errors.New("edge unavailable")
}
return transport.DialEdge(ctx, edgeAddr, token, logger)
})
fakeSleeper := func(context.Context, time.Duration) {}
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
defer func() { _ = app.Stop(context.Background()) }()
select {
case <-regCh:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for registration after >10 unlimited attempts")
}
count := atomic.LoadInt32(&dialCount)
if count != failFirst+1 {
t.Fatalf("expected %d dials, got %d", failFirst+1, count)
}
if count <= 10 {
t.Fatalf("unlimited policy must retry past 10 attempts, got %d", count)
}
}
// TestSupervisorFiniteExhaustionExitsOne verifies that a finite reconnect policy
// makes exactly max_attempts initial-connect attempts and then terminates the
// node with exit code 1. SDD S17.
func TestSupervisorFiniteExhaustionExitsOne(t *testing.T) {
useTempCwd(t)
_, _, addr := freeAddrForBootstrap(t) // no server listening
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
return nil, errors.New("edge unavailable")
})
fakeSleeper := func(context.Context, time.Duration) {}
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 3, IntervalSec: 0},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
defer func() { _ = app.Stop(context.Background()) }()
select {
case sig := <-app.Wait():
if sig.ExitCode != 1 {
t.Fatalf("expected exit code 1, got %d", sig.ExitCode)
}
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for finite exhaustion shutdown")
}
if count := atomic.LoadInt32(&dialCount); count != 3 {
t.Fatalf("expected exactly 3 initial-connect attempts, got %d", count)
}
}
// TestSupervisorFatalFailureDoesNotRetry verifies that an authoritative
// registration rejection is non-retryable: the supervisor exits with code 1
// after a single attempt without any retry. SDD S17.
func TestSupervisorFatalFailureDoesNotRetry(t *testing.T) {
useTempCwd(t)
host, port, addr := freeAddrForBootstrap(t)
serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(serverCancel)
startRejectingEdge(t, serverCtx, host, port, "bad-token")
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
return transport.DialEdge(ctx, edgeAddr, token, logger)
})
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 5, IntervalSec: 0},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
defer func() { _ = app.Stop(context.Background()) }()
select {
case sig := <-app.Wait():
if sig.ExitCode != 1 {
t.Fatalf("expected exit code 1 on fatal failure, got %d", sig.ExitCode)
}
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for fatal shutdown")
}
if count := atomic.LoadInt32(&dialCount); count != 1 {
t.Fatalf("fatal failure must not retry: expected 1 dial, got %d", count)
}
}
// TestSupervisorShutdownCancelsRetry verifies that a local shutdown cancels a
// pending retry sleep, ends cleanly with no error, and runs no additional
// attempt. SDD S17.
func TestSupervisorShutdownCancelsRetry(t *testing.T) {
useTempCwd(t)
_, _, addr := freeAddrForBootstrap(t) // no server listening
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(_ context.Context, _, _ string, _ *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
return nil, errors.New("edge unavailable")
})
sleeping := make(chan struct{}, 1)
fakeSleeper := func(ctx context.Context, _ time.Duration) {
select {
case sleeping <- struct{}{}:
default:
}
<-ctx.Done() // block until the supervisor is cancelled by shutdown
}
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer), bootstrap.WithSleeper(fakeSleeper)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
// Wait until the supervisor has dialed once and parked in the retry sleep.
select {
case <-sleeping:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for supervisor to enter retry sleep")
}
countAtSleep := atomic.LoadInt32(&dialCount)
if countAtSleep != 1 {
t.Fatalf("expected exactly 1 dial before parking in sleep, got %d", countAtSleep)
}
// Local shutdown must cancel the sleep and stop cleanly with no error.
stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer stopCancel()
if err := app.Stop(stopCtx); err != nil {
t.Fatalf("app stop must be clean, got: %v", err)
}
if count := atomic.LoadInt32(&dialCount); count != countAtSleep {
t.Fatalf("shutdown must not trigger extra dial: had %d, got %d", countAtSleep, count)
}
}
// TestSupervisorShutdownWinsInFlightDialFatal fixes the ordering between a
// local stop and a dialer that observes cancellation but returns a fatal result
// afterward. Cancellation must produce a clean stop, no exit 1, and no retry.
func TestSupervisorShutdownWinsInFlightDialFatal(t *testing.T) {
useTempCwd(t)
_, _, addr := freeAddrForBootstrap(t)
started := make(chan struct{}, 1)
var dialCount int32
dialer := bootstrap.DialFunc(func(ctx context.Context, _, _ string, logger *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
started <- struct{}{}
<-ctx.Done()
return transport.DialEdge(context.Background(), "127.0.0.1:0", "tok", logger)
})
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
select {
case <-started:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for in-flight dial")
}
stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer stopCancel()
if err := app.Stop(stopCtx); err != nil {
t.Fatalf("app stop must be clean, got: %v", err)
}
if count := atomic.LoadInt32(&dialCount); count != 1 {
t.Fatalf("shutdown must not retry fatal late result: got %d dials", count)
}
select {
case sig := <-app.Wait():
t.Fatalf("shutdown must not request failure exit, got %+v", sig)
case <-time.After(100 * time.Millisecond):
}
}
// TestSupervisorShutdownClosesInFlightDialSuccess verifies that a dialer which
// observes cancellation and then returns a successful registered session cannot
// install or leak that late owner. The one registered session is closed by the
// post-dial cancellation fence and no failure exit or retry is emitted.
func TestSupervisorShutdownClosesInFlightDialSuccess(t *testing.T) {
useTempCwd(t)
host, port, addr := freeAddrForBootstrap(t)
serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(serverCancel)
regCh, disconnectCh := startAcceptingEdge(t, serverCtx, host, port, "late-success-node")
started := make(chan struct{}, 1)
var dialCount int32
dialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
atomic.AddInt32(&dialCount, 1)
started <- struct{}{}
<-ctx.Done()
lateCtx, lateCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer lateCancel()
return transport.DialEdge(lateCtx, edgeAddr, token, logger)
})
cfg := &config.NodeConfig{
Transport: config.TransportConf{EdgeAddr: addr, Token: "tok"},
Reconnect: config.ReconnectConf{MaxAttempts: 0, IntervalSec: 1},
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(dialer)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
if err := app.Start(startCtx); err != nil {
t.Fatalf("app start: %v", err)
}
select {
case <-started:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for in-flight dial")
}
stopCtx, stopCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer stopCancel()
stopDone := make(chan error, 1)
go func() { stopDone <- app.Stop(stopCtx) }()
select {
case <-regCh:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for late successful registration")
}
select {
case err := <-stopDone:
if err != nil {
t.Fatalf("app stop must be clean, got: %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for app stop")
}
select {
case <-disconnectCh:
case <-time.After(5 * time.Second):
t.Fatal("late successful session leaked after supervisor shutdown")
}
if count := atomic.LoadInt32(&dialCount); count != 1 {
t.Fatalf("shutdown must not retry late success: got %d dials", count)
}
select {
case <-regCh:
t.Fatal("unexpected second registration after shutdown")
case sig := <-app.Wait():
t.Fatalf("shutdown must not request failure exit, got %+v", sig)
case <-time.After(100 * time.Millisecond):
}
}
// TestReconnectSupervisorReestablishesSession verifies that after the edge-side
// client closes the connection, the supervisor reconnects and re-registers.
func TestReconnectSupervisorReestablishesSession(t *testing.T) {
@ -69,8 +556,11 @@ func TestReconnectSupervisorReestablishesSession(t *testing.T) {
serverCtx, serverCancel := context.WithTimeout(context.Background(), 10*time.Second)
t.Cleanup(serverCancel)
// regCh fires each time a RegisterRequest is handled.
// regCh fires each time a RegisterRequest is handled; readyCh fires when a
// connection completes its ready handshake so the test can close it only after
// the session is fully established.
regCh := make(chan struct{}, 8)
readyCh := make(chan struct{}, 8)
acceptedCh := make(chan *toki.TcpClient, 8)
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
@ -86,6 +576,7 @@ func TestReconnectSupervisorReestablishesSession(t *testing.T) {
}, nil
},
)
bindEdgeReadyResponder(client, readyCh)
acceptedCh <- client
return client
})
@ -100,7 +591,19 @@ func TestReconnectSupervisorReestablishesSession(t *testing.T) {
Logging: config.LoggingConf{Level: "error"},
}
app := fx.New(bootstrap.Module(cfg), fx.NopLogger)
// Wrap the real dialer so the test can wait for the initial session to be
// fully established before inducing a disconnect (OnStart is non-blocking).
established := make(chan struct{}, 1)
var establishOnce sync.Once
fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
if err == nil {
establishOnce.Do(func() { established <- struct{}{} })
}
return res, err
})
app := fx.New(bootstrap.Module(cfg, bootstrap.WithDialer(fakeDialer)), fx.NopLogger)
startCtx, startCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer startCancel()
@ -109,18 +612,30 @@ func TestReconnectSupervisorReestablishesSession(t *testing.T) {
}
defer func() { _ = app.Stop(context.Background()) }()
// Wait for first registration.
// Wait for first registration and full establishment.
select {
case <-regCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first registration")
}
select {
case <-established:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial session establishment")
}
var firstClient *toki.TcpClient
select {
case firstClient = <-acceptedCh:
case <-time.After(time.Second):
t.Fatal("timeout waiting for first accepted client")
}
// Wait for the ready handshake before closing so the disconnect is induced on a
// fully-established session, not mid-SignalReady.
select {
case <-readyCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial ready handshake")
}
// Force edge-side close → remote disconnect on node.
_ = firstClient.Close()
@ -143,6 +658,7 @@ func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
t.Cleanup(serverCancel)
acceptedCh := make(chan *toki.TcpClient, 4)
readyCh := make(chan struct{}, 4)
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
@ -155,6 +671,7 @@ func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
}, nil
},
)
bindEdgeReadyResponder(client, readyCh)
acceptedCh <- client
return client
})
@ -164,10 +681,17 @@ func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
t.Cleanup(func() { server.Stop() })
// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
// established fires once the initial session is fully established so the test
// can induce the disconnect without racing the initial registration.
established := make(chan struct{}, 1)
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
if atomic.AddInt32(&dialCount, 1) == 1 {
return transport.DialEdge(ctx, edgeAddr, token, logger)
res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
if err == nil {
established <- struct{}{}
}
return res, err
}
return nil, errors.New("fake reconnect failure")
})
@ -191,13 +715,25 @@ func TestReconnectSupervisorExhaustsRetries(t *testing.T) {
}
defer func() { _ = app.Stop(context.Background()) }()
// Wait for first connection then close it from the edge side.
// Wait for the initial session to establish, then close it from the edge side.
select {
case <-established:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial session establishment")
}
var firstClient *toki.TcpClient
select {
case firstClient = <-acceptedCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first connection")
}
// Close only after the ready handshake so the first session counts as an
// established connection, not a mid-SignalReady failure.
select {
case <-readyCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial ready handshake")
}
_ = firstClient.Close()
// Supervisor exhausts retries → Shutdown(ExitCode(1)).
@ -227,6 +763,7 @@ func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
t.Cleanup(serverCancel)
acceptedCh := make(chan *toki.TcpClient, 4)
readyCh := make(chan struct{}, 4)
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
client := toki.NewTcpClient(conn, 30, 10, bootstrapEdgeParserMap())
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
@ -239,6 +776,7 @@ func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
}, nil
},
)
bindEdgeReadyResponder(client, readyCh)
acceptedCh <- client
return client
})
@ -248,10 +786,17 @@ func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
t.Cleanup(func() { server.Stop() })
// Fake dialer: first call uses real DialEdge; subsequent calls fail immediately.
// established fires once the initial session is fully established so the test
// can induce the disconnect without racing the initial registration.
established := make(chan struct{}, 1)
var dialCount int32
fakeDialer := bootstrap.DialFunc(func(ctx context.Context, edgeAddr, token string, logger *zap.Logger) (*transport.RegisterResult, error) {
if atomic.AddInt32(&dialCount, 1) == 1 {
return transport.DialEdge(ctx, edgeAddr, token, logger)
res, err := transport.DialEdge(ctx, edgeAddr, token, logger)
if err == nil {
established <- struct{}{}
}
return res, err
}
return nil, errors.New("fake reconnect failure")
})
@ -284,13 +829,25 @@ func TestReconnectSupervisorPolicyTimingAndLimit(t *testing.T) {
}
defer func() { _ = app.Stop(context.Background()) }()
// Wait for first connection then close it from the edge side.
// Wait for the initial session to establish, then close it from the edge side.
select {
case <-established:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial session establishment")
}
var firstClient *toki.TcpClient
select {
case firstClient = <-acceptedCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for first connection")
}
// Close only after the ready handshake so the first session counts as an
// established connection, not a mid-SignalReady failure.
select {
case <-readyCh:
case <-time.After(3 * time.Second):
t.Fatal("timeout waiting for initial ready handshake")
}
_ = firstClient.Close()
// Supervisor exhausts all 10 retries → Shutdown(ExitCode(1)).

View file

@ -2,7 +2,6 @@ package bootstrap
import (
"context"
"fmt"
"sync"
"time"
@ -14,10 +13,11 @@ import (
"iop/packages/go/observability"
)
// runtimeSupervisor owns the node's single active edge connection and the
// reconnect loop. It is created once per fx lifecycle and drives
// connect/reconnect/shutdown. The mutex guards current across the reconnect
// goroutine and the OnStop hook.
// runtimeSupervisor owns the node's single active edge connection across its
// entire lifetime. A single supervisor goroutine serially performs the initial
// dial, waits for the established session to end, and reconnects, so at most one
// in-flight dial and one active session exist at any time. The mutex guards
// current across the supervisor goroutine and the OnStop hook.
type runtimeSupervisor struct {
cfg *config.NodeConfig
logger *zap.Logger
@ -28,25 +28,22 @@ type runtimeSupervisor struct {
mu sync.Mutex
current *runtimeOwner
cancelSupervisor context.CancelFunc
done chan struct{}
}
// start establishes the initial connection, launches the reconnect supervisor
// goroutine, and starts the metrics server. It maps a connect failure to the
// same "bootstrap: %w" error the OnStart hook returned before extraction.
func (s *runtimeSupervisor) start(ctx context.Context) error {
owner, err := connectRuntime(ctx, s.cfg, s.logger, s.dialer)
if err != nil {
return fmt.Errorf("bootstrap: %w", err)
}
s.mu.Lock()
s.current = owner
s.mu.Unlock()
// start launches the metrics server and the single supervisor goroutine, then
// returns immediately. Establishing the initial connection is the supervisor's
// responsibility, so a transient Edge outage at startup no longer fails the fx
// OnStart hook (SDD S16).
func (s *runtimeSupervisor) start(_ context.Context) error {
supCtx, cancel := context.WithCancel(context.Background())
s.cancelSupervisor = cancel
s.done = make(chan struct{})
go s.run(supCtx, owner.sess)
go func() {
defer close(s.done)
s.run(supCtx)
}()
go func() {
if err := observability.ServeMetrics(s.cfg.Metrics.Port); err != nil {
@ -56,84 +53,107 @@ func (s *runtimeSupervisor) start(ctx context.Context) error {
return nil
}
// run waits for the active session to end, then drives a reconnect. It exits on
// supervisor cancellation, a local shutdown, or when reconnect returns nil
// (cancelled or exhausted).
func (s *runtimeSupervisor) run(supCtx context.Context, sess *transport.Session) {
// run drives the full connectivity lifecycle: connect (initial), wait for the
// active session to end, then reconnect, repeating until the supervisor is
// cancelled, a fatal error occurs, finite attempts are exhausted, or a local
// shutdown ends the session.
func (s *runtimeSupervisor) run(supCtx context.Context) {
owner := s.connect(supCtx, true)
if owner == nil {
return
}
for {
select {
case <-supCtx.Done():
return
case <-sess.Done():
case <-owner.sess.Done():
}
// Guard against context cancellation racing with disconnect.
// Guard against supervisor cancellation racing with disconnect.
select {
case <-supCtx.Done():
return
default:
}
if sess.IsLocalShutdown() {
if owner.sess.IsLocalShutdown() {
return
}
newOwner := s.reconnect(supCtx)
if newOwner == nil {
owner = s.connect(supCtx, false)
if owner == nil {
return
}
sess = newOwner.sess
}
}
// reconnect runs the bounded reconnect loop. On success it swaps in the new
// owner and returns it. On supervisor cancellation it returns nil without
// changing state. On exhaustion it closes the current owner, requests node
// shutdown, and returns nil. The attempt/interval defaults, sleep, dial timeout,
// and logging match the prior inline loop.
func (s *runtimeSupervisor) reconnect(supCtx context.Context) *runtimeOwner {
// connect runs the bounded connect/retry loop shared by the initial dial and
// established-session reconnect. On success it installs the new owner and
// returns it. It returns nil after supervisor cancellation, a fatal
// (non-retryable) failure, or finite attempt exhaustion; the latter two request
// a node shutdown with exit code 1 before returning.
//
// The reconnect path sleeps before every attempt; the initial path issues its
// first dial immediately for fast startup and sleeps only before retries. A
// finite policy (max_attempts>0) makes exactly max_attempts attempts before
// exhausting; an unlimited policy (max_attempts=0) retries until cancelled.
func (s *runtimeSupervisor) connect(supCtx context.Context, initial bool) *runtimeOwner {
maxAttempts := s.cfg.Reconnect.MaxAttempts
if maxAttempts <= 0 {
maxAttempts = 10
}
intervalSec := s.cfg.Reconnect.IntervalSec
if intervalSec < 0 {
intervalSec = 10
}
unlimited := maxAttempts == 0
for attempt := 1; attempt <= maxAttempts; attempt++ {
select {
case <-supCtx.Done():
return nil
default:
for attempt := 1; unlimited || attempt <= maxAttempts; attempt++ {
if attempt > 1 || !initial {
s.sleeper(supCtx, time.Duration(intervalSec)*time.Second)
}
s.sleeper(supCtx, time.Duration(intervalSec)*time.Second)
if supCtx.Err() != nil {
return nil
}
s.logger.Info("reconnecting to edge",
s.logger.Info("connecting to edge",
zap.Bool("initial", initial),
zap.Int("attempt", attempt),
zap.Int("max_attempts", maxAttempts),
zap.Bool("unlimited", unlimited),
zap.Int("interval_sec", intervalSec),
)
dialCtx, dialCancel := context.WithTimeout(supCtx, 30*time.Second)
o, dialErr := connectRuntime(dialCtx, s.cfg, s.logger, s.dialer)
owner, err := connectRuntime(dialCtx, s.cfg, s.logger, s.dialer)
dialCancel()
if dialErr == nil {
s.swapOwner(o)
return o
// Local shutdown wins when a dial completes concurrently with supervisor
// cancellation. A dialer may observe cancellation and still return a late
// success or fatal result; never install that owner or turn it into exit 1.
if supCtx.Err() != nil {
if owner != nil {
owner.close()
}
return nil
}
if err == nil {
s.swapOwner(owner)
return owner
}
s.logger.Warn("reconnect attempt failed",
if transport.IsFatalConnectError(err) {
s.fatal(err)
return nil
}
s.logger.Warn("connect attempt failed",
zap.Bool("initial", initial),
zap.Int("attempt", attempt),
zap.Int("max_attempts", maxAttempts),
zap.Error(dialErr),
zap.Error(err),
)
}
// A local shutdown that lands on the final finite attempt falls out of the
// loop; treat it as cancellation rather than exhaustion.
if supCtx.Err() != nil {
return nil
}
s.exhaust(maxAttempts)
return nil
}
@ -149,9 +169,28 @@ func (s *runtimeSupervisor) swapOwner(owner *runtimeOwner) {
}
}
// fatal clears the current owner, logs the non-retryable failure, and requests a
// node shutdown with a non-zero exit code.
func (s *runtimeSupervisor) fatal(err error) {
s.clearCurrent()
s.logger.Error("connect failed with non-retryable error, shutting down node",
zap.Error(err),
)
_ = s.shutdowner.Shutdown(fx.ExitCode(1))
}
// exhaust clears the current owner, logs the exhaustion, and requests a node
// shutdown with a non-zero exit code.
func (s *runtimeSupervisor) exhaust(maxAttempts int) {
s.clearCurrent()
s.logger.Error("reconnect exhausted, shutting down node",
zap.Int("max_attempts", maxAttempts),
)
_ = s.shutdowner.Shutdown(fx.ExitCode(1))
}
// clearCurrent detaches and closes the current owner, if any.
func (s *runtimeSupervisor) clearCurrent() {
s.mu.Lock()
prev := s.current
s.current = nil
@ -159,22 +198,17 @@ func (s *runtimeSupervisor) exhaust(maxAttempts int) {
if prev != nil {
prev.close()
}
s.logger.Error("reconnect exhausted, shutting down node",
zap.Int("max_attempts", maxAttempts),
)
_ = s.shutdowner.Shutdown(fx.ExitCode(1))
}
// stop cancels the supervisor goroutine and closes the current connection.
// stop cancels the supervisor goroutine, waits for it to exit, and closes the
// current connection. Waiting guarantees no retry sleep or dial survives the
// shutdown and no additional attempt runs afterward.
func (s *runtimeSupervisor) stop() {
if s.cancelSupervisor != nil {
s.cancelSupervisor()
}
s.mu.Lock()
owner := s.current
s.current = nil
s.mu.Unlock()
if owner != nil {
owner.close()
if s.done != nil {
<-s.done
}
s.clearCurrent()
}

View file

@ -2,6 +2,7 @@ package transport
import (
"context"
"errors"
"fmt"
"net"
"strconv"
@ -13,6 +14,60 @@ import (
iop "iop/proto/gen/iop"
)
// connectFailureClass classifies a DialEdge failure so the connectivity
// supervisor can distinguish transient failures it should retry from fatal ones
// that must terminate the Node.
type connectFailureClass int
const (
// connectFailureRetryable marks a transient failure (dial refused, network
// unreachable, timeout, transient register transport failure) that the
// supervisor should retry under the reconnect policy.
connectFailureRetryable connectFailureClass = iota
// connectFailureFatal marks a non-retryable failure (invalid address or
// authoritative registration rejection) that retrying cannot recover.
connectFailureFatal
)
// ConnectError wraps a DialEdge failure with a retry classification while
// preserving the underlying cause for errors.Is/errors.As and any transport
// status it carries. The supervisor consumes the classification through
// IsFatalConnectError rather than comparing error strings.
type ConnectError struct {
class connectFailureClass
cause error
}
// Error returns the wrapped cause's message unchanged.
func (e *ConnectError) Error() string {
if e.cause == nil {
return "transport: connect error"
}
return e.cause.Error()
}
// Unwrap exposes the original cause so errors.Is/errors.As reach it.
func (e *ConnectError) Unwrap() error { return e.cause }
// Retryable reports whether the failure is a transient one the supervisor
// should retry.
func (e *ConnectError) Retryable() bool { return e.class == connectFailureRetryable }
func wrapConnectError(class connectFailureClass, cause error) *ConnectError {
return &ConnectError{class: class, cause: cause}
}
// IsFatalConnectError reports whether err (or any wrapped cause) is a fatal,
// non-retryable connect failure. Errors that are not classified ConnectErrors
// are treated as retryable transient failures.
func IsFatalConnectError(err error) bool {
var ce *ConnectError
if errors.As(err, &ce) {
return !ce.Retryable()
}
return false
}
const (
heartbeatIntervalSec = 2
// heartbeatWaitSec is kept above heartbeatIntervalSec as defence in depth:
@ -56,28 +111,39 @@ type RegisterResult struct {
func DialEdge(ctx context.Context, addr, token string, logger *zap.Logger) (*RegisterResult, error) {
host, portStr, err := net.SplitHostPort(addr)
if err != nil {
return nil, fmt.Errorf("transport: invalid addr %q: %w", addr, err)
// Local address preparation error: retrying cannot recover it.
return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid addr %q: %w", addr, err))
}
port, err := strconv.Atoi(portStr)
if err != nil {
return nil, fmt.Errorf("transport: invalid port %q: %w", portStr, err)
return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, err))
}
if port < 1 || port > 65535 {
// A numeric port outside the TCP range is still a local configuration
// error. Preserve a standard range cause so callers can inspect it with
// errors.Is/errors.As instead of relying on the message.
rangeErr := &strconv.NumError{Func: "Atoi", Num: portStr, Err: strconv.ErrRange}
return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: invalid port %q: %w", portStr, rangeErr))
}
dialer := net.Dialer{KeepAlive: 15 * time.Second}
conn, err := dialer.DialContext(ctx, "tcp", net.JoinHostPort(host, strconv.Itoa(port)))
if err != nil {
return nil, fmt.Errorf("transport: dial edge %s: %w", addr, err)
// Edge unavailable / network transient failure: retryable.
return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: dial edge %s: %w", addr, err))
}
client := toki.NewTcpClient(&writeDeadlineConn{Conn: conn, timeout: tcpWriteTimeout}, heartbeatIntervalSec, heartbeatWaitSec, nodeParserMap())
resp, err := registerWithEdge(ctx, client, token, logger)
if err != nil {
_ = client.Close()
return nil, fmt.Errorf("transport: register: %w", err)
// Register request transport failure or context cancellation: retryable.
return nil, wrapConnectError(connectFailureRetryable, fmt.Errorf("transport: register: %w", err))
}
if !resp.GetAccepted() {
_ = client.Close()
return nil, fmt.Errorf("transport: register rejected: %s", resp.GetReason())
// Authoritative registration rejection (bad token/credential): fatal.
return nil, wrapConnectError(connectFailureFatal, fmt.Errorf("transport: register rejected: %s", resp.GetReason()))
}
sess := newSession(client, logger, resp.GetNodeId(), resp.GetAlias())

View file

@ -1,9 +1,19 @@
package transport
import (
"context"
"errors"
"net"
"strconv"
"strings"
"testing"
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
iop "iop/proto/gen/iop"
)
type recordingConn struct {
@ -27,6 +37,166 @@ type noopAddr string
func (a noopAddr) Network() string { return string(a) }
func (a noopAddr) String() string { return string(a) }
// closedLocalAddr returns a loopback address that has no listener, so a dial to
// it is refused deterministically.
func closedLocalAddr(t *testing.T) string {
t.Helper()
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
addr := l.Addr().String()
_ = l.Close()
return addr
}
func rejectingEdgeParserMap() toki.ParserMap {
return toki.ParserMap{
toki.TypeNameOf(&iop.RegisterRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.RegisterRequest{}
return m, proto.Unmarshal(b, m)
},
}
}
// startRejectingEdge starts a mock edge that answers RegisterRequest with
// Accepted=false and the given reason.
func startRejectingEdge(t *testing.T, ctx context.Context, reason string) string {
t.Helper()
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
addr := l.Addr().String()
host, portStr, _ := net.SplitHostPort(addr)
port, _ := strconv.Atoi(portStr)
_ = l.Close()
server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
client := toki.NewTcpClient(conn, 30, 10, rejectingEdgeParserMap())
toki.AddRequestListenerTyped[*iop.RegisterRequest, *iop.RegisterResponse](
&client.Communicator,
func(_ *iop.RegisterRequest) (*iop.RegisterResponse, error) {
return &iop.RegisterResponse{Accepted: false, Reason: reason}, nil
},
)
return client
})
if err := server.Start(ctx); err != nil {
t.Fatalf("start rejecting edge: %v", err)
}
t.Cleanup(func() { server.Stop() })
return addr
}
// TestDialEdgeClassifiesConnectFailures verifies that DialEdge classifies its
// failures into fatal (invalid address, authoritative registration rejection)
// and retryable (Edge unavailable) while preserving the underlying cause. REFACTOR-2.
func TestDialEdgeClassifiesConnectFailures(t *testing.T) {
logger := zap.NewNop()
t.Run("invalid address is fatal", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_, err := DialEdge(ctx, "missing-port", "tok", logger)
if err == nil {
t.Fatal("expected error for invalid address")
}
if !IsFatalConnectError(err) {
t.Fatalf("invalid address must be fatal, got retryable: %v", err)
}
var ce *ConnectError
if !errors.As(err, &ce) || ce.Retryable() {
t.Fatalf("expected non-retryable ConnectError, got %v", err)
}
})
t.Run("non-numeric port is fatal", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_, err := DialEdge(ctx, "127.0.0.1:not-a-port", "tok", logger)
if err == nil {
t.Fatal("expected error for invalid port")
}
if !IsFatalConnectError(err) {
t.Fatalf("invalid port must be fatal, got retryable: %v", err)
}
var numErr *strconv.NumError
if !errors.As(err, &numErr) || !errors.Is(err, strconv.ErrSyntax) {
t.Fatalf("expected preserved strconv syntax cause, got %v", err)
}
})
for _, port := range []string{"0", "-1", "65536"} {
port := port
t.Run("out-of-range port "+port+" is fatal", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
_, err := DialEdge(ctx, net.JoinHostPort("127.0.0.1", port), "tok", logger)
if err == nil {
t.Fatalf("expected error for out-of-range port %s", port)
}
if !IsFatalConnectError(err) {
t.Fatalf("out-of-range port %s must be fatal, got retryable: %v", port, err)
}
var ce *ConnectError
if !errors.As(err, &ce) || ce.Retryable() {
t.Fatalf("expected non-retryable ConnectError, got %v", err)
}
var numErr *strconv.NumError
if !errors.As(err, &numErr) || !errors.Is(err, strconv.ErrRange) {
t.Fatalf("expected preserved strconv range cause, got %v", err)
}
if numErr.Num != port {
t.Fatalf("range cause port = %q, want %q", numErr.Num, port)
}
})
}
t.Run("connection refused is retryable and preserves cause", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
_, err := DialEdge(ctx, closedLocalAddr(t), "tok", logger)
if err == nil {
t.Fatal("expected error for refused connection")
}
if IsFatalConnectError(err) {
t.Fatalf("refused connection must be retryable, got fatal: %v", err)
}
var ce *ConnectError
if !errors.As(err, &ce) || !ce.Retryable() {
t.Fatalf("expected retryable ConnectError, got %v", err)
}
// The underlying network error must remain reachable through the wrap.
var opErr *net.OpError
if !errors.As(err, &opErr) {
t.Fatalf("expected wrapped *net.OpError cause, got %v", err)
}
})
t.Run("registration rejected is fatal and preserves reason", func(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
addr := startRejectingEdge(t, ctx, "bad-token")
_, err := DialEdge(ctx, addr, "tok", logger)
if err == nil {
t.Fatal("expected error for rejected registration")
}
if !IsFatalConnectError(err) {
t.Fatalf("rejected registration must be fatal, got retryable: %v", err)
}
if !strings.Contains(err.Error(), "bad-token") {
t.Fatalf("expected rejection reason preserved in error, got %v", err)
}
})
t.Run("unclassified error defaults to retryable", func(t *testing.T) {
if IsFatalConnectError(errors.New("plain")) {
t.Fatal("unclassified error must not be fatal")
}
})
}
func TestWriteDeadlineConnSetsAndClearsDeadline(t *testing.T) {
base := &recordingConn{}
conn := &writeDeadlineConn{Conn: base, timeout: 25 * time.Millisecond}

View file

@ -33,6 +33,10 @@ func nodeParserMap() toki.ParserMap {
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.NodeConfigRefreshRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.NodeConfigRefreshRequest{}
return m, proto.Unmarshal(b, m)

View file

@ -2,7 +2,9 @@ package transport
import (
"context"
"fmt"
"sync"
"time"
toki "git.toki-labs.com/toki/proto-socket/go"
"go.uber.org/zap"
@ -154,6 +156,28 @@ func (s *Session) SetHandler(h Handler) {
s.mu.Unlock()
}
// SignalReady tells edge this node has applied the config from RegisterResponse
// and installed its message handler, so edge may now open dispatch eligibility
// and pump any waiters stranded while the node was offline. It MUST be called
// after SetHandler: edge dispatches run/tunnel requests in response to this
// signal, and a request arriving before the handler is installed would be
// dropped. A non-ready ack (stale connection) or a transport error is returned so
// the caller tears the session down and lets the supervisor reconnect.
func (s *Session) SignalReady(timeout time.Duration) error {
resp, err := toki.SendRequestTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
&s.client.Communicator,
&iop.NodeReadyRequest{NodeId: s.nodeID},
timeout,
)
if err != nil {
return err
}
if !resp.GetReady() {
return fmt.Errorf("edge rejected ready signal: %s", resp.GetReason())
}
return nil
}
// NodeID returns the session's node ID.
func (s *Session) NodeID() string {
return s.nodeID

View file

@ -287,3 +287,55 @@ func TestSessionProviderTunnelRequest_NilAndErrHandler(t *testing.T) {
sess.Close()
}
}
// TestSessionSignalReady verifies the node→edge dispatch-ready handshake:
// SignalReady sends a NodeReadyRequest carrying the session's node id and returns
// nil only when the edge acks ready, and an error when the edge rejects the
// signal (stale connection) so the caller can tear down and reconnect.
func TestSessionSignalReady(t *testing.T) {
newReadyPipe := func(t *testing.T, ready bool, reason string) *transport.Session {
t.Helper()
edgeConn, nodeConn := net.Pipe()
edgeParser := toki.ParserMap{
toki.TypeNameOf(&iop.NodeReadyRequest{}): func(b []byte) (proto.Message, error) {
m := &iop.NodeReadyRequest{}
return m, proto.Unmarshal(b, m)
},
}
nodeParser := toki.ParserMap{
toki.TypeNameOf(&iop.NodeReadyResponse{}): func(b []byte) (proto.Message, error) {
m := &iop.NodeReadyResponse{}
return m, proto.Unmarshal(b, m)
},
}
edgeSide := toki.NewTcpClient(edgeConn, 0, 0, edgeParser)
nodeSide := toki.NewTcpClient(nodeConn, 0, 0, nodeParser)
t.Cleanup(func() { edgeSide.Close(); nodeSide.Close() })
toki.AddRequestListenerTyped[*iop.NodeReadyRequest, *iop.NodeReadyResponse](
&edgeSide.Communicator,
func(req *iop.NodeReadyRequest) (*iop.NodeReadyResponse, error) {
if req.GetNodeId() != "node-ready-test" {
return &iop.NodeReadyResponse{Ready: false, Reason: "unexpected node id"}, nil
}
return &iop.NodeReadyResponse{Ready: ready, Reason: reason}, nil
},
)
return transport.ExportNewSession(nodeSide, zap.NewNop(), "node-ready-test", "alias")
}
t.Run("ready ack succeeds", func(t *testing.T) {
sess := newReadyPipe(t, true, "")
if err := sess.SignalReady(2 * time.Second); err != nil {
t.Fatalf("SignalReady on ready ack: %v", err)
}
})
t.Run("non-ready ack errors", func(t *testing.T) {
sess := newReadyPipe(t, false, "superseded")
err := sess.SignalReady(2 * time.Second)
if err == nil {
t.Fatal("SignalReady must error when edge rejects the ready signal")
}
})
}

View file

@ -0,0 +1,334 @@
#!/usr/bin/env bash
# edge-node-reconnect-diagnostic.sh - repo-internal diagnostic for edge/node reconnect
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
echo "[diagnostic] Starting edge-node-reconnect-diagnostic (repo-internal)..."
TMP_DIR=$(mktemp -d /tmp/iop-reconnect-diag-XXXXXX)
kill_process_tree() {
local pid="$1"
local child
if [ -z "$pid" ]; then
return 0
fi
if command -v pgrep >/dev/null 2>&1; then
while IFS= read -r child; do
if [ -n "$child" ]; then
kill_process_tree "$child"
fi
done <<EOF_CHILDREN
$(pgrep -P "$pid" 2>/dev/null || true)
EOF_CHILDREN
fi
kill "$pid" 2>/dev/null || true
}
cleanup() {
set +e
echo "[diagnostic] Cleaning up..."
kill_process_tree "${NODE_PID:-}"
kill_process_tree "${EDGE_PID:-}"
rm -rf "$TMP_DIR" 2>/dev/null || true
}
trap cleanup EXIT
EDGE_CONFIG="$TMP_DIR/edge.yaml"
NODE_CONFIG="$TMP_DIR/node.yaml"
EDGE_OUT="$TMP_DIR/edge.log"
NODE_OUT="$TMP_DIR/node.log"
CONSOLE_FIFO="$TMP_DIR/console.fifo"
# Registration wait ceiling; honor the dev override used by the verification command.
BIND_TIMEOUT="${IOP_DEV_RECONNECT_BIND_TIMEOUT:-45}"
PORT=$((30000 + RANDOM % 10000))
BOOTSTRAP_PORT=$((20000 + RANDOM % 10000))
EDGE_METRICS_PORT=$((40000 + RANDOM % 10000))
NODE_METRICS_PORT=$((50000 + RANDOM % 10000))
# Create fake-cli.sh. The Node's own [node-message] transcript below is the
# node-side payload record used for the Edge relay comparison.
MOCK_CLI="$TMP_DIR/fake-cli.sh"
cat <<'EOF' > "$MOCK_CLI"
#!/usr/bin/env sh
while IFS= read -r line; do
token=$(printf '%s\n' "$line" | sed -nE 's/.*(IOP_E2E_[A-Z0-9_]+).*/\1/p' | head -n 1)
if [ -n "$token" ]; then
printf '%s\n%s_TAIL\n' "$token" "$token"
else
printf 'IOP_E2E_UNKNOWN\n'
fi
done
EOF
chmod +x "$MOCK_CLI"
# Create edge config
cat <<EOF > "$EDGE_CONFIG"
server:
listen: "127.0.0.1:$PORT"
metrics:
port: $EDGE_METRICS_PORT
bootstrap:
listen: "127.0.0.1:$BOOTSTRAP_PORT"
artifact_dir: "$TMP_DIR/artifacts"
nodes:
- id: test-node
alias: test-node
token: test-token
adapters:
cli:
enabled: true
profiles:
fake-cli:
command: "$MOCK_CLI"
persistent: true
response_idle_timeout_ms: 1000
console:
adapter: cli
target: fake-cli
session_id: default
EOF
# Create node config
cat <<EOF > "$NODE_CONFIG"
transport:
edge_addr: "127.0.0.1:$PORT"
token: test-token
reconnect:
interval_sec: 1
max_attempts: 0
metrics:
port: $NODE_METRICS_PORT
node:
id: test-node
alias: test-node
EOF
# Make console FIFO
mkfifo "$CONSOLE_FIFO"
# 1. Start edge
echo "[diagnostic] Starting edge.sh..."
IOP_EDGE_CONFIG="$EDGE_CONFIG" "$REPO_ROOT/scripts/dev/edge.sh" < "$CONSOLE_FIFO" > "$EDGE_OUT" 2>&1 &
EDGE_PID=$!
# Open FIFO for writing (blocks until edge reads it)
exec 3> "$CONSOLE_FIFO"
# 2. Start node
echo "[diagnostic] Starting node.sh..."
IOP_NODE_CONFIG="$NODE_CONFIG" "$REPO_ROOT/scripts/dev/node.sh" > "$NODE_OUT" 2>&1 &
NODE_PID=$!
# Wait for node registration
echo "[diagnostic] Awaiting node registration..."
deadline=$((SECONDS + BIND_TIMEOUT))
while ! grep -q 'connected reason="registered"' "$EDGE_OUT" 2>/dev/null; do
if (( SECONDS >= deadline )); then
echo "[diagnostic] Timeout waiting for node registration" >&2
cat "$EDGE_OUT" >&2
exit 1
fi
sleep 0.5
done
echo "[diagnostic] Node registered"
# Wait for registration to settle
sleep 1
# Send /nodes command
echo "/nodes" >&3
sleep 0.5
# Send message 1
echo "Convert token IOP_E2E_HELLO_BASIC and reply only with converted token" >&3
# Wait for message 1 complete
deadline=$((SECONDS + 15))
while ! grep -q "complete run_id=" "$EDGE_OUT" 2>/dev/null; do
if (( SECONDS >= deadline )); then
echo "[diagnostic] Timeout waiting for message 1 completion" >&2
exit 1
fi
sleep 0.5
done
echo "[diagnostic] Message 1 completed"
# Send message 2
# Save baseline of edge log count to check next completion
EDGE_BASELINE=$(wc -l < "$EDGE_OUT" | tr -d ' ')
echo "Convert token IOP_E2E_HELLO_FORMAL and reply only with converted token" >&3
# Wait for message 2 complete (new lines only)
deadline=$((SECONDS + 15))
while ! tail -n +"$((EDGE_BASELINE + 1))" "$EDGE_OUT" | grep -q "complete run_id=" 2>/dev/null; do
if (( SECONDS >= deadline )); then
echo "[diagnostic] Timeout waiting for message 2 completion" >&2
exit 1
fi
sleep 0.5
done
echo "[diagnostic] Message 2 completed"
# Send commands
echo "/capabilities" >&3
sleep 0.2
echo "/transport" >&3
sleep 0.2
echo "/sessions" >&3
sleep 0.2
echo "/terminate-session" >&3
sleep 0.5
# Restart Node
echo "[diagnostic] Killing node for reconnect test..."
kill_process_tree "$NODE_PID"
sleep 2
EDGE_BASELINE=$(wc -l < "$EDGE_OUT" | tr -d ' ')
echo "[diagnostic] Restarting node..."
IOP_NODE_CONFIG="$NODE_CONFIG" "$REPO_ROOT/scripts/dev/node.sh" >> "$NODE_OUT" 2>&1 &
NODE_PID=$!
# Wait for re-registration (new lines only)
deadline=$((SECONDS + BIND_TIMEOUT))
while ! tail -n +"$((EDGE_BASELINE + 1))" "$EDGE_OUT" | grep -E 'connected reason="registered"|node ready' 2>/dev/null; do
if (( SECONDS >= deadline )); then
echo "[diagnostic] Timeout waiting for node reconnect" >&2
exit 1
fi
sleep 0.5
done
echo "[diagnostic] Node reconnected"
# Wait for reconnect to settle
sleep 1
# Send message 3 post-reconnect
EDGE_BASELINE=$(wc -l < "$EDGE_OUT" | tr -d ' ')
echo "Convert token IOP_E2E_PING_BASIC and reply only with converted token" >&3
# Wait for message 3 complete
deadline=$((SECONDS + 10))
while ! tail -n +"$((EDGE_BASELINE + 1))" "$EDGE_OUT" | grep -q "complete run_id=" 2>/dev/null; do
if (( SECONDS >= deadline )); then
echo "[diagnostic] Timeout waiting for message 3 completion" >&2
exit 1
fi
sleep 0.5
done
echo "[diagnostic] Message 3 completed"
# Exit edge console
echo "/exit" >&3
exec 3>&-
# Wait for edge to finish
wait $EDGE_PID || true
echo "=== EDGE LOG ==="
cat "$EDGE_OUT"
echo "=== NODE LOG ==="
cat "$NODE_OUT"
# ---- Fail-fast verification of the direct two-process transcript ----
# Every branch below exits non-zero on a missing/extra token, a Node-vs-Edge
# payload mismatch, a mis-ordered or duplicated terminal event, a missing
# reconnect, or a missing command response. There is no silent skip: the previous
# version matched run_id on the [node0-msg] lines (which carry none), so its whole
# loop `continue`d and it printed PASS having compared nothing.
echo "[diagnostic] Verifying payload sequence, terminal ordering, and command responses..."
fail() {
printf '[diagnostic] VALIDATION FAILED: %b\n' "$*" >&2
exit 1
}
# Runs are dispatched in this order: two before the reconnect, one after. Each run
# streams exactly the extracted token then that token with a _TAIL suffix.
EXPECTED_TOKENS=(IOP_E2E_HELLO_BASIC IOP_E2E_HELLO_FORMAL IOP_E2E_PING_BASIC)
# Edge-observed run ids, in order, from the authoritative start events.
mapfile -t RUN_IDS < <(grep -oE '\[node0-evt\] start run_id=[A-Za-z0-9_-]+' "$EDGE_OUT" | sed 's/.*run_id=//')
if [ "${#RUN_IDS[@]}" -ne "${#EXPECTED_TOKENS[@]}" ]; then
fail "expected ${#EXPECTED_TOKENS[@]} runs, edge start events show ${#RUN_IDS[@]}: ${RUN_IDS[*]:-<none>}"
fi
# Reconnect must have occurred: at least the initial + one reconnect registration.
CONNECTED_COUNT=$(grep -c 'connected reason="registered"' "$EDGE_OUT" || true)
if [ "${CONNECTED_COUNT:-0}" -lt 2 ]; then
fail "expected >=2 connected(registered) events (initial + reconnect), got ${CONNECTED_COUNT:-0}"
fi
for i in "${!RUN_IDS[@]}"; do
rid="${RUN_IDS[$i]}"
token="${EXPECTED_TOKENS[$i]}"
echo "[diagnostic] Checking run $((i + 1)) run_id=$rid token=$token"
start_ln=$(grep -nE "\[node0-evt\] start run_id=${rid}\$" "$EDGE_OUT" | head -n1 | cut -d: -f1 || true)
[ -n "$start_ln" ] || fail "run $rid: missing edge start event"
# The terminal event must appear exactly once for this run.
complete_ln=$(grep -nE "\[node0-evt\] complete run_id=${rid}( |\$)" "$EDGE_OUT" | cut -d: -f1 || true)
complete_count=$(printf '%s\n' "$complete_ln" | grep -c . || true)
[ "$complete_count" -eq 1 ] || fail "run $rid: expected exactly one complete event, got $complete_count"
[ "$complete_ln" -gt "$start_ln" ] || fail "run $rid: complete event precedes start"
# Payload tokens strictly between this run's start and complete, in order.
run_payloads=$(sed -n "${start_ln},${complete_ln}p" "$EDGE_OUT" \
| grep -oE '\[node0-msg\] IOP_E2E_[A-Z0-9_]+' | sed 's/.*\] //' || true)
want_payloads=$(printf '%s\n%s_TAIL' "$token" "$token")
if [ "$run_payloads" != "$want_payloads" ]; then
fail "run $rid ($token): edge payload sequence mismatch\n got: [$run_payloads]\n want: [$want_payloads]"
fi
# The terminal event must come strictly after the last payload line.
last_msg_rel=$(sed -n "${start_ln},${complete_ln}p" "$EDGE_OUT" | grep -nE '\[node0-msg\]' | tail -n1 | cut -d: -f1 || true)
[ -n "$last_msg_rel" ] || fail "run $rid: no payload lines in trace"
last_msg_abs=$((start_ln + last_msg_rel - 1))
[ "$complete_ln" -gt "$last_msg_abs" ] || fail "run $rid: complete event is not after the last payload"
# The Node process logs the same run with [node-event] start/complete bounds.
# Capture both the prefixed first payload line and any continuation lines
# emitted by a multiline adapter delta, then compare the sequence directly
# with the Edge relay for this exact run id.
node_start_ln=$(grep -nE "^\[node-event\] start run_id=${rid}\$" "$NODE_OUT" | cut -d: -f1 || true)
node_start_count=$(printf '%s\n' "$node_start_ln" | grep -c . || true)
[ "$node_start_count" -eq 1 ] || fail "run $rid: expected one Node start event, got $node_start_count"
node_complete_ln=$(grep -nE "^\[node-event\] complete run_id=${rid}( |\$)" "$NODE_OUT" | cut -d: -f1 || true)
node_complete_count=$(printf '%s\n' "$node_complete_ln" | grep -c . || true)
[ "$node_complete_count" -eq 1 ] || fail "run $rid: expected one Node complete event, got $node_complete_count"
[ "$node_complete_ln" -gt "$node_start_ln" ] || fail "run $rid: Node complete event precedes start"
node_payloads=$(sed -n "${node_start_ln},${node_complete_ln}p" "$NODE_OUT" \
| sed -nE '/^\[node-message\] IOP_E2E_[A-Z0-9_]+$/{s/^\[node-message\] //;p;b}; /^IOP_E2E_[A-Z0-9_]+$/p' || true)
if [ "${IOP_DEV_RECONNECT_INJECT_FAIL:-}" != "" ] && [ "$i" -eq "$((${#RUN_IDS[@]} - 1))" ]; then
# Test-only fault injection: drop the final Node payload from the last
# run so the actual per-run Node-vs-Edge comparison must fail-fast.
node_payloads=$(printf '%s\n' "$node_payloads" | sed '$d')
fi
if [ "$node_payloads" != "$run_payloads" ]; then
fail "run $rid: Node-vs-Edge payload sequence mismatch\n node: [$node_payloads]\n edge: [$run_payloads]"
fi
node_last_msg_rel=$(sed -n "${node_start_ln},${node_complete_ln}p" "$NODE_OUT" \
| grep -nE '^\[node-message\] |^IOP_E2E_' | tail -n1 | cut -d: -f1 || true)
[ -n "$node_last_msg_rel" ] || fail "run $rid: no Node payload lines in trace"
node_last_msg_abs=$((node_start_ln + node_last_msg_rel - 1))
[ "$node_complete_ln" -gt "$node_last_msg_abs" ] || fail "run $rid: Node complete event is not after the last payload"
done
# Required command responses in the edge transcript (secret-safe: fixed markers only).
require_cmd() {
grep -qE "$2" "$EDGE_OUT" || fail "missing $1 command response"
}
require_cmd "/nodes" '^ node0 = test-node'
require_cmd "/capabilities" '\[node0-capabilities\]'
require_cmd "/transport" '\[node0-transport\]'
require_cmd "/sessions" '\[node0-sessions\]'
require_cmd "/terminate-session" 'terminated session .* node=node0'
echo "[diagnostic] PASS: ${#RUN_IDS[@]} runs verified — payload sequence, one terminal after the last payload, Node==Edge; reconnect observed; all five command responses present."