mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 19:43:15 +00:00
229 lines
8.1 KiB
Python
229 lines
8.1 KiB
Python
"""Middleware that keeps Slack's assistant status current during agent work.
|
|
|
|
Slack's ``assistant.threads.setStatus`` indicator expires after two minutes
|
|
if no message is sent. This middleware refreshes the indicator while model and
|
|
tool calls are actively running, then clears it when the run exits without
|
|
relying on the model to post a final Slack reply.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any, TypeVar
|
|
|
|
from langchain.agents.middleware.types import (
|
|
AgentMiddleware,
|
|
AgentState,
|
|
ModelRequest,
|
|
ModelResponse,
|
|
)
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.config import get_config
|
|
from langgraph.prebuilt.tool_node import ToolCallRequest
|
|
from langgraph.runtime import Runtime
|
|
from langgraph.types import Command
|
|
|
|
from ..utils.slack import (
|
|
DEFAULT_ASSISTANT_STATUS,
|
|
DEFAULT_LOADING_MESSAGES,
|
|
set_slack_assistant_status,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_HEARTBEAT_INTERVAL_SECONDS = 60.0
|
|
_MAX_HEARTBEAT_SECONDS = 60 * 60
|
|
|
|
_T = TypeVar("_T")
|
|
|
|
|
|
# Tool-name -> human-readable status. Keep in sync with the tool list in
|
|
# agent/server.py, agent/reviewer.py, and the deepagents built-ins (read_file,
|
|
# write_file, edit_file, execute, glob, grep, task).
|
|
_TOOL_STATUS: dict[str, str] = {
|
|
"read_file": "reading files...",
|
|
"write_file": "editing files...",
|
|
"edit_file": "editing files...",
|
|
"execute": "running commands...",
|
|
"glob": "scanning the repo...",
|
|
"grep": "searching the codebase...",
|
|
"task": "delegating to a subagent...",
|
|
"web_search": "searching the web...",
|
|
"fetch_url": "fetching a URL...",
|
|
"http_request": "making an HTTP request...",
|
|
"request_pr_review": "requesting a PR review...",
|
|
"slack_read_thread_messages": "reading Slack history...",
|
|
"slack_thread_reply": "drafting a Slack reply...",
|
|
"linear_comment": "commenting on Linear...",
|
|
"linear_create_issue": "creating a Linear issue...",
|
|
"linear_get_issue": "checking Linear...",
|
|
"linear_get_issue_comments": "checking Linear...",
|
|
"linear_list_teams": "checking Linear...",
|
|
"linear_update_issue": "updating Linear...",
|
|
"linear_delete_issue": "updating Linear...",
|
|
"add_finding": "recording review findings...",
|
|
"update_finding": "updating review findings...",
|
|
"list_findings": "checking review findings...",
|
|
"publish_review": "publishing the review...",
|
|
}
|
|
|
|
|
|
def _slack_thread_from_config() -> tuple[str, str] | None:
|
|
config = get_config()
|
|
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
|
slack_thread = configurable.get("slack_thread") if isinstance(configurable, dict) else None
|
|
if not isinstance(slack_thread, dict):
|
|
return None
|
|
|
|
channel_id = slack_thread.get("channel_id")
|
|
thread_ts = slack_thread.get("thread_ts")
|
|
if not isinstance(channel_id, str) or not isinstance(thread_ts, str):
|
|
return None
|
|
if not channel_id or not thread_ts:
|
|
return None
|
|
return channel_id, thread_ts
|
|
|
|
|
|
def _tool_call_name(tool_call: object) -> str | None:
|
|
if isinstance(tool_call, dict):
|
|
name = tool_call.get("name")
|
|
else:
|
|
name = getattr(tool_call, "name", None)
|
|
return name if isinstance(name, str) and name else None
|
|
|
|
|
|
def _status_from_recent_tool_calls(messages: list[Any]) -> str:
|
|
"""Pick a status string based on the last assistant message's tool calls."""
|
|
for msg in reversed(messages):
|
|
tool_calls = getattr(msg, "tool_calls", None)
|
|
if not tool_calls:
|
|
continue
|
|
# Use the first tool call's name; if the agent fans out, this is fine
|
|
# as a single-line indicator.
|
|
name = _tool_call_name(tool_calls[0])
|
|
if isinstance(name, str) and name in _TOOL_STATUS:
|
|
return _TOOL_STATUS[name]
|
|
return DEFAULT_ASSISTANT_STATUS
|
|
return DEFAULT_ASSISTANT_STATUS
|
|
|
|
|
|
async def _set_status(channel_id: str, thread_ts: str, status: str) -> None:
|
|
try:
|
|
await set_slack_assistant_status(
|
|
channel_id,
|
|
thread_ts,
|
|
status=status,
|
|
loading_messages=list(DEFAULT_LOADING_MESSAGES) if status else None,
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to update Slack assistant status")
|
|
|
|
|
|
class SlackAssistantStatusMiddleware(AgentMiddleware):
|
|
"""Maintain Slack's assistant status for Slack-triggered agent runs."""
|
|
|
|
state_schema = AgentState
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
heartbeat_interval_seconds: float = _HEARTBEAT_INTERVAL_SECONDS,
|
|
max_heartbeat_seconds: float = _MAX_HEARTBEAT_SECONDS,
|
|
) -> None:
|
|
self._heartbeat_interval_seconds = heartbeat_interval_seconds
|
|
self._max_heartbeat_seconds = max_heartbeat_seconds
|
|
|
|
async def abefore_agent(
|
|
self,
|
|
state: AgentState, # noqa: ARG002
|
|
runtime: Runtime, # noqa: ARG002
|
|
) -> dict[str, Any] | None:
|
|
await self._try_set(DEFAULT_ASSISTANT_STATUS)
|
|
return None
|
|
|
|
async def aafter_agent(
|
|
self,
|
|
state: AgentState, # noqa: ARG002
|
|
runtime: Runtime, # noqa: ARG002
|
|
) -> dict[str, Any] | None:
|
|
# Slack auto-clears on bot replies. This explicit clear covers soft
|
|
# exits where the agent stops without posting a message.
|
|
await self._try_set("")
|
|
return None
|
|
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
|
) -> ModelResponse:
|
|
messages = request.state.get("messages", []) if isinstance(request.state, dict) else []
|
|
status = _status_from_recent_tool_calls(messages)
|
|
return await self._run_with_heartbeat(status, handler(request))
|
|
|
|
async def awrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
|
|
) -> ToolMessage | Command:
|
|
name = _tool_call_name(request.tool_call)
|
|
status = _TOOL_STATUS.get(name or "", DEFAULT_ASSISTANT_STATUS)
|
|
return await self._run_with_heartbeat(status, handler(request))
|
|
|
|
def wrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], ModelResponse],
|
|
) -> ModelResponse:
|
|
return handler(request)
|
|
|
|
def wrap_tool_call(
|
|
self,
|
|
request: ToolCallRequest,
|
|
handler: Callable[[ToolCallRequest], ToolMessage | Command],
|
|
) -> ToolMessage | Command:
|
|
return handler(request)
|
|
|
|
async def _try_set(self, status: str) -> None:
|
|
try:
|
|
slack_thread = _slack_thread_from_config()
|
|
if slack_thread is None:
|
|
return
|
|
channel_id, thread_ts = slack_thread
|
|
await _set_status(channel_id, thread_ts, status)
|
|
except Exception:
|
|
logger.exception("Failed to read Slack thread config")
|
|
|
|
async def _run_with_heartbeat(self, status: str, awaitable: Awaitable[_T]) -> _T:
|
|
try:
|
|
slack_thread = _slack_thread_from_config()
|
|
except Exception:
|
|
logger.exception("Failed to read Slack thread config")
|
|
return await awaitable
|
|
if slack_thread is None:
|
|
return await awaitable
|
|
|
|
channel_id, thread_ts = slack_thread
|
|
await _set_status(channel_id, thread_ts, status)
|
|
heartbeat = asyncio.create_task(self._heartbeat(channel_id, thread_ts, status))
|
|
try:
|
|
return await awaitable
|
|
finally:
|
|
heartbeat.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await heartbeat
|
|
|
|
async def _heartbeat(self, channel_id: str, thread_ts: str, status: str) -> None:
|
|
loop = asyncio.get_running_loop()
|
|
started_at = loop.time()
|
|
while True:
|
|
await asyncio.sleep(self._heartbeat_interval_seconds)
|
|
if loop.time() - started_at >= self._max_heartbeat_seconds:
|
|
logger.info(
|
|
"Stopping Slack assistant status heartbeat after %.0f seconds",
|
|
self._max_heartbeat_seconds,
|
|
)
|
|
return
|
|
await _set_status(channel_id, thread_ts, status)
|