From ecb584cbe0a8b4a15ed1b7649da58cc199cffd74 Mon Sep 17 00:00:00 2001
From: Andrew Nguonly <andrewnguonly@users.noreply.github.com>
Date: Fri, 1 Nov 2024 13:32:49 -0700
Subject: [PATCH] sdk: Add `action` query param to cancel Run methods (#2284)

---
 libs/sdk-js/src/client.ts           |  4 ++++
 libs/sdk-js/src/schema.ts           |  2 ++
 libs/sdk-py/langgraph_sdk/client.py | 33 +++++++++++++++++++++++------
 libs/sdk-py/langgraph_sdk/schema.py |  7 ++++++
 4 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/libs/sdk-js/src/client.ts b/libs/sdk-js/src/client.ts
index 32e75017b..6139e8c41 100644
--- a/libs/sdk-js/src/client.ts
+++ b/libs/sdk-js/src/client.ts
@@ -1,6 +1,7 @@
 import {
   Assistant,
   AssistantGraph,
+  CancelAction,
   Config,
   DefaultValues,
   GraphSchema,
@@ -928,17 +929,20 @@ export class RunsClient extends BaseClient {
    * @param threadId The ID of the thread.
    * @param runId The ID of the run.
    * @param wait Whether to block when canceling
+   * @param action Action to take when cancelling the run. Possible values are `interrupt` or `rollback`. Default is `interrupt`.
    * @returns
    */
   async cancel(
     threadId: string,
     runId: string,
     wait: boolean = false,
+    action: CancelAction = "interrupt",
   ): Promise<void> {
     return this.fetch<void>(`/threads/${threadId}/runs/${runId}/cancel`, {
       method: "POST",
       params: {
         wait: wait ? "1" : "0",
+        action: action,
       },
     });
   }
diff --git a/libs/sdk-js/src/schema.ts b/libs/sdk-js/src/schema.ts
index c2de251f6..86a668d01 100644
--- a/libs/sdk-js/src/schema.ts
+++ b/libs/sdk-js/src/schema.ts
@@ -14,6 +14,8 @@ type ThreadStatus = "idle" | "busy" | "interrupted" | "error";
 
 type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
 
+export type CancelAction = "interrupt" | "rollback";
+
 export interface Config {
   /**
    * Tags for this call and any sub-calls (eg. a Chain calling an LLM).
diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py
index 8118d6a5c..a057d452f 100644
--- a/libs/sdk-py/langgraph_sdk/client.py
+++ b/libs/sdk-py/langgraph_sdk/client.py
@@ -34,6 +34,7 @@ from langgraph_sdk.schema import (
     All,
     Assistant,
     AssistantVersion,
+    CancelAction,
     Checkpoint,
     Config,
     Cron,
@@ -1712,13 +1713,22 @@ class RunsClient:
 
         return await self.http.get(f"/threads/{thread_id}/runs/{run_id}")
 
-    async def cancel(self, thread_id: str, run_id: str, *, wait: bool = False) -> None:
+    async def cancel(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        wait: bool = False,
+        action: CancelAction = "interrupt",
+    ) -> None:
         """Get a run.
 
         Args:
             thread_id: The thread ID to cancel.
             run_id: The run ID to cancek.
             wait: Whether to wait until run has completed.
+            action: Action to take when cancelling the run. Possible values
+                are `interrupt` or `rollback`. Default is `interrupt`.
 
         Returns:
             None
@@ -1728,12 +1738,13 @@ class RunsClient:
             await client.runs.cancel(
                 thread_id="thread_id_to_cancel",
                 run_id="run_id_to_cancel",
-                wait=True
+                wait=True,
+                action="interrupt"
             )
 
         """  # noqa: E501
         return await self.http.post(
-            f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}",
+            f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}&action={action}",
             json=None,
         )
 
@@ -3792,13 +3803,22 @@ class SyncRunsClient:
 
         return self.http.get(f"/threads/{thread_id}/runs/{run_id}")
 
-    def cancel(self, thread_id: str, run_id: str, *, wait: bool = False) -> None:
+    def cancel(
+        self,
+        thread_id: str,
+        run_id: str,
+        *,
+        wait: bool = False,
+        action: CancelAction = "interrupt",
+    ) -> None:
         """Get a run.
 
         Args:
             thread_id: The thread ID to cancel.
             run_id: The run ID to cancek.
             wait: Whether to wait until run has completed.
+            action: Action to take when cancelling the run. Possible values
+                are `interrupt` or `rollback`. Default is `interrupt`.
 
         Returns:
             None
@@ -3808,12 +3828,13 @@ class SyncRunsClient:
             client.runs.cancel(
                 thread_id="thread_id_to_cancel",
                 run_id="run_id_to_cancel",
-                wait=True
+                wait=True,
+                action="interrupt"
             )
 
         """  # noqa: E501
         return self.http.post(
-            f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}",
+            f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}&action={action}",
             json=None,
         )
 
diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py
index 43e218669..9583a1c1b 100644
--- a/libs/sdk-py/langgraph_sdk/schema.py
+++ b/libs/sdk-py/langgraph_sdk/schema.py
@@ -78,6 +78,13 @@ Specifies behavior if the thread doesn't exist:
 - "reject": Reject the operation if the thread doesn't exist.
 """
 
+CancelAction = Literal["interrupt", "rollback"]
+"""
+Action to take when cancelling the run.
+- "interrupt": Simply cancel the run.
+- "rollback": Cancel the run. Then delete the run and associated checkpoints.
+"""
+
 
 class Config(TypedDict, total=False):
     """Configuration options for a call."""