Cancel "zombie" thread runs
Occasionally, thread runs remain in an "in_progress" state even after
delivering the response. These "zombie" runs obstruct the addition of
new messages to the thread and must be manually cancelled once the run
step has been retrieved.
Change-Id: I967da84bba66ed8b489b5fdc79001b62b8a9ec5d
Signed-off-by: Patrizio <patrizio.gelosi@amarulasolutions.com>
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
index a313d5a..b56aa43 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/UriResourceLocatorStateful.java
@@ -36,4 +36,8 @@
public static String runStepsUri(String threadId, String runId) {
return runRetrieveUri(threadId, runId) + "/steps";
}
+
+ public static String runCancelUri(String threadId, String runId) {
+ return runRetrieveUri(threadId, runId) + "/cancel";
+ }
}
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
index ca3bc3d..f0eb5f8 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptClientStateful.java
@@ -67,13 +67,16 @@
requestBody = chatGptThreadMessage.getAddMessageRequestBody();
log.debug("ChatGPT request body: {}", requestBody);
- return getResponseContentStateful(threadId, chatGptRun);
+ ChatGptResponseContent chatGptResponseContent = getResponseContentStateful(threadId, chatGptRun);
+ chatGptRun.cancelRun();
+
+ return chatGptResponseContent;
}
private ChatGptResponseContent getResponseContentStateful(String threadId, ChatGptRun chatGptRun) {
return switch (chatGptRun.getFirstStepDetails().getType()) {
case TYPE_MESSAGE_CREATION -> retrieveThreadMessage(threadId, chatGptRun);
- case TYPE_TOOL_CALLS -> getResponseContent(chatGptRun.getFirstStep());
+ case TYPE_TOOL_CALLS -> getResponseContent(chatGptRun.getFirstStepToolCalls());
default -> throw new IllegalStateException("Unexpected Step Type in stateful ChatGpt response: " +
chatGptRun);
};
diff --git a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
index 31d537e..1fa4b0d 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/chatgpt/mode/stateful/client/api/chatgpt/ChatGptRun.java
@@ -28,6 +28,7 @@
"cancelling"
));
public static final String COMPLETED_STATUS = "completed";
+ public static final String CANCELLED_STATUS = "cancelled";
private final ChatGptHttpClient httpClient = new ChatGptHttpClient();
private final ChangeSetData changeSetData;
@@ -100,13 +101,35 @@
}
public ChatGptResponseMessage getFirstStepDetails() {
- return stepResponse.getData().get(0).getStepDetails();
+ return getFirstStep().getStepDetails();
}
- public List<ChatGptToolCall> getFirstStep() {
+ public List<ChatGptToolCall> getFirstStepToolCalls() {
return getFirstStepDetails().getToolCalls();
}
+ public void cancelRun() {
+ if (getFirstStep().getStatus().equals(COMPLETED_STATUS)) return;
+
+ Request cancelRequest = getCancelRequest();
+ log.debug("ChatGPT Cancel Run request: {}", cancelRequest);
+ try {
+ String fullResponse = httpClient.execute(cancelRequest);
+ log.debug("ChatGPT Cancel Run Full response: {}", fullResponse);
+ ChatGptResponse response = getGson().fromJson(fullResponse, ChatGptResponse.class);
+ if (!response.getStatus().equals(CANCELLED_STATUS)) {
+ log.error("Unable to cancel run. Run cancel response: {}", fullResponse);
+ }
+ }
+ catch (Exception e) {
+ log.error("Error cancelling run", e);
+ }
+ }
+
+ private ChatGptRunStepsResponse getFirstStep() {
+ return stepResponse.getData().get(0);
+ }
+
private Request runCreateRequest() {
URI uri = URI.create(config.getGptDomain() + UriResourceLocatorStateful.runsUri(threadId));
log.debug("ChatGPT Create Run request URI: {}", uri);
@@ -135,6 +158,14 @@
return getRunPollRequest(uri);
}
+ private Request getCancelRequest() {
+ URI uri = URI.create(config.getGptDomain()
+ + UriResourceLocatorStateful.runCancelUri(threadId, runResponse.getId()));
+ log.debug("ChatGPT Run Cancel request URI: {}", uri);
+
+ return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), new Object());
+ }
+
private Request getRunPollRequest(URI uri) {
return httpClient.createRequestFromJson(uri.toString(), config.getGptToken(), null);
}
diff --git a/src/test/resources/__files/chatGptResponseRequestMessageStateful.json b/src/test/resources/__files/chatGptResponseRequestMessageStateful.json
index 812caa4..37709e9 100644
--- a/src/test/resources/__files/chatGptResponseRequestMessageStateful.json
+++ b/src/test/resources/__files/chatGptResponseRequestMessageStateful.json
@@ -4,6 +4,7 @@
{
"id": "step_UKMPnSinQy6XjSWO7SFleu1v",
"object": "thread.run.step",
+ "status": "completed",
"step_details": {
"type": "message_creation",
"message_creation": {
diff --git a/src/test/resources/__files/chatGptResponseRequestStateful.json b/src/test/resources/__files/chatGptResponseRequestStateful.json
index b66217b..a3c26e4 100644
--- a/src/test/resources/__files/chatGptResponseRequestStateful.json
+++ b/src/test/resources/__files/chatGptResponseRequestStateful.json
@@ -4,6 +4,7 @@
{
"id": "step_UKMPny6XSinQjSWOhSFleu1v",
"object": "thread.run.step",
+ "status": "completed",
"step_details": {
"type": "tool_calls",
"tool_calls": [
diff --git a/src/test/resources/__files/chatGptRunStepsResponse.json b/src/test/resources/__files/chatGptRunStepsResponse.json
index b9b6044..3b92beb 100644
--- a/src/test/resources/__files/chatGptRunStepsResponse.json
+++ b/src/test/resources/__files/chatGptRunStepsResponse.json
@@ -4,6 +4,7 @@
{
"id": "step_TEST_STEP_ID",
"object": "thread.run.step",
+ "status": "completed",
"created_at": 1714983398.0,
"run_id": "run_TEST_RUN_ID",
"assistant_id": "asst_TEST_ASSISTANT_ID",
