diff --git a/src/features/context-injector/collector.test.ts b/src/features/context-injector/collector.test.ts new file mode 100644 index 0000000..52f4c05 --- /dev/null +++ b/src/features/context-injector/collector.test.ts @@ -0,0 +1,330 @@ +import { describe, it, expect, beforeEach } from "bun:test" +import { ContextCollector } from "./collector" +import type { ContextPriority, ContextSourceType } from "./types" + +describe("ContextCollector", () => { + let collector: ContextCollector + + beforeEach(() => { + collector = new ContextCollector() + }) + + describe("register", () => { + it("registers context for a session", () => { + // #given + const sessionID = "ses_test1" + const options = { + id: "ulw-context", + source: "keyword-detector" as ContextSourceType, + content: "Ultrawork mode activated", + } + + // #when + collector.register(sessionID, options) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.hasContent).toBe(true) + expect(pending.entries).toHaveLength(1) + expect(pending.entries[0].content).toBe("Ultrawork mode activated") + }) + + it("assigns default priority of 'normal' when not specified", () => { + // #given + const sessionID = "ses_test2" + + // #when + collector.register(sessionID, { + id: "test", + source: "keyword-detector", + content: "test content", + }) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.entries[0].priority).toBe("normal") + }) + + it("uses specified priority", () => { + // #given + const sessionID = "ses_test3" + + // #when + collector.register(sessionID, { + id: "critical-context", + source: "keyword-detector", + content: "critical content", + priority: "critical", + }) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.entries[0].priority).toBe("critical") + }) + + it("deduplicates by source + id combination", () => { + // #given + const sessionID = "ses_test4" + const options = { + id: "ulw-context", + source: "keyword-detector" as ContextSourceType, + content: "First content", + } + + // #when + collector.register(sessionID, options) + collector.register(sessionID, { ...options, content: "Updated content" }) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.entries).toHaveLength(1) + expect(pending.entries[0].content).toBe("Updated content") + }) + + it("allows same id from different sources", () => { + // #given + const sessionID = "ses_test5" + + // #when + collector.register(sessionID, { + id: "context-1", + source: "keyword-detector", + content: "From keyword-detector", + }) + collector.register(sessionID, { + id: "context-1", + source: "rules-injector", + content: "From rules-injector", + }) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.entries).toHaveLength(2) + }) + }) + + describe("getPending", () => { + it("returns empty result for session with no context", () => { + // #given + const sessionID = "ses_empty" + + // #when + const pending = collector.getPending(sessionID) + + // #then + expect(pending.hasContent).toBe(false) + expect(pending.entries).toHaveLength(0) + expect(pending.merged).toBe("") + }) + + it("merges multiple contexts with separator", () => { + // #given + const sessionID = "ses_merge" + collector.register(sessionID, { + id: "ctx-1", + source: "keyword-detector", + content: "First context", + }) + collector.register(sessionID, { + id: "ctx-2", + source: "rules-injector", + content: "Second context", + }) + + // #when + const pending = collector.getPending(sessionID) + + // #then + expect(pending.hasContent).toBe(true) + expect(pending.merged).toContain("First context") + expect(pending.merged).toContain("Second context") + }) + + it("orders contexts by priority (critical > high > normal > low)", () => { + // #given + const sessionID = "ses_priority" + collector.register(sessionID, { + id: "low", + source: "custom", + content: "LOW", + priority: "low", + }) + collector.register(sessionID, { + id: "critical", + source: "custom", + content: "CRITICAL", + priority: "critical", + }) + collector.register(sessionID, { + id: "normal", + source: "custom", + content: "NORMAL", + priority: "normal", + }) + collector.register(sessionID, { + id: "high", + source: "custom", + content: "HIGH", + priority: "high", + }) + + // #when + const pending = collector.getPending(sessionID) + + // #then + const order = pending.entries.map((e) => e.priority) + expect(order).toEqual(["critical", "high", "normal", "low"]) + }) + + it("maintains registration order within same priority", () => { + // #given + const sessionID = "ses_order" + collector.register(sessionID, { + id: "first", + source: "custom", + content: "First", + priority: "normal", + }) + collector.register(sessionID, { + id: "second", + source: "custom", + content: "Second", + priority: "normal", + }) + collector.register(sessionID, { + id: "third", + source: "custom", + content: "Third", + priority: "normal", + }) + + // #when + const pending = collector.getPending(sessionID) + + // #then + const ids = pending.entries.map((e) => e.id) + expect(ids).toEqual(["first", "second", "third"]) + }) + }) + + describe("consume", () => { + it("clears pending context for session", () => { + // #given + const sessionID = "ses_consume" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "test", + }) + + // #when + collector.consume(sessionID) + + // #then + const pending = collector.getPending(sessionID) + expect(pending.hasContent).toBe(false) + }) + + it("returns the consumed context", () => { + // #given + const sessionID = "ses_consume_return" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "test content", + }) + + // #when + const consumed = collector.consume(sessionID) + + // #then + expect(consumed.hasContent).toBe(true) + expect(consumed.entries[0].content).toBe("test content") + }) + + it("does not affect other sessions", () => { + // #given + const session1 = "ses_1" + const session2 = "ses_2" + collector.register(session1, { + id: "ctx", + source: "keyword-detector", + content: "session 1", + }) + collector.register(session2, { + id: "ctx", + source: "keyword-detector", + content: "session 2", + }) + + // #when + collector.consume(session1) + + // #then + expect(collector.getPending(session1).hasContent).toBe(false) + expect(collector.getPending(session2).hasContent).toBe(true) + }) + }) + + describe("clear", () => { + it("removes all context for a session", () => { + // #given + const sessionID = "ses_clear" + collector.register(sessionID, { + id: "ctx-1", + source: "keyword-detector", + content: "test 1", + }) + collector.register(sessionID, { + id: "ctx-2", + source: "rules-injector", + content: "test 2", + }) + + // #when + collector.clear(sessionID) + + // #then + expect(collector.getPending(sessionID).hasContent).toBe(false) + }) + }) + + describe("hasPending", () => { + it("returns true when session has pending context", () => { + // #given + const sessionID = "ses_has" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "test", + }) + + // #when / #then + expect(collector.hasPending(sessionID)).toBe(true) + }) + + it("returns false when session has no pending context", () => { + // #given + const sessionID = "ses_empty" + + // #when / #then + expect(collector.hasPending(sessionID)).toBe(false) + }) + + it("returns false after consume", () => { + // #given + const sessionID = "ses_after_consume" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "test", + }) + + // #when + collector.consume(sessionID) + + // #then + expect(collector.hasPending(sessionID)).toBe(false) + }) + }) +}) diff --git a/src/features/context-injector/collector.ts b/src/features/context-injector/collector.ts new file mode 100644 index 0000000..af60e41 --- /dev/null +++ b/src/features/context-injector/collector.ts @@ -0,0 +1,85 @@ +import type { + ContextEntry, + ContextPriority, + PendingContext, + RegisterContextOptions, +} from "./types" + +const PRIORITY_ORDER: Record = { + critical: 0, + high: 1, + normal: 2, + low: 3, +} + +const CONTEXT_SEPARATOR = "\n\n---\n\n" + +export class ContextCollector { + private sessions: Map> = new Map() + + register(sessionID: string, options: RegisterContextOptions): void { + if (!this.sessions.has(sessionID)) { + this.sessions.set(sessionID, new Map()) + } + + const sessionMap = this.sessions.get(sessionID)! + const key = `${options.source}:${options.id}` + + const entry: ContextEntry = { + id: options.id, + source: options.source, + content: options.content, + priority: options.priority ?? "normal", + timestamp: Date.now(), + metadata: options.metadata, + } + + sessionMap.set(key, entry) + } + + getPending(sessionID: string): PendingContext { + const sessionMap = this.sessions.get(sessionID) + + if (!sessionMap || sessionMap.size === 0) { + return { + merged: "", + entries: [], + hasContent: false, + } + } + + const entries = this.sortEntries([...sessionMap.values()]) + const merged = entries.map((e) => e.content).join(CONTEXT_SEPARATOR) + + return { + merged, + entries, + hasContent: entries.length > 0, + } + } + + consume(sessionID: string): PendingContext { + const pending = this.getPending(sessionID) + this.clear(sessionID) + return pending + } + + clear(sessionID: string): void { + this.sessions.delete(sessionID) + } + + hasPending(sessionID: string): boolean { + const sessionMap = this.sessions.get(sessionID) + return sessionMap !== undefined && sessionMap.size > 0 + } + + private sortEntries(entries: ContextEntry[]): ContextEntry[] { + return entries.sort((a, b) => { + const priorityDiff = PRIORITY_ORDER[a.priority] - PRIORITY_ORDER[b.priority] + if (priorityDiff !== 0) return priorityDiff + return a.timestamp - b.timestamp + }) + } +} + +export const contextCollector = new ContextCollector() diff --git a/src/features/context-injector/index.ts b/src/features/context-injector/index.ts new file mode 100644 index 0000000..9ab1330 --- /dev/null +++ b/src/features/context-injector/index.ts @@ -0,0 +1,12 @@ +export { ContextCollector, contextCollector } from "./collector" +export { injectPendingContext, createContextInjectorHook } from "./injector" +export type { + ContextSourceType, + ContextPriority, + ContextEntry, + RegisterContextOptions, + PendingContext, + MessageContext, + OutputParts, + InjectionStrategy, +} from "./types" diff --git a/src/features/context-injector/injector.test.ts b/src/features/context-injector/injector.test.ts new file mode 100644 index 0000000..6dc5b5a --- /dev/null +++ b/src/features/context-injector/injector.test.ts @@ -0,0 +1,172 @@ +import { describe, it, expect, beforeEach } from "bun:test" +import { ContextCollector } from "./collector" +import { injectPendingContext, createContextInjectorHook } from "./injector" + +describe("injectPendingContext", () => { + let collector: ContextCollector + + beforeEach(() => { + collector = new ContextCollector() + }) + + describe("when parts have text content", () => { + it("prepends context to first text part", () => { + // #given + const sessionID = "ses_inject1" + collector.register(sessionID, { + id: "ulw", + source: "keyword-detector", + content: "Ultrawork mode activated", + }) + const parts = [{ type: "text", text: "User message" }] + + // #when + const result = injectPendingContext(collector, sessionID, parts) + + // #then + expect(result.injected).toBe(true) + expect(parts[0].text).toContain("Ultrawork mode activated") + expect(parts[0].text).toContain("User message") + }) + + it("uses separator between context and original message", () => { + // #given + const sessionID = "ses_inject2" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context content", + }) + const parts = [{ type: "text", text: "Original message" }] + + // #when + injectPendingContext(collector, sessionID, parts) + + // #then + expect(parts[0].text).toBe("Context content\n\n---\n\nOriginal message") + }) + + it("consumes context after injection", () => { + // #given + const sessionID = "ses_inject3" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context", + }) + const parts = [{ type: "text", text: "Message" }] + + // #when + injectPendingContext(collector, sessionID, parts) + + // #then + expect(collector.hasPending(sessionID)).toBe(false) + }) + + it("returns injected=false when no pending context", () => { + // #given + const sessionID = "ses_empty" + const parts = [{ type: "text", text: "Message" }] + + // #when + const result = injectPendingContext(collector, sessionID, parts) + + // #then + expect(result.injected).toBe(false) + expect(parts[0].text).toBe("Message") + }) + }) + + describe("when parts have no text content", () => { + it("does not inject and preserves context", () => { + // #given + const sessionID = "ses_notext" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context", + }) + const parts = [{ type: "image", url: "https://example.com/img.png" }] + + // #when + const result = injectPendingContext(collector, sessionID, parts) + + // #then + expect(result.injected).toBe(false) + expect(collector.hasPending(sessionID)).toBe(true) + }) + }) + + describe("with multiple text parts", () => { + it("injects into first text part only", () => { + // #given + const sessionID = "ses_multi" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Context", + }) + const parts = [ + { type: "text", text: "First" }, + { type: "text", text: "Second" }, + ] + + // #when + injectPendingContext(collector, sessionID, parts) + + // #then + expect(parts[0].text).toContain("Context") + expect(parts[1].text).toBe("Second") + }) + }) +}) + +describe("createContextInjectorHook", () => { + let collector: ContextCollector + + beforeEach(() => { + collector = new ContextCollector() + }) + + describe("chat.message handler", () => { + it("injects pending context into output parts", async () => { + // #given + const hook = createContextInjectorHook(collector) + const sessionID = "ses_hook1" + collector.register(sessionID, { + id: "ctx", + source: "keyword-detector", + content: "Hook context", + }) + const input = { sessionID } + const output = { + message: {}, + parts: [{ type: "text", text: "User message" }], + } + + // #when + await hook["chat.message"](input, output) + + // #then + expect(output.parts[0].text).toContain("Hook context") + expect(output.parts[0].text).toContain("User message") + }) + + it("does nothing when no pending context", async () => { + // #given + const hook = createContextInjectorHook(collector) + const sessionID = "ses_hook2" + const input = { sessionID } + const output = { + message: {}, + parts: [{ type: "text", text: "User message" }], + } + + // #when + await hook["chat.message"](input, output) + + // #then + expect(output.parts[0].text).toBe("User message") + }) + }) +}) diff --git a/src/features/context-injector/injector.ts b/src/features/context-injector/injector.ts new file mode 100644 index 0000000..2ffb915 --- /dev/null +++ b/src/features/context-injector/injector.ts @@ -0,0 +1,61 @@ +import type { ContextCollector } from "./collector" + +const MESSAGE_SEPARATOR = "\n\n---\n\n" + +interface OutputPart { + type: string + text?: string + [key: string]: unknown +} + +interface InjectionResult { + injected: boolean + contextLength: number +} + +export function injectPendingContext( + collector: ContextCollector, + sessionID: string, + parts: OutputPart[] +): InjectionResult { + if (!collector.hasPending(sessionID)) { + return { injected: false, contextLength: 0 } + } + + const textPartIndex = parts.findIndex((p) => p.type === "text" && p.text !== undefined) + if (textPartIndex === -1) { + return { injected: false, contextLength: 0 } + } + + const pending = collector.consume(sessionID) + const originalText = parts[textPartIndex].text ?? "" + parts[textPartIndex].text = `${pending.merged}${MESSAGE_SEPARATOR}${originalText}` + + return { + injected: true, + contextLength: pending.merged.length, + } +} + +interface ChatMessageInput { + sessionID: string + agent?: string + model?: { providerID: string; modelID: string } + messageID?: string +} + +interface ChatMessageOutput { + message: Record + parts: OutputPart[] +} + +export function createContextInjectorHook(collector: ContextCollector) { + return { + "chat.message": async ( + input: ChatMessageInput, + output: ChatMessageOutput + ): Promise => { + injectPendingContext(collector, input.sessionID, output.parts) + }, + } +} diff --git a/src/features/context-injector/types.ts b/src/features/context-injector/types.ts new file mode 100644 index 0000000..c203be9 --- /dev/null +++ b/src/features/context-injector/types.ts @@ -0,0 +1,91 @@ +/** + * Source identifier for context injection + * Each source registers context that will be merged and injected together + */ +export type ContextSourceType = + | "keyword-detector" + | "rules-injector" + | "directory-agents" + | "directory-readme" + | "custom" + +/** + * Priority levels for context ordering + * Higher priority contexts appear first in the merged output + */ +export type ContextPriority = "critical" | "high" | "normal" | "low" + +/** + * A single context entry registered by a source + */ +export interface ContextEntry { + /** Unique identifier for this entry within the source */ + id: string + /** The source that registered this context */ + source: ContextSourceType + /** The actual context content to inject */ + content: string + /** Priority for ordering (default: normal) */ + priority: ContextPriority + /** Timestamp when registered */ + timestamp: number + /** Optional metadata for debugging/logging */ + metadata?: Record +} + +/** + * Options for registering context + */ +export interface RegisterContextOptions { + /** Unique ID for this context entry (used for deduplication) */ + id: string + /** Source identifier */ + source: ContextSourceType + /** The content to inject */ + content: string + /** Priority for ordering (default: normal) */ + priority?: ContextPriority + /** Optional metadata */ + metadata?: Record +} + +/** + * Result of getting pending context for a session + */ +export interface PendingContext { + /** Merged context string, ready for injection */ + merged: string + /** Individual entries that were merged */ + entries: ContextEntry[] + /** Whether there's any content to inject */ + hasContent: boolean +} + +/** + * Message context from the original user message + * Used when injecting to match the message format + */ +export interface MessageContext { + agent?: string + model?: { + providerID?: string + modelID?: string + } + path?: { + cwd?: string + root?: string + } + tools?: Record +} + +/** + * Output parts from chat.message hook + */ +export interface OutputParts { + parts: Array<{ type: string; text?: string; [key: string]: unknown }> +} + +/** + * Injection strategy + */ +export type InjectionStrategy = "prepend-parts" | "storage" | "auto" diff --git a/src/hooks/keyword-detector/index.ts b/src/hooks/keyword-detector/index.ts index 0db3ec0..e7a7886 100644 --- a/src/hooks/keyword-detector/index.ts +++ b/src/hooks/keyword-detector/index.ts @@ -1,13 +1,12 @@ import type { PluginInput } from "@opencode-ai/plugin" import { detectKeywordsWithType, extractPromptText } from "./detector" import { log } from "../../shared" -import { injectHookMessage } from "../../features/hook-message-injector" +import { contextCollector } from "../../features/context-injector" export * from "./detector" export * from "./constants" export * from "./types" -const sessionFirstMessageProcessed = new Set() const sessionUltraworkNotified = new Set() export function createKeywordDetectorHook(ctx: PluginInput) { @@ -24,9 +23,6 @@ export function createKeywordDetectorHook(ctx: PluginInput) { parts: Array<{ type: string; text?: string; [key: string]: unknown }> } ): Promise => { - const isFirstMessage = !sessionFirstMessageProcessed.has(input.sessionID) - sessionFirstMessageProcessed.add(input.sessionID) - const promptText = extractPromptText(output.parts) const detectedKeywords = detectKeywordsWithType(promptText) const messages = detectedKeywords.map((k) => k.message) @@ -39,50 +35,36 @@ export function createKeywordDetectorHook(ctx: PluginInput) { if (hasUltrawork && !sessionUltraworkNotified.has(input.sessionID)) { sessionUltraworkNotified.add(input.sessionID) log(`[keyword-detector] Ultrawork mode activated`, { sessionID: input.sessionID }) - - ctx.client.tui.showToast({ - body: { - title: "Ultrawork Mode Activated", - message: "Maximum precision engaged. All agents at your disposal.", - variant: "success" as const, - duration: 3000, - }, - }).catch((err) => log(`[keyword-detector] Failed to show toast`, { error: err, sessionID: input.sessionID })) + + ctx.client.tui + .showToast({ + body: { + title: "Ultrawork Mode Activated", + message: "Maximum precision engaged. All agents at your disposal.", + variant: "success" as const, + duration: 3000, + }, + }) + .catch((err) => + log(`[keyword-detector] Failed to show toast`, { error: err, sessionID: input.sessionID }) + ) } const context = messages.join("\n") - // First message: transform parts directly (for title generation compatibility) - if (isFirstMessage) { - log(`Keywords detected on first message, transforming parts directly`, { sessionID: input.sessionID, keywordCount: messages.length }) - const idx = output.parts.findIndex((p) => p.type === "text" && p.text) - if (idx >= 0) { - output.parts[idx].text = `${context}\n\n---\n\n${output.parts[idx].text ?? ""}` - } - return + for (const keyword of detectedKeywords) { + contextCollector.register(input.sessionID, { + id: `keyword-${keyword.type}`, + source: "keyword-detector", + content: keyword.message, + priority: keyword.type === "ultrawork" ? "critical" : "high", + }) } - // Subsequent messages: inject as separate message - log(`Keywords detected: ${messages.length}`, { sessionID: input.sessionID }) - - const message = output.message as { - agent?: string - model?: { modelID?: string; providerID?: string } - path?: { cwd?: string; root?: string } - tools?: Record - } - - log(`[keyword-detector] Injecting context for ${messages.length} keywords`, { sessionID: input.sessionID, contextLength: context.length }) - const success = injectHookMessage(input.sessionID, context, { - agent: message.agent, - model: message.model, - path: message.path, - tools: message.tools, + log(`[keyword-detector] Registered ${messages.length} keyword contexts`, { + sessionID: input.sessionID, + contextLength: context.length, }) - - if (success) { - log("Keyword context injected", { sessionID: input.sessionID }) - } }, } } diff --git a/src/index.ts b/src/index.ts index 0175ac8..56acf13 100644 --- a/src/index.ts +++ b/src/index.ts @@ -27,6 +27,10 @@ import { createAutoSlashCommandHook, createEditErrorRecoveryHook, } from "./hooks"; +import { + contextCollector, + createContextInjectorHook, +} from "./features/context-injector"; import { createGoogleAntigravityAuthPlugin } from "./auth/antigravity"; import { discoverUserClaudeSkills, @@ -130,6 +134,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { const keywordDetector = isHookEnabled("keyword-detector") ? createKeywordDetectorHook(ctx) : null; + const contextInjector = createContextInjectorHook(contextCollector); const agentUsageReminder = isHookEnabled("agent-usage-reminder") ? createAgentUsageReminderHook(ctx) : null; @@ -243,6 +248,7 @@ const OhMyOpenCodePlugin: Plugin = async (ctx) => { "chat.message": async (input, output) => { await claudeCodeHooks["chat.message"]?.(input, output); await keywordDetector?.["chat.message"]?.(input, output); + await contextInjector["chat.message"]?.(input, output); await autoSlashCommand?.["chat.message"]?.(input, output); if (ralphLoop) {