@probelabs/probe 0.6.0-rc232 → 0.6.0-rc234

package/build/agent/dsl/sandbox-experiment.mjs

@@ -0,0 +1,309 @@
+/**
+ * Quick experiment to verify SandboxJS capabilities for our DSL runtime.
+ *
+ * Tests:
+ * 1. compileAsync() with host async functions as globals
+ * 2. Error propagation from sandbox to host
+ * 3. Tick limits
+ * 4. Sandbox.audit() for introspection
+ * 5. map() concurrency pattern
+ * 6. Nested async calls (callback inside map that calls async)
+ */
+
+import SandboxModule from '@nyariv/sandboxjs';
+const Sandbox = SandboxModule.default || SandboxModule;
+
+async function test(name, fn) {
+  try {
+    const result = await fn();
+    console.log(`PASS: ${name}`, result !== undefined ? `→ ${JSON.stringify(result)}` : '');
+  } catch (e) {
+    console.log(`FAIL: ${name} → ${e.message}`);
+  }
+}
+
+// Test 1: Basic async function as global
+await test('Host async function as global', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      fetchData: async (query) => {
+        return { results: [`result for: ${query}`] };
+      }
+    }
+  });
+  const exec = s.compileAsync(`
+    const data = await fetchData("test query");
+    return data.results[0];
+  `);
+  const result = await exec().run();
+  if (result !== 'result for: test query') throw new Error(`Expected 'result for: test query', got '${result}'`);
+  return result;
+});
+
+// Test 2: Multiple sequential async calls
+await test('Multiple sequential async calls', async () => {
+  const callLog = [];
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      step: async (n) => {
+        callLog.push(n);
+        return `done-${n}`;
+      }
+    }
+  });
+  const exec = s.compileAsync(`
+    const a = await step(1);
+    const b = await step(2);
+    const c = await step(3);
+    return a + "," + b + "," + c;
+  `);
+  const result = await exec().run();
+  if (callLog.join(',') !== '1,2,3') throw new Error(`Wrong call order: ${callLog}`);
+  return result;
+});
+
+// Test 3: Error propagation
+await test('Error propagation from async global', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      failingTool: async () => {
+        throw new Error('Tool failed!');
+      }
+    }
+  });
+  const exec = s.compileAsync(`
+    const result = await failingTool();
+    return result;
+  `);
+  try {
+    await exec().run();
+    throw new Error('Should have thrown');
+  } catch (e) {
+    if (!e.message.includes('Tool failed')) throw new Error(`Wrong error: ${e.message}`);
+    return 'Error correctly propagated';
+  }
+});
+
+// Test 4: Sandbox.audit()
+await test('Sandbox.audit() reports accessed globals', async () => {
+  const audit = Sandbox.audit(`
+    const x = myFunc("test");
+    const y = otherFunc(x);
+    return y;
+  `);
+  return audit;
+});
+
+// Test 5: Code without await (what LLM would write, before our transform)
+await test('Async global called WITHOUT await', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      fetchData: async (query) => {
+        return { results: [`result for: ${query}`] };
+      }
+    }
+  });
+  // Without await - should return a Promise object, not resolved value
+  const exec = s.compileAsync(`
+    const data = fetchData("test");
+    return data;
+  `);
+  const result = await exec().run();
+  // Check if it's a Promise (unresolved) or the actual value
+  const isPromise = result && typeof result.then === 'function';
+  return { isPromise, type: typeof result, value: isPromise ? 'Promise (unresolved)' : result };
+});
+
+// Test 6: Custom throw for pause-like mechanism (just to verify throws work)
+await test('Custom throw propagation', async () => {
+  class PauseSignal {
+    constructor(value) { this.value = value; this.isPause = true; }
+  }
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      pause: (value) => { throw new PauseSignal(value); }
+    }
+  });
+  const exec = s.compileAsync(`
+    const x = 42;
+    pause({ result: x });
+    return "should not reach here";
+  `);
+  try {
+    await exec().run();
+    throw new Error('Should have thrown');
+  } catch (e) {
+    if (e.isPause) {
+      return { paused: true, value: e.value };
+    }
+    return { paused: false, error: e.message };
+  }
+});
+
+// Test 7: for...of loop with async
+await test('for...of with async calls', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      process: async (item) => item * 2  // Note: 'process' as a name - might conflict
+    }
+  });
+  const exec = s.compileAsync(`
+    const items = [1, 2, 3, 4, 5];
+    const results = [];
+    for (const item of items) {
+      results.push(await process(item));
+    }
+    return results;
+  `);
+  const result = await exec().run();
+  const expected = [2, 4, 6, 8, 10];
+  if (JSON.stringify(result) !== JSON.stringify(expected)) throw new Error(`Got ${JSON.stringify(result)}`);
+  return result;
+});
+
+// Test 8: Passing scope variables
+await test('Scope variables accessible in sandbox', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      transform: async (x) => x.toUpperCase()
+    }
+  });
+  const exec = s.compileAsync(`
+    const result = await transform(inputData);
+    return result;
+  `);
+  const result = await exec({ inputData: 'hello world' }).run();
+  if (result !== 'HELLO WORLD') throw new Error(`Got '${result}'`);
+  return result;
+});
+
+// Test 9: Arrow function callback with async inside
+await test('Arrow function with async call inside', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      processItem: async (item) => item * 10
+    }
+  });
+  const exec = s.compileAsync(`
+    const items = [1, 2, 3];
+    const fn = async (item) => {
+      const result = await processItem(item);
+      return result;
+    };
+    const results = [];
+    for (const item of items) {
+      results.push(await fn(item));
+    }
+    return results;
+  `);
+  const result = await exec().run();
+  if (JSON.stringify(result) !== '[10,20,30]') throw new Error(`Got ${JSON.stringify(result)}`);
+  return result;
+});
+
+// Test 10: map() as a custom global with concurrency
+await test('Custom map() with concurrency control', async () => {
+  let concurrent = 0;
+  let maxConcurrent = 0;
+
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      processItem: async (item) => {
+        concurrent++;
+        maxConcurrent = Math.max(maxConcurrent, concurrent);
+        await new Promise(r => setTimeout(r, 50)); // simulate work
+        concurrent--;
+        return item * 2;
+      },
+      map: async (items, fn) => {
+        const concurrency = 3;
+        const results = [];
+        const executing = new Set();
+        for (const item of items) {
+          const p = fn(item).then(result => {
+            executing.delete(p);
+            return result;
+          });
+          executing.add(p);
+          results.push(p);
+          if (executing.size >= concurrency) {
+            await Promise.race(executing);
+          }
+        }
+        return Promise.all(results);
+      }
+    }
+  });
+  const exec = s.compileAsync(`
+    const items = [1, 2, 3, 4, 5, 6, 7, 8];
+    const results = await map(items, async (item) => {
+      return await processItem(item);
+    });
+    return results;
+  `);
+  const result = await exec().run();
+  const expected = [2, 4, 6, 8, 10, 12, 14, 16];
+  if (JSON.stringify(result) !== JSON.stringify(expected)) throw new Error(`Got ${JSON.stringify(result)}`);
+  return { result, maxConcurrent };
+});
+
+// Test 11: map() called WITHOUT async/await in the callback (what LLM would write)
+await test('map() where LLM writes sync-looking callback', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      processItem: async (item) => item * 2,
+      map: async (items, fn) => {
+        // fn might return a promise even if not declared async
+        const results = [];
+        for (const item of items) {
+          const result = await fn(item);
+          results.push(result);
+        }
+        return results;
+      }
+    }
+  });
+  // LLM writes this - no async, no await in callback
+  const exec = s.compileAsync(`
+    const items = [1, 2, 3];
+    const results = await map(items, (item) => {
+      return processItem(item);
+    });
+    return results;
+  `);
+  const result = await exec().run();
+  if (JSON.stringify(result) !== '[2,4,6]') throw new Error(`Got ${JSON.stringify(result)}`);
+  return result;
+});
+
+// Test 12: Verify blocked globals are truly inaccessible
+await test('Blocked globals not accessible', async () => {
+  const s = new Sandbox({
+    globals: {
+      ...Sandbox.SAFE_GLOBALS,
+      // Deliberately NOT including: require, process, setTimeout, fetch
+    }
+  });
+  const exec = s.compileAsync(`
+    try {
+      const x = setTimeout;
+      return "FAIL: setTimeout accessible";
+    } catch(e) {
+      return "PASS: setTimeout blocked";
+    }
+  `);
+  const result = await exec().run();
+  return result;
+});
+
+console.log('\n--- Experiment complete ---');

