fix(team-mode): gate status injection by keyword

This commit is contained in:
YeonGyu-Kim
2026-05-15 15:35:35 +09:00
parent 1e7a7a22af
commit 15e7330ff0
3 changed files with 67 additions and 11 deletions
@@ -3,7 +3,7 @@ import { describe, expect, it } from "bun:test"
import { TeamModeConfigSchema } from "../../config/schema/team-mode" import { TeamModeConfigSchema } from "../../config/schema/team-mode"
import { createTeamModeStatusInjector } from "./hook" import { createTeamModeStatusInjector } from "./hook"
function createOutput(sessionID: string): { function createOutput(sessionID: string, text = "original message"): {
messages: Array<{ messages: Array<{
info: { role: string; sessionID: string } info: { role: string; sessionID: string }
parts: Array<{ type: string; text?: string; synthetic?: boolean }> parts: Array<{ type: string; text?: string; synthetic?: boolean }>
@@ -16,7 +16,7 @@ function createOutput(sessionID: string): {
role: "user", role: "user",
sessionID, sessionID,
}, },
parts: [{ type: "text", text: "original message" }], parts: [{ type: "text", text }],
}, },
], ],
} }
@@ -26,7 +26,7 @@ describe("createTeamModeStatusInjector", () => {
it("injects a one-time team mode enabled message before the latest user message", async () => { it("injects a one-time team mode enabled message before the latest user message", async () => {
// given // given
const hook = createTeamModeStatusInjector(TeamModeConfigSchema.parse({ enabled: true })) const hook = createTeamModeStatusInjector(TeamModeConfigSchema.parse({ enabled: true }))
const output = createOutput("session-team-mode") const output = createOutput("session-team-mode", "team mode please")
// when // when
await hook["experimental.chat.messages.transform"]?.( await hook["experimental.chat.messages.transform"]?.(
@@ -49,14 +49,14 @@ describe("createTeamModeStatusInjector", () => {
}, },
], ],
}) })
expect(output.messages[1]?.parts[0]?.text).toBe("original message") expect(output.messages[1]?.parts[0]?.text).toBe("team mode please")
}) })
it("does not inject again when the team mode status was already added", async () => { it("does not inject again when the team mode status was already added", async () => {
// given // given
const hook = createTeamModeStatusInjector(TeamModeConfigSchema.parse({ enabled: true })) const hook = createTeamModeStatusInjector(TeamModeConfigSchema.parse({ enabled: true }))
const firstOutput = createOutput("session-team-mode") const firstOutput = createOutput("session-team-mode", "team mode please")
const secondOutput = createOutput("session-team-mode") const secondOutput = createOutput("session-team-mode", "team mode please")
// when // when
await hook["experimental.chat.messages.transform"]?.( await hook["experimental.chat.messages.transform"]?.(
@@ -94,4 +94,39 @@ describe("createTeamModeStatusInjector", () => {
expect(output.messages).toHaveLength(1) expect(output.messages).toHaveLength(1)
expect(output.messages[0]?.parts[0]?.text).toBe("original message") expect(output.messages[0]?.parts[0]?.text).toBe("original message")
}) })
it("does not inject team mode status for punctuation-only prompts", async () => {
// given
const hook = createTeamModeStatusInjector(TeamModeConfigSchema.parse({ enabled: true }))
const output = createOutput("session-team-mode", ".")
// when
await hook["experimental.chat.messages.transform"]?.(
{ sessionID: "session-team-mode" },
output,
)
// then
expect(output.messages).toHaveLength(1)
expect(output.messages[0]?.parts[0]?.text).toBe(".")
})
it("does not inject team mode status when the team keyword is disabled", async () => {
// given
const hook = createTeamModeStatusInjector(
TeamModeConfigSchema.parse({ enabled: true }),
{ disabled_keywords: ["team"] },
)
const output = createOutput("session-team-mode", "team mode please")
// when
await hook["experimental.chat.messages.transform"]?.(
{ sessionID: "session-team-mode" },
output,
)
// then
expect(output.messages).toHaveLength(1)
expect(output.messages[0]?.parts[0]?.text).toBe("team mode please")
})
}) })
+25 -4
View File
@@ -1,4 +1,6 @@
import type { TeamModeConfig } from "../../config/schema/team-mode" import type { TeamModeConfig } from "../../config/schema/team-mode"
import type { KeywordDetectorConfig } from "../../config/schema/keyword-detector"
import { detectKeywordsWithType, extractPromptText } from "../keyword-detector/detector"
type TransformPart = { type TransformPart = {
type: string type: string
@@ -72,6 +74,25 @@ function hasInjectedTeamModeStatus(messages: MessageWithParts[]): boolean {
) )
} }
function latestUserMessageRequestsTeamMode(
messages: MessageWithParts[],
userMessageIndex: number,
keywordDetectorConfig?: KeywordDetectorConfig,
): boolean {
const message = messages[userMessageIndex]
if (message === undefined) {
return false
}
const promptText = extractPromptText(message.parts)
return detectKeywordsWithType(
promptText,
undefined,
undefined,
keywordDetectorConfig?.disabled_keywords,
).some((keyword) => keyword.type === "team")
}
function buildTeamModeStatusContent(): string { function buildTeamModeStatusContent(): string {
return `${TEAM_MODE_STATUS_MARKER} return `${TEAM_MODE_STATUS_MARKER}
Team mode is ENABLED for this session. Team mode is ENABLED for this session.
@@ -93,6 +114,7 @@ function createInjectedMessage(sessionID: string): MessageWithParts {
export function createTeamModeStatusInjector( export function createTeamModeStatusInjector(
config: TeamModeConfig, config: TeamModeConfig,
keywordDetectorConfig?: KeywordDetectorConfig,
): TeamModeStatusInjectorHook { ): TeamModeStatusInjectorHook {
return { return {
"experimental.chat.messages.transform": async ( "experimental.chat.messages.transform": async (
@@ -113,13 +135,12 @@ export function createTeamModeStatusInjector(
} }
const lastUserMessageIndex = findLastUserMessageIndex(output.messages) const lastUserMessageIndex = findLastUserMessageIndex(output.messages)
const injectedMessage = createInjectedMessage(sessionID) if (!latestUserMessageRequestsTeamMode(output.messages, lastUserMessageIndex, keywordDetectorConfig)) {
if (lastUserMessageIndex === -1) {
output.messages.unshift(injectedMessage)
return return
} }
const injectedMessage = createInjectedMessage(sessionID)
output.messages.splice(lastUserMessageIndex, 0, injectedMessage) output.messages.splice(lastUserMessageIndex, 0, injectedMessage)
}, },
} }
+1 -1
View File
@@ -74,7 +74,7 @@ export function createTransformHooks(args: {
const teamModeStatusInjector = teamModeConfig?.enabled const teamModeStatusInjector = teamModeConfig?.enabled
? safeCreateHook( ? safeCreateHook(
"team-mode-status-injector", "team-mode-status-injector",
() => createTeamModeStatusInjector(teamModeConfig), () => createTeamModeStatusInjector(teamModeConfig, pluginConfig.keyword_detector),
{ enabled: safeHookEnabled }, { enabled: safeHookEnabled },
) )
: null : null