diff --git a/backend/apps/workflows/cloud/credential_readiness.py b/backend/apps/workflows/cloud/credential_readiness.py new file mode 100644 index 00000000..13953efb --- /dev/null +++ b/backend/apps/workflows/cloud/credential_readiness.py @@ -0,0 +1,43 @@ +"""Whether this account has an AI connection the cloud could actually run with. + +A cloud run signs its LLM calls with the user's OWN subscription, handed up by +`credential_lease`. Only a rotating OAuth connection can be handed up: an API key has no +refresh token, and the runner cannot mint one. So an account whose only provider is a +Gemini or OpenAI key can never run in the cloud, and the honest moment to say so is +before the user schedules anything, not at 9am when the run refuses. +""" +from __future__ import annotations + +from typing import List, Literal, Optional + +from pydantic import BaseModel, ConfigDict +from typeguard import typechecked + +from backend.apps.nine_router import credential_store + +CONNECT_HINT = ( + "Cloud runs sign in with your own Claude or ChatGPT subscription, so connect one in " + "Settings to run a workflow in the cloud. An API key alone can't be used up there." +) + + +class CredentialReadiness(BaseModel): + model_config = ConfigDict(validate_assignment=True) + + # ready: something is leasable or already lent. none_eligible: only API keys, or nothing at all. + state: Literal["ready", "none_eligible"] + connection_ids: List[str] = [] + # Written for the user, present only when they cannot proceed. + reason: Optional[str] = None + + @property + def ok(self) -> bool: + return self.state == "ready" + + +@typechecked +def cloud_credential_readiness() -> CredentialReadiness: + ids = credential_store.list_oauth_connection_ids() + if not ids: + return CredentialReadiness(state="none_eligible", reason=CONNECT_HINT) + return CredentialReadiness(state="ready", connection_ids=ids) diff --git a/backend/apps/workflows/cloud/handover.py b/backend/apps/workflows/cloud/handover.py index 0d92caae..7027f0c7 100644 --- a/backend/apps/workflows/cloud/handover.py +++ b/backend/apps/workflows/cloud/handover.py @@ -15,8 +15,10 @@ from typing import Optional from pydantic import BaseModel, ConfigDict from typeguard import typechecked +from backend.apps.nine_router import credential_lease, credential_store from backend.apps.workflows import scheduler, storage from backend.apps.workflows.cloud import client as cloud +from backend.apps.workflows.cloud.credential_readiness import cloud_credential_readiness from backend.apps.workflows.cloud.definition import cloud_definition, definition_signature from backend.apps.workflows.cloud.portable_context import portable_context from backend.apps.workflows.cloud.schedule import ScheduleSupported, to_cloud_schedule, wire @@ -44,6 +46,42 @@ class TargetOutcome(BaseModel): message: Optional[str] = None +LEASE_FAILED = ( + "Couldn't lend your AI account to the cloud, so nothing was scheduled there. " + "This workflow still runs on this device. Try again in a moment." +) +LEASE_STRANDED = ( + "We couldn't confirm whether your AI account reached the cloud, so nothing was scheduled " + "there. Open Settings and reconnect the provider before trying again." +) + + +@typechecked +async def lend_credential_for_cloud() -> TargetOutcome: + """Make sure the cloud holds a credential it can sign this user's runs with. + + Already-lent is the common case (one lease covers every cloud workflow), so this is a + no-op after the first one. + """ + readiness = cloud_credential_readiness() + if not readiness.ok: + return TargetOutcome(ok=False, message=readiness.reason) + + for connection_id in readiness.connection_ids: + outcome = await credential_lease.lease_to_cloud(connection_id) + if outcome.status in ("leased", "not_rotatable"): + # not_rotatable here means the local copy has already been stripped, i.e. the cloud has it. + return TargetOutcome(ok=True) + if outcome.status == "not_signed_in": + return TargetOutcome(ok=False, message=SIGN_IN_MESSAGE) + if outcome.status == "ownership_unknown": + logger.error("credential lease outcome unknown: %s", outcome.detail) + return TargetOutcome(ok=False, message=LEASE_STRANDED) + logger.info("credential lease for %s failed: %s %s", connection_id, outcome.status, outcome.detail) + + return TargetOutcome(ok=False, message=LEASE_FAILED) + + @typechecked async def hand_to_cloud(wf: Workflow, enabled: bool) -> TargetOutcome: mapping = to_cloud_schedule(wf.schedule) @@ -51,6 +89,11 @@ async def hand_to_cloud(wf: Workflow, enabled: bool) -> TargetOutcome: return TargetOutcome(ok=False, message=mapping.reason) if enabled and not scheduler.is_schedule_configured(wf.schedule): return TargetOutcome(ok=False, message="Finish setting up the schedule before choosing where it runs.") + # Lend the credential BEFORE the workflow goes up. The other order parks a workflow in the cloud + # that cannot sign a single call, and the user only finds out when 9am comes and goes. + lent = await lend_credential_for_cloud() + if not lent.ok: + return TargetOutcome(ok=False, message=lent.message) definition = cloud_definition(wf) context = portable_context().as_body() try: @@ -104,10 +147,30 @@ async def take_back(wf: Workflow, enabled: bool) -> TargetOutcome: wf.next_run_at = scheduler.compute_next_fire(wf) if wf.schedule.enabled else None wf.updated_at = datetime.now() storage.save_workflow(wf) + await p_reclaim_credential_if_last(wf.id) scheduler.kick() return TargetOutcome(ok=True) +@typechecked +async def p_reclaim_credential_if_last(leaving_id: str) -> None: + """Bring custody home once nothing is left in the cloud that needs it. + + Reclaiming while another cloud workflow is still scheduled would break that one, so the + last one out turns off the lights. Best-effort: a failure here leaves the credential + lent, which still works, rather than failing a toggle the user already got. + """ + if any( + w.id != leaving_id and w.execution_target == "cloud" + for w in storage.list_workflows() + ): + return + for connection_id in credential_store.list_oauth_connection_ids(): + outcome = await credential_lease.release_to_device(connection_id) + if outcome.status not in ("released", "no_such_connection"): + logger.info("could not reclaim %s: %s %s", connection_id, outcome.status, outcome.detail) + + @typechecked async def release_before_removing(wf: Workflow) -> TargetOutcome: """Take the cloud copy down before a workflow disappears from this machine. diff --git a/backend/apps/workflows/cloud/status.py b/backend/apps/workflows/cloud/status.py index 48f6ce67..3ff4dcfc 100644 --- a/backend/apps/workflows/cloud/status.py +++ b/backend/apps/workflows/cloud/status.py @@ -15,6 +15,7 @@ from typeguard import typechecked from backend.apps.workflows import storage from backend.apps.workflows.cloud import client as cloud +from backend.apps.workflows.cloud.credential_readiness import CredentialReadiness, cloud_credential_readiness from backend.apps.workflows.cloud.definition import cloud_definition, definition_signature from backend.apps.workflows.cloud.portable_context import portable_context from backend.apps.workflows.cloud.schedule import ScheduleSupported, to_cloud_schedule, wire @@ -56,6 +57,9 @@ class CloudStatusReady(CloudStatusBase): # None when this control plane cannot tell us whether the runner could do the job. capability: Optional[cloud.CloudCapability] = None hosted: Optional[HostedState] = None + # Whether this account owns an AI connection the cloud could sign runs with. Read locally, + # because it is our 9router db that knows, not the control plane. + credential: CredentialReadiness CloudStatus = Union[CloudStatusReady, CloudStatusSignedOut, CloudStatusUnknown] @@ -138,5 +142,6 @@ async def compute_status(wf: Workflow) -> CloudStatus: usage=pre.usage, capability=pre.capability, hosted=hosted, + credential=cloud_credential_readiness(), **shared, ) diff --git a/backend/config/entity_references.py b/backend/config/entity_references.py index cc7c2400..b1fe0833 100644 --- a/backend/config/entity_references.py +++ b/backend/config/entity_references.py @@ -30,6 +30,8 @@ class EntityKind(str, Enum): OUTPUT = "output" WORKSPACE = "workspace" CLOUD_WORKFLOW = "cloud_workflow" + # A provider login in 9router's own db, not one of our JSON records. + PROVIDER_CONNECTION = "provider_connection" class EntityStore(BaseModel): @@ -64,6 +66,7 @@ ENTITY_STORES: List[EntityStore] = [ EntityStore(kind=EntityKind.WORKSPACE, module="backend.apps.outputs.outputs", lookup="read_workspace"), # The one referent that does not live on this machine. preflight asks the cloud whether it still has the row; a miss renders as "nothing is running this", never as a silent blank. EntityStore(kind=EntityKind.CLOUD_WORKFLOW, module="backend.apps.workflows.cloud.client", lookup="preflight"), + EntityStore(kind=EntityKind.PROVIDER_CONNECTION, module="backend.apps.nine_router.credential_store", lookup="read_credential"), ] CROSS_ENTITY_REFERENCES: List[EntityReference] = [ @@ -100,6 +103,7 @@ CROSS_ENTITY_REFERENCES: List[EntityReference] = [ EntityReference(module="backend.apps.workflows.models", model="AskRunBody", field="run_id", target=EntityKind.WORKFLOW_RUN), EntityReference(module="backend.apps.workflows.models", model="MissedRun", field="workflow_id", target=EntityKind.WORKFLOW), EntityReference(module="backend.apps.workflows.models", model="Workflow", field="cloud_workflow_id", target=EntityKind.CLOUD_WORKFLOW), + EntityReference(module="backend.apps.workflows.cloud.credential_readiness", model="CredentialReadiness", field="connection_ids", target=EntityKind.PROVIDER_CONNECTION), EntityReference(module="backend.apps.workflows.models", model="Workflow", field="dashboard_id", target=EntityKind.DASHBOARD), EntityReference(module="backend.apps.workflows.models", model="Workflow", field="edit_agent_session_id", target=EntityKind.SESSION), EntityReference(module="backend.apps.workflows.models", model="Workflow", field="last_run_id", target=EntityKind.WORKFLOW_RUN), diff --git a/backend/tests/test_cloud_credential_wiring.py b/backend/tests/test_cloud_credential_wiring.py new file mode 100644 index 00000000..f608c319 --- /dev/null +++ b/backend/tests/test_cloud_credential_wiring.py @@ -0,0 +1,180 @@ +"""A cloud workflow may never exist without a credential the cloud can sign it with. + +Before this wiring, `lease_to_cloud` had zero callers outside its own unit tests. Every part +worked and nothing joined them, so every cloud run in existence died at dispatch with +`no_cloud_credential` and the user found out at 9am. These tests pin the join. +""" +import pytest + +from backend.apps.nine_router.credential_lease import LeaseOutcome +from backend.apps.workflows.cloud import credential_readiness, handover + + +@pytest.fixture +def p_oauth(monkeypatch): + def set_ids(ids): + monkeypatch.setattr( + credential_readiness.credential_store, "list_oauth_connection_ids", lambda: list(ids) + ) + monkeypatch.setattr( + handover.credential_store, "list_oauth_connection_ids", lambda: list(ids) + ) + return set_ids + + +@pytest.fixture +def p_lease(monkeypatch): + calls = [] + + def set_result(*statuses): + seq = list(statuses) + + async def fake(connection_id: str) -> LeaseOutcome: + calls.append(connection_id) + return LeaseOutcome(status=seq.pop(0) if seq else "cloud_rejected") + + monkeypatch.setattr(handover.credential_lease, "lease_to_cloud", fake) + return calls + + return set_result + + +def test_an_account_with_only_api_keys_cannot_use_cloud_runs(p_oauth): + # Gemini AI Studio and a raw OpenAI key are apikey rows: no refresh token, nothing to lend. + p_oauth([]) + r = credential_readiness.cloud_credential_readiness() + assert r.ok is False + assert r.state == "none_eligible" + assert "Claude or ChatGPT" in (r.reason or ""), "must name what to connect, not just refuse" + assert "API key" in (r.reason or ""), "the API-key user needs to know why theirs will not do" + + +def test_an_oauth_connection_reads_as_ready(p_oauth): + p_oauth(["conn-claude"]) + r = credential_readiness.cloud_credential_readiness() + assert r.ok is True + assert r.connection_ids == ["conn-claude"] + assert r.reason is None + + +@pytest.mark.asyncio +async def test_lending_succeeds_on_the_first_usable_connection(p_oauth, p_lease): + p_oauth(["conn-a", "conn-b"]) + calls = p_lease("leased") + out = await handover.lend_credential_for_cloud() + assert out.ok is True + assert calls == ["conn-a"], "one lease covers every cloud workflow; do not lend them all" + + +@pytest.mark.asyncio +async def test_an_already_stripped_connection_counts_as_lent(p_oauth, p_lease): + # not_rotatable means the local refresh token is already gone, i.e. the cloud has it. + p_oauth(["conn-a"]) + p_lease("not_rotatable") + assert (await handover.lend_credential_for_cloud()).ok is True + + +@pytest.mark.asyncio +async def test_a_refused_connection_falls_through_to_the_next(p_oauth, p_lease): + p_oauth(["conn-dead", "conn-good"]) + calls = p_lease("cloud_rejected", "leased") + out = await handover.lend_credential_for_cloud() + assert out.ok is True + assert calls == ["conn-dead", "conn-good"] + + +@pytest.mark.asyncio +async def test_every_connection_failing_refuses_with_words_a_user_can_act_on(p_oauth, p_lease): + p_oauth(["conn-a"]) + p_lease("cloud_rejected") + out = await handover.lend_credential_for_cloud() + assert out.ok is False + assert "still runs on this device" in (out.message or ""), "say what DID happen, not just what failed" + + +@pytest.mark.asyncio +async def test_signed_out_says_sign_in_rather_than_a_lease_error(p_oauth, p_lease): + p_oauth(["conn-a"]) + p_lease("not_signed_in") + out = await handover.lend_credential_for_cloud() + assert out.ok is False + assert out.message == handover.SIGN_IN_MESSAGE + + +@pytest.mark.asyncio +async def test_an_unknown_lease_outcome_tells_the_user_to_reconnect(p_oauth, p_lease): + # The token is off this device and we cannot prove the cloud took it. Silence here strands them. + p_oauth(["conn-a"]) + p_lease("ownership_unknown") + out = await handover.lend_credential_for_cloud() + assert out.ok is False + assert "reconnect" in (out.message or "").lower() + + +@pytest.mark.asyncio +async def test_no_eligible_connection_refuses_before_anything_is_lent(p_oauth, p_lease): + p_oauth([]) + calls = p_lease("leased") + out = await handover.lend_credential_for_cloud() + assert out.ok is False + assert calls == [], "nothing to lend, so nothing should have been attempted" + assert "Claude or ChatGPT" in (out.message or "") + + +# The join itself. Everything above passes even if hand_to_cloud never calls any of it, which is +# exactly the shape of the bug being fixed: the parts all worked and nothing wired them together. + +from backend.apps.workflows import storage +from backend.apps.workflows.cloud import client as cloud +from backend.apps.workflows.models import ScheduleConfig, Workflow, WorkflowStep + +pytestmark = pytest.mark.usefixtures("isolated_workflows_data") + + +def p_wf() -> Workflow: + wf = Workflow( + title="Morning digest", + steps=[WorkflowStep(text="summarize the news")], + schedule=ScheduleConfig( + enabled=True, repeat_unit="day", repeat_every=1, hour=9, minute=0, timezone="UTC" + ), + ) + storage.save_workflow(wf) + return wf + + +@pytest.mark.asyncio +async def test_a_workflow_never_reaches_the_cloud_without_a_credential(monkeypatch, p_oauth): + """The 9am bug, pinned: no lendable account means the push must not happen at all.""" + p_oauth([]) + wf = p_wf() + talked: list = [] + + async def p_call(method, path, body=None): + talked.append(path) + raise AssertionError("pushed a workflow the cloud could never run") + + monkeypatch.setattr(cloud, "p_call", p_call) + + out = await handover.hand_to_cloud(wf, enabled=True) + assert out.ok is False + assert talked == [], "the credential check has to come BEFORE the push" + assert wf.execution_target == "device", "a refused handover leaves it running here" + assert "Claude or ChatGPT" in (out.message or "") + + +@pytest.mark.asyncio +async def test_a_failed_lease_leaves_the_workflow_on_this_device(monkeypatch, p_oauth, p_lease): + p_oauth(["conn-a"]) + p_lease("cloud_rejected") + wf = p_wf() + + async def p_call(method, path, body=None): + raise AssertionError("pushed despite the lease failing") + + monkeypatch.setattr(cloud, "p_call", p_call) + + out = await handover.hand_to_cloud(wf, enabled=True) + assert out.ok is False + assert wf.execution_target == "device" + assert storage.get_workflow(wf.id).execution_target == "device", "and it stayed that way on disk" diff --git a/backend/tests/test_cloud_workflow_target.py b/backend/tests/test_cloud_workflow_target.py index 9d56949f..d38849ff 100644 --- a/backend/tests/test_cloud_workflow_target.py +++ b/backend/tests/test_cloud_workflow_target.py @@ -17,6 +17,24 @@ from backend.apps.workflows.models import ScheduleConfig, Workflow, WorkflowStep pytestmark = pytest.mark.usefixtures("isolated_workflows_data") +@pytest.fixture(autouse=True) +def p_credential_already_lent(monkeypatch): + """These tests are about the timer, not credential custody. Without this they would reach the + real lease, which reads settings this harness never signs in, and every handover would refuse + with a sign-in message instead of the answer under test. Custody has its own file: + test_cloud_credential_wiring.py.""" + from backend.apps.workflows.cloud import handover + + async def lent(): + return handover.TargetOutcome(ok=True) + + async def reclaimed(leaving_id: str) -> None: + return None + + monkeypatch.setattr(handover, "lend_credential_for_cloud", lent) + monkeypatch.setattr(handover, "p_reclaim_credential_if_last", reclaimed) + + def p_sched(**overrides) -> ScheduleConfig: base = dict(enabled=True, repeat_unit="day", repeat_every=1, hour=9, minute=0, timezone="UTC") base.update(overrides) diff --git a/frontend/src/app/pages/Workflows/app/CloudRunSection.tsx b/frontend/src/app/pages/Workflows/app/CloudRunSection.tsx index 13736877..5175cfed 100644 --- a/frontend/src/app/pages/Workflows/app/CloudRunSection.tsx +++ b/frontend/src/app/pages/Workflows/app/CloudRunSection.tsx @@ -108,6 +108,9 @@ const CloudRunSection: React.FC<{ workflow: Workflow; cloud: CloudStatusHandle } {availability.kind === 'blocked' && availability.action === 'plans' && ( )} + {availability.kind === 'blocked' && availability.action === 'connect' && ( + + )} @@ -139,6 +142,9 @@ const CloudRunSection: React.FC<{ workflow: Workflow; cloud: CloudStatusHandle } {availability.action === 'plans' && ( )} + {availability.action === 'connect' && ( + + )} )} diff --git a/frontend/src/app/pages/Workflows/app/cloudApi.ts b/frontend/src/app/pages/Workflows/app/cloudApi.ts index 8589f904..ec2d3e9f 100644 --- a/frontend/src/app/pages/Workflows/app/cloudApi.ts +++ b/frontend/src/app/pages/Workflows/app/cloudApi.ts @@ -35,6 +35,13 @@ interface CloudStatusShared { schedule_reason: string | null; } +/** Whether an AI account exists that the cloud could sign runs with. An API key alone cannot. */ +export interface CloudCredential { + state: 'ready' | 'none_eligible'; + connection_ids: string[]; + reason: string | null; +} + export interface CloudStatusReady extends CloudStatusShared { state: 'ready'; plan: string | null; @@ -43,6 +50,7 @@ export interface CloudStatusReady extends CloudStatusShared { /** Null when the control plane could not tell us; create re-checks either way. */ capability: CloudCapability | null; hosted: HostedState | null; + credential: CloudCredential; } export interface CloudStatusSignedOut extends CloudStatusShared { diff --git a/frontend/src/app/pages/Workflows/app/cloudAvailability.ts b/frontend/src/app/pages/Workflows/app/cloudAvailability.ts index ec08107f..b65995ef 100644 --- a/frontend/src/app/pages/Workflows/app/cloudAvailability.ts +++ b/frontend/src/app/pages/Workflows/app/cloudAvailability.ts @@ -11,7 +11,7 @@ export type CloudProbe = export type CloudAvailability = | { kind: 'checking' } | { kind: 'unknown'; detail: string | null } - | { kind: 'blocked'; reason: string; action: 'sign_in' | 'plans' | null } + | { kind: 'blocked'; reason: string; action: 'sign_in' | 'plans' | 'connect' | null } | { kind: 'available' }; const PLAN_REQUIRED = 'Cloud runs come with Pro and up. On this plan, workflows run on this device.'; @@ -53,6 +53,12 @@ export function cloudAvailability(probe: CloudProbe): CloudAvailability { if (status.capability && !status.capability.ok && status.capability.reason) { return { kind: 'blocked', reason: status.capability.reason, action: null }; } + // Before the plan, deliberately. Someone whose only provider is an API key cannot run in the + // cloud at any price, so leading them to the pricing page would sell them a thing that still + // would not work. + if (status.credential && status.credential.state !== 'ready' && status.credential.reason) { + return { kind: 'blocked', reason: status.credential.reason, action: 'connect' }; + } return blockedForAccount(status) ?? { kind: 'available' }; }