-
Notifications
You must be signed in to change notification settings - Fork 865
Tool injector guardrail #1297
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Tool injector guardrail #1297
Changes from all commits
166cf0d
2b612ab
9c5e5ac
f221689
803b9ac
b774a35
2c58e25
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1 +1,5 @@ | ||
| npm run format:check || (npm run format && return 1) | ||
| #!/usr/bin/env sh | ||
| . "$(dirname -- "$0")/_/husky.sh" | ||
|
|
||
| # Check formatting; if it fails, format and fail the commit so dev can re-stage | ||
| npm run format:check || (npm run format && echo "Prettier formatted files. Please re-stage and commit again." && exit 1) |
| Original file line number | Diff line number | Diff line change | ||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,247 @@ | ||||||||||||||||
| import { | ||||||||||||||||
| HookEventType, | ||||||||||||||||
| PluginContext, | ||||||||||||||||
| PluginHandler, | ||||||||||||||||
| PluginParameters, | ||||||||||||||||
| } from '../types'; | ||||||||||||||||
| import { getCurrentContentPart, setCurrentContentPart } from '../utils'; | ||||||||||||||||
|
|
||||||||||||||||
| const addPrefixToCompletion = ( | ||||||||||||||||
| context: PluginContext, | ||||||||||||||||
| prefix: string, | ||||||||||||||||
| eventType: HookEventType | ||||||||||||||||
| ): Record<string, any> => { | ||||||||||||||||
| const transformedData: Record<string, any> = { | ||||||||||||||||
| request: { json: null }, | ||||||||||||||||
| response: { json: null }, | ||||||||||||||||
| }; | ||||||||||||||||
|
|
||||||||||||||||
| const { content, textArray } = getCurrentContentPart(context, eventType); | ||||||||||||||||
| if (!content) { | ||||||||||||||||
| return transformedData; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| const updatedTexts = ( | ||||||||||||||||
| Array.isArray(textArray) ? textArray : [String(textArray)] | ||||||||||||||||
| ).map((text, index) => (index === 0 ? `${prefix}${text ?? ''}` : text)); | ||||||||||||||||
|
|
||||||||||||||||
| setCurrentContentPart(context, eventType, transformedData, updatedTexts); | ||||||||||||||||
| return transformedData; | ||||||||||||||||
| }; | ||||||||||||||||
|
|
||||||||||||||||
| const addPrefixToChatCompletion = ( | ||||||||||||||||
| context: PluginContext, | ||||||||||||||||
| prefix: string, | ||||||||||||||||
| applyToRole: string = 'user', | ||||||||||||||||
| addToExisting: boolean = true, | ||||||||||||||||
| onlyIfEmpty: boolean = false, | ||||||||||||||||
| eventType: HookEventType | ||||||||||||||||
| ): Record<string, any> => { | ||||||||||||||||
| const json = context.request.json; | ||||||||||||||||
| const updatedJson = { ...json }; | ||||||||||||||||
| const messages = Array.isArray(json.messages) ? [...json.messages] : []; | ||||||||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🐛 Bug Fix Issue: Potential null pointer exception when
Suggested change
|
||||||||||||||||
|
|
||||||||||||||||
| // Find the target role message | ||||||||||||||||
| const targetIndex = messages.findIndex((msg) => msg.role === applyToRole); | ||||||||||||||||
|
|
||||||||||||||||
| // Helper to build a message content with the prefix in both chatComplete and messages formats | ||||||||||||||||
| const buildPrefixedContent = (existing: any): any => { | ||||||||||||||||
| if (existing == null || typeof existing === 'string') { | ||||||||||||||||
| return `${prefix}${existing ?? ''}`; | ||||||||||||||||
| } | ||||||||||||||||
| if (Array.isArray(existing)) { | ||||||||||||||||
| if (existing.length > 0 && existing[0]?.type === 'text') { | ||||||||||||||||
| const cloned = existing.map((item) => ({ ...item })); | ||||||||||||||||
| cloned[0].text = `${prefix}${cloned[0]?.text ?? ''}`; | ||||||||||||||||
| return cloned; | ||||||||||||||||
| } | ||||||||||||||||
| return [{ type: 'text', text: prefix }, ...existing]; | ||||||||||||||||
| } | ||||||||||||||||
| return `${prefix}${String(existing)}`; | ||||||||||||||||
| }; | ||||||||||||||||
|
|
||||||||||||||||
| // If the target role exists | ||||||||||||||||
| if (targetIndex !== -1) { | ||||||||||||||||
| const targetMsg = messages[targetIndex]; | ||||||||||||||||
| const content = targetMsg?.content; | ||||||||||||||||
|
|
||||||||||||||||
| const isEmptyContent = | ||||||||||||||||
| (typeof content === 'string' && content.trim().length === 0) || | ||||||||||||||||
| (Array.isArray(content) && content.length === 0); | ||||||||||||||||
|
Comment on lines
+68
to
+70
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🐛 Bug Fix Issue: Logical OR operator creates incorrect boolean condition - will always be true
Suggested change
|
||||||||||||||||
|
|
||||||||||||||||
| if (onlyIfEmpty && !isEmptyContent) { | ||||||||||||||||
| // Respect onlyIfEmpty by skipping modification when non-empty | ||||||||||||||||
| return { | ||||||||||||||||
| request: { json: updatedJson }, | ||||||||||||||||
| response: { json: null }, | ||||||||||||||||
| }; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| if (addToExisting) { | ||||||||||||||||
| // If this is the last message, leverage utils to ensure messages route compatibility | ||||||||||||||||
| if (targetIndex === messages.length - 1) { | ||||||||||||||||
| const transformedData: Record<string, any> = { | ||||||||||||||||
| request: { json: null }, | ||||||||||||||||
| response: { json: null }, | ||||||||||||||||
| }; | ||||||||||||||||
| const { content: currentContent, textArray } = getCurrentContentPart( | ||||||||||||||||
| context, | ||||||||||||||||
| eventType | ||||||||||||||||
| ); | ||||||||||||||||
| if (currentContent) { | ||||||||||||||||
| const updatedTexts = ( | ||||||||||||||||
| Array.isArray(textArray) ? textArray : [String(textArray)] | ||||||||||||||||
| ).map((text, idx) => (idx === 0 ? `${prefix}${text ?? ''}` : text)); | ||||||||||||||||
| setCurrentContentPart( | ||||||||||||||||
| context, | ||||||||||||||||
| eventType, | ||||||||||||||||
| transformedData, | ||||||||||||||||
| updatedTexts | ||||||||||||||||
| ); | ||||||||||||||||
| } | ||||||||||||||||
| return transformedData; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| // Otherwise, modify the specific message inline | ||||||||||||||||
| messages[targetIndex] = { | ||||||||||||||||
| ...targetMsg, | ||||||||||||||||
| content: buildPrefixedContent(targetMsg.content), | ||||||||||||||||
| }; | ||||||||||||||||
| } else { | ||||||||||||||||
| // Create new message with prefix before the existing one | ||||||||||||||||
| const newMessage = { | ||||||||||||||||
| role: applyToRole, | ||||||||||||||||
| content: | ||||||||||||||||
| context.requestType === 'messages' | ||||||||||||||||
| ? [{ type: 'text', text: prefix }] | ||||||||||||||||
| : prefix, | ||||||||||||||||
| }; | ||||||||||||||||
| messages.splice(targetIndex, 0, newMessage); | ||||||||||||||||
| } | ||||||||||||||||
| } else { | ||||||||||||||||
| // No message with target role exists, create one | ||||||||||||||||
| const newMessage = { | ||||||||||||||||
| role: applyToRole, | ||||||||||||||||
| content: | ||||||||||||||||
| context.requestType === 'messages' | ||||||||||||||||
| ? [{ type: 'text', text: prefix }] | ||||||||||||||||
| : prefix, | ||||||||||||||||
| }; | ||||||||||||||||
|
|
||||||||||||||||
| if (applyToRole === 'system') { | ||||||||||||||||
| messages.unshift(newMessage); | ||||||||||||||||
| } else { | ||||||||||||||||
| messages.push(newMessage); | ||||||||||||||||
| } | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| updatedJson.messages = messages; | ||||||||||||||||
|
|
||||||||||||||||
| return { | ||||||||||||||||
| request: { | ||||||||||||||||
| json: updatedJson, | ||||||||||||||||
| }, | ||||||||||||||||
| response: { | ||||||||||||||||
| json: null, | ||||||||||||||||
| }, | ||||||||||||||||
| }; | ||||||||||||||||
| }; | ||||||||||||||||
|
|
||||||||||||||||
| export const handler: PluginHandler = async ( | ||||||||||||||||
| context: PluginContext, | ||||||||||||||||
| parameters: PluginParameters, | ||||||||||||||||
| eventType: HookEventType | ||||||||||||||||
| ) => { | ||||||||||||||||
| let error = null; | ||||||||||||||||
| let verdict = true; // Always allow the request to continue | ||||||||||||||||
| let data = null; | ||||||||||||||||
| const transformedData: Record<string, any> = { | ||||||||||||||||
| request: { | ||||||||||||||||
| json: null, | ||||||||||||||||
| }, | ||||||||||||||||
| response: { | ||||||||||||||||
| json: null, | ||||||||||||||||
| }, | ||||||||||||||||
| }; | ||||||||||||||||
| let transformed = false; | ||||||||||||||||
|
|
||||||||||||||||
| try { | ||||||||||||||||
| // Only process before request and only for completion/chat completion/messages | ||||||||||||||||
| if ( | ||||||||||||||||
| eventType !== 'beforeRequestHook' || | ||||||||||||||||
| !['complete', 'chatComplete', 'messages'].includes( | ||||||||||||||||
| context.requestType || '' | ||||||||||||||||
| ) | ||||||||||||||||
| ) { | ||||||||||||||||
| return { | ||||||||||||||||
| error: null, | ||||||||||||||||
| verdict: true, | ||||||||||||||||
| data: null, | ||||||||||||||||
| transformedData, | ||||||||||||||||
| transformed, | ||||||||||||||||
| }; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| // Get prefix from parameters | ||||||||||||||||
| const prefix = parameters.prefix; | ||||||||||||||||
| if (!prefix || typeof prefix !== 'string') { | ||||||||||||||||
| return { | ||||||||||||||||
| error: { message: 'Prefix parameter is required and must be a string' }, | ||||||||||||||||
| verdict: true, | ||||||||||||||||
| data: null, | ||||||||||||||||
| transformedData, | ||||||||||||||||
| transformed, | ||||||||||||||||
| }; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| // Check if request JSON exists | ||||||||||||||||
| if (!context.request?.json) { | ||||||||||||||||
| return { | ||||||||||||||||
| error: { message: 'Request JSON is empty or missing' }, | ||||||||||||||||
| verdict: true, | ||||||||||||||||
| data: null, | ||||||||||||||||
| transformedData, | ||||||||||||||||
| transformed, | ||||||||||||||||
| }; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| let newTransformedData; | ||||||||||||||||
|
|
||||||||||||||||
| if ( | ||||||||||||||||
| context.requestType && | ||||||||||||||||
| ['chatComplete', 'messages'].includes(context.requestType) | ||||||||||||||||
| ) { | ||||||||||||||||
| // Handle chat completion | ||||||||||||||||
| newTransformedData = addPrefixToChatCompletion( | ||||||||||||||||
| context, | ||||||||||||||||
| prefix, | ||||||||||||||||
| parameters.applyToRole || 'user', | ||||||||||||||||
| parameters.addToExisting !== false, // default to true | ||||||||||||||||
| parameters.onlyIfEmpty === true, // default to false | ||||||||||||||||
| eventType | ||||||||||||||||
| ); | ||||||||||||||||
| } else { | ||||||||||||||||
| // Handle regular completion | ||||||||||||||||
| newTransformedData = addPrefixToCompletion(context, prefix, eventType); | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| Object.assign(transformedData, newTransformedData); | ||||||||||||||||
| transformed = true; | ||||||||||||||||
|
|
||||||||||||||||
| data = { | ||||||||||||||||
| prefix: prefix, | ||||||||||||||||
| requestType: context.requestType, | ||||||||||||||||
| applyToRole: parameters.applyToRole || 'user', | ||||||||||||||||
| addToExisting: parameters.addToExisting !== false, | ||||||||||||||||
| onlyIfEmpty: parameters.onlyIfEmpty === true, | ||||||||||||||||
| }; | ||||||||||||||||
| } catch (e: any) { | ||||||||||||||||
| delete e.stack; | ||||||||||||||||
| error = { | ||||||||||||||||
| message: `Error in addPrefix plugin: ${e.message || 'Unknown error'}`, | ||||||||||||||||
| originalError: e, | ||||||||||||||||
| }; | ||||||||||||||||
| } | ||||||||||||||||
|
|
||||||||||||||||
| return { error, verdict, data, transformedData, transformed }; | ||||||||||||||||
| }; | ||||||||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why was this changed?