diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt index f9117725..73e38109 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt @@ -195,7 +195,6 @@ object AgentFactory { extraBehavior: String? = null, toolOverrides: Set? = null ): AIAgent { - val selectedTools = toolOverrides ?: SubagentTool.entries.toSet() return AIAgent.Companion( promptExecutor = createExecutor(provider), @@ -214,7 +213,7 @@ object AgentFactory { # Tool Usage Policy - You may call multiple tools in a single turn. If calls are independent, run them in parallel. If calls depend on earlier results, run sequentially. - - Prefer specialized tools over bash: use Read for file content, Grep for search, Edit/Write for code changes. Use Bash only for true shell operations. + - Prefer specialized tools over bash: use Read for file content, IntelliJSearch for search, Edit/Write for code changes. Use Bash only for true shell operations. - Never use Bash to "echo" thoughts or communicate. All communication happens in your response text. - Never guess parameters. If a required argument is unknown, first gather it via an appropriate tool. - Respect approval gates: Edit/Write operations require confirmation hooks and may be denied. @@ -223,7 +222,7 @@ object AgentFactory { # Tool Routing Rules - If the user asks about how to use a library, best practices, API semantics, configuration, or conventions: prefer ResolveLibraryId followed by GetLibraryDocs to gather authoritative guidance before proposing changes. - - Use IntelliJSearch/Grep to locate symbols and examples before editing files. + - Use IntelliJSearch to locate symbols and examples before editing files. # Collaboration as Subagent - Assume a parent agent orchestrates overall strategy. Focus on execution quality and clear, minimal output. @@ -231,7 +230,7 @@ object AgentFactory { # Good Practices - Be precise and cite evidence (paths/lines) for findings and changes. - - Batch independent reads/greps/web queries in parallel for speed. + - Batch independent reads/searches/web queries in parallel for speed. - Validate at boundaries (user input, external APIs); trust internal code guarantees. """.trimIndent() ) @@ -272,18 +271,18 @@ object AgentFactory { ${getEnvironmentInfo(project)}${behaviorSection(extraBehavior)} # Tool Usage Policy (Read-only) - - Use Read to examine files; Grep to search patterns; WebSearch for external context; ResolveLibraryId/GetLibraryDocs for dependencies; TodoWrite to record findings. - - You may call multiple tools in parallel when independent (e.g., read multiple files, run several greps at once). + - Use Read to examine files; IntelliJSearch to search patterns; WebSearch for external context; ResolveLibraryId/GetLibraryDocs for dependencies; TodoWrite to record findings. + - You may call multiple tools in parallel when independent (e.g., read multiple files, run several searches at once). - Do not use Edit or Write. Avoid destructive Bash. Use Bash only for safe, read-only operations if strictly necessary. - Never guess parameters; first gather precise paths/patterns. # Tool Routing Rules - For library usage and best practices: ResolveLibraryId then GetLibraryDocs to retrieve relevant docs prior to summarizing advice. - - Use IntelliJSearch/Grep for code navigation and symbol discovery. + - Use IntelliJSearch for code navigation and symbol discovery. # Exploration Workflow - Initial scan: read key config/entry files in parallel; map structure. - - Deep dive: run targeted greps and reads in parallel for related components. + - Deep dive: run targeted searches and reads in parallel for related components. - Summarize: synthesize findings with concrete references; use TodoWrite to capture a brief outline of insights. # Good Practices diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt index 8d14fa76..a417ce95 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt @@ -189,7 +189,6 @@ internal class SingleRunStrategyProvider : AgentRunStrategyProvider { edge(nodeShowCompressionLoading forwardTo nodeCompressHistory) edge(nodeCompressHistory forwardTo nodeResetCompressionLoading) edge(nodeResetCompressionLoading forwardTo nodeSendCompressedHistory) - edge(nodeSendToolResult forwardTo nodeFinish onEmptyOutput { true }) edge(nodeSendToolResult forwardTo nodeFinish onSingleAssistantResponse { true }) edge(nodeSendToolResult forwardTo nodeExecuteTool onMultipleToolCalls { true }) @@ -326,17 +325,3 @@ private infix fun AIAgentEdgeBuilderIntermediate .onCondition { messages -> block(messages[0]) } .transformed { it[0].content } } - -@EdgeTransformationDslMarker -private infix fun AIAgentEdgeBuilderIntermediate, OutgoingInput>.onEmptyOutput( - block: suspend (Message.Response) -> Boolean -): AIAgentEdgeBuilderIntermediate { - return onIsInstance(List::class) - .transformed { response -> - response.filter { item -> item is Message.Response && item !is Message.Reasoning } - .filterIsInstance() - } - .onCondition { it.isEmpty() } - .onCondition { messages -> block(messages[0]) } - .transformed { it[0].content } -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt index ab0607e1..b5b95a69 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt @@ -4,21 +4,24 @@ import ai.koog.agents.core.tools.Tool import ai.koog.agents.core.tools.annotations.LLMDescription import ai.koog.agents.ext.tool.shell.ShellCommandConfirmation import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.ModalityState +import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.application.runWriteAction +import com.intellij.openapi.vfs.VirtualFileManager import ee.carlrobert.codegpt.agent.AgentToolOutputNotifier import ee.carlrobert.codegpt.agent.ToolRunContext import ee.carlrobert.codegpt.settings.ProxyAISettingsService import ee.carlrobert.codegpt.tokens.truncateToolResult -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.Job -import kotlinx.coroutines.currentCoroutineContext -import kotlinx.coroutines.launch -import kotlinx.coroutines.withContext -import kotlinx.coroutines.withTimeoutOrNull +import kotlinx.coroutines.* +import kotlinx.coroutines.channels.Channel +import kotlinx.coroutines.channels.ReceiveChannel +import kotlinx.coroutines.selects.select import kotlinx.serialization.SerialName import kotlinx.serialization.Serializable import java.io.IOException import java.nio.file.Paths import java.util.* +import java.util.concurrent.atomic.AtomicBoolean import kotlin.coroutines.cancellation.CancellationException fun interface BashCommandConfirmationHandler { @@ -281,6 +284,7 @@ class BashTool( return withContext(Dispatchers.IO) { val stdoutBuilder = StringBuilder() val stderrBuilder = StringBuilder() + val activityChannel = Channel(Channel.CONFLATED) val shellCommand = buildShellCommand(args.command) val process = ProcessBuilder(shellCommand) @@ -289,12 +293,13 @@ class BashTool( redirectErrorStream(false) } .start() + closeStdin(process) try { val job = currentCoroutineContext()[Job] job?.invokeOnCompletion { cause -> if (cause is CancellationException && process.isAlive) { - runCatching { process.destroyForcibly() } + terminateProcess(process) } } val stdoutJob = launch { @@ -303,6 +308,7 @@ class BashTool( lines.forEach { line -> stdoutBuilder.appendLine(line) publisher.toolOutput(toolId, line, false) + activityChannel.trySend(Unit) } } } catch (_: IOException) { @@ -316,6 +322,7 @@ class BashTool( lines.forEach { line -> stderrBuilder.appendLine(line) publisher.toolOutput(toolId, line, true) + activityChannel.trySend(Unit) } } } catch (_: IOException) { @@ -323,16 +330,15 @@ class BashTool( } } - val timeoutMs = args.timeout.coerceIn(1, 600_000) - val isCompleted = withTimeoutOrNull(timeoutMs.toLong()) { - withContext(Dispatchers.IO) { - process.waitFor() - } - } != null - - if (!isCompleted) { - process.destroyForcibly() - } + val timedOut = AtomicBoolean(false) + val exitDeferred = async(Dispatchers.IO) { process.waitFor() } + waitForExitOrIdleTimeout( + process = process, + timeoutMs = args.timeout, + activityChannel = activityChannel, + timedOut = timedOut, + exitDeferred = exitDeferred + ) stdoutJob.join() stderrJob.join() @@ -340,18 +346,24 @@ class BashTool( val combinedOutput = buildCombinedOutput( stdoutBuilder.toString().trimEnd(), stderrBuilder.toString().trimEnd(), - if (!isCompleted) "Command timed out after ${timeoutMs}ms" else null + if (timedOut.get()) "Command timed out after ${args.timeout}ms of inactivity" else null ) + runInEdt(ModalityState.defaultModalityState()) { + runWriteAction { + VirtualFileManager.getInstance().syncRefresh() + } + } + Result( args.command, - if (isCompleted) process.exitValue() else null, + if (!timedOut.get()) process.exitValue() else null, combinedOutput, null ) } finally { if (process.isAlive) { - process.destroyForcibly() + terminateProcess(process) } } } @@ -369,6 +381,63 @@ class BashTool( }.trimEnd() } + private suspend fun waitForExitOrIdleTimeout( + process: Process, + timeoutMs: Int, + activityChannel: ReceiveChannel, + timedOut: AtomicBoolean, + exitDeferred: Deferred + ) { + var done = false + while (!done) { + val event = withTimeoutOrNull(timeoutMs.toLong()) { + select { + exitDeferred.onAwait { WaitEvent.EXIT } + activityChannel.onReceive { WaitEvent.ACTIVITY } + } + } + when (event) { + WaitEvent.EXIT -> done = true + WaitEvent.ACTIVITY -> { /* reset idle timer */ + } + + null -> { + timedOut.set(true) + terminateProcess(process) + done = true + } + } + } + } + + private enum class WaitEvent { + EXIT, + ACTIVITY + } + + private fun destroyProcessTree(process: Process) { + val handle = process.toHandle() + handle.descendants().forEach { child -> + runCatching { child.destroyForcibly() } + } + runCatching { handle.destroyForcibly() } + } + + private fun closeProcessStreams(process: Process) { + runCatching { process.inputStream.close() } + runCatching { process.errorStream.close() } + runCatching { process.outputStream.close() } + } + + private fun closeStdin(process: Process) { + runCatching { process.outputStream.close() } + } + + private fun terminateProcess(process: Process) { + destroyProcessTree(process) + closeProcessStreams(process) + } + override fun encodeResultToString(result: Result): String = with(result) { val raw = buildString { appendLine("Command: $command") @@ -431,7 +500,20 @@ class BashTool( private fun shouldBlockByIgnore(command: String): Boolean { val svc = settingsService ?: return false - val readers = setOf("cat", "grep", "rg", "sed", "awk", "head", "tail", "less", "more", "wc", "stat", "file") + val readers = setOf( + "cat", + "grep", + "rg", + "sed", + "awk", + "head", + "tail", + "less", + "more", + "wc", + "stat", + "file" + ) val tokens = tokenize(command) val paths = mutableListOf() var lastWasReader = false @@ -474,7 +556,9 @@ class BashTool( if (c == '\'' || c == '"') { quote = c } else if (c.isWhitespace()) { - if (sb.isNotEmpty()) { out.add(sb.toString()); sb.setLength(0) } + if (sb.isNotEmpty()) { + out.add(sb.toString()); sb.setLength(0) + } } else { sb.append(c) } @@ -488,7 +572,9 @@ class BashTool( } private fun looksLikePath(token: String): Boolean = - token.startsWith("/") || token.startsWith("./") || token.startsWith("../") || token.contains('/') + token.startsWith("/") || token.startsWith("./") || token.startsWith("../") || token.contains( + '/' + ) private fun toAbsolute(token: String): String { return try { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt index 81ee8946..5d81c035 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt @@ -10,7 +10,6 @@ import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import ee.carlrobert.codegpt.EncodingManager import ee.carlrobert.codegpt.agent.* -import ee.carlrobert.codegpt.conversations.message.TokenUsage import ee.carlrobert.codegpt.settings.ProxyAISettingsService import ee.carlrobert.codegpt.settings.ProxyAISubagent import ee.carlrobert.codegpt.settings.agents.SubagentDefaults @@ -77,16 +76,6 @@ class TaskTool( val totalTokens: Long = 0 ) - internal data class InternalResult( - val agentType: String, - val description: String, - val prompt: String, - val output: String, - val executionTime: Long, - val totalTokens: Long = 0, - val tokenUsage: TokenUsage? = null - ) - override suspend fun execute(args: Args): Result { val startTime = System.currentTimeMillis() val parentId = ToolRunContext.getToolId(sessionId) @@ -126,9 +115,9 @@ class TaskTool( approveToolCall = approvalHandler, onAgentToolCallStarting = toolCallBridge::onToolCallStarting, onAgentToolCallCompleted = toolCallBridge::onToolCallCompleted, + onCreditsAvailable = events::onCreditsAvailable, extraBehavior = extraBehavior, toolOverrides = toolOverrides, - onCreditsAvailable = events::onCreditsAvailable ) } @@ -183,7 +172,6 @@ private fun buildTaskDescription(project: Project): String { Built-in agent types: - general-purpose: General-purpose agent for researching complex questions, searching for code, and executing multi-step tasks. - explore: Fast agent specialized for exploring codebases. - - plan: Software architect agent for designing implementation plans. Custom subagents: - You may also pass the exact name of a configured subagent (as shown in ProxyAI Settings > Subagents) in the subagent_type field. @@ -194,19 +182,19 @@ private fun buildTaskDescription(project: Project): String { """.trimIndent() ) - val customs = runCatching { + val subagents = runCatching { project.service().getSubagents() .filterNot { SubagentDefaults.isBuiltInId(it.id) } }.getOrNull()?.takeIf { it.isNotEmpty() } - if (customs != null) { + if (subagents != null) { appendLine() appendLine("Configured subagents available:") - customs.forEach { sa -> - val title = sa.title.trim() - if (title.isEmpty()) return@forEach - val desc = sa.objective.trim() + subagents + .filter { it.title.trim().isNotBlank() } + .forEach { sa -> append("- ") - append(title) + append(sa.title) + val desc = sa.objective.trim() if (desc.isNotBlank()) { append(": ") append(desc.take(140)) @@ -225,7 +213,7 @@ private data class ConfiguredSubagent( private fun isBuiltInAgentType(value: String): Boolean { return when (value.lowercase().trim()) { - "general-purpose", "explore", "plan" -> true + "general-purpose", "explore" -> true else -> false } }