feat: update codex formatting and prevent terminal echo from being captured as agent messages (#200)

35C4n0r · web-flow · commit 702b10b653e2 · 2026-03-05T17:54:32.000+05:30
diff --git a/lib/msgfmt/agent_readiness.go b/lib/msgfmt/agent_readiness.go
@@ -45,7 +45,7 @@ func isOpencodeAgentReadyForInitialPrompt(message string) bool {
 
 func isCodexAgentReadyForInitialPrompt(message string) bool {
 	message = trimEmptyLines(message)
-	messageWithoutInputBox := removeCodexInputBox(message)
+	messageWithoutInputBox := removeCodexMessageBox(message)
 	return len(messageWithoutInputBox) != len(message)
 }
 
diff --git a/lib/msgfmt/message_box.go b/lib/msgfmt/message_box.go
@@ -51,12 +51,11 @@ func removeMessageBox(msg string) string {
 	return strings.Join(lines, "\n")
 }
 
-func removeCodexInputBox(msg string) string {
+func removeCodexMessageBox(msg string) string {
 	lines := strings.Split(msg, "\n")
-	// Remove the input box, we need to match the exact pattern, because thinking follows the same pattern of ▌ followed by text
-	if len(lines) >= 2 && strings.Contains(lines[len(lines)-2], "▌ Ask Codex to do anything") {
-		idx := len(lines) - 2
-		lines = append(lines[:idx], lines[idx+1:]...)
+	if len(lines) >= 3 && strings.Contains(lines[len(lines)-3], "›") {
+		idx := len(lines) - 3
+		lines = append(lines[:idx], lines[idx+2])
 	}
 	return strings.Join(lines, "\n")
 }
diff --git a/lib/msgfmt/msgfmt.go b/lib/msgfmt/msgfmt.go
@@ -43,8 +43,8 @@ func IndexSubslice[T comparable](s, sub []T) int {
 // Return the runes, the lines, and the rune to line location mapping.
 func normalizeAndGetRuneLineMapping(msgRaw string) ([]rune, []string, []int) {
 	msgLines := strings.Split(msgRaw, "\n")
-	msgRuneLineLocations := []int{}
-	runes := []rune{}
+	var msgRuneLineLocations []int
+	var runes []rune
 	for lineIdx, line := range msgLines {
 		for _, r := range line {
 			if !strings.ContainsRune(WhiteSpaceChars, r) {
@@ -256,7 +256,7 @@ func formatGenericMessage(message string, userInput string, agentType AgentType)
 
 func formatCodexMessage(message string, userInput string) string {
 	message = RemoveUserInput(message, userInput, AgentTypeCodex)
-	message = removeCodexInputBox(message)
+	message = removeCodexMessageBox(message)
 	message = trimEmptyLines(message)
 	return message
 }
diff --git a/lib/msgfmt/testdata/format/codex/first_message/expected.txt b/lib/msgfmt/testdata/format/codex/first_message/expected.txt
@@ -7,4 +7,4 @@
  /diff - show git diff (including untracked files)
  /prompts - show example prompts
 
- ⏎ send   Shift+⏎ newline   Ctrl+C quit
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/format/codex/first_message/msg.txt b/lib/msgfmt/testdata/format/codex/first_message/msg.txt
@@ -7,5 +7,6 @@
  /diff - show git diff (including untracked files)
  /prompts - show example prompts
 
-▌ Ask Codex to do anything
- ⏎ send   Shift+⏎ newline   Ctrl+C quit
+› A for Apple
+
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/format/codex/multi-line-input/expected.txt b/lib/msgfmt/testdata/format/codex/multi-line-input/expected.txt
@@ -9,4 +9,4 @@ The `formatCodexMessage` function is defined in:
 
 `lib/msgfmt/msgfmt.go` (around line 219)
 
- Ctrl+C again to quit   12437 tokens used   96% context left 
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/format/codex/multi-line-input/msg.txt b/lib/msgfmt/testdata/format/codex/multi-line-input/msg.txt
@@ -18,5 +18,6 @@ The `formatCodexMessage` function is defined in:
 
 `lib/msgfmt/msgfmt.go` (around line 219)
 
-▌ Ask Codex to do anything                                                                                                                                                                                
- Ctrl+C again to quit   12437 tokens used   96% context left 
+› Ask Codex to do anything
+
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/format/codex/second_message/expected.txt b/lib/msgfmt/testdata/format/codex/second_message/expected.txt
@@ -7,4 +7,4 @@
 codex
 There are 2 untracked files (`.env` and `forge.yaml`).
 
- ⏎ send   Ctrl+J newline   Ctrl+C quit   18724 tokens used   96% context left
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/format/codex/second_message/msg.txt b/lib/msgfmt/testdata/format/codex/second_message/msg.txt
@@ -10,5 +10,6 @@ How many untracked files are there?
 codex
 There are 2 untracked files (`.env` and `forge.yaml`).
 
-▌ Ask Codex to do anything
- ⏎ send   Ctrl+J newline   Ctrl+C quit   18724 tokens used   96% context left
+› Ask Codex to do anything
+
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/msgfmt/testdata/initialization/codex/ready/msg.txt b/lib/msgfmt/testdata/initialization/codex/ready/msg.txt
@@ -7,5 +7,6 @@
  /diff - show git diff (including untracked files)
  /prompts - show example prompts
 
-▌ Ask Codex to do anything
- ⏎ send   Shift+⏎ newline   Ctrl+C quit
+› Refactor this Code...
+
+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…
diff --git a/lib/screentracker/pty_conversation.go b/lib/screentracker/pty_conversation.go
@@ -124,6 +124,9 @@ type PTYConversation struct {
 	// Set under lock in the snapshot loop when signaling, cleared under
 	// lock in the send loop after sendMessage returns.
 	sendingMessage bool
+	// writingMessage is true while writeStabilize is executing.
+	// When true, updateLastAgentMessageLocked skips updates to avoid capturing terminal echo.
+	writingMessage bool
 	// stableSignal is used by the snapshot loop to signal the send loop
 	// when the agent is stable and there are items in the outbound queue.
 	stableSignal chan struct{}
@@ -177,6 +180,7 @@ func NewPTY(ctx context.Context, cfg PTYConversationConfig, emitter Emitter) *PT
 		dirty:                         false,
 		userSentMessageAfterLoadState: false,
 		loadStateStatus:               LoadStatePending,
+		writingMessage:                false,
 	}
 	if c.cfg.ReadyForInitialPrompt == nil {
 		c.cfg.ReadyForInitialPrompt = func(string) bool { return true }
@@ -295,6 +299,9 @@ func (c *PTYConversation) lastMessage(role ConversationRole) ConversationMessage
 
 // caller MUST hold c.lock
 func (c *PTYConversation) updateLastAgentMessageLocked(screen string, timestamp time.Time) {
+	if c.writingMessage {
+		return
+	}
 	agentMessage := screenDiff(c.screenBeforeLastUserMessage, screen, c.cfg.AgentType)
 	lastUserMessage := c.lastMessage(ConversationRoleUser)
 	var toolCalls []string
@@ -380,23 +387,17 @@ func (c *PTYConversation) sendMessage(ctx context.Context, messageParts ...Messa
 	screenBeforeMessage := c.cfg.AgentIO.ReadScreen()
 	now := c.cfg.Clock.Now()
 	c.updateLastAgentMessageLocked(screenBeforeMessage, now)
+	c.writingMessage = true
 	c.lock.Unlock()
 
 	if err := c.writeStabilize(ctx, messageParts...); err != nil {
+		c.lock.Lock()
+		defer c.lock.Unlock()
+		c.writingMessage = false
 		return xerrors.Errorf("failed to send message: %w", err)
 	}
 
 	c.lock.Lock()
-	// Re-apply the pre-send agent message from the screen captured before
-	// the write. While the lock was released during writeStabilize, the
-	// snapshot loop continued taking snapshots and calling
-	// updateLastAgentMessageLocked with whatever was on screen at each
-	// tick (typically echoed user input or intermediate terminal state).
-	// Those updates corrupt the agent message for this turn. Restoring it
-	// here ensures the conversation history is correct. The next line sets
-	// screenBeforeLastUserMessage so the *next* agent message will be
-	// diffed relative to the pre-send screen.
-	c.updateLastAgentMessageLocked(screenBeforeMessage, now)
 	c.screenBeforeLastUserMessage = screenBeforeMessage
 	c.messages = append(c.messages, ConversationMessage{
 		Id:      len(c.messages),
@@ -405,7 +406,7 @@ func (c *PTYConversation) sendMessage(ctx context.Context, messageParts ...Messa
 		Time:    now,
 	})
 	c.userSentMessageAfterLoadState = true
-
+	c.writingMessage = false
 	c.lock.Unlock()
 	return nil
 }

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,7 @@ func isOpencodeAgentReadyForInitialPrompt(message string) bool {`
`45`	`45`
`46`	`46`	`func isCodexAgentReadyForInitialPrompt(message string) bool {`
`47`	`47`	`message = trimEmptyLines(message)`
`48`		`- messageWithoutInputBox := removeCodexInputBox(message)`
	`48`	`+ messageWithoutInputBox := removeCodexMessageBox(message)`
`49`	`49`	`return len(messageWithoutInputBox) != len(message)`
`50`	`50`	`}`
`51`	`51`
Original file line number	Diff line number	Diff line change
`@@ -51,12 +51,11 @@ func removeMessageBox(msg string) string {`
`51`	`51`	`return strings.Join(lines, "\n")`
`52`	`52`	`}`
`53`	`53`
`54`		`-func removeCodexInputBox(msg string) string {`
	`54`	`+func removeCodexMessageBox(msg string) string {`
`55`	`55`	`lines := strings.Split(msg, "\n")`
`56`		`- // Remove the input box, we need to match the exact pattern, because thinking follows the same pattern of ▌ followed by text`
`57`		`- if len(lines) >= 2 && strings.Contains(lines[len(lines)-2], "▌ Ask Codex to do anything") {`
`58`		`- idx := len(lines) - 2`
`59`		`- lines = append(lines[:idx], lines[idx+1:]...)`
	`56`	`+ if len(lines) >= 3 && strings.Contains(lines[len(lines)-3], "›") {`
	`57`	`+ idx := len(lines) - 3`
	`58`	`+ lines = append(lines[:idx], lines[idx+2])`
`60`	`59`	`}`
`61`	`60`	`return strings.Join(lines, "\n")`
`62`	`61`	`}`
Original file line number	Diff line number	Diff line change
@@ -9,4 +9,4 @@ The `formatCodexMessage` function is defined in:
`9`	`9`
`10`	`10`	`lib/msgfmt/msgfmt.go` (around line 219)
`11`	`11`
`12`		`- Ctrl+C again to quit 12437 tokens used 96% context left`
	`12`	`+gpt-5.1-codex-max default · 100% left · ~/Documents/work/coder_o…`