From 7c2e2fe1fac9a6905e7934fcf067e7c667b3b627 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sat, 16 May 2026 17:12:48 +0900 Subject: [PATCH] fix(background-agent): redact task registry views --- src/features/background-agent/manager.test.ts | 38 ++++++++++++ .../background-agent/task-registry.ts | 59 ++++++++++++++++--- 2 files changed, 90 insertions(+), 7 deletions(-) diff --git a/src/features/background-agent/manager.test.ts b/src/features/background-agent/manager.test.ts index 030b3670d..d1fca8183 100644 --- a/src/features/background-agent/manager.test.ts +++ b/src/features/background-agent/manager.test.ts @@ -6822,6 +6822,44 @@ describe("BackgroundManager regression fixes - resume and aborted notification", secondManager.shutdown() }) + test("should redact active task prompts resolved from an earlier plugin manager instance", () => { + //#given + const firstManager = createBackgroundManager() + const secondManager = createBackgroundManager() + const task: BackgroundTask = { + id: "task-cross-manager-active-redaction", + parentSessionId: "parent-session", + parentMessageId: "msg-1", + description: "cross manager active redaction", + prompt: "secret prompt", + agent: "explore", + status: "pending", + queuedAt: new Date(), + } + + //#when + ;(cast<{ addTask: (task: BackgroundTask) => void }>(firstManager)).addTask(task) + task.sessionId = "session-cross-manager-active-redaction" + task.status = "running" + task.startedAt = new Date() + task.progress = { + lastUpdate: new Date(), + toolCalls: 1, + countedToolPartIDs: new Set(["part-1"]), + } + + //#then + const localTask = firstManager.getTask(task.id) + const registeredTask = secondManager.getTask(task.id) + expect(localTask?.prompt).toBe("secret prompt") + expect(registeredTask?.sessionId).toBe(task.sessionId) + expect(registeredTask?.prompt).toBe("[redacted]") + expect(registeredTask?.progress?.countedToolPartIDs).toEqual(new Set(["part-1"])) + + firstManager.shutdown() + secondManager.shutdown() + }) + test("should resolve archived completed task from an earlier plugin manager instance", () => { //#given const firstManager = createBackgroundManager() diff --git a/src/features/background-agent/task-registry.ts b/src/features/background-agent/task-registry.ts index ae01889a9..8a027e667 100644 --- a/src/features/background-agent/task-registry.ts +++ b/src/features/background-agent/task-registry.ts @@ -4,7 +4,7 @@ const MAX_COMPLETED_TASK_REGISTRY_SIZE = 100 const REGISTRY_KEY = "__omoBackgroundTaskRegistry" type BackgroundTaskRegistry = { - activeTasks: Map + activeTasks: Map BackgroundTask> completedTasks: Map } @@ -22,28 +22,67 @@ const TERMINAL_TASK_STATUSES = new Set([ function getRegistry(): BackgroundTaskRegistry { const registryGlobal = globalThis as GlobalWithBackgroundTaskRegistry registryGlobal[REGISTRY_KEY] ??= { - activeTasks: new Map(), + activeTasks: new Map BackgroundTask>(), completedTasks: new Map(), } - return registryGlobal[REGISTRY_KEY] + const registry = registryGlobal[REGISTRY_KEY] + return registry } -function cloneCompletedTask(task: BackgroundTask): BackgroundTask { +function cloneProgress(progress: BackgroundTask["progress"]): BackgroundTask["progress"] { + if (!progress) { + return undefined + } + + return { + ...progress, + countedToolPartIDs: progress.countedToolPartIDs ? new Set(progress.countedToolPartIDs) : undefined, + } +} + +function cloneAttempts(attempts: BackgroundTask["attempts"]): BackgroundTask["attempts"] { + if (!attempts) { + return undefined + } + + return attempts.map((attempt) => ({ ...attempt })) +} + +function cloneRegisteredTask(task: BackgroundTask): BackgroundTask { return { id: task.id, + rootSessionId: task.rootSessionId, parentSessionId: task.parentSessionId, parentMessageId: task.parentMessageId, + teamRunId: task.teamRunId, description: task.description, prompt: "[redacted]", agent: task.agent, + spawnDepth: task.spawnDepth, sessionId: task.sessionId, status: task.status, queuedAt: task.queuedAt, startedAt: task.startedAt, completedAt: task.completedAt, + result: task.result, + progress: cloneProgress(task.progress), + parentModel: task.parentModel, model: task.model, + fallbackChain: task.fallbackChain, + attemptCount: task.attemptCount, + concurrencyKey: task.concurrencyKey, + concurrencyGroup: task.concurrencyGroup, + parentAgent: task.parentAgent, + parentTools: task.parentTools, + isUnstableAgent: task.isUnstableAgent, error: task.error, category: task.category, + retryNotification: task.retryNotification ? { ...task.retryNotification } : undefined, + attempts: cloneAttempts(task.attempts), + currentAttemptID: task.currentAttemptID, + lastMsgCount: task.lastMsgCount, + stablePolls: task.stablePolls, + consecutiveMissedPolls: task.consecutiveMissedPolls, } } @@ -60,7 +99,7 @@ function trimCompletedTasks(registry: BackgroundTaskRegistry): void { export function rememberBackgroundTask(task: BackgroundTask): void { const registry = getRegistry() registry.completedTasks.delete(task.id) - registry.activeTasks.set(task.id, task) + registry.activeTasks.set(task.id, () => cloneRegisteredTask(task)) } export function archiveBackgroundTask(task: BackgroundTask): void { @@ -70,13 +109,19 @@ export function archiveBackgroundTask(task: BackgroundTask): void { if (!task.sessionId || !TERMINAL_TASK_STATUSES.has(task.status)) { return } - registry.completedTasks.set(task.id, cloneCompletedTask(task)) + registry.completedTasks.set(task.id, cloneRegisteredTask(task)) trimCompletedTasks(registry) } export function getRegisteredBackgroundTask(taskID: string): BackgroundTask | undefined { const registry = getRegistry() - return registry.activeTasks.get(taskID) ?? registry.completedTasks.get(taskID) + const activeTask = registry.activeTasks.get(taskID) + if (activeTask) { + return activeTask() + } + + const completedTask = registry.completedTasks.get(taskID) + return completedTask ? cloneRegisteredTask(completedTask) : undefined } export function forgetBackgroundTask(taskID: string): void {