mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-31 04:09:44 +02:00
105 lines
4.0 KiB
Python
105 lines
4.0 KiB
Python
"""Edit a prior user message: fork a branch at it and run the turn again from there. Split out of
|
|
Messaging so each file keeps one job; self.* resolves across the AgentManager MRO as before."""
|
|
|
|
import asyncio
|
|
from uuid import uuid4
|
|
|
|
from typeguard import typechecked
|
|
|
|
from backend.apps.agents.core.models import Message, MessageBranch
|
|
from backend.apps.agents.core.ws_manager import ws_manager
|
|
from backend.apps.agents.manager.session.session_store import snapshot_session_now
|
|
from backend.apps.agents.manager.AgentManagerProtocol import AgentManagerProtocol
|
|
|
|
|
|
class EditMessage(AgentManagerProtocol):
|
|
@typechecked
|
|
async def edit_message(self, session_id: str, message_id: str, new_content: str):
|
|
"""Edit a prior user message, creating a new branch (fork)."""
|
|
session = self.sessions.get(session_id)
|
|
if not session:
|
|
raise ValueError(f"Session {session_id} not found")
|
|
|
|
existing = self.tasks.get(session_id)
|
|
if existing and not existing.done():
|
|
existing.cancel()
|
|
try:
|
|
await existing
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
target_msg = None
|
|
for i, msg in enumerate(session.messages):
|
|
if msg.id == message_id:
|
|
target_msg = msg
|
|
break
|
|
|
|
if not target_msg or target_msg.role != "user":
|
|
raise ValueError("Can only edit user messages")
|
|
|
|
fork_point_id = message_id
|
|
fork_parent_branch = target_msg.branch_id
|
|
|
|
msg_branch = session.branches.get(target_msg.branch_id)
|
|
if msg_branch and msg_branch.fork_point_message_id:
|
|
branch_user_msgs = [
|
|
m for m in session.messages
|
|
if m.branch_id == target_msg.branch_id and m.role == "user"
|
|
]
|
|
if branch_user_msgs and branch_user_msgs[0].id == message_id:
|
|
fork_point_id = msg_branch.fork_point_message_id
|
|
fork_parent_branch = msg_branch.parent_branch_id or "main"
|
|
|
|
new_branch_id = uuid4().hex
|
|
new_branch = MessageBranch(
|
|
id=new_branch_id,
|
|
parent_branch_id=fork_parent_branch,
|
|
fork_point_message_id=fork_point_id,
|
|
)
|
|
session.branches[new_branch_id] = new_branch
|
|
session.active_branch_id = new_branch_id
|
|
session.needs_fresh_session = True
|
|
|
|
|
|
edited_msg = Message(
|
|
role="user",
|
|
content=new_content,
|
|
branch_id=new_branch_id,
|
|
parent_id=target_msg.parent_id,
|
|
images=target_msg.images,
|
|
context_paths=target_msg.context_paths,
|
|
forced_tools=target_msg.forced_tools,
|
|
attached_skills=target_msg.attached_skills,
|
|
)
|
|
session.messages.append(edited_msg)
|
|
# Same status-before-snapshot rule as send_message: a stale terminal status on disk hides a mid-turn dirty death from the crash detector.
|
|
session.status = "running"
|
|
snapshot_session_now(session)
|
|
|
|
await ws_manager.send_to_session(session_id, "agent:message", {
|
|
"session_id": session_id,
|
|
"message": edited_msg.model_dump(mode="json"),
|
|
})
|
|
await ws_manager.send_to_session(session_id, "agent:branch_created", {
|
|
"session_id": session_id,
|
|
"branch": new_branch.model_dump(mode="json"),
|
|
"active_branch_id": new_branch_id,
|
|
})
|
|
|
|
session.status = "running"
|
|
await ws_manager.send_to_session(session_id, "agent:status", {
|
|
"session_id": session_id,
|
|
"status": "running",
|
|
"session": session.model_dump(mode="json"),
|
|
})
|
|
|
|
task = asyncio.create_task(self.run_agent_loop(
|
|
session_id, new_content,
|
|
images=target_msg.images,
|
|
context_paths=target_msg.context_paths,
|
|
forced_tools=target_msg.forced_tools,
|
|
attached_skills=target_msg.attached_skills,
|
|
fork_session=True,
|
|
))
|
|
self.register_turn_task(session_id, task)
|