diff --git a/products/tasks/backend/facade/api.py b/products/tasks/backend/facade/api.py index a7199e03d551..16312f0a3c64 100644 --- a/products/tasks/backend/facade/api.py +++ b/products/tasks/backend/facade/api.py @@ -94,6 +94,13 @@ WIZARD_PR_READY_EMAIL_FEATURE_FLAG = "wizard-cloud-run-pr-ready-email-enabled" +# Runtime posture for a setup-wizard cloud run, applied in create_wizard_cloud_run. The model is +# pinned because these runs route to the unbilled `onboarding` gateway product, which allowlists a +# narrow model set; the string form avoids pulling the temporal RuntimeAdapter enum onto this path. +WIZARD_CLOUD_RUN_RUNTIME_ADAPTER = "claude" +WIZARD_CLOUD_RUN_MODEL = "claude-sonnet-5" +WIZARD_CLOUD_RUN_AI_STAGE = "wizard_pr_agent" + __all__ = [ "CODE_INVITE_INVALID_CODE", "CODE_INVITE_NOT_REDEEMABLE", @@ -900,6 +907,10 @@ def create_wizard_cloud_run( The PR head branch is generated here (not by the agent) so the GitHub PR webhook can bind the opened PR back to this run by branch + repository — wizard PRs are bot-authored, which the agent-side PR attribution cannot match. + + The model is pinned rather than left to the agent's default because these runs bill to nobody: + they route to the unbilled ``onboarding`` gateway product, whose model allowlist is narrow, and + PostHog absorbs the cost. Keep the pin inside that allowlist or the run fails at the gateway. """ head_branch = generate_wizard_head_branch() prompt = build_wizard_pr_agent_prompt(head_branch) @@ -916,6 +927,9 @@ def create_wizard_cloud_run( wizard_config={}, wizard_head_branch=head_branch, posthog_mcp_scopes="read_only", + runtime_adapter=WIZARD_CLOUD_RUN_RUNTIME_ADAPTER, + model=WIZARD_CLOUD_RUN_MODEL, + ai_stage=WIZARD_CLOUD_RUN_AI_STAGE, # The agent server boots idle; this is the message that actually kicks it off once ready # (delivered by forward_pending_user_message). Without it the run stalls after "Started agent". pending_user_message=prompt, @@ -1769,6 +1783,16 @@ def _sync_automation_schedule(automation: TaskAutomation) -> None: "loop_trigger_id", "trigger_context", "config_snapshot", + # The run's model posture, chosen at creation by the server-owned caller and read back out + # of state when the run dispatches. It decides what the run costs, and for a run routed to + # an unbilled gateway product (create_wizard_cloud_run pins claude-sonnet-5 for the + # `onboarding` product) it is the only thing keeping the run off the more expensive models + # that product still allowlists. Every writer is server-side, so nothing legitimate PATCHes + # these. + "runtime_adapter", + "provider", + "model", + "reasoning_effort", } ) diff --git a/products/tasks/backend/presentation/serializers.py b/products/tasks/backend/presentation/serializers.py index 2d8b56bfbc05..d9c6ff731153 100644 --- a/products/tasks/backend/presentation/serializers.py +++ b/products/tasks/backend/presentation/serializers.py @@ -608,6 +608,11 @@ def validate_origin_product(self, value): # would route the task's run logs into PostHog's internal Logs project # (run_log_mirror) and inherit scout visibility semantics. raise serializers.ValidationError("origin_product 'signals_scout' is reserved for signals scout runs") + if value == tasks_facade.TaskOriginProduct.ONBOARDING: + # This origin routes the run's LLM traffic to the unbilled `onboarding` gateway + # product, so a forged one would be free model access. Only create_wizard_cloud_run + # sets it, behind its own rate limits and daily cap. + raise serializers.ValidationError("origin_product 'onboarding' is reserved for setup wizard cloud runs") return value def validate_repository(self, value): diff --git a/products/tasks/backend/tests/test_api.py b/products/tasks/backend/tests/test_api.py index fb5fa4ab19ec..5ad5cd54dc5b 100644 --- a/products/tasks/backend/tests/test_api.py +++ b/products/tasks/backend/tests/test_api.py @@ -1267,6 +1267,7 @@ def test_create_task_with_hogdesk_origin_product(self): [ ("image_builder",), ("experiments",), + ("onboarding",), ] ) def test_create_task_rejects_internal_origin(self, origin: str): @@ -4385,6 +4386,10 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish "pending_dispatch": {"workflow_id_prefix": "review-real", "create_pr": True}, "pending_external_followups": pending_external_followups, "pending_external_followups_generation": 7, + "runtime_adapter": "claude", + "provider": "anthropic", + "model": "claude-sonnet-5", + "reasoning_effort": "low", }, ) @@ -4394,7 +4399,8 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish # (which would mint a write-scoped wizard token into the sandbox), change rollout # decisions, change Modal resume snapshot metadata, repoint the run at another # team's Temporal workflow, or steer an orphan re-dispatch (workflow ID prefix / MCP - # scopes) via pending_dispatch. Non-protected keys still merge. + # scopes) via pending_dispatch, or repoint the run at a costlier model (which for a run + # routed to an unbilled gateway product is free spend). Non-protected keys still merge. response = self.client.patch( f"/api/projects/@current/tasks/{task.id}/runs/{run.id}/", { @@ -4424,6 +4430,10 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish } ], "pending_external_followups_generation": 999, + "runtime_adapter": "codex", + "provider": "openai", + "model": "claude-opus-4-8", + "reasoning_effort": "high", "scratch": "ok", } }, @@ -4450,6 +4460,10 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish assert run.state["pending_dispatch"] == {"workflow_id_prefix": "review-real", "create_pr": True} assert run.state["pending_external_followups"] == pending_external_followups assert run.state["pending_external_followups_generation"] == 7 + assert run.state["runtime_adapter"] == "claude" + assert run.state["provider"] == "anthropic" + assert run.state["model"] == "claude-sonnet-5" + assert run.state["reasoning_effort"] == "low" assert run.state["scratch"] == "ok" # non-protected keys still merge # Nor can a caller remove a protected key to force a fallback or unguarded path. @@ -4470,6 +4484,10 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish "pending_dispatch", "pending_external_followups", "pending_external_followups_generation", + "runtime_adapter", + "provider", + "model", + "reasoning_effort", "scratch", ], }, @@ -4489,6 +4507,13 @@ def test_patch_cannot_mutate_protected_credential_state_keys(self, _mock_publish assert run.state["pending_dispatch"] == {"workflow_id_prefix": "review-real", "create_pr": True} assert run.state["pending_external_followups"] == pending_external_followups assert run.state["pending_external_followups_generation"] == 7 + # Dropping the model posture is as good as repointing it: the processing context reads these + # back with .get(), so an absent key silently falls back to the runtime's default rather than + # the pin the server chose. + assert run.state["runtime_adapter"] == "claude" # protected key survives removal + assert run.state["provider"] == "anthropic" # protected key survives removal + assert run.state["model"] == "claude-sonnet-5" # protected key survives removal + assert run.state["reasoning_effort"] == "low" # protected key survives removal assert "scratch" not in run.state # non-protected key removed @patch("products.tasks.backend.facade.api.signal_workflow_completion") diff --git a/products/tasks/backend/tests/test_facade.py b/products/tasks/backend/tests/test_facade.py index 53af1386f636..62b1aa491caa 100644 --- a/products/tasks/backend/tests/test_facade.py +++ b/products/tasks/backend/tests/test_facade.py @@ -586,6 +586,23 @@ def test_create_wizard_cloud_run_seeds_pending_user_message(self, _mock_workflow # and the run never opens a PR. Wizard runs must pin the overlap boot off. self.assertIs(run.state.get("overlap_clone_boot_enabled"), False) + @patch("products.tasks.backend.temporal.client.execute_task_processing_workflow") + def test_create_wizard_cloud_run_pins_its_model(self, _mock_workflow): + Integration.objects.create(team=self.team, kind="github", config={}) + created = facade.create_wizard_cloud_run( + team=self.team, + user_id=self.user.id, + repository="acme-co/web", + ) + run = TaskRun.objects.get(task_id=created.task_id) + # Wizard runs route to the unbilled `onboarding` gateway product, which allowlists only + # these models. Dropping the pin puts the run back on the agent-server's premium default, + # which that product rejects, so every wizard cloud run would 403 at the gateway. Changing + # the pin means changing the allowlist in services/llm-gateway too. + self.assertEqual(run.state.get("runtime_adapter"), "claude") + self.assertEqual(run.state.get("model"), "claude-sonnet-5") + self.assertEqual(run.state.get("ai_stage"), "wizard_pr_agent") + class TestRecentWizardCloudRunTimes(TestCase): organization: ClassVar[Organization] diff --git a/services/llm-gateway/README.md b/services/llm-gateway/README.md index 013e8d51e06c..eb864517810c 100644 --- a/services/llm-gateway/README.md +++ b/services/llm-gateway/README.md @@ -234,6 +234,7 @@ OAuth access is permitted only for products with an explicit `allowed_applicatio | `ci` | API key only | All | CI / e2e test runs | | `posthog_code` | OAuth only | Restricted set | Desktop coding agent | | `background_agents` | OAuth only | Restricted set | Cloud background agents | +| `onboarding` | OAuth only | claude-sonnet-5 | Unbilled setup wizard cloud run | | `wizard` | API key + OAuth | All | Max AI assistant | | `django` | API key only | All | Server-side Django calls | | `growth` | API key only | All | Growth team | @@ -334,5 +335,6 @@ response = client.chat.completions.create( ``` `ai_product` and `$ai_billable` are derived from the product config (`products/config.py`): -the route sets `ai_product` from the `product` arg, and `$ai_billable` from that product's -`billable` flag. Set `billable=True` on the product config to bill its generations. +the route sets `ai_product` from the `product` arg, and `$ai_billable` from whether that +product has a `credit_bucket`. Set `credit_bucket` on the product config to bill its +generations into that bucket; leave it `None` to keep them unbilled. diff --git a/services/llm-gateway/src/llm_gateway/config.py b/services/llm-gateway/src/llm_gateway/config.py index 3f5088d16f16..962f02b73c6c 100644 --- a/services/llm-gateway/src/llm_gateway/config.py +++ b/services/llm-gateway/src/llm_gateway/config.py @@ -30,6 +30,7 @@ class UserCostLimit(BaseModel, frozen=True): "wizard": ProductCostLimit(limit_usd=10000.0, window_seconds=86400), "posthog_code": ProductCostLimit(limit_usd=5000.0, window_seconds=3600), "background_agents": ProductCostLimit(limit_usd=1000.0, window_seconds=3600), + "onboarding": ProductCostLimit(limit_usd=1000.0, window_seconds=3600), "django": ProductCostLimit(limit_usd=5000.0, window_seconds=86400), "custom_image_scans": ProductCostLimit(limit_usd=1000.0, window_seconds=86400), "signals": ProductCostLimit(limit_usd=25000.0, window_seconds=86400), @@ -61,6 +62,19 @@ class UserCostLimit(BaseModel, frozen=True): sustained_limit_usd=10000.0, sustained_window_seconds=2592000, ), + # Nobody is billed for onboarding (credit_bucket=None), so this bounds blast radius rather than + # spend: the route's server-credential marker proves a token was minted server-side, not that it + # belongs to a wizard run, and INTERNAL_SCOPES in posthog/temporal/oauth.py grants that marker to + # every task run. Sized to stay clear of real onboarding rather than to be tight, since cutting a + # user off mid-setup is worse than the unbilled spend: half of DEFAULT_USER_COST_LIMIT, and well + # under the comparable agentic product (background_agents, $500/week burst). Staff bypass this + # entirely via is_usage_unlimited, so internal runs are never capped by it. + "onboarding": UserCostLimit( + burst_limit_usd=50.0, + burst_window_seconds=86400, + sustained_limit_usd=500.0, + sustained_window_seconds=2592000, + ), } FREE_PLAN_COST_LIMIT = UserCostLimit( diff --git a/services/llm-gateway/src/llm_gateway/products/config.py b/services/llm-gateway/src/llm_gateway/products/config.py index 84c777fab530..28b228406d32 100644 --- a/services/llm-gateway/src/llm_gateway/products/config.py +++ b/services/llm-gateway/src/llm_gateway/products/config.py @@ -88,6 +88,18 @@ class ProductConfig: } ) +# Products whose requires_server_credential applies right away rather than waiting for +# posthog_code_model_gate_enabled. The flag exists so products that already shipped accepting plain +# Code OAuth tokens keep working until the Code billing cutover. A product that never had such a +# permissive period has nothing to stay compatible with, and leaving it flag-gated would ship an +# unbilled route open to any Code OAuth token for as long as the flag is off. +UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS: Final[frozenset[str]] = frozenset( + { + "custom_image_scans", + "onboarding", + } +) + PRODUCTS: Final[dict[str, ProductConfig]] = { "llm_gateway": ProductConfig( allowed_application_ids=None, @@ -139,6 +151,19 @@ class ProductConfig: credit_bucket=None, requires_server_credential=True, ), + # The setup wizard's cloud run (Task.OriginProduct.ONBOARDING). Unbilled like + # background_agents, and this one runs before the user has decided to buy anything. + # Two gates keep the free route shut: Django refuses `onboarding` as a caller-supplied + # task origin, and the agent-server only routes here for a run carrying the protected + # `wizard_config` state key. Models stay narrow because a free bucket shouldn't reach + # the whole fleet; claude-opus-4-8 is only the SDK's fallback for the pinned sonnet. + "onboarding": ProductConfig( + allowed_application_ids=frozenset({POSTHOG_CODE_US_APP_ID, POSTHOG_CODE_EU_APP_ID, POSTHOG_CODE_DEV_APP_ID}), + allowed_models=frozenset({"claude-sonnet-5", "claude-opus-4-8"}) | BEDROCK_MODELS, + allow_api_keys=False, + credit_bucket=None, + requires_server_credential=True, + ), "slack_app": ProductConfig( allowed_application_ids=frozenset({POSTHOG_CODE_US_APP_ID, POSTHOG_CODE_EU_APP_ID, POSTHOG_CODE_DEV_APP_ID}), allowed_models=_POSTHOG_CODE_AGENT_MODELS | BEDROCK_MODELS, @@ -418,13 +443,14 @@ def check_product_access( # and route around the posthog_code free-tier model gate. Require the internal marker that # only server-minted tokens carry. OAuth-only: personal API keys reach the gateway with an # explicit, feature-gated llm_gateway:read scope (a `*` PAK is rejected at auth), so the - # shared server-side gateway key still works here. Gated behind the same flag as the - # free-tier gate so it stays inert until the Code billing cutover. + # shared server-side gateway key still works here. Products that shipped before this check + # existed stay behind the free-tier flag so they keep working until the Code billing cutover; + # the rest enforce it now, per UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS. if ( config.requires_server_credential and is_oauth and INTERNAL_RUN_SCOPE not in (scopes or []) - and (settings.posthog_code_model_gate_enabled or resolved_product == "custom_image_scans") + and (settings.posthog_code_model_gate_enabled or resolved_product in UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS) ): return False, f"Product '{product}' requires a server-minted credential" diff --git a/services/llm-gateway/tests/test_product_config.py b/services/llm-gateway/tests/test_product_config.py index cb3bf30ef5ee..b25ac055d167 100644 --- a/services/llm-gateway/tests/test_product_config.py +++ b/services/llm-gateway/tests/test_product_config.py @@ -16,6 +16,7 @@ PRODUCTS, TWIG_EU_APP_ID, TWIG_US_APP_ID, + UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS, WIZARD_EU_APP_ID, WIZARD_US_APP_ID, check_free_tier_model_access, @@ -552,9 +553,9 @@ def test_denied_model_alias_variant_is_also_denied(self): class TestServerCredentialRequirement: """The internal products that share the PostHog Desktop OAuth app (background_agents, signals, - slack_app, conversations) must accept only server-minted tokens — those carrying the internal - `internal_run:read` marker. Otherwise a user's own Desktop OAuth token could route around the - posthog_code free-tier gate through these products to premium models.""" + slack_app, conversations, onboarding) must accept only server-minted tokens — those carrying the + internal `internal_run:read` marker. Otherwise a user's own Desktop OAuth token could route around + the posthog_code free-tier gate through these products to premium models.""" _MARKER_SCOPES = ["llm_gateway:read", "task:write", "internal_run:read"] @@ -567,7 +568,7 @@ def gate_enabled(self, monkeypatch: pytest.MonkeyPatch): yield get_settings.cache_clear() - @pytest.mark.parametrize("product", ["background_agents", "signals", "slack_app", "conversations"]) + @pytest.mark.parametrize("product", ["background_agents", "signals", "slack_app", "conversations", "onboarding"]) def test_oauth_without_marker_is_rejected(self, product: str): # a desktop Code token (wildcard scope, no internal marker); claude-sonnet-5 is in every # sibling's model list, so the rejection is unambiguously the missing server credential @@ -577,7 +578,7 @@ def test_oauth_without_marker_is_rejected(self, product: str): assert allowed is False assert error is not None and "server-minted" in error - @pytest.mark.parametrize("product", ["background_agents", "signals", "slack_app", "conversations"]) + @pytest.mark.parametrize("product", ["background_agents", "signals", "slack_app", "conversations", "onboarding"]) def test_oauth_with_marker_is_allowed(self, product: str): allowed, error = check_product_access( product, "oauth_access_token", POSTHOG_CODE_US_APP_ID, "claude-sonnet-5", scopes=self._MARKER_SCOPES @@ -613,6 +614,33 @@ def test_gate_disabled_leaves_sibling_access_unchanged(self, monkeypatch: pytest assert allowed is True assert error is None + # The rest of this class runs with the gate forced on, which is the state in which the + # requirement was already known to hold. These cover the default state, where the products that + # never shipped without the check have to enforce it anyway. Spelled out rather than derived + # from UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS: parameterizing over the set under test would + # make dropping a product from it delete its own coverage instead of failing. + @pytest.mark.parametrize("product", ["custom_image_scans", "onboarding"]) + def test_gate_disabled_still_refuses_unconditional_products(self, product: str, monkeypatch: pytest.MonkeyPatch): + from llm_gateway.config import get_settings + + monkeypatch.delenv("LLM_GATEWAY_POSTHOG_CODE_MODEL_GATE_ENABLED", raising=False) + get_settings.cache_clear() + allowed, error = check_product_access(product, "oauth_access_token", POSTHOG_CODE_US_APP_ID, None, scopes=["*"]) + assert allowed is False + assert error is not None and "server-minted" in error + + @pytest.mark.parametrize("product", ["custom_image_scans", "onboarding"]) + def test_gate_disabled_still_admits_server_minted_tokens(self, product: str, monkeypatch: pytest.MonkeyPatch): + from llm_gateway.config import get_settings + + monkeypatch.delenv("LLM_GATEWAY_POSTHOG_CODE_MODEL_GATE_ENABLED", raising=False) + get_settings.cache_clear() + allowed, error = check_product_access( + product, "oauth_access_token", POSTHOG_CODE_US_APP_ID, None, scopes=self._MARKER_SCOPES + ) + assert allowed is True + assert error is None + _CODE_APP_IDS = frozenset({POSTHOG_CODE_DEV_APP_ID, POSTHOG_CODE_EU_APP_ID, POSTHOG_CODE_US_APP_ID}) _CODE_APP_PRODUCTS = [ @@ -635,6 +663,12 @@ def test_internal_code_app_products_require_a_server_credential(self, product: s "posthog_code free-tier model gate" ) + def test_unconditional_products_are_the_ones_enforcing_without_the_flag(self): + # Pairs with the two flag-off tests above, which name their products literally. If a + # product is added here without flag-off coverage, or removed from here while still + # expected to enforce, exactly one of the two sides fails. + assert UNCONDITIONAL_SERVER_CREDENTIAL_PRODUCTS == frozenset({"custom_image_scans", "onboarding"}) + def test_posthog_code_is_the_only_code_app_product_open_to_user_tokens(self): # desktop users hold marker-less Code tokens; requiring the marker on the # user-facing product would lock them all out. Membership is asserted so a