fix(agent): Close stdin and stdout separately to fix pty output loss (#6862)

Fixes #6656
Closes #6840

agent/agent.go

@@ -1115,7 +1115,8 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
 		go func() {
 			for win := range windowSize {
 				resizeErr := ptty.Resize(uint16(win.Height), uint16(win.Width))
-				if resizeErr != nil {
+				// If the pty is closed, then command has exited, no need to log.
+				if resizeErr != nil && !errors.Is(resizeErr, pty.ErrClosed) {
 					a.logger.Warn(ctx, "failed to resize tty", slog.Error(resizeErr))
 				}
 			}
@@ -1131,19 +1132,32 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
 		// output being lost. To avoid this, we wait for the output copy to
 		// start before waiting for the command to exit. This ensures that the
 		// output copy goroutine will be scheduled before calling close on the
-		// pty. There is still a risk of data loss if a command produces a lot
-		// of output, see TestAgent_Session_TTY_HugeOutputIsNotLost (skipped).
+		// pty. This shouldn't be needed because of `pty.Dup()` below, but it
+		// may not be supported on all platforms.
 		outputCopyStarted := make(chan struct{})
-		ptyOutput := func() io.Reader {
+		ptyOutput := func() io.ReadCloser {
 			defer close(outputCopyStarted)
-			return ptty.Output()
+			// Try to dup so we can separate stdin and stdout closure.
+			// Once the original pty is closed, the dup will return
+			// input/output error once the buffered data has been read.
+			stdout, err := ptty.Dup()
+			if err == nil {
+				return stdout
+			}
+			// If we can't dup, we shouldn't close
+			// the fd since it's tied to stdin.
+			return readNopCloser{ptty.Output()}
 		}
 		wg.Add(1)
 		go func() {
 			// Ensure data is flushed to session on command exit, if we
 			// close the session too soon, we might lose data.
 			defer wg.Done()
-			_, _ = io.Copy(session, ptyOutput())
+
+			stdout := ptyOutput()
+			defer stdout.Close()
+
+			_, _ = io.Copy(session, stdout)
 		}()
 		<-outputCopyStarted
 
@@ -1176,6 +1190,11 @@ func (a *agent) handleSSHSession(session ssh.Session) (retErr error) {
 	return cmd.Wait()
 }
 
+type readNopCloser struct{ io.Reader }
+
+// Close implements io.Closer.
+func (readNopCloser) Close() error { return nil }
+
 func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, msg codersdk.WorkspaceAgentReconnectingPTYInit, conn net.Conn) (retErr error) {
 	defer conn.Close()
 

agent/agent_test.go

@@ -350,15 +350,8 @@ func TestAgent_Session_TTY_Hushlogin(t *testing.T) {
 
 func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
 	t.Parallel()
-	if runtime.GOOS == "windows" {
-		// This might be our implementation, or ConPTY itself.
-		// It's difficult to find extensive tests for it, so
-		// it seems like it could be either.
-		t.Skip("ConPTY appears to be inconsistent on Windows.")
-	}
-
 	// This test is here to prevent regressions where quickly executing
-	// commands (with TTY) don't flush their output to the SSH session.
+	// commands (with TTY) don't sync their output to the SSH session.
 	//
 	// See: https://github.com/coder/coder/issues/6656
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -404,20 +397,13 @@ func TestAgent_Session_TTY_FastCommandHasOutput(t *testing.T) {
 
 func TestAgent_Session_TTY_HugeOutputIsNotLost(t *testing.T) {
 	t.Parallel()
-	if runtime.GOOS == "windows" {
-		// This might be our implementation, or ConPTY itself.
-		// It's difficult to find extensive tests for it, so
-		// it seems like it could be either.
-		t.Skip("ConPTY appears to be inconsistent on Windows.")
-	}
-	t.Skip("This test proves we have a bug where parts of large output on a PTY can be lost after the command exits, skipped to avoid test failures.")
 
-	// This test is here to prevent prove we have a bug where quickly executing
-	// commands (with TTY) don't flush their output to the SSH session. This is
-	// due to the pty being closed before all the output has been copied, but
-	// protecting against this requires a non-trivial rewrite of the output
-	// processing (or figuring out a way to put the pty in a mode where this
-	// does not happen).
+	// This test is here to prevent regressions where a command (with or
+	// without) a large amount of output would not be fully copied to the
+	// SSH session. On unix systems, this was fixed by duplicating the file
+	// descriptor of the PTY master and using it for copying the output.
+	//
+	// See: https://github.com/coder/coder/issues/6656
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled