From 290ad29d26696c9dac8cb5ab50ca75e0b3d903fc Mon Sep 17 00:00:00 2001
From: Jack Williams <jack@inngest.com>
Date: Thu, 21 Nov 2024 17:52:40 +0000
Subject: [PATCH] Allow editing `step.ai.infer()` (#752)

## Summary
<!-- Succinctly describe your change, providing context, what you've
changed, and why. -->

Allows editing `step.ai.infer()` steps.

## Checklist
<!-- Tick these items off as you progress. -->
<!-- If an item isn't applicable, ideally please strikeout the item by
wrapping it in "~~"" and suffix it with "N/A My reason for skipping
this." -->
<!-- e.g. "- [ ] ~~Added tests~~ N/A Only touches docs" -->

- [ ] ~Added a [docs PR](https://github.com/inngest/website) that
references this PR~ N/A
- [ ] ~Added unit/integration tests~ N/A
- [x] Added changesets if applicable
---
 .changeset/late-planes-travel.md              |  5 ++++
 .../inngest/src/components/execution/v1.ts    | 28 +++++++++++++++----
 2 files changed, 28 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/late-planes-travel.md

diff --git a/.changeset/late-planes-travel.md b/.changeset/late-planes-travel.md
new file mode 100644
index 000000000..315d88405
--- /dev/null
+++ b/.changeset/late-planes-travel.md
@@ -0,0 +1,5 @@
+---
+"inngest": patch
+---
+
+Allow editing `step.ai.infer()`
diff --git a/packages/inngest/src/components/execution/v1.ts b/packages/inngest/src/components/execution/v1.ts
index 09f7004e8..61b348f85 100644
--- a/packages/inngest/src/components/execution/v1.ts
+++ b/packages/inngest/src/components/execution/v1.ts
@@ -870,14 +870,32 @@ class V1InngestExecution extends InngestExecution implements IInngestExecution {
       let fnArgs = [...args];
 
       if (
-        opId.op === StepOpCode.StepPlanned &&
         typeof stepState?.input !== "undefined" &&
         Array.isArray(stepState.input)
       ) {
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-        fnArgs = [...args.slice(0, 2), ...stepState.input];
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-        extraOpts = { input: [...stepState.input] };
+        switch (opId.op) {
+          // `step.run()` has its function input affected
+          case StepOpCode.StepPlanned: {
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            fnArgs = [...args.slice(0, 2), ...stepState.input];
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+            extraOpts = { input: [...stepState.input] };
+            break;
+          }
+
+          // `step.ai.infer()` has its body affected
+          case StepOpCode.AiGateway: {
+            extraOpts = {
+              body: {
+                ...(typeof opId.opts?.body === "object"
+                  ? { ...opId.opts.body }
+                  : {}),
+                ...stepState.input[0],
+              },
+            };
+            break;
+          }
+        }
       }
 
       const step: FoundStep = {