Просмотр исходного кода

fix(codemode): reject Promise.all immediately (#35604)

Kit Langton 1 месяц назад
Родитель
Сommit
df471872de
2 измененных файлов с 58 добавлено и 6 удалено
  1. 27 5
      packages/codemode/src/interpreter/runtime.ts
  2. 31 1
      packages/codemode/test/promise.test.ts

+ 27 - 5
packages/codemode/src/interpreter/runtime.ts

@@ -2229,14 +2229,36 @@ class Interpreter<R> {
     switch (ref.name) {
       case "all": {
         // Mark every promise element observed up-front (Promise.all handles all of its
-        // members' failures, as in JS), then join in index order; the first failure rejects
-        // the whole call while unrelated in-flight members keep running.
-        const settles = items.map((item) =>
-          item instanceof SandboxPromise ? this.settlePromise(item, node) : Effect.succeed(item),
+        // members' failures, as in JS), race their settlements for fail-fast rejection, and
+        // preserve input order when they all fulfill. Rejected calls keep draining siblings.
+        const observations = items.map((item, index) =>
+          item instanceof SandboxPromise
+            ? Effect.map(this.observePromise(item), (exit) => ({ index, item, exit }))
+            : Effect.succeed({ index, item: undefined, exit: Exit.succeed(item) }),
         )
         return Effect.gen(function* () {
+          const remaining = [...observations]
           const values: Array<unknown> = []
-          for (const settle of settles) values.push(yield* settle)
+          values.length = items.length
+          while (remaining.length > 0) {
+            const winner = yield* Effect.raceAll(remaining)
+            const position = remaining.indexOf(observations[winner.index])
+            if (position >= 0) remaining.splice(position, 1)
+            if (Exit.isSuccess(winner.exit)) {
+              values[winner.index] = winner.exit.value
+              continue
+            }
+            yield* self.createPromise(
+              Effect.asVoid(
+                Effect.forEach(
+                  items,
+                  (item) => (item instanceof SandboxPromise ? self.observePromise(item) : Effect.void),
+                  { concurrency: "unbounded" },
+                ),
+              ),
+            )
+            return yield* self.unwrapPromiseExit(winner.item, winner.exit, node)
+          }
           return values
         })
       }

+ 31 - 1
packages/codemode/test/promise.test.ts

@@ -48,6 +48,14 @@ const failingTool = Tool.make({
   run: () => Effect.fail(toolError("Lookup refused")),
 })
 
+const completedTool = (trace: Trace) =>
+  Tool.make({
+    description: "Return the number of completed sleepy calls",
+    input: Schema.Struct({}),
+    output: Schema.Number,
+    run: () => Effect.succeed(trace.completed),
+  })
+
 const run = (
   code: string,
   options: { trace?: Trace; limits?: CodeMode.ExecutionLimits } = {},
@@ -55,7 +63,7 @@ const run = (
   const trace = options.trace ?? makeTrace()
   return Effect.runPromise(
     CodeMode.execute({
-      tools: { host: { sleepy: sleepyTool(trace), fail: failingTool } },
+      tools: { host: { sleepy: sleepyTool(trace), fail: failingTool, completed: completedTool(trace) } },
       code,
       ...(options.limits ? { limits: options.limits } : {}),
     }),
@@ -338,6 +346,28 @@ describe("Promise.all over arbitrary arrays", () => {
     ).toBe("Lookup refused")
   })
 
+  test("rejects before an earlier slow promise fulfills", async () => {
+    const trace = makeTrace()
+    expect(
+      await value(
+        `
+          try {
+            await Promise.all([
+              tools.host.sleepy({ id: 1, ms: 100 }),
+              tools.host.fail({}),
+            ])
+            return -1
+          } catch {
+            return await tools.host.completed({})
+          }
+        `,
+        { trace },
+      ),
+    ).toBe(0)
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
+
   test("a non-collection argument is a clear error", async () => {
     const diagnostic = await error(`return await Promise.all(42)`)
     expect(diagnostic.message).toContain("Promise.all expects an array")