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>

1 files changed, 433 insertions, 114 deletions

diff --git a/internal/server/handlers/basehandler.go b/internal/server/handlers/basehandler.go
index f6ab3db..1ac159f 100644
--- a/internal/server/handlers/basehandler.go
+++ b/internal/server/handlers/basehandler.go
@@ -3,8 +3,6 @@ package handlers
 import (
 	"bytes"
 	"context"
-	"encoding/base64"
-	"errors"
 	"fmt"
 	"io"
 	"strconv"
@@ -19,26 +17,76 @@ import (
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/io/pool"
 	"github.com/mimecast/dtail/internal/lcontext"
-	"github.com/mimecast/dtail/internal/mapr/server"
+	maprserver "github.com/mimecast/dtail/internal/mapr/server"
 	"github.com/mimecast/dtail/internal/protocol"
 	user "github.com/mimecast/dtail/internal/user/server"
 )
 
 type handleCommandCb func(context.Context, lcontext.LContext, int, []string, string)
 
+// commandCancelKeyType is a private key type for stashing a per-command
+// context.CancelFunc inside a context.Context. It is used to hand the cancel
+// ownership from handleCommand (which creates the context) to the command
+// completion callback (handleUserCommand.commandFinished), which is the only
+// place that knows when the asynchronous command is actually done.
+type commandCancelKeyType struct{}
+
+var commandCancelKey commandCancelKeyType
+
+// withCommandCancel returns a derived context that carries the per-command
+// cancel func. See cancelCommandContext for the matching consumer.
+func withCommandCancel(ctx context.Context, cancel context.CancelFunc) context.Context {
+	if cancel == nil {
+		return ctx
+	}
+	return context.WithValue(ctx, commandCancelKey, cancel)
+}
+
+// cancelCommandContext invokes the per-command cancel func stashed on ctx (if
+// any) exactly once. It is a no-op when ctx carries no cancel (for example in
+// the session-command path where the session state owns the cancel).
+func cancelCommandContext(ctx context.Context) {
+	cancel, ok := ctx.Value(commandCancelKey).(context.CancelFunc)
+	if !ok || cancel == nil {
+		return
+	}
+	cancel()
+}
+
 type baseHandler struct {
-	done             *internal.Done
-	handleCommandCb  handleCommandCb
-	lines            chan *line.Line
-	aggregate        *server.Aggregate
+	done            *internal.Done
+	handleCommandCb handleCommandCb
+	lines           chan *line.Line
+
+	// aggregate is written by handleMapCommand on the command-dispatch
+	// goroutine and read concurrently by Shutdown, Aggregate, and
+	// resetSessionAggregates. Using atomic.Pointer eliminates the data race
+	// without requiring h.mutex to be held around every access site.
+	aggregate atomic.Pointer[maprserver.Aggregate]
+
 	maprMessages     chan string
 	serverMessages   chan string
 	hostname         string
 	user             *user.User
 	ackCloseReceived chan struct{}
+	ackCloseOnce     sync.Once
 	activeCommands   int32
-	readBuf          bytes.Buffer
-	writeBuf         bytes.Buffer
+	codec            protocolCodec
+
+	// readBuf holds the formatted protocol message currently being sent to
+	// the client. It is only touched by Read (single session output
+	// goroutine) and retains any bytes that did not fit into the caller's
+	// buffer, so messages larger than one Read are delivered across multiple
+	// calls instead of being truncated (see Read/drainReadBuf).
+	readBuf  bytes.Buffer
+	writeBuf bytes.Buffer
+
+	// maxCommandFrameSize is the maximum number of bytes that may be buffered
+	// between two ';' delimiters. When a frame grows beyond this limit the
+	// Write method closes the session immediately to prevent a malicious or
+	// misbehaving client from exhausting server memory. The value is set at
+	// construction time from ServerConfig.MaxCommandFrameSize.
+	maxCommandFrameSize int
 
 	// Some global options + sync primitives required.
 	once       sync.Once
@@ -46,10 +94,30 @@ type baseHandler struct {
 	quiet      bool
 	plain      bool
 	serverless bool
+
+	output outputManager
+
+	activeGeneration func() uint64
+}
+
+// getAggregate returns the current output MapReduce aggregate atomically.
+func (h *baseHandler) getAggregate() *maprserver.Aggregate {
+	return h.aggregate.Load()
 }
 
-// Shutdown the handler.
+// setAggregate stores a output MapReduce aggregate atomically.
+func (h *baseHandler) setAggregate(ta *maprserver.Aggregate) {
+	h.aggregate.Store(ta)
+}
+
+// Shutdown the handler. Uses atomic accessors to read aggregate pointers so
+// the reads are race-free with concurrent writes from handleMapCommand.
 func (h *baseHandler) Shutdown() {
+	// Shutdown output aggregate if present.
+	if ta := h.getAggregate(); ta != nil {
+		dlog.Server.Info(h.user, "Shutting down output aggregate")
+		ta.Shutdown()
+	}
 	h.done.Shutdown()
 }
 
@@ -59,73 +127,143 @@ func (h *baseHandler) Done() <-chan struct{} {
 }
 
 // Read is to send data to the dtail client via Reader interface.
+//
+// A formatted protocol message can be larger than p (io.Copy drives this
+// reader with a 32KB buffer while MaxLineLength allows lines up to 1MB), so
+// each Read drains any bytes left over from a previous call first and every
+// message path keeps its unsent remainder in readBuf across calls. Dropping
+// the remainder would truncate long lines and lose the trailing message
+// delimiter, desyncing the client-side parser. This mirrors the remainder
+// buffer used by the output path (outputManager.tryRead).
 func (h *baseHandler) Read(p []byte) (n int, err error) {
-	defer h.readBuf.Reset()
+	if h.readBuf.Len() > 0 {
+		return h.drainReadBuf(p), nil
+	}
 
-	select {
-	case message := <-h.serverMessages:
-		if len(message) > 0 && message[0] == '.' {
-			// Handle hidden message (don't display to the user)
-			h.readBuf.WriteString(message)
-			h.readBuf.WriteByte(protocol.MessageDelimiter)
-			n = copy(p, h.readBuf.Bytes())
-			return
+	for {
+		if n, handled := h.output.tryRead(p, h.user, h.shouldDropGeneration); handled {
+			if n == 0 {
+				continue
+			}
+			return n, nil
 		}
 
-		if h.serverless {
-			return
+		pollInterval := time.Second
+		if h.output.enabled() {
+			// Output reads require tighter wake-ups so we can continue draining the output channel.
+			pollInterval = h.output.resolvedReadRetryInterval()
 		}
+		poll := time.After(pollInterval)
 
-		// Handle normal server message (display to the user)
-		h.readBuf.WriteString("SERVER")
-		h.readBuf.WriteString(protocol.FieldDelimiter)
-		h.readBuf.WriteString(h.hostname)
-		h.readBuf.WriteString(protocol.FieldDelimiter)
-		h.readBuf.WriteString(message)
-		h.readBuf.WriteByte(protocol.MessageDelimiter)
-		n = copy(p, h.readBuf.Bytes())
-
-	case message := <-h.maprMessages:
-		// Send mapreduce-aggregated data as a message.
-		h.readBuf.WriteString("AGGREGATE")
-		h.readBuf.WriteString(protocol.FieldDelimiter)
-		h.readBuf.WriteString(h.hostname)
-		h.readBuf.WriteString(protocol.FieldDelimiter)
-		h.readBuf.WriteString(message)
-		h.readBuf.WriteByte(protocol.MessageDelimiter)
-		n = copy(p, h.readBuf.Bytes())
-
-	case line := <-h.lines:
-		if !h.plain {
-			h.readBuf.WriteString("REMOTE")
+		select {
+		case message := <-h.serverMessages:
+			generation, decodedMessage := decodeGeneratedMessage(message)
+			if h.shouldDropGeneration(generation) {
+				continue
+			}
+			message = decodedMessage
+			if len(message) > 0 && message[0] == '.' {
+				// Handle hidden message (don't display to the user)
+				h.readBuf.WriteString(message)
+				h.readBuf.WriteByte(protocol.MessageDelimiter)
+				n = h.drainReadBuf(p)
+				return
+			}
+
+			if h.serverless {
+				return
+			}
+
+			// Skip empty server messages when in plain mode
+			if h.plain && (message == "" || message == "\n") {
+				return
+			}
+
+			// Handle normal server message (display to the user).
+			formatServerMessage(&h.readBuf, h.hostname, message, h.plain)
+			n = h.drainReadBuf(p)
+			return
+
+		case message := <-h.maprMessages:
+			generation, decodedMessage := decodeGeneratedMessage(message)
+			if h.shouldDropGeneration(generation) {
+				continue
+			}
+			message = decodedMessage
+			// Send mapreduce-aggregated data as a message. The leading
+			// AggregateMessageID field lets the mapr client tell aggregate
+			// data apart from plain server acks that happen to start with 'A'.
+			h.readBuf.WriteString(protocol.AggregateMessageID)
 			h.readBuf.WriteString(protocol.FieldDelimiter)
 			h.readBuf.WriteString(h.hostname)
 			h.readBuf.WriteString(protocol.FieldDelimiter)
-			h.readBuf.WriteString(fmt.Sprintf("%3d", line.TransmittedPerc))
-			h.readBuf.WriteString(protocol.FieldDelimiter)
-			h.readBuf.WriteString(fmt.Sprintf("%v", line.Count))
-			h.readBuf.WriteString(protocol.FieldDelimiter)
-			h.readBuf.WriteString(line.SourceID)
-			h.readBuf.WriteString(protocol.FieldDelimiter)
-		}
-		h.readBuf.WriteString(line.Content.String())
-		h.readBuf.WriteByte(protocol.MessageDelimiter)
-		n = copy(p, h.readBuf.Bytes())
-		pool.RecycleBytesBuffer(line.Content)
-		line.Recycle()
+			h.readBuf.WriteString(message)
+			h.readBuf.WriteByte(protocol.MessageDelimiter)
+			n = h.drainReadBuf(p)
+			return
+
+		case line := <-h.lines:
+			if line == nil {
+				continue
+			}
+			if h.shouldDropGeneration(line.Generation) {
+				pool.RecycleBytesBuffer(line.Content)
+				line.Recycle()
+				continue
+			}
+			if h.plain {
+				h.readBuf.Write(line.Content.Bytes())
+				h.readBuf.WriteByte(protocol.MessageDelimiter)
+			} else {
+				formatRemoteLine(
+					&h.readBuf,
+					h.hostname,
+					fmt.Sprintf("%3d", line.TransmittedPerc),
+					line.Count,
+					line.SourceID,
+					line.Content.Bytes(),
+				)
+			}
+			n = h.drainReadBuf(p)
+			pool.RecycleBytesBuffer(line.Content)
+			line.Recycle()
+			return
 
-	case <-time.After(time.Second):
-		select {
 		case <-h.done.Done():
 			err = io.EOF
 			return
-		default:
+
+		case <-poll:
+			// Wake periodically so output mode transitions don't leave this read blocked forever.
+			select {
+			case <-h.done.Done():
+				err = io.EOF
+				return
+			default:
+			}
+			return
 		}
 	}
-	return
+}
+
+// drainReadBuf copies as many buffered message bytes as fit into p and keeps
+// the remainder in readBuf for subsequent Read calls. bytes.Buffer.Read
+// consumes exactly the bytes it returns, so nothing is ever discarded; its
+// io.EOF (only possible on an empty buffer) is deliberately not propagated
+// because an empty buffer here simply means there is nothing left to drain.
+func (h *baseHandler) drainReadBuf(p []byte) int {
+	n, _ := h.readBuf.Read(p)
+	return n
 }
 
 // Write is to receive data from the dtail client via Writer interface.
+// Each byte is accumulated in writeBuf until a ';' delimiter arrives, at which
+// point the buffered frame is dispatched as a command and the buffer is reset.
+//
+// To prevent a client from exhausting server memory with an unterminated frame,
+// the buffer length is checked against maxCommandFrameSize on every append. When
+// the limit is exceeded the session is shut down and io.ErrClosedPipe is returned
+// so the SSH layer tears down the connection.
 func (h *baseHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
 		switch b {
@@ -134,6 +272,19 @@ func (h *baseHandler) Write(p []byte) (n int, err error) {
 			h.writeBuf.Reset()
 		default:
 			h.writeBuf.WriteByte(b)
+			// Guard against unbounded frame growth: a client could send bytes
+			// without ever emitting a ';' delimiter and grow the buffer
+			// indefinitely. Reject and close when the configurable limit is hit.
+			if h.maxCommandFrameSize > 0 && h.writeBuf.Len() > h.maxCommandFrameSize {
+				dlog.Server.Error(h.user,
+					"command frame exceeds maximum size, closing session",
+					"frameSize", h.writeBuf.Len(),
+					"limit", h.maxCommandFrameSize,
+				)
+				h.writeBuf.Reset()
+				h.done.Shutdown()
+				return len(p), io.ErrClosedPipe
+			}
 		}
 	}
 	n = len(p)
@@ -153,79 +304,145 @@ func (h *baseHandler) handleCommand(commandStr string) {
 		h.sendln(h.serverMessages, dlog.Server.Error(h.user, err))
 		return
 	}
-	ctx, cancel := context.WithCancel(context.Background())
-	go func() {
-		<-h.done.Done()
+	ctx, cancel := h.newCommandContext(context.Background())
+	// Cancel ownership is transferred to the command completion callback
+	// (see cancelCommandContext + handleUserCommand.commandFinished) so the
+	// per-command context and its watcher goroutine are released once the
+	// (possibly asynchronous) command has finished. If dispatch fails before
+	// the callback is ever invoked we must cancel here to avoid a leak.
+	ctx = withCommandCancel(ctx, cancel)
+
+	if err := h.dispatchCommand(ctx, args, argc); err != nil {
 		cancel()
-	}()
+		h.sendln(h.serverMessages, dlog.Server.Error(h.user, err))
+	}
+}
+
+func (h *baseHandler) dispatchCommand(ctx context.Context, args []string, argc int) error {
+	// Strip and apply a leading "timeout N <cmd>..." prefix. The client emits
+	// this when --timeout>0 (see internal/session/spec.go queryCommands); it
+	// caps how long the server collects data for that read command before its
+	// context is canceled. Handling it here covers both the legacy command
+	// stream and the SESSION dispatch path, which both funnel through here.
+	ctx, args, argc, err := applyCommandTimeout(ctx, args, argc)
+	if err != nil {
+		return err
+	}
 
-	parts := strings.Split(args[0], ":")
+	parts := strings.SplitN(args[0], ":", 2)
 	commandName := parts[0]
 
 	// Either no options or empty options provided.
 	if len(parts) == 1 || len(parts[1]) == 0 {
 		h.handleCommandCb(ctx, lcontext.LContext{}, argc, args, commandName)
-		return
+		return nil
 	}
 
-	options, ltx, err := config.DeserializeOptions(parts[1:])
+	options, ltx, err := config.DeserializeOptions([]string{parts[1]})
 	if err != nil {
-		h.sendln(h.serverMessages, dlog.Server.Error(h.user, err))
-		return
+		return err
 	}
 	h.handleOptions(options)
 	h.handleCommandCb(ctx, ltx, argc, args, commandName)
+	return nil
 }
 
-func (h *baseHandler) handleProtocolVersion(args []string) ([]string, int, string, error) {
-	argc := len(args)
-	var add string
-
-	if argc <= 2 || args[0] != "protocol" {
-		return args, argc, add, errors.New("unable to determine protocol version")
+// maxCommandTimeoutSeconds caps the "timeout N <cmd>" prefix value. 24h is far
+// beyond any realistic collection window yet nowhere near the int64 overflow
+// point of time.Duration (~292 years in nanoseconds), so it doubles as an
+// overflow guard for the multiplication in applyCommandTimeout.
+const maxCommandTimeoutSeconds = 24 * 60 * 60
+
+// applyCommandTimeout detects a leading "timeout N <cmd>..." command prefix (as
+// emitted by the client when --timeout>0) and returns a context that is
+// canceled after N seconds together with the remaining command (the prefix
+// stripped). This restores the original server-side deadline semantics: "Max
+// time dtail server will collect data until disconnection". When no timeout
+// prefix is present, or N<=0, the context and args are returned unchanged so
+// the --timeout 0 / unset case behaves exactly as before.
+//
+// The timeout child cancel is chained onto the per-command cancel already
+// stashed on ctx (if any) so cancelCommandContext, invoked once the command
+// finishes, releases both the parent cancel and the timeout timer. In the
+// session-dispatch path ctx carries no per-command cancel, so the returned
+// context is the sole owner and its cancel still fires on command completion.
+func applyCommandTimeout(ctx context.Context, args []string, argc int) (context.Context, []string, int, error) {
+	if argc < 3 || args[0] != "timeout" {
+		return ctx, args, argc, nil
 	}
 
-	if args[1] != protocol.ProtocolCompat {
-		clientCompat, _ := strconv.Atoi(args[1])
-		serverCompat, _ := strconv.Atoi(protocol.ProtocolCompat)
-		if clientCompat <= 3 {
-			// Protocol version 3 or lower expect a newline as message separator
-			// One day (after 2 major versions) this exception may be removed!
-			add = "\n"
-		}
+	seconds, err := strconv.Atoi(args[1])
+	if err != nil {
+		return ctx, args, argc, fmt.Errorf("invalid timeout value %q: %w", args[1], err)
+	}
+	// Reject absurd values rather than clamp: an out-of-range N is a client
+	// mistake, and erroring (like the non-numeric case above) surfaces it
+	// instead of silently substituting a different deadline. This also guards
+	// against int64 overflow in time.Duration(seconds)*time.Second below, which
+	// for a huge N would wrap to a negative (already-elapsed) deadline and
+	// cancel the read immediately.
+	if seconds > maxCommandTimeoutSeconds {
+		return ctx, args, argc, fmt.Errorf("timeout value %d exceeds maximum of %d seconds",
+			seconds, maxCommandTimeoutSeconds)
+	}
+	if seconds <= 0 {
+		return ctx, args[2:], argc - 2, nil
+	}
 
-		toUpdate := "client"
-		if clientCompat > serverCompat {
-			toUpdate = "server"
+	timeoutCtx, cancel := context.WithTimeout(ctx, time.Duration(seconds)*time.Second)
+	parentCancel, _ := ctx.Value(commandCancelKey).(context.CancelFunc)
+	combined := func() {
+		cancel()
+		if parentCancel != nil {
+			parentCancel()
 		}
-		err := fmt.Errorf("the DTail server protocol version '%s' does not match "+
-			"client protocol version '%s', please update DTail %s",
-			protocol.ProtocolCompat, args[1], toUpdate)
-		return args, argc, add, err
 	}
 
-	return args[2:], argc - 2, add, nil
+	return withCommandCancel(timeoutCtx, combined), args[2:], argc - 2, nil
+}
+
+func (h *baseHandler) handleProtocolVersion(args []string) ([]string, int, string, error) {
+	return h.codec.handleProtocolVersion(args)
 }
 
 func (h *baseHandler) handleBase64(args []string, argc int) ([]string, int, error) {
-	err := errors.New("unable to decode client message, DTail server and client " +
-		"versions may not be compatible")
-	if argc != 2 || args[0] != "base64" {
-		return args, argc, err
-	}
+	return h.codec.handleBase64(args, argc)
+}
 
-	decoded, err := base64.StdEncoding.DecodeString(args[1])
-	if err != nil {
-		return args, argc, err
+func (h *baseHandler) handleRawCommand(ctx context.Context, command string) error {
+	args := strings.Fields(command)
+	if len(args) == 0 {
+		return fmt.Errorf("empty command")
 	}
-	decodedStr := string(decoded)
+	return h.dispatchCommand(ctx, args, len(args))
+}
 
-	args = strings.Split(decodedStr, " ")
-	argc = len(decodedStr)
-	dlog.Server.Trace(h.user, "Base64 decoded received command",
-		decodedStr, argc, args)
+// newCommandContext creates a cancellable context for a single command
+// invocation. The caller owns the returned cancel func and MUST invoke it
+// exactly once (typically via defer or through the per-command cancel
+// stashed on the context, see withCommandCancel/cancelCommandContext).
+// Failing to cancel leaks both the context and the watcher goroutine
+// spawned below, because the watcher only returns when the handler is shut
+// down; on long-lived sessions (:reload, continuous/scheduled workloads)
+// those leaks accumulate per command.
+//
+// The watcher goroutine doubles as a defensive safety net: even if a
+// caller forgets to cancel, handler shutdown still drains it by cancelling
+// the context via <-h.done.Done().
+func (h *baseHandler) newCommandContext(parent context.Context) (context.Context, context.CancelFunc) {
+	if parent == nil {
+		parent = context.Background()
+	}
 
-	return args, argc, nil
+	ctx, cancel := context.WithCancel(parent)
+	go func() {
+		select {
+		case <-h.done.Done():
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+	return ctx, cancel
 }
 
 func (h *baseHandler) handleAckCommand(argc int, args []string) {
@@ -237,11 +454,9 @@ func (h *baseHandler) handleAckCommand(argc int, args []string) {
 		return
 	}
 	if args[1] == "close" && args[2] == "connection" {
-		select {
-		case <-h.ackCloseReceived:
-		default:
+		h.ackCloseOnce.Do(func() {
 			close(h.ackCloseReceived)
-		}
+		})
 	}
 }
 
@@ -278,24 +493,76 @@ func (h *baseHandler) sendln(ch chan<- string, message string) {
 	h.send(ch, message+"\n")
 }
 
+func (h *baseHandler) shouldDropGeneration(generation uint64) bool {
+	if generation == 0 || h.activeGeneration == nil {
+		return false
+	}
+
+	activeGeneration := h.activeGeneration()
+	if activeGeneration == 0 {
+		return false
+	}
+
+	return activeGeneration != generation
+}
+
 func (h *baseHandler) flush() {
 	dlog.Server.Trace(h.user, "flush()")
 	numUnsentMessages := func() int {
-		return len(h.lines) + len(h.serverMessages) + len(h.maprMessages)
+		lineCount := len(h.lines)
+		serverCount := len(h.serverMessages)
+		maprCount := len(h.maprMessages)
+		outputCount := h.output.channelLen()
+		dlog.Server.Trace(h.user, "flush", "lines", lineCount, "server", serverCount, "mapr", maprCount, "output", outputCount)
+		return lineCount + serverCount + maprCount + outputCount
+	}
+
+	// Use atomic accessors to avoid a data race with handleMapCommand, which
+	// may be concurrently writing aggregate pointers on another goroutine.
+	maxWait := time.Second
+	if h.output.enabled() || h.getAggregate() != nil {
+		maxWait = 3 * time.Second
+	}
+	if h.serverless && maxWait < 5*time.Second {
+		maxWait = 5 * time.Second
 	}
-	for i := 0; i < 10; i++ {
-		if numUnsentMessages() == 0 {
+
+	deadline := time.Now().Add(maxWait)
+	for i := 0; ; i++ {
+		unsent := numUnsentMessages()
+		if unsent == 0 {
 			dlog.Server.Debug(h.user, "ALL lines sent", fmt.Sprintf("%p", h))
 			return
 		}
-		dlog.Server.Debug(h.user, "Still lines to be sent")
+		if time.Now().After(deadline) {
+			dlog.Server.Warn(h.user, "Some lines remain unsent", unsent)
+			return
+		}
+		dlog.Server.Debug(h.user, "Still lines to be sent", "iteration", i, "unsent", unsent, "deadline", deadline.Sub(time.Now()))
 		time.Sleep(time.Millisecond * 10)
 	}
-	dlog.Server.Warn(h.user, "Some lines remain unsent", numUnsentMessages())
 }
 
 func (h *baseHandler) shutdown() {
-	dlog.Server.Debug(h.user, "shutdown()")
+	// Log current state at shutdown
+	activeCommands := atomic.LoadInt32(&h.activeCommands)
+	dlog.Server.Info(h.user, "shutdown() called", "activeCommands", activeCommands, "outputMode", h.output.enabled())
+
+	// In output mode, ensure all data is flushed before shutdown
+	if h.output.enabled() {
+		h.flushOutput()
+	}
+
+	// Shutdown the aggregate BEFORE flush to ensure MapReduce data is available.
+	// Use the atomic accessor to avoid a data race with handleMapCommand which
+	// may be concurrently storing the aggregate pointer on another goroutine.
+	if ta := h.getAggregate(); ta != nil {
+		dlog.Server.Info(h.user, "Shutting down output aggregate in shutdown()")
+		ta.Shutdown()
+		// Give time for serialization to complete.
+		time.Sleep(100 * time.Millisecond)
+	}
+
 	h.flush()
 
 	go func() {
@@ -322,3 +589,55 @@ func (h *baseHandler) decrementActiveCommands() int32 {
 	atomic.AddInt32(&h.activeCommands, -1)
 	return atomic.LoadInt32(&h.activeCommands)
 }
+
+// EnableDirectOutput enables output mode for direct line processing. It is an
+// atomic check-and-enable: the return value is true when this call switched
+// output mode on and false when it was already active (in which case the
+// existing output state is left untouched).
+func (h *baseHandler) EnableDirectOutput() bool {
+	return h.output.enable()
+}
+
+// DirectOutputActive returns true if output mode is enabled
+func (h *baseHandler) DirectOutputActive() bool {
+	return h.output.enabled()
+}
+
+// HasOutputEOF returns true when a output EOF channel exists.
+func (h *baseHandler) HasOutputEOF() bool {
+	return h.output.hasEOF()
+}
+
+// OutputEpoch returns the current output handshake epoch. Capture it before
+// checking the pending-work count and pass it to SignalOutputEOF so a stale
+// "batch over" decision cannot EOF a batch that joined in between.
+func (h *baseHandler) OutputEpoch() uint64 {
+	return h.output.currentEpoch()
+}
+
+// SignalOutputEOF closes the output EOF channel once, unless the handshake
+// epoch has advanced past the given captured value (i.e. another command
+// joined the output session since), in which case the stale signal is dropped.
+func (h *baseHandler) SignalOutputEOF(epoch uint64) {
+	h.output.signalEOF(epoch)
+}
+
+// flushOutput ensures all output channel data is processed
+func (h *baseHandler) flushOutput() {
+	h.output.flush(h.user)
+}
+
+// GetOutputChannel returns the output lines channel for direct writing
+func (h *baseHandler) GetOutputChannel() chan []byte {
+	return h.output.channel()
+}
+
+// OutputChannelLen returns current output channel buffered size.
+func (h *baseHandler) OutputChannelLen() int {
+	return h.output.channelLen()
+}
+
+// WaitForOutputEOFAck waits until output reader acknowledges EOF or timeout.
+func (h *baseHandler) WaitForOutputEOFAck(timeout time.Duration) bool {
+	return h.output.waitForEOFAck(timeout)
+}