mirror of
https://github.com/carlrobertoh/ProxyAI.git
synced 2026-10-20 11:34:25 +00:00
feat: agent chat history
This commit is contained in:
Carl-Robert Linnupuu
parent
01757c4af2
commit
b456b53752
30 changed files with 2149 additions and 171 deletions
|
|
@ -28,6 +28,7 @@ import javax.swing.JPanel;
|
|||
public class ChatToolWindowScrollablePanel extends ScrollablePanel {
|
||||
|
||||
private final Map<UUID, JPanel> 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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String, ArrayDeque<MessageWithContext>>,
|
||||
): AIAgentService<MessageWithContext, String, out AIAgent<MessageWithContext, String>>
|
||||
}
|
||||
|
||||
@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<MessageWithContext, String, out AIAgent<MessageWithContext, String>>,
|
||||
val provider: ServiceType,
|
||||
val modelId: String,
|
||||
val selectedServerIds: Set<String>,
|
||||
val events: AgentEvents
|
||||
)
|
||||
|
||||
private val sessionJobs = ConcurrentHashMap<String, Job>()
|
||||
private val pendingMessages = ConcurrentHashMap<String, ArrayDeque<MessageWithContext>>()
|
||||
private val sessionAgents = ConcurrentHashMap<String, AIAgent<MessageWithContext, String>>()
|
||||
private val sessionRuntimes = ConcurrentHashMap<String, SessionRuntime>()
|
||||
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<AgentCheckpointHistoryService>()
|
||||
|
||||
private val _queuedMessageProcessed = MutableSharedFlow<String>()
|
||||
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<AgentToolWindowContentManager>().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<AgentToolWindowContentManager>()
|
||||
.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<ModelSelectionService>().getServiceForFeature(FeatureType.AGENT)
|
||||
val previousCheckpoint = runBlocking { getCheckpoint(sessionId) }
|
||||
val contentManager = project.service<AgentToolWindowContentManager>()
|
||||
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<AgentMcpContextService>().clear(sessionId)
|
||||
}
|
||||
|
||||
|
|
@ -138,4 +200,74 @@ class AgentService(private val project: Project) {
|
|||
project.service<AgentMcpContextService>()
|
||||
.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<AgentToolWindowContentManager>().setResumeCheckpointRef(sessionId, ref)
|
||||
}
|
||||
}
|
||||
|
||||
private fun ensureSessionRuntime(
|
||||
sessionId: String,
|
||||
provider: ServiceType,
|
||||
events: AgentEvents
|
||||
): SessionRuntime {
|
||||
val modelId = service<ModelSelectionService>().getAgentModel().id
|
||||
val selectedServerIds = project.service<AgentMcpContextService>()
|
||||
.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
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<TagDetails>,
|
||||
@Transient
|
||||
val tags: List<TagDetails> = emptyList(),
|
||||
val uiVisible: Boolean = true,
|
||||
val uiText: String = text
|
||||
) {
|
||||
|
||||
@Transient
|
||||
val id: UUID = UUID.randomUUID()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,13 +0,0 @@
|
|||
package ee.carlrobert.codegpt.agent
|
||||
|
||||
import kotlinx.serialization.Serializable
|
||||
|
||||
@Serializable
|
||||
data class MessageWithContextSnapshot(
|
||||
val text: String,
|
||||
val tagNames: List<String>
|
||||
)
|
||||
|
||||
fun MessageWithContext.toSnapshot(): MessageWithContextSnapshot {
|
||||
return MessageWithContextSnapshot(text, tags.map { it.name })
|
||||
}
|
||||
|
|
@ -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<MessageWithContextSnapshot>()
|
||||
|
||||
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<String, ArrayDeque<MessageWithContext>>,
|
||||
): AIAgent<MessageWithContext, String> {
|
||||
): GraphAIAgentService<MessageWithContext, String> {
|
||||
val modelSelection =
|
||||
service<ModelSelectionService>().getModelSelectionForFeature(FeatureType.AGENT)
|
||||
val skills = project.service<SkillDiscoveryService>().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<ModelSelectionService>().getAgentModel()
|
||||
val agent = AIAgent(
|
||||
return AIAgentService<MessageWithContext, String>(
|
||||
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<String, JsonPrimitive> {
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -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<String> = emptySet()) {
|
||||
READ("Read"),
|
||||
|
|
@ -22,34 +23,6 @@ enum class ToolName(val id: String, val aliases: Set<String> = 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<TArgs, TResult>(
|
||||
|
|
@ -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<Any>, payload)
|
||||
}.getOrNull()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<ToolDescriptor>
|
||||
): Flow<StreamFrame> {
|
||||
var attempt = 1
|
||||
var refinedPrompt = prompt
|
||||
var hasReceivedData: Boolean
|
||||
|
||||
fun createStream(attemptNum: Int): Flow<StreamFrame> {
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<ToolCall>()
|
||||
val toolResults = LinkedHashMap<String, String>()
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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<AgentHistoryThreadSummary>? = null
|
||||
|
||||
suspend fun listThreads(limit: Int? = null): List<AgentHistoryThreadSummary> =
|
||||
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<PromptMessage>): 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<AgentHistoryThreadSummary> {
|
||||
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<AgentHistoryThreadSummary> {
|
||||
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<AgentHistoryThreadSummary>,
|
||||
query: String
|
||||
): List<AgentHistoryThreadSummary> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
package ee.carlrobert.codegpt.agent.history
|
||||
|
||||
data class AgentHistoryPage(
|
||||
val items: List<AgentHistoryThreadSummary>,
|
||||
val hasMore: Boolean,
|
||||
val total: Int
|
||||
)
|
||||
|
|
@ -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 = ""
|
||||
)
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
package ee.carlrobert.codegpt.agent.history
|
||||
|
||||
data class CheckpointRef(
|
||||
val agentId: String,
|
||||
val checkpointId: String
|
||||
)
|
||||
|
|
@ -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<MessageWithContext>,
|
||||
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<MessageWithContext>,
|
||||
historyCompressionConfig: HistoryCompressionConfig,
|
||||
events: AgentEvents,
|
||||
|
|
@ -67,29 +65,6 @@ internal class SingleRunStrategyProvider : AgentRunStrategyProvider {
|
|||
): AIAgentGraphStrategy<MessageWithContext, String> = strategy("single_run") {
|
||||
val nodeCallLLM by node<MessageWithContext, List<Message.Response>> { 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<StreamFrame>()
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -71,7 +71,6 @@ abstract class BaseTool<Args : Any, Result : Any>(
|
|||
BEFORE_TOOL_USE,
|
||||
beforePayload,
|
||||
name,
|
||||
toolId,
|
||||
sessionId
|
||||
)
|
||||
|
||||
|
|
@ -111,7 +110,6 @@ abstract class BaseTool<Args : Any, Result : Any>(
|
|||
HookEventType.AFTER_TOOL_USE,
|
||||
afterPayload,
|
||||
name,
|
||||
toolId,
|
||||
sessionId
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,7 +18,6 @@ class HookManager(
|
|||
event: HookEventType,
|
||||
payload: Map<String, Any?>,
|
||||
toolName: String? = null,
|
||||
toolId: String? = null,
|
||||
sessionId: String? = null
|
||||
): List<HookExecutionResult> {
|
||||
val settings = project.service<ProxyAISettingsService>().getSettings()
|
||||
|
|
@ -81,10 +80,9 @@ class HookManager(
|
|||
event: HookEventType,
|
||||
payload: Map<String, Any?>,
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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<VirtualFile> = emptyList(),
|
||||
val createdAt: Long = System.currentTimeMillis(),
|
||||
var lastActiveAt: Long = System.currentTimeMillis()
|
||||
|
|
|
|||
|
|
@ -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<AgentService>().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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<AgentHistoryThreadSummary>, 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<AgentHistoryThreadSummary>()
|
||||
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<AgentHistoryThreadSummary> {
|
||||
|
||||
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<out AgentHistoryThreadSummary>,
|
||||
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) + "…"
|
||||
}
|
||||
|
|
@ -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<AgentCheckpointHistoryService>()
|
||||
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<AgentHistoryThreadSummary>, 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<AgentToolWindowContentManager>()
|
||||
.openCheckpointConversation(thread, conversation)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private fun welcomeMessage(): String {
|
||||
val name = GeneralSettings.getCurrentState().displayName
|
||||
return """
|
||||
|
|
|
|||
|
|
@ -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<EditorColorsManager>().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) {
|
||||
|
|
|
|||
|
|
@ -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 :
|
||||
|
|
|
|||
|
|
@ -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<String>()
|
||||
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<ModelSelectionService>().getAgentModel()
|
||||
}
|
||||
|
||||
private data class SessionSnapshot(
|
||||
val runtimeAgentId: String?,
|
||||
val resumeCheckpointRef: CheckpointRef?
|
||||
)
|
||||
}
|
||||
|
||||
private class RecordingRuntimeFactory(
|
||||
private val model: LLModel
|
||||
) : AgentRuntimeFactory {
|
||||
|
||||
val createdServices = mutableListOf<FakeManagedAgentService>()
|
||||
|
||||
override fun create(
|
||||
project: Project,
|
||||
checkpointStorage: JVMFilePersistenceStorageProvider,
|
||||
provider: ServiceType,
|
||||
events: AgentEvents,
|
||||
sessionId: String,
|
||||
pendingMessages: ConcurrentHashMap<String, ArrayDeque<MessageWithContext>>
|
||||
): AIAgentService<MessageWithContext, String, out AIAgent<MessageWithContext, String>> {
|
||||
val service = FakeManagedAgentService(checkpointStorage, model)
|
||||
createdServices.add(service)
|
||||
return service
|
||||
}
|
||||
}
|
||||
|
||||
private class FakeManagedAgentService(
|
||||
private val checkpointStorage: JVMFilePersistenceStorageProvider,
|
||||
model: LLModel
|
||||
) : AIAgentService<MessageWithContext, String, FakeManagedAgent> {
|
||||
|
||||
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<String, FakeManagedAgent>()
|
||||
val createdAgentIds = mutableListOf<String>()
|
||||
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<FakeManagedAgent> = agents.values.toList()
|
||||
|
||||
override suspend fun listActiveAgents(): List<FakeManagedAgent> = emptyList()
|
||||
|
||||
override suspend fun listInactiveAgents(): List<FakeManagedAgent> = agents.values.toList()
|
||||
|
||||
override suspend fun listFinishedAgents(): List<FakeManagedAgent> = 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<MessageWithContext, String> {
|
||||
|
||||
private var state: AIAgent.Companion.State<String> = AIAgent.Companion.State.NotStarted()
|
||||
|
||||
override suspend fun getState(): AIAgent.Companion.State<String> = 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<ai.koog.agents.core.tools.ToolDescriptor>
|
||||
): List<Message.Response> {
|
||||
throw UnsupportedOperationException("NoopPromptExecutor should not be used in tests")
|
||||
}
|
||||
|
||||
override fun executeStreaming(
|
||||
prompt: Prompt,
|
||||
model: LLModel,
|
||||
tools: List<ai.koog.agents.core.tools.ToolDescriptor>
|
||||
): Flow<StreamFrame> {
|
||||
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
|
||||
}
|
||||
|
|
@ -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<String, String>
|
||||
) {
|
||||
assertThat(conversation.messages)
|
||||
.extracting("prompt", "response")
|
||||
.containsExactly(tuple(expectedTurn.first, expectedTurn.second))
|
||||
}
|
||||
|
||||
private fun checkpoint(history: List<Message>): 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
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -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<Message>,
|
||||
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
|
||||
)
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue