Merge pull request #655 from thecodacus/streaming-fixed

fix: Add Code Streaming Sampling for Performance Optimization

Author: thecodacus

app/lib/stores/workbench.ts

@@ -16,6 +16,7 @@ import * as nodePath from 'node:path';
 import { extractRelativePath } from '~/utils/diff';
 import { description } from '~/lib/persistence';
 import Cookies from 'js-cookie';
+import { createSampler } from '~/utils/sampler';
 
 export interface ArtifactState {
   id: string;
@@ -280,7 +281,7 @@ export class WorkbenchStore {
 
   runAction(data: ActionCallbackData, isStreaming: boolean = false) {
     if (isStreaming) {
-      this._runAction(data, isStreaming);
+      this.actionStreamSampler(data, isStreaming);
     } else {
       this.addToExecutionQueue(() => this._runAction(data, isStreaming));
     }
@@ -296,7 +297,8 @@ export class WorkbenchStore {
 
     const action = artifact.runner.actions.get()[data.actionId];
 
-    if (action.executed) {
+
+    if (!action || action.executed) {
       return;
     }
 
@@ -329,6 +331,10 @@ export class WorkbenchStore {
     }
   }
 
+  actionStreamSampler = createSampler(async (data: ActionCallbackData, isStreaming: boolean = false) => {
+    return await this._runAction(data, isStreaming);
+  }, 100); // TODO: remove this magic number to have it configurable
+
   #getArtifact(id: string) {
     const artifacts = this.artifacts.get();
     return artifacts[id];

app/utils/sampler.ts

@@ -0,0 +1,49 @@
+/**
+ * Creates a function that samples calls at regular intervals and captures trailing calls.
+ * - Drops calls that occur between sampling intervals
+ * - Takes one call per sampling interval if available
+ * - Captures the last call if no call was made during the interval
+ *
+ * @param fn The function to sample
+ * @param sampleInterval How often to sample calls (in ms)
+ * @returns The sampled function
+ */
+export function createSampler<T extends (...args: any[]) => any>(fn: T, sampleInterval: number): T {
+  let lastArgs: Parameters<T> | null = null;
+  let lastTime = 0;
+  let timeout: NodeJS.Timeout | null = null;
+
+  // Create a function with the same type as the input function
+  const sampled = function (this: any, ...args: Parameters<T>) {
+    const now = Date.now();
+    lastArgs = args;
+
+    // If we're within the sample interval, just store the args
+    if (now - lastTime < sampleInterval) {
+      // Set up trailing call if not already set
+      if (!timeout) {
+        timeout = setTimeout(
+          () => {
+            timeout = null;
+            lastTime = Date.now();
+
+            if (lastArgs) {
+              fn.apply(this, lastArgs);
+              lastArgs = null;
+            }
+          },
+          sampleInterval - (now - lastTime),
+        );
+      }
+
+      return;
+    }
+
+    // If we're outside the interval, execute immediately
+    lastTime = now;
+    fn.apply(this, args);
+    lastArgs = null;
+  } as T;
+
+  return sampled;
+}