feat: DTail fork — server/client feature development

Squashed development of the snonux/dtail fork's product code (internal/, cmd/)
since diverging from mimecast/dtail. Major areas:

- Read/output path: the former "turbo" channel-less path is now the single,
  default server-side read/output path for cat/grep/tail and MapReduce; the old
  channel-based path and its config/env toggles were removed.
- MapReduce: single aggregate implementation (server + serverless) fed directly
  by a processor pipeline, with input-exhausted finalization via the shutdown
  coordinator; high-concurrency and data-race fixes.
- Journal source reads (journal:unit.service) via journalctl, Linux-gated behind
  a journal-v1 capability.
- Auth-key fast reconnect: in-memory per-user public-key cache with TTL/max-keys,
  registered over an authenticated session (AUTHKEY), checked before
  authorized_keys.
- Interactive query reload (--interactive-query) with SESSION START/UPDATE
  generation boundaries and capability negotiation.
- Client-side deadlines: --timeout / --shutdownAfter as context deadlines;
  follow shutdown handling.
- Client logging: diagnostics-only daily log by default, opt-in payload tee via
  --log-payload.
- Numerous correctness fixes (buffer-pool double-recycle races, EOF-sentinel
  leaks, glob-expansion cap, TOCTOU in CSV parsing) with accompanying unit tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

5 files changed, 1502 insertions, 70 deletions

diff --git a/internal/clients/connectors/connector.go b/internal/clients/connectors/connector.go
index 3ab6a08..00de32e 100644
--- a/internal/clients/connectors/connector.go
+++ b/internal/clients/connectors/connector.go
@@ -2,8 +2,10 @@ package connectors
 
 import (
 	"context"
+	"time"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	sessionspec "github.com/mimecast/dtail/internal/session"
 )
 
 // Connector interface.
@@ -14,4 +16,20 @@ type Connector interface {
 	Server() string
 	// Handler for the connection.
 	Handler() handlers.Handler
+	// SupportsQueryUpdates reports whether the connected server advertised
+	// runtime query replacement support within the given timeout.
+	SupportsQueryUpdates(timeout time.Duration) bool
+	// ApplySessionSpec starts or updates the interactive session workload on an
+	// already connected server when query updates are supported.
+	ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error
+	// ApplySessionSpecWithGeneration starts or updates the interactive session
+	// workload using the provided committed generation as the base for the
+	// session update command.
+	ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error
+	// CommittedSession returns the last session spec and generation that the
+	// server acknowledged for this connection.
+	CommittedSession() (sessionspec.Spec, uint64, bool)
+	// RestoreCommittedSession resets the local committed session snapshot without
+	// advancing the generation.
+	RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool)
 }
diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
index 5c3d455..b98f815 100644
--- a/internal/clients/connectors/serverconnection.go
+++ b/internal/clients/connectors/serverconnection.go
@@ -2,20 +2,38 @@ package connectors
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"io"
+	"net"
+	"os"
+	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
-	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
+	sessionspec "github.com/mimecast/dtail/internal/session"
 	"github.com/mimecast/dtail/internal/ssh/client"
 
 	"golang.org/x/crypto/ssh"
 )
 
