Say what reload is about to do before doing it, on every surface

The rebase onto the Run/Result model left AfterReply registered by nobody, so
the "Reloading config" line that used to go out before the reload went out
after it, folded into the result, and a reload that restarted the ssh server
cut the line off with the session. Run now says what is about to happen and a
second step, Then, does it and says how it went.

- Command.Then replaces AfterReply: Run's Result is rendered, Then runs, its Result follows and decides the outcome
- over ssh in JSON only Then's outcome travels, one document for the program that asked
- the control socket runs Then before its one reply and carries both lines as text, or the outcome as JSON
- the info API runs Then before answering; nothing it serves sets one
- reload: Run promises, Then reloads and reports, without the prefix runReload used to add
- the output golden loses the prefix from reload's JSON reasons

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

commands.go

@@ -127,7 +127,12 @@ func attachCommands(l *slog.Logger, c *config.C, reg *commands.Registry, f *Inte
 	reg.Register(&commands.Command{
 		Name:             commands.Reload,
 		ShortDescription: "Reloads configuration from disk, same as sending HUP to the process",
+		// The line goes out before the reload, as it always has: a reload that restarts
+		// the ssh server ends the session that asked, and this is what it sees.
 		Run: func(fs any, a []string) (commands.Result, error) {
+			return commands.Did("Reloading config"), nil
+		},
+		Then: func() commands.Result {
 			return runReload(l, c)
 		},
 	})
@@ -633,14 +638,13 @@ func textDeviceInfo(d deviceInfo, w io.Writer) error {
 	return err
 }
 
-// runReload reloads the config and reports how it went, under the "Reloading config"
-// line the command has always printed first. A reload that failed is a refusal, exit 1
-// over ssh; one that merely needs a restart is not.
-func runReload(l *slog.Logger, c *config.C) (commands.Result, error) {
+// runReload reloads the config and reports how it went. A reload that failed is a
+// refusal, exit 1 over ssh; one that merely needs a restart is not.
+func runReload(l *slog.Logger, c *config.C) commands.Result {
 	res := c.ReloadConfig()
 	res.Log(l)
 	if res.Err != nil || res.Failed {
-		return commands.Failed("Reloading config\n%s", res.String()), nil
+		return commands.Failed("%s", res.String())
 	}
-	return commands.Did("Reloading config\n%s", res.String()), nil
+	return commands.Did("%s", res.String())
 }

commands/command.go

@@ -62,11 +62,14 @@ type Command struct {
 	// rendered as a sentence, since that always has a human form.
 	JSONOnly bool
 
-	// AfterReply runs once the surface has delivered Run's result, and only then. It is
-	// for a command whose action would cut off its own reply: reload restarts the ssh
-	// server, which ends the session that asked, so the sentence Run returns has to be
-	// on the wire before the reload begins. Run promises; AfterReply acts.
-	AfterReply func()
+	// Then, when set, is the second half of a command that acts: Run says what is about
+	// to happen and Then does it and says how it went, with its Result rendered after
+	// Run's and deciding the outcome. It exists for a command whose action can cut off
+	// its own reply: reload restarts the ssh server when the sshd keys changed, which
+	// ends the session that asked, so Run's line has to be on the wire before the reload
+	// begins. A surface that answers in one piece, the control socket, runs Then before
+	// it answers and carries both.
+	Then func() Result
 }
 
 // Failure is the error a surface gets for a refusal the person has already been shown:
@@ -171,39 +174,52 @@ func (r *Registry) execLine(c *Command, args []string, w Output) error {
 		return err
 	}
 
-	if res.Empty() {
-		runAfterReply(c)
-		return nil
-	}
-
-	// The line wins when it asked for a format. Otherwise the surface that owns the
-	// registry decides, from the command and its Result, which is how ssh keeps printing
-	// what it always has while another surface can choose differently.
-	f := format
-	if f == FormatUnset {
-		f = r.defaultFormatFor(c, res)
+	// Run's reply goes out first, and then the action, if the command has one, whose
+	// Result is the outcome. In JSON only the outcome travels: the promise is for a
+	// person to read, and a program expects one document.
+	if c.Then == nil || r.formatFor(c, res, format) == FormatText {
+		if err := r.renderLine(c, res, format, w); err != nil {
+			return err
+		}
 	}
-
-	if err := Render(c, res, f, w.Out); err != nil {
-		if errors.Is(err, ErrNoTextRendering) {
-			_, _ = fmt.Fprintln(w.Err, err.Error())
+	if c.Then != nil {
+		res = c.Then()
+		if err := r.renderLine(c, res, format, w); err != nil {
+			return err
 		}
-		return err
 	}
 
-	runAfterReply(c)
 	if !res.OK {
 		return &Failure{Msg: res.Reason}
 	}
 	return nil
 }
 
-// runAfterReply performs the command's deferred action, if it has one. It is called only
-// once the reply has been written, which is the whole point of AfterReply.
-func runAfterReply(c *Command) {
-	if c.AfterReply != nil {
-		c.AfterReply()
+// formatFor is the format a Result is rendered in. The line wins when it asked for one;
+// otherwise the surface that owns the registry decides, from the command and its
+// Result, which is how ssh keeps printing what it always has while another surface can
+// choose differently.
+func (r *Registry) formatFor(c *Command, res Result, format Format) Format {
+	if format == FormatUnset {
+		return r.defaultFormatFor(c, res)
 	}
+	return format
+}
+
+// renderLine writes one Result of a command line to w; see formatFor. An empty Result
+// writes nothing.
+func (r *Registry) renderLine(c *Command, res Result, format Format, w Output) error {
+	if res.Empty() {
+		return nil
+	}
+
+	if err := Render(c, res, r.formatFor(c, res, format), w.Out); err != nil {
+		if errors.Is(err, ErrNoTextRendering) {
+			_, _ = fmt.Fprintln(w.Err, err.Error())
+		}
+		return err
+	}
+	return nil
 }
 
 // ParseArgs parses args with fl and returns the operands.

commands/registry_test.go

@@ -387,11 +387,13 @@ func TestCloneHelpSeesCloneCommands(t *testing.T) {
 	assert.Contains(t, out.String(), "logout")
 }
 
-// TestAfterReplyRunsAfterTheReply pins the order AfterReply exists for: the reply is on
-// the writer before the action begins.
-func TestAfterReplyRunsAfterTheReply(t *testing.T) {
+// TestThenRunsAfterTheReply pins the order Then exists for: Run's line is on the writer
+// before the action begins, Then's Result follows it, and Then's outcome is the
+// command's.
+func TestThenRunsAfterTheReply(t *testing.T) {
 	out := &bytes.Buffer{}
 	seen := "never ran"
+	outcome := Did("acted")
 
 	r := New()
 	r.Register(&Command{
@@ -400,12 +402,26 @@ func TestAfterReplyRunsAfterTheReply(t *testing.T) {
 		Run: func(fs any, a []string) (Result, error) {
 			return Did("about to act"), nil
 		},
-		AfterReply: func() { seen = out.String() },
+		Then: func() Result {
+			seen = out.String()
+			return outcome
+		},
 	})
 
 	require.NoError(t, r.DispatchArgs([]string{"act"}, MergedOutput(out)))
-	assert.Contains(t, seen, "about to act")
-	assert.Equal(t, out.String(), seen)
+	assert.Equal(t, "about to act\n", seen, "the line was written before the action")
+	assert.Equal(t, "about to act\nacted\n", out.String())
+
+	out.Reset()
+	outcome = Failed("could not act")
+	var f *Failure
+	require.ErrorAs(t, r.DispatchArgs([]string{"act"}, MergedOutput(out)), &f)
+	assert.Equal(t, "about to act\ncould not act\n", out.String())
+
+	// In JSON only the outcome travels: one document for the program that asked.
+	out.Reset()
+	require.ErrorAs(t, r.DispatchArgs([]string{"act", "-json"}, MergedOutput(out)), &f)
+	assert.JSONEq(t, `{"ok":false,"reason":"could not act"}`, out.String())
 }
 
 // TestWriteCommandHelpIsTheOneRendering pins the layout both `help <command>` on the

e2e/testdata/commands_output.golden

@@ -1337,13 +1337,13 @@ Config reload failed, the new config did not load: no config files found at
 exit status 1
 
 $ reload -json
-{"ok":false,"reason":"Reloading config\nConfig reload failed, the new config did not load: no config files found at "}
+{"ok":false,"reason":"Config reload failed, the new config did not load: no config files found at "}
 exit status 1
 
 $ reload -pretty
 {
     "ok": false,
-    "reason": "Reloading config\nConfig reload failed, the new config did not load: no config files found at "
+    "reason": "Config reload failed, the new config did not load: no config files found at "
 }
 exit status 1
 

localapi/info.go

@@ -258,14 +258,10 @@ func (s *InfoServer) handleCommand(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if cmd.AfterReply != nil {
-		// The reply sits in the server's buffer until the handler returns or flushes;
-		// AfterReply's contract is that the surface has delivered it, so flush first.
-		// writeJSON sets Content-Length, so a flushed reply is complete to its reader.
-		defer func() {
-			_ = http.NewResponseController(w).Flush()
-			cmd.AfterReply()
-		}()
+	// A command that acts does so before this surface's one reply, whose body is the
+	// outcome; see Command.Then.
+	if cmd.Then != nil {
+		res = cmd.Then()
 	}
 
 	if res.Empty() {

localapi/info_test.go

@@ -428,46 +428,6 @@ func TestWireInfoKeepsRunningOnABadReload(t *testing.T) {
 	assert.Equal(t, http.StatusOK, status)
 }
 
-// TestInfoAfterReplyRunsOnceTheResponseIsOut: the hook blocks until the test releases
-// it, so the only way the response can arrive is if it left the server before the hook
-// began. The reply carries a Content-Length, so a flushed reply is complete to its reader.
-func TestInfoAfterReplyRunsOnceTheResponseIsOut(t *testing.T) {
-	release := make(chan struct{})
-	acted := make(chan struct{})
-
-	reg := commands.New()
-	reg.Register(&commands.Command{
-		Name:             "act",
-		ShortDescription: "replies, then acts",
-		Identity:         true,
-		Run: func(fs any, a []string) (commands.Result, error) {
-			return commands.Did("about to act"), nil
-		},
-		AfterReply: func() {
-			<-release
-			close(acted)
-		},
-	})
-	_, path := serveInfo(t, reg, 0o600)
-
-	status, body := test.HTTPGet(t, path, "/v1/act")
-	assert.Equal(t, http.StatusOK, status)
-	assert.JSONEq(t, `{"ok":true,"reason":"about to act"}`, body)
-
-	select {
-	case <-acted:
-		t.Fatal("AfterReply ran before the response was delivered")
-	default:
-	}
-
-	close(release)
-	select {
-	case <-acted:
-	case <-time.After(2 * time.Second):
-		t.Fatal("AfterReply never ran")
-	}
-}
-
 // rawInfoRequest opens the socket and writes req as it is, so a test can send what an
 // HTTP client would not: a body that never arrives. It returns the connection.
 func rawInfoRequest(t *testing.T, path, req string) net.Conn {

localapi/server.go

@@ -207,10 +207,7 @@ func (s *Server) accept(ln net.Listener, sem chan struct{}) {
 func (s *Server) handle(conn net.Conn) {
 	defer conn.Close()
 
-	var (
-		args    []string
-		replied bool
-	)
+	var args []string
 	defer func() {
 		if r := recover(); r != nil {
 			s.l.Error("Recovered from panic serving a local control command",
@@ -218,11 +215,7 @@ func (s *Server) handle(conn net.Conn) {
 				"args", args,
 				"stack", string(debug.Stack()),
 			)
-			// A panic in a command's AfterReply arrives here after the answer has gone
-			// out; the client has its reply, so it is not sent a second one.
-			if !replied {
-				s.reply(conn, &Response{Version: ProtocolVersion, Error: "command failed"})
-			}
+			s.reply(conn, &Response{Version: ProtocolVersion, Error: "command failed"})
 		}
 	}()
 
@@ -259,25 +252,17 @@ func (s *Server) handle(conn net.Conn) {
 	_ = conn.SetReadDeadline(time.Time{})
 	_ = conn.SetWriteDeadline(time.Now().Add(writeTimeout))
 
-	res, after := s.serve(req)
-	s.reply(conn, res)
-	replied = true
-
-	// A command that would cut off its own reply acts only now, once the answer is out.
-	if after != nil {
-		after()
-	}
+	s.reply(conn, s.serve(req))
 }
 
-// serve runs one request and describes the outcome. after, when not nil, is the
-// command's AfterReply, which the caller runs once the response has been written.
+// serve runs one request and describes the outcome.
 //
 // Nothing bounds a command's run; a context would have to reach every Run. What is
 // bounded is the silence: a command still running at slowAfter is named in the log, with
 // the slot it is holding, so a socket that has gone "busy" can be traced to the commands
 // that wedged it. The timer is stopped however the command ends, a panic included, so
 // the line is only ever written about a command that is still running.
-func (s *Server) serve(req *Request) (res *Response, after func()) {
+func (s *Server) serve(req *Request) *Response {
 	slow := time.AfterFunc(s.slowAfter, func() {
 		s.l.Warn("Local control command is still running and holding a connection slot",
 			"command", req.Command,
@@ -288,66 +273,79 @@ func (s *Server) serve(req *Request) (res *Response, after func()) {
 	})
 	defer slow.Stop()
 
-	res = &Response{Version: ProtocolVersion}
+	res := &Response{Version: ProtocolVersion}
 	reg := s.registry()
 
 	if req.Describe {
 		schema, err := json.Marshal(reg.Schema())
 		if err != nil {
 			res.Error = err.Error()
-			return res, nil
+			return res
 		}
 
 		res.Result = schema
-		return res, nil
+		return res
 	}
 
 	if req.Command == "" {
 		res.Error = "no command given"
-		return res, nil
+		return res
 	}
 
 	params, err := stringParams(req.Params)
 	if err != nil {
 		res.Error = err.Error()
-		return res, nil
+		return res
 	}
 
 	format, err := parseFormat(req.Format)
 	if err != nil {
 		res.Error = err.Error()
-		return res, nil
+		return res
 	}
 
 	cmd, err := reg.Resolve(commands.Name(req.Command))
 	if err != nil {
 		s.l.Debug("Local control command refused", "command", req.Command, "error", err)
 		res.Error = err.Error()
-		return res, nil
+		return res
 	}
 
 	r, err := commands.Invoke(commands.Request{Command: cmd, Params: params, Args: req.Args})
 	if err != nil {
 		s.l.Debug("Local control command failed", "command", req.Command, "error", err)
 		res.Error = err.Error()
-		return res, nil
+		return res
 	}
 
-	after = cmd.AfterReply
+	// A command that acts does so now, before the one reply this surface gives, and the
+	// reply carries both halves: the line Run promised and the outcome Then produced.
+	// In JSON only the outcome travels, since the promise is for a person to read.
+	var preface []byte
+	if cmd.Then != nil {
+		if format == commands.FormatText && !r.Empty() {
+			buf := &bytes.Buffer{}
+			if err := commands.Render(cmd, r, format, buf); err == nil {
+				preface = buf.Bytes()
+			}
+		}
+		r = cmd.Then()
+	}
 
 	// A command with nothing to say gets an empty, successful reply.
 	if r.Empty() {
-		return res, after
+		res.Text = preface
+		return res
 	}
 
-	return render(res, cmd, r, format), after
+	return render(res, cmd, r, format, preface)
 }
 
 // render turns a command's Result into the response, in the format the caller asked for.
 // Both formats come from commands.Render, which is where ssh and the control socket
 // meet: the JSON a program reads here is byte for byte what -json prints at an ssh
 // prompt, and the text is the command's own rendering, which only the daemon can run.
-func render(res *Response, cmd *commands.Command, r commands.Result, format commands.Format) *Response {
+func render(res *Response, cmd *commands.Command, r commands.Result, format commands.Format, preface []byte) *Response {
 	buf := &bytes.Buffer{}
 	if err := commands.Render(cmd, r, format, buf); err != nil {
 		res.Error = err.Error()
@@ -378,7 +376,7 @@ func render(res *Response, cmd *commands.Command, r commands.Result, format comm
 	}
 
 	// The bytes the command wrote, as they are; see Response.Text.
-	res.Text = buf.Bytes()
+	res.Text = append(preface, buf.Bytes()...)
 	return res
 }
 

localapi/server_test.go

@@ -858,42 +858,34 @@ func TestCommandRegisteredAfterStartIsServed(t *testing.T) {
 	assert.JSONEq(t, `"late"`, string(res.Result))
 }
 
-// TestAfterReplyRunsOnceTheResponseIsOut: the hook blocks until the test releases it, so
-// the only way a response can arrive is if it was written before the hook began.
-func TestAfterReplyRunsOnceTheResponseIsOut(t *testing.T) {
-	release := make(chan struct{})
-	acted := make(chan struct{})
-
+// TestThenIsAnsweredInTheOneReply: this surface answers in one piece, so a command that
+// acts does so before the reply, and the reply carries Run's line and Then's outcome in
+// text, and the outcome alone in JSON.
+func TestThenIsAnsweredInTheOneReply(t *testing.T) {
 	reg := commands.New()
+	acted := false
 	reg.Register(&commands.Command{
 		Name:             "act",
 		ShortDescription: "replies, then acts",
 		Run: func(fs any, a []string) (commands.Result, error) {
 			return commands.Did("about to act"), nil
 		},
-		AfterReply: func() {
-			<-release
-			close(acted)
+		Then: func() commands.Result {
+			acted = true
+			return commands.Did("acted")
 		},
 	})
 	_, path := serve(t, reg)
 
-	res, err := (&Client{Path: path}).Invoke("act", nil, nil, "")
+	res, err := (&Client{Path: path}).Invoke("act", nil, nil, "text")
 	require.NoError(t, err)
-	require.Empty(t, res.Error)
-
-	select {
-	case <-acted:
-		t.Fatal("AfterReply ran before the response was delivered")
-	default:
-	}
+	assert.True(t, acted)
+	assert.Empty(t, res.Error)
+	assert.Equal(t, "about to act\nacted\n", string(res.Text))
 
-	close(release)
-	select {
-	case <-acted:
-	case <-time.After(2 * time.Second):
-		t.Fatal("AfterReply never ran")
-	}
+	res, err = (&Client{Path: path}).Invoke("act", nil, nil, "")
+	require.NoError(t, err)
+	assert.JSONEq(t, `{"ok":true,"reason":"acted"}`, string(res.Result))
 }
 
 // TestRebindRefusesDuringShutdown: a reload that lands after the daemon's context is

ssh_reload_test.go

@@ -3,7 +3,6 @@ package nebula
 import (
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/slackhq/nebula/config"
@@ -32,10 +31,9 @@ func TestReloadCommandReportsTheResult(t *testing.T) {
 	// reload reports what happened. A reload that failed is a refusal, exit 1 over ssh;
 	// one that merely needs a restart is not.
 	reload := func(failed bool) string {
-		res, err := runReload(l, c)
-		require.NoError(t, err)
+		res := runReload(l, c)
 		assert.Equal(t, !failed, res.OK)
-		return strings.TrimPrefix(res.Reason, "Reloading config\n")
+		return res.Reason
 	}
 
 	write("a: 2\nb: 1")