fix local codex review

Author: seratch

packages/agents-core/src/run.ts

@@ -1348,33 +1348,37 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
     };
     const session = effectiveOptions.session;
     const resumingFromState = input instanceof RunState;
-    let sessionOriginalInput: string | AgentInputItem[] | undefined;
-    if (resumingFromState) {
-      if (session) {
-        sessionOriginalInput = [] as AgentInputItem[];
-      }
-    } else {
-      sessionOriginalInput = input as string | AgentInputItem[];
-    }
+    let sessionInputItemsToPersist: AgentInputItem[] | undefined =
+      session && resumingFromState ? [] : undefined;
 
     let preparedInput: typeof input = input;
     if (!(preparedInput instanceof RunState)) {
-      preparedInput = await prepareInputItemsWithSession(
+      const prepared = await prepareInputItemsWithSession(
         preparedInput,
         session,
         sessionInputCallback,
       );
+      preparedInput = prepared.preparedInput;
+      if (session) {
+        const items = prepared.sessionItems ?? [];
+        // Clone the items that will be persisted so later mutations (filters, hooks) cannot desync history.
+        sessionInputItemsToPersist = items.map((item) => structuredClone(item));
+      }
     }
 
     let ensureStreamInputPersisted: (() => Promise<void>) | undefined;
-    if (session && typeof sessionOriginalInput !== 'undefined') {
+    if (
+      session &&
+      sessionInputItemsToPersist &&
+      sessionInputItemsToPersist.length > 0
+    ) {
       let persisted = false;
       ensureStreamInputPersisted = async () => {
         if (persisted) {
           return;
         }
         persisted = true;
-        await saveStreamInputToSession(session, sessionOriginalInput);
+        await saveStreamInputToSession(session, sessionInputItemsToPersist);
       };
     }
 
@@ -1393,7 +1397,7 @@ export class Runner extends RunHooks<any, AgentOutputType<unknown>> {
         preparedInput,
         effectiveOptions,
       );
-      await saveToSession(session, sessionOriginalInput, runResult);
+      await saveToSession(session, sessionInputItemsToPersist, runResult);
       return runResult;
     };
 

packages/agents-core/src/runImplementation.ts

@@ -1737,13 +1737,13 @@ export function extractOutputItemsFromRunItems(
  */
 export async function saveToSession(
   session: Session | undefined,
-  originalInput: string | AgentInputItem[] | undefined,
+  sessionInputItems: AgentInputItem[] | undefined,
   result: RunResult<any, any>,
 ): Promise<void> {
   if (!session) {
     return;
   }
-  const inputItems = originalInput ? toInputItemList(originalInput) : [];
+  const inputItems = sessionInputItems ?? [];
   const state = result.state;
   const alreadyPersisted = state._currentTurnPersistedItemCount ?? 0;
   // Persist only the portion of _generatedItems that has not yet been stored for this turn.
@@ -1771,19 +1771,15 @@ export async function saveToSession(
  */
 export async function saveStreamInputToSession(
   session: Session | undefined,
-  originalInput: string | AgentInputItem[] | undefined,
+  sessionInputItems: AgentInputItem[] | undefined,
 ): Promise<void> {
   if (!session) {
     return;
   }
-  if (!originalInput) {
+  if (!sessionInputItems || sessionInputItems.length === 0) {
     return;
   }
-  const itemsToSave = toInputItemList(originalInput);
-  if (itemsToSave.length === 0) {
-    return;
-  }
-  await session.addItems(itemsToSave);
+  await session.addItems(sessionInputItems);
 }
 
 /**
@@ -1814,13 +1810,21 @@ export async function saveStreamResultToSession(
  * @internal
  * If a session is provided, expands the input with session history; otherwise returns the input.
  */
+export type PreparedInputWithSessionResult = {
+  preparedInput: string | AgentInputItem[];
+  sessionItems?: AgentInputItem[];
+};
+
 export async function prepareInputItemsWithSession(
   input: string | AgentInputItem[],
   session?: Session,
   sessionInputCallback?: SessionInputCallback,
-): Promise<string | AgentInputItem[]> {
+): Promise<PreparedInputWithSessionResult> {
   if (!session) {
-    return input;
+    return {
+      preparedInput: input,
+      sessionItems: undefined,
+    };
   }
 
   const history = await session.getItems();
@@ -1839,7 +1843,10 @@ export async function prepareInputItemsWithSession(
   }
 
   if (!sessionInputCallback) {
-    return [...history, ...newInputItems];
+    return {
+      preparedInput: [...history, ...newInputItems],
+      sessionItems: newInputItems,
+    };
   }
 
   // Delegate history reconciliation to the user-supplied callback. It must return a concrete list
@@ -1851,5 +1858,26 @@ export async function prepareInputItemsWithSession(
     );
   }
 
-  return combined;
+  const historySet = new Set(history);
+  const newInputSet = new Set(newInputItems);
+  const appended: AgentInputItem[] = [];
+  for (const item of combined) {
+    if (historySet.has(item) || !newInputSet.has(item)) {
+      if (!historySet.has(item) && !newInputSet.has(item)) {
+        appended.push(item);
+      }
+      continue;
+    }
+    appended.push(item);
+  }
+
+  if (appended.length === 0 && combined.length > history.length) {
+    // When callbacks replace every new item with fresh objects, fall back to the tail slice.
+    appended.push(...combined.slice(history.length));
+  }
+
+  return {
+    preparedInput: combined,
+    sessionItems: appended.length > 0 ? appended : [],
+  };
 }

packages/agents-core/test/run.stream.test.ts

@@ -714,7 +714,16 @@ describe('Runner.run (streaming)', () => {
     await result.completed;
 
     expect(saveInputSpy).toHaveBeenCalledTimes(1);
-    expect(saveInputSpy).toHaveBeenCalledWith(session, 'hello world');
+    const [sessionArg, persistedItems] = saveInputSpy.mock.calls[0];
+    expect(sessionArg).toBe(session);
+    if (!Array.isArray(persistedItems)) {
+      throw new Error('Expected persisted session items to be an array.');
+    }
+    expect(persistedItems).toHaveLength(1);
+    expect(persistedItems[0]).toMatchObject({
+      role: 'user',
+      content: 'hello world',
+    });
   });
 
   it('skips persisting streaming input when an input guardrail triggers', async () => {

packages/agents-core/test/run.test.ts

@@ -775,6 +775,44 @@ describe('Runner.run', () => {
         expect(getFirstTextContent(recordedInput[1])).toBe('Fresh input');
       });
 
+      it('persists transformed session input from callback', async () => {
+        const model = new RecordingModel([
+          {
+            ...TEST_MODEL_RESPONSE_BASIC,
+            output: [fakeModelMessage('session response')],
+          },
+        ]);
+        const agent = new Agent({ name: 'SessionTransform', model });
+        const session = new MemorySession();
+        const runner = new Runner();
+        const original = 'Sensitive payload';
+        const redacted = '[redacted]';
+
+        await runner.run(agent, original, {
+          session,
+          sessionInputCallback: (history, newItems) => {
+            expect(history).toHaveLength(0);
+            if (newItems[0] && typeof newItems[0] === 'object') {
+              (newItems[0] as protocol.UserMessageItem).content = redacted;
+            }
+            return history.concat(newItems);
+          },
+        });
+
+        const recordedInput = model.lastRequest?.input as AgentInputItem[];
+        expect(recordedInput[recordedInput.length - 1]).toMatchObject({
+          role: 'user',
+          content: redacted,
+        });
+
+        expect(session.added).toHaveLength(1);
+        const persistedTurn = session.added[0];
+        expect(persistedTurn[0]).toMatchObject({
+          role: 'user',
+          content: redacted,
+        });
+      });
+
       it('throws when session input callback returns invalid data', async () => {
         const model = new RecordingModel([
           {

packages/agents-core/test/runImplementation.test.ts

@@ -32,6 +32,7 @@ import {
   streamStepItemsToRunResult,
   saveToSession,
   executeInterruptedToolsAndSideEffects,
+  toInputItemList,
 } from '../src/runImplementation';
 import {
   FunctionTool,
@@ -195,7 +196,11 @@ describe('saveToSession', () => {
     state._generatedItems = [new ToolApprovalItem(functionCall, textAgent)];
 
     const preApprovalResult = new RunResult(state);
-    await saveToSession(session, state._originalInput, preApprovalResult);
+    await saveToSession(
+      session,
+      toInputItemList(state._originalInput),
+      preApprovalResult,
+    );
 
     expect(session.items).toEqual([
       {