diff --git a/src/core/callable-compiler.ts b/src/core/callable-compiler.ts index 9a0d6726..1bd1b630 100644 --- a/src/core/callable-compiler.ts +++ b/src/core/callable-compiler.ts @@ -286,6 +286,211 @@ function resolveTuple( }; } +type WireRelation = 'disjoint' | 'same-decoding' | 'ambiguous'; + +function combineWireRelations(relations: readonly WireRelation[]): WireRelation { + if (relations.includes('ambiguous')) { + return 'ambiguous'; + } + return relations.includes('same-decoding') ? 'same-decoding' : 'disjoint'; +} + +type WireCategory = 'null' | 'boolean' | 'number' | 'string' | 'array' | 'object' | 'unknown'; + +function acceptsWireCategory(value: ValueContract, category: WireCategory): boolean { + if (category === 'unknown' || value.kind === 'unsupported') { + return true; + } + if (value.kind === 'union') { + return value.options.some(option => acceptsWireCategory(option, category)); + } + const expected: WireCategory = (() => { + switch (value.kind) { + case 'null': + return 'null'; + case 'boolean': + return 'boolean'; + case 'integer': + case 'float': + return 'number'; + case 'string': + return 'string'; + case 'sequence': + case 'tuple': + return 'array'; + case 'bytes': + case 'record': + case 'ndarray-float16': + case 'torch-float16': + return 'object'; + } + })(); + return expected === category; +} + +interface TaggedEnvelopeShape { + readonly required: Readonly>; + readonly optional?: Readonly>; +} + +function taggedEnvelopeShapes(value: ValueContract): readonly TaggedEnvelopeShape[] { + switch (value.kind) { + case 'bytes': + return [ + { required: { __type__: 'string', encoding: 'string', data: 'string' } }, + { required: { __tywrap_bytes__: 'boolean', b64: 'string' } }, + ]; + case 'ndarray-float16': + // Legacy envelopes may omit the version, shape, and dtype. + return [ + { + required: { __tywrap__: 'string', encoding: 'string', b64: 'string' }, + optional: { codecVersion: 'number' }, + }, + { + required: { __tywrap__: 'string', encoding: 'string', data: 'unknown' }, + optional: { codecVersion: 'number' }, + }, + ]; + case 'torch-float16': + return [ + { + required: { __tywrap__: 'string', encoding: 'string', value: 'object' }, + optional: { + codecVersion: 'number', + device: 'string', + sourceDtype: 'string', + sourceDevice: 'string', + }, + }, + ]; + default: + return []; + } +} + +function recordCanMatchEnvelope( + record: Extract, + envelope: TaggedEnvelopeShape +): boolean { + const fields = new Map(record.fields.map(field => [field.name, field] as const)); + // Current decoders accept extra keys on marker envelopes. + for (const [name, category] of Object.entries(envelope.required)) { + const field = fields.get(name); + if (field && !acceptsWireCategory(field.value, category)) { + return false; + } + if (record.additionalValues && !acceptsWireCategory(record.additionalValues, category)) { + return false; + } + } + // An optional marker field proves disjointness only when the record requires it. + for (const [name, category] of Object.entries(envelope.optional ?? {})) { + const field = fields.get(name); + if (field?.required && !acceptsWireCategory(field.value, category)) { + return false; + } + if ( + field?.required && + record.additionalValues && + !acceptsWireCategory(record.additionalValues, category) + ) { + return false; + } + } + return true; +} + +function taggedRecordRelation( + record: Extract, + tagged: ValueContract +): WireRelation { + return taggedEnvelopeShapes(tagged).some(envelope => recordCanMatchEnvelope(record, envelope)) + ? 'ambiguous' + : 'disjoint'; +} + +function recordWireRelation(left: ValueContract, right: ValueContract): WireRelation { + if (left.kind !== 'record' || right.kind !== 'record') { + return 'disjoint'; + } + const leftFields = new Map(left.fields.map(field => [field.name, field] as const)); + const rightFields = new Map(right.fields.map(field => [field.name, field] as const)); + const relations: WireRelation[] = []; + for (const field of left.fields) { + const otherField = rightFields.get(field.name); + const other = otherField?.value ?? right.additionalValues; + if (other) { + const relation = valueWireRelation(field.value, other); + if (relation === 'disjoint' && (field.required || otherField?.required)) { + return 'disjoint'; + } + relations.push(relation); + } + } + for (const field of right.fields) { + if (leftFields.has(field.name)) { + continue; + } + if (left.additionalValues) { + const relation = valueWireRelation(left.additionalValues, field.value); + if (relation === 'disjoint' && field.required) { + return 'disjoint'; + } + relations.push(relation); + } + } + if (left.additionalValues && right.additionalValues) { + relations.push(valueWireRelation(left.additionalValues, right.additionalValues)); + } + return relations.includes('ambiguous') ? 'ambiguous' : 'same-decoding'; +} + +function valueWireRelation(left: ValueContract, right: ValueContract): WireRelation { + if (left.kind === 'unsupported' || right.kind === 'unsupported') { + return 'ambiguous'; + } + if (left.kind === 'union') { + return combineWireRelations(left.options.map(option => valueWireRelation(option, right))); + } + if (right.kind === 'union') { + return combineWireRelations(right.options.map(option => valueWireRelation(left, option))); + } + if (left.kind === 'record' && right.kind === 'record') { + return recordWireRelation(left, right); + } + if (left.kind === 'record' && taggedEnvelopeShapes(right).length > 0) { + return taggedRecordRelation(left, right); + } + if (right.kind === 'record' && taggedEnvelopeShapes(left).length > 0) { + return taggedRecordRelation(right, left); + } + if (left.kind === 'sequence' && right.kind === 'sequence') { + return valueWireRelation(left.item, right.item) === 'ambiguous' ? 'ambiguous' : 'same-decoding'; + } + if (left.kind === 'tuple' && right.kind === 'tuple') { + if (left.items.length !== right.items.length) { + return 'disjoint'; + } + const relations = left.items.map((item, index) => valueWireRelation(item, right.items[index]!)); + return relations.includes('disjoint') ? 'disjoint' : combineWireRelations(relations); + } + if (left.kind === 'sequence' && right.kind === 'tuple') { + const relations = right.items.map(item => valueWireRelation(left.item, item)); + return relations.includes('disjoint') ? 'disjoint' : combineWireRelations(relations); + } + if (left.kind === 'tuple' && right.kind === 'sequence') { + return valueWireRelation(right, left); + } + if ( + (left.kind === 'integer' && right.kind === 'float') || + (left.kind === 'float' && right.kind === 'integer') + ) { + return 'same-decoding'; + } + return left.kind === right.kind ? 'same-decoding' : 'disjoint'; +} + function resolveUnion( options: readonly PythonType[], request: ValueConversionRequest, @@ -320,6 +525,17 @@ function resolveUnion( guidance: 'Provide two explicit value alternatives.', }; } + for (let index = 0; index < values.length; index += 1) { + for (let earlier = 0; earlier < index; earlier += 1) { + if (valueWireRelation(values[earlier]!.value, values[index]!.value) === 'ambiguous') { + return { + status: 'unsupported', + reason: 'Union alternatives can share a wire value but decode differently.', + guidance: 'Use a disjoint tagged record or split the union.', + }; + } + } + } return { status: 'supported', value: { diff --git a/test/callable-compiler.test.ts b/test/callable-compiler.test.ts index e154674f..0ea44529 100644 --- a/test/callable-compiler.test.ts +++ b/test/callable-compiler.test.ts @@ -14,6 +14,7 @@ import { CodeGenerator } from '../src/core/generator.js'; import { parseAnnotationToPythonType } from '../src/core/annotation-parser.js'; import { validateIrContract } from '../src/core/ir-contract.js'; import { BridgeValidationError } from '../src/runtime/errors.js'; +import { BridgeCodec } from '../src/runtime/bridge-codec.js'; import { HttpBridge } from '../src/runtime/http.js'; import { clearRuntimeBridge, setRuntimeBridge } from 'tywrap/runtime'; import type { PythonModule, PythonType } from '../src/types/index.js'; @@ -1701,6 +1702,337 @@ describe('compileContract', () => { }); }); + it('rejects unions whose marker wire can decode as a different value', async () => { + const annotations = [ + 'bytes | dict[str, str]', + 'list[bytes] | list[dict[str, str]]', + 'tuple[bytes] | tuple[dict[str, str]]', + 'dict[str, bytes] | dict[str, dict[str, str]]', + ]; + for (const annotation of annotations) { + expect( + DEFAULT_VALUE_CONVERSION.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType(annotation), + }) + ).toMatchObject({ + status: 'unsupported', + reason: 'Union alternatives can share a wire value but decode differently.', + }); + } + for (const annotation of ['int | float', 'bytes | str', 'bytes | dict[str, int]']) { + expect( + DEFAULT_VALUE_CONVERSION.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType(annotation), + }).status + ).toBe('supported'); + } + const tagged: typeof DEFAULT_VALUE_CONVERSION = { + revision: DEFAULT_VALUE_CONVERSION.revision, + resolve(request) { + if (request.logicalType.kind === 'custom' && request.logicalType.name === 'IntTag') { + return { + status: 'supported', + value: { + kind: 'record', + wire: 'json', + decodedAs: 'object', + fields: [ + { + name: 'tag', + required: true, + value: { + kind: 'integer', + wire: 'json', + decodedAs: 'number', + constraint: 'safe-integer', + }, + }, + { + name: 'payload', + required: true, + value: { kind: 'bytes', wire: 'base64-envelope', decodedAs: 'Uint8Array' }, + }, + ], + }, + }; + } + if (request.logicalType.kind === 'custom' && request.logicalType.name === 'StringTag') { + return { + status: 'supported', + value: { + kind: 'record', + wire: 'json', + decodedAs: 'object', + fields: [ + { + name: 'tag', + required: true, + value: { kind: 'string', wire: 'json', decodedAs: 'string' }, + }, + { + name: 'payload', + required: true, + value: { + kind: 'record', + wire: 'json', + decodedAs: 'object', + fields: [], + additionalValues: { kind: 'string', wire: 'json', decodedAs: 'string' }, + }, + }, + ], + }, + }; + } + if ( + request.logicalType.kind === 'custom' && + (request.logicalType.name === 'ExtraRecord' || + request.logicalType.name === 'ExtraRecordValues') + ) { + return { + status: 'supported', + value: { + kind: 'record', + wire: 'json', + decodedAs: 'object', + fields: [ + { + name: '__type__', + required: true, + value: { kind: 'string', wire: 'json', decodedAs: 'string' }, + }, + { + name: 'encoding', + required: true, + value: { kind: 'string', wire: 'json', decodedAs: 'string' }, + }, + { + name: 'data', + required: true, + value: { kind: 'string', wire: 'json', decodedAs: 'string' }, + }, + { + name: 'extra', + required: true, + value: { + kind: 'integer', + wire: 'json', + decodedAs: 'number', + constraint: 'safe-integer', + }, + }, + ], + additionalValues: + request.logicalType.name === 'ExtraRecordValues' + ? { + kind: 'union', + wire: 'selected-option', + decodedAs: 'selected-option', + options: [ + { kind: 'string', wire: 'json', decodedAs: 'string' }, + { + kind: 'integer', + wire: 'json', + decodedAs: 'number', + constraint: 'safe-integer', + }, + ], + } + : undefined, + }, + }; + } + if ( + request.logicalType.kind === 'custom' && + (request.logicalType.name === 'NdarrayRecord' || + request.logicalType.name === 'NdarrayVersionRecord' || + request.logicalType.name === 'TorchRecord') + ) { + const stringValue = { kind: 'string', wire: 'json', decodedAs: 'string' } as const; + const integerValue = { + kind: 'integer', + wire: 'json', + decodedAs: 'number', + constraint: 'safe-integer', + } as const; + const scientific = request.logicalType.name !== 'TorchRecord'; + return { + status: 'supported', + value: { + kind: 'record', + wire: 'json', + decodedAs: 'object', + fields: [ + { name: '__tywrap__', required: true, value: stringValue }, + { name: 'encoding', required: true, value: stringValue }, + { + name: scientific ? 'data' : 'value', + required: true, + value: scientific + ? { kind: 'sequence', wire: 'json', decodedAs: 'array', item: integerValue } + : { kind: 'record', wire: 'json', decodedAs: 'object', fields: [] }, + }, + { + name: 'codecVersion', + required: request.logicalType.name === 'NdarrayVersionRecord', + value: stringValue, + }, + ...(scientific + ? [ + { + name: 'shape', + required: true, + value: { + kind: 'sequence' as const, + wire: 'json' as const, + decodedAs: 'array' as const, + item: integerValue, + }, + }, + { name: 'dtype', required: true, value: stringValue }, + ] + : [{ name: 'device', required: false, value: integerValue }]), + ], + }, + }; + } + return DEFAULT_VALUE_CONVERSION.resolve({ ...request, resolveNested: tagged.resolve }); + }, + }; + expect( + tagged.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType('IntTag | StringTag'), + }).status + ).toBe('supported'); + for (const annotation of ['bytes | ExtraRecord', 'bytes | ExtraRecordValues']) { + expect( + tagged.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType(annotation), + }) + ).toMatchObject({ + status: 'unsupported', + reason: 'Union alternatives can share a wire value but decode differently.', + }); + } + for (const annotation of [ + 'numpy.NDArray[numpy.float16] | NdarrayRecord', + 'torch.Tensor[torch.float16] | TorchRecord', + ]) { + expect( + tagged.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType(annotation), + }) + ).toMatchObject({ + status: 'unsupported', + reason: 'Union alternatives can share a wire value but decode differently.', + }); + } + expect( + tagged.resolve({ + direction: 'output', + path: '$.returns', + logicalType: parseAnnotationToPythonType( + 'numpy.NDArray[numpy.float16] | NdarrayVersionRecord' + ), + }).status + ).toBe('supported'); + const decoded = await new BridgeCodec().decodeResponseAsync( + JSON.stringify({ + id: 1, + protocol: 'tywrap/1', + result: { __type__: 'bytes', encoding: 'base64', data: 'eA==', extra: 1 }, + }) + ); + expect(decoded).toEqual(Uint8Array.from([120])); + const ndarrayDecoded = await new BridgeCodec().decodeResponseAsync( + JSON.stringify({ + id: 2, + protocol: 'tywrap/1', + result: { + __tywrap__: 'ndarray', + encoding: 'json', + data: [1], + shape: [1], + dtype: 'float16', + }, + }) + ); + expect(ndarrayDecoded).toEqual([1]); + const torchDecoded = await new BridgeCodec().decodeResponseAsync( + JSON.stringify({ + id: 3, + protocol: 'tywrap/1', + result: { + __tywrap__: 'torch.tensor', + encoding: 'ndarray', + value: { + __tywrap__: 'ndarray', + encoding: 'json', + data: [1], + shape: [1], + dtype: 'float16', + }, + shape: [1], + dtype: 'torch.float16', + }, + }) + ); + expect(torchDecoded).toMatchObject({ data: [1], shape: [1], dtype: 'torch.float16' }); + + const source = rawIr.functions[1]!; + const ir = validateIrContract( + { + ...rawIr, + functions: [ + { + ...source, + name: 'ambiguous_wire', + qualname: 'fixture.ambiguous_wire', + returns: 'bytes | dict[str, str]', + }, + ], + }, + 'ambiguous union wire contract' + ); + expect(ir.ok).toBe(true); + if (!ir.ok) { + return; + } + const module = { + ...moduleModel, + functions: [{ ...moduleModel.functions[1]!, name: 'ambiguous_wire' }], + classes: [], + }; + const compiled = compileContract(ir.contract, { + module, + generator: new CodeGenerator(), + conversion: DEFAULT_VALUE_CONVERSION, + capabilities: DEFAULT_CALLABLE_CAPABILITIES, + }); + expect(compiled.callables[0]?.result.resolution.status).toBe('unsupported'); + expect(compiled.diagnostics).toContainEqual( + expect.objectContaining({ + severity: 'error', + code: 'conversion-unsupported', + path: '$.functions[0].returns', + }) + ); + expect(compiled.generated.declaration).toContain('ambiguousWire(): Promise'); + expect(compiled.generated.typescript).toContain( + 'const __validateambiguous_wireResult = createReturnValidator({"kind":"any"}' + ); + }); + it('requires returned dataclass fields even when their constructor has defaults', () => { const point = rawIr.classes[0]!; const ir = validateIrContract(