diff --git a/posthog/temporal/ai/__init__.py b/posthog/temporal/ai/__init__.py index 720d3a9122dc..45f107b263a1 100644 --- a/posthog/temporal/ai/__init__.py +++ b/posthog/temporal/ai/__init__.py @@ -15,6 +15,7 @@ from posthog.temporal.ai.slack_app.posthog_code_slack_mention import PostHogCodeSlackMentionWorkflow from posthog.temporal.ai.slack_app.posthog_code_slack_mention_command import PostHogCodeSlackMentionCommandWorkflow from posthog.temporal.ai.slack_app.posthog_slack_inbox_onboarding import PostHogSlackInboxOnboardingWorkflow +from posthog.temporal.ai.slack_app.slack_app_mention import SlackAppMentionWorkflow from .llm_traces_summaries.summarize_traces import ( SummarizeLLMTracesInputs, @@ -35,6 +36,7 @@ # workflows. POSTHOG_CODE_SLACK_WORKFLOWS = [ PostHogCodeSlackMentionWorkflow, + SlackAppMentionWorkflow, PostHogCodeSlackMentionCommandWorkflow, PostHogCodeSlackTerminateTaskWorkflow, PostHogSlackInboxOnboardingWorkflow, diff --git a/posthog/temporal/ai/slack_app/__init__.py b/posthog/temporal/ai/slack_app/__init__.py index 456852b300f6..ffdcebbdb005 100644 --- a/posthog/temporal/ai/slack_app/__init__.py +++ b/posthog/temporal/ai/slack_app/__init__.py @@ -23,6 +23,7 @@ forward_posthog_code_followup_activity, handle_posthog_code_rules_command_activity, handle_posthog_code_slack_mention_command_activity, + mark_slack_app_message_processing_activity, post_posthog_code_authorship_timeout_activity, post_posthog_code_internal_error_activity, post_posthog_code_no_repos_activity, @@ -34,12 +35,16 @@ run_posthog_slack_inbox_onboarding_activity, ) from posthog.temporal.ai.slack_app.types import ( + SLACK_APP_PROCESSING_REACTION, + SLACK_APP_QUEUED_REACTION, + MarkSlackAppMessageProcessingInput, PostHogCodeRepoCascadeOutcome, PostHogCodeRulesCommandResult, PostHogCodeSlackMentionCommandResult, PostHogCodeSlackMentionCommandWorkflowInputs, PostHogCodeSlackMentionWorkflowInputs, PostHogSlackInboxOnboardingInputs, + SlackAppMentionWorkflowInputs, SlackRepoSelectionOutcome, ) @@ -64,13 +69,17 @@ handle_posthog_code_rules_command_activity, create_posthog_code_routing_rule_activity, handle_posthog_code_slack_mention_command_activity, + mark_slack_app_message_processing_activity, run_posthog_slack_inbox_onboarding_activity, ] __all__ = [ "CLASSIFIER_THREAD_HISTORY_MESSAGES", + "MarkSlackAppMessageProcessingInput", "POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE", "POSTHOG_CODE_SLACK_RULES_ADD_PICKER_GUIDANCE", + "SLACK_APP_PROCESSING_REACTION", + "SLACK_APP_QUEUED_REACTION", "PostHogCodeRepoCascadeOutcome", "PostHogCodeRulesCommandResult", "PostHogCodeSlackMentionCommandResult", @@ -78,6 +87,7 @@ "PostHogCodeSlackMentionWorkflowInputs", "PostHogSlackInboxOnboardingInputs", "SLACK_APP_ACTIVITIES", + "SlackAppMentionWorkflowInputs", "SlackRepoSelectionOutcome", "block_posthog_code_task_if_no_personal_github_activity", "cascade_posthog_code_repository_activity", @@ -93,6 +103,7 @@ "forward_posthog_code_followup_activity", "handle_posthog_code_rules_command_activity", "handle_posthog_code_slack_mention_command_activity", + "mark_slack_app_message_processing_activity", "post_posthog_code_authorship_timeout_activity", "post_posthog_code_internal_error_activity", "post_posthog_code_no_repos_activity", diff --git a/posthog/temporal/ai/slack_app/activities/__init__.py b/posthog/temporal/ai/slack_app/activities/__init__.py index 02dc8ba6b76c..1be05794807d 100644 --- a/posthog/temporal/ai/slack_app/activities/__init__.py +++ b/posthog/temporal/ai/slack_app/activities/__init__.py @@ -10,6 +10,7 @@ POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE, POSTHOG_CODE_SLACK_RULES_ADD_PICKER_GUIDANCE, block_posthog_code_task_if_no_personal_github_activity, + mark_slack_app_message_processing_activity, post_posthog_code_authorship_timeout_activity, post_posthog_code_internal_error_activity, post_posthog_code_no_repos_activity, @@ -60,6 +61,7 @@ "forward_posthog_code_followup_activity", "handle_posthog_code_rules_command_activity", "handle_posthog_code_slack_mention_command_activity", + "mark_slack_app_message_processing_activity", "post_posthog_code_authorship_timeout_activity", "post_posthog_code_internal_error_activity", "post_posthog_code_no_repos_activity", diff --git a/posthog/temporal/ai/slack_app/activities/messaging.py b/posthog/temporal/ai/slack_app/activities/messaging.py index ef571b9d3b53..cf370ab687f9 100644 --- a/posthog/temporal/ai/slack_app/activities/messaging.py +++ b/posthog/temporal/ai/slack_app/activities/messaging.py @@ -4,7 +4,11 @@ import structlog from temporalio import activity -from posthog.temporal.ai.slack_app.types import PostHogCodeSlackMentionWorkflowInputs, coerce_mention_workflow_inputs +from posthog.temporal.ai.slack_app.types import ( + MarkSlackAppMessageProcessingInput, + PostHogCodeSlackMentionWorkflowInputs, + coerce_mention_workflow_inputs, +) from posthog.temporal.common.utils import close_db_connections if TYPE_CHECKING: @@ -428,3 +432,35 @@ def post_posthog_code_internal_error_activity( thread_ts=thread_ts, text="Sorry, I hit an internal error while processing that request. Please try again.", ) + + +@activity.defn +@close_db_connections +def mark_slack_app_message_processing_activity(input: MarkSlackAppMessageProcessingInput) -> None: + """Swap the queued :hourglass: reaction for :eyes: when the conversation + queue starts processing a message. + + Purely cosmetic UX feedback: never raises, so a Slack hiccup can't stall + the conversation queue behind retries of a reaction. + """ + from posthog.models.integration import Integration, SlackIntegration + from posthog.temporal.ai.slack_app.helpers import swap_reaction + from posthog.temporal.ai.slack_app.types import SLACK_APP_PROCESSING_REACTION, SLACK_APP_QUEUED_REACTION + + try: + integration = Integration.objects.get( + id=input.integration_id, + kind="slack", + integration_id=input.slack_team_id, + ) + slack = SlackIntegration(integration) + swap_reaction( + slack.client, input.channel, input.message_ts, SLACK_APP_QUEUED_REACTION, SLACK_APP_PROCESSING_REACTION + ) + except Exception as e: + logger.warning( + "slack_app_processing_reaction_failed", + channel=input.channel, + message_ts=input.message_ts, + error=str(e), + ) diff --git a/posthog/temporal/ai/slack_app/activities/task_creation.py b/posthog/temporal/ai/slack_app/activities/task_creation.py index d845c76caa8f..f0e87a5aac18 100644 --- a/posthog/temporal/ai/slack_app/activities/task_creation.py +++ b/posthog/temporal/ai/slack_app/activities/task_creation.py @@ -563,12 +563,14 @@ def forward_posthog_code_followup_activity( slack = SlackIntegration(integration) followup_user_text_prefix: str | None = None + actor_user = mapping.task.created_by if slack_user_id != mapping.mentioning_slack_user_id: - # The follow-up is from a different Slack user than the one who started the - # thread. Try to resolve them to a PostHog user with access to the same team - # — if so, let them participate; the message is still relayed in the original - # author's name (their sandbox token, their identity to the agent), with the - # actual sender's name prefixed onto the text so the agent sees who spoke. + # Follow-up from someone other than the original mentioner. Resolve them to a + # PostHog user with access to the same team; if they qualify, they participate + # under their own identity — the sandbox JWT and the PostHog MCP OAuth token + # are both rebound to them via send_refresh_session below, so their actions + # (insights, dashboards, etc.) attribute to them rather than the task creator. + # The actor's name is still prefixed onto the text so the agent sees who spoke. resolved = resolve_slack_user(slack, integration, slack_user_id, channel, thread_ts) if not resolved: logger.info( @@ -584,6 +586,7 @@ def forward_posthog_code_followup_activity( # into the LLM-forwarded prefix when both name and slack_email are absent. actor_name = resolved.user.get_full_name() or resolved.slack_email or resolved.user.email followup_user_text_prefix = f"{actor_name}: " + actor_user = resolved.user logger.info( "posthog_code_followup_cross_user_authorized", channel=channel, @@ -682,14 +685,38 @@ def forward_posthog_code_followup_activity( if user_message_ts: safe_react(slack.client, channel, user_message_ts, "eyes") + # Per-message identity applies only to queue-dispatched messages + # (slack-app-queue-workflow flag); the legacy per-message workflow keeps + # every credential bound to the task creator, exactly as before. + identity_user = actor_user if inputs.per_message_identity else mapping.task.created_by + auth_token = None - created_by = mapping.task.created_by - if created_by and created_by.id: - distinct_id = created_by.distinct_id or f"user_{created_by.id}" + if identity_user and identity_user.id: + distinct_id = identity_user.distinct_id or f"user_{identity_user.id}" auth_token = tasks_facade.create_sandbox_connection_token( - task_run.id, user_id=created_by.id, distinct_id=distinct_id + task_run.id, user_id=identity_user.id, distinct_id=distinct_id ) + # Rebind the sandbox to the message's actor *before* sending it, so + # this turn's actions attribute to whoever actually spoke, not + # whoever spoke last. The tasks layer owns which credentials that + # covers, tracks the sandbox's current identity, bypasses its refresh + # rate limits on any transition (including switching *back* to the + # task creator), and is best-effort by contract — a rebind failure is + # logged there and never blocks the message. The belt-and-braces + # except covers facade-level surprises for the same reason: deliver + # the message under the previous identity rather than dropping it. + if inputs.per_message_identity: + try: + tasks_facade.rebind_sandbox_identity_for_user(task_run.id, identity_user.id, auth_token=auth_token) + except Exception: + logger.exception( + "slack_app_followup_identity_rebind_failed", + channel=channel, + thread_ts=thread_ts, + actor_user_id=identity_user.id, + ) + result = tasks_facade.send_user_message(task_run.id, user_text, auth_token=auth_token, timeout=90) if not result.success and result.retryable and result.status_code != 504: result = tasks_facade.send_user_message(task_run.id, user_text, auth_token=auth_token, timeout=90) diff --git a/posthog/temporal/ai/slack_app/helpers.py b/posthog/temporal/ai/slack_app/helpers/__init__.py similarity index 85% rename from posthog/temporal/ai/slack_app/helpers.py rename to posthog/temporal/ai/slack_app/helpers/__init__.py index 3f5606961a07..cc347ce33f3a 100644 --- a/posthog/temporal/ai/slack_app/helpers.py +++ b/posthog/temporal/ai/slack_app/helpers/__init__.py @@ -62,3 +62,12 @@ def safe_react(client: Any, channel: str, timestamp: str, name: str) -> None: pass else: raise + + +def swap_reaction(client: Any, channel: str, timestamp: str, remove: str, add: str) -> None: + """Replace one reaction with another; a missing old reaction is a no-op.""" + try: + client.reactions_remove(channel=channel, timestamp=timestamp, name=remove) + except Exception: + pass + safe_react(client, channel, timestamp, add) diff --git a/posthog/temporal/ai/slack_app/helpers/process_mention_message.py b/posthog/temporal/ai/slack_app/helpers/process_mention_message.py new file mode 100644 index 000000000000..72a21de6925e --- /dev/null +++ b/posthog/temporal/ai/slack_app/helpers/process_mention_message.py @@ -0,0 +1,406 @@ +"""Shared per-message mention orchestration. + +This is workflow-side code: it runs inside a Temporal workflow and must stay +deterministic — no I/O, no Django, everything external goes through +activities. Both the per-message ``PostHogCodeSlackMentionWorkflow`` and the +per-conversation ``SlackAppMentionWorkflow`` drive their messages through +``process_mention_message``. Activity-side utilities live in this package's +``__init__``. +""" + +from datetime import timedelta +from typing import Any + +from temporalio import workflow +from temporalio.common import RetryPolicy + +from posthog.temporal.ai.slack_app import ( + POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE, + PostHogCodeSlackMentionWorkflowInputs, + block_posthog_code_task_if_no_personal_github_activity, + cascade_posthog_code_repository_activity, + classify_posthog_code_task_needs_repo_activity, + classify_untagged_followup_activity, + collect_posthog_code_thread_messages_activity, + create_posthog_code_task_for_repo_activity, + discover_posthog_code_repository_via_agent_activity, + enforce_posthog_code_billing_quota_activity, + forward_posthog_code_followup_activity, + post_posthog_code_authorship_timeout_activity, + post_posthog_code_internal_error_activity, + post_posthog_code_picker_timeout_activity, + post_posthog_code_repo_picker_activity, + resolve_posthog_code_authorship_activity, + resolve_posthog_code_slack_user_activity, +) + +POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS = 10 * 60 +POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES = 15 + + +class MentionSignalState: + """Interactive-signal state for one mention's processing. + + Owned by whichever workflow class hosts the processing (the per-message + mention workflow or the per-conversation queue workflow); its signal + handlers write here and ``process_mention_message`` waits on it. The queue + workflow resets it between messages, so a click on a stale picker from an + earlier message can, in rare cases, resolve the currently pending one — + accepted limitation. + """ + + def __init__(self) -> None: + self.selected_repo: str | None + self.repo_selection_resolved: bool + self.authorship_resolved: bool + self.reset() + + def reset(self) -> None: + self.selected_repo = None + self.repo_selection_resolved = False + self.authorship_resolved = False + + def resolve_repo_selection(self, repository: str | None) -> None: + """First click wins; later picker clicks for the same wait are ignored.""" + if not self.repo_selection_resolved: + self.repo_selection_resolved = True + self.selected_repo = repository + + def confirm_authorship(self) -> None: + self.authorship_resolved = True + + +class MentionSignalHandlersMixin: + """The interactive Slack signals, shared by both workflow classes. + + The webhook signals by name against an untyped workflow handle, so the + handler names are a cross-class contract — inheriting them from one place + keeps the two dispatch modes from drifting apart. + """ + + def __init__(self) -> None: + super().__init__() + self._signals = MentionSignalState() + + @workflow.signal + async def repo_selected(self, repository: str) -> None: + self._signals.resolve_repo_selection(repository) + + @workflow.signal + async def no_repo_needed(self) -> None: + self._signals.resolve_repo_selection(None) + + @workflow.signal + async def authorship_confirmed(self) -> None: + self._signals.confirm_authorship() + + +async def _resolve_authorship( + inputs: PostHogCodeSlackMentionWorkflowInputs, + signals: MentionSignalState, + channel: str, + thread_ts: str, + slack_user_id: str, + user_id: int, + repository: str, +) -> bool: + """Return True if the workflow must stop (blocked or timed out); False to proceed.""" + status = await _execute_posthog_code_activity( + resolve_posthog_code_authorship_activity, + inputs, + channel, + thread_ts, + slack_user_id, + user_id, + workflow.info().workflow_id, + repository, + ) + if status == "proceed": + return False + if status == "awaiting_confirmation": + try: + await workflow.wait_condition( + lambda: signals.authorship_resolved, + timeout=timedelta(minutes=POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES), + ) + except TimeoutError: + await _execute_posthog_code_activity( + post_posthog_code_authorship_timeout_activity, inputs, channel, thread_ts + ) + return True + return False + return True + + +async def process_mention_message( + inputs: PostHogCodeSlackMentionWorkflowInputs, + signals: MentionSignalState, +) -> None: + """Full per-message orchestration: quota gate -> followup classify/forward -> + thread collect -> repo cascade/picker -> authorship gate -> create task. + + Catches its own exceptions (posts an internal-error reply); never raises, + so a queue workflow calling it in a loop survives a poisoned message. + """ + event = inputs.event + channel = event.get("channel") + thread_ts = event.get("thread_ts") or event.get("ts") + slack_user_id = event.get("user") + + if not channel or not thread_ts or not slack_user_id: + return + + try: + # Gate every workflow entry on the team's AI-credits quota before any + # other activity runs. Webhook-level short-circuit catches the common + # case (see products/slack_app/backend/api.py); this is the defense in + # depth that also covers replays, manual workflow starts, and the race + # where the webhook saw "not limited" but Redis flipped before we got + # here. + blocked = await _execute_posthog_code_activity( + enforce_posthog_code_billing_quota_activity, + inputs, + channel, + thread_ts, + slack_user_id, + ) + if blocked: + return + + # Untagged thread replies face the Haiku classifier before any + # forward. The webhook handler punted on this so its 3-second ack + # budget stays unencumbered; here we run it under Temporal's retry + # policy. Drop on chitchat or any failure (default-deny). + if inputs.untagged_followup: + should_forward = await _execute_posthog_code_activity( + classify_untagged_followup_activity, + inputs, + channel, + thread_ts, + slack_user_id, + event.get("text", ""), + ) + if not should_forward: + return + + followup_handled = await _execute_posthog_code_activity( + forward_posthog_code_followup_activity, + inputs, + channel, + thread_ts, + slack_user_id, + event.get("text", ""), + event.get("ts"), + ) + if followup_handled: + return + + # Untagged thread replies must not fall through to the new-task path. + # The user never @mentioned us — they only typed in a thread that + # used to have an active task. If the mapping is gone by the time we + # got here, the right behaviour is to do nothing. + if inputs.untagged_followup: + return + + # New starts carry ``user_id`` from routing-time resolution and skip + # the activity. Legacy histories started before the field existed + # deserialize with ``user_id=None`` and replay through the activity so + # the recorded command stream still matches. Drop this fallback (and + # make ``user_id`` required on inputs) once the workflow history + # retention window has elapsed. + if inputs.user_id is not None: + user_id = inputs.user_id + else: + user_id = await _execute_posthog_code_activity( + resolve_posthog_code_slack_user_activity, inputs, channel, thread_ts, slack_user_id + ) + if not user_id: + return + + thread_messages = await _execute_posthog_code_activity( + collect_posthog_code_thread_messages_activity, + inputs, + channel, + thread_ts, + ) + if not thread_messages: + return + + repository: str | None + # Set only on the ambiguous path that runs the discovery sandbox + repo_research_task_id: str | None = None + repo_research_run_id: str | None = None + + cascade = await _execute_posthog_code_activity( + cascade_posthog_code_repository_activity, + inputs, + event.get("text", ""), + user_id, + ) + + if cascade.mode == "auto": + repository = cascade.repository + elif cascade.mode == "no_repo": + # Cascade only emits `no_repo` when neither the team nor the + # mentioning user has any GitHub install. Classify first so + # non-coding asks ("how do I configure retention?") still + # answer with no repo; coding asks surface the connect-personal- + # GitHub prompt instead of silently no-op'ing. + repository = None + needs_repo = await _execute_posthog_code_activity( + classify_posthog_code_task_needs_repo_activity, + event.get("text", ""), + thread_messages, + ) + if needs_repo: + blocked = await _execute_posthog_code_activity( + block_posthog_code_task_if_no_personal_github_activity, + inputs, + channel, + thread_ts, + user_id, + ) + if blocked: + return + elif cascade.mode == "needs_user_github": + # Team has GitHub, but the mentioning user hasn't connected their + # personal install. Fire the gate so they get the Connect button + # instead of a silently no-repo task. + await _execute_posthog_code_activity( + block_posthog_code_task_if_no_personal_github_activity, + inputs, + channel, + thread_ts, + user_id, + ) + return + else: + # Multiple candidates and no explicit mention. Cheap Haiku + # check first to skip the agent entirely for analytics/config + # questions; otherwise hand off to the discovery agent. + needs_repo = await _execute_posthog_code_activity( + classify_posthog_code_task_needs_repo_activity, + event.get("text", ""), + thread_messages, + ) + if not needs_repo: + repository = None + else: + outcome = await _execute_posthog_code_agent_activity( + discover_posthog_code_repository_via_agent_activity, + inputs, + channel, + event, + thread_messages, + user_id, + ) + repo_research_task_id = outcome.repo_research_task_id + repo_research_run_id = outcome.repo_research_run_id + + if outcome.status == "found": + repository = outcome.repository + elif outcome.status == "no_match": + repository = None + else: + # Agent crashed/timed out/hallucinated — italicize its reason + # above the picker guidance so the user sees why. + picker_guidance = f"_{outcome.reason}_\n\n{POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE}" + await _execute_posthog_code_activity( + post_posthog_code_repo_picker_activity, + inputs, + channel, + thread_ts, + slack_user_id, + event, + workflow.info().workflow_id, + picker_guidance, + True, + user_id, + ) + try: + await workflow.wait_condition( + lambda: signals.repo_selection_resolved, + timeout=timedelta(minutes=POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES), + ) + except TimeoutError: + await _execute_posthog_code_activity( + post_posthog_code_picker_timeout_activity, inputs, channel, thread_ts + ) + return + repository = signals.selected_repo + if repository: + if workflow.patched("posthog-code-authorship-confirm-2026-06"): + if await _resolve_authorship(inputs, signals, channel, thread_ts, slack_user_id, user_id, repository): + return + elif await _gate_on_personal_github(inputs, channel, thread_ts, user_id): + return + await _execute_posthog_code_activity( + create_posthog_code_task_for_repo_activity, + inputs, + channel, + thread_ts, + slack_user_id, + user_id, + event, + thread_messages, + repository, + repo_research_task_id, + repo_research_run_id, + ) + except Exception as exc: + workflow.logger.exception( + "posthog_code_workflow_unhandled_exception", + extra={ + "channel": channel, + "thread_ts": thread_ts, + "error": str(exc), + "error_type": type(exc).__name__, + }, + ) + await _execute_posthog_code_activity( + post_posthog_code_internal_error_activity, + inputs, + channel, + thread_ts, + ) + + +async def _gate_on_personal_github( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + user_id: int, +) -> bool: + """Return True when the workflow must abort because the mentioner has no personal GitHub.""" + return await _execute_posthog_code_activity( + block_posthog_code_task_if_no_personal_github_activity, + inputs, + channel, + thread_ts, + user_id, + ) + + +async def _execute_posthog_code_activity(activity_fn: Any, *args: Any) -> Any: + return await workflow.execute_activity( + activity_fn, + args=args, + start_to_close_timeout=timedelta(seconds=POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS), + retry_policy=RetryPolicy(maximum_attempts=3), + ) + + +async def _execute_posthog_code_agent_activity(activity_fn: Any, *args: Any) -> Any: + """Wrapper for the discovery-agent activity. + + No retries: a hung agent shouldn't block the Slack thread for tens of + minutes — the activity catches its own exceptions and returns + `status='failed'` so the workflow falls through to the picker. + """ + return await workflow.execute_activity( + activity_fn, + args=args, + start_to_close_timeout=timedelta(seconds=POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS), + heartbeat_timeout=timedelta(minutes=5), + retry_policy=RetryPolicy(maximum_attempts=1), + ) diff --git a/posthog/temporal/ai/slack_app/posthog_code_slack_mention.py b/posthog/temporal/ai/slack_app/posthog_code_slack_mention.py index 882ca077f2c7..16b9546ba05a 100644 --- a/posthog/temporal/ai/slack_app/posthog_code_slack_mention.py +++ b/posthog/temporal/ai/slack_app/posthog_code_slack_mention.py @@ -1,361 +1,29 @@ # Workflows in this module run on the max-ai temporal task queue. import json -from datetime import timedelta -from typing import Any from temporalio import workflow -from temporalio.common import RetryPolicy -from posthog.temporal.ai.slack_app import ( - POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE, - PostHogCodeSlackMentionWorkflowInputs, - block_posthog_code_task_if_no_personal_github_activity, - cascade_posthog_code_repository_activity, - classify_posthog_code_task_needs_repo_activity, - classify_untagged_followup_activity, - collect_posthog_code_thread_messages_activity, - create_posthog_code_task_for_repo_activity, - discover_posthog_code_repository_via_agent_activity, - enforce_posthog_code_billing_quota_activity, - forward_posthog_code_followup_activity, - post_posthog_code_authorship_timeout_activity, - post_posthog_code_internal_error_activity, - post_posthog_code_picker_timeout_activity, - post_posthog_code_repo_picker_activity, - resolve_posthog_code_authorship_activity, - resolve_posthog_code_slack_user_activity, +from posthog.temporal.ai.slack_app.helpers.process_mention_message import ( + MentionSignalHandlersMixin, + process_mention_message, ) +from posthog.temporal.ai.slack_app.types import PostHogCodeSlackMentionWorkflowInputs from posthog.temporal.common.base import PostHogWorkflow -POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS = 10 * 60 -POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES = 15 - @workflow.defn(name="posthog-code-slack-mention-processing") -class PostHogCodeSlackMentionWorkflow(PostHogWorkflow): - def __init__(self) -> None: - self._selected_repo: str | None = None - self._repo_selection_resolved = False - self._authorship_resolved = False - - @workflow.signal - async def repo_selected(self, repository: str) -> None: - if not self._repo_selection_resolved: - self._repo_selection_resolved = True - self._selected_repo = repository - - @workflow.signal - async def no_repo_needed(self) -> None: - if not self._repo_selection_resolved: - self._repo_selection_resolved = True - self._selected_repo = None - - @workflow.signal - async def authorship_confirmed(self) -> None: - self._authorship_resolved = True +class PostHogCodeSlackMentionWorkflow(MentionSignalHandlersMixin, PostHogWorkflow): + """One workflow per message — the pre-queue dispatch mode. When the + ``slack-app-queue-workflow`` flag is on, dispatch goes to the + per-conversation ``SlackAppMentionWorkflow`` instead; both drive messages + through the shared ``process_mention_message`` orchestration. + """ @staticmethod def parse_inputs(inputs: list[str]) -> PostHogCodeSlackMentionWorkflowInputs: loaded = json.loads(inputs[0]) return PostHogCodeSlackMentionWorkflowInputs(**loaded) - async def _resolve_authorship( - self, - inputs: PostHogCodeSlackMentionWorkflowInputs, - channel: str, - thread_ts: str, - slack_user_id: str, - user_id: int, - repository: str, - ) -> bool: - """Return True if the workflow must stop (blocked or timed out); False to proceed.""" - status = await _execute_posthog_code_activity( - resolve_posthog_code_authorship_activity, - inputs, - channel, - thread_ts, - slack_user_id, - user_id, - workflow.info().workflow_id, - repository, - ) - if status == "proceed": - return False - if status == "awaiting_confirmation": - try: - await workflow.wait_condition( - lambda: self._authorship_resolved, - timeout=timedelta(minutes=POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES), - ) - except TimeoutError: - await _execute_posthog_code_activity( - post_posthog_code_authorship_timeout_activity, inputs, channel, thread_ts - ) - return True - return False - return True - @workflow.run async def run(self, inputs: PostHogCodeSlackMentionWorkflowInputs) -> None: - event = inputs.event - channel = event.get("channel") - thread_ts = event.get("thread_ts") or event.get("ts") - slack_user_id = event.get("user") - - if not channel or not thread_ts or not slack_user_id: - return - - try: - # Gate every workflow entry on the team's AI-credits quota before any - # other activity runs. Webhook-level short-circuit catches the common - # case (see products/slack_app/backend/api.py); this is the defense in - # depth that also covers replays, manual workflow starts, and the race - # where the webhook saw "not limited" but Redis flipped before we got - # here. - blocked = await _execute_posthog_code_activity( - enforce_posthog_code_billing_quota_activity, - inputs, - channel, - thread_ts, - slack_user_id, - ) - if blocked: - return - - # Untagged thread replies face the Haiku classifier before any - # forward. The webhook handler punted on this so its 3-second ack - # budget stays unencumbered; here we run it under Temporal's retry - # policy. Drop on chitchat or any failure (default-deny). - if inputs.untagged_followup: - should_forward = await _execute_posthog_code_activity( - classify_untagged_followup_activity, - inputs, - channel, - thread_ts, - slack_user_id, - event.get("text", ""), - ) - if not should_forward: - return - - followup_handled = await _execute_posthog_code_activity( - forward_posthog_code_followup_activity, - inputs, - channel, - thread_ts, - slack_user_id, - event.get("text", ""), - event.get("ts"), - ) - if followup_handled: - return - - # Untagged thread replies must not fall through to the new-task path. - # The user never @mentioned us — they only typed in a thread that - # used to have an active task. If the mapping is gone by the time we - # got here, the right behaviour is to do nothing. - if inputs.untagged_followup: - return - - # New starts carry ``user_id`` from routing-time resolution and skip - # the activity. Legacy histories started before the field existed - # deserialize with ``user_id=None`` and replay through the activity so - # the recorded command stream still matches. Drop this fallback (and - # make ``user_id`` required on inputs) once the workflow history - # retention window has elapsed. - if inputs.user_id is not None: - user_id = inputs.user_id - else: - user_id = await _execute_posthog_code_activity( - resolve_posthog_code_slack_user_activity, inputs, channel, thread_ts, slack_user_id - ) - if not user_id: - return - - thread_messages = await _execute_posthog_code_activity( - collect_posthog_code_thread_messages_activity, - inputs, - channel, - thread_ts, - ) - if not thread_messages: - return - - repository: str | None - # Set only on the ambiguous path that runs the discovery sandbox - repo_research_task_id: str | None = None - repo_research_run_id: str | None = None - - cascade = await _execute_posthog_code_activity( - cascade_posthog_code_repository_activity, - inputs, - event.get("text", ""), - user_id, - ) - - if cascade.mode == "auto": - repository = cascade.repository - elif cascade.mode == "no_repo": - # Cascade only emits `no_repo` when neither the team nor the - # mentioning user has any GitHub install. Classify first so - # non-coding asks ("how do I configure retention?") still - # answer with no repo; coding asks surface the connect-personal- - # GitHub prompt instead of silently no-op'ing. - repository = None - needs_repo = await _execute_posthog_code_activity( - classify_posthog_code_task_needs_repo_activity, - event.get("text", ""), - thread_messages, - ) - if needs_repo: - blocked = await _execute_posthog_code_activity( - block_posthog_code_task_if_no_personal_github_activity, - inputs, - channel, - thread_ts, - user_id, - ) - if blocked: - return - elif cascade.mode == "needs_user_github": - # Team has GitHub, but the mentioning user hasn't connected their - # personal install. Fire the gate so they get the Connect button - # instead of a silently no-repo task. - await _execute_posthog_code_activity( - block_posthog_code_task_if_no_personal_github_activity, - inputs, - channel, - thread_ts, - user_id, - ) - return - else: - # Multiple candidates and no explicit mention. Cheap Haiku - # check first to skip the agent entirely for analytics/config - # questions; otherwise hand off to the discovery agent. - needs_repo = await _execute_posthog_code_activity( - classify_posthog_code_task_needs_repo_activity, - event.get("text", ""), - thread_messages, - ) - if not needs_repo: - repository = None - else: - outcome = await _execute_posthog_code_agent_activity( - discover_posthog_code_repository_via_agent_activity, - inputs, - channel, - event, - thread_messages, - user_id, - ) - repo_research_task_id = outcome.repo_research_task_id - repo_research_run_id = outcome.repo_research_run_id - - if outcome.status == "found": - repository = outcome.repository - elif outcome.status == "no_match": - repository = None - else: - # Agent crashed/timed out/hallucinated — italicize its reason - # above the picker guidance so the user sees why. - picker_guidance = f"_{outcome.reason}_\n\n{POSTHOG_CODE_SLACK_MENTION_PICKER_GUIDANCE}" - await _execute_posthog_code_activity( - post_posthog_code_repo_picker_activity, - inputs, - channel, - thread_ts, - slack_user_id, - event, - workflow.info().workflow_id, - picker_guidance, - True, - user_id, - ) - try: - await workflow.wait_condition( - lambda: self._repo_selection_resolved, - timeout=timedelta(minutes=POSTHOG_CODE_SLACK_PICKER_TIMEOUT_MINUTES), - ) - except TimeoutError: - await _execute_posthog_code_activity( - post_posthog_code_picker_timeout_activity, inputs, channel, thread_ts - ) - return - repository = self._selected_repo - if repository: - if workflow.patched("posthog-code-authorship-confirm-2026-06"): - if await self._resolve_authorship(inputs, channel, thread_ts, slack_user_id, user_id, repository): - return - elif await _gate_on_personal_github(inputs, channel, thread_ts, user_id): - return - await _execute_posthog_code_activity( - create_posthog_code_task_for_repo_activity, - inputs, - channel, - thread_ts, - slack_user_id, - user_id, - event, - thread_messages, - repository, - repo_research_task_id, - repo_research_run_id, - ) - except Exception as exc: - workflow.logger.exception( - "posthog_code_workflow_unhandled_exception", - extra={ - "channel": channel, - "thread_ts": thread_ts, - "error": str(exc), - "error_type": type(exc).__name__, - }, - ) - await _execute_posthog_code_activity( - post_posthog_code_internal_error_activity, - inputs, - channel, - thread_ts, - ) - - -async def _gate_on_personal_github( - inputs: PostHogCodeSlackMentionWorkflowInputs, - channel: str, - thread_ts: str, - user_id: int, -) -> bool: - """Return True when the workflow must abort because the mentioner has no personal GitHub.""" - return await _execute_posthog_code_activity( - block_posthog_code_task_if_no_personal_github_activity, - inputs, - channel, - thread_ts, - user_id, - ) - - -async def _execute_posthog_code_activity(activity_fn: Any, *args: Any) -> Any: - return await workflow.execute_activity( - activity_fn, - args=args, - start_to_close_timeout=timedelta(seconds=POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS), - retry_policy=RetryPolicy(maximum_attempts=3), - ) - - -async def _execute_posthog_code_agent_activity(activity_fn: Any, *args: Any) -> Any: - """Wrapper for the discovery-agent activity. - - No retries: a hung agent shouldn't block the Slack thread for tens of - minutes — the activity catches its own exceptions and returns - `status='failed'` so the workflow falls through to the picker. - """ - return await workflow.execute_activity( - activity_fn, - args=args, - start_to_close_timeout=timedelta(seconds=POSTHOG_CODE_SLACK_MENTION_TIMEOUT_SECONDS), - heartbeat_timeout=timedelta(minutes=5), - retry_policy=RetryPolicy(maximum_attempts=1), - ) + await process_mention_message(inputs, self._signals) diff --git a/posthog/temporal/ai/slack_app/slack_app_mention.py b/posthog/temporal/ai/slack_app/slack_app_mention.py new file mode 100644 index 000000000000..8da1b910f888 --- /dev/null +++ b/posthog/temporal/ai/slack_app/slack_app_mention.py @@ -0,0 +1,152 @@ +# Workflows in this module run on the max-ai temporal task queue. +import json +from datetime import timedelta + +from temporalio import exceptions, workflow +from temporalio.common import RetryPolicy + +from posthog.temporal.ai.slack_app import derive_mention_workflow_id, mark_slack_app_message_processing_activity +from posthog.temporal.ai.slack_app.helpers.process_mention_message import ( + MentionSignalHandlersMixin, + process_mention_message, +) +from posthog.temporal.ai.slack_app.types import ( + MarkSlackAppMessageProcessingInput, + PostHogCodeSlackMentionWorkflowInputs, + SlackAppMentionWorkflowInputs, +) +from posthog.temporal.common.base import PostHogWorkflow + +SLACK_APP_MENTION_IDLE_TIMEOUT_SECONDS = 30 +# Dedup keys carried across continue_as_new. Bounded so the carry-over payload +# stays small; old keys only matter for Slack retries, which arrive within +# minutes of the original event. +SLACK_APP_MENTION_MAX_PROCESSED_KEYS = 200 + + +def derive_slack_app_mention_workflow_id(inputs: PostHogCodeSlackMentionWorkflowInputs) -> str | None: + """Conversation-scoped workflow ID: one per thread (or DM thread). + + Anchored on the thread root ts — the same anchor the rest of the pipeline + uses — so every message in a conversation resolves to the same workflow. + Returns None when the event lacks a channel or ts; callers fall back to + the per-message workflow. + """ + event = inputs.event + channel = event.get("channel") + anchor = event.get("thread_ts") or event.get("ts") + if not channel or not anchor: + return None + return f"slack-app-mention-{inputs.slack_team_id}:{channel}:{anchor}" + + +@workflow.defn(name="slack-app-mention") +class SlackAppMentionWorkflow(MentionSignalHandlersMixin, PostHogWorkflow): + """Per-conversation queue over the mention pipeline. + + The per-message ``PostHogCodeSlackMentionWorkflow`` races when several + messages land in one thread. This workflow serializes them: the webhook + signal-with-starts one instance per conversation, messages queue up as + ``new_message`` signals, and the loop feeds them one at a time through the + shared ``process_mention_message`` orchestration. After the idle timeout + with an empty queue the workflow completes; the next message simply starts + a fresh instance, which finds the conversation's task via + ``SlackThreadTaskMapping`` and continues in followup mode. + """ + + def __init__(self) -> None: + super().__init__() + self._queue: list[PostHogCodeSlackMentionWorkflowInputs] = [] + # Insertion-ordered so the continue_as_new carry-over stays deterministic + # across replays (set iteration order is not). + self._seen_keys: dict[str, None] = {} + + @workflow.signal + async def new_message(self, message: PostHogCodeSlackMentionWorkflowInputs) -> None: + # The per-message workflow ID doubles as the message's identity, so the + # dedup rule stays single-sourced with derive_mention_workflow_id. + key = derive_mention_workflow_id(message) + if key in self._seen_keys: + return + self._seen_keys[key] = None + self._queue.append(message) + + @staticmethod + def parse_inputs(inputs: list[str]) -> SlackAppMentionWorkflowInputs: + loaded = json.loads(inputs[0]) + loaded["pending_messages"] = [ + PostHogCodeSlackMentionWorkflowInputs(**message) for message in loaded.get("pending_messages", []) + ] + return SlackAppMentionWorkflowInputs(**loaded) + + async def _mark_processing(self, message: PostHogCodeSlackMentionWorkflowInputs) -> None: + """Swap the dispatch-time :hourglass: for :eyes: as the message leaves + the queue. Mentions and DMs only — untagged thread followups get no + dispatch reaction, and most are dropped by the chitchat classifier. + """ + channel = message.event.get("channel") + message_ts = message.event.get("ts") + if message.untagged_followup or not channel or not message_ts: + return + try: + await workflow.execute_activity( + mark_slack_app_message_processing_activity, + MarkSlackAppMessageProcessingInput( + integration_id=message.integration_id, + slack_team_id=message.slack_team_id, + channel=channel, + message_ts=message_ts, + ), + start_to_close_timeout=timedelta(seconds=10), + retry_policy=RetryPolicy(maximum_attempts=1), + ) + except exceptions.ActivityError: + # The activity swallows Slack errors itself; this only fires on a + # timeout. Cosmetic either way — never stall the queue for it. + workflow.logger.warning("slack_app_processing_reaction_activity_failed") + + @workflow.run + async def run(self, inputs: SlackAppMentionWorkflowInputs) -> None: + for key in inputs.processed_event_keys: + self._seen_keys[key] = None + for message in inputs.pending_messages: + self._seen_keys.setdefault(derive_mention_workflow_id(message), None) + self._queue.append(message) + + while True: + try: + await workflow.wait_condition( + lambda: bool(self._queue), + timeout=timedelta(seconds=SLACK_APP_MENTION_IDLE_TIMEOUT_SECONDS), + ) + except TimeoutError: + # Idle exit. Drain in-flight signal handlers and re-check: a + # message landing between the timeout and the return must not + # be dropped. A signal racing the completion command makes the + # server fail the workflow task and replay, landing here again + # with the queue non-empty. Once the workflow has fully closed, + # the webhook's signal-with-start spawns a fresh instance. + await workflow.wait_condition(workflow.all_handlers_finished) + if self._queue: + continue + return + + message = self._queue.pop(0) + # Being processed by the queue workflow is what per-message + # identity *means* — stamping here (not at dispatch) makes it + # impossible for a future dispatcher into this workflow to forget. + message.per_message_identity = True + self._signals.reset() + await self._mark_processing(message) + # Never raises: internal errors are posted back to the thread, so + # one poisoned message can't wedge the conversation's queue. + await process_mention_message(message, self._signals) + + if workflow.info().is_continue_as_new_suggested(): + await workflow.wait_condition(workflow.all_handlers_finished) + workflow.continue_as_new( + SlackAppMentionWorkflowInputs( + pending_messages=list(self._queue), + processed_event_keys=list(self._seen_keys)[-SLACK_APP_MENTION_MAX_PROCESSED_KEYS:], + ) + ) diff --git a/posthog/temporal/ai/slack_app/types.py b/posthog/temporal/ai/slack_app/types.py index 663177287a5c..209f3f351c3c 100644 --- a/posthog/temporal/ai/slack_app/types.py +++ b/posthog/temporal/ai/slack_app/types.py @@ -5,9 +5,11 @@ creating an import cycle with the workflow modules. """ -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields from typing import Any, Literal +from pydantic import BaseModel + @dataclass class PostHogSlackInboxOnboardingInputs: @@ -33,6 +35,12 @@ class PostHogCodeSlackMentionWorkflowInputs: # cleanup), we must NOT fall through to the new-task path — the user never # tagged us, so kicking off a brand-new agent run would be wrong. untagged_followup: bool = False + # True when the message was dispatched to the per-conversation queue + # workflow (slack-app-queue-workflow flag). Gates per-message identity: + # the sandbox JWT and credential rebinds follow the message's actor + # instead of the task creator. The legacy per-message workflow leaves + # this False and keeps creator-bound credentials throughout. + per_message_identity: bool = False def coerce_mention_workflow_inputs(inputs: object) -> PostHogCodeSlackMentionWorkflowInputs: @@ -64,6 +72,41 @@ def coerce_mention_workflow_inputs(inputs: object) -> PostHogCodeSlackMentionWor ) +@dataclass +class SlackAppMentionWorkflowInputs: + """Conversation-level inputs for the per-thread queue workflow. + + One workflow instance covers one Slack conversation (channel thread or DM + thread), identified entirely by its workflow ID; individual messages + arrive as ``new_message`` signals carrying + ``PostHogCodeSlackMentionWorkflowInputs``. These fields exist only to + carry state across ``continue_as_new`` — fresh starts leave them empty. + """ + + pending_messages: list[PostHogCodeSlackMentionWorkflowInputs] = field(default_factory=list) + processed_event_keys: list[str] = field(default_factory=list) + + +# The queue-ack reaction contract: the webhook adds the queued reaction at +# dispatch and the queue workflow swaps it for the processing one. Both sides +# must agree, so the names live here rather than as literals at each call site. +SLACK_APP_QUEUED_REACTION = "hourglass" +SLACK_APP_PROCESSING_REACTION = "eyes" + + +class MarkSlackAppMessageProcessingInput(BaseModel): + """Single-argument input for ``mark_slack_app_message_processing_activity``. + + New Slack-app activities take one pydantic model instead of positional + arguments so the payload can grow fields without signature churn. + """ + + integration_id: int + slack_team_id: str + channel: str + message_ts: str + + @dataclass class PostHogCodeSlackMentionCommandWorkflowInputs: event: dict[str, Any] diff --git a/posthog/temporal/tests/ai/slack_app/test_slack_app_mention_workflow.py b/posthog/temporal/tests/ai/slack_app/test_slack_app_mention_workflow.py new file mode 100644 index 000000000000..e629f41420dd --- /dev/null +++ b/posthog/temporal/tests/ai/slack_app/test_slack_app_mention_workflow.py @@ -0,0 +1,408 @@ +import os +import uuid +import asyncio +from typing import Any, Literal + +import pytest + +from temporalio import activity +from temporalio.common import WorkflowIDConflictPolicy, WorkflowIDReusePolicy +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import UnsandboxedWorkflowRunner, Worker + +from posthog.temporal.ai.slack_app import derive_mention_workflow_id +from posthog.temporal.ai.slack_app.slack_app_mention import SlackAppMentionWorkflow +from posthog.temporal.ai.slack_app.types import ( + MarkSlackAppMessageProcessingInput, + PostHogCodeRepoCascadeOutcome, + PostHogCodeSlackMentionWorkflowInputs, + SlackAppMentionWorkflowInputs, + SlackRepoSelectionOutcome, +) + + +def _message( + ts: str, + *, + event_id: str | None = None, + untagged: bool = False, +) -> PostHogCodeSlackMentionWorkflowInputs: + return PostHogCodeSlackMentionWorkflowInputs( + event={"channel": "C1", "ts": ts, "thread_ts": "100.0", "user": "U1", "text": "fix the bug"}, + integration_id=1, + slack_team_id="T1", + slack_event_id=event_id, + user_id=42, + untagged_followup=untagged, + ) + + +class _Recorder: + def __init__(self) -> None: + # (ts, repository) per create-task call, in execution order. + self.created: list[tuple[str, str | None]] = [] + # per_message_identity as seen by the first activity of each message. + self.identity_flags: list[bool] = [] + # ts per hourglass->eyes reaction swap, in execution order. + self.processing_marked: list[str] = [] + # ts per forwarded followup, in execution order. + self.forwarded: list[str] = [] + # ts -> forward result; missing means False (no existing task, fall through to new-task path). + self.forward_results: dict[str, bool] = {} + # ts -> cascade mode; missing means "auto" with a fixed repository. + self.cascade_modes: dict[str, Literal["auto", "no_repo", "agent_needed", "needs_user_github"]] = {} + # ts -> gate the create-task fake blocks on, to hold a message mid-processing. + self.create_gates: dict[str, asyncio.Event] = {} + self.create_reached: dict[str, asyncio.Event] = {} + self.picker_posted = asyncio.Event() + self.picker_workflow_id: str | None = None + + +def _fake_activities(rec: _Recorder) -> list: + @activity.defn(name="enforce_posthog_code_billing_quota_activity") + async def quota( + inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str, slack_user_id: str + ) -> bool: + rec.identity_flags.append(inputs.per_message_identity) + return False + + @activity.defn(name="classify_untagged_followup_activity") + async def classify_followup( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + slack_user_id: str, + event_text: str, + ) -> bool: + return True + + @activity.defn(name="forward_posthog_code_followup_activity") + async def forward( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + slack_user_id: str, + event_text: str, + user_message_ts: str | None, + ) -> bool: + ts = inputs.event["ts"] + if rec.forward_results.get(ts, False): + rec.forwarded.append(ts) + return True + return False + + @activity.defn(name="collect_posthog_code_thread_messages_activity") + async def collect( + inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str + ) -> list[dict[str, str]]: + return [{"user": "U1", "text": inputs.event["text"]}] + + @activity.defn(name="cascade_posthog_code_repository_activity") + async def cascade( + inputs: PostHogCodeSlackMentionWorkflowInputs, event_text: str, user_id: int | None = None + ) -> PostHogCodeRepoCascadeOutcome: + mode = rec.cascade_modes.get(inputs.event["ts"], "auto") + repository = "org/auto-repo" if mode == "auto" else None + return PostHogCodeRepoCascadeOutcome(mode=mode, repository=repository, reason="test") + + @activity.defn(name="classify_posthog_code_task_needs_repo_activity") + async def needs_repo(event_text: str, thread_messages: list[dict[str, str]]) -> bool: + return True + + @activity.defn(name="discover_posthog_code_repository_via_agent_activity") + async def discover( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + event: dict[str, Any], + thread_messages: list[dict[str, str]], + user_id: int, + ) -> SlackRepoSelectionOutcome: + return SlackRepoSelectionOutcome(status="failed", repository=None, reason="agent crashed") + + @activity.defn(name="post_posthog_code_repo_picker_activity") + async def post_picker( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + slack_user_id: str, + event: dict[str, Any], + workflow_id: str, + guidance: str, + allow_no_repo: bool, + user_id: int | None = None, + ) -> None: + rec.picker_workflow_id = workflow_id + rec.picker_posted.set() + + @activity.defn(name="resolve_posthog_code_authorship_activity") + async def resolve_authorship( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + slack_user_id: str, + user_id: int, + workflow_id: str, + repository: str, + ) -> str: + return "proceed" + + @activity.defn(name="block_posthog_code_task_if_no_personal_github_activity") + async def block_github( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + user_id: int, + allow_bot_prs: bool = False, + ) -> bool: + return False + + @activity.defn(name="create_posthog_code_task_for_repo_activity") + async def create_task( + inputs: PostHogCodeSlackMentionWorkflowInputs, + channel: str, + thread_ts: str, + slack_user_id: str, + user_id: int, + event: dict[str, Any], + thread_messages: list[dict[str, str]], + repository: str | None, + repo_research_task_id: str | None = None, + repo_research_run_id: str | None = None, + ) -> None: + ts = inputs.event["ts"] + reached = rec.create_reached.get(ts) + if reached: + reached.set() + gate = rec.create_gates.get(ts) + if gate: + await gate.wait() + rec.created.append((ts, repository)) + + @activity.defn(name="post_posthog_code_picker_timeout_activity") + async def picker_timeout(inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str) -> None: + return None + + @activity.defn(name="post_posthog_code_authorship_timeout_activity") + async def authorship_timeout(inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str) -> None: + return None + + @activity.defn(name="post_posthog_code_internal_error_activity") + async def internal_error(inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str) -> None: + return None + + @activity.defn(name="resolve_posthog_code_slack_user_activity") + async def resolve_user( + inputs: PostHogCodeSlackMentionWorkflowInputs, channel: str, thread_ts: str, slack_user_id: str + ) -> int | None: + return 42 + + @activity.defn(name="mark_slack_app_message_processing_activity") + async def mark_processing(input: MarkSlackAppMessageProcessingInput) -> None: + rec.processing_marked.append(input.message_ts) + + return [ + mark_processing, + quota, + classify_followup, + forward, + collect, + cascade, + needs_repo, + discover, + post_picker, + resolve_authorship, + block_github, + create_task, + picker_timeout, + authorship_timeout, + internal_error, + resolve_user, + ] + + +async def _signal_with_start(env, task_queue: str, workflow_id: str, message: PostHogCodeSlackMentionWorkflowInputs): + """Mirror the production dispatch shape from api._start_mention_workflow.""" + return await env.client.start_workflow( + SlackAppMentionWorkflow.run, + SlackAppMentionWorkflowInputs(), + id=workflow_id, + task_queue=task_queue, + id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, + id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, + start_signal="new_message", + start_signal_args=[message], + ) + + +class _Harness: + """One time-skipping environment + worker per test. + + Time only skips while the test awaits the workflow result, so the idle + timer (and the 15-minute picker timer) stay parked while the test delivers + signals in real time — no sleeps, no timer races. + """ + + def __init__(self, rec: _Recorder) -> None: + self.rec = rec + self.task_queue = str(uuid.uuid4()) + + async def __aenter__(self): + # Escape hatch for networks where the SDK's temporal.download fetch is + # blocked: point at a pre-downloaded temporal-test-server binary (from + # the temporalio/sdk-java GitHub releases). Unset, the SDK downloads + # and caches the binary itself. + self._env_cm = await WorkflowEnvironment.start_time_skipping( + test_server_existing_path=os.environ.get("TEMPORAL_TEST_SERVER_PATH") + ) + self.env = await self._env_cm.__aenter__() + self._worker_cm = Worker( + self.env.client, + task_queue=self.task_queue, + workflows=[SlackAppMentionWorkflow], + activities=_fake_activities(self.rec), + workflow_runner=UnsandboxedWorkflowRunner(), + ) + await self._worker_cm.__aenter__() + return self + + async def __aexit__(self, *exc_info): + await self._worker_cm.__aexit__(*exc_info) + await self._env_cm.__aexit__(*exc_info) + + +@pytest.mark.asyncio +async def test_queued_messages_process_serially_in_arrival_order(): + rec = _Recorder() + first, second, third = _message("1.1"), _message("1.2"), _message("1.3") + rec.create_reached["1.1"] = asyncio.Event() + rec.create_gates["1.1"] = asyncio.Event() + + async with _Harness(rec) as h: + handle = await _signal_with_start(h.env, h.task_queue, f"wf-{uuid.uuid4()}", first) + # Hold the first message inside its create-task activity, queue two + # more behind it, then release. FIFO must be preserved. + await asyncio.wait_for(rec.create_reached["1.1"].wait(), timeout=30) + await handle.signal(SlackAppMentionWorkflow.new_message, second) + await handle.signal(SlackAppMentionWorkflow.new_message, third) + rec.create_gates["1.1"].set() + await asyncio.wait_for(handle.result(), timeout=30) + + assert rec.created == [("1.1", "org/auto-repo"), ("1.2", "org/auto-repo"), ("1.3", "org/auto-repo")] + # Each mention gets its hourglass swapped for eyes as it leaves the queue. + assert rec.processing_marked == ["1.1", "1.2", "1.3"] + # The queue workflow stamps per-message identity on every message it + # processes — the payload arrives with the field unset. + assert rec.identity_flags == [True, True, True] + + +@pytest.mark.asyncio +async def test_signal_with_start_on_running_workflow_signals_same_run(): + rec = _Recorder() + first, second = _message("1.1"), _message("1.2") + rec.create_reached["1.1"] = asyncio.Event() + rec.create_gates["1.1"] = asyncio.Event() + + async with _Harness(rec) as h: + workflow_id = f"wf-{uuid.uuid4()}" + handle = await _signal_with_start(h.env, h.task_queue, workflow_id, first) + # Issue a second signal-with-start (the exact production dispatch call) + # while the first run is mid-message. It must NOT start a second + # execution — the server delivers the signal to the running one. + await asyncio.wait_for(rec.create_reached["1.1"].wait(), timeout=30) + handle_two = await _signal_with_start(h.env, h.task_queue, workflow_id, second) + rec.create_gates["1.1"].set() + await asyncio.wait_for(handle.result(), timeout=30) + assert handle_two.first_execution_run_id == handle.first_execution_run_id + + assert rec.created == [("1.1", "org/auto-repo"), ("1.2", "org/auto-repo")] + + +@pytest.mark.asyncio +async def test_duplicate_slack_event_id_is_processed_once(): + rec = _Recorder() + message = _message("1.1", event_id="Ev123") + + async with _Harness(rec) as h: + handle = await _signal_with_start(h.env, h.task_queue, f"wf-{uuid.uuid4()}", message) + await handle.signal(SlackAppMentionWorkflow.new_message, message) + await asyncio.wait_for(handle.result(), timeout=30) + + assert rec.created == [("1.1", "org/auto-repo")] + + +@pytest.mark.asyncio +async def test_untagged_followup_forwards_without_task_creation_or_reactions(): + rec = _Recorder() + rec.forward_results["1.1"] = True + + async with _Harness(rec) as h: + handle = await _signal_with_start(h.env, h.task_queue, f"wf-{uuid.uuid4()}", _message("1.1", untagged=True)) + await asyncio.wait_for(handle.result(), timeout=30) + + assert rec.forwarded == ["1.1"] + assert rec.created == [] + # Untagged followups were never addressed to the bot: no reaction swap. + assert rec.processing_marked == [] + + +@pytest.mark.asyncio +async def test_idle_exit_then_signal_with_start_processes_in_fresh_run(): + rec = _Recorder() + workflow_id = f"wf-{uuid.uuid4()}" + + async with _Harness(rec) as h: + handle = await _signal_with_start(h.env, h.task_queue, workflow_id, _message("1.1")) + await asyncio.wait_for(handle.result(), timeout=30) + # First run has idled out and completed; the production dispatch shape + # must start a fresh run under the same conversation ID. + handle_two = await _signal_with_start(h.env, h.task_queue, workflow_id, _message("2.1")) + await asyncio.wait_for(handle_two.result(), timeout=30) + assert handle_two.result_run_id != handle.result_run_id + + assert rec.created == [("1.1", "org/auto-repo"), ("2.1", "org/auto-repo")] + + +@pytest.mark.asyncio +async def test_repo_picker_signal_resolves_and_queue_continues(): + rec = _Recorder() + rec.cascade_modes["1.1"] = "agent_needed" + + async with _Harness(rec) as h: + workflow_id = f"wf-{uuid.uuid4()}" + handle = await _signal_with_start(h.env, h.task_queue, workflow_id, _message("1.1")) + # First message falls through discovery to the picker and blocks the + # queue; a second message queues up behind it in the meantime. + await asyncio.wait_for(rec.picker_posted.wait(), timeout=30) + await handle.signal(SlackAppMentionWorkflow.new_message, _message("1.2")) + await handle.signal(SlackAppMentionWorkflow.repo_selected, "org/picked") + await asyncio.wait_for(handle.result(), timeout=30) + + # The picker message must carry the conversation workflow ID — it is what + # the interactivity webhook uses to route the click back as a signal. + assert rec.picker_workflow_id == workflow_id + + assert rec.created == [("1.1", "org/picked"), ("1.2", "org/auto-repo")] + + +@pytest.mark.asyncio +async def test_continue_as_new_carry_over_processes_pending_and_dedups_seen(): + rec = _Recorder() + pending = _message("1.1", event_id="Ev-pending") + already_seen = _message("1.2", event_id="Ev-seen") + + async with _Harness(rec) as h: + # Start with post-continue_as_new-shaped inputs: one carried pending + # message and one already-processed key. + handle = await h.env.client.start_workflow( + SlackAppMentionWorkflow.run, + SlackAppMentionWorkflowInputs( + pending_messages=[pending], + processed_event_keys=[derive_mention_workflow_id(already_seen)], + ), + id=f"wf-{uuid.uuid4()}", + task_queue=h.task_queue, + ) + await handle.signal(SlackAppMentionWorkflow.new_message, already_seen) + await asyncio.wait_for(handle.result(), timeout=30) + + assert rec.created == [("1.1", "org/auto-repo")] diff --git a/products/slack_app/backend/api.py b/products/slack_app/backend/api.py index 89f744c51022..95013e5fbaca 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -37,16 +37,23 @@ from posthog.models.user import User from posthog.models.user_integration import UserGitHubIntegration, UserIntegration from posthog.temporal.ai.slack_app import ( + SLACK_APP_QUEUED_REACTION, PostHogCodeSlackMentionCommandWorkflowInputs, PostHogCodeSlackMentionWorkflowInputs, + SlackAppMentionWorkflowInputs, derive_mention_workflow_id, ) +from posthog.temporal.ai.slack_app.helpers import safe_react from posthog.temporal.ai.slack_app.posthog_code_slack_interactivity import ( PostHogCodeSlackInteractivityInputs, PostHogCodeSlackTerminateTaskWorkflow, ) from posthog.temporal.ai.slack_app.posthog_code_slack_mention import PostHogCodeSlackMentionWorkflow from posthog.temporal.ai.slack_app.posthog_code_slack_mention_command import PostHogCodeSlackMentionCommandWorkflow +from posthog.temporal.ai.slack_app.slack_app_mention import ( + SlackAppMentionWorkflow, + derive_slack_app_mention_workflow_id, +) from posthog.temporal.common.client import sync_connect from posthog.user_permissions import UserPermissions from posthog.utils import get_instance_region @@ -56,6 +63,7 @@ is_slack_app_assistant_enabled, is_slack_app_bot_prs_enabled, is_slack_app_oauth_enabled, + is_slack_app_queue_workflow_enabled, is_slack_app_untagged_thread_followups_enabled, ) from products.slack_app.backend.models import SlackChannel, SlackThreadTaskMapping @@ -1417,7 +1425,16 @@ def _start_posthog_code_workflow( event: dict, event_id: str | None, workflow_id: str | None = None, + start_signal: str | None = None, + start_signal_args: list[Any] | None = None, ) -> None: + """Start a Slack-app workflow, optionally as a signal-with-start. + + With ``start_signal`` set the operation is atomic on the server: a running + execution gets the signal, a finished (or never-started) one is started + with the signal as its first event — how the queue workflow guarantees a + message lands exactly once in its conversation's queue. + """ if workflow_id is None: fallback = event_id if event_id else f"{event.get('channel', '')}:{event.get('ts', '')}" workflow_id = f"{id_prefix}-{slack_team_id}:{fallback}" @@ -1430,10 +1447,32 @@ def _start_posthog_code_workflow( task_queue=settings.TASKS_TASK_QUEUE, id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, + # None / [] match the SDK defaults, so a plain start stays a plain start. + start_signal=start_signal, + start_signal_args=start_signal_args or [], ) ) +def _react_message_queued(integration: Integration, event: dict) -> None: + """Ack an enqueued mention/DM with the queued reaction. The queue workflow + swaps it for the processing one when the message's turn starts. Best-effort + — a reaction failure must never fail the dispatch that already succeeded.""" + message_ts = event.get("ts") + channel = event.get("channel") + if not message_ts or not channel: + return + try: + safe_react(SlackIntegration(integration).client, channel, message_ts, SLACK_APP_QUEUED_REACTION) + except Exception as e: + logger.warning( + "slack_app_queued_reaction_failed", + channel=channel, + message_ts=message_ts, + error=str(e), + ) + + _ASSISTANT_CONTEXT_TTL_SECONDS = 60 * 60 _ASSISTANT_SUGGESTED_PROMPTS = [ {"title": "Fix a bug", "message": "Open a PR to fix a bug in my connected repo"}, @@ -2413,6 +2452,33 @@ def _start_mention_workflow( user_id=posthog_user.id if posthog_user else None, untagged_followup=untagged_followup, ) + # Deriving the conversation ID is free; check it before paying the remote + # flag evaluation. When it can't be derived (missing channel/ts) we fall + # through to the per-message workflow purely so flag-on and flag-off + # behavior stay identical for such events. + queue_workflow_id = derive_slack_app_mention_workflow_id(workflow_inputs) + if queue_workflow_id is not None and is_slack_app_queue_workflow_enabled(integration, slack_team_id): + # Note: under the queue workflow the slack_mention_workflow_id the + # task-creation activity persists (derived per message) has no + # Temporal execution behind it, so the debug-tool Temporal link + # dangles. Threading the real workflow id through needs an activity + # signature change — deferred to a follow-up. + _start_posthog_code_workflow( + SlackAppMentionWorkflow, + SlackAppMentionWorkflowInputs(), + id_prefix="slack-app-mention", + slack_team_id=slack_team_id, + event=event, + event_id=event_id, + workflow_id=queue_workflow_id, + start_signal="new_message", + start_signal_args=[workflow_inputs], + ) + # Untagged followups stay reaction-free: they were never addressed + # to the bot, and most get dropped by the chitchat classifier. + if not untagged_followup: + _react_message_queued(integration, event) + return ROUTE_HANDLED_LOCALLY # Use derive_mention_workflow_id as the single source of truth: the workflow persists the same # value as slack_mention_workflow_id, so dispatch and the debug-tool Temporal link stay consistent _start_posthog_code_workflow( diff --git a/products/slack_app/backend/feature_flags.py b/products/slack_app/backend/feature_flags.py index 78524aaec344..6f5403e75223 100644 --- a/products/slack_app/backend/feature_flags.py +++ b/products/slack_app/backend/feature_flags.py @@ -33,6 +33,7 @@ SLACK_APP_AGENT_DESIGN_FLAG = "slack-app-agent-design" SLACK_APP_ASSISTANT_FLAG = "slack-app-assistant" SLACK_APP_BOT_PRS_FLAG = "slack-app-bot-prs" +SLACK_APP_QUEUE_WORKFLOW_FLAG = "slack-app-queue-workflow" UNTAGGED_THREAD_FOLLOWUPS_FLAG = "posthog-slack-app-untagged-thread-followups" @@ -136,6 +137,31 @@ def is_slack_app_untagged_thread_followups_enabled(integration: Integration, sla return False +def is_slack_app_queue_workflow_enabled(integration: Integration, slack_team_id: str) -> bool: + """Gate for the per-conversation queue workflow: when on, mention/followup/DM + dispatch signal-with-starts one ``SlackAppMentionWorkflow`` per thread that + processes messages serially, instead of one workflow per message. Keyed on + the Slack workspace + PostHog org.""" + try: + return bool( + posthoganalytics.feature_enabled( + SLACK_APP_QUEUE_WORKFLOW_FLAG, + f"slack_workspace:{slack_team_id}", + groups={"organization": str(integration.team.organization_id)}, + person_properties=_region_properties(), + only_evaluate_locally=False, + send_feature_flag_events=False, + ) + ) + except Exception: + logger.exception( + "slack_app_queue_workflow_flag_check_failed", + slack_team_id=slack_team_id, + integration_id=integration.id, + ) + return False + + def is_slack_app_bot_prs_enabled(team: Team) -> bool: organization_id = str(team.organization_id) project_id = str(team.id) diff --git a/products/slack_app/backend/tests/test_followup_forwarding.py b/products/slack_app/backend/tests/test_followup_forwarding.py index 75139413a62c..b1d032727112 100644 --- a/products/slack_app/backend/tests/test_followup_forwarding.py +++ b/products/slack_app/backend/tests/test_followup_forwarding.py @@ -26,11 +26,14 @@ from products.slack_app.backend.models import SlackThreadTaskMapping -def _make_inputs(integration_id: int, slack_team_id: str = "T_SLACK") -> PostHogCodeSlackMentionWorkflowInputs: +def _make_inputs( + integration_id: int, slack_team_id: str = "T_SLACK", per_message_identity: bool = False +) -> PostHogCodeSlackMentionWorkflowInputs: return PostHogCodeSlackMentionWorkflowInputs( event={"channel": "C123", "ts": "1234.5678", "user": "U_ALICE", "text": "<@BOT> do something"}, integration_id=integration_id, slack_team_id=slack_team_id, + per_message_identity=per_message_identity, ) @@ -705,6 +708,7 @@ def test_unauthorized_actor_returns_true_with_resolver_feedback(self, mock_slack mock_resolve.assert_called_once_with(mock_slack_instance, self.integration, "U_BOB", "C123", "1234.5678") mock_slack_instance.client.chat_postMessage.assert_not_called() + @patch("products.tasks.backend.facade.api.rebind_sandbox_identity_for_user") @patch( "products.tasks.backend.logic.services.connection_token.create_sandbox_connection_token", return_value="jwt-token", @@ -713,11 +717,13 @@ def test_unauthorized_actor_returns_true_with_resolver_feedback(self, mock_slack @patch("products.slack_app.backend.api.resolve_slack_user") @patch("posthog.models.integration.SlackIntegration") def test_cross_user_followup_authorized_prefixes_actor_name( - self, mock_slack_cls, mock_resolve, mock_send, mock_token + self, mock_slack_cls, mock_resolve, mock_send, mock_token, mock_rebind ): - # A second user in the same PostHog org and team should be allowed to chip in - # on the thread; their message is forwarded under the original author's identity - # but their name is prepended so the agent knows who actually spoke. + # A second user in the same PostHog org and team chips in on the thread: + # we re-bind the sandbox JWT and the PostHog MCP credentials to *them*, so + # subsequent agent writes (insights, dashboards) attribute to the live + # actor rather than the long-gone task creator. Their name is also + # prepended onto the text so the agent knows who spoke. self._create_mapping(mentioning_user="U_ALICE") bob = User.objects.create(email="bob@test.com", first_name="Bob") mock_slack_instance = MagicMock() @@ -725,12 +731,51 @@ def test_cross_user_followup_authorized_prefixes_actor_name( mock_resolve.return_value = SlackUserContext(user=bob, slack_email="bob@test.com") mock_send.return_value = _command_result(success=True, status_code=200) + inputs = _make_inputs(self.integration.id, per_message_identity=True) + result = forward_posthog_code_followup_activity( + inputs, "C123", "1234.5678", "U_BOB", "<@BOT> please retry the build", "1234.5679" + ) + + assert result is True + # JWT is minted for Bob (the live actor), not Alice (the original creator). + # The facade forwards user_id/distinct_id positionally to the underlying mint, + # so the mock sees args = (run, user_id, distinct_id). + assert mock_token.call_args.args[1] == bob.id + # The sandbox identity is rebound to Bob before the follow-up is + # delivered, so the agent acts as Bob for this turn. The tasks layer + # detects the identity transition and bypasses its refresh rate limits. + mock_rebind.assert_called_once_with(self.task_run.id, bob.id, auth_token="jwt-token") + + @patch("products.tasks.backend.facade.api.rebind_sandbox_identity_for_user") + @patch( + "products.tasks.backend.logic.services.connection_token.create_sandbox_connection_token", + return_value="jwt-token", + ) + @patch("products.tasks.backend.logic.services.agent_command.send_user_message") + @patch("products.slack_app.backend.api.resolve_slack_user") + @patch("posthog.models.integration.SlackIntegration") + def test_legacy_dispatch_keeps_creator_identity( + self, mock_slack_cls, mock_resolve, mock_send, mock_token, mock_rebind + ): + # Messages dispatched by the legacy per-message workflow (queue flag + # off) must behave exactly as before: no MCP or GitHub rebind, and the + # sandbox JWT stays bound to the task creator even for a cross-user + # follow-up. Bob still participates — only via the creator's identity. + self._create_mapping(mentioning_user="U_ALICE") + bob = User.objects.create(email="bob-legacy@test.com", first_name="Bob") + mock_slack_instance = MagicMock() + mock_slack_cls.return_value = mock_slack_instance + mock_resolve.return_value = SlackUserContext(user=bob, slack_email="bob-legacy@test.com") + mock_send.return_value = _command_result(success=True, status_code=200) + inputs = _make_inputs(self.integration.id) result = forward_posthog_code_followup_activity( inputs, "C123", "1234.5678", "U_BOB", "<@BOT> please retry the build", "1234.5679" ) assert result is True + assert mock_token.call_args.args[1] == self.task.created_by_id + mock_rebind.assert_not_called() mock_send.assert_called_once_with( self.task_run, "Bob: please retry the build", auth_token="jwt-token", timeout=90 ) @@ -763,6 +808,37 @@ def test_cross_user_followup_falls_back_to_email_when_no_full_name( mock_send.assert_called_once_with(self.task_run, "bob@test.com: ping", auth_token="jwt-token", timeout=90) + @patch("products.tasks.backend.facade.api.rebind_sandbox_identity_for_user") + @patch( + "products.tasks.backend.logic.services.connection_token.create_sandbox_connection_token", + return_value="jwt-token", + ) + @patch("products.tasks.backend.logic.services.agent_command.send_user_message") + @patch("products.slack_app.backend.api.resolve_slack_user") + @patch("posthog.models.integration.SlackIntegration") + def test_cross_user_followup_proceeds_when_identity_rebind_raises( + self, mock_slack_cls, mock_resolve, mock_send, mock_token, mock_rebind + ): + # The facade rebind is best-effort by contract, but even if it raises + # unexpectedly we still deliver the user's message rather than dropping + # it on the floor. The cost is that this turn may still attribute to + # the previous identity — an accepted downside vs. losing the input. + self._create_mapping(mentioning_user="U_ALICE") + bob = User.objects.create(email="bob@test.com", first_name="Bob") + mock_slack_cls.return_value = MagicMock() + mock_resolve.return_value = SlackUserContext(user=bob, slack_email="bob@test.com") + mock_send.return_value = _command_result(success=True, status_code=200) + mock_rebind.side_effect = RuntimeError("refresh failed") + + inputs = _make_inputs(self.integration.id, per_message_identity=True) + result = forward_posthog_code_followup_activity( + inputs, "C123", "1234.5678", "U_BOB", "<@BOT> ping", "1234.5679" + ) + + assert result is True + mock_rebind.assert_called_once() + mock_send.assert_called_once_with(self.task_run, "Bob: ping", auth_token="jwt-token", timeout=90) + @patch("products.slack_app.backend.api.resolve_slack_user", return_value=None) @patch("posthog.models.integration.SlackIntegration") def test_cross_user_followup_unmapped_user_delegates_feedback_to_resolver(self, mock_slack_cls, mock_resolve): @@ -825,13 +901,14 @@ def test_sandbox_not_ready_returns_true_with_message(self, mock_slack_cls): call_kwargs = mock_slack_instance.client.chat_postMessage.call_args.kwargs assert "still starting up" in call_kwargs["text"] + @patch("products.tasks.backend.facade.api.rebind_sandbox_identity_for_user") @patch( "products.tasks.backend.logic.services.connection_token.create_sandbox_connection_token", return_value="jwt-token", ) @patch("products.tasks.backend.logic.services.agent_command.send_user_message") @patch("posthog.models.integration.SlackIntegration") - def test_successful_forwarding(self, mock_slack_cls, mock_send, mock_token): + def test_successful_forwarding(self, mock_slack_cls, mock_send, mock_token, mock_rebind): mapping = self._create_mapping() mock_slack_instance = MagicMock() mock_slack_cls.return_value = mock_slack_instance @@ -841,13 +918,17 @@ def test_successful_forwarding(self, mock_slack_cls, mock_send, mock_token): data={"result": {"assistant_message": "thanks"}}, ) - inputs = _make_inputs(self.integration.id) + inputs = _make_inputs(self.integration.id, per_message_identity=True) result = forward_posthog_code_followup_activity( inputs, "C123", "1234.5678", "U_ALICE", "<@BOT> do something", "1234.5679" ) assert result is True mock_token.assert_called_once() + # Every message rebinds the sandbox identity to its actor — here the + # task creator. The tasks layer rate-limits same-identity refreshes, + # so this stays cheap; what matters is the actor is always the one authed. + mock_rebind.assert_called_once_with(self.task_run.id, self.task.created_by_id, auth_token="jwt-token") mock_send.assert_called_once_with(self.task_run, "do something", auth_token="jwt-token", timeout=90) # Agent is now working on the message, so the :eyes: reaction stays up — it is # not swapped to :hedgehog: until the task genuinely completes. diff --git a/products/slack_app/backend/tests/test_posthog_code_event_handler.py b/products/slack_app/backend/tests/test_posthog_code_event_handler.py index f3baf258671c..0178e6ba598b 100644 --- a/products/slack_app/backend/tests/test_posthog_code_event_handler.py +++ b/products/slack_app/backend/tests/test_posthog_code_event_handler.py @@ -1253,3 +1253,78 @@ def test_slack_error_is_swallowed(self): with enabled_p, slack as slack_cls: slack_cls.return_value.client.chat_postMessage.side_effect = Exception("slack down") send_assistant_install_welcome(self.integration) # must not raise + + +class TestQueueWorkflowDispatch(TestCase): + def setUp(self): + from django.utils import timezone + + from posthog.helpers.slack_scopes import REQUIRED_SLACK_SCOPES + + cache.clear() + self.factory = RequestFactory() + self.organization = Organization.objects.create(name="Test Org") + self.team = Team.objects.create(organization=self.organization, name="Test Team") + self.user = User.objects.create(email="dev@example.com", distinct_id="queue-user-1") + OrganizationMembership.objects.create(organization=self.organization, user=self.user) + self.user.current_organization = self.organization + self.user.current_team = self.team + self.user.save() + self.integration = Integration.objects.create( + team=self.team, + kind="slack", + integration_id="T12345", + config={"scope": ",".join(sorted(REQUIRED_SLACK_SCOPES))}, + sensitive_config={"access_token": "xoxb-posthog-code-test"}, + ) + SlackUserProfileCache.objects.create( + integration=self.integration, + slack_user_id="U123", + email="dev@example.com", + display_name="Dev", + real_name="Dev User", + refreshed_at=timezone.now(), + ) + + @parameterized.expand( + [ + ("top_level_mention_anchors_on_ts", None, "1234.5678"), + ("threaded_mention_anchors_on_thread_root", "1000.0001", "1000.0001"), + ] + ) + @patch("products.slack_app.backend.api.SlackIntegration") + @patch("products.slack_app.backend.api.is_slack_app_queue_workflow_enabled", return_value=True) + @patch("products.slack_app.backend.api.asyncio.run") + @patch("products.slack_app.backend.api.sync_connect") + @override_settings(DEBUG=False, CLOUD_DEPLOYMENT="US") + def test_flag_on_signal_with_starts_conversation_workflow( + self, _name, thread_ts, expected_anchor, mock_sync_connect, mock_asyncio_run, mock_flag, mock_slack + ): + # With the flag on, every message in a conversation must land in ONE + # per-thread workflow via signal-with-start — the conversation ID + # anchors on the thread root so followups reach the same instance. + from posthog.temporal.ai.slack_app.slack_app_mention import SlackAppMentionWorkflow + + from products.slack_app.backend.api import ROUTE_HANDLED_LOCALLY, route_posthog_code_event_to_relevant_region + + mock_slack.return_value.missing_scopes.return_value = set() + event = {"type": "app_mention", "channel": "C001", "user": "U123", "ts": "1234.5678"} + if thread_ts: + event["thread_ts"] = thread_ts + request = self.factory.post("/slack/event-callback/", HTTP_HOST="us.posthog.com") + + result = route_posthog_code_event_to_relevant_region(request, event, "T12345") + + assert result == ROUTE_HANDLED_LOCALLY + mock_sync_connect.return_value.start_workflow.assert_called_once() + call = mock_sync_connect.return_value.start_workflow.call_args + assert call.args[0] == SlackAppMentionWorkflow.run + assert call.kwargs["id"] == f"slack-app-mention-T12345:C001:{expected_anchor}" + assert call.kwargs["start_signal"] == "new_message" + message = call.kwargs["start_signal_args"][0] + assert message.user_id == self.user.id + assert message.event["ts"] == "1234.5678" + # The enqueued message is acked with an hourglass on the message itself. + mock_slack.return_value.client.reactions_add.assert_called_once_with( + channel="C001", timestamp="1234.5678", name="hourglass" + ) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index d2f055348235..2608fe15a0a4 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -175,6 +175,7 @@ "presign_task_run_artifact", "read_task_run_artifact", "read_task_run_logs", + "rebind_sandbox_identity_for_user", "redeem_code_invite", "redispatch_task_run", "refresh_team_code_workstreams", @@ -4581,3 +4582,41 @@ def forward_thread_message( message.forwarded_run = run message.save(update_fields=["forwarded_to_agent_at", "forwarded_by", "forwarded_run"]) return "ok", _thread_message_to_dto(message) + + +def rebind_sandbox_identity_for_user( + run_id: str | UUID, + user_id: int, + *, + auth_token: str | None = None, +) -> None: + """Rebind a run's live sandbox to ``user_id``: mint a PostHog MCP OAuth + token for them and push fresh MCP configs. + + Used by Slack follow-ups so a teammate taking over the conversation acts + as themselves — insights, dashboards, and other PostHog writes attribute + to the live actor rather than the task creator. Best-effort by contract: + failures are logged, never raised, so a rebind problem can't block the + message that triggered it. + """ + from products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox import ( # noqa: PLC0415 — keep sandbox deps off the api import path + refresh_sandbox_mcp_for_user, + ) + + try: + run = TaskRun.objects.select_related("task").get(id=run_id) + user = User.objects.get(id=user_id) + except Exception as e: + logger.warning( + "Sandbox identity rebind skipped: could not load run or user", + extra={"run_id": str(run_id), "user_id": user_id, "error": str(e)}, + ) + return + + try: + # Mirrors the scope Slack-originated runs boot with. The effective + # per-run scope isn't persisted in run state (follow-up); a non-Slack + # caller of this facade must revisit this constant. + refresh_sandbox_mcp_for_user(run, user, scopes="full", auth_token=auth_token) + except Exception: + logger.exception("Sandbox MCP identity rebind failed", extra={"run_id": str(run_id), "user_id": user_id}) diff --git a/products/tasks/backend/temporal/oauth.py b/products/tasks/backend/temporal/oauth.py index 2010d5d4ceef..e1cb842e8201 100644 --- a/products/tasks/backend/temporal/oauth.py +++ b/products/tasks/backend/temporal/oauth.py @@ -27,12 +27,15 @@ def _oauth_application_for_task(task: Task) -> SandboxOAuthApplication: return "array" -def create_oauth_access_token(task: Task, *, scopes: PosthogMcpScopes = "read_only") -> str: +def create_oauth_access_token(task: Task, *, scopes: PosthogMcpScopes = "read_only", user=None) -> str: """Create an OAuth access token for the task's sandbox app, scoped to the task's team. - OAuth tokens auto-expire after 6 hours, so no cleanup is needed. + Minted for ``user`` when given — per-message identity rebinds mint for the + live Slack actor — otherwise for the task creator. OAuth tokens + auto-expire after 6 hours, so no cleanup is needed. """ - if not task.created_by: + target = user or task.created_by + if not target: raise TaskInvalidStateError( f"Task {task.id} has no created_by user", {"task_id": task.id}, @@ -40,7 +43,7 @@ def create_oauth_access_token(task: Task, *, scopes: PosthogMcpScopes = "read_on ) return create_oauth_access_token_for_user( - task.created_by, + target, task.team_id, scopes=scopes, application=_oauth_application_for_task(task), diff --git a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py index fcb2d4467994..13095d6b75d9 100644 --- a/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py +++ b/products/tasks/backend/temporal/process_task/activities/get_sandbox_for_repository.py @@ -219,7 +219,6 @@ def get_sandbox_for_repository(input: GetSandboxForRepositoryInput) -> GetSandbo environment_variables["LLM_GATEWAY_URL"] = settings.SANDBOX_LLM_GATEWAY_URL environment_variables.update(get_git_identity_env_vars(task, ctx.state)) - run_state = parse_run_state(ctx.state) # Set resume run ID independently of snapshot so conversation history diff --git a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py index 6aed39ff3aaa..c5b03649bdf2 100644 --- a/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/send_followup_to_sandbox.py @@ -3,7 +3,7 @@ import threading import contextvars from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any import structlog from temporalio import activity @@ -26,14 +26,20 @@ from products.tasks.backend.redis import get_tasks_stream_redis_sync, run_uses_dedicated_stream from products.tasks.backend.temporal.oauth import create_oauth_access_token from products.tasks.backend.temporal.process_task.utils import ( + get_last_sandbox_identity, get_sandbox_ph_mcp_configs, get_user_mcp_server_configs, mark_mcp_token_issued, + mark_sandbox_identity, + sandbox_identity_scope, should_refresh_mcp_token, ) from ee.hogai.sandbox import STOP_REASON_END_TURN, TURN_COMPLETE_METHOD +if TYPE_CHECKING: + from posthog.models.user import User + logger = structlog.get_logger(__name__) REFRESH_RETRY_DELAY_SECONDS = 0.5 @@ -213,24 +219,65 @@ def _refresh_sandbox_mcp( scopes: PosthogMcpScopes, auth_token: str | None, ) -> None: - """Mint a fresh OAuth token and push updated MCP configs to the sandbox. + """Best-effort MCP refresh for the web-layer follow-up flow. + + Scopes the OAuth token to the task creator; the web-layer follow-up signal + has no per-turn actor concept. The Slack follow-up path calls + :func:`refresh_sandbox_mcp_for_user` directly so the live mentioner's + identity is used. + """ + created_by = task_run.task.created_by + if not created_by: + logger.info("refresh_mcp_skipped_no_user", run_id=str(task_run.id)) + return + refresh_sandbox_mcp_for_user(task_run, created_by, scopes=scopes, auth_token=auth_token) - Best-effort: retries once on failure, then logs and returns. Never raises - — a failed refresh should not block an otherwise-valid follow-up. - Skipped entirely if a token was issued for this run within the last - MCP_TOKEN_REFRESH_INTERVAL_SECONDS — the in-sandbox token is still fresh. +def refresh_sandbox_mcp_for_user( + task_run: TaskRun, + user: "User", + *, + scopes: PosthogMcpScopes, + auth_token: str | None, +) -> None: + """Mint a fresh OAuth token scoped to ``user`` and push updated MCP configs + to the sandbox via ``send_refresh_session``. + + Retries once on failure, then logs and returns. Never raises — a failed + refresh should not block an otherwise-valid follow-up. + + Rate-limited per sandbox and user: skipped if the same user already had a + token issued to this sandbox within the + ``MCP_TOKEN_REFRESH_INTERVAL_SECONDS`` window. The rate limit never + applies to an identity transition — when ``user`` differs from the + identity the sandbox currently holds (defaulting to the task creator), + the refresh always goes through so the actor swap can't be silently + skipped. + + Note the session's MCP server set is *replaced* on refresh: the personal + MCP Store servers follow ``user``, so a swap uninstalls the previous + actor's servers and installs the new actor's for the rest of their turn. """ run_id = str(task_run.id) - if not should_refresh_mcp_token(run_id): - logger.info("refresh_mcp_skipped_within_interval", run_id=run_id) + task = task_run.task + scope = sandbox_identity_scope(run_id, task_run.state) + current_identity = get_last_sandbox_identity(scope, "mcp") or task.created_by_id + identity_changed = user.id != current_identity + if not identity_changed and not should_refresh_mcp_token(scope, user.id): + logger.info("refresh_mcp_skipped_within_interval", run_id=run_id, user_id=user.id) return + if identity_changed: + logger.info( + "refresh_mcp_identity_transition", + run_id=run_id, + from_user_id=current_identity, + to_user_id=user.id, + ) - task = task_run.task try: - access_token = create_oauth_access_token(task, scopes=scopes) + access_token = create_oauth_access_token(task, scopes=scopes, user=user) except Exception as e: - logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, error=str(e)) + logger.warning("refresh_mcp_token_mint_failed", run_id=run_id, user_id=user.id, error=str(e)) return mcp_configs = get_sandbox_ph_mcp_configs( @@ -240,56 +287,57 @@ def _refresh_sandbox_mcp( interaction_origin=(task_run.state or {}).get("interaction_origin"), task_id=str(task_run.task_id), ) - if task.created_by_id: - user_mcp_configs = get_user_mcp_server_configs( - token=access_token, - team_id=task_run.team_id, - user_id=task.created_by_id, - interaction_origin=(task_run.state or {}).get("interaction_origin"), - ) - if user_mcp_configs: - mcp_configs = mcp_configs + user_mcp_configs + user_mcp_configs = get_user_mcp_server_configs( + token=access_token, + team_id=task_run.team_id, + user_id=user.id, + interaction_origin=(task_run.state or {}).get("interaction_origin"), + ) + if user_mcp_configs: + mcp_configs = mcp_configs + user_mcp_configs if not mcp_configs: - logger.info("refresh_mcp_skipped_no_configs", run_id=run_id) + # Nothing to push means nothing can diverge — record both marks so + # hosts with no resolvable MCP URL and no installs don't repeat the + # mint above on every message (an identity transition would otherwise + # bypass the rate-limit mark forever). + mark_mcp_token_issued(scope, user.id) + mark_sandbox_identity(scope, "mcp", user.id) + logger.info("refresh_mcp_skipped_no_configs", run_id=run_id, user_id=user.id) return mcp_servers = [config.to_dict() for config in mcp_configs] - result = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if result.success: - mark_mcp_token_issued(run_id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=1) - return - - logger.info( - "refresh_mcp_retrying", - run_id=run_id, - error=result.error, - status_code=result.status_code, - ) - time.sleep(REFRESH_RETRY_DELAY_SECONDS) - retry: CommandResult = send_refresh_session( - task_run, - mcp_servers, - auth_token=auth_token, - timeout=REFRESH_TIMEOUT_SECONDS, - ) - if retry.success: - mark_mcp_token_issued(run_id) - logger.info("refresh_mcp_delivered", run_id=run_id, attempts=2) - return + result: CommandResult | None = None + for attempt in (1, 2): + result = send_refresh_session( + task_run, + mcp_servers, + auth_token=auth_token, + timeout=REFRESH_TIMEOUT_SECONDS, + ) + if result.success: + mark_mcp_token_issued(scope, user.id) + mark_sandbox_identity(scope, "mcp", user.id) + logger.info("refresh_mcp_delivered", run_id=run_id, user_id=user.id, attempts=attempt) + return + if attempt == 1: + logger.info( + "refresh_mcp_retrying", + run_id=run_id, + user_id=user.id, + error=result.error, + status_code=result.status_code, + ) + time.sleep(REFRESH_RETRY_DELAY_SECONDS) + assert result is not None logger.warning( "refresh_mcp_failed", run_id=run_id, - error=retry.error, - status_code=retry.status_code, + user_id=user.id, + error=result.error, + status_code=result.status_code, ) diff --git a/products/tasks/backend/temporal/process_task/activities/start_agent_server.py b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py index 60f9dffafee2..c1d0a25d3e45 100644 --- a/products/tasks/backend/temporal/process_task/activities/start_agent_server.py +++ b/products/tasks/backend/temporal/process_task/activities/start_agent_server.py @@ -301,10 +301,10 @@ def _invoke_start_agent_server( wait_for_health=wait_for_health, ) - # Mark startup-time token issuance so follow-ups within the next - # 30m window skip the redundant refresh. - if params.mcp_configs: - mark_mcp_token_issued(ctx.run_id) + # Mark startup-time token issuance so the creator's follow-ups within + # the refresh window skip a redundant refresh_session round-trip. + if params.mcp_configs and ctx.task_created_by_id: + mark_mcp_token_issued(sandbox.id, ctx.task_created_by_id) except Exception as e: if params.agentsh_domains is not None: _emit_agentsh_log_tail(ctx, sandbox) diff --git a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py index 08ede70f60f2..ee575aa66433 100644 --- a/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py +++ b/products/tasks/backend/temporal/process_task/tests/test_send_followup_to_sandbox.py @@ -11,12 +11,16 @@ SEND_FOLLOWUP_MAX_ATTEMPTS, SendFollowupToSandboxInput, _refresh_sandbox_mcp, + refresh_sandbox_mcp_for_user, send_followup_to_sandbox, ) from products.tasks.backend.temporal.process_task.utils import ( McpServerConfig, _mcp_token_issued_cache_key, + _sandbox_identity_cache_key, + get_last_sandbox_identity, mark_mcp_token_issued, + mark_sandbox_identity, ) pytestmark = pytest.mark.django_db @@ -25,10 +29,19 @@ @pytest.fixture(autouse=True) def _clear_mcp_token_cache(): """Ensure each test starts with no recorded token issuances so the - refresh gate doesn't carry state between tests.""" - cache.delete(_mcp_token_issued_cache_key("run-1")) + refresh gate doesn't carry state between tests. The rate-limit cache key + is scoped to ``(run_id, user_id)`` so cross-user follow-ups don't get + silently blocked by a same-run prior refresh — clear both shapes.""" + keys = [ + _mcp_token_issued_cache_key("run-1", 42), + _mcp_token_issued_cache_key("run-1", 43), + _sandbox_identity_cache_key("run-1", "mcp"), + ] + for key in keys: + cache.delete(key) yield - cache.delete(_mcp_token_issued_cache_key("run-1")) + for key in keys: + cache.delete(key) def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: @@ -43,6 +56,7 @@ def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConf def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: dict | None = None) -> MagicMock: task = MagicMock() task.created_by_id = created_by_id + task.created_by = MagicMock(id=created_by_id) if created_by_id is not None else None task_run = MagicMock() task_run.id = "run-1" task_run.team_id = team_id @@ -55,15 +69,36 @@ def _make_task_run_mock(team_id: int = 7, created_by_id: int | None = 42, state: return task_run -class TestRefreshSandboxMcp: - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") +def _make_user_mock(user_id: int = 42) -> MagicMock: + user = MagicMock() + user.id = user_id + return user + + +_OAUTH_MINT_PATH = ( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token" +) +_PH_CONFIGS_PATH = ( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" +) +_USER_CONFIGS_PATH = ( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" +) +_SEND_REFRESH_PATH = ( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session" +) +_TIME_SLEEP_PATH = "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep" + + +class TestRefreshSandboxMcpForUser: + """Covers the actor-parameterized helper that the Slack and web follow-up + paths share. The legacy ``_refresh_sandbox_mcp`` is now a thin wrapper that + fills in ``task.created_by`` and delegates here.""" + + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] @@ -71,9 +106,10 @@ def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_c mock_send_refresh.return_value = CommandResult(success=True, status_code=200) task_run = _make_task_run_mock() - _refresh_sandbox_mcp(task_run, "read_only", auth_token="jwt") + user = _make_user_mock(user_id=42) + refresh_sandbox_mcp_for_user(task_run, user, scopes="read_only", auth_token="jwt") - mock_oauth.assert_called_once_with(task_run.task, scopes="read_only") + mock_oauth.assert_called_once_with(task_run.task, scopes="read_only", user=user) mock_ph_configs.assert_called_once_with( token="fresh-token", project_id=7, scopes="read_only", interaction_origin=None, task_id="task-1" ) @@ -86,15 +122,11 @@ def test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_c mcp_servers = mock_send_refresh.call_args.args[1] assert mcp_servers == [_make_mcp_config(token="fresh-token").to_dict()] - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_TIME_SLEEP_PATH) + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_retries_once_on_first_failure( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep ): @@ -106,20 +138,16 @@ def test_retries_once_on_first_failure( CommandResult(success=True, status_code=200), ] - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) assert mock_send_refresh.call_count == 2 mock_sleep.assert_called_once_with(REFRESH_RETRY_DELAY_SECONDS) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_TIME_SLEEP_PATH) + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_two_failures_are_non_fatal( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): @@ -129,37 +157,29 @@ def test_two_failures_are_non_fatal( mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") # Must not raise. - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) assert mock_send_refresh.call_count == 2 - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_token_mint_failure_is_non_fatal_and_skips_send( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.side_effect = RuntimeError("oauth service down") - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) mock_ph_configs.assert_not_called() mock_user_configs.assert_not_called() mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_skips_send_when_no_mcp_configs_resolved( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): @@ -167,38 +187,14 @@ def test_skips_send_when_no_mcp_configs_resolved( mock_ph_configs.return_value = [] mock_user_configs.return_value = [] - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") - def test_user_mcp_configs_skipped_when_no_creator( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh - ): - mock_oauth.return_value = "fresh-token" - mock_ph_configs.return_value = [_make_mcp_config()] - mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - - _refresh_sandbox_mcp(_make_task_run_mock(created_by_id=None), "read_only", auth_token=None) - - mock_user_configs.assert_not_called() - mock_send_refresh.assert_called_once() - - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_scopes_propagate_to_oauth_and_configs( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): @@ -207,57 +203,108 @@ def test_scopes_propagate_to_oauth_and_configs( mock_user_configs.return_value = [] mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - _refresh_sandbox_mcp(_make_task_run_mock(), "full", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="full", auth_token=None) - mock_oauth.assert_called_once_with(mock_oauth.call_args.args[0], scopes="full") + mock_oauth.assert_called_once() + assert mock_oauth.call_args.kwargs["scopes"] == "full" mock_ph_configs.assert_called_once_with( token="fresh-token", project_id=7, scopes="full", interaction_origin=None, task_id="task-1" ) + @patch(_SEND_REFRESH_PATH) + @patch(_OAUTH_MINT_PATH) + def test_legacy_wrapper_skips_when_no_creator(self, mock_oauth, mock_send_refresh): + """``_refresh_sandbox_mcp`` (the no-actor wrapper used by the web-layer + signal path) short-circuits when the task has no ``created_by``. The + prior implementation would try to mint and log a warning; the new + wrapper just skips, since the helper requires an explicit user.""" + _refresh_sandbox_mcp(_make_task_run_mock(created_by_id=None), "read_only", auth_token=None) + + mock_oauth.assert_not_called() + mock_send_refresh.assert_not_called() + class TestRefreshIntervalGate: """Refreshes within MCP_TOKEN_REFRESH_INTERVAL_SECONDS of a previous - successful issuance must be skipped without minting a new token or - contacting the sandbox.""" + successful issuance for the *same actor* must be skipped without minting a + new token or contacting the sandbox. The rate-limit key is scoped to + ``(run_id, user_id)`` so a cross-user follow-up isn't silently blocked by + a prior same-run refresh under a different identity.""" - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_SEND_REFRESH_PATH) + @patch(_OAUTH_MINT_PATH) def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): - mark_mcp_token_issued("run-1") + mark_mcp_token_issued("run-1", 42) - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) mock_oauth.assert_not_called() mock_send_refresh.assert_not_called() - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_SEND_REFRESH_PATH) + @patch(_OAUTH_MINT_PATH) + def test_identity_transition_bypasses_rate_limit(self, mock_oauth, mock_send_refresh): + """When the actor differs from the identity the sandbox currently + holds (default: the task creator), the refresh must go through even if + this actor's own rate-limit window is warm — an identity swap is never + silently skipped.""" + mark_mcp_token_issued("run-1", 43) + mock_oauth.return_value = "fresh-token" + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + with ( + patch(_PH_CONFIGS_PATH, return_value=[_make_mcp_config()]), + patch(_USER_CONFIGS_PATH, return_value=[]), + ): + refresh_sandbox_mcp_for_user( + _make_task_run_mock(), _make_user_mock(user_id=43), scopes="read_only", auth_token=None + ) + + mock_oauth.assert_called_once() + mock_send_refresh.assert_called_once() + assert get_last_sandbox_identity("run-1", "mcp") == 43 + + @patch(_SEND_REFRESH_PATH) + @patch(_OAUTH_MINT_PATH) + def test_switching_back_to_creator_refreshes(self, mock_oauth, mock_send_refresh): + """Ping-pong threads: after a teammate took over the sandbox identity, + a message from the task creator must rebind the MCP back to them — + "actor == creator" alone is not proof the sandbox is authed as the + creator.""" + mark_sandbox_identity("run-1", "mcp", 43) + mark_mcp_token_issued("run-1", 42) + mock_oauth.return_value = "fresh-token" + mock_send_refresh.return_value = CommandResult(success=True, status_code=200) + with ( + patch(_PH_CONFIGS_PATH, return_value=[_make_mcp_config()]), + patch(_USER_CONFIGS_PATH, return_value=[]), + ): + refresh_sandbox_mcp_for_user( + _make_task_run_mock(), _make_user_mock(user_id=42), scopes="read_only", auth_token=None + ) + + mock_send_refresh.assert_called_once() + assert get_last_sandbox_identity("run-1", "mcp") == 42 + + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] mock_user_configs.return_value = [] mock_send_refresh.return_value = CommandResult(success=True, status_code=200) - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) # Cache entry now exists → next refresh within the interval is gated. - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_TIME_SLEEP_PATH) + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_marks_after_successful_retry( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): @@ -269,19 +316,15 @@ def test_marks_after_successful_retry( CommandResult(success=True, status_code=200), ] - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) - assert cache.get(_mcp_token_issued_cache_key("run-1")) is True + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.time.sleep") - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.send_refresh_session") - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_user_mcp_server_configs" - ) - @patch( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.get_sandbox_ph_mcp_configs" - ) - @patch("products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token") + @patch(_TIME_SLEEP_PATH) + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_MINT_PATH) def test_does_not_mark_after_two_failures( self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): @@ -290,10 +333,10 @@ def test_does_not_mark_after_two_failures( mock_user_configs.return_value = [] mock_send_refresh.return_value = CommandResult(success=False, status_code=502, error="down") - _refresh_sandbox_mcp(_make_task_run_mock(), "read_only", auth_token=None) + refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) # Cache stays empty so the next follow-up retries the dispatch. - assert cache.get(_mcp_token_issued_cache_key("run-1")) is None + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is None class TestSendFollowupActivityRefreshOrdering: diff --git a/products/tasks/backend/temporal/process_task/utils.py b/products/tasks/backend/temporal/process_task/utils.py index 2b6974f1a83e..5f1e270613ce 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -3,7 +3,7 @@ import logging from dataclasses import dataclass, field from enum import StrEnum -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Literal, Optional from urllib.parse import urlparse from django.conf import settings @@ -340,23 +340,67 @@ def get_sandbox_snapshot_metadata(snapshot: SandboxSnapshot) -> SnapshotMetadata MCP_TOKEN_REFRESH_INTERVAL_SECONDS = TOKEN_EXPIRATION_SECONDS / 2 # 3 hours -def _mcp_token_issued_cache_key(run_id: str) -> str: - return f"posthog_ai:task-run-mcp-token-issued:{run_id}" +def sandbox_identity_scope(run_id: str, state: dict[str, Any] | None) -> str: + """Cache scope for the marks describing what a run's sandbox holds. + + The freshness and identity marks below describe the state of a *sandbox*, + so they key on the sandbox id: a replacement sandbox (fresh provision, + snapshot restore, mid-run workflow retry) starts unmarked by construction + and therefore defaults back to the boot-time creator identity — nothing + ever needs clearing. Falls back to the run id for runs that haven't + recorded a sandbox id in their state yet. + """ + return (state or {}).get("sandbox_id") or run_id + + +def _mcp_token_issued_cache_key(scope: str, user_id: int) -> str: + return f"posthog_ai:sandbox-mcp-token-issued:{scope}:{user_id}" + + +def mark_mcp_token_issued(scope: str, user_id: int) -> None: + """Record that a fresh MCP token for ``user_id`` was issued to the sandbox. + + ``scope`` comes from ``sandbox_identity_scope``. The entry self-expires + after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so ``should_refresh_mcp_token`` + returns True again past that window. + """ + get_tasks_cache().set(_mcp_token_issued_cache_key(scope, user_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + + +def should_refresh_mcp_token(scope: str, user_id: int) -> bool: + """True when no MCP token for ``user_id`` was issued to the sandbox within + the last MCP_TOKEN_REFRESH_INTERVAL_SECONDS window.""" + return get_tasks_cache().get(_mcp_token_issued_cache_key(scope, user_id)) is None + + +# How long a sandbox's swapped identity is remembered — comfortably past any +# plausible sandbox lifetime. On eviction the identity is assumed to be the +# boot-time one (the task creator / the task's own integration). +SANDBOX_IDENTITY_TTL_SECONDS = 7 * 24 * 60 * 60 + +SandboxIdentityKind = Literal["mcp"] + + +def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: + return f"posthog_ai:sandbox-{kind}-identity:{scope}" -def mark_mcp_token_issued(run_id: str) -> None: - """Record that a fresh MCP token was issued to the sandbox for this run. +def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: + """Record which identity the sandbox currently holds for a credential kind. - The cache entry self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so - `should_refresh_mcp_token` returns True again past that window. + ``mcp`` stores the user id the OAuth token was minted for. Written on + every successful rebind so identity transitions (a different Slack actor + taking over the thread) are detected and never silently skipped by the + per-credential freshness rate limits. ``scope`` comes from + ``sandbox_identity_scope``, so a replacement sandbox starts unmarked. """ - get_tasks_cache().set(_mcp_token_issued_cache_key(run_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + get_tasks_cache().set(_sandbox_identity_cache_key(scope, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) -def should_refresh_mcp_token(run_id: str) -> bool: - """Return True if no MCP token has been issued for this run within the - last MCP_TOKEN_REFRESH_INTERVAL_SECONDS window.""" - return get_tasks_cache().get(_mcp_token_issued_cache_key(run_id)) is None +def get_last_sandbox_identity(scope: str, kind: SandboxIdentityKind) -> int | str | None: + """Return the identity the sandbox's credentials were last bound to for a + kind, or None when unknown (never swapped, or the entry was evicted).""" + return get_tasks_cache().get(_sandbox_identity_cache_key(scope, kind)) @dataclass(frozen=True)