fix(tests): update ScopedStorage tests for type safety

Author: sroussey

packages/knowledge-base/src/knowledge-base/KnowledgeBaseSchema.ts

@@ -56,5 +56,7 @@ export function knowledgeBaseTableNames(kbId: string): {
  * Checks whether a KnowledgeBaseRecord uses shared-table mode.
  */
 export function isSharedTableMode(record: KnowledgeBaseRecord): boolean {
-  return record.document_table === SHARED_DOCUMENT_TABLE && record.chunk_table === SHARED_CHUNK_TABLE;
+  return (
+    record.document_table === SHARED_DOCUMENT_TABLE && record.chunk_table === SHARED_CHUNK_TABLE
+  );
 }

packages/knowledge-base/src/knowledge-base/ScopedTabularStorage.ts

@@ -32,8 +32,7 @@ export class ScopedTabularStorage<
   Entity = any,
   PrimaryKey = any,
   InsertType = any,
-> implements ITabularStorage<Schema, PrimaryKeyNames, Entity, PrimaryKey, InsertType>
-{
+> implements ITabularStorage<Schema, PrimaryKeyNames, Entity, PrimaryKey, InsertType> {
   protected readonly inner: AnyTabularStorage;
   protected readonly kbId: string;
   protected readonly events = new EventEmitter<TabularEventListeners<PrimaryKey, Entity>>();
@@ -104,10 +103,7 @@ export class ScopedTabularStorage<
     const pageSize = 1000;
     let offset = 0;
     while (true) {
-      const page = await this.inner.query(
-        { kb_id: this.kbId } as any,
-        { offset, limit: pageSize }
-      );
+      const page = await this.inner.query({ kb_id: this.kbId } as any, { offset, limit: pageSize });
       if (!page || page.length === 0) break;
       count += page.length;
       if (page.length < pageSize) break;

packages/knowledge-base/src/knowledge-base/ScopedVectorStorage.ts

@@ -19,13 +19,13 @@ import { ScopedTabularStorage } from "./ScopedTabularStorage";
  * post-filtering results by `kb_id`.
  */
 export class ScopedVectorStorage<
-    Metadata extends Record<string, unknown> | undefined,
-    Schema extends DataPortSchemaObject,
-    Entity = any,
-    PrimaryKeyNames extends ReadonlyArray<keyof Schema["properties"]> = ReadonlyArray<
-      keyof Schema["properties"]
-    >,
-  >
+  Metadata extends Record<string, unknown> | undefined,
+  Schema extends DataPortSchemaObject,
+  Entity = any,
+  PrimaryKeyNames extends ReadonlyArray<keyof Schema["properties"]> = ReadonlyArray<
+    keyof Schema["properties"]
+  >,
+>
   extends ScopedTabularStorage<Schema, PrimaryKeyNames, Entity>
   implements IVectorStorage<Metadata, Schema, Entity, PrimaryKeyNames>
 {
@@ -47,19 +47,12 @@ export class ScopedVectorStorage<
     topK: number | undefined,
     overfetchLimit: number | undefined
   ): (Entity & { score: number })[] {
-    const filtered = results
-      .filter((r: any) => r.kb_id === this.kbId)
-      .slice(0, topK);
+    const filtered = results.filter((r: any) => r.kb_id === this.kbId).slice(0, topK);
 
     // Only warn when cross-KB filtering is likely the culprit: the inner search
     // returned exactly as many results as we requested (overfetch limit) and
     // filtering still left us short of topK.
-    if (
-      topK &&
-      overfetchLimit &&
-      results.length >= overfetchLimit &&
-      filtered.length < topK
-    ) {
+    if (topK && overfetchLimit && results.length >= overfetchLimit && filtered.length < topK) {
       console.warn(
         `ScopedVectorStorage: search returned ${filtered.length}/${topK} results after ` +
           `kb_id filtering. Consider increasing overFetchMultiplier (currently ${this.overFetchMultiplier}).`

packages/knowledge-base/src/knowledge-base/SharedTableSchemas.ts

@@ -78,13 +78,14 @@ export type SharedChunkPrimaryKey = typeof SharedChunkPrimaryKey;
  * Index definitions for efficient KB-scoped queries on shared document table.
  */
 export const SharedDocumentIndexes = [["kb_id"]] as const satisfies readonly (
-  keyof any | readonly (keyof any)[]
+  | keyof any
+  | readonly (keyof any)[]
 )[];
 
 /**
  * Index definitions for efficient KB-scoped queries on shared chunk table.
  */
-export const SharedChunkIndexes = [
-  ["kb_id"],
-  ["kb_id", "doc_id"],
-] as const satisfies readonly (keyof any | readonly (keyof any)[])[]; 
+export const SharedChunkIndexes = [["kb_id"], ["kb_id", "doc_id"]] as const satisfies readonly (
+  | keyof any
+  | readonly (keyof any)[]
+)[];

packages/storage/src/tabular/BaseTabularStorage.ts

@@ -4,8 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { DataPortSchemaObject, FromSchema, TypedArraySchemaOptions } from "@workglow/util/schema";
 import { createServiceToken, EventEmitter, makeFingerprint } from "@workglow/util";
+import { DataPortSchemaObject, FromSchema, TypedArraySchemaOptions } from "@workglow/util/schema";
 import {
   AnyTabularStorage,
   AutoGeneratedKeys,
@@ -70,7 +70,7 @@ export abstract class BaseTabularStorage<
   protected primaryKeySchema: DataPortSchemaObject;
   protected valueSchema: DataPortSchemaObject;
 
-  /** Name of the auto-generated key column (only first primary key column can be auto-generated) */
+  /** Name of the auto-generated key column (at most one primary key column may be auto-generated) */
   protected autoGeneratedKeyName: keyof Entity | null = null;
   /** Strategy for generating the auto-generated key */
   protected autoGeneratedKeyStrategy: KeyGenerationStrategy | null = null;
@@ -165,8 +165,8 @@ export abstract class BaseTabularStorage<
       }
     }
 
-    // Detect and validate auto-generated keys
-    // Only the first primary key column can be auto-generated
+    // Detect and validate auto-generated keys (at most one PK column; any PK position is allowed).
+    // Composite keys often put a scope column first (e.g. kb_id) and auto-generate a second id.
     const autoGeneratedKeys: string[] = [];
     for (const key of primaryKeyNames) {
       const keyStr = String(key);
@@ -181,20 +181,12 @@ export abstract class BaseTabularStorage<
     if (autoGeneratedKeys.length > 1) {
       throw new Error(
         `Multiple auto-generated keys detected: ${autoGeneratedKeys.join(", ")}. ` +
-          `Only the first primary key column can be auto-generated.`
+          `At most one primary key column can be auto-generated.`
       );
     }
 
     if (autoGeneratedKeys.length > 0) {
       const autoGenKeyName = autoGeneratedKeys[0];
-      const firstPrimaryKey = String(primaryKeyNames[0]);
-
-      if (autoGenKeyName !== firstPrimaryKey) {
-        throw new Error(
-          `Auto-generated key "${autoGenKeyName}" must be the first primary key column. ` +
-            `Current first primary key is "${firstPrimaryKey}".`
-        );
-      }
 
       this.autoGeneratedKeyName = autoGenKeyName as keyof Entity;
       this.autoGeneratedKeyStrategy = this.determineGenerationStrategy(

packages/task-graph/src/task-graph/TaskGraph.ts

@@ -67,8 +67,10 @@ export interface TaskGraphRunConfig {
   enforceEntitlements?: boolean;
 }
 
-export interface TaskGraphRunReactiveConfig
-  extends Omit<TaskGraphRunConfig, "enforceEntitlements" | "timeout"> {
+export interface TaskGraphRunReactiveConfig extends Omit<
+  TaskGraphRunConfig,
+  "enforceEntitlements" | "timeout"
+> {
   /** Optional service registry to use for this task graph */
   registry?: ServiceRegistry;
 }

packages/task-graph/src/task/EntitlementEnforcer.ts

@@ -105,11 +105,7 @@ export function createPolicyEnforcer(
 
     async checkTask(task: ITask): Promise<readonly TaskEntitlement[]> {
       const entitlements = task.entitlements();
-      return resolveAsks(
-        entitlements,
-        (task.constructor as typeof Task).type,
-        task.id
-      );
+      return resolveAsks(entitlements, (task.constructor as typeof Task).type, task.id);
     },
   };
 }

packages/task-graph/src/task/EntitlementPolicy.ts

@@ -4,7 +4,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { EntitlementGrant, EntitlementId, TaskEntitlement, TaskEntitlements } from "./TaskEntitlements";
+import type {
+  EntitlementGrant,
+  EntitlementId,
+  TaskEntitlement,
+  TaskEntitlements,
+} from "./TaskEntitlements";
 import { entitlementCovers, grantCoversResources } from "./TaskEntitlements";
 
 // ========================================================================

packages/tasks/src/task/mcp/McpPromptGetTask.ts

@@ -180,9 +180,7 @@ export class McpPromptGetTask extends Task<
 
   public static override entitlements(): TaskEntitlements {
     return {
-      entitlements: [
-        { id: Entitlements.MCP_PROMPT_GET, reason: "Gets prompts from MCP servers" },
-      ],
+      entitlements: [{ id: Entitlements.MCP_PROMPT_GET, reason: "Gets prompts from MCP servers" }],
     };
   }
 

packages/test/src/test/rag/ScopedStorage.test.ts

@@ -4,12 +4,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage";
 import {
   createKnowledgeBase,
+  getGlobalKnowledgeBaseRepository,
   isSharedTableMode,
   registerKnowledgeBase,
-  getGlobalKnowledgeBaseRepository,
   ScopedTabularStorage,
   ScopedVectorStorage,
   SharedChunkIndexes,
@@ -19,11 +18,15 @@ import {
   SharedDocumentPrimaryKey,
   SharedDocumentStorageSchema,
 } from "@workglow/knowledge-base";
+import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage";
 import { uuid4 } from "@workglow/util";
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 
 describe("ScopedTabularStorage", () => {
-  let sharedStorage: InstanceType<typeof InMemoryTabularStorage>;
+  let sharedStorage: InMemoryTabularStorage<
+    typeof SharedDocumentStorageSchema,
+    typeof SharedDocumentPrimaryKey
+  >;
   let scopeA: ScopedTabularStorage<any, any>;
   let scopeB: ScopedTabularStorage<any, any>;
 
@@ -244,7 +247,10 @@ describe("ScopedTabularStorage", () => {
 });
 
 describe("ScopedVectorStorage", () => {
-  let sharedStorage: InstanceType<typeof InMemoryVectorStorage>;
+  let sharedStorage: InMemoryVectorStorage<
+    typeof SharedChunkVectorStorageSchema,
+    typeof SharedChunkPrimaryKey
+  >;
   let scopeA: ScopedVectorStorage<any, any>;
   let scopeB: ScopedVectorStorage<any, any>;
 
@@ -261,7 +267,7 @@ describe("ScopedVectorStorage", () => {
   });
 
   afterEach(() => {
-    sharedStorage.destroy();
+    sharedStorage?.destroy?.();
   });
 
   test("getVectorDimensions delegates to inner", () => {