TUI

cli/run.go

@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -20,11 +21,11 @@ import (
 	"syscall"
 	"time"
 
+	sdkv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/v1"
+	tea "github.com/charmbracelet/bubbletea"
 	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 	"google.golang.org/protobuf/proto"
-
-	sdkv1 "buf.build/gen/go/stealthrocket/dispatch-proto/protocolbuffers/go/dispatch/sdk/v1"
 )
 
 var (
@@ -81,24 +82,33 @@ previous run.`, defaultEndpoint),
 
 			prefixWidth := max(len("dispatch"), len(arg0))
 
-			if Verbose {
+			// Enable the TUI if this is an interactive session and
+			// stdout/stderr aren't redirected.
+			var tui *TUI
+			var logWriter io.Writer = os.Stderr
+			if isTerminal(os.Stdin) && isTerminal(os.Stdout) && isTerminal(os.Stderr) {
+				tui = &TUI{}
+				logWriter = tui
+			}
+
+			if Verbose || tui != nil {
 				prefix := []byte(pad("dispatch", prefixWidth) + " | ")
 				if Color {
 					prefix = []byte("\033[32m" + pad("dispatch", prefixWidth) + " \033[90m|\033[0m ")
 				}
-				// Print Dispatch logs with a prefix in verbose mode.
-				slog.SetDefault(slog.New(&prefixHandler{
-					Handler: slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug}),
-					stream:  os.Stderr,
-					prefix:  prefix,
-				}))
+				slog.SetDefault(slog.New(
+					slog.NewTextHandler(&prefixLogWriter{
+						stream: logWriter,
+						prefix: prefix,
+					}, &slog.HandlerOptions{Level: slog.LevelDebug}),
+				))
 			}
 
 			if BridgeSession == "" {
 				BridgeSession = uuid.New().String()
 			}
 
-			if Verbose {
+			if Verbose || tui != nil {
 				slog.Info("starting session", "session_id", BridgeSession)
 			} else {
 				dialog(`Starting Dispatch session: %v
@@ -117,7 +127,7 @@ Run 'dispatch help run' to learn about Dispatch sessions.`, BridgeSession)
 			// to disambiguate Dispatch logs from the local application's logs.
 			var stdout io.ReadCloser
 			var stderr io.ReadCloser
-			if Verbose {
+			if Verbose || tui != nil {
 				var err error
 				stdout, err = cmd.StdoutPipe()
 				if err != nil {
@@ -131,8 +141,8 @@ Run 'dispatch help run' to learn about Dispatch sessions.`, BridgeSession)
 				}
 				defer stderr.Close()
 			} else {
-				cmd.Stdout = os.Stdout
-				cmd.Stderr = os.Stderr
+				cmd.Stdout = logWriter
+				cmd.Stderr = logWriter
 			}
 
 			// Pass on environment variables to the local application.
