From f6bcc2b34ae33df8b870d41aeef9f3a8d62d9699 Mon Sep 17 00:00:00 2001
From: Simon Leee <simonlee@tw.ibm.com>
Date: Fri, 26 Apr 2024 23:30:01 +0800
Subject: [PATCH] CustomAgent : add timeout for big agent Write : check endLine
 value for empty file

---
 docs/agent/custom-ai-agent.md                                | 5 +++--
 .../unitmesh/devti/language/compiler/exec/WriteInsCommand.kt | 2 +-
 .../kotlin/cc/unitmesh/devti/agent/CustomAgentExecutor.kt    | 3 ++-
 .../cc/unitmesh/devti/agent/model/CustomAgentConfig.kt       | 1 +
 .../unitmesh/devti/custom/schema/autodev-custom-agent.json   | 3 +++
 5 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/docs/agent/custom-ai-agent.md b/docs/agent/custom-ai-agent.md
index 4576adbad3..bcb2c91652 100644
--- a/docs/agent/custom-ai-agent.md
+++ b/docs/agent/custom-ai-agent.md
@@ -50,9 +50,10 @@ Demo:
   },
   {
     "name": "DevInInsert",
-    "description": "Update",
+    "description": "Update，並指定20秒的timeout時間",
     "url": "http://127.0.0.1:8765/api/agent/devins-sample",
-    "responseAction": "DevIns"
+    "responseAction": "DevIns",
+    "defaultTimeout": 20
   }
 ]
 ```
diff --git a/exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/exec/WriteInsCommand.kt b/exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/exec/WriteInsCommand.kt
index 83056f057c..a51b786b56 100644
--- a/exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/exec/WriteInsCommand.kt
+++ b/exts/devins-lang/src/main/kotlin/cc/unitmesh/devti/language/compiler/exec/WriteInsCommand.kt
@@ -87,7 +87,7 @@ class WriteInsCommand(val myProject: Project, val argument: String, val content:
         } ?: return "$DEVINS_ERROR: File not found: $argument"
 
         val startLine = range?.startLine ?: 0
-        val endLine = range?.endLine ?: document.lineCount
+        val endLine = if (document.lineCount == 0) 1 else range?.endLine ?: document.lineCount
 
         try {
             val startOffset = document.getLineStartOffset(startLine)
diff --git a/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentExecutor.kt b/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentExecutor.kt
index 69f2bcc992..ecf0b42259 100644
--- a/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentExecutor.kt
+++ b/src/main/kotlin/cc/unitmesh/devti/agent/CustomAgentExecutor.kt
@@ -10,6 +10,7 @@ import cc.unitmesh.devti.llms.custom.updateCustomFormat
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.project.Project
+import java.util.concurrent.TimeUnit
 import kotlinx.coroutines.flow.Flow
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
@@ -50,7 +51,7 @@ class CustomAgentExecutor(val project: Project) : CustomSSEProcessor(project) {
             }
         }
 
-        client = client.newBuilder().build()
+        client = client.newBuilder().connectTimeout(agent.defaultTimeout, TimeUnit.SECONDS).readTimeout(agent.defaultTimeout, TimeUnit.SECONDS).build()
         val call = client.newCall(builder.url(agent.url).post(body).build())
 
         return when (agent.responseAction) {
diff --git a/src/main/kotlin/cc/unitmesh/devti/agent/model/CustomAgentConfig.kt b/src/main/kotlin/cc/unitmesh/devti/agent/model/CustomAgentConfig.kt
index 2d15a17ad4..47a902310d 100644
--- a/src/main/kotlin/cc/unitmesh/devti/agent/model/CustomAgentConfig.kt
+++ b/src/main/kotlin/cc/unitmesh/devti/agent/model/CustomAgentConfig.kt
@@ -55,6 +55,7 @@ data class CustomAgentConfig(
     val transition: List<CustomFlowTransition> = emptyList(),
     val interactive: InteractionType = InteractionType.ChatPanel,
     val auth: CustomAgentAuth? = null,
+    val defaultTimeout: Long = 10,
 ) {
     var state: CustomAgentState = CustomAgentState.START
 
diff --git a/src/main/resources/cc/unitmesh/devti/custom/schema/autodev-custom-agent.json b/src/main/resources/cc/unitmesh/devti/custom/schema/autodev-custom-agent.json
index bbf708e9bc..b5e9223fa9 100644
--- a/src/main/resources/cc/unitmesh/devti/custom/schema/autodev-custom-agent.json
+++ b/src/main/resources/cc/unitmesh/devti/custom/schema/autodev-custom-agent.json
@@ -53,6 +53,9 @@
           "WebView",
           "DevIns"
         ]
+      },
+      "defaultTimeout": {
+        "type": "number"
       }
     },
     "required": [