Skip to content

Commit 3f32f19

Browse files
carlrobertohcarlrobertoh
committed
feat: agent timeline
1 parent 24e2a5f commit 3f32f19

31 files changed

+3768
-439
lines changed

‎src/main/kotlin/ee/carlrobert/codegpt/agent/AgentEvents.kt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
package ee.carlrobert.codegpt.agent
22

33
import ai.koog.prompt.executor.clients.LLMClientException
4+
import ee.carlrobert.codegpt.agent.history.CheckpointRef
45
import ee.carlrobert.codegpt.agent.tools.AskUserQuestionTool
56
import ee.carlrobert.codegpt.conversations.message.TokenUsage
67
import ee.carlrobert.codegpt.settings.service.ServiceType
78
import ee.carlrobert.codegpt.toolwindow.agent.AgentCreditsEvent
89
import ee.carlrobert.codegpt.toolwindow.agent.ui.approval.ToolApprovalRequest
10+
import java.util.UUID
911

1012
interface AgentEvents {
1113
fun onTextReceived(text: String) {}
@@ -22,6 +24,7 @@ interface AgentEvents {
2224
}
2325

2426
fun onRetry(attempt: Int, maxAttempts: Int, reason: String? = null) {}
27+
fun onRunCheckpointUpdated(runMessageId: UUID, ref: CheckpointRef?) {}
2528
fun onQueuedMessagesResolved()
2629
fun onTokenUsageAvailable(tokenUsage: Long) {}
2730
fun onCreditsAvailable(event: AgentCreditsEvent) {}

‎src/main/kotlin/ee/carlrobert/codegpt/agent/AgentService.kt‎

Lines changed: 50 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,12 @@ import ee.carlrobert.codegpt.ui.textarea.header.tag.McpTagDetails
1919
import kotlinx.coroutines.*
2020
import kotlinx.coroutines.flow.MutableSharedFlow
2121
import kotlinx.coroutines.flow.asSharedFlow
22+
import kotlinx.datetime.Clock
23+
import kotlinx.serialization.json.JsonNull
2224
import java.util.*
2325
import java.util.concurrent.ConcurrentHashMap
2426
import kotlin.io.path.Path
27+
import ai.koog.prompt.message.Message as PromptMessage
2528

2629
internal fun interface AgentRuntimeFactory {
2730
fun create(
@@ -53,16 +56,17 @@ class AgentService(private val project: Project) {
5356
private val pendingMessages = ConcurrentHashMap<String, ArrayDeque<MessageWithContext>>()
5457
private val sessionAgents = ConcurrentHashMap<String, AIAgent<MessageWithContext, String>>()
5558
private val sessionRuntimes = ConcurrentHashMap<String, SessionRuntime>()
56-
internal var runtimeFactory: AgentRuntimeFactory = AgentRuntimeFactory { p, storage, provider, events, sid, pending ->
57-
ProxyAIAgent.createService(
58-
project = p,
59-
checkpointStorage = storage,
60-
provider = provider,
61-
events = events,
62-
sessionId = sid,
63-
pendingMessages = pending
64-
)
65-
}
59+
internal var runtimeFactory: AgentRuntimeFactory =
60+
AgentRuntimeFactory { p, storage, provider, events, sid, pending ->
61+
ProxyAIAgent.createService(
62+
project = p,
63+
checkpointStorage = storage,
64+
provider = provider,
65+
events = events,
66+
sessionId = sid,
67+
pendingMessages = pending
68+
)
69+
}
6670
private val checkpointStorage =
6771
JVMFilePersistenceStorageProvider(Path(project.basePath ?: "", ".proxyai"))
6872
private val historyService = project.service<AgentCheckpointHistoryService>()
@@ -132,7 +136,8 @@ class AgentService(private val project: Project) {
132136
} catch (ex: Exception) {
133137
logger.error(ex)
134138
} finally {
135-
refreshSessionResumeCheckpoint(sessionId, agent.id)
139+
val ref = refreshSessionResumeCheckpoint(sessionId, agent.id)
140+
events.onRunCheckpointUpdated(message.id, ref)
136141
sessionAgents.remove(sessionId, agent)
137142
runCatching { runtime.service.removeAgentWithId(agent.id) }
138143
.onFailure { ex ->
@@ -201,7 +206,39 @@ class AgentService(private val project: Project) {
201206
.update(sessionId, conversationId, selectedServerIds)
202207
}
203208

204-
private suspend fun refreshSessionResumeCheckpoint(sessionId: String, agentId: String) {
209+
suspend fun createSeedCheckpointFromHistory(history: List<PromptMessage>): CheckpointRef? =
210+
withContext(Dispatchers.IO) {
211+
if (history.isEmpty()) {
212+
return@withContext null
213+
}
214+
215+
val agentId = UUID.randomUUID().toString()
216+
val checkpointId = UUID.randomUUID().toString()
217+
val checkpoint = AgentCheckpointData(
218+
checkpointId = checkpointId,
219+
createdAt = Clock.System.now(),
220+
nodePath = "$agentId/single_run/nodeExecuteTool",
221+
lastInput = JsonNull,
222+
messageHistory = history,
223+
version = 0
224+
)
225+
226+
runCatching {
227+
checkpointStorage.saveCheckpoint(agentId, checkpoint)
228+
CheckpointRef(agentId, checkpointId)
229+
}.onFailure { ex ->
230+
logger.warn(
231+
"Agent checkpoints: failed to create seed checkpoint from history " +
232+
"agentId=$agentId error=${ex.message}",
233+
ex
234+
)
235+
}.getOrNull()
236+
}
237+
238+
private suspend fun refreshSessionResumeCheckpoint(
239+
sessionId: String,
240+
agentId: String
241+
): CheckpointRef? {
205242
val ref = runCatching {
206243
historyService.loadLatestResumeCheckpoint(agentId)
207244
?.let { CheckpointRef(agentId, it.checkpointId) }
@@ -216,6 +253,7 @@ class AgentService(private val project: Project) {
216253
if (ref != null) {
217254
project.service<AgentToolWindowContentManager>().setResumeCheckpointRef(sessionId, ref)
218255
}
256+
return ref
219257
}
220258

221259
private fun ensureSessionRuntime(

‎src/main/kotlin/ee/carlrobert/codegpt/agent/ProxyAIAgent.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,8 @@ object ProxyAIAgent {
171171
}
172172

173173
onNodeExecutionCompleted { ctx ->
174+
if (stream) return@onNodeExecutionCompleted
175+
174176
(ctx.output as? List<*>)?.forEach { msg ->
175177
(msg as? Message.Assistant)?.let {
176178
events.onTextReceived(it.content)

‎src/main/kotlin/ee/carlrobert/codegpt/agent/ToolSpecs.kt‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -169,25 +169,20 @@ object ToolSpecs {
169169

170170
fun find(toolName: String): ToolSpec<*, *>? = specsByName[toolName.lowercase()]
171171

172-
fun approvalTypeFor(toolName: String): ToolApprovalType {
173-
return find(toolName)?.approvalType ?: ToolApprovalType.GENERIC
174-
}
172+
fun approvalTypeFor(toolName: String): ToolApprovalType =
173+
find(toolName)?.approvalType ?: ToolApprovalType.GENERIC
175174

176175
fun decodeArgsOrNull(
177176
json: Json,
178177
toolName: String,
179178
payload: String
180-
): Any? {
181-
return decodeOrNull(json, find(toolName)?.argsSerializer, payload)
182-
}
179+
) = decodeOrNull(json, find(toolName)?.argsSerializer, payload)
183180

184181
fun decodeResultOrNull(
185182
json: Json,
186183
toolName: String,
187184
payload: String
188-
): Any? {
189-
return decodeOrNull(json, find(toolName)?.resultSerializer, payload)
190-
}
185+
) = decodeOrNull(json, find(toolName)?.resultSerializer, payload)
191186

192187
@Suppress("UNCHECKED_CAST")
193188
private fun decodeOrNull(

‎src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointConversationMapper.kt‎

Lines changed: 35 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@ package ee.carlrobert.codegpt.agent.history
33
import ai.koog.agents.snapshot.feature.AgentCheckpointData
44
import ee.carlrobert.codegpt.conversations.Conversation
55
import ee.carlrobert.codegpt.conversations.message.Message
6+
import ee.carlrobert.codegpt.util.StringUtil.stripThinkingBlocks
67
import ee.carlrobert.llm.client.openai.completion.response.ToolCall
78
import ee.carlrobert.llm.client.openai.completion.response.ToolFunctionResponse
8-
import ai.koog.prompt.message.Message as PromptMessage
99

1010
object AgentCheckpointConversationMapper {
1111

@@ -14,87 +14,69 @@ object AgentCheckpointConversationMapper {
1414
projectInstructions: String?
1515
): Conversation {
1616
val conversation = Conversation()
17-
val history = checkpoint.messageHistory.filterNot { it is PromptMessage.System }
17+
val turns = AgentCheckpointTurnSequencer.toVisibleTurns(
18+
history = checkpoint.messageHistory,
19+
projectInstructions = projectInstructions
20+
)
1821

19-
var currentPrompt: String? = null
20-
val response = StringBuilder()
21-
val toolCalls = mutableListOf<ToolCall>()
22-
val toolResults = LinkedHashMap<String, String>()
23-
var syntheticToolIdIndex = 0
22+
turns.forEach { turn ->
23+
val response = StringBuilder()
24+
val toolCalls = mutableListOf<ToolCall>()
25+
val toolResults = LinkedHashMap<String, String>()
26+
var syntheticToolIdIndex = 0
2427

25-
fun flushTurn() {
26-
val prompt = currentPrompt?.trim().orEmpty()
27-
if (prompt.isBlank()) {
28-
return
29-
}
30-
31-
val uiMessage = Message(prompt)
32-
uiMessage.response = response.toString().trim()
33-
if (toolCalls.isNotEmpty()) {
34-
uiMessage.toolCalls = ArrayList(toolCalls)
35-
}
36-
if (toolResults.isNotEmpty()) {
37-
uiMessage.toolCallResults = LinkedHashMap(toolResults)
38-
}
39-
conversation.addMessage(uiMessage)
40-
currentPrompt = null
41-
response.setLength(0)
42-
toolCalls.clear()
43-
toolResults.clear()
44-
}
28+
turn.events.forEach { event ->
29+
when (event) {
30+
is AgentCheckpointTurnSequencer.TurnEvent.Assistant -> {
31+
appendAssistant(response, event.content)
32+
}
4533

46-
history.forEach { msg ->
47-
when (msg) {
48-
is PromptMessage.User -> {
49-
val text = msg.content.trim()
50-
if (shouldHideInAgentToolWindow(msg, projectInstructions)) {
51-
return@forEach
34+
is AgentCheckpointTurnSequencer.TurnEvent.Reasoning -> {
35+
appendAssistant(response, event.content)
5236
}
53-
flushTurn()
54-
currentPrompt = text
55-
}
5637

57-
is PromptMessage.Assistant -> appendAssistant(response, msg.content)
58-
is PromptMessage.Reasoning -> appendAssistant(response, msg.content)
59-
is PromptMessage.Tool.Call -> {
60-
if (currentPrompt != null) {
61-
val callId =
62-
msg.id?.takeIf { it.isNotBlank() }
63-
?: "tool-call-${++syntheticToolIdIndex}"
38+
is AgentCheckpointTurnSequencer.TurnEvent.ToolCall -> {
39+
val callId = event.id?.takeIf { it.isNotBlank() }
40+
?: "tool-call-${++syntheticToolIdIndex}"
6441
toolCalls.add(
6542
ToolCall(
6643
null,
6744
callId,
6845
"function",
69-
ToolFunctionResponse(msg.tool, msg.content.trim())
46+
ToolFunctionResponse(event.tool, event.content.trim())
7047
)
7148
)
7249
}
73-
}
7450

75-
is PromptMessage.Tool.Result -> {
76-
if (currentPrompt != null) {
77-
val callId = msg.id?.takeIf { it.isNotBlank() }
51+
is AgentCheckpointTurnSequencer.TurnEvent.ToolResult -> {
52+
val callId = event.id?.takeIf { it.isNotBlank() }
7853
?: toolCalls.lastOrNull()?.id
7954
?: "tool-call-${++syntheticToolIdIndex}"
8055
val prior = toolResults[callId]
81-
val merged = if (prior.isNullOrBlank()) msg.content.trim() else {
82-
"$prior\n\n${msg.content.trim()}"
56+
val merged = if (prior.isNullOrBlank()) event.content.trim() else {
57+
"$prior\n\n${event.content.trim()}"
8358
}
8459
toolResults[callId] = merged
8560
}
8661
}
62+
}
8763

88-
else -> Unit
64+
val uiMessage = Message(turn.prompt)
65+
uiMessage.response = response.toString().trim()
66+
if (toolCalls.isNotEmpty()) {
67+
uiMessage.toolCalls = ArrayList(toolCalls)
8968
}
69+
if (toolResults.isNotEmpty()) {
70+
uiMessage.toolCallResults = LinkedHashMap(toolResults)
71+
}
72+
conversation.addMessage(uiMessage)
9073
}
9174

92-
flushTurn()
9375
return conversation
9476
}
9577

9678
private fun appendAssistant(sb: StringBuilder, content: String) {
97-
val text = content.trim()
79+
val text = content.stripThinkingBlocks()
9880
if (text.isBlank()) {
9981
return
10082
}

‎src/main/kotlin/ee/carlrobert/codegpt/agent/history/AgentCheckpointHistoryService.kt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,13 @@ class AgentCheckpointHistoryService(project: Project) {
9696
?: checkpoints.maxByOrNull { it.createdAt }
9797
}
9898

99+
suspend fun listCheckpoints(agentId: String): List<AgentCheckpointData> =
100+
withContext(Dispatchers.IO) {
101+
storage.getCheckpoints(agentId)
102+
.filterNot { it.isTombstone() }
103+
.sortedByDescending { it.createdAt }
104+
}
105+
99106
private suspend fun buildSummary(agentId: String): AgentHistoryThreadSummary? {
100107
val checkpoints = storage.getCheckpoints(agentId)
101108
.filterNot { it.isTombstone() }

0 commit comments

Comments
 (0)