fix: let waitFor accept TaskBuilders, fix type issues (#7)

Author: AlCalzone

src/Task.test.ts

@@ -1637,11 +1637,18 @@ test("The type-safe wrapper works correctly", async (t) => {
 	const scheduler = new TaskScheduler();
 	scheduler.start();
 
+	const inner = {
+		priority: TaskPriority.Normal,
+		task: async function* () {
+			return 42;
+		},
+	};
+
 	const task = scheduler.queueTask({
 		priority: TaskPriority.Normal,
 		task: async function* () {
 			const foo: string = yield* waitFor(Promise.resolve("bar"));
-			const baz: number = yield* waitFor(Promise.resolve(42));
+			const baz = yield* waitFor(inner);
 			return [foo, baz] as const;
 		},
 	});

src/Task.ts

@@ -78,9 +78,9 @@ export interface TaskBuilder<
 	 * Yielded Promises should not spawn new tasks. If they do, the spawned tasks MUST have a higher priority than the parent task.
 	 */
 	task: () => AsyncGenerator<
-		| (() => Promise<TInner> | TaskBuilder<TReturn, TaskTag, TInner>)
-		| (() => TaskBuilder<TReturn, TaskTag, TInner>)
-		| TaskBuilder<TReturn, TaskTag, TInner>
+		| (() => Promise<unknown> | TaskBuilder<unknown, TaskTag, unknown>)
+		| (() => TaskBuilder<unknown, TaskTag, unknown>)
+		| TaskBuilder<unknown, TaskTag, unknown>
 		| undefined,
 		TReturn,
 		TInner

src/index.ts

@@ -9,4 +9,4 @@ export {
 	type TaskStepResult,
 	type TaskConcurrencyGroup,
 } from "./Task.js";
-export { waitFor } from "./utils.js";
+export { waitFor, type WaitForReturnType } from "./utils.js";

src/utils.ts

@@ -1,3 +1,5 @@
+import { type TaskBuilder } from "./Task.js";
+
 /** Returns a timestamp with nano-second precision */
 export function highResTimestamp(): number {
 	if (typeof process !== "undefined") {
@@ -31,9 +33,12 @@ export function evalOrStatic<T>(
 	return typeof fnOrConst === "function" ? fnOrConst(...args) : fnOrConst;
 }
 
-export async function* waitFor<T>(
-	promise: Promise<T>,
-): AsyncGenerator<() => Promise<T>, T> {
-	const result = yield () => promise;
+export type WaitForReturnType<T> =
+	T extends Promise<infer U> ? U : T extends TaskBuilder<infer U> ? U : never;
+
+export async function* waitFor<
+	T extends Promise<unknown> | TaskBuilder<unknown>,
+>(thing: T): AsyncGenerator<() => T, WaitForReturnType<T>> {
+	const result = yield () => thing;
 	return result;
 }