feat(auth): carry the acting org's restoreRules entitlement from whoami

Author: theCodeDrift

packages/cli/src/auth/identity.ts

@@ -1,5 +1,5 @@
 import { getToken } from "./token";
-import { resolveOrgSubject } from "./org";
+import { resolveActingOrg } from "./org";
 import { CLIError } from "../util/cli-error";
 import { resolveRepositoryUrl } from "../util/git-remote";
 import { getCliPrefix } from "../util/package-manager";
@@ -9,10 +9,15 @@ export interface Identity {
   token: string;
   /**
    * Org subject to send on write calls: the current org's Taskless UUID
-   * (preferred) or the deprecated numeric `orgId` claim. See `resolveOrgSubject`.
+   * (preferred) or the deprecated numeric `orgId` claim. See `resolveActingOrg`.
    */
   orgSubject: string | number;
   repositoryUrl: string;
+  /**
+   * Whether the acting org's plan includes rule recovery; `undefined` when
+   * unknown. Chooses which recovery to suggest, never whether to try one.
+   */
+  restoreRules?: boolean;
 }
 
 /**
@@ -21,6 +26,8 @@ export interface Identity {
  *   repo's remotes, falling back to the token's canonical id claim, and finally
  *   to the nil-UUID `NIL_ORG_ID` so a subject is always present
  * - repositoryUrl: inferred from `git remote get-url origin`
+ * - restoreRules: the matched org's plan entitlement, from the same `whoami`
+ *   call; absent when unknown
  *
  * Throws a `CLIError` carrying a stable `CLIErrorCode` if auth is missing
  * (`AUTH_REQUIRED`) or the repository URL cannot be resolved, in which case
@@ -39,9 +46,16 @@ export async function resolveIdentity(cwd: string): Promise<Identity> {
   }
 
   const repositoryUrl = await resolveRepositoryUrl(cwd);
-  const orgSubject = await resolveOrgSubject(cwd, token);
+  const org = await resolveActingOrg(cwd, token);
 
-  return { token, orgSubject, repositoryUrl };
+  return {
+    token,
+    orgSubject: org.subject,
+    repositoryUrl,
+    ...(org.restoreRules === undefined
+      ? {}
+      : { restoreRules: org.restoreRules }),
+  };
 }
 
 /**
@@ -52,7 +66,7 @@ export async function resolveIdentity(cwd: string): Promise<Identity> {
  * prose. It is shared by every caller so the mapping cannot drift between
  * them.
  *
- * `resolveOrgSubject` is deliberately absent from the list of expected
+ * `resolveActingOrg` is deliberately absent from the list of expected
  * failures: it cannot throw. `fetchWhoami` swallows every network and HTTP
  * error and returns `undefined`, and `decodeOrgId` falls back to the nil-UUID
  * `NIL_ORG_ID`, so a repo with no matching org resolves a subject rather than

packages/cli/src/auth/org.ts

@@ -65,26 +65,50 @@ export async function resolveCurrentOrg(
   return selectOrgForOwners(ownerUrls, orgs);
 }
 
+/** The organization the CLI acts as, and what its plan is known to grant. */
+export interface ActingOrg {
+  /** The org subject to send on write calls. See `resolveActingOrg`. */
+  subject: string | number;
+  /**
+   * The matched organization's `entitlements.restoreRules`, or `undefined`
+   * when it is unknown: whoami failed, the organization carried no boolean
+   * for it, or no organization matched and the subject came from the token.
+   * A hint for which recovery to suggest, never a gate. The service refuses
+   * a recovery the plan lacks on its own.
+   */
+  restoreRules?: boolean;
+}
+
 /**
- * The org subject to send on write calls. Prefers the current org's Taskless
- * UUID (`id`), resolved by matching the repo's remotes against `whoami`; falls
+ * The org to act as. `subject` prefers the current org's Taskless UUID
+ * (`id`), resolved by matching the repo's remotes against `whoami`; falls
  * back to the token's canonical id (or its deprecated numeric `orgId` claim)
  * when whoami is unavailable or no org owns the repo. A new client thus routes
  * multi-org users correctly, while older single-org behaviour is preserved via
  * the claim.
  *
- * Never `undefined`: when neither a matched org nor a token claim resolves,
- * returns the nil-UUID `NIL_ORG_ID` so the canonical subject is always known
- * (unattributed usage is sent under one stable, known org id).
+ * `subject` is never `undefined`: when neither a matched org nor a token claim
+ * resolves, it is the nil-UUID `NIL_ORG_ID` so the canonical subject is always
+ * known (unattributed usage is sent under one stable, known org id).
+ *
+ * `restoreRules` rides along from the same whoami call, so knowing the plan
+ * costs no request of its own.
  */
