fix: address PR review findings

Author: elkaix

packages/kosong/src/providers/pythinker-schema.ts

@@ -157,24 +157,56 @@ function foldRootAnyOf(root: Record<string, unknown>): void {
   delete root['anyOf'];
 
   const rootProperties = root['properties'];
-  const merged: Record<string, unknown> = isRecord(rootProperties)
-    ? rootProperties
-    : {};
+  const alternativesByName = new Map<string, unknown[]>();
+  if (isRecord(rootProperties)) {
+    for (const [name, property] of Object.entries(rootProperties)) {
+      alternativesByName.set(name, [cloneJsonValue(property)]);
+    }
+  }
   for (const branch of branches) {
     const branchProperties = branch['properties'];
     if (!isRecord(branchProperties)) continue;
     for (const [name, property] of Object.entries(branchProperties)) {
-      if (!hasOwn(merged, name)) {
-        merged[name] = cloneJsonValue(property);
-      }
+      addRootPropertyAlternative(alternativesByName, name, property);
     }
   }
-  if (Object.keys(merged).length > 0) {
+
+  if (alternativesByName.size > 0) {
+    const merged: Record<string, unknown> = {};
+    for (const [name, alternatives] of alternativesByName) {
+      merged[name] = alternatives.length === 1 ? alternatives[0] : { anyOf: alternatives };
+    }
     root['properties'] = merged;
   }
   root['type'] = 'object';
 }
 
+/**
+ * Record one branch's schema for a merged root property.
+ *
+ * Root `anyOf` branches are alternatives, so two branches declaring the same
+ * property with different schemas (e.g. `value` as a string in one branch, an
+ * integer in another) must both stay representable — keeping only the first
+ * one seen would silently narrow what the tool actually accepts. Identical
+ * schemas collapse to one; differing schemas fold into an `anyOf` on the
+ * merged property.
+ */
+function addRootPropertyAlternative(
+  alternativesByName: Map<string, unknown[]>,
+  name: string,
+  property: unknown,
+): void {
+  const cloned = cloneJsonValue(property);
+  const alternatives = alternativesByName.get(name);
+  if (!alternatives) {
+    alternativesByName.set(name, [cloned]);
+    return;
+  }
+  if (!alternatives.some((existing) => deepEqualJson(existing, cloned))) {
+    alternatives.push(cloned);
+  }
+}
+
 function hasUnresolvedDefinitionRef(node: unknown, bucketKey: string): boolean {
   if (Array.isArray(node)) {
     return node.some((child) => hasUnresolvedDefinitionRef(child, bucketKey));
@@ -337,8 +369,13 @@ const ANYOF_PARENT_KEEP_KEYS = new Set([
  * so schemas that are perfectly valid elsewhere are rejected on this wire.
  *
  * Distributing is lossless: `P ∧ (B₁ ∨ B₂)` and `(P ∧ B₁) ∨ (P ∧ B₂)` accept
- * exactly the same instances. A branch that already declares a keyword keeps
- * its own, which is the narrower of the two.
+ * exactly the same instances — *if* a branch that already declares the same
+ * keyword is merged conjunctively with the parent's value rather than simply
+ * overriding it. `required` is the one keyword this function merges that way
+ * (parent and branch field lists are unioned, since both are actually
+ * required). Every other overlapping keyword still keeps the branch's own
+ * value: a full conjunctive merge for arbitrary keywords (`properties`,
+ * `items`, …) is out of scope for this compatibility normalizer.
  */
 function distributeAnyOfParentKeywords(node: Record<string, unknown>): void {
   const branches = node['anyOf'];
@@ -355,6 +392,8 @@ function distributeAnyOfParentKeywords(node: Record<string, unknown>): void {
     for (const key of inherited) {
       if (!hasOwn(branch, key)) {
         branch[key] = cloneJsonValue(node[key]);
+      } else if (key === 'required') {
+        branch[key] = mergeRequired(node[key], branch[key]);
       }
     }
   }
@@ -363,6 +402,26 @@ function distributeAnyOfParentKeywords(node: Record<string, unknown>): void {
   }
 }
 
+/**
+ * Union two `required` field lists.
+ *
+ * A parent's `required` and a branch's own `required` are both mandatory —
+ * dropping the parent's list when the branch already has one would silently
+ * accept objects missing a field the parent demanded.
+ */
+function mergeRequired(parentValue: unknown, branchValue: unknown): unknown {
+  if (!Array.isArray(parentValue) || !Array.isArray(branchValue)) {
+    return branchValue;
+  }
+  const merged = [...branchValue];
+  for (const name of parentValue) {
+    if (!merged.includes(name)) {
+      merged.push(name);
+    }
+  }
+  return merged;
+}
+
 function visitChildSchemas(node: Record<string, unknown>, visit: (schema: unknown) => void): void {
   for (const { key, kind } of CHILD_SCHEMA_SLOTS) {
     const value = node[key];
@@ -574,6 +633,21 @@ function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === 'object' && value !== null && !Array.isArray(value);
 }
 
+function deepEqualJson(a: unknown, b: unknown): boolean {
+  if (a === b) {
+    return true;
+  }
+  if (Array.isArray(a) && Array.isArray(b)) {
+    return a.length === b.length && a.every((item, index) => deepEqualJson(item, b[index]));
+  }
+  if (isRecord(a) && isRecord(b)) {
+    const aKeys = Object.keys(a);
+    const bKeys = Object.keys(b);
+    return aKeys.length === bKeys.length && aKeys.every((key) => hasOwn(b, key) && deepEqualJson(a[key], b[key]));
+  }
+  return false;
+}
+
 function hasOwn(obj: Record<string, unknown>, key: string): boolean {
   return Object.prototype.hasOwnProperty.call(obj, key);
 }

packages/kosong/test/providers/pythinker-schema.test.ts

@@ -612,6 +612,40 @@ describe('normalizePythinkerToolSchema', () => {
     });
   });
 
+  it('folds a root anyOf property declared differently across branches into an alternative', () => {
+    const result = normalizePythinkerToolSchema({
+      anyOf: [
+        { type: 'object', properties: { value: { type: 'string' } } },
+        { type: 'object', properties: { value: { type: 'integer' } } },
+      ],
+    });
+
+    expect(result).toEqual({
+      type: 'object',
+      properties: {
+        value: { anyOf: [{ type: 'string' }, { type: 'integer' }] },
+      },
+    });
+  });
+
+  it('collapses a root anyOf property that is identical across every branch', () => {
+    const result = normalizePythinkerToolSchema({
+      anyOf: [
+        { type: 'object', properties: { mode: { type: 'string' }, task_id: { type: 'string' } } },
+        { type: 'object', properties: { mode: { type: 'string' }, shell_id: { type: 'string' } } },
+      ],
+    });
+
+    expect(result).toEqual({
+      type: 'object',
+      properties: {
+        mode: { type: 'string' },
+        task_id: { type: 'string' },
+        shell_id: { type: 'string' },
+      },
+    });
+  });
+
   it('distributes into nested anyOf nodes and keeps narrower branch keywords', () => {
     const result = normalizePythinkerToolSchema({
       type: 'object',
@@ -637,6 +671,31 @@ describe('normalizePythinkerToolSchema', () => {
     });
   });
 
+  it('unions a parent required list with a branch that already declares its own', () => {
+    const result = normalizePythinkerToolSchema({
+      type: 'object',
+      properties: {
+        variant: {
+          type: 'object',
+          required: ['common'],
+          anyOf: [{ required: ['variant'] }, { required: ['common'] }],
+        },
+      },
+    });
+
+    expect(result).toEqual({
+      type: 'object',
+      properties: {
+        variant: {
+          anyOf: [
+            { type: 'object', required: ['variant', 'common'] },
+            { type: 'object', required: ['common'] },
+          ],
+        },
+      },
+    });
+  });
+
   it('leaves anyOf nodes alone when the parent only carries metadata', () => {
     const schema = {
       type: 'object',
@@ -651,23 +710,28 @@ describe('normalizePythinkerToolSchema', () => {
     expect(normalizePythinkerToolSchema(schema)).toEqual(schema);
   });
 
-  it('keeps cyclic $ref and definition buckets on the parent of an anyOf', () => {
+  it('keeps $ref and $defs on the anyOf parent instead of distributing them into branches', () => {
+    // $ref must be genuinely cyclic (self-referential) to survive derefJsonSchema
+    // and still be present by the time distributeAnyOfParentKeywords runs.
     const result = normalizePythinkerToolSchema({
       type: 'object',
       properties: {
         node: {
-          anyOf: [{ $ref: '#/$defs/Node' }, { type: 'null' }],
-        },
-      },
-      $defs: {
-        Node: {
-          type: 'object',
-          properties: { next: { $ref: '#/$defs/Node' } },
+          $ref: '#/properties/node',
+          $defs: { Extra: { type: 'string' } },
+          anyOf: [{ type: 'null' }],
         },
       },
     });
 
-    expect(result['$defs']).toBeDefined();
+    const node = (result['properties'] as Record<string, unknown>)['node'] as Record<string, unknown>;
+    expect(node['$ref']).toBe('#/properties/node');
+    expect(node['$defs']).toEqual({ Extra: { type: 'string' } });
+    const branches = node['anyOf'] as Record<string, unknown>[];
+    for (const branch of branches) {
+      expect(branch).not.toHaveProperty('$ref');
+      expect(branch).not.toHaveProperty('$defs');
+    }
   });
 
   it('dereferences and normalizes local definition buckets', () => {