fix(agent): store parsed values and enforce object refinements in context setters

.changeset/tool-context-store-parsed.md

@@ -0,0 +1,5 @@
+---
+"@openrouter/agent": patch
+---
+
+`setContext` and `setSharedContext` now store the schema's parsed output for each updated field, so field transforms, defaults and coercions apply to stored context, and they validate the merged context against the full schema before storing it, so object-level refinements are enforced at write time. An update that fails validation throws and leaves the stored context unchanged.

packages/agent/src/lib/tool-context.ts

@@ -108,29 +108,32 @@ export class ToolContextStore {
 
 /**
  * Validate a partial update against a schema's shape, filtering to known keys
- * and validating each key individually. Returns the filtered partial.
+ * and parsing each key individually. Returns the parsed partial, so field
+ * transforms, defaults and coercions apply to the stored value. The merged
+ * result (existing values plus the update) is then validated as a whole so
+ * object-level refinements are enforced before the update is stored.
  */
 function validatePartialAgainstSchema(
   partial: Record<string, unknown>,
   schema: $ZodObject<$ZodShape>,
+  existing: Record<string, unknown>,
 ): Record<string, unknown> {
-  const schemaKeys = Object.keys(schema._zod.def.shape);
-  const filteredPartial: Record<string, unknown> = {};
-  for (const [key, value] of Object.entries(partial)) {
-    if (schemaKeys.includes(key)) {
-      filteredPartial[key] = value;
-    }
-  }
-
   const shape = schema._zod.def.shape;
-  for (const [key, value] of Object.entries(filteredPartial)) {
+  const schemaKeys = Object.keys(shape);
+  const parsedPartial: Record<string, unknown> = {};
+  for (const [key, value] of Object.entries(partial)) {
     const keySchema = shape[key];
-    if (keySchema) {
-      z4.parse(keySchema, value);
+    if (schemaKeys.includes(key) && keySchema) {
+      parsedPartial[key] = z4.parse(keySchema, value);
     }
   }
 
-  return filteredPartial;
+  z4.parse(schema, {
+    ...existing,
+    ...parsedPartial,
+  });
+
+  return parsedPartial;
 }
 
 /**
@@ -238,11 +241,12 @@ export function buildToolExecuteContext<
       if (!store || !schema) {
         return;
       }
-      const filteredPartial = validatePartialAgainstSchema(
+      const parsedPartial = validatePartialAgainstSchema(
         partial as Record<string, unknown>,
         schema,
+        store.getToolContext(toolName),
       );
-      store.mergeToolContext(toolName, filteredPartial);
+      store.mergeToolContext(toolName, parsedPartial);
     },
 
     get shared(): Readonly<TShared> {
@@ -254,11 +258,12 @@ export function buildToolExecuteContext<
       if (!store || !sharedSchema) {
         return;
       }
-      const filteredPartial = validatePartialAgainstSchema(
+      const parsedPartial = validatePartialAgainstSchema(
         partial as Record<string, unknown>,
         sharedSchema,
+        store.getToolContext(SHARED_CONTEXT_KEY),
       );
-      store.mergeToolContext(SHARED_CONTEXT_KEY, filteredPartial);
+      store.mergeToolContext(SHARED_CONTEXT_KEY, parsedPartial);
     },
   };
 

packages/agent/tests/unit/shared-context.test.ts

@@ -336,3 +336,115 @@ describe('SHARED_CONTEXT_KEY', () => {
 });
 
 //#endregion
+
+//#region Context setters store schema-validated values
+
+describe('context setters store schema-validated values', () => {
+  const turnContext: TurnContext = {
+    numberOfTurns: 1,
+  };
+
+  it('setSharedContext stores the parsed output of field transforms', () => {
+    const store = new ToolContextStore({
+      [SHARED_CONTEXT_KEY]: {
+        level: 0,
+      },
+    });
+    const sharedSchema = z4.object({
+      level: z4.number().transform((n) => Math.min(n, 0)),
+    });
+
+    const ctx = buildToolExecuteContext<
+      'my_tool',
+      Record<string, unknown>,
+      {
+        level: number;
+      }
+    >(turnContext, store, 'my_tool', undefined, sharedSchema);
+
+    ctx.setSharedContext({
+      level: 1,
+    });
+
+    expect(ctx.shared).toEqual({
+      level: 0,
+    });
+  });
+
+  it('setContext stores the parsed output of field transforms', () => {
+    const store = new ToolContextStore({
+      my_tool: {
+        mode: 'read',
+      },
+    });
+    const schema = z4.object({
+      mode: z4.string().transform((s) => s.trim().toLowerCase()),
+    });
+
+    const ctx = buildToolExecuteContext<
+      'my_tool',
+      {
+        mode: string;
+      }
+    >(turnContext, store, 'my_tool', schema);
+
+    ctx.setContext({
+      mode: '  WRITE ',
+    });
+
+    expect(ctx.local).toEqual({
+      mode: 'write',
+    });
+  });
+
+  it('setSharedContext rejects an update that breaks an object-level refinement', () => {
+    const initial = {
+      role: 'guest',
+      access: 'read',
+    };
+    const store = new ToolContextStore({
+      [SHARED_CONTEXT_KEY]: initial,
+    });
+    const sharedSchema = z4
+      .object({
+        role: z4.enum([
+          'guest',
+          'staff',
+        ]),
+        access: z4.enum([
+          'read',
+          'write',
+        ]),
+      })
+      .refine((v) => v.access !== 'write' || v.role === 'staff');
+
+    const ctx = buildToolExecuteContext<
+      'my_tool',
+      Record<string, unknown>,
+      {
+        role: 'guest' | 'staff';
+        access: 'read' | 'write';
+      }
+    >(turnContext, store, 'my_tool', undefined, sharedSchema);
+
+    expect(() =>
+      ctx.setSharedContext({
+        access: 'write',
+      }),
+    ).toThrow();
+    expect(ctx.shared).toEqual(initial);
+
+    ctx.setSharedContext({
+      role: 'staff',
+    });
+    ctx.setSharedContext({
+      access: 'write',
+    });
+    expect(ctx.shared).toEqual({
+      role: 'staff',
+      access: 'write',
+    });
+  });
+});
+
+//#endregion