package/build/agent/dsl/transformer.js

@@ -0,0 +1,156 @@
+/**
+ * AST Transformer - Auto-injects await before async tool calls.
+ *
+ * The LLM writes synchronous-looking code. This transformer:
+ * 1. Parses the code into an AST
+ * 2. Finds all CallExpressions where the callee is a known async tool function
+ * 3. Inserts `await` before those calls in the source
+ * 4. Marks arrow functions containing async calls as `async`
+ * 5. Wraps the whole program in an async IIFE
+ *
+ * Uses offset-based string insertion (not AST regeneration) to preserve
+ * the original code structure as much as possible.
+ */
+
+import * as acorn from 'acorn';
+import * as walk from 'acorn-walk';
+
+/**
+ * Transform DSL code by injecting await and async wrappers.
+ *
+ * @param {string} code - The sync-looking DSL code
+ * @param {Set<string>} asyncFunctionNames - Names of functions that are async (tool functions)
+ * @returns {string} Transformed code with await injected, wrapped in async IIFE
+ */
+export function transformDSL(code, asyncFunctionNames) {
+  let ast;
+  try {
+    ast = acorn.parse(code, {
+      ecmaVersion: 2022,
+      sourceType: 'script',
+      allowReturnOutsideFunction: true,
+    });
+  } catch (e) {
+    throw new Error(`Transform parse error: ${e.message}`);
+  }
+
+  // Collect insertions: { offset, text } sorted by offset descending
+  // We insert from end to start so offsets don't shift
+  const insertions = [];
+
+  // Track which arrow/function expressions need to be marked async
+  const functionsNeedingAsync = new Set();
+
+  // Find the enclosing function for a given node position
+  function findEnclosingFunction(node) {
+    // Walk the AST to find parent functions
+    // We'll use a different approach: collect all functions and their ranges
+    return null; // Handled by the parent tracking below
+  }
+
+  // First pass: collect all function scopes with their ranges
+  const functionScopes = [];
+  walk.full(ast, (node) => {
+    if (node.type === 'ArrowFunctionExpression' || node.type === 'FunctionExpression') {
+      functionScopes.push(node);
+    }
+  });
+
+  // Second pass: find async calls and determine what needs transformation
+  walk.full(ast, (node) => {
+    if (node.type !== 'CallExpression') return;
+
+    const calleeName = getCalleeName(node);
+    if (!calleeName || !asyncFunctionNames.has(calleeName)) return;
+
+    // This call needs await. Check if it's already awaited.
+    // (It shouldn't be since we block AwaitExpression in the validator,
+    // but be defensive.)
+
+    // Insert 'await ' before the call expression
+    insertions.push({ offset: node.start, text: 'await ' });
+
+    // Find the enclosing function (if any) and mark it as needing async
+    for (const fn of functionScopes) {
+      if (fn.body.start <= node.start && fn.body.end >= node.end) {
+        functionsNeedingAsync.add(fn);
+      }
+    }
+  });
+
+  // Also check: if 'map' is called with a callback that contains async calls,
+  // mark that callback as async. The callback is typically the second argument.
+  walk.full(ast, (node) => {
+    if (node.type !== 'CallExpression') return;
+    const calleeName = getCalleeName(node);
+    if (calleeName !== 'map' || node.arguments.length < 2) return;
+
+    const callback = node.arguments[1];
+    if (callback.type === 'ArrowFunctionExpression' || callback.type === 'FunctionExpression') {
+      // Check if this callback contains any async tool calls
+      let hasAsyncCall = false;
+      walk.full(callback, (inner) => {
+        if (inner.type === 'CallExpression') {
+          const innerName = getCalleeName(inner);
+          if (innerName && asyncFunctionNames.has(innerName)) {
+            hasAsyncCall = true;
+          }
+        }
+      });
+      if (hasAsyncCall) {
+        functionsNeedingAsync.add(callback);
+      }
+    }
+  });
+
+  // Third pass: inject loop guards (__checkLoop()) into while/for loops
+  walk.full(ast, (node) => {
+    if (node.type === 'WhileStatement' || node.type === 'ForStatement' || node.type === 'ForOfStatement' || node.type === 'ForInStatement') {
+      // Insert __checkLoop(); at the start of the loop body
+      const body = node.body;
+      if (body.type === 'BlockStatement' && body.body.length > 0) {
+        // Insert after the opening brace
+        insertions.push({ offset: body.start + 1, text: ' __checkLoop();' });
+      }
+    }
+  });
+
+  // Build insertions for async markers on functions
+  for (const fn of functionsNeedingAsync) {
+    // Insert 'async ' before the function
+    // For arrow functions: `(x) => ...` → `async (x) => ...`
+    // For function expressions: `function(x) { ... }` → `async function(x) { ... }`
+    insertions.push({ offset: fn.start, text: 'async ' });
+  }
+
+  // Sort insertions by offset descending (apply from end to preserve offsets)
+  insertions.sort((a, b) => b.offset - a.offset);
+
+  // Apply insertions to the source code
+  let transformed = code;
+  for (const ins of insertions) {
+    transformed = transformed.slice(0, ins.offset) + ins.text + transformed.slice(ins.offset);
+  }
+
+  // Wrap in async IIFE with return so SandboxJS awaits the result
+  transformed = `return (async () => {\n${transformed}\n})()`;
+
+  return transformed;
+}
+
+/**
+ * Extract the function name from a CallExpression callee.
+ * Handles: `foo()` → 'foo', `obj.foo()` → 'foo' (for member access)
+ *
+ * @param {import('acorn').Node} callExpr
+ * @returns {string|null}
+ */
+function getCalleeName(callExpr) {
+  const callee = callExpr.callee;
+  if (callee.type === 'Identifier') {
+    return callee.name;
+  }
+  // For member expressions like mcp_server.tool(), get the full dotted name
+  // But our tools use flat names like mcp_github_create_issue, so Identifier is sufficient
+  return null;
+}

