From 02cb913eccab72e503c469b78e8190aa7ec9cc50 Mon Sep 17 00:00:00 2001 From: Derek <294233080+derek201759@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:13:48 -0700 Subject: [PATCH] feat(plan-execution): add token consumption tracking for plan steps - add total_input_tokens, total_output_tokens, total_tokens, total_duration_ms columns to plan_executions table - extend PlanStepOutput with token fields and serialization - update PlanExecutionListener and PlanExecutor to populate token data - expose plan-level token totals in PlansViewModel - update UI to display token summary and duration - maintain backward compatibility by treating missing values as null BREAKING CHANGE: PlanStepOutput now includes token fields; legacy rows will have null values. --- .../io/askimo/ui/plan/PlanDetailView.kt | 110 +++++++++++++++--- .../io/askimo/ui/plan/PlanHistorySidePanel.kt | 18 ++- .../io/askimo/ui/plan/PlansViewModel.kt | 42 ++++++- .../io/askimo/core/db/DatabaseManager.kt | 22 ++++ .../askimo/core/plan/PlanExecutionListener.kt | 38 +++++- .../io/askimo/core/plan/PlanExecutor.kt | 17 ++- .../kotlin/io/askimo/core/plan/PlanService.kt | 29 ++++- .../io/askimo/core/plan/PlanStepEvent.kt | 6 + .../askimo/core/plan/domain/PlanExecution.kt | 42 ++++++- .../repository/PlanExecutionRepository.kt | 73 +++++++++--- 10 files changed, 343 insertions(+), 54 deletions(-) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanDetailView.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanDetailView.kt index 853bd4d1..08ca6aa9 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanDetailView.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanDetailView.kt @@ -59,9 +59,12 @@ import androidx.compose.runtime.remember import androidx.compose.runtime.rememberCoroutineScope import androidx.compose.runtime.setValue import androidx.compose.ui.Alignment +import androidx.compose.ui.ExperimentalComposeUiApi import androidx.compose.ui.Modifier import androidx.compose.ui.graphics.Color +import androidx.compose.ui.input.pointer.PointerEventType import androidx.compose.ui.input.pointer.PointerIcon +import androidx.compose.ui.input.pointer.onPointerEvent import androidx.compose.ui.input.pointer.pointerHoverIcon import androidx.compose.ui.platform.LocalClipboardManager import androidx.compose.ui.semantics.Role @@ -71,6 +74,7 @@ import androidx.compose.ui.text.font.FontWeight import androidx.compose.ui.text.input.KeyboardType import androidx.compose.ui.unit.dp import androidx.compose.ui.unit.sp +import io.askimo.core.i18n.LocalizationManager import io.askimo.core.plan.PlanStepEvent import io.askimo.core.plan.domain.PlanInput import io.askimo.core.util.TimeUtil @@ -639,6 +643,7 @@ private fun agenticStepProgressPanel( } } +@OptIn(ExperimentalComposeUiApi::class) @Composable private fun agenticStepRow( event: PlanStepEvent, @@ -647,6 +652,7 @@ private fun agenticStepRow( val clipboardManager = LocalClipboardManager.current val coroutineScope = rememberCoroutineScope() var showCopyFeedback by remember { mutableStateOf(false) } + var isHovered by remember { mutableStateOf(false) } var outputExpanded by remember { mutableStateOf(true) } @@ -669,7 +675,11 @@ private fun agenticStepRow( null } - Column(modifier = Modifier.fillMaxWidth().padding(vertical = 2.dp)) { + Column( + modifier = Modifier.fillMaxWidth().padding(vertical = 2.dp) + .onPointerEvent(PointerEventType.Enter) { isHovered = true } + .onPointerEvent(PointerEventType.Exit) { isHovered = false }, + ) { Row( modifier = Modifier.fillMaxWidth(), horizontalArrangement = Arrangement.spacedBy(Spacing.medium), @@ -725,22 +735,38 @@ private fun agenticStepRow( ) when (event) { is PlanStepEvent.Completed -> { - Surface( - shape = MaterialTheme.shapes.extraSmall, - color = MaterialTheme.colorScheme.secondaryContainer, + Row( + horizontalArrangement = Arrangement.spacedBy(Spacing.extraSmall), + verticalAlignment = Alignment.CenterVertically, ) { - Text( - text = TimeUtil.formatDurationMs( - durationMs = event.durationMs, - hourLabel = stringResource("plans.steps.duration.hour"), - minuteLabel = stringResource("plans.steps.duration.minute"), - secondLabel = stringResource("plans.steps.duration.second"), - lessThanOne = stringResource("plans.steps.duration.less.than.one.second"), - ), - style = MaterialTheme.typography.labelSmall, - color = MaterialTheme.colorScheme.onSecondaryContainer, - modifier = Modifier.padding(horizontal = 4.dp, vertical = 1.dp), - ) + Surface( + shape = MaterialTheme.shapes.extraSmall, + color = MaterialTheme.colorScheme.secondaryContainer, + ) { + Text( + text = TimeUtil.formatDurationMs( + durationMs = event.durationMs, + hourLabel = stringResource("plans.steps.duration.hour"), + minuteLabel = stringResource("plans.steps.duration.minute"), + secondLabel = stringResource("plans.steps.duration.second"), + lessThanOne = stringResource("plans.steps.duration.less.than.one.second"), + ), + style = MaterialTheme.typography.labelSmall, + color = MaterialTheme.colorScheme.onSecondaryContainer, + modifier = Modifier.padding(horizontal = 4.dp, vertical = 1.dp), + ) + } + stepTokenUsageLabel(event)?.let { usage -> + if (isHovered) { + themedTooltip(text = stringResource("message.token.usage.tooltip")) { + Text( + text = usage, + style = MaterialTheme.typography.labelSmall, + color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.5f), + ) + } + } + } } } @@ -877,6 +903,43 @@ private fun agenticStepRow( } } +private fun stepTokenUsageLabel(event: PlanStepEvent.Completed): String? { + val total = event.totalTokens + val input = event.inputTokens + val output = event.outputTokens + if (total == null || total <= 0) return null + + return if (input != null && output != null) { + "${LocalizationManager.formatNumber(total)} tokens (${LocalizationManager.formatNumber(input)} in / ${LocalizationManager.formatNumber(output)} out)" + } else { + "${LocalizationManager.formatNumber(total)} tokens" + } +} + +@Composable +private fun planTotalSummaryLabel(viewModel: PlansViewModel): String? { + val total = viewModel.planTotalTokens + val input = viewModel.planTotalInputTokens + val output = viewModel.planTotalOutputTokens + val durationMs = viewModel.planTotalDurationMs + if (total == null || total <= 0) return null + + return buildString { + if (input != null && output != null) { + append("${LocalizationManager.formatNumber(total)} tokens (${LocalizationManager.formatNumber(input)} in / ${LocalizationManager.formatNumber(output)} out)") + } else { + append("${LocalizationManager.formatNumber(total)} tokens") + } + if (durationMs != null && durationMs > 0) { + val durationLabel = when { + durationMs < 1_000 -> "${durationMs}ms" + else -> "${"%.1f".format(durationMs / 1000.0)}s" + } + append(" · $durationLabel") + } + } +} + /** * Inline panel rendered while the plan executor is paused waiting for the user to answer * an interactive [io.askimo.core.plan.domain.WorkflowNode.Ask] step. @@ -1033,9 +1096,11 @@ private fun resultPanel( horizontalArrangement = Arrangement.SpaceBetween, verticalAlignment = Alignment.CenterVertically, ) { + // Title + plan totals on hover Row( horizontalArrangement = Arrangement.spacedBy(Spacing.small), verticalAlignment = Alignment.CenterVertically, + modifier = Modifier.weight(1f), ) { Icon( if (isPinned) Icons.Default.PushPin else Icons.Default.CheckCircle, @@ -1049,7 +1114,20 @@ private fun resultPanel( fontWeight = FontWeight.SemiBold, color = MaterialTheme.colorScheme.onSurface, ) + // Plan-level totals — only for the live (non-pinned) result panel + if (!isPinned) { + planTotalSummaryLabel(viewModel)?.let { summary -> + themedTooltip(text = stringResource("message.token.usage.tooltip")) { + Text( + text = summary, + style = MaterialTheme.typography.labelSmall, + color = MaterialTheme.colorScheme.onSurfaceVariant.copy(alpha = 0.5f), + ) + } + } + } } + Row( horizontalArrangement = Arrangement.spacedBy(Spacing.extraSmall), verticalAlignment = Alignment.CenterVertically, diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanHistorySidePanel.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanHistorySidePanel.kt index 134ef79c..f0cfc5e0 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanHistorySidePanel.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlanHistorySidePanel.kt @@ -55,6 +55,7 @@ import androidx.compose.ui.input.pointer.pointerHoverIcon import androidx.compose.ui.input.pointer.pointerInput import androidx.compose.ui.text.font.FontWeight import androidx.compose.ui.unit.dp +import io.askimo.core.i18n.LocalizationManager import io.askimo.core.plan.domain.PlanExecution import io.askimo.core.plan.domain.PlanExecutionStatus import io.askimo.core.util.TimeUtil @@ -282,7 +283,22 @@ private fun planHistoryItem( val interactionSource = remember { MutableInteractionSource() } val isHovered by interactionSource.collectIsHoveredAsState() - themedTooltip(text = stringResource("plans.history.restore")) { + val tooltipText = buildString { + append(stringResource("plans.history.restore")) + val parts = mutableListOf() + execution.totalTokens?.let { parts += "Tokens: ${LocalizationManager.formatNumber(it)}" } + execution.totalInputTokens?.let { parts += "In: ${LocalizationManager.formatNumber(it)}" } + execution.totalOutputTokens?.let { parts += "Out: ${LocalizationManager.formatNumber(it)}" } + execution.totalDurationMs?.let { ms -> + parts += if (ms >= 1_000) "Time: ${"%.1f".format(ms / 1000.0)}s" else "Time: ${ms}ms" + } + if (parts.isNotEmpty()) { + append("\n") + append(parts.joinToString(" · ")) + } + } + + themedTooltip(text = tooltipText) { Surface( modifier = modifier .hoverable(interactionSource) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlansViewModel.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlansViewModel.kt index 6fbe1f16..4aef24ae 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlansViewModel.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/plan/PlansViewModel.kt @@ -91,6 +91,16 @@ class PlansViewModel( var runError by mutableStateOf(null) private set + /** Plan-level token and duration summary for the current (or restored) run. */ + var planTotalInputTokens by mutableStateOf(null) + private set + var planTotalOutputTokens by mutableStateOf(null) + private set + var planTotalTokens by mutableStateOf(null) + private set + var planTotalDurationMs by mutableStateOf(null) + private set + var exportError by mutableStateOf(null) private set @@ -359,7 +369,16 @@ class PlansViewModel( }.fold( onSuccess = { result -> result.fold( - onSuccess = { runResult = it }, + onSuccess = { planResult -> + runResult = planResult + val saved = withContext(Dispatchers.IO) { + planService.getExecutions(plan.id).firstOrNull { it.id == planResult.executionId } + } + planTotalInputTokens = saved?.totalInputTokens + planTotalOutputTokens = saved?.totalOutputTokens + planTotalTokens = saved?.totalTokens + planTotalDurationMs = saved?.totalDurationMs + }, onFailure = { runError = it.message ?: "Plan failed" }, ) }, @@ -392,6 +411,10 @@ class PlansViewModel( followUpError = null pendingQuestion = null pendingAnswerText = "" + planTotalInputTokens = null + planTotalOutputTokens = null + planTotalTokens = null + planTotalDurationMs = null selectedPlan?.id?.let { planStateCache.remove(it) } } @@ -537,13 +560,16 @@ class PlansViewModel( // Reconstruct step progress from persisted step outputs so the UI shows // the breakdown from the previous run when the user clicks a history entry. - stepProgress = execution.stepOutputs.map { (stepName, output) -> + stepProgress = execution.stepOutputs.map { step -> PlanStepEvent.Completed( planId = execution.planId, - stepName = stepName, + stepName = step.stepName, executionId = execution.id, - output = output, - durationMs = 0L, + output = step.output, + inputTokens = step.inputTokens, + outputTokens = step.outputTokens, + totalTokens = step.totalTokens, + durationMs = step.durationMs ?: 0L, ) } @@ -564,6 +590,12 @@ class PlansViewModel( runError = null } } + + // Restore plan-level totals from the persisted execution record + planTotalInputTokens = execution.totalInputTokens + planTotalOutputTokens = execution.totalOutputTokens + planTotalTokens = execution.totalTokens + planTotalDurationMs = execution.totalDurationMs } /** Opens the editor to create a brand-new plan with a starter template. */ diff --git a/shared/src/main/kotlin/io/askimo/core/db/DatabaseManager.kt b/shared/src/main/kotlin/io/askimo/core/db/DatabaseManager.kt index 411a1ed4..31b98841 100644 --- a/shared/src/main/kotlin/io/askimo/core/db/DatabaseManager.kt +++ b/shared/src/main/kotlin/io/askimo/core/db/DatabaseManager.kt @@ -655,6 +655,28 @@ class DatabaseManager private constructor( } catch (_: Exception) { // Column already exists — safe to ignore. } + + // Migration: add token usage columns for plan executions. + try { + stmt.executeUpdate("ALTER TABLE plan_executions ADD COLUMN total_input_tokens INTEGER") + } catch (_: Exception) { + // Column already exists — safe to ignore. + } + try { + stmt.executeUpdate("ALTER TABLE plan_executions ADD COLUMN total_output_tokens INTEGER") + } catch (_: Exception) { + // Column already exists — safe to ignore. + } + try { + stmt.executeUpdate("ALTER TABLE plan_executions ADD COLUMN total_tokens INTEGER") + } catch (_: Exception) { + // Column already exists — safe to ignore. + } + try { + stmt.executeUpdate("ALTER TABLE plan_executions ADD COLUMN total_duration_ms INTEGER") + } catch (_: Exception) { + // Column already exists — safe to ignore. + } } } diff --git a/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutionListener.kt b/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutionListener.kt index 8d40156e..8f984716 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutionListener.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutionListener.kt @@ -10,6 +10,7 @@ import dev.langchain4j.agentic.observability.AgentRequest import dev.langchain4j.agentic.observability.AgentResponse import io.askimo.core.event.EventBus import io.askimo.core.logging.logger +import io.askimo.core.plan.domain.PlanStepOutput import kotlinx.coroutines.runBlocking import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.CopyOnWriteArrayList @@ -38,10 +39,9 @@ internal class PlanExecutionListener( /** * Outputs of completed steps, in completion order. - * Key = stepName (agent name / outputKey), value = output string. * Used by [PlanExecutor] to derive the final plan output. */ - val stepOutputs = CopyOnWriteArrayList>() + val stepOutputs = CopyOnWriteArrayList() override fun inheritedBySubagents(): Boolean = true @@ -70,11 +70,38 @@ internal class PlanExecutionListener( val durationMs = elapsed(startTimes.remove(key(agentResponse.agentId(), stepName))) val output = agentResponse.output() val outputStr = output?.toString() ?: "" - log.debug("[{}] ✔ step '{}' done in {}ms | output: {}", planId, stepName, durationMs, truncate(outputStr)) + + val tokenUsage = agentResponse.chatResponse() + ?.metadata() + ?.tokenUsage() + + val inputTokens = tokenUsage?.inputTokenCount() + val outputTokens = tokenUsage?.outputTokenCount() + val totalTokens = tokenUsage?.totalTokenCount() + + log.debug( + "[{}] ✔ step '{}' done in {}ms | tokens(total/in/out)={}/{}/{} | output: {}", + planId, + stepName, + durationMs, + totalTokens ?: "-", + inputTokens ?: "-", + outputTokens ?: "-", + truncate(outputStr), + ) // Record non-blank outputs so the executor can find the final result. if (outputStr.isNotBlank()) { - stepOutputs.add(stepName to outputStr) + stepOutputs.add( + PlanStepOutput( + stepName = stepName, + output = outputStr, + inputTokens = inputTokens, + outputTokens = outputTokens, + totalTokens = totalTokens, + durationMs = durationMs, + ), + ) } post( @@ -83,6 +110,9 @@ internal class PlanExecutionListener( stepName = stepName, executionId = executionId, output = output, + inputTokens = inputTokens, + outputTokens = outputTokens, + totalTokens = totalTokens, durationMs = durationMs, ), ) diff --git a/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutor.kt b/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutor.kt index d74608eb..d88361c4 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutor.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/PlanExecutor.kt @@ -13,6 +13,7 @@ import io.askimo.core.event.EventBus import io.askimo.core.logging.logger import io.askimo.core.plan.domain.PlanDef import io.askimo.core.plan.domain.PlanStep +import io.askimo.core.plan.domain.PlanStepOutput import io.askimo.core.plan.domain.WorkflowNode import kotlinx.coroutines.runBlocking import java.util.concurrent.Executors @@ -55,15 +56,14 @@ class PlanExecutor(private val chatModel: ChatModel) { * @param executionId Optional [PlanExecution] record id — passed to [PlanExecutionListener] * so every [PlanStepEvent] is correlated back to the DB record. * @param onStepOutputs Optional callback invoked after execution with the ordered list of - * (stepName → output) pairs for all completed steps. Use this to - * persist intermediate step results alongside the final output. + * completed step outputs (including nullable usage metadata). * @return The final output string produced by the last step in the workflow. */ fun execute( plan: PlanDef, inputs: Map, executionId: String = "", - onStepOutputs: ((List>) -> Unit)? = null, + onStepOutputs: ((List) -> Unit)? = null, ): String { validateInputs(plan, inputs) @@ -124,7 +124,12 @@ class PlanExecutor(private val chatModel: ChatModel) { } } - private fun executeInternal(plan: PlanDef, inputs: Map, executionId: String, onStepOutputs: ((List>) -> Unit)? = null): String { + private fun executeInternal( + plan: PlanDef, + inputs: Map, + executionId: String, + onStepOutputs: ((List) -> Unit)? = null, + ): String { log.debug("Executing plan '{}' (execution={}) with inputs: {}", plan.id, executionId, inputs.keys) // Collect interactive answers first by traversing the workflow for Ask nodes. @@ -143,8 +148,8 @@ class PlanExecutor(private val chatModel: ChatModel) { rootAgent?.invoke(scopeInputs) val lastStepId = lastLeafStepId(plan.workflow) - val output = listener.stepOutputs.lastOrNull { (stepName, _) -> stepName == lastStepId }?.second - ?: listener.stepOutputs.lastOrNull()?.second + val output = listener.stepOutputs.lastOrNull { it.stepName == lastStepId }?.output + ?: listener.stepOutputs.lastOrNull()?.output ?: "" logExecutionSummary(plan, output) diff --git a/shared/src/main/kotlin/io/askimo/core/plan/PlanService.kt b/shared/src/main/kotlin/io/askimo/core/plan/PlanService.kt index 26c59fa2..53820c70 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/PlanService.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/PlanService.kt @@ -12,6 +12,7 @@ import io.askimo.core.logging.logger import io.askimo.core.plan.domain.PlanDef import io.askimo.core.plan.domain.PlanExecution import io.askimo.core.plan.domain.PlanExecutionStatus +import io.askimo.core.plan.domain.PlanStepOutput import io.askimo.core.plan.repository.PlanDefRepository import io.askimo.core.plan.repository.PlanExecutionRepository import java.time.Instant @@ -119,20 +120,38 @@ class PlanService( val executor = PlanExecutor(chatModel) - var stepOutputs: List> = emptyList() + var stepOutputs: List = emptyList() return runCatching { executor.execute(plan, inputs, execution.id) { stepOutputs = it } }.fold( onSuccess = { output -> + val totalInputTokens = stepOutputs.mapNotNull { it.inputTokens }.takeIf { it.isNotEmpty() }?.sum() + val totalOutputTokens = stepOutputs.mapNotNull { it.outputTokens }.takeIf { it.isNotEmpty() }?.sum() + val totalTokens = stepOutputs.mapNotNull { it.totalTokens }.takeIf { it.isNotEmpty() }?.sum() + val totalDurationMs = stepOutputs.mapNotNull { it.durationMs }.takeIf { it.isNotEmpty() }?.sum() + planExecutionRepository.update( execution.copy( status = PlanExecutionStatus.COMPLETED, output = output.takeIf { it.isNotBlank() }, stepOutputs = stepOutputs, + totalInputTokens = totalInputTokens, + totalOutputTokens = totalOutputTokens, + totalTokens = totalTokens, + totalDurationMs = totalDurationMs, ), ) - log.info("Plan '{}' (execution={}) completed successfully", plan.id, execution.id) + log.info( + "Plan '{}' (execution={}) completed — steps={}, tokens(total/in/out)={}/{}/{}, duration={}ms", + plan.id, + execution.id, + stepOutputs.size, + totalTokens ?: "-", + totalInputTokens ?: "-", + totalOutputTokens ?: "-", + totalDurationMs ?: "-", + ) Result.success(PlanRunResult(executionId = execution.id, output = output)) }, onFailure = { e -> @@ -194,9 +213,9 @@ class PlanService( // Multi-step plan — show each intermediate step's output for full context. // The last entry equals priorOutput so we show it separately as "Final result". append("This result was produced by a multi-step plan. Here are the intermediate step outputs:\n\n") - stepOutputs.dropLast(1).forEachIndexed { index, (stepName, output) -> - append("### Step ${index + 1}: $stepName\n\n") - append(output) + stepOutputs.dropLast(1).forEachIndexed { index, step -> + append("### Step ${index + 1}: ${step.stepName}\n\n") + append(step.output) append("\n\n") } append("### Final result\n\n") diff --git a/shared/src/main/kotlin/io/askimo/core/plan/PlanStepEvent.kt b/shared/src/main/kotlin/io/askimo/core/plan/PlanStepEvent.kt index 86cab75b..916a3d5e 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/PlanStepEvent.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/PlanStepEvent.kt @@ -45,6 +45,9 @@ sealed class PlanStepEvent : Event { * Fired after a step agent returns successfully. * * @param output The raw output object returned by the agent (usually a String). + * @param inputTokens Prompt/input token count for this step (nullable when unavailable). + * @param outputTokens Completion/output token count for this step (nullable when unavailable). + * @param totalTokens Total token count for this step (nullable when unavailable). * @param durationMs Wall-clock time for this step in milliseconds. */ data class Completed( @@ -52,6 +55,9 @@ sealed class PlanStepEvent : Event { override val stepName: String, override val executionId: String, val output: Any?, + val inputTokens: Int? = null, + val outputTokens: Int? = null, + val totalTokens: Int? = null, val durationMs: Long = 0L, override val timestamp: Instant = Instant.now(), ) : PlanStepEvent() { diff --git a/shared/src/main/kotlin/io/askimo/core/plan/domain/PlanExecution.kt b/shared/src/main/kotlin/io/askimo/core/plan/domain/PlanExecution.kt index 5ddb6b65..23c04e01 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/domain/PlanExecution.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/domain/PlanExecution.kt @@ -5,6 +5,7 @@ package io.askimo.core.plan.domain import io.askimo.core.db.sqliteInstant +import kotlinx.serialization.Serializable import org.jetbrains.exposed.v1.core.Table import java.time.Instant @@ -17,6 +18,17 @@ enum class PlanExecutionStatus { FAILED, } +/** Persisted output payload for one completed plan step. */ +@Serializable +data class PlanStepOutput( + val stepName: String, + val output: String, + val inputTokens: Int? = null, + val outputTokens: Int? = null, + val totalTokens: Int? = null, + val durationMs: Long? = null, +) + /** * Persisted record of one run of a [PlanDef]. * @@ -31,8 +43,12 @@ enum class PlanExecutionStatus { * @param runCount How many times this execution has been re-run (starts at 1). * @param sessionId Linked ChatSession id that holds the message history. * @param output The final AI-generated result text; populated on COMPLETED. - * @param stepOutputs Ordered list of (stepName → output) pairs for every completed step; - * populated on COMPLETED so users can inspect intermediate results. + * @param stepOutputs Ordered list of completed step outputs, including optional + * token usage and duration per step. + * @param totalInputTokens Sum of input tokens across all completed steps; null when unavailable. + * @param totalOutputTokens Sum of output tokens across all completed steps; null when unavailable. + * @param totalTokens Sum of total tokens across all completed steps; null when unavailable. + * @param totalDurationMs Sum of wall-clock duration across all steps in milliseconds; null when unavailable. * @param errorMessage Populated when [status] is [PlanExecutionStatus.FAILED]. * @param createdAt When the execution record was first created. * @param updatedAt When the execution record was last modified. @@ -46,7 +62,11 @@ data class PlanExecution( val runCount: Int = 1, val sessionId: String? = null, val output: String? = null, - val stepOutputs: List> = emptyList(), + val stepOutputs: List = emptyList(), + val totalInputTokens: Int? = null, + val totalOutputTokens: Int? = null, + val totalTokens: Int? = null, + val totalDurationMs: Long? = null, val errorMessage: String? = null, val createdAt: Instant = Instant.now(), val updatedAt: Instant = Instant.now(), @@ -73,10 +93,24 @@ object PlanExecutionsTable : Table("plan_executions") { val output = text("output").nullable() /** - * JSON-encoded list of step outputs: `[{"step":"stepName","output":"..."},...]`. + * JSON-encoded list of step outputs: + * `[{"stepName":"step-a","output":"...","inputTokens":12,"outputTokens":34,"totalTokens":46,"durationMs":1000}, ...]`. * Null for executions created before this column was added. */ val stepOutputs = text("step_outputs").nullable() + + /** Sum of input tokens across all steps; null for old rows or when no token data available. */ + val totalInputTokens = integer("total_input_tokens").nullable() + + /** Sum of output tokens across all steps; null for old rows or when no token data available. */ + val totalOutputTokens = integer("total_output_tokens").nullable() + + /** Sum of total tokens across all steps; null for old rows or when no token data available. */ + val totalTokens = integer("total_tokens").nullable() + + /** Sum of wall-clock duration across all steps in milliseconds; null for old rows. */ + val totalDurationMs = long("total_duration_ms").nullable() + val errorMessage = text("error_message").nullable() val createdAt = sqliteInstant("created_at") val updatedAt = sqliteInstant("updated_at") diff --git a/shared/src/main/kotlin/io/askimo/core/plan/repository/PlanExecutionRepository.kt b/shared/src/main/kotlin/io/askimo/core/plan/repository/PlanExecutionRepository.kt index 4baf8179..79cb4260 100644 --- a/shared/src/main/kotlin/io/askimo/core/plan/repository/PlanExecutionRepository.kt +++ b/shared/src/main/kotlin/io/askimo/core/plan/repository/PlanExecutionRepository.kt @@ -10,6 +10,15 @@ import io.askimo.core.logging.logger import io.askimo.core.plan.domain.PlanExecution import io.askimo.core.plan.domain.PlanExecutionStatus import io.askimo.core.plan.domain.PlanExecutionsTable +import io.askimo.core.plan.domain.PlanStepOutput +import io.askimo.core.util.JsonUtils.json +import kotlinx.serialization.json.JsonArray +import kotlinx.serialization.json.JsonObject +import kotlinx.serialization.json.contentOrNull +import kotlinx.serialization.json.intOrNull +import kotlinx.serialization.json.jsonArray +import kotlinx.serialization.json.jsonPrimitive +import kotlinx.serialization.json.longOrNull import org.jetbrains.exposed.v1.core.ResultRow import org.jetbrains.exposed.v1.core.SortOrder import org.jetbrains.exposed.v1.core.eq @@ -24,9 +33,8 @@ import java.util.UUID /** * Repository for persisting and querying [PlanExecution] records. * - * Inputs are stored as a simple `key=value` text (one per line) to avoid a JSON - * dependency at this layer. The format is intentionally human-readable so the - * SQLite file can be inspected without tooling. + * Inputs are stored as a simple `key=value` text (one per line). + * Step outputs are stored as JSON for extensibility. */ class PlanExecutionRepository internal constructor( databaseManager: DatabaseManager = DatabaseManager.getInstance(), @@ -55,6 +63,10 @@ class PlanExecutionRepository internal constructor( it[sessionId] = record.sessionId it[output] = record.output it[stepOutputs] = encodeStepOutputs(record.stepOutputs) + it[totalInputTokens] = record.totalInputTokens + it[totalOutputTokens] = record.totalOutputTokens + it[totalTokens] = record.totalTokens + it[totalDurationMs] = record.totalDurationMs it[errorMessage] = record.errorMessage it[createdAt] = record.createdAt it[updatedAt] = record.updatedAt @@ -78,6 +90,10 @@ class PlanExecutionRepository internal constructor( it[sessionId] = record.sessionId it[output] = record.output it[stepOutputs] = encodeStepOutputs(record.stepOutputs) + it[totalInputTokens] = record.totalInputTokens + it[totalOutputTokens] = record.totalOutputTokens + it[totalTokens] = record.totalTokens + it[totalDurationMs] = record.totalDurationMs it[errorMessage] = record.errorMessage it[runCount] = record.runCount it[updatedAt] = record.updatedAt @@ -149,6 +165,10 @@ class PlanExecutionRepository internal constructor( sessionId = this[PlanExecutionsTable.sessionId], output = this[PlanExecutionsTable.output], stepOutputs = decodeStepOutputs(this[PlanExecutionsTable.stepOutputs]), + totalInputTokens = this[PlanExecutionsTable.totalInputTokens], + totalOutputTokens = this[PlanExecutionsTable.totalOutputTokens], + totalTokens = this[PlanExecutionsTable.totalTokens], + totalDurationMs = this[PlanExecutionsTable.totalDurationMs], errorMessage = this[PlanExecutionsTable.errorMessage], createdAt = this[PlanExecutionsTable.createdAt], updatedAt = this[PlanExecutionsTable.updatedAt], @@ -168,23 +188,50 @@ class PlanExecutionRepository internal constructor( } } - /** - * Encodes step outputs as `stepName|output` lines (one per step). - * Step names are kebab-case identifiers and cannot contain `|` or newlines. - */ - private fun encodeStepOutputs(steps: List>): String? { + /** Encodes structured step outputs as a JSON array. */ + private fun encodeStepOutputs(steps: List): String? { if (steps.isEmpty()) return null - return steps.joinToString("\n") { (name, output) -> - "$name|${output.replace("\n", "\\n")}" - } + return json.encodeToString(steps) } - private fun decodeStepOutputs(raw: String?): List> { + private fun decodeStepOutputs(raw: String?): List { if (raw.isNullOrBlank()) return emptyList() + + // New format: JSON array of PlanStepOutput. + runCatching { + json.decodeFromString>(raw) + }.getOrNull()?.let { return it } + + // Compatibility format: JSON with legacy keys {"step":"...","output":"..."}. + val legacySteps = runCatching { + val element = json.parseToJsonElement(raw) + if (element !is JsonArray) return@runCatching null + element.jsonArray.mapNotNull { item -> + val obj = item as? JsonObject ?: return@mapNotNull null + val stepName = obj["stepName"]?.jsonPrimitive?.contentOrNull + ?: obj["step"]?.jsonPrimitive?.contentOrNull + ?: return@mapNotNull null + val output = obj["output"]?.jsonPrimitive?.contentOrNull ?: "" + PlanStepOutput( + stepName = stepName, + output = output, + inputTokens = obj["inputTokens"]?.jsonPrimitive?.intOrNull, + outputTokens = obj["outputTokens"]?.jsonPrimitive?.intOrNull, + totalTokens = obj["totalTokens"]?.jsonPrimitive?.intOrNull, + durationMs = obj["durationMs"]?.jsonPrimitive?.longOrNull, + ) + } + }.getOrNull() + if (legacySteps != null) return legacySteps + + // Legacy line format: stepName|output return raw.lines().mapNotNull { line -> val sep = line.indexOf('|') if (sep < 0) return@mapNotNull null - line.substring(0, sep) to line.substring(sep + 1).replace("\\n", "\n") + PlanStepOutput( + stepName = line.substring(0, sep), + output = line.substring(sep + 1).replace("\\n", "\n"), + ) } } }