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:
Rajat Ahuja 2025-08-06 17:52:35 -04:00 committed by GitHub
parent 7b174dd34b
commit 3bea3da169
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
31 changed files with 2587 additions and 1494 deletions

View File

@ -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)

View File

@ -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 = [

View File

@ -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__":

View File

@ -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:

View File

@ -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",

View File

@ -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)

115
src/crud/webhook.py Normal file
View File

@ -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")

View File

@ -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:

View File

@ -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)

View File

@ -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(

View File

@ -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",
}

View File

@ -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()

View File

@ -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

67
src/deriver/utils.py Normal file
View File

@ -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}")

View File

@ -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

View File

@ -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

97
src/routers/webhooks.py Normal file
View File

@ -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)

View File

@ -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

View File

@ -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):

View File

@ -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

21
src/webhooks/README.md Normal file
View File

@ -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.

81
src/webhooks/events.py Normal file
View File

@ -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,
)

View File

@ -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()

View File

@ -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()

View File

@ -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()

View File

@ -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,

View File

@ -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

View File

@ -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

View File

@ -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)

View File

@ -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"]

2404
uv.lock

File diff suppressed because it is too large Load Diff