feat(explorer): add on_completion hook to client

roaga · roaga · commit 5d6d25af66db · 2025-12-03T18:48:26.000-08:00
diff --git a/src/sentry/seer/endpoints/seer_rpc.py b/src/sentry/seer/endpoints/seer_rpc.py
@@ -86,6 +86,7 @@
     rpc_get_trace_for_transaction,
     rpc_get_transactions_for_project,
 )
+from sentry.seer.explorer.on_completion_hook import call_on_completion_hook
 from sentry.seer.explorer.tools import (
     execute_table_query,
     execute_timeseries_query,
@@ -1039,6 +1040,7 @@ def check_repository_integrations_status(*, repository_integrations: list[dict[s
     "get_trace_item_attributes": get_trace_item_attributes,
     "get_repository_definition": get_repository_definition,
     "call_custom_tool": call_custom_tool,
+    "call_on_completion_hook": call_on_completion_hook,
     "get_log_attributes_for_trace": get_log_attributes_for_trace,
     "get_metric_attributes_for_trace": get_metric_attributes_for_trace,
     #
diff --git a/src/sentry/seer/explorer/client.py b/src/sentry/seer/explorer/client.py
@@ -18,6 +18,10 @@
     poll_until_done,
 )
 from sentry.seer.explorer.custom_tool_utils import ExplorerTool, extract_tool_schema
+from sentry.seer.explorer.on_completion_hook import (
+    ExplorerOnCompletionHook,
+    extract_hook_definition,
+)
 from sentry.seer.models import SeerPermissionError
 from sentry.seer.signed_seer_api import sign_with_seer_secret
 from sentry.users.models.user import User
@@ -97,6 +101,22 @@ def execute(cls, organization, params: DeploymentStatusParams) -> str:
             custom_tools=[DeploymentStatusTool]
         )
         run_id = client.start_run("Check if payment-service is deployed in production")
+
+        # WITH ON-COMPLETION HOOK
+        from sentry.seer.explorer.on_completion_hook import ExplorerOnCompletionHook
+
+        class NotifyOnComplete(ExplorerOnCompletionHook):
+            @classmethod
+            def execute(cls, organization: Organization, run_id: int) -> None:
+                # Called when the agent completes (regardless of status)
+                send_notification(organization, f"Explorer run {run_id} completed")
+
+        client = SeerExplorerClient(
+            organization,
+            user,
+            on_completion=NotifyOnComplete
+        )
+        run_id = client.start_run("Analyze this issue")
     ```
 
         Args:
@@ -105,6 +125,7 @@ def execute(cls, organization, params: DeploymentStatusParams) -> str:
             category_key: Optional category key for filtering/grouping runs (e.g., "bug-fixer", "trace-analyzer"). Must be provided together with category_value. Makes it easy to retrieve runs for your feature later.
             category_value: Optional category value for filtering/grouping runs (e.g., issue ID, trace ID). Must be provided together with category_key. Makes it easy to retrieve a specific run for your feature later.
             custom_tools: Optional list of `ExplorerTool` classes to make available as tools to the agent. Each tool must inherit from ExplorerTool, define a params_model (Pydantic BaseModel), and implement execute(). Tools are automatically given access to the organization context. Tool classes must be module-level (not nested classes).
+            on_completion_hook: Optional `ExplorerOnCompletionHook` class to call when the agent completes. The hook's execute() method receives the organization and run ID. This is called whether or not the agent was successful. Hook classes must be module-level (not nested classes).
             intelligence_level: Optionally set the intelligence level of the agent. Higher intelligence gives better result quality at the cost of significantly higher latency and cost.
             is_interactive: Enable full interactive, human-like features of the agent. Only enable if you support *all* available interactions in Seer. An example use of this is the explorer chat in Sentry UI.
     """
@@ -116,12 +137,14 @@ def __init__(
         category_key: str | None = None,
         category_value: str | None = None,
         custom_tools: list[type[ExplorerTool[Any]]] | None = None,
+        on_completion_hook: type[ExplorerOnCompletionHook] | None = None,
         intelligence_level: Literal["low", "medium", "high"] = "medium",
         is_interactive: bool = False,
     ):
         self.organization = organization
         self.user = user
         self.custom_tools = custom_tools or []
+        self.on_completion_hook = on_completion_hook
         self.intelligence_level = intelligence_level
         self.category_key = category_key
         self.category_value = category_value
@@ -188,6 +211,10 @@ def start_run(
                 extract_tool_schema(tool).dict() for tool in self.custom_tools
             ]
 
+        # Add on-completion hook if provided
+        if self.on_completion_hook:
+            payload["on_completion_hook"] = extract_hook_definition(self.on_completion_hook).dict()
+
         if self.category_key and self.category_value:
             payload["category_key"] = self.category_key
             payload["category_value"] = self.category_value
diff --git a/src/sentry/seer/explorer/on_completion_hook.py b/src/sentry/seer/explorer/on_completion_hook.py
@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import importlib
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+from sentry.models.organization import Organization
+
+
+class OnCompletionHookDefinition(BaseModel):
+    """Definition of an on-completion hook to pass to Seer."""
+
+    module_path: str
+
+
+class ExplorerOnCompletionHook(ABC):
+    """Base class for Explorer on-completion hooks.
+
+    Hooks are called when an Explorer agent run completes (regardless of status).
+
+    Example:
+        class MyCompletionHook(ExplorerOnCompletionHook):
+            @classmethod
+            def execute(cls, organization: Organization, run_id: int) -> None:
+                # Do something when the run completes
+                notify_user(organization, run_id)
+
+        # Pass to client
+        client = SeerExplorerClient(
+            organization,
+            user,
+            on_completion_hook=MyCompletionHook
+        )
+    """
+
+    @classmethod
+    @abstractmethod
+    def execute(cls, organization: Organization, run_id: int) -> None:
+        """Execute the hook when the agent completes.
+
+        Args:
+            organization: The organization context
+            run_id: The ID of the completed run
+        """
+        ...
+
+    @classmethod
+    def get_module_path(cls) -> str:
+        """Get the full module path for this hook class."""
+        if not hasattr(cls, "__module__") or not hasattr(cls, "__name__"):
+            raise ValueError(f"Hook class {cls} must have __module__ and __name__ attributes")
+        if not cls.__module__ or not cls.__name__:
+            raise ValueError(f"Hook class {cls} has empty __module__ or __name__")
+        return f"{cls.__module__}.{cls.__name__}"
+
+
+def extract_hook_definition(
+    hook_class: type[ExplorerOnCompletionHook],
+) -> OnCompletionHookDefinition:
+    """Extract hook definition from an ExplorerOnCompletionHook class."""
+    # Enforce module-level classes only (no nested classes)
+    if "." in hook_class.__qualname__:
+        raise ValueError(
+            f"Hook class {hook_class.__name__} must be a module-level class. "
+            f"Nested classes are not supported. (qualname: {hook_class.__qualname__})"
+        )
+
+    return OnCompletionHookDefinition(module_path=hook_class.get_module_path())
+
+
+def call_on_completion_hook(
+    *,
+    module_path: str,
+    organization_id: int,
+    run_id: int,
+    allowed_prefixes: tuple[str, ...] = ("sentry.",),
+) -> None:
+    """Dynamically import and call an on-completion hook class.
+
+    Args:
+        module_path: Full module path to the hook class (e.g., "sentry.api.MyHook")
+        organization_id: Organization ID to load and pass to the hook
+        run_id: The run ID that completed
+        allowed_prefixes: Tuple of allowed module path prefixes for security
+    """
+    # Only allow imports from approved package prefixes
+    if not any(module_path.startswith(prefix) for prefix in allowed_prefixes):
+        raise ValueError(
+            f"Module path must start with one of {allowed_prefixes}, got: {module_path}"
+        )
+
+    # Load the organization
+    try:
+        organization = Organization.objects.get(id=organization_id)
+    except Organization.DoesNotExist:
+        raise ValueError(f"Organization with id {organization_id} does not exist")
+
+    # Split module path and class name
+    parts = module_path.rsplit(".", 1)
+    if len(parts) != 2:
+        raise ValueError(f"Invalid module path: {module_path}")
+
+    module_name, class_name = parts
+
+    # Import the hook class
+    try:
+        module = importlib.import_module(module_name)
+        hook_class = getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        raise ValueError(f"Could not import {module_path}: {e}")
+
+    # Validate it's an ExplorerOnCompletionHook subclass
+    if not isinstance(hook_class, type) or not issubclass(hook_class, ExplorerOnCompletionHook):
+        raise ValueError(
+            f"{module_path} must be a class that inherits from ExplorerOnCompletionHook"
+        )
+
+    # Execute the hook
+    hook_class.execute(organization, run_id)
diff --git a/tests/sentry/seer/explorer/test_on_completion_hook.py b/tests/sentry/seer/explorer/test_on_completion_hook.py
@@ -0,0 +1,87 @@
+import pytest
+
+from sentry.models.organization import Organization
+from sentry.seer.explorer.on_completion_hook import (
+    ExplorerOnCompletionHook,
+    OnCompletionHookDefinition,
+    call_on_completion_hook,
+    extract_hook_definition,
+)
+from sentry.testutils.cases import TestCase
+
+
+# Test hook class (defined at module level as required)
+class SampleCompletionHook(ExplorerOnCompletionHook):
+    @classmethod
+    def execute(cls, organization: Organization, run_id: int) -> None:
+        # Side effect: write to organization options so we can verify execution
+        organization.update_option("test_hook_run_id", run_id)
+
+
+class OnCompletionHookTest(TestCase):
+    def test_extract_hook_definition(self):
+        """Test extracting hook definition from a hook class."""
+        hook_def = extract_hook_definition(SampleCompletionHook)
+
+        assert isinstance(hook_def, OnCompletionHookDefinition)
+        assert hook_def.module_path.endswith("test_on_completion_hook.SampleCompletionHook")
+
+    def test_extract_hook_definition_nested_class_raises(self):
+        """Test that nested classes are rejected."""
+
+        class OuterClass:
+            class NestedHook(ExplorerOnCompletionHook):
+                @classmethod
+                def execute(cls, organization: Organization, run_id: int) -> None:
+                    pass
+
+        with pytest.raises(ValueError) as cm:
+            extract_hook_definition(OuterClass.NestedHook)
+        assert "module-level class" in str(cm.value)
+
+    def test_call_on_completion_hook_success(self):
+        """Test calling a completion hook successfully."""
+        module_path = "tests.sentry.seer.explorer.test_on_completion_hook.SampleCompletionHook"
+
+        call_on_completion_hook(
+            module_path=module_path,
+            organization_id=self.organization.id,
+            run_id=12345,
+            allowed_prefixes=("sentry.", "tests.sentry."),
+        )
+
+        # Verify side effect: hook wrote run_id to organization options
+        assert self.organization.get_option("test_hook_run_id") == 12345
+
+    def test_call_on_completion_hook_security_restriction(self):
+        """Test that module path must start with allowed prefix."""
+        with pytest.raises(ValueError) as cm:
+            call_on_completion_hook(
+                module_path="malicious.module.Hook",
+                organization_id=self.organization.id,
+                run_id=123,
+                allowed_prefixes=("sentry.",),
+            )
+        assert "must start with one of" in str(cm.value)
+
+    def test_call_on_completion_hook_invalid_module(self):
+        """Test calling a non-existent hook module."""
+        with pytest.raises(ValueError) as cm:
+            call_on_completion_hook(
+                module_path="sentry.nonexistent.module.Hook",
+                organization_id=self.organization.id,
+                run_id=123,
+            )
+        assert "Could not import" in str(cm.value)
+
+    def test_call_on_completion_hook_not_a_hook_class(self):
+        """Test calling something that isn't an ExplorerOnCompletionHook."""
+        # BaseModel is importable but not an ExplorerOnCompletionHook
+        with pytest.raises(ValueError) as cm:
+            call_on_completion_hook(
+                module_path="pydantic.BaseModel",
+                organization_id=self.organization.id,
+                run_id=123,
+                allowed_prefixes=("pydantic.",),
+            )
+        assert "must be a class that inherits from ExplorerOnCompletionHook" in str(cm.value)
