From 9ba9f906c523bfec05d2f204fbed9d347588dacf Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 4 Jan 2026 23:50:27 +0900 Subject: [PATCH] feat(context-injector): implement messages transform hook for context injection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Implement `createContextInjectorMessagesTransformHook` for messages transform hook - Refactor existing `chat.message` handler to be a no-op (context injection moved to transform) - Add comprehensive test suite for the new hook (4 test cases) - Update exports to expose new hook function 🤖 GENERATED WITH ASSISTANCE OF [OhMyOpenCode](https://github.com/code-yeongyu/oh-my-opencode) --- src/features/context-injector/index.ts | 6 +- .../context-injector/injector.test.ts | 128 +++++++++++++++++- src/features/context-injector/injector.ts | 117 +++++++++++++++- src/index.ts | 5 + 4 files changed, 245 insertions(+), 11 deletions(-) diff --git a/src/features/context-injector/index.ts b/src/features/context-injector/index.ts index 9ab1330..51e3045 100644 --- a/src/features/context-injector/index.ts +++ b/src/features/context-injector/index.ts @@ -1,5 +1,9 @@ export { ContextCollector, contextCollector } from "./collector" -export { injectPendingContext, createContextInjectorHook } from "./injector" +export { + injectPendingContext, + createContextInjectorHook, + createContextInjectorMessagesTransformHook, +} from "./injector" export type { ContextSourceType, ContextPriority, diff --git a/src/features/context-injector/injector.test.ts b/src/features/context-injector/injector.test.ts index 6dc5b5a..ae0924d 100644 --- a/src/features/context-injector/injector.test.ts +++ b/src/features/context-injector/injector.test.ts @@ -1,6 +1,10 @@ import { describe, it, expect, beforeEach } from "bun:test" import { ContextCollector } from "./collector" -import { injectPendingContext, createContextInjectorHook } from "./injector" +import { + injectPendingContext, + createContextInjectorHook, + createContextInjectorMessagesTransformHook, +} from "./injector" describe("injectPendingContext", () => { let collector: ContextCollector @@ -129,7 +133,7 @@ describe("createContextInjectorHook", () => { }) describe("chat.message handler", () => { - it("injects pending context into output parts", async () => { + it("is a no-op (context injection moved to messages transform)", async () => { // #given const hook = createContextInjectorHook(collector) const sessionID = "ses_hook1" @@ -148,8 +152,8 @@ describe("createContextInjectorHook", () => { await hook["chat.message"](input, output) // #then - expect(output.parts[0].text).toContain("Hook context") - expect(output.parts[0].text).toContain("User message") + expect(output.parts[0].text).toBe("User message") + expect(collector.hasPending(sessionID)).toBe(true) }) it("does nothing when no pending context", async () => { @@ -170,3 +174,119 @@ describe("createContextInjectorHook", () => { }) }) }) + +describe("createContextInjectorMessagesTransformHook", () => { + let collector: ContextCollector + + beforeEach(() => { + collector = new ContextCollector() + }) + + const createMockMessage = ( + role: "user" | "assistant", + text: string, + sessionID: string + ) => ({ + info: { + id: `msg_${Date.now()}_${Math.random()}`, + sessionID, + role, + time: { created: Date.now() }, + agent: "Sisyphus", + model: { providerID: "test", modelID: "test" }, + path: { cwd: "/", root: "/" }, + }, + parts: [ + { + id: `part_${Date.now()}`, + sessionID, + messageID: `msg_${Date.now()}`, + type: "text" as const, + text, + }, + ], + }) + + it("inserts synthetic message before last user message", async () => { + // #given + const hook = createContextInjectorMessagesTransformHook(collector) + const sessionID = "ses_transform1" + collector.register(sessionID, { + id: "ulw", + source: "keyword-detector", + content: "Ultrawork context", + }) + const messages = [ + createMockMessage("user", "First message", sessionID), + createMockMessage("assistant", "Response", sessionID), + createMockMessage("user", "Second message", sessionID), + ] + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const output = { messages } as any + + // #when + await hook["experimental.chat.messages.transform"]!({}, output) + + // #then + expect(output.messages.length).toBe(4) + expect(output.messages[2].parts[0].text).toBe("Ultrawork context") + expect(output.messages[2].parts[0].synthetic).toBe(true) + expect(output.messages[3].parts[0].text).toBe("Second message") + }) + + it("does nothing when no pending context", async () => { + // #given + const hook = createContextInjectorMessagesTransformHook(collector) + const sessionID = "ses_transform2" + const messages = [createMockMessage("user", "Message", sessionID)] + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const output = { messages } as any + + // #when + await hook["experimental.chat.messages.transform"]!({}, output) + + // #then + expect(output.messages.length).toBe(1) + }) + + it("does nothing when no user messages", async () => { + // #given + const hook = createContextInjectorMessagesTransformHook(collector) + const sessionID = "ses_transform3" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context", + }) + const messages = [createMockMessage("assistant", "Response", sessionID)] + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const output = { messages } as any + + // #when + await hook["experimental.chat.messages.transform"]!({}, output) + + // #then + expect(output.messages.length).toBe(1) + expect(collector.hasPending(sessionID)).toBe(true) + }) + + it("consumes context after injection", async () => { + // #given + const hook = createContextInjectorMessagesTransformHook(collector) + const sessionID = "ses_transform4" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context", + }) + const messages = [createMockMessage("user", "Message", sessionID)] + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const output = { messages } as any + + // #when + await hook["experimental.chat.messages.transform"]!({}, output) + + // #then + expect(collector.hasPending(sessionID)).toBe(false) + }) +}) diff --git a/src/features/context-injector/injector.ts b/src/features/context-injector/injector.ts index 2ffb915..4ed1ebe 100644 --- a/src/features/context-injector/injector.ts +++ b/src/features/context-injector/injector.ts @@ -1,6 +1,6 @@ import type { ContextCollector } from "./collector" - -const MESSAGE_SEPARATOR = "\n\n---\n\n" +import type { Message, Part } from "@opencode-ai/sdk" +import { log } from "../../shared" interface OutputPart { type: string @@ -29,7 +29,7 @@ export function injectPendingContext( const pending = collector.consume(sessionID) const originalText = parts[textPartIndex].text ?? "" - parts[textPartIndex].text = `${pending.merged}${MESSAGE_SEPARATOR}${originalText}` + parts[textPartIndex].text = `${pending.merged}\n\n---\n\n${originalText}` return { injected: true, @@ -52,10 +52,115 @@ interface ChatMessageOutput { export function createContextInjectorHook(collector: ContextCollector) { return { "chat.message": async ( - input: ChatMessageInput, - output: ChatMessageOutput + _input: ChatMessageInput, + _output: ChatMessageOutput ): Promise => { - injectPendingContext(collector, input.sessionID, output.parts) + void collector + }, + } +} + +interface MessageWithParts { + info: Message + parts: Part[] +} + +type MessagesTransformHook = { + "experimental.chat.messages.transform"?: ( + input: Record, + output: { messages: MessageWithParts[] } + ) => Promise +} + +export function createContextInjectorMessagesTransformHook( + collector: ContextCollector +): MessagesTransformHook { + return { + "experimental.chat.messages.transform": async (_input, output) => { + const { messages } = output + if (messages.length === 0) { + log("[context-injector] messages.transform: no messages") + return + } + + const lastMessage = messages[messages.length - 1] + const sessionID = (lastMessage.info as unknown as { sessionID?: string }).sessionID + if (!sessionID) { + log("[context-injector] messages.transform: no sessionID on last message") + return + } + + const hasPending = collector.hasPending(sessionID) + log("[context-injector] messages.transform check", { + sessionID, + hasPending, + messageCount: messages.length, + }) + + if (!hasPending) return + + let lastUserMessageIndex = -1 + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].info.role === "user") { + lastUserMessageIndex = i + break + } + } + + if (lastUserMessageIndex === -1) { + log("[context-injector] messages.transform: no user message found") + return + } + + const pending = collector.consume(sessionID) + if (!pending.hasContent) { + log("[context-injector] messages.transform: pending was empty") + return + } + + const refMessage = messages[lastUserMessageIndex] + const refInfo = refMessage.info as unknown as { + sessionID?: string + agent?: string + model?: { providerID?: string; modelID?: string } + path?: { cwd?: string; root?: string } + } + + const syntheticMessageId = `synthetic_ctx_${Date.now()}` + const syntheticPartId = `synthetic_ctx_part_${Date.now()}` + const now = Date.now() + + const syntheticMessage: MessageWithParts = { + info: { + id: syntheticMessageId, + sessionID: sessionID, + role: "user", + time: { created: now }, + agent: refInfo.agent ?? "Sisyphus", + model: refInfo.model ?? { providerID: "unknown", modelID: "unknown" }, + path: refInfo.path ?? { cwd: "/", root: "/" }, + } as unknown as Message, + parts: [ + { + id: syntheticPartId, + sessionID: sessionID, + messageID: syntheticMessageId, + type: "text", + text: pending.merged, + synthetic: true, + time: { start: now, end: now }, + } as Part, + ], + } + + messages.splice(lastUserMessageIndex, 0, syntheticMessage) + + log("[context-injector] Injected synthetic message", { + sessionID, + insertIndex: lastUserMessageIndex, + contextLength: pending.merged.length, + newMessageCount: messages.length, + }) }, } } diff --git a/src/index.ts b/src/index.ts index 9a8798f..9e188e6 100644 --- a/src/index.ts +++ b/src/index.ts @@ -30,6 +30,7 @@ import { import { contextCollector, createContextInjectorHook, + createContextInjectorMessagesTransformHook, } from "./features/context-injector"; import { createGoogleAntigravityAuthPlugin } from "./auth/antigravity"; import { @@ -138,6 +139,8 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { ? createKeywordDetectorHook(ctx) : null; const contextInjector = createContextInjectorHook(contextCollector); + const contextInjectorMessagesTransform = + createContextInjectorMessagesTransformHook(contextCollector); const agentUsageReminder = isHookEnabled("agent-usage-reminder") ? createAgentUsageReminderHook(ctx) : null; @@ -315,6 +318,8 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { input: Record, output: { messages: Array<{ info: unknown; parts: unknown[] }> } ) => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + await contextInjectorMessagesTransform?.["experimental.chat.messages.transform"]?.(input, output as any); await thinkingBlockValidator?.[ "experimental.chat.messages.transform" // eslint-disable-next-line @typescript-eslint/no-explicit-any