-export async function resolveOrgSubject(
+export async function resolveActingOrg(
   cwd: string,
   token: string
-): Promise<string | number> {
+): Promise<ActingOrg> {
   const whoami = await fetchWhoami(token);
   if (whoami && whoami.orgs.length > 0) {
     const org = await resolveCurrentOrg(cwd, whoami.orgs);
-    if (org) return org.id;
+    if (org) {
+      // Read defensively: the schema says an absent `entitlements` means
+      // unknown, and an older service may not send it at all.
+      const restoreRules: unknown = org.entitlements?.restoreRules;
+      return typeof restoreRules === "boolean"
+        ? { subject: org.id, restoreRules }
+        : { subject: org.id };
+    }
   }
-  return decodeOrgId(token) ?? NIL_ORG_ID;
+  return { subject: decodeOrgId(token) ?? NIL_ORG_ID };
 }

packages/cli/src/rules/plan-check.ts

@@ -1,5 +1,5 @@
 import { getToken } from "../auth/token";
-import { resolveOrgSubject } from "../auth/org";
+import { resolveActingOrg } from "../auth/org";
 import { reconcileRules } from "../api/v2";
 import { resolveRepositoryUrl } from "../util/git-remote";
 import { getCliPrefix } from "../util/package-manager";
@@ -191,9 +191,9 @@ export async function planCheck(
     };
   }
 
-  const orgSubject = await resolveOrgSubject(cwd, token);
+  const org = await resolveActingOrg(cwd, token);
   const outcome = await reconcileRules(token, {
-    orgId: orgSubject,
+    orgId: org.subject,
     repositoryUrl,
     rules: report.rules.map(({ ruleId, files }) => ({ ruleId, files })),
   });

packages/cli/test/error-envelope.test.ts

@@ -6,7 +6,7 @@ import { join, resolve } from "node:path";
 import { promisify } from "node:util";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
 
-import { resolveOrgSubject } from "../src/auth/org";
+import { resolveActingOrg } from "../src/auth/org";
 import { cliRejectionToResult } from "./support/spawn-cli";
 
 const execFileAsync = promisify(execFile);
