diff --git a/src/features/claude-code-session-state/state.ts b/src/features/claude-code-session-state/state.ts index 751ab83e..a864b75d 100644 --- a/src/features/claude-code-session-state/state.ts +++ b/src/features/claude-code-session-state/state.ts @@ -9,3 +9,23 @@ export function setMainSession(id: string | undefined) { export function getMainSessionID(): string | undefined { return mainSessionID } + +const sessionAgentMap = new Map() + +export function setSessionAgent(sessionID: string, agent: string): void { + if (!sessionAgentMap.has(sessionID)) { + sessionAgentMap.set(sessionID, agent) + } +} + +export function updateSessionAgent(sessionID: string, agent: string): void { + sessionAgentMap.set(sessionID, agent) +} + +export function getSessionAgent(sessionID: string): string | undefined { + return sessionAgentMap.get(sessionID) +} + +export function clearSessionAgent(sessionID: string): void { + sessionAgentMap.delete(sessionID) +} diff --git a/src/features/hook-message-injector/index.ts b/src/features/hook-message-injector/index.ts index 2262a0b3..fcb0624d 100644 --- a/src/features/hook-message-injector/index.ts +++ b/src/features/hook-message-injector/index.ts @@ -1,4 +1,4 @@ -export { injectHookMessage, findNearestMessageWithFields } from "./injector" +export { injectHookMessage, findNearestMessageWithFields, findFirstMessageWithAgent } from "./injector" export type { StoredMessage } from "./injector" export type { MessageMeta, OriginalMessageContext, TextPart } from "./types" export { MESSAGE_STORAGE } from "./constants" diff --git a/src/features/hook-message-injector/injector.ts b/src/features/hook-message-injector/injector.ts index acc2c46a..f5d87018 100644 --- a/src/features/hook-message-injector/injector.ts +++ b/src/features/hook-message-injector/injector.ts @@ -48,6 +48,35 @@ export function findNearestMessageWithFields(messageDir: string): StoredMessage return null } +/** + * Finds the FIRST (oldest) message in the session with agent field. + * This is used to get the original agent that started the session, + * avoiding issues where newer messages may have a different agent + * due to OpenCode's internal agent switching. + */ +export function findFirstMessageWithAgent(messageDir: string): string | null { + try { + const files = readdirSync(messageDir) + .filter((f) => f.endsWith(".json")) + .sort() // Oldest first (no reverse) + + for (const file of files) { + try { + const content = readFileSync(join(messageDir, file), "utf-8") + const msg = JSON.parse(content) as StoredMessage + if (msg.agent) { + return msg.agent + } + } catch { + continue + } + } + } catch { + return null + } + return null +} + function generateMessageId(): string { const timestamp = Date.now().toString(16) const random = Math.random().toString(36).substring(2, 14) diff --git a/src/hooks/prometheus-md-only/index.ts b/src/hooks/prometheus-md-only/index.ts index d5839e81..c562e39e 100644 --- a/src/hooks/prometheus-md-only/index.ts +++ b/src/hooks/prometheus-md-only/index.ts @@ -2,7 +2,8 @@ import type { PluginInput } from "@opencode-ai/plugin" import { existsSync, readdirSync } from "node:fs" import { join, resolve, relative, isAbsolute } from "node:path" import { HOOK_NAME, PROMETHEUS_AGENTS, ALLOWED_EXTENSIONS, ALLOWED_PATH_PREFIX, BLOCKED_TOOLS, PLANNING_CONSULT_WARNING } from "./constants" -import { findNearestMessageWithFields, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { findNearestMessageWithFields, findFirstMessageWithAgent, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { getSessionAgent } from "../../features/claude-code-session-state" import { log } from "../../shared/logger" export * from "./constants" @@ -61,10 +62,14 @@ function getMessageDir(sessionID: string): string | null { const TASK_TOOLS = ["sisyphus_task", "task", "call_omo_agent"] -function getAgentFromSession(sessionID: string): string | undefined { +function getAgentFromMessageFiles(sessionID: string): string | undefined { const messageDir = getMessageDir(sessionID) if (!messageDir) return undefined - return findNearestMessageWithFields(messageDir)?.agent + return findFirstMessageWithAgent(messageDir) ?? findNearestMessageWithFields(messageDir)?.agent +} + +function getAgentFromSession(sessionID: string): string | undefined { + return getSessionAgent(sessionID) ?? getAgentFromMessageFiles(sessionID) } export function createPrometheusMdOnlyHook(ctx: PluginInput) { diff --git a/src/index.ts b/src/index.ts index 22e75cba..995c5ea4 100644 --- a/src/index.ts +++ b/src/index.ts @@ -49,6 +49,8 @@ import { getSystemMcpServerNames } from "./features/claude-code-mcp-loader"; import { setMainSession, getMainSessionID, + setSessionAgent, + clearSessionAgent, } from "./features/claude-code-session-state"; import { builtinTools, @@ -428,11 +430,22 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { setMainSession(undefined); } if (sessionInfo?.id) { + clearSessionAgent(sessionInfo.id); await skillMcpManager.disconnectSession(sessionInfo.id); await lspManager.cleanupTempDirectoryClients(); } } + if (event.type === "message.updated") { + const info = props?.info as Record | undefined; + const sessionID = info?.sessionID as string | undefined; + const agent = info?.agent as string | undefined; + const role = info?.role as string | undefined; + if (sessionID && agent && role === "user") { + setSessionAgent(sessionID, agent); + } + } + if (event.type === "session.error") { const sessionID = props?.sessionID as string | undefined; const error = props?.error; diff --git a/src/tools/background-task/tools.ts b/src/tools/background-task/tools.ts index 3df7b053..3a2eeae5 100644 --- a/src/tools/background-task/tools.ts +++ b/src/tools/background-task/tools.ts @@ -4,7 +4,9 @@ import { join } from "node:path" import type { BackgroundManager, BackgroundTask } from "../../features/background-agent" import type { BackgroundTaskArgs, BackgroundOutputArgs, BackgroundCancelArgs } from "./types" import { BACKGROUND_TASK_DESCRIPTION, BACKGROUND_OUTPUT_DESCRIPTION, BACKGROUND_CANCEL_DESCRIPTION } from "./constants" -import { findNearestMessageWithFields, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { findNearestMessageWithFields, findFirstMessageWithAgent, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { getSessionAgent } from "../../features/claude-code-session-state" +import { log } from "../../shared/logger" type OpencodeClient = PluginInput["client"] @@ -63,6 +65,19 @@ export function createBackgroundTask(manager: BackgroundManager): ToolDefinition try { const messageDir = getMessageDir(ctx.sessionID) const prevMessage = messageDir ? findNearestMessageWithFields(messageDir) : null + const firstMessageAgent = messageDir ? findFirstMessageWithAgent(messageDir) : null + const sessionAgent = getSessionAgent(ctx.sessionID) + const parentAgent = ctx.agent ?? sessionAgent ?? firstMessageAgent ?? prevMessage?.agent + + log("[background_task] parentAgent resolution", { + sessionID: ctx.sessionID, + ctxAgent: ctx.agent, + sessionAgent, + firstMessageAgent, + prevMessageAgent: prevMessage?.agent, + resolvedParentAgent: parentAgent, + }) + const parentModel = prevMessage?.model?.providerID && prevMessage?.model?.modelID ? { providerID: prevMessage.model.providerID, modelID: prevMessage.model.modelID } : undefined @@ -74,7 +89,7 @@ export function createBackgroundTask(manager: BackgroundManager): ToolDefinition parentSessionID: ctx.sessionID, parentMessageID: ctx.messageID, parentModel, - parentAgent: ctx.agent ?? prevMessage?.agent, + parentAgent, }) ctx.metadata?.({ diff --git a/src/tools/call-omo-agent/tools.ts b/src/tools/call-omo-agent/tools.ts index b30e2286..c9ea3814 100644 --- a/src/tools/call-omo-agent/tools.ts +++ b/src/tools/call-omo-agent/tools.ts @@ -1,8 +1,26 @@ import { tool, type PluginInput, type ToolDefinition } from "@opencode-ai/plugin" +import { existsSync, readdirSync } from "node:fs" +import { join } from "node:path" import { ALLOWED_AGENTS, CALL_OMO_AGENT_DESCRIPTION } from "./constants" import type { CallOmoAgentArgs } from "./types" import type { BackgroundManager } from "../../features/background-agent" import { log } from "../../shared/logger" +import { findFirstMessageWithAgent, findNearestMessageWithFields, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { getSessionAgent } from "../../features/claude-code-session-state" + +function getMessageDir(sessionID: string): string | null { + if (!existsSync(MESSAGE_STORAGE)) return null + + const directPath = join(MESSAGE_STORAGE, sessionID) + if (existsSync(directPath)) return directPath + + for (const dir of readdirSync(MESSAGE_STORAGE)) { + const sessionPath = join(MESSAGE_STORAGE, dir, sessionID) + if (existsSync(sessionPath)) return sessionPath + } + + return null +} type ToolContextWithMetadata = { sessionID: string @@ -60,12 +78,29 @@ async function executeBackground( manager: BackgroundManager ): Promise { try { + const messageDir = getMessageDir(toolContext.sessionID) + const prevMessage = messageDir ? findNearestMessageWithFields(messageDir) : null + const firstMessageAgent = messageDir ? findFirstMessageWithAgent(messageDir) : null + const sessionAgent = getSessionAgent(toolContext.sessionID) + const parentAgent = toolContext.agent ?? sessionAgent ?? firstMessageAgent ?? prevMessage?.agent + + log("[call_omo_agent] parentAgent resolution", { + sessionID: toolContext.sessionID, + messageDir, + ctxAgent: toolContext.agent, + sessionAgent, + firstMessageAgent, + prevMessageAgent: prevMessage?.agent, + resolvedParentAgent: parentAgent, + }) + const task = await manager.launch({ description: args.description, prompt: args.prompt, agent: args.subagent_type, parentSessionID: toolContext.sessionID, parentMessageID: toolContext.messageID, + parentAgent, }) toolContext.metadata?.({ diff --git a/src/tools/sisyphus-task/tools.ts b/src/tools/sisyphus-task/tools.ts index ca4534fe..c082f8b1 100644 --- a/src/tools/sisyphus-task/tools.ts +++ b/src/tools/sisyphus-task/tools.ts @@ -5,11 +5,12 @@ import type { BackgroundManager } from "../../features/background-agent" import type { SisyphusTaskArgs } from "./types" import type { CategoryConfig, CategoriesConfig, GitMasterConfig } from "../../config/schema" import { SISYPHUS_TASK_DESCRIPTION, DEFAULT_CATEGORIES, CATEGORY_PROMPT_APPENDS } from "./constants" -import { findNearestMessageWithFields, MESSAGE_STORAGE } from "../../features/hook-message-injector" +import { findNearestMessageWithFields, findFirstMessageWithAgent, MESSAGE_STORAGE } from "../../features/hook-message-injector" import { resolveMultipleSkills } from "../../features/opencode-skill-loader/skill-content" import { createBuiltinSkills } from "../../features/builtin-skills/skills" import { getTaskToastManager } from "../../features/task-toast-manager" -import { subagentSessions } from "../../features/claude-code-session-state" +import { subagentSessions, getSessionAgent } from "../../features/claude-code-session-state" +import { log } from "../../shared/logger" type OpencodeClient = PluginInput["client"] @@ -147,7 +148,19 @@ export function createSisyphusTask(options: SisyphusTaskToolOptions): ToolDefini const messageDir = getMessageDir(ctx.sessionID) const prevMessage = messageDir ? findNearestMessageWithFields(messageDir) : null - const parentAgent = ctx.agent ?? prevMessage?.agent + const firstMessageAgent = messageDir ? findFirstMessageWithAgent(messageDir) : null + const sessionAgent = getSessionAgent(ctx.sessionID) + const parentAgent = ctx.agent ?? sessionAgent ?? firstMessageAgent ?? prevMessage?.agent + + log("[sisyphus_task] parentAgent resolution", { + sessionID: ctx.sessionID, + messageDir, + ctxAgent: ctx.agent, + sessionAgent, + firstMessageAgent, + prevMessageAgent: prevMessage?.agent, + resolvedParentAgent: parentAgent, + }) const parentModel = prevMessage?.model?.providerID && prevMessage?.model?.modelID ? { providerID: prevMessage.model.providerID, modelID: prevMessage.model.modelID } : undefined