vrr/ProxyAI
2
0
Fork 0
mirror of https://github.com/carlrobertoh/ProxyAI.git synced 2026-05-19 16:28:46 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: refine agent tools behavior and strategy routing

Browse source
This commit is contained in:
Carl-Robert Linnupuu 2026-01-15 11:41:31 +00:00
parent e558a27ada
commit b00593b9bd
4 changed files with 125 additions and 67 deletions
Download patch file Download diff file Expand all files Collapse all files

15
src/main/kotlin/ee/carlrobert/codegpt/agent/AgentFactory.kt
View file

@ -195,7 +195,6 @@ object AgentFactory {
extraBehavior: String? = null,
toolOverrides: Set<SubagentTool>? = null
): AIAgent<String, String> {
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

15
src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt
View file

@ -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 <IncomingOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate
.onCondition { messages -> block(messages[0]) }
.transformed { it[0].content }
}
@EdgeTransformationDslMarker
private infix fun <IncomingOutput, OutgoingInput> AIAgentEdgeBuilderIntermediate<IncomingOutput, List<Message.Response>, OutgoingInput>.onEmptyOutput(
block: suspend (Message.Response) -> Boolean
): AIAgentEdgeBuilderIntermediate<IncomingOutput, String, OutgoingInput> {
return onIsInstance(List::class)
.transformed { response ->
response.filter { item -> item is Message.Response && item !is Message.Reasoning }
.filterIsInstance<Message.Response>()
}
.onCondition { it.isEmpty() }
.onCondition { messages -> block(messages[0]) }
.transformed { it[0].content }
}

132
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt
View file

@ -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<Unit>(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<Unit>,
timedOut: AtomicBoolean,
exitDeferred: Deferred<Int>
) {
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<String>()
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 {

30
src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt
View file

@ -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<ProxyAISettingsService>().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
}
}