diff --git a/actions/setup/js/assign_milestone.cjs b/actions/setup/js/assign_milestone.cjs
index b10b7c3c960..1ffbc46c73b 100644
--- a/actions/setup/js/assign_milestone.cjs
+++ b/actions/setup/js/assign_milestone.cjs
@@ -8,9 +8,7 @@
 const { getErrorMessage } = require("./error_helpers.cjs");
 const { logStagedPreviewInfo } = require("./staged_preview.cjs");
 const { createAuthenticatedGitHubClient } = require("./handler_auth.cjs");
-
-/** @type {string} Safe output type handled by this module */
-const HANDLER_TYPE = "assign_milestone";
+const { loadTemporaryIdMapFromResolved, resolveRepoIssueTarget } = require("./temporary_id.cjs");
 
 /**
  * Main handler factory for assign_milestone
@@ -55,24 +53,37 @@ async function main(config = {}) {
 
     processedCount++;
 
-    const item = message;
+    // Build temporary ID map from resolved IDs
+    const temporaryIdMap = loadTemporaryIdMapFromResolved(resolvedTemporaryIds);
+
+    // Determine target issue number, with temporary ID support
+    const resolvedTarget = resolveRepoIssueTarget(message.issue_number, temporaryIdMap, context.repo.owner, context.repo.repo);
 
-    const issueNumber = Number(item.issue_number);
-    const milestoneNumber = Number(item.milestone_number);
+    if (resolvedTarget.wasTemporaryId && !resolvedTarget.resolved) {
+      core.info(`Deferring assign_milestone: unresolved temporary ID (${message.issue_number})`);
+      return {
+        success: false,
+        deferred: true,
+        error: resolvedTarget.errorMessage || `Unresolved temporary ID: ${message.issue_number}`,
+      };
+    }
 
-    if (isNaN(issueNumber) || issueNumber <= 0) {
-      core.error(`Invalid issue_number: ${item.issue_number}`);
+    if (resolvedTarget.errorMessage || !resolvedTarget.resolved) {
+      core.error(`Invalid issue_number: ${message.issue_number}`);
       return {
         success: false,
-        error: `Invalid issue_number: ${item.issue_number}`,
+        error: `Invalid issue_number: ${message.issue_number}`,
       };
     }
 
+    const issueNumber = resolvedTarget.resolved.number;
+    const milestoneNumber = Number(message.milestone_number);
+
     if (isNaN(milestoneNumber) || milestoneNumber <= 0) {
-      core.error(`Invalid milestone_number: ${item.milestone_number}`);
+      core.error(`Invalid milestone_number: ${message.milestone_number}`);
       return {
         success: false,
-        error: `Invalid milestone_number: ${item.milestone_number}`,
+        error: `Invalid milestone_number: ${message.milestone_number}`,
       };
     }
 
@@ -121,20 +132,20 @@ async function main(config = {}) {
     }
 
     // Assign the milestone to the issue
-    try {
-      // If in staged mode, preview without executing
-      if (isStaged) {
-        logStagedPreviewInfo(`Would assign milestone #${milestoneNumber} to issue #${issueNumber}`);
-        return {
-          success: true,
-          staged: true,
-          previewInfo: {
-            issue_number: issueNumber,
-            milestone_number: milestoneNumber,
-          },
-        };
-      }
+    // If in staged mode, preview without executing
+    if (isStaged) {
+      logStagedPreviewInfo(`Would assign milestone #${milestoneNumber} to issue #${issueNumber}`);
+      return {
+        success: true,
+        staged: true,
+        previewInfo: {
+          issue_number: issueNumber,
+          milestone_number: milestoneNumber,
+        },
+      };
+    }
 
+    try {
       await githubClient.rest.issues.update({
         owner: context.repo.owner,
         repo: context.repo.repo,
diff --git a/actions/setup/js/assign_milestone.test.cjs b/actions/setup/js/assign_milestone.test.cjs
index 2a05d251c9d..88c79140fb2 100644
--- a/actions/setup/js/assign_milestone.test.cjs
+++ b/actions/setup/js/assign_milestone.test.cjs
@@ -209,4 +209,172 @@ describe("assign_milestone (Handler Factory Architecture)", () => {
     expect(result.success).toBe(false);
     expect(result.error).toContain("Invalid milestone_number");
   });
+
+  it("should return staged preview without calling API when staged mode is enabled", async () => {
+    const origStaged = process.env.GH_AW_SAFE_OUTPUTS_STAGED;
+    process.env.GH_AW_SAFE_OUTPUTS_STAGED = "true";
+
+    const { main } = require("./assign_milestone.cjs");
+    const stagedHandler = await main({ max: 10 });
+
+    const message = {
+      type: "assign_milestone",
+      issue_number: 42,
+      milestone_number: 5,
+    };
+
+    const result = await stagedHandler(message, {});
+
+    expect(result.success).toBe(true);
+    expect(result.staged).toBe(true);
+    expect(result.previewInfo.issue_number).toBe(42);
+    expect(result.previewInfo.milestone_number).toBe(5);
+    expect(mockGithub.rest.issues.update).not.toHaveBeenCalled();
+
+    process.env.GH_AW_SAFE_OUTPUTS_STAGED = origStaged ?? "";
+  });
+
+  it("should handle failure when fetching milestones for allowed list validation", async () => {
+    const { main } = require("./assign_milestone.cjs");
+    const handlerWithAllowed = await main({
+      max: 10,
+      allowed: ["v1.0"],
+    });
+
+    mockGithub.rest.issues.listMilestones.mockRejectedValue(new Error("Network error"));
+
+    const result = await handlerWithAllowed({ issue_number: 42, milestone_number: 5 }, {});
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("Failed to fetch milestones for validation");
+    expect(result.error).toContain("Network error");
+  });
+
+  it("should return error when milestone not found in repository", async () => {
+    const { main } = require("./assign_milestone.cjs");
+    const handlerWithAllowed = await main({
+      max: 10,
+      allowed: ["v1.0"],
+    });
+
+    mockGithub.rest.issues.listMilestones.mockResolvedValue({
+      data: [{ number: 7, title: "v2.0" }],
+    });
+
+    const result = await handlerWithAllowed({ issue_number: 42, milestone_number: 99 }, {});
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("not found in repository");
+  });
+
+  it("should allow milestone matched by number string in allowed list", async () => {
+    const { main } = require("./assign_milestone.cjs");
+    const handlerWithAllowed = await main({
+      max: 10,
+      allowed: ["5"], // Allowed by number as string
+    });
+
+    mockGithub.rest.issues.listMilestones.mockResolvedValue({
+      data: [{ number: 5, title: "v1.0" }],
+    });
+    mockGithub.rest.issues.update.mockResolvedValue({});
+
+    const result = await handlerWithAllowed({ issue_number: 42, milestone_number: 5 }, {});
+
+    expect(result.success).toBe(true);
+  });
+
+  it("should use cached milestones on second call", async () => {
+    const { main } = require("./assign_milestone.cjs");
+    const handlerWithAllowed = await main({
+      max: 10,
+      allowed: ["v1.0"],
+    });
+
+    mockGithub.rest.issues.listMilestones.mockResolvedValue({
+      data: [{ number: 5, title: "v1.0" }],
+    });
+    mockGithub.rest.issues.update.mockResolvedValue({});
+
+    await handlerWithAllowed({ issue_number: 42, milestone_number: 5 }, {});
+    await handlerWithAllowed({ issue_number: 43, milestone_number: 5 }, {});
+
+    // Should only fetch milestones once (cached after first call)
+    expect(mockGithub.rest.issues.listMilestones).toHaveBeenCalledTimes(1);
+  });
+
+  it("should handle zero issue_number as invalid", async () => {
+    const result = await handler({ issue_number: 0, milestone_number: 5 }, {});
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("Invalid issue_number");
+  });
+
+  it("should handle zero milestone_number as invalid", async () => {
+    const result = await handler({ issue_number: 42, milestone_number: 0 }, {});
+
+    expect(result.success).toBe(false);
+    expect(result.error).toContain("Invalid milestone_number");
+  });
+
+  describe("temporary ID resolution", () => {
+    it("should resolve temporary ID in issue_number to real issue number", async () => {
+      mockGithub.rest.issues.update.mockResolvedValue({});
+
+      const result = await handler(
+        {
+          type: "assign_milestone",
+          issue_number: "aw_issue1",
+          milestone_number: 5,
+        },
+        { aw_issue1: { repo: "test-owner/test-repo", number: 42 } }
+      );
+
+      expect(result.success).toBe(true);
+      expect(result.issue_number).toBe(42);
+      expect(mockGithub.rest.issues.update).toHaveBeenCalledWith({
+        owner: "test-owner",
+        repo: "test-repo",
+        issue_number: 42,
+        milestone: 5,
+      });
+    });
+
+    it("should defer when issue_number is an unresolved temporary ID", async () => {
+      const result = await handler(
+        {
+          type: "assign_milestone",
+          issue_number: "aw_issue1",
+          milestone_number: 5,
+        },
+        {}
+      );
+
+      expect(result.success).toBe(false);
+      expect(result.deferred).toBe(true);
+      expect(result.error).toContain("aw_issue1");
+    });
+
+    it("should resolve temporary ID with hash prefix in issue_number", async () => {
+      mockGithub.rest.issues.update.mockResolvedValue({});
+
+      const result = await handler(
+        {
+          type: "assign_milestone",
+          issue_number: "#aw_issue1",
+          milestone_number: 5,
+        },
+        { aw_issue1: { repo: "test-owner/test-repo", number: 99 } }
+      );
+
+      expect(result.success).toBe(true);
+      expect(result.issue_number).toBe(99);
+      expect(mockGithub.rest.issues.update).toHaveBeenCalledWith({
+        owner: "test-owner",
+        repo: "test-repo",
+        issue_number: 99,
+        milestone: 5,
+      });
+    });
+  });
 });