open-swe/agent/middleware/refresh_slack_status.py

"""Middleware that keeps Slack's assistant status current during agent work.

Slack's ``assistants.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)