diff --git a/src/index.ts b/src/index.ts index baf4f95903..1ef3f6ebb9 100644 --- a/src/index.ts +++ b/src/index.ts @@ -100,6 +100,7 @@ import { isOpenCodeVersionAtLeast, OPENCODE_NATIVE_AGENTS_INJECTION_VERSION, injectServerAuthIntoClient, + createDirectoryBoundClient, } from "./shared"; import { loadPluginConfig } from "./plugin-config"; import { createModelCacheState } from "./plugin-state"; @@ -110,10 +111,16 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { directory: ctx.directory, }); injectServerAuthIntoClient(ctx.client); + + const ctxWithDirectory = { + ...ctx, + client: createDirectoryBoundClient(ctx.client, ctx.directory), + }; + // Start background tmux check immediately startTmuxCheck(); - const pluginConfig = loadPluginConfig(ctx.directory, ctx); + const pluginConfig = loadPluginConfig(ctxWithDirectory.directory, ctxWithDirectory); const disabledHooks = new Set(pluginConfig.disabled_hooks ?? []); const firstMessageVariantGate = createFirstMessageVariantGate(); @@ -129,15 +136,15 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const modelCacheState = createModelCacheState(); const contextWindowMonitor = isHookEnabled("context-window-monitor") - ? createContextWindowMonitorHook(ctx) + ? createContextWindowMonitorHook(ctxWithDirectory) : null; const preemptiveCompaction = isHookEnabled("preemptive-compaction") && pluginConfig.experimental?.preemptive_compaction - ? createPreemptiveCompactionHook(ctx) + ? createPreemptiveCompactionHook(ctxWithDirectory) : null; const sessionRecovery = isHookEnabled("session-recovery") - ? createSessionRecoveryHook(ctx, { + ? createSessionRecoveryHook(ctxWithDirectory, { experimental: pluginConfig.experimental, }) : null; @@ -156,7 +163,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { allPlugins: externalNotifier.allPlugins, }); } else { - sessionNotification = createSessionNotification(ctx); + sessionNotification = createSessionNotification(ctxWithDirectory); } } @@ -164,7 +171,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { ? createCommentCheckerHooks(pluginConfig.comment_checker) : null; const toolOutputTruncator = isHookEnabled("tool-output-truncator") - ? createToolOutputTruncatorHook(ctx, { + ? createToolOutputTruncatorHook(ctxWithDirectory, { experimental: pluginConfig.experimental, }) : null; @@ -185,20 +192,20 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { }, ); } else { - directoryAgentsInjector = createDirectoryAgentsInjectorHook(ctx); + directoryAgentsInjector = createDirectoryAgentsInjectorHook(ctxWithDirectory); } } const directoryReadmeInjector = isHookEnabled("directory-readme-injector") - ? createDirectoryReadmeInjectorHook(ctx) + ? createDirectoryReadmeInjectorHook(ctxWithDirectory) : null; const emptyTaskResponseDetector = isHookEnabled( "empty-task-response-detector", ) - ? createEmptyTaskResponseDetectorHook(ctx) + ? createEmptyTaskResponseDetectorHook(ctxWithDirectory) : null; const thinkMode = isHookEnabled("think-mode") ? createThinkModeHook() : null; const claudeCodeHooks = createClaudeCodeHooksHook( - ctx, + ctxWithDirectory, { disabledHooks: (pluginConfig.claude_code?.hooks ?? true) ? undefined : true, @@ -209,33 +216,33 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const anthropicContextWindowLimitRecovery = isHookEnabled( "anthropic-context-window-limit-recovery", ) - ? createAnthropicContextWindowLimitRecoveryHook(ctx, { + ? createAnthropicContextWindowLimitRecoveryHook(ctxWithDirectory, { experimental: pluginConfig.experimental, }) : null; const rulesInjector = isHookEnabled("rules-injector") - ? createRulesInjectorHook(ctx) + ? createRulesInjectorHook(ctxWithDirectory) : null; const autoUpdateChecker = isHookEnabled("auto-update-checker") - ? createAutoUpdateCheckerHook(ctx, { + ? createAutoUpdateCheckerHook(ctxWithDirectory, { showStartupToast: isHookEnabled("startup-toast"), isSisyphusEnabled: pluginConfig.sisyphus_agent?.disabled !== true, autoUpdate: pluginConfig.auto_update ?? true, }) : null; const keywordDetector = isHookEnabled("keyword-detector") - ? createKeywordDetectorHook(ctx, contextCollector) + ? createKeywordDetectorHook(ctxWithDirectory, contextCollector) : null; const contextInjectorMessagesTransform = createContextInjectorMessagesTransformHook(contextCollector); const agentUsageReminder = isHookEnabled("agent-usage-reminder") - ? createAgentUsageReminderHook(ctx) + ? createAgentUsageReminderHook(ctxWithDirectory) : null; const nonInteractiveEnv = isHookEnabled("non-interactive-env") - ? createNonInteractiveEnvHook(ctx) + ? createNonInteractiveEnvHook(ctxWithDirectory) : null; const interactiveBashSession = isHookEnabled("interactive-bash-session") - ? createInteractiveBashSessionHook(ctx) + ? createInteractiveBashSessionHook(ctxWithDirectory) : null; const thinkingBlockValidator = isHookEnabled("thinking-block-validator") @@ -243,34 +250,34 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { : null; const categorySkillReminder = isHookEnabled("category-skill-reminder") - ? createCategorySkillReminderHook(ctx) + ? createCategorySkillReminderHook(ctxWithDirectory) : null; const ralphLoop = isHookEnabled("ralph-loop") - ? createRalphLoopHook(ctx, { + ? createRalphLoopHook(ctxWithDirectory, { config: pluginConfig.ralph_loop, checkSessionExists: async (sessionId) => sessionExists(sessionId), }) : null; const editErrorRecovery = isHookEnabled("edit-error-recovery") - ? createEditErrorRecoveryHook(ctx) + ? createEditErrorRecoveryHook(ctxWithDirectory) : null; const delegateTaskRetry = isHookEnabled("delegate-task-retry") - ? createDelegateTaskRetryHook(ctx) + ? createDelegateTaskRetryHook(ctxWithDirectory) : null; const startWork = isHookEnabled("start-work") - ? createStartWorkHook(ctx) + ? createStartWorkHook(ctxWithDirectory) : null; const prometheusMdOnly = isHookEnabled("prometheus-md-only") - ? createPrometheusMdOnlyHook(ctx) + ? createPrometheusMdOnlyHook(ctxWithDirectory) : null; const sisyphusJuniorNotepad = isHookEnabled("sisyphus-junior-notepad") - ? createSisyphusJuniorNotepadHook(ctx) + ? createSisyphusJuniorNotepadHook(ctxWithDirectory) : null; const tasksTodowriteDisabler = isHookEnabled("tasks-todowrite-disabler") @@ -282,15 +289,15 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const questionLabelTruncator = createQuestionLabelTruncatorHook(); const subagentQuestionBlocker = createSubagentQuestionBlockerHook(); const writeExistingFileGuard = isHookEnabled("write-existing-file-guard") - ? createWriteExistingFileGuardHook(ctx) + ? createWriteExistingFileGuardHook(ctxWithDirectory) : null; const taskResumeInfo = createTaskResumeInfoHook(); - const tmuxSessionManager = new TmuxSessionManager(ctx, tmuxConfig); + const tmuxSessionManager = new TmuxSessionManager(ctxWithDirectory, tmuxConfig); const backgroundManager = new BackgroundManager( - ctx, + ctxWithDirectory, pluginConfig.background_task, { tmuxConfig, @@ -321,13 +328,13 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { ); const atlasHook = isHookEnabled("atlas") - ? createAtlasHook(ctx, { directory: ctx.directory, backgroundManager }) + ? createAtlasHook(ctxWithDirectory, { directory: ctx.directory, backgroundManager }) : null; - initTaskToastManager(ctx.client); + initTaskToastManager(ctxWithDirectory.client); const stopContinuationGuard = isHookEnabled("stop-continuation-guard") - ? createStopContinuationGuardHook(ctx) + ? createStopContinuationGuardHook(ctxWithDirectory) : null; const compactionContextInjector = isHookEnabled("compaction-context-injector") @@ -335,7 +342,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { : null; const todoContinuationEnforcer = isHookEnabled("todo-continuation-enforcer") - ? createTodoContinuationEnforcer(ctx, { + ? createTodoContinuationEnforcer(ctxWithDirectory, { backgroundManager, isContinuationStopped: stopContinuationGuard?.isStopped, }) @@ -348,7 +355,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { client: { session: { messages: async (args) => { - const result = await ctx.client.session.messages(args); + const result = await ctxWithDirectory.client.session.messages(args); if (Array.isArray(result)) return result; if ( typeof result === "object" && @@ -361,7 +368,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { return []; }, prompt: async (args) => { - await ctx.client.session.prompt(args); + await ctxWithDirectory.client.session.prompt(args); }, }, }, @@ -383,19 +390,22 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const backgroundNotificationHook = isHookEnabled("background-notification") ? createBackgroundNotificationHook(backgroundManager) : null; - const backgroundTools = createBackgroundTools(backgroundManager, ctx.client); + const backgroundTools = createBackgroundTools( + backgroundManager, + ctxWithDirectory.client, + ); - const callOmoAgent = createCallOmoAgent(ctx, backgroundManager); + const callOmoAgent = createCallOmoAgent(ctxWithDirectory, backgroundManager); const isMultimodalLookerEnabled = !(pluginConfig.disabled_agents ?? []).some( (agent) => agent.toLowerCase() === "multimodal-looker", ); - const lookAt = isMultimodalLookerEnabled ? createLookAt(ctx) : null; + const lookAt = isMultimodalLookerEnabled ? createLookAt(ctxWithDirectory) : null; const browserProvider = pluginConfig.browser_automation_engine?.provider ?? "playwright"; const disabledSkills = new Set(pluginConfig.disabled_skills ?? []); const delegateTask = createDelegateTask({ manager: backgroundManager, - client: ctx.client, + client: ctxWithDirectory.client, directory: ctx.directory, userCategories: pluginConfig.categories, gitMasterConfig: pluginConfig.git_master, @@ -472,7 +482,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { : null; const configHandler = createConfigHandler({ - ctx: { directory: ctx.directory, client: ctx.client }, + ctx: { directory: ctxWithDirectory.directory, client: ctxWithDirectory.client }, pluginConfig, modelCacheState, }); @@ -480,10 +490,10 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const taskSystemEnabled = pluginConfig.experimental?.task_system ?? false; const taskToolsRecord: Record = taskSystemEnabled ? { - task_create: createTaskCreateTool(pluginConfig, ctx), + task_create: createTaskCreateTool(pluginConfig, ctxWithDirectory), task_get: createTaskGetTool(pluginConfig), task_list: createTaskList(pluginConfig), - task_update: createTaskUpdateTool(pluginConfig, ctx), + task_update: createTaskUpdateTool(pluginConfig, ctxWithDirectory), } : {}; @@ -538,7 +548,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { await startWork?.["chat.message"]?.(input, output); if (!hasConnectedProvidersCache()) { - ctx.client.tui + ctxWithDirectory.client.tui .showToast({ body: { title: "⚠️ Provider Cache Missing", @@ -710,11 +720,11 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { sessionID === getMainSessionID() && !stopContinuationGuard?.isStopped(sessionID) ) { - await ctx.client.session + await ctxWithDirectory.client.session .prompt({ path: { id: sessionID }, body: { parts: [{ type: "text", text: "continue" }] }, - query: { directory: ctx.directory }, + query: { directory: ctxWithDirectory.directory }, }) .catch(() => {}); } diff --git a/src/shared/index.ts b/src/shared/index.ts index 93163dcc08..703257b06a 100644 --- a/src/shared/index.ts +++ b/src/shared/index.ts @@ -41,3 +41,4 @@ export * from "./tmux" export * from "./model-suggestion-retry" export * from "./opencode-server-auth" export * from "./port-utils" +export * from "./session-directory-injection" diff --git a/src/shared/session-directory-injection.test.ts b/src/shared/session-directory-injection.test.ts new file mode 100644 index 0000000000..cf3946c572 --- /dev/null +++ b/src/shared/session-directory-injection.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it } from "bun:test" +import { + createDirectoryBoundClient, + withDirectoryArgs, + wrapSessionWithDirectory, +} from "./session-directory-injection" + +describe("withDirectoryArgs", () => { + it("adds query.directory when args is undefined", () => { + // given + const directory = "/repo/worktree" + + // when + const next = withDirectoryArgs(undefined, directory) + + // then + expect(next).toEqual({ query: { directory } }) + }) + + it("merges query and preserves existing keys", () => { + // given + const directory = "/repo/worktree" + + // when + const next = withDirectoryArgs( + { query: { limit: 10 } }, + directory, + ) + + // then + expect(next).toEqual({ query: { limit: 10, directory } }) + }) + + it("preserves caller-provided query.directory", () => { + // given + const input = { query: { directory: "/explicit", limit: 5 } } + + // when + const next = withDirectoryArgs(input, "/ignored") + + // then + expect(next).toBe(input) + }) + + it("overwrites malformed query with a directory-only object", () => { + // given + const directory = "/repo/worktree" + const input = { query: "not-an-object" as unknown as Record } + + // when + const next = withDirectoryArgs(input, directory) + + // then + expect(next).toEqual({ query: { directory } }) + }) +}) + +describe("wrapSessionWithDirectory", () => { + it("injects query.directory only for allowlisted methods", () => { + // given + const calls: Record = {} + const session = { + get: (args?: unknown) => { + calls.get = [args] + return args + }, + status: (args?: unknown) => { + calls.status = [args] + return args + }, + } + + const wrapped = wrapSessionWithDirectory(session, "/repo/worktree", new Set(["get"])) + + // when + wrapped.get({ query: { limit: 1 } }) + wrapped.status({ query: { limit: 1 } }) + + // then + expect(calls.get[0]).toEqual({ query: { limit: 1, directory: "/repo/worktree" } }) + expect(calls.status[0]).toEqual({ query: { limit: 1 } }) + }) + + it("does not override caller-provided query.directory", () => { + // given + let seen: unknown + const session = { + get: (args?: unknown) => { + seen = args + return args + }, + } + + const wrapped = wrapSessionWithDirectory(session, "/repo/worktree", new Set(["get"])) + + // when + wrapped.get({ query: { directory: "/explicit" } }) + + // then + expect(seen).toEqual({ query: { directory: "/explicit" } }) + }) +}) + +describe("createDirectoryBoundClient", () => { + it("returns the original client when session is missing", () => { + // given + const client: { foo: number; session?: unknown } = { foo: 1 } + + // when + const next = createDirectoryBoundClient(client, "/repo/worktree") + + // then + expect(next).toBe(client) + }) + + it("does not mutate the original client and wraps session methods", () => { + // given + let seen: unknown + const client = { + session: { + get: (args?: unknown) => { + seen = args + return args + }, + }, + other: 123, + } + + // when + const next = createDirectoryBoundClient(client, "/repo/worktree") + + // then + expect(next).not.toBe(client) + expect(next.other).toBe(123) + expect(next.session).not.toBe(client.session) + + next.session.get({ query: { limit: 3 } }) + expect(seen).toEqual({ query: { limit: 3, directory: "/repo/worktree" } }) + }) +}) diff --git a/src/shared/session-directory-injection.ts b/src/shared/session-directory-injection.ts new file mode 100644 index 0000000000..a52e2bed95 --- /dev/null +++ b/src/shared/session-directory-injection.ts @@ -0,0 +1,89 @@ +export type DirectoryQuery = { + query?: { directory?: string } & Record +} + +export function withDirectoryArgs( + args: DirectoryQuery | undefined, + directory: string, +): DirectoryQuery { + if (!args) return { query: { directory } } + + const query = + typeof args.query === "object" && args.query ? args.query : undefined + + // Preserve caller-provided query.directory (even if it differs). + if (query && "directory" in query) return args + + return { ...args, query: { ...query, directory } } +} + +const DEFAULT_SESSION_METHODS_NEEDING_DIRECTORY = new Set([ + "get", + "messages", + "prompt", + "summarize", + "create", + "abort", + "status", + "todo", + "children", + "promptAsync", + "message", + "diff", + "fork", + "init", + "command", + "shell", + "revert", + "unrevert", + "share", + "unshare", + "update", + "delete", + "list", +]) + +export function wrapSessionWithDirectory( + session: T, + directory: string, + methodsNeedingDirectory: ReadonlySet = + DEFAULT_SESSION_METHODS_NEEDING_DIRECTORY, +): T { + return new Proxy(session, { + get(target, prop) { + const value = target[prop as keyof T] + if (typeof value !== "function") return value + + const key = typeof prop === "string" ? prop : undefined + if (!key || !methodsNeedingDirectory.has(key)) { + return value.bind(target) + } + + return (args?: DirectoryQuery) => { + const next = withDirectoryArgs(args, directory) + const fn = value as (input: DirectoryQuery) => unknown + return fn.call(target, next) + } + }, + }) +} + +function cloneWithPrototype(input: T): T { + const clone = Object.create(Object.getPrototypeOf(input)) as T + Object.assign(clone, input) + return clone +} + +export function createDirectoryBoundClient( + client: T, + directory: string, +): T { + if (!client.session || typeof client.session !== "object") return client + + const session = client.session as object + const wrappedSession = wrapSessionWithDirectory(session, directory) + + const next = cloneWithPrototype(client) + ;(next as { session: unknown }).session = wrappedSession + return next +} diff --git a/tsconfig.test.json b/tsconfig.test.json new file mode 100644 index 0000000000..253ec471b9 --- /dev/null +++ b/tsconfig.test.json @@ -0,0 +1,8 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "noEmit": true + }, + "include": ["src/**/*.test.ts"], + "exclude": ["node_modules", "dist", "script"] +}