fix(telegram): preserve code block when editing approval messages (#21)

* fix(telegram): preserve code block when editing approval messages

Edits on approval prompts (resolved by button, timed out, or
cancelled) were appending a status suffix to a stored rendered HTML
string. Under some conditions the appended suffix caused Telegram to
drop the <pre><code class="language-json"> formatting on the args
block, so approved/denied prompts lost the pretty-printed JSON code
block.

Switch msgMap to store the full ApprovalRequest instead of a
pre-rendered text string. Every edit path (CancelApproval, inflight
timeout, callback-resolved, callback-timed-out) now re-renders the
body by calling FormatApprovalMessage(am.req) fresh and appending the
status suffix to the result. This guarantees the <pre><code> block is
always intact in the final body we send to editMessageText.

* fix(telegram): pre-load msgMap to avoid edit race with cancelOnChannels

broker.Resolve synchronously calls cancelOnChannels -> tc.CancelApproval
on the same Telegram channel that just resolved the request. That cleanup
path did LoadAndDelete on msgMap and issued its own "(resolved via another
channel)" edit. Control then returned to handleCallback, which could no
longer find the entry in msgMap and fell back to editing with
cq.Message.Text - Telegram's plain-text extraction, where <pre>/<code>
tags are already stripped. The second edit clobbered the first, leaving
the user with plain-text JSON after tapping Allow/Deny.

Fix: handleCallback now LoadAndDelete's the msgMap entry into a local
'am' variable before calling broker.Resolve. This takes ownership of the
entry so CancelApproval becomes a no-op, and handleCallback's edit
becomes the single last-write-wins render with the <pre><code> block
intact and the "Allowed (once)/(request timed out)" label appended.

Add regression test TestHandleCallbackAllowOncePreservesCodeBlockOnEdit
that simulates an MCP approval with ToolArgs, taps Allow, and asserts
the final editMessageText payload still contains
<pre><code class="language-json"> and the status label.

Author: nnemirovsky

internal/telegram/approval.go

@@ -101,11 +101,15 @@ func NewTelegramChannel(cfg ChannelConfig) (*TelegramChannel, error) {
 	return tc, nil
 }
 
-// approvalMsg stores the Telegram message ID and original text for an
-// approval request so that timeout/cancel edits can preserve context.
+// approvalMsg stores the Telegram message ID and the original approval
+// request so that timeout/cancel/resolved edits can re-render the full
+// HTML-formatted body (including the <pre><code> args block) without
+// worrying about whether the stored text survived appending, parse
+// mode changes, or Telegram's plain-text extraction in callback
+// queries.
 type approvalMsg struct {
 	messageID int
-	text      string
+	req       channel.ApprovalRequest
 }
 
 // SetBroker sets the broker reference for resolving approval requests.
@@ -155,15 +159,14 @@ func (tc *TelegramChannel) sendApprovalMessage(req channel.ApprovalRequest) {
 		}
 		return
 	}