@@ -290,7 +290,7 @@ describe("standardized error envelope (--json)", () => {
  * rewording any of these strings fails a test instead of silently telling a
  * machine consumer to log in when the real problem is the git remote.
  *
- * Origin 4, a `resolveOrgSubject` failure, is absent on purpose and is
+ * Origin 4, a `resolveActingOrg` failure, is absent on purpose and is
  * covered by its own test below: it cannot throw.
  */
 describe("resolveIdentity failure codes (--json)", () => {
@@ -501,14 +501,14 @@ describe("resolveIdentity failure codes (--json)", () => {
 });
 
 /**
- * Origin 4 of #181: a `resolveOrgSubject` failure. There is no such failure
+ * Origin 4 of #181: a `resolveActingOrg` failure. There is no such failure
  * to give a code to. `fetchWhoami` returns `undefined` on any network or HTTP
  * error, and `decodeOrgId` falls back to the nil-UUID `NIL_ORG_ID`, so the
  * step resolves a subject rather than throwing. This test holds that shape in
- * place: if `resolveOrgSubject` ever grows a throw, it needs a code of its
+ * place: if `resolveActingOrg` ever grows a throw, it needs a code of its
  * own and this test says so by failing.
  */
-describe("resolveOrgSubject", () => {
+describe("resolveActingOrg", () => {
   let cwd: string;
 
   beforeEach(async () => {
@@ -525,7 +525,7 @@ describe("resolveOrgSubject", () => {
     process.env.TASKLESS_API_URL = "http://127.0.0.1:1/cli";
     try {
       await expect(
-        resolveOrgSubject(cwd, "not-a-real-token")
+        resolveActingOrg(cwd, "not-a-real-token")
       ).resolves.toBeDefined();
     } finally {
       if (previous === undefined) {

packages/cli/test/org.test.ts

@@ -6,7 +6,7 @@ import { promisify } from "node:util";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { fetchWhoami } from "../src/auth/whoami";
 import {
-  resolveOrgSubject,
+  resolveActingOrg,
   selectOrgForOwners,
   type WhoamiOrg,
 } from "../src/auth/org";
@@ -39,6 +39,15 @@ const org = (id: string, url: string, source = "github"): WhoamiOrg =>
     url,
   }) as WhoamiOrg;
 
+const entitled = (restoreRules: boolean): WhoamiOrg => ({
+  ...org("uuid-acme", "https://github.com/acme"),
+  entitlements: {
+    remoteGeneration: true,
+    runtimeSignatures: true,
+    restoreRules,
+  },
+});
+
 describe("selectOrgForOwners", () => {
   const orgs = [
     org("uuid-acme", "https://github.com/acme"),
@@ -78,7 +87,7 @@ describe("selectOrgForOwners", () => {
   });
 });
 
-describe("resolveOrgSubject", () => {
+describe("resolveActingOrg", () => {
   let directory: string;
   const token = makeJwt({ orgId: 4242 });
 
@@ -100,25 +109,84 @@ describe("resolveOrgSubject", () => {
     mockedFetchWhoami.mockResolvedValue(
       whoamiWith([org("uuid-acme", "https://github.com/acme")])
     );
-    expect(await resolveOrgSubject(directory, token)).toBe("uuid-acme");
+    await expect(resolveActingOrg(directory, token)).resolves.toMatchObject({
+      subject: "uuid-acme",
+    });
   });
 
   it("falls back to the numeric claim when no whoami org matches the repo", async () => {
     mockedFetchWhoami.mockResolvedValue(
       whoamiWith([org("uuid-other", "https://github.com/other")])
     );
-    expect(await resolveOrgSubject(directory, token)).toBe(4242);
+    await expect(resolveActingOrg(directory, token)).resolves.toMatchObject({
+      subject: 4242,
+    });
   });
 
   it("falls back to the numeric claim when whoami is unavailable", async () => {
     mockedFetchWhoami.mockResolvedValue(WHOAMI_UNAVAILABLE);
-    expect(await resolveOrgSubject(directory, token)).toBe(4242);
+    await expect(resolveActingOrg(directory, token)).resolves.toMatchObject({
+      subject: 4242,
+    });
   });
 
   it("falls back to the nil-UUID when whoami is unavailable and the token has no org claim", async () => {
     mockedFetchWhoami.mockResolvedValue(WHOAMI_UNAVAILABLE);
-    expect(await resolveOrgSubject(directory, makeJwt({ sub: "u" }))).toBe(
-      "00000000-0000-0000-0000-000000000000"
+    await expect(
+      resolveActingOrg(directory, makeJwt({ sub: "u" }))
+    ).resolves.toMatchObject({
+      subject: "00000000-0000-0000-0000-000000000000",
+    });
+  });
+
+  it.each([true, false])(
+    "carries the matched org's restoreRules (%s) from the same whoami call",
+    async (restoreRules) => {
+      mockedFetchWhoami.mockResolvedValue(whoamiWith([entitled(restoreRules)]));
+      expect(await resolveActingOrg(directory, token)).toEqual({
+        subject: "uuid-acme",
+        restoreRules,
+      });
+      expect(mockedFetchWhoami).toHaveBeenCalledTimes(1);
+    }
+  );
+
+  it("leaves restoreRules unknown when the matched org has no entitlements", async () => {
+    mockedFetchWhoami.mockResolvedValue(
+      whoamiWith([org("uuid-acme", "https://github.com/acme")])
     );
+    expect(await resolveActingOrg(directory, token)).toEqual({
+      subject: "uuid-acme",
+    });
+  });
+
+  it("leaves restoreRules unknown when entitlements carries no boolean for it", async () => {
+    mockedFetchWhoami.mockResolvedValue(
+      whoamiWith([
+        {
+          ...org("uuid-acme", "https://github.com/acme"),
+          entitlements: { remoteGeneration: true } as WhoamiOrg["entitlements"],
+        },
+      ])
+    );
+    expect(await resolveActingOrg(directory, token)).toEqual({
+      subject: "uuid-acme",
+    });
+  });
+
+  it("leaves restoreRules unknown when no org matches, even if another org has it", async () => {
+    mockedFetchWhoami.mockResolvedValue(
+      whoamiWith([
+        { ...entitled(false), url: "https://github.com/other" } as WhoamiOrg,
+      ])
+    );
+    expect(await resolveActingOrg(directory, token)).toEqual({
+      subject: 4242,
+    });
+  });
+
+  it("leaves restoreRules unknown when whoami is unavailable", async () => {
+    mockedFetchWhoami.mockResolvedValue(WHOAMI_UNAVAILABLE);
+    expect(await resolveActingOrg(directory, token)).toEqual({ subject: 4242 });
   });
 });