From 970e780ef71eb351a8b32941d7c4515bce0c0f9c Mon Sep 17 00:00:00 2001
From: Ian Macartney <ianhmacartney@gmail.com>
Date: Wed, 12 Nov 2025 16:05:55 -0800
Subject: [PATCH] make step types match ctx better

---
 src/client/step.ts            |  3 ++-
 src/client/workflowContext.ts | 28 ++++++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/client/step.ts b/src/client/step.ts
index 9fa2a7e..2726b3a 100644
--- a/src/client/step.ts
+++ b/src/client/step.ts
@@ -8,6 +8,7 @@ import {
   createFunctionHandle,
   type FunctionReference,
   type FunctionType,
+  type FunctionVisibility,
   type GenericDataModel,
   type GenericMutationCtx,
 } from "convex/server";
@@ -31,7 +32,7 @@ export type StepRequest = {
     | {
         kind: "function";
         functionType: FunctionType;
-        function: FunctionReference<FunctionType, "internal">;
+        function: FunctionReference<FunctionType, FunctionVisibility>;
         args: unknown;
       }
     | {
diff --git a/src/client/workflowContext.ts b/src/client/workflowContext.ts
index ea6a1aa..4dd07e9 100644
--- a/src/client/workflowContext.ts
+++ b/src/client/workflowContext.ts
@@ -6,6 +6,7 @@ import type {
   FunctionReference,
   FunctionReturnType,
   FunctionType,
+  FunctionVisibility,
 } from "convex/server";
 import type { Validator } from "convex/values";
 import type { EventId, SchedulerOptions, WorkflowId } from "../types.js";
@@ -33,10 +34,9 @@ export type WorkflowCtx = {
    * @param args - The arguments to the query function.
    * @param opts - Options for scheduling and naming the query.
    */
-  runQuery<Query extends FunctionReference<"query", "internal">>(
+  runQuery<Query extends FunctionReference<"query", FunctionVisibility>>(
     query: Query,
-    args: FunctionArgs<Query>,
-    opts?: RunOptions,
+    ...args: OptionalRestArgs<RunOptions, Query>
   ): Promise<FunctionReturnType<Query>>;
 
   /**
@@ -46,10 +46,11 @@ export type WorkflowCtx = {
    * @param args - The arguments to the mutation function.
    * @param opts - Options for scheduling and naming the mutation.
    */
-  runMutation<Mutation extends FunctionReference<"mutation", "internal">>(
+  runMutation<
+    Mutation extends FunctionReference<"mutation", FunctionVisibility>,
+  >(
     mutation: Mutation,
-    args: FunctionArgs<Mutation>,
-    opts?: RunOptions,
+    ...args: OptionalRestArgs<RunOptions, Mutation>
   ): Promise<FunctionReturnType<Mutation>>;
 
   /**
@@ -59,10 +60,9 @@ export type WorkflowCtx = {
    * @param args - The arguments to the action function.
    * @param opts - Options for retrying, scheduling and naming the action.
    */
-  runAction<Action extends FunctionReference<"action", "internal">>(
+  runAction<Action extends FunctionReference<"action", FunctionVisibility>>(
     action: Action,
-    args: FunctionArgs<Action>,
-    opts?: RunOptions & RetryOption,
+    ...args: OptionalRestArgs<RunOptions & RetryOption, Action>
   ): Promise<FunctionReturnType<Action>>;
 
   /**
@@ -99,6 +99,14 @@ export type WorkflowCtx = {
   ): Promise<T>;
 };
 
+export type OptionalRestArgs<
+  Opts,
+  FuncRef extends FunctionReference<FunctionType, FunctionVisibility>,
+> =
+  FuncRef["_args"] extends Record<string, never>
+    ? [args?: Record<string, never>, opts?: Opts]
+    : [args: FuncRef["_args"], opts?: Opts];
+
 export function createWorkflowCtx(
   workflowId: WorkflowId,
   sender: BaseChannel<StepRequest>,
@@ -150,7 +158,7 @@ export function createWorkflowCtx(
 }
 
 async function runFunction<
-  F extends FunctionReference<FunctionType, "internal">,
+  F extends FunctionReference<FunctionType, FunctionVisibility>,
 >(
   sender: BaseChannel<StepRequest>,
   functionType: FunctionType,