From 65f5d24620e76867062091a528d66d0f0837a4dc Mon Sep 17 00:00:00 2001 From: _Kerman Date: Thu, 9 Jul 2026 15:23:07 +0800 Subject: [PATCH] fix: align v2 media and task compaction handling --- .../agent-core-v2/docs/di-scope-domains.puml | 1 + .../agent-core-v2/src/_base/utils/tokens.ts | 21 +++-- .../src/agent/task/taskService.ts | 41 ++++++++- .../src/agent/task/tools/task-list.ts | 2 +- .../test/_base/utils/tokens.test.ts | 61 +++++++++++++ .../agent-core-v2/test/task/service.test.ts | 85 ++++++++++++++++++- 6 files changed, 199 insertions(+), 12 deletions(-) create mode 100644 packages/agent-core-v2/test/_base/utils/tokens.test.ts diff --git a/packages/agent-core-v2/docs/di-scope-domains.puml b/packages/agent-core-v2/docs/di-scope-domains.puml index 0ca7df5c7..11a6201eb 100644 --- a/packages/agent-core-v2/docs/di-scope-domains.puml +++ b/packages/agent-core-v2/docs/di-scope-domains.puml @@ -285,6 +285,7 @@ task --> wireRecord #34495E task --> telemetry #34495E task --> prompt #34495E task --> contextMemory #34495E +task --> contextInjector #34495E task --> config #34495E task --> storage #34495E task --> session_context #34495E diff --git a/packages/agent-core-v2/src/_base/utils/tokens.ts b/packages/agent-core-v2/src/_base/utils/tokens.ts index 5a5229af0..a45b24afd 100644 --- a/packages/agent-core-v2/src/_base/utils/tokens.ts +++ b/packages/agent-core-v2/src/_base/utils/tokens.ts @@ -64,11 +64,22 @@ export function estimateTokensForContentParts(parts: readonly ContentPart[]): nu return total; } +export const MEDIA_TOKEN_ESTIMATE = 2000; + export function estimateTokensForContentPart(part: ContentPart): number { - if (part.type === 'text') { - return estimateTokens(part.text); - } else if (part.type === 'think') { - return estimateTokens(part.think); + switch (part.type) { + case 'text': + return estimateTokens(part.text); + case 'think': + return estimateTokens(part.think); + case 'image_url': + case 'audio_url': + case 'video_url': + return MEDIA_TOKEN_ESTIMATE; + default: { + const exhaustive: never = part; + void exhaustive; + return 0; + } } - return 0; } diff --git a/packages/agent-core-v2/src/agent/task/taskService.ts b/packages/agent-core-v2/src/agent/task/taskService.ts index d0a9e34d6..336c8b1f9 100644 --- a/packages/agent-core-v2/src/agent/task/taskService.ts +++ b/packages/agent-core-v2/src/agent/task/taskService.ts @@ -7,8 +7,9 @@ * limits through `config`, records lifecycle and broadcasts through `wire` * (`task.started` / `task.terminated` Ops into `TaskModel`, plus the matching * signals), restores ghosts through a single `wire.onRestored` handler (wire - * replay -> disk load -> reconcile, in that order), and delivers terminal - * notifications through `contextMemory`. Bound at Agent scope. + * replay -> disk load -> reconcile, in that order), delivers terminal + * notifications through `contextMemory`, and re-surfaces active tasks through + * `contextInjector` after compaction. Bound at Agent scope. */ import { randomBytes } from 'node:crypto'; @@ -21,6 +22,7 @@ import { Disposable } from '#/_base/di/lifecycle'; import { escapeXml, escapeXmlAttr } from '#/_base/utils/xml-escape'; import { IEventBus } from '#/app/event/eventBus'; import type { TaskOrigin } from '#/agent/contextMemory/types'; +import { IAgentContextInjectorService } from '#/agent/contextInjector/contextInjector'; import { ITaskService, type ITaskHandle, TERMINAL_TASK_STATES } from '#/app/task/task'; import { TERMINAL_STATUSES, @@ -55,7 +57,7 @@ import { import { LEGACY_BACKGROUND_SECTION, TASK_SECTION, type AgentTaskConfig } from './configSection'; import { AgentTaskPersistence } from './persist'; import { TaskModel, taskStarted, taskTerminated } from './taskOps'; -import '#/agent/task/tools/task-list'; +import { formatTaskList } from '#/agent/task/tools/task-list'; import '#/agent/task/tools/task-output'; import '#/agent/task/tools/task-stop'; @@ -147,6 +149,11 @@ const SIGTERM_GRACE_MS = 5_000; const TASK_ID_ALPHABET = '0123456789abcdefghijklmnopqrstuvwxyz'; const USER_INTERRUPT_REASON = 'Interrupted by user'; const NOTIFICATION_FALLBACK_PREVIEW_BYTES = 3_000; +const ACTIVE_BACKGROUND_TASK_INJECTION_VARIANT = 'background_task_status'; +const ACTIVE_BACKGROUND_TASK_GUIDANCE = [ + 'The conversation was compacted, so the earlier messages that started these background tasks are gone - but the tasks are still running from before.', + 'Do not start duplicates. Use TaskOutput to fetch a task result, TaskList to list them, and TaskStop to cancel one.', +].join(' '); export function isAgentTaskTerminal(status: AgentTaskStatus): boolean { return TERMINAL_STATUSES.has(status); @@ -184,6 +191,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { private readonly scheduledNotificationKeys = new Set(); private readonly deliveredNotificationKeys = new Set(); private readonly persistence: AgentTaskPersistence; + private activeTaskReminderPending = false; constructor( @ITelemetryService private readonly telemetry: ITelemetryService, @@ -197,6 +205,7 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { @IAgentWireRecordService wireRecord: IAgentWireRecordService, @IAgentWireService private readonly wire: IWireService, @IEventBus private readonly eventBus: IEventBus, + @IAgentContextInjectorService injector: IAgentContextInjectorService, ) { super(); this.persistence = new AgentTaskPersistence( @@ -208,6 +217,9 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { this._register(this.wire.onRestored(() => this.restoreAfterReplay())); this._register( this.eventBus.subscribe('context.spliced', (e) => { + if (isCompactionSplice(e)) { + this.activeTaskReminderPending = true; + } for (const message of e.messages) { if (isTaskOrigin(message.origin)) { this.markDeliveredNotification(message.origin); @@ -215,6 +227,11 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { } }), ); + this._register( + injector.register(ACTIVE_BACKGROUND_TASK_INJECTION_VARIANT, () => + this.activeBackgroundTaskReminder(), + ), + ); this._register( wireRecord.hooks.onRestoredRecord.register( 'task-delivered-notifications', @@ -239,6 +256,14 @@ export class AgentTaskService extends Disposable implements IAgentTaskService { await this.reconcile(); } + private activeBackgroundTaskReminder(): string | undefined { + if (!this.activeTaskReminderPending) return undefined; + this.activeTaskReminderPending = false; + const tasks = this.list(true); + if (tasks.length === 0) return undefined; + return `${ACTIVE_BACKGROUND_TASK_GUIDANCE}\n\n${formatTaskList(tasks, true)}`; + } + private restoreGhostsFromWire(): void { for (const [taskId, info] of this.wire.getModel(TaskModel)) { if (this.tasks.has(taskId)) continue; @@ -1119,6 +1144,16 @@ function shouldListTask(info: AgentTaskInfo, activeOnly: boolean): boolean { return info.detached !== false; } +function isCompactionSplice(splice: { + readonly deleteCount: number; + readonly messages: readonly { readonly origin?: { readonly kind: string } | undefined }[]; +}): boolean { + return ( + splice.deleteCount > 0 && + splice.messages.some((message) => message.origin?.kind === 'compaction_summary') + ); +} + function newerRestoredTask( existing: AgentTaskInfo, loaded: AgentTaskInfo, diff --git a/packages/agent-core-v2/src/agent/task/tools/task-list.ts b/packages/agent-core-v2/src/agent/task/tools/task-list.ts index 4cdd2147b..1e17f23af 100644 --- a/packages/agent-core-v2/src/agent/task/tools/task-list.ts +++ b/packages/agent-core-v2/src/agent/task/tools/task-list.ts @@ -36,7 +36,7 @@ export type TaskListInput = z.infer; // ── Implementation ─────────────────────────────────────────────────── -function formatTaskList(tasks: readonly AgentTaskInfo[], activeOnly: boolean): string { +export function formatTaskList(tasks: readonly AgentTaskInfo[], activeOnly: boolean): string { // `active_only=false` mixes in terminal/lost tasks, so the count is no // longer purely "active" — use a neutral label to avoid mislabeling them. const label = activeOnly ? 'active_background_tasks' : 'background_tasks'; diff --git a/packages/agent-core-v2/test/_base/utils/tokens.test.ts b/packages/agent-core-v2/test/_base/utils/tokens.test.ts new file mode 100644 index 000000000..4af1e9689 --- /dev/null +++ b/packages/agent-core-v2/test/_base/utils/tokens.test.ts @@ -0,0 +1,61 @@ +/** + * Scenario: token estimation for rich content parts. + * Responsibilities: media parts contribute bounded non-zero estimates to + * content-part and whole-message estimates. Wiring: pure utility functions, no + * collaborators. Run with: + * `vitest run --config packages/agent-core-v2/vitest.config.ts test/_base/utils/tokens.test.ts`. + */ + +import type { ContentPart } from '#/app/llmProtocol/message'; +import { describe, expect, it } from 'vitest'; + +import { + estimateTokensForContentPart, + estimateTokensForMessage, + MEDIA_TOKEN_ESTIMATE, +} from '#/_base/utils/tokens'; + +describe('token estimates for media content parts', () => { + const imagePart: ContentPart = { + type: 'image_url', + imageUrl: { url: 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAAB' }, + }; + const audioPart: ContentPart = { + type: 'audio_url', + audioUrl: { url: 'data:audio/mp3;base64,AAAA' }, + }; + const videoPart: ContentPart = { + type: 'video_url', + videoUrl: { url: 'data:video/mp4;base64,AAAA' }, + }; + + it('counts image parts with the fixed media estimate', () => { + expect(estimateTokensForContentPart(imagePart)).toBe(MEDIA_TOKEN_ESTIMATE); + expect(MEDIA_TOKEN_ESTIMATE).toBeGreaterThan(100); + }); + + it('counts audio and video parts as non-zero media', () => { + expect(estimateTokensForContentPart(audioPart)).toBe(MEDIA_TOKEN_ESTIMATE); + expect(estimateTokensForContentPart(videoPart)).toBe(MEDIA_TOKEN_ESTIMATE); + }); + + it('keeps large data URLs bounded instead of counting base64 as text', () => { + const part: ContentPart = { + type: 'image_url', + imageUrl: { url: `data:image/png;base64,${'A'.repeat(4_000_000)}` }, + }; + + expect(estimateTokensForContentPart(part)).toBe(MEDIA_TOKEN_ESTIMATE); + expect(estimateTokensForContentPart(part)).toBeLessThan(50_000); + }); + + it('includes media when estimating a whole message', () => { + const estimate = estimateTokensForMessage({ + role: 'user', + content: [{ type: 'text', text: 'see screenshot' }, imagePart], + toolCalls: [], + }); + + expect(estimate).toBeGreaterThan(100); + }); +}); diff --git a/packages/agent-core-v2/test/task/service.test.ts b/packages/agent-core-v2/test/task/service.test.ts index a8553e7ca..0c90887d4 100644 --- a/packages/agent-core-v2/test/task/service.test.ts +++ b/packages/agent-core-v2/test/task/service.test.ts @@ -5,6 +5,11 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; import { DisposableStore, toDisposable } from '#/_base/di/lifecycle'; import { TestInstantiationService } from '#/_base/di/test'; +import { + IAgentContextInjectorService, + type ContextInjectionContext, + type ContextInjectionProvider, +} from '#/agent/contextInjector/contextInjector'; import { IAgentTaskService, type AgentTask, @@ -16,6 +21,7 @@ import { ProcessTask } from '#/os/backends/node-local/tools/process-task'; import type { IProcess } from '#/session/process/processRunner'; import { IConfigRegistry, IConfigService } from '#/app/config/config'; import { IAgentContextMemoryService } from '#/agent/contextMemory/contextMemory'; +import type { ContextMessage } from '#/agent/contextMemory/types'; import { IAgentPromptService } from '#/agent/prompt/prompt'; import { ISessionContext } from '#/session/sessionContext/sessionContext'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; @@ -26,6 +32,7 @@ import { IAgentWireRecordService } from '#/agent/wireRecord/wireRecord'; import { IAgentWireService } from '#/wire/tokens'; import type { IWireService } from '#/wire/wireService'; import { IEventBus } from '#/app/event/eventBus'; +import { EventBusService } from '#/app/event/eventBusService'; import { ITaskService } from '#/app/task/task'; import { stubContextMemory, stubWireRecord } from '../contextMemory/stubs'; @@ -58,15 +65,24 @@ function stubWireService(): IWireService { describe('AgentTaskService', () => { let disposables: DisposableStore; let ix: TestInstantiationService; + let eventBus: EventBusService; + let injectionProviders: Map; beforeEach(() => { disposables = new DisposableStore(); ix = disposables.add(new TestInstantiationService()); + eventBus = disposables.add(new EventBusService()); + injectionProviders = new Map(); ix.stub(IAgentWireRecordService, stubWireRecord()); ix.stub(IAgentWireService, stubWireService()); - ix.stub(IEventBus, { - publish: () => {}, - subscribe: () => toDisposable(() => {}), + ix.stub(IEventBus, eventBus); + ix.stub(IAgentContextInjectorService, { + register: (name, provider) => { + injectionProviders.set(name, provider); + return toDisposable(() => { + injectionProviders.delete(name); + }); + }, }); ix.stub(ITaskService, { run: () => { @@ -128,6 +144,69 @@ describe('AgentTaskService', () => { await svc.stop(id); }); + function compactionSummary(text: string): ContextMessage { + return { + role: 'user', + content: [{ type: 'text', text }], + toolCalls: [], + origin: { kind: 'compaction_summary' }, + }; + } + + function publishCompactionSplice(): void { + eventBus.publish({ + type: 'context.spliced', + start: 0, + deleteCount: 2, + messages: [compactionSummary('Compacted summary.')], + }); + } + + async function backgroundTaskReminder( + context: ContextInjectionContext = { + injectedPositions: [], + lastInjectedAt: null, + isNewTurn: false, + }, + ): Promise { + const provider = injectionProviders.get('background_task_status'); + expect(provider).toBeDefined(); + const content = await provider!(context); + return typeof content === 'string' ? content : undefined; + } + + it('injects active background task status when compaction dropped the original launch context', async () => { + const svc = ix.get(IAgentTaskService); + const taskId = svc.registerTask(fakeProcessTask()); + + expect(await backgroundTaskReminder()).toBeUndefined(); + + publishCompactionSplice(); + + const reminder = await backgroundTaskReminder(); + expect(reminder).toContain('The conversation was compacted'); + expect(reminder).toContain('active_background_tasks: 1'); + expect(reminder).toContain(taskId); + expect(reminder).toContain('TaskOutput'); + expect(reminder).toContain('TaskList'); + expect(reminder).toContain('TaskStop'); + expect(await backgroundTaskReminder()).toBeUndefined(); + + await svc.stop(taskId); + }); + + it('does not carry post-compaction task reminder eligibility forward when no task is active', async () => { + const svc = ix.get(IAgentTaskService); + publishCompactionSplice(); + + expect(await backgroundTaskReminder()).toBeUndefined(); + + const taskId = svc.registerTask(fakeProcessTask()); + expect(await backgroundTaskReminder()).toBeUndefined(); + + await svc.stop(taskId); + }); + // ── Output ceiling for shell (process) tasks ───────────────────────── // // A single foreground shell command that streams more output than the