root/goodbuddy
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

feat: separate system task token usage and add usage refresh

Browse Source 
Add a System tasks tab to run history that splits background usage (supervisor, auto-supervision, magic notes, knowledge base, delegation) from conversation usage, with a manual refresh button and automatic refresh when usage is recorded. Record previously dropped usage from supervisor reviews, image regeneration, knowledge graph extraction, and remote embeddings.
This commit is contained in:
mesalogo
2026-09-29 10:38:37 +08:00
parent 3c9229ac7a
commit 9a9170441e
24 changed files with 991 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/assistant/assistant-database.test.ts
+57 -18
View File
@@ -6900,16 +6900,19 @@ describe('AssistantDatabase', () => {
cacheWrite: 12
})
const usageTotals = {
callCount: 1,
input: 125,
output: 25,
cacheRead: 40,
cacheWrite: 12,
cacheInput: 177,
totalTokens: 150
}
expect(database.getTokenUsageSummary()).toEqual({
totals: {
callCount: 1,
input: 125,
output: 25,
cacheRead: 40,
cacheWrite: 12,
cacheInput: 177,
totalTokens: 150
},
totals: usageTotals,
conversationTotals: usageTotals,
systemTotals: expect.objectContaining({ callCount: 0, totalTokens: 0 }),
records: [
expect.objectContaining({
requestId: taskId,
@@ -7525,18 +7528,54 @@ describe('AssistantDatabase', () => {
expect(database.listTasks()).toEqual([])
expect(database.listArtifacts(project.id)).toEqual([])
expect(database.listMagicNotes()).toEqual([])
const emptyTotals = {
callCount: 0,
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
cacheInput: 0,
totalTokens: 0
}
expect(database.getTokenUsageSummary()).toEqual({
totals: {
callCount: 0,
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
cacheInput: 0,
totalTokens: 0
},
totals: emptyTotals,
conversationTotals: emptyTotals,
systemTotals: emptyTotals,
records: []
})
database.close()
})
it('separates system task usage from conversation usage', async () => {
const database = await createDatabase()
const usage = { runtime: 'model', provider: 'openai', model: 'gpt', input: 10, output: 5, cacheRead: 0, cacheWrite: 0 }
const conversationTask = '00000000-0000-4000-8000-000000000401'
database.createTask({ id: conversationTask, conversationId: 'chat-1', title: 'Chat', instructions: 'Reply', workMode: 'ask' })
database.upsertModelUsageCall({ ...usage, requestId: conversationTask, callId: 'chat' })
const systemTasks = [
['00000000-0000-4000-8000-000000000402', 'heartbeat:a', 'heartbeat'],
['00000000-0000-4000-8000-000000000403', 'supervision:b', 'supervision'],
['00000000-0000-4000-8000-000000000404', undefined, 'magic-notes']
] as const
for (const [id, conversationId] of systemTasks) {
database.createTask({ id, conversationId, title: 'System', instructions: 'Run', workMode: 'ask', origin: 'assistant', visible: false })
database.upsertModelUsageCall({ ...usage, requestId: id, callId: 'system' })
}
database.recordSystemModelUsage({ ...usage, runtime: 'embedding', source: 'knowledge', bucket: 'embedding', output: 0 })
database.recordSystemModelUsage({ ...usage, runtime: 'embedding', source: 'knowledge', bucket: 'embedding', output: 0 })
const summary = database.getTokenUsageSummary()
expect(summary.conversationTotals).toMatchObject({ callCount: 1, input: 10, output: 5 })
expect(summary.systemTotals).toMatchObject({ callCount: 5, input: 50, output: 15 })
expect(summary.totals).toMatchObject({ callCount: 6, input: 60, output: 20 })
expect(summary.records.find(record => record.requestId === conversationTask)?.systemSource).toBeUndefined()
for (const [id, , source] of systemTasks) {
expect(summary.records.find(record => record.requestId === id)?.systemSource).toBe(source)
}
const knowledge = summary.records.filter(record => record.systemSource === 'knowledge')
expect(knowledge).toHaveLength(1)
expect(knowledge[0]).toMatchObject({ callCount: 2, input: 20, output: 0 })
expect(database.listTasks().map(task => task.id)).toEqual([conversationTask])
database.close()
})
})
src/main/assistant/assistant-database.ts
+165 -75
View File
@@ -58,9 +58,13 @@ import type {
ProjectChannel,
ProjectCreateInput,
ScheduleCreateInput,
SystemModelUsageInput,
TokenUsageRecord,
TokenUsageSummary
TokenUsageSummary,
TokenUsageSystemSource,
TokenUsageTotals
} from '../../shared/assistant-contracts'
import { tokenUsageSystemSources } from '../../shared/assistant-contracts'
import type { StoredSupervisionResult, SupervisionCandidate } from './supervisor-service'
import { supervisionEntitySchema, supervisionResultViewSchema, type SupervisionTarget, type SupervisionGraphRequest, type SupervisionRunRequest, type SupervisionActivity } from '../../shared/supervision-contracts'
import type { AgentEvent } from '../../shared/contracts'
@@ -466,6 +470,7 @@ type TokenUsageRecordRow = {
project_name: string | null
conversation_id: string | null
conversation_title: string | null
system_source: string | null
runtime: string
provider: string
model: string
@@ -1824,6 +1829,7 @@ export class AssistantDatabase {
private subagentProgress?: SubagentProgressStorage
private channelEventWrites = 0
private channelOutboxWrites = 0
private modelUsageNotificationQueued = false
private readonly executionStatsCache = new Map<
string,
{ value: ExecutionStats; cachedAt: number; dataVersion: number; totalChanges: number }
@@ -1834,6 +1840,7 @@ export class AssistantDatabase {
private readonly options: {
onMagicNotesChanged?: () => void
onMagicTodosChanged?: () => void
onModelUsageChanged?: () => void
} = {}
) {
this.noteStorage = new MagicNoteStorage(databasePath)
@@ -5577,53 +5584,121 @@ export class AssistantDatabase {
now,
now
)
this.notifyModelUsageChanged()
}
/**
* Records usage for background work that has no request of its own, such as
* knowledge graph extraction or remote embeddings. Calls are grouped into one
* hidden task per source, bucket and UTC day so the foreign key to tasks holds.
*/
recordSystemModelUsage(input: SystemModelUsageInput): void {
const day = new Date().toISOString().slice(0, 10)
const taskId = `${input.source}:${input.bucket}:${day}`
const now = new Date().toISOString()
this.requireDatabase()
.prepare(
`INSERT OR IGNORE INTO tasks
(id, project_id, conversation_id, title, instructions, origin,
status, priority, work_mode, created_at, started_at, completed_at,
visible)
VALUES (?, NULL, ?, ?, ?, 'assistant', 'completed', 0, 'ask', ?, ?, ?, 0)`
)
.run(
taskId,
taskId,
input.bucket === 'embedding' ? '知识库向量化' : '知识图谱抽取',
'后台知识库模型用量汇总',
now,
now,
now
)
this.upsertModelUsageCall({
...input,
requestId: taskId,
callId: `system-call:${randomUUID()}`
})
}
private notifyModelUsageChanged(): void {
if (this.modelUsageNotificationQueued || !this.options.onModelUsageChanged) {
return
}
this.modelUsageNotificationQueued = true
// Coalesce bursts of per-call upserts into one renderer refresh.
const timer = setTimeout(() => {
this.modelUsageNotificationQueued = false
this.options.onModelUsageChanged?.()
}, 500)
timer.unref?.()
}
getTokenUsageSummary(): TokenUsageSummary {
const rows = this.requireDatabase()
.prepare(
`SELECT
usage.request_id,
tasks.project_id,
projects.name AS project_name,
tasks.conversation_id,
conversations.title AS conversation_title,
usage.runtime,
usage.provider,
usage.model,
COUNT(*) AS call_count,
SUM(usage.input_tokens) AS input_tokens,
SUM(usage.output_tokens) AS output_tokens,
SUM(usage.cache_read_tokens) AS cache_read_tokens,
SUM(usage.cache_write_tokens) AS cache_write_tokens,
SUM(
usage.input_tokens +
CASE
WHEN LOWER(usage.provider) LIKE '%anthropic%'
THEN usage.cache_read_tokens + usage.cache_write_tokens
ELSE 0
END
) AS cache_input_tokens
FROM model_usage_calls usage
JOIN tasks ON tasks.id = usage.request_id
LEFT JOIN projects ON projects.id = tasks.project_id
LEFT JOIN conversations
ON conversations.id = tasks.conversation_id
GROUP BY
usage.request_id,
tasks.project_id,
projects.name,
tasks.conversation_id,
conversations.title,
usage.runtime,
usage.provider,
usage.model
ORDER BY MAX(usage.updated_at) DESC
LIMIT 500`
)
.all() as TokenUsageRecordRow[]
const records: TokenUsageRecord[] = rows.map((row) => ({
const systemSourceSql = `CASE
WHEN tasks.conversation_id LIKE 'heartbeat:%' THEN 'heartbeat'
WHEN tasks.conversation_id LIKE 'supervision:%' THEN 'supervision'
WHEN tasks.conversation_id LIKE 'knowledge:%' THEN 'knowledge'
WHEN tasks.conversation_id LIKE 'delegation:%' THEN 'delegation'
WHEN tasks.origin = 'assistant' AND tasks.conversation_id IS NULL
THEN 'magic-notes'
END`
const database = this.requireDatabase()
// Conversation and system rows are limited separately so frequent chat
// usage cannot hide background usage (or the reverse).
const selectRows = (system: boolean): TokenUsageRecordRow[] =>
database
.prepare(
`SELECT
usage.request_id,
tasks.project_id,
projects.name AS project_name,
tasks.conversation_id,
conversations.title AS conversation_title,
${systemSourceSql} AS system_source,
usage.runtime,
usage.provider,
usage.model,
COUNT(*) AS call_count,
SUM(usage.input_tokens) AS input_tokens,
SUM(usage.output_tokens) AS output_tokens,
SUM(usage.cache_read_tokens) AS cache_read_tokens,
SUM(usage.cache_write_tokens) AS cache_write_tokens,
SUM(
usage.input_tokens +
CASE
WHEN LOWER(usage.provider) LIKE '%anthropic%'
THEN usage.cache_read_tokens + usage.cache_write_tokens
ELSE 0
END
) AS cache_input_tokens
FROM model_usage_calls usage
JOIN tasks ON tasks.id = usage.request_id
LEFT JOIN projects ON projects.id = tasks.project_id
LEFT JOIN conversations
ON conversations.id = tasks.conversation_id
WHERE (${systemSourceSql}) IS ${system ? 'NOT ' : ''}NULL
GROUP BY
usage.request_id,
tasks.project_id,
projects.name,
tasks.conversation_id,
conversations.title,
usage.runtime,
usage.provider,
usage.model
ORDER BY MAX(usage.updated_at) DESC
LIMIT 500`
)
.all() as TokenUsageRecordRow[]
const knownSources = new Set<string>(tokenUsageSystemSources)
const records: TokenUsageRecord[] = [
...selectRows(false),
...selectRows(true)
].map((row) => ({
requestId: row.request_id,
...(row.system_source && knownSources.has(row.system_source)
? { systemSource: row.system_source as TokenUsageSystemSource }
: {}),
projectId: row.project_id ?? undefined,
projectName: row.project_name ?? undefined,
conversationId: row.conversation_id ?? undefined,
@@ -5639,28 +5714,8 @@ export class AssistantDatabase {
cacheInput: row.cache_input_tokens,
totalTokens: row.input_tokens + row.output_tokens
}))
const totalRow = this.requireDatabase()
.prepare(
`SELECT
COUNT(*) AS call_count,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(cache_write_tokens), 0) AS cache_write_tokens,
COALESCE(
SUM(
input_tokens +
CASE
WHEN LOWER(provider) LIKE '%anthropic%'
THEN cache_read_tokens + cache_write_tokens
ELSE 0
END
),
0
) AS cache_input_tokens
FROM model_usage_calls`
)
.get() as {
type TotalRow = {
is_system: number | null
call_count: number
input_tokens: number
output_tokens: number
@@ -5668,16 +5723,51 @@ export class AssistantDatabase {
cache_write_tokens: number
cache_input_tokens: number
}
const totals = {
callCount: totalRow.call_count,
input: totalRow.input_tokens,
output: totalRow.output_tokens,
cacheRead: totalRow.cache_read_tokens,
cacheWrite: totalRow.cache_write_tokens,
cacheInput: totalRow.cache_input_tokens,
totalTokens: totalRow.input_tokens + totalRow.output_tokens
// Usage rows always reference a task (foreign key), so the join is lossless.
const totalRows = database
.prepare(
`SELECT
(${systemSourceSql}) IS NOT NULL AS is_system,
COUNT(*) AS call_count,
COALESCE(SUM(usage.input_tokens), 0) AS input_tokens,
COALESCE(SUM(usage.output_tokens), 0) AS output_tokens,
COALESCE(SUM(usage.cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(usage.cache_write_tokens), 0) AS cache_write_tokens,
COALESCE(
SUM(
usage.input_tokens +
CASE
WHEN LOWER(usage.provider) LIKE '%anthropic%'
THEN usage.cache_read_tokens + usage.cache_write_tokens
ELSE 0
END
),
0
) AS cache_input_tokens
FROM model_usage_calls usage
LEFT JOIN tasks ON tasks.id = usage.request_id
GROUP BY is_system`
)
.all() as TotalRow[]
const toTotals = (rows: TotalRow[]): TokenUsageTotals => {
const sum = (key: Exclude<keyof TotalRow, 'is_system'>): number =>
rows.reduce((total, row) => total + row[key], 0)
return {
callCount: sum('call_count'),
input: sum('input_tokens'),
output: sum('output_tokens'),
cacheRead: sum('cache_read_tokens'),
cacheWrite: sum('cache_write_tokens'),
cacheInput: sum('cache_input_tokens'),
totalTokens: sum('input_tokens') + sum('output_tokens')
}
}
return {
totals: toTotals(totalRows),
conversationTotals: toTotals(totalRows.filter((row) => !row.is_system)),
systemTotals: toTotals(totalRows.filter((row) => row.is_system)),
records
}
return { totals, records }
}
updateTaskStatus(
src/main/assistant/supervision-production.test.ts
+18
View File
@@ -131,6 +131,24 @@ function fixture(content = 'Atlas and Beacon are separate projects.') {
return { db, request, respond, service }
}
it('records supervisor model usage under a hidden supervision task', async () => {
const f = fixture()
const run: AgentRuntime['run'] = async function* (input) {
yield { type: 'text', requestId: input.requestId, delta: JSON.stringify(f.respond(input.prompt)) }
yield { type: 'model-usage', requestId: input.requestId, callId: 'call-1', runtime: 'model', provider: 'openai',
model: 'gpt-review', inputTokens: 40, outputTokens: 8, cacheReadTokens: 0, cacheWriteTokens: 0 }
yield { type: 'done', requestId: input.requestId }
}
const service = createProductionSupervisorService(f.db, async () => ({ supervisorModelConcurrency: 1,
supervisionReview: { pageSize: 10, batchCharacters: 1000, batchMessages: 10, executionSeconds: 300, responseKiB: 1024 } }),
async () => ({ runtimeId: 'model', capability: 'chat', run } as AgentRuntime), new SupervisionModelPool())
await service.run(f.request)
const summary = f.db.getTokenUsageSummary()
expect(summary.records).toEqual([expect.objectContaining({ systemSource: 'supervision', model: 'gpt-review', input: 40, output: 8 })])
expect(summary.systemTotals).toMatchObject({ callCount: 1, input: 40 })
expect(f.db.listTasks()).toEqual([])
})
it('production factory normalizes two absent identities, then reuses exact candidate UUIDs without extra calls', async () => {
const f = fixture()
const first = await f.service().run(f.request)
src/main/assistant/supervision-production.ts
+22 -5
View File
@@ -2,7 +2,7 @@ import { randomUUID } from 'node:crypto'
import type { ApplicationSettings } from '../../shared/application-settings-contracts'
import { defaultSupervisionTimeoutSeconds, defaultSupervisorModelConcurrency } from '../../shared/application-settings-contracts'
import { supervisionReviewSettingsSchema } from '../../shared/supervision-review-contracts'
import type { AgentRuntime } from '../agent/runtime'
import type { AgentRuntime, RuntimeModelUsageEvent } from '../agent/runtime'
import type { AssistantDatabase } from './assistant-database'
import type { SupervisionModelPool } from './supervision-model-pool'
import { SupervisorService } from './supervisor-service'
@@ -11,7 +11,11 @@ export function createProductionSupervisorService(
database: AssistantDatabase,
getSettings: () => Promise<Partial<ApplicationSettings> | undefined>,
resolveRuntime: () => Promise<AgentRuntime>,
pool: SupervisionModelPool
pool: SupervisionModelPool,
persistUsage: (event: RuntimeModelUsageEvent) => void = event => database.upsertModelUsageCall({
requestId: event.requestId, callId: event.callId, runtime: event.runtime, provider: event.provider, model: event.model,
input: event.inputTokens, output: event.outputTokens, cacheRead: event.cacheReadTokens, cacheWrite: event.cacheWriteTokens
})
): SupervisorService {
return new SupervisorService({ collect: async () => { throw new Error('Paged review collector required') } }, {
summarize: async request => {
@@ -30,9 +34,13 @@ export function createProductionSupervisorService(
}, timeoutSeconds * 1000)
const responseKiB = settings?.supervisionReview?.responseKiB ?? 1024
let output = '', completed = false
const conversationId = `supervision:${randomUUID()}`
const requestId = randomUUID()
const conversationId = `supervision:${requestId}`
// Hidden task row so model usage can be attributed to supervision.
database.createTask({ id: requestId, conversationId, title: '监督者回顾', instructions: '根据有界证据整理监督者回顾',
workMode: 'ask', origin: 'assistant', visible: false })
try {
for await (const event of runtime.run({ requestId: randomUUID(), conversationId, workMode: 'ask',
for await (const event of runtime.run({ requestId, conversationId, workMode: 'ask',
prompt: [request.systemInstruction, 'OUTPUT CONTRACT:', request.outputContract,
'REVIEW TIME RANGE:', JSON.stringify(request.request.timeRange), 'KNOWN ENTITIES:', JSON.stringify(request.candidates.map((candidate, index) => ({
candidateRef: `known_${index + 1}`, label: candidate.label, description: candidate.description
@@ -48,14 +56,23 @@ export function createProductionSupervisorService(
modelSignal.throwIfAborted()
}
}
if (event.type === 'model-usage') persistUsage(event)
if (event.type === 'tool') throw new Error('监督者只允许只读模型摘要，不允许工具调用')
if (event.type === 'error') throw new Error(event.message)
if (event.type === 'done') completed = true
}
modelSignal.throwIfAborted()
if (!completed) throw new Error('监督者模型未报告完成')
database.updateTaskStatus(requestId, 'completed')
return output
} catch (error) { throw timeoutError ?? error }
} catch (error) {
const failure = timeoutError ?? error
try {
database.updateTaskStatus(requestId, modelSignal.aborted ? 'cancelled' : 'failed',
failure instanceof Error ? failure.message.slice(0, 2_000) : '监督者模型调用失败')
} catch { /* status bookkeeping must not mask the model failure */ }
throw failure
}
finally { clearTimeout(timeout); await runtime.releaseConversation?.(conversationId) }
}, request.signal)
}