Cortex-Inara/cortex/routers/google_chat.py

import asyncio
import logging
from fastapi import APIRouter, HTTPException, Request, Response
from google.auth.transport import requests as google_requests
from google.oauth2 import id_token
from auth_utils import get_user_channels
from context_loader import load_context
from llm_client import complete
from persona import set_context
from session_logger import log_turn
from session_store import load as load_session, save as save_session
from config import settings

logger = logging.getLogger(__name__)
router = APIRouter()

# Workspace Add-on Chat apps: JWT is issued by accounts.google.com.
# (Legacy standalone Chat bots used chat@system.gserviceaccount.com — different format.)
_GOOGLE_ISSUER = "https://accounts.google.com"


def _msg(text: str) -> dict:
    """Wrap a text reply in the Workspace Add-on hostAppDataAction format.

    Standalone Chat apps use {"text": "..."} directly, but Workspace Add-on
    Chat apps require the hostAppDataAction wrapper for Google Chat to render
    the response as a bot message.
    """
    return {
        "hostAppDataAction": {
            "chatDataAction": {
                "createMessageAction": {
                    "message": {"text": text}
                }
            }
        }
    }


def _verify_system_id_token(token: str, audience: str) -> None:
    """Verify the systemIdToken from authorizationEventObject.

    For Workspace Add-on Chat apps Google sends the token in the request body
    at body["authorizationEventObject"]["systemIdToken"], not in the
    Authorization header.

    Claims verified:
      iss  = "https://accounts.google.com"
      aud  = the per-user audience from channels.json (the endpoint URL)
    """
    try:
        claims = id_token.verify_oauth2_token(
            token,
            google_requests.Request(),
            audience=audience,
        )
    except Exception as exc:
        logger.warning("Google Chat JWT verification failed: %s", exc)
        raise HTTPException(status_code=401, detail="Invalid token")
    if claims.get("iss") != _GOOGLE_ISSUER:
        logger.warning("Google Chat JWT wrong issuer: %s", claims.get("iss"))
        raise HTTPException(status_code=401, detail="Wrong issuer")


@router.post("/channels/google-chat/{username}")
async def receive(username: str, request: Request):
    channels = get_user_channels(username)
    cfg = channels.get("google_chat")
    if not cfg:
        logger.warning("Google Chat: no channel config for user %r", username)
        raise HTTPException(status_code=404, detail="Channel not configured for this user")

    persona_name = cfg.get("persona", "inara")
    audience     = cfg.get("audience", "")
    backend      = cfg.get("backend", settings.primary_backend)
    timeout      = cfg.get("timeout", 25)

    set_context(username, persona_name)

    body = await request.json()

    # Verify the systemIdToken embedded in the request body
    if audience:
        token = body.get("authorizationEventObject", {}).get("systemIdToken", "")
        if not token:
            logger.warning("Google Chat: missing systemIdToken for %s", username)
            raise HTTPException(status_code=401, detail="Missing token")
        _verify_system_id_token(token, audience)

    chat = body.get("chat", {})

    # Event type is inferred from which payload key is present — there is no
    # top-level "type" field in the Workspace Add-on event format.
    if "addedToSpacePayload" in chat:
        space_type = chat["addedToSpacePayload"].get("space", {}).get("type", "")
        if space_type == "DM":
            return _msg(f"✨ Hello! I'm {persona_name.capitalize()}. What can I help you with?")
        return _msg(f"✨ Hello! I'm {persona_name.capitalize()}. Send me a message and I'll do my best to help.")

    if "removedFromSpacePayload" in chat:
        return Response(status_code=200)

    if "messagePayload" not in chat:
        logger.info("Google Chat: unhandled event keys: %s", list(chat.keys()))
        return Response(status_code=200)

    payload        = chat["messagePayload"]
    message        = payload.get("message", {})
    space          = payload.get("space", {})
    user           = chat.get("user", {})

    # argumentText strips @BotName mentions in Spaces; fall back to full text in DMs
    user_text      = (message.get("argumentText") or message.get("text", "")).strip()
    sender_display = user.get("displayName", "User")
    space_name     = space.get("name", "unknown")
    space_type     = space.get("type", "")

    logger.info("Google Chat message from %s in %s (%s): %r",
                sender_display, space_name, space_type, user_text[:80])

    if not user_text:
        logger.warning("Google Chat: empty user_text, ignoring")
        return Response(status_code=200)

    session_id    = f"gc_{username}_{space_name.replace('/', '_')}"
    system_prompt = load_context(settings.default_tier)
    history       = load_session(session_id)
    history.append({"role": "user", "content": user_text})

    try:
        response_text, actual_backend = await asyncio.wait_for(
            complete(
                system_prompt=system_prompt,
                messages=history,
                model=backend,
            ),
            timeout=timeout,
        )
    except asyncio.TimeoutError:
        logger.warning("Google Chat request timed out for session %s", session_id)
        return _msg("⏳ Still thinking — this is taking a bit longer than usual. Try again in a moment.")
    except Exception as e:
        logger.error("Google Chat error for session %s: %s", session_id, e)
        return _msg("⚠️ Something went wrong on my end. Try again shortly.")

    logger.info("Google Chat LLM responded via %s (%d chars)", actual_backend, len(response_text))
    history.append({"role": "assistant", "content": response_text})
    save_session(session_id, history)
    log_turn(session_id, user_text, response_text)

    return _msg(response_text)