sdk/go: keep a handle's relay connection across calls

docs/sdk.md

@@ -156,6 +156,16 @@ TLS handshakes share the existing dial timeout/cancellation budget. The agent
 connection then uses HTTP/1.1 `Upgrade: silkd` and remains a bidirectional
 stream; the guest protocol and port-forwarding frames do not change.
 
+Data-plane calls share a handle's relay connection: after a call the SDK
+keeps the connection for 30 seconds (`WithKeepAlive` tunes the window; 0
+dials per call) and the next call on that handle sends its request on it, so
+a busy handle pays the dial, upgrade and TLS handshake once. A kept
+connection counts as live for `idle_hibernate_seconds` until it closes, so
+keep the window below that setting; `Close` and `Hibernate` drop it at once.
+Long-lived streams (`Watch`, `OpenPty`, `DialPort`, an LSP session) take a
+connection of their own, and a guest whose silkd predates the back-to-back
+protocol gets one connection per call as before.
+
 An explicit scheme in an owner, redirect, peer, or `Attach` address wins;
 a bare address inherits the entry client's scheme. Trust settings are shared
 across these connections. The SDK does not translate private addresses to

e2e/cmd/rpcbench/main.go

@@ -1,8 +1,9 @@
 // rpcbench measures the one-connection-per-RPC overhead on a live node and
-// what a pre-dialed spare connection would buy (H-4's decision data): mode A
-// dials+upgrades per RPC like the SDK does today; mode B keeps one dialed
-// connection ahead, hiding the handshake behind the previous call. Run by
-// hand against a claimed sandbox:
+// what the alternatives buy: mode A dials+upgrades per RPC by hand; mode C
+// is the SDK's own path, one kept connection serving RPCs back to back; mode
+// B keeps one hand-dialed connection ahead, hiding the handshake behind the
+// previous call (H-4's decision data). Run by hand against a claimed
+// sandbox:
 //
 //	rpcbench -addr <node> -token <api token> -template <ref> -n 200
 package main
@@ -70,7 +71,20 @@ func run(addr, token, template string, n int) error {
 		}
 		a = append(a, time.Since(start))
 	}
-	report("A dial-per-RPC (today)", a)
+	report("A dial-per-RPC", a)
+
+	if _, err := sb.Stat(ctx, "/"); err != nil {
+		return err
+	}
+	c := make([]time.Duration, 0, n)
+	for range n {
+		start := time.Now()
+		if _, err := sb.Stat(ctx, "/"); err != nil {
+			return err
+		}
+		c = append(c, time.Since(start))
+	}
+	report("C SDK keep-alive", c)
 
 	spare := make(chan net.Conn, 1)
 	errs := make(chan error, 1)
@@ -104,11 +118,11 @@ func run(addr, token, template string, n int) error {
 		b = append(b, time.Since(start))
 	}
 	cancel()
-	report("B pre-dialed spare   ", b)
+	report("B pre-dialed spare", b)
 	return nil
 }
 
