feat: webhooks (#168)
* feat: webhooks * feat: Enhance webhook security and typing, fix validation and encryption bugs * fix: lint / types * fix: rm files * fix: rm mcp * fix: pydantic issue with TypedDict in python version <= 3.11 * fix: pre-commit hook for test coverage * fix: simplify API -- store url on workspace * fix: redo architecture * fix: webhook body * fix: make workspace optional * fix: comments * refactor: add webhook secret * fix: CR comments * feat: use deriver for webhooks * use key-value approach * feat: add work unit key to deriver * fix: add work unit key to webhooks * fix: tests * fix: cr comments #2 * fix: endpoint structure; make webhook delivery into a function; add tests; other general comments * chore: change webhook secret, fix test event and workspace_id, use async with * feat: implement queue.empty and backfill * fix: unique constraint * refactor: queue to use outerjoin and remove skip locked; also fix publish queue.empty * fix: tests * fix: migration - make columns non-nullable
This commit is contained in:
parent
7b174dd34b
commit
3bea3da169
|
|
@ -0,0 +1,185 @@
|
|||
"""add webhooks table
|
||||
|
||||
Revision ID: 88b0fb10906f
|
||||
Revises: 05486ce795d5
|
||||
Create Date: 2025-07-25 16:12:11.015327
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from migrations.utils import (
|
||||
column_exists,
|
||||
constraint_exists,
|
||||
index_exists,
|
||||
table_exists,
|
||||
)
|
||||
from src.config import settings
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "88b0fb10906f"
|
||||
down_revision: str | None = "05486ce795d5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = settings.DB.SCHEMA
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# 1. Add webhook_endpoints table
|
||||
op.create_table(
|
||||
"webhook_endpoints",
|
||||
sa.Column(
|
||||
"id",
|
||||
sa.TEXT(),
|
||||
primary_key=True,
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column(
|
||||
"workspace_name",
|
||||
sa.TEXT(),
|
||||
sa.ForeignKey("workspaces.name"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("url", sa.TEXT(), nullable=False),
|
||||
sa.Column(
|
||||
"created_at",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.func.now(),
|
||||
),
|
||||
sa.CheckConstraint("length(url) <= 2048", name="webhook_endpoint_url_length"),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
op.create_index(
|
||||
op.f("idx_webhook_endpoints_workspace_lookup"),
|
||||
"webhook_endpoints",
|
||||
["workspace_name"],
|
||||
unique=False,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# 2. Add columns to queue table
|
||||
op.add_column(
|
||||
"queue",
|
||||
sa.Column("task_type", sa.TEXT(), nullable=True),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
op.add_column(
|
||||
"queue",
|
||||
sa.Column("work_unit_key", sa.Text(), nullable=True),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# 2.5 Backfill task_type and work_unit_key for existing queue items (batched)
|
||||
op.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
rows_updated INT;
|
||||
BEGIN
|
||||
LOOP
|
||||
UPDATE {schema}.queue
|
||||
SET
|
||||
task_type = COALESCE(payload->>'task_type', 'representation'),
|
||||
work_unit_key =
|
||||
COALESCE(payload->>'task_type', 'representation') || ':' ||
|
||||
COALESCE(payload->>'workspace_name', 'None') || ':' ||
|
||||
COALESCE(payload->>'session_name', 'None') || ':' ||
|
||||
COALESCE(payload->>'sender_name', 'None') || ':' ||
|
||||
COALESCE(payload->>'target_name', 'None')
|
||||
WHERE id IN (
|
||||
SELECT id FROM {schema}.queue
|
||||
WHERE task_type IS NULL OR work_unit_key IS NULL
|
||||
LIMIT 1000
|
||||
);
|
||||
|
||||
GET DIAGNOSTICS rows_updated = ROW_COUNT;
|
||||
EXIT WHEN rows_updated = 0;
|
||||
END LOOP;
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
# Make both columns non-nullable
|
||||
op.alter_column("queue", "task_type", nullable=False, schema=schema)
|
||||
op.alter_column("queue", "work_unit_key", nullable=False, schema=schema)
|
||||
|
||||
# 3. Alter active queue sessions table
|
||||
op.add_column(
|
||||
"active_queue_sessions",
|
||||
sa.Column("work_unit_key", sa.Text(), index=True),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Add unique constraint for work_unit_key
|
||||
op.create_unique_constraint(
|
||||
"unique_work_unit_key",
|
||||
"active_queue_sessions",
|
||||
["work_unit_key"],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
|
||||
if constraint_exists(
|
||||
"active_queue_sessions", "unique_active_queue_session", "unique", inspector
|
||||
):
|
||||
op.drop_constraint(
|
||||
"unique_active_queue_session", "active_queue_sessions", schema=schema
|
||||
)
|
||||
|
||||
if column_exists("active_queue_sessions", "session_id", inspector):
|
||||
op.drop_column("active_queue_sessions", "session_id", schema=schema)
|
||||
|
||||
if column_exists("active_queue_sessions", "sender_name", inspector):
|
||||
op.drop_column("active_queue_sessions", "sender_name", schema=schema)
|
||||
|
||||
if column_exists("active_queue_sessions", "target_name", inspector):
|
||||
op.drop_column("active_queue_sessions", "target_name", schema=schema)
|
||||
|
||||
if column_exists("active_queue_sessions", "task_type", inspector):
|
||||
op.drop_column("active_queue_sessions", "task_type", schema=schema)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
|
||||
if table_exists("webhook_endpoints", inspector):
|
||||
if index_exists(
|
||||
"webhook_endpoints", "idx_webhook_endpoints_workspace_lookup", inspector
|
||||
):
|
||||
op.drop_index(
|
||||
op.f("idx_webhook_endpoints_workspace_lookup"),
|
||||
table_name="webhook_endpoints",
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
op.drop_table("webhook_endpoints", schema=schema)
|
||||
|
||||
if column_exists("queue", "task_type", inspector):
|
||||
op.drop_column("queue", "task_type", schema=schema)
|
||||
|
||||
if column_exists("queue", "work_unit_key", inspector):
|
||||
op.drop_column("queue", "work_unit_key", schema=schema)
|
||||
|
||||
# Drop unique constraint first if it exists
|
||||
if constraint_exists(
|
||||
"active_queue_sessions", "unique_work_unit_key", "unique", inspector
|
||||
):
|
||||
op.drop_constraint(
|
||||
"unique_work_unit_key", "active_queue_sessions", schema=schema
|
||||
)
|
||||
|
||||
if column_exists("active_queue_sessions", "work_unit_key", inspector):
|
||||
op.drop_column("active_queue_sessions", "work_unit_key", schema=schema)
|
||||
|
||||
if column_exists("active_queue_sessions", "work_unit_data", inspector):
|
||||
op.drop_column("active_queue_sessions", "work_unit_data", schema=schema)
|
||||
|
|
@ -28,6 +28,7 @@ dependencies = [
|
|||
"pydantic-settings>=2.10.1",
|
||||
"google-generativeai>=0.8.5",
|
||||
"pdfplumber>=0.11.7",
|
||||
"typing-extensions>=4.11.0",
|
||||
]
|
||||
[tool.uv]
|
||||
dev-dependencies = [
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ def main():
|
|||
print(f"Generated JWT secret: {secret}")
|
||||
print("\nAdd this to your .env file as:")
|
||||
print(f"AUTH_JWT_SECRET={secret}")
|
||||
print(f"or as WEBHOOK_SECRET={secret}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ class TomlConfigSettingsSource(PydanticBaseSettingsSource):
|
|||
"DERIVER": "deriver",
|
||||
"DIALECTIC": "dialectic",
|
||||
"SUMMARY": "summary",
|
||||
"WEBHOOK": "webhook",
|
||||
"": "app", # For AppSettings with no prefix
|
||||
}
|
||||
|
||||
|
|
@ -232,6 +233,13 @@ class SummarySettings(HonchoSettings):
|
|||
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=512, gt=0, le=2000)] = 512
|
||||
|
||||
|
||||
class WebhookSettings(HonchoSettings):
|
||||
model_config = SettingsConfigDict(env_prefix="WEBHOOK_", extra="ignore") # pyright: ignore
|
||||
|
||||
SECRET: str | None = None # Must be set if configuring webhooks
|
||||
MAX_WORKSPACE_LIMIT: int = 10
|
||||
|
||||
|
||||
class AppSettings(HonchoSettings):
|
||||
# No env_prefix for app-level settings
|
||||
model_config = SettingsConfigDict( # pyright: ignore
|
||||
|
|
@ -262,6 +270,7 @@ class AppSettings(HonchoSettings):
|
|||
DERIVER: DeriverSettings = Field(default_factory=DeriverSettings)
|
||||
DIALECTIC: DialecticSettings = Field(default_factory=DialecticSettings)
|
||||
SUMMARY: SummarySettings = Field(default_factory=SummarySettings)
|
||||
WEBHOOK: WebhookSettings = Field(default_factory=WebhookSettings)
|
||||
|
||||
@field_validator("LOG_LEVEL")
|
||||
def validate_log_level(cls, v: str) -> str:
|
||||
|
|
|
|||
|
|
@ -36,6 +36,11 @@ from .session import (
|
|||
set_peers_for_session,
|
||||
update_session,
|
||||
)
|
||||
from .webhook import (
|
||||
delete_webhook_endpoint,
|
||||
get_or_create_webhook_endpoint,
|
||||
list_webhook_endpoints,
|
||||
)
|
||||
from .workspace import get_all_workspaces, get_or_create_workspace, update_workspace
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -80,6 +85,10 @@ __all__ = [
|
|||
"set_peers_for_session",
|
||||
"get_peer_config",
|
||||
"set_peer_config",
|
||||
# Webhook
|
||||
"get_or_create_webhook_endpoint",
|
||||
"delete_webhook_endpoint",
|
||||
"list_webhook_endpoints",
|
||||
# Workspace
|
||||
"get_or_create_workspace",
|
||||
"get_all_workspaces",
|
||||
|
|
|
|||
|
|
@ -58,7 +58,6 @@ def _build_queue_status_query(
|
|||
"""Build SQL query for queue status with validation and aggregation."""
|
||||
sender_name_expr = models.QueueItem.payload["sender_name"].astext
|
||||
target_name_expr = models.QueueItem.payload["target_name"].astext
|
||||
task_type_expr = models.QueueItem.payload["task_type"].astext
|
||||
|
||||
# Define conditions for cleaner window functions
|
||||
is_completed = models.QueueItem.processed
|
||||
|
|
@ -94,10 +93,7 @@ def _build_queue_status_query(
|
|||
|
||||
stmt = stmt.outerjoin(
|
||||
models.ActiveQueueSession,
|
||||
(models.QueueItem.session_id == models.ActiveQueueSession.session_id)
|
||||
& (sender_name_expr == models.ActiveQueueSession.sender_name)
|
||||
& (target_name_expr == models.ActiveQueueSession.target_name)
|
||||
& (task_type_expr == models.ActiveQueueSession.task_type),
|
||||
models.QueueItem.work_unit_key == models.ActiveQueueSession.work_unit_key,
|
||||
)
|
||||
|
||||
stmt = stmt.join(models.Session, models.QueueItem.session_id == models.Session.id)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,115 @@
|
|||
from logging import getLogger
|
||||
|
||||
from sqlalchemy import Select, select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models, schemas
|
||||
from src.config import settings
|
||||
from src.crud.workspace import get_workspace
|
||||
from src.exceptions import ResourceNotFoundException
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
async def get_or_create_webhook_endpoint(
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
webhook: schemas.WebhookEndpointCreate,
|
||||
) -> schemas.WebhookEndpoint:
|
||||
"""
|
||||
Get or create a webhook endpoint, optionally for a workspace.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
webhook: Webhook endpoint creation schema
|
||||
|
||||
Returns:
|
||||
The webhook endpoint
|
||||
|
||||
Raises:
|
||||
ResourceNotFoundException: If the workspace is specified and does not exist
|
||||
"""
|
||||
# Verify workspace exists
|
||||
await get_workspace(db, workspace_name=workspace_name)
|
||||
|
||||
stmt = select(models.WebhookEndpoint).where(
|
||||
models.WebhookEndpoint.workspace_name == workspace_name,
|
||||
)
|
||||
result = await db.execute(stmt)
|
||||
endpoints = result.scalars().all()
|
||||
|
||||
# No more than WORKSPACE_LIMIT webhooks per workspace
|
||||
if len(endpoints) >= settings.WEBHOOK.MAX_WORKSPACE_LIMIT:
|
||||
raise ValueError(
|
||||
f"Maximum number of webhook endpoints ({settings.WEBHOOK.MAX_WORKSPACE_LIMIT}) reached for this workspace."
|
||||
)
|
||||
|
||||
# Check if webhook already exists for this workspace
|
||||
for endpoint in endpoints:
|
||||
if endpoint.url == webhook.url:
|
||||
return schemas.WebhookEndpoint.model_validate(endpoint)
|
||||
|
||||
# Create new webhook endpoint
|
||||
webhook_endpoint = models.WebhookEndpoint(
|
||||
workspace_name=workspace_name,
|
||||
url=webhook.url,
|
||||
)
|
||||
db.add(webhook_endpoint)
|
||||
await db.commit()
|
||||
await db.refresh(webhook_endpoint)
|
||||
|
||||
logger.info(f"Webhook endpoint created: {webhook.url}")
|
||||
return schemas.WebhookEndpoint.model_validate(webhook_endpoint)
|
||||
|
||||
|
||||
async def list_webhook_endpoints(
|
||||
db: AsyncSession, workspace_name: str
|
||||
) -> Select[tuple[models.WebhookEndpoint]]:
|
||||
"""
|
||||
List all webhook endpoints, optionally filtered by workspace.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
workspace_name: Name of the workspace (optional)
|
||||
|
||||
Returns:
|
||||
List of webhook endpoints
|
||||
"""
|
||||
# Verify workspace exists
|
||||
await get_workspace(db, workspace_name)
|
||||
|
||||
return select(models.WebhookEndpoint).where(
|
||||
models.WebhookEndpoint.workspace_name == workspace_name
|
||||
)
|
||||
|
||||
|
||||
async def delete_webhook_endpoint(
|
||||
db: AsyncSession, workspace_name: str, endpoint_id: str
|
||||
) -> None:
|
||||
"""
|
||||
Delete a webhook endpoint.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
endpoint_id: ID of the webhook endpoint
|
||||
|
||||
Raises:
|
||||
ResourceNotFoundException: If the webhook endpoint is not found
|
||||
"""
|
||||
# Verify webhook endpoint exists
|
||||
stmt = select(models.WebhookEndpoint).where(
|
||||
models.WebhookEndpoint.id == endpoint_id,
|
||||
models.WebhookEndpoint.workspace_name == workspace_name,
|
||||
)
|
||||
result = await db.execute(stmt)
|
||||
endpoint = result.scalar_one_or_none()
|
||||
|
||||
if not endpoint:
|
||||
raise ResourceNotFoundException(
|
||||
f"Webhook endpoint {endpoint_id} not found for workspace {workspace_name}"
|
||||
)
|
||||
|
||||
await db.delete(endpoint)
|
||||
await db.commit()
|
||||
|
||||
logger.info(f"Webhook endpoint {endpoint_id} deleted")
|
||||
|
|
@ -6,7 +6,7 @@ from sqlalchemy.exc import IntegrityError
|
|||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models, schemas
|
||||
from src.exceptions import ConflictException
|
||||
from src.exceptions import ConflictException, ResourceNotFoundException
|
||||
from src.utils.filter import apply_filter
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
|
@ -77,6 +77,34 @@ async def get_all_workspaces(
|
|||
return stmt
|
||||
|
||||
|
||||
async def get_workspace(
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
) -> models.Workspace:
|
||||
"""
|
||||
Get an existing workspace.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
workspace_name: Name of the workspace
|
||||
|
||||
Returns:
|
||||
The workspace if found or created
|
||||
|
||||
Raises:
|
||||
ResourceNotFoundException: If the workspace does not exist
|
||||
"""
|
||||
# Try to get the existing peer
|
||||
stmt = select(models.Workspace).where(models.Workspace.name == workspace_name)
|
||||
result = await db.execute(stmt)
|
||||
existing_workspace = result.scalar_one_or_none()
|
||||
|
||||
if existing_workspace is not None:
|
||||
return existing_workspace
|
||||
|
||||
raise ResourceNotFoundException(f"Workspace {workspace_name} not found")
|
||||
|
||||
|
||||
async def update_workspace(
|
||||
db: AsyncSession, workspace_name: str, workspace: schemas.WorkspaceUpdate
|
||||
) -> models.Workspace:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from pydantic import ValidationError
|
|||
from rich.console import Console
|
||||
|
||||
from .deriver import Deriver
|
||||
from .queue_payload import RepresentationPayload, SummaryPayload
|
||||
from .queue_payload import RepresentationPayload, SummaryPayload, WebhookPayload
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
|
||||
|
|
@ -15,14 +15,15 @@ console = Console(markup=True)
|
|||
deriver = Deriver()
|
||||
|
||||
|
||||
async def process_item(payload: dict[str, Any]) -> None:
|
||||
async def process_item(task_type: str, payload: dict[str, Any]) -> None:
|
||||
# Validate payload structure and types before processing
|
||||
try:
|
||||
task_type = payload.get("task_type")
|
||||
if task_type == "representation":
|
||||
validated_payload = RepresentationPayload(**payload)
|
||||
elif task_type == "summary":
|
||||
validated_payload = SummaryPayload(**payload)
|
||||
elif task_type == "webhook":
|
||||
validated_payload = WebhookPayload(**payload)
|
||||
else:
|
||||
raise ValueError(f"Invalid task_type: {task_type}")
|
||||
except ValidationError as e:
|
||||
|
|
@ -30,10 +31,20 @@ async def process_item(payload: dict[str, Any]) -> None:
|
|||
raise ValueError(f"Invalid payload structure: {str(e)}") from e
|
||||
|
||||
logger.debug(
|
||||
"process_item received payload for message %s in session %s, task type %s",
|
||||
validated_payload.message_id,
|
||||
validated_payload.session_name,
|
||||
validated_payload.task_type,
|
||||
"process_item received payload for task type %s ",
|
||||
task_type,
|
||||
)
|
||||
await deriver.process_message(validated_payload)
|
||||
logger.debug("Finished processing message %s", validated_payload.message_id)
|
||||
|
||||
if task_type == "webhook":
|
||||
if not isinstance(validated_payload, WebhookPayload):
|
||||
raise ValueError(f"Expected WebhookPayload, got {type(validated_payload)}")
|
||||
await deriver.process_webhook(validated_payload)
|
||||
logger.debug("Finished processing webhook %s", validated_payload.event_type)
|
||||
else:
|
||||
if not isinstance(validated_payload, RepresentationPayload | SummaryPayload):
|
||||
raise ValueError(
|
||||
f"Expected DeriverQueuePayload, got {type(validated_payload)}"
|
||||
)
|
||||
deriver_payload = validated_payload
|
||||
await deriver.process_message(task_type, deriver_payload)
|
||||
logger.debug("Finished processing message %s", deriver_payload.message_id)
|
||||
|
|
|
|||
|
|
@ -36,9 +36,15 @@ from src.utils.shared_models import (
|
|||
ReasoningResponseWithThinking,
|
||||
UnifiedObservation,
|
||||
)
|
||||
from src.webhooks import webhook_delivery
|
||||
|
||||
from .prompts import critical_analysis_prompt
|
||||
from .queue_payload import DeriverQueuePayload, RepresentationPayload, SummaryPayload
|
||||
from .queue_payload import (
|
||||
DeriverQueuePayload,
|
||||
RepresentationPayload,
|
||||
SummaryPayload,
|
||||
WebhookPayload,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
|
||||
|
|
@ -77,9 +83,18 @@ async def critical_analysis_call(
|
|||
class Deriver:
|
||||
"""Deriver class for processing messages and extracting insights."""
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def process_webhook(
|
||||
self,
|
||||
payload: WebhookPayload,
|
||||
) -> None:
|
||||
async with tracked_db() as db:
|
||||
await webhook_delivery.deliver_webhook(db, payload)
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def process_message(
|
||||
self,
|
||||
task_type: str,
|
||||
payload: DeriverQueuePayload,
|
||||
) -> None:
|
||||
"""
|
||||
|
|
@ -96,10 +111,18 @@ class Deriver:
|
|||
|
||||
# Open a DB session only for the duration of the processing call
|
||||
async with tracked_db("deriver") as db:
|
||||
if payload.task_type == "summary":
|
||||
if task_type == "summary":
|
||||
if not isinstance(payload, SummaryPayload):
|
||||
raise ValueError(f"Expected SummaryPayload, got {type(payload)}")
|
||||
await self.process_summary_task(db, payload)
|
||||
else:
|
||||
elif task_type == "representation":
|
||||
if not isinstance(payload, RepresentationPayload):
|
||||
raise ValueError(
|
||||
f"Expected RepresentationPayload, got {type(payload)}"
|
||||
)
|
||||
await self.process_representation_task(db, payload)
|
||||
else:
|
||||
raise ValueError(f"Unknown task type: {task_type}")
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def process_summary_task(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
from src import crud, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import tracked_db
|
||||
from src.deriver.utils import get_work_unit_key
|
||||
from src.exceptions import ValidationException
|
||||
from src.models import QueueItem
|
||||
|
||||
|
|
@ -152,8 +153,12 @@ def create_representation_record(
|
|||
task_type="representation",
|
||||
)
|
||||
return {
|
||||
"work_unit_key": get_work_unit_key(
|
||||
task_type="representation", payload=processed_payload
|
||||
),
|
||||
"payload": processed_payload,
|
||||
"session_id": session_id,
|
||||
"task_type": "representation",
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -180,8 +185,12 @@ def create_summary_record(
|
|||
message_seq_in_session=message_seq_in_session,
|
||||
)
|
||||
return {
|
||||
"work_unit_key": get_work_unit_key(
|
||||
task_type="summary", payload=processed_payload
|
||||
),
|
||||
"payload": processed_payload,
|
||||
"session_id": session_id,
|
||||
"task_type": "summary",
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,8 +2,7 @@ import asyncio
|
|||
import signal
|
||||
from asyncio import Task
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from logging import getLogger
|
||||
|
||||
import sentry_sdk
|
||||
|
|
@ -15,6 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
from sqlalchemy.sql import func
|
||||
|
||||
from src.config import settings
|
||||
from src.models import QueueItem
|
||||
|
||||
from .. import models
|
||||
from ..dependencies import tracked_db
|
||||
|
|
@ -25,33 +25,11 @@ logger = getLogger(__name__)
|
|||
load_dotenv(override=True)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkUnit:
|
||||
"""
|
||||
Represents a unit of work in the queue system.
|
||||
|
||||
A work unit is uniquely identified by the combination of session_id,
|
||||
sender_name, target_name, and task_type. This allows multiple workers
|
||||
to process different work units from the same session in parallel.
|
||||
|
||||
For summary tasks, sender_name and target_name are None since summary
|
||||
tasks don't have these fields and should be processed sequentially per session.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
sender_name: str | None
|
||||
target_name: str | None
|
||||
task_type: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"({self.session_id}, {self.sender_name}, {self.target_name}, {self.task_type})"
|
||||
|
||||
|
||||
class QueueManager:
|
||||
def __init__(self):
|
||||
self.shutdown_event: asyncio.Event = asyncio.Event()
|
||||
self.active_tasks: set[asyncio.Task[None]] = set()
|
||||
self.owned_work_units: set[WorkUnit] = set()
|
||||
self.owned_work_units: set[str] = set()
|
||||
self.queue_empty_flag: asyncio.Event = asyncio.Event()
|
||||
|
||||
# Initialize from settings
|
||||
|
|
@ -70,21 +48,21 @@ class QueueManager:
|
|||
integrations=[AsyncioIntegration()],
|
||||
)
|
||||
|
||||
def add_task(self, task: asyncio.Task[None]):
|
||||
def add_task(self, task: asyncio.Task[None]) -> None:
|
||||
"""Track a new task"""
|
||||
self.active_tasks.add(task)
|
||||
task.add_done_callback(self.active_tasks.discard)
|
||||
|
||||
def track_work_unit(self, work_unit: WorkUnit):
|
||||
def track_work_unit(self, work_unit_key: str) -> None:
|
||||
"""Track a new work unit owned by this process"""
|
||||
self.owned_work_units.add(work_unit)
|
||||
self.owned_work_units.add(work_unit_key)
|
||||
|
||||
def untrack_work_unit(self, work_unit: WorkUnit):
|
||||
def untrack_work_unit(self, work_unit_key: str) -> None:
|
||||
"""Remove a work unit from tracking"""
|
||||
self.owned_work_units.discard(work_unit)
|
||||
self.owned_work_units.discard(work_unit_key)
|
||||
|
||||
async def initialize(self):
|
||||
"""Setup signal handlers and start the main polling loop"""
|
||||
async def initialize(self) -> None:
|
||||
"""Setup signal handlers, initialize client, and start the main polling loop"""
|
||||
logger.debug(f"Initializing QueueManager with {self.workers} workers")
|
||||
|
||||
# Set up signal handlers
|
||||
|
|
@ -103,7 +81,7 @@ class QueueManager:
|
|||
finally:
|
||||
await self.cleanup()
|
||||
|
||||
async def shutdown(self, sig: signal.Signals):
|
||||
async def shutdown(self, sig: signal.Signals) -> None:
|
||||
"""Handle graceful shutdown"""
|
||||
logger.info(f"Received exit signal {sig.name}...")
|
||||
self.shutdown_event.set()
|
||||
|
|
@ -114,24 +92,17 @@ class QueueManager:
|
|||
)
|
||||
await asyncio.gather(*self.active_tasks, return_exceptions=True)
|
||||
|
||||
async def cleanup(self):
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up owned work units"""
|
||||
if self.owned_work_units:
|
||||
logger.info(f"Cleaning up {len(self.owned_work_units)} owned work units...")
|
||||
try:
|
||||
# Use the tracked_db dependency for transaction safety
|
||||
async with tracked_db("queue_cleanup") as db:
|
||||
for work_unit in self.owned_work_units:
|
||||
for work_unit_key in self.owned_work_units:
|
||||
await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.session_id
|
||||
== work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name
|
||||
== work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name
|
||||
== work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type
|
||||
== work_unit.task_type,
|
||||
models.ActiveQueueSession.work_unit_key == work_unit_key
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
|
@ -145,13 +116,13 @@ class QueueManager:
|
|||
# Polling and Scheduling #
|
||||
##########################
|
||||
|
||||
async def get_available_work_units(self, db: AsyncSession) -> Sequence[WorkUnit]:
|
||||
async def get_available_work_units(self, db: AsyncSession) -> Sequence[str]:
|
||||
"""
|
||||
Get available work units that aren't being processed.
|
||||
Returns a list of WorkUnit objects.
|
||||
Returns a list of work unit keys.
|
||||
"""
|
||||
# Clean up stale work units
|
||||
five_minutes_ago = datetime.now(UTC) - timedelta(
|
||||
five_minutes_ago = datetime.now(timezone.utc) - timedelta(
|
||||
minutes=settings.DERIVER.STALE_SESSION_TIMEOUT_MINUTES
|
||||
)
|
||||
await db.execute(
|
||||
|
|
@ -160,64 +131,24 @@ class QueueManager:
|
|||
)
|
||||
)
|
||||
|
||||
# Create the JSON path expressions once to ensure they're identical in SELECT and GROUP BY
|
||||
sender_name_expr = models.QueueItem.payload["sender_name"].astext
|
||||
target_name_expr = models.QueueItem.payload["target_name"].astext
|
||||
task_type_expr = models.QueueItem.payload["task_type"].astext
|
||||
|
||||
# Get available work units by extracting sender_name, target_name, task_type from payload
|
||||
# We need to join with ActiveQueueSession to find units that aren't already being processed
|
||||
result = await db.execute(
|
||||
select(
|
||||
models.QueueItem.session_id,
|
||||
sender_name_expr.label("sender_name"),
|
||||
target_name_expr.label("target_name"),
|
||||
task_type_expr.label("task_type"),
|
||||
)
|
||||
query = (
|
||||
select(models.QueueItem.work_unit_key)
|
||||
.outerjoin(
|
||||
models.ActiveQueueSession,
|
||||
(models.QueueItem.session_id == models.ActiveQueueSession.session_id)
|
||||
& (
|
||||
(sender_name_expr == models.ActiveQueueSession.sender_name)
|
||||
| (
|
||||
sender_name_expr.is_(None)
|
||||
& models.ActiveQueueSession.sender_name.is_(None)
|
||||
)
|
||||
)
|
||||
& (
|
||||
(target_name_expr == models.ActiveQueueSession.target_name)
|
||||
| (
|
||||
target_name_expr.is_(None)
|
||||
& models.ActiveQueueSession.target_name.is_(None)
|
||||
)
|
||||
)
|
||||
& (task_type_expr == models.ActiveQueueSession.task_type),
|
||||
models.QueueItem.work_unit_key
|
||||
== models.ActiveQueueSession.work_unit_key,
|
||||
)
|
||||
.where(~models.QueueItem.processed)
|
||||
.where(
|
||||
models.ActiveQueueSession.id.is_(None)
|
||||
) # Only work units not in active_queue_sessions
|
||||
.group_by(
|
||||
models.QueueItem.session_id,
|
||||
sender_name_expr,
|
||||
target_name_expr,
|
||||
task_type_expr,
|
||||
)
|
||||
.limit(self.workers) # Process multiple work units in parallel
|
||||
.where(models.QueueItem.work_unit_key.isnot(None))
|
||||
.where(models.ActiveQueueSession.work_unit_key.is_(None))
|
||||
.distinct()
|
||||
.limit(self.workers)
|
||||
)
|
||||
|
||||
rows = result.fetchall()
|
||||
return [
|
||||
WorkUnit(
|
||||
session_id=row.session_id,
|
||||
sender_name=row.sender_name,
|
||||
target_name=row.target_name,
|
||||
task_type=row.task_type,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
result = await db.execute(query)
|
||||
return result.scalars().all()
|
||||
|
||||
async def polling_loop(self):
|
||||
async def polling_loop(self) -> None:
|
||||
"""Main polling loop to find and process new work units"""
|
||||
logger.debug("Starting polling loop")
|
||||
try:
|
||||
|
|
@ -242,13 +173,10 @@ class QueueManager:
|
|||
if new_work_units and not self.shutdown_event.is_set():
|
||||
for work_unit in new_work_units:
|
||||
try:
|
||||
# Try to claim the work unit
|
||||
# Try to claim the work unit using work_unit_key
|
||||
await db.execute(
|
||||
insert(models.ActiveQueueSession).values(
|
||||
session_id=work_unit.session_id,
|
||||
sender_name=work_unit.sender_name,
|
||||
target_name=work_unit.target_name,
|
||||
task_type=work_unit.task_type,
|
||||
work_unit_key=work_unit
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
|
@ -265,12 +193,14 @@ class QueueManager:
|
|||
self.process_work_unit(work_unit)
|
||||
)
|
||||
self.add_task(task)
|
||||
|
||||
except IntegrityError:
|
||||
# Rollback the failed transaction to clear the error state
|
||||
await db.rollback()
|
||||
logger.debug(
|
||||
f"Failed to claim work unit {work_unit}, already owned"
|
||||
f"Failed to claim work unit {work_unit}, already owned by another worker"
|
||||
)
|
||||
# If we couldn't claim any work units, avoid tight loop
|
||||
else:
|
||||
self.queue_empty_flag.set()
|
||||
await asyncio.sleep(
|
||||
|
|
@ -292,32 +222,34 @@ class QueueManager:
|
|||
######################
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def process_work_unit(self, work_unit: WorkUnit):
|
||||
"""Process all messages for a specific work unit"""
|
||||
logger.debug(f"Starting to process work unit {work_unit}")
|
||||
# Use the tracked_db dependency for transaction safety
|
||||
async def process_work_unit(self, work_unit_key: str):
|
||||
"""Process all messages for a specific work unit by routing to the correct handler."""
|
||||
logger.debug(f"Starting to process work unit {work_unit_key}")
|
||||
async with (
|
||||
self.semaphore,
|
||||
tracked_db("queue_process_work_unit") as db,
|
||||
): # Hold the semaphore for the entire work unit duration
|
||||
message_count = 0
|
||||
try:
|
||||
message_count = 0
|
||||
while not self.shutdown_event.is_set():
|
||||
message = await self.get_next_message(db, work_unit)
|
||||
message = await self.get_next_message(db, work_unit_key)
|
||||
if not message:
|
||||
logger.debug(f"No more messages for work unit {work_unit}")
|
||||
logger.debug(f"No more messages for work unit {work_unit_key}")
|
||||
|
||||
break
|
||||
|
||||
message_count += 1
|
||||
try:
|
||||
logger.info(
|
||||
f"Processing message {message.payload['message_id']} from work unit {work_unit}"
|
||||
f"Processing item for task type {message.task_type} with id {message.id} from work unit {work_unit_key}"
|
||||
)
|
||||
await process_item(message.task_type, message.payload)
|
||||
logger.debug(
|
||||
f"Successfully processed queue item for task type {message.task_type} with id {message.id}"
|
||||
)
|
||||
await process_item(message.payload)
|
||||
logger.debug(f"Successfully processed message {message.id}")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing message {message.id}: {str(e)}",
|
||||
f"Error processing queue item for task type {message.task_type} with id {message.id}: {str(e)}",
|
||||
exc_info=True,
|
||||
)
|
||||
if settings.SENTRY.ENABLED:
|
||||
|
|
@ -326,69 +258,83 @@ class QueueManager:
|
|||
# Prevent malformed messages from stalling queue indefinitely
|
||||
message.processed = True
|
||||
await db.commit()
|
||||
logger.debug(f"Marked message {message.id} as processed")
|
||||
|
||||
if self.shutdown_event.is_set():
|
||||
logger.debug(
|
||||
f"Shutdown requested, stopping processing for work unit {work_unit}"
|
||||
f"Shutdown requested, stopping processing for work unit {work_unit_key}"
|
||||
)
|
||||
break
|
||||
|
||||
# Update last_updated timestamp to show this work unit is still being processed
|
||||
await db.execute(
|
||||
update(models.ActiveQueueSession)
|
||||
.where(
|
||||
models.ActiveQueueSession.session_id
|
||||
== work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name
|
||||
== work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name
|
||||
== work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type == work_unit.task_type,
|
||||
)
|
||||
.where(models.ActiveQueueSession.work_unit_key == work_unit_key)
|
||||
.values(last_updated=func.now())
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
logger.debug(
|
||||
f"Completed processing work unit {work_unit}, processed {message_count} messages"
|
||||
f"Completed processing work unit {work_unit_key}, processed {message_count} messages"
|
||||
)
|
||||
finally:
|
||||
# Remove work unit from active_queue_sessions when done
|
||||
logger.debug(f"Removing work unit {work_unit} from active sessions")
|
||||
await db.execute(
|
||||
logger.debug(f"Removing work unit {work_unit_key} from active sessions")
|
||||
delete_result = await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.session_id == work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name == work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name == work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type == work_unit.task_type,
|
||||
models.ActiveQueueSession.work_unit_key == work_unit_key
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
self.untrack_work_unit(work_unit)
|
||||
|
||||
# Only publish webhook if we actually removed an active session
|
||||
if delete_result.rowcount > 0 and message_count > 0:
|
||||
try:
|
||||
from src.deriver.utils import parse_work_unit_key
|
||||
from src.webhooks.events import (
|
||||
QueueEmptyEvent,
|
||||
publish_webhook_event,
|
||||
)
|
||||
|
||||
parsed_key = parse_work_unit_key(work_unit_key)
|
||||
if parsed_key["task_type"] in ["representation", "summary"]:
|
||||
logger.info(
|
||||
f"Publishing queue.empty event for {work_unit_key}"
|
||||
)
|
||||
await publish_webhook_event(
|
||||
QueueEmptyEvent(
|
||||
workspace_id=parsed_key["workspace_name"],
|
||||
queue_type=parsed_key["task_type"],
|
||||
session_id=parsed_key["session_name"],
|
||||
sender_name=parsed_key["sender_name"],
|
||||
observer_name=parsed_key["target_name"],
|
||||
)
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
f"Skipping queue.empty event for webhook work unit {work_unit_key}"
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Error triggering queue_empty webhook")
|
||||
else:
|
||||
logger.debug(
|
||||
f"Work unit {work_unit_key} already cleaned up by another worker, skipping webhook"
|
||||
)
|
||||
|
||||
self.untrack_work_unit(work_unit_key)
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def get_next_message(self, db: AsyncSession, work_unit: WorkUnit):
|
||||
"""Get the next unprocessed message for a specific work unit"""
|
||||
async def get_next_message(
|
||||
self, db: AsyncSession, work_unit_key: str
|
||||
) -> QueueItem | None:
|
||||
"""Get the next unprocessed message for a specific work unit."""
|
||||
|
||||
query = (
|
||||
select(models.QueueItem)
|
||||
.where(models.QueueItem.session_id == work_unit.session_id)
|
||||
.where(models.QueueItem.payload["task_type"].astext == work_unit.task_type)
|
||||
.where(models.QueueItem.work_unit_key == work_unit_key)
|
||||
.where(~models.QueueItem.processed)
|
||||
.order_by(models.QueueItem.id)
|
||||
.with_for_update(skip_locked=True)
|
||||
.limit(1)
|
||||
)
|
||||
|
||||
# For summary tasks, sender_name and target_name don't exist in payload
|
||||
# For other tasks, filter by sender_name and target_name
|
||||
if work_unit.task_type != "summary":
|
||||
query = query.where(
|
||||
models.QueueItem.payload["sender_name"].astext == work_unit.sender_name
|
||||
).where(
|
||||
models.QueueItem.payload["target_name"].astext == work_unit.target_name
|
||||
)
|
||||
|
||||
result = await db.execute(query)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
|
|
|||
|
|
@ -7,17 +7,16 @@ from pydantic import BaseModel, ConfigDict
|
|||
class BasePayload(BaseModel):
|
||||
"""Base payload with common fields."""
|
||||
|
||||
workspace_name: str
|
||||
session_name: str
|
||||
message_id: int
|
||||
|
||||
model_config = ConfigDict(extra="forbid") # pyright: ignore
|
||||
model_config = ConfigDict(extra="forbid") # pyright: ignore[reportUnannotatedClassAttribute]
|
||||
|
||||
|
||||
class RepresentationPayload(BasePayload):
|
||||
"""Payload for representation tasks."""
|
||||
|
||||
task_type: Literal["representation"] = "representation"
|
||||
workspace_name: str
|
||||
session_name: str
|
||||
message_id: int
|
||||
content: str
|
||||
sender_name: str
|
||||
target_name: str
|
||||
|
|
@ -28,11 +27,34 @@ class SummaryPayload(BasePayload):
|
|||
"""Payload for summary tasks."""
|
||||
|
||||
task_type: Literal["summary"] = "summary"
|
||||
workspace_name: str
|
||||
session_name: str
|
||||
message_id: int
|
||||
message_seq_in_session: int
|
||||
|
||||
|
||||
# Union type for the actual payload
|
||||
class WebhookPayload(BasePayload):
|
||||
"""Payload for webhook delivery tasks."""
|
||||
|
||||
task_type: Literal["webhook"] = "webhook"
|
||||
workspace_name: str
|
||||
event_type: str
|
||||
data: dict[str, Any]
|
||||
|
||||
|
||||
# Union type for all possible queue payloads
|
||||
DeriverQueuePayload = RepresentationPayload | SummaryPayload
|
||||
QueuePayload = DeriverQueuePayload | WebhookPayload
|
||||
|
||||
|
||||
def create_webhook_payload(
|
||||
workspace_name: str,
|
||||
event_type: str,
|
||||
data: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
return WebhookPayload(
|
||||
workspace_name=workspace_name, event_type=event_type, data=data
|
||||
).model_dump(mode="json")
|
||||
|
||||
|
||||
def create_payload(
|
||||
|
|
@ -112,6 +134,7 @@ def create_payload(
|
|||
# Convert back to dict for compatibility with JSON serialization
|
||||
# mode='json' ensures datetime is converted to ISO string
|
||||
payload = validated_payload.model_dump(mode="json")
|
||||
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to create valid payload: {str(e)}") from e
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,67 @@
|
|||
from typing_extensions import Any, TypedDict
|
||||
|
||||
|
||||
class ParsedWorkUnit(TypedDict):
|
||||
task_type: str
|
||||
workspace_name: str
|
||||
session_name: str | None
|
||||
sender_name: str | None
|
||||
target_name: str | None
|
||||
|
||||
|
||||
def get_work_unit_key(task_type: str, payload: dict[str, Any]) -> str:
|
||||
"""
|
||||
Generate a work unit key for a given task type, workspace name, and event type.
|
||||
"""
|
||||
workspace_name = payload.get("workspace_name")
|
||||
if not workspace_name:
|
||||
raise ValueError("workspace_name is required to generate a work_unit_key")
|
||||
|
||||
if task_type in ["representation", "summary"]:
|
||||
sender_name = payload.get("sender_name", "None")
|
||||
target_name = payload.get("target_name", "None")
|
||||
session_name = payload.get("session_name", "None")
|
||||
return (
|
||||
f"{task_type}:{workspace_name}:{session_name}:{sender_name}:{target_name}"
|
||||
)
|
||||
|
||||
if task_type == "webhook":
|
||||
return f"webhook:{workspace_name}"
|
||||
|
||||
raise ValueError(f"Invalid task type: {task_type}")
|
||||
|
||||
|
||||
def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit:
|
||||
"""
|
||||
Parse a work unit key to extract its components.
|
||||
"""
|
||||
parts = work_unit_key.split(":")
|
||||
task_type = parts[0]
|
||||
|
||||
if task_type in ["representation", "summary"]:
|
||||
if len(parts) != 5:
|
||||
raise ValueError(
|
||||
f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
|
||||
)
|
||||
return {
|
||||
"task_type": task_type,
|
||||
"workspace_name": parts[1],
|
||||
"session_name": parts[2],
|
||||
"sender_name": parts[3],
|
||||
"target_name": parts[4],
|
||||
}
|
||||
|
||||
if task_type == "webhook":
|
||||
if len(parts) != 2:
|
||||
raise ValueError(
|
||||
f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
|
||||
)
|
||||
return {
|
||||
"task_type": task_type,
|
||||
"workspace_name": parts[1],
|
||||
"session_name": None,
|
||||
"sender_name": None,
|
||||
"target_name": None,
|
||||
}
|
||||
|
||||
raise ValueError(f"Invalid task type in work_unit_key: {task_type}")
|
||||
|
|
@ -24,6 +24,7 @@ from src.routers import (
|
|||
messages,
|
||||
peers,
|
||||
sessions,
|
||||
webhooks,
|
||||
workspaces,
|
||||
)
|
||||
from src.security import create_admin_jwt
|
||||
|
|
@ -101,6 +102,7 @@ if SENTRY_ENABLED:
|
|||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI):
|
||||
# Lifespan events are now handled by the respective services
|
||||
yield
|
||||
await engine.dispose()
|
||||
|
||||
|
|
@ -150,6 +152,7 @@ app.include_router(peers.router, prefix="/v2")
|
|||
app.include_router(sessions.router, prefix="/v2")
|
||||
app.include_router(messages.router, prefix="/v2")
|
||||
app.include_router(keys.router, prefix="/v2")
|
||||
app.include_router(webhooks.router, prefix="/v2")
|
||||
|
||||
|
||||
# Global exception handlers
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import datetime
|
||||
from logging import getLogger
|
||||
from typing import Any, final
|
||||
from typing import Any, Literal, final
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
|
@ -83,6 +83,7 @@ class Workspace(Base):
|
|||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True, unique=True)
|
||||
peers = relationship("Peer", back_populates="workspace")
|
||||
webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace")
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
|
|
@ -358,6 +359,9 @@ class Document(Base):
|
|||
)
|
||||
|
||||
|
||||
TaskType = Literal["webhook", "summary", "representation"]
|
||||
|
||||
|
||||
@final
|
||||
class QueueItem(Base):
|
||||
__tablename__: str = "queue"
|
||||
|
|
@ -367,6 +371,9 @@ class QueueItem(Base):
|
|||
session_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("sessions.id"), index=True, nullable=True
|
||||
)
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
|
||||
task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False)
|
||||
payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
|
||||
processed: Mapped[bool] = mapped_column(Boolean, default=False)
|
||||
|
||||
|
|
@ -376,25 +383,35 @@ class ActiveQueueSession(Base):
|
|||
__tablename__: str = "active_queue_sessions"
|
||||
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
session_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("sessions.id"), nullable=True
|
||||
)
|
||||
sender_name: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
target_name: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
task_type: Mapped[str] = mapped_column(TEXT)
|
||||
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True, index=True)
|
||||
|
||||
last_updated: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=func.now(), onupdate=func.now()
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"session_id",
|
||||
"sender_name",
|
||||
"target_name",
|
||||
"task_type",
|
||||
name="unique_active_queue_session",
|
||||
),
|
||||
|
||||
@final
|
||||
class WebhookEndpoint(Base):
|
||||
__tablename__: str = "webhook_endpoints"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True, nullable=False
|
||||
)
|
||||
url: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=func.now()
|
||||
)
|
||||
|
||||
workspace = relationship("Workspace", back_populates="webhook_endpoints")
|
||||
|
||||
__table_args__ = (
|
||||
CheckConstraint("length(url) <= 2048", name="webhook_endpoint_url_length"),
|
||||
Index("idx_webhook_endpoints_workspace_lookup", "workspace_name"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"WebhookEndpoint(id={self.id}, workspace_name={self.workspace_name}, url={self.url})"
|
||||
|
||||
|
||||
@final
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
import logging
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, Path
|
||||
from fastapi_pagination import Page
|
||||
from fastapi_pagination.ext.sqlalchemy import apaginate
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import schemas
|
||||
from src.config import settings
|
||||
from src.crud import webhook as crud
|
||||
from src.dependencies import db
|
||||
from src.exceptions import AuthenticationException, ConflictException
|
||||
from src.security import JWTParams, require_auth
|
||||
from src.webhooks.events import (
|
||||
TestEvent,
|
||||
publish_webhook_event,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/workspaces/{workspace_id}/webhooks",
|
||||
tags=["webhooks"],
|
||||
)
|
||||
|
||||
|
||||
@router.post("", response_model=schemas.WebhookEndpoint)
|
||||
async def get_or_create_webhook_endpoint(
|
||||
workspace_id: str = Path(..., description="Workspace ID"),
|
||||
webhook: schemas.WebhookEndpointCreate = Body(
|
||||
..., description="Webhook endpoint parameters"
|
||||
),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db: AsyncSession = db,
|
||||
) -> schemas.WebhookEndpoint:
|
||||
"""
|
||||
Get or create a webhook endpoint URL.
|
||||
"""
|
||||
if not jwt_params.ad and jwt_params.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
|
||||
try:
|
||||
return await crud.get_or_create_webhook_endpoint(
|
||||
db, workspace_id, webhook=webhook
|
||||
)
|
||||
except ValueError as e:
|
||||
raise ConflictException(
|
||||
f"Maximum number of webhook endpoints ({settings.WEBHOOK.MAX_WORKSPACE_LIMIT}) reached for this workspace."
|
||||
) from e
|
||||
|
||||
|
||||
@router.get("", response_model=Page[schemas.WebhookEndpoint])
|
||||
async def list_webhook_endpoints(
|
||||
workspace_id: str = Path(..., description="Workspace ID"),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db: AsyncSession = db,
|
||||
) -> Page[schemas.WebhookEndpoint]:
|
||||
"""
|
||||
List all webhook endpoints, optionally filtered by workspace.
|
||||
"""
|
||||
if not jwt_params.ad and jwt_params.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
|
||||
stmt = await crud.list_webhook_endpoints(db, workspace_id)
|
||||
return await apaginate(db, stmt)
|
||||
|
||||
|
||||
@router.delete("/{endpoint_id}", response_model=None)
|
||||
async def delete_webhook_endpoint(
|
||||
workspace_id: str = Path(..., description="Workspace ID"),
|
||||
endpoint_id: str = Path(..., description="Webhook endpoint ID"),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db: AsyncSession = db,
|
||||
) -> None:
|
||||
"""
|
||||
Delete a specific webhook endpoint.
|
||||
"""
|
||||
|
||||
if not jwt_params.ad and jwt_params.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
|
||||
await crud.delete_webhook_endpoint(db, workspace_id, endpoint_id)
|
||||
|
||||
|
||||
@router.get("/test")
|
||||
async def test_emit(
|
||||
workspace_id: str = Path(..., description="Workspace ID"),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
) -> None:
|
||||
"""
|
||||
Test publishing a webhook event.
|
||||
"""
|
||||
if not jwt_params.ad and jwt_params.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unable to publish test webhook")
|
||||
|
||||
event = TestEvent(workspace_id=workspace_id)
|
||||
await publish_webhook_event(event)
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
# pyright: reportUnannotatedClassAttribute=false # pyright: ignore
|
||||
import datetime
|
||||
import ipaddress
|
||||
from typing import Annotated, Any, Self
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import tiktoken
|
||||
from pydantic import (
|
||||
|
|
@ -8,6 +10,7 @@ from pydantic import (
|
|||
ConfigDict,
|
||||
Field,
|
||||
PrivateAttr,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
|
|
@ -351,3 +354,44 @@ class DeriverStatus(BaseModel):
|
|||
sessions: dict[str, SessionDeriverStatus] | None = Field(
|
||||
default=None, description="Per-session status when not filtered by session"
|
||||
)
|
||||
|
||||
|
||||
# Webhook endpoint schemas
|
||||
class WebhookEndpointBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class WebhookEndpointCreate(WebhookEndpointBase):
|
||||
url: str
|
||||
|
||||
@field_validator("url")
|
||||
@classmethod
|
||||
def validate_webhook_url(cls, v: str) -> str:
|
||||
parsed = urlparse(v)
|
||||
|
||||
if not all([parsed.scheme, parsed.netloc]):
|
||||
raise ValueError("Invalid URL format")
|
||||
|
||||
# Only allow HTTP/HTTPS
|
||||
if parsed.scheme not in ["http", "https"]:
|
||||
raise ValueError("Only HTTP and HTTPS URLs are allowed")
|
||||
|
||||
# Block private/internal addresses
|
||||
if parsed.hostname:
|
||||
try:
|
||||
ip_address = ipaddress.ip_address(parsed.hostname)
|
||||
if ip_address.is_private:
|
||||
raise ValueError("Private IP addresses are not allowed")
|
||||
except ValueError: # Not an IP address, might be a hostname
|
||||
pass
|
||||
|
||||
return v
|
||||
|
||||
|
||||
class WebhookEndpoint(WebhookEndpointBase):
|
||||
id: str
|
||||
workspace_name: str | None = Field(serialization_alias="workspace_id")
|
||||
url: str
|
||||
created_at: datetime.datetime
|
||||
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True) # pyright: ignore
|
||||
|
|
|
|||
|
|
@ -6,9 +6,9 @@ from __future__ import annotations
|
|||
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
||||
class ReasoningLevel(str, Enum):
|
||||
|
|
|
|||
|
|
@ -3,11 +3,11 @@ import datetime
|
|||
import logging
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import TypedDict
|
||||
|
||||
from mirascope import llm
|
||||
from sqlalchemy import update
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from src.config import settings
|
||||
from src.dependencies import tracked_db
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
# Webhooks
|
||||
|
||||
Webhooks are used to deliver event notifications to user-configured URLs.
|
||||
|
||||
## System Architecture
|
||||
|
||||
The webhooks system consists of several key components:
|
||||
|
||||
* **API Endpoints (`routers/webhooks.py`):** Provides endpoints for users to create, list, delete, and test their webhook subscriptions.
|
||||
* **Event Publishing (`webhooks/events.py`):** Defines the event types and allows us to publish new events to the processing queue.
|
||||
* **Webhook Delivery (`webhooks/webhook_delivery.py`):** Contains the logic for sending the webhook to the subscriber's URL.
|
||||
|
||||
## Event Flow
|
||||
|
||||
1. An event is triggered within the application by calling `publish_webhook_event` with a defined event payload.
|
||||
2. This function creates a `QueueItem` and stores it in the database.
|
||||
3. The `QueueManager` background process polls the database for new items.
|
||||
4. When a new event is found, it is passed to the `deliver_webhook` function.
|
||||
5. `deliver_webhook` fetches all subscriber URLs for the event's workspace, signs the payload with a secret key, and sends an HTTP POST request to each URL.
|
||||
|
||||
Note that the webhooks require the *deriver* process to be running to facilitate the delivery of the webhook.
|
||||
|
|
@ -0,0 +1,81 @@
|
|||
import logging
|
||||
from enum import Enum
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from src.dependencies import tracked_db
|
||||
from src.deriver.queue_payload import create_webhook_payload
|
||||
from src.deriver.utils import get_work_unit_key
|
||||
from src.models import QueueItem
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WebhookEventType(str, Enum):
|
||||
QUEUE_EMPTY = "queue.empty"
|
||||
TEST = "test.event"
|
||||
|
||||
|
||||
class BaseWebhookEvent(BaseModel):
|
||||
"""Base class for all webhook events."""
|
||||
|
||||
workspace_id: str
|
||||
|
||||
|
||||
class QueueEmptyEvent(BaseWebhookEvent):
|
||||
"""Webhook event for when a queue becomes empty."""
|
||||
|
||||
type: Literal[WebhookEventType.QUEUE_EMPTY] = WebhookEventType.QUEUE_EMPTY
|
||||
queue_type: str
|
||||
session_id: str | None = None
|
||||
sender_name: str | None = None
|
||||
observer_name: str | None = None
|
||||
|
||||
|
||||
class TestEvent(BaseWebhookEvent):
|
||||
"""Webhook event for testing."""
|
||||
|
||||
type: Literal[WebhookEventType.TEST] = WebhookEventType.TEST
|
||||
|
||||
|
||||
# Union type for all webhook events
|
||||
WebhookEvent = QueueEmptyEvent | TestEvent
|
||||
|
||||
|
||||
async def publish_webhook_event(event: WebhookEvent) -> None:
|
||||
"""
|
||||
Add a webhook event to our DB queue.
|
||||
|
||||
Args:
|
||||
event: The webhook event to publish.
|
||||
"""
|
||||
try:
|
||||
payload = create_webhook_payload(
|
||||
workspace_name=event.workspace_id,
|
||||
event_type=event.type.value,
|
||||
data=event.model_dump(mode="json", exclude={"type"}),
|
||||
)
|
||||
|
||||
async with tracked_db("publish_webhook_event") as db:
|
||||
queue_item = QueueItem(
|
||||
work_unit_key=get_work_unit_key(
|
||||
"webhook", {"workspace_name": event.workspace_id}
|
||||
),
|
||||
payload=payload,
|
||||
session_id=None,
|
||||
task_type="webhook",
|
||||
)
|
||||
db.add(queue_item)
|
||||
await db.commit()
|
||||
logger.debug(
|
||||
"Published webhook event '%s' for workspace '%s'",
|
||||
event.type,
|
||||
event.workspace_id,
|
||||
)
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Failed to publish webhook event %s",
|
||||
event.type,
|
||||
)
|
||||
|
|
@ -0,0 +1,103 @@
|
|||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src.config import settings
|
||||
from src.crud.webhook import list_webhook_endpoints
|
||||
from src.deriver.queue_payload import WebhookPayload
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def deliver_webhook(db: AsyncSession, payload: WebhookPayload) -> None:
|
||||
"""
|
||||
Deliver a single webhook event to its configured endpoints.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
try:
|
||||
webhook_urls = await _get_webhook_urls(db, payload.workspace_name)
|
||||
if not webhook_urls:
|
||||
logger.info(
|
||||
f"No webhook endpoints for workspace {payload.workspace_name}, skipping."
|
||||
)
|
||||
return
|
||||
|
||||
event_payload = {
|
||||
"type": payload.event_type,
|
||||
"data": payload.data,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
event_json = json.dumps(
|
||||
event_payload, separators=(",", ":"), sort_keys=True
|
||||
)
|
||||
|
||||
try:
|
||||
signature = _generate_webhook_signature(event_json)
|
||||
except ValueError:
|
||||
logger.exception("Failed to generate webhook signature")
|
||||
return
|
||||
|
||||
tasks = [
|
||||
client.post(
|
||||
url=url,
|
||||
content=event_json,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Honcho-Signature": signature,
|
||||
},
|
||||
)
|
||||
for url in webhook_urls
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for url, result in zip(webhook_urls, results, strict=False):
|
||||
if isinstance(result, httpx.Response):
|
||||
if 200 <= result.status_code < 300:
|
||||
logger.info(
|
||||
f"Successfully delivered webhook {payload.event_type} to {url}"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed delivery for {payload.event_type} to {url}. Status: {result.status_code}"
|
||||
)
|
||||
else:
|
||||
logger.error(
|
||||
f"Failed delivery for {payload.event_type} to {url}. Exception: {result}"
|
||||
)
|
||||
|
||||
except httpx.RequestError:
|
||||
logger.exception(f"Error sending webhook for {payload.workspace_name}.")
|
||||
except Exception:
|
||||
logger.exception("Unexpected error delivering webhook.")
|
||||
|
||||
|
||||
async def _get_webhook_urls(db: AsyncSession, workspace_name: str) -> list[str]:
|
||||
"""
|
||||
Get all webhook endpoint URLs for a workspace.
|
||||
"""
|
||||
try:
|
||||
endpoints = await list_webhook_endpoints(db, workspace_name)
|
||||
result = await db.execute(endpoints)
|
||||
return [endpoint.url for endpoint in result.scalars().all()]
|
||||
except Exception:
|
||||
logger.exception(f"Error fetching endpoints for {workspace_name}")
|
||||
return []
|
||||
|
||||
|
||||
def _generate_webhook_signature(payload: str) -> str:
|
||||
"""
|
||||
Generate HMAC-SHA256 signature for webhook payload using WEBHOOK_SECRET.
|
||||
"""
|
||||
webhook_secret = settings.WEBHOOK.SECRET
|
||||
if not webhook_secret:
|
||||
raise ValueError("WEBHOOK_SECRET not found - cannot sign webhook")
|
||||
|
||||
return hmac.new(
|
||||
webhook_secret.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256
|
||||
).hexdigest()
|
||||
|
|
@ -18,13 +18,14 @@ import os
|
|||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, TypedDict
|
||||
from typing import Any
|
||||
|
||||
import tiktoken
|
||||
from anthropic import AsyncAnthropic
|
||||
from dotenv import load_dotenv
|
||||
from honcho import Honcho
|
||||
from honcho.session import SessionPeerConfig
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from src import crud, models, schemas
|
||||
from src.deriver.queue_payload import create_payload
|
||||
from src.deriver.utils import get_work_unit_key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -18,7 +19,9 @@ def mock_deriver_process(monkeypatch: pytest.MonkeyPatch) -> Callable[..., Any]:
|
|||
"""Mock the deriver process_message method to avoid actual LLM calls"""
|
||||
from src.deriver.deriver import Deriver
|
||||
|
||||
async def mock_process_message(_self: Any, _payload: dict[str, Any]) -> None:
|
||||
async def mock_process_message(
|
||||
_self: Any, _task_type: str, _payload: dict[str, Any]
|
||||
) -> None:
|
||||
# Simulate processing without making actual LLM calls
|
||||
pass
|
||||
|
||||
|
|
@ -174,8 +177,14 @@ async def add_queue_items(
|
|||
"""Add queue items to the database and return them"""
|
||||
queue_items: list[models.QueueItem] = []
|
||||
for payload in payloads:
|
||||
# Generate work_unit_key from the payload
|
||||
task_type = payload.get("task_type", "unknown")
|
||||
work_unit_key = get_work_unit_key(task_type, payload)
|
||||
|
||||
queue_item = models.QueueItem(
|
||||
session_id=session_id,
|
||||
task_type=task_type,
|
||||
work_unit_key=work_unit_key,
|
||||
payload=payload,
|
||||
processed=False,
|
||||
)
|
||||
|
|
@ -249,17 +258,11 @@ async def create_active_queue_session(db_session: AsyncSession) -> Callable[...,
|
|||
"""Helper function to create active queue sessions for testing work unit tracking"""
|
||||
|
||||
async def _create_active_session(
|
||||
session_id: str,
|
||||
sender_name: str | None = None,
|
||||
target_name: str | None = None,
|
||||
task_type: str = "representation",
|
||||
work_unit_key: str,
|
||||
) -> models.ActiveQueueSession:
|
||||
"""Create an active queue session"""
|
||||
active_session = models.ActiveQueueSession(
|
||||
session_id=session_id,
|
||||
sender_name=sender_name,
|
||||
target_name=target_name,
|
||||
task_type=task_type,
|
||||
work_unit_key=work_unit_key,
|
||||
)
|
||||
db_session.add(active_session)
|
||||
await db_session.commit()
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@ from typing import Any
|
|||
import pytest
|
||||
|
||||
from src import models
|
||||
from src.deriver.queue_manager import WorkUnit
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -40,39 +39,43 @@ class TestDeriverProcessing:
|
|||
# The mock should be in place and return a predefined response
|
||||
# This ensures no actual LLM calls are made during testing
|
||||
|
||||
async def test_work_unit_creation(
|
||||
async def test_work_unit_key_generation(
|
||||
self,
|
||||
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
|
||||
):
|
||||
"""Test that WorkUnit objects can be created correctly"""
|
||||
"""Test that work unit keys are generated correctly"""
|
||||
from src.deriver.utils import get_work_unit_key
|
||||
|
||||
session, peers = sample_session_with_peers
|
||||
peer1, peer2, _ = peers
|
||||
|
||||
# Create a WorkUnit for representation task
|
||||
work_unit = WorkUnit(
|
||||
session_id=session.id,
|
||||
sender_name=peer1.name,
|
||||
target_name=peer2.name,
|
||||
task_type="representation",
|
||||
# Create a payload for representation task
|
||||
representation_payload = {
|
||||
"workspace_name": "workspace1",
|
||||
"session_name": session.name,
|
||||
"sender_name": peer1.name,
|
||||
"target_name": peer2.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
|
||||
# Generate work unit key for representation
|
||||
work_unit_key = get_work_unit_key("representation", representation_payload)
|
||||
expected_key = (
|
||||
f"representation:workspace1:{session.name}:{peer1.name}:{peer2.name}"
|
||||
)
|
||||
assert work_unit_key == expected_key
|
||||
|
||||
assert work_unit.session_id == session.id
|
||||
assert work_unit.sender_name == peer1.name
|
||||
assert work_unit.target_name == peer2.name
|
||||
assert work_unit.task_type == "representation"
|
||||
# Create a payload for summary task (sender_name and target_name should be None)
|
||||
summary_payload = {
|
||||
"workspace_name": "workspace1",
|
||||
"session_name": session.name,
|
||||
"task_type": "summary",
|
||||
}
|
||||
|
||||
# Create a WorkUnit for summary task (sender_name and target_name should be None)
|
||||
summary_work_unit = WorkUnit(
|
||||
session_id=session.id,
|
||||
sender_name=None,
|
||||
target_name=None,
|
||||
task_type="summary",
|
||||
)
|
||||
|
||||
assert summary_work_unit.session_id == session.id
|
||||
assert summary_work_unit.sender_name is None
|
||||
assert summary_work_unit.target_name is None
|
||||
assert summary_work_unit.task_type == "summary"
|
||||
# Generate work unit key for summary
|
||||
summary_work_unit_key = get_work_unit_key("summary", summary_payload)
|
||||
expected_summary_key = f"summary:workspace1:{session.name}:None:None"
|
||||
assert summary_work_unit_key == expected_summary_key
|
||||
|
||||
async def test_mock_queue_manager(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -70,16 +70,21 @@ class TestQueueOperations:
|
|||
session, peers = sample_session_with_peers
|
||||
peer1, peer2, _ = peers
|
||||
|
||||
# Create an active queue session
|
||||
active_session = await create_active_queue_session(
|
||||
session_id=session.id,
|
||||
sender_name=peer1.name,
|
||||
target_name=peer2.name,
|
||||
task_type="representation",
|
||||
# Create a work unit key for a representation task
|
||||
work_unit_key = (
|
||||
f"representation:workspace1:{session.name}:{peer1.name}:{peer2.name}"
|
||||
)
|
||||
|
||||
# Create an active queue session
|
||||
active_session = await create_active_queue_session(work_unit_key=work_unit_key)
|
||||
|
||||
assert active_session is not None
|
||||
assert active_session.session_id == session.id
|
||||
assert active_session.sender_name == peer1.name
|
||||
assert active_session.target_name == peer2.name
|
||||
assert active_session.task_type == "representation"
|
||||
assert active_session.work_unit_key == work_unit_key
|
||||
|
||||
# Verify the work unit key has the expected format
|
||||
parts = work_unit_key.split(":")
|
||||
assert parts[0] == "representation"
|
||||
assert parts[1] == "workspace1"
|
||||
assert parts[2] == session.name
|
||||
assert parts[3] == peer1.name
|
||||
assert parts[4] == peer2.name
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from src import models
|
||||
from src.deriver.consumer import process_item
|
||||
from src.deriver.queue_manager import QueueManager, WorkUnit
|
||||
from src.deriver.queue_manager import QueueManager
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -35,11 +35,11 @@ class TestQueueProcessing:
|
|||
|
||||
# Check that all work units have the expected structure
|
||||
for work_unit in work_units:
|
||||
assert isinstance(work_unit, WorkUnit)
|
||||
assert work_unit.task_type in ["representation", "summary"]
|
||||
assert isinstance(work_unit, str)
|
||||
assert work_unit.split(":")[0] in ["representation", "summary"]
|
||||
|
||||
# The test is mainly verifying that get_available_work_units works without errors
|
||||
# and returns properly structured WorkUnit objects
|
||||
# and returns properly structured work unit key strings
|
||||
|
||||
async def test_work_unit_claiming(
|
||||
self,
|
||||
|
|
@ -60,10 +60,7 @@ class TestQueueProcessing:
|
|||
# Claim a work unit by creating an ActiveQueueSession entry
|
||||
work_unit = work_units[0]
|
||||
active_session = models.ActiveQueueSession(
|
||||
session_id=work_unit.session_id,
|
||||
sender_name=work_unit.sender_name,
|
||||
target_name=work_unit.target_name,
|
||||
task_type=work_unit.task_type,
|
||||
work_unit_key=work_unit,
|
||||
)
|
||||
db_session.add(active_session)
|
||||
await db_session.commit()
|
||||
|
|
@ -108,45 +105,37 @@ class TestQueueProcessing:
|
|||
queue_item = sample_queue_items[0]
|
||||
|
||||
# This should not raise an exception since the deriver is mocked
|
||||
await process_item(queue_item.payload)
|
||||
await process_item(queue_item.task_type, queue_item.payload)
|
||||
|
||||
# The mock should have been called
|
||||
# Note: We can't easily verify this since we're mocking the class method directly
|
||||
# In a real test, we might want to mock at a different level
|
||||
|
||||
async def test_work_unit_string_representation(
|
||||
async def test_work_unit_key_format(
|
||||
self, sample_session_with_peers: tuple[models.Session, list[models.Peer]]
|
||||
):
|
||||
"""Test that WorkUnit string representation works correctly"""
|
||||
"""Test that work unit keys have the correct format"""
|
||||
session, peers = sample_session_with_peers
|
||||
peer1, peer2, _ = peers
|
||||
|
||||
# Create a representation work unit
|
||||
work_unit = WorkUnit(
|
||||
session_id=session.id,
|
||||
sender_name=peer1.name,
|
||||
target_name=peer2.name,
|
||||
task_type="representation",
|
||||
# Create a representation work unit key
|
||||
# Format: task_type:workspace:session:sender:target
|
||||
work_unit_key = (
|
||||
f"representation:workspace1:{session.name}:{peer1.name}:{peer2.name}"
|
||||
)
|
||||
|
||||
# Convert to string
|
||||
work_unit_str = str(work_unit)
|
||||
# Check that the key contains the expected information
|
||||
assert session.name in work_unit_key
|
||||
assert peer1.name in work_unit_key
|
||||
assert peer2.name in work_unit_key
|
||||
assert "representation" in work_unit_key
|
||||
assert "workspace1" in work_unit_key
|
||||
|
||||
# Check that the string contains the expected information
|
||||
assert session.id in work_unit_str
|
||||
assert peer1.name in work_unit_str
|
||||
assert peer2.name in work_unit_str
|
||||
assert "representation" in work_unit_str
|
||||
# Create a summary work unit key
|
||||
# Summary work units use None for sender/target
|
||||
summary_work_unit_key = f"summary:workspace1:{session.name}:None:None"
|
||||
|
||||
# Create a summary work unit
|
||||
summary_work_unit = WorkUnit(
|
||||
session_id=session.id,
|
||||
sender_name=None,
|
||||
target_name=None,
|
||||
task_type="summary",
|
||||
)
|
||||
|
||||
summary_str = str(summary_work_unit)
|
||||
assert session.id in summary_str
|
||||
assert "None" in summary_str
|
||||
assert "summary" in summary_str
|
||||
assert session.name in summary_work_unit_key
|
||||
assert "None" in summary_work_unit_key
|
||||
assert "summary" in summary_work_unit_key
|
||||
assert "workspace1" in summary_work_unit_key
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from fastapi.testclient import TestClient
|
|||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import models
|
||||
from src.deriver.utils import get_work_unit_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -144,18 +145,23 @@ class TestDeriverStatusEndpoint:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(session)
|
||||
# Add queue items
|
||||
queue_items = [
|
||||
models.QueueItem(
|
||||
queue_items: list[models.QueueItem] = []
|
||||
for _ in range(5):
|
||||
payload = {
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "representation",
|
||||
"workspace_name": workspace.name,
|
||||
"session_name": session.name,
|
||||
}
|
||||
queue_item = models.QueueItem(
|
||||
session_id=session.id,
|
||||
payload={
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "derive",
|
||||
},
|
||||
task_type="representation",
|
||||
work_unit_key=get_work_unit_key("representation", payload),
|
||||
payload=payload,
|
||||
processed=False,
|
||||
)
|
||||
for _ in range(5)
|
||||
]
|
||||
queue_items.append(queue_item)
|
||||
db_session.add_all(queue_items)
|
||||
await db_session.commit()
|
||||
# Test without parameters
|
||||
|
|
@ -213,18 +219,23 @@ class TestDeriverStatusEndpoint:
|
|||
await db_session.refresh(s)
|
||||
# Add queue items to different sessions
|
||||
for i, session in enumerate(sessions):
|
||||
queue_items = [
|
||||
models.QueueItem(
|
||||
queue_items: list[models.QueueItem] = []
|
||||
for _ in range(i + 1): # 1,2,3 items respectively
|
||||
payload = {
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "representation",
|
||||
"workspace_name": workspace.name,
|
||||
"session_name": session.name,
|
||||
}
|
||||
queue_item = models.QueueItem(
|
||||
session_id=session.id,
|
||||
payload={
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "derive",
|
||||
},
|
||||
task_type="representation",
|
||||
work_unit_key=get_work_unit_key("representation", payload),
|
||||
payload=payload,
|
||||
processed=False,
|
||||
)
|
||||
for _ in range(i + 1) # 1,2,3 items respectively
|
||||
]
|
||||
queue_items.append(queue_item)
|
||||
db_session.add_all(queue_items)
|
||||
await db_session.commit()
|
||||
response = client.get(
|
||||
|
|
@ -271,13 +282,18 @@ class TestDeriverStatusEndpoint:
|
|||
db_session.add(session)
|
||||
await db_session.commit()
|
||||
await db_session.refresh(session)
|
||||
payload = {
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "representation",
|
||||
"workspace_name": workspace.name,
|
||||
"session_name": session.name,
|
||||
}
|
||||
queue_item = models.QueueItem(
|
||||
session_id=session.id,
|
||||
payload={
|
||||
"sender_name": peer.name,
|
||||
"target_name": peer.name,
|
||||
"task_type": "derive",
|
||||
},
|
||||
task_type="representation",
|
||||
work_unit_key=get_work_unit_key("representation", payload),
|
||||
payload=payload,
|
||||
processed=False,
|
||||
)
|
||||
db_session.add(queue_item)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,285 @@
|
|||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from src.config import settings
|
||||
from src.models import Peer, Workspace
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_webhook_endpoint(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "http://example.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
response_json = response.json()
|
||||
assert response_json["url"] == "http://example.com/webhook"
|
||||
assert "id" in response_json
|
||||
assert response_json["workspace_id"] == test_workspace.name
|
||||
assert "created_at" in response_json
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_webhook_endpoint_invalid_url(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "192.168.1.1/webhook",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert "Invalid URL format" in error["msg"]
|
||||
assert error["type"] == "value_error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_webhook_endpoint_missing_workspace(client: TestClient):
|
||||
response = client.post(
|
||||
"/v2/workspaces/nonexistent-workspace/webhooks",
|
||||
json={
|
||||
"url": "http://example.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
assert response.json() == {"detail": "Workspace nonexistent-workspace not found"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_webhook_endpoints_with_data(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
initial_count = len(list_response.json()["items"])
|
||||
|
||||
# Create first endpoint
|
||||
response1 = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "http://example1.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response1.status_code == 200
|
||||
|
||||
# Create second endpoint
|
||||
response2 = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "http://example2.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response2.status_code == 200
|
||||
|
||||
# List endpoints
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
assert list_response.status_code == 200
|
||||
response_data = list_response.json()
|
||||
endpoints = response_data["items"]
|
||||
assert len(endpoints) == initial_count + 2
|
||||
|
||||
# Verify both endpoints are returned
|
||||
endpoint_urls = [ep["url"] for ep in endpoints]
|
||||
assert "http://example1.com/webhook" in endpoint_urls
|
||||
assert "http://example2.com/webhook" in endpoint_urls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_webhook_endpoints_missing_workspace(client: TestClient):
|
||||
response = client.get("/v2/workspaces/nonexistent-workspace/webhooks")
|
||||
assert response.status_code == 404
|
||||
assert response.json() == {"detail": "Workspace nonexistent-workspace not found"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_webhook_endpoint(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Create webhook endpoint
|
||||
create_response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "http://example.com/webhook",
|
||||
},
|
||||
)
|
||||
assert create_response.status_code == 200
|
||||
endpoint = create_response.json()
|
||||
endpoint_id = endpoint["id"]
|
||||
|
||||
# Delete webhook endpoint
|
||||
delete_response = client.delete(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks/{endpoint_id}"
|
||||
)
|
||||
assert delete_response.status_code == 200
|
||||
|
||||
# Verify endpoint is deleted
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
assert list_response.status_code == 200
|
||||
response_data = list_response.json()
|
||||
assert response_data["items"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_webhook_endpoint_not_found(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
response = client.delete(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks/nonexistent-id"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
assert "not found" in response.json()["detail"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_endpoints_per_workspace(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
"""Test that workspaces can have multiple webhook endpoints"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Create multiple endpoints
|
||||
urls = [
|
||||
"http://app1.com/webhook",
|
||||
"http://app2.com/webhook",
|
||||
"http://app3.com/webhook",
|
||||
]
|
||||
|
||||
initial_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
initial_count = len(initial_response.json()["items"])
|
||||
|
||||
created_endpoints: list[Any] = []
|
||||
for url in urls:
|
||||
response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={"url": url},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
created_endpoints.append(response.json())
|
||||
|
||||
# List all endpoints
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
assert list_response.status_code == 200
|
||||
response_data = list_response.json()
|
||||
endpoints = response_data["items"]
|
||||
assert len(endpoints) == initial_count + 3
|
||||
|
||||
# Verify all URLs are present
|
||||
returned_urls = [ep["url"] for ep in endpoints]
|
||||
for url in urls:
|
||||
assert url in returned_urls
|
||||
|
||||
# Delete one endpoint
|
||||
delete_response = client.delete(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks/{created_endpoints[0]['id']}"
|
||||
)
|
||||
assert delete_response.status_code == 200
|
||||
|
||||
# Verify only 2 endpoints remain
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
assert list_response.status_code == 200
|
||||
response_data = list_response.json()
|
||||
endpoints = response_data["items"]
|
||||
assert len(endpoints) == initial_count + 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_duplicate_webhook_endpoint(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
url = "http://example.com/duplicate"
|
||||
|
||||
# Create the endpoint first
|
||||
response1 = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={"url": url},
|
||||
)
|
||||
assert response1.status_code == 200
|
||||
|
||||
# Try to create it again
|
||||
response2 = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={"url": url},
|
||||
)
|
||||
assert response2.status_code == 200
|
||||
assert response1.json() == response2.json()
|
||||
|
||||
# Verify only one endpoint exists
|
||||
list_response = client.get(f"/v2/workspaces/{test_workspace.name}/webhooks")
|
||||
assert list_response.status_code == 200
|
||||
response_data = list_response.json()
|
||||
endpoints = response_data["items"]
|
||||
assert len(endpoints) == 1
|
||||
assert endpoints[0]["url"] == url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_webhook_endpoints_per_workspace(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
test_workspace, _ = sample_data
|
||||
limit: int = settings.WEBHOOK.MAX_WORKSPACE_LIMIT
|
||||
|
||||
# Create endpoints up to the limit
|
||||
for i in range(limit):
|
||||
response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": f"http://example{i}.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Try to create one more
|
||||
response = client.post(
|
||||
f"/v2/workspaces/{test_workspace.name}/webhooks",
|
||||
json={
|
||||
"url": "http://extra.com/webhook",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 409
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_endpoint_in_different_workspaces(
|
||||
client: TestClient, sample_data: tuple[Workspace, Peer]
|
||||
):
|
||||
ws1, _ = sample_data
|
||||
ws2_response = client.post("/v2/workspaces", json={"name": "workspace-2"})
|
||||
assert ws2_response.status_code == 200
|
||||
ws2 = ws2_response.json()
|
||||
|
||||
url = "http://example.com/shared"
|
||||
|
||||
# Create endpoint in workspace 1
|
||||
response1 = client.post(
|
||||
f"/v2/workspaces/{ws1.name}/webhooks",
|
||||
json={"url": url},
|
||||
)
|
||||
assert response1.status_code == 200
|
||||
|
||||
# Create same endpoint in workspace 2
|
||||
response2 = client.post(
|
||||
f"/v2/workspaces/{ws2['id']}/webhooks",
|
||||
json={"url": url},
|
||||
)
|
||||
assert response2.status_code == 200
|
||||
|
||||
# Verify they are different resources
|
||||
assert response1.json()["id"] != response2.json()["id"]
|
||||
assert response1.json()["workspace_id"] == ws1.name
|
||||
assert response2.json()["workspace_id"] == ws2["id"]
|
||||
Loading…
Reference in New Issue