diff --git a/src/features/background-agent/manager.ts b/src/features/background-agent/manager.ts index b2d4a40a..3a680dd6 100644 --- a/src/features/background-agent/manager.ts +++ b/src/features/background-agent/manager.ts @@ -7,6 +7,7 @@ import type { } from "./types" import { TaskHistory } from "./task-history" import { log, getAgentToolRestrictions, promptWithModelSuggestionRetry } from "../../shared" +import { setSessionTools } from "../../shared/session-tools-store" import { ConcurrencyManager } from "./concurrency" import type { BackgroundTaskConfig, TmuxConfig } from "../../config/schema" import { isInsideTmux } from "../../shared/tmux" @@ -141,6 +142,7 @@ export class BackgroundManager { parentMessageID: input.parentMessageID, parentModel: input.parentModel, parentAgent: input.parentAgent, + parentTools: input.parentTools, model: input.model, category: input.category, } @@ -328,12 +330,16 @@ export class BackgroundManager { ...(launchModel ? { model: launchModel } : {}), ...(launchVariant ? { variant: launchVariant } : {}), system: input.skillContent, - tools: { - ...getAgentToolRestrictions(input.agent), - task: false, - call_omo_agent: true, - question: false, - }, + tools: (() => { + const tools = { + ...getAgentToolRestrictions(input.agent), + task: false, + call_omo_agent: true, + question: false, + } + setSessionTools(sessionID, tools) + return tools + })(), parts: [{ type: "text", text: input.prompt }], }, }).catch((error) => { @@ -535,6 +541,9 @@ export class BackgroundManager { existingTask.parentMessageID = input.parentMessageID existingTask.parentModel = input.parentModel existingTask.parentAgent = input.parentAgent + if (input.parentTools) { + existingTask.parentTools = input.parentTools + } // Reset startedAt on resume to prevent immediate completion // The MIN_IDLE_TIME_MS check uses startedAt, so resumed tasks need fresh timing existingTask.startedAt = new Date() @@ -588,12 +597,16 @@ export class BackgroundManager { agent: existingTask.agent, ...(resumeModel ? { model: resumeModel } : {}), ...(resumeVariant ? { variant: resumeVariant } : {}), - tools: { - ...getAgentToolRestrictions(existingTask.agent), - task: false, - call_omo_agent: true, - question: false, - }, + tools: (() => { + const tools = { + ...getAgentToolRestrictions(existingTask.agent), + task: false, + call_omo_agent: true, + question: false, + } + setSessionTools(existingTask.sessionID!, tools) + return tools + })(), parts: [{ type: "text", text: input.prompt }], }, }).catch((error) => { @@ -1252,6 +1265,7 @@ Use \`background_output(task_id="${task.id}")\` to retrieve this result when rea noReply: !allComplete, ...(agent !== undefined ? { agent } : {}), ...(model !== undefined ? { model } : {}), + ...(task.parentTools ? { tools: task.parentTools } : {}), parts: [{ type: "text", text: notification }], }, }) diff --git a/src/features/background-agent/notify-parent-session.ts b/src/features/background-agent/notify-parent-session.ts index ea28f025..da6a531e 100644 --- a/src/features/background-agent/notify-parent-session.ts +++ b/src/features/background-agent/notify-parent-session.ts @@ -148,6 +148,7 @@ export async function notifyParentSession(args: { noReply: !allComplete, ...(agent !== undefined ? { agent } : {}), ...(model !== undefined ? { model } : {}), + ...(task.parentTools ? { tools: task.parentTools } : {}), parts: [{ type: "text", text: notification }], }, }) diff --git a/src/features/background-agent/parent-session-notifier.ts b/src/features/background-agent/parent-session-notifier.ts index 28fb3a37..a0b228fc 100644 --- a/src/features/background-agent/parent-session-notifier.ts +++ b/src/features/background-agent/parent-session-notifier.ts @@ -71,6 +71,7 @@ export async function notifyParentSession( noReply: !allComplete, ...(agent !== undefined ? { agent } : {}), ...(model !== undefined ? { model } : {}), + ...(task.parentTools ? { tools: task.parentTools } : {}), parts: [{ type: "text", text: notification }], }, }) diff --git a/src/features/background-agent/spawner/task-factory.ts b/src/features/background-agent/spawner/task-factory.ts index 7af433d4..e7788be7 100644 --- a/src/features/background-agent/spawner/task-factory.ts +++ b/src/features/background-agent/spawner/task-factory.ts @@ -13,6 +13,7 @@ export function createTask(input: LaunchInput): BackgroundTask { parentMessageID: input.parentMessageID, parentModel: input.parentModel, parentAgent: input.parentAgent, + parentTools: input.parentTools, model: input.model, } } diff --git a/src/features/background-agent/spawner/task-resumer.ts b/src/features/background-agent/spawner/task-resumer.ts index 4a517a65..7c7d5d2a 100644 --- a/src/features/background-agent/spawner/task-resumer.ts +++ b/src/features/background-agent/spawner/task-resumer.ts @@ -1,5 +1,6 @@ import type { BackgroundTask, ResumeInput } from "../types" import { log, getAgentToolRestrictions } from "../../../shared" +import { setSessionTools } from "../../../shared/session-tools-store" import type { SpawnerContext } from "./spawner-context" import { subagentSessions } from "../../claude-code-session-state" import { getTaskToastManager } from "../../task-toast-manager" @@ -35,6 +36,9 @@ export async function resumeTask( task.parentMessageID = input.parentMessageID task.parentModel = input.parentModel task.parentAgent = input.parentAgent + if (input.parentTools) { + task.parentTools = input.parentTools + } task.startedAt = new Date() task.progress = { @@ -75,12 +79,16 @@ export async function resumeTask( agent: task.agent, ...(resumeModel ? { model: resumeModel } : {}), ...(resumeVariant ? { variant: resumeVariant } : {}), - tools: { - ...getAgentToolRestrictions(task.agent), - task: false, - call_omo_agent: true, - question: false, - }, + tools: (() => { + const tools = { + ...getAgentToolRestrictions(task.agent), + task: false, + call_omo_agent: true, + question: false, + } + setSessionTools(task.sessionID!, tools) + return tools + })(), parts: [{ type: "text", text: input.prompt }], }, }) diff --git a/src/features/background-agent/spawner/task-starter.ts b/src/features/background-agent/spawner/task-starter.ts index 4dfb48d1..a904a20b 100644 --- a/src/features/background-agent/spawner/task-starter.ts +++ b/src/features/background-agent/spawner/task-starter.ts @@ -1,5 +1,6 @@ import type { QueueItem } from "../constants" import { log, getAgentToolRestrictions, promptWithModelSuggestionRetry } from "../../../shared" +import { setSessionTools } from "../../../shared/session-tools-store" import { subagentSessions } from "../../claude-code-session-state" import { getTaskToastManager } from "../../task-toast-manager" import { createBackgroundSession } from "./background-session-creator" @@ -79,12 +80,16 @@ export async function startTask(item: QueueItem, ctx: SpawnerContext): Promise { + const tools = { + ...getAgentToolRestrictions(input.agent), + task: false, + call_omo_agent: true, + question: false, + } + setSessionTools(sessionID, tools) + return tools + })(), parts: [{ type: "text", text: input.prompt }], }, }).catch((error: unknown) => { diff --git a/src/features/background-agent/types.ts b/src/features/background-agent/types.ts index fd5309dc..ea56186e 100644 --- a/src/features/background-agent/types.ts +++ b/src/features/background-agent/types.ts @@ -37,6 +37,8 @@ export interface BackgroundTask { concurrencyGroup?: string /** Parent session's agent name for notification */ parentAgent?: string + /** Parent session's tool restrictions for notification prompts */ + parentTools?: Record /** Marks if the task was launched from an unstable agent/category */ isUnstableAgent?: boolean /** Category used for this task (e.g., 'quick', 'visual-engineering') */ @@ -56,6 +58,7 @@ export interface LaunchInput { parentMessageID: string parentModel?: { providerID: string; modelID: string } parentAgent?: string + parentTools?: Record model?: { providerID: string; modelID: string; variant?: string } isUnstableAgent?: boolean skills?: string[] @@ -70,4 +73,5 @@ export interface ResumeInput { parentMessageID: string parentModel?: { providerID: string; modelID: string } parentAgent?: string + parentTools?: Record } diff --git a/src/shared/session-tools-store.test.ts b/src/shared/session-tools-store.test.ts new file mode 100644 index 00000000..1e6f8c1b --- /dev/null +++ b/src/shared/session-tools-store.test.ts @@ -0,0 +1,72 @@ +import { describe, test, expect, beforeEach } from "bun:test" +import { setSessionTools, getSessionTools, clearSessionTools } from "./session-tools-store" + +describe("session-tools-store", () => { + beforeEach(() => { + clearSessionTools() + }) + + test("returns undefined for unknown session", () => { + //#given + const sessionID = "ses_unknown" + + //#when + const result = getSessionTools(sessionID) + + //#then + expect(result).toBeUndefined() + }) + + test("stores and retrieves tools for a session", () => { + //#given + const sessionID = "ses_abc123" + const tools = { question: false, task: true, call_omo_agent: true } + + //#when + setSessionTools(sessionID, tools) + const result = getSessionTools(sessionID) + + //#then + expect(result).toEqual({ question: false, task: true, call_omo_agent: true }) + }) + + test("overwrites existing tools for same session", () => { + //#given + const sessionID = "ses_abc123" + setSessionTools(sessionID, { question: false }) + + //#when + setSessionTools(sessionID, { question: true, task: false }) + const result = getSessionTools(sessionID) + + //#then + expect(result).toEqual({ question: true, task: false }) + }) + + test("clearSessionTools removes all entries", () => { + //#given + setSessionTools("ses_1", { question: false }) + setSessionTools("ses_2", { task: true }) + + //#when + clearSessionTools() + + //#then + expect(getSessionTools("ses_1")).toBeUndefined() + expect(getSessionTools("ses_2")).toBeUndefined() + }) + + test("returns a copy, not a reference", () => { + //#given + const sessionID = "ses_abc123" + const tools = { question: false } + setSessionTools(sessionID, tools) + + //#when + const result = getSessionTools(sessionID)! + result.question = true + + //#then + expect(getSessionTools(sessionID)).toEqual({ question: false }) + }) +}) diff --git a/src/shared/session-tools-store.ts b/src/shared/session-tools-store.ts new file mode 100644 index 00000000..f717488e --- /dev/null +++ b/src/shared/session-tools-store.ts @@ -0,0 +1,14 @@ +const store = new Map>() + +export function setSessionTools(sessionID: string, tools: Record): void { + store.set(sessionID, { ...tools }) +} + +export function getSessionTools(sessionID: string): Record | undefined { + const tools = store.get(sessionID) + return tools ? { ...tools } : undefined +} + +export function clearSessionTools(): void { + store.clear() +} diff --git a/src/tools/call-omo-agent/background-agent-executor.ts b/src/tools/call-omo-agent/background-agent-executor.ts index 35cb5fe1..7babb43c 100644 --- a/src/tools/call-omo-agent/background-agent-executor.ts +++ b/src/tools/call-omo-agent/background-agent-executor.ts @@ -5,6 +5,7 @@ import { log } from "../../shared" import type { CallOmoAgentArgs } from "./types" import type { ToolContextWithMetadata } from "./tool-context-with-metadata" import { getMessageDir } from "./message-storage-directory" +import { getSessionTools } from "../../shared/session-tools-store" export async function executeBackgroundAgent( args: CallOmoAgentArgs, @@ -36,6 +37,7 @@ export async function executeBackgroundAgent( parentSessionID: toolContext.sessionID, parentMessageID: toolContext.messageID, parentAgent, + parentTools: getSessionTools(toolContext.sessionID), }) const waitStart = Date.now() diff --git a/src/tools/call-omo-agent/background-executor.ts b/src/tools/call-omo-agent/background-executor.ts index a89fff40..5751664a 100644 --- a/src/tools/call-omo-agent/background-executor.ts +++ b/src/tools/call-omo-agent/background-executor.ts @@ -5,6 +5,7 @@ import { consumeNewMessages } from "../../shared/session-cursor" import { findFirstMessageWithAgent, findNearestMessageWithFields } from "../../features/hook-message-injector" import { getSessionAgent } from "../../features/claude-code-session-state" import { getMessageDir } from "./message-dir" +import { getSessionTools } from "../../shared/session-tools-store" export async function executeBackground( args: CallOmoAgentArgs, @@ -41,6 +42,7 @@ export async function executeBackground( parentSessionID: toolContext.sessionID, parentMessageID: toolContext.messageID, parentAgent, + parentTools: getSessionTools(toolContext.sessionID), }) const WAIT_FOR_SESSION_INTERVAL_MS = 50 diff --git a/src/tools/delegate-task/background-continuation.ts b/src/tools/delegate-task/background-continuation.ts index 71e20970..a6e38293 100644 --- a/src/tools/delegate-task/background-continuation.ts +++ b/src/tools/delegate-task/background-continuation.ts @@ -2,6 +2,7 @@ import type { DelegateTaskArgs, ToolContextWithMetadata } from "./types" import type { ExecutorContext, ParentContext } from "./executor-types" import { storeToolMetadata } from "../../features/tool-metadata-store" import { formatDetailedError } from "./error-formatting" +import { getSessionTools } from "../../shared/session-tools-store" export async function executeBackgroundContinuation( args: DelegateTaskArgs, @@ -19,6 +20,7 @@ export async function executeBackgroundContinuation( parentMessageID: parentContext.messageID, parentModel: parentContext.model, parentAgent: parentContext.agent, + parentTools: getSessionTools(parentContext.sessionID), }) const bgContMeta = { diff --git a/src/tools/delegate-task/background-task.ts b/src/tools/delegate-task/background-task.ts index 35d9af9a..e724695b 100644 --- a/src/tools/delegate-task/background-task.ts +++ b/src/tools/delegate-task/background-task.ts @@ -3,6 +3,7 @@ import type { ExecutorContext, ParentContext } from "./executor-types" import { getTimingConfig } from "./timing" import { storeToolMetadata } from "../../features/tool-metadata-store" import { formatDetailedError } from "./error-formatting" +import { getSessionTools } from "../../shared/session-tools-store" export async function executeBackgroundTask( args: DelegateTaskArgs, @@ -24,6 +25,7 @@ export async function executeBackgroundTask( parentMessageID: parentContext.messageID, parentModel: parentContext.model, parentAgent: parentContext.agent, + parentTools: getSessionTools(parentContext.sessionID), model: categoryModel, skills: args.load_skills.length > 0 ? args.load_skills : undefined, skillContent: systemContent, diff --git a/src/tools/delegate-task/sync-continuation.ts b/src/tools/delegate-task/sync-continuation.ts index 95f6baec..72355982 100644 --- a/src/tools/delegate-task/sync-continuation.ts +++ b/src/tools/delegate-task/sync-continuation.ts @@ -9,6 +9,7 @@ import { promptWithModelSuggestionRetry } from "../../shared/model-suggestion-re import { findNearestMessageWithFields } from "../../features/hook-message-injector" import { formatDuration } from "./time-formatter" import { syncContinuationDeps, type SyncContinuationDeps } from "./sync-continuation-deps" +import { setSessionTools } from "../../shared/session-tools-store" export async function executeSyncContinuation( args: DelegateTaskArgs, @@ -77,6 +78,13 @@ export async function executeSyncContinuation( } const allowTask = isPlanFamily(resumeAgent) + const tools = { + ...(resumeAgent ? getAgentToolRestrictions(resumeAgent) : {}), + task: allowTask, + call_omo_agent: true, + question: false, + } + setSessionTools(args.session_id!, tools) await promptWithModelSuggestionRetry(client, { path: { id: args.session_id! }, @@ -84,12 +92,7 @@ export async function executeSyncContinuation( ...(resumeAgent !== undefined ? { agent: resumeAgent } : {}), ...(resumeModel !== undefined ? { model: resumeModel } : {}), ...(resumeVariant !== undefined ? { variant: resumeVariant } : {}), - tools: { - ...(resumeAgent ? getAgentToolRestrictions(resumeAgent) : {}), - task: allowTask, - call_omo_agent: true, // Intentionally overrides restrictions - continuation context needs delegation capability even for restricted agents - question: false, - }, + tools, parts: [{ type: "text", text: args.prompt }], }, }) diff --git a/src/tools/delegate-task/sync-prompt-sender.ts b/src/tools/delegate-task/sync-prompt-sender.ts index 815fe380..34659b81 100644 --- a/src/tools/delegate-task/sync-prompt-sender.ts +++ b/src/tools/delegate-task/sync-prompt-sender.ts @@ -3,6 +3,7 @@ import { isPlanFamily } from "./constants" import { promptWithModelSuggestionRetry } from "../../shared/model-suggestion-retry" import { formatDetailedError } from "./error-formatting" import { getAgentToolRestrictions } from "../../shared/agent-tool-restrictions" +import { setSessionTools } from "../../shared/session-tools-store" export async function sendSyncPrompt( client: OpencodeClient, @@ -18,17 +19,19 @@ export async function sendSyncPrompt( ): Promise { try { const allowTask = isPlanFamily(input.agentToUse) + const tools = { + task: allowTask, + call_omo_agent: true, + question: false, + ...getAgentToolRestrictions(input.agentToUse), + } + setSessionTools(input.sessionID, tools) await promptWithModelSuggestionRetry(client, { path: { id: input.sessionID }, body: { agent: input.agentToUse, system: input.systemContent, - tools: { - task: allowTask, - call_omo_agent: true, - question: false, - ...getAgentToolRestrictions(input.agentToUse), - }, + tools, parts: [{ type: "text", text: input.args.prompt }], ...(input.categoryModel ? { model: { providerID: input.categoryModel.providerID, modelID: input.categoryModel.modelID } } : {}), ...(input.categoryModel?.variant ? { variant: input.categoryModel.variant } : {}), diff --git a/src/tools/delegate-task/unstable-agent-task.ts b/src/tools/delegate-task/unstable-agent-task.ts index ae97153b..9e0bf853 100644 --- a/src/tools/delegate-task/unstable-agent-task.ts +++ b/src/tools/delegate-task/unstable-agent-task.ts @@ -4,6 +4,7 @@ import { getTimingConfig } from "./timing" import { storeToolMetadata } from "../../features/tool-metadata-store" import { formatDuration } from "./time-formatter" import { formatDetailedError } from "./error-formatting" +import { getSessionTools } from "../../shared/session-tools-store" export async function executeUnstableAgentTask( args: DelegateTaskArgs, @@ -26,6 +27,7 @@ export async function executeUnstableAgentTask( parentMessageID: parentContext.messageID, parentModel: parentContext.model, parentAgent: parentContext.agent, + parentTools: getSessionTools(parentContext.sessionID), model: categoryModel, skills: args.load_skills.length > 0 ? args.load_skills : undefined, skillContent: systemContent,