feat(codemode): support custom async iterators (#38141)

This commit is contained in:
Aiden Cline 2026-07-21 12:18:46 -05:00 committed by GitHub
commit 6f4b9504e5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 636 additions and 51 deletions

View file

@ -36,7 +36,7 @@ export type StatementResult =
export type MemberReference = {
target: SafeObject | Array<unknown> | CodeModeRegExp | CodeModeURL
key: string | number
key: PropertyKey
}
export class CodeModeFunction {
@ -61,6 +61,12 @@ export class ComputedValue {
export class PromiseNamespace {}
export class SymbolNamespace {}
export const AsyncIteratorSymbol: unique symbol = Symbol("codemode.async-iterator")
export const IteratorSymbol: unique symbol = Symbol("codemode.iterator")
export const IteratorSymbols = [AsyncIteratorSymbol, IteratorSymbol] as const
export type PromiseMethodName = "all" | "allSettled" | "race" | "any" | "resolve" | "reject"
export class PromiseMethodReference {

View file

@ -13,6 +13,7 @@ import {
PromiseMethodReference,
PromiseNamespace,
SearchFunction,
SymbolNamespace,
UriFunction,
} from "./model.js"
import { ToolReference } from "../tool-runtime.js"
@ -34,6 +35,7 @@ export const isRuntimeReference = (value: unknown): boolean =>
value instanceof SearchFunction ||
value instanceof PromiseCapabilityFunction ||
value instanceof ErrorConstructorReference ||
value instanceof SymbolNamespace ||
isCodeModeValue(value)
function* childValues(value: object): Generator<unknown> {
@ -113,7 +115,8 @@ export const typeofValue = (value: unknown): string => {
value instanceof PromiseInstanceMethodReference ||
value instanceof PromiseNamespace ||
value instanceof PromiseCapabilityFunction ||
value instanceof ErrorConstructorReference
value instanceof ErrorConstructorReference ||
value instanceof SymbolNamespace
)
return "function"
if (value instanceof UriFunction || value instanceof SearchFunction) return "function"

View file

@ -1,7 +1,8 @@
import { Cause, Effect } from "effect"
import { Cause, Effect, Exit } from "effect"
import { isBlockedMember, ToolReference, ToolRuntimeError, type SafeObject } from "../tool-runtime.js"
import {
type AstNode,
AsyncIteratorSymbol,
asNode,
type Binding,
CodeModeFunction,
@ -19,6 +20,8 @@ import {
IntrinsicReference,
InterpreterRuntimeError,
isRecord,
IteratorSymbol,
IteratorSymbols,
JsonMethodReference,
type MemberReference,
OptionalShortCircuit,
@ -30,6 +33,7 @@ import {
ProgramThrow,
type ProgramNode,
SearchFunction,
SymbolNamespace,
type StatementResult,
supportedSyntaxMessage,
unsupportedSyntax,
@ -211,6 +215,12 @@ const loopDeclaration = (left: AstNode, statement: "for...of" | "for...in") => {
}
}
type CustomIterator = {
iterator: SafeObject
next: unknown
asynchronous: boolean
}
export class Interpreter<R> {
private scopes: ScopeStack
private readonly invokeTool: (path: ReadonlyArray<string>, args: Array<unknown>) => Effect.Effect<unknown, unknown, R>
@ -241,6 +251,7 @@ export class Interpreter<R> {
globalScope.set("tools", { mutable: false, value: new ToolReference([]) })
globalScope.set("search", { mutable: false, value: new SearchFunction() })
globalScope.set("Promise", { mutable: false, value: new PromiseNamespace() })
globalScope.set("Symbol", { mutable: false, value: new SymbolNamespace() })
globalScope.set("undefined", { mutable: false, value: undefined })
globalScope.set("Object", { mutable: false, value: new GlobalNamespace("Object") })
globalScope.set("Math", { mutable: false, value: new GlobalNamespace("Math") })
@ -636,9 +647,10 @@ export class Interpreter<R> {
const body = getNode(node, "body")
const iterable = spreadItems(right)
if (iterable === undefined) {
const iterator = iterable === undefined && awaiting ? yield* self.customIterator(right, node) : undefined
if (iterable === undefined && iterator === undefined) {
throw new InterpreterRuntimeError(
`${awaiting ? "for await...of" : "for...of"} requires an array, string, Map, Set, or URLSearchParams value.`,
`${awaiting ? "for await...of" : "for...of"} requires an array, string, Map, Set, or URLSearchParams${awaiting ? ", or custom iterator" : ""} value.`,
node,
)
}
@ -657,19 +669,14 @@ export class Interpreter<R> {
throw new InterpreterRuntimeError("Unsupported for...of binding.", left)
}
for (const value of iterable) {
const resolved = awaiting
? value instanceof CodeModePromise
? yield* self.settlePromise(value)
: yield* Effect.as(Effect.yieldNow, value)
: value
const result = yield* Effect.gen(function* () {
const evaluateBody = (value: unknown) =>
Effect.gen(function* () {
if (declared) {
self.scopes.push()
if (declared.lexical) self.predeclarePattern(declared.pattern, declared.mutable, left)
yield* self.declarePattern(declared.pattern, resolved, declared.mutable, left, declared.lexical)
yield* self.declarePattern(declared.pattern, value, declared.mutable, left, declared.lexical)
} else if (assignment) {
yield* self.assignPattern(assignment, resolved, left)
yield* self.assignPattern(assignment, value, left)
}
return yield* self.evaluateStatement(body)
}).pipe(
@ -680,22 +687,48 @@ export class Interpreter<R> {
),
)
if (iterable !== undefined) {
for (const value of iterable) {
const result = yield* evaluateBody(awaiting ? yield* self.awaitValue(value) : value)
if (result.kind === "return") return result
if (result.kind === "break") {
if (result.label !== undefined && !labels?.has(result.label)) return result
return { kind: "none" } satisfies StatementResult
}
if (result.kind === "continue" && result.label !== undefined && !labels?.has(result.label)) return result
}
return { kind: "none" } satisfies StatementResult
}
if (iterator === undefined) throw new InterpreterRuntimeError("Custom iterator is unavailable.", node)
while (true) {
const step = yield* self.nextIteratorResult(iterator, node)
if (step.done) return { kind: "none" } satisfies StatementResult
const bodyExit = yield* Effect.exit(evaluateBody(step.value))
if (!Exit.isSuccess(bodyExit)) {
// Process interruption must remain prompt; user cleanup cannot extend a timeout.
if (!Cause.hasInterruptsOnly(bodyExit.cause)) yield* Effect.exit(self.closeIterator(iterator, node))
return yield* Effect.failCause(bodyExit.cause)
}
const result = bodyExit.value
if (result.kind === "return") {
yield* self.closeIterator(iterator, node)
return result
}
if (result.kind === "break") {
yield* self.closeIterator(iterator, node)
if (result.label !== undefined && !labels?.has(result.label)) return result
return { kind: "none" } satisfies StatementResult
}
if (result.kind === "continue") {
if (result.label !== undefined && !labels?.has(result.label)) return result
continue
if (result.kind === "continue" && result.label !== undefined && !labels?.has(result.label)) {
yield* self.closeIterator(iterator, node)
return result
}
}
return { kind: "none" } satisfies StatementResult
}).pipe(
Effect.ensuring(
Effect.sync(() => {
@ -705,6 +738,101 @@ export class Interpreter<R> {
)
}
private awaitValue(value: unknown): Effect.Effect<unknown, unknown, R> {
return value instanceof CodeModePromise ? this.settlePromise(value) : Effect.as(Effect.yieldNow, value)
}
private customIterator(value: unknown, node: AstNode) {
if (!isRecord(value) || isRuntimeReference(value)) return Effect.succeed(undefined)
const asyncMethod = Reflect.get(value, AsyncIteratorSymbol)
const method = asyncMethod ?? Reflect.get(value, IteratorSymbol)
if (method === undefined || method === null) return Effect.succeed(undefined)
const self = this
return Effect.map(
this.invokeCallable(this.requireIteratorMethod(method, "Iterator method", node), [], node),
(iterator) => {
const object = self.requireIteratorObject(iterator, "Iterator method result", node)
return {
iterator: object,
next: self.requireIteratorMethod(object.next, "Iterator next", node),
asynchronous: asyncMethod !== undefined && asyncMethod !== null,
}
},
)
}
private nextIteratorResult(iterator: CustomIterator, node: AstNode) {
const self = this
return Effect.gen(function* () {
if (iterator.asynchronous) {
const object = self.requireIteratorObject(
yield* self.awaitValue(yield* self.invokeCallable(iterator.next, [], node)),
"Iterator next() result",
node,
)
return { done: Boolean(object.done), value: object.value }
}
const called = yield* Effect.exit(self.invokeCallable(iterator.next, [], node))
if (!Exit.isSuccess(called)) {
yield* Effect.yieldNow
return yield* Effect.failCause(called.cause)
}
const captured = yield* Effect.exit(
Effect.sync(() => {
const object = self.requireIteratorObject(called.value, "Iterator next() result", node)
return { done: Boolean(object.done), value: object.value }
}),
)
if (!Exit.isSuccess(captured)) {
yield* Effect.yieldNow
return yield* Effect.failCause(captured.cause)
}
return { done: captured.value.done, value: yield* self.awaitValue(captured.value.value) }
})
}
private closeIterator(iterator: CustomIterator, node: AstNode): Effect.Effect<void, unknown, R> {
const close = iterator.iterator.return
if (close === undefined || close === null) return iterator.asynchronous ? Effect.void : Effect.yieldNow
const self = this
return Effect.gen(function* () {
const method = self.requireIteratorMethod(close, "Iterator return", node)
if (iterator.asynchronous) {
self.requireIteratorObject(
yield* self.awaitValue(yield* self.invokeCallable(method, [], node)),
"Iterator return() result",
node,
)
return
}
const called = yield* Effect.exit(self.invokeCallable(method, [], node))
if (!Exit.isSuccess(called)) {
yield* Effect.yieldNow
return yield* Effect.failCause(called.cause)
}
const captured = yield* Effect.exit(
Effect.sync(() => self.requireIteratorObject(called.value, "Iterator return() result", node).value),
)
if (!Exit.isSuccess(captured)) {
yield* Effect.yieldNow
return yield* Effect.failCause(captured.cause)
}
yield* self.awaitValue(captured.value)
})
}
private requireIteratorObject(value: unknown, context: string, node: AstNode): SafeObject {
if (isRecord(value) && !isRuntimeReference(value)) return value
throw new InterpreterRuntimeError(`${context} must be an object.`, node).as("TypeError")
}
private requireIteratorMethod(value: unknown, context: string, node: AstNode): unknown {
if (typeofValue(value) === "function") return value
throw new InterpreterRuntimeError(`${context} must be a function.`, node).as("TypeError")
}
private enumerableKeys(value: unknown): Array<string> | undefined {
if (value instanceof ToolReference) {
return [...this.toolKeys(value.path)]
@ -922,7 +1050,7 @@ export class Interpreter<R> {
)
}
const consumed = new Set<string>()
const consumed = new Set<PropertyKey>()
for (const propertyValue of getArray(pattern, "properties")) {
const property = asNode(propertyValue, "properties")
@ -931,6 +1059,10 @@ export class Interpreter<R> {
for (const [key, item] of Object.entries(value as SafeObject)) {
if (!consumed.has(key) && !isBlockedMember(key)) rest[key] = item
}
for (const symbol of IteratorSymbols) {
if (!consumed.has(symbol) && Object.hasOwn(value, symbol))
Reflect.set(rest, symbol, Reflect.get(value, symbol))
}
yield* self.declarePattern(getNode(property, "argument"), rest, mutable, property, initialize)
continue
}
@ -939,7 +1071,7 @@ export class Interpreter<R> {
if (isBlockedMember(String(key))) {
throw new InterpreterRuntimeError(`Property '${String(key)}' is not available.`, property)
}
consumed.add(String(key))
consumed.add(typeof key === "symbol" ? key : String(key))
yield* self.declarePattern(
getNode(property, "value"),
self.destructuringPropertyValue(value as SafeObject | Array<unknown>, key),
@ -1002,7 +1134,7 @@ export class Interpreter<R> {
}
const source = value as SafeObject | Array<unknown>
const consumed = new Set<string>()
const consumed = new Set<PropertyKey>()
for (const propertyValue of getArray(pattern, "properties")) {
const property = asNode(propertyValue, "properties")
if (property.type === "RestElement") {
@ -1010,6 +1142,10 @@ export class Interpreter<R> {
for (const [key, item] of Object.entries(source)) {
if (!consumed.has(key) && !isBlockedMember(key)) rest[key] = item
}
for (const symbol of IteratorSymbols) {
if (!consumed.has(symbol) && Object.hasOwn(source, symbol))
Reflect.set(rest, symbol, Reflect.get(source, symbol))
}
yield* self.assignPattern(getNode(property, "argument"), rest, property)
continue
}
@ -1017,7 +1153,7 @@ export class Interpreter<R> {
if (isBlockedMember(String(key))) {
throw new InterpreterRuntimeError(`Property '${String(key)}' is not available.`, property)
}
consumed.add(String(key))
consumed.add(typeof key === "symbol" ? key : String(key))
yield* self.assignPattern(getNode(property, "value"), self.destructuringPropertyValue(source, key), property)
}
return
@ -1044,7 +1180,7 @@ export class Interpreter<R> {
})
}
private destructuringPropertyKey(property: AstNode): Effect.Effect<string | number, unknown, R> {
private destructuringPropertyKey(property: AstNode): Effect.Effect<PropertyKey, unknown, R> {
if (property.type !== "Property" || getString(property, "kind") !== "init") {
throw new InterpreterRuntimeError("Unsupported object destructuring property.", property)
}
@ -1055,12 +1191,12 @@ export class Interpreter<R> {
return Effect.succeed(keyNode.type === "Identifier" ? getString(keyNode, "name") : String(keyNode.value))
}
private destructuringPropertyValue(source: SafeObject | Array<unknown>, key: string | number): unknown {
if (!Array.isArray(source)) return source[String(key)]
private destructuringPropertyValue(source: SafeObject | Array<unknown>, key: PropertyKey): unknown {
if (!Array.isArray(source)) return Reflect.get(source, key)
if (key === "length") return source.length
if (typeof key === "number") return source[key]
if (Object.hasOwn(source, key)) return (source as Record<string, unknown> & Array<unknown>)[key]
if (arrayMethods.has(key)) return new IntrinsicReference(source, key)
if (Object.hasOwn(source, key)) return Reflect.get(source, key)
if (typeof key === "string" && arrayMethods.has(key)) return new IntrinsicReference(source, key)
return undefined
}
@ -1690,6 +1826,12 @@ export class Interpreter<R> {
if (callable instanceof PromiseNamespace) {
throw new InterpreterRuntimeError("Constructor Promise requires 'new'.", node).as("TypeError")
}
if (callable instanceof SymbolNamespace) {
throw new InterpreterRuntimeError(
"Symbol is not callable; only Symbol.asyncIterator and Symbol.iterator are available.",
node,
).as("TypeError")
}
if (callable instanceof PromiseCapabilityFunction) {
callable.settle(args[0])
return undefined
@ -1800,6 +1942,9 @@ export class Interpreter<R> {
if (isBlockedMember(key)) throw new InterpreterRuntimeError(`Property '${key}' is not available.`, property)
objectValue[key] = value
}
for (const symbol of IteratorSymbols) {
if (Object.hasOwn(spread, symbol)) Reflect.set(objectValue, symbol, Reflect.get(spread, symbol))
}
continue
}
@ -1830,7 +1975,7 @@ export class Interpreter<R> {
if (isBlockedMember(String(key))) {
throw new InterpreterRuntimeError(`Property '${String(key)}' is not available.`, keyNode)
}
objectValue[String(key)] = yield* self.evaluateExpression(valueNode)
Reflect.set(objectValue, key, yield* self.evaluateExpression(valueNode))
}
return objectValue
@ -1955,6 +2100,12 @@ export class Interpreter<R> {
)
}
if (objectValue instanceof SymbolNamespace) {
if (key === "asyncIterator") return new ComputedValue(AsyncIteratorSymbol)
if (key === "iterator") return new ComputedValue(IteratorSymbol)
return new ComputedValue(undefined)
}
if (objectValue instanceof GlobalNamespace) {
if (typeof key === "string" && isBlockedMember(key)) {
throw new InterpreterRuntimeError(`${objectValue.name}.${key} is not available.`, propertyNode)
@ -1976,7 +2127,7 @@ export class Interpreter<R> {
if (typeof objectValue === "string") {
if (key === "length") return new ComputedValue(objectValue.length)
const index = parseArrayIndex(key)
const index = typeof key === "symbol" ? undefined : parseArrayIndex(key)
if (index !== undefined) return new ComputedValue(objectValue[index])
if (typeof key === "string" && stringMethods.has(key)) return new IntrinsicReference(objectValue, key)
return new ComputedValue(undefined)
@ -2072,7 +2223,7 @@ export class Interpreter<R> {
if (Array.isArray(objectValue)) {
if (operation === "delete") return { target: objectValue, key }
const index = parseArrayIndex(key)
const index = typeof key === "symbol" ? undefined : parseArrayIndex(key)
if (key !== "length" && !(typeof key === "string" && arrayMethods.has(key)) && index === undefined) {
if (typeof key === "string" && Object.hasOwn(objectValue, key)) {
return new ComputedValue((objectValue as Record<string, unknown> & Array<unknown>)[key])
@ -2103,13 +2254,13 @@ export class Interpreter<R> {
if (Array.isArray(reference.target)) {
if (reference.key === "length") return reference.target.length
if (typeof reference.key === "string") return new IntrinsicReference(reference.target, reference.key)
return reference.target[reference.key]
return Reflect.get(reference.target, reference.key)
}
if (reference.target instanceof CodeModeRegExp) return reference.target.lastIndex
if (reference.target instanceof CodeModeURL) {
return (reference.target.url as unknown as Record<string, unknown>)[String(reference.key)]
return Reflect.get(reference.target.url, reference.key)
}
return reference.target[String(reference.key)]
return Reflect.get(reference.target, reference.key)
})
}
@ -2171,22 +2322,22 @@ export class Interpreter<R> {
throw new InterpreterRuntimeError("Array methods cannot be assigned.", node)
}
}
const key = Array.isArray(reference.target) ? reference.key : String(reference.key)
const key = reference.key
const { write, next, result } = yield* compute(self.readReferenceValue(reference, key))
if (write) self.assignToReference(reference, key, next, node)
return result
})
}
private readReferenceValue(reference: MemberReference, key: number | string): unknown {
private readReferenceValue(reference: MemberReference, key: PropertyKey): unknown {
if (reference.target instanceof CodeModeURL) {
return (reference.target.url as unknown as Record<string, unknown>)[key]
return Reflect.get(reference.target.url, key)
}
if (reference.target instanceof CodeModeRegExp) return reference.target.lastIndex
return (reference.target as Record<PropertyKey, unknown>)[key]
return Reflect.get(reference.target, key)
}
private assignToReference(reference: MemberReference, key: number | string, next: unknown, node: AstNode): void {
private assignToReference(reference: MemberReference, key: PropertyKey, next: unknown, node: AstNode): void {
if (Array.isArray(reference.target)) {
const target = reference.target
if (typeof key !== "number" || parseArrayIndex(key) === undefined) {
@ -2219,16 +2370,19 @@ export class Interpreter<R> {
return
}
const target = reference.target as SafeObject
const objectKey = key as string
rejectCircularInsertion(target, next, "Object assignment result", node)
target[objectKey] = next
Reflect.set(target, key, next)
}
private toPropertyKey(value: unknown, node: AstNode): string | number {
private toPropertyKey(value: unknown, node: AstNode): PropertyKey {
if (typeof value === "string" || typeof value === "number") {
return value
}
if (value === AsyncIteratorSymbol || value === IteratorSymbol) return value
throw new InterpreterRuntimeError("Property key must be a string or number.", node)
throw new InterpreterRuntimeError(
"Property key must be a string or number, or Symbol.asyncIterator/Symbol.iterator.",
node,
)
}
}

View file

@ -1,4 +1,10 @@
import { type AstNode, InterpreterRuntimeError } from "../interpreter/model.js"
import {
type AstNode,
AsyncIteratorSymbol,
InterpreterRuntimeError,
IteratorSymbol,
IteratorSymbols,
} from "../interpreter/model.js"
import { containsOpaqueReference } from "../interpreter/references.js"
import { isBlockedMember } from "../tool-runtime.js"
import { isCodeModeValue, CodeModeMap, CodeModePromise, CodeModeSet, CodeModeURLSearchParams } from "../values.js"
@ -46,7 +52,10 @@ export const invokeObjectMethod = (name: string, args: Array<unknown>, node: Ast
case "entries":
return Object.entries(requireObject()).map(([key, item]) => [key, item])
case "hasOwn":
return Object.hasOwn(requireObject(), String(args[1]))
return Object.hasOwn(
requireObject(),
args[1] === AsyncIteratorSymbol || args[1] === IteratorSymbol ? args[1] : String(args[1]),
)
case "is":
if (containsOpaqueReference(args[0]) || containsOpaqueReference(args[1])) {
throw new InterpreterRuntimeError("Object.is requires data values.", node, "InvalidDataValue")
@ -64,6 +73,9 @@ export const invokeObjectMethod = (name: string, args: Array<unknown>, node: Ast
throw new InterpreterRuntimeError("Object.assign expects data objects.", node)
}
for (const [key, item] of Object.entries(source)) guardedSet(out, key, item)
for (const symbol of IteratorSymbols) {
if (Object.hasOwn(source, symbol)) Reflect.set(out, symbol, Reflect.get(source, symbol))
}
}
return out
}