+// SSHSettings provides the connection settings needed by ServerConnection.
+type SSHSettings interface {
+	SSHPort() int
+	SSHConnectTimeout() time.Duration
+}
+
+const (
+	defaultSSHConnectTimeout = 2 * time.Second
+	defaultSSHPort           = 2222
+	defaultCapabilityWait    = 250 * time.Millisecond
+)
+
 // ServerConnection represents a connection to a single remote dtail server via
 // SSH protocol.
 type ServerConnection struct {
@@ -28,30 +46,61 @@ type ServerConnection struct {
 	config          *ssh.ClientConfig
 	handler         handlers.Handler
 	commands        []string
+	sessionSpec     sessionspec.Spec
+	sessionState    committedSessionState
+	interactive     bool
+	authKeyPath     string
+	authKeyDisabled bool
 	hostKeyCallback client.HostKeyCallback
-	throttlingDone  bool
+	// throttleReleased ensures the throttle slot is returned to throttleCh
+	// exactly once, even if both the early-release path in handle() and the
+	// deferred cleanup in Start() execute concurrently or the same goroutine
+	// hits both paths.  sync.Once is safe for concurrent callers whereas the
+	// previous bool guard was not synchronized.
+	throttleReleased sync.Once
 }
 
+var _ Connector = (*ServerConnection)(nil)
+
 // NewServerConnection returns a new DTail SSH server connection.
 func NewServerConnection(server string, userName string,
 	authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback,
-	handler handlers.Handler, commands []string) *ServerConnection {
+	handler handlers.Handler, commands []string, sessionSpec sessionspec.Spec,
+	interactive bool, authKeyPath string, authKeyDisabled bool, settings SSHSettings) *ServerConnection {
 
 	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
+	sshConnectTimeout := defaultSSHConnectTimeout
+	defaultPort := defaultSSHPort
+	if settings != nil {
+		sshConnectTimeout = settings.SSHConnectTimeout()
+		defaultPort = settings.SSHPort()
+	}
+	if sshConnectTimeout <= 0 {
+		sshConnectTimeout = defaultSSHConnectTimeout
+	}
+	if defaultPort <= 0 {
+		defaultPort = defaultSSHPort
+	}
+
 	c := ServerConnection{
 		hostKeyCallback: hostKeyCallback,
 		server:          server,
 		handler:         handler,
 		commands:        commands,
+		sessionSpec:     sessionSpec,
+		interactive:     interactive,
+		authKeyPath:     resolveAuthKeyPath(authKeyPath),
+		authKeyDisabled: authKeyDisabled,
 		config: &ssh.ClientConfig{
-			User:            userName,
-			Auth:            authMethods,
-			HostKeyCallback: hostKeyCallback.Wrap(),
-			Timeout:         time.Second * 2,
+			User:    userName,
+			Auth:    authMethods,
+			Timeout: sshConnectTimeout,
+			// HostKeyCallback is assigned per-handshake in dial() so the
+			// callback can honour the handshake's context (see dial()).
 		},
 	}
 
-	c.initServerPort()
+	c.initServerPort(defaultPort)
 	return &c
 }
 
@@ -61,12 +110,42 @@ func (c *ServerConnection) Server() string { return c.server }
 // Handler returns the handler used for the connection.
 func (c *ServerConnection) Handler() handlers.Handler { return c.handler }
 
+// SupportsQueryUpdates reports whether the remote server advertised the
+// runtime query replacement capability. Older servers simply time out and
+// return false here without affecting the legacy command path.
+func (c *ServerConnection) SupportsQueryUpdates(timeout time.Duration) bool {
+	return supportsQueryUpdates(c.handler, timeout)
+}
+
+// ApplySessionSpec starts or updates the interactive session state on the
+// existing SSH connection when runtime query updates are supported.
+func (c *ServerConnection) ApplySessionSpec(spec sessionspec.Spec, timeout time.Duration) error {
+	return applySessionSpec(c.server, c.handler, &c.sessionState, spec, timeout)
+}
+
+// ApplySessionSpecWithGeneration starts or updates the interactive session
+// state using an explicit committed generation as the base for the update.
+func (c *ServerConnection) ApplySessionSpecWithGeneration(spec sessionspec.Spec, generation uint64, timeout time.Duration) error {
+	return applySessionSpecWithGeneration(c.server, c.handler, &c.sessionState, spec, generation, false, timeout)
+}
+
+// CommittedSession returns the last server-acknowledged session state.
+func (c *ServerConnection) CommittedSession() (sessionspec.Spec, uint64, bool) {
+	return c.sessionState.snapshot()
+}
+
+// RestoreCommittedSession resets the local session snapshot without advancing
+// the generation.
+func (c *ServerConnection) RestoreCommittedSession(spec sessionspec.Spec, generation uint64, committed bool) {
+	c.sessionState.restore(spec, generation, committed)
+}
+
 // Attempt to parse the server port address from the provided server FQDN.
-func (c *ServerConnection) initServerPort() {
+func (c *ServerConnection) initServerPort(defaultPort int) {
 	parts := strings.Split(c.server, ":")
 	if len(parts) == 1 {
 		c.hostname = c.server
-		c.port = config.Common.SSHPort
+		c.port = defaultPort
 		return
 	}
 
@@ -99,12 +178,15 @@ func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
 
 	go func() {
 		defer func() {
-			if !c.throttlingDone {
-				dlog.Client.Debug(c.server, "Unthrottling connection (1)",
+			// Release the throttle slot on the way out regardless of which
+			// code path already attempted it.  throttleReleased.Do guarantees
+			// the drain happens exactly once even if handle() already released
+			// the slot early (the fast path when the session is fully up).
+			c.throttleReleased.Do(func() {
+				dlog.Client.Debug(c.server, "Unthrottling connection (cleanup)",
 					len(throttleCh), cap(throttleCh))
-				c.throttlingDone = true
 				<-throttleCh
-			}
+			})
 			cancel()
 		}()
 
@@ -133,10 +215,33 @@ func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc,
 	address := fmt.Sprintf("%s:%d", c.hostname, c.port)
 	dlog.Client.Debug(c.server, "Dialing into the connection", address)
 
-	client, err := ssh.Dial("tcp", address, c.config)
+	// Use context-aware dialing to enable proper cancellation during connection establishment.
+	// TCP KeepAlive (30s) prevents silent connection failures on long-lived connections.
+	dialer := &net.Dialer{
+		Timeout:   c.config.Timeout, // Use the SSH config timeout (2 seconds)
+		KeepAlive: 30 * time.Second, // Standard Go default for connection health monitoring
+	}
+
+	// Establish TCP connection with context support for cancellation
+	conn, err := dialer.DialContext(ctx, "tcp", address)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to dial TCP connection to %s: %w", address, err)
 	}
+
+	// Perform SSH handshake over the established TCP connection. Build a
+	// per-handshake ssh.ClientConfig so the host-key callback is bound to
+	// ctx and unblocks cleanly if the handshake is cancelled (e.g. when the
+	// user aborts before responding to the unknown-host prompt).
+	handshakeConfig := *c.config
+	handshakeConfig.HostKeyCallback = c.hostKeyCallback.Wrap(ctx)
+	sshConn, chans, reqs, err := ssh.NewClientConn(conn, address, &handshakeConfig)
+	if err != nil {
+		conn.Close()
+		return fmt.Errorf("SSH handshake failed for %s: %w", address, err)
+	}
+
+	// Create SSH client from the connection components
+	client := ssh.NewClient(sshConn, chans, reqs)
 	defer client.Close()
 
 	return c.session(ctx, cancel, client, throttleCh)
@@ -149,7 +254,7 @@ func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFun
 	dlog.Client.Debug(c.server, "Creating SSH session")
 	session, err := client.NewSession()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to create SSH session for %s: %w", c.server, err)
 	}
 	defer session.Close()
 	return c.handle(ctx, cancel, session, throttleCh)
@@ -161,14 +266,14 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 	dlog.Client.Debug(c.server, "Creating handler for SSH session")
 	stdinPipe, err := session.StdinPipe()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get SSH session stdin pipe for %s: %w", c.server, err)
 	}
 	stdoutPipe, err := session.StdoutPipe()
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to get SSH session stdout pipe for %s: %w", c.server, err)
 	}
 	if err := session.Shell(); err != nil {
-		return err
+		return fmt.Errorf("failed to start SSH shell for %s: %w", c.server, err)
 	}
 
 	go func() {
@@ -191,22 +296,105 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 		}
 	}()
 