package/build/agent/dsl/trigger-test.mjs

@@ -0,0 +1,159 @@
+#!/usr/bin/env node
+/**
+ * Trigger test: verifies that the agent picks execute_plan for the right queries.
+ *
+ * Runs the real ProbeAgent with enableExecutePlan=true and observes which tools
+ * get called for different types of questions. This tests the tool-selection
+ * logic end-to-end — the system prompt, tool descriptions, and LLM decision-making.
+ *
+ * Usage:
+ *   node npm/src/agent/dsl/trigger-test.mjs
+ *
+ * Requires: GOOGLE_API_KEY or GOOGLE_GENERATIVE_AI_API_KEY in .env
+ */
+
+import { ProbeAgent } from '../ProbeAgent.js';
+import { config } from 'dotenv';
+import { resolve, dirname } from 'path';
+import { fileURLToPath } from 'url';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const projectRoot = resolve(__dirname, '../../../..');
+
+config({ path: resolve(projectRoot, '.env') });
+
+// Check for API key
+const apiKey = process.env.GOOGLE_GENERATIVE_AI_API_KEY || process.env.GOOGLE_API_KEY;
+if (!apiKey) {
+  console.error('ERROR: No Google API key found. Set GOOGLE_API_KEY or GOOGLE_GENERATIVE_AI_API_KEY');
+  process.exit(1);
+}
+
+// ── Test definitions ──
+// Each test has a query and an expected tool choice
+const tests = [
+  // ── Should trigger execute_plan ──
+  {
+    name: 'Aggregate question (all patterns)',
+    query: 'Find ALL error handling patterns across the entire codebase and give me a comprehensive summary covering every module.',
+    expectTool: 'execute_plan',
+    reason: 'Aggregate question needing full data coverage + "ALL" + "comprehensive" + "every module"',
+  },
+  {
+    name: 'Multi-topic bulk scan',
+    query: 'Search for authentication, authorization, and session management patterns. Analyze each topic across the full codebase and produce a security report.',
+    expectTool: 'execute_plan',
+    reason: 'Multiple topics + full codebase scan + synthesis',
+  },
+  {
+    name: 'Open-ended discovery',
+    query: 'What are all the different testing approaches used in this codebase? Give me a complete inventory.',
+    expectTool: 'execute_plan',
+    reason: 'Open-ended, needs discovery + comprehensive scan',
+  },
+
+  // ── Should NOT trigger execute_plan ──
+  {
+    name: 'Simple search (specific function)',
+    query: 'How does the validateDSL function work?',
+    expectTool: 'search',
+    reason: 'Specific function lookup, 1-2 tool calls',
+  },
+  {
+    name: 'Simple search (single concept)',
+    query: 'What is the timeout configuration for the DSL runtime?',
+    expectTool: 'search',
+    reason: 'Narrow question, single concept',
+  },
+];
+
+// ── Test runner ──
+let testNum = 0;
+let passed = 0;
+let failed = 0;
+
+async function runTriggerTest(test) {
+  testNum++;
+  console.log(`\n${'─'.repeat(70)}`);
+  console.log(`▶ Test ${testNum}: ${test.name}`);
+  console.log(`  Query: "${test.query.substring(0, 100)}${test.query.length > 100 ? '...' : ''}"`);
+  console.log(`  Expected tool: ${test.expectTool}`);
+
+  const toolCalls = [];
+
+  const agent = new ProbeAgent({
+    path: projectRoot,
+    provider: 'google',
+    model: 'gemini-2.5-flash',
+    enableExecutePlan: true,
+    maxIterations: 3, // Only need first few iterations to see what tool gets picked
+  });
+
+  // Listen for tool call events
+  agent.events.on('toolCall', (event) => {
+    if (event.status === 'started') {
+      toolCalls.push(event.name);
+      console.log(`  [tool] ${event.name}`);
+    }
+  });
+
+  await agent.initialize();
+
+  const start = Date.now();
+  try {
+    await agent.answer(test.query);
+  } catch (e) {
+    // May hit maxIterations limit — that's fine, we just want tool selection
+    if (!e.message?.includes('iteration') && !e.message?.includes('cancelled')) {
+      console.log(`  [warn] Agent error: ${e.message?.substring(0, 150)}`);
+    }
+  }
+
+  const elapsed = Date.now() - start;
+  const firstMeaningfulTool = toolCalls.find(t =>
+    t === 'execute_plan' || t === 'analyze_all' || t === 'search' || t === 'query'
+  );
+
+  console.log(`  All tool calls: [${toolCalls.join(', ')}]`);
+  console.log(`  First meaningful tool: ${firstMeaningfulTool || '(none)'}`);
+
+  const toolMatch = firstMeaningfulTool === test.expectTool;
+
+  if (toolMatch) {
+    console.log(`  ✓ PASSED — picked ${firstMeaningfulTool} as expected (${elapsed}ms)`);
+    passed++;
+  } else {
+    console.log(`  ✗ FAILED — expected ${test.expectTool}, got ${firstMeaningfulTool || '(none)'} (${elapsed}ms)`);
+    console.log(`  Reason it should use ${test.expectTool}: ${test.reason}`);
+    failed++;
+  }
+
+  try {
+    await agent.close();
+  } catch (e) {
+    // ignore cleanup errors
+  }
+}
+
+async function main() {
+  console.log('═'.repeat(70));
+  console.log('  Execute Plan Trigger Tests — Tool Selection Verification');
+  console.log('═'.repeat(70));
+  console.log(`\nRunning with: enableExecutePlan=true, provider=google, model=gemini-2.5-flash`);
+  console.log(`Project root: ${projectRoot}`);
+
+  for (const test of tests) {
+    await runTriggerTest(test);
+  }
+
+  console.log(`\n${'═'.repeat(70)}`);
+  console.log(`  Results: ${passed} passed, ${failed} failed, ${testNum} total`);
+  console.log('═'.repeat(70));
+
+  process.exit(failed > 0 ? 1 : 0);
+}
+
+main().catch(e => {
+  console.error('Fatal error:', e);
+  process.exit(1);
+});