refactor(iam): add candidate inline policy swaps

src/core/executionRole.tsx

@@ -1,9 +1,4 @@
-import {
-  CreateRoleCommand,
-  GetRoleCommand,
-  PutRolePolicyCommand,
-  type IAMClient,
-} from "@aws-sdk/client-iam";
+import { CreateRoleCommand, GetRoleCommand, type IAMClient } from "@aws-sdk/client-iam";
 
 // Default harness execution role provisioning.
 //
@@ -15,7 +10,13 @@ import {
 // role is reused and its inline policy refreshed — so repeated creates of the
 // same harness name converge on one role.
 
-const POLICY_NAME = "AgentCoreHarnessExecutionPolicy";
+export const HARNESS_EXECUTION_POLICY_NAME = "AgentCoreHarnessExecutionPolicy";
+
+export type HarnessExecutionRole = {
+  roleArn: string;
+  roleName: string;
+  policyDocument: string;
+};
 
 // executionRoleName derives the default role's name from the harness name. IAM
 // role names cap at 64 characters; harness names are alphanumeric/underscore so
@@ -230,14 +231,13 @@ function accountIdFromRoleArn(arn: string): string {
   return accountId;
 }
 
-// ensureDefaultExecutionRole returns the ARN of the default execution role for
-// `harnessName`, creating the role if it doesn't exist and (re)attaching the
-// default inline policy either way.
+// ensureDefaultExecutionRole creates or reuses the role and returns the complete
+// policy document that HarnessClient applies transactionally around creation.
 export async function ensureDefaultExecutionRole(
   iam: IAMClient,
   harnessName: string,
   region: string,
-): Promise<string> {
+): Promise<HarnessExecutionRole> {
   const roleName = executionRoleName(harnessName);
 
   let roleArn: string;
@@ -256,13 +256,9 @@ export async function ensureDefaultExecutionRole(
     roleArn = created.Role!.Arn!;
   }
 
-  await iam.send(
-    new PutRolePolicyCommand({
-      RoleName: roleName,
-      PolicyName: POLICY_NAME,
-      PolicyDocument: executionPolicy(region, accountIdFromRoleArn(roleArn), harnessName),
-    }),
-  );
-
-  return roleArn;
+  return {
+    roleArn,
+    roleName,
+    policyDocument: executionPolicy(region, accountIdFromRoleArn(roleArn), harnessName),
+  };
 }

src/core/harness.tsx

@@ -38,7 +38,8 @@ import {
 import type { CoreHarnessClient, CreateHarnessInput } from "../handlers/harness/types";
 import type { AwsClients, CoreOptions } from "./types";
 import { abortable } from "./abortable";
-import { ensureDefaultExecutionRole } from "./executionRole";
+import { ensureDefaultExecutionRole, HARNESS_EXECUTION_POLICY_NAME } from "./executionRole";
+import { InlinePolicySwap } from "./inlinePolicySwap";
 import { toClientConfig } from "./utils";
 
 // HarnessClient implements the harness-facing operations on top of the shared AWS
@@ -119,15 +120,21 @@ export class HarnessClient implements CoreHarnessClient {
     // create the harness with it. IAM is eventually consistent — a role created
     // moments ago may not yet be assumable by the AgentCore service principal —
     // so retry the create while the service reports the role as unusable.
-    const defaultRoleArn = await ensureDefaultExecutionRole(
+    const iam = this.clients.iam({ region: options.region });
+    const role = await ensureDefaultExecutionRole(
       // IAM is a global service; the region only selects the endpoint, and the
       // agentcore endpoint override must not leak onto it.
-      this.clients.iam({ region: options.region }),
+      iam,
       input.harnessName!,
       options.region,
     );
-    return retryWhileRoleUnassumable(() =>
-      control.send(new CreateHarnessCommand({ ...request, executionRoleArn: defaultRoleArn })),
+    return new InlinePolicySwap(iam, {
+      roleName: role.roleName,
+      policyNamePrefix: HARNESS_EXECUTION_POLICY_NAME,
+    }).run(role.policyDocument, () =>
+      retryWhileRoleUnassumable(() =>
+        control.send(new CreateHarnessCommand({ ...request, executionRoleArn: role.roleArn })),
+      ),
     );
   }
 

src/core/inlinePolicySwap.test.ts

@@ -0,0 +1,239 @@
+import { expect, test } from "bun:test";
+import {
+  DeleteRolePolicyCommand,
+  ListRolePoliciesCommand,
+  PutRolePolicyCommand,
+  type IAMClient,
+} from "@aws-sdk/client-iam";
+import { InlinePolicySwap } from "./inlinePolicySwap";
+
+const ROLE_NAME = "AgentCoreHarness-orders";
+const POLICY_PREFIX = "AgentCoreHarnessExecutionPolicy";
+const POLICY_DOCUMENT = '{"Version":"2012-10-17","Statement":[]}';
+
+type SentCommand = DeleteRolePolicyCommand | ListRolePoliciesCommand | PutRolePolicyCommand;
+
+function iamClient(send: (command: SentCommand) => Promise<Record<string, unknown>>): IAMClient {
+  return { send } as unknown as IAMClient;
+}
+
+function policySwap(iam: IAMClient): InlinePolicySwap {
+  return new InlinePolicySwap(iam, {
+    roleName: ROLE_NAME,
+    policyNamePrefix: POLICY_PREFIX,
+  });
+}
+
+test("stages a complete candidate and removes previous family policies after success", async () => {
+  const sent: SentCommand[] = [];
+  const oldCandidate = `${POLICY_PREFIX}-${"a".repeat(32)}`;
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return {
+        PolicyNames: [POLICY_PREFIX, oldCandidate, "CustomerPolicy"],
+        IsTruncated: false,
+      };
+    }
+    return {};
+  });
+
+  const result = await policySwap(iam).run(POLICY_DOCUMENT, async () => {
+    expect(sent.at(-1)).toBeInstanceOf(PutRolePolicyCommand);
+    return "created";
+  });
+
+  const candidateName = InlinePolicySwap.candidatePolicyName(POLICY_PREFIX, POLICY_DOCUMENT);
+  expect(result).toBe("created");
+  expect(sent.map((command) => command.constructor.name)).toEqual([
+    "ListRolePoliciesCommand",
+    "PutRolePolicyCommand",
+    "DeleteRolePolicyCommand",
+    "DeleteRolePolicyCommand",
+  ]);
+  expect((sent[1] as PutRolePolicyCommand).input.PolicyName).toBe(candidateName);
+  expect(
+    sent
+      .filter((command) => command instanceof DeleteRolePolicyCommand)
+      .map((command) => command.input.PolicyName),
+  ).toEqual([POLICY_PREFIX, oldCandidate]);
+});
+
+test("removes a new candidate and preserves previous policies when the operation fails", async () => {
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [POLICY_PREFIX], IsTruncated: false };
+    }
+    return {};
+  });
+
+  await expect(
+    policySwap(iam).run(POLICY_DOCUMENT, async () => {
+      throw new Error("deployment failed");
+    }),
+  ).rejects.toThrow("deployment failed");
+
+  expect(
+    sent
+      .filter((command) => command instanceof DeleteRolePolicyCommand)
+      .map((command) => command.input.PolicyName),
+  ).toEqual([InlinePolicySwap.candidatePolicyName(POLICY_PREFIX, POLICY_DOCUMENT)]);
+});
+
+test("keeps a pre-existing candidate when the operation fails", async () => {
+  const candidateName = InlinePolicySwap.candidatePolicyName(POLICY_PREFIX, POLICY_DOCUMENT);
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [candidateName], IsTruncated: false };
+    }
+    return {};
+  });
+
+  await expect(
+    policySwap(iam).run(POLICY_DOCUMENT, async () => {
+      throw new Error("deployment failed");
+    }),
+  ).rejects.toThrow("deployment failed");
+
+  expect(sent.some((command) => command instanceof DeleteRolePolicyCommand)).toBeFalse();
+});
+
+test("removes a first candidate when the first operation fails", async () => {
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [], IsTruncated: false };
+    }
+    return {};
+  });
+
+  await expect(
+    policySwap(iam).run(POLICY_DOCUMENT, async () => {
+      throw new Error("deployment failed");
+    }),
+  ).rejects.toThrow("deployment failed");
+
+  expect(sent.map((command) => command.constructor.name)).toEqual([
+    "ListRolePoliciesCommand",
+    "PutRolePolicyCommand",
+    "DeleteRolePolicyCommand",
+  ]);
+});
+
+test("preserves both operation and rollback failures", async () => {
+  const iam = iamClient(async (command) => {
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [POLICY_PREFIX], IsTruncated: false };
+    }
+    if (command instanceof DeleteRolePolicyCommand) throw new Error("rollback failed");
+    return {};
+  });
+
+  const error = await policySwap(iam)
+    .run(POLICY_DOCUMENT, async () => {
+      throw new Error("deployment failed");
+    })
+    .catch((caught) => caught);
+
+  expect(error).toBeInstanceOf(AggregateError);
+  expect((error as AggregateError).errors.map((cause) => (cause as Error).message)).toEqual([
+    "deployment failed",
+    "One or more IAM policies could not be removed",
+  ]);
+});
+
+test("keeps the candidate when previous-policy cleanup fails after success", async () => {
+  const candidateName = InlinePolicySwap.candidatePolicyName(POLICY_PREFIX, POLICY_DOCUMENT);
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [POLICY_PREFIX], IsTruncated: false };
+    }
+    if (command instanceof DeleteRolePolicyCommand && command.input.PolicyName === POLICY_PREFIX) {
+      throw new Error("cleanup failed");
+    }
+    return {};
+  });
+
+  await expect(policySwap(iam).run(POLICY_DOCUMENT, async () => "created")).rejects.toThrow(
+    "Operation succeeded but previous IAM policies could not be removed",
+  );
+
+  expect((sent[1] as PutRolePolicyCommand).input.PolicyName).toBe(candidateName);
+  expect(
+    sent
+      .filter((command) => command instanceof DeleteRolePolicyCommand)
+      .map((command) => command.input.PolicyName),
+  ).toEqual([POLICY_PREFIX]);
+});
+
+test("lists every policy page before staging the candidate", async () => {
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return command.input.Marker
+        ? { PolicyNames: [`${POLICY_PREFIX}-${"b".repeat(32)}`], IsTruncated: false }
+        : { PolicyNames: [POLICY_PREFIX], IsTruncated: true, Marker: "page-2" };
+    }
+    return {};
+  });
+
+  await policySwap(iam).run(POLICY_DOCUMENT, async () => "created");
+
+  expect(sent.slice(0, 3).map((command) => command.constructor.name)).toEqual([
+    "ListRolePoliciesCommand",
+    "ListRolePoliciesCommand",
+    "PutRolePolicyCommand",
+  ]);
+});
+
+test("removes previous family policies after a successful operation with no candidate", async () => {
+  const sent: SentCommand[] = [];
+  const oldCandidate = `${POLICY_PREFIX}-${"a".repeat(32)}`;
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return {
+        PolicyNames: [POLICY_PREFIX, oldCandidate, "CustomerPolicy"],
+        IsTruncated: false,
+      };
+    }
+    return {};
+  });
+
+  await policySwap(iam).run(undefined, async () => {
+    expect(sent.some((command) => command instanceof PutRolePolicyCommand)).toBeFalse();
+  });
+
+  expect(
+    sent
+      .filter((command) => command instanceof DeleteRolePolicyCommand)
+      .map((command) => command.input.PolicyName),
+  ).toEqual([POLICY_PREFIX, oldCandidate]);
+});
+
+test("preserves previous family policies after a failed operation with no candidate", async () => {
+  const sent: SentCommand[] = [];
+  const iam = iamClient(async (command) => {
+    sent.push(command);
+    if (command instanceof ListRolePoliciesCommand) {
+      return { PolicyNames: [POLICY_PREFIX], IsTruncated: false };
+    }
+    return {};
+  });
+
+  await expect(
+    policySwap(iam).run(undefined, async () => {
+      throw new Error("deployment failed");
+    }),
+  ).rejects.toThrow("deployment failed");
+
+  expect(sent.map((command) => command.constructor.name)).toEqual(["ListRolePoliciesCommand"]);
+});