Skip to content

Commit 6fbc61f

Browse files
tjb-techtjb-tech
committed
fix(ohmo): support channel slash commands and stabilize tests
1 parent a359e02 commit 6fbc61f

8 files changed

Lines changed: 841 additions & 174 deletions

File tree

‎ohmo/gateway/bridge.py‎

Lines changed: 124 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ def __init__(self, *, bus: MessageBus, runtime_pool: OhmoSessionRuntimePool) ->
5656
self._bus = bus
5757
self._runtime_pool = runtime_pool
5858
self._running = False
59+
self._session_tasks: dict[str, asyncio.Task[None]] = {}
60+
self._session_cancel_reasons: dict[str, str] = {}
5961

6062
async def run(self) -> None:
6163
self._running = True
@@ -76,63 +78,137 @@ async def run(self) -> None:
7678
session_key,
7779
_content_snippet(message.content),
7880
)
79-
try:
80-
reply = ""
81-
async for update in self._runtime_pool.stream_message(message, session_key):
82-
if update.kind == "final":
83-
reply = update.text
84-
continue
85-
if not update.text:
86-
continue
87-
logger.info(
88-
"ohmo outbound update channel=%s chat_id=%s session_key=%s kind=%s content=%r",
89-
message.channel,
90-
message.chat_id,
91-
session_key,
92-
update.kind,
93-
_content_snippet(update.text),
94-
)
95-
await self._bus.publish_outbound(
96-
OutboundMessage(
97-
channel=message.channel,
98-
chat_id=message.chat_id,
99-
content=update.text,
100-
metadata=update.metadata,
101-
)
102-
)
103-
except Exception as exc: # pragma: no cover - gateway failure path
104-
logger.exception(
105-
"ohmo gateway failed to process inbound message channel=%s chat_id=%s sender_id=%s session_key=%s content=%r",
106-
message.channel,
107-
message.chat_id,
108-
message.sender_id,
109-
session_key,
110-
_content_snippet(message.content),
111-
)
112-
reply = _format_gateway_error(exc)
113-
if not reply:
81+
if message.content.strip() == "/stop":
82+
await self._handle_stop(message, session_key)
83+
continue
84+
await self._interrupt_session(
85+
session_key,
86+
reason="replaced by a newer user message",
87+
notify=OutboundMessage(
88+
channel=message.channel,
89+
chat_id=message.chat_id,
90+
content="⏹️ 已停止上一条正在处理的任务，继续看你的最新消息。",
91+
metadata={"_progress": True, "_session_key": session_key},
92+
),
93+
)
94+
task = asyncio.create_task(
95+
self._process_message(message, session_key),
96+
name=f"ohmo-session:{session_key}",
97+
)
98+
self._session_tasks[session_key] = task
99+
task.add_done_callback(lambda finished, key=session_key: self._cleanup_task(key, finished))
100+
101+
def stop(self) -> None:
102+
self._running = False
103+
for session_key, task in list(self._session_tasks.items()):
104+
self._session_cancel_reasons[session_key] = "gateway stopping"
105+
task.cancel()
106+
107+
async def _handle_stop(self, message, session_key: str) -> None:
108+
stopped = await self._interrupt_session(
109+
session_key,
110+
reason="stopped by user command",
111+
)
112+
content = "⏹️ 已停止当前正在运行的任务。" if stopped else "当前没有正在运行的任务。"
113+
await self._bus.publish_outbound(
114+
OutboundMessage(
115+
channel=message.channel,
116+
chat_id=message.chat_id,
117+
content=content,
118+
metadata={"_session_key": session_key},
119+
)
120+
)
121+
122+
async def _interrupt_session(
123+
self,
124+
session_key: str,
125+
*,
126+
reason: str,
127+
notify: OutboundMessage | None = None,
128+
) -> bool:
129+
task = self._session_tasks.get(session_key)
130+
if task is None or task.done():
131+
return False
132+
self._session_cancel_reasons[session_key] = reason
133+
task.cancel()
134+
if notify is not None:
135+
await self._bus.publish_outbound(notify)
136+
try:
137+
await asyncio.wait_for(asyncio.shield(task), timeout=3.0)
138+
except (asyncio.CancelledError, asyncio.TimeoutError):
139+
pass
140+
return True
141+
142+
async def _process_message(self, message, session_key: str) -> None:
143+
try:
144+
reply = ""
145+
async for update in self._runtime_pool.stream_message(message, session_key):
146+
if update.kind == "final":
147+
reply = update.text
148+
continue
149+
if not update.text:
150+
continue
114151
logger.info(
115-
"ohmo inbound finished without final reply channel=%s chat_id=%s session_key=%s",
152+
"ohmo outbound update channel=%s chat_id=%s session_key=%s kind=%s content=%r",
116153
message.channel,
117154
message.chat_id,
118155
session_key,
156+
update.kind,
157+
_content_snippet(update.text),
119158
)
120-
continue
159+
await self._bus.publish_outbound(
160+
OutboundMessage(
161+
channel=message.channel,
162+
chat_id=message.chat_id,
163+
content=update.text,
164+
metadata=update.metadata,
165+
)
166+
)
167+
except asyncio.CancelledError:
121168
logger.info(
122-
"ohmo outbound final channel=%s chat_id=%s session_key=%s content=%r",
169+
"ohmo session interrupted channel=%s chat_id=%s session_key=%s reason=%s",
123170
message.channel,
124171
message.chat_id,
125172
session_key,
126-
_content_snippet(reply),
173+
self._session_cancel_reasons.get(session_key, "cancelled"),
127174
)
128-
await self._bus.publish_outbound(
129-
OutboundMessage(
130-
channel=message.channel,
131-
chat_id=message.chat_id,
132-
content=reply,
133-
metadata={"_session_key": session_key},
134-
)
175+
raise
176+
except Exception as exc: # pragma: no cover - gateway failure path
177+
logger.exception(
178+
"ohmo gateway failed to process inbound message channel=%s chat_id=%s sender_id=%s session_key=%s content=%r",
179+
message.channel,
180+
message.chat_id,
181+
message.sender_id,
182+
session_key,
183+
_content_snippet(message.content),
135184
)
185+
reply = _format_gateway_error(exc)
186+
if not reply:
187+
logger.info(
188+
"ohmo inbound finished without final reply channel=%s chat_id=%s session_key=%s",
189+
message.channel,
190+
message.chat_id,
191+
session_key,
192+
)
193+
return
194+
logger.info(
195+
"ohmo outbound final channel=%s chat_id=%s session_key=%s content=%r",
196+
message.channel,
197+
message.chat_id,
198+
session_key,
199+
_content_snippet(reply),
200+
)
201+
await self._bus.publish_outbound(
202+
OutboundMessage(
203+
channel=message.channel,
204+
chat_id=message.chat_id,
205+
content=reply,
206+
metadata={"_session_key": session_key},
207+
)
208+
)
136209

137-
def stop(self) -> None:
138-
self._running = False
210+
def _cleanup_task(self, session_key: str, task: asyncio.Task[None]) -> None:
211+
current = self._session_tasks.get(session_key)
212+
if current is task:
213+
self._session_tasks.pop(session_key, None)
214+
self._session_cancel_reasons.pop(session_key, None)

0 commit comments

Comments
 (0)