Fix some responses from terminal sometimes leaking into shell on after kitten exit

Always do a roundtrip at kitten exit, except for special purpose
kittens. This slows down exit by one round trip time (capped at 2
seconds), however it ensures that we never get terminal response leak.

Fixes #9839

tools/tui/loop/api.go

@@ -58,11 +58,17 @@ type Loop struct {
 	style_ctx                              style.Context
 	atomic_update_active                   bool
 	pointer_shapes                         []PointerShape
-	waiting_for_capabilities_response      bool
 
 	// Queried capabilities from terminal
 	TerminalCapabilities TerminalCapabilities
 
+	// Set this to true to avoid doing a query response loop to the terminal at
+	// exit. This loop is needed for most kittens to ensure that in-flight
+	// responses such as in-band resize notifications, color queries, kitty
+	// keyboard events, etc. are not leaked to the shell. For some special
+	// purpose uses of the loop, this is not appropriate, hence this setting.
+	NoRoundtripToTerminalOnExit bool
+
 	// Suspend the loop restoring terminal state, and run the provided function. When it returns terminal state is
 	// put back to what it was before suspending unless the function returns an error or an error occurs saving/restoring state.
 	SuspendAndRun func(func() error) error
@@ -333,7 +339,7 @@ func (self *Loop) Run() (err error) {
 			os.Stderr.WriteString(err.Error())
 			os.Stderr.WriteString("\n")
 			if is_terminal {
-				if term, err := tty.OpenControllingTerm(tty.SetRaw); err == nil {
+				if term, err := tty.OpenControllingTerm(tty.SetRaw); err != nil {
 					defer term.RestoreAndClose()
 					term.DebugPrintln(err.Error())
 					fmt.Println("Press any key to exit.\r")
@@ -568,12 +574,7 @@ func (self *Loop) CurrentPointerShape() (ans PointerShape, has_shape bool) {
 // callback will be called once the query response is received. This
 // function should be called as early as possible ideally in OnInitialize.
 func (self *Loop) QueryCapabilities() {
-	if !self.waiting_for_capabilities_response {
-		self.waiting_for_capabilities_response = true
-		self.StartAtomicUpdate()
-		self.QueueWriteString("\x1b[?u\x1b[?996n\x1b[c")
-		self.EndAtomicUpdate()
-	}
+	self.QueueWriteString("\x1b[?u\x1b[?996n\x1b[c")
 }
 
 type Alignment int

tools/tui/loop/read.go

@@ -115,11 +115,17 @@ func has_da1_response(s string) bool {
 	return pat.FindString(s) != ""
 }
 
-func read_until_primary_device_attributes_response(term *tty.Term, initial_bytes []byte, timeout time.Duration) {
-	s := strings.Builder{}
-	if initial_bytes != nil {
-		s.Write(initial_bytes)
+func do_roundtrip_to_terminal(term *tty.Term, timeout time.Duration) {
+	// ask for primary device attributes
+	for {
+		if err := term.WriteAllString("\x1b[c"); err != nil && !is_temporary_error(err) {
+			return
+		} else {
+			break
+		}
 	}
+	s := strings.Builder{}
+	s.Grow(256)
 	received := make(chan error)
 	go func() {
 		defer func() {
@@ -128,17 +134,19 @@ func read_until_primary_device_attributes_response(term *tty.Term, initial_bytes
 				received <- fmt.Errorf("%s", text)
 			}
 		}()
-		buf := make([]byte, 1024)
-		n, err := read_ignoring_temporary_errors(term, buf)
-		if n > 0 {
-			s.Write(buf[:n])
-			if has_da1_response(s.String()) {
-				received <- nil
-				return
+		for {
+			buf := make([]byte, 1024)
+			n, err := read_ignoring_temporary_errors(term, buf)
+			if n > 0 {
+				s.Write(buf[:n])
+				if has_da1_response(s.String()) {
+					received <- nil
+					return
+				}
+			}
+			if err != nil {
+				received <- err
 			}
-		}
-		if err != nil {
-			received <- err
 		}
 	}()
 	select {

tools/tui/loop/run.go

@@ -87,38 +87,72 @@ func (self *Loop) update_screen_size() error {
 }
 
 func (self *Loop) handle_csi(raw []byte) (err error) {
+	if len(raw) == 0 {
+		return nil
+	}
 	csi := string(raw)
-	if strings.HasSuffix(csi, "t") && strings.HasPrefix(csi, "48;") {
-		if parts := strings.Split(csi[3:len(csi)-1], ";"); len(parts) > 3 {
-			var parsed [4]int
-			ok := true
-			for i, x := range parts {
-				x, _, _ = strings.Cut(x, ":")
-				if parsed[i], err = strconv.Atoi(x); err != nil {
-					ok = false
-					break
+	switch raw[len(raw)-1] {
+	case 't':
+		if strings.HasSuffix(csi, "t") {
+			if parts := strings.Split(csi[3:len(csi)-1], ";"); len(parts) > 3 {
+				var parsed [4]int
+				ok := true
+				for i, x := range parts {
+					x, _, _ = strings.Cut(x, ":")
+					if parsed[i], err = strconv.Atoi(x); err != nil {
+						ok = false
+						break
+					}
 				}
-			}
-			if ok {
-				self.seen_inband_resize = true
-				old_size := self.screen_size
-				s := &self.screen_size
-				s.updated = true
-				s.HeightCells, s.WidthCells = uint(parsed[0]), uint(parsed[1])
-				s.HeightPx, s.WidthPx = uint(parsed[2]), uint(parsed[3])
-				s.CellWidth = s.WidthPx / s.WidthCells
-				s.CellHeight = s.HeightPx / s.HeightCells
-				if self.OnResize != nil {
-					return self.OnResize(old_size, self.screen_size)
+				if ok {
+					self.seen_inband_resize = true
+					old_size := self.screen_size
+					s := &self.screen_size
+					s.updated = true
+					s.HeightCells, s.WidthCells = uint(parsed[0]), uint(parsed[1])
+					s.HeightPx, s.WidthPx = uint(parsed[2]), uint(parsed[3])
+					s.CellWidth = s.WidthPx / s.WidthCells
+					s.CellHeight = s.HeightPx / s.HeightCells
+					if self.OnResize != nil {
+						return self.OnResize(old_size, self.screen_size)
+					}
+					return nil
 				}
-				return nil
 			}
 		}
-	} else if csi == "I" || csi == "O" {
-		if self.OnFocusChange != nil {
-			return self.OnFocusChange(csi == "I")
+	case 'I', 'O':
+		if len(csi) == 1 {
+			if self.OnFocusChange != nil {
+				return self.OnFocusChange(csi == "I")
+			}
 		}
 		return nil
+	case 'c':
+		if strings.HasPrefix(csi, "?") {
+			if self.OnCapabilitiesReceived != nil {
+				if err = self.OnCapabilitiesReceived(self.TerminalCapabilities); err != nil {
+					return err
+				}
+			}
+		}
+	case 'u':
+		if strings.HasPrefix(csi, "?") {
+			self.TerminalCapabilities.KeyboardProtocol = true
+			self.TerminalCapabilities.KeyboardProtocolResponseReceived = true
+		}
+	case 'n':
+		if strings.HasPrefix(csi, "?997;") {
+			switch csi[len(csi)-2] {
+			case '1':
+				self.TerminalCapabilities.ColorPreference = DARK_COLOR_PREFERENCE
+			case '2':
+				self.TerminalCapabilities.ColorPreference = LIGHT_COLOR_PREFERENCE
+			}
+			self.TerminalCapabilities.ColorPreferenceResponseReceived = true
+			if self.OnColorSchemeChange != nil {
+				return self.OnColorSchemeChange(self.TerminalCapabilities.ColorPreference)
+			}
+		}
 	}
 	ke := KeyEventFromCSI(csi)
 	if ke != nil {
@@ -131,38 +165,6 @@ func (self *Loop) handle_csi(raw []byte) (err error) {
 			return self.handle_mouse_event(me)
 		}
 	}
-	if self.waiting_for_capabilities_response {
-		if strings.HasPrefix(csi, "?") && strings.HasSuffix(csi, "c") {
-			self.waiting_for_capabilities_response = false
-			if self.OnCapabilitiesReceived != nil {
-				if err = self.OnCapabilitiesReceived(self.TerminalCapabilities); err != nil {
-					return err
-				}
-			}
-		} else if strings.HasPrefix(csi, "?997;") && strings.HasSuffix(csi, "n") {
-			switch csi[len(csi)-2] {
-			case '1':
-				self.TerminalCapabilities.ColorPreference = DARK_COLOR_PREFERENCE
-			case '2':
-				self.TerminalCapabilities.ColorPreference = LIGHT_COLOR_PREFERENCE
-			}
-			self.TerminalCapabilities.ColorPreferenceResponseReceived = true
-		} else if strings.HasPrefix(csi, "?") && strings.HasSuffix(csi, "u") {
-			self.TerminalCapabilities.KeyboardProtocol = true
-			self.TerminalCapabilities.KeyboardProtocolResponseReceived = true
-		}
-	} else if self.terminal_options.color_scheme_change_notification && strings.HasPrefix(csi, "?997;") && strings.HasSuffix(csi, "n") {
-		switch csi[len(csi)-2] {
-		case '1':
-			self.TerminalCapabilities.ColorPreference = DARK_COLOR_PREFERENCE
-		case '2':
-			self.TerminalCapabilities.ColorPreference = LIGHT_COLOR_PREFERENCE
-		}
-		self.TerminalCapabilities.ColorPreferenceResponseReceived = true
-		if self.OnColorSchemeChange != nil {
-			return self.OnColorSchemeChange(self.TerminalCapabilities.ColorPreference)
-		}
-	}
 	if self.OnEscapeCode != nil {
 		return self.OnEscapeCode(CSI, raw)
 	}
@@ -446,19 +448,16 @@ func (self *Loop) run() (err error) {
 		// wait for tty reader to exit cleanly
 		for range tty_read_channel {
 		}
-		if !self.waiting_for_capabilities_response {
-			close(tty_leftover_read_channel)
-			return
-		}
-		var pending_bytes []byte
 		select {
-		case msg, ok := <-tty_leftover_read_channel:
-			if ok {
-				pending_bytes = msg
-			}
+		case <-tty_leftover_read_channel:
 		default:
 		}
-		read_until_primary_device_attributes_response(controlling_term, pending_bytes, 2*time.Second)
+		if !self.NoRoundtripToTerminalOnExit {
+			// ensure that any terminal responses such as kitty keyboard events,
+			// color scheme changes, in-band resize notifications, etc. do not
+			// bleed into the shell.
+			do_roundtrip_to_terminal(controlling_term, 2*time.Second)
+		}
 	}
 
 	defer func() {