From 24801027353c6904a1d7fb7c684cee9a7d36b54c Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Wed, 8 Jul 2026 16:43:45 +0200 Subject: [PATCH 1/7] refactor(slack-app): extract shared mention orchestration into helpers package Move the per-message orchestration out of PostHogCodeSlackMentionWorkflow into helpers/process_mention_message.py: process_mention_message plus the MentionSignalState holder and a MentionSignalHandlersMixin carrying the interactive signal handlers. The workflow class becomes a thin shell with an unchanged defn name, signals, and command sequence, so in-flight histories replay identically. helpers.py becomes the helpers package __init__ so activity-side imports are unchanged. No behavior change. --- .../{helpers.py => helpers/__init__.py} | 0 .../helpers/process_mention_message.py | 406 ++++++++++++++++++ .../slack_app/posthog_code_slack_mention.py | 354 +-------------- 3 files changed, 417 insertions(+), 343 deletions(-) rename posthog/temporal/ai/slack_app/{helpers.py => helpers/__init__.py} (100%) create mode 100644 posthog/temporal/ai/slack_app/helpers/process_mention_message.py diff --git a/posthog/temporal/ai/slack_app/helpers.py b/posthog/temporal/ai/slack_app/helpers/__init__.py similarity index 100% rename from posthog/temporal/ai/slack_app/helpers.py rename to posthog/temporal/ai/slack_app/helpers/__init__.py 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) From 66855c7a2be7be7f2da100d40f9594ecd6ad1d7c Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Wed, 8 Jul 2026 16:51:29 +0200 Subject: [PATCH 2/7] feat(slack-app): queue slack messages per conversation behind feature flag One SlackAppMentionWorkflow per thread/DM serializes message processing: the webhook signal-with-starts the conversation workflow, messages queue as signals carrying their own routing-time-resolved user, and the loop feeds them one at a time through the shared mention orchestration. The workflow idles out after 30s with an empty queue; the next message starts a fresh instance that picks the conversation back up via SlackThreadTaskMapping. Gated by the slack-app-queue-workflow flag; off, dispatch stays one workflow per message. --- posthog/temporal/ai/__init__.py | 2 + posthog/temporal/ai/slack_app/__init__.py | 2 + .../ai/slack_app/slack_app_mention.py | 116 ++++++ posthog/temporal/ai/slack_app/types.py | 17 +- .../test_slack_app_mention_workflow.py | 390 ++++++++++++++++++ products/slack_app/backend/api.py | 39 ++ products/slack_app/backend/feature_flags.py | 26 ++ .../tests/test_posthog_code_event_handler.py | 69 ++++ 8 files changed, 660 insertions(+), 1 deletion(-) create mode 100644 posthog/temporal/ai/slack_app/slack_app_mention.py create mode 100644 posthog/temporal/tests/ai/slack_app/test_slack_app_mention_workflow.py 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..e360506aa3bc 100644 --- a/posthog/temporal/ai/slack_app/__init__.py +++ b/posthog/temporal/ai/slack_app/__init__.py @@ -40,6 +40,7 @@ PostHogCodeSlackMentionCommandWorkflowInputs, PostHogCodeSlackMentionWorkflowInputs, PostHogSlackInboxOnboardingInputs, + SlackAppMentionWorkflowInputs, SlackRepoSelectionOutcome, ) @@ -78,6 +79,7 @@ "PostHogCodeSlackMentionWorkflowInputs", "PostHogSlackInboxOnboardingInputs", "SLACK_APP_ACTIVITIES", + "SlackAppMentionWorkflowInputs", "SlackRepoSelectionOutcome", "block_posthog_code_task_if_no_personal_github_activity", "cascade_posthog_code_repository_activity", 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..ec18d62c22be --- /dev/null +++ b/posthog/temporal/ai/slack_app/slack_app_mention.py @@ -0,0 +1,116 @@ +# Workflows in this module run on the max-ai temporal task queue. +import json +from datetime import timedelta + +from temporalio import workflow + +from posthog.temporal.ai.slack_app import derive_mention_workflow_id +from posthog.temporal.ai.slack_app.helpers.process_mention_message import ( + MentionSignalHandlersMixin, + process_mention_message, +) +from posthog.temporal.ai.slack_app.types import 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) + + @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) + self._signals.reset() + # 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..938f1f10c467 100644 --- a/posthog/temporal/ai/slack_app/types.py +++ b/posthog/temporal/ai/slack_app/types.py @@ -5,7 +5,7 @@ creating an import cycle with the workflow modules. """ -from dataclasses import dataclass, fields +from dataclasses import dataclass, field, fields from typing import Any, Literal @@ -64,6 +64,21 @@ 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) + + @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..4c87c47d61d6 --- /dev/null +++ b/posthog/temporal/tests/ai/slack_app/test_slack_app_mention_workflow.py @@ -0,0 +1,390 @@ +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 ( + 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]] = [] + # 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: + 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 + + return [ + 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")] + + +@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(): + 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 == [] + + +@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..331c15e01685 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -39,6 +39,7 @@ from posthog.temporal.ai.slack_app import ( PostHogCodeSlackMentionCommandWorkflowInputs, PostHogCodeSlackMentionWorkflowInputs, + SlackAppMentionWorkflowInputs, derive_mention_workflow_id, ) from posthog.temporal.ai.slack_app.posthog_code_slack_interactivity import ( @@ -47,6 +48,10 @@ ) 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 +61,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 +1423,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,6 +1445,7 @@ def _start_posthog_code_workflow( task_queue=settings.TASKS_TASK_QUEUE, id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, + **({"start_signal": start_signal, "start_signal_args": start_signal_args or []} if start_signal else {}), ) ) @@ -2413,6 +2429,29 @@ 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], + ) + 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_posthog_code_event_handler.py b/products/slack_app/backend/tests/test_posthog_code_event_handler.py index f3baf258671c..920a545a2ecd 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,72 @@ 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.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 + ): + # 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 + + 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" From bfbd9ea2b777db642b820ff39370670b0f937292 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Wed, 8 Jul 2026 16:52:07 +0200 Subject: [PATCH 3/7] feat(slack-app): react hourglass on queued messages, eyes when processing Mentions and DMs get an :hourglass: reaction when they land in the conversation queue and an :eyes: swap when their turn starts. Untagged thread followups stay reaction-free since most are dropped by the chitchat classifier. The reaction activity takes a single pydantic input model; reactions are best-effort and never stall the queue. --- posthog/temporal/ai/slack_app/__init__.py | 9 +++++ .../ai/slack_app/activities/__init__.py | 2 + .../ai/slack_app/activities/messaging.py | 38 ++++++++++++++++++- .../temporal/ai/slack_app/helpers/__init__.py | 9 +++++ .../ai/slack_app/slack_app_mention.py | 38 +++++++++++++++++-- posthog/temporal/ai/slack_app/types.py | 22 +++++++++++ .../test_slack_app_mention_workflow.py | 14 ++++++- products/slack_app/backend/api.py | 25 ++++++++++++ .../tests/test_posthog_code_event_handler.py | 8 +++- 9 files changed, 159 insertions(+), 6 deletions(-) diff --git a/posthog/temporal/ai/slack_app/__init__.py b/posthog/temporal/ai/slack_app/__init__.py index e360506aa3bc..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,6 +35,9 @@ 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, @@ -65,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", @@ -95,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/helpers/__init__.py b/posthog/temporal/ai/slack_app/helpers/__init__.py index 3f5606961a07..cc347ce33f3a 100644 --- a/posthog/temporal/ai/slack_app/helpers/__init__.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/slack_app_mention.py b/posthog/temporal/ai/slack_app/slack_app_mention.py index ec18d62c22be..bbdfafe75ffb 100644 --- a/posthog/temporal/ai/slack_app/slack_app_mention.py +++ b/posthog/temporal/ai/slack_app/slack_app_mention.py @@ -2,14 +2,19 @@ import json from datetime import timedelta -from temporalio import workflow +from temporalio import exceptions, workflow +from temporalio.common import RetryPolicy -from posthog.temporal.ai.slack_app import derive_mention_workflow_id +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 PostHogCodeSlackMentionWorkflowInputs, SlackAppMentionWorkflowInputs +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 @@ -74,6 +79,32 @@ def parse_inputs(inputs: list[str]) -> SlackAppMentionWorkflowInputs: ] 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: @@ -102,6 +133,7 @@ async def run(self, inputs: SlackAppMentionWorkflowInputs) -> None: message = self._queue.pop(0) 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) diff --git a/posthog/temporal/ai/slack_app/types.py b/posthog/temporal/ai/slack_app/types.py index 938f1f10c467..a59e02a92aa1 100644 --- a/posthog/temporal/ai/slack_app/types.py +++ b/posthog/temporal/ai/slack_app/types.py @@ -8,6 +8,8 @@ from dataclasses import dataclass, field, fields from typing import Any, Literal +from pydantic import BaseModel + @dataclass class PostHogSlackInboxOnboardingInputs: @@ -79,6 +81,26 @@ class SlackAppMentionWorkflowInputs: 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 index 4c87c47d61d6..54fc9c6d2681 100644 --- 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 @@ -13,6 +13,7 @@ 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, @@ -40,6 +41,8 @@ class _Recorder: def __init__(self) -> None: # (ts, repository) per create-task call, in execution order. self.created: list[tuple[str, str | None]] = [] + # 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). @@ -190,7 +193,12 @@ async def resolve_user( ) -> 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, @@ -277,6 +285,8 @@ async def test_queued_messages_process_serially_in_arrival_order(): 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"] @pytest.mark.asyncio @@ -315,7 +325,7 @@ async def test_duplicate_slack_event_id_is_processed_once(): @pytest.mark.asyncio -async def test_untagged_followup_forwards_without_task_creation(): +async def test_untagged_followup_forwards_without_task_creation_or_reactions(): rec = _Recorder() rec.forward_results["1.1"] = True @@ -325,6 +335,8 @@ async def test_untagged_followup_forwards_without_task_creation(): 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 diff --git a/products/slack_app/backend/api.py b/products/slack_app/backend/api.py index 331c15e01685..4ce9c229fffa 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -37,11 +37,13 @@ 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, @@ -1450,6 +1452,25 @@ def _start_posthog_code_workflow( ) +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"}, @@ -2451,6 +2472,10 @@ def _start_mention_workflow( 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 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 920a545a2ecd..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 @@ -1292,12 +1292,13 @@ def setUp(self): ("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 + 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 @@ -1306,6 +1307,7 @@ def test_flag_on_signal_with_starts_conversation_workflow( 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 @@ -1322,3 +1324,7 @@ def test_flag_on_signal_with_starts_conversation_workflow( 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" + ) From 9de847ca75640c05919237e60e24be56f3b6f2c2 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Wed, 8 Jul 2026 17:29:08 +0200 Subject: [PATCH 4/7] fix(slack-app): pass start_signal kwargs directly for mypy overload match --- products/slack_app/backend/api.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/products/slack_app/backend/api.py b/products/slack_app/backend/api.py index 4ce9c229fffa..95013e5fbaca 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -1447,7 +1447,9 @@ def _start_posthog_code_workflow( task_queue=settings.TASKS_TASK_QUEUE, id_conflict_policy=WorkflowIDConflictPolicy.USE_EXISTING, id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE, - **({"start_signal": start_signal, "start_signal_args": start_signal_args or []} if start_signal else {}), + # None / [] match the SDK defaults, so a plain start stays a plain start. + start_signal=start_signal, + start_signal_args=start_signal_args or [], ) ) From 32e0ed75f7bfc3f4613b1ed34dd6d4dad9d4f7ee Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Fri, 10 Jul 2026 15:00:23 +0200 Subject: [PATCH 5/7] feat(tasks): rebind sandbox MCP identity to a given user mid-run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds rebind_sandbox_identity_for_user to the tasks facade: mint a PostHog MCP OAuth token for the given user and push fresh MCP configs into the run's live sandbox, so subsequent agent actions attribute to that user. Best-effort by contract — failures are logged, never raised. The sandbox's current identity is tracked per run in the tasks cache so transitions always bypass the token-freshness rate limit (including switching back to the task creator in ping-pong threads), and every sandbox-creation path clears the marks — a fresh sandbox boots with creator credentials, so a mark surviving a mid-run workflow retry or dead-restore fallback would block the actor's next rebind. --- products/tasks/backend/facade/api.py | 36 ++ products/tasks/backend/temporal/oauth.py | 5 +- .../create_sandbox_from_snapshot.py | 6 + .../activities/get_sandbox_for_repository.py | 6 + .../activities/provision_sandbox.py | 12 + .../activities/send_followup_to_sandbox.py | 152 ++++++--- .../tests/test_send_followup_to_sandbox.py | 322 +++++++++++------- .../backend/temporal/process_task/utils.py | 47 ++- 8 files changed, 401 insertions(+), 185 deletions(-) diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index d2f055348235..6afcf757c698 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -176,6 +176,7 @@ "read_task_run_artifact", "read_task_run_logs", "redeem_code_invite", + "rebind_sandbox_identity_for_user", "redispatch_task_run", "refresh_team_code_workstreams", "relay_task_run_message", @@ -4581,3 +4582,38 @@ 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: + 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..bfcf6ed1fccb 100644 --- a/products/tasks/backend/temporal/oauth.py +++ b/products/tasks/backend/temporal/oauth.py @@ -18,10 +18,11 @@ "create_oauth_access_token", "create_oauth_access_token_for_user", "create_wizard_oauth_access_token", + "oauth_application_for_task", ] -def _oauth_application_for_task(task: Task) -> SandboxOAuthApplication: +def oauth_application_for_task(task: Task) -> SandboxOAuthApplication: if task.origin_product == Task.OriginProduct.POSTHOG_AI: return "posthog_ai" return "array" @@ -43,7 +44,7 @@ def create_oauth_access_token(task: Task, *, scopes: PosthogMcpScopes = "read_on task.created_by, task.team_id, scopes=scopes, - application=_oauth_application_for_task(task), + application=oauth_application_for_task(task), ) diff --git a/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py b/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py index 7ef1ae7fd3a2..6a8c8063018d 100644 --- a/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py +++ b/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py @@ -17,6 +17,7 @@ from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.utils import ( build_sandbox_environment_variables, + clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_github_token, get_sandbox_name_for_task, @@ -113,6 +114,11 @@ def create_sandbox_from_snapshot(input: CreateSandboxFromSnapshotInput) -> Creat sandbox_environment=sandbox_env, ) environment_variables.update(get_git_identity_env_vars(task, ctx.state)) + # A brand-new sandbox boots with the task's own (creator) credentials; + # forget any per-message identity swap recorded against this run_id by + # a previous sandbox (mid-run workflow retry, dead-restore fallback), + # or the marks would diverge from what this sandbox actually holds. + clear_sandbox_identities(str(ctx.run_id)) config = SandboxConfig( name=get_sandbox_name_for_task(ctx.task_id), 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..f42605b12e0c 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 @@ -26,6 +26,7 @@ from products.tasks.backend.temporal.oauth import create_oauth_access_token from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.utils import ( + clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_api_url, get_sandbox_github_token, @@ -219,6 +220,11 @@ 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)) + # A brand-new sandbox boots with the task's own (creator) credentials; + # forget any per-message identity swap recorded against this run_id by + # a previous sandbox (mid-run workflow retry, dead-restore fallback), + # or the marks would diverge from what this sandbox actually holds. + clear_sandbox_identities(str(ctx.run_id)) run_state = parse_run_state(ctx.state) diff --git a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py index 0b825ae5acd4..bd21359fdead 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -28,6 +28,7 @@ from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.sandbox_credentials import set_git_remote_token from products.tasks.backend.temporal.process_task.utils import ( + clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_api_url, get_sandbox_github_token, @@ -252,6 +253,11 @@ def _build_environment_variables( environment_variables.update(NETWORK_RESTRICTED_AGENT_ENV) environment_variables.update(get_git_identity_env_vars(task, ctx.state)) + # A brand-new sandbox boots with the task's own (creator) credentials; + # forget any per-message identity swap recorded against this run_id by + # a previous sandbox (mid-run workflow retry, dead-restore fallback), + # or the marks would diverge from what this sandbox actually holds. + clear_sandbox_identities(str(ctx.run_id)) run_state = parse_run_state(ctx.state) if run_state.resume_from_run_id: @@ -651,6 +657,12 @@ def inject_fresh_tokens_on_resume(input: InjectFreshTokensOnResumeInput) -> None ): task = _load_task(ctx) + # Resume re-applies the boot-time (task creator) credentials below, so + # forget any per-message identity swap a Slack actor made before the + # snapshot — a stale mark would block their next rebind as a + # same-identity no-op while the refresh loop pulls the other way. + clear_sandbox_identities(ctx.run_id) + github_token = "" if ctx.has_github_credentials: try: 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..7da844b13044 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 @@ -24,16 +24,21 @@ from products.tasks.backend.logic.stream.redis_stream import get_task_run_stream_key from products.tasks.backend.models import TaskRun 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.oauth import create_oauth_access_token_for_user, oauth_application_for_task 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, 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 +218,69 @@ 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 ``(run_id, user_id)``: skipped if the same user already + had a token issued 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 + current_identity = get_last_sandbox_identity(run_id, "mcp") or task.created_by_id + identity_changed = user.id != current_identity + rate_limit_key = f"{run_id}:{user.id}" + if not identity_changed and not should_refresh_mcp_token(rate_limit_key): + 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_for_user( + user, + task_run.team_id, + scopes=scopes, + application=oauth_application_for_task(task), + ) 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 +290,54 @@ 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) + # Mark the window anyway: on hosts where no MCP URL resolves and the + # user has no installs, every message would otherwise repeat the + # mint-and-discard above. + mark_mcp_token_issued(rate_limit_key) + 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): + if attempt > 1: + logger.info( + "refresh_mcp_retrying", + run_id=run_id, + user_id=user.id, + error=result.error if result else None, + status_code=result.status_code if result else None, + ) + time.sleep(REFRESH_RETRY_DELAY_SECONDS) + result = send_refresh_session( + task_run, + mcp_servers, + auth_token=auth_token, + timeout=REFRESH_TIMEOUT_SECONDS, + ) + if result.success: + mark_mcp_token_issued(rate_limit_key) + mark_sandbox_identity(run_id, "mcp", user.id) + logger.info("refresh_mcp_delivered", run_id=run_id, user_id=user.id, attempts=attempt) + return logger.warning( "refresh_mcp_failed", run_id=run_id, - error=retry.error, - status_code=retry.status_code, + user_id=user.id, + error=result.error if result else None, + status_code=result.status_code if result else None, ) 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..28f6648aad88 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, + clear_sandbox_identities, + get_last_sandbox_identity, mark_mcp_token_issued, + mark_sandbox_identity, ) pytestmark = pytest.mark.django_db @@ -25,10 +29,21 @@ @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"), + _mcp_token_issued_cache_key("run-1:42"), + _mcp_token_issued_cache_key("run-1:43"), + ] + for key in keys: + cache.delete(key) + clear_sandbox_identities("run-1") yield - cache.delete(_mcp_token_issued_cache_key("run-1")) + for key in keys: + cache.delete(key) + clear_sandbox_identities("run-1") def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: @@ -43,6 +58,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,25 +71,52 @@ 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 test_success_path_single_call(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): +def _make_user_mock(user_id: int = 42) -> MagicMock: + user = MagicMock() + user.id = user_id + return user + + +_OAUTH_FOR_USER_PATH = "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_user" +_OAUTH_APP_PATH = ( + "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.oauth_application_for_task" +) +_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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) + def test_success_path_single_call( + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh + ): mock_oauth.return_value = "fresh-token" + mock_oauth_app.return_value = "array" mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] mock_user_configs.return_value = [] 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(user, 7, scopes="read_only", application="array") mock_ph_configs.assert_called_once_with( token="fresh-token", project_id=7, scopes="read_only", interaction_origin=None, task_id="task-1" ) @@ -86,17 +129,14 @@ 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_retries_once_on_first_failure( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -106,22 +146,19 @@ 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_two_failures_are_non_fatal( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -129,137 +166,165 @@ 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_token_mint_failure_is_non_fatal_and_skips_send( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_oauth_app, 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_skips_send_when_no_mcp_configs_resolved( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.return_value = "fresh-token" 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_scopes_propagate_to_oauth_and_configs( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.return_value = "fresh-token" + mock_oauth_app.return_value = "array" 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(), "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_with(mock_oauth.call_args.args[0], 7, scopes="full", application="array") 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_FOR_USER_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_FOR_USER_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") - def test_marks_after_successful_refresh(self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh): + @patch(_SEND_REFRESH_PATH) + @patch(_OAUTH_FOR_USER_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=[]), + patch(_OAUTH_APP_PATH, return_value="array"), + ): + 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_FOR_USER_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=[]), + patch(_OAUTH_APP_PATH, return_value="array"), + ): + 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) + def test_marks_after_successful_refresh( + self, mock_oauth, mock_oauth_app, 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 - - @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") + assert cache.get(_mcp_token_issued_cache_key("run-1:42")) is True + + @patch(_TIME_SLEEP_PATH) + @patch(_SEND_REFRESH_PATH) + @patch(_USER_CONFIGS_PATH) + @patch(_PH_CONFIGS_PATH) + @patch(_OAUTH_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_marks_after_successful_retry( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -269,31 +334,28 @@ 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_APP_PATH) + @patch(_OAUTH_FOR_USER_PATH) def test_does_not_mark_after_two_failures( - self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): 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=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..f817ac70413a 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 @@ -359,6 +359,51 @@ def should_refresh_mcp_token(run_id: str) -> bool: return get_tasks_cache().get(_mcp_token_issued_cache_key(run_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"] + +SANDBOX_IDENTITY_KINDS: tuple[SandboxIdentityKind, ...] = ("mcp",) + + +def _sandbox_identity_cache_key(run_id: str, kind: SandboxIdentityKind) -> str: + return f"posthog_ai:task-run-{kind}-identity:{run_id}" + + +def mark_sandbox_identity(run_id: str, kind: SandboxIdentityKind, value: int | str) -> None: + """Record which identity the sandbox currently holds for a credential kind. + + ``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. + """ + get_tasks_cache().set(_sandbox_identity_cache_key(run_id, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) + + +def get_last_sandbox_identity(run_id: 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(run_id, kind)) + + +def clear_sandbox_identities(run_id: str) -> None: + """Forget a run's swapped identities. + + Called when a sandbox is restored from a snapshot: the resume path + re-applies the boot-time (task creator) credentials, so any remembered + swap would diverge from what the sandbox actually holds — blocking the + actor's next rebind as a same-identity no-op while the refresh loop pulls + the other way. + """ + cache = get_tasks_cache() + for kind in SANDBOX_IDENTITY_KINDS: + cache.delete(_sandbox_identity_cache_key(run_id, kind)) + + @dataclass(frozen=True) class McpServerConfig: """Configuration for a remote MCP server matching the ACP McpServer schema. From ad4731e916d41f0315ef4e725c9812d73b9bb5c8 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Fri, 10 Jul 2026 15:00:25 +0200 Subject: [PATCH 6/7] feat(slack-app): bind sandbox MCP identity to each message's actor Messages dispatched to the per-conversation queue workflow carry per_message_identity: the sandbox connection JWT is minted for the message's actor and the facade rebind runs before delivery, so insights, dashboards, and other PostHog writes attribute to whoever actually spoke. Rebind failures are best-effort and never drop the message. The legacy per-message workflow leaves the field False and keeps every credential bound to the task creator, exactly as before. --- .../ai/slack_app/activities/task_creation.py | 45 +++++++-- posthog/temporal/ai/slack_app/types.py | 6 ++ products/slack_app/backend/api.py | 3 + .../backend/tests/test_followup_forwarding.py | 98 +++++++++++++++++-- .../tests/test_posthog_code_event_handler.py | 3 + 5 files changed, 138 insertions(+), 17 deletions(-) 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/types.py b/posthog/temporal/ai/slack_app/types.py index a59e02a92aa1..209f3f351c3c 100644 --- a/posthog/temporal/ai/slack_app/types.py +++ b/posthog/temporal/ai/slack_app/types.py @@ -35,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: diff --git a/products/slack_app/backend/api.py b/products/slack_app/backend/api.py index 95013e5fbaca..7ec74b2f3894 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -2458,6 +2458,9 @@ def _start_mention_workflow( # 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): + # Queue-dispatched messages carry per-message identity: the sandbox + # JWT and credential rebinds follow each message's actor. + workflow_inputs.per_message_identity = True # 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 diff --git a/products/slack_app/backend/tests/test_followup_forwarding.py b/products/slack_app/backend/tests/test_followup_forwarding.py index 75139413a62c..e6eb4aaf3650 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 ) @@ -742,6 +787,7 @@ def test_cross_user_followup_authorized_prefixes_actor_name( ] assert not post_calls + @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", @@ -750,7 +796,7 @@ def test_cross_user_followup_authorized_prefixes_actor_name( @patch("products.slack_app.backend.api.resolve_slack_user") @patch("posthog.models.integration.SlackIntegration") def test_cross_user_followup_falls_back_to_email_when_no_full_name( - self, mock_slack_cls, mock_resolve, mock_send, mock_token + self, mock_slack_cls, mock_resolve, mock_send, mock_token, mock_rebind ): self._create_mapping(mentioning_user="U_ALICE") bob = User.objects.create(email="bob@test.com") # no full name @@ -763,6 +809,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 +902,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 +919,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 0178e6ba598b..f3026d45e96a 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 @@ -1324,6 +1324,9 @@ def test_flag_on_signal_with_starts_conversation_workflow( message = call.kwargs["start_signal_args"][0] assert message.user_id == self.user.id assert message.event["ts"] == "1234.5678" + # Queue-dispatched messages carry per-message identity so the sandbox + # JWT and MCP/GitHub rebinds follow each message's actor. + assert message.per_message_identity is True # 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" From a08508da6d6d92521135e36bfc878b51c6506e91 Mon Sep 17 00:00:00 2001 From: Vojta Bartos Date: Mon, 13 Jul 2026 13:10:26 +0200 Subject: [PATCH 7/7] refactor(tasks): key sandbox identity marks by sandbox, not run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The freshness and identity marks describe what a sandbox holds, so key them by sandbox id: a replacement sandbox (fresh provision, snapshot restore, mid-run workflow retry) starts unmarked by construction and defaults back to the boot-time creator identity — deleting all the creation-path clear calls and the clear helper outright. Make the cache helpers own the composite key as (scope, user_id) signatures; the boot-time mark in start_agent_server had silently kept the old single-run shape after the rekey, orphaning the write and losing the post-boot refresh suppression. Stamp per_message_identity inside the queue workflow loop instead of at dispatch, so no future dispatcher into the workflow can forget it. Mint the rebind token via create_oauth_access_token with a user override rather than re-composing its body, and record both marks on the no-configs bail so hosts with no resolvable MCP URL stop re-minting on every cross-user message. --- .../ai/slack_app/slack_app_mention.py | 4 + .../test_slack_app_mention_workflow.py | 6 ++ products/slack_app/backend/api.py | 3 - .../backend/tests/test_followup_forwarding.py | 3 +- .../tests/test_posthog_code_event_handler.py | 3 - products/tasks/backend/facade/api.py | 5 +- products/tasks/backend/temporal/oauth.py | 16 ++-- .../create_sandbox_from_snapshot.py | 6 -- .../activities/get_sandbox_for_repository.py | 7 -- .../activities/provision_sandbox.py | 12 --- .../activities/send_followup_to_sandbox.py | 66 ++++++------- .../activities/start_agent_server.py | 8 +- .../tests/test_send_followup_to_sandbox.py | 93 ++++++++----------- .../backend/temporal/process_task/utils.py | 67 +++++++------ 14 files changed, 131 insertions(+), 168 deletions(-) diff --git a/posthog/temporal/ai/slack_app/slack_app_mention.py b/posthog/temporal/ai/slack_app/slack_app_mention.py index bbdfafe75ffb..8da1b910f888 100644 --- a/posthog/temporal/ai/slack_app/slack_app_mention.py +++ b/posthog/temporal/ai/slack_app/slack_app_mention.py @@ -132,6 +132,10 @@ async def run(self, inputs: SlackAppMentionWorkflowInputs) -> None: 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 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 index 54fc9c6d2681..e629f41420dd 100644 --- 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 @@ -41,6 +41,8 @@ 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. @@ -61,6 +63,7 @@ def _fake_activities(rec: _Recorder) -> list: 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") @@ -287,6 +290,9 @@ async def test_queued_messages_process_serially_in_arrival_order(): 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 diff --git a/products/slack_app/backend/api.py b/products/slack_app/backend/api.py index 7ec74b2f3894..95013e5fbaca 100644 --- a/products/slack_app/backend/api.py +++ b/products/slack_app/backend/api.py @@ -2458,9 +2458,6 @@ def _start_mention_workflow( # 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): - # Queue-dispatched messages carry per-message identity: the sandbox - # JWT and credential rebinds follow each message's actor. - workflow_inputs.per_message_identity = True # 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 diff --git a/products/slack_app/backend/tests/test_followup_forwarding.py b/products/slack_app/backend/tests/test_followup_forwarding.py index e6eb4aaf3650..b1d032727112 100644 --- a/products/slack_app/backend/tests/test_followup_forwarding.py +++ b/products/slack_app/backend/tests/test_followup_forwarding.py @@ -787,7 +787,6 @@ def test_legacy_dispatch_keeps_creator_identity( ] assert not post_calls - @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", @@ -796,7 +795,7 @@ def test_legacy_dispatch_keeps_creator_identity( @patch("products.slack_app.backend.api.resolve_slack_user") @patch("posthog.models.integration.SlackIntegration") def test_cross_user_followup_falls_back_to_email_when_no_full_name( - self, mock_slack_cls, mock_resolve, mock_send, mock_token, mock_rebind + self, mock_slack_cls, mock_resolve, mock_send, mock_token ): self._create_mapping(mentioning_user="U_ALICE") bob = User.objects.create(email="bob@test.com") # no full name 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 f3026d45e96a..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 @@ -1324,9 +1324,6 @@ def test_flag_on_signal_with_starts_conversation_workflow( message = call.kwargs["start_signal_args"][0] assert message.user_id == self.user.id assert message.event["ts"] == "1234.5678" - # Queue-dispatched messages carry per-message identity so the sandbox - # JWT and MCP/GitHub rebinds follow each message's actor. - assert message.per_message_identity is True # 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 6afcf757c698..2608fe15a0a4 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -175,8 +175,8 @@ "presign_task_run_artifact", "read_task_run_artifact", "read_task_run_logs", - "redeem_code_invite", "rebind_sandbox_identity_for_user", + "redeem_code_invite", "redispatch_task_run", "refresh_team_code_workstreams", "relay_task_run_message", @@ -4614,6 +4614,9 @@ def rebind_sandbox_identity_for_user( 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 bfcf6ed1fccb..e1cb842e8201 100644 --- a/products/tasks/backend/temporal/oauth.py +++ b/products/tasks/backend/temporal/oauth.py @@ -18,22 +18,24 @@ "create_oauth_access_token", "create_oauth_access_token_for_user", "create_wizard_oauth_access_token", - "oauth_application_for_task", ] -def oauth_application_for_task(task: Task) -> SandboxOAuthApplication: +def _oauth_application_for_task(task: Task) -> SandboxOAuthApplication: if task.origin_product == Task.OriginProduct.POSTHOG_AI: return "posthog_ai" 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}, @@ -41,10 +43,10 @@ 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), + application=_oauth_application_for_task(task), ) diff --git a/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py b/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py index 6a8c8063018d..7ef1ae7fd3a2 100644 --- a/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py +++ b/products/tasks/backend/temporal/process_task/activities/create_sandbox_from_snapshot.py @@ -17,7 +17,6 @@ from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.utils import ( build_sandbox_environment_variables, - clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_github_token, get_sandbox_name_for_task, @@ -114,11 +113,6 @@ def create_sandbox_from_snapshot(input: CreateSandboxFromSnapshotInput) -> Creat sandbox_environment=sandbox_env, ) environment_variables.update(get_git_identity_env_vars(task, ctx.state)) - # A brand-new sandbox boots with the task's own (creator) credentials; - # forget any per-message identity swap recorded against this run_id by - # a previous sandbox (mid-run workflow retry, dead-restore fallback), - # or the marks would diverge from what this sandbox actually holds. - clear_sandbox_identities(str(ctx.run_id)) config = SandboxConfig( name=get_sandbox_name_for_task(ctx.task_id), 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 f42605b12e0c..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 @@ -26,7 +26,6 @@ from products.tasks.backend.temporal.oauth import create_oauth_access_token from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.utils import ( - clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_api_url, get_sandbox_github_token, @@ -220,12 +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)) - # A brand-new sandbox boots with the task's own (creator) credentials; - # forget any per-message identity swap recorded against this run_id by - # a previous sandbox (mid-run workflow retry, dead-restore fallback), - # or the marks would diverge from what this sandbox actually holds. - clear_sandbox_identities(str(ctx.run_id)) - 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/provision_sandbox.py b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py index bd21359fdead..0b825ae5acd4 100644 --- a/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py +++ b/products/tasks/backend/temporal/process_task/activities/provision_sandbox.py @@ -28,7 +28,6 @@ from products.tasks.backend.temporal.observability import emit_agent_log, log_activity_execution from products.tasks.backend.temporal.process_task.sandbox_credentials import set_git_remote_token from products.tasks.backend.temporal.process_task.utils import ( - clear_sandbox_identities, get_git_identity_env_vars, get_sandbox_api_url, get_sandbox_github_token, @@ -253,11 +252,6 @@ def _build_environment_variables( environment_variables.update(NETWORK_RESTRICTED_AGENT_ENV) environment_variables.update(get_git_identity_env_vars(task, ctx.state)) - # A brand-new sandbox boots with the task's own (creator) credentials; - # forget any per-message identity swap recorded against this run_id by - # a previous sandbox (mid-run workflow retry, dead-restore fallback), - # or the marks would diverge from what this sandbox actually holds. - clear_sandbox_identities(str(ctx.run_id)) run_state = parse_run_state(ctx.state) if run_state.resume_from_run_id: @@ -657,12 +651,6 @@ def inject_fresh_tokens_on_resume(input: InjectFreshTokensOnResumeInput) -> None ): task = _load_task(ctx) - # Resume re-applies the boot-time (task creator) credentials below, so - # forget any per-message identity swap a Slack actor made before the - # snapshot — a stale mark would block their next rebind as a - # same-identity no-op while the refresh loop pulls the other way. - clear_sandbox_identities(ctx.run_id) - github_token = "" if ctx.has_github_credentials: try: 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 7da844b13044..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 @@ -24,13 +24,14 @@ from products.tasks.backend.logic.stream.redis_stream import get_task_run_stream_key from products.tasks.backend.models import TaskRun 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_for_user, oauth_application_for_task +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, ) @@ -245,12 +246,13 @@ def refresh_sandbox_mcp_for_user( Retries once on failure, then logs and returns. Never raises — a failed refresh should not block an otherwise-valid follow-up. - Rate-limited per ``(run_id, user_id)``: skipped if the same user already - had a token issued 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. + 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 @@ -258,10 +260,10 @@ def refresh_sandbox_mcp_for_user( """ run_id = str(task_run.id) task = task_run.task - current_identity = get_last_sandbox_identity(run_id, "mcp") or task.created_by_id + 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 - rate_limit_key = f"{run_id}:{user.id}" - if not identity_changed and not should_refresh_mcp_token(rate_limit_key): + 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: @@ -273,12 +275,7 @@ def refresh_sandbox_mcp_for_user( ) try: - access_token = create_oauth_access_token_for_user( - user, - task_run.team_id, - scopes=scopes, - application=oauth_application_for_task(task), - ) + 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, user_id=user.id, error=str(e)) return @@ -300,10 +297,12 @@ def refresh_sandbox_mcp_for_user( mcp_configs = mcp_configs + user_mcp_configs if not mcp_configs: - # Mark the window anyway: on hosts where no MCP URL resolves and the - # user has no installs, every message would otherwise repeat the - # mint-and-discard above. - mark_mcp_token_issued(rate_limit_key) + # 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 @@ -311,15 +310,6 @@ def refresh_sandbox_mcp_for_user( result: CommandResult | None = None for attempt in (1, 2): - if attempt > 1: - logger.info( - "refresh_mcp_retrying", - run_id=run_id, - user_id=user.id, - error=result.error if result else None, - status_code=result.status_code if result else None, - ) - time.sleep(REFRESH_RETRY_DELAY_SECONDS) result = send_refresh_session( task_run, mcp_servers, @@ -327,17 +317,27 @@ def refresh_sandbox_mcp_for_user( timeout=REFRESH_TIMEOUT_SECONDS, ) if result.success: - mark_mcp_token_issued(rate_limit_key) - mark_sandbox_identity(run_id, "mcp", user.id) + 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, user_id=user.id, - error=result.error if result else None, - status_code=result.status_code if result else None, + 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 28f6648aad88..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 @@ -17,7 +17,7 @@ from products.tasks.backend.temporal.process_task.utils import ( McpServerConfig, _mcp_token_issued_cache_key, - clear_sandbox_identities, + _sandbox_identity_cache_key, get_last_sandbox_identity, mark_mcp_token_issued, mark_sandbox_identity, @@ -33,17 +33,15 @@ def _clear_mcp_token_cache(): 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"), - _mcp_token_issued_cache_key("run-1:42"), - _mcp_token_issued_cache_key("run-1:43"), + _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) - clear_sandbox_identities("run-1") yield for key in keys: cache.delete(key) - clear_sandbox_identities("run-1") def _make_mcp_config(name: str = "posthog", token: str = "tok") -> McpServerConfig: @@ -77,9 +75,8 @@ def _make_user_mock(user_id: int = 42) -> MagicMock: return user -_OAUTH_FOR_USER_PATH = "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.create_oauth_access_token_for_user" -_OAUTH_APP_PATH = ( - "products.tasks.backend.temporal.process_task.activities.send_followup_to_sandbox.oauth_application_for_task" +_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" @@ -101,13 +98,9 @@ class TestRefreshSandboxMcpForUser: @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) - def test_success_path_single_call( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh - ): + @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_oauth_app.return_value = "array" mock_ph_configs.return_value = [_make_mcp_config(token="fresh-token")] mock_user_configs.return_value = [] mock_send_refresh.return_value = CommandResult(success=True, status_code=200) @@ -116,7 +109,7 @@ def test_success_path_single_call( 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(user, 7, scopes="read_only", application="array") + 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" ) @@ -133,10 +126,9 @@ def test_success_path_single_call( @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_retries_once_on_first_failure( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -155,10 +147,9 @@ def test_retries_once_on_first_failure( @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_two_failures_are_non_fatal( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -173,10 +164,9 @@ def test_two_failures_are_non_fatal( @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_token_mint_failure_is_non_fatal_and_skips_send( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.side_effect = RuntimeError("oauth service down") @@ -189,10 +179,9 @@ def test_token_mint_failure_is_non_fatal_and_skips_send( @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_skips_send_when_no_mcp_configs_resolved( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [] @@ -205,26 +194,25 @@ def test_skips_send_when_no_mcp_configs_resolved( @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_scopes_propagate_to_oauth_and_configs( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh ): mock_oauth.return_value = "fresh-token" - mock_oauth_app.return_value = "array" 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_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], 7, scopes="full", application="array") + 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_FOR_USER_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 @@ -244,9 +232,9 @@ class TestRefreshIntervalGate: a prior same-run refresh under a different identity.""" @patch(_SEND_REFRESH_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh): - mark_mcp_token_issued("run-1:42") + mark_mcp_token_issued("run-1", 42) refresh_sandbox_mcp_for_user(_make_task_run_mock(), _make_user_mock(), scopes="read_only", auth_token=None) @@ -254,19 +242,18 @@ def test_skipped_when_token_recently_issued(self, mock_oauth, mock_send_refresh) mock_send_refresh.assert_not_called() @patch(_SEND_REFRESH_PATH) - @patch(_OAUTH_FOR_USER_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") + 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=[]), - patch(_OAUTH_APP_PATH, return_value="array"), ): refresh_sandbox_mcp_for_user( _make_task_run_mock(), _make_user_mock(user_id=43), scopes="read_only", auth_token=None @@ -277,20 +264,19 @@ def test_identity_transition_bypasses_rate_limit(self, mock_oauth, mock_send_ref assert get_last_sandbox_identity("run-1", "mcp") == 43 @patch(_SEND_REFRESH_PATH) - @patch(_OAUTH_FOR_USER_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") + 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=[]), - patch(_OAUTH_APP_PATH, return_value="array"), ): refresh_sandbox_mcp_for_user( _make_task_run_mock(), _make_user_mock(user_id=42), scopes="read_only", auth_token=None @@ -302,11 +288,8 @@ def test_switching_back_to_creator_refreshes(self, mock_oauth, mock_send_refresh @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) - def test_marks_after_successful_refresh( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh - ): + @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 = [] @@ -315,16 +298,15 @@ def test_marks_after_successful_refresh( 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:42")) is True + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True @patch(_TIME_SLEEP_PATH) @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_marks_after_successful_retry( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -336,16 +318,15 @@ def test_marks_after_successful_retry( 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:42")) is True + assert cache.get(_mcp_token_issued_cache_key("run-1", 42)) is True @patch(_TIME_SLEEP_PATH) @patch(_SEND_REFRESH_PATH) @patch(_USER_CONFIGS_PATH) @patch(_PH_CONFIGS_PATH) - @patch(_OAUTH_APP_PATH) - @patch(_OAUTH_FOR_USER_PATH) + @patch(_OAUTH_MINT_PATH) def test_does_not_mark_after_two_failures( - self, mock_oauth, mock_oauth_app, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep + self, mock_oauth, mock_ph_configs, mock_user_configs, mock_send_refresh, _mock_sleep ): mock_oauth.return_value = "fresh-token" mock_ph_configs.return_value = [_make_mcp_config()] @@ -355,7 +336,7 @@ def test_does_not_mark_after_two_failures( 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:42")) 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 f817ac70413a..5f1e270613ce 100644 --- a/products/tasks/backend/temporal/process_task/utils.py +++ b/products/tasks/backend/temporal/process_task/utils.py @@ -340,23 +340,37 @@ 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(run_id: str) -> None: - """Record that a fresh MCP token was issued to the sandbox for this run. - The cache entry self-expires after MCP_TOKEN_REFRESH_INTERVAL_SECONDS, so - `should_refresh_mcp_token` returns True again past that window. +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(run_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_SECONDS) + get_tasks_cache().set(_mcp_token_issued_cache_key(scope, user_id), True, timeout=MCP_TOKEN_REFRESH_INTERVAL_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 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 @@ -366,42 +380,27 @@ def should_refresh_mcp_token(run_id: str) -> bool: SandboxIdentityKind = Literal["mcp"] -SANDBOX_IDENTITY_KINDS: tuple[SandboxIdentityKind, ...] = ("mcp",) - -def _sandbox_identity_cache_key(run_id: str, kind: SandboxIdentityKind) -> str: - return f"posthog_ai:task-run-{kind}-identity:{run_id}" +def _sandbox_identity_cache_key(scope: str, kind: SandboxIdentityKind) -> str: + return f"posthog_ai:sandbox-{kind}-identity:{scope}" -def mark_sandbox_identity(run_id: str, kind: SandboxIdentityKind, value: int | str) -> None: +def mark_sandbox_identity(scope: str, kind: SandboxIdentityKind, value: int | str) -> None: """Record which identity the sandbox currently holds for a credential kind. ``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. + per-credential freshness rate limits. ``scope`` comes from + ``sandbox_identity_scope``, so a replacement sandbox starts unmarked. """ - get_tasks_cache().set(_sandbox_identity_cache_key(run_id, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) + get_tasks_cache().set(_sandbox_identity_cache_key(scope, kind), value, timeout=SANDBOX_IDENTITY_TTL_SECONDS) -def get_last_sandbox_identity(run_id: str, kind: SandboxIdentityKind) -> int | str | 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(run_id, kind)) - - -def clear_sandbox_identities(run_id: str) -> None: - """Forget a run's swapped identities. - - Called when a sandbox is restored from a snapshot: the resume path - re-applies the boot-time (task creator) credentials, so any remembered - swap would diverge from what the sandbox actually holds — blocking the - actor's next rebind as a same-identity no-op while the refresh loop pulls - the other way. - """ - cache = get_tasks_cache() - for kind in SANDBOX_IDENTITY_KINDS: - cache.delete(_sandbox_identity_cache_key(run_id, kind)) + return get_tasks_cache().get(_sandbox_identity_cache_key(scope, kind)) @dataclass(frozen=True)