-	originalText := FormatApprovalMessage(req)
-	tc.msgMap.Store(req.ID, approvalMsg{messageID: sent.MessageID, text: originalText})
+	tc.msgMap.Store(req.ID, approvalMsg{messageID: sent.MessageID, req: req})
 
 	// If the request timed out while the Telegram API call was in
 	// flight, update the message immediately so the user does not
 	// see a stale prompt.
 	if tc.broker != nil && tc.broker.WasTimedOut(req.ID) {
 		edit := tgbotapi.NewEditMessageText(tc.chatID, sent.MessageID,
-			originalText+"\n\n(request timed out)")
+			FormatApprovalMessage(req)+"\n\n(request timed out)")
 		edit.ParseMode = tgbotapi.ModeHTML
 		if _, editErr := tc.api.Send(edit); editErr == nil {
 			tc.broker.ClearTimedOut(req.ID)
@@ -186,7 +189,8 @@ func (tc *TelegramChannel) CancelApproval(id string) error {
 	} else if tc.broker != nil && tc.broker.IsClosed() {
 		reason = "(proxy shutting down)"
 	}
-	edit := tgbotapi.NewEditMessageText(tc.chatID, am.messageID, am.text+"\n\n"+reason)
+	edit := tgbotapi.NewEditMessageText(tc.chatID, am.messageID,
+		FormatApprovalMessage(am.req)+"\n\n"+reason)
 	edit.ParseMode = tgbotapi.ModeHTML
 	_, _ = tc.api.Send(edit)
 	return nil
@@ -300,6 +304,22 @@ func (tc *TelegramChannel) handleCallback(cq *tgbotapi.CallbackQuery) {
 		return
 	}
 
+	// Take ownership of the msgMap entry BEFORE calling broker.Resolve.
+	// Resolve synchronously calls cancelOnChannels, which invokes
+	// tc.CancelApproval on this same Telegram channel. CancelApproval
+	// does its own LoadAndDelete and issues a "(resolved via another
+	// channel)" edit. If we let that happen first, the msgMap entry is
+	// gone by the time handleCallback runs its own edit, and we fall
+	// back to cq.Message.Text (Telegram's plain-text extraction where
+	// <pre><code> tags are already stripped). Pre-loading makes
+	// CancelApproval a no-op so handleCallback owns the final render.
+	var am approvalMsg
+	haveAM := false
+	if val, ok := tc.msgMap.LoadAndDelete(reqID); ok {
+		am = val.(approvalMsg)
+		haveAM = true
+	}
+
 	resolved := false
 	if tc.broker != nil {
 		resolved = tc.broker.Resolve(reqID, resp)
@@ -309,32 +329,37 @@ func (tc *TelegramChannel) handleCallback(cq *tgbotapi.CallbackQuery) {
 		callback := tgbotapi.NewCallback(cq.ID, label)
 		_, _ = tc.api.Request(callback)
 
-		// Use the original stored text (with HTML markup) rather than
-		// cq.Message.Text, which is Telegram's plain-text extraction
-		// and loses our <pre><code> formatting on tool arguments.
-		originalText := cq.Message.Text
-		if val, ok := tc.msgMap.Load(reqID); ok {
-			originalText = val.(approvalMsg).text
+		// Re-render the full HTML body from the stored request. This
+		// preserves the <pre><code> args block because we are not
+		// appending to or reusing a string that already contains HTML
+		// tags. cq.Message.Text is Telegram's plain-text extraction and
+		// loses the code block formatting entirely, so we fall back to
+		// a minimal plain message only if the msgMap entry is gone.
+		var body string
+		if haveAM {
+			body = fmt.Sprintf("%s\n\n%s at %s",
+				FormatApprovalMessage(am.req), label, time.Now().UTC().Format("15:04:05"))
+		} else {
+			body = fmt.Sprintf("%s\n\n%s at %s",
+				cq.Message.Text, label, time.Now().UTC().Format("15:04:05"))
 		}
-		edit := tgbotapi.NewEditMessageText(tc.chatID, cq.Message.MessageID,
-			fmt.Sprintf("%s\n\n%s at %s", originalText, label, time.Now().UTC().Format("15:04:05")))
+		edit := tgbotapi.NewEditMessageText(tc.chatID, cq.Message.MessageID, body)
 		edit.ParseMode = tgbotapi.ModeHTML
 		_, _ = tc.api.Send(edit)
-		tc.msgMap.Delete(reqID)
 	} else if tc.broker != nil && tc.broker.WasTimedOut(reqID) {
 		tc.broker.ClearTimedOut(reqID)
 		callback := tgbotapi.NewCallback(cq.ID, "Request timed out")
 		_, _ = tc.api.Request(callback)
 
-		originalText := cq.Message.Text
-		if val, ok := tc.msgMap.Load(reqID); ok {
-			originalText = val.(approvalMsg).text
+		var body string
+		if haveAM {
+			body = FormatApprovalMessage(am.req) + "\n\n(request timed out)"
+		} else {
+			body = cq.Message.Text + "\n\n(request timed out)"
 		}
-		edit := tgbotapi.NewEditMessageText(tc.chatID, cq.Message.MessageID,
-			originalText+"\n\n(request timed out)")
+		edit := tgbotapi.NewEditMessageText(tc.chatID, cq.Message.MessageID, body)
 		edit.ParseMode = tgbotapi.ModeHTML
 		_, _ = tc.api.Send(edit)
-		tc.msgMap.Delete(reqID)
 
 	} else {
 		// Request was already resolved by a previous callback (e.g. double-tap

internal/telegram/approval_test.go

@@ -459,7 +459,10 @@ func TestCancelApproval(t *testing.T) {
 	tc.SetBroker(broker)
 
 	// Store a message ID mapping.
-	tc.msgMap.Store("req_test", approvalMsg{messageID: 42, text: "test approval message"})
+	tc.msgMap.Store("req_test", approvalMsg{
+		messageID: 42,
+		req:       channel.ApprovalRequest{ID: "req_test", Destination: "test.example.com", Port: 443, Protocol: "https"},
+	})
 
 	err := tc.CancelApproval("req_test")
 	if err != nil {
@@ -508,7 +511,10 @@ func TestCancelApprovalShowsShutdownReason(t *testing.T) {
 	tc.SetBroker(broker)
 	broker.CancelAll() // Mark broker as closed.
 
-	tc.msgMap.Store("req_shutdown", approvalMsg{messageID: 43, text: "shutdown test message"})
+	tc.msgMap.Store("req_shutdown", approvalMsg{
+		messageID: 43,
+		req:       channel.ApprovalRequest{ID: "req_shutdown", Destination: "shutdown.example.com", Port: 443, Protocol: "https"},
+	})
 	_ = tc.CancelApproval("req_shutdown")
 
 	time.Sleep(50 * time.Millisecond)
@@ -808,6 +814,90 @@ func TestHandleCallbackUnknownAction(t *testing.T) {
 	})
 }
 
+// TestHandleCallbackAllowOncePreservesCodeBlockOnEdit reproduces the bug
+// where the <pre><code> args block was lost from the final edit after the
+// user tapped Allow. The root cause was that broker.Resolve synchronously
+// calls cancelOnChannels -> tc.CancelApproval on the same Telegram channel,
+// which LoadAndDelete'd the msgMap entry and issued its own edit. Control
+// then returned to handleCallback, which could not find the entry in
+// msgMap and fell back to editing with cq.Message.Text (Telegram's
+// plain-text extraction, where <pre>/<code> tags are already stripped).
+// That second edit clobbered the first and lost the code block. The fix
+// is for handleCallback to take ownership of the msgMap entry BEFORE
+// calling broker.Resolve so CancelApproval becomes a no-op.
+func TestHandleCallbackAllowOncePreservesCodeBlockOnEdit(t *testing.T) {
+	mock := newMockTelegramAPI(t)
+	s := newTestStore(t)
+	tc := newTestTelegramChannel(t, mock, s)
+
+	broker := channel.NewBroker([]channel.Channel{tc})
+	tc.SetBroker(broker)
+
+	const toolArgs = `{"owner":"me","repo":"x"}`
+	done := make(chan channel.Response, 1)
+	go func() {
+		resp, _ := broker.Request(
+			"github__list_branches", 0, "mcp", 5*time.Second,
+			channel.WithToolArgs(toolArgs),
+		)
+		done <- resp
+	}()
+
+	waitForPending(t, broker, 1)
+	reqs := broker.PendingRequests()
+	if len(reqs) != 1 {
+		t.Fatalf("expected 1 pending request, got %d", len(reqs))
+	}
+	req := reqs[0]
+
+	// Wait for sendApprovalMessage goroutine to populate msgMap (it runs
+	// async after broadcast). The mock API returns immediately so this is
+	// typically fast.
+	deadline := time.Now().Add(500 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		if _, ok := tc.msgMap.Load(req.ID); ok {
+			break
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	if _, ok := tc.msgMap.Load(req.ID); !ok {
+		t.Fatal("msgMap entry was not populated by sendApprovalMessage")
+	}
+
+	// Simulate user tapping "Allow". Message.Text is what Telegram returns
+	// in callback queries: the sent HTML text with all formatting tags
+	// stripped. The JSON content is still present (newlines preserved)
+	// but the <pre><code> wrapping is gone.
+	tc.handleCallback(&tgbotapi.CallbackQuery{
+		ID: "cb_allow",
+		Message: &tgbotapi.Message{
+			MessageID: 101,
+			Chat:      &tgbotapi.Chat{ID: 12345},
+			Text:      "OpenClaw wants to call tool:\n\ngithub__list_branches\n\nArguments:\n{\n  \"owner\": \"me\",\n  \"repo\": \"x\"\n}\n\nAllow this tool call?",
+		},
+		Data: req.ID + "|allow_once",
+	})
+
+	if resp := <-done; resp != channel.ResponseAllowOnce {
+		t.Errorf("expected AllowOnce, got %v", resp)
+	}
+
+	// The FINAL edit (last-write-wins) must still contain the code block
+	// and the "Allowed (once)" status label. Multiple edits may fire
+	// (cancelOnChannels also triggers one), so we check the last entry.
+	edits := mock.getEditedMessages()
+	if len(edits) == 0 {
+		t.Fatal("expected at least one edit")
+	}
+	last := edits[len(edits)-1]
+	if !strings.Contains(last.Text, `<pre><code class="language-json">`) {
+		t.Errorf("final edit lost <pre><code> block.\nfinal edit text:\n%s", last.Text)
+	}
+	if !strings.Contains(last.Text, "Allowed (once)") {
+		t.Errorf("final edit missing 'Allowed (once)' label.\nfinal edit text:\n%s", last.Text)
+	}
+}
+
 func TestHandleCallbackAlreadyResolved(t *testing.T) {
 	mock := newMockTelegramAPI(t)
 	s := newTestStore(t)