Skip to content

feat: add Telegram message editing support#185

Closed
Marenz wants to merge 1 commit intospacedriveapp:mainfrom
Marenz:feat/telegram-edit-message
Closed

feat: add Telegram message editing support#185
Marenz wants to merge 1 commit intospacedriveapp:mainfrom
Marenz:feat/telegram-edit-message

Conversation

@Marenz
Copy link
Collaborator

@Marenz Marenz commented Feb 24, 2026

This PR adds Telegram message editing support via editMessageText API. See: https://github.com/Marenz/spacebot/pull/new/feat/telegram-edit-message

Note

Summary

This PR introduces message editing capability to Spacebot, enabling agents to modify previously sent Telegram messages. The implementation adds a new edit_message tool that allows agents to specify a message ID and new content, with Telegram-specific handling via the editMessageText API and graceful no-op fallbacks for other messaging platforms.

Key changes:
A new EditMessageTool in src/tools/edit_message.rs that sends OutboundResponse::EditMessage events. A new OutboundResponse::EditMessage variant carrying message_id and text fields. Integration of an edited_flag through the channel event loop to track when a message edit occurs. Platform-specific implementations: Telegram uses editMessageText with automatic HTML fallback to plain text; Discord, Slack, Twitch, and Webhook adapters handle edits gracefully without errors. Updates to src/agent/channel.rs to thread the edited_flag through the agent turn and result handling.

Written by Tembo for commit 436fe39e. This will update automatically on new commits.

result: std::result::Result<String, rig::completion::PromptError>,
skip_flag: &crate::tools::SkipFlag,
replied_flag: &crate::tools::RepliedFlag,
edited_flag: &crate::tools::EditedFlag,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

edited_flag gets threaded through here, but handle_agent_result still only checks skip_flag/replied_flag. If the model calls edit_message without reply, you'll likely still hit the fallback text send and double-post. Consider loading edited_flag and treating it like replied_flag for fallback suppression.

StreamEnd,
/// Edit a previously sent message by ID.
/// Telegram: edits the text of a message using editMessageText.
/// Other platforms: falls back to sending a new message (no-op).
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment says other platforms "falls back to sending a new message", but adapters currently treat EditMessage as a no-op. Suggest tightening this doc to match behavior.

Suggested change
/// Other platforms: falls back to sending a new message (no-op).
/// Other platforms: no-op (message unchanged).

Comment on lines +478 to +480
let html = markdown_to_telegram_html(&text);

if let Err(html_error) = self
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: there’s a whitespace-only line after markdown_to_telegram_html(&text); (shows up in the diff).

Suggested change
let html = markdown_to_telegram_html(&text);
if let Err(html_error) = self
let html = markdown_to_telegram_html(&text);
if let Err(html_error) = self

Also, this branch duplicates the StreamChunk edit logic above (markdown->HTML->plain fallback). Might be worth extracting a small helper to keep behavior consistent.

