Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/config/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ const stdioServerSchema = z.object({
maxRestarts: z.number().int().min(0).default(5),
/** Initial backoff in ms, doubles up to 30s. */
restartBackoffMs: z.number().int().min(0).default(500),
/** Per-upstream stderr log budget in bytes per minute (0 = unlimited). */
stderrLogBytesPerMinute: z.number().int().min(0).default(10_000),
/** Per-server description-shrink override. */
shrink: z
.object({
Expand Down Expand Up @@ -144,6 +146,8 @@ export const tooltrimConfigSchema = z.object({
observability: observabilitySchema,
policy: policySchema,
logLevel: z.enum(["trace", "debug", "info", "warn", "error", "fatal", "silent"]).default("info"),
/** Timeout in ms for upstream listTools calls (default 30s). */
upstreamTimeoutMs: z.number().int().min(1000).default(30_000),
});

export type StdioServerConfig = z.infer<typeof stdioServerSchema>;
Expand Down
19 changes: 18 additions & 1 deletion src/core/aggregator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,9 @@ export class Aggregator {
private readonly promptRoute = new Map<string, { upstreamId: string; original: string }>();
/** uri -> upstreamId for resources (URIs are unique upstream-side; we keep first wins on collision). */
private readonly resourceRoute = new Map<string, string>();
/** Short-lived cache for collectTools() to debounce redundant fan-out. */
private toolsCache: { tools: unknown[]; ts: number } | null = null;
private readonly TOOLS_CACHE_TTL_MS = 2000;

constructor(deps: AggregatorDeps) {
this.deps = deps;
Expand Down Expand Up @@ -210,12 +213,24 @@ export class Aggregator {
* and return the merged list.
*/
async collectTools(): Promise<unknown[]> {
// Return cached result if fresh enough to avoid redundant upstream fan-out.
if (this.toolsCache && Date.now() - this.toolsCache.ts < this.TOOLS_CACHE_TTL_MS) {
return this.toolsCache.tools;
}

this.toolRoute.clear();
const out: Record<string, unknown>[] = [];
const timeoutMs = this.deps.cfg.upstreamTimeoutMs ?? 30_000;

for (const [id, conn] of this.deps.upstream.connections) {
if (conn.status !== "connected" || !conn.capabilities?.tools) continue;
let timer: ReturnType<typeof setTimeout> | undefined;
try {
const result = await conn.client.listTools();
const timeout = new Promise<never>((_, reject) => {
timer = setTimeout(() => reject(new Error(`upstream "${id}" tools/list timed out after ${timeoutMs}ms`)), timeoutMs);
});
const result = await Promise.race([conn.client.listTools(), timeout]);
clearTimeout(timer);
for (const t of result.tools ?? []) {
const namespaced = this.namespace(id, t.name);
if (!this.deps.filter.isAllowed(namespaced, "tool")) continue;
Expand Down Expand Up @@ -246,9 +261,11 @@ export class Aggregator {
});
}
} catch (err) {
clearTimeout(timer);
this.log.warn({ id, err: errMsg(err) }, "upstream tools/list failed");
}
}
this.toolsCache = { tools: out, ts: Date.now() };
return out;
}

Expand Down
9 changes: 7 additions & 2 deletions src/core/filter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ export class ToolFilter {
private readonly applyTools: boolean;
private readonly applyResources: boolean;
private readonly applyPrompts: boolean;
/** Pre-compiled matchers for batch glob evaluation. */
private readonly allowMatcher: ((name: string) => boolean) | null;
private readonly denyMatcher: ((name: string) => boolean) | null;

constructor(opts: FilterOptions) {
this.allow = opts.allow;
Expand All @@ -34,6 +37,8 @@ export class ToolFilter {
this.applyTools = apply.tools ?? true;
this.applyResources = apply.resources ?? true;
this.applyPrompts = apply.prompts ?? true;
this.allowMatcher = this.allow.length > 0 ? micromatch.matcher(this.allow) : null;

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pnpm typecheck fails here - @types/micromatch types matcher() as taking string, not string[]:

Suggested changes:

this.allowMatcher = this.allow.length > 0 ? micromatch.matcher(this.allow) : null;
this.denyMatcher = this.deny.length > 0 ? micromatch.matcher(this.deny) : null;

to

this.allowMatcher = this.allow.length > 0 ? micromatch.matcher(this.allow as string | string[]) : null;
this.denyMatcher = this.deny.length > 0 ? micromatch.matcher(this.deny as string | string[]) : null;

this.denyMatcher = this.deny.length > 0 ? micromatch.matcher(this.deny) : null;
}

static fromConfig(cfg: TooltrimConfig): ToolFilter {
Expand All @@ -49,10 +54,10 @@ export class ToolFilter {
if (kind === "resource" && !this.applyResources) return true;
if (kind === "prompt" && !this.applyPrompts) return true;

if (this.allow.length > 0 && !micromatch.isMatch(namespacedName, this.allow)) {
if (this.allowMatcher && !this.allowMatcher(namespacedName)) {
return false;
}
if (this.deny.length > 0 && micromatch.isMatch(namespacedName, this.deny)) {
if (this.denyMatcher && this.denyMatcher(namespacedName)) {
return false;
}
return true;
Expand Down
4 changes: 2 additions & 2 deletions src/core/shrinker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -180,10 +180,10 @@ export class Shrinker {
.join(" ");
s = s.replace(/\s+/g, " ").trim();

// 6. capitalise first letter for readability
// 8. capitalise first letter for readability
if (s.length > 0) s = s[0]!.toUpperCase() + s.slice(1);

// 7. truncate at the first sentence boundary past `maxChars`
// 9. truncate at the first sentence boundary past `maxChars`
if (s.length > maxChars) {
const sliceEnd = this.findSentenceEnd(s, maxChars);
s = s.slice(0, sliceEnd).trimEnd();
Expand Down
8 changes: 7 additions & 1 deletion src/observability/audit.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ export class AuditLogger {
private readonly enabled: boolean;
private readonly filePath: string;
private dirEnsured = false;
/** Serializes appendFile calls to prevent interleaved NDJSON lines. */
private writeQueue = Promise.resolve();

constructor(enabled: boolean, filePath: string) {
this.enabled = enabled;
Expand All @@ -41,6 +43,10 @@ export class AuditLogger {
this.dirEnsured = true;
}
const line = JSON.stringify({ ts: new Date().toISOString(), ...ev });
await appendFile(this.filePath, line + "\n", "utf8");
// Serialize writes to prevent interleaving under concurrent requests.
this.writeQueue = this.writeQueue.then(() =>
appendFile(this.filePath, line + "\n", "utf8"),
);
return this.writeQueue;
}
}
15 changes: 15 additions & 0 deletions src/types/micromatch.d.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
declare module "micromatch" {
function micromatch(
patterns: string | string[],
options?: Record<string, unknown>,
): (str: string) => boolean;

namespace micromatch {
function matcher(
patterns: string | string[],
options?: Record<string, unknown>,
): (str: string) => boolean;
}

export = micromatch;
}
14 changes: 13 additions & 1 deletion src/upstream/manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -161,8 +161,20 @@ export class UpstreamManager {
cwd: cfg.cwd,
stderr: "pipe",
});
let stderrBytesThisMinute = 0;
let stderrMinuteStart = Date.now();
const stderrBudget = cfg.stderrLogBytesPerMinute ?? 10_000;

transport.stderr?.on("data", (chunk: Buffer) => {
this.log.debug({ id, chunk: chunk.toString("utf8").trim() }, "upstream stderr");
const now = Date.now();
if (now - stderrMinuteStart > 60_000) {
stderrBytesThisMinute = 0;
stderrMinuteStart = now;
}
if (stderrBudget === 0 || stderrBytesThisMinute < stderrBudget) {
this.log.debug({ id, chunk: chunk.toString("utf8").trim() }, "upstream stderr");
stderrBytesThisMinute += chunk.length;
}
});
await client.connect(transport);
}
Expand Down
Loading