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

@@ -0,0 +1,5 @@
+---
+"inngest": patch
+---
+
+Allow editing `step.ai.infer()`

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 = {