pub struct EditMessageTool {
response_tx: mpsc::Sender<OutboundResponse>,
conversation_id: String,
conversation_logger: ConversationLogger,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

EditMessageTool stores conversation_logger/channel_id, but call() doesn’t use them yet. If they’re intended for future (logging edits / mention conversion), maybe add a quick note or wire it up; otherwise consider dropping them from the struct/ctor to keep the tool minimal.

- Add OutboundResponse::EditMessage variant for editing previously sent messages
- Implement editMessageText API in Telegram adapter
- Add edit_message tool for agents to edit messages by ID
- Add fallback handlers in other messaging adapters (Discord, Slack, Twitch, webchat, webhook)

The edit_message tool allows the agent to edit a Telegram message by providing
the message_id and new content. On other platforms, calls are logged but
have no effect.
@Marenz Marenz force-pushed the feat/telegram-edit-message branch from 436fe39 to 8588afd Compare February 26, 2026 10:14
@coderabbitai
Copy link

coderabbitai bot commented Feb 26, 2026

Walkthrough

The changes introduce a message editing capability to the agent framework. A new EditMessage variant is added to the OutboundResponse enum, with an EditedFlag tracking mechanism threaded through agent execution. The EditMessageTool implements the editing functionality, while messaging adapters handle the new variant with platform-specific behavior (Telegram executes edits with HTML fallback; other platforms log and no-op).

Changes

Cohort / File(s) Summary
Message Editing Core
src/lib.rs, src/tools.rs, src/tools/edit_message.rs
Introduces EditMessage variant to OutboundResponse, creates EditMessageTool with Arc-based EditedFlag for shared state, exports public types (EditedFlag, EditMessageArgs, EditMessageError, EditMessageOutput, EditMessageTool, new_edited_flag), and extends add_channel_tools signature to accept and thread EditedFlag through tool initialization.
Agent Execution Flow
src/agent/channel.rs
Extends run_agent_turn to return EditedFlag alongside skip and replied flags; updates all call sites and handle_agent_result to destructure and propagate edited_flag through agent turn execution and result handling paths.
Messaging Adapters
src/messaging/discord.rs, src/messaging/slack.rs, src/messaging/telegram.rs, src/messaging/twitch.rs, src/messaging/webchat.rs, src/messaging/webhook.rs
Adds EditMessage variant handling across all adapters: Discord, Slack, Twitch, WebChat, and Webhook treat it as no-op with debug logs; Telegram implements actual message editing with HTML formatting and plain-text fallback retry logic.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 3
✅ Passed checks (3 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately and concisely describes the main feature being added: Telegram message editing support.
Description check ✅ Passed The description is comprehensive and directly related to the changeset, clearly explaining the purpose, key components, and implementation details.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
src/agent/channel.rs (1)

795-806: ⚠️ Potential issue | 🟡 Minor

This call-site region still needs rustfmt cleanup.

CI reports formatting drift in this area; run cargo fmt --all to unblock the check.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/agent/channel.rs` around lines 795 - 806, The highlighted call-site is
poorly formatted and failing CI's rustfmt check; run rustfmt (e.g., cargo fmt
--all) to reformat this block so the multi-line await call to run_agent_turn and
the subsequent handle_agent_result await follow the project's rustfmt style.
Focus on reformatting the expression that assigns (result, skip_flag,
replied_flag, edited_flag) from self.run_agent_turn(...) and the trailing
self.handle_agent_result(...) .await; ensuring proper line breaks and
indentation around run_agent_turn and handle_agent_result to match the rest of
the file.
♻️ Duplicate comments (4)
src/lib.rs (1)

389-392: ⚠️ Potential issue | 🟡 Minor

Update the EditMessage docs to match actual adapter behavior.

The current text says “sending a new message (no-op)”, but adapters here treat non-Telegram edits as no-op (message unchanged).

💡 Suggested patch
-    /// Other platforms: falls back to sending a new message (no-op).
+    /// Other platforms: no-op (message unchanged).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/lib.rs` around lines 389 - 392, The docs for the EditMessage enum variant
are inaccurate: change the comment for EditMessage to state that Telegram uses
editMessageText while non-Telegram adapters treat edits as a no-op (the original
message remains unchanged) rather than falling back to sending a new message;
update the docstring on the EditMessage variant (EditMessage) to reflect this
exact behavior and clarify that only Telegram supports editing via
editMessageText and other adapters ignore the edit.
src/messaging/telegram.rs (1)

474-494: ⚠️ Potential issue | 🟡 Minor

This block still fails rustfmt in CI.

Please run cargo fmt --all (or format this section) to fix the current formatting failure.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/messaging/telegram.rs` around lines 474 - 494, The code block for
handling OutboundResponse::EditMessage is not rustfmt-compliant; run `cargo fmt
--all` or reformat the block so method chains and arguments follow rustfmt
conventions (adjust indentation and line breaks) for
markdown_to_telegram_html(&text) and the calls to
self.bot.edit_message_text(chat_id, MessageId(msg_id),
&html).parse_mode(ParseMode::Html).send().await and the fallback
self.bot.edit_message_text(chat_id, MessageId(msg_id), &text).send().await;
ensure trailing commas/spacing match rustfmt so the file containing
OutboundResponse::EditMessage compiles cleanly with cargo fmt.
src/agent/channel.rs (1)

1127-1236: ⚠️ Potential issue | 🟠 Major

edited_flag is threaded through but never used in fallback gating.

If a turn edits a message without calling reply, this path can still emit fallback Text, which defeats edit-only behavior and can double-post.

💡 Suggested patch
                 let skipped = skip_flag.load(std::sync::atomic::Ordering::Relaxed);
                 let replied = replied_flag.load(std::sync::atomic::Ordering::Relaxed);
+                let edited = edited_flag.load(std::sync::atomic::Ordering::Relaxed);
@@
-                } else if replied {
-                    tracing::debug!(channel_id = %self.id, "channel turn replied via tool (fallback suppressed)");
+                } else if replied || edited {
+                    tracing::debug!(
+                        channel_id = %self.id,
+                        replied,
+                        edited,
+                        "channel turn handled via tool (fallback suppressed)"
+                    );
                 } else if is_retrigger {
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/agent/channel.rs` around lines 1127 - 1236, The code threads edited_flag
but never uses it to suppress fallback Text, so turns that only edit a message
can still emit a new Text (double-post); load edited_flag (e.g., let edited =
edited_flag.load(std::sync::atomic::Ordering::Relaxed)) alongside skipped and
replied and treat edited like replied for gating: where the code checks replied
(and in combined checks like skipped && is_retrigger), add edited into the
conditions that suppress sending fallback OutboundResponse::Text (i.e., skip
sending when edited is true and log appropriately), and ensure any places that
log or extract reply-from-tool consider edited to avoid emitting fallback text
for edit-only turns.
src/tools/edit_message.rs (1)

30-31: ⚠️ Potential issue | 🟠 Major

Remove or use dead fields to fix CI (conversation_logger, channel_id).

conversation_logger and channel_id are stored but never read, and CI is already failing on this. Either wire them into logging/auditing now or remove them from the struct/constructor.

Proposed fix (remove dead fields)
 use crate::conversation::ConversationLogger;
-use crate::{ChannelId, OutboundResponse};
+use crate::OutboundResponse;
@@
 pub struct EditMessageTool {
     response_tx: mpsc::Sender<OutboundResponse>,
     conversation_id: String,
-    conversation_logger: ConversationLogger,
-    channel_id: ChannelId,
     edited_flag: EditedFlag,
 }
@@
     pub fn new(
         response_tx: mpsc::Sender<OutboundResponse>,
         conversation_id: impl Into<String>,
-        conversation_logger: ConversationLogger,
-        channel_id: ChannelId,
         edited_flag: EditedFlag,
     ) -> Self {
         Self {
             response_tx,
             conversation_id: conversation_id.into(),
-            conversation_logger,
-            channel_id,
             edited_flag,
         }
     }
 }

Also update the call site in src/tools.rs Line 264-269 accordingly.

Also applies to: 36-49

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/tools/edit_message.rs` around lines 30 - 31, The struct in
edit_message.rs contains unused fields conversation_logger and channel_id;
remove these fields from the struct definition and from its
constructor/signature (e.g., EditMessage::new or the struct literal) and delete
any assignments to them so they are no longer stored, then update the call site
in src/tools.rs where the EditMessage is constructed to stop passing those two
arguments (also remove any now-unused imports or variables related to
conversation_logger/channel_id). Ensure no other code expects those fields
(adjust type usages/signatures accordingly).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/messaging/webchat.rs`:
- Around line 95-96: The match arm handling OutboundResponse::EditMessage
currently returns early without any logging; add a log entry before the return
to record unsupported edit events so they can be audited. Locate the match over
OutboundResponse (the arm matching OutboundResponse::EditMessage { .. }) in
webchat.rs and call the appropriate logger (e.g., debug!/info!/warn! or
process_logger.error depending on project conventions) with a clear message
including relevant context (such as message id, user id, or the outbound
response debug representation) and then return Ok(()) as before.

In `@src/messaging/webhook.rs`:
- Around line 195-199: The handler in src/messaging/webhook.rs currently returns
early and drops OutboundResponse::EditMessage silently; update the match arm
that covers OutboundResponse::EditMessage (alongside
Reaction/RemoveReaction/Status) to emit a debug log before returning (e.g.,
using the existing logger instance or tracing::debug!) that includes the variant
name and any relevant payload identifiers so EditMessage handling is traceable;
keep the early return behavior but ensure the log message runs prior to
returning from the match arm.

In `@src/tools.rs`:
- Line 81: The re-export line in src/tools.rs that exposes EditMessage symbols
(EditedFlag, EditMessageArgs, EditMessageError, EditMessageOutput,
EditMessageTool, new_edited_flag) is misformatted and failing rustfmt; reflow
the export into rustfmt-compliant style (or run cargo fmt --all) so the pub use
edit_message::{...} line is wrapped/commas/spacing adjusted to satisfy rustfmt
and CI.
- Around line 263-271: add_channel_tools registers a per-turn EditMessageTool
but remove_channel_tools never removes it, which can leak per-turn state; in
remove_channel_tools remove the same EditMessageTool instance from the tool
handle (use the same identifying keys used when adding —
conversation_id/channel_id/edited_flag or the tool name/type) by calling the
corresponding handle removal API (e.g. handle.remove_tool(...) or
handle.remove_tool_by_name("EditMessageTool")) and await/propagate the result so
the per-turn tool is unregistered during teardown.

In `@src/tools/edit_message.rs`:
- Around line 123-125: The code sets edited_flag via
self.edited_flag.store(true, Ordering::Relaxed) immediately after enqueuing a
channel edit, which can suppress fallbacks even if no platform edit actually
occurred; change the logic so edited_flag is only set to true after a confirmed
successful platform edit (e.g., after the channel edit call returns success or
the platform-specific edit method indicates success). Update the same pattern
around the other occurrence (the block referenced at 132-137) to check the edit
result or channel response before setting edited_flag, and for non-Telegram
adapters ensure you do not toggle edited_flag when the adapter reports no-op or
failure so fallback flows remain possible.

---

Outside diff comments:
In `@src/agent/channel.rs`:
- Around line 795-806: The highlighted call-site is poorly formatted and failing
CI's rustfmt check; run rustfmt (e.g., cargo fmt --all) to reformat this block
so the multi-line await call to run_agent_turn and the subsequent
handle_agent_result await follow the project's rustfmt style. Focus on
reformatting the expression that assigns (result, skip_flag, replied_flag,
edited_flag) from self.run_agent_turn(...) and the trailing
self.handle_agent_result(...) .await; ensuring proper line breaks and
indentation around run_agent_turn and handle_agent_result to match the rest of
the file.

---

Duplicate comments:
In `@src/agent/channel.rs`:
- Around line 1127-1236: The code threads edited_flag but never uses it to
suppress fallback Text, so turns that only edit a message can still emit a new
Text (double-post); load edited_flag (e.g., let edited =
edited_flag.load(std::sync::atomic::Ordering::Relaxed)) alongside skipped and
replied and treat edited like replied for gating: where the code checks replied
(and in combined checks like skipped && is_retrigger), add edited into the
conditions that suppress sending fallback OutboundResponse::Text (i.e., skip
sending when edited is true and log appropriately), and ensure any places that
log or extract reply-from-tool consider edited to avoid emitting fallback text
for edit-only turns.

In `@src/lib.rs`:
- Around line 389-392: The docs for the EditMessage enum variant are inaccurate:
change the comment for EditMessage to state that Telegram uses editMessageText
while non-Telegram adapters treat edits as a no-op (the original message remains
unchanged) rather than falling back to sending a new message; update the
docstring on the EditMessage variant (EditMessage) to reflect this exact
behavior and clarify that only Telegram supports editing via editMessageText and
other adapters ignore the edit.

In `@src/messaging/telegram.rs`:
- Around line 474-494: The code block for handling OutboundResponse::EditMessage
is not rustfmt-compliant; run `cargo fmt --all` or reformat the block so method
chains and arguments follow rustfmt conventions (adjust indentation and line
breaks) for markdown_to_telegram_html(&text) and the calls to
self.bot.edit_message_text(chat_id, MessageId(msg_id),
&html).parse_mode(ParseMode::Html).send().await and the fallback
self.bot.edit_message_text(chat_id, MessageId(msg_id), &text).send().await;
ensure trailing commas/spacing match rustfmt so the file containing
OutboundResponse::EditMessage compiles cleanly with cargo fmt.

In `@src/tools/edit_message.rs`:
- Around line 30-31: The struct in edit_message.rs contains unused fields
conversation_logger and channel_id; remove these fields from the struct
definition and from its constructor/signature (e.g., EditMessage::new or the
struct literal) and delete any assignments to them so they are no longer stored,
then update the call site in src/tools.rs where the EditMessage is constructed
to stop passing those two arguments (also remove any now-unused imports or
variables related to conversation_logger/channel_id). Ensure no other code
expects those fields (adjust type usages/signatures accordingly).

ℹ️ Review info

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d0d58b7 and 8588afd.

📒 Files selected for processing (10)
  • src/agent/channel.rs
  • src/lib.rs
  • src/messaging/discord.rs
  • src/messaging/slack.rs
  • src/messaging/telegram.rs
  • src/messaging/twitch.rs
  • src/messaging/webchat.rs
  • src/messaging/webhook.rs
  • src/tools.rs
  • src/tools/edit_message.rs

Comment on lines +95 to +96
| OutboundResponse::Status(_)
| OutboundResponse::EditMessage { .. } => return Ok(()),
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Log unsupported EditMessage events before returning.

This branch no-ops correctly, but without a log it’s hard to audit edit tool activity on webchat.

💡 Suggested patch
-            | OutboundResponse::Status(_)
-            | OutboundResponse::EditMessage { .. } => return Ok(()),
+            | OutboundResponse::Status(_) => return Ok(()),
+            OutboundResponse::EditMessage { message_id, .. } => {
+                tracing::debug!(
+                    conversation_id = %message.conversation_id,
+                    %message_id,
+                    "EditMessage not supported on webchat, ignoring"
+                );
+                return Ok(());
+            }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
| OutboundResponse::Status(_)
| OutboundResponse::EditMessage { .. } => return Ok(()),
| OutboundResponse::Status(_) => return Ok(()),
OutboundResponse::EditMessage { message_id, .. } => {
tracing::debug!(
conversation_id = %message.conversation_id,
%message_id,
"EditMessage not supported on webchat, ignoring"
);
return Ok(());
}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/messaging/webchat.rs` around lines 95 - 96, The match arm handling
OutboundResponse::EditMessage currently returns early without any logging; add a
log entry before the return to record unsupported edit events so they can be
audited. Locate the match over OutboundResponse (the arm matching
OutboundResponse::EditMessage { .. }) in webchat.rs and call the appropriate
logger (e.g., debug!/info!/warn! or process_logger.error depending on project
conventions) with a clear message including relevant context (such as message
id, user id, or the outbound response debug representation) and then return
Ok(()) as before.

Comment on lines 195 to +199
// Reactions, status updates, and remove-reaction aren't meaningful over webhook
OutboundResponse::Reaction(_)
| OutboundResponse::RemoveReaction(_)
| OutboundResponse::Status(_) => return Ok(()),
| OutboundResponse::Status(_)
| OutboundResponse::EditMessage { .. } => return Ok(()),
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add logging for ignored EditMessage webhook responses.

EditMessage is currently dropped silently; adding a debug log will preserve traceability of tool behavior.

💡 Suggested patch
-            OutboundResponse::Reaction(_)
-            | OutboundResponse::RemoveReaction(_)
-            | OutboundResponse::Status(_)
-            | OutboundResponse::EditMessage { .. } => return Ok(()),
+            OutboundResponse::Reaction(_)
+            | OutboundResponse::RemoveReaction(_)
+            | OutboundResponse::Status(_) => return Ok(()),
+            OutboundResponse::EditMessage { message_id, .. } => {
+                tracing::debug!(
+                    conversation_id = %message.conversation_id,
+                    %message_id,
+                    "EditMessage not supported over webhook, ignoring"
+                );
+                return Ok(());
+            }
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/messaging/webhook.rs` around lines 195 - 199, The handler in
src/messaging/webhook.rs currently returns early and drops
OutboundResponse::EditMessage silently; update the match arm that covers
OutboundResponse::EditMessage (alongside Reaction/RemoveReaction/Status) to emit
a debug log before returning (e.g., using the existing logger instance or
tracing::debug!) that includes the variant name and any relevant payload
identifiers so EditMessage handling is traceable; keep the early return behavior
but ensure the log message runs prior to returning from the match arm.

pub use react::{ReactArgs, ReactError, ReactOutput, ReactTool};
pub use read_skill::{ReadSkillArgs, ReadSkillError, ReadSkillOutput, ReadSkillTool};
pub use reply::{RepliedFlag, ReplyArgs, ReplyError, ReplyOutput, ReplyTool, new_replied_flag};
pub use edit_message::{EditedFlag, EditMessageArgs, EditMessageError, EditMessageOutput, EditMessageTool, new_edited_flag};
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Rustfmt CI is failing on the new re-export formatting.

Line 81 is currently formatted in a way that fails the formatting check. Please run cargo fmt --all (or reflow this export) to unblock CI.

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/tools.rs` at line 81, The re-export line in src/tools.rs that exposes
EditMessage symbols (EditedFlag, EditMessageArgs, EditMessageError,
EditMessageOutput, EditMessageTool, new_edited_flag) is misformatted and failing
rustfmt; reflow the export into rustfmt-compliant style (or run cargo fmt --all)
so the pub use edit_message::{...} line is wrapped/commas/spacing adjusted to
satisfy rustfmt and CI.

Comment on lines +263 to +271
handle
.add_tool(EditMessageTool::new(
response_tx.clone(),
conversation_id.clone(),
state.conversation_logger.clone(),
state.channel_id.clone(),
edited_flag.clone(),
))
.await?;
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Per-turn EditMessageTool is registered but not removed at turn teardown.

Line 263 adds EditMessageTool in add_channel_tools(), but remove_channel_tools() never removes it. This can leak per-turn state and may fail later re-registration depending on ToolServer semantics.

Proposed fix
 pub async fn remove_channel_tools(
     handle: &ToolServerHandle,
 ) -> Result<(), rig::tool::server::ToolServerError> {
     handle.remove_tool(ReplyTool::NAME).await?;
+    handle.remove_tool(EditMessageTool::NAME).await?;
     handle.remove_tool(BranchTool::NAME).await?;
     handle.remove_tool(SpawnWorkerTool::NAME).await?;
     handle.remove_tool(RouteTool::NAME).await?;
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/tools.rs` around lines 263 - 271, add_channel_tools registers a per-turn
EditMessageTool but remove_channel_tools never removes it, which can leak
per-turn state; in remove_channel_tools remove the same EditMessageTool instance
from the tool handle (use the same identifying keys used when adding —
conversation_id/channel_id/edited_flag or the tool name/type) by calling the
corresponding handle removal API (e.g. handle.remove_tool(...) or
handle.remove_tool_by_name("EditMessageTool")) and await/propagate the result so
the per-turn tool is unregistered during teardown.

Comment on lines +123 to +125
// Mark the turn as handled so handle_agent_result skips any fallback
self.edited_flag.store(true, Ordering::Relaxed);

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

edited_flag is set even when no real edit may occur.

Line 123 marks the turn as edited after channel enqueue, not after a confirmed platform edit. On non-Telegram adapters (or downstream edit failures), this can still suppress fallback flows while nothing was actually edited.

Proposed fix
-        // Mark the turn as handled so handle_agent_result skips any fallback
-        self.edited_flag.store(true, Ordering::Relaxed);
+        // Only mark edited for Telegram conversation targets.
+        // Other adapters currently no-op EditMessage.
+        let is_telegram = self.conversation_id.starts_with("telegram:");
+        if is_telegram {
+            self.edited_flag.store(true, Ordering::Relaxed);
+        }

         Ok(EditMessageOutput {
-            success: true,
+            success: is_telegram,
             conversation_id: self.conversation_id.clone(),
             message_id: args.message_id,
             content: args.content,
         })

Also applies to: 132-137

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@src/tools/edit_message.rs` around lines 123 - 125, The code sets edited_flag
via self.edited_flag.store(true, Ordering::Relaxed) immediately after enqueuing
a channel edit, which can suppress fallbacks even if no platform edit actually
occurred; change the logic so edited_flag is only set to true after a confirmed
successful platform edit (e.g., after the channel edit call returns success or
the platform-specific edit method indicates success). Update the same pattern
around the other occurrence (the block referenced at 132-137) to check the edit
result or channel response before setting edited_flag, and for non-Telegram
adapters ensure you do not toggle edited_flag when the adapter reports no-op or
failure so fallback flows remain possible.

@Marenz
Copy link
Collaborator Author

Marenz commented Feb 26, 2026

I am gonna recreate this when I have more time to do this properly

@Marenz Marenz closed this Feb 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant