From b456b537525eb0bcbbff100a89c6185f783d4bff Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Fri, 6 Feb 2026 16:03:07 +0000 Subject: [PATCH] feat: agent chat history --- .../ui/ChatToolWindowScrollablePanel.java | 12 + .../carlrobert/codegpt/agent/AgentService.kt | 178 +++++++-- .../codegpt/agent/MessageWithContext.kt | 9 +- .../agent/MessageWithContextSnapshot.kt | 13 - .../carlrobert/codegpt/agent/ProxyAIAgent.kt | 60 ++- .../ee/carlrobert/codegpt/agent/ToolSpecs.kt | 59 +-- .../agent/clients/RetryingPromptExecutor.kt | 8 +- .../AgentCheckpointConversationMapper.kt | 106 +++++ .../history/AgentCheckpointHistoryService.kt | 235 +++++++++++ .../history/AgentHistoryMessageVisibility.kt | 36 ++ .../codegpt/agent/history/AgentHistoryPage.kt | 7 + .../history/AgentHistoryThreadSummary.kt | 20 + .../codegpt/agent/history/CheckpointRef.kt | 6 + .../strategy/SingleRunStrategyProvider.kt | 50 +-- .../codegpt/agent/tools/BaseTool.kt | 2 - .../codegpt/agent/tools/BashTool.kt | 9 +- .../codegpt/agent/tools/EditTool.kt | 3 - .../codegpt/agent/tools/ReadTool.kt | 5 +- .../codegpt/agent/tools/TaskTool.kt | 12 +- .../codegpt/settings/hooks/HookManager.kt | 4 +- .../codegpt/toolwindow/agent/AgentSession.kt | 3 + .../agent/AgentToolWindowContentManager.kt | 45 ++- .../agent/AgentToolWindowTabPanel.kt | 127 +++++- .../agent/history/AgentHistoryListPanel.kt | 375 ++++++++++++++++++ .../agent/ui/AgentToolWindowLandingPanel.kt | 99 +++++ .../codegpt/ui/textarea/PromptTextField.kt | 33 +- .../ui/textarea/header/tag/TagDetails.kt | 8 +- .../agent/AgentServiceIntegrationTest.kt | 324 +++++++++++++++ .../AgentCheckpointConversationMapperTest.kt | 171 ++++++++ .../AgentCheckpointHistoryServiceTest.kt | 301 ++++++++++++++ 30 files changed, 2149 insertions(+), 171 deletions(-) delete mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContextSnapshot.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryService.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryMessageVisibility.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryPage.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/agent/history/CheckpointRef.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/history/AgentHistoryListPanel.kt create mode 100644 src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt create mode 100644 src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt create mode 100644 src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java index 899b4f52..9d4865d6 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/ChatToolWindowScrollablePanel.java @@ -28,6 +28,7 @@ import javax.swing.JPanel; public class ChatToolWindowScrollablePanel extends ScrollablePanel { private final Map visibleMessagePanels = new HashMap<>(); + private boolean landingViewVisible; public ChatToolWindowScrollablePanel() { super(new VerticalStackLayout()); @@ -36,6 +37,7 @@ public class ChatToolWindowScrollablePanel extends ScrollablePanel { public void displayLandingView(JComponent landingView) { clearAll(); add(landingView); + landingViewVisible = true; if (ModelSelectionService.getInstance().getServiceForFeature(FeatureType.CHAT) == ServiceType.PROXYAI && !CredentialsStore.INSTANCE.isCredentialSet(CredentialKey.CodeGptApiKey.INSTANCE)) { @@ -73,6 +75,9 @@ public class ChatToolWindowScrollablePanel extends ScrollablePanel { } public JPanel addMessage(UUID messageId) { + if (landingViewVisible) { + clearAll(); + } var messageWrapper = new JPanel(); messageWrapper.setLayout(new BoxLayout(messageWrapper, BoxLayout.PAGE_AXIS)); add(messageWrapper); @@ -86,8 +91,15 @@ public class ChatToolWindowScrollablePanel extends ScrollablePanel { visibleMessagePanels.remove(messageId); } + public void clearLandingViewIfVisible() { + if (landingViewVisible) { + clearAll(); + } + } + public void clearAll() { visibleMessagePanels.clear(); + landingViewVisible = false; removeAll(); update(); } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt index b90005fe..f76b82ea 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt.agent import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.AIAgentService import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider import ai.koog.prompt.executor.clients.LLMClientException @@ -8,8 +9,11 @@ import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.project.Project +import ee.carlrobert.codegpt.agent.history.AgentCheckpointHistoryService +import ee.carlrobert.codegpt.agent.history.CheckpointRef import ee.carlrobert.codegpt.settings.service.FeatureType import ee.carlrobert.codegpt.settings.service.ModelSelectionService +import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails import kotlinx.coroutines.* @@ -19,6 +23,17 @@ import java.util.* import java.util.concurrent.ConcurrentHashMap import kotlin.io.path.Path +internal fun interface AgentRuntimeFactory { + fun create( + project: Project, + checkpointStorage: JVMFilePersistenceStorageProvider, + provider: ServiceType, + events: AgentEvents, + sessionId: String, + pendingMessages: ConcurrentHashMap>, + ): AIAgentService> +} + @Service(Service.Level.PROJECT) class AgentService(private val project: Project) { @@ -26,11 +41,31 @@ class AgentService(private val project: Project) { private val logger = thisLogger() } + private data class SessionRuntime( + val service: AIAgentService>, + val provider: ServiceType, + val modelId: String, + val selectedServerIds: Set, + val events: AgentEvents + ) + private val sessionJobs = ConcurrentHashMap() private val pendingMessages = ConcurrentHashMap>() private val sessionAgents = ConcurrentHashMap>() + private val sessionRuntimes = ConcurrentHashMap() + internal var runtimeFactory: AgentRuntimeFactory = AgentRuntimeFactory { p, storage, provider, events, sid, pending -> + ProxyAIAgent.createService( + project = p, + checkpointStorage = storage, + provider = provider, + events = events, + sessionId = sid, + pendingMessages = pending + ) + } private val checkpointStorage = JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai")) + private val historyService = project.service() private val _queuedMessageProcessed = MutableSharedFlow() val queuedMessageProcessed = _queuedMessageProcessed.asSharedFlow() @@ -40,19 +75,28 @@ class AgentService(private val project: Project) { } suspend fun getCheckpoint(sessionId: String): AgentCheckpointData? { - val prevAgentId = sessionAgents[sessionId]?.id ?: return null - return runCatching { - checkpointStorage.getCheckpoints(prevAgentId) - .filter { it.nodePath != "tombstone" } - .maxByOrNull { it.createdAt } - }.onFailure { ex -> - val sessionInfo = sessionId.let { " session=$it" } - logger.error( - "Agent checkpoints: failed to load for$sessionInfo agentId=$prevAgentId " + - "error=${ex.message}", - ex - ) - }.getOrNull() + val session = + project.service().getSession(sessionId) ?: return null + val ref = session.resumeCheckpointRef + + if (ref != null) { + val checkpoint = runCatching { historyService.loadResumeCheckpoint(ref) } + .onFailure { ex -> logCheckpointLoadFailure(sessionId, ref.agentId, ex) } + .getOrNull() + if (checkpoint != null) { + return checkpoint + } + } + + val runtimeAgentId = session.runtimeAgentId ?: return null + val latest = runCatching { historyService.loadLatestResumeCheckpoint(runtimeAgentId) } + .onFailure { ex -> logCheckpointLoadFailure(sessionId, runtimeAgentId, ex) } + .getOrNull() + ?: return null + + project.service() + .setResumeCheckpointRef(sessionId, CheckpointRef(runtimeAgentId, latest.checkpointId)) + return latest } fun submitMessage(message: MessageWithContext, events: AgentEvents, sessionId: String) { @@ -64,17 +108,19 @@ class AgentService(private val project: Project) { } val provider = service().getServiceForFeature(FeatureType.AGENT) - val previousCheckpoint = runBlocking { getCheckpoint(sessionId) } + val contentManager = project.service() + val runtimeAgentId = contentManager.getSession(sessionId)?.runtimeAgentId + val runtime = ensureSessionRuntime(sessionId, provider, events) - val agent = ProxyAIAgent.create( - project, - checkpointStorage, - previousCheckpoint, - provider, - events, - sessionId, - pendingMessages - ) + val agent = runCatching { + runBlocking { + runtime.service.createAgent(id = runtimeAgentId) + } + }.onFailure { ex -> + logger.warn("Failed to create managed agent for session=$sessionId", ex) + }.getOrNull() ?: return + + contentManager.setRuntimeAgentId(sessionId, agent.id) sessionAgents[sessionId] = agent sessionJobs[sessionId] = CoroutineScope(Dispatchers.IO).launch { try { @@ -86,6 +132,15 @@ class AgentService(private val project: Project) { } catch (ex: Exception) { logger.error(ex) } finally { + refreshSessionResumeCheckpoint(sessionId, agent.id) + sessionAgents.remove(sessionId, agent) + runCatching { runtime.service.removeAgentWithId(agent.id) } + .onFailure { ex -> + logger.warn( + "Failed removing managed agent for session=$sessionId agentId=${agent.id}", + ex + ) + } sessionJobs.remove(sessionId) } } @@ -100,6 +155,13 @@ class AgentService(private val project: Project) { cancelCurrentRun(sessionId) pendingMessages.remove(sessionId) sessionAgents.remove(sessionId) + sessionRuntimes.remove(sessionId)?.let { runtime -> + runCatching { + runBlocking { runtime.service.closeAll() } + }.onFailure { ex -> + logger.warn("Failed closing managed agent service for session=$sessionId", ex) + } + } project.service().clear(sessionId) } @@ -138,4 +200,74 @@ class AgentService(private val project: Project) { project.service() .update(sessionId, conversationId, selectedServerIds) } + + private suspend fun refreshSessionResumeCheckpoint(sessionId: String, agentId: String) { + val ref = runCatching { + historyService.loadLatestResumeCheckpoint(agentId) + ?.let { CheckpointRef(agentId, it.checkpointId) } + }.onFailure { ex -> + logger.warn( + "Agent checkpoints: failed to refresh session checkpoint session=$sessionId " + + "agentId=$agentId error=${ex.message}", + ex + ) + }.getOrNull() + + if (ref != null) { + project.service().setResumeCheckpointRef(sessionId, ref) + } + } + + private fun ensureSessionRuntime( + sessionId: String, + provider: ServiceType, + events: AgentEvents + ): SessionRuntime { + val modelId = service().getAgentModel().id + val selectedServerIds = project.service() + .get(sessionId) + ?.selectedServerIds + ?: emptySet() + val existing = sessionRuntimes[sessionId] + if (existing != null && + existing.provider == provider && + existing.modelId == modelId && + existing.selectedServerIds == selectedServerIds && + existing.events === events + ) { + return existing + } + + existing?.let { stale -> + runCatching { runBlocking { stale.service.closeAll() } } + .onFailure { ex -> + logger.warn("Failed closing stale managed service for session=$sessionId", ex) + } + } + + val created = SessionRuntime( + service = runtimeFactory.create( + project = project, + checkpointStorage = checkpointStorage, + provider = provider, + events = events, + sessionId = sessionId, + pendingMessages = pendingMessages + ), + provider = provider, + modelId = modelId, + selectedServerIds = selectedServerIds, + events = events + ) + sessionRuntimes[sessionId] = created + return created + } + + private fun logCheckpointLoadFailure(sessionId: String, agentId: String, ex: Throwable) { + val sessionInfo = sessionId.let { " session=$it" } + logger.error( + "Agent checkpoints: failed to load for$sessionInfo agentId=$agentId error=${ex.message}", + ex + ) + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContext.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContext.kt index 3dd0dfa1..b2f5527f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContext.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContext.kt @@ -1,14 +1,19 @@ package ee.carlrobert.codegpt.agent import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails -import java.util.UUID +import kotlinx.serialization.Serializable +import kotlinx.serialization.Transient +import java.util.* +@Serializable class MessageWithContext( val text: String, - val tags: List, + @Transient + val tags: List = emptyList(), val uiVisible: Boolean = true, val uiText: String = text ) { + @Transient val id: UUID = UUID.randomUUID() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContextSnapshot.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContextSnapshot.kt deleted file mode 100644 index 44311af5..00000000 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/MessageWithContextSnapshot.kt +++ /dev/null @@ -1,13 +0,0 @@ -package ee.carlrobert.codegpt.agent - -import kotlinx.serialization.Serializable - -@Serializable -data class MessageWithContextSnapshot( - val text: String, - val tagNames: List -) - -fun MessageWithContext.toSnapshot(): MessageWithContextSnapshot { - return MessageWithContextSnapshot(text, tags.map { it.name }) -} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt index bd5ff36b..b4099d11 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt.agent -import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.AIAgentService +import ai.koog.agents.core.agent.GraphAIAgentService import ai.koog.agents.core.agent.config.AIAgentConfig import ai.koog.agents.core.agent.context.RollbackStrategy import ai.koog.agents.core.tools.Tool @@ -11,13 +12,12 @@ import ai.koog.agents.features.eventHandler.feature.handleEvents import ai.koog.agents.features.tokenizer.feature.MessageTokenizer import ai.koog.agents.features.tracing.feature.Tracing import ai.koog.agents.features.tracing.writer.TraceFeatureMessageLogWriter -import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.agents.snapshot.feature.Persistence -import ai.koog.agents.snapshot.feature.persistence import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider import ai.koog.prompt.dsl.prompt import ai.koog.prompt.llm.LLModel import ai.koog.prompt.message.Message +import ai.koog.prompt.message.RequestMetaInfo import ai.koog.prompt.streaming.StreamFrame import ai.koog.prompt.tokenizer.Tokenizer import com.intellij.openapi.components.service @@ -40,12 +40,13 @@ import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.BashPayload import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalType import io.github.oshai.kotlinlogging.KotlinLogging +import kotlinx.datetime.Clock import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive import java.nio.file.Files import java.nio.file.Paths import java.util.* import java.util.concurrent.ConcurrentHashMap -import kotlin.reflect.typeOf data class ToolError(val message: String) @@ -54,9 +55,8 @@ object ProxyAIAgent { private const val INSTRUCTION_FILE_NAME = "PROXYAI.md" private val logger = KotlinLogging.logger { } - private val messageWithContextSnapshotType = typeOf() - private fun searchForInstructions(projectPath: String?): String? { + internal fun loadProjectInstructions(projectPath: String?): String? { if (projectPath == null) return null val projectRoot = Paths.get(projectPath) if (!Files.exists(projectRoot)) return null @@ -76,32 +76,29 @@ object ProxyAIAgent { return null } - fun create( + fun createService( project: Project, checkpointStorage: JVMFilePersistenceStorageProvider, - previousCheckpoint: AgentCheckpointData?, provider: ServiceType, events: AgentEvents, sessionId: String, pendingMessages: ConcurrentHashMap>, - ): AIAgent { + ): GraphAIAgentService { val modelSelection = service().getModelSelectionForFeature(FeatureType.AGENT) val skills = project.service().listSkills() val stream = provider != ServiceType.CUSTOM_OPENAI - val projectInstructions = searchForInstructions(project.basePath) + val projectInstructions = loadProjectInstructions(project.basePath) val executor = AgentFactory.createExecutor(provider, events) val pendingMessageQueue = pendingMessages.getOrPut(sessionId) { ArrayDeque() } val hookManager = HookManager(project) val toolRegistry = createToolRegistry(project, events, sessionId, provider, hookManager) val agentModel = service().getAgentModel() - val agent = AIAgent( + return AIAgentService( promptExecutor = executor, strategy = SingleRunStrategyProvider().build( project, executor, - projectInstructions, - previousCheckpoint, pendingMessageQueue, HistoryCompressionConfig( isLimitExceeded = buildHistoryTooBigPredicate(computeAvailableInput(agentModel)), @@ -122,6 +119,20 @@ object ProxyAIAgent { skills ) ) + + projectInstructions?.let { + message( + Message.User( + content = it, + metaInfo = RequestMetaInfo( + Clock.System.now(), + JsonObject(buildMap { + put("cacheable", JsonPrimitive(true)) + }) + ) + ) + ) + } }, model = agentModel, maxAgentIterations = 100 @@ -135,7 +146,8 @@ object ProxyAIAgent { } install(Persistence) { storage = checkpointStorage - rollbackStrategy = RollbackStrategy.Default + enableAutomaticPersistence = true + rollbackStrategy = RollbackStrategy.MessageHistoryOnly } install(MessageTokenizer) { tokenizer = object : Tokenizer { @@ -159,25 +171,6 @@ object ProxyAIAgent { } onNodeExecutionCompleted { ctx -> - val input = ctx.input - val (lastInput, lastInputType) = if (input is MessageWithContext) { - val snapshot = input.toSnapshot() - snapshot to messageWithContextSnapshotType - } else { - input to ctx.inputType - } - val checkpoint = ctx.context.persistence().createCheckpoint( - agentContext = ctx.context, - nodePath = ctx.context.executionInfo.path(), - lastInput = lastInput, - lastInputType = lastInputType, - checkpointId = ctx.context.runId, - version = 0L - ) - checkpoint?.checkpointId ?: return@onNodeExecutionCompleted - - if (stream) return@onNodeExecutionCompleted - (ctx.output as? List<*>)?.forEach { msg -> (msg as? Message.Assistant)?.let { events.onTextReceived(it.content) @@ -264,7 +257,6 @@ object ProxyAIAgent { } } } - return agent } private fun createToolRegistry( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/ToolSpecs.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/ToolSpecs.kt index fd754af4..5669c515 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/ToolSpecs.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/ToolSpecs.kt @@ -4,6 +4,7 @@ import ee.carlrobert.codegpt.agent.tools.* import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalType import kotlinx.serialization.KSerializer import kotlinx.serialization.builtins.serializer +import kotlinx.serialization.json.Json enum class ToolName(val id: String, val aliases: Set = emptySet()) { READ("Read"), @@ -22,34 +23,6 @@ enum class ToolName(val id: String, val aliases: Set = emptySet()) { ASK_USER_QUESTION("AskUserQuestion"), TODO_WRITE("TodoWrite", setOf("TodoWriteTool")), EXIT("Exit"); - - companion object { - val HOOK_AWARE = setOf( - READ, - WRITE, - EDIT, - BASH, - TASK, - INTELLIJ_SEARCH, - WEB_SEARCH, - MCP, - RESOLVE_LIBRARY_ID, - GET_LIBRARY_DOCS, - LOAD_SKILL, - ASK_USER_QUESTION, - TODO_WRITE, - BASH_OUTPUT, - KILL_SHELL - ) - - fun isHookAware(toolId: String): Boolean { - return fromId(toolId) in HOOK_AWARE - } - - fun fromId(toolId: String): ToolName? { - return entries.find { it.id == toolId } - } - } } data class ToolSpec( @@ -191,4 +164,34 @@ object ToolSpecs { fun approvalTypeFor(toolName: String): ToolApprovalType { return find(toolName)?.approvalType ?: ToolApprovalType.GENERIC } + + fun decodeArgsOrNull( + json: Json, + toolName: String, + payload: String + ): Any? { + return decodeOrNull(json, find(toolName)?.argsSerializer, payload) + } + + fun decodeResultOrNull( + json: Json, + toolName: String, + payload: String + ): Any? { + return decodeOrNull(json, find(toolName)?.resultSerializer, payload) + } + + @Suppress("UNCHECKED_CAST") + private fun decodeOrNull( + json: Json, + serializer: KSerializer<*>?, + payload: String + ): Any? { + if (serializer == null || payload.isBlank()) { + return null + } + return runCatching { + json.decodeFromString(serializer as KSerializer, payload) + }.getOrNull() + } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt index e2bc6f1e..e90ac518 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/clients/RetryingPromptExecutor.kt @@ -17,6 +17,7 @@ import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.catch import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.onEach +import kotlinx.serialization.SerializationException import kotlin.math.pow import kotlin.random.Random import kotlin.time.Duration @@ -53,12 +54,13 @@ class RetryingPromptExecutor( tools: List ): Flow { var attempt = 1 + var refinedPrompt = prompt var hasReceivedData: Boolean fun createStream(attemptNum: Int): Flow { hasReceivedData = false - return delegate.executeStreaming(prompt, model, tools) + return delegate.executeStreaming(refinedPrompt, model, tools) .onEach { hasReceivedData = true } .catch { error -> val shouldRetry = !hasReceivedData && attemptNum < retryPolicy.maxAttempts @@ -111,6 +113,10 @@ class RetryingPromptExecutor( } else { throw ce } + } catch (_: SerializationException) { + refinedPrompt = prompt.copy(messages = prompt.messages.dropLast(1)) + currentAttempt++ + continue } } } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt new file mode 100644 index 00000000..6c74363a --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt @@ -0,0 +1,106 @@ +package ee.carlrobert.codegpt.agent.history + +import ai.koog.agents.snapshot.feature.AgentCheckpointData +import ee.carlrobert.codegpt.conversations.Conversation +import ee.carlrobert.codegpt.conversations.message.Message +import ee.carlrobert.llm.client.openai.completion.response.ToolCall +import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse +import ai.koog.prompt.message.Message as PromptMessage + +object AgentCheckpointConversationMapper { + + fun toConversation( + checkpoint: AgentCheckpointData, + projectInstructions: String? + ): Conversation { + val conversation = Conversation() + val history = checkpoint.messageHistory.filterNot { it is PromptMessage.System } + + var currentPrompt: String? = null + val response = StringBuilder() + val toolCalls = mutableListOf() + val toolResults = LinkedHashMap() + var syntheticToolIdIndex = 0 + + fun flushTurn() { + val prompt = currentPrompt?.trim().orEmpty() + if (prompt.isBlank()) { + return + } + + val uiMessage = Message(prompt) + uiMessage.response = response.toString().trim() + if (toolCalls.isNotEmpty()) { + uiMessage.toolCalls = ArrayList(toolCalls) + } + if (toolResults.isNotEmpty()) { + uiMessage.toolCallResults = LinkedHashMap(toolResults) + } + conversation.addMessage(uiMessage) + currentPrompt = null + response.setLength(0) + toolCalls.clear() + toolResults.clear() + } + + history.forEach { msg -> + when (msg) { + is PromptMessage.User -> { + val text = msg.content.trim() + if (shouldHideInAgentToolWindow(msg, projectInstructions)) { + return@forEach + } + flushTurn() + currentPrompt = text + } + + is PromptMessage.Assistant -> appendAssistant(response, msg.content) + is PromptMessage.Reasoning -> appendAssistant(response, msg.content) + is PromptMessage.Tool.Call -> { + if (currentPrompt != null) { + val callId = + msg.id?.takeIf { it.isNotBlank() } + ?: "tool-call-${++syntheticToolIdIndex}" + toolCalls.add( + ToolCall( + null, + callId, + "function", + ToolFunctionResponse(msg.tool, msg.content.trim()) + ) + ) + } + } + + is PromptMessage.Tool.Result -> { + if (currentPrompt != null) { + val callId = msg.id?.takeIf { it.isNotBlank() } + ?: toolCalls.lastOrNull()?.id + ?: "tool-call-${++syntheticToolIdIndex}" + val prior = toolResults[callId] + val merged = if (prior.isNullOrBlank()) msg.content.trim() else { + "$prior\n\n${msg.content.trim()}" + } + toolResults[callId] = merged + } + } + + else -> Unit + } + } + + flushTurn() + return conversation + } + + private fun appendAssistant(sb: StringBuilder, content: String) { + val text = content.trim() + if (text.isBlank()) { + return + } + if (sb.isNotEmpty()) { + sb.append("\n\n") + } + sb.append(text) + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryService.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryService.kt new file mode 100644 index 00000000..f2167efb --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryService.kt @@ -0,0 +1,235 @@ +package ee.carlrobert.codegpt.agent.history + +import ai.koog.agents.snapshot.feature.AgentCheckpointData +import ai.koog.agents.snapshot.feature.isTombstone +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import com.intellij.openapi.components.Service +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.project.Project +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.withContext +import kotlinx.serialization.json.Json +import kotlinx.serialization.json.jsonObject +import kotlinx.serialization.json.jsonPrimitive +import java.nio.file.Files +import kotlin.io.path.Path +import ai.koog.prompt.message.Message as PromptMessage + +@Service(Service.Level.PROJECT) +class AgentCheckpointHistoryService(project: Project) { + + companion object { + private val logger = thisLogger() + } + + private val root = Path(project.basePath ?: "", ".proxyai") + private val checkpointsRoot = root.resolve("checkpoints") + private val storage = JVMFilePersistenceStorageProvider(root) + + @Volatile + private var cachedSummaries: List? = null + + suspend fun listThreads(limit: Int? = null): List = + withContext(Dispatchers.IO) { + val summaries = loadThreadSummaries() + + return@withContext if (limit == null) summaries else summaries.take(limit) + } + + suspend fun listThreadsPage( + query: String, + offset: Int, + limit: Int, + refresh: Boolean = false + ): AgentHistoryPage = withContext(Dispatchers.IO) { + val safeOffset = offset.coerceAtLeast(0) + val safeLimit = limit.coerceAtLeast(1) + val summaries = getCachedSummaries(refresh) + val filtered = filterSummaries(summaries, query) + if (safeOffset >= filtered.size) { + return@withContext AgentHistoryPage(emptyList(), false, filtered.size) + } + val endIndex = (safeOffset + safeLimit).coerceAtMost(filtered.size) + AgentHistoryPage( + items = filtered.subList(safeOffset, endIndex), + hasMore = endIndex < filtered.size, + total = filtered.size + ) + } + + suspend fun loadCheckpoint(ref: CheckpointRef): AgentCheckpointData? = + withContext(Dispatchers.IO) { + val checkpoints = storage.getCheckpoints(ref.agentId) + .filterNot { it.isTombstone() } + checkpoints.firstOrNull { it.checkpointId == ref.checkpointId } + } + + suspend fun loadResumeCheckpoint(ref: CheckpointRef): AgentCheckpointData? = + withContext(Dispatchers.IO) { + val checkpoints = storage.getCheckpoints(ref.agentId) + .filterNot { it.isTombstone() } + .sortedByDescending { it.createdAt } + if (checkpoints.isEmpty()) return@withContext null + + val preferred = checkpoints.firstOrNull { it.checkpointId == ref.checkpointId } + if (preferred == null) { + return@withContext checkpoints.firstOrNull { it.isResumableNode() } + } + if (preferred.isResumableNode()) { + return@withContext preferred + } + + checkpoints + .asSequence() + .filter { it.createdAt <= preferred.createdAt } + .firstOrNull { it.isResumableNode() } + ?: checkpoints.firstOrNull { it.isResumableNode() } + } + + suspend fun loadLatestResumeCheckpoint(agentId: String): AgentCheckpointData? = + withContext(Dispatchers.IO) { + val checkpoints = storage.getCheckpoints(agentId) + .filterNot { it.isTombstone() } + checkpoints + .filter { it.isResumableNode() } + .maxByOrNull { it.createdAt } + ?: checkpoints.maxByOrNull { it.createdAt } + } + + private suspend fun buildSummary(agentId: String): AgentHistoryThreadSummary? { + val checkpoints = storage.getCheckpoints(agentId) + .filterNot { it.isTombstone() } + if (checkpoints.isEmpty()) { + return null + } + + val latest = checkpoints.maxBy { it.createdAt } + val status = when (latest.nodePath.substringAfterLast('/')) { + "__finish__" -> ThreadStatus.COMPLETED + "__start__" -> ThreadStatus.UNKNOWN + else -> ThreadStatus.PARTIAL + } + + val title = extractThreadTitle(latest.messageHistory) + + val lastAssistant = latest.messageHistory + .asReversed() + .firstNotNullOfOrNull { msg -> + when (msg) { + is PromptMessage.Assistant -> toSingleLine(msg.content).takeIf { it.isNotBlank() } + is PromptMessage.Reasoning -> toSingleLine(msg.content).takeIf { it.isNotBlank() } + else -> null + } + } ?: "" + + return AgentHistoryThreadSummary( + agentId = agentId, + latest = CheckpointRef(agentId, latest.checkpointId), + latestCreatedAt = latest.createdAt, + status = status, + runCount = checkpoints.size, + title = title.take(80), + preview = lastAssistant.take(160) + ) + } + + private fun extractThreadTitle(history: List): String { + val todoWriteTitle = history + .asReversed() + .firstNotNullOfOrNull { message -> + if (message !is PromptMessage.Tool.Call || !isTodoWriteTool(message.tool)) { + return@firstNotNullOfOrNull null + } + extractTodoWriteTitle(message.content) + } + + if (!todoWriteTitle.isNullOrBlank()) { + return toSingleLine(todoWriteTitle) + } + + return history + .asReversed() + .firstNotNullOfOrNull { message -> + val user = message as? PromptMessage.User ?: return@firstNotNullOfOrNull null + if (shouldHideInAgentToolWindow(user)) { + return@firstNotNullOfOrNull null + } + user.content + .let(::toSingleLine) + .takeIf { it.isNotBlank() } + } ?: "Recovered conversation" + } + + private fun isTodoWriteTool(toolName: String): Boolean { + return toolName.equals("TodoWrite", ignoreCase = true) || + toolName.equals("TodoWriteTool", ignoreCase = true) + } + + private fun extractTodoWriteTitle(rawArgs: String): String? { + val args = runCatching { + Json.parseToJsonElement(rawArgs).jsonObject + }.getOrNull() ?: return null + + return args["title"]?.jsonPrimitive?.content?.trim()?.takeIf { it.isNotBlank() } + } + + private fun toSingleLine(value: String): String { + return value.replace("\\s+".toRegex(), " ").trim() + } + + private fun AgentCheckpointData.isResumableNode(): Boolean { + return nodePath.substringAfterLast('/') != "__finish__" + } + + private suspend fun loadThreadSummaries(): List { + if (!Files.exists(checkpointsRoot)) { + return emptyList() + } + + val agentIds = withContext(Dispatchers.IO) { + Files.list(checkpointsRoot) + }.use { stream -> + stream + .filter { Files.isDirectory(it) } + .map { it.fileName.toString() } + .toList() + } + + return agentIds.mapNotNull { agentId -> + runCatching { buildSummary(agentId) } + .onFailure { logger.warn("Failed to load checkpoints for agentId=$agentId", it) } + .getOrNull() + }.sortedByDescending { it.latestCreatedAt } + } + + private suspend fun getCachedSummaries(refresh: Boolean): List { + if (!refresh) { + synchronized(this) { + cachedSummaries?.let { return it } + } + } + + val loaded = loadThreadSummaries() + synchronized(this) { + if (refresh || cachedSummaries == null) { + cachedSummaries = loaded + } + return cachedSummaries ?: loaded + } + } + + private fun filterSummaries( + summaries: List, + query: String + ): List { + val normalized = query.trim() + if (normalized.isBlank()) { + return summaries + } + return summaries.filter { item -> + item.title.contains(normalized, ignoreCase = true) || + item.preview.contains(normalized, ignoreCase = true) || + item.agentId.contains(normalized, ignoreCase = true) + } + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryMessageVisibility.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryMessageVisibility.kt new file mode 100644 index 00000000..c06bc760 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryMessageVisibility.kt @@ -0,0 +1,36 @@ +package ee.carlrobert.codegpt.agent.history + +import ai.koog.prompt.message.Message as PromptMessage +import kotlinx.serialization.json.booleanOrNull +import kotlinx.serialization.json.contentOrNull +import kotlinx.serialization.json.jsonPrimitive + +private const val CACHEABLE_METADATA_KEY = "cacheable" + +internal fun shouldHideInAgentToolWindow( + message: PromptMessage.User, + projectInstructions: String? = null +): Boolean { + return isCacheableInstructionMessage(message) || + isProjectInstructionsMessage(message.content, projectInstructions) +} + +private fun isCacheableInstructionMessage(message: PromptMessage.User): Boolean { + val cacheable = message.metaInfo.metadata + ?.get(CACHEABLE_METADATA_KEY) + ?.jsonPrimitive + ?: return false + + return cacheable.booleanOrNull ?: (cacheable.contentOrNull?.equals("true", ignoreCase = true) == true) +} + +private fun isProjectInstructionsMessage(text: String, projectInstructions: String?): Boolean { + if (projectInstructions.isNullOrBlank()) { + return false + } + return normalize(text) == normalize(projectInstructions) +} + +private fun normalize(value: String): String { + return value.replace("\\s+".toRegex(), " ").trim() +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryPage.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryPage.kt new file mode 100644 index 00000000..e56add08 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryPage.kt @@ -0,0 +1,7 @@ +package ee.carlrobert.codegpt.agent.history + +data class AgentHistoryPage( + val items: List, + val hasMore: Boolean, + val total: Int +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt new file mode 100644 index 00000000..8950bb2e --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentHistoryThreadSummary.kt @@ -0,0 +1,20 @@ +package ee.carlrobert.codegpt.agent.history + +import kotlinx.datetime.Instant + +enum class ThreadStatus { + COMPLETED, + PARTIAL, + UNKNOWN +} + +data class AgentHistoryThreadSummary( + val agentId: String, + val latest: CheckpointRef, + val latestCreatedAt: Instant, + val status: ThreadStatus, + val runCount: Int, + val title: String, + val preview: String, + val searchText: String = "" +) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/history/CheckpointRef.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/CheckpointRef.kt new file mode 100644 index 00000000..4ea0baf7 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/history/CheckpointRef.kt @@ -0,0 +1,6 @@ +package ee.carlrobert.codegpt.agent.history + +data class CheckpointRef( + val agentId: String, + val checkpointId: String +) 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 ce495681..420165dc 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/strategy/SingleRunStrategyProvider.kt @@ -11,7 +11,6 @@ import ai.koog.agents.core.dsl.extension.* import ai.koog.agents.core.environment.ReceivedToolResult import ai.koog.agents.core.environment.result import ai.koog.agents.features.tokenizer.feature.tokenizer -import ai.koog.agents.snapshot.feature.AgentCheckpointData import ai.koog.prompt.dsl.Prompt import ai.koog.prompt.executor.model.PromptExecutor import ai.koog.prompt.message.Message @@ -29,6 +28,9 @@ import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsEvent import ee.carlrobert.codegpt.ui.textarea.TagProcessorFactory import ee.carlrobert.codegpt.ui.textarea.header.tag.TagDetails +import kotlinx.serialization.SerializationException +import kotlinx.serialization.json.Json +import kotlinx.serialization.json.jsonObject import java.util.* import ee.carlrobert.codegpt.conversations.message.Message as ChatMessage @@ -36,8 +38,6 @@ internal interface AgentRunStrategyProvider { fun build( project: Project, executor: PromptExecutor, - projectInstructions: String?, - previousCheckpoint: AgentCheckpointData?, pendingMessageQueue: ArrayDeque, historyCompressionConfig: HistoryCompressionConfig, events: AgentEvents, @@ -56,8 +56,6 @@ internal class SingleRunStrategyProvider : AgentRunStrategyProvider { override fun build( project: Project, executor: PromptExecutor, - projectInstructions: String?, - previousCheckpoint: AgentCheckpointData?, pendingMessageQueue: ArrayDeque, historyCompressionConfig: HistoryCompressionConfig, events: AgentEvents, @@ -67,29 +65,6 @@ internal class SingleRunStrategyProvider : AgentRunStrategyProvider { ): AIAgentGraphStrategy = strategy("single_run") { val nodeCallLLM by node> { message -> llm.writeSession { - if (previousCheckpoint == null) { - projectInstructions?.let { - appendPrompt { - user(it) - } - } - } else { - val filteredMessageHistory = previousCheckpoint.messageHistory - .filter { it !is Message.System } - .toMutableList() - .apply { - if (lastOrNull() is Message.Tool.Call) { - removeLast() - } - } - - for (previousMessage in filteredMessageHistory) { - if (previousMessage !is Message.System) { - appendPrompt { message(previousMessage) } - } - } - } - if (message.tags.isNotEmpty()) { val context = buildTagContext(project, message.tags) if (context.isNotBlank()) { @@ -206,7 +181,24 @@ private suspend fun AIAgentLLMWriteSession.requestResponses( val responses = if (stream) { val streamFrames = mutableListOf() requestLLMStreaming().collect { streamFrames.add(it) } - streamFrames.toMessageResponses() + streamFrames + .toMessageResponses() + .map { + if (it is Message.Tool.Call) { + try { + // validate json + Json.parseToJsonElement(it.content).jsonObject + it + } catch (_: SerializationException) { + // allows agent to retry the request + it.copy(parts = listOf(it.parts[0].copy(text = "{}"))) + } catch (e: Exception) { + throw e + } + } else { + it + } + } } else { val preparedPrompt = config.missingToolsConversionStrategy.convertPrompt(prompt, tools) executor.execute(preparedPrompt, model, tools) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt index 27aa6d55..4d0c7e6f 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BaseTool.kt @@ -71,7 +71,6 @@ abstract class BaseTool( BEFORE_TOOL_USE, beforePayload, name, - toolId, sessionId ) @@ -111,7 +110,6 @@ abstract class BaseTool( HookEventType.AFTER_TOOL_USE, afterPayload, name, - toolId, sessionId ) 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 a7d9f076..040dca86 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/BashTool.kt @@ -249,8 +249,7 @@ class BashTool( val deniedReason = hookManager.checkHooksForDenial( HookEventType.BEFORE_BASH_EXECUTION, prePayload, - "Bash", - toolId, + name, sessionId ) if (deniedReason != null) { @@ -291,8 +290,7 @@ class BashTool( hookManager.executeHooksForEvent( HookEventType.AFTER_BASH_EXECUTION, postPayload, - "Bash", - toolId, + name, sessionId ) @@ -309,8 +307,7 @@ class BashTool( hookManager.executeHooksForEvent( HookEventType.AFTER_BASH_EXECUTION, errorPayload, - "Bash", - toolId, + name, sessionId ) Result( diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt index b7f9c17c..417d4b8c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/EditTool.kt @@ -10,7 +10,6 @@ import com.intellij.openapi.project.Project import com.intellij.openapi.util.TextRange import com.intellij.openapi.vfs.LocalFileSystem import com.intellij.openapi.vfs.VirtualFileManager -import ee.carlrobert.codegpt.agent.ToolRunContext import ee.carlrobert.codegpt.settings.ProxyAISettingsService import ee.carlrobert.codegpt.settings.hooks.HookEventType import ee.carlrobert.codegpt.settings.hooks.HookManager @@ -280,7 +279,6 @@ class EditTool( editLocations = editLocations ) - val toolId = sessionId?.let { id -> ToolRunContext.getToolId(id) } val payload = mapOf( "file_path" to args.filePath, "replacements_made" to replacementsMade, @@ -295,7 +293,6 @@ class EditTool( HookEventType.AFTER_FILE_EDIT, payload, "Edit", - toolId, sessionId ) if (deniedReason != null) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt index 313a7b51..7e1856e0 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/ReadTool.kt @@ -8,7 +8,6 @@ import com.intellij.openapi.fileEditor.FileDocumentManager import com.intellij.openapi.fileTypes.FileTypeManager import com.intellij.openapi.project.Project import com.intellij.openapi.vfs.VirtualFileManager -import ee.carlrobert.codegpt.agent.ToolRunContext import ee.carlrobert.codegpt.settings.ProxyAISettingsService import ee.carlrobert.codegpt.settings.ToolPermissionPolicy import ee.carlrobert.codegpt.settings.hooks.HookEventType @@ -216,7 +215,6 @@ class ReadTool( } if (result is Result.Success) { - val toolId = sessionId?.let { id -> ToolRunContext.getToolId(id) } val payload = mapOf( "file_path" to args.filePath, "content" to result.content, @@ -225,8 +223,7 @@ class ReadTool( val deniedReason = hookManager.checkHooksForDenial( HookEventType.BEFORE_READ_FILE, payload, - "Read", - toolId, + name, sessionId ) if (deniedReason != null) { 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 ff3106de..3ba79b21 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/agent/tools/TaskTool.kt @@ -113,7 +113,6 @@ class TaskTool( HookEventType.SUBAGENT_START, startPayload, "Task", - ToolRunContext.getToolId(sessionId), sessionId ) } @@ -188,7 +187,10 @@ class TaskTool( } private fun configuredSubagent(subagentType: String): ConfiguredSubagent? { - return if (isBuiltInAgentType(subagentType)) null else findConfiguredSubagent(project, subagentType) + return if (isBuiltInAgentType(subagentType)) null else findConfiguredSubagent( + project, + subagentType + ) } private suspend fun emitStopHook(args: Args, output: String, duration: Long) { @@ -201,8 +203,7 @@ class TaskTool( "result" to output, "duration" to duration ), - "Task", - ToolRunContext.getToolId(sessionId), + name, sessionId ) } @@ -301,11 +302,14 @@ class TaskTool( details.contains("Unexpected response: no tool calls and no content", true) -> "Subagent '$subagentType' received an invalid model response: no content " + "and no tool calls were returned." + details.equals("Something went wrong", true) || details.equals("Something went wrong.", true) -> "Subagent '$subagentType' received a generic provider error without details." + details.isNotBlank() -> "Subagent '$subagentType' failed in the LLM client: $details" + else -> "Subagent '$subagentType' failed in the LLM client without an error message." } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt index 0968f9d8..46987027 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/hooks/HookManager.kt @@ -18,7 +18,6 @@ class HookManager( event: HookEventType, payload: Map, toolName: String? = null, - toolId: String? = null, sessionId: String? = null ): List { val settings = project.service().getSettings() @@ -81,10 +80,9 @@ class HookManager( event: HookEventType, payload: Map, toolName: String? = null, - toolId: String? = null, sessionId: String? = null ): String? { - val results = executeHooksForEvent(event, payload, toolName, toolId, sessionId) + val results = executeHooksForEvent(event, payload, toolName, sessionId) results.forEach { result -> when (result) { is HookExecutionResult.Denied -> return result.reason diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSession.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSession.kt index c686884d..4ad157a6 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSession.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentSession.kt @@ -1,6 +1,7 @@ package ee.carlrobert.codegpt.toolwindow.agent import com.intellij.openapi.vfs.VirtualFile +import ee.carlrobert.codegpt.agent.history.CheckpointRef import ee.carlrobert.codegpt.conversations.Conversation import ee.carlrobert.codegpt.settings.service.ServiceType @@ -13,6 +14,8 @@ data class AgentSession( val conversation: Conversation, var displayName: String = "", var serviceType: ServiceType? = null, + var runtimeAgentId: String? = null, + var resumeCheckpointRef: CheckpointRef? = null, val referencedFiles: List = emptyList(), val createdAt: Long = System.currentTimeMillis(), var lastActiveAt: Long = System.currentTimeMillis() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowContentManager.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowContentManager.kt index 24fbcf35..32290392 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowContentManager.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowContentManager.kt @@ -5,9 +5,12 @@ import com.intellij.openapi.components.Service import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import com.intellij.openapi.util.Disposer +import com.intellij.openapi.wm.ToolWindowManager import ee.carlrobert.codegpt.agent.AgentService +import ee.carlrobert.codegpt.agent.history.AgentHistoryThreadSummary +import ee.carlrobert.codegpt.agent.history.CheckpointRef import ee.carlrobert.codegpt.conversations.Conversation -import java.util.UUID +import java.util.* @Service(Service.Level.PROJECT) class AgentToolWindowContentManager(private val project: Project) : Disposable { @@ -43,6 +46,35 @@ class AgentToolWindowContentManager(private val project: Project) : Disposable { return tabPanel } + fun openCheckpointConversation( + thread: AgentHistoryThreadSummary, + conversation: Conversation + ): AgentToolWindowTabPanel { + selectAgentTab() + val recoveredName = thread.title.ifBlank { "Recovered ${thread.agentId.take(8)}" } + val session = AgentSession( + sessionId = UUID.randomUUID().toString(), + conversation = conversation, + displayName = recoveredName, + runtimeAgentId = thread.agentId, + resumeCheckpointRef = thread.latest + ) + val panel = createNewAgentTab(session, select = true) + return panel + } + + private fun selectAgentTab() { + val toolWindow = ToolWindowManager.getInstance(project).getToolWindow("ProxyAI") ?: return + toolWindow.show() + + val contentManager = toolWindow.contentManager + val agentContent = contentManager.contents.firstOrNull { it.tabName == "Agent" } + ?: contentManager.getContent(0) + if (agentContent != null) { + contentManager.setSelectedContent(agentContent) + } + } + fun getActiveTabPanel(): AgentToolWindowTabPanel? { return tabbedPane.tryFindActiveTabPanel().orElse(null) } @@ -81,8 +113,19 @@ class AgentToolWindowContentManager(private val project: Project) : Disposable { project.service().removeSession(sessionId) } + @Synchronized fun getSession(sessionId: String): AgentSession? = activeSessions[sessionId] + @Synchronized + fun setRuntimeAgentId(sessionId: String, agentId: String?) { + activeSessions[sessionId]?.runtimeAgentId = agentId + } + + @Synchronized + fun setResumeCheckpointRef(sessionId: String, ref: CheckpointRef?) { + activeSessions[sessionId]?.resumeCheckpointRef = ref + } + companion object { fun getInstance(project: Project): AgentToolWindowContentManager { return project.service() diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt index 7d2630c3..caa9707c 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/AgentToolWindowTabPanel.kt @@ -15,6 +15,7 @@ import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.agent.AgentService import ee.carlrobert.codegpt.agent.AgentToolOutputNotifier import ee.carlrobert.codegpt.agent.MessageWithContext +import ee.carlrobert.codegpt.agent.ToolSpecs import ee.carlrobert.codegpt.agent.ToolRunContext import ee.carlrobert.codegpt.agent.rollback.RollbackService import ee.carlrobert.codegpt.conversations.Conversation @@ -26,6 +27,7 @@ import ee.carlrobert.codegpt.settings.service.ModelSelectionService import ee.carlrobert.codegpt.toolwindow.agent.ui.AgentToolWindowLandingPanel import ee.carlrobert.codegpt.toolwindow.agent.ui.RollbackPanel import ee.carlrobert.codegpt.toolwindow.agent.ui.TodoListPanel +import ee.carlrobert.codegpt.toolwindow.agent.ui.ToolCallCard import ee.carlrobert.codegpt.toolwindow.chat.MessageBuilder import ee.carlrobert.codegpt.toolwindow.chat.editor.actions.CopyAction import ee.carlrobert.codegpt.toolwindow.chat.structure.data.PsiStructureRepository @@ -42,6 +44,7 @@ import ee.carlrobert.codegpt.ui.textarea.header.tag.TagManager import ee.carlrobert.codegpt.util.EditorUtil import ee.carlrobert.codegpt.util.coroutines.CoroutineDispatchers import kotlinx.coroutines.launch +import kotlinx.serialization.json.Json import javax.swing.Box import javax.swing.BoxLayout import javax.swing.JComponent @@ -51,6 +54,11 @@ class AgentToolWindowTabPanel( private val project: Project, private val agentSession: AgentSession ) : BorderLayoutPanel(), Disposable { + private val replayJson = Json { + ignoreUnknownKeys = true + isLenient = true + explicitNulls = false + } private val scrollablePanel = ChatToolWindowScrollablePanel() private val tagManager = TagManager() @@ -106,13 +114,13 @@ class AgentToolWindowTabPanel( sessionIdProvider = { sessionId }, conversationIdProvider = { conversation.id } ) - private lateinit var rollbackPanel: RollbackPanel + private var rollbackPanel: RollbackPanel private val todoListPanel = TodoListPanel() private val projectMessageBusConnection = project.messageBus.connect() private val appMessageBusConnection = ApplicationManager.getApplication().messageBus.connect() private val rollbackService = RollbackService.getInstance(project) - private val agentEventHandler = AgentEventHandler( + private val eventHandler = AgentEventHandler( project = project, sessionId = sessionId, agentApprovalManager = AgentApprovalManager(project), @@ -150,10 +158,12 @@ class AgentToolWindowTabPanel( if (conversation.messages.isEmpty()) { displayLandingView() + } else { + displayRecoveredConversation() } userInputPanel.setStopEnabled(false) - Disposer.register(this, agentEventHandler) + Disposer.register(this, eventHandler) } private fun setupMessageBusSubscriptions() { @@ -172,7 +182,7 @@ class AgentToolWindowTabPanel( object : AgentToolOutputNotifier { override fun toolOutput(toolId: String, text: String, isError: Boolean) { val namespacedToolId = "${sessionId}:${toolId}" - agentEventHandler.handleToolOutput(namespacedToolId, text, isError) + eventHandler.handleToolOutput(namespacedToolId, text, isError) } } ) @@ -220,6 +230,7 @@ class AgentToolWindowTabPanel( private fun handleSubmit(text: String) { if (text.isBlank()) return + scrollablePanel.clearLandingViewIfVisible() agentSession.serviceType = ModelSelectionService.getInstance().getServiceForFeature(FeatureType.AGENT) @@ -233,7 +244,7 @@ class AgentToolWindowTabPanel( agentService.submitMessage( MessageWithContext(text, userInputPanel.getSelectedTags()), - agentEventHandler, + eventHandler, sessionId ) return @@ -271,8 +282,8 @@ class AgentToolWindowTabPanel( messagePanel.add(responsePanel) scrollablePanel.update() - agentEventHandler.resetForNewSubmission() - agentEventHandler.setCurrentResponseBody(responseBody) + eventHandler.resetForNewSubmission() + eventHandler.setCurrentResponseBody(responseBody) loadingLabel.text = CodeGPTBundle.get("toolwindow.chat.loading") loadingLabel.isVisible = true @@ -281,7 +292,7 @@ class AgentToolWindowTabPanel( agentService.clearPendingMessages(sessionId) userInputPanel.setStopEnabled(true) - agentService.submitMessage(message, agentEventHandler, sessionId) + agentService.submitMessage(message, eventHandler, sessionId) } private fun handleCancel() { @@ -308,6 +319,102 @@ class AgentToolWindowTabPanel( scrollablePanel.displayLandingView(createLandingView()) } + private fun displayRecoveredConversation() { + scrollablePanel.clearAll() + conversation.messages.forEach { message -> + val prompt = message.prompt.orEmpty() + val wrapper = scrollablePanel.addMessage(message.id) + val userPanel = UserMessagePanel(project, message, this) + userPanel.addCopyAction { CopyAction.copyToClipboard(prompt) } + wrapper.add(userPanel) + + val responseBody = ChatMessageResponseBody( + project, + false, + false, + false, + false, + false, + this + ) + addRecoveredToolCards(responseBody, message) + responseBody.withResponse(message.response.orEmpty()) + val responsePanel = ResponseMessagePanel().apply { + setResponseContent(responseBody) + } + wrapper.add(responsePanel) + } + + scrollablePanel.update() + scrollablePanel.scrollToBottom() + } + + private fun addRecoveredToolCards(responseBody: ChatMessageResponseBody, message: Message) { + val toolCalls = message.toolCalls ?: return + val toolCallResults = message.toolCallResults ?: emptyMap() + + toolCalls.forEach { toolCall -> + val toolName = toolCall.function.name ?: return@forEach + if (toolName == "TodoWrite") { + return@forEach + } + + val rawArgs = toolCall.function.arguments.orEmpty() + val args = parseRecoveredToolArgs(toolName, rawArgs) + val card = createRecoveredToolCard(toolName, args, rawArgs) + responseBody.addToolStatusPanel(card) + + val rawResult = toolCallResults[toolCall.id] ?: return@forEach + val parsedResult = parseRecoveredToolResult(toolName, rawResult) + val success = inferRecoveredToolSuccess(parsedResult, rawResult) + card.complete(success, parsedResult ?: rawResult) + } + } + + private fun createRecoveredToolCard( + toolName: String, + args: Any, + rawArgs: String + ): ToolCallCard { + return try { + ToolCallCard(project, toolName, args) + } catch (_: Exception) { + val fallbackName = "Recovered $toolName" + val fallbackArgs = rawArgs.ifBlank { "(no arguments)" } + ToolCallCard(project, fallbackName, fallbackArgs) + } + } + + private fun parseRecoveredToolArgs(toolName: String, rawArgs: String): Any { + val payload = rawArgs.trim() + if (payload.isBlank()) return "" + + return ToolSpecs.decodeArgsOrNull( + json = replayJson, + toolName = toolName, + payload = payload + ) ?: payload + } + + private fun parseRecoveredToolResult(toolName: String, rawResult: String): Any? { + val payload = rawResult.trim() + if (payload.isBlank()) return null + + return ToolSpecs.decodeResultOrNull( + json = replayJson, + toolName = toolName, + payload = payload + ) ?: payload + } + + private fun inferRecoveredToolSuccess(parsedResult: Any?, rawResult: String): Boolean { + if (parsedResult != null) { + return parsedResult::class.simpleName != "Error" + } + return !rawResult.contains("failed", ignoreCase = true) && + !rawResult.contains("error", ignoreCase = true) + } + private fun createLandingView(): AgentToolWindowLandingPanel { return AgentToolWindowLandingPanel(project) } @@ -376,8 +483,8 @@ class AgentToolWindowTabPanel( responsePanel.setResponseContent(responseBody) messagePanel.add(responsePanel) - agentEventHandler.resetForNewSubmission() - agentEventHandler.setCurrentResponseBody(responseBody) + eventHandler.resetForNewSubmission() + eventHandler.setCurrentResponseBody(responseBody) scrollablePanel.update() } diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/history/AgentHistoryListPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/history/AgentHistoryListPanel.kt new file mode 100644 index 00000000..f2b7e1f7 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/history/AgentHistoryListPanel.kt @@ -0,0 +1,375 @@ +package ee.carlrobert.codegpt.toolwindow.agent.history + +import com.intellij.ui.DocumentAdapter +import com.intellij.ui.SearchTextField +import com.intellij.ui.components.JBLabel +import com.intellij.ui.components.JBList +import com.intellij.ui.components.JBScrollPane +import com.intellij.util.ui.JBFont +import com.intellij.util.ui.JBUI +import com.intellij.util.ui.UIUtil +import com.intellij.util.ui.components.BorderLayoutPanel +import ee.carlrobert.codegpt.agent.history.AgentHistoryThreadSummary +import ee.carlrobert.codegpt.agent.history.CheckpointRef +import java.awt.BorderLayout +import java.awt.Component +import java.awt.Cursor +import java.awt.Dimension +import java.awt.event.MouseAdapter +import java.awt.event.MouseEvent +import java.time.Instant +import java.time.ZoneId +import java.time.format.DateTimeFormatter +import javax.swing.* +import javax.swing.event.DocumentEvent + +class AgentHistoryListPanel( + private val defaultLimit: Int? = null +) : BorderLayoutPanel() { + + companion object { + val TIMESTAMP_FORMAT: DateTimeFormatter = DateTimeFormatter.ofPattern("MMM d, h:mm a") + private const val MIN_VISIBLE_ROWS = 2 + } + + var onOpen: ((AgentHistoryThreadSummary) -> Unit)? = null + var onLoadPage: ((query: String, offset: Int, limit: Int, onResult: (List, Boolean, Int) -> Unit) -> Unit)? = + null + + private val searchField = SearchTextField() + private val statusLabel = JBLabel().apply { + font = JBFont.small() + foreground = UIUtil.getContextHelpForeground() + border = JBUI.Borders.empty(0, 8, 4, 8) + } + private val listModel = DefaultListModel() + private val historyList = JBList(listModel).apply { + selectionMode = ListSelectionModel.SINGLE_SELECTION + cellRenderer = AgentHistoryCellRenderer() + fixedCellHeight = 54 + visibleRowCount = defaultLimit ?: 5 + border = null + } + private val listScrollPane = JBScrollPane(historyList).apply { + horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER + verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED + border = null + viewport.background = UIUtil.getListBackground() + } + private val emptyPanel = BorderLayoutPanel().apply { + border = JBUI.Borders.empty(12, 8) + addToCenter( + JBLabel("No conversations found", SwingConstants.LEFT).apply { + foreground = UIUtil.getContextHelpForeground() + } + ) + } + + private val pageSize = 50 + private var hasMore = false + private var isLoading = false + private var resetRequested = false + private var currentQuery = "" + private var currentTotal = 0 + private var requestId = 0 + + init { + setupUI() + setupListeners() + updateStatus() + } + + fun reload() { + requestPage(reset = true) + } + + private fun setupUI() { + searchField.textEditor.emptyText.text = "Search checkpoint conversations" + searchField.border = JBUI.Borders.empty(4, 0, 2, 0) + historyList.border = JBUI.Borders.emptyTop(4) + + val topPanel = BorderLayoutPanel() + .addToTop(searchField) + .addToBottom(statusLabel) + addToTop(topPanel) + addToCenter(listScrollPane) + } + + private fun setupListeners() { + searchField.addDocumentListener(object : DocumentAdapter() { + override fun textChanged(e: DocumentEvent) { + val nextQuery = searchField.text.trim() + if (nextQuery == currentQuery && !resetRequested) { + return + } + currentQuery = nextQuery + requestPage(reset = true) + } + }) + + listScrollPane.verticalScrollBar.addAdjustmentListener { + if (!it.valueIsAdjusting && shouldLoadNextPage()) { + requestPage(reset = false) + } + } + + historyList.addMouseListener(object : MouseAdapter() { + override fun mouseClicked(e: MouseEvent) { + if (e.button == MouseEvent.BUTTON1 && e.clickCount == 1) { + openAt(e) + } + } + + override fun mouseExited(e: MouseEvent) { + historyList.cursor = Cursor.getDefaultCursor() + } + }) + + historyList.addMouseMotionListener(object : MouseAdapter() { + override fun mouseMoved(e: MouseEvent) { + val hasItem = indexAt(e) >= 0 + historyList.cursor = if (hasItem) { + Cursor.getPredefinedCursor(Cursor.HAND_CURSOR) + } else { + Cursor.getDefaultCursor() + } + } + }) + + historyList.inputMap.put(KeyStroke.getKeyStroke("ENTER"), "openThread") + historyList.actionMap.put("openThread", object : AbstractAction() { + override fun actionPerformed(e: java.awt.event.ActionEvent?) { + openSelected() + } + }) + } + + private fun requestPage(reset: Boolean) { + val loader = onLoadPage ?: return + if (isLoading) { + if (reset) { + resetRequested = true + } + return + } + if (!reset && !hasMore) { + return + } + + val previousSelection = historyList.selectedValue?.latest + val querySnapshot = currentQuery + val offset = if (reset) 0 else listModel.size + val requestSnapshot = ++requestId + + if (reset) { + listModel.clear() + hasMore = false + currentTotal = 0 + updateListViewportSize(0) + updateCenterView(isEmpty = false) + } + + isLoading = true + updateStatus() + + loader(querySnapshot, offset, pageSize) { pageItems, pageHasMore, pageTotal -> + if (requestSnapshot != requestId) { + return@loader + } + + if (querySnapshot != currentQuery) { + isLoading = false + if (resetRequested) { + resetRequested = false + } + requestPage(reset = true) + return@loader + } + + if (reset) { + listModel.clear() + } + + pageItems.forEach(listModel::addElement) + hasMore = pageHasMore + currentTotal = pageTotal + isLoading = false + + restoreSelection(previousSelection) + updateListViewportSize(listModel.size) + updateCenterView(isEmpty = listModel.isEmpty) + updateStatus() + + if (resetRequested) { + resetRequested = false + requestPage(reset = true) + return@loader + } + + if (shouldLoadNextPage()) { + requestPage(reset = false) + } + } + } + + private fun restoreSelection(selected: CheckpointRef?) { + if (listModel.isEmpty) { + return + } + + if (selected != null) { + val idx = (0 until listModel.size()).firstOrNull { + val item = listModel.getElementAt(it) + item.latest.agentId == selected.agentId && item.latest.checkpointId == selected.checkpointId + } + if (idx != null) { + historyList.selectedIndex = idx + historyList.ensureIndexIsVisible(idx) + return + } + } + + historyList.selectedIndex = 0 + } + + private fun shouldLoadNextPage(): Boolean { + if (!hasMore || isLoading) { + return false + } + val scrollbar = listScrollPane.verticalScrollBar + val threshold = historyList.fixedCellHeight.takeIf { it > 0 } ?: 54 + return scrollbar.value + scrollbar.visibleAmount >= scrollbar.maximum - threshold + } + + private fun updateListViewportSize(itemCount: Int) { + val maxRows = defaultLimit ?: 5 + val visibleRows = itemCount.coerceIn(MIN_VISIBLE_ROWS, maxRows) + val rowHeight = historyList.fixedCellHeight.takeIf { it > 0 } ?: 54 + val height = rowHeight * visibleRows + listScrollPane.minimumSize = Dimension(0, rowHeight * MIN_VISIBLE_ROWS) + listScrollPane.preferredSize = Dimension(0, height) + } + + private fun updateStatus() { + val loaded = listModel.size + val queryActive = currentQuery.isNotBlank() + statusLabel.text = when { + isLoading && loaded == 0 -> "Loading..." + currentTotal == 0 && queryActive -> "No results" + currentTotal == 0 -> "No threads" + isLoading -> "Showing $loaded of $currentTotal..." + else -> "Showing $loaded of $currentTotal" + } + } + + private fun updateCenterView(isEmpty: Boolean) { + val current = if (componentCount > 1) getComponent(1) else null + val expected = if (isEmpty) emptyPanel else listScrollPane + if (current === expected) return + if (current != null) remove(current) + addToCenter(expected) + revalidate() + repaint() + } + + private fun openSelected() { + historyList.selectedValue?.let { onOpen?.invoke(it) } + } + + private fun openAt(e: MouseEvent) { + val index = indexAt(e) + if (index < 0) { + return + } + historyList.selectedIndex = index + openSelected() + } + + private fun indexAt(e: MouseEvent): Int { + val index = historyList.locationToIndex(e.point) + if (index < 0) { + return -1 + } + val bounds = historyList.getCellBounds(index, index) ?: return -1 + return if (bounds.contains(e.point)) index else -1 + } + + private class AgentHistoryCellRenderer : JPanel(BorderLayout()), + ListCellRenderer { + + private val titleLabel = JBLabel().apply { + font = JBFont.label().asBold() + } + private val previewLabel = JBLabel().apply { + font = JBFont.small() + } + private val metaLabel = JBLabel().apply { + font = JBFont.small() + } + + init { + isOpaque = true + border = JBUI.Borders.empty(4, 8) + + val header = JPanel(BorderLayout()).apply { + isOpaque = false + add(titleLabel, BorderLayout.CENTER) + add(metaLabel, BorderLayout.EAST) + } + + val body = JPanel(BorderLayout()).apply { + isOpaque = false + add(previewLabel, BorderLayout.CENTER) + } + + val content = JPanel().apply { + isOpaque = false + layout = BoxLayout(this, BoxLayout.Y_AXIS) + add(header) + add(body) + } + + add(content, BorderLayout.CENTER) + } + + override fun getListCellRendererComponent( + list: JList, + value: AgentHistoryThreadSummary, + index: Int, + isSelected: Boolean, + cellHasFocus: Boolean + ): Component { + titleLabel.text = + ellipsize(value.title.ifBlank { "Recovered ${value.agentId.take(8)}" }, 44) + previewLabel.text = ellipsize( + value.preview.ifBlank { "No assistant response in checkpoint" }, + 56 + ) + metaLabel.text = formatTimestamp(value.latestCreatedAt.toString()) + + background = + if (isSelected) UIUtil.getListSelectionBackground(true) else UIUtil.getListBackground() + titleLabel.foreground = + if (isSelected) UIUtil.getListSelectionForeground(true) else UIUtil.getLabelForeground() + previewLabel.foreground = + if (isSelected) UIUtil.getListSelectionForeground(true) else UIUtil.getContextHelpForeground() + metaLabel.foreground = + if (isSelected) UIUtil.getListSelectionForeground(true) else UIUtil.getContextHelpForeground() + + return this + } + } +} + +private fun formatTimestamp(raw: String): String { + return runCatching { + val instant = Instant.parse(raw) + instant.atZone(ZoneId.systemDefault()).format(AgentHistoryListPanel.TIMESTAMP_FORMAT) + }.getOrElse { raw } +} + +private fun ellipsize(value: String, maxChars: Int): String { + val cleaned = value.replace("\\s+".toRegex(), " ").trim() + if (cleaned.length <= maxChars) return cleaned + return cleaned.take(maxChars - 1) + "…" +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt index d9d1a425..f2a5a18b 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/toolwindow/agent/ui/AgentToolWindowLandingPanel.kt @@ -1,6 +1,9 @@ package ee.carlrobert.codegpt.toolwindow.agent.ui +import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.runInEdt import com.intellij.openapi.command.WriteCommandAction +import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.options.ShowSettingsUtil @@ -13,11 +16,18 @@ import com.intellij.ui.components.ActionLink import com.intellij.ui.components.JBLabel import com.intellij.util.ui.JBUI import com.intellij.util.ui.components.BorderLayoutPanel +import ee.carlrobert.codegpt.agent.ProxyAIAgent.loadProjectInstructions +import ee.carlrobert.codegpt.agent.history.AgentCheckpointConversationMapper +import ee.carlrobert.codegpt.agent.history.AgentCheckpointHistoryService +import ee.carlrobert.codegpt.agent.history.AgentHistoryThreadSummary import ee.carlrobert.codegpt.settings.GeneralSettings import ee.carlrobert.codegpt.settings.agents.SubagentsConfigurable import ee.carlrobert.codegpt.tokens.TokenComputationService +import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager +import ee.carlrobert.codegpt.toolwindow.agent.history.AgentHistoryListPanel import ee.carlrobert.codegpt.toolwindow.ui.ResponseMessagePanel import ee.carlrobert.codegpt.ui.UIUtil.createTextPane +import kotlinx.coroutines.runBlocking import java.awt.BorderLayout import java.awt.Color import java.awt.Desktop @@ -37,6 +47,10 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag private val logger = thisLogger() } + private val historyService = project.service() + private val historyListPanel = AgentHistoryListPanel(defaultLimit = 5) + private var refreshHistory = true + private fun createLabel(text: String) = JBLabel(text) private fun createLink(text: String, onClick: () -> Unit) = ActionLink(text) { onClick() } @@ -47,7 +61,12 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag } init { + historyListPanel.onOpen = { thread -> openCheckpointThread(thread) } + historyListPanel.onLoadPage = { query, offset, limit, onResult -> + loadHistoryPage(query, offset, limit, onResult) + } addContent(buildContent()) + loadHistory() } private fun buildContent(): JPanel { @@ -59,11 +78,20 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag } private fun centerPanel(): JPanel { + val panel = BorderLayoutPanel() + panel.addToTop(topSectionsPanel()) + panel.addToCenter(previousChatsPanel()) + return panel + } + + private fun topSectionsPanel(): JPanel { val panel = JPanel() panel.layout = BoxLayout(panel, BoxLayout.Y_AXIS) + panel.alignmentX = LEFT_ALIGNMENT panel.add(actionsListPanel()) panel.add(Box.createVerticalStrut(12)) panel.add(projectInfoPanel()) + panel.add(Box.createVerticalStrut(12)) return panel } @@ -210,10 +238,25 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag container.add(Box.createVerticalStrut(2)) container.add(bottomRow) } + panel.add(container) return panel } + private fun previousChatsPanel(): JPanel { + val panel = BorderLayoutPanel().apply { + border = JBUI.Borders.empty() + } + panel.addToTop(JPanel().apply { + layout = BoxLayout(this, BoxLayout.Y_AXIS) + alignmentX = LEFT_ALIGNMENT + isOpaque = false + add(createLabel("Previous Chats")) + add(Box.createVerticalStrut(6)) + }) + panel.addToCenter(historyListPanel) + return panel + } private fun findProxyAiFile() = project.basePath?.let { @@ -240,10 +283,66 @@ class AgentToolWindowLandingPanel(private val project: Project) : ResponseMessag private fun refresh() { removeAll() addContent(buildContent()) + loadHistory() revalidate() repaint() } + private fun loadHistory() { + refreshHistory = true + historyListPanel.reload() + } + + private fun loadHistoryPage( + query: String, + offset: Int, + limit: Int, + onResult: (List, Boolean, Int) -> Unit + ) { + val shouldRefresh = offset == 0 && refreshHistory + ApplicationManager.getApplication().executeOnPooledThread { + val page = runCatching { + runBlocking { + historyService.listThreadsPage( + query = query, + offset = offset, + limit = limit, + refresh = shouldRefresh + ) + } + } + .onFailure { logger.warn("Failed to load checkpoint history", it) } + .getOrNull() + + runInEdt { + if (page != null) { + refreshHistory = false + } + onResult(page?.items.orEmpty(), page?.hasMore == true, page?.total ?: 0) + } + } + } + + private fun openCheckpointThread(thread: AgentHistoryThreadSummary) { + ApplicationManager.getApplication().executeOnPooledThread { + val conversation = runCatching { + val checkpoint = runBlocking { historyService.loadCheckpoint(thread.latest) } + ?: return@executeOnPooledThread + AgentCheckpointConversationMapper.toConversation( + checkpoint = checkpoint, + projectInstructions = loadProjectInstructions(project.basePath) + ) + }.onFailure { + logger.warn("Failed to open checkpoint thread ${thread.agentId}", it) + }.getOrNull() ?: return@executeOnPooledThread + + ApplicationManager.getApplication().invokeLater { + project.service() + .openCheckpointConversation(thread, conversation) + } + } + } + private fun welcomeMessage(): String { val name = GeneralSettings.getCurrentState().displayName return """ diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt index f6b3dae6..db506cde 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/PromptTextField.kt @@ -30,6 +30,7 @@ import com.intellij.openapi.wm.ToolWindow import com.intellij.openapi.wm.ToolWindowManager import com.intellij.ui.EditorTextField import com.intellij.ui.JBColor +import com.intellij.util.IncorrectOperationException import com.intellij.util.ui.JBUI import ee.carlrobert.codegpt.CodeGPTBundle import ee.carlrobert.codegpt.CodeGPTKeys.IS_PROMPT_TEXT_FIELD_DOCUMENT @@ -64,6 +65,8 @@ class PromptTextField( private val coroutineScope = CoroutineScope(Dispatchers.Main + SupervisorJob()) private val lookupManager = PromptTextFieldLookupManager(project, onLookupAdded) private val searchManager = SearchManager(project, tagManager, featureType) + @Volatile + private var isFieldDisposed = false private var mouseClickListener: MouseAdapter? = null private var mouseMotionListener: MouseMotionAdapter? = null @@ -85,6 +88,9 @@ class PromptTextField( } override fun onEditorAdded(editor: Editor) { + if (isFieldDisposed) { + return + } IdeEventQueue.getInstance().addDispatcher( PromptTextFieldEventDispatcher(dispatcherId, onBackSpace) { event -> val shown = lookup?.let { it.isShown && !it.isLookupDisposed } == true @@ -279,6 +285,9 @@ class PromptTextField( override fun createEditor(): EditorEx { val editorEx = super.createEditor() + if (isFieldDisposed) { + return editorEx + } editorEx.settings.isUseSoftWraps = true editorEx.backgroundColor = service().globalScheme.defaultBackground setupDocumentListener(editorEx) @@ -295,6 +304,7 @@ class PromptTextField( } override fun dispose() { + isFieldDisposed = true showSuggestionsJob?.cancel() lastSearchResults = null clearPlaceholders() @@ -452,13 +462,24 @@ class PromptTextField( } private fun setupDocumentListener(editor: EditorEx) { - editor.document.addDocumentListener(object : DocumentListener { - override fun documentChanged(event: DocumentEvent) { - adjustHeight(editor) - onTextChanged(event.document.text) - handleDocumentChange(event) + if (isFieldDisposed) { + return + } + + try { + editor.document.addDocumentListener(object : DocumentListener { + override fun documentChanged(event: DocumentEvent) { + adjustHeight(editor) + onTextChanged(event.document.text) + handleDocumentChange(event) + } + }, this) + } catch (e: IncorrectOperationException) { + if (!isFieldDisposed) { + throw e } - }, this) + logger.debug("Skipping document listener registration for disposed PromptTextField", e) + } } private fun handleDocumentChange(event: DocumentEvent) { diff --git a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt index 9677afb6..fe2e3324 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/ui/textarea/header/tag/TagDetails.kt @@ -11,12 +11,16 @@ import ee.carlrobert.codegpt.mcp.McpTool import ee.carlrobert.codegpt.settings.prompts.PersonaDetails import ee.carlrobert.codegpt.ui.DocumentationDetails import git4idea.GitCommit +import kotlinx.serialization.Contextual +import kotlinx.serialization.Serializable import java.util.* import javax.swing.Icon +@Serializable sealed class TagDetails( val name: String, val icon: Icon? = null, + @Contextual val id: UUID = UUID.randomUUID(), val createdOn: Long = System.currentTimeMillis(), val isRemovable: Boolean = true @@ -125,7 +129,7 @@ class EditorSelectionTagDetails( data class DocumentationTagDetails(var documentationDetails: DocumentationDetails) : TagDetails(documentationDetails.name, AllIcons.Toolwindows.Documentation) { - override fun getTooltipText(): String? = documentationDetails.url + override fun getTooltipText(): String = documentationDetails.url } data class PersonaTagDetails(var personaDetails: PersonaDetails) : @@ -135,7 +139,7 @@ data class PersonaTagDetails(var personaDetails: PersonaDetails) : data class GitCommitTagDetails(var gitCommit: GitCommit) : TagDetails(gitCommit.id.asString().take(6), AllIcons.Vcs.CommitNode) { - override fun getTooltipText(): String? = gitCommit.fullMessage + override fun getTooltipText(): String = gitCommit.fullMessage } class CurrentGitChangesTagDetails : diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt new file mode 100644 index 00000000..94403fe8 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/AgentServiceIntegrationTest.kt @@ -0,0 +1,324 @@ +package ee.carlrobert.codegpt.agent + +import ai.koog.agents.core.agent.AIAgent +import ai.koog.agents.core.agent.AIAgentService +import ai.koog.agents.core.agent.config.AIAgentConfig +import ai.koog.agents.core.tools.ToolRegistry +import ai.koog.agents.snapshot.feature.AgentCheckpointData +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import ai.koog.prompt.dsl.Prompt +import ai.koog.prompt.dsl.prompt +import ai.koog.prompt.executor.model.PromptExecutor +import ai.koog.prompt.llm.LLModel +import ai.koog.prompt.message.Message +import ai.koog.prompt.message.RequestMetaInfo +import ai.koog.prompt.message.ResponseMetaInfo +import ai.koog.prompt.streaming.StreamFrame +import com.intellij.openapi.components.service +import com.intellij.openapi.project.Project +import ee.carlrobert.codegpt.agent.history.CheckpointRef +import ee.carlrobert.codegpt.conversations.Conversation +import ee.carlrobert.codegpt.settings.service.ModelSelectionService +import ee.carlrobert.codegpt.settings.service.ServiceType +import ee.carlrobert.codegpt.toolwindow.agent.AgentSession +import ee.carlrobert.codegpt.toolwindow.agent.AgentToolWindowContentManager +import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.runBlocking +import kotlinx.datetime.Clock +import kotlinx.serialization.json.JsonNull +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.groups.Tuple.tuple +import testsupport.IntegrationTest +import java.util.* +import java.util.concurrent.ConcurrentHashMap +import kotlin.io.path.Path + +class AgentServiceIntegrationTest : IntegrationTest() { + + private lateinit var agentService: AgentService + private lateinit var contentManager: AgentToolWindowContentManager + private lateinit var originalRuntimeFactory: AgentRuntimeFactory + private val createdSessions = mutableListOf() + private val noopEvents = object : AgentEvents { + override fun onQueuedMessagesResolved() = Unit + } + private val runTimeoutMillis = 5_000L + + override fun setUp() { + super.setUp() + agentService = project.service() + contentManager = project.service() + originalRuntimeFactory = agentService.runtimeFactory + } + + override fun tearDown() { + agentService.runtimeFactory = originalRuntimeFactory + createdSessions.forEach { sessionId -> + runCatching { contentManager.removeSession(sessionId) } + } + createdSessions.clear() + super.tearDown() + } + + fun testSubmitMessageReusesRuntimeAndUpdatesResumeCheckpoint() { + val sessionId = createSession("agent-runtime-reuse") + val factory = RecordingRuntimeFactory(agentModel()) + agentService.runtimeFactory = factory + val firstMessage = MessageWithContext("First request") + val secondMessage = MessageWithContext("Second request") + + submitAndAwait(sessionId, firstMessage) + val firstSnapshot = snapshot(sessionId) + submitAndAwait(sessionId, secondMessage) + val secondSnapshot = snapshot(sessionId) + + val runtimeAgentId = firstSnapshot.runtimeAgentId + val managedService = factory.createdServices.single() + assertThat(runtimeAgentId).isNotBlank() + assertThat(factory.createdServices).hasSize(1) + assertThat(listOf(firstSnapshot, secondSnapshot)) + .extracting("runtimeAgentId", "resumeCheckpointRef.agentId") + .containsExactly( + tuple(runtimeAgentId, runtimeAgentId), + tuple(runtimeAgentId, runtimeAgentId) + ) + assertThat(firstSnapshot.resumeCheckpointRef?.checkpointId).isNotBlank() + assertThat(secondSnapshot.resumeCheckpointRef?.checkpointId).isNotBlank() + assertThat(secondSnapshot.resumeCheckpointRef?.checkpointId) + .isNotEqualTo(firstSnapshot.resumeCheckpointRef?.checkpointId) + assertThat(managedService.createdAgentIds).hasSize(2) + assertThat(managedService.createdAgentIds.last()).isEqualTo(runtimeAgentId) + } + + fun testSubmitMessageRebuildsRuntimeWhenMcpSelectionChanges() { + val sessionId = createSession("agent-runtime-mcp") + val factory = RecordingRuntimeFactory(agentModel()) + agentService.runtimeFactory = factory + val withoutMcp = MessageWithContext("Without MCP") + val mcpTag = McpTagDetails(serverId = "server-1", serverName = "Server 1") + val withMcp = MessageWithContext("With MCP", tags = listOf(mcpTag)) + + submitAndAwait(sessionId, withoutMcp) + submitAndAwait(sessionId, withMcp) + + assertThat(factory.createdServices).hasSize(2) + assertThat(factory.createdServices) + .extracting("closeAllCalls") + .containsExactly(1, 0) + } + + fun testGetCheckpointFallsBackToRuntimeAgentIdWhenResumeRefIsStale() { + val sessionId = createSession("agent-checkpoint-fallback") + val runtimeAgentId = "runtime-agent-${UUID.randomUUID()}" + val checkpointId = "checkpoint-${UUID.randomUUID()}" + val checkpointStorage = JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai")) + val staleRef = CheckpointRef(agentId = "missing-agent", checkpointId = "missing-checkpoint") + runBlocking { + checkpointStorage.saveCheckpoint( + runtimeAgentId, + AgentCheckpointData( + checkpointId = checkpointId, + createdAt = Clock.System.now(), + nodePath = "$runtimeAgentId/single_run/nodeExecuteTool", + lastInput = JsonNull, + messageHistory = listOf( + Message.User("hello", RequestMetaInfo.Empty), + Message.Assistant("world", ResponseMetaInfo.Empty), + ), + version = 0, + ) + ) + } + contentManager.setRuntimeAgentId(sessionId, runtimeAgentId) + contentManager.setResumeCheckpointRef(sessionId, staleRef) + + val checkpoint = runBlocking { agentService.getCheckpoint(sessionId) } + + assertThat(checkpoint).isNotNull + assertThat(checkpoint!!.checkpointId).isEqualTo(checkpointId) + assertThat(contentManager.getSession(sessionId)!!.resumeCheckpointRef) + .isEqualTo(CheckpointRef(runtimeAgentId, checkpointId)) + } + + private fun submitAndAwait( + sessionId: String, + message: MessageWithContext + ) { + agentService.submitMessage(message, noopEvents, sessionId) + awaitSessionToFinish(sessionId) + } + + private fun awaitSessionToFinish(sessionId: String) { + val deadline = System.currentTimeMillis() + runTimeoutMillis + while (agentService.isSessionRunning(sessionId) && System.currentTimeMillis() < deadline) { + Thread.sleep(20) + } + assertThat(agentService.isSessionRunning(sessionId)).isFalse + } + + private fun createSession(prefix: String): String { + val sessionId = "$prefix-${UUID.randomUUID()}" + contentManager.createNewAgentTab( + AgentSession( + sessionId = sessionId, + conversation = Conversation() + ), + select = false + ) + createdSessions.add(sessionId) + return sessionId + } + + private fun snapshot(sessionId: String): SessionSnapshot { + val session = requireNotNull(contentManager.getSession(sessionId)) + return SessionSnapshot( + runtimeAgentId = session.runtimeAgentId, + resumeCheckpointRef = session.resumeCheckpointRef + ) + } + + private fun agentModel(): LLModel { + return project.service().getAgentModel() + } + + private data class SessionSnapshot( + val runtimeAgentId: String?, + val resumeCheckpointRef: CheckpointRef? + ) +} + +private class RecordingRuntimeFactory( + private val model: LLModel +) : AgentRuntimeFactory { + + val createdServices = mutableListOf() + + override fun create( + project: Project, + checkpointStorage: JVMFilePersistenceStorageProvider, + provider: ServiceType, + events: AgentEvents, + sessionId: String, + pendingMessages: ConcurrentHashMap> + ): AIAgentService> { + val service = FakeManagedAgentService(checkpointStorage, model) + createdServices.add(service) + return service + } +} + +private class FakeManagedAgentService( + private val checkpointStorage: JVMFilePersistenceStorageProvider, + model: LLModel +) : AIAgentService { + + override val promptExecutor: PromptExecutor = NoopPromptExecutor + override val agentConfig: AIAgentConfig = AIAgentConfig( + prompt = prompt("fake-service") {}, + model = model, + maxAgentIterations = 1 + ) + override val toolRegistry: ToolRegistry = ToolRegistry.EMPTY + + private val agents = linkedMapOf() + val createdAgentIds = mutableListOf() + var closeAllCalls: Int = 0 + + override suspend fun createAgent( + id: String?, + additionalToolRegistry: ToolRegistry, + agentConfig: AIAgentConfig, + clock: Clock + ): FakeManagedAgent { + val actualId = id ?: UUID.randomUUID().toString() + val agent = FakeManagedAgent(actualId, agentConfig, checkpointStorage) + agents[actualId] = agent + createdAgentIds.add(actualId) + return agent + } + + override suspend fun removeAgent(agent: FakeManagedAgent): Boolean { + return removeAgentWithId(agent.id) + } + + override suspend fun removeAgentWithId(id: String): Boolean { + return agents.remove(id) != null + } + + override suspend fun agentById(id: String): FakeManagedAgent? = agents[id] + + override suspend fun listAllAgents(): List = agents.values.toList() + + override suspend fun listActiveAgents(): List = emptyList() + + override suspend fun listInactiveAgents(): List = agents.values.toList() + + override suspend fun listFinishedAgents(): List = agents.values.toList() + + override suspend fun closeAll() { + closeAllCalls += 1 + super.closeAll() + } +} + +private class FakeManagedAgent( + override val id: String, + override val agentConfig: AIAgentConfig, + private val checkpointStorage: JVMFilePersistenceStorageProvider +) : AIAgent { + + private var state: AIAgent.Companion.State = AIAgent.Companion.State.NotStarted() + + override suspend fun getState(): AIAgent.Companion.State = state + + override suspend fun run(agentInput: MessageWithContext): String { + state = AIAgent.Companion.State.Starting() + + val latest = checkpointStorage.getLatestCheckpoint(id) + val checkpoint = AgentCheckpointData( + checkpointId = UUID.randomUUID().toString(), + createdAt = Clock.System.now(), + nodePath = "$id/single_run/nodeExecuteTool", + lastInput = JsonNull, + messageHistory = listOf( + Message.User(agentInput.text, RequestMetaInfo.Empty), + Message.Assistant("ok", ResponseMetaInfo.Empty) + ), + version = (latest?.version ?: -1) + 1 + ) + checkpointStorage.saveCheckpoint(id, checkpoint) + + state = AIAgent.Companion.State.Finished("ok") + return "ok" + } + + override suspend fun close() = Unit +} + +private object NoopPromptExecutor : PromptExecutor { + override suspend fun execute( + prompt: Prompt, + model: LLModel, + tools: List + ): List { + throw UnsupportedOperationException("NoopPromptExecutor should not be used in tests") + } + + override fun executeStreaming( + prompt: Prompt, + model: LLModel, + tools: List + ): Flow { + throw UnsupportedOperationException("NoopPromptExecutor should not be used in tests") + } + + override suspend fun moderate( + prompt: Prompt, + model: LLModel + ): ai.koog.prompt.dsl.ModerationResult { + throw UnsupportedOperationException("NoopPromptExecutor should not be used in tests") + } + + override fun close() = Unit +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt new file mode 100644 index 00000000..074970ea --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapperTest.kt @@ -0,0 +1,171 @@ +package ee.carlrobert.codegpt.agent.history + +import ai.koog.agents.snapshot.feature.AgentCheckpointData +import ai.koog.prompt.message.Message +import ai.koog.prompt.message.RequestMetaInfo +import ai.koog.prompt.message.ResponseMetaInfo +import ee.carlrobert.codegpt.conversations.Conversation +import kotlinx.datetime.Clock +import kotlinx.datetime.Instant +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.groups.Tuple.tuple +import kotlin.test.Test + +class AgentCheckpointConversationMapperTest { + private val actualUserTurn = "Actual user prompt" to "Actual response" + + @Test + fun `maps user and assistant turns into conversation messages`() { + val checkpoint = checkpoint( + listOf( + Message.System("System", RequestMetaInfo.Empty), + Message.User("First prompt", RequestMetaInfo.Empty), + Message.Assistant("First response", ResponseMetaInfo.Empty), + Message.User("Second prompt", RequestMetaInfo.Empty), + Message.Assistant("Second response", ResponseMetaInfo.Empty) + ) + ) + + val conversation = AgentCheckpointConversationMapper.toConversation(checkpoint, null) + + assertThat(conversation.messages) + .extracting("prompt", "response") + .containsExactly( + tuple("First prompt", "First response"), + tuple("Second prompt", "Second response") + ) + } + + @Test + fun `maps tool calls and results into structured message fields`() { + val checkpoint = checkpoint( + listOf( + Message.User("Run the command", RequestMetaInfo.Empty), + Message.Assistant("Working on it", ResponseMetaInfo.Empty), + Message.Tool.Call( + id = "tool-1", + tool = "Bash", + content = """{"command":"ls"}""", + metaInfo = ResponseMetaInfo.Empty + ), + Message.Tool.Result( + id = "tool-1", + tool = "Bash", + content = "file1.txt", + metaInfo = RequestMetaInfo.Empty + ) + ) + ) + + val conversation = AgentCheckpointConversationMapper.toConversation(checkpoint, null) + val mapped = conversation.messages.single() + + assertThat(mapped.response).isEqualTo("Working on it") + assertThat(mapped.toolCalls.orEmpty()) + .extracting("function.name", "function.arguments") + .containsExactly(tuple("Bash", """{"command":"ls"}""")) + assertThat(mapped.toolCallResults).containsEntry("tool-1", "file1.txt") + } + + @Test + fun `filters project instructions from user turns`() { + val projectInstructions = "Follow project instructions carefully." + val checkpoint = checkpoint( + listOf( + Message.User(projectInstructions, RequestMetaInfo.Empty), + Message.User("Actual user prompt", RequestMetaInfo.Empty), + Message.Assistant("Actual response", ResponseMetaInfo.Empty) + ) + ) + + val conversation = AgentCheckpointConversationMapper.toConversation( + checkpoint = checkpoint, + projectInstructions = projectInstructions + ) + + assertSingleTurn(conversation, actualUserTurn) + } + + @Test + fun `filters cacheable user messages from user turns`() { + val checkpoint = checkpoint( + listOf( + Message.User( + content = "Project-level instruction payload", + metaInfo = cacheableMetaInfo() + ), + Message.User("Actual user prompt", RequestMetaInfo.Empty), + Message.Assistant("Actual response", ResponseMetaInfo.Empty) + ) + ) + + val conversation = AgentCheckpointConversationMapper.toConversation( + checkpoint = checkpoint, + projectInstructions = null + ) + + assertSingleTurn(conversation, actualUserTurn) + } + + @Test + fun `merges multiple tool results for the same tool call`() { + val checkpoint = checkpoint( + listOf( + Message.User("Run the command", RequestMetaInfo.Empty), + Message.Tool.Call( + id = "tool-1", + tool = "Bash", + content = """{"command":"ls"}""", + metaInfo = ResponseMetaInfo.Empty + ), + Message.Tool.Result( + id = "tool-1", + tool = "Bash", + content = "file1.txt", + metaInfo = RequestMetaInfo.Empty + ), + Message.Tool.Result( + id = "tool-1", + tool = "Bash", + content = "file2.txt", + metaInfo = RequestMetaInfo.Empty + ) + ) + ) + + val conversation = AgentCheckpointConversationMapper.toConversation(checkpoint, null) + + assertThat(conversation.messages.single().toolCallResults) + .containsEntry("tool-1", "file1.txt\n\nfile2.txt") + } + + private fun cacheableMetaInfo(): RequestMetaInfo { + return RequestMetaInfo( + timestamp = Clock.System.now(), + metadata = JsonObject(mapOf("cacheable" to JsonPrimitive(true))) + ) + } + + private fun assertSingleTurn( + conversation: Conversation, + expectedTurn: Pair + ) { + assertThat(conversation.messages) + .extracting("prompt", "response") + .containsExactly(tuple(expectedTurn.first, expectedTurn.second)) + } + + private fun checkpoint(history: List): AgentCheckpointData { + return AgentCheckpointData( + checkpointId = "checkpoint-1", + createdAt = Instant.parse("2026-02-04T00:00:00Z"), + nodePath = "agent/single_run/nodeExecuteTool", + lastInput = JsonNull, + messageHistory = history, + version = 0 + ) + } +} diff --git a/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt new file mode 100644 index 00000000..e9cae785 --- /dev/null +++ b/src/test/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryServiceTest.kt @@ -0,0 +1,301 @@ +package ee.carlrobert.codegpt.agent.history + +import ai.koog.agents.snapshot.feature.AgentCheckpointData +import ai.koog.agents.snapshot.providers.file.JVMFilePersistenceStorageProvider +import ai.koog.prompt.message.Message +import ai.koog.prompt.message.RequestMetaInfo +import ai.koog.prompt.message.ResponseMetaInfo +import com.intellij.openapi.components.service +import kotlinx.coroutines.runBlocking +import kotlinx.datetime.Clock +import kotlinx.datetime.Instant +import kotlinx.serialization.json.JsonNull +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.JsonPrimitive +import org.assertj.core.api.Assertions.assertThat +import org.assertj.core.groups.Tuple.tuple +import testsupport.IntegrationTest +import kotlin.io.path.Path +import java.util.UUID + +class AgentCheckpointHistoryServiceTest : IntegrationTest() { + + fun testListsThreadsAndPicksLatestNonTombstoneCheckpoint() { + runBlocking { + val storage = storage() + val agentAId = "agent-a-${UUID.randomUUID()}" + val agentBId = "agent-b-${UUID.randomUUID()}" + + storage.saveCheckpoint( + agentAId, + checkpoint( + checkpointId = "a-1", + createdAt = Instant.parse("2026-02-04T00:00:00Z"), + nodePath = "$agentAId/single_run/__finish__", + history = listOf( + Message.User("Initial prompt", RequestMetaInfo.Empty), + Message.Assistant("Initial response", ResponseMetaInfo.Empty) + ) + ) + ) + storage.saveCheckpoint( + agentAId, + checkpoint( + checkpointId = "a-2", + createdAt = Instant.parse("2026-02-04T00:05:00Z"), + nodePath = "$agentAId/single_run/nodeExecuteTool", + history = listOf( + Message.User("Latest prompt", RequestMetaInfo.Empty), + Message.Assistant("Latest response", ResponseMetaInfo.Empty) + ) + ) + ) + storage.saveCheckpoint( + agentAId, + checkpoint( + checkpointId = "a-3-tombstone", + createdAt = Instant.parse("2026-02-04T00:06:00Z"), + nodePath = "$agentAId/single_run/__finish__", + history = emptyList(), + tombstone = true + ) + ) + storage.saveCheckpoint( + agentBId, + checkpoint( + checkpointId = "b-1", + createdAt = Instant.parse("2026-02-04T00:01:00Z"), + nodePath = "$agentBId/single_run/__start__", + history = listOf(Message.User("Thread B", RequestMetaInfo.Empty)) + ) + ) + + val scopedThreads = historyService() + .listThreads() + .filter { it.agentId == agentAId || it.agentId == agentBId } + + assertThat(scopedThreads) + .extracting( + "agentId", + "latest.checkpointId", + "status", + "runCount", + "title", + "preview" + ) + .containsExactlyInAnyOrder( + tuple( + agentAId, + "a-2", + ThreadStatus.PARTIAL, + 2, + "Latest prompt", + "Latest response" + ), + tuple(agentBId, "b-1", ThreadStatus.UNKNOWN, 1, "Thread B", "") + ) + } + } + + fun testLoadsCheckpointByCheckpointReference() { + runBlocking { + val storage = storage() + val agentId = "agent-z-${UUID.randomUUID()}" + val persisted = checkpoint( + checkpointId = "a-9", + createdAt = Instant.parse("2026-02-04T00:15:00Z"), + nodePath = "$agentId/single_run/__finish__", + history = listOf(Message.User("Prompt", RequestMetaInfo.Empty)) + ) + storage.saveCheckpoint(agentId, persisted) + + val loaded = historyService().loadCheckpoint(CheckpointRef(agentId, "a-9")) + + assertThat(loaded).isNotNull + assertThat(loaded!!.checkpointId).isEqualTo("a-9") + assertThat(loaded.messageHistory).hasSize(1) + } + } + + fun testLoadResumeCheckpointFallsBackWhenLatestIsFinishNode() { + runBlocking { + val storage = storage() + val agentId = "agent-r-${UUID.randomUUID()}" + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "r-1", + createdAt = Instant.parse("2026-02-04T01:00:00Z"), + nodePath = "$agentId/single_run/nodeExecuteTool", + history = listOf(Message.User("First", RequestMetaInfo.Empty)) + ) + ) + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "r-2", + createdAt = Instant.parse("2026-02-04T01:01:00Z"), + nodePath = "$agentId/single_run/__finish__", + history = listOf(Message.Assistant("Done", ResponseMetaInfo.Empty)) + ) + ) + + val resume = historyService().loadResumeCheckpoint(CheckpointRef(agentId, "r-2")) + + assertThat(resume).isNotNull + assertThat(resume!!.checkpointId).isEqualTo("r-1") + } + } + + fun testLoadLatestResumeCheckpointPrefersResumableNodeWhenLatestIsFinish() { + runBlocking { + val storage = storage() + val agentId = "agent-latest-resume-${UUID.randomUUID()}" + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "latest-resumable", + createdAt = Instant.parse("2026-02-04T01:00:00Z"), + nodePath = "$agentId/single_run/nodeExecuteTool", + history = listOf(Message.User("First", RequestMetaInfo.Empty)) + ) + ) + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "latest-finish", + createdAt = Instant.parse("2026-02-04T01:01:00Z"), + nodePath = "$agentId/single_run/__finish__", + history = listOf(Message.Assistant("Done", ResponseMetaInfo.Empty)) + ) + ) + + val latestResume = historyService().loadLatestResumeCheckpoint(agentId) + + assertThat(latestResume).isNotNull + assertThat(latestResume!!.checkpointId).isEqualTo("latest-resumable") + } + } + + fun testPrefersLatestTodoWriteTitleForThreadTitle() { + runBlocking { + val storage = storage() + val agentId = "agent-todo-${UUID.randomUUID()}" + + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "todo-1", + createdAt = Instant.parse("2026-02-04T01:02:00Z"), + nodePath = "$agentId/single_run/nodeExecuteTool", + history = listOf( + Message.User("Please debug the crash", RequestMetaInfo.Empty), + Message.Tool.Call( + id = "tool-1", + tool = "TodoWrite", + content = """{"title":"Fix crash flow","todos":[]}""", + metaInfo = ResponseMetaInfo.Empty + ), + Message.User("Also handle retries", RequestMetaInfo.Empty) + ) + ) + ) + + val thread = historyService().listThreads().first { it.agentId == agentId } + + assertThat(thread.title).isEqualTo("Fix crash flow") + } + } + + fun testFallsBackToLatestUserMessageWhenNoTodoWriteTitleExists() { + runBlocking { + val storage = storage() + val agentId = "agent-user-${UUID.randomUUID()}" + + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "user-1", + createdAt = Instant.parse("2026-02-04T01:03:00Z"), + nodePath = "$agentId/single_run/nodeExecuteTool", + history = listOf( + Message.User("First request", RequestMetaInfo.Empty), + Message.Tool.Call( + id = "tool-2", + tool = "TodoWrite", + content = """{"title":" ","todos":[]}""", + metaInfo = ResponseMetaInfo.Empty + ), + Message.User("Latest request", RequestMetaInfo.Empty) + ) + ) + ) + + val thread = historyService().listThreads().first { it.agentId == agentId } + + assertThat(thread.title).isEqualTo("Latest request") + } + } + + fun testSkipsCacheableUserMessagesWhenResolvingThreadTitle() { + runBlocking { + val storage = storage() + val agentId = "agent-cacheable-${UUID.randomUUID()}" + + storage.saveCheckpoint( + agentId, + checkpoint( + checkpointId = "cacheable-1", + createdAt = Instant.parse("2026-02-04T01:04:00Z"), + nodePath = "$agentId/single_run/nodeExecuteTool", + history = listOf( + Message.User( + "Project-level instruction payload", + RequestMetaInfo( + timestamp = Clock.System.now(), + metadata = JsonObject(mapOf("cacheable" to JsonPrimitive(true))) + ) + ), + Message.User("Real user request", RequestMetaInfo.Empty) + ) + ) + ) + + val thread = historyService().listThreads().first { it.agentId == agentId } + + assertThat(thread.title).isEqualTo("Real user request") + } + } + + private fun storage(): JVMFilePersistenceStorageProvider { + return JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai")) + } + + private fun historyService(): AgentCheckpointHistoryService { + return project.service() + } + + private fun checkpoint( + checkpointId: String, + createdAt: Instant, + nodePath: String, + history: List, + tombstone: Boolean = false + ): AgentCheckpointData { + val properties = if (tombstone) { + mapOf("tombstone" to JsonPrimitive(true)) + } else { + null + } + + return AgentCheckpointData( + checkpointId = checkpointId, + createdAt = createdAt, + nodePath = nodePath, + lastInput = JsonNull, + messageHistory = history, + version = 0, + properties = properties + ) + } +}