Agent Mentions in Conversations - Implementation Plan¶

┌─────────────────────────────────────────────────────────────────┐
│                    UNIFIED AGENT CONSUMER                        │
├─────────────────────────────────────────────────────────────────┤
│  WebSocket Path: /ws/agent-chat/                                │
│                                                                  │
│  Query Parameters:                                               │
│    - corpus_id: Optional[str]   (GraphQL ID)                    │
│    - document_id: Optional[str] (GraphQL ID)                    │
│    - conversation_id: Optional[str] (GraphQL ID)                │
│    - agent_id: Optional[str]    (GraphQL ID, defaults to global)│
│                                                                  │
│  Agent Selection Logic:                                          │
│    if agent_id provided:                                         │
│        agent = AgentConfiguration.get(agent_id)                 │
│    elif document_id:                                             │
│        agent = get_default_document_agent()  # GLOBAL scope     │
│    elif corpus_id:                                               │
│        agent = get_default_corpus_agent()    # GLOBAL scope     │
│    else:                                                         │
│        raise Error("No context provided")                        │
└─────────────────────────────────────────────────────────────────┘

# Default Corpus Agent
AgentConfiguration.objects.create(
    name="Corpus Assistant",
    slug="default-corpus-agent",
    scope="GLOBAL",
    is_active=True,
    system_instructions=settings.DEFAULT_CORPUS_AGENT_INSTRUCTIONS,
    available_tools=["similarity_search", "get_document_summary"],
)

# Default Document Agent
AgentConfiguration.objects.create(
    name="Document Assistant",
    slug="default-document-agent",
    scope="GLOBAL",
    is_active=True,
    system_instructions=settings.DEFAULT_DOCUMENT_AGENT_INSTRUCTIONS,
    available_tools=["similarity_search", "load_document_text"],
)

// useAgentChat.ts - ALL shared WebSocket/chat logic
export function useAgentChat(options: {
  corpusId?: string;
  documentId?: string;
  conversationId?: string;
  agentId?: string;  // Optional - uses default if not specified
}) {
  // WebSocket connection
  // Message state management
  // Streaming logic
  // Approval flow
  // Source/timeline handling

  return {
    messages,
    sendMessage,
    isConnected,
    isStreaming,
    pendingApproval,
    sendApprovalDecision,
    // ...
  };
}

// CorpusChat.tsx becomes ~200 lines (UI only)
// ChatTray.tsx becomes ~200 lines (UI only)

from django.conf import settings
from django.db import migrations

DEFAULT_CORPUS_AGENT_INSTRUCTIONS = """
You are a helpful corpus assistant. You have access to all documents in this corpus
and can search for relevant information, summarize content, and answer questions
about the corpus contents.
"""

DEFAULT_DOCUMENT_AGENT_INSTRUCTIONS = """
You are a helpful document assistant. You have access to the full text and structure
of this document and can answer questions, extract information, and explain content.
"""

def create_default_agents(apps, schema_editor):
    AgentConfiguration = apps.get_model("agents", "AgentConfiguration")

    # Default Corpus Agent
    AgentConfiguration.objects.get_or_create(
        slug="default-corpus-agent",
        defaults={
            "name": "Corpus Assistant",
            "scope": "GLOBAL",
            "is_active": True,
            "system_instructions": DEFAULT_CORPUS_AGENT_INSTRUCTIONS,
            "available_tools": ["similarity_search", "get_document_summary"],
            "description": "Default assistant for corpus-level conversations",
        }
    )

    # Default Document Agent
    AgentConfiguration.objects.get_or_create(
        slug="default-document-agent",
        defaults={
            "name": "Document Assistant",
            "scope": "GLOBAL",
            "is_active": True,
            "system_instructions": DEFAULT_DOCUMENT_AGENT_INSTRUCTIONS,
            "available_tools": ["similarity_search", "load_document_text"],
            "description": "Default assistant for document-level conversations",
        }
    )

def remove_default_agents(apps, schema_editor):
    AgentConfiguration = apps.get_model("agents", "AgentConfiguration")
    AgentConfiguration.objects.filter(
        slug__in=["default-corpus-agent", "default-document-agent"]
    ).delete()

class Migration(migrations.Migration):
    dependencies = [
        ("agents", "XXXX_previous_migration"),  # Update with actual
    ]
    operations = [
        migrations.RunPython(create_default_agents, remove_default_agents),
    ]

class AgentConfiguration(BaseOCModel):
    # Existing fields...

    slug = models.CharField(
        max_length=128,
        db_index=True,
        unique=True,  # Global uniqueness for now
        null=True,
        blank=True,
        help_text="URL-friendly identifier for mentions (e.g., 'research-assistant')"
    )

"""
Unified WebSocket consumer for all agent conversations.

Replaced the legacy corpus_conversation.py, document_conversation.py,
and standalone_document_conversation.py (all removed).
"""
import json
import logging
from typing import Optional
from channels.generic.websocket import AsyncWebsocketConsumer
from django.conf import settings

logger = logging.getLogger(__name__)


class UnifiedAgentConsumer(AsyncWebsocketConsumer):
    """
    Single WebSocket consumer for all agent chat contexts.

    Query Parameters:
        corpus_id: Optional GraphQL ID for corpus context
        document_id: Optional GraphQL ID for document context
        conversation_id: Optional GraphQL ID for existing conversation
        agent_id: Optional GraphQL ID for specific agent (uses default if omitted)

    Agent Selection:
        1. If agent_id provided → use that agent
        2. If document_id provided → use default-document-agent
        3. If corpus_id provided → use default-corpus-agent
        4. Otherwise → error
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.corpus_id: Optional[str] = None
        self.document_id: Optional[str] = None
        self.conversation_id: Optional[str] = None
        self.agent_id: Optional[str] = None
        self.agent = None
        self.room_group_name: Optional[str] = None

    async def connect(self):
        """Authenticate, extract context, and initialize agent."""
        # Extract query parameters
        query_string = self.scope.get("query_string", b"").decode()
        params = dict(p.split("=") for p in query_string.split("&") if "=" in p)

        self.corpus_id = params.get("corpus_id")
        self.document_id = params.get("document_id")
        self.conversation_id = params.get("conversation_id")
        self.agent_id = params.get("agent_id")

        # Authenticate
        user = self.scope.get("user")
        if not user or not user.is_authenticated:
            await self.close(code=4001)
            return

        # Validate context and permissions
        if not await self._validate_context(user):
            await self.close(code=4003)
            return

        # Resolve agent
        self.agent = await self._resolve_agent()
        if not self.agent:
            await self.close(code=4004)
            return

        # Set up room group
        self.room_group_name = self._get_room_group_name()
        await self.channel_layer.group_add(self.room_group_name, self.channel_name)

        await self.accept()
        logger.info(f"Agent consumer connected: {self.room_group_name}")

    async def disconnect(self, close_code):
        if self.room_group_name:
            await self.channel_layer.group_discard(
                self.room_group_name, self.channel_name
            )

    async def receive(self, text_data):
        """Handle incoming chat messages."""
        data = json.loads(text_data)
        message_type = data.get("type")

        if message_type == "chat_message":
            await self._handle_chat_message(data)
        elif message_type == "approval":
            await self._handle_approval(data)
        elif message_type == "cancel":
            await self._handle_cancel(data)

    async def _validate_context(self, user) -> bool:
        """Validate user has access to the specified context."""
        from asgiref.sync import sync_to_async

        if self.corpus_id:
            from opencontractserver.corpuses.models import Corpus
            try:
                await sync_to_async(
                    Corpus.objects.visible_to_user(user).get
                )(pk=self._decode_id(self.corpus_id))
            except Corpus.DoesNotExist:
                return False

        if self.document_id:
            from opencontractserver.documents.models import Document
            try:
                await sync_to_async(
                    Document.objects.visible_to_user(user).get
                )(pk=self._decode_id(self.document_id))
            except Document.DoesNotExist:
                return False

        return True

    async def _resolve_agent(self):
        """Resolve which agent to use based on context."""
        from asgiref.sync import sync_to_async
        from opencontractserver.agents.models import AgentConfiguration

        # Priority 1: Explicit agent_id
        if self.agent_id:
            try:
                return await sync_to_async(
                    AgentConfiguration.objects.get
                )(pk=self._decode_id(self.agent_id), is_active=True)
            except AgentConfiguration.DoesNotExist:
                logger.error(f"Specified agent not found: {self.agent_id}")
                return None

        # Priority 2: Document context → default document agent
        if self.document_id:
            try:
                return await sync_to_async(
                    AgentConfiguration.objects.get
                )(slug="default-document-agent", is_active=True)
            except AgentConfiguration.DoesNotExist:
                logger.error("Default document agent not found")
                return None

        # Priority 3: Corpus context → default corpus agent
        if self.corpus_id:
            try:
                return await sync_to_async(
                    AgentConfiguration.objects.get
                )(slug="default-corpus-agent", is_active=True)
            except AgentConfiguration.DoesNotExist:
                logger.error("Default corpus agent not found")
                return None

        logger.error("No context provided for agent resolution")
        return None

    def _get_room_group_name(self) -> str:
        """Generate unique room group name based on context."""
        if self.conversation_id:
            return f"conversation_{self.conversation_id}"
        if self.document_id:
            return f"document_{self.document_id}_{self.scope['user'].id}"
        if self.corpus_id:
            return f"corpus_{self.corpus_id}_{self.scope['user'].id}"
        return f"agent_{self.agent.id}_{self.scope['user'].id}"

    def _decode_id(self, global_id: str) -> int:
        """Decode GraphQL global ID to database PK."""
        from graphene.relay import Node
        try:
            _, pk = Node.from_global_id(global_id)
            return int(pk)
        except:
            return int(global_id)  # Fallback to raw ID

    async def _handle_chat_message(self, data):
        """Process chat message and stream agent response."""
        from opencontractserver.llms import agents as agent_factory
        from asgiref.sync import sync_to_async

        message = data.get("message", "")
        user = self.scope["user"]

        # Build agent based on context
        if self.document_id:
            agent = await sync_to_async(agent_factory.for_document_sync)(
                document_pk=self._decode_id(self.document_id),
                user_id=user.id,
                corpus_pk=self._decode_id(self.corpus_id) if self.corpus_id else None,
                system_instructions=self.agent.system_instructions,
                tools=self.agent.available_tools,
            )
        elif self.corpus_id:
            agent = await sync_to_async(agent_factory.for_corpus_sync)(
                corpus=self._decode_id(self.corpus_id),
                user_id=user.id,
                system_instructions=self.agent.system_instructions,
                tools=self.agent.available_tools,
            )
        else:
            await self.send(json.dumps({
                "type": "ERROR",
                "error": "No valid context for agent",
            }))
            return

        # Send start event
        await self.send(json.dumps({
            "type": "ASYNC_START",
            "agent_name": self.agent.name,
            "agent_id": str(self.agent.pk),
        }))

        # Stream response
        try:
            async for event in agent.stream(message):
                if hasattr(event, 'content') and event.content:
                    await self.send(json.dumps({
                        "type": "ASYNC_CONTENT",
                        "content": event.content,
                    }))

                if hasattr(event, 'tool_call'):
                    await self.send(json.dumps({
                        "type": "TOOL_CALL",
                        "tool": event.tool_call.name,
                        "args": event.tool_call.args,
                    }))

                if hasattr(event, 'approval_needed') and event.approval_needed:
                    await self.send(json.dumps({
                        "type": "APPROVAL_NEEDED",
                        "tool": event.tool_name,
                        "args": event.tool_args,
                    }))

            # Send completion
            await self.send(json.dumps({
                "type": "ASYNC_FINISH",
            }))

        except Exception as e:
            logger.exception(f"Agent response error: {e}")
            await self.send(json.dumps({
                "type": "ERROR",
                "error": str(e),
            }))

    async def _handle_approval(self, data):
        """Handle tool approval decisions."""
        # Implementation depends on agent framework approval handling
        pass

    async def _handle_cancel(self, data):
        """Handle cancellation requests."""
        # Implementation for cancelling ongoing generation
        pass

# Replace three separate routes with one unified route
from config.websocket.consumers.unified_agent_conversation import UnifiedAgentConsumer

websocket_urlpatterns = [
    # NEW: Unified agent consumer
    re_path(r"ws/agent-chat/$", UnifiedAgentConsumer.as_asgi()),

    # DEPRECATED: Keep old routes for backwards compatibility during transition
    # These can be removed after frontend migration
    re_path(r"ws/corpus-chat/(?P<corpus_id>\d+)/$", UnifiedAgentConsumer.as_asgi()),
    re_path(r"ws/document-chat/(?P<document_id>\d+)/$", UnifiedAgentConsumer.as_asgi()),

    # ... other routes (threads, etc.)
]

/**
 * Unified hook for agent chat WebSocket connections.
 *
 * Replaces duplicated logic in:
 * - CorpusChat.tsx (corpus-level chat)
 * - ChatTray.tsx (document-level chat)
 *
 * Both components now become thin UI wrappers around this hook.
 */
import { useState, useEffect, useRef, useCallback } from "react";

export interface AgentChatOptions {
  corpusId?: string;
  documentId?: string;
  conversationId?: string;
  agentId?: string;  // Optional - uses default if not specified
}

export interface AgentMessage {
  id: string;
  content: string;
  role: "user" | "assistant";
  timestamp: Date;
  agentName?: string;
  agentId?: string;
  sources?: AgentSource[];
  isStreaming?: boolean;
}

export interface AgentSource {
  annotationId: string;
  documentId: string;
  text: string;
  relevance: number;
}

export interface ApprovalRequest {
  tool: string;
  args: Record<string, unknown>;
}

export function useAgentChat(options: AgentChatOptions) {
  const { corpusId, documentId, conversationId, agentId } = options;

  // Connection state
  const wsRef = useRef<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [connectionError, setConnectionError] = useState<string | null>(null);

  // Message state
  const [messages, setMessages] = useState<AgentMessage[]>([]);
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingContent, setStreamingContent] = useState("");

  // Approval flow
  const [pendingApproval, setPendingApproval] = useState<ApprovalRequest | null>(null);

  // Sources/timeline
  const [sources, setSources] = useState<AgentSource[]>([]);

  // Build WebSocket URL
  const wsUrl = useCallback(() => {
    const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
    const params = new URLSearchParams();
    if (corpusId) params.set("corpus_id", corpusId);
    if (documentId) params.set("document_id", documentId);
    if (conversationId) params.set("conversation_id", conversationId);
    if (agentId) params.set("agent_id", agentId);
    return `${protocol}//${window.location.host}/ws/agent-chat/?${params}`;
  }, [corpusId, documentId, conversationId, agentId]);

  // Connect to WebSocket
  useEffect(() => {
    // Need at least corpus or document context
    if (!corpusId && !documentId) return;

    const ws = new WebSocket(wsUrl());
    wsRef.current = ws;

    ws.onopen = () => {
      setIsConnected(true);
      setConnectionError(null);
    };

    ws.onclose = (event) => {
      setIsConnected(false);
      if (event.code !== 1000) {
        setConnectionError(`Connection closed: ${event.code}`);
      }
    };

    ws.onerror = () => {
      setConnectionError("WebSocket error");
    };

    ws.onmessage = (event) => {
      const data = JSON.parse(event.data);
      handleMessage(data);
    };

    return () => {
      ws.close();
    };
  }, [corpusId, documentId, conversationId, agentId, wsUrl]);

  // Handle incoming WebSocket messages
  const handleMessage = useCallback((data: any) => {
    switch (data.type) {
      case "ASYNC_START":
        setIsStreaming(true);
        setStreamingContent("");
        break;

      case "ASYNC_CONTENT":
        setStreamingContent(prev => prev + data.content);
        break;

      case "ASYNC_FINISH":
        setIsStreaming(false);
        // Add completed message
        setMessages(prev => [...prev, {
          id: Date.now().toString(),
          content: streamingContent,
          role: "assistant",
          timestamp: new Date(),
          agentName: data.agent_name,
        }]);
        setStreamingContent("");
        break;

      case "APPROVAL_NEEDED":
        setPendingApproval({
          tool: data.tool,
          args: data.args,
        });
        break;

      case "SOURCES":
        setSources(data.sources);
        break;

      case "ERROR":
        setIsStreaming(false);
        setConnectionError(data.error);
        break;
    }
  }, [streamingContent]);

  // Send user message
  const sendMessage = useCallback((content: string) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setConnectionError("Not connected");
      return;
    }

    // Add user message to state
    setMessages(prev => [...prev, {
      id: Date.now().toString(),
      content,
      role: "user",
      timestamp: new Date(),
    }]);

    // Send to WebSocket
    wsRef.current.send(JSON.stringify({
      type: "chat_message",
      message: content,
    }));
  }, []);

  // Send approval decision
  const sendApprovalDecision = useCallback((approved: boolean) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;

    wsRef.current.send(JSON.stringify({
      type: "approval",
      approved,
    }));
    setPendingApproval(null);
  }, []);

  // Cancel current generation
  const cancelGeneration = useCallback(() => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) return;

    wsRef.current.send(JSON.stringify({
      type: "cancel",
    }));
    setIsStreaming(false);
  }, []);

  // Clear messages
  const clearMessages = useCallback(() => {
    setMessages([]);
    setSources([]);
  }, []);

  return {
    // Connection
    isConnected,
    connectionError,

    // Messages
    messages,
    isStreaming,
    streamingContent,

    // Approval
    pendingApproval,
    sendApprovalDecision,

    // Sources
    sources,

    // Actions
    sendMessage,
    cancelGeneration,
    clearMessages,
  };
}

/**
 * Corpus-level chat component.
 *
 * Uses useAgentChat hook for all WebSocket/message logic.
 * This component only handles corpus-specific UI.
 */
import { useAgentChat } from "../../hooks/useAgentChat";
import { ChatUI } from "../shared/ChatUI";

interface CorpusChatProps {
  corpusId: string;
}

export function CorpusChat({ corpusId }: CorpusChatProps) {
  const chat = useAgentChat({ corpusId });

  return (
    <ChatUI
      messages={chat.messages}
      isStreaming={chat.isStreaming}
      streamingContent={chat.streamingContent}
      sources={chat.sources}
      pendingApproval={chat.pendingApproval}
      isConnected={chat.isConnected}
      connectionError={chat.connectionError}
      onSendMessage={chat.sendMessage}
      onApprovalDecision={chat.sendApprovalDecision}
      onCancel={chat.cancelGeneration}
      onClear={chat.clearMessages}
      placeholder="Ask about this corpus..."
    />
  );
}

/**
 * Document-level chat component (in right tray).
 *
 * Uses useAgentChat hook for all WebSocket/message logic.
 * This component only handles document-specific UI.
 */
import { useAgentChat } from "../../../../hooks/useAgentChat";
import { ChatUI } from "../../../shared/ChatUI";

interface ChatTrayProps {
  documentId: string;
  corpusId?: string;  // Optional corpus context
}

export function ChatTray({ documentId, corpusId }: ChatTrayProps) {
  const chat = useAgentChat({ documentId, corpusId });

  return (
    <ChatUI
      messages={chat.messages}
      isStreaming={chat.isStreaming}
      streamingContent={chat.streamingContent}
      sources={chat.sources}
      pendingApproval={chat.pendingApproval}
      isConnected={chat.isConnected}
      connectionError={chat.connectionError}
      onSendMessage={chat.sendMessage}
      onApprovalDecision={chat.sendApprovalDecision}
      onCancel={chat.cancelGeneration}
      onClear={chat.clearMessages}
      placeholder="Ask about this document..."
    />
  );
}

/**
 * Shared chat UI component used by both CorpusChat and ChatTray.
 */
import React, { useState, useRef, useEffect } from "react";
import styled from "styled-components";
import { AgentMessage, AgentSource, ApprovalRequest } from "../../hooks/useAgentChat";
import { MessageComposer } from "../threads/MessageComposer";
import { ApprovalDialog } from "./ApprovalDialog";
import { SourcesPanel } from "./SourcesPanel";
import { StreamingMessage } from "./StreamingMessage";

interface ChatUIProps {
  messages: AgentMessage[];
  isStreaming: boolean;
  streamingContent: string;
  sources: AgentSource[];
  pendingApproval: ApprovalRequest | null;
  isConnected: boolean;
  connectionError: string | null;
  onSendMessage: (content: string) => void;
  onApprovalDecision: (approved: boolean) => void;
  onCancel: () => void;
  onClear: () => void;
  placeholder?: string;
}

export function ChatUI({
  messages,
  isStreaming,
  streamingContent,
  sources,
  pendingApproval,
  isConnected,
  connectionError,
  onSendMessage,
  onApprovalDecision,
  onCancel,
  onClear,
  placeholder,
}: ChatUIProps) {
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // Auto-scroll on new messages
  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [messages, streamingContent]);

  return (
    <Container>
      {connectionError && (
        <ErrorBanner>{connectionError}</ErrorBanner>
      )}

      <MessagesContainer>
        {messages.map(message => (
          <MessageBubble key={message.id} role={message.role}>
            {message.content}
          </MessageBubble>
        ))}

        {isStreaming && (
          <StreamingMessage content={streamingContent} onCancel={onCancel} />
        )}

        <div ref={messagesEndRef} />
      </MessagesContainer>

      {sources.length > 0 && (
        <SourcesPanel sources={sources} />
      )}

      {pendingApproval && (
        <ApprovalDialog
          tool={pendingApproval.tool}
          args={pendingApproval.args}
          onDecision={onApprovalDecision}
        />
      )}

      <MessageComposer
        onSend={onSendMessage}
        disabled={!isConnected || isStreaming}
        placeholder={placeholder}
      />
    </Container>
  );
}

// Styled components...
const Container = styled.div`/* ... */`;
const MessagesContainer = styled.div`/* ... */`;
const MessageBubble = styled.div<{ role: string }>`/* ... */`;
const ErrorBanner = styled.div`/* ... */`;

class AgentConfiguration(BaseOCModel):
    # Existing fields...

    slug = models.CharField(
        max_length=128,
        db_index=True,
        null=True,
        blank=True,
        help_text="URL-friendly identifier for mentions (e.g., 'research-assistant')"
    )

    class Meta:
        constraints = [
            # Existing constraints...
            models.UniqueConstraint(
                fields=["corpus", "slug"],
                name="unique_agent_slug_per_corpus",
                condition=Q(scope="CORPUS"),
            ),
            models.UniqueConstraint(
                fields=["slug"],
                name="unique_global_agent_slug",
                condition=Q(scope="GLOBAL"),
            ),
        ]

def parse_mentions_from_content(markdown_content: str) -> dict[str, set[str]]:
    mentioned = {
        "users": set(),
        "documents": set(),
        "annotations": set(),
        "corpuses": set(),
        "agents": set(),  # NEW
    }

    # ... existing parsing ...

    # Agent: /agents/{agentSlug} or /c/{corpusIdent}/agents/{agentSlug}
    if path.startswith("/agents/"):
        parts = path.split("/")
        if len(parts) >= 3:
            agent_slug = parts[2]
            mentioned["agents"].add(agent_slug)
            logger.debug(f"Found agent mention: {agent_slug}")

class ChatMessage(BaseOCModel, HasEmbeddingMixin):
    # Existing fields...

    mentioned_agents = models.ManyToManyField(
        "agents.AgentConfiguration",
        related_name="mentioned_in_messages",
        blank=True,
        help_text="Agents mentioned in this message that should respond",
    )

@transaction.atomic
def link_message_to_resources(
    chat_message, mentioned_ids: dict[str, set[str]]
) -> dict[str, int]:
    # ... existing logic ...

    result = {
        "documents_linked": 0,
        "annotations_linked": 0,
        "users_mentioned": len(mentioned_ids.get("users", set())),
        "corpuses_mentioned": len(mentioned_ids.get("corpuses", set())),
        "agents_linked": 0,  # NEW
    }

    # Link agents (ManyToMany)
    if mentioned_ids.get("agents"):
        from opencontractserver.agents.models import AgentConfiguration

        # Get conversation's corpus for scoped agent lookup
        corpus = chat_message.conversation.chat_with_corpus

        # Find agents by slug - global OR matching corpus
        agent_slugs = list(mentioned_ids["agents"])
        if corpus:
            agents_qs = AgentConfiguration.objects.filter(
                Q(slug__in=agent_slugs, scope="GLOBAL", is_active=True) |
                Q(slug__in=agent_slugs, scope="CORPUS", corpus=corpus, is_active=True)
            )
        else:
            agents_qs = AgentConfiguration.objects.filter(
                slug__in=agent_slugs, scope="GLOBAL", is_active=True
            )

        chat_message.mentioned_agents.set(agents_qs)
        result["agents_linked"] = agents_qs.count()

    return result

# Add to Query class
search_agents_for_mention = DjangoConnectionField(
    AgentConfigurationType,
    text_search=graphene.String(
        description="Search query to find agents by name or description"
    ),
    corpus_id=graphene.ID(
        description="Corpus ID to scope agent search (includes global + corpus agents)"
    ),
)

@graphql_ratelimit_dynamic(get_rate=get_user_tier_rate("READ_LIGHT"))
def resolve_search_agents_for_mention(
    self, info, text_search=None, corpus_id=None, **kwargs
):
    """
    Search agents for @ mention autocomplete.

    Returns:
    - All active global agents
    - Corpus-specific agents for the provided corpus (if user has access)

    SECURITY: Filters by visibility - users only see agents they can mention.
    """
    from opencontractserver.agents.models import AgentConfiguration

    user = info.context.user

    # Anonymous users cannot mention
    if user.is_anonymous:
        return AgentConfiguration.objects.none()

    # Build base queryset using visible_to_user
    qs = AgentConfiguration.objects.visible_to_user(user).filter(is_active=True)

    # If corpus_id provided, further filter to global + that corpus
    if corpus_id:
        _, corpus_pk = from_global_id(corpus_id)
        qs = qs.filter(
            Q(scope="GLOBAL") | Q(scope="CORPUS", corpus_id=corpus_pk)
        )

    # Apply text search
    if text_search:
        qs = qs.filter(
            Q(name__icontains=text_search) |
            Q(description__icontains=text_search) |
            Q(slug__icontains=text_search)
        )

    return qs.order_by("scope", "name")  # Global first, then corpus, alphabetical

class AgentConfigurationType(AnnotatePermissionsForReadMixin, DjangoObjectType):
    # Existing fields...

    mention_format = graphene.String(
        description="The @ mention format for this agent"
    )

    def resolve_mention_format(self, info):
        return f"@agent:{self.slug}"

"""
Celery tasks for agent response generation in threads.
"""
import logging
from celery import shared_task
from channels.layers import get_channel_layer
from asgiref.sync import async_to_sync

from opencontractserver.agents.models import AgentConfiguration
from opencontractserver.conversations.models import ChatMessage, Conversation

logger = logging.getLogger(__name__)


@shared_task(bind=True, max_retries=3)
def generate_agent_response(
    self,
    message_id: int,
    agent_config_id: int,
    user_id: int,
):
    """
    Generate an agent response to a message that mentioned the agent.

    This task:
    1. Creates a placeholder LLM message in 'in_progress' state
    2. Loads the agent configuration and builds the agent
    3. Gathers thread context (previous messages)
    4. Generates response using the agent
    5. Updates the message with final content
    6. Broadcasts updates via WebSocket
    """
    from django.contrib.auth import get_user_model
    from opencontractserver.llms import agents as agent_factory

    User = get_user_model()

    try:
        user = User.objects.get(pk=user_id)
        source_message = ChatMessage.objects.get(pk=message_id)
        agent_config = AgentConfiguration.objects.get(pk=agent_config_id)
        conversation = source_message.conversation
        corpus = conversation.chat_with_corpus

        if not corpus:
            logger.error(f"Cannot generate agent response: no corpus for conversation {conversation.pk}")
            return

        # Create placeholder message
        llm_message = ChatMessage.objects.create(
            conversation=conversation,
            msg_type="LLM",
            content="",
            state="in_progress",
            agent_configuration=agent_config,
            parent_message=source_message,  # Thread as reply
            creator=user,
        )

        # Broadcast start event
        channel_layer = get_channel_layer()
        async_to_sync(channel_layer.group_send)(
            f"thread_{conversation.pk}",
            {
                "type": "agent.response.start",
                "message_id": llm_message.pk,
                "agent_id": agent_config.pk,
                "agent_name": agent_config.name,
            }
        )

        # Build agent with custom system instructions
        agent = agent_factory.for_corpus_sync(
            corpus=corpus.pk,
            user_id=user_id,
            system_instructions=agent_config.system_instructions,
            tools=agent_config.available_tools,
        )

        # Build prompt with thread context
        thread_context = _build_thread_context(conversation, source_message)
        prompt = f"{thread_context}\n\nUser: {source_message.content}"

        # Generate response (streaming to WebSocket)
        full_response = ""
        sources = []

        for event in agent.stream_sync(prompt):
            if hasattr(event, 'content') and event.content:
                full_response += event.content
                # Broadcast content chunk
                async_to_sync(channel_layer.group_send)(
                    f"thread_{conversation.pk}",
                    {
                        "type": "agent.response.chunk",
                        "message_id": llm_message.pk,
                        "content": event.content,
                    }
                )

            if hasattr(event, 'sources') and event.sources:
                sources.extend(event.sources)

        # Update message with final content
        llm_message.content = full_response
        llm_message.state = "completed"
        llm_message.data = {
            "sources": [s.to_dict() for s in sources],
        }
        llm_message.save()

        # Link source annotations
        if sources:
            from opencontractserver.annotations.models import Annotation
            annotation_ids = [s.annotation_id for s in sources if s.annotation_id]
            llm_message.source_annotations.set(
                Annotation.objects.filter(pk__in=annotation_ids)
            )

        # Broadcast completion
        async_to_sync(channel_layer.group_send)(
            f"thread_{conversation.pk}",
            {
                "type": "agent.response.complete",
                "message_id": llm_message.pk,
                "content": full_response,
                "sources": [s.to_dict() for s in sources],
            }
        )

    except Exception as e:
        logger.exception(f"Agent response generation failed: {e}")
        # Update message to error state if it was created
        if 'llm_message' in locals():
            llm_message.state = "error"
            llm_message.data = {"error": str(e)}
            llm_message.save()

            # Broadcast error
            async_to_sync(channel_layer.group_send)(
                f"thread_{conversation.pk}",
                {
                    "type": "agent.response.error",
                    "message_id": llm_message.pk,
                    "error": str(e),
                }
            )
        raise


def _build_thread_context(conversation: Conversation, up_to_message: ChatMessage) -> str:
    """Build conversation context from thread messages."""
    messages = ChatMessage.objects.filter(
        conversation=conversation,
        deleted_at__isnull=True,
    ).exclude(pk=up_to_message.pk).order_by("created_at")[:20]  # Limit context

    context_parts = []
    for msg in messages:
        role = "Assistant" if msg.msg_type == "LLM" else "User"
        context_parts.append(f"{role}: {msg.content}")

    return "\n".join(context_parts)

# In CreateThreadMessageMutation.mutate()

# After mention parsing:
if result.get("agents_linked", 0) > 0:
    from opencontractserver.tasks.agent_tasks import generate_agent_response

    # Trigger response for each mentioned agent
    for agent in chat_message.mentioned_agents.all():
        generate_agent_response.delay(
            message_id=chat_message.pk,
            agent_config_id=agent.pk,
            user_id=user.id,
        )

"""
WebSocket consumer for thread conversations with agent responses.
"""
import json
import logging
from channels.generic.websocket import AsyncWebsocketConsumer

logger = logging.getLogger(__name__)


class ThreadConversationConsumer(AsyncWebsocketConsumer):
    """
    WebSocket consumer for real-time updates in thread conversations.

    Handles:
    - New messages in thread
    - Agent response streaming
    - Typing indicators
    """

    async def connect(self):
        self.conversation_id = self.scope["url_route"]["kwargs"]["conversation_id"]
        self.room_group_name = f"thread_{self.conversation_id}"

        # Verify user has access to conversation
        if not await self._user_can_access_conversation():
            await self.close(code=4003)
            return

        # Join room group
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        await self.accept()

    async def disconnect(self, close_code):
        await self.channel_layer.group_discard(
            self.room_group_name,
            self.channel_name
        )

    async def receive(self, text_data):
        """Handle incoming messages (e.g., typing indicators)."""
        data = json.loads(text_data)
        msg_type = data.get("type")

        if msg_type == "typing":
            await self.channel_layer.group_send(
                self.room_group_name,
                {
                    "type": "user.typing",
                    "user_id": self.scope["user"].id,
                    "username": self.scope["user"].username,
                }
            )

    # Event handlers for agent responses

    async def agent_response_start(self, event):
        await self.send(json.dumps({
            "type": "AGENT_RESPONSE_START",
            "message_id": event["message_id"],
            "agent_id": event["agent_id"],
            "agent_name": event["agent_name"],
        }))

    async def agent_response_chunk(self, event):
        await self.send(json.dumps({
            "type": "AGENT_RESPONSE_CHUNK",
            "message_id": event["message_id"],
            "content": event["content"],
        }))

    async def agent_response_complete(self, event):
        await self.send(json.dumps({
            "type": "AGENT_RESPONSE_COMPLETE",
            "message_id": event["message_id"],
            "content": event["content"],
            "sources": event.get("sources", []),
        }))

    async def agent_response_error(self, event):
        await self.send(json.dumps({
            "type": "AGENT_RESPONSE_ERROR",
            "message_id": event["message_id"],
            "error": event["error"],
        }))

    async def user_typing(self, event):
        await self.send(json.dumps({
            "type": "USER_TYPING",
            "user_id": event["user_id"],
            "username": event["username"],
        }))

    async def _user_can_access_conversation(self) -> bool:
        """Check if user has access to the conversation."""
        from opencontractserver.conversations.models import Conversation
        from asgiref.sync import sync_to_async

        user = self.scope.get("user")
        if not user or not user.is_authenticated:
            return False

        try:
            conversation = await sync_to_async(
                Conversation.objects.visible_to_user(user).get
            )(pk=self.conversation_id)
            return True
        except Conversation.DoesNotExist:
            return False

# Add to websocket_urlpatterns
from config.websocket.consumers.thread_conversation import ThreadConversationConsumer

websocket_urlpatterns = [
    # Existing routes...
    re_path(
        r"ws/threads/(?P<conversation_id>\d+)/$",
        ThreadConversationConsumer.as_asgi()
    ),
]

import React, { forwardRef, useEffect, useImperativeHandle, useState } from "react";
import styled from "styled-components";
import { Bot } from "lucide-react";
import { color } from "../../theme/colors";
import { spacing } from "../../theme/spacing";

// ... styled components similar to ResourceMentionPicker ...

export interface MentionAgent {
  id: string;
  slug: string;
  name: string;
  description: string;
  scope: "GLOBAL" | "CORPUS";
  badgeConfig: {
    icon?: string;
    color?: string;
    label?: string;
  };
}

export interface AgentMentionPickerProps {
  agents: MentionAgent[];
  onSelect: (agent: MentionAgent) => void;
  selectedIndex: number;
}

export const AgentMentionPicker = forwardRef<
  AgentMentionPickerRef,
  AgentMentionPickerProps
>(({ agents, onSelect, selectedIndex }, ref) => {
  // Implementation similar to ResourceMentionPicker
  // with agent-specific display (icon from badgeConfig, scope indicator)
});

export const SEARCH_AGENTS_FOR_MENTION = gql`
  query SearchAgentsForMention($textSearch: String, $corpusId: ID) {
    searchAgentsForMention(textSearch: $textSearch, corpusId: $corpusId) {
      edges {
        node {
          id
          slug
          name
          description
          scope
          badgeConfig
          avatarUrl
          corpus {
            id
            title
          }
        }
      }
    }
  }
`;

import { useQuery } from "@apollo/client";
import { useMemo } from "react";
import { SEARCH_AGENTS_FOR_MENTION } from "../graphql/queries";

export function useAgentMentionSearch(
  searchText: string,
  corpusId: string | null
) {
  const { data, loading, error } = useQuery(SEARCH_AGENTS_FOR_MENTION, {
    variables: { textSearch: searchText, corpusId },
    skip: searchText.length < 1,
    fetchPolicy: "cache-and-network",
  });

  const agents = useMemo(() => {
    if (!data?.searchAgentsForMention?.edges) return [];
    return data.searchAgentsForMention.edges.map((edge: any) => edge.node);
  }, [data]);

  return { agents, loading, error };
}

// In the mention extension configuration
const mentionSuggestion = {
  char: "@",
  items: async ({ query }: { query: string }) => {
    // Detect mention type from query
    if (query.startsWith("agent:")) {
      const searchText = query.slice(6); // Remove "agent:" prefix
      const agents = await fetchAgentsForMention(searchText, corpusId);
      return agents.map(a => ({ ...a, mentionType: "agent" }));
    }

    if (query.startsWith("corpus:") || query.startsWith("document:")) {
      // Existing resource mention logic
      return fetchResourcesForMention(query);
    }

    // Default: user mentions
    return fetchUsersForMention(query);
  },

  render: () => ({
    // Render appropriate picker based on mentionType
  }),
};

export function MentionChip({ mention }: { mention: ParsedMention }) {
  if (mention.type === "agent") {
    return (
      <AgentChip
        href={`/agents/${mention.slug}`}
        color={mention.badgeConfig?.color || "#8B5CF6"}
      >
        <Bot size={12} />
        <span>{mention.name}</span>
      </AgentChip>
    );
  }
  // ... existing corpus/document chip rendering
}

import { useEffect, useRef, useCallback, useState } from "react";
import { useAtomValue } from "jotai";
import { currentThreadIdAtom } from "../atoms/threadAtoms";

interface AgentResponseState {
  messageId: number;
  agentId: number;
  agentName: string;
  content: string;
  isComplete: boolean;
  error?: string;
}

export function useThreadWebSocket() {
  const threadId = useAtomValue(currentThreadIdAtom);
  const wsRef = useRef<WebSocket | null>(null);
  const [agentResponses, setAgentResponses] = useState<Map<number, AgentResponseState>>(new Map());

  useEffect(() => {
    if (!threadId) return;

    const wsUrl = `${window.location.protocol === 'https:' ? 'wss:' : 'ws:'}//${window.location.host}/ws/threads/${threadId}/`;
    const ws = new WebSocket(wsUrl);
    wsRef.current = ws;

    ws.onmessage = (event) => {
      const data = JSON.parse(event.data);

      switch (data.type) {
        case "AGENT_RESPONSE_START":
          setAgentResponses(prev => new Map(prev).set(data.message_id, {
            messageId: data.message_id,
            agentId: data.agent_id,
            agentName: data.agent_name,
            content: "",
            isComplete: false,
          }));
          break;

        case "AGENT_RESPONSE_CHUNK":
          setAgentResponses(prev => {
            const updated = new Map(prev);
            const existing = updated.get(data.message_id);
            if (existing) {
              updated.set(data.message_id, {
                ...existing,
                content: existing.content + data.content,
              });
            }
            return updated;
          });
          break;

        case "AGENT_RESPONSE_COMPLETE":
          setAgentResponses(prev => {
            const updated = new Map(prev);
            const existing = updated.get(data.message_id);
            if (existing) {
              updated.set(data.message_id, {
                ...existing,
                content: data.content,
                isComplete: true,
              });
            }
            return updated;
          });
          break;

        case "AGENT_RESPONSE_ERROR":
          setAgentResponses(prev => {
            const updated = new Map(prev);
            const existing = updated.get(data.message_id);
            if (existing) {
              updated.set(data.message_id, {
                ...existing,
                isComplete: true,
                error: data.error,
              });
            }
            return updated;
          });
          break;
      }
    };

    return () => {
      ws.close();
    };
  }, [threadId]);

  return { agentResponses };
}

function ThreadDetail() {
  const { agentResponses } = useThreadWebSocket();

  // ... existing logic ...

  return (
    <div>
      {messages.map(message => (
        <MessageItem key={message.id} message={message} />
      ))}

      {/* Show streaming agent responses */}
      {Array.from(agentResponses.values())
        .filter(r => !r.isComplete)
        .map(response => (
          <StreamingAgentMessage
            key={response.messageId}
            agentName={response.agentName}
            content={response.content}
          />
        ))}
    </div>
  );
}

"""
Tests for agent mention parsing, autocomplete, and response generation.
"""

class AgentMentionParsingTests(TestCase):
    """Test agent mention parsing from markdown content."""

    def test_parse_agent_mention_from_markdown(self):
        """Test parsing @agent: mentions from markdown links."""
        pass

    def test_agent_mention_linked_to_message(self):
        """Test that mentioned agents are linked to ChatMessage."""
        pass


class AgentMentionAutocompleteTests(TestCase):
    """Test GraphQL agent search for mentions."""

    def test_search_returns_global_agents(self):
        """Global agents visible to all authenticated users."""
        pass

    def test_search_returns_corpus_agents_for_accessible_corpus(self):
        """Corpus agents only visible if user has corpus access."""
        pass

    def test_search_filters_inactive_agents(self):
        """Inactive agents not returned in search."""
        pass


class AgentResponseGenerationTests(TransactionTestCase):
    """Test agent response generation task."""

    def test_agent_response_created_for_mention(self):
        """Test that mentioning an agent triggers response generation."""
        pass

    def test_multiple_agents_respond_in_sequence(self):
        """Test that multiple mentioned agents each generate responses."""
        pass