fix: correct workflow token counters#718
Open
ArshAnan wants to merge 4 commits intopensarai:canaryfrom
Open
Conversation
… and cache metrics - Added `onStepFinish` and `onCacheMetrics` callbacks to the `PentestSwarmInput` interface for better tracking of agent performance and resource usage. - Updated related components to propagate these new callbacks, improving the overall observability of the pentesting process. - Included new dependencies in `package.json` and `bun.lock` for enhanced functionality. Made-with: Cursor
- Added `preserveLargerTokenUsage` option to `WriteExecutionMetricsInput` to allow for retaining larger token usage values across sessions. - Implemented `preserveLargerTokenUsage` logic in `writeExecutionMetrics` to conditionally merge token usage. - Updated `addTokenUsage` method in `AgentContextValue` to accept an optional total token count. - Enhanced `accumulateTokenUsage` function to handle total tokens, ensuring accurate tracking during workflow execution. - Adjusted `OperatorDashboard` to reset token usage on new sessions and accommodate the new total tokens feature. Co-authored-by: Cursor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #671.
This PR fixes token counter behavior for
/pentestworkflow runs. The bug had three related symptoms:/pentestlaunched workflow-managed agents, the footer token counters could stay at zero because usage from those workflow agents was not reliably flowing back into the same counter path used by normal operator messages.inputTokens + outputTokens, even when the AI SDK/provider reported an explicittotalTokensvalue.This branch fixes those paths in a few places:
src/core/agents/offSecAgent/tools/types.tsonStepFinishandonCacheMetricsto the tool context so tools that spawn internal agents can forward usage and cache metrics back to the parent operator run.src/core/agents/offSecAgent/offensiveSecurityAgent.tsonStepFinishandonCacheMetricscallbacks into tool creation.run_pentest_workflowreport token usage from nested workflow agents instead of isolating that usage inside the tool call.src/core/agents/offSecAgent/tools/runPentestWorkflow.tsctx.onStepFinishandctx.onCacheMetricsinto the deterministic pentest workflow./pentesttool call back into the operator dashboard’s token accounting.src/core/workflows/pentest.tsonStepFinishso each discovery/swarm agent step contributes to workflow metrics and still forwards to the parent UI callback.preserveLargerTokenUsage: truewhen writing workflow metrics so a workflow-local total cannot overwrite larger totals already recorded by the operator UI.src/tui/context/agent.tsxaddTokenUsage(input, output)to accept an optionaltotal.totalTokensis provided by the SDK, the footer now uses that reported value instead of always deriving total as input + output.src/tui/components/operator-dashboard/logic.tsaccumulateTokenUsageto accept an optional step-leveltotalTokens.src/tui/components/operator-dashboard/index.tsxevent.usage.totalTokensthrough both the local ref accumulator and the sharedAgentProvidercounter.src/core/session/execution-metrics.tspreserveLargerTokenUsagetowriteExecutionMetrics.inputTokens,outputTokens, andtotalTokensvalues instead of replacing them with smaller workflow-local totals.src/tui/components/operator-dashboard/logic.test.tstotalTokens.src/core/session/execution-metrics.test.tsTogether, these changes make workflow token usage follow the same accounting path as normal operator usage, reset cleanly for new sessions, hydrate correctly for resumed sessions, and avoid metric overwrites from workflow-local bookkeeping.
The attached screenshot shows
/pentestrunning with populated footer counters for input, output, cached, and total tokens.How did you verify your code works?
Ran focused tests for token accumulation and execution metric persistence:
bun test src/tui/components/operator-dashboard/logic.test.ts src/core/session/execution-metrics.test.tsRan TS Checks
Ran the full test suite