From a41e49c60790b154788962540cce5a4e7582e288 Mon Sep 17 00:00:00 2001
From: nightcityblade <nightcityblade@gmail.com>
Date: Tue, 21 Jul 2026 23:36:34 +0800
Subject: [PATCH] fix(langgraph): preserve context during regeneration

Signed-off-by: nightcityblade <nightcityblade@gmail.com>
---
 .../typescript/src/__tests__/agent.test.ts    | 30 +++++++++++++++++++
 .../langgraph/typescript/src/agent.ts         | 18 +++++++++++
 2 files changed, 48 insertions(+)

diff --git a/integrations/langgraph/typescript/src/__tests__/agent.test.ts b/integrations/langgraph/typescript/src/__tests__/agent.test.ts
index a07bed9b3d..e6a56bbfab 100644
--- a/integrations/langgraph/typescript/src/__tests__/agent.test.ts
+++ b/integrations/langgraph/typescript/src/__tests__/agent.test.ts
@@ -148,6 +148,36 @@ async function runPrepareStream(
   ]);
 }
 
+describe("prepareRegenerateStream payload", () => {
+  it("preserves context-schema values when regenerating", async () => {
+    const { agent, capturedPayload } = buildMockedAgent();
+    (agent as any).assistant = MOCK_ASSISTANT;
+    (agent as any).activeRun.schemaKeys = {
+      config: ["model"],
+      context: ["model"],
+    };
+    (agent as any).getCheckpointByMessage = vi.fn().mockResolvedValue({
+      values: { messages: [] },
+      checkpoint: { checkpoint_id: "checkpoint-1" },
+      next: [],
+    });
+    (agent as any).client.threads.updateState.mockResolvedValue({
+      checkpoint: { checkpoint_id: "fork-1" },
+    });
+
+    await agent.prepareRegenerateStream(
+      {
+        threadId: "thread-1",
+        messageCheckpoint: { id: "message-1", type: "human", content: "hi" },
+        forwardedProps: { config: { configurable: { model: "gpt-5" } } },
+      } as any,
+      ["values"],
+    );
+
+    expect(capturedPayload.value?.context).toEqual({ model: "gpt-5" });
+  });
+});
+
 // ─── Part A: prepareStream payload shape ─────────────────────────────────────
 
 describe("prepareStream payload partitioning", () => {
diff --git a/integrations/langgraph/typescript/src/agent.ts b/integrations/langgraph/typescript/src/agent.ts
index 393380405a..7c3ab5ae7f 100644
--- a/integrations/langgraph/typescript/src/agent.ts
+++ b/integrations/langgraph/typescript/src/agent.ts
@@ -478,6 +478,23 @@ export class LangGraphAgent extends AbstractAgent {
       });
     }
 
+    const contextSchemaKeys = new Set(
+      this.activeRun!.schemaKeys?.context ?? [],
+    );
+    const payloadContext = Object.fromEntries(
+      Object.entries(payloadConfig?.configurable ?? {}).filter(([key]) =>
+        contextSchemaKeys.has(key),
+      ),
+    );
+    const hasContext = Object.keys(payloadContext).length > 0;
+    if (hasContext && payloadConfig) {
+      const { configurable: _configurable, ...configWithoutConfigurable } =
+        payloadConfig;
+      payloadConfig = Object.keys(configWithoutConfigurable).length
+        ? configWithoutConfigurable
+        : undefined;
+    }
+
     const payload = {
       ...(input.forwardedProps ?? {}),
       input: this.langGraphDefaultMergeState(
@@ -489,6 +506,7 @@ export class LangGraphAgent extends AbstractAgent {
       checkpointId: fork.checkpoint.checkpoint_id!,
       streamMode,
       config: payloadConfig,
+      ...(hasContext ? { context: payloadContext } : {}),
     };
     return {
       streamResponse: this.client.runs.stream(