diff --git a/core/config/default.ts b/core/config/default.ts index 07d96277bf4..7a7783ebc52 100644 --- a/core/config/default.ts +++ b/core/config/default.ts @@ -2,6 +2,7 @@ import { AssistantUnrolled, ConfigYaml, ModelConfig, + graniteCodeModelSlugs, } from "@continuedev/config-yaml"; export const defaultContextProvidersVsCode = [ @@ -24,7 +25,7 @@ export const defaultConfig: ConfigYaml = { name: "Local Assistant", version: "1.0.0", schema: "v1", - models: [], + models: graniteCodeModelSlugs.map((slug) => ({ uses: slug })), context: defaultContextProvidersVsCode, }; diff --git a/core/granite/config/virtualReference.ts b/core/granite/config/virtualReference.ts new file mode 100644 index 00000000000..2476202817f --- /dev/null +++ b/core/granite/config/virtualReference.ts @@ -0,0 +1,176 @@ +import { graniteCodeModelSlugs } from "@continuedev/config-yaml"; +import fs from "fs"; +import * as YAML from "yaml"; +import { IDE } from "../.."; +import { getConfigYamlPath } from "../../util/paths"; +import { parseUsesSlug } from "../utils/parseUsesSlug"; + +const keyOrder = ["name", "version", "schema", "models"]; + +function extractSlugsFromComment(comment: string | undefined | null): string[] { + const res: string[] = []; + + if (!comment) { + return res; + } + + const lines = comment.split("\n"); + for (const line of lines) { + let slug = parseUsesSlug(line); + if (slug) { + res.push(slug); + } + } + return res; +} + +/** + * Determines which model slugs need to be added. + * + * This function identifies which model slugs from the predefined `graniteCodeModelSlugs` array are not currently in use within the provided YAML sequence of models. It considers two sources of "in-use" slugs: + * 1. An active model entry (e.g. `- uses: $granite-code/models/chat`). + * 2. Any comment within the `models` block (checked via a deep search). + * + * @param models The YAML sequence of models to check against. + * @returns An array of strings representing the model slugs that are not in use. + */ +function slugsToAdd(models: YAML.YAMLSeq): string[] { + const graniteCodeModelSlugsSet = new Set(graniteCodeModelSlugs); + + // Trailing comments of models section + const slugs = extractSlugsFromComment(models.comment); + for (const slug of slugs) { + graniteCodeModelSlugsSet.delete(slug); + } + + // Models in use + for (const model of models.items) { + if (YAML.isMap(model)) { + const usesSlug = model.get("uses"); + if (usesSlug && typeof usesSlug === "string") { + graniteCodeModelSlugsSet.delete(usesSlug); + } + } + } + + // Commented out models + const extractedSlugs = extractYamlCommentsDFS(models); + extractedSlugs.forEach((slug) => graniteCodeModelSlugsSet.delete(slug)); + + return [...graniteCodeModelSlugsSet]; +} + +/** + * Recursively extracts "slugs" from all comments within a YAML document structure. + * + * This function performs a depth-first search (DFS) through a given YAML Map (`{...}`) or + * Sequence (`[...]`). It inspects the `comment` (inline) and `commentBefore` (on the preceding line) + * properties of every key, value, and item. It then uses a helper function, `extractSlugsFromComment`, + * to parse slugs from those comment strings. + * + * The final result is a deduplicated array of all unique slugs found throughout the entire + * nested structure. + * + * @param {YAML.YAMLMap | YAML.YAMLSeq} yamlNode - The root YAML node (Map or Sequence) to start the search from. + * @returns {string[]} A unique, flat array of all slugs extracted from the comments. + */ +function extractYamlCommentsDFS( + yamlNode: YAML.YAMLMap | YAML.YAMLSeq, +): string[] { + let extractedSlugs: string[] = []; + + function pushSlugs(node: YAML.Node) { + extractedSlugs.push(...extractSlugsFromComment(node.comment)); + extractedSlugs.push(...extractSlugsFromComment(node.commentBefore)); + } + + // Comments of the current node + pushSlugs(yamlNode); + + for (const item of yamlNode.items) { + if (YAML.isPair(item)) { + // Key Section + // If it's not a scalar, ignore it + if (YAML.isScalar(item.key)) { + pushSlugs(item.key); + } + + // Value Section + // Only consider Scalar, Map, and Seq + if (YAML.isScalar(item.value)) { + pushSlugs(item.value); + } else if (YAML.isMap(item.value)) { + extractedSlugs.push(...extractYamlCommentsDFS(item.value)); + } else if (YAML.isSeq(item.value)) { + extractedSlugs.push(...extractYamlCommentsDFS(item.value)); + } + } else if (YAML.isMap(item)) { + extractedSlugs.push(...extractYamlCommentsDFS(item)); + } else if (YAML.isSeq(item)) { + extractedSlugs.push(...extractYamlCommentsDFS(item)); + } + } + + return [...new Set(extractedSlugs)]; +} + +/** + * Validates and injects missing virtual references into the global configuration YAML file. + * + * This function reads the global YAML config, ensures a `models` sequence (an array) exists, + * and then populates it with any required virtual model references ("slugs") that are missing. + * The changes are then written back to the file. + * + * It operates safely by: + * - Aborting without changes if the YAML file has parsing errors. + * - Catching any exceptions during file I/O or processing and displaying a warning toast in the IDE. + * + * @param {IDE} ide - The IDE instance, used to display a warning notification if the process fails. + * @param {string} configFilePath - The path to the config YAML file to validate. If not provided, the default path will be used. + */ +export async function validateVirtualReferences( + ide: IDE, + configFilePath?: string, +): Promise { + try { + const filePath = configFilePath || getConfigYamlPath("vscode"); + const rawContent = (await fs.promises.readFile(filePath)).toString(); + const yamlDoc = YAML.parseDocument(rawContent); + + // Don't do anything, if there is a parsing error + if (yamlDoc.errors.length !== 0) { + return; + } + + // Add an empty models section if there is no models section + if (yamlDoc.get("models") === undefined) { + yamlDoc.add(new YAML.Pair("models", [])); + } + + const models = yamlDoc.get("models"); + + // Models should be a Seq + if (YAML.isSeq(models)) { + // Change to block style if models section is empty + if (models.items.length === 0) { + models.flow = false; + } + + const missingRefs = slugsToAdd(models); + // Only inject missing references if there are any + if (missingRefs.length > 0) { + for (const missingRef of missingRefs) { + const node = yamlDoc.createNode({ uses: missingRef }); + models.add(node); + } + const stringYaml = yamlDoc.toString(); + await fs.promises.writeFile(filePath, stringYaml); + } + } + } catch { + void ide.showToast( + "warning", + "Failed to inject missing virtual references", + ); + } +} diff --git a/core/granite/config/virtualReference.vitest.ts b/core/granite/config/virtualReference.vitest.ts new file mode 100644 index 00000000000..95cc7a4e5a8 --- /dev/null +++ b/core/granite/config/virtualReference.vitest.ts @@ -0,0 +1,166 @@ +import { graniteCodeModelSlugs } from "@continuedev/config-yaml"; +import path from "path"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { testIde } from "../../test/fixtures"; +import { + addToTestDir, + setUpTestDir, + tearDownTestDir, + TEST_DIR_PATH, +} from "../../test/testDir"; +import { localPathToUri } from "../../util/pathToUri"; +import { validateVirtualReferences } from "./virtualReference"; + +describe("Test validateVirtualReferences", () => { + beforeAll(async () => { + setUpTestDir(); + }); + + afterAll(async () => { + tearDownTestDir(); + }); + + it("injects all references when the models section is empty", async () => { + const injectedReferences = graniteCodeModelSlugs + .map((slug) => ` - uses: ${slug}`) + .join("\n"); + const rawConfig = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: []`; + const expectedResult = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: +${injectedReferences} +`; + addToTestDir([["config.yaml", rawConfig]]); + + const configUri = localPathToUri(path.join(TEST_DIR_PATH, "config.yaml")); + await validateVirtualReferences( + testIde, + path.join(TEST_DIR_PATH, "config.yaml"), + ); + + const res = await testIde.readFile(configUri); + + expect(res).toBe(expectedResult); + }); + + it("don't inject the references that are commented out", async () => { + const commentedOutRef = graniteCodeModelSlugs[0]; + const injectedReferences = graniteCodeModelSlugs + .slice(1) + .map((slug) => ` - uses: ${slug}`) + .join("\n"); + const rawConfig = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: + # - uses: ${commentedOutRef} + - name: Local Chat Model + provider: ollama + model: granite3.3:8b + roles: + - chat`; + + const expectedResult = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: + # - uses: ${commentedOutRef} + - name: Local Chat Model + provider: ollama + model: granite3.3:8b + roles: + - chat +${injectedReferences} +`; + + addToTestDir([["config.yaml", rawConfig]]); + + const configUri = localPathToUri(path.join(TEST_DIR_PATH, "config.yaml")); + await validateVirtualReferences( + testIde, + path.join(TEST_DIR_PATH, "config.yaml"), + ); + + const res = await testIde.readFile(configUri); + + expect(res).toBe(expectedResult); + }); + + it("don't inject the references that are commented out (flow style)", async () => { + const commentedOutRef = graniteCodeModelSlugs[0]; + const injectedReferences = graniteCodeModelSlugs + .slice(1) + .map((slug) => ` { uses: ${slug} }`) + .join(",\n"); + const rawConfig = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: [ + # {uses: ${commentedOutRef}, override: { name: new name }}, + { name: Local Chat Model, + provider: ollama, + model: granite3.3:8b, + roles: [chat], + } +]`; + + const expectedResult = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: + [ + # {uses: ${commentedOutRef}, override: { name: new name }}, + { + name: Local Chat Model, + provider: ollama, + model: granite3.3:8b, + roles: [ chat ] + }, +${injectedReferences} + ] +`; + + addToTestDir([["config.yaml", rawConfig]]); + + const configUri = localPathToUri(path.join(TEST_DIR_PATH, "config.yaml")); + await validateVirtualReferences( + testIde, + path.join(TEST_DIR_PATH, "config.yaml"), + ); + + const res = await testIde.readFile(configUri); + + expect(res).toBe(expectedResult); + }); + + it("don't rewtire the config file if there are no missing references", async () => { + const commentedOutRefsWithIncreasingIndent = graniteCodeModelSlugs + .map((slug, ind) => `#${" ".repeat(ind + 1)}{ uses: ${slug} }`) + .join(",\n"); + const rawConfig = `name: Local Assistant +version: 1.0.0 +schema: v1 +models: [ +${commentedOutRefsWithIncreasingIndent} +]`; + + // The expected result should keep the increasing indentation because we didn't rewrite the file + const expectedResult = rawConfig; + + addToTestDir([["config.yaml", rawConfig]]); + + const configUri = localPathToUri(path.join(TEST_DIR_PATH, "config.yaml")); + await validateVirtualReferences( + testIde, + path.join(TEST_DIR_PATH, "config.yaml"), + ); + + const res = await testIde.readFile(configUri); + + expect(res).toBe(expectedResult); + }); +}); diff --git a/extensions/vscode/src/extension/parseUsesSlug.ts b/core/granite/utils/parseUsesSlug.ts similarity index 97% rename from extensions/vscode/src/extension/parseUsesSlug.ts rename to core/granite/utils/parseUsesSlug.ts index 2e2f32fe6fc..1cd089399c5 100644 --- a/extensions/vscode/src/extension/parseUsesSlug.ts +++ b/core/granite/utils/parseUsesSlug.ts @@ -30,6 +30,7 @@ export function findSlug( } else { // If not quoted, remove any trailing comment slug = slug.replace(/\s*#.*$/, "").trim(); + slug = slug.match(/^[^\s,}]*/)![0]; } if (!slug || !slug.startsWith("$")) { diff --git a/extensions/vscode/src/extension/parseUsesSlug.vitest.ts b/core/granite/utils/parseUsesSlug.vitest.ts similarity index 100% rename from extensions/vscode/src/extension/parseUsesSlug.vitest.ts rename to core/granite/utils/parseUsesSlug.vitest.ts diff --git a/extensions/vscode/src/extension/ConfigYamlDefinitionProvider.ts b/extensions/vscode/src/extension/ConfigYamlDefinitionProvider.ts index 29086591d02..4aef615563f 100644 --- a/extensions/vscode/src/extension/ConfigYamlDefinitionProvider.ts +++ b/extensions/vscode/src/extension/ConfigYamlDefinitionProvider.ts @@ -1,6 +1,6 @@ import * as vscode from "vscode"; -import { parseUsesSlug } from "./parseUsesSlug"; +import { parseUsesSlug } from "core/granite/utils/parseUsesSlug"; import { getVirtualConfigUri } from "./VirtualConfigUris"; export function registerConfigYamlDefinitionProvider(): vscode.Disposable { diff --git a/extensions/vscode/src/extension/ConfigYamlDiagnosticsProvider.ts b/extensions/vscode/src/extension/ConfigYamlDiagnosticsProvider.ts index 73cd5289648..2de440a4278 100644 --- a/extensions/vscode/src/extension/ConfigYamlDiagnosticsProvider.ts +++ b/extensions/vscode/src/extension/ConfigYamlDiagnosticsProvider.ts @@ -1,7 +1,7 @@ import { isSupportedSlug } from "core/config/yaml/VirtualConfigYamlSupport"; import * as vscode from "vscode"; -import { findSlug } from "./parseUsesSlug"; +import { findSlug } from "core/granite/utils/parseUsesSlug"; export function registerConfigYamlDiagnosticsProvider(): vscode.Disposable { const diagnosticCollection = vscode.languages.createDiagnosticCollection( diff --git a/extensions/vscode/src/extension/ConfigYamlHoverProvider.ts b/extensions/vscode/src/extension/ConfigYamlHoverProvider.ts index 44c18399818..09216d04f2b 100644 --- a/extensions/vscode/src/extension/ConfigYamlHoverProvider.ts +++ b/extensions/vscode/src/extension/ConfigYamlHoverProvider.ts @@ -3,7 +3,7 @@ import * as vscode from "vscode"; import { getExtensionVersion } from "../util/util"; -import { parseUsesSlug } from "./parseUsesSlug"; +import { parseUsesSlug } from "core/granite/utils/parseUsesSlug"; export function registerConfigYamlHoverProvider(): vscode.Disposable { return vscode.languages.registerHoverProvider( diff --git a/extensions/vscode/src/extension/VsCodeExtension.ts b/extensions/vscode/src/extension/VsCodeExtension.ts index 77e1d4bdadf..09a3c5b2afb 100644 --- a/extensions/vscode/src/extension/VsCodeExtension.ts +++ b/extensions/vscode/src/extension/VsCodeExtension.ts @@ -47,6 +47,7 @@ import { registerConfigYamlHoverProvider } from "./ConfigYamlHoverProvider"; import { registerVirtualConfigDocumentProvider } from "./VirtualConfigYamlDocumentProvider"; import { VsCodeMessenger } from "./VsCodeMessenger"; +import { validateVirtualReferences } from "core/granite/config/virtualReference"; import setupNextEditWindowManager, { NextEditWindowManager, } from "../activation/NextEditWindowManager"; @@ -424,6 +425,8 @@ export class VsCodeExtension { ), ); + void validateVirtualReferences(this.ide); + const linkProvider = vscode.languages.registerDocumentLinkProvider( { language: "yaml" }, new ConfigYamlDocumentLinkProvider(),