modifySchemaContent only partially applies schema updates

[ogerhardt]

Describe the bug

modifySchemaContent appears to apply simple updates (for example enum) but not complex updates (for example allOf with if/then, or full object replacement).

Use Case: I publish a generic schema to a registry, then specialize it at runtime from a VSCode extension.

Generic shape:

$schema: "https://schema.example/schema.json"
myType: <any string>
myProperties:
  <any additional property>

Desired specialized shapes:

$schema: "https://schema.example/schema.json"
myType: "first"
myProperties:
  myFirstProperty: ""

$schema: "https://schema.example/schema.json"
myType: "second"
myProperties:
  mySecondProperty: ""

For this, I add an enum node to myType and an allOf node to the whole object to define the available properties under myProperties with if/then conditions (see test code below).

When I split the modification into multiple calls, the enum update works, but the allOf update does not. If I replace the full object definition, neither works.

The JSON schema itself works (for example through registerCustomSchemaProvider).

Expected Behavior

The schema content modification works if the whole object is replaced.

Current Behavior

The schem content modification only works for the enum property.

LLM Danger Zone

Since I do not know the internals of your implementation, I used an LLM to help inspect this (i.e. double check that the problem is not on my side), so treat that part with a huuuuge grain of salt. I also do not know whether this level of schema modification is intended to be supported.

