update tests to properly carry cell brands

mathpirate · mathpirate · commit 5879f1d3cbe7 · 2025-11-04T18:34:51.000-08:00
diff --git a/packages/schema-generator/src/formatters/common-tools-formatter.ts b/packages/schema-generator/src/formatters/common-tools-formatter.ts
@@ -388,16 +388,16 @@ export class CommonToolsFormatter implements TypeFormatter {
       | undefined,
     targetKind: WrapperKind,
   ): ts.TypeReferenceNode | undefined {
-    if (resolvedWrapper?.kind === targetKind) {
-      return resolvedWrapper.node;
-    }
     if (
       originalNode &&
       ts.isTypeReferenceNode(originalNode) &&
       originalNode.typeArguments
     ) {
       return originalNode;
     }
+    if (resolvedWrapper?.kind === targetKind) {
+      return resolvedWrapper.node;
+    }
     return undefined;
   }
 
diff --git a/packages/schema-generator/test/fixtures/schema/alias-edge-cases.input.ts b/packages/schema-generator/test/fixtures/schema/alias-edge-cases.input.ts
@@ -1,9 +1,9 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
 
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 
diff --git a/packages/schema-generator/test/fixtures/schema/aliases-of-aliases-cell-array.input.ts b/packages/schema-generator/test/fixtures/schema/aliases-of-aliases-cell-array.input.ts
@@ -1,4 +1,4 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/fixtures/schema/array-of-cell-string.input.ts b/packages/schema-generator/test/fixtures/schema/array-of-cell-string.input.ts
@@ -1,4 +1,4 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/fixtures/schema/cell-array-types.input.ts b/packages/schema-generator/test/fixtures/schema/cell-array-types.input.ts
@@ -1,4 +1,4 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/fixtures/schema/cell-brand-variants.expected.json b/packages/schema-generator/test/fixtures/schema/cell-brand-variants.expected.json
@@ -1,32 +1,49 @@
 {
-  "properties": {
-    "comparableValue": {
+  "$defs": {
+    "MyComparableCell": {
       "asCell": true,
       "type": "number"
     },
-    "readonlyValue": {
+    "MyOpaque": {
+      "asOpaque": true,
+      "type": "boolean"
+    },
+    "MyReadonlyCell": {
       "asCell": true,
       "type": "string"
     },
-    "writeonlyValue": {
+    "MyStream": {
+      "asStream": true,
+      "type": "number"
+    },
+    "MyWriteonlyCell": {
       "asCell": true,
       "type": "boolean"
+    }
+  },
+  "properties": {
+    "comparableValue": {
+      "$ref": "#/$defs/MyComparableCell"
+    },
+    "opaqueValue": {
+      "$ref": "#/$defs/MyOpaque"
+    },
+    "readonlyValue": {
+      "$ref": "#/$defs/MyReadonlyCell"
     },
     "streamValue": {
-      "asStream": true,
-      "type": "number"
+      "$ref": "#/$defs/MyStream"
     },
-    "opaqueValue": {
-      "asOpaque": true,
-      "type": "boolean"
+    "writeonlyValue": {
+      "$ref": "#/$defs/MyWriteonlyCell"
     }
   },
   "required": [
     "comparableValue",
+    "opaqueValue",
     "readonlyValue",
-    "writeonlyValue",
     "streamValue",
-    "opaqueValue"
+    "writeonlyValue"
   ],
   "type": "object"
 }
diff --git a/packages/schema-generator/test/fixtures/schema/cell-type.input.ts b/packages/schema-generator/test/fixtures/schema/cell-type.input.ts
@@ -1,4 +1,4 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/fixtures/schema/default-type.input.ts b/packages/schema-generator/test/fixtures/schema/default-type.input.ts
@@ -41,7 +41,7 @@ interface ComplexDefault {
 }
 
 // Test Default with Cell types
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/fixtures/schema/descriptions-wrappers.input.ts b/packages/schema-generator/test/fixtures/schema/descriptions-wrappers.input.ts
@@ -1,9 +1,9 @@
 type Default<T, V extends T = T> = T;
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 
diff --git a/packages/schema-generator/test/fixtures/schema/opaque-opaqueref-stream.input.ts b/packages/schema-generator/test/fixtures/schema/opaque-opaqueref-stream.input.ts
@@ -2,7 +2,7 @@
 // This mimics the structure of BuiltInLLMState where cancelGeneration: Stream<void>
 // becomes Opaque<OpaqueRef<Stream<void>>> when returned inside OpaqueRef<BuiltInLLMState>
 
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   send(event: T): void;
 }
 
@@ -12,7 +12,10 @@ interface OpaqueRefMethods<T> {
 }
 
 // OpaqueRef<T> is an intersection of OpaqueRefMethods<T> and T
-type OpaqueRef<T> = OpaqueRefMethods<T> & (
+type OpaqueRef<T> =
+  & BrandedCell<T, "opaque">
+  & OpaqueRefMethods<T>
+  & (
   T extends Array<infer U> ? Array<OpaqueRef<U>>
   : T extends object ? { [K in keyof T]: OpaqueRef<T[K]> }
   : T
diff --git a/packages/schema-generator/test/fixtures/schema/stream-of-cell-number.input.ts b/packages/schema-generator/test/fixtures/schema/stream-of-cell-number.input.ts
@@ -1,8 +1,8 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 interface SchemaRoot {
diff --git a/packages/schema-generator/test/fixtures/schema/stream-type.input.ts b/packages/schema-generator/test/fixtures/schema/stream-type.input.ts
@@ -1,4 +1,4 @@
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 interface SchemaRoot {
diff --git a/packages/schema-generator/test/fixtures/schema/type-aliases.input.ts b/packages/schema-generator/test/fixtures/schema/type-aliases.input.ts
@@ -1,8 +1,8 @@
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 type Default<T, V extends T = T> = T;
diff --git a/packages/schema-generator/test/fixtures/schema/wrapper-aliases.input.ts b/packages/schema-generator/test/fixtures/schema/wrapper-aliases.input.ts
@@ -1,11 +1,11 @@
 // Test wrapper type aliases - both direct and aliased wrapper types
 // Using proper interface definitions that match actual CommonTools wrappers
 type Default<T, V = T> = T;
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
-interface Stream<T> {
+interface Stream<T> extends BrandedCell<T, "stream"> {
   subscribe(cb: (v: T) => void): void;
 }
 
diff --git a/packages/schema-generator/test/fixtures/schema/wrappers-nested.input.ts b/packages/schema-generator/test/fixtures/schema/wrappers-nested.input.ts
@@ -1,5 +1,5 @@
 type Default<T, V extends T = T> = T;
-interface Cell<T> {
+interface Cell<T> extends BrandedCell<T, "cell"> {
   get(): T;
   set(v: T): void;
 }
diff --git a/packages/schema-generator/test/schema/cell-type.test.ts b/packages/schema-generator/test/schema/cell-type.test.ts
@@ -6,7 +6,7 @@ import { asObjectSchema, getTypeFromCode } from "../utils.ts";
 describe("Schema: Cell types", () => {
   it("handles Cell<string>", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(value: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(value: T): void; }
       interface X { name: Cell<string>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -21,7 +21,7 @@ describe("Schema: Cell types", () => {
 
   it("handles Cell<Array<{id:string}>>", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(value: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(value: T): void; }
       interface X { users: Cell<Array<{ id: string }>>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -39,8 +39,8 @@ describe("Schema: Cell types", () => {
 
   it("handles Stream<Cell<number>>", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(value: T): void; }
-      interface Stream<T> { subscribe(cb: (v: T) => void): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(value: T): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v: T) => void): void; }
       interface X { value: Stream<Cell<number>>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -56,8 +56,8 @@ describe("Schema: Cell types", () => {
 
   it("disallows Cell<Stream<T>> and suggests boxing", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(value: T): void; }
-      interface Stream<T> { subscribe(cb: (v: T) => void): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(value: T): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v: T) => void): void; }
       interface X { invalid: Cell<Stream<number>>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
diff --git a/packages/schema-generator/test/schema/nested-wrappers.test.ts b/packages/schema-generator/test/schema/nested-wrappers.test.ts
@@ -23,7 +23,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
   it("Cell<Default<string,'default'>>", async () => {
     const code = `
       interface Default<T, V> {}
-      interface Cell<T> { get(): T; set(v: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
       interface X { value: Cell<Default<string, "default">>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -39,7 +39,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
   it("Stream<Default<string,'initial'>>", async () => {
     const code = `
       interface Default<T, V> {}
-      interface Stream<T> { subscribe(cb: (v:T) => void): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v:T) => void): void; }
       interface X { events: Stream<Default<string, "initial">>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -55,7 +55,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
   it("Stream<Default<string[], ['a']>> yields array schema with default", async () => {
     const code = `
       interface Default<T, V> {}
-      interface Stream<T> { subscribe(cb: (v:T) => void): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v:T) => void): void; }
       interface X { events: Stream<Default<string[], ["a"]>>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -72,7 +72,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
 
   it("array of Cell<string>", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(v: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
       interface X { items: Array<Cell<string>>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -91,7 +91,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
 
   it("Cell<string[]>", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(v: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
       interface X { tags: Cell<string[]>; }
     `;
     const { type, checker } = await getTypeFromCode(code, "X");
@@ -108,7 +108,7 @@ describe("Schema: Nested wrappers (Cell, Stream, Default)", () => {
   it("complex nesting: Cell<Default<string,'d'>> and Default<string[], ['a','b']>", async () => {
     const code = `
       interface Default<T, V> {}
-      interface Cell<T> { get(): T; set(v: T): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
       interface X {
         cellOfDefault: Cell<Default<string, "d">>;
         defaultArray: Default<string[], ["a", "b"]>;
diff --git a/packages/schema-generator/test/schema/type-aliases-and-shared.test.ts b/packages/schema-generator/test/schema/type-aliases-and-shared.test.ts
@@ -6,8 +6,8 @@ import { asObjectSchema, getTypeFromCode } from "../utils.ts";
 describe("Schema: Type aliases and shared types", () => {
   it("handles basic Cell/Stream/Default aliases", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(v: T): void; }
-      interface Stream<T> { subscribe(cb: (v:T) => void): void; }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v:T) => void): void; }
       interface Default<T, V> {}
       type MyCell<T> = Cell<T>;
       type StringCell = Cell<string>;
diff --git a/packages/schema-generator/test/schema/type-to-schema.test.ts b/packages/schema-generator/test/schema/type-to-schema.test.ts
@@ -6,8 +6,8 @@ import { asObjectSchema, getTypeFromCode } from "../utils.ts";
 describe("Schema: type-to-schema parity", () => {
   it("generates schemas for inputs/outputs structures", async () => {
     const code = `
-      interface Cell<T> { get(): T; set(v: T): void; }
-      interface Stream<T> { subscribe(cb: (v:T)=>void): void }
+      interface Cell<T> extends BrandedCell<T, "cell"> { get(): T; set(v: T): void; }
+      interface Stream<T> extends BrandedCell<T, "stream"> { subscribe(cb: (v:T)=>void): void }
 
       interface UpdaterInput { newValues: string[]; }
       interface RecipeInput { values: Cell<string[]>; }

Original file line number	Diff line number	Diff line change
`@@ -1,9 +1,9 @@`
`1`		`-interface Cell<T> {`
	`1`	`+interface Cell<T> extends BrandedCell<T, "cell"> {`
`2`	`2`	`get(): T;`
`3`	`3`	`set(v: T): void;`
`4`	`4`	`}`
`5`	`5`
`6`		`-interface Stream<T> {`
	`6`	`+interface Stream<T> extends BrandedCell<T, "stream"> {`
`7`	`7`	`subscribe(cb: (v: T) => void): void;`
`8`	`8`	`}`
`9`	`9`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-interface Cell<T> {`
	`1`	`+interface Cell<T> extends BrandedCell<T, "cell"> {`
`2`	`2`	`get(): T;`
`3`	`3`	`set(v: T): void;`
`4`	`4`	`}`
Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ interface ComplexDefault {`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`// Test Default with Cell types`
`44`		`-interface Cell<T> {`
	`44`	`+interface Cell<T> extends BrandedCell<T, "cell"> {`
`45`	`45`	`get(): T;`
`46`	`46`	`set(v: T): void;`
`47`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,9 +1,9 @@`
`1`	`1`	`type Default<T, V extends T = T> = T;`
`2`		`-interface Cell<T> {`
	`2`	`+interface Cell<T> extends BrandedCell<T, "cell"> {`
`3`	`3`	`get(): T;`
`4`	`4`	`set(v: T): void;`
`5`	`5`	`}`
`6`		`-interface Stream<T> {`
	`6`	`+interface Stream<T> extends BrandedCell<T, "stream"> {`
`7`	`7`	`subscribe(cb: (v: T) => void): void;`
`8`	`8`	`}`
`9`	`9`