From b2a78c849d3776d9c4652c5ee8c2983d16c77368 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Mon, 9 Feb 2026 19:15:54 -0500 Subject: [PATCH 1/9] feat: add new middleware package including fallback, filesystem, and retry, and update core AI modules to support it. --- js/plugins/middleware/.npmignore | 3 + js/plugins/middleware/LICENSE | 203 +++++++++++ js/plugins/middleware/README.md | 3 + js/plugins/middleware/examples/fallback.ts | 42 +++ js/plugins/middleware/examples/filesystem.ts | 39 +++ js/plugins/middleware/examples/retry.ts | 42 +++ js/plugins/middleware/package.json | 51 +++ js/plugins/middleware/src/fallback.ts | 138 ++++++++ js/plugins/middleware/src/filesystem.ts | 124 +++++++ js/plugins/middleware/src/index.ts | 19 ++ js/plugins/middleware/src/retry.ts | 169 ++++++++++ js/plugins/middleware/tests/fallback_test.ts | 193 +++++++++++ .../middleware/tests/filesystem_test.ts | 142 ++++++++ js/plugins/middleware/tests/retry_test.ts | 314 ++++++++++++++++++ js/plugins/middleware/tsconfig.json | 4 + js/plugins/middleware/tsup.config.ts | 22 ++ js/plugins/middleware/typedoc.json | 3 + 17 files changed, 1511 insertions(+) create mode 100644 js/plugins/middleware/.npmignore create mode 100644 js/plugins/middleware/LICENSE create mode 100644 js/plugins/middleware/README.md create mode 100644 js/plugins/middleware/examples/fallback.ts create mode 100644 js/plugins/middleware/examples/filesystem.ts create mode 100644 js/plugins/middleware/examples/retry.ts create mode 100644 js/plugins/middleware/package.json create mode 100644 js/plugins/middleware/src/fallback.ts create mode 100644 js/plugins/middleware/src/filesystem.ts create mode 100644 js/plugins/middleware/src/index.ts create mode 100644 js/plugins/middleware/src/retry.ts create mode 100644 js/plugins/middleware/tests/fallback_test.ts create mode 100644 js/plugins/middleware/tests/filesystem_test.ts create mode 100644 js/plugins/middleware/tests/retry_test.ts create mode 100644 js/plugins/middleware/tsconfig.json create mode 100644 js/plugins/middleware/tsup.config.ts create mode 100644 js/plugins/middleware/typedoc.json diff --git a/js/plugins/middleware/.npmignore b/js/plugins/middleware/.npmignore new file mode 100644 index 0000000000..5ea8bc610d --- /dev/null +++ b/js/plugins/middleware/.npmignore @@ -0,0 +1,3 @@ +node_modules +tsconfig.json +tsup.config.ts \ No newline at end of file diff --git a/js/plugins/middleware/LICENSE b/js/plugins/middleware/LICENSE new file mode 100644 index 0000000000..26a870243c --- /dev/null +++ b/js/plugins/middleware/LICENSE @@ -0,0 +1,203 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + \ No newline at end of file diff --git a/js/plugins/middleware/README.md b/js/plugins/middleware/README.md new file mode 100644 index 0000000000..bd07280cda --- /dev/null +++ b/js/plugins/middleware/README.md @@ -0,0 +1,3 @@ +# Genkit Middleware + +This package provides middleware for Genkit. diff --git a/js/plugins/middleware/examples/fallback.ts b/js/plugins/middleware/examples/fallback.ts new file mode 100644 index 0000000000..27899fce44 --- /dev/null +++ b/js/plugins/middleware/examples/fallback.ts @@ -0,0 +1,42 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { googleAI } from '@genkit-ai/google-genai'; +import { genkit } from 'genkit'; +import { fallback } from '../src/index.js'; + +const ai = genkit({ + plugins: [googleAI()], +}); + +async function main() { + const { text } = await ai.generate({ + model: googleAI.model('gemini-3-pro-preview'), + prompt: 'Tell me a joke about fallbacks.', + use: [ + fallback({ + models: [googleAI.model('gemini-3-flash-preview')], + statuses: ['RESOURCE_EXHAUSTED', 'UNAVAILABLE', 'DEADLINE_EXCEEDED'], + onError: (err) => { + console.error(`Fallback triggered due to error: ${err.message}`); + }, + }), + ], + }); + console.log(text); +} + +main().catch(console.error); diff --git a/js/plugins/middleware/examples/filesystem.ts b/js/plugins/middleware/examples/filesystem.ts new file mode 100644 index 0000000000..1e8ab20bab --- /dev/null +++ b/js/plugins/middleware/examples/filesystem.ts @@ -0,0 +1,39 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { googleAI } from '@genkit-ai/google-genai'; +import { genkit } from 'genkit'; +import { filesystem } from '../src/index.js'; + +const ai = genkit({ + plugins: [googleAI()], +}); + +async function main() { + const { text } = await ai.generate({ + model: googleAI.model('gemini-3-flash-preview'), + prompt: + 'Can you list the files in my temporary directory and read whatever is in `hello.txt`?', + use: [ + filesystem({ + rootDirectory: process.env.TEMP_DIR || '/tmp', + }), + ], + }); + console.log(text); +} + +main().catch(console.error); diff --git a/js/plugins/middleware/examples/retry.ts b/js/plugins/middleware/examples/retry.ts new file mode 100644 index 0000000000..44c46475eb --- /dev/null +++ b/js/plugins/middleware/examples/retry.ts @@ -0,0 +1,42 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { googleAI } from '@genkit-ai/google-genai'; +import { genkit } from 'genkit'; +import { retry } from '../src/index.js'; // @genkit-ai/middleware + +const ai = genkit({ + plugins: [googleAI()], +}); + +async function main() { + const { text } = await ai.generate({ + model: googleAI.model('gemini-3-pro-preview'), + prompt: 'Tell me a joke about retries.', + use: [ + retry({ + maxRetries: 3, + initialDelayMs: 500, + onError: (err, attempt) => { + console.error(`Retry attempt ${attempt} failed: ${err.message}`); + }, + }), + ], + }); + console.log(text); +} + +main().catch(console.error); diff --git a/js/plugins/middleware/package.json b/js/plugins/middleware/package.json new file mode 100644 index 0000000000..8301d1b1a6 --- /dev/null +++ b/js/plugins/middleware/package.json @@ -0,0 +1,51 @@ +{ + "name": "@genkit-ai/middleware", + "description": "Middleware for Genkit", + "keywords": [ + "genkit", + "genkit-plugin", + "google cloud", + "google ai", + "ai", + "genai", + "generative-ai" + ], + "version": "1.28.0", + "type": "commonjs", + "scripts": { + "check": "tsc", + "compile": "tsup-node", + "build:clean": "rimraf ./lib", + "build": "npm-run-all build:clean check compile", + "build:watch": "tsup-node --watch", + "test": "tsx --test tests/**/*.ts" + }, + "repository": { + "type": "git", + "url": "https://github.com/firebase/genkit.git", + "directory": "js/plugins/middleware" + }, + "author": "genkit", + "license": "Apache-2.0", + "peerDependencies": { + "genkit": "workspace:^" + }, + "devDependencies": { + "@genkit-ai/google-genai": "workspace:^", + "@types/node": "^20.11.16", + "npm-run-all": "^4.1.5", + "rimraf": "^6.0.1", + "tsup": "^8.0.2", + "tsx": "^4.7.0", + "typescript": "^4.9.0" + }, + "types": "./lib/index.d.ts", + "exports": { + ".": { + "types": "./lib/index.d.ts", + "import": "./lib/index.mjs", + "require": "./lib/index.js", + "default": "./lib/index.js" + } + } +} diff --git a/js/plugins/middleware/src/fallback.ts b/js/plugins/middleware/src/fallback.ts new file mode 100644 index 0000000000..54c1771961 --- /dev/null +++ b/js/plugins/middleware/src/fallback.ts @@ -0,0 +1,138 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + GenkitError, + ModelReferenceSchema, + generateMiddleware, + modelRef, + z, + type GenerateMiddleware, + type ModelReference, + type StatusName, +} from 'genkit'; + +const DEFAULT_FALLBACK_STATUSES: StatusName[] = [ + 'UNAVAILABLE', + 'DEADLINE_EXCEEDED', + 'RESOURCE_EXHAUSTED', + 'ABORTED', + 'INTERNAL', + 'NOT_FOUND', + 'UNIMPLEMENTED', +]; + +export const FallbackOptionsSchema = z + .object({ + /** + * An array of models to try in order. + */ + models: z.array(z.union([z.string(), ModelReferenceSchema])), + /** + * An array of `StatusName` values that should trigger a fallback. + * @default ['UNAVAILABLE', 'DEADLINE_EXCEEDED', 'RESOURCE_EXHAUSTED', 'ABORTED', 'INTERNAL', 'NOT_FOUND', 'UNIMPLEMENTED'] + */ + statuses: z.array(z.string()).optional(), + }) + .passthrough(); + +export interface FallbackOptions extends z.infer { + /** + * A callback to be executed on each fallback attempt. + */ + onError?: (error: Error) => void; +} + +/** + * Creates a middleware that falls back to a different model on specific error statuses. + * + * ```ts + * const { text } = await ai.generate({ + * model: googleAI.model('gemini-2.5-pro'), + * prompt: 'You are a helpful AI assistant named Walt, say hello', + * use: [ + * fallback({ + * models: [googleAI.model('gemini-2.5-flash')], + * statuses: ['RESOURCE_EXHAUSTED'], + * }), + * ], + * }); + * ``` + */ +export const fallback: GenerateMiddleware = + generateMiddleware( + { + name: 'fallback', + configSchema: FallbackOptionsSchema, + }, + (options?: FallbackOptions) => { + const { + models = [], + statuses = DEFAULT_FALLBACK_STATUSES, + onError, + } = options || {}; + + return { + generate: async (req, ctx, next) => { + try { + return await next(req, ctx); + } catch (e) { + if ( + e instanceof GenkitError && + statuses.includes(e.status as StatusName) + ) { + onError?.(e); + let lastError: any = e; + for (const model of models) { + const normalizedModel = normalizeModel(model); + try { + return await next( + { + ...req, + model: normalizedModel.name, + config: normalizedModel.config ?? req.config, + }, + ctx + ); + } catch (e2) { + lastError = e2; + if ( + e2 instanceof GenkitError && + statuses.includes(e2.status as StatusName) + ) { + onError?.(e2); + continue; + } + throw e2; + } + } + throw lastError; + } + throw e; + } + }, + }; + } + ); + +function normalizeModel( + model: string | z.infer +): ModelReference { + if (typeof model === 'string') { + return modelRef({ name: model }); + } + return modelRef({ name: model.name, config: model.config }); +} diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts new file mode 100644 index 0000000000..91315c91d7 --- /dev/null +++ b/js/plugins/middleware/src/filesystem.ts @@ -0,0 +1,124 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as fs from 'fs/promises'; +import { generateMiddleware, z, type GenerateMiddleware } from 'genkit'; +import { tool } from 'genkit/beta'; +import * as path from 'path'; + +export const FilesystemOptionsSchema = z.object({ + rootDirectory: z + .string() + .describe( + 'The root directory to which all filesystem operations are restricted.' + ), +}); + +export type FilesystemOptions = z.infer; + +/** + * Creates a middleware that grants the LLM basic readonly access to the filesystem. + * Injects `list_files` and `read_file` tools restricted to the provided `rootDirectory`. + */ +export const filesystem: GenerateMiddleware = + generateMiddleware( + { + name: 'filesystem', + configSchema: FilesystemOptionsSchema, + }, + (options) => { + if (!options?.rootDirectory) { + throw new Error( + 'filesystem middleware requires a rootDirectory option' + ); + } + const rootDir = path.resolve(options.rootDirectory); + + function resolvePath(requestedPath: string) { + const p = path.resolve(rootDir, requestedPath); + // Ensure the resolved path starts with the rootDir and a path separator + // to prevent directory traversal attacks (e.g. rootDir is /a/b, requested is ../b_secret) + if (!p.startsWith(rootDir + path.sep) && p !== rootDir) { + throw new Error('Access denied: Path is outside of root directory.'); + } + return p; + } + + const listFiles = tool( + { + name: 'list_files', + description: + 'Lists files and directories in a given path. Returns a list of strings.', + inputSchema: z.object({ + dirPath: z + .string() + .describe('Directory path relative to root.') + .default(''), + recursive: z + .boolean() + .describe('Whether to list files recursively.') + .default(false), + }), + outputSchema: z.array(z.string()), + }, + async (input) => { + const targetDir = resolvePath(input.dirPath); + + async function list( + dir: string, + recursive: boolean, + base: string = '' + ) { + const results: string[] = []; + const entries = await fs.readdir(dir, { withFileTypes: true }); + for (const entry of entries) { + const relativePath = path.join(base, entry.name); + results.push(relativePath); + if (entry.isDirectory() && recursive) { + const subResults = await list( + path.join(dir, entry.name), + true, + relativePath + ); + results.push(...subResults); + } + } + return results; + } + return await list(targetDir, input.recursive); + } + ); + + const readFile = tool( + { + name: 'read_file', + description: 'Reads the contents of a file', + inputSchema: z.object({ + filePath: z.string().describe('File path relative to root.'), + }), + outputSchema: z.string(), + }, + async (input) => { + const targetFile = resolvePath(input.filePath); + return await fs.readFile(targetFile, 'utf8'); + } + ); + + return { + tools: [listFiles, readFile], + }; + } + ); diff --git a/js/plugins/middleware/src/index.ts b/js/plugins/middleware/src/index.ts new file mode 100644 index 0000000000..6dbbaac587 --- /dev/null +++ b/js/plugins/middleware/src/index.ts @@ -0,0 +1,19 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +export { fallback } from './fallback.js'; +export { filesystem } from './filesystem.js'; +export { retry } from './retry.js'; diff --git a/js/plugins/middleware/src/retry.ts b/js/plugins/middleware/src/retry.ts new file mode 100644 index 0000000000..161fbf835f --- /dev/null +++ b/js/plugins/middleware/src/retry.ts @@ -0,0 +1,169 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + GenkitError, + generateMiddleware, + z, + type GenerateMiddleware, + type StatusName, +} from 'genkit'; + +export const RetryOptionsSchema = z + .object({ + /** + * The maximum number of times to retry a failed request. + * @default 3 + */ + maxRetries: z.number().optional(), + /** + * An array of `StatusName` values that should trigger a retry. + * @default ['UNAVAILABLE', 'DEADLINE_EXCEEDED', 'RESOURCE_EXHAUSTED', 'ABORTED', 'INTERNAL'] + */ + statuses: z.array(z.string()).optional(), + /** + * The initial delay between retries, in milliseconds. + * @default 1000 + */ + initialDelayMs: z.number().optional(), + /** + * The maximum delay between retries, in milliseconds. + * @default 60000 + */ + maxDelayMs: z.number().optional(), + /** + * The factor by which the delay increases after each retry (exponential backoff). + * @default 2 + */ + backoffFactor: z.number().optional(), + /** + * Whether to disable jitter on the delay. Jitter adds a random factor to the + * delay to help prevent a "thundering herd" of clients all retrying at the + * same time. + * @default false + */ + noJitter: z.boolean().optional(), + }) + .passthrough(); + +export interface RetryOptions extends z.infer { + /** + * A callback to be executed on each retry attempt. + */ + onError?: (error: Error, attempt: number) => void; +} + +const DEFAULT_RETRY_STATUSES: StatusName[] = [ + 'UNAVAILABLE', + 'DEADLINE_EXCEEDED', + 'RESOURCE_EXHAUSTED', + 'ABORTED', + 'INTERNAL', +]; + +let __setTimeout: ( + callback: (...args: any[]) => void, + ms?: number +) => NodeJS.Timeout = setTimeout; + +/** + * FOR TESTING ONLY. + * @internal + */ +export const TEST_ONLY = { + setRetryTimeout( + impl: (callback: (...args: any[]) => void, ms?: number) => NodeJS.Timeout + ) { + __setTimeout = impl; + }, +}; + +/** + * Creates a middleware that retries requests on specific error statuses. + * + * ```ts + * const { text } = await ai.generate({ + * model: googleAI.model('gemini-2.5-pro'), + * prompt: 'You are a helpful AI assistant named Walt, say hello', + * use: [ + * retry({ + * maxRetries: 2, + * initialDelayMs: 1000, + * backoffFactor: 2, + * }), + * ], + * }); + * ``` + */ +export const retry: GenerateMiddleware = + generateMiddleware( + { + name: 'retry', + configSchema: RetryOptionsSchema, + }, + (options?: RetryOptions) => { + const { + maxRetries = 3, + statuses = DEFAULT_RETRY_STATUSES, + initialDelayMs = 1000, + maxDelayMs = 60000, + backoffFactor = 2, + noJitter = false, + onError, + } = options || {}; + + return { + model: async (req, ctx, next) => { + let lastError: any; + let currentDelay = initialDelayMs; + for (let i = 0; i <= maxRetries; i++) { + try { + return await next(req, ctx); + } catch (e) { + lastError = e; + const error = e as Error; + if (i < maxRetries) { + let shouldRetry = false; + if (error instanceof GenkitError) { + if ((statuses as string[]).includes(error.status)) { + shouldRetry = true; + } + } else { + shouldRetry = true; + } + + if (shouldRetry) { + onError?.(error, i + 1); + let delay = currentDelay; + if (!noJitter) { + delay = delay + 1000 * Math.pow(2, i) * Math.random(); + } + await new Promise((resolve) => __setTimeout(resolve, delay)); + currentDelay = Math.min( + currentDelay * backoffFactor, + maxDelayMs + ); + continue; + } + } + throw error; + } + } + throw lastError; + }, + }; + } + ); diff --git a/js/plugins/middleware/tests/fallback_test.ts b/js/plugins/middleware/tests/fallback_test.ts new file mode 100644 index 0000000000..009ad53d16 --- /dev/null +++ b/js/plugins/middleware/tests/fallback_test.ts @@ -0,0 +1,193 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import { GenkitError, genkit } from 'genkit'; +import { describe, it } from 'node:test'; +import { fallback } from '../src/fallback.js'; + +describe('fallback', () => { + it('should not fallback on success', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + let pmFallbackRequestCount = 0; + const pmFallback = ai.defineModel( + { name: 'programmableModelFallback' }, + async (req) => { + pmFallbackRequestCount++; + return { + message: { role: 'model', content: [{ text: 'fallback success' }] }, + }; + } + ); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [fallback({ models: [pmFallback] })], + }); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(pmFallbackRequestCount, 0); + assert.strictEqual(result.text, 'success'); + }); + + it('should call onError callback', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'UNAVAILABLE', message: 'test' }); + }); + + let pmFallbackRequestCount = 0; + ai.defineModel({ name: 'programmableModelFallback' }, async (req) => { + pmFallbackRequestCount++; + throw new GenkitError({ + status: 'RESOURCE_EXHAUSTED', + message: 'test2', + }); + }); + + let errorCount = 0; + let lastError: Error | undefined; + + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [ + fallback({ + models: ['programmableModelFallback'], + onError: (err) => { + errorCount++; + lastError = err; + }, + }), + ], + }), + /RESOURCE_EXHAUSTED: test2/ + ); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(pmFallbackRequestCount, 1); + assert.strictEqual(errorCount, 2); + assert.ok(lastError); + assert.strictEqual(lastError!.message, 'RESOURCE_EXHAUSTED: test2'); + }); + + it('should fallback on a fallbackable error', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'RESOURCE_EXHAUSTED', message: 'test' }); + }); + + let pmFallbackRequestCount = 0; + ai.defineModel({ name: 'programmableModelFallback' }, async (req) => { + pmFallbackRequestCount++; + return { + message: { role: 'model', content: [{ text: 'fallback success' }] }, + }; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [ + fallback({ + models: ['programmableModelFallback'], + statuses: ['RESOURCE_EXHAUSTED'], + }), + ], + }); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(pmFallbackRequestCount, 1); + assert.strictEqual(result.text, 'fallback success'); + }); + + it('should throw after all fallbacks fail', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'UNAVAILABLE', message: 'test' }); + }); + + let pmFallbackRequestCount = 0; + ai.defineModel({ name: 'programmableModelFallback' }, async (req) => { + pmFallbackRequestCount++; + throw new GenkitError({ status: 'UNAVAILABLE', message: 'test2' }); + }); + + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [ + fallback({ + models: ['programmableModelFallback'], + }), + ], + }), + /UNAVAILABLE: test2/ + ); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(pmFallbackRequestCount, 1); + }); + + it('should not fallback on non-fallbackable error', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'INVALID_ARGUMENT', message: 'test' }); + }); + + let pmFallbackRequestCount = 0; + ai.defineModel({ name: 'programmableModelFallback' }, async (req) => { + pmFallbackRequestCount++; + return { + message: { role: 'model', content: [{ text: 'fallback success' }] }, + }; + }); + + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [ + fallback({ + models: ['programmableModelFallback'], + statuses: ['RESOURCE_EXHAUSTED'], + }), + ], + }), + /INVALID_ARGUMENT: test/ + ); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(pmFallbackRequestCount, 0); + }); +}); diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts new file mode 100644 index 0000000000..e0d1b8b0ca --- /dev/null +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -0,0 +1,142 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import * as fs from 'fs/promises'; +import { afterEach, beforeEach, describe, it } from 'node:test'; +import * as os from 'os'; +import * as path from 'path'; +import { filesystem } from '../src/filesystem.js'; + +describe('filesystem middleware', () => { + let tempDir: string; + let fakeGenerateAPI: any = {}; + + beforeEach(async () => { + tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'genkit-filesystem-test-') + ); + await fs.mkdir(path.join(tempDir, 'sub')); + await fs.writeFile(path.join(tempDir, 'file1.txt'), 'hello world'); + await fs.writeFile(path.join(tempDir, 'sub', 'file2.txt'), 'sub file'); + }); + + afterEach(async () => { + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + it('fails if rootDirectory is not provided', () => { + assert.throws( + () => filesystem.instantiate({} as any, fakeGenerateAPI), + /requires a rootDirectory option/ + ); + }); + + it('injects tools', () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + assert.ok(mw.tools); + assert.strictEqual(mw.tools.length, 2); + assert.strictEqual(mw.tools[0].__action.name, 'list_files'); + assert.strictEqual(mw.tools[1].__action.name, 'read_file'); + }); + + describe('list_files', () => { + it('lists files in root directory', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); + const { result } = await listFiles!.run( + { dirPath: '', recursive: false }, + {} as any + ); + assert.ok(result.includes('file1.txt')); + assert.ok(result.includes('sub')); + assert.ok(!result.includes(path.join('sub', 'file2.txt'))); + }); + + it('lists files recursively', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); + const { result } = await listFiles!.run( + { dirPath: '', recursive: true }, + {} as any + ); + assert.ok(result.includes('file1.txt')); + assert.ok(result.includes('sub')); + assert.ok(result.includes(path.join('sub', 'file2.txt'))); + }); + + it('rejects listing outside root directory', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); + await assert.rejects( + listFiles!.run({ dirPath: '../', recursive: false }, {} as any), + /Access denied/ + ); + }); + }); + + describe('read_file', () => { + it('reads a file in root directory', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); + const { result } = await readFile!.run( + { filePath: 'file1.txt' }, + {} as any + ); + assert.strictEqual(result, 'hello world'); + }); + + it('reads a file in sub directory', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); + const { result } = await readFile!.run( + { filePath: 'sub/file2.txt' }, + {} as any + ); + assert.strictEqual(result, 'sub file'); + }); + + it('rejects reading outside root directory', async () => { + const mw = filesystem.instantiate( + { rootDirectory: tempDir }, + fakeGenerateAPI + ); + const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); + await assert.rejects( + readFile!.run({ filePath: '../etc/passwd' }, {} as any), + /Access denied/ + ); + }); + }); +}); diff --git a/js/plugins/middleware/tests/retry_test.ts b/js/plugins/middleware/tests/retry_test.ts new file mode 100644 index 0000000000..cb722e1eeb --- /dev/null +++ b/js/plugins/middleware/tests/retry_test.ts @@ -0,0 +1,314 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import { GenkitError, genkit } from 'genkit'; +import { describe, it } from 'node:test'; +import { TEST_ONLY, retry } from '../src/retry.js'; + +describe('retry', () => { + it('should not retry on success', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [retry()], + }); + + assert.strictEqual(requestCount, 1); + assert.strictEqual(result.text, 'success'); + }); + + it('should retry on a retryable GenkitError', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 3) { + throw new GenkitError({ status: 'UNAVAILABLE', message: 'test' }); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + TEST_ONLY.setRetryTimeout((callback, ms) => { + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [retry({ maxRetries: 3 })], + }); + + assert.strictEqual(requestCount, 3); + assert.strictEqual(result.text, 'success'); + }); + + it('should retry on a non-GenkitError', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 2) { + throw new Error('generic error'); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + TEST_ONLY.setRetryTimeout((callback, ms) => { + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [retry({ maxRetries: 2 })], + }); + + assert.strictEqual(requestCount, 2); + assert.strictEqual(result.text, 'success'); + }); + + it('should throw after exhausting retries', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'UNAVAILABLE', message: 'test' }); + }); + + TEST_ONLY.setRetryTimeout((callback, ms) => { + callback(); + return 0 as any; + }); + + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [retry({ maxRetries: 2 })], + }), + /UNAVAILABLE: test/ + ); + + assert.strictEqual(requestCount, 3); + }); + + it('should call onError callback', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new Error('test error'); + }); + + TEST_ONLY.setRetryTimeout((callback, ms) => { + callback(); + return 0 as any; + }); + + let errorCount = 0; + let lastError: Error | undefined; + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [ + retry({ + maxRetries: 2, + onError: (err, attempt) => { + errorCount++; + lastError = err; + assert.strictEqual(attempt, errorCount); + }, + }), + ], + }), + /test error/ + ); + + assert.strictEqual(requestCount, 3); + assert.strictEqual(errorCount, 2); + assert.ok(lastError); + assert.strictEqual(lastError!.message, 'test error'); + }); + + it('should not retry on non-retryable status', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + throw new GenkitError({ status: 'INVALID_ARGUMENT', message: 'test' }); + }); + + await assert.rejects( + ai.generate({ + model: pm, + prompt: 'test', + use: [retry({ maxRetries: 2 })], + }), + /INVALID_ARGUMENT: test/ + ); + + assert.strictEqual(requestCount, 1); + }); + + it('should respect initial delay', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 2) { + throw new Error('generic error'); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + let totalDelay = 0; + TEST_ONLY.setRetryTimeout((callback, ms) => { + totalDelay += ms!; + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [retry({ maxRetries: 2, initialDelayMs: 50, noJitter: true })], + }); + + assert.strictEqual(requestCount, 2); + assert.strictEqual(result.text, 'success'); + assert.strictEqual(totalDelay, 50); + }); + + it('should respect backoff factor', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 3) { + throw new Error('generic error'); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + let totalDelay = 0; + TEST_ONLY.setRetryTimeout((callback, ms) => { + totalDelay += ms!; + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [ + retry({ + maxRetries: 3, + initialDelayMs: 20, + backoffFactor: 2, + noJitter: true, + }), + ], + }); + + assert.strictEqual(requestCount, 3); + assert.strictEqual(result.text, 'success'); + assert.strictEqual(totalDelay, 20 + 40); + }); + + it('should apply jitter', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 2) { + throw new Error('generic error'); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + let totalDelay = 0; + TEST_ONLY.setRetryTimeout((callback, ms) => { + totalDelay += ms!; + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [ + retry({ + maxRetries: 2, + initialDelayMs: 50, + noJitter: false, // do jitter + }), + ], + }); + + assert.strictEqual(requestCount, 2); + assert.strictEqual(result.text, 'success'); + assert.ok(totalDelay >= 50); + assert.ok(totalDelay <= 1050); + }); + + it('should respect max delay', async () => { + let requestCount = 0; + const ai = genkit({}); + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + requestCount++; + if (requestCount < 3) { + throw new Error('generic error'); + } + return { message: { role: 'model', content: [{ text: 'success' }] } }; + }); + + let totalDelay = 0; + TEST_ONLY.setRetryTimeout((callback, ms) => { + totalDelay += ms!; + callback(); + return 0 as any; + }); + + const result = await ai.generate({ + model: pm, + prompt: 'test', + use: [ + retry({ + maxRetries: 3, + initialDelayMs: 50, + maxDelayMs: 60, + backoffFactor: 2, + noJitter: true, + }), + ], + }); + + assert.strictEqual(requestCount, 3); + assert.strictEqual(result.text, 'success'); + assert.strictEqual(totalDelay, 50 + 60); + }); +}); diff --git a/js/plugins/middleware/tsconfig.json b/js/plugins/middleware/tsconfig.json new file mode 100644 index 0000000000..596e2cf729 --- /dev/null +++ b/js/plugins/middleware/tsconfig.json @@ -0,0 +1,4 @@ +{ + "extends": "../../tsconfig.json", + "include": ["src"] +} diff --git a/js/plugins/middleware/tsup.config.ts b/js/plugins/middleware/tsup.config.ts new file mode 100644 index 0000000000..c36c8e05fa --- /dev/null +++ b/js/plugins/middleware/tsup.config.ts @@ -0,0 +1,22 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { defineConfig, type Options } from 'tsup'; +import { defaultOptions } from '../../tsup.common'; + +export default defineConfig({ + ...(defaultOptions as Options), +}); diff --git a/js/plugins/middleware/typedoc.json b/js/plugins/middleware/typedoc.json new file mode 100644 index 0000000000..35fed2c958 --- /dev/null +++ b/js/plugins/middleware/typedoc.json @@ -0,0 +1,3 @@ +{ + "entryPoints": ["src/index.ts"] +} From e5182044da4420220d471beeb6ad1f01c8710c65 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Mon, 9 Feb 2026 21:56:48 -0500 Subject: [PATCH 2/9] feat: Update middleware to pass AI context for model resolution and directly invoke fallback models. --- js/plugins/middleware/src/fallback.ts | 25 +++++++++++++----------- js/pnpm-lock.yaml | 28 +++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/js/plugins/middleware/src/fallback.ts b/js/plugins/middleware/src/fallback.ts index 54c1771961..a5ae7ff9a5 100644 --- a/js/plugins/middleware/src/fallback.ts +++ b/js/plugins/middleware/src/fallback.ts @@ -18,12 +18,12 @@ import { GenkitError, ModelReferenceSchema, generateMiddleware, - modelRef, z, type GenerateMiddleware, - type ModelReference, type StatusName, } from 'genkit'; +import { ModelAction } from 'genkit/model'; +import { Registry } from 'genkit/registry'; const DEFAULT_FALLBACK_STATUSES: StatusName[] = [ 'UNAVAILABLE', @@ -78,7 +78,7 @@ export const fallback: GenerateMiddleware = name: 'fallback', configSchema: FallbackOptionsSchema, }, - (options?: FallbackOptions) => { + (options: FallbackOptions | undefined, ai) => { const { models = [], statuses = DEFAULT_FALLBACK_STATUSES, @@ -86,7 +86,7 @@ export const fallback: GenerateMiddleware = } = options || {}; return { - generate: async (req, ctx, next) => { + model: async (req, ctx, next) => { try { return await next(req, ctx); } catch (e) { @@ -97,12 +97,11 @@ export const fallback: GenerateMiddleware = onError?.(e); let lastError: any = e; for (const model of models) { - const normalizedModel = normalizeModel(model); + const normalizedModel = await resolveModel(ai.registry, model); try { - return await next( + return await normalizedModel.model( { ...req, - model: normalizedModel.name, config: normalizedModel.config ?? req.config, }, ctx @@ -128,11 +127,15 @@ export const fallback: GenerateMiddleware = } ); -function normalizeModel( +async function resolveModel( + registry: Registry, model: string | z.infer -): ModelReference { +): Promise<{ model: ModelAction; config?: any }> { if (typeof model === 'string') { - return modelRef({ name: model }); + return { model: await registry.lookupAction(`/model/${model}`) }; } - return modelRef({ name: model.name, config: model.config }); + return { + model: await registry.lookupAction(`/model/${model.name}`), + config: model.config, + }; } diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index b365e14875..ae57b3f734 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -853,6 +853,34 @@ importers: specifier: ^5.3.0 version: 5.8.3 + plugins/middleware: + dependencies: + genkit: + specifier: workspace:^ + version: link:../../genkit + devDependencies: + '@genkit-ai/google-genai': + specifier: workspace:^ + version: link:../google-genai + '@types/node': + specifier: ^20.11.16 + version: 20.19.1 + npm-run-all: + specifier: ^4.1.5 + version: 4.1.5 + rimraf: + specifier: ^6.0.1 + version: 6.0.1 + tsup: + specifier: ^8.0.2 + version: 8.5.0(postcss@8.4.47)(tsx@4.20.3)(typescript@4.9.5)(yaml@2.8.0) + tsx: + specifier: ^4.7.0 + version: 4.20.3 + typescript: + specifier: ^4.9.0 + version: 4.9.5 + plugins/next: devDependencies: '@jest/globals': From f58f2d120a316d8015a01f5463f918c98858191b Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 10 Feb 2026 10:05:32 -0500 Subject: [PATCH 3/9] feat(middleware): return file metadata in list_files output --- js/plugins/middleware/src/filesystem.ts | 13 +++++--- .../middleware/tests/filesystem_test.ts | 33 +++++++++++++++---- 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts index 91315c91d7..6f99cce33e 100644 --- a/js/plugins/middleware/src/filesystem.ts +++ b/js/plugins/middleware/src/filesystem.ts @@ -61,7 +61,7 @@ export const filesystem: GenerateMiddleware = { name: 'list_files', description: - 'Lists files and directories in a given path. Returns a list of strings.', + 'Lists files and directories in a given path. Returns a list of objects with path and type.', inputSchema: z.object({ dirPath: z .string() @@ -72,7 +72,9 @@ export const filesystem: GenerateMiddleware = .describe('Whether to list files recursively.') .default(false), }), - outputSchema: z.array(z.string()), + outputSchema: z.array( + z.object({ path: z.string(), isDirectory: z.boolean() }) + ), }, async (input) => { const targetDir = resolvePath(input.dirPath); @@ -82,11 +84,14 @@ export const filesystem: GenerateMiddleware = recursive: boolean, base: string = '' ) { - const results: string[] = []; + const results: { path: string; isDirectory: boolean }[] = []; const entries = await fs.readdir(dir, { withFileTypes: true }); for (const entry of entries) { const relativePath = path.join(base, entry.name); - results.push(relativePath); + results.push({ + path: relativePath, + isDirectory: entry.isDirectory(), + }); if (entry.isDirectory() && recursive) { const subResults = await list( path.join(dir, entry.name), diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts index e0d1b8b0ca..6942868c47 100644 --- a/js/plugins/middleware/tests/filesystem_test.ts +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -67,9 +67,19 @@ describe('filesystem middleware', () => { { dirPath: '', recursive: false }, {} as any ); - assert.ok(result.includes('file1.txt')); - assert.ok(result.includes('sub')); - assert.ok(!result.includes(path.join('sub', 'file2.txt'))); + assert.ok( + (result as any[]).find( + (r) => r.path === 'file1.txt' && r.isDirectory === false + ) + ); + assert.ok( + (result as any[]).find((r) => r.path === 'sub' && r.isDirectory === true) + ); + assert.ok( + !(result as any[]).find( + (r) => r.path === path.join('sub', 'file2.txt') + ) + ); }); it('lists files recursively', async () => { @@ -82,9 +92,20 @@ describe('filesystem middleware', () => { { dirPath: '', recursive: true }, {} as any ); - assert.ok(result.includes('file1.txt')); - assert.ok(result.includes('sub')); - assert.ok(result.includes(path.join('sub', 'file2.txt'))); + assert.ok( + (result as any[]).find( + (r) => r.path === 'file1.txt' && r.isDirectory === false + ) + ); + assert.ok( + (result as any[]).find((r) => r.path === 'sub' && r.isDirectory === true) + ); + assert.ok( + (result as any[]).find( + (r) => + r.path === path.join('sub', 'file2.txt') && r.isDirectory === false + ) + ); }); it('rejects listing outside root directory', async () => { From d86e65a52565ec90e3937a7719ad5b61ce17d305 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 10 Feb 2026 19:12:03 -0500 Subject: [PATCH 4/9] feat(middleware): gracefully handle errors in filesystem tools Intercept errors thrown by filesystem tools (e.g., `readFile`, `listFiles`) and convert them into user messages injected into the conversation context. This prevents the execution from crashing on errors like ENOENT and allows the model to react to the failure details. - Add middleware logic to catch tool errors and queue them as user messages. - Inject queued error messages into the prompt during the generate phase. - Add robustness tests to verify error handling and message flow. --- js/plugins/middleware/src/filesystem.ts | 35 +++++++++- .../middleware/tests/filesystem_test.ts | 68 +++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts index 6f99cce33e..2076842f8d 100644 --- a/js/plugins/middleware/src/filesystem.ts +++ b/js/plugins/middleware/src/filesystem.ts @@ -15,7 +15,7 @@ */ import * as fs from 'fs/promises'; -import { generateMiddleware, z, type GenerateMiddleware } from 'genkit'; +import { generateMiddleware, MessageData, z, type GenerateMiddleware } from 'genkit'; import { tool } from 'genkit/beta'; import * as path from 'path'; @@ -122,8 +122,41 @@ export const filesystem: GenerateMiddleware = } ); + const filesystemTools = [ + listFiles.__action.name, + readFile.__action.name, + ]; + const messageQueue: MessageData[] = []; + return { tools: [listFiles, readFile], + tool: async (req, ctx, next) => { + try { + return await next(req, ctx); + } catch (e: any) { + if (filesystemTools.includes(req.toolRequest.name)) { + messageQueue.push({ + role: 'user', + content: [ + { + text: `Tool '${req.toolRequest.name}' failed: ${ + e.message || String(e) + }`, + }, + ], + }); + return; + } + throw e; + } + }, + generate: async (req, ctx, next) => { + if (messageQueue.length > 0) { + req.messages.push(...messageQueue); + messageQueue.length = 0; + } + return await next(req, ctx); + }, }; } ); diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts index 6942868c47..5582458f03 100644 --- a/js/plugins/middleware/tests/filesystem_test.ts +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -16,6 +16,7 @@ import * as assert from 'assert'; import * as fs from 'fs/promises'; +import { genkit } from 'genkit'; import { afterEach, beforeEach, describe, it } from 'node:test'; import * as os from 'os'; import * as path from 'path'; @@ -160,4 +161,71 @@ describe('filesystem middleware', () => { ); }); }); + + describe('robustness', () => { + it('should handle tool errors gracefully by injecting user message', async () => { + const ai = genkit({}); + let turn = 0; + const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { + turn++; + if (turn === 1) { + return { + message: { + role: 'model', + content: [ + { + toolRequest: { + name: 'read_file', + ref: '123', + input: { filePath: 'nonexistent' }, + }, + }, + ], + }, + }; + } + return { + message: { role: 'model', content: [{ text: 'done' }] }, + }; + }); + + const result = (await ai.generate({ + model: pm, + prompt: 'start', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const messages = result.messages; + const lastModelIndex = messages.findLastIndex( + (m: any) => m.role === 'model' && m.content[0].toolRequest + ); + const injectedUserIndex = messages.findIndex( + (m: any) => + m.role === 'user' && + m.content[0].text.includes("Tool 'read_file' failed") + ); + + assert.ok( + injectedUserIndex > lastModelIndex, + 'User message should appear after tool request' + ); + + const userMsg = messages[injectedUserIndex]; + assert.match( + userMsg.content[0].text, + /Tool 'read_file' failed: .*ENOENT.*/, + 'Error message should contain underlying error details' + ); + + const roles = messages.map((m: any) => m.role); + assert.deepStrictEqual(roles, ['user', 'model', 'user', 'model']); + + const toolMsg = messages.find((m: any) => m.role === 'tool'); + assert.strictEqual( + toolMsg, + undefined, + 'Tool message should not be present' + ); + }); + }); }); From c8c5adf06e28876faa34afac99aed59f8748b2f4 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Tue, 10 Feb 2026 22:54:00 -0500 Subject: [PATCH 5/9] feat(middleware): add image reading support to filesystem plugin - Update `readFile` tool to detect mime types using the newly added `mime` dependency. - Handle image files by reading them as base64 and injecting them as media parts into the conversation. - Wrap text file content in `` XML tags for clearer context delineation. - Refactor message queue logic to append tool outputs and errors to the current user message turn instead of creating new turns. - Update filesystem example configuration with `maxTurns`. --- js/plugins/middleware/examples/filesystem.ts | 1 + js/plugins/middleware/package.json | 3 + js/plugins/middleware/src/filesystem.ts | 83 +++++-- .../middleware/tests/filesystem_image_test.ts | 101 +++++++++ .../middleware/tests/filesystem_test.ts | 211 ++++++++++-------- js/pnpm-lock.yaml | 10 + 6 files changed, 297 insertions(+), 112 deletions(-) create mode 100644 js/plugins/middleware/tests/filesystem_image_test.ts diff --git a/js/plugins/middleware/examples/filesystem.ts b/js/plugins/middleware/examples/filesystem.ts index 1e8ab20bab..8f234db147 100644 --- a/js/plugins/middleware/examples/filesystem.ts +++ b/js/plugins/middleware/examples/filesystem.ts @@ -32,6 +32,7 @@ async function main() { rootDirectory: process.env.TEMP_DIR || '/tmp', }), ], + maxTurns: 10, }); console.log(text); } diff --git a/js/plugins/middleware/package.json b/js/plugins/middleware/package.json index 8301d1b1a6..5839c8940c 100644 --- a/js/plugins/middleware/package.json +++ b/js/plugins/middleware/package.json @@ -47,5 +47,8 @@ "require": "./lib/index.js", "default": "./lib/index.js" } + }, + "dependencies": { + "mime": "^4.1.0" } } diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts index 2076842f8d..4be41d65a1 100644 --- a/js/plugins/middleware/src/filesystem.ts +++ b/js/plugins/middleware/src/filesystem.ts @@ -15,8 +15,15 @@ */ import * as fs from 'fs/promises'; -import { generateMiddleware, MessageData, z, type GenerateMiddleware } from 'genkit'; +import { + generateMiddleware, + MessageData, + Part, + z, + type GenerateMiddleware, +} from 'genkit'; import { tool } from 'genkit/beta'; +import mime from 'mime'; import * as path from 'path'; export const FilesystemOptionsSchema = z.object({ @@ -57,6 +64,8 @@ export const filesystem: GenerateMiddleware = return p; } + const messageQueue: MessageData[] = []; + const listFiles = tool( { name: 'list_files', @@ -118,33 +127,69 @@ export const filesystem: GenerateMiddleware = }, async (input) => { const targetFile = resolvePath(input.filePath); - return await fs.readFile(targetFile, 'utf8'); + const ext = path.extname(targetFile).toLowerCase(); + const mimeType = mime.getType(ext); + const isImage = mimeType?.startsWith('image/'); + + const parts: Part[] = []; + + if (isImage && mimeType) { + const buffer = await fs.readFile(targetFile); + const base64 = buffer.toString('base64'); + + parts.push({ text: `\n\nread_file media ${input.filePath}` }); + parts.push({ + media: { + url: `data:${mimeType};base64,${base64}`, + contentType: mimeType, + }, + }); + } else { + const content = await fs.readFile(targetFile, 'utf8'); + parts.push({ + text: `\n${content}\n`, + }); + } + + if ( + messageQueue.length > 0 && + messageQueue[messageQueue.length - 1].role === 'user' + ) { + messageQueue[messageQueue.length - 1].content.push(...parts); + } else { + messageQueue.push({ role: 'user', content: parts }); + } + + return `File ${input.filePath} read successfully, see contents below`; } ); - const filesystemTools = [ - listFiles.__action.name, - readFile.__action.name, - ]; - const messageQueue: MessageData[] = []; + const filesystemTools = [listFiles, readFile]; + const filesystemToolNames = filesystemTools.map((t) => t.__action.name); return { - tools: [listFiles, readFile], + tools: filesystemTools, tool: async (req, ctx, next) => { try { return await next(req, ctx); } catch (e: any) { - if (filesystemTools.includes(req.toolRequest.name)) { - messageQueue.push({ - role: 'user', - content: [ - { - text: `Tool '${req.toolRequest.name}' failed: ${ - e.message || String(e) - }`, - }, - ], - }); + if (filesystemToolNames.includes(req.toolRequest.name)) { + const errorPart = { + text: `Tool '${req.toolRequest.name}' failed: ${ + e.message || String(e) + }`, + }; + if ( + messageQueue.length > 0 && + messageQueue[messageQueue.length - 1].role === 'user' + ) { + messageQueue[messageQueue.length - 1].content.push(errorPart); + } else { + messageQueue.push({ + role: 'user', + content: [errorPart], + }); + } return; } throw e; diff --git a/js/plugins/middleware/tests/filesystem_image_test.ts b/js/plugins/middleware/tests/filesystem_image_test.ts new file mode 100644 index 0000000000..d887945d1e --- /dev/null +++ b/js/plugins/middleware/tests/filesystem_image_test.ts @@ -0,0 +1,101 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import * as fs from 'fs/promises'; +import { genkit } from 'genkit'; +import { afterEach, beforeEach, describe, it } from 'node:test'; +import * as os from 'os'; +import * as path from 'path'; +import { filesystem } from '../src/filesystem.js'; + +describe('filesystem middleware image support', () => { + let tempDir: string; + + beforeEach(async () => { + tempDir = await fs.mkdtemp( + path.join(os.tmpdir(), 'genkit-filesystem-test-') + ); + await fs.writeFile(path.join(tempDir, 'image.png'), 'fake image content'); + await fs.writeFile(path.join(tempDir, 'unknown.xyz'), 'unknown content'); + }); + + afterEach(async () => { + await fs.rm(tempDir, { recursive: true, force: true }); + }); + + function createToolModel(ai: any, toolName: string, input: any) { + let turn = 0; + return ai.defineModel( + { name: `pm-${toolName}-${Math.random()}` }, + async () => { + turn++; + if (turn === 1) { + return { + message: { + role: 'model', + content: [{ toolRequest: { name: toolName, input } }], + }, + }; + } + return { message: { role: 'model', content: [{ text: 'done' }] } }; + } + ); + } + + it('reads an image file as media', async () => { + const ai = genkit({}); + const pm = createToolModel(ai, 'read_file', { filePath: 'image.png' }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + assert.match(toolMsg.content[0].toolResponse.output, /read successfully/); + + const userMsg = result.messages.find( + (m: any) => m.role === 'user' && m.content.some((c: any) => c.media) + ); + assert.ok(userMsg); + const mediaPart = userMsg.content.find((c: any) => c.media); + assert.ok(mediaPart); + assert.strictEqual(mediaPart.media.contentType, 'image/png'); + assert.ok(mediaPart.media.url.startsWith('data:image/png;base64,')); + }); + + it('reads unknown file as text', async () => { + const ai = genkit({}); + const pm = createToolModel(ai, 'read_file', { filePath: 'unknown.xyz' }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && + m.content.some((c: any) => c.text && c.text.includes(' c.media)); + }); +}); diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts index 5582458f03..e68fa98b2c 100644 --- a/js/plugins/middleware/tests/filesystem_test.ts +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -39,6 +39,25 @@ describe('filesystem middleware', () => { await fs.rm(tempDir, { recursive: true, force: true }); }); + function createToolModel(ai: any, toolName: string, input: any) { + let turn = 0; + return ai.defineModel( + { name: `pm-${toolName}-${Math.random()}` }, + async () => { + turn++; + if (turn === 1) { + return { + message: { + role: 'model', + content: [{ toolRequest: { name: toolName, input } }], + }, + }; + } + return { message: { role: 'model', content: [{ text: 'done' }] } }; + } + ); + } + it('fails if rootDirectory is not provided', () => { assert.throws( () => filesystem.instantiate({} as any, fakeGenerateAPI), @@ -59,135 +78,141 @@ describe('filesystem middleware', () => { describe('list_files', () => { it('lists files in root directory', async () => { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); - const { result } = await listFiles!.run( - { dirPath: '', recursive: false }, - {} as any - ); + const ai = genkit({}); + const pm = createToolModel(ai, 'list_files', { dirPath: '' }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + const output = toolMsg.content[0].toolResponse.output; assert.ok( - (result as any[]).find( - (r) => r.path === 'file1.txt' && r.isDirectory === false - ) + output.find((r: any) => r.path === 'file1.txt' && !r.isDirectory) ); + assert.ok(output.find((r: any) => r.path === 'sub' && r.isDirectory)); assert.ok( - (result as any[]).find((r) => r.path === 'sub' && r.isDirectory === true) - ); - assert.ok( - !(result as any[]).find( - (r) => r.path === path.join('sub', 'file2.txt') - ) + !output.find((r: any) => r.path === path.join('sub', 'file2.txt')) ); }); it('lists files recursively', async () => { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); - const { result } = await listFiles!.run( - { dirPath: '', recursive: true }, - {} as any - ); - assert.ok( - (result as any[]).find( - (r) => r.path === 'file1.txt' && r.isDirectory === false - ) - ); + const ai = genkit({}); + const pm = createToolModel(ai, 'list_files', { + dirPath: '', + recursive: true, + }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + const output = toolMsg.content[0].toolResponse.output; assert.ok( - (result as any[]).find((r) => r.path === 'sub' && r.isDirectory === true) + output.find((r: any) => r.path === 'file1.txt' && !r.isDirectory) ); + assert.ok(output.find((r: any) => r.path === 'sub' && r.isDirectory)); assert.ok( - (result as any[]).find( - (r) => - r.path === path.join('sub', 'file2.txt') && r.isDirectory === false + output.find( + (r: any) => r.path === path.join('sub', 'file2.txt') && !r.isDirectory ) ); }); it('rejects listing outside root directory', async () => { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const listFiles = mw.tools!.find((t) => t.__action.name === 'list_files'); - await assert.rejects( - listFiles!.run({ dirPath: '../', recursive: false }, {} as any), - /Access denied/ + const ai = genkit({}); + const pm = createToolModel(ai, 'list_files', { dirPath: '../' }); + + // The middleware catches errors and injects user message. + // So verify that user message contains access denied error. + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && m.content[0].text.includes('Access denied') ); + assert.ok(userMsg); }); }); describe('read_file', () => { it('reads a file in root directory', async () => { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); - const { result } = await readFile!.run( - { filePath: 'file1.txt' }, - {} as any + const ai = genkit({}); + const pm = createToolModel(ai, 'read_file', { filePath: 'file1.txt' }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + assert.match(toolMsg.content[0].toolResponse.output, /read successfully/); + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && m.content[0].text.includes(' { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); - const { result } = await readFile!.run( - { filePath: 'sub/file2.txt' }, - {} as any + const ai = genkit({}); + const pm = createToolModel(ai, 'read_file', { + filePath: 'sub/file2.txt', + }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + assert.match(toolMsg.content[0].toolResponse.output, /read successfully/); + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && m.content[0].text.includes(' { - const mw = filesystem.instantiate( - { rootDirectory: tempDir }, - fakeGenerateAPI - ); - const readFile = mw.tools!.find((t) => t.__action.name === 'read_file'); - await assert.rejects( - readFile!.run({ filePath: '../etc/passwd' }, {} as any), - /Access denied/ + const ai = genkit({}); + const pm = createToolModel(ai, 'read_file', { + filePath: '../etc/passwd', + }); + + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && m.content[0].text.includes('Access denied') ); + assert.ok(userMsg); }); }); describe('robustness', () => { it('should handle tool errors gracefully by injecting user message', async () => { const ai = genkit({}); - let turn = 0; - const pm = ai.defineModel({ name: 'programmableModel' }, async (req) => { - turn++; - if (turn === 1) { - return { - message: { - role: 'model', - content: [ - { - toolRequest: { - name: 'read_file', - ref: '123', - input: { filePath: 'nonexistent' }, - }, - }, - ], - }, - }; - } - return { - message: { role: 'model', content: [{ text: 'done' }] }, - }; - }); + const pm = createToolModel(ai, 'read_file', { filePath: 'nonexistent' }); const result = (await ai.generate({ model: pm, diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index ae57b3f734..1fce804630 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -858,6 +858,9 @@ importers: genkit: specifier: workspace:^ version: link:../../genkit + mime: + specifier: ^4.1.0 + version: 4.1.0 devDependencies: '@genkit-ai/google-genai': specifier: workspace:^ @@ -7205,6 +7208,11 @@ packages: engines: {node: '>=10.0.0'} hasBin: true + mime@4.1.0: + resolution: {integrity: sha512-X5ju04+cAzsojXKes0B/S4tcYtFAJ6tTMuSPBEn9CPGlrWr8Fiw7qYeLT0XyH80HSoAoqWCaz+MWKh22P7G1cw==} + engines: {node: '>=16'} + hasBin: true + mimic-fn@2.1.0: resolution: {integrity: sha512-OqbOk5oEQeAZ8WXWydlu9HJjz9WVdEIvamMCcXmuqUYjTknH/sqsWvhQ3vgwKFRR1HpjvNBKQ37nbJgYzGqGcg==} engines: {node: '>=6'} @@ -15251,6 +15259,8 @@ snapshots: mime@3.0.0: optional: true + mime@4.1.0: {} + mimic-fn@2.1.0: {} minimatch@10.0.1: From 2033544b1bf5c006f642b8768704dd9fa4c87a06 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 11 Feb 2026 09:12:52 -0500 Subject: [PATCH 6/9] refactor --- js/plugins/middleware/src/filesystem.ts | 107 +----------------- .../middleware/src/filesystem/list_files.ts | 70 ++++++++++++ .../middleware/src/filesystem/read_file.ts | 74 ++++++++++++ 3 files changed, 149 insertions(+), 102 deletions(-) create mode 100644 js/plugins/middleware/src/filesystem/list_files.ts create mode 100644 js/plugins/middleware/src/filesystem/read_file.ts diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts index 4be41d65a1..724e9f49a5 100644 --- a/js/plugins/middleware/src/filesystem.ts +++ b/js/plugins/middleware/src/filesystem.ts @@ -14,17 +14,15 @@ * limitations under the License. */ -import * as fs from 'fs/promises'; import { generateMiddleware, MessageData, - Part, z, type GenerateMiddleware, } from 'genkit'; -import { tool } from 'genkit/beta'; -import mime from 'mime'; import * as path from 'path'; +import { defineListFileTool } from './filesystem/list_files'; +import { defineReadFileTool } from './filesystem/read_file'; export const FilesystemOptionsSchema = z.object({ rootDirectory: z @@ -66,105 +64,10 @@ export const filesystem: GenerateMiddleware = const messageQueue: MessageData[] = []; - const listFiles = tool( - { - name: 'list_files', - description: - 'Lists files and directories in a given path. Returns a list of objects with path and type.', - inputSchema: z.object({ - dirPath: z - .string() - .describe('Directory path relative to root.') - .default(''), - recursive: z - .boolean() - .describe('Whether to list files recursively.') - .default(false), - }), - outputSchema: z.array( - z.object({ path: z.string(), isDirectory: z.boolean() }) - ), - }, - async (input) => { - const targetDir = resolvePath(input.dirPath); - - async function list( - dir: string, - recursive: boolean, - base: string = '' - ) { - const results: { path: string; isDirectory: boolean }[] = []; - const entries = await fs.readdir(dir, { withFileTypes: true }); - for (const entry of entries) { - const relativePath = path.join(base, entry.name); - results.push({ - path: relativePath, - isDirectory: entry.isDirectory(), - }); - if (entry.isDirectory() && recursive) { - const subResults = await list( - path.join(dir, entry.name), - true, - relativePath - ); - results.push(...subResults); - } - } - return results; - } - return await list(targetDir, input.recursive); - } - ); - - const readFile = tool( - { - name: 'read_file', - description: 'Reads the contents of a file', - inputSchema: z.object({ - filePath: z.string().describe('File path relative to root.'), - }), - outputSchema: z.string(), - }, - async (input) => { - const targetFile = resolvePath(input.filePath); - const ext = path.extname(targetFile).toLowerCase(); - const mimeType = mime.getType(ext); - const isImage = mimeType?.startsWith('image/'); - - const parts: Part[] = []; - - if (isImage && mimeType) { - const buffer = await fs.readFile(targetFile); - const base64 = buffer.toString('base64'); - - parts.push({ text: `\n\nread_file media ${input.filePath}` }); - parts.push({ - media: { - url: `data:${mimeType};base64,${base64}`, - contentType: mimeType, - }, - }); - } else { - const content = await fs.readFile(targetFile, 'utf8'); - parts.push({ - text: `\n${content}\n`, - }); - } - - if ( - messageQueue.length > 0 && - messageQueue[messageQueue.length - 1].role === 'user' - ) { - messageQueue[messageQueue.length - 1].content.push(...parts); - } else { - messageQueue.push({ role: 'user', content: parts }); - } - - return `File ${input.filePath} read successfully, see contents below`; - } - ); + const listFilesTool = defineListFileTool(resolvePath); + const readFileTool = defineReadFileTool(messageQueue, resolvePath); - const filesystemTools = [listFiles, readFile]; + const filesystemTools = [listFilesTool, readFileTool]; const filesystemToolNames = filesystemTools.map((t) => t.__action.name); return { diff --git a/js/plugins/middleware/src/filesystem/list_files.ts b/js/plugins/middleware/src/filesystem/list_files.ts new file mode 100644 index 0000000000..c697ba693b --- /dev/null +++ b/js/plugins/middleware/src/filesystem/list_files.ts @@ -0,0 +1,70 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as fs from 'fs/promises'; +import { z } from 'genkit'; +import { tool } from 'genkit/beta'; +import * as path from 'path'; + +export function defineListFileTool( + resolvePath: (requestedPath: string) => string +) { + return tool( + { + name: 'list_files', + description: + 'Lists files and directories in a given path. Returns a list of objects with path and type.', + inputSchema: z.object({ + dirPath: z + .string() + .describe('Directory path relative to root.') + .default(''), + recursive: z + .boolean() + .describe('Whether to list files recursively.') + .default(false), + }), + outputSchema: z.array( + z.object({ path: z.string(), isDirectory: z.boolean() }) + ), + }, + async (input) => { + const targetDir = resolvePath(input.dirPath); + + async function list(dir: string, recursive: boolean, base: string = '') { + const results: { path: string; isDirectory: boolean }[] = []; + const entries = await fs.readdir(dir, { withFileTypes: true }); + for (const entry of entries) { + const relativePath = path.join(base, entry.name); + results.push({ + path: relativePath, + isDirectory: entry.isDirectory(), + }); + if (entry.isDirectory() && recursive) { + const subResults = await list( + path.join(dir, entry.name), + true, + relativePath + ); + results.push(...subResults); + } + } + return results; + } + return await list(targetDir, input.recursive); + } + ); +} diff --git a/js/plugins/middleware/src/filesystem/read_file.ts b/js/plugins/middleware/src/filesystem/read_file.ts new file mode 100644 index 0000000000..e99a2be451 --- /dev/null +++ b/js/plugins/middleware/src/filesystem/read_file.ts @@ -0,0 +1,74 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as fs from 'fs/promises'; +import { MessageData, Part, z } from 'genkit'; +import { tool } from 'genkit/beta'; +import mime from 'mime'; +import * as path from 'path'; + +export function defineReadFileTool( + messageQueue: MessageData[], + resolvePath: (requestedPath: string) => string +) { + return tool( + { + name: 'read_file', + description: 'Reads the contents of a file', + inputSchema: z.object({ + filePath: z.string().describe('File path relative to root.'), + }), + outputSchema: z.string(), + }, + async (input) => { + const targetFile = resolvePath(input.filePath); + const ext = path.extname(targetFile).toLowerCase(); + const mimeType = mime.getType(ext); + const isImage = mimeType?.startsWith('image/'); + + const parts: Part[] = []; + + if (isImage && mimeType) { + const buffer = await fs.readFile(targetFile); + const base64 = buffer.toString('base64'); + + parts.push({ text: `\n\nread_file result ${mimeType} ${input.filePath}` }); + parts.push({ + media: { + url: `data:${mimeType};base64,${base64}`, + contentType: mimeType, + }, + }); + } else { + const content = await fs.readFile(targetFile, 'utf8'); + parts.push({ + text: `\n${content}\n`, + }); + } + + if ( + messageQueue.length > 0 && + messageQueue[messageQueue.length - 1].role === 'user' + ) { + messageQueue[messageQueue.length - 1].content.push(...parts); + } else { + messageQueue.push({ role: 'user', content: parts }); + } + + return `File ${input.filePath} read successfully, see contents below`; + } + ); +} From 93e0d541b983555170676ca10be39fb3f09f46f6 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 11 Feb 2026 09:41:45 -0500 Subject: [PATCH 7/9] feat(middleware): add write and search/replace tools to filesystem --- js/plugins/middleware/src/filesystem.ts | 15 +- .../src/filesystem/search_and_replace.ts | 122 ++++++++ .../middleware/src/filesystem/write_file.ts | 42 +++ .../middleware/tests/filesystem_test.ts | 266 +++++++++++++++++- 4 files changed, 441 insertions(+), 4 deletions(-) create mode 100644 js/plugins/middleware/src/filesystem/search_and_replace.ts create mode 100644 js/plugins/middleware/src/filesystem/write_file.ts diff --git a/js/plugins/middleware/src/filesystem.ts b/js/plugins/middleware/src/filesystem.ts index 724e9f49a5..8a8fdc141c 100644 --- a/js/plugins/middleware/src/filesystem.ts +++ b/js/plugins/middleware/src/filesystem.ts @@ -23,6 +23,8 @@ import { import * as path from 'path'; import { defineListFileTool } from './filesystem/list_files'; import { defineReadFileTool } from './filesystem/read_file'; +import { defineSearchAndReplaceTool } from './filesystem/search_and_replace'; +import { defineWriteFileTool } from './filesystem/write_file'; export const FilesystemOptionsSchema = z.object({ rootDirectory: z @@ -35,8 +37,8 @@ export const FilesystemOptionsSchema = z.object({ export type FilesystemOptions = z.infer; /** - * Creates a middleware that grants the LLM basic readonly access to the filesystem. - * Injects `list_files` and `read_file` tools restricted to the provided `rootDirectory`. + * Creates a middleware that grants the LLM access to the filesystem. + * Injects `list_files`, `read_file`, `write_file`, and `search_and_replace` tools restricted to the provided `rootDirectory`. */ export const filesystem: GenerateMiddleware = generateMiddleware( @@ -66,8 +68,15 @@ export const filesystem: GenerateMiddleware = const listFilesTool = defineListFileTool(resolvePath); const readFileTool = defineReadFileTool(messageQueue, resolvePath); + const writeFileTool = defineWriteFileTool(resolvePath); + const searchAndReplaceTool = defineSearchAndReplaceTool(resolvePath); - const filesystemTools = [listFilesTool, readFileTool]; + const filesystemTools = [ + listFilesTool, + readFileTool, + writeFileTool, + searchAndReplaceTool, + ]; const filesystemToolNames = filesystemTools.map((t) => t.__action.name); return { diff --git a/js/plugins/middleware/src/filesystem/search_and_replace.ts b/js/plugins/middleware/src/filesystem/search_and_replace.ts new file mode 100644 index 0000000000..dc31feeb8c --- /dev/null +++ b/js/plugins/middleware/src/filesystem/search_and_replace.ts @@ -0,0 +1,122 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as fs from 'fs/promises'; +import { z } from 'genkit'; +import { tool } from 'genkit/beta'; + +export function defineSearchAndReplaceTool( + resolvePath: (requestedPath: string) => string +) { + return tool( + { + name: 'search_and_replace', + description: + 'Replaces text in a file using search and replace blocks. ', + inputSchema: z.object({ + filePath: z.string().describe('File path relative to root.'), + edits: z.array( + z + .string() + .describe( + 'A search and replace block string in the format:\n' + + '<<<<<<< SEARCH\n[search content]\n=======\n[replace content]\n>>>>>>> REPLACE' + ) + ), + }), + outputSchema: z.string(), + }, + async (input) => { + const targetFile = resolvePath(input.filePath); + let content = await fs.readFile(targetFile, 'utf8'); + + for (const editBlock of input.edits) { + const startMarker = '<<<<<<< SEARCH\n'; + const endMarker = '\n>>>>>>> REPLACE'; + const separator = '\n=======\n'; + + if ( + !editBlock.startsWith(startMarker) || + !editBlock.endsWith(endMarker) + ) { + throw new Error( + 'Invalid edit block format. Block must start with "<<<<<<< SEARCH\\n" and end with "\\n>>>>>>> REPLACE"' + ); + } + + const innerContent = editBlock.substring( + startMarker.length, + editBlock.length - endMarker.length + ); + + // Find all possible separator positions + const separatorIndices: number[] = []; + let pos = innerContent.indexOf(separator); + while (pos !== -1) { + separatorIndices.push(pos); + pos = innerContent.indexOf(separator, pos + 1); + } + + if (separatorIndices.length === 0) { + throw new Error( + 'Invalid edit block format. Missing separator "\\n=======\\n"' + ); + } + + // Try to find a split that matches the content + let match: { search: string; replace: string } | null = null; + let matchCount = 0; + + for (const splitIndex of separatorIndices) { + const search = innerContent.substring(0, splitIndex); + const replace = innerContent.substring( + splitIndex + separator.length + ); + + if (content.includes(search)) { + // If we already have a match, only replace it if this one is longer (more specific) + if (!match || search.length > match.search.length) { + match = { search, replace }; + } + matchCount++; + } + } + + if (matchCount === 0) { + throw new Error( + `Search content not found in file ${input.filePath}. ` + + `Make sure the search block matches the file content exactly, ` + + `including whitespace and indentation.` + ); + } + + if (matchCount > 1) { + // If multiple splits match, prefer the one with the longest search string. + // This assumes that a longer match is more specific and likely what the user intended + // if they included markers in their search block. + } + + // Apply the replacement (first occurrence only) + if (match) { + content = content.replace(match.search, match.replace); + } + } + + await fs.writeFile(targetFile, content, 'utf8'); + return `Successfully applied ${input.edits.length} edit(s) to ${input.filePath}.`; + } + ); +} diff --git a/js/plugins/middleware/src/filesystem/write_file.ts b/js/plugins/middleware/src/filesystem/write_file.ts new file mode 100644 index 0000000000..31ec7b2c7d --- /dev/null +++ b/js/plugins/middleware/src/filesystem/write_file.ts @@ -0,0 +1,42 @@ +/** + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as fs from 'fs/promises'; +import { z } from 'genkit'; +import { tool } from 'genkit/beta'; +import * as path from 'path'; + +export function defineWriteFileTool( + resolvePath: (requestedPath: string) => string +) { + return tool( + { + name: 'write_file', + description: 'Writes content to a file, overwriting it if it exists.', + inputSchema: z.object({ + filePath: z.string().describe('File path relative to root.'), + content: z.string().describe('Content to write to the file.'), + }), + outputSchema: z.string(), + }, + async (input) => { + const targetFile = resolvePath(input.filePath); + await fs.mkdir(path.dirname(targetFile), { recursive: true }); + await fs.writeFile(targetFile, input.content, 'utf8'); + return `File ${input.filePath} written successfully.`; + } + ); +} diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts index e68fa98b2c..ff35bc7b79 100644 --- a/js/plugins/middleware/tests/filesystem_test.ts +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -71,9 +71,11 @@ describe('filesystem middleware', () => { fakeGenerateAPI ); assert.ok(mw.tools); - assert.strictEqual(mw.tools.length, 2); + assert.strictEqual(mw.tools.length, 4); assert.strictEqual(mw.tools[0].__action.name, 'list_files'); assert.strictEqual(mw.tools[1].__action.name, 'read_file'); + assert.strictEqual(mw.tools[2].__action.name, 'write_file'); + assert.strictEqual(mw.tools[3].__action.name, 'search_and_replace'); }); describe('list_files', () => { @@ -209,6 +211,268 @@ describe('filesystem middleware', () => { }); }); + describe('write_file', () => { + it('writes a new file', async () => { + const ai = genkit({}); + const pm = createToolModel(ai, 'write_file', { + filePath: 'new.txt', + content: 'new content', + }); + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg); + assert.match( + toolMsg.content[0].toolResponse.output, + /written successfully/ + ); + + const content = await fs.readFile(path.join(tempDir, 'new.txt'), 'utf8'); + assert.strictEqual(content, 'new content'); + }); + + it('creates directories if needed', async () => { + const ai = genkit({}); + const pm = createToolModel(ai, 'write_file', { + filePath: 'deep/nested/file.txt', + content: 'nested content', + }); + await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + }); + + const content = await fs.readFile( + path.join(tempDir, 'deep/nested/file.txt'), + 'utf8' + ); + assert.strictEqual(content, 'nested content'); + }); + }); + + describe('search_and_replace', () => { + it('replaces content', async () => { + const ai = genkit({}); + const editBlock = `<<<<<<< SEARCH +hello world +======= +hello universe +>>>>>>> REPLACE`; + const pm = createToolModel(ai, 'search_and_replace', { + filePath: 'file1.txt', + edits: [editBlock], + }); + + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + if (!toolMsg) { + const errorMsg = result.messages.find( + (m: any) => + m.role === 'user' && m.content[0].text.includes('failed') + ); + if (errorMsg) { + throw new Error( + `Tool failed unexpectedly: ${errorMsg.content[0].text}` + ); + } + } + assert.ok(toolMsg); + assert.match( + toolMsg.content[0].toolResponse.output, + /Successfully applied/ + ); + + const content = await fs.readFile( + path.join(tempDir, 'file1.txt'), + 'utf8' + ); + assert.strictEqual(content, 'hello universe'); + }); + + it('fails if search content not found', async () => { + const ai = genkit({}); + const editBlock = `<<<<<<< SEARCH +nonexistent +======= +replace +>>>>>>> REPLACE`; + const pm = createToolModel(ai, 'search_and_replace', { + filePath: 'file1.txt', + edits: [editBlock], + }); + + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const userMsg = result.messages.find( + (m: any) => + m.role === 'user' && + m.content[0].text.includes('Search content not found') + ); + if (!userMsg) { + console.log( + 'Messages received:', + JSON.stringify(result.messages, null, 2) + ); + } + assert.ok(userMsg); + }); + + it('handles tricky search/replace cases', async () => { + const cases = [ + { + name: 'marker in search', + initial: 'line1\n=======\nline2', + block: `<<<<<<< SEARCH +line1 +======= +line2 +======= +replacement +>>>>>>> REPLACE`, + expected: 'replacement', + }, + { + name: 'marker in replace', + initial: 'original', + block: `<<<<<<< SEARCH +original +======= +new +======= +line +>>>>>>> REPLACE`, + expected: 'new\n=======\nline', + }, + { + name: 'start marker in search', + initial: '<<<<<<< SEARCH\ncontent', + block: `<<<<<<< SEARCH +<<<<<<< SEARCH +content +======= +replaced +>>>>>>> REPLACE`, + expected: 'replaced', + }, + { + name: 'start marker in replace', + initial: 'content', + block: `<<<<<<< SEARCH +content +======= +<<<<<<< SEARCH +new +>>>>>>> REPLACE`, + expected: '<<<<<<< SEARCH\nnew', + }, + { + name: 'end marker in search', + initial: 'content\n>>>>>>> REPLACE', + block: `<<<<<<< SEARCH +content +>>>>>>> REPLACE +======= +replaced +>>>>>>> REPLACE`, + expected: 'replaced', + }, + { + name: 'end marker in replace', + initial: 'content', + block: `<<<<<<< SEARCH +content +======= +new +>>>>>>> REPLACE +>>>>>>> REPLACE`, + expected: 'new\n>>>>>>> REPLACE', + }, + { + name: 'multiple markers greedy search', + initial: 'part1\n=======\npart2', + block: `<<<<<<< SEARCH +part1 +======= +part2 +======= +replacement +>>>>>>> REPLACE`, + expected: 'replacement', + }, + { + name: 'ambiguous separators preferring longest match', + initial: 'A\n=======\nB', + // search: A\n=======\nB -> replace: C\n=======\nD + // block structure: A = B = C = D (where = is separator) + // splits: + // 1. S=A, R=B=C=D. (Match A? Yes) + // 2. S=A=B, R=C=D. (Match A=B? Yes) + // 3. S=A=B=C, R=D. (Match A=B=C? No) + // Winner: 2. + block: `<<<<<<< SEARCH +A +======= +B +======= +C +======= +D +>>>>>>> REPLACE`, + expected: 'C\n=======\nD', + }, + ]; + + for (const c of cases) { + const ai = genkit({}); + const filename = `tricky-${c.name.replace(/\s+/g, '-')}.txt`; + await fs.writeFile(path.join(tempDir, filename), c.initial); + + const pm = createToolModel(ai, 'search_and_replace', { + filePath: filename, + edits: [c.block], + }); + + const result = (await ai.generate({ + model: pm, + prompt: 'test', + use: [filesystem({ rootDirectory: tempDir })], + })) as any; + + const toolMsg = result.messages.find((m: any) => m.role === 'tool'); + assert.ok(toolMsg, `Tool execution failed for case: ${c.name}`); + assert.match( + toolMsg.content[0].toolResponse.output, + /Successfully applied/, + `Tool output mismatch for case: ${c.name}` + ); + + const newContent = await fs.readFile( + path.join(tempDir, filename), + 'utf8' + ); + assert.strictEqual( + newContent, + c.expected, + `Content mismatch for case: ${c.name}` + ); + } + }); + }); + describe('robustness', () => { it('should handle tool errors gracefully by injecting user message', async () => { const ai = genkit({}); From 22fbe937d3d8c2a0a84b0cb4b0218ea2d7450fa0 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 11 Feb 2026 10:33:13 -0500 Subject: [PATCH 8/9] fix mime --- js/plugins/middleware/package.json | 3 +- .../middleware/src/filesystem/list_files.ts | 4 +-- .../middleware/src/filesystem/read_file.ts | 10 +++--- .../src/filesystem/search_and_replace.ts | 4 +-- .../middleware/src/filesystem/write_file.ts | 4 +-- js/pnpm-lock.yaml | 35 ++++++++++++------- 6 files changed, 35 insertions(+), 25 deletions(-) diff --git a/js/plugins/middleware/package.json b/js/plugins/middleware/package.json index 5839c8940c..059f38ebd7 100644 --- a/js/plugins/middleware/package.json +++ b/js/plugins/middleware/package.json @@ -32,6 +32,7 @@ }, "devDependencies": { "@genkit-ai/google-genai": "workspace:^", + "@types/mime-types": "^3.0.1", "@types/node": "^20.11.16", "npm-run-all": "^4.1.5", "rimraf": "^6.0.1", @@ -49,6 +50,6 @@ } }, "dependencies": { - "mime": "^4.1.0" + "mime-types": "^3.0.2" } } diff --git a/js/plugins/middleware/src/filesystem/list_files.ts b/js/plugins/middleware/src/filesystem/list_files.ts index c697ba693b..9ea1376406 100644 --- a/js/plugins/middleware/src/filesystem/list_files.ts +++ b/js/plugins/middleware/src/filesystem/list_files.ts @@ -15,13 +15,13 @@ */ import * as fs from 'fs/promises'; -import { z } from 'genkit'; +import { ToolAction, z } from 'genkit'; import { tool } from 'genkit/beta'; import * as path from 'path'; export function defineListFileTool( resolvePath: (requestedPath: string) => string -) { +): ToolAction { return tool( { name: 'list_files', diff --git a/js/plugins/middleware/src/filesystem/read_file.ts b/js/plugins/middleware/src/filesystem/read_file.ts index e99a2be451..011a93b3c0 100644 --- a/js/plugins/middleware/src/filesystem/read_file.ts +++ b/js/plugins/middleware/src/filesystem/read_file.ts @@ -15,15 +15,15 @@ */ import * as fs from 'fs/promises'; -import { MessageData, Part, z } from 'genkit'; +import { MessageData, Part, ToolAction, z } from 'genkit'; import { tool } from 'genkit/beta'; -import mime from 'mime'; +import * as mime from 'mime-types'; import * as path from 'path'; export function defineReadFileTool( messageQueue: MessageData[], resolvePath: (requestedPath: string) => string -) { +): ToolAction { return tool( { name: 'read_file', @@ -36,8 +36,8 @@ export function defineReadFileTool( async (input) => { const targetFile = resolvePath(input.filePath); const ext = path.extname(targetFile).toLowerCase(); - const mimeType = mime.getType(ext); - const isImage = mimeType?.startsWith('image/'); + const mimeType = mime.lookup(ext); + const isImage = mimeType != false && mimeType?.startsWith('image/'); const parts: Part[] = []; diff --git a/js/plugins/middleware/src/filesystem/search_and_replace.ts b/js/plugins/middleware/src/filesystem/search_and_replace.ts index dc31feeb8c..952cbcc31a 100644 --- a/js/plugins/middleware/src/filesystem/search_and_replace.ts +++ b/js/plugins/middleware/src/filesystem/search_and_replace.ts @@ -15,12 +15,12 @@ */ import * as fs from 'fs/promises'; -import { z } from 'genkit'; +import { ToolAction, z } from 'genkit'; import { tool } from 'genkit/beta'; export function defineSearchAndReplaceTool( resolvePath: (requestedPath: string) => string -) { +): ToolAction { return tool( { name: 'search_and_replace', diff --git a/js/plugins/middleware/src/filesystem/write_file.ts b/js/plugins/middleware/src/filesystem/write_file.ts index 31ec7b2c7d..886467a8b4 100644 --- a/js/plugins/middleware/src/filesystem/write_file.ts +++ b/js/plugins/middleware/src/filesystem/write_file.ts @@ -15,13 +15,13 @@ */ import * as fs from 'fs/promises'; -import { z } from 'genkit'; +import { ToolAction, z } from 'genkit'; import { tool } from 'genkit/beta'; import * as path from 'path'; export function defineWriteFileTool( resolvePath: (requestedPath: string) => string -) { +): ToolAction { return tool( { name: 'write_file', diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index 1fce804630..170041cfe9 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -858,13 +858,16 @@ importers: genkit: specifier: workspace:^ version: link:../../genkit - mime: - specifier: ^4.1.0 - version: 4.1.0 + mime-types: + specifier: ^3.0.2 + version: 3.0.2 devDependencies: '@genkit-ai/google-genai': specifier: workspace:^ version: link:../google-genai + '@types/mime-types': + specifier: ^3.0.1 + version: 3.0.1 '@types/node': specifier: ^20.11.16 version: 20.19.1 @@ -4437,6 +4440,9 @@ packages: '@types/memcached@2.2.10': resolution: {integrity: sha512-AM9smvZN55Gzs2wRrqeMHVP7KE8KWgCJO/XL5yCly2xF6EKa4YlbpK+cLSAH4NG/Ah64HrlegmGqW8kYws7Vxg==} + '@types/mime-types@3.0.1': + resolution: {integrity: sha512-xRMsfuQbnRq1Ef+C+RKaENOxXX87Ygl38W1vDfPHRku02TgQr+Qd8iivLtAMcR0KF5/29xlnFihkTlbqFrGOVQ==} + '@types/mime@1.3.5': resolution: {integrity: sha512-/pyBZWSLD2n0dcHE3hq8s8ZvcETHtEuF+3E7XVt0Ig2nvsVQXdghHVcEkIWjy9A0wKfTn97a/PSDYohKIlnP/w==} @@ -7193,6 +7199,10 @@ packages: resolution: {integrity: sha512-xRc4oEhT6eaBpU1XF7AjpOFD+xQmXNB5OVKwp4tqCuBpHLS/ZbBDrc07mYTDqVMg6PfxUjjNp85O6Cd2Z/5HWA==} engines: {node: '>= 0.6'} + mime-types@3.0.2: + resolution: {integrity: sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==} + engines: {node: '>=18'} + mime@1.6.0: resolution: {integrity: sha512-x0Vn8spI+wuJ1O6S7gnbaQg8Pxh4NNHb7KSINmEWKiPE4RKOplvijn+NkmYmmRgP68mc70j2EbeTFRsrswaQeg==} engines: {node: '>=4'} @@ -7208,11 +7218,6 @@ packages: engines: {node: '>=10.0.0'} hasBin: true - mime@4.1.0: - resolution: {integrity: sha512-X5ju04+cAzsojXKes0B/S4tcYtFAJ6tTMuSPBEn9CPGlrWr8Fiw7qYeLT0XyH80HSoAoqWCaz+MWKh22P7G1cw==} - engines: {node: '>=16'} - hasBin: true - mimic-fn@2.1.0: resolution: {integrity: sha512-OqbOk5oEQeAZ8WXWydlu9HJjz9WVdEIvamMCcXmuqUYjTknH/sqsWvhQ3vgwKFRR1HpjvNBKQ37nbJgYzGqGcg==} engines: {node: '>=6'} @@ -11765,6 +11770,8 @@ snapshots: dependencies: '@types/node': 20.19.1 + '@types/mime-types@3.0.1': {} + '@types/mime@1.3.5': {} '@types/mime@3.0.4': {} @@ -11897,7 +11904,7 @@ snapshots: accepts@2.0.0: dependencies: - mime-types: 3.0.1 + mime-types: 3.0.2 negotiator: 1.0.0 acorn-import-attributes@1.9.5(acorn@8.11.3): @@ -15252,6 +15259,10 @@ snapshots: dependencies: mime-db: 1.54.0 + mime-types@3.0.2: + dependencies: + mime-db: 1.54.0 + mime@1.6.0: {} mime@2.6.0: {} @@ -15259,8 +15270,6 @@ snapshots: mime@3.0.0: optional: true - mime@4.1.0: {} - mimic-fn@2.1.0: {} minimatch@10.0.1: @@ -16318,7 +16327,7 @@ snapshots: etag: 1.8.1 fresh: 2.0.0 http-errors: 2.0.0 - mime-types: 3.0.1 + mime-types: 3.0.2 ms: 2.1.3 on-finished: 2.4.1 range-parser: 1.2.1 @@ -17018,7 +17027,7 @@ snapshots: dependencies: content-type: 1.0.5 media-typer: 1.1.0 - mime-types: 3.0.1 + mime-types: 3.0.2 typed-array-buffer@1.0.3: dependencies: From 4f80aaa57d06cfda15a558f9476d9ed9f98eb8cb Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Wed, 11 Feb 2026 12:28:33 -0500 Subject: [PATCH 9/9] fmt --- js/plugins/middleware/src/filesystem/read_file.ts | 4 +++- js/plugins/middleware/src/filesystem/search_and_replace.ts | 7 ++----- js/plugins/middleware/tests/filesystem_test.ts | 3 +-- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/js/plugins/middleware/src/filesystem/read_file.ts b/js/plugins/middleware/src/filesystem/read_file.ts index 011a93b3c0..4bef98ac2d 100644 --- a/js/plugins/middleware/src/filesystem/read_file.ts +++ b/js/plugins/middleware/src/filesystem/read_file.ts @@ -45,7 +45,9 @@ export function defineReadFileTool( const buffer = await fs.readFile(targetFile); const base64 = buffer.toString('base64'); - parts.push({ text: `\n\nread_file result ${mimeType} ${input.filePath}` }); + parts.push({ + text: `\n\nread_file result ${mimeType} ${input.filePath}`, + }); parts.push({ media: { url: `data:${mimeType};base64,${base64}`, diff --git a/js/plugins/middleware/src/filesystem/search_and_replace.ts b/js/plugins/middleware/src/filesystem/search_and_replace.ts index 952cbcc31a..a64eb826a6 100644 --- a/js/plugins/middleware/src/filesystem/search_and_replace.ts +++ b/js/plugins/middleware/src/filesystem/search_and_replace.ts @@ -24,8 +24,7 @@ export function defineSearchAndReplaceTool( return tool( { name: 'search_and_replace', - description: - 'Replaces text in a file using search and replace blocks. ', + description: 'Replaces text in a file using search and replace blocks. ', inputSchema: z.object({ filePath: z.string().describe('File path relative to root.'), edits: z.array( @@ -82,9 +81,7 @@ export function defineSearchAndReplaceTool( for (const splitIndex of separatorIndices) { const search = innerContent.substring(0, splitIndex); - const replace = innerContent.substring( - splitIndex + separator.length - ); + const replace = innerContent.substring(splitIndex + separator.length); if (content.includes(search)) { // If we already have a match, only replace it if this one is longer (more specific) diff --git a/js/plugins/middleware/tests/filesystem_test.ts b/js/plugins/middleware/tests/filesystem_test.ts index ff35bc7b79..22839bf38d 100644 --- a/js/plugins/middleware/tests/filesystem_test.ts +++ b/js/plugins/middleware/tests/filesystem_test.ts @@ -277,8 +277,7 @@ hello universe const toolMsg = result.messages.find((m: any) => m.role === 'tool'); if (!toolMsg) { const errorMsg = result.messages.find( - (m: any) => - m.role === 'user' && m.content[0].text.includes('failed') + (m: any) => m.role === 'user' && m.content[0].text.includes('failed') ); if (errorMsg) { throw new Error(