mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
refactor(automations): remove phase-6 chat watch (watch service/routes + chat_message action)
This commit is contained in:
CREDO23
parent
5da624399d
commit
c66b2f0e0e
17 changed files with 1 additions and 959 deletions
|
|
@ -2,4 +2,4 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from . import agent_task, chat_message # noqa: F401
|
||||
from . import agent_task # noqa: F401
|
||||
|
|
|
|||
|
|
@ -1,15 +0,0 @@
|
|||
"""``chat_message`` action: post one turn into an existing chat thread.
|
||||
|
||||
Imports ``definition`` for its side-effect (self-registration on the actions
|
||||
registry) and re-exports ``build_handler`` for direct consumers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .factory import build_handler
|
||||
from .params import ChatMessageActionParams
|
||||
|
||||
__all__ = ["ChatMessageActionParams", "build_handler"]
|
||||
|
||||
# Side-effect: register on the actions store.
|
||||
from . import definition # noqa: F401
|
||||
|
|
@ -1,18 +0,0 @@
|
|||
"""``chat_message`` ``ActionDefinition`` registration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from ...store import register_action
|
||||
from ...types import ActionDefinition
|
||||
from .factory import build_handler
|
||||
from .params import ChatMessageActionParams
|
||||
|
||||
CHAT_MESSAGE_ACTION = ActionDefinition(
|
||||
type="chat_message",
|
||||
name="Chat message",
|
||||
description="Post a message into an existing chat thread and run one durable, persisted agent turn.",
|
||||
params_model=ChatMessageActionParams,
|
||||
build_handler=build_handler,
|
||||
)
|
||||
|
||||
register_action(CHAT_MESSAGE_ACTION)
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
"""Bind ``ActionContext`` to a callable that runs one ``chat_message`` step."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ...types import ActionContext, ActionHandler
|
||||
from .invoke import run_chat_message
|
||||
from .params import ChatMessageActionParams
|
||||
|
||||
|
||||
def build_handler(ctx: ActionContext) -> ActionHandler:
|
||||
"""Return a handler closure that validates params and posts the turn."""
|
||||
|
||||
async def handle(params: dict[str, Any]) -> dict[str, Any]:
|
||||
validated = ChatMessageActionParams.model_validate(params)
|
||||
return await run_chat_message(
|
||||
ctx=ctx,
|
||||
thread_id=validated.thread_id,
|
||||
message=validated.message,
|
||||
)
|
||||
|
||||
return handle
|
||||
|
|
@ -1,59 +0,0 @@
|
|||
"""Run one ``chat_message`` step by draining ``stream_new_chat``.
|
||||
|
||||
Reuses the interactive chat turn: it persists the messages and advances the
|
||||
durable checkpointer for the thread, so a scheduled tick shares the same
|
||||
conversation memory as a user turn. The SSE frames have no client, so drain them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.auth.context import AuthContext
|
||||
from app.db import ChatVisibility, User
|
||||
from app.services.chat_session_state_service import get_session_state
|
||||
from app.tasks.chat.streaming.flows.new_chat.orchestrator import stream_new_chat
|
||||
|
||||
from ...types import ActionContext
|
||||
|
||||
|
||||
async def run_chat_message(
|
||||
*,
|
||||
ctx: ActionContext,
|
||||
thread_id: int,
|
||||
message: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Post ``message`` into ``thread_id`` and run one durable agent turn.
|
||||
|
||||
Skips when a turn is already in flight on the thread, so a slow tick can't
|
||||
overlap the next scheduled fire.
|
||||
"""
|
||||
state = await get_session_state(ctx.session, thread_id)
|
||||
if state is not None and state.ai_responding_to_user_id is not None:
|
||||
return {"thread_id": thread_id, "frames": 0, "skipped": "in_flight"}
|
||||
|
||||
user_id = str(ctx.creator_user_id) if ctx.creator_user_id else None
|
||||
auth_context: AuthContext | None = None
|
||||
if ctx.creator_user_id:
|
||||
user = await ctx.session.get(User, ctx.creator_user_id)
|
||||
if user is not None:
|
||||
auth_context = AuthContext.system(user, source="automation")
|
||||
|
||||
llm_config_id = ctx.chat_model_id if ctx.chat_model_id is not None else -1
|
||||
request_id = f"automation:{ctx.run_id}:{ctx.step_id}"
|
||||
|
||||
frames = 0
|
||||
async for _sse in stream_new_chat(
|
||||
user_query=message,
|
||||
workspace_id=ctx.workspace_id,
|
||||
chat_id=thread_id,
|
||||
user_id=user_id,
|
||||
llm_config_id=llm_config_id,
|
||||
needs_history_bootstrap=False,
|
||||
thread_visibility=ChatVisibility.PRIVATE,
|
||||
auth_context=auth_context,
|
||||
request_id=request_id,
|
||||
):
|
||||
frames += 1
|
||||
|
||||
return {"thread_id": thread_id, "frames": frames, "skipped": None}
|
||||
|
|
@ -1,21 +0,0 @@
|
|||
"""``ChatMessageActionParams`` — params for the ``chat_message`` action type."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class ChatMessageActionParams(BaseModel):
|
||||
"""Post one turn into an existing chat thread from an automation step."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
thread_id: int = Field(
|
||||
...,
|
||||
description="NewChatThread id (the LangGraph thread) to post the turn into.",
|
||||
)
|
||||
message: str = Field(
|
||||
...,
|
||||
min_length=1,
|
||||
description="Message sent as the user turn; rendered at execute time.",
|
||||
)
|
||||
|
|
@ -5,14 +5,10 @@ from __future__ import annotations
|
|||
from fastapi import APIRouter
|
||||
|
||||
from .automation import router as automation_router
|
||||
from .chat_watch import router as watch_router
|
||||
from .run import router as run_router
|
||||
from .trigger import router as trigger_router
|
||||
|
||||
router = APIRouter()
|
||||
# Before automation_router so the literal ``/automations/watches`` path is not
|
||||
# shadowed by the ``/automations/{automation_id}`` route.
|
||||
router.include_router(watch_router)
|
||||
router.include_router(automation_router)
|
||||
router.include_router(trigger_router)
|
||||
router.include_router(run_router)
|
||||
|
|
|
|||
|
|
@ -1,50 +0,0 @@
|
|||
"""HTTP routes for chat watches (automations bound to a chat thread).
|
||||
|
||||
A watch is an automation whose plan re-posts a question into a chat on a
|
||||
schedule. These routes let the chat UI show watch state and controls:
|
||||
|
||||
* ``GET /automations/watches`` — the watches bound to a thread (is-watched + list).
|
||||
* ``POST /automations/{automation_id}/run`` — run a watch now (manual refresh).
|
||||
|
||||
Stopping a watch is a plain ``DELETE /automations/{automation_id}``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from fastapi import APIRouter, Depends, Query, status
|
||||
|
||||
from app.automations.schemas.api import AutomationList, AutomationSummary, RunSummary
|
||||
from app.automations.services import AutomationService, get_automation_service
|
||||
from app.automations.services.chat_watch import find_watches_for_thread, run_watch_now
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.get("/automations/watches", response_model=AutomationList)
|
||||
async def list_watches(
|
||||
workspace_id: int = Query(...),
|
||||
thread_id: int = Query(...),
|
||||
service: AutomationService = Depends(get_automation_service),
|
||||
) -> AutomationList:
|
||||
"""List the watches bound to a chat thread (empty when it isn't watched)."""
|
||||
watches = await find_watches_for_thread(
|
||||
service, workspace_id=workspace_id, thread_id=thread_id
|
||||
)
|
||||
return AutomationList(
|
||||
items=[AutomationSummary.model_validate(a) for a in watches],
|
||||
total=len(watches),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/automations/{automation_id}/run",
|
||||
response_model=RunSummary,
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def run_watch(
|
||||
automation_id: int,
|
||||
service: AutomationService = Depends(get_automation_service),
|
||||
) -> RunSummary:
|
||||
"""Enqueue an immediate run of a watch (manual refresh)."""
|
||||
run = await run_watch_now(service, automation_id=automation_id)
|
||||
return RunSummary.model_validate(run)
|
||||
|
|
@ -1,152 +0,0 @@
|
|||
"""A chat watch is an automation bound to a chat: a ``schedule`` trigger firing
|
||||
a single ``chat_message`` step that re-posts the question into the same thread.
|
||||
|
||||
Starting a watch creates that automation; stopping deletes it. Whether a chat
|
||||
is watched is derived from the plan (``plan_targets_thread``), not a stored flag.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from app.automations.persistence.enums.trigger_type import TriggerType
|
||||
from app.automations.persistence.models.automation import Automation
|
||||
from app.automations.persistence.models.run import AutomationRun
|
||||
from app.automations.persistence.models.trigger import AutomationTrigger
|
||||
from app.automations.schemas.api import AutomationCreate, TriggerCreate
|
||||
from app.automations.schemas.definition import AutomationDefinition, PlanStep
|
||||
from app.automations.services.automation import AutomationService
|
||||
|
||||
WATCH_ACTION_TYPE = "chat_message"
|
||||
_WATCH_STEP_ID = "watch"
|
||||
_NAME_MAX = 200
|
||||
# Watches per chat are few; one generous page finds them all.
|
||||
_WATCH_SCAN_LIMIT = 500
|
||||
|
||||
|
||||
def _derive_name(message: str) -> str:
|
||||
"""A short, human name for the watch from the question it re-asks."""
|
||||
condensed = " ".join(message.split())
|
||||
label = condensed[:80].rstrip()
|
||||
return f"Watch: {label}" if label else "Watch"
|
||||
|
||||
|
||||
def _build_watch_payload(
|
||||
*,
|
||||
workspace_id: int,
|
||||
thread_id: int,
|
||||
message: str,
|
||||
cron: str,
|
||||
timezone: str,
|
||||
name: str | None,
|
||||
description: str | None,
|
||||
) -> AutomationCreate:
|
||||
watch_name = (name or _derive_name(message))[:_NAME_MAX]
|
||||
return AutomationCreate(
|
||||
workspace_id=workspace_id,
|
||||
name=watch_name,
|
||||
description=description,
|
||||
definition=AutomationDefinition(
|
||||
name=watch_name,
|
||||
plan=[
|
||||
PlanStep(
|
||||
step_id=_WATCH_STEP_ID,
|
||||
action=WATCH_ACTION_TYPE,
|
||||
params={"thread_id": thread_id, "message": message},
|
||||
)
|
||||
],
|
||||
),
|
||||
triggers=[
|
||||
TriggerCreate(
|
||||
type=TriggerType.SCHEDULE,
|
||||
params={"cron": cron, "timezone": timezone},
|
||||
enabled=True,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def create_watch(
|
||||
service: AutomationService,
|
||||
*,
|
||||
workspace_id: int,
|
||||
thread_id: int,
|
||||
message: str,
|
||||
cron: str,
|
||||
timezone: str,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
) -> Automation:
|
||||
"""Bind a schedule + chat_message watch automation to ``thread_id``."""
|
||||
payload = _build_watch_payload(
|
||||
workspace_id=workspace_id,
|
||||
thread_id=thread_id,
|
||||
message=message,
|
||||
cron=cron,
|
||||
timezone=timezone,
|
||||
name=name,
|
||||
description=description,
|
||||
)
|
||||
return await service.create(payload)
|
||||
|
||||
|
||||
async def stop_watch(service: AutomationService, *, automation_id: int) -> None:
|
||||
"""Stop a watch by deleting its automation; the chat reverts to normal."""
|
||||
await service.delete(automation_id)
|
||||
|
||||
|
||||
def plan_targets_thread(definition: dict[str, Any] | None, thread_id: int) -> bool:
|
||||
"""Whether an automation's plan re-posts into ``thread_id`` (i.e. is a watch).
|
||||
|
||||
Pure predicate over the persisted ``definition`` JSON so callers can filter
|
||||
a workspace's automations without a DB round-trip per row.
|
||||
"""
|
||||
for step in (definition or {}).get("plan", []):
|
||||
if step.get("action") != WATCH_ACTION_TYPE:
|
||||
continue
|
||||
if (step.get("params") or {}).get("thread_id") == thread_id:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def find_watches_for_thread(
|
||||
service: AutomationService,
|
||||
*,
|
||||
workspace_id: int,
|
||||
thread_id: int,
|
||||
) -> list[Automation]:
|
||||
"""Return the workspace automations that watch ``thread_id`` (i.e. its watches)."""
|
||||
automations, _total = await service.list(
|
||||
workspace_id=workspace_id, limit=_WATCH_SCAN_LIMIT, offset=0
|
||||
)
|
||||
return [a for a in automations if plan_targets_thread(a.definition, thread_id)]
|
||||
|
||||
|
||||
def schedule_trigger(automation: Automation) -> AutomationTrigger | None:
|
||||
"""The automation's ``schedule`` trigger row, or ``None``."""
|
||||
for trigger in automation.triggers:
|
||||
if trigger.type == TriggerType.SCHEDULE:
|
||||
return trigger
|
||||
return None
|
||||
|
||||
|
||||
async def run_watch_now(
|
||||
service: AutomationService,
|
||||
*,
|
||||
automation_id: int,
|
||||
) -> AutomationRun:
|
||||
"""Enqueue an immediate run of a watch (a manual refresh)."""
|
||||
# Lazy: launch_run pulls in the automation task graph, which imports
|
||||
# multi_agent_chat; this module is imported from there, so a top-level
|
||||
# import would cycle.
|
||||
from app.automations.dispatch import launch as launch_mod
|
||||
|
||||
automation = await service.get(automation_id)
|
||||
trigger = schedule_trigger(automation)
|
||||
if trigger is None:
|
||||
raise HTTPException(
|
||||
status_code=422, detail="watch has no schedule trigger to run"
|
||||
)
|
||||
return await launch_mod.launch_run(session=service.session, trigger=trigger)
|
||||
|
|
@ -1,118 +0,0 @@
|
|||
"""Integration tests for the ``chat_message`` action's concurrency guard.
|
||||
|
||||
The guard reads real ``ChatSessionState`` (seeded via the production
|
||||
``set_ai_responding`` helper) to decide whether a scheduled watch tick may run.
|
||||
This proves the skip decision against real persistence — the part a fake
|
||||
``get_session_state`` could silently get wrong.
|
||||
|
||||
The turn's actual work (``stream_new_chat``) resolves an LLM from the DB and
|
||||
runs the agent; that end-to-end drain is a running-stack concern. Here it is
|
||||
replaced by a spy so we can assert the guard's two outcomes: skip without
|
||||
streaming when a turn is in flight, and stream (with system auth) when not.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.automations.actions.builtin.chat_message import invoke as invoke_mod
|
||||
from app.automations.actions.types import ActionContext
|
||||
from app.db import ChatVisibility, NewChatThread, User, Workspace
|
||||
from app.services.chat_session_state_service import (
|
||||
clear_ai_responding,
|
||||
set_ai_responding,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def thread(
|
||||
db_session: AsyncSession, db_user: User, db_workspace: Workspace
|
||||
) -> NewChatThread:
|
||||
row = NewChatThread(
|
||||
title="Watched chat",
|
||||
workspace_id=db_workspace.id,
|
||||
created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
db_session.add(row)
|
||||
await db_session.flush()
|
||||
return row
|
||||
|
||||
|
||||
def _ctx(session: AsyncSession, *, workspace_id: int, creator_id) -> ActionContext:
|
||||
return ActionContext(
|
||||
session=session,
|
||||
run_id=1,
|
||||
step_id="watch",
|
||||
workspace_id=workspace_id,
|
||||
creator_user_id=creator_id,
|
||||
chat_model_id=1,
|
||||
)
|
||||
|
||||
|
||||
async def test_skips_without_streaming_when_a_turn_is_in_flight(
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
await set_ai_responding(db_session, thread.id, db_user.id)
|
||||
|
||||
streamed = {"called": False}
|
||||
|
||||
async def _spy(**_kwargs: Any):
|
||||
streamed["called"] = True
|
||||
yield "data: x\n\n"
|
||||
|
||||
monkeypatch.setattr(invoke_mod, "stream_new_chat", _spy)
|
||||
|
||||
out = await invoke_mod.run_chat_message(
|
||||
ctx=_ctx(db_session, workspace_id=db_workspace.id, creator_id=db_user.id),
|
||||
thread_id=thread.id,
|
||||
message="tick",
|
||||
)
|
||||
|
||||
assert out["skipped"] == "in_flight"
|
||||
assert out["frames"] == 0
|
||||
assert streamed["called"] is False
|
||||
|
||||
|
||||
async def test_streams_under_system_auth_when_no_turn_in_flight(
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Ensure the thread has no in-flight turn.
|
||||
await clear_ai_responding(db_session, thread.id)
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _spy(**kwargs: Any):
|
||||
captured.update(kwargs)
|
||||
for frame in ("data: a\n\n", "data: b\n\n"):
|
||||
yield frame
|
||||
|
||||
monkeypatch.setattr(invoke_mod, "stream_new_chat", _spy)
|
||||
|
||||
out = await invoke_mod.run_chat_message(
|
||||
ctx=_ctx(db_session, workspace_id=db_workspace.id, creator_id=db_user.id),
|
||||
thread_id=thread.id,
|
||||
message="what changed?",
|
||||
)
|
||||
|
||||
assert out["skipped"] is None
|
||||
assert out["frames"] == 2
|
||||
assert captured["chat_id"] == thread.id
|
||||
assert captured["user_query"] == "what changed?"
|
||||
assert captured["user_id"] == str(db_user.id)
|
||||
auth = captured["auth_context"]
|
||||
assert auth is not None and auth.method == "system" and auth.source == "automation"
|
||||
|
|
@ -1,136 +0,0 @@
|
|||
"""Integration tests for the watch HTTP routes (real app, DB, auth).
|
||||
|
||||
Watches are seeded through the real service and read back through the mounted
|
||||
endpoints, exercising query validation, auth, and response mapping. Run-now is
|
||||
asserted without a broker by spying the Celery enqueue while still persisting a
|
||||
real PENDING run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.auth.context import AuthContext
|
||||
from app.automations.persistence.enums.run_status import RunStatus
|
||||
from app.automations.persistence.models.run import AutomationRun
|
||||
from app.automations.services.automation import AutomationService
|
||||
from app.automations.services.chat_watch import create_watch
|
||||
from app.db import ChatVisibility, NewChatThread, User, Workspace
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
BASE = "/api/v1/automations"
|
||||
_CRON = "0 9 * * 1-5"
|
||||
_TZ = "UTC"
|
||||
|
||||
|
||||
def _billable(workspace: Workspace) -> None:
|
||||
workspace.chat_model_id = 1
|
||||
workspace.image_gen_model_id = 1
|
||||
workspace.vision_model_id = 1
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def thread(
|
||||
db_session: AsyncSession, db_user: User, db_workspace: Workspace
|
||||
) -> NewChatThread:
|
||||
row = NewChatThread(
|
||||
title="Watched chat",
|
||||
workspace_id=db_workspace.id,
|
||||
created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
db_session.add(row)
|
||||
await db_session.flush()
|
||||
return row
|
||||
|
||||
|
||||
async def _seed_watch(
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
*,
|
||||
workspace_id: int,
|
||||
thread_id: int,
|
||||
):
|
||||
service = AutomationService(session=db_session, auth=AuthContext.session(db_user))
|
||||
return await create_watch(
|
||||
service,
|
||||
workspace_id=workspace_id,
|
||||
thread_id=thread_id,
|
||||
message="what changed?",
|
||||
cron=_CRON,
|
||||
timezone=_TZ,
|
||||
)
|
||||
|
||||
|
||||
async def test_list_watches_returns_the_threads_watch(
|
||||
client: httpx.AsyncClient,
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
watch = await _seed_watch(
|
||||
db_session, db_user, workspace_id=db_workspace.id, thread_id=thread.id
|
||||
)
|
||||
|
||||
resp = await client.get(
|
||||
f"{BASE}/watches",
|
||||
params={"workspace_id": db_workspace.id, "thread_id": thread.id},
|
||||
)
|
||||
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert body["total"] == 1
|
||||
assert body["items"][0]["id"] == watch.id
|
||||
|
||||
|
||||
async def test_list_watches_empty_when_thread_not_watched(
|
||||
client: httpx.AsyncClient,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
resp = await client.get(
|
||||
f"{BASE}/watches",
|
||||
params={"workspace_id": db_workspace.id, "thread_id": thread.id},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert resp.json()["total"] == 0
|
||||
|
||||
|
||||
async def test_run_watch_now_enqueues_and_persists_pending_run(
|
||||
client: httpx.AsyncClient,
|
||||
enqueue_spy: list[dict],
|
||||
db_session: AsyncSession,
|
||||
db_user: User,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
watch = await _seed_watch(
|
||||
db_session, db_user, workspace_id=db_workspace.id, thread_id=thread.id
|
||||
)
|
||||
|
||||
resp = await client.post(f"{BASE}/{watch.id}/run")
|
||||
|
||||
assert resp.status_code == 202
|
||||
assert len(enqueue_spy) == 1
|
||||
|
||||
runs = (
|
||||
(
|
||||
await db_session.execute(
|
||||
select(AutomationRun).where(AutomationRun.automation_id == watch.id)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
assert len(runs) == 1
|
||||
assert runs[0].status == RunStatus.PENDING
|
||||
|
|
@ -5,9 +5,6 @@ behavior runs against real Postgres and rolls back at test end:
|
|||
|
||||
* ``client`` — httpx over ASGI with ``get_async_session``/``get_auth_context``
|
||||
overridden to the test session + owner.
|
||||
* ``tools_use_test_session`` — the watch tools open their own
|
||||
``async_session_maker()``; repoint it at ``db_session`` so they join the
|
||||
test transaction (and can see the fixture's uncommitted membership/models).
|
||||
* ``enqueue_spy`` — capture ``automation_run_execute.apply_async`` so run-now
|
||||
can be asserted without a Redis broker.
|
||||
"""
|
||||
|
|
@ -15,7 +12,6 @@ behavior runs against real Postgres and rolls back at test end:
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -55,24 +51,6 @@ async def client(
|
|||
app.dependency_overrides.update(previous)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tools_use_test_session(monkeypatch, db_session: AsyncSession) -> None:
|
||||
"""Make the watch tools' ``async_session_maker()`` yield the test session."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def _session_cm():
|
||||
yield db_session # owned by the outer fixture; do not close
|
||||
|
||||
from app.agents.chat.multi_agent_chat.subagents.builtins.scraping.tools import (
|
||||
refresh_watch,
|
||||
start_watch,
|
||||
stop_watch,
|
||||
)
|
||||
|
||||
for module in (start_watch, stop_watch, refresh_watch):
|
||||
monkeypatch.setattr(module, "async_session_maker", _session_cm)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def enqueue_spy(monkeypatch) -> list[dict]:
|
||||
"""Capture Celery enqueues so run-now needs no broker."""
|
||||
|
|
|
|||
|
|
@ -1,217 +0,0 @@
|
|||
"""Integration tests for the chat watch service against real Postgres.
|
||||
|
||||
A watch is an automation bound to a chat. These exercise the real
|
||||
``AutomationService`` (RBAC, the automation model-billing gate, and the
|
||||
definition JSON round-trip) rather than a fake, so the behaviors the unit
|
||||
layer can't see are proven end to end:
|
||||
|
||||
* creating a watch persists a ``schedule`` + ``chat_message`` automation and
|
||||
it reads back as the thread's watch;
|
||||
* the billing gate rejects a watch when the workspace's models aren't
|
||||
billable (Auto/free) — the exact failure a fake service would hide;
|
||||
* finding watches is thread-scoped; stopping one deletes it;
|
||||
* run-now refuses an automation that has no schedule trigger.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.auth.context import AuthContext
|
||||
from app.automations.persistence.enums.trigger_type import TriggerType
|
||||
from app.automations.persistence.models.automation import Automation
|
||||
from app.automations.persistence.models.trigger import AutomationTrigger
|
||||
from app.automations.services.automation import AutomationService
|
||||
from app.automations.services.chat_watch import (
|
||||
create_watch,
|
||||
find_watches_for_thread,
|
||||
run_watch_now,
|
||||
stop_watch,
|
||||
)
|
||||
from app.db import ChatVisibility, NewChatThread, User, Workspace
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
_CRON = "0 9 * * 1-5"
|
||||
_TZ = "UTC"
|
||||
|
||||
|
||||
def _billable(workspace: Workspace) -> None:
|
||||
"""Give the workspace BYOK (positive id) models so automations are billable."""
|
||||
workspace.chat_model_id = 1
|
||||
workspace.image_gen_model_id = 1
|
||||
workspace.vision_model_id = 1
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def service(db_session: AsyncSession, db_user: User) -> AutomationService:
|
||||
return AutomationService(session=db_session, auth=AuthContext.session(db_user))
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def thread(
|
||||
db_session: AsyncSession, db_user: User, db_workspace: Workspace
|
||||
) -> NewChatThread:
|
||||
row = NewChatThread(
|
||||
title="Watched chat",
|
||||
workspace_id=db_workspace.id,
|
||||
created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
db_session.add(row)
|
||||
await db_session.flush()
|
||||
return row
|
||||
|
||||
|
||||
async def test_create_watch_persists_and_reads_back_as_the_threads_watch(
|
||||
service: AutomationService,
|
||||
db_session: AsyncSession,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
thread_id = thread.id
|
||||
|
||||
created = await create_watch(
|
||||
service,
|
||||
workspace_id=db_workspace.id,
|
||||
thread_id=thread_id,
|
||||
message="what changed on the pricing page?",
|
||||
cron=_CRON,
|
||||
timezone=_TZ,
|
||||
)
|
||||
|
||||
assert created.id is not None
|
||||
plan = created.definition["plan"]
|
||||
assert len(plan) == 1
|
||||
assert plan[0]["action"] == "chat_message"
|
||||
assert plan[0]["params"] == {
|
||||
"thread_id": thread_id,
|
||||
"message": "what changed on the pricing page?",
|
||||
}
|
||||
schedule = [t for t in created.triggers if t.type == TriggerType.SCHEDULE]
|
||||
assert len(schedule) == 1
|
||||
assert schedule[0].params["cron"] == _CRON
|
||||
|
||||
found = await find_watches_for_thread(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_id
|
||||
)
|
||||
assert [a.id for a in found] == [created.id]
|
||||
|
||||
|
||||
async def test_create_watch_rejected_when_workspace_models_not_billable(
|
||||
service: AutomationService,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
# Default workspace has no model prefs (Auto/None) -> not billable.
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await create_watch(
|
||||
service,
|
||||
workspace_id=db_workspace.id,
|
||||
thread_id=thread.id,
|
||||
message="track it",
|
||||
cron=_CRON,
|
||||
timezone=_TZ,
|
||||
)
|
||||
assert exc.value.status_code == 422
|
||||
|
||||
|
||||
async def test_find_watches_is_scoped_to_the_thread(
|
||||
service: AutomationService,
|
||||
db_session: AsyncSession,
|
||||
db_workspace: Workspace,
|
||||
db_user: User,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
|
||||
thread_a = NewChatThread(
|
||||
title="A", workspace_id=db_workspace.id, created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
thread_b = NewChatThread(
|
||||
title="B", workspace_id=db_workspace.id, created_by_id=db_user.id,
|
||||
visibility=ChatVisibility.PRIVATE,
|
||||
)
|
||||
db_session.add_all([thread_a, thread_b])
|
||||
await db_session.flush()
|
||||
|
||||
watch_a = await create_watch(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_a.id,
|
||||
message="a", cron=_CRON, timezone=_TZ,
|
||||
)
|
||||
await create_watch(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_b.id,
|
||||
message="b", cron=_CRON, timezone=_TZ,
|
||||
)
|
||||
|
||||
found = await find_watches_for_thread(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_a.id
|
||||
)
|
||||
assert [a.id for a in found] == [watch_a.id]
|
||||
|
||||
|
||||
async def test_stop_watch_deletes_and_leaves_the_thread_unwatched(
|
||||
service: AutomationService,
|
||||
db_session: AsyncSession,
|
||||
db_workspace: Workspace,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
_billable(db_workspace)
|
||||
await db_session.flush()
|
||||
thread_id = thread.id
|
||||
|
||||
created = await create_watch(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_id,
|
||||
message="m", cron=_CRON, timezone=_TZ,
|
||||
)
|
||||
|
||||
await stop_watch(service, automation_id=created.id)
|
||||
|
||||
found = await find_watches_for_thread(
|
||||
service, workspace_id=db_workspace.id, thread_id=thread_id
|
||||
)
|
||||
assert found == []
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await service.get(created.id)
|
||||
assert exc.value.status_code == 404
|
||||
|
||||
|
||||
async def test_run_watch_now_rejects_automation_without_schedule_trigger(
|
||||
service: AutomationService,
|
||||
db_session: AsyncSession,
|
||||
db_workspace: Workspace,
|
||||
db_user: User,
|
||||
thread: NewChatThread,
|
||||
) -> None:
|
||||
# Seed a watch-shaped automation whose only trigger is non-schedule.
|
||||
automation = Automation(
|
||||
workspace_id=db_workspace.id,
|
||||
created_by_user_id=db_user.id,
|
||||
name="Watch: broken",
|
||||
definition={
|
||||
"name": "Watch: broken",
|
||||
"plan": [
|
||||
{
|
||||
"step_id": "watch",
|
||||
"action": "chat_message",
|
||||
"params": {"thread_id": thread.id, "message": "m"},
|
||||
}
|
||||
],
|
||||
},
|
||||
version=1,
|
||||
)
|
||||
automation.triggers.append(
|
||||
AutomationTrigger(type=TriggerType.EVENT, params={}, enabled=True)
|
||||
)
|
||||
db_session.add(automation)
|
||||
await db_session.flush()
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await run_watch_now(service, automation_id=automation.id)
|
||||
assert exc.value.status_code == 422
|
||||
|
|
@ -1,48 +0,0 @@
|
|||
"""Pure-logic unit tests for the ``chat_message`` action.
|
||||
|
||||
Only the infra-free contract lives here: the action self-registers, and its
|
||||
params model requires ``thread_id`` + ``message`` and forbids extras. The
|
||||
run-time behavior (the in-flight guard against real ``ChatSessionState`` and
|
||||
streaming under system auth) is proven in
|
||||
``tests/integration/automations/actions/builtin/chat_message/test_chat_message.py``,
|
||||
against real persistence rather than a faked ``get_session_state``/``stream_new_chat``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def test_chat_message_action_is_registered_after_package_import() -> None:
|
||||
import app.automations # noqa: F401 (force side-effect registrations)
|
||||
from app.automations.actions.store import get_action
|
||||
|
||||
definition = get_action("chat_message")
|
||||
|
||||
assert definition is not None
|
||||
assert definition.type == "chat_message"
|
||||
|
||||
|
||||
def test_params_require_thread_id_and_message_and_forbid_extra() -> None:
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.automations.actions.builtin.chat_message.params import (
|
||||
ChatMessageActionParams,
|
||||
)
|
||||
|
||||
ok = ChatMessageActionParams.model_validate({"thread_id": 7, "message": "hi"})
|
||||
assert ok.thread_id == 7
|
||||
assert ok.message == "hi"
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessageActionParams.model_validate({"thread_id": 7, "message": ""})
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessageActionParams.model_validate({"thread_id": 7})
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessageActionParams.model_validate(
|
||||
{"thread_id": 7, "message": "hi", "surprise": True}
|
||||
)
|
||||
|
|
@ -1,75 +0,0 @@
|
|||
"""Pure-logic unit tests for the chat watch service.
|
||||
|
||||
Only the deterministic, infra-free pieces live here: the plan predicate that
|
||||
decides whether an automation watches a thread, and the schedule-trigger
|
||||
picker. Everything that touches ``AutomationService``/Postgres (create, find,
|
||||
stop, run-now) is proven against real infra in
|
||||
``tests/integration/automations/services/test_chat_watch.py`` — faking the
|
||||
service we own would only assert wiring and hide the model-billing gate,
|
||||
RBAC, and JSON round-trip that those calls really go through.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from app.automations.persistence.enums.trigger_type import TriggerType
|
||||
from app.automations.persistence.models.trigger import AutomationTrigger
|
||||
from app.automations.services.chat_watch import (
|
||||
WATCH_ACTION_TYPE,
|
||||
plan_targets_thread,
|
||||
schedule_trigger,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
def _definition_with_step(action: str, params: dict[str, Any]) -> dict[str, Any]:
|
||||
return {"plan": [{"step_id": "s", "action": action, "params": params}]}
|
||||
|
||||
|
||||
def test_plan_targets_thread_true_for_matching_chat_message_step() -> None:
|
||||
definition = _definition_with_step("chat_message", {"thread_id": 55, "message": "x"})
|
||||
assert plan_targets_thread(definition, 55) is True
|
||||
assert WATCH_ACTION_TYPE == "chat_message"
|
||||
|
||||
|
||||
def test_plan_targets_thread_false_for_other_thread_or_action() -> None:
|
||||
assert (
|
||||
plan_targets_thread(
|
||||
_definition_with_step("chat_message", {"thread_id": 999}), 55
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert (
|
||||
plan_targets_thread(
|
||||
_definition_with_step("agent_task", {"thread_id": 55}), 55
|
||||
)
|
||||
is False
|
||||
)
|
||||
assert plan_targets_thread({}, 55) is False
|
||||
assert plan_targets_thread(None, 55) is False
|
||||
|
||||
|
||||
class _Automation:
|
||||
"""Minimal stand-in carrying only the ``triggers`` list the picker reads."""
|
||||
|
||||
def __init__(self, triggers: list[AutomationTrigger]) -> None:
|
||||
self.triggers = triggers
|
||||
|
||||
|
||||
def _trigger(type_: TriggerType) -> AutomationTrigger:
|
||||
return AutomationTrigger(type=type_, params={}, enabled=True)
|
||||
|
||||
|
||||
def test_schedule_trigger_picks_the_schedule_row() -> None:
|
||||
schedule = _trigger(TriggerType.SCHEDULE)
|
||||
automation = _Automation([_trigger(TriggerType.EVENT), schedule])
|
||||
assert schedule_trigger(automation) is schedule
|
||||
|
||||
|
||||
def test_schedule_trigger_returns_none_without_a_schedule() -> None:
|
||||
automation = _Automation([_trigger(TriggerType.EVENT)])
|
||||
assert schedule_trigger(automation) is None
|
||||
Loading…
Add table
Add a link
Reference in a new issue