versioned-components fixes

packages/hub/src/app/ai/SquigglePlaygroundForWorkflow.tsx

@@ -42,7 +42,7 @@ export function SquigglePlaygroundForWorkflow({
     <squiggle.components.SquigglePlayground
       height={height}
       defaultCode={defaultCode}
-      linker={llmLinker}
+      linker={llmLinker as any}
       renderExtraControls={() => (
         <div className="flex h-full items-center justify-end gap-2">
           <SquigglePlaygroundVersionPicker

packages/hub/src/squiggle/linker.ts

@@ -1,7 +1,8 @@
 import { z } from "zod";
 
-import { SqLinker, SqModule } from "@quri/squiggle-lang";
 import {
+  type GuardedSquigglePackages,
+  type SquigglePackages,
   versionedSquigglePackages,
   versionSupportsSqProjectV2,
 } from "@quri/versioned-squiggle-components";
@@ -35,38 +36,40 @@ export function serializeSourceId({ owner, slug }: ParsedSourceId): string {
   return `${PREFIX}:${owner}/${slug}`;
 }
 
-const linker: SqLinker = {
-  resolve(name) {
-    return name;
-  },
-  async loadModule(sourceId: string) {
-    const { owner, slug } = parseSourceId(sourceId);
-
-    const data = await fetch(
-      `/api/get-source?${new URLSearchParams({ owner, slug })}`
-    ).then((res) => res.json());
-
-    const parsed = z.object({ code: z.string() }).safeParse(data);
-    if (!parsed.success) {
-      throw new Error(`Failed to fetch source for ${sourceId}`);
-    }
+async function fetchSource(sourceId: string) {
+  const { owner, slug } = parseSourceId(sourceId);
 
-    return new SqModule({
-      name: sourceId,
-      code: parsed.data.code,
-    });
-  },
-};
+  const data = await fetch(
+    `/api/get-source?${new URLSearchParams({ owner, slug })}`
+  ).then((res) => res.json());
+  const parsed = z.object({ code: z.string() }).safeParse(data);
+  if (!parsed.success) {
+    throw new Error(`Failed to fetch source for ${sourceId}`);
+  }
 
-const oldLinker: OldSqLinker = {
-  resolve: (name) => name,
-  loadSource: async (sourceId: string) => {
-    return (await linker.loadModule(sourceId)).code;
-  },
-};
+  return parsed.data.code;
+}
 
-// export directly from versioned-components?
-type SquigglePackages = Awaited<ReturnType<typeof versionedSquigglePackages>>;
+function getModernLinker(
+  squiggle: GuardedSquigglePackages<typeof versionSupportsSqProjectV2>
+) {
+  const linker: ReturnType<typeof squiggle.lang.makeSelfContainedLinker> = {
+    resolve(name: string) {
+      return name;
+    },
+    async loadModule(sourceId: string) {
+      const code = await fetchSource(sourceId);
+
+      return new squiggle.lang.SqModule({
+        name: sourceId,
+        code,
+      }) as any;
+      // Different SqModule types are not compatible, so the union of `SqModule | SqModule` doesn't work.
+      // But this should be fine if we're careful because we instantiate the SqModule for `squiggle.lang`.
+    },
+  };
+  return linker;
+}
 
 // this type extraction is awkward but it works
 type OldSqLinker = NonNullable<
@@ -77,11 +80,18 @@ type OldSqLinker = NonNullable<
   >["linker"]
 >;
 
+const oldLinker: OldSqLinker = {
+  resolve: (name: string) => name,
+  loadSource: async (sourceId: string) => {
+    return await fetchSource(sourceId);
+  },
+};
+
 export function getHubLinker(
   squiggle: Awaited<ReturnType<typeof versionedSquigglePackages>>
 ) {
-  if (versionSupportsSqProjectV2.plain(squiggle.version)) {
-    return linker;
+  if (versionSupportsSqProjectV2.object(squiggle)) {
+    return getModernLinker(squiggle);
   } else {
     return oldLinker;
   }
@@ -91,7 +101,9 @@ export function sqProjectWithHubLinker(
   squiggle: Awaited<ReturnType<typeof versionedSquigglePackages>>
 ) {
   if (versionSupportsSqProjectV2.object(squiggle)) {
-    return new squiggle.lang.SqProject({ linker });
+    return new squiggle.lang.SqProject({
+      linker: getModernLinker(squiggle) as any,
+    });
   } else {
     return new squiggle.lang.SqProject({ linker: oldLinker });
   }

packages/versioned-components/src/index.ts

@@ -1,6 +1,9 @@
 export { SquigglePlaygroundVersionPicker } from "./SquigglePlaygroundVersionPicker.js";
 export { SquiggleVersionShower } from "./SquiggleVersionShower.js";
-export { versionedSquigglePackages } from "./versionedSquigglePackages.js";
+export {
+  type SquigglePackages,
+  versionedSquigglePackages,
+} from "./versionedSquigglePackages.js";
 
 export {
   checkSquiggleVersion,
@@ -20,3 +23,8 @@ export {
 } from "./predicates.js";
 
 export { useAdjustSquiggleVersion } from "./hooks.js";
+
+export type {
+  AntiGuardedSquigglePackages,
+  GuardedSquigglePackages,
+} from "./predicates.js";

packages/versioned-components/src/predicates.ts

@@ -115,3 +115,15 @@ export const versionSupportsSqProjectV2 = excludeVersions([
   "0.9.4",
   "0.9.5",
 ]);
+
+// Useful if you need to write a function that takes `squiggle` object that was previously guarded.
+export type GuardedSquigglePackages<
+  Guard extends CompositeGuard<SquiggleVersion>,
+> = SquigglePackages<Guard extends CompositeGuard<infer T> ? T : never>;
+
+// Useful if you need to write a function that takes `squiggle` object that was previously guarded, and the guard check fails.
+export type AntiGuardedSquigglePackages<
+  Guard extends CompositeGuard<SquiggleVersion>,
+> = SquigglePackages<
+  Guard extends CompositeGuard<infer T> ? Exclude<SquiggleVersion, T> : never
+>;

packages/versioned-components/src/versionedSquiggleComponents.ts

@@ -3,7 +3,9 @@
 // squiggle-components are not RSC-friendly
 import { FC } from "react";
 import { type SqLinker as SqLinkerV1 } from "squiggle-lang-0.9.5";
-import { type SqLinker as SqLinkerV2 } from "@quri/squiggle-lang";
+
+import { type SqLinker as SqLinkerDev } from "@quri/squiggle-lang";
+
 import { SquiggleLangPackageTypes } from "./versionedSquiggleLang.js";
 import { SquiggleVersion } from "./versions.js";
 
@@ -12,6 +14,13 @@ import { SquiggleVersion } from "./versions.js";
  * This makes them less type-safe, but matching squiggle-lang object versions with components verisons is almost impossible without this.
  */
 
+type SqLinkerV2 = Omit<SqLinkerDev, "loadModule"> & {
+  // Relax the type of `loadModule` - various versions of `SqModule` return type
+  // are not compatible, even if they're identical
+  // (because of "Types have separate declarations of a private property '_ast'" error.)
+  loadModule: (sourceId: string, hash?: string) => Promise<unknown>;
+};
+
 type AnySqProject = InstanceType<
   SquiggleLangPackageTypes[SquiggleVersion]["SqProject"]
 >;