From f20e1aa0d0d8c54ddb68a346b9f88715f56dce84 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sat, 14 Feb 2026 14:30:22 +0900 Subject: [PATCH] feat: store tool restrictions in session-tools-store at prompt-send sites Call setSessionTools(sessionID, tools) before every prompt dispatch so the tools object is captured and available for later retrieval when background tasks complete. --- .../background-agent/spawner/task-resumer.ts | 20 +++++++++++++------ .../background-agent/spawner/task-starter.ts | 17 ++++++++++------ src/tools/delegate-task/sync-continuation.ts | 15 ++++++++------ src/tools/delegate-task/sync-prompt-sender.ts | 15 ++++++++------ 4 files changed, 43 insertions(+), 24 deletions(-) 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/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 } : {}),