bug

Author: roninjin10

src/SmithersCtx.ts

@@ -1,5 +1,14 @@
 import type { OutputKey } from "./OutputKey";
 import type { OutputAccessor, InferOutputEntry } from "./OutputAccessor";
+import type { z } from "zod";
+
+/**
+ * Reverse-lookup: given Schema and a value type V, find the key K where Schema[K] extends V.
+ * Used to narrow return types when passing Zod schema objects directly.
+ */
+type SchemaKeyForValue<Schema, V> = {
+  [K in keyof Schema & string]: Schema[K] extends V ? K : never;
+}[keyof Schema & string];
 
 export interface SmithersCtx<Schema> {
   runId: string;
@@ -8,22 +17,49 @@ export interface SmithersCtx<Schema> {
   input: Schema extends { input: infer T } ? T : Record<string, unknown>;
   outputs: OutputAccessor<Schema>;
 
+  // Overload: pass Zod schema value directly → narrowed return type
+  output<V extends z.ZodTypeAny>(
+    table: V,
+    key: OutputKey,
+  ): SchemaKeyForValue<Schema, V> extends never
+    ? InferOutputEntry<V>
+    : InferOutputEntry<V>;
+
+  // Overload: pass string key → narrowed via K
   output<K extends keyof Schema & string>(
-    table: K | Schema[K],
+    table: K,
     key: OutputKey,
   ): InferOutputEntry<Schema[K]>;
 
+  // Overload: pass Zod schema value directly → narrowed return type
+  outputMaybe<V extends z.ZodTypeAny>(
+    table: V,
+    key: OutputKey,
+  ): SchemaKeyForValue<Schema, V> extends never
+    ? InferOutputEntry<V> | undefined
+    : InferOutputEntry<V> | undefined;
+
+  // Overload: pass string key → narrowed via K
   outputMaybe<K extends keyof Schema & string>(
-    table: K | Schema[K],
+    table: K,
     key: OutputKey,
   ): InferOutputEntry<Schema[K]> | undefined;
 
+  // Overload: pass Zod schema value directly → narrowed return type
+  latest<V extends z.ZodTypeAny>(
+    table: V,
+    nodeId: string,
+  ): SchemaKeyForValue<Schema, V> extends never
+    ? InferOutputEntry<V> | undefined
+    : InferOutputEntry<V> | undefined;
+
+  // Overload: pass string key → narrowed via K
   latest<K extends keyof Schema & string>(
-    table: K | Schema[K],
+    table: K,
     nodeId: string,
   ): InferOutputEntry<Schema[K]> | undefined;
 
-  latestArray(value: unknown, schema: import("zod").ZodType): any[];
+  latestArray(value: unknown, schema: z.ZodType): any[];
 
   iterationCount(table: any, nodeId: string): number;
 }