With the following adjustments in src/languageservice/services/yamlSchemaService.ts, it appears to work in my extension. That said, I do not know whether this is a correct fix, introduces side effects, or conflicts with intended behavior.

  async resolveSchemaContent(
    schemaToResolve: UnresolvedSchema,
    schemaURL: string,
    dependencies: SchemaDependencies
  ): Promise<ResolvedSchema> {
    [...]
-    let schema = raw as JSONSchema;
+    // LLM "fix" and comment
+    // Resolve against a clone to avoid mutating cached unresolved schema content.
+    // Otherwise prior ref expansions can leave stale inlined sections across
+    let schema = _cloneSchema(raw as JSONSchema, new Map()) as JSONSchema;

  public async saveSchema(schemaId: string, schemaContent: JSONSchema): Promise<void> {
    const id = normalizeId(schemaId);
-    this.getOrAddSchemaHandle(id, schemaContent);
+    // LLM "fix" and comment
+    // Replace the handle so resolved/unresolved caches are rebuilt from updated source content.
+    this.schemasById[id] = new SchemaHandle(this, id, schemaContent);
+    this.cachedSchemaForResource = undefined;
    this.schemaPriorityMapping.set(id, new Set<SchemaPriority>().add(SchemaPriority.Settings));
    return Promise.resolve(undefined);
  }

  public async saveSchema(schemaId: string, schemaContent: JSONSchema): Promise<void> {
    const id = normalizeId(schemaId);
-    this.getOrAddSchemaHandle(id, schemaContent);
+    // LLM "fix" and comment
+    // Replace the handle so resolved/unresolved caches are rebuilt from updated source content.
+    this.schemasById[id] = new SchemaHandle(this, id, schemaContent);
+    this.cachedSchemaForResource = undefined;
    this.schemaPriorityMapping.set(id, new Set<SchemaPriority>().add(SchemaPriority.Settings));
    return Promise.resolve(undefined);
  }

  public async addContent(additions: SchemaAdditions): Promise<void> {
-    const schema = await this.getResolvedSchema(additions.schema);
+    const schemaHandle = this.getOrAddSchemaHandle(normalizeId(additions.schema));
+    const unresolvedSchema = await Promise.resolve(schemaHandle.getUnresolvedSchema());
-    if (schema) {
+    if (unresolvedSchema?.schema) {
-      const resolvedSchemaLocation = this.resolveJSONSchemaToSection(schema.schema, additions.path);
+      const resolvedSchemaLocation = this.resolveJSONSchemaToSection(unresolvedSchema.schema, additions.path);

      if (typeof resolvedSchemaLocation === 'object') {
        resolvedSchemaLocation[additions.key] = additions.content;
      }
-      await this.saveSchema(additions.schema, schema.schema);
+      await this.saveSchema(additions.schema, unresolvedSchema.schema);
    }

Steps to Reproduce

I tried writing a unit test, but I am not fully confident it reflects the real runtime behavior. In unit tests, only part of the "fix" seems necessary. In manual testing with my VSCode extension, all parts seem necessary. So there appears to be a gap between test behavior and runtime behavior.

I fiddled around in test/autoCompletionFix.test.ts.

  function getCompletionLabels(completion: CompletionList): string[] {
    return completion.items.map((item) => String(item.label));
  }

  function setup(content: string): TextDocument {
    const testTextDocument = setupSchemaIDTextDocument(content);
    yamlSettings.documents = new TextDocumentTestManager();
    (yamlSettings.documents as TextDocumentTestManager).set(testTextDocument);
    return testTextDocument;
  }

  function act(testDocument: TextDocument, line: number, character: number): Promise<CompletionList> {
    return languageHandler.completionHandler({
      position: Position.create(line, character),
      textDocument: testDocument,
    });
  }

  const schemaUri = 'https://schema.example/schema.json';
  const initialSchema: JSONSchema = {
    $schema: 'https://json-schema.org/draft/2020-12/schema',
    $id: 'https://schema.example/schema.json',
    $ref: '#/$defs/myobject',
    $defs: {
      myobject: {
        type: 'object',
        required: ['$schema', 'myType'],
        additionalProperties: false,
        properties: {
          $schema: {
            type: 'string',
          },
          myType: {
            type: 'string',
          },
          myProperties: {
            type: 'object',
            additionalProperties: true,
          },
        },
      },
    },
  };

  const modifiedMyObject = {
    type: 'object',
    required: ['$schema', 'myType'],
    additionalProperties: false,
    properties: {
      $schema: {
        type: 'string',
      },
      myType: {
        type: 'string',
        enum: ['first', 'second'],
      },
      myProperties: {
        type: 'object',
        additionalProperties: true,
      },
    },
    allOf: [
      {
        if: {
          properties: {
            myType: {
              const: 'first',
            },
          },
        },
        then: {
          required: ['$schema', 'myType', 'myProperties'],
          properties: {
            myProperties: {
              type: 'object',
              additionalProperties: false,
              required: ['myFirstProperty'],
              properties: {
                myFirstProperty: {
                  type: 'string',
                },
              },
            },
          },
        },
      },
      {
        if: {
          properties: {
            myType: {
              const: 'second',
            },
          },
        },
        then: {
          required: ['$schema', 'myType'],
          properties: {
            myProperties: {
              type: 'object',
              additionalProperties: false,
              required: [],
              properties: {
                mySecondProperty: {
                  type: 'string',
                },
              },
            },
          },
        },
      },
    ],
  };

  it('schema modification should keep myType completions when enum is replaced', async () => {
    schemaProvider.addSchemaWithUri(SCHEMA_ID, schemaUri, initialSchema);

    const content = `$schema: https://schema.example/schema.json
myType: 
`;
    const textDocument = await setup(content);

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs/myobject/properties/myType',
      key: 'enum',
      content: ['first', 'second'],
    });

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs/myobject',
      key: 'allOf',
      content: modifiedMyObject.allOf,
    });

    const typeCompletion = await act(textDocument, 1, 8);
    const typeLabels = getCompletionLabels(typeCompletion);
    expect(typeLabels).to.include.members(['first', 'second']);
  });

  it('schema modification should keep myProperties completions when allOf is replaced', async () => {
    schemaProvider.addSchemaWithUri(SCHEMA_ID, schemaUri, initialSchema);

    const content = `$schema: https://schema.example/schema.json
myType: first
myProperties:
  
`;
    const textDocument = await setup(content);

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs/myobject/properties/myType',
      key: 'enum',
      content: ['first', 'second'],
    });

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs/myobject',
      key: 'allOf',
      content: modifiedMyObject.allOf,
    });

    const propertiesCompletion = await act(textDocument, 3, 2);
    const propertiesLabels = getCompletionLabels(propertiesCompletion);
    expect(propertiesLabels).to.include('myFirstProperty');
  });

  it('schema modification should keep myType completions when $defs/myobject is replaced', async () => {
    schemaProvider.addSchemaWithUri(SCHEMA_ID, schemaUri, initialSchema);

    const content = `$schema: https://schema.example/schema.json
myType: 
`;
    const textDocument = setup(content);

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs',
      key: 'myobject',
      content: modifiedMyObject,
    });

    const typeCompletion = await act(textDocument, 1, 8);
    const typeLabels = getCompletionLabels(typeCompletion);
    expect(typeLabels).to.include.members(['first', 'second']);
  });

  it('schema modification should keep myProperties completions when $defs/myobject is replaced', async () => {
    schemaProvider.addSchemaWithUri(SCHEMA_ID, schemaUri, initialSchema);

    const content = `$schema: https://schema.example/schema.json
myType: first
myProperties:
  
`;
    const textDocument = setup(content);

    languageService.modifySchemaContent({
      action: MODIFICATION_ACTIONS.add,
      schema: schemaUri,
      path: '$defs',
      key: 'myobject',
      content: modifiedMyObject,
    });

    const propertiesCompletion = await act(textDocument, 3, 2);
    const propertiesLabels = getCompletionLabels(propertiesCompletion);
    expect(propertiesLabels).to.include('myFirstProperty');
  });

Environment

• Windows
• Mac
• Linux