feat: implement patchGlobals method for AxCodeSession and update related functionality

Author: dosco

docs/AXAGENT.md

@@ -731,7 +731,7 @@ const analyzer = agent('query:string -> answer:string', {
 });
 ```
 
-Updates from this callback are merged into current inputs (unknown keys are ignored), then synchronized into runtime `inputs.<field>` and existing non-colliding top-level aliases before code execution.
+Updates from this callback are merged into current inputs (unknown keys are ignored), then synchronized into runtime `inputs.<field>` and existing non-colliding top-level aliases via `AxCodeSession.patchGlobals(...)` before code execution. This host-side sync does not run through the Actor's `execute(code)` path.
 
 ### Actor/Responder Forward Options
 
@@ -922,11 +922,37 @@ class MyBrowserInterpreter implements AxCodeRuntime {
   }
 
   createSession(globals?: Record<string, unknown>): AxCodeSession {
+    const scope = { ...globals };
+    const isPlainObject = (
+      value: unknown
+    ): value is Record<string, unknown> => {
+      if (!value || typeof value !== 'object' || Array.isArray(value)) {
+        return false;
+      }
+      const proto = Object.getPrototypeOf(value);
+      return proto === Object.prototype || proto === null;
+    };
+
     // Set up your execution environment with globals
     return {
       async execute(code: string) {
         // Execute code and return result
       },
+      async patchGlobals(nextGlobals: Record<string, unknown>) {
+        for (const [key, value] of Object.entries(nextGlobals)) {
+          const current = scope[key];
+          if (isPlainObject(current) && isPlainObject(value)) {
+            for (const existingKey of Object.keys(current)) {
+              if (!(existingKey in value)) {
+                delete current[existingKey];
+              }
+            }
+            Object.assign(current, value);
+            continue;
+          }
+          scope[key] = value;
+        }
+      },
       close() {
         // Clean up resources
       },
@@ -935,6 +961,8 @@ class MyBrowserInterpreter implements AxCodeRuntime {
 }
 ```
 
+When patching object-valued globals such as `inputs`, reconcile the existing object in place instead of blindly replacing the reference. That keeps previously saved references in the runtime session aligned with later host-side updates.
+
 The `globals` object passed to `createSession` includes:
 - All context field values (by field name)
 - `llmQuery` function (supports both single and batched queries)
@@ -1020,7 +1048,14 @@ interface AxCodeRuntime {
 
 ```typescript
 interface AxCodeSession {
-  execute(code: string, options?: { signal?: AbortSignal });
+  execute(
+    code: string,
+    options?: { signal?: AbortSignal; reservedNames?: readonly string[] }
+  ): Promise<unknown>;
+  patchGlobals(
+    globals: Record<string, unknown>,
+    options?: { signal?: AbortSignal }
+  ): Promise<void>;
   close(): void;
 }
 ```

src/ax/ai/google-gemini/agent-flow.test.ts

@@ -1,8 +1,7 @@
 import { describe, expect, it } from 'vitest';
-
-import { AxMockAIService } from '../mock/api.js';
 import { agent, s } from '../../index.js';
 import type { AxCodeRuntime } from '../../prompts/rlm.js';
+import { AxMockAIService } from '../mock/api.js';
 
 const makeModelUsage = () => ({
   ai: 'mock',
@@ -119,6 +118,7 @@ describe('Agent Split Architecture Flow', () => {
             }
             return `executed: ${code}`;
           },
+          patchGlobals: async () => {},
           close: () => {},
         };
       },

src/ax/funcs/jsRuntime.integration.test.ts

@@ -698,4 +698,94 @@ describe('AxJSRuntime integration', () => {
       session.close();
     }
   });
+
+  it('patchGlobals overwrites scalars and reconciles plain objects in place', async () => {
+    const runtime = new AxJSRuntime({ outputMode: 'return' });
+    const session = runtime.createSession({
+      query: 'initial',
+      tags: ['old'],
+      inputs: { query: 'initial', stale: 'old' },
+    });
+
+    try {
+      await session.execute('globalThis.savedInputs = inputs');
+      await session.patchGlobals({
+        query: 'updated',
+        tags: ['fresh'],
+        note: undefined,
+        inputs: { query: 'updated' },
+        helper: {
+          describe: async () => 'patched helper',
+        },
+      });
+
+      const result = await session.execute(`
+        ({
+          query,
+          tags,
+          note,
+          savedInputsSame: globalThis.savedInputs === inputs,
+          savedInputsQuery: globalThis.savedInputs.query,
+          inputsKeys: Object.keys(inputs).sort(),
+          helperResult: await helper.describe()
+        })
+      `);
+
+      expect(result).toEqual({
+        query: 'updated',
+        tags: ['fresh'],
+        note: undefined,
+        savedInputsSame: true,
+        savedInputsQuery: 'updated',
+        inputsKeys: ['query'],
+        helperResult: 'patched helper',
+      });
+    } finally {
+      session.close();
+    }
+  });
+
+  it('patchGlobals survives worker reset and recreation', async () => {
+    const runtime = new AxJSRuntime({ outputMode: 'return', timeout: 100 });
+    const session = runtime.createSession({
+      query: 'initial',
+      inputs: { query: 'initial' },
+    });
+
+    try {
+      await session.patchGlobals({
+        query: 'updated',
+        inputs: { query: 'updated' },
+      });
+
+      await expect(session.execute('while(true){}')).rejects.toThrow(
+        'Execution timed out'
+      );
+
+      const result = await session.execute(
+        '({ query, inputQuery: inputs.query })'
+      );
+      expect(result).toEqual({ query: 'updated', inputQuery: 'updated' });
+    } finally {
+      session.close();
+    }
+  });
+
+  it('patchGlobals respects abort signals', async () => {
+    const runtime = new AxJSRuntime({ outputMode: 'return' });
+    const session = runtime.createSession();
+    const controller = new AbortController();
+    controller.abort('stop patching');
+
+    try {
+      await expect(
+        session.patchGlobals(
+          { query: 'updated' },
+          { signal: controller.signal }
+        )
+      ).rejects.toThrow('Aborted: stop patching');
+    } finally {
+      session.close();
+    }
+  });
 });