-	// Send all commands to client.
-	for _, command := range c.commands {
-		dlog.Client.Debug(command)
-		if err := c.handler.SendMessage(command); err != nil {
-			dlog.Client.Debug(err)
-		}
+	if c.authKeyDisabled {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration because auth-key is disabled")
+	} else {
+		c.sendAuthKeyRegistrationCommand()
+	}
+
+	if err := dispatchInitialCommands(c.server, c.handler, c.commands, c.interactive, c.sessionSpec, &c.sessionState); err != nil {
+		c.handler.Shutdown()
+		return err
 	}
 
-	if !c.throttlingDone {
-		dlog.Client.Debug(c.server, "Unthrottling connection (2)",
+	// Release the throttle slot as soon as the session is fully established so
+	// the next pending connection can proceed without waiting for this session
+	// to finish.  throttleReleased.Do is idempotent: if the deferred cleanup
+	// in Start() fires first (e.g. on a dial error path that never reaches
+	// here), the slot is still returned exactly once.
+	c.throttleReleased.Do(func() {
+		dlog.Client.Debug(c.server, "Unthrottling connection (session up)",
 			len(throttleCh), cap(throttleCh))
-		c.throttlingDone = true
 		<-throttleCh
-	}
+	})
 
 	<-ctx.Done()
 	c.handler.Shutdown()
 	return nil
 }
+
+// resolveAuthKeyPath returns the effective auth-key path. When the provided
+// path is non-empty it is used as-is. Otherwise the function falls back to
+// $HOME/.ssh/id_rsa. If HOME is also empty it returns "" so that the AUTHKEY
+// registration step (sendAuthKeyRegistrationCommand) will skip gracefully
+// instead of trying to open a path that the SSH library cannot expand.
+func resolveAuthKeyPath(authKeyPath string) string {
+	if strings.TrimSpace(authKeyPath) != "" {
+		return authKeyPath
+	}
+	homeDir := os.Getenv("HOME")
+	if homeDir == "" {
+		return ""
+	}
+	return filepath.Join(homeDir, ".ssh", "id_rsa")
+}
+
+func (c *ServerConnection) sendAuthKeyRegistrationCommand() {
+	authKeyPubPath := c.authKeyPath + ".pub"
+	authKeyPubBytes, err := os.ReadFile(authKeyPubPath)
+	if err != nil {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, unable to read public key", authKeyPubPath, err)
+		return
+	}
+
+	authKeyBase64, err := extractAuthKeyBase64(authKeyPubBytes)
+	if err != nil {
+		dlog.Client.Debug(c.server, "Skipping AUTHKEY registration, invalid public key file", authKeyPubPath, err)
+		return
+	}
+
+	if err := c.handler.SendMessage("AUTHKEY " + authKeyBase64); err != nil {
+		dlog.Client.Debug(c.server, "Unable to send AUTHKEY registration command", err)
+		return
+	}
+	dlog.Client.Debug(c.server, "Sent AUTHKEY registration command", authKeyPubPath)
+}
+
+func extractAuthKeyBase64(authKeyPubBytes []byte) (string, error) {
+	authKeyPubContent := string(authKeyPubBytes)
+	for _, line := range strings.Split(authKeyPubContent, "\n") {
+		trimmedLine := strings.TrimSpace(line)
+		if trimmedLine == "" || strings.HasPrefix(trimmedLine, "#") {
+			continue
+		}
+
+		fields := strings.Fields(trimmedLine)
+		if len(fields) < 2 {
+			return "", fmt.Errorf("expected authorized key format '<type> <base64-key> [comment]'")
+		}
+
+		authKeyBase64 := strings.TrimSpace(fields[1])
+		if _, err := base64.StdEncoding.DecodeString(authKeyBase64); err != nil {
+			return "", fmt.Errorf("invalid base64 public key: %w", err)
+		}
+
+		return authKeyBase64, nil
+	}
+
+	return "", fmt.Errorf("no public key found")
+}
+
+func supportsQueryUpdates(handler handlers.Handler, timeout time.Duration) bool {
+	if handler == nil {
+		return false
+	}
+
+	if timeout <= 0 {
+		timeout = defaultCapabilityWait
+	}
+	if !handler.WaitForCapabilities(timeout) {
+		return false
+	}
+
+	return handler.HasCapability(protocol.CapabilityQueryUpdateV1)
+}
diff --git a/internal/clients/connectors/serverconnection_test.go b/internal/clients/connectors/serverconnection_test.go
new file mode 100644
index 0000000..4e22dc7
--- /dev/null
+++ b/internal/clients/connectors/serverconnection_test.go
@@ -0,0 +1,880 @@
+package connectors
+
+import (
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/protocol"
+	sessionspec "github.com/mimecast/dtail/internal/session"
+
+	"golang.org/x/crypto/ssh"
+)
+
+// TestResolveAuthKeyPathNoLiteralPath is a regression test for the bug
+// described in task l6: when authKeyPath is empty and HOME is also unset,
+// resolveAuthKeyPath must return "" instead of a mangled path like
+// "/.ssh/id_rsa" or the literal "~/.ssh/id_rsa" that the SSH stack cannot use.
+func TestResolveAuthKeyPathNoLiteralPath(t *testing.T) {
+	// Unset HOME so the environment fallback is also empty.
+	t.Setenv("HOME", "")
+
+	got := resolveAuthKeyPath("")
+	if got != "" {
+		t.Fatalf("resolveAuthKeyPath(\"\") with empty HOME = %q; want \"\"", got)
+	}
+}
+
+// TestResolveAuthKeyPathExplicitPathPassedThrough verifies that a non-empty
+// explicit auth key path is returned unchanged.
+func TestResolveAuthKeyPathExplicitPathPassedThrough(t *testing.T) {
+	got := resolveAuthKeyPath("/custom/key")
+	if got != "/custom/key" {
+		t.Fatalf("resolveAuthKeyPath(\"/custom/key\") = %q; want \"/custom/key\"", got)
+	}
+}
+
+// TestResolveAuthKeyPathFallsBackToHome verifies that when authKeyPath is empty
+// but HOME is set, the function returns the expected default path.
+func TestResolveAuthKeyPathFallsBackToHome(t *testing.T) {
+	t.Setenv("HOME", "/home/testuser")
+
+	got := resolveAuthKeyPath("")
+	want := "/home/testuser/.ssh/id_rsa"
+	if got != want {
+		t.Fatalf("resolveAuthKeyPath(\"\") = %q; want %q", got, want)
+	}
+}
+
+func TestExtractAuthKeyBase64(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	t.Run("valid authorized key line", func(t *testing.T) {
+		pubKey := []byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n")
+
+		got, err := extractAuthKeyBase64(pubKey)
+		if err != nil {
+			t.Fatalf("Expected valid key, got error: %v", err)
+		}
+		if got != "AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA" {
+			t.Fatalf("Unexpected base64 payload: %s", got)
+		}
+	})
+
+	t.Run("invalid key format", func(t *testing.T) {
+		_, err := extractAuthKeyBase64([]byte("not-a-valid-authorized-key-line"))
+		if err == nil {
+			t.Fatalf("Expected parse error for invalid key format")
+		}
+	})
+
+	t.Run("invalid base64 payload", func(t *testing.T) {
+		_, err := extractAuthKeyBase64([]byte("ssh-ed25519 !!! not-valid\n"))
+		if err == nil {
+			t.Fatalf("Expected error for invalid base64 payload")
+		}
+	})
+}
+
+func TestSendAuthKeyRegistrationCommand(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	tempDir := t.TempDir()
+	privateKeyPath := filepath.Join(tempDir, "id_rsa")
+	publicKeyPath := privateKeyPath + ".pub"
+	if err := os.WriteFile(publicKeyPath,
+		[]byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA user@host\n"), 0600); err != nil {
+		t.Fatalf("Unable to write public key test file: %v", err)
+	}
+
+	handler := &mockHandler{}
+	conn := &ServerConnection{
+		server:      "srv1",
+		handler:     handler,
+		authKeyPath: privateKeyPath,
+	}
+
+	conn.sendAuthKeyRegistrationCommand()
+
+	if len(handler.commands) != 1 {
+		t.Fatalf("Expected one AUTHKEY command, got %d", len(handler.commands))
+	}
+	expected := "AUTHKEY AAAAC3NzaC1lZDI1NTE5AAAAIAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
+	if handler.commands[0] != expected {
+		t.Fatalf("Unexpected AUTHKEY command.\nexpected: %s\ngot:      %s", expected, handler.commands[0])
+	}
+}
+
+func TestNewServerConnectionUsesInjectedSettings(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := NewServerConnection(
+		"srv1",
+		"user",
+		nil,
+		testHostKeyCallback{},
+		&mockHandler{},
+		nil,
+		sessionspec.Spec{},
+		false,
+		"",
+		false,
+		testSSHSettings{port: 3022, timeout: 5 * time.Second},
+	)
+
+	if conn.hostname != "srv1" {
+		t.Fatalf("Expected hostname srv1, got %q", conn.hostname)
+	}
+	if conn.port != 3022 {
+		t.Fatalf("Expected injected port 3022, got %d", conn.port)
+	}
+	if conn.config.Timeout != 5*time.Second {
+		t.Fatalf("Expected injected timeout 5s, got %v", conn.config.Timeout)
+	}
+}
+
+func TestNewServerConnectionFallsBackToDefaults(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := NewServerConnection(
+		"srv1",
+		"user",
+		nil,
+		testHostKeyCallback{},
+		&mockHandler{},
+		nil,
+		sessionspec.Spec{},
+		false,
+		"",
+		false,
+		testSSHSettings{},
+	)
+
+	if conn.port != defaultSSHPort {
+		t.Fatalf("Expected default port %d, got %d", defaultSSHPort, conn.port)
+	}
+	if conn.config.Timeout != defaultSSHConnectTimeout {
+		t.Fatalf("Expected default timeout %v, got %v", defaultSSHConnectTimeout, conn.config.Timeout)
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdates(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+		},
+	}
+
+	if !conn.SupportsQueryUpdates(10 * time.Millisecond) {
+		t.Fatalf("expected query-update capability to be detected")
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdatesFallsBackForOlderServers(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{},
+	}
+
+	if conn.SupportsQueryUpdates(5 * time.Millisecond) {
+		t.Fatalf("expected old-server fallback when no capability is advertised")
+	}
+}
+
+func TestServerConnectionSupportsQueryUpdatesRequiresCapabilityFlag(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{
+			waitForCapabilities: true,
+		},
+	}
+
+	if conn.SupportsQueryUpdates(10 * time.Millisecond) {
+		t.Fatalf("expected capability wait success alone to be insufficient")
+	}
+}
+
+func TestServerConnectionApplySessionSpecStart(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		server: "srv1",
+		handler: &mockHandler{
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+			sessionAcks: []handlers.SessionAck{{
+				Action:     "start",
+				Generation: 1,
+			}},
+		},
+	}
+
+	spec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	if err := conn.ApplySessionSpec(spec, 10*time.Millisecond); err != nil {
+		t.Fatalf("ApplySessionSpec() error = %v", err)
+	}
+
+	mock := conn.handler.(*mockHandler)
+	if len(mock.commands) != 1 {
+		t.Fatalf("expected one session command, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 1 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecUpdateUsesNextGeneration(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 4},
+			{Action: "update", Generation: 5},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	updateSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("update ApplySessionSpec() error = %v", err)
+	}
+	if len(mock.commands) != 2 {
+		t.Fatalf("expected two session commands, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 5 || committedSpec.Regex != "WARN" {
+		t.Fatalf("unexpected committed session after update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecReappliesPreviousSpecForRollback(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 4},
+			{Action: "update", Generation: 5},
+			{Action: "update", Generation: 6},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	updateSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(updateSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("update ApplySessionSpec() error = %v", err)
+	}
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("rollback ApplySessionSpec() error = %v", err)
+	}
+	if len(mock.commands) != 3 {
+		t.Fatalf("expected three session commands, got %d", len(mock.commands))
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 6 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session after rollback: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecFallsBackForUnsupportedServer(t *testing.T) {
+	resetClientLogger(t)
+
+	conn := &ServerConnection{
+		handler: &mockHandler{},
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"}, 5*time.Millisecond)
+	if !errors.Is(err, ErrSessionUnsupported) {
+		t.Fatalf("expected ErrSessionUnsupported, got %v", err)
+	}
+}
+
+func TestRequireJournalCapability(t *testing.T) {
+	tests := []struct {
+		name                string
+		spec                sessionspec.Spec
+		waitForCapabilities bool
+		capabilities        map[string]bool
+		wantErr             error
+		wantServerError     bool
+	}{
+		{
+			name: "journal file with journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityJournalV1: true,
+			},
+		},
+		{
+			name: "journal file without journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+			wantErr:         ErrJournalUnsupported,
+			wantServerError: true,
+		},
+		{
+			name: "journal file without capabilities advertisement",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"journal:ssh.service"},
+			},
+			wantErr:         ErrJournalUnsupported,
+			wantServerError: true,
+		},
+		{
+			name: "regular file without journal capability",
+			spec: sessionspec.Spec{
+				Mode:  omode.CatClient,
+				Files: []string{"/var/log/app.log"},
+			},
+			waitForCapabilities: true,
+			capabilities: map[string]bool{
+				protocol.CapabilityQueryUpdateV1: true,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			handler := &mockHandler{
+				waitForCapabilities: tc.waitForCapabilities,
+				capabilities:        tc.capabilities,
+			}
+
+			err := requireJournalCapability("srv1", handler, tc.spec, 10*time.Millisecond)
+			if !errors.Is(err, tc.wantErr) {
+				t.Fatalf("requireJournalCapability() error = %v, want %v", err, tc.wantErr)
+			}
+			if got := handler.serverError != ""; got != tc.wantServerError {
+				t.Fatalf("server error recorded = %v, want %v", got, tc.wantServerError)
+			}
+			if tc.wantServerError && !strings.Contains(handler.serverError, protocol.CapabilityJournalV1) {
+				t.Fatalf("server error %q does not mention %s", handler.serverError, protocol.CapabilityJournalV1)
+			}
+		})
+	}
+}
+
+func TestDispatchInitialCommandsRejectsJournalWithoutCapability(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	spec := sessionspec.Spec{
+		Mode:  omode.CatClient,
+		Files: []string{"journal:ssh.service"},
+	}
+
+	err := dispatchInitialCommands("srv1", handler, []string{"cat: journal:ssh.service ."}, false, spec, &committedSessionState{})
+	if !errors.Is(err, ErrJournalUnsupported) {
+		t.Fatalf("expected ErrJournalUnsupported, got %v", err)
+	}
+	if len(handler.commands) != 0 {
+		t.Fatalf("expected no commands to be sent, got %#v", handler.commands)
+	}
+	if handler.Status() != 1 {
+		t.Fatalf("handler status = %d, want 1", handler.Status())
+	}
+}
+
+func TestDispatchInitialCommandsRejectsInteractiveJournalWithoutCapability(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	spec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"journal:ssh.service"},
+	}
+
+	err := dispatchInitialCommands("srv1", handler, []string{"tail: journal:ssh.service ."}, true, spec, &committedSessionState{})
+	if !errors.Is(err, ErrJournalUnsupported) {
+		t.Fatalf("expected ErrJournalUnsupported, got %v", err)
+	}
+	if len(handler.commands) != 0 {
+		t.Fatalf("expected no commands to be sent, got %#v", handler.commands)
+	}
+	if handler.Status() != 1 {
+		t.Fatalf("handler status = %d, want 1", handler.Status())
+	}
+}
+
+func TestServerConnectionApplySessionSpecPreservesCommittedStateOnRejectedUpdate(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "start", Generation: 2},
+			{Action: "error", Error: "bad reload"},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	startSpec := sessionspec.Spec{Mode: omode.TailClient, Regex: "ERROR"}
+	if err := conn.ApplySessionSpec(startSpec, 10*time.Millisecond); err != nil {
+		t.Fatalf("start ApplySessionSpec() error = %v", err)
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{Mode: omode.TailClient, Regex: "WARN"}, 10*time.Millisecond)
+	if !errors.Is(err, ErrSessionRejected) {
+		t.Fatalf("expected ErrSessionRejected, got %v", err)
+	}
+	if committedSpec, generation, ok := conn.CommittedSession(); !ok || generation != 2 || committedSpec.Regex != "ERROR" {
+		t.Fatalf("unexpected committed session after rejected update: spec=%#v generation=%d ok=%v", committedSpec, generation, ok)
+	}
+}
+
+func TestServerConnectionApplySessionSpecRejectsUnexpectedAck(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+		sessionAcks: []handlers.SessionAck{
+			{Action: "update", Generation: 1},
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}, 10*time.Millisecond)
+	if !errors.Is(err, ErrUnexpectedSessionAck) {
+		t.Fatalf("expected ErrUnexpectedSessionAck, got %v", err)
+	}
+	if _, _, ok := conn.CommittedSession(); ok {
+		t.Fatalf("unexpected committed session after mismatched ack")
+	}
+}
+
+func TestServerConnectionApplySessionSpecTimesOutWaitingForAck(t *testing.T) {
+	resetClientLogger(t)
+
+	mock := &mockHandler{
+		waitForCapabilities: true,
+		capabilities: map[string]bool{
+			protocol.CapabilityQueryUpdateV1: true,
+		},
+	}
+	conn := &ServerConnection{
+		server:  "srv1",
+		handler: mock,
+	}
+
+	err := conn.ApplySessionSpec(sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}, 10*time.Millisecond)
+	if !errors.Is(err, ErrSessionAckTimeout) {
+		t.Fatalf("expected ErrSessionAckTimeout, got %v", err)
+	}
+	if len(mock.commands) != 1 {
+		t.Fatalf("expected session command to be sent before timeout, got %d", len(mock.commands))
+	}
+	if _, _, ok := conn.CommittedSession(); ok {
+		t.Fatalf("unexpected committed session after missing ack")
+	}
+}
+
+func TestApplySessionSpecSerializesConcurrentBootstrapAndReload(t *testing.T) {
+	resetClientLogger(t)
+
+	handler := newBlockingSessionHandler()
+	state := &committedSessionState{}
+
+	initialSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "ERROR",
+	}
+	reloadSpec := sessionspec.Spec{
+		Mode:  omode.TailClient,
+		Files: []string{"/var/log/app.log"},
+		Regex: "WARN",
+	}
+
+	initialErrCh := make(chan error, 1)
+	go func() {
+		initialErrCh <- dispatchInitialCommands("srv1", handler, nil, true, initialSpec, state)
+	}()
+
+	firstCommand := <-handler.commandsCh
+	if !strings.HasPrefix(firstCommand, "SESSION START ") {
+		t.Fatalf("expected initial SESSION START command, got %q", firstCommand)
+	}
+
+	reloadErrCh := make(chan error, 1)
+	go func() {
+		reloadErrCh <- applySessionSpec("srv1", handler, state, reloadSpec, 50*time.Millisecond)
+	}()
+
+	select {
+	case command := <-handler.commandsCh:
+		t.Fatalf("unexpected concurrent session command before bootstrap ack: %q", command)