-// statRPC: the protocol is one RPC per connection.
+// statRPC drives one hand-dialed connection through one RPC and closes it.
 func statRPC(conn net.Conn) error {
 	defer func() { _ = conn.Close() }()
 	sc := silkd.NewConn(conn)
@@ -172,7 +186,7 @@ func dialAgent(ctx context.Context, addr, id, token string) (net.Conn, error) {
 func report(label string, samples []time.Duration) {
 	slices.Sort(samples)
 	pct := func(p float64) time.Duration { return samples[int(p*float64(len(samples)-1))] }
-	fmt.Printf("%s  n=%d p50=%.2fms p90=%.2fms p99=%.2fms\n",
+	fmt.Printf("%-22s n=%d p50=%.2fms p90=%.2fms p99=%.2fms\n",
 		label, len(samples), ms(pct(0.50)), ms(pct(0.90)), ms(pct(0.99)))
 }
 

mcp/server_test.go

@@ -11,8 +11,19 @@ import (
 	"net/http/httptest"
 	"strings"
 	"testing"
+
+	"github.com/cocoonstack/sandbox/protocol/wire"
+	"github.com/cocoonstack/sandbox/sdk/go/silkd/silkdtest"
 )
 
+var infoFrame = func() string {
+	frame, err := wire.EncodeResponse(&wire.InfoResp{Version: "test", Proto: wire.KeepAliveProto})
+	if err != nil {
+		panic(err)
+	}
+	return string(frame) + "\n"
+}()
+
 func TestServeSpeaksMCP(t *testing.T) {
 	srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) bool {
 		if !strings.HasSuffix(r.URL.Path, "/checkpoint") {
@@ -66,20 +77,9 @@ func TestServeSpeaksMCP(t *testing.T) {
 }
 
 func TestExecKeepsOutputWhenTheGuestDrops(t *testing.T) {
-	srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) bool {
-		if !strings.HasSuffix(r.URL.Path, "/agent") {
-			return false
-		}
-		conn, _, err := http.NewResponseController(w).Hijack()
-		if err != nil {
-			t.Errorf("hijack: %v", err)
-			return true
-		}
-		_, _ = io.WriteString(conn, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: silkd\r\nConnection: Upgrade\r\n\r\n"+
-			`{"type":"started","pid":7}`+"\n"+`{"type":"stdout","data":"cGFydGlhbA=="}`+"\n")
-		_ = conn.Close()
-		return true
-	})
+	srv := newTestServer(t, agentRoute(t, func(string) (string, bool) {
+		return `{"type":"started","pid":7}` + "\n" + `{"type":"stdout","data":"cGFydGlhbA=="}` + "\n", true
+	}))
 	replies := serveLines(t, srv,
 		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"create_sandbox","arguments":{}}}`,
 		`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"exec","arguments":{"sandbox_id":"sb_1","command":"yes"}}}`,
@@ -93,36 +93,12 @@ func TestExecKeepsOutputWhenTheGuestDrops(t *testing.T) {
 
 func TestReadFileStopsAtTheCap(t *testing.T) {
 	chunk := `{"type":"data","data":"` + base64.StdEncoding.EncodeToString(make([]byte, 256<<10)) + `"}` + "\n"
-	srv := newTestServer(t, func(w http.ResponseWriter, r *http.Request) bool {
-		if !strings.HasSuffix(r.URL.Path, "/agent") {
-			return false
-		}
-		conn, _, err := http.NewResponseController(w).Hijack()
-		if err != nil {
-			t.Errorf("hijack: %v", err)
-			return true
-		}
-		defer conn.Close()
-		br := bufio.NewReader(conn)
-		if _, err = io.WriteString(conn, "HTTP/1.1 101 Switching Protocols\r\nUpgrade: silkd\r\nConnection: Upgrade\r\n\r\n"); err != nil {
-			return true
-		}
-		req, err := br.ReadString('\n')
-		if err != nil {
-			return true
-		}
+	srv := newTestServer(t, agentRoute(t, func(req string) (string, bool) {
 		if strings.Contains(req, `"op":"fs_stat"`) {
-			_, _ = io.WriteString(conn, `{"type":"stat","info":{"kind":"file","size":0}}`+"\n")
-			return true
+			return `{"type":"stat","info":{"kind":"file","size":0}}` + "\n", false
 		}
-		for range 8 {
-			if _, err := io.WriteString(conn, chunk); err != nil {
-				return true
-			}
-		}
-		_, _ = io.WriteString(conn, `{"type":"done"}`+"\n")
-		return true
-	})
+		return strings.Repeat(chunk, 8) + `{"type":"done"}` + "\n", false
+	}))
 	replies := serveLines(t, srv,
 		`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"create_sandbox","arguments":{}}}`,
 		`{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"read_file","arguments":{"sandbox_id":"sb_1","path":"/dev/zero"}}}`,
@@ -184,6 +160,35 @@ func newTestServer(t *testing.T, extra func(http.ResponseWriter, *http.Request)
 	return srv
 }
 
+func agentRoute(t *testing.T, reply func(req string) (frames string, drop bool)) func(http.ResponseWriter, *http.Request) bool {
+	t.Helper()
+	return func(w http.ResponseWriter, r *http.Request) bool {
+		if !strings.HasSuffix(r.URL.Path, "/agent") {
+			return false
+		}
+		conn, err := silkdtest.Upgrade(w)
+		if err != nil {
+			t.Errorf("upgrade: %v", err)
+			return true
+		}
+		defer conn.Close()
+		br := bufio.NewReader(conn)
+		for {
+			req, err := br.ReadString('\n')
+			if err != nil {
+				return true
+			}
+			frames, drop := infoFrame, false
+			if !strings.Contains(req, `"op":"info"`) {
+				frames, drop = reply(req)
+			}
+			if _, err := io.WriteString(conn, frames); err != nil || drop {
+				return true
+			}
+		}
+	}
+}
+
 func serveLines(t *testing.T, srv *server, lines ...string) map[int]map[string]any {
 	t.Helper()
 	var out bytes.Buffer

protocol/wire/frame.go

@@ -20,6 +20,8 @@ const (
 	// ProtoVersion is stamped into every request as "v"; silkd ignores
 	// unknown fields, which is the forward-compatibility story.
 	ProtoVersion = 1
+	// KeepAliveProto is the InfoResp.Proto from which silkd serves RPCs back to back on one connection.
+	KeepAliveProto = 2
 	// MaxFrame mirrors silkd's frame cap.
 	MaxFrame = 8 << 20
 	// BulkChunk mirrors silkd's BULK_CHUNK, the payload of one bulk data frame.

sdk/go/client.go

@@ -129,6 +129,7 @@ type Client struct {
 	apiToken  string
 	hc        *http.Client
 	tlsConfig *tls.Config
+	keepAlive time.Duration
 }
 
 // Connect returns a client for a sandboxd node.
@@ -142,7 +143,7 @@ func Connect(addr string, opts ...ClientOption) (*Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	c := &Client{addr: strings.TrimPrefix(u.String(), "http://"), scheme: u.Scheme, hc: &http.Client{}}
+	c := &Client{addr: strings.TrimPrefix(u.String(), "http://"), scheme: u.Scheme, hc: &http.Client{}, keepAlive: keepAliveIdle}
 	for _, opt := range opts {
 		opt(c)
 	}

sdk/go/keepalive.go

@@ -0,0 +1,175 @@
+package sandbox
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"sync"
+	"sync/atomic"
+	"syscall"
+	"time"
+
+	"github.com/cocoonstack/sandbox/protocol/wire"
+	"github.com/cocoonstack/sandbox/sdk/go/silkd"
+)
+
+const (
+	keepAliveIdle  = 30 * time.Second
+	keepAliveConns = 8
+)
+
+// WithKeepAlive bounds how long a handle keeps an idle relay connection for
+// its next call (default 30s; 0 dials per call). An open connection holds
+// the sandbox's idle-hibernate clock, so keep it below the deployment's
+// idle_hibernate_seconds.
+func WithKeepAlive(idle time.Duration) ClientOption {
+	return func(c *Client) { c.keepAlive = idle }
+}
+
+// agentPool parks a handle's idle relay connections between calls.
+type agentPool struct {
+	mu   sync.Mutex
+	idle []*agentConn
+}
+
+// take returns a parked connection whose peer is still there, or nil.
+func (p *agentPool) take() *agentConn {
+	for {
+		p.mu.Lock()
+		n := len(p.idle)
+		if n == 0 {
+			p.mu.Unlock()
+			return nil
+		}
+		c := p.idle[n-1]
+		p.idle = p.idle[:n-1]
+		p.mu.Unlock()
+		if c.timer.Stop() && c.quiet() {
+			return c
+		}
+		_ = c.Close()
+	}
+}
+
+// park keeps c for the next call and closes it once idle has passed.
+func (p *agentPool) park(c *agentConn, idle time.Duration) {
+	p.mu.Lock()
+	if len(p.idle) >= keepAliveConns {
+		p.mu.Unlock()
+		_ = c.Close()
+		return
+	}
+	if c.timer == nil {
+		c.timer = time.AfterFunc(idle, func() { p.evict(c) })
+	} else {
+		c.timer.Reset(idle)
+	}
+	p.idle = append(p.idle, c)
+	p.mu.Unlock()
+}
+
+func (p *agentPool) drain() {
+	p.mu.Lock()
+	idle := p.idle
+	p.idle = nil
+	p.mu.Unlock()
+	for _, c := range idle {
+		c.timer.Stop()
+		_ = c.Close()
+	}
+}
+
+func (p *agentPool) evict(c *agentConn) {
+	p.mu.Lock()
+	if i := slices.Index(p.idle, c); i >= 0 {
+		p.idle = slices.Delete(p.idle, i, i+1)
+	}
+	p.mu.Unlock()
+	_ = c.Close()
+}
+
+// agentConn is one relayed silkd connection, with the TCP socket's raw handle for the liveness peek.
+type agentConn struct {
+	*silkd.Conn
+	sock   syscall.RawConn
+	peeked bool
+	peekFn func(uintptr) bool
+	timer  *time.Timer
+}
+
+func newAgentConn(raw *upgradedConn) *agentConn {
+	c := &agentConn{Conn: silkd.NewConn(raw)}
+	if sc, ok := raw.tcp.(syscall.Conn); ok {
+		c.sock, _ = sc.SyscallConn()
+	}
+	c.peekFn = c.peek
+	return c
+}
+
+// quiet reports whether the parked connection's peer has neither hung up nor spoken.
+func (c *agentConn) quiet() bool {
+	return c.sock != nil && c.sock.Read(c.peekFn) == nil && c.peeked
+}
+
+// sent sends req on c and hands it back, closing it on failure.
+func (c *agentConn) sent(req wire.Request) (*agentConn, error) {
+	if err := c.Send(req); err != nil {
+		_ = c.Close()
+		return nil, err
+	}
+	return c, nil
+}
+
+// probeProto pipelines info ahead of req and reads the daemon's answer.
+func (c *agentConn) probeProto(ctx context.Context, req wire.Request) (uint32, error) {
+	if err := c.Send(wire.Info{}); err != nil {
+		return 0, err
+	}
+	if err := c.Send(req); err != nil {
+		return 0, err
+	}
+	stop := context.AfterFunc(ctx, func() { _ = c.Close() })
+	defer stop()
+	info, err := expect[wire.InfoResp](ctx, c.Conn)
+	if err != nil {
+		return 0, err
+	}
+	return max(info.Proto, 1), nil
+}
+
+// lease is one RPC's hold on a relay connection: close drops it, reuse parks it for the next call.
+type lease struct {
+	s    *Sandbox
+	conn *agentConn
+	stop func() bool
+	over atomic.Bool
+}
+
+// done ends the RPC: a terminal frame, error frames included, parks the connection; any other failure drops it.
+func (l *lease) done(err error) error {
+	if _, frame := errors.AsType[*wire.ErrorResp](err); err == nil || frame {
+		l.reuse()
+	} else {
+		l.close()
+	}
+	return err
+}
+
+func (l *lease) reuse() {
+	if !l.over.CompareAndSwap(false, true) {
+		return
+	}
+	if !l.stop() || !canProbe || l.s.c.keepAlive <= 0 || l.s.proto.Load() < wire.KeepAliveProto {
+		_ = l.conn.Close()
+		return
+	}
+	l.s.pool.park(l.conn, l.s.c.keepAlive)
+}
+
+func (l *lease) close() {
+	if !l.over.CompareAndSwap(false, true) {
+		return
+	}
+	l.stop()
+	_ = l.conn.Close()
+}