@@ -176,6 +186,28 @@ Run 'dispatch help run' to learn about Dispatch sessions.`, BridgeSession)
 				}
 			}()
 
+			// Initialize the TUI.
+			if tui != nil {
+				p := tea.NewProgram(tui,
+					tea.WithContext(ctx),
+					tea.WithoutSignalHandler(),
+					tea.WithoutCatchPanics(),
+					tea.WithMouseCellMotion())
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+
+					if _, err := p.Run(); err != nil && !errors.Is(err, tea.ErrProgramKilled) {
+						panic(err)
+					}
+					// Quitting the TUI sends an implicit interrupt.
+					select {
+					case signals <- syscall.SIGINT:
+					default:
+					}
+				}()
+			}
+
 			bridgeSessionURL := fmt.Sprintf("%s/sessions/%s", DispatchBridgeUrl, BridgeSession)
 
 			// Poll for work in the background.
@@ -212,7 +244,7 @@ Run 'dispatch help run' to learn about Dispatch sessions.`, BridgeSession)
 					go func() {
 						defer wg.Done()
 
-						err := invoke(ctx, httpClient, bridgeSessionURL, requestID, res)
+						err := invoke(ctx, httpClient, bridgeSessionURL, requestID, res, tui)
 						res.Body.Close()
 						if err != nil {
 							if ctx.Err() == nil {
@@ -237,14 +269,14 @@ Run 'dispatch help run' to learn about Dispatch sessions.`, BridgeSession)
 				return fmt.Errorf("failed to start %s: %v", strings.Join(args, " "), err)
 			}
 
-			if Verbose {
+			if Verbose || tui != nil {
 				prefix := []byte(pad(arg0, prefixWidth) + " | ")
 				suffix := []byte("\n")
 				if Color {
 					prefix = []byte("\033[35m" + pad(arg0, prefixWidth) + " \033[90m|\033[0m ")
 				}
-				go printPrefixedLines(os.Stderr, stdout, prefix, suffix)
-				go printPrefixedLines(os.Stderr, stderr, prefix, suffix)
+				go printPrefixedLines(logWriter, stdout, prefix, suffix)
+				go printPrefixedLines(logWriter, stderr, prefix, suffix)
 			}
 
 			err = cmd.Wait()
@@ -317,7 +349,28 @@ func poll(ctx context.Context, client *http.Client, url string) (string, *http.R
 	return requestID, res, nil
 }
 
-func invoke(ctx context.Context, client *http.Client, url, requestID string, bridgeGetRes *http.Response) error {
+// FunctionCallObserver observes function call requests and responses.
+//
+// The observer may be invoked concurrently from many goroutines.
+type FunctionCallObserver interface {
+	// ObserveRequest observes a RunRequest as it passes from the API through
+	// the CLI to the local application.
+	ObserveRequest(*sdkv1.RunRequest)
+
+	// ObserveResponse observes a response to the RunRequest.
+	//
+	// If the RunResponse is nil, it means the local application did not return
+	// a valid response. If the http.Response is not nil, it means an HTTP
+	// response was generated, but it wasn't a valid RunResponse. The error may
+	// be present if there was either an error making the HTTP request, or parsing
+	// the response.
+	//
+	// ObserveResponse always comes after a call to ObserveRequest for any given
+	// RunRequest.
+	ObserveResponse(*sdkv1.RunRequest, error, *http.Response, *sdkv1.RunResponse)
+}
+
+func invoke(ctx context.Context, client *http.Client, url, requestID string, bridgeGetRes *http.Response, observer FunctionCallObserver) error {
 	slog.Debug("sending request to local application", "endpoint", LocalEndpoint, "request_id", requestID)
 
 	// Extract the nested request header/body.
@@ -354,6 +407,9 @@ func invoke(ctx context.Context, client *http.Client, url, requestID string, bri
 	case *sdkv1.RunRequest_PollResult:
 		slog.Debug("resuming function", "function", runRequest.Function, "call_results", len(d.PollResult.Results), "request_id", requestID)
 	}
+	if observer != nil {
+		observer.ObserveRequest(&runRequest)
+	}
 
 	// The RequestURI field must be cleared for client.Do() to
 	// accept the request below.
@@ -365,6 +421,9 @@ func invoke(ctx context.Context, client *http.Client, url, requestID string, bri
 	endpointReq.URL.Host = LocalEndpoint
 	endpointRes, err := client.Do(endpointReq)
 	if err != nil {
+		if observer != nil {
+			observer.ObserveResponse(&runRequest, err, nil, nil)
+		}
 		return fmt.Errorf("failed to contact local application endpoint (%s): %v. Please check that -e,--endpoint is correct.", LocalEndpoint, err)
 	}
 
@@ -376,6 +435,9 @@ func invoke(ctx context.Context, client *http.Client, url, requestID string, bri
 	_, err = io.Copy(endpointResBody, endpointRes.Body)
 	endpointRes.Body.Close()
 	if err != nil {
+		if observer != nil {
+			observer.ObserveResponse(&runRequest, err, endpointRes, nil)
+		}
 		return fmt.Errorf("failed to read response from local application endpoint (%s): %v", LocalEndpoint, err)
 	}
 	endpointRes.Body = io.NopCloser(endpointResBody)
@@ -385,6 +447,9 @@ func invoke(ctx context.Context, client *http.Client, url, requestID string, bri
 	if endpointRes.StatusCode == http.StatusOK && endpointRes.Header.Get("Content-Type") == "application/proto" {
 		var runResponse sdkv1.RunResponse
 		if err := proto.Unmarshal(endpointResBody.Bytes(), &runResponse); err != nil {
+			if observer != nil {
+				observer.ObserveResponse(&runRequest, err, endpointRes, nil)
+			}
 			return fmt.Errorf("invalid response from local application endpoint (%s): %v", LocalEndpoint, err)
 		}
 		switch runResponse.Status {
@@ -402,9 +467,15 @@ func invoke(ctx context.Context, client *http.Client, url, requestID string, bri
 		default:
 			slog.Debug("function call failed", "function", runRequest.Function, "status", runResponse.Status, "request_id", requestID)
 		}
+		if observer != nil {
+			observer.ObserveResponse(&runRequest, nil, endpointRes, &runResponse)
+		}
 	} else {
 		// The response might indicate some other issue, e.g. it could be a 404 if the function can't be found
 		slog.Debug("function call failed", "function", runRequest.Function, "http_status", endpointRes.StatusCode, "request_id", requestID)
+		if observer != nil {
+			observer.ObserveResponse(&runRequest, nil, endpointRes, nil)
+		}
 	}
 
 	// Use io.Pipe to convert the response writer into an io.Reader.
@@ -484,22 +555,20 @@ func withoutEnv(env []string, prefixes ...string) []string {
 	})
 }
 
-type prefixHandler struct {
-	slog.Handler
+type prefixLogWriter struct {
 	stream io.Writer
 	prefix []byte
-	suffix []byte
 }
 
-func (h *prefixHandler) Handle(ctx context.Context, r slog.Record) error {
-	if _, err := h.stream.Write(h.prefix); err != nil {
-		return err
+func (p *prefixLogWriter) Write(b []byte) (int, error) {
+	var buffer bytes.Buffer
+	if _, err := buffer.Write(p.prefix); err != nil {
+		return 0, err
 	}
-	if err := h.Handler.Handle(ctx, r); err != nil {
-		return err
+	if _, err := buffer.Write(b); err != nil {
+		return 0, err
 	}
-	_, err := h.stream.Write(h.suffix)
-	return err
+	return p.stream.Write(buffer.Bytes())
 }
 
 func printPrefixedLines(w io.Writer, r io.Reader, prefix, suffix []byte) {