@adhd/agent-plugin-budget 0.2.3 → 0.2.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (4) hide show
  1. package/CHANGELOG.md +11 -0
  2. package/index.cjs +33 -27
  3. package/index.js +120 -92
  4. package/package.json +4 -3
package/CHANGELOG.md CHANGED
@@ -1,3 +1,14 @@
1
+ ## 0.2.4 (2026-10-06)
2
+
3
+ ### 🩹 Fixes
4
+
5
+ - **agent-plugin-budget:** enforce global-scope caps via the drizzle handle + scoped cost ([6a2da60d](https://github.com/PseudoSky/adhd/commit/6a2da60d))
6
+ - **agent-mcp:** pin wildcard externals to real ranges ([9869db26](https://github.com/PseudoSky/adhd/commit/9869db26))
7
+
8
+ ### ❤️ Thank You
9
+
10
+ - pseudosky
11
+
1
12
  ## 0.2.2 (2026-09-26)
2
13
 
3
14
  ### 🚀 Features
package/index.cjs CHANGED
@@ -1,27 +1,33 @@
1
- Object.defineProperties(exports,{__esModule:{value:!0},[Symbol.toStringTag]:{value:`Module`}});let e=require("zod"),t=require("@adhd/agent-base-types");function n(e){let t=e.match(/^P(?:(\d+)Y)?(?:(\d+)M)?(?:(\d+)D)?(?:T(?:(\d+)H)?(?:(\d+)M)?(?:(\d+(?:\.\d+)?)S)?)?$/);if(!t)throw Error(`invalid ISO 8601 duration: ${e}`);let[,n,r,i,a,o,s]=t,c=0;return n&&(c+=parseInt(n)*365.25*864e5),r&&(c+=parseInt(r)*30.44*864e5),i&&(c+=parseInt(i)*864e5),a&&(c+=parseInt(a)*36e5),o&&(c+=parseInt(o)*6e4),s&&(c+=parseFloat(s)*1e3),Math.round(c)}var r=[`context`,`inputTokens`,`outputTokens`,`calls`,`wallClock`,`modelMs`,`cost`,`toolCalls`,`errors`,`consecutiveErrors`,`responseSize`],i=new Set([`toolCalls`,`errors`,`consecutiveErrors`]),a=e.z.object({field:e.z.enum(r),maximum:e.z.number().min(0).optional(),contextWindowFraction:e.z.number().min(0).max(1).optional(),window:e.z.string().optional(),scope:e.z.enum([`task`,`session`,`agent`,`global`]).optional(),mode:e.z.enum([`warning`,`block`]).optional(),message:e.z.string().optional()}).superRefine((t,n)=>{t.field===`context`?(t.maximum!==void 0==(t.contextWindowFraction!==void 0)&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`contextWindowFraction`],message:`cap field 'context' requires exactly one of 'maximum' or 'contextWindowFraction'`}),t.window!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`window`],message:`cap field 'context' rejects 'window' — a windowed peak is meaningless; windowed cumulative volume is expressible via 'inputTokens'/'outputTokens' + 'window'`})):((t.field===`errors`||t.field===`consecutiveErrors`)&&(t.window!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`window`],message:`cap field '${t.field}' rejects 'window' — windowed error budgets are not expressible this wave`}),t.scope!==void 0&&t.scope!==`task`&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`scope`],message:`cap field '${t.field}' is counted per-task in memory (no task_usage column) — only 'task' scope is expressible this wave`})),t.window!==void 0&&(t.scope===void 0||t.scope===`task`)&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`scope`],message:`cap field '${t.field}' carries 'window' but resolves to 'task' scope, where windowed caps are a silent no-op (no task-level window query); set an explicit 'scope': 'session' | 'agent' | 'global'`}),t.contextWindowFraction!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`contextWindowFraction`],message:`'contextWindowFraction' is only valid on cap field 'context' (got '${t.field}')`}),t.maximum===void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`maximum`],message:`cap field '${t.field}' requires a 'maximum'`}))}),o=e.z.object({caps:e.z.array(a).optional(),mode:e.z.enum([`warning`,`block`]).optional(),costPerInputToken:e.z.number().min(0).optional(),costPerOutputToken:e.z.number().min(0).optional(),costPerCacheReadToken:e.z.number().min(0).optional(),costPerCacheWriteToken:e.z.number().min(0).optional(),scope:e.z.enum([`task`,`session`,`agent`,`global`]).optional()}),s=o.partial().superRefine((t,n)=>{let r=t.caps??[];for(let[t,i]of r.entries())i.field===`context`&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`caps`,t,`field`],message:`cap field 'context' is only valid at model scope — a tool-scoped 'context' cap is a silent no-op (enforced on the model path, invisible to tool overrides); place it in 'defaults'/'agent'/'provider' instead`})}),c=e.z.object({defaults:o.optional(),agent:e.z.object({default:o.optional(),overrides:e.z.record(e.z.string(),o.partial()).optional().default({})}).optional(),provider:e.z.object({default:o.optional(),overrides:e.z.record(e.z.string(),o.partial()).optional().default({})}).optional(),tool:e.z.object({default:s.optional(),overrides:e.z.record(e.z.string(),s).optional().default({})}).optional()}),l=e.z.object({}).passthrough(),u={maxInputTokens:{field:`inputTokens`},maxOutputTokens:{field:`outputTokens`},maxModelCalls:{field:`calls`},maxWallClockMs:{field:`wallClock`},maxModelMs:{field:`modelMs`},maxCostUSD:{field:`cost`},maxTokensPer24h:{field:`inputTokens`,window:`PT24H`},maxCalls:{field:`toolCalls`}};function d(e){let t=[],n={};for(let[r,i]of Object.entries(e)){let e=u[r];if(e&&typeof i==`number`){let n={field:e.field,maximum:i};e.window&&(n.window=e.window),t.push(n)}else(r===`scope`||r===`mode`||r===`costPerInputToken`||r===`costPerOutputToken`||r===`costPerCacheReadToken`||r===`costPerCacheWriteToken`||r===`message`)&&(n[r]=i)}t.length>0&&(n.caps=t);let r=n.scope,i=r===`task`||r===`session`||r===`agent`||r===`global`?r:void 0;for(let e of t)e.window!==void 0&&e.scope===void 0&&(e.scope=i??`global`);return o.parse(n)}var f=e=>`legacy 'tokens' budget cap detected (${e}): cap field 'tokens' is removed; use 'context' (peak request input) or 'inputTokens'/'outputTokens' (windowed volume)`;function p(e){let t=e??{};if(typeof t.maxTotalTokens==`number`)throw Error(f(`flat 'maxTotalTokens'`));let n=e=>{if(typeof e!=`object`||!e)return;let t=e;if(typeof t.maxTotalTokens==`number`)throw Error(f(`structured 'maxTotalTokens'`));let n=t.caps;if(Array.isArray(n)){for(let e of n)if(typeof e==`object`&&e&&e.field===`tokens`)throw Error(f(`caps[].field === 'tokens'`))}},r=e=>{if(typeof e!=`object`||!e)return;let t=e;n(t),`default`in t&&n(t.default);let r=t.overrides;if(typeof r==`object`&&r)for(let e of Object.values(r))n(e)};n(t.defaults),r(t.agent),r(t.provider),r(t.tool)}function m(e){p(e);let t=e;if(t.defaults!==void 0||t.agent!==void 0||t.provider!==void 0||t.tool!==void 0){let t=c.parse(e);return{defaults:t.defaults??o.parse({}),agent:t.agent??{overrides:{}},provider:t.provider??{overrides:{}},tool:t.tool??{overrides:{}}}}return{defaults:d(t),agent:{overrides:{}},provider:{overrides:{}},tool:{overrides:{}}}}function h(e,t,n,r){return{isEnforcementError:!0,code:`BUDGET_EXCEEDED`,message:r??`${e} limit is ${t}, current value is ${Math.round(n)}`}}function g(e,t,n){return{isToolWarning:!0,toolName:e,callId:t,message:n}}function _(e,t){let n=0;for(let t of e){n+=t.content?.length??0;for(let e of t.toolCalls??[])n+=JSON.stringify(e.arguments??{}).length;for(let e of t.toolResults??[])n+=JSON.stringify(e.result??null).length}for(let e of t)n+=JSON.stringify({name:e.name,description:e.description,inputSchema:e.inputSchema}).length;return Math.ceil(n/4)}var v=class{constructor(e,t,n=0,r=0,i=0,a=0){this.db=e,this.cfg=t,this.costPerInput=n,this.costPerOutput=r,this.costPerCacheRead=i,this.costPerCacheWrite=a,this.name=`agent-mcp-budget`,this.accumulators=new Map}install(e){this.hooks=e,e.register(`task:start`,e=>{try{this.onTaskStart(e)}catch{}}),e.register(`pre:model_request`,e=>{try{this.onPreModelRequest(e)}catch{}}),e.register(`post:model_response`,e=>{try{this.onPostModelResponse(e)}catch{}}),e.register(`post:tool_call`,e=>{try{this.onPostToolCall(e)}catch{}}),e.register(`task:completed`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.register(`task:failed`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.register(`task:cancelled`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.registerEnforcement(`pre:model_request`,e=>this.enforcePreModel(e)),e.registerEnforcement(`pre:tool_call`,e=>this.enforcePreTool(e)),e.register(`transform:tool_result`,e=>{try{this.enforceResponseSize(e)}catch{}})}onTaskStart(e){let{taskId:t,sessionId:n,agentName:r}=e.executionContext,i=e.executionContext.agentDefinition?.provider?.type??`unknown`;this.accumulators.set(t,{taskId:t,sessionId:n??void 0,agentName:r,providerType:i,startedAtMs:Date.now(),uncachedInputTokens:0,cacheReadTokens:0,cacheCreationTokens:0,inputTokens:0,outputTokens:0,peakContextTokens:0,modelCalls:0,totalModelMs:0,toolCalls:new Map,errors:0,consecutiveErrors:0})}onPreModelRequest(e){let t=this.accumulators.get(e.executionContext.taskId);t&&(t.modelCallStartMs=Date.now())}onPostModelResponse(e){let t=this.accumulators.get(e.executionContext.taskId);if(!t)return;let n=e.tokenUsage;if(n){let e=n.inputTokens??0,r=n.cacheReadTokens??0,i=n.cacheCreationTokens??0,a=n.uncachedInputTokens!==void 0||n.cacheReadTokens!==void 0||n.cacheCreationTokens!==void 0?n.uncachedInputTokens??Math.max(0,e-r-i):e;t.uncachedInputTokens+=a,t.cacheReadTokens+=r,t.cacheCreationTokens+=i,t.inputTokens+=a+r+i,t.outputTokens+=n.outputTokens??0,t.peakContextTokens=Math.max(t.peakContextTokens,e)}t.modelCalls+=1,t.modelCallStartMs!==void 0&&(t.totalModelMs+=Date.now()-t.modelCallStartMs,t.modelCallStartMs=void 0)}onTerminal(e){this.accumulators.delete(e)}onPostToolCall(e){let t=this.accumulators.get(e.executionContext.taskId);t&&(e.isError?(t.errors+=1,t.consecutiveErrors+=1):t.consecutiveErrors=0)}mergeDim(e){let t={caps:[]};for(let n of e)n&&(t={caps:[...t.caps??[],...n.caps??[]],mode:n.mode??t.mode,costPerInputToken:n.costPerInputToken??t.costPerInputToken,costPerOutputToken:n.costPerOutputToken??t.costPerOutputToken,costPerCacheReadToken:n.costPerCacheReadToken??t.costPerCacheReadToken,costPerCacheWriteToken:n.costPerCacheWriteToken??t.costPerCacheWriteToken,scope:n.scope??t.scope});return t}resolveCaps(e,t,n){let r=this.cfg.defaults,i=this.cfg.agent,a=this.cfg.provider,o=this.cfg.tool;if(n){let e=o?.overrides?.[n],t=this.mergeDim([r,o?.default,e]);return{caps:t.caps??[],mode:t.mode,scope:t.scope}}let s=i?.overrides?.[e],c=a?.overrides?.[t],l=this.mergeDim([r,i?.default,s,a?.default,c]);return{caps:l.caps??[],mode:l.mode,scope:l.scope}}queryScopeTotals(e,t,n,r){let i=this.accumulators.get(e),a=i?{inputTokens:i.inputTokens,outputTokens:i.outputTokens,modelCalls:i.modelCalls,peakContextTokens:i.peakContextTokens}:{inputTokens:0,outputTokens:0,modelCalls:0,peakContextTokens:0};if(r===`task`||!this.db)return a;try{let i=this.db,o;if(r===`session`&&t?o=i.prepare(`SELECT
2
- COALESCE(SUM(tu.input_tokens), 0) AS input,
3
- COALESCE(SUM(tu.output_tokens), 0) AS output,
4
- COALESCE(SUM(tu.model_calls), 0) AS calls,
5
- COALESCE(MAX(tu.peak_context_tokens), 0) AS peak
6
- FROM task_usage tu
7
- JOIN tasks t ON tu.task_id = t.id
8
- WHERE t.session_id = ? AND tu.task_id != ?`).get(t,e):r===`agent`?o=i.prepare(`SELECT
9
- COALESCE(SUM(input_tokens), 0) AS input,
10
- COALESCE(SUM(output_tokens), 0) AS output,
11
- COALESCE(SUM(model_calls), 0) AS calls,
12
- COALESCE(MAX(peak_context_tokens), 0) AS peak
13
- FROM task_usage
14
- WHERE agent_name = ? AND task_id != ?`).get(n,e):r===`global`&&(o=i.prepare(`SELECT
15
- COALESCE(SUM(input_tokens), 0) AS input,
16
- COALESCE(SUM(output_tokens), 0) AS output,
17
- COALESCE(SUM(model_calls), 0) AS calls,
18
- COALESCE(MAX(peak_context_tokens), 0) AS peak
19
- FROM task_usage
20
- WHERE task_id != ?`).get(e)),o)return{inputTokens:(o.input??0)+a.inputTokens,outputTokens:(o.output??0)+a.outputTokens,modelCalls:(o.calls??0)+a.modelCalls,peakContextTokens:Math.max(o.peak??0,a.peakContextTokens)}}catch{}return a}queryWindowTokens(e,t,n,r){if(!this.db)return 0;try{let i=this.db,a=new Date(Date.now()-n).toISOString(),o;if(e===`session`){let e=r?` AND tu.task_id != ?`:``,n=[t,a];r&&n.push(r),o=i.prepare(`SELECT COALESCE(SUM(tu.input_tokens + tu.output_tokens), 0) AS total
21
- FROM task_usage tu
22
- JOIN tasks t ON tu.task_id = t.id
23
- WHERE t.session_id = ? AND tu.created_at >= ?${e}`).get(...n)}else if(e===`agent`){let e=r?` AND task_id != ?`:``,n=[t,a];r&&n.push(r),o=i.prepare(`SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
24
- FROM task_usage
25
- WHERE agent_name = ? AND created_at >= ?${e}`).get(...n)}else if(e===`global`){let e=r?` AND task_id != ?`:``,t=[a];r&&t.push(r),o=i.prepare(`SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
26
- FROM task_usage
27
- WHERE created_at >= ?${e}`).get(...t)}return o?.total??0}catch{return 0}}buildSnapshot(e,t,r,i,a,o,s=0){let c={};c.inputTokens=t.inputTokens,c.outputTokens=t.outputTokens,c.calls=t.modelCalls,c.wallClock=Date.now()-t.startedAtMs,c.modelMs=t.totalModelMs,c.context=Math.max(t.peakContextTokens,s),c.errors=t.errors,c.consecutiveErrors=t.consecutiveErrors,c.cost=t.uncachedInputTokens*this.costPerInput+t.cacheReadTokens*this.costPerCacheRead+t.cacheCreationTokens*this.costPerCacheWrite+t.outputTokens*this.costPerOutput;let l=new Set,u=new Map;for(let t of e){let e=t.scope??o??`task`;if(e!==`task`&&l.add(e),t.window){let r=`${e}:${t.window}`;u.has(r)||u.set(r,{scope:e,windowMs:n(t.window)})}}for(let e of l){let t=this.queryScopeTotals(r,i,a,e);c[`${e}:inputTokens`]=t.inputTokens,c[`${e}:outputTokens`]=t.outputTokens,c[`${e}:calls`]=t.modelCalls,c[`${e}:context`]=Math.max(t.peakContextTokens,s)}for(let[e,{scope:t,windowMs:n}]of u){let o=``;t===`session`?o=i??``:t===`agent`&&(o=a??``),c[e]=this.queryWindowTokens(t,o,n,r)}return c}getSnapshotValue(e,t,n){let r=t.scope??n??`task`,i;i=t.window||r===`task`?``:`${r}:`;let a;switch(t.field){case`inputTokens`:a=e[`${i}inputTokens`]??e.inputTokens;break;case`outputTokens`:a=e[`${i}outputTokens`]??e.outputTokens;break;case`context`:a=e[`${i}context`]??e.context;break;case`calls`:a=e[`${i}calls`]??e.calls;break;case`wallClock`:a=e.wallClock;break;case`modelMs`:a=e.modelMs;break;case`cost`:a=e.cost;break;case`toolCalls`:a=0;break;case`errors`:a=e.errors;break;case`consecutiveErrors`:a=e.consecutiveErrors;break;default:a=0}return t.window&&(a+=e[`${r}:${t.window}`]??0),a}async emitBudgetEvent(e,t,n,r,i,a){this.hooks&&await this.hooks.emit(e,{executionContext:t,field:n.field,maximum:i,current:r,message:a??n.message??`${n.field} limit is ${i}, current value is ${Math.round(r)}`})}resolveCapLimit(e,n){if(e.maximum!==void 0)return e.maximum;let r=n.agentDefinition.provider;return Math.floor((0,t.contextWindowFor)(r.model)*(e.contextWindowFraction??0))}async evaluateCap(e,t,n,r,i){let a=this.resolveCapLimit(e,n),o=this.getSnapshotValue(t,e,i);if(!(o<a)){if((e.mode??r??`warning`)===`warning`){await this.emitBudgetEvent(`budget:warning`,n,e,o,a);return}throw await this.emitBudgetEvent(`budget:block`,n,e,o,a),h(e.field,a,o,e.message)}}async enforcePreModel(e){let{taskId:t,sessionId:n,agentName:r}=e.executionContext,a=e.executionContext.agentDefinition?.provider?.type??`unknown`,o=this.accumulators.get(t);if(!o)return;let{caps:s,mode:c,scope:l}=this.resolveCaps(r,a),u=s.filter(e=>!i.has(e.field));if(u.length===0)return;let d=this.buildSnapshot(u,o,t,n,r,l,_(e.messages,e.tools));for(let t of u)await this.evaluateCap(t,d,e.executionContext,c,l)}async enforcePreTool(e){let{toolName:t,callId:n,executionContext:r}=e,{caps:i,mode:a,scope:o}=this.resolveCaps(r.agentName,``,t),s=this.accumulators.get(r.taskId);if(!s)return;let c=i.filter(e=>e.field!==`context`);if(c.length===0)return;let l=s.toolCalls.get(t)??0,u=this.buildSnapshot(c,s,r.taskId,r.sessionId,r.agentName,o);for(let e of c){let i=this.resolveCapLimit(e,r),s=e.field===`toolCalls`?l:this.getSnapshotValue(u,e,o);if(s>=i){let o=e.message??`tool "${t}": ${e.field} limit is ${i}, current value is ${Math.round(s)}`;throw(e.mode??a??`warning`)===`warning`?(await this.emitBudgetEvent(`budget:warning`,r,e,s,i,o),g(t,n,o)):(await this.emitBudgetEvent(`budget:block`,r,e,s,i,o),h(`tool:${t}:${e.field}`,i,s,e.message))}}s.toolCalls.set(t,l+1)}enforceResponseSize(e){let{toolName:t,result:n}=e;if(typeof n!=`object`||!n)return;let{caps:r,mode:i}=this.resolveCaps(``,``,t),a=r.filter(e=>e.field===`responseSize`);if(a.length===0)return;let o=n,s=o.content;if(!Array.isArray(s))return;let c=0;for(let e of s)if(typeof e==`object`&&e){let t=e;t.type===`text`&&(c+=(t.text??``).length)}for(let t of a){let n=this.resolveCapLimit(t,e.executionContext);if(!(c<=n)){if((t.mode??i??`warning`)===`block`)o.content=[{type:`text`,text:t.message??`Response size (${c} chars) exceeds limit of ${n}. Use offset/limit or shell paging tools instead.`}],e.isError=!0;else{let e=n,r=[];for(let t of s){if(typeof t!=`object`||!t){r.push(t);continue}let n=t;if(n.type!==`text`){r.push(t);continue}let i=n.text??``;if(i.length<=e)r.push(t),e-=i.length;else{r.push({type:`text`,text:i.slice(0,e)});break}}r.push({type:`text`,text:`\n\n[truncated: response was ${c} chars, limited to ${t.maximum}. ${t.message??`Use offset/limit or shell paging tools for full content.`}]`}),o.content=r}break}}}},y=({db:e,config:t})=>{let n=m(t),r=n.defaults?.costPerInputToken??0;return new v(e,n,r,n.defaults?.costPerOutputToken??0,n.defaults?.costPerCacheReadToken??r,n.defaults?.costPerCacheWriteToken??r)};exports.configSchema=l,exports.createPlugin=y,exports.default=y,exports.pluginConfigSchema=c;
1
+ Object.defineProperties(exports,{__esModule:{value:!0},[Symbol.toStringTag]:{value:`Module`}});let e=require("zod"),t=require("drizzle-orm"),n=require("@adhd/agent-base-types");function r(e){let t=e.match(/^P(?:(\d+)Y)?(?:(\d+)M)?(?:(\d+)D)?(?:T(?:(\d+)H)?(?:(\d+)M)?(?:(\d+(?:\.\d+)?)S)?)?$/);if(!t)throw Error(`invalid ISO 8601 duration: ${e}`);let[,n,r,i,a,o,s]=t,c=0;return n&&(c+=parseInt(n)*365.25*864e5),r&&(c+=parseInt(r)*30.44*864e5),i&&(c+=parseInt(i)*864e5),a&&(c+=parseInt(a)*36e5),o&&(c+=parseInt(o)*6e4),s&&(c+=parseFloat(s)*1e3),Math.round(c)}var i=[`context`,`inputTokens`,`outputTokens`,`calls`,`wallClock`,`modelMs`,`cost`,`toolCalls`,`errors`,`consecutiveErrors`,`responseSize`],a=new Set([`toolCalls`,`errors`,`consecutiveErrors`]),o=e.z.object({field:e.z.enum(i),maximum:e.z.number().min(0).optional(),contextWindowFraction:e.z.number().min(0).max(1).optional(),window:e.z.string().optional(),scope:e.z.enum([`task`,`session`,`agent`,`global`]).optional(),mode:e.z.enum([`warning`,`block`]).optional(),message:e.z.string().optional()}).superRefine((t,n)=>{t.field===`context`?(t.maximum!==void 0==(t.contextWindowFraction!==void 0)&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`contextWindowFraction`],message:`cap field 'context' requires exactly one of 'maximum' or 'contextWindowFraction'`}),t.window!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`window`],message:`cap field 'context' rejects 'window' — a windowed peak is meaningless; windowed cumulative volume is expressible via 'inputTokens'/'outputTokens' + 'window'`})):((t.field===`errors`||t.field===`consecutiveErrors`)&&(t.window!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`window`],message:`cap field '${t.field}' rejects 'window' — windowed error budgets are not expressible this wave`}),t.scope!==void 0&&t.scope!==`task`&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`scope`],message:`cap field '${t.field}' is counted per-task in memory (no task_usage column) — only 'task' scope is expressible this wave`})),t.window!==void 0&&(t.scope===void 0||t.scope===`task`)&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`scope`],message:`cap field '${t.field}' carries 'window' but resolves to 'task' scope, where windowed caps are a silent no-op (no task-level window query); set an explicit 'scope': 'session' | 'agent' | 'global'`}),t.contextWindowFraction!==void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`contextWindowFraction`],message:`'contextWindowFraction' is only valid on cap field 'context' (got '${t.field}')`}),t.maximum===void 0&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`maximum`],message:`cap field '${t.field}' requires a 'maximum'`}))}),s=e.z.object({caps:e.z.array(o).optional(),mode:e.z.enum([`warning`,`block`]).optional(),costPerInputToken:e.z.number().min(0).optional(),costPerOutputToken:e.z.number().min(0).optional(),costPerCacheReadToken:e.z.number().min(0).optional(),costPerCacheWriteToken:e.z.number().min(0).optional(),scope:e.z.enum([`task`,`session`,`agent`,`global`]).optional()}),c=s.partial().superRefine((t,n)=>{let r=t.caps??[];for(let[t,i]of r.entries())i.field===`context`&&n.addIssue({code:e.z.ZodIssueCode.custom,path:[`caps`,t,`field`],message:`cap field 'context' is only valid at model scope — a tool-scoped 'context' cap is a silent no-op (enforced on the model path, invisible to tool overrides); place it in 'defaults'/'agent'/'provider' instead`})}),l=e.z.object({defaults:s.optional(),agent:e.z.object({default:s.optional(),overrides:e.z.record(e.z.string(),s.partial()).optional().default({})}).optional(),provider:e.z.object({default:s.optional(),overrides:e.z.record(e.z.string(),s.partial()).optional().default({})}).optional(),tool:e.z.object({default:c.optional(),overrides:e.z.record(e.z.string(),c).optional().default({})}).optional()}),u=e.z.object({}).passthrough(),d={maxInputTokens:{field:`inputTokens`},maxOutputTokens:{field:`outputTokens`},maxModelCalls:{field:`calls`},maxWallClockMs:{field:`wallClock`},maxModelMs:{field:`modelMs`},maxCostUSD:{field:`cost`},maxTokensPer24h:{field:`inputTokens`,window:`PT24H`},maxCalls:{field:`toolCalls`}};function f(e){let t=[],n={};for(let[r,i]of Object.entries(e)){let e=d[r];if(e&&typeof i==`number`){let n={field:e.field,maximum:i};e.window&&(n.window=e.window),t.push(n)}else(r===`scope`||r===`mode`||r===`costPerInputToken`||r===`costPerOutputToken`||r===`costPerCacheReadToken`||r===`costPerCacheWriteToken`||r===`message`)&&(n[r]=i)}t.length>0&&(n.caps=t);let r=n.scope,i=r===`task`||r===`session`||r===`agent`||r===`global`?r:void 0;for(let e of t)e.window!==void 0&&e.scope===void 0&&(e.scope=i??`global`);return s.parse(n)}var p=e=>`legacy 'tokens' budget cap detected (${e}): cap field 'tokens' is removed; use 'context' (peak request input) or 'inputTokens'/'outputTokens' (windowed volume)`;function m(e){let t=e??{};if(typeof t.maxTotalTokens==`number`)throw Error(p(`flat 'maxTotalTokens'`));let n=e=>{if(typeof e!=`object`||!e)return;let t=e;if(typeof t.maxTotalTokens==`number`)throw Error(p(`structured 'maxTotalTokens'`));let n=t.caps;if(Array.isArray(n)){for(let e of n)if(typeof e==`object`&&e&&e.field===`tokens`)throw Error(p(`caps[].field === 'tokens'`))}},r=e=>{if(typeof e!=`object`||!e)return;let t=e;n(t),`default`in t&&n(t.default);let r=t.overrides;if(typeof r==`object`&&r)for(let e of Object.values(r))n(e)};n(t.defaults),r(t.agent),r(t.provider),r(t.tool)}function h(e){m(e);let t=e;if(t.defaults!==void 0||t.agent!==void 0||t.provider!==void 0||t.tool!==void 0){let t=l.parse(e);return{defaults:t.defaults??s.parse({}),agent:t.agent??{overrides:{}},provider:t.provider??{overrides:{}},tool:t.tool??{overrides:{}}}}return{defaults:f(t),agent:{overrides:{}},provider:{overrides:{}},tool:{overrides:{}}}}function g(e,t,n,r){return{isEnforcementError:!0,code:`BUDGET_EXCEEDED`,message:r??`${e} limit is ${t}, current value is ${Math.round(n)}`}}function _(e,t,n){return{isToolWarning:!0,toolName:e,callId:t,message:n}}function v(e,t){let n=0;for(let t of e){n+=t.content?.length??0;for(let e of t.toolCalls??[])n+=JSON.stringify(e.arguments??{}).length;for(let e of t.toolResults??[])n+=JSON.stringify(e.result??null).length}for(let e of t)n+=JSON.stringify({name:e.name,description:e.description,inputSchema:e.inputSchema}).length;return Math.ceil(n/4)}var y=class{constructor(e,t,n=0,r=0,i=0,a=0){this.db=e,this.cfg=t,this.costPerInput=n,this.costPerOutput=r,this.costPerCacheRead=i,this.costPerCacheWrite=a,this.name=`agent-mcp-budget`,this.accumulators=new Map}install(e){this.hooks=e,e.register(`task:start`,e=>{try{this.onTaskStart(e)}catch{}}),e.register(`pre:model_request`,e=>{try{this.onPreModelRequest(e)}catch{}}),e.register(`post:model_response`,e=>{try{this.onPostModelResponse(e)}catch{}}),e.register(`post:tool_call`,e=>{try{this.onPostToolCall(e)}catch{}}),e.register(`task:completed`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.register(`task:failed`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.register(`task:cancelled`,e=>{try{this.onTerminal(e.executionContext.taskId)}catch{}}),e.registerEnforcement(`pre:model_request`,e=>this.enforcePreModel(e)),e.registerEnforcement(`pre:tool_call`,e=>this.enforcePreTool(e)),e.register(`transform:tool_result`,e=>{try{this.enforceResponseSize(e)}catch{}})}onTaskStart(e){let{taskId:t,sessionId:n,agentName:r}=e.executionContext,i=e.executionContext.agentDefinition?.provider?.type??`unknown`;this.accumulators.set(t,{taskId:t,sessionId:n??void 0,agentName:r,providerType:i,startedAtMs:Date.now(),uncachedInputTokens:0,cacheReadTokens:0,cacheCreationTokens:0,inputTokens:0,outputTokens:0,peakContextTokens:0,modelCalls:0,totalModelMs:0,toolCalls:new Map,errors:0,consecutiveErrors:0})}onPreModelRequest(e){let t=this.accumulators.get(e.executionContext.taskId);t&&(t.modelCallStartMs=Date.now())}onPostModelResponse(e){let t=this.accumulators.get(e.executionContext.taskId);if(!t)return;let n=e.tokenUsage;if(n){let e=n.inputTokens??0,r=n.cacheReadTokens??0,i=n.cacheCreationTokens??0,a=n.uncachedInputTokens!==void 0||n.cacheReadTokens!==void 0||n.cacheCreationTokens!==void 0?n.uncachedInputTokens??Math.max(0,e-r-i):e;t.uncachedInputTokens+=a,t.cacheReadTokens+=r,t.cacheCreationTokens+=i,t.inputTokens+=a+r+i,t.outputTokens+=n.outputTokens??0,t.peakContextTokens=Math.max(t.peakContextTokens,e)}t.modelCalls+=1,t.modelCallStartMs!==void 0&&(t.totalModelMs+=Date.now()-t.modelCallStartMs,t.modelCallStartMs=void 0)}onTerminal(e){this.accumulators.delete(e)}onPostToolCall(e){let t=this.accumulators.get(e.executionContext.taskId);t&&(e.isError?(t.errors+=1,t.consecutiveErrors+=1):t.consecutiveErrors=0)}mergeDim(e){let t={caps:[]};for(let n of e)n&&(t={caps:[...t.caps??[],...n.caps??[]],mode:n.mode??t.mode,costPerInputToken:n.costPerInputToken??t.costPerInputToken,costPerOutputToken:n.costPerOutputToken??t.costPerOutputToken,costPerCacheReadToken:n.costPerCacheReadToken??t.costPerCacheReadToken,costPerCacheWriteToken:n.costPerCacheWriteToken??t.costPerCacheWriteToken,scope:n.scope??t.scope});return t}resolveCaps(e,t,n){let r=this.cfg.defaults,i=this.cfg.agent,a=this.cfg.provider,o=this.cfg.tool;if(n){let e=o?.overrides?.[n],t=this.mergeDim([r,o?.default,e]);return{caps:t.caps??[],mode:t.mode,scope:t.scope}}let s=i?.overrides?.[e],c=a?.overrides?.[t],l=this.mergeDim([r,i?.default,s,a?.default,c]);return{caps:l.caps??[],mode:l.mode,scope:l.scope}}queryScopeTotals(e,n,r,i){let a=this.accumulators.get(e),o=a?{inputTokens:a.inputTokens,outputTokens:a.outputTokens,modelCalls:a.modelCalls,peakContextTokens:a.peakContextTokens,uncachedInputTokens:a.uncachedInputTokens,cacheReadTokens:a.cacheReadTokens,cacheCreationTokens:a.cacheCreationTokens}:{inputTokens:0,outputTokens:0,modelCalls:0,peakContextTokens:0,uncachedInputTokens:0,cacheReadTokens:0,cacheCreationTokens:0};if(i===`task`||!this.db)return o;try{let a=this.db,s=e=>t.sql`
2
+ SELECT
3
+ COALESCE(SUM(tu.input_tokens), 0) AS input,
4
+ COALESCE(SUM(tu.output_tokens), 0) AS output,
5
+ COALESCE(SUM(tu.model_calls), 0) AS calls,
6
+ COALESCE(MAX(tu.peak_context_tokens), 0) AS peak,
7
+ COALESCE(SUM(COALESCE(
8
+ tu.uncached_input_tokens,
9
+ tu.input_tokens - COALESCE(tu.cache_read_input_tokens, 0)
10
+ - COALESCE(tu.cache_creation_input_tokens, 0)
11
+ )), 0) AS uncached,
12
+ COALESCE(SUM(tu.cache_read_input_tokens), 0) AS cache_read,
13
+ COALESCE(SUM(tu.cache_creation_input_tokens), 0) AS cache_write
14
+ ${e}`,c;if(i===`session`&&n?c=a.get(s(t.sql`FROM task_usage tu
15
+ JOIN tasks t ON tu.task_id = t.id
16
+ WHERE t.session_id = ${n} AND tu.task_id != ${e}`)):i===`agent`?c=a.get(s(t.sql`FROM task_usage tu
17
+ WHERE tu.agent_name = ${r} AND tu.task_id != ${e}`)):i===`global`&&(c=a.get(s(t.sql`FROM task_usage tu WHERE tu.task_id != ${e}`))),c)return{inputTokens:(c.input??0)+o.inputTokens,outputTokens:(c.output??0)+o.outputTokens,modelCalls:(c.calls??0)+o.modelCalls,peakContextTokens:Math.max(c.peak??0,o.peakContextTokens),uncachedInputTokens:(c.uncached??0)+o.uncachedInputTokens,cacheReadTokens:(c.cache_read??0)+o.cacheReadTokens,cacheCreationTokens:(c.cache_write??0)+o.cacheCreationTokens}}catch{}return o}queryWindowTokens(e,n,r,i){if(!this.db)return 0;try{let a=this.db,o=new Date(Date.now()-r).toISOString(),s;return e===`session`?s=a.get(t.sql`
18
+ SELECT COALESCE(SUM(tu.input_tokens + tu.output_tokens), 0) AS total
19
+ FROM task_usage tu
20
+ JOIN tasks t ON tu.task_id = t.id
21
+ WHERE t.session_id = ${n} AND tu.created_at >= ${o}
22
+ ${i?t.sql` AND tu.task_id != ${i}`:t.sql``}
23
+ `):e===`agent`?s=a.get(t.sql`
24
+ SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
25
+ FROM task_usage
26
+ WHERE agent_name = ${n} AND created_at >= ${o}
27
+ ${i?t.sql` AND task_id != ${i}`:t.sql``}
28
+ `):e===`global`&&(s=a.get(t.sql`
29
+ SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
30
+ FROM task_usage
31
+ WHERE created_at >= ${o}
32
+ ${i?t.sql` AND task_id != ${i}`:t.sql``}
33
+ `)),s?.total??0}catch{return 0}}costFromTotals(e){return e.uncachedInputTokens*this.costPerInput+e.cacheReadTokens*this.costPerCacheRead+e.cacheCreationTokens*this.costPerCacheWrite+e.outputTokens*this.costPerOutput}buildSnapshot(e,t,n,i,a,o,s=0){let c={};c.inputTokens=t.inputTokens,c.outputTokens=t.outputTokens,c.calls=t.modelCalls,c.wallClock=Date.now()-t.startedAtMs,c.modelMs=t.totalModelMs,c.context=Math.max(t.peakContextTokens,s),c.errors=t.errors,c.consecutiveErrors=t.consecutiveErrors,c.cost=this.costFromTotals(t);let l=new Set,u=new Map;for(let t of e){let e=t.scope??o??`task`;if(e!==`task`&&l.add(e),t.window){let n=`${e}:${t.window}`;u.has(n)||u.set(n,{scope:e,windowMs:r(t.window)})}}for(let e of l){let t=this.queryScopeTotals(n,i,a,e);c[`${e}:inputTokens`]=t.inputTokens,c[`${e}:outputTokens`]=t.outputTokens,c[`${e}:calls`]=t.modelCalls,c[`${e}:context`]=Math.max(t.peakContextTokens,s),c[`${e}:cost`]=this.costFromTotals(t)}for(let[e,{scope:t,windowMs:r}]of u){let o=``;t===`session`?o=i??``:t===`agent`&&(o=a??``),c[e]=this.queryWindowTokens(t,o,r,n)}return c}getSnapshotValue(e,t,n){let r=t.scope??n??`task`,i;i=t.window||r===`task`?``:`${r}:`;let a;switch(t.field){case`inputTokens`:a=e[`${i}inputTokens`]??e.inputTokens;break;case`outputTokens`:a=e[`${i}outputTokens`]??e.outputTokens;break;case`context`:a=e[`${i}context`]??e.context;break;case`calls`:a=e[`${i}calls`]??e.calls;break;case`wallClock`:a=e.wallClock;break;case`modelMs`:a=e.modelMs;break;case`cost`:a=e[`${i}cost`]??e.cost;break;case`toolCalls`:a=0;break;case`errors`:a=e.errors;break;case`consecutiveErrors`:a=e.consecutiveErrors;break;default:a=0}return t.window&&(a+=e[`${r}:${t.window}`]??0),a}async emitBudgetEvent(e,t,n,r,i,a){this.hooks&&await this.hooks.emit(e,{executionContext:t,field:n.field,maximum:i,current:r,message:a??n.message??`${n.field} limit is ${i}, current value is ${Math.round(r)}`})}resolveCapLimit(e,t){if(e.maximum!==void 0)return e.maximum;let r=t.agentDefinition.provider;return Math.floor((0,n.contextWindowFor)(r.model)*(e.contextWindowFraction??0))}async evaluateCap(e,t,n,r,i){let a=this.resolveCapLimit(e,n),o=this.getSnapshotValue(t,e,i);if(!(o<a)){if((e.mode??r??`warning`)===`warning`){await this.emitBudgetEvent(`budget:warning`,n,e,o,a);return}throw await this.emitBudgetEvent(`budget:block`,n,e,o,a),g(e.field,a,o,e.message)}}async enforcePreModel(e){let{taskId:t,sessionId:n,agentName:r}=e.executionContext,i=e.executionContext.agentDefinition?.provider?.type??`unknown`,o=this.accumulators.get(t);if(!o)return;let{caps:s,mode:c,scope:l}=this.resolveCaps(r,i),u=s.filter(e=>!a.has(e.field));if(u.length===0)return;let d=this.buildSnapshot(u,o,t,n,r,l,v(e.messages,e.tools));for(let t of u)await this.evaluateCap(t,d,e.executionContext,c,l)}async enforcePreTool(e){let{toolName:t,callId:n,executionContext:r}=e,{caps:i,mode:a,scope:o}=this.resolveCaps(r.agentName,``,t),s=this.accumulators.get(r.taskId);if(!s)return;let c=i.filter(e=>e.field!==`context`);if(c.length===0)return;let l=s.toolCalls.get(t)??0,u=this.buildSnapshot(c,s,r.taskId,r.sessionId,r.agentName,o);for(let e of c){let i=this.resolveCapLimit(e,r),s=e.field===`toolCalls`?l:this.getSnapshotValue(u,e,o);if(s>=i){let o=e.message??`tool "${t}": ${e.field} limit is ${i}, current value is ${Math.round(s)}`;throw(e.mode??a??`warning`)===`warning`?(await this.emitBudgetEvent(`budget:warning`,r,e,s,i,o),_(t,n,o)):(await this.emitBudgetEvent(`budget:block`,r,e,s,i,o),g(`tool:${t}:${e.field}`,i,s,e.message))}}s.toolCalls.set(t,l+1)}enforceResponseSize(e){let{toolName:t,result:n}=e;if(typeof n!=`object`||!n)return;let{caps:r,mode:i}=this.resolveCaps(``,``,t),a=r.filter(e=>e.field===`responseSize`);if(a.length===0)return;let o=n,s=o.content;if(!Array.isArray(s))return;let c=0;for(let e of s)if(typeof e==`object`&&e){let t=e;t.type===`text`&&(c+=(t.text??``).length)}for(let t of a){let n=this.resolveCapLimit(t,e.executionContext);if(!(c<=n)){if((t.mode??i??`warning`)===`block`)o.content=[{type:`text`,text:t.message??`Response size (${c} chars) exceeds limit of ${n}. Use offset/limit or shell paging tools instead.`}],e.isError=!0;else{let e=n,r=[];for(let t of s){if(typeof t!=`object`||!t){r.push(t);continue}let n=t;if(n.type!==`text`){r.push(t);continue}let i=n.text??``;if(i.length<=e)r.push(t),e-=i.length;else{r.push({type:`text`,text:i.slice(0,e)});break}}r.push({type:`text`,text:`\n\n[truncated: response was ${c} chars, limited to ${t.maximum}. ${t.message??`Use offset/limit or shell paging tools for full content.`}]`}),o.content=r}break}}}},b=({db:e,config:t})=>{let n=h(t),r=n.defaults?.costPerInputToken??0;return new y(e,n,r,n.defaults?.costPerOutputToken??0,n.defaults?.costPerCacheReadToken??r,n.defaults?.costPerCacheWriteToken??r)};exports.configSchema=u,exports.createPlugin=b,exports.default=b,exports.pluginConfigSchema=l;
package/index.js CHANGED
@@ -1,13 +1,14 @@
1
1
  import { z as e } from "zod";
2
- import { contextWindowFor as t } from "@adhd/agent-base-types";
2
+ import { sql as t } from "drizzle-orm";
3
+ import { contextWindowFor as n } from "@adhd/agent-base-types";
3
4
  //#region src/index.ts
4
- function n(e) {
5
+ function r(e) {
5
6
  let t = e.match(/^P(?:(\d+)Y)?(?:(\d+)M)?(?:(\d+)D)?(?:T(?:(\d+)H)?(?:(\d+)M)?(?:(\d+(?:\.\d+)?)S)?)?$/);
6
7
  if (!t) throw Error(`invalid ISO 8601 duration: ${e}`);
7
8
  let [, n, r, i, a, o, s] = t, c = 0;
8
9
  return n && (c += parseInt(n) * 365.25 * 864e5), r && (c += parseInt(r) * 30.44 * 864e5), i && (c += parseInt(i) * 864e5), a && (c += parseInt(a) * 36e5), o && (c += parseInt(o) * 6e4), s && (c += parseFloat(s) * 1e3), Math.round(c);
9
10
  }
10
- var r = [
11
+ var i = [
11
12
  "context",
12
13
  "inputTokens",
13
14
  "outputTokens",
@@ -19,12 +20,12 @@ var r = [
19
20
  "errors",
20
21
  "consecutiveErrors",
21
22
  "responseSize"
22
- ], i = /* @__PURE__ */ new Set([
23
+ ], a = /* @__PURE__ */ new Set([
23
24
  "toolCalls",
24
25
  "errors",
25
26
  "consecutiveErrors"
26
- ]), a = e.object({
27
- field: e.enum(r),
27
+ ]), o = e.object({
28
+ field: e.enum(i),
28
29
  maximum: e.number().min(0).optional(),
29
30
  contextWindowFraction: e.number().min(0).max(1).optional(),
30
31
  window: e.string().optional(),
@@ -66,8 +67,8 @@ var r = [
66
67
  path: ["maximum"],
67
68
  message: `cap field '${t.field}' requires a 'maximum'`
68
69
  }));
69
- }), o = e.object({
70
- caps: e.array(a).optional(),
70
+ }), s = e.object({
71
+ caps: e.array(o).optional(),
71
72
  mode: e.enum(["warning", "block"]).optional(),
72
73
  costPerInputToken: e.number().min(0).optional(),
73
74
  costPerOutputToken: e.number().min(0).optional(),
@@ -79,7 +80,7 @@ var r = [
79
80
  "agent",
80
81
  "global"
81
82
  ]).optional()
82
- }), s = o.partial().superRefine((t, n) => {
83
+ }), c = s.partial().superRefine((t, n) => {
83
84
  let r = t.caps ?? [];
84
85
  for (let [t, i] of r.entries()) i.field === "context" && n.addIssue({
85
86
  code: e.ZodIssueCode.custom,
@@ -90,21 +91,21 @@ var r = [
90
91
  ],
91
92
  message: "cap field 'context' is only valid at model scope — a tool-scoped 'context' cap is a silent no-op (enforced on the model path, invisible to tool overrides); place it in 'defaults'/'agent'/'provider' instead"
92
93
  });
93
- }), c = e.object({
94
- defaults: o.optional(),
94
+ }), l = e.object({
95
+ defaults: s.optional(),
95
96
  agent: e.object({
96
- default: o.optional(),
97
- overrides: e.record(e.string(), o.partial()).optional().default({})
97
+ default: s.optional(),
98
+ overrides: e.record(e.string(), s.partial()).optional().default({})
98
99
  }).optional(),
99
100
  provider: e.object({
100
- default: o.optional(),
101
- overrides: e.record(e.string(), o.partial()).optional().default({})
101
+ default: s.optional(),
102
+ overrides: e.record(e.string(), s.partial()).optional().default({})
102
103
  }).optional(),
103
104
  tool: e.object({
104
- default: s.optional(),
105
- overrides: e.record(e.string(), s).optional().default({})
105
+ default: c.optional(),
106
+ overrides: e.record(e.string(), c).optional().default({})
106
107
  }).optional()
107
- }), l = e.object({}).passthrough(), u = {
108
+ }), u = e.object({}).passthrough(), d = {
108
109
  maxInputTokens: { field: "inputTokens" },
109
110
  maxOutputTokens: { field: "outputTokens" },
110
111
  maxModelCalls: { field: "calls" },
@@ -117,10 +118,10 @@ var r = [
117
118
  },
118
119
  maxCalls: { field: "toolCalls" }
119
120
  };
120
- function d(e) {
121
+ function f(e) {
121
122
  let t = [], n = {};
122
123
  for (let [r, i] of Object.entries(e)) {
123
- let e = u[r];
124
+ let e = d[r];
124
125
  if (e && typeof i == "number") {
125
126
  let n = {
126
127
  field: e.field,
@@ -132,19 +133,19 @@ function d(e) {
132
133
  t.length > 0 && (n.caps = t);
133
134
  let r = n.scope, i = r === "task" || r === "session" || r === "agent" || r === "global" ? r : void 0;
134
135
  for (let e of t) e.window !== void 0 && e.scope === void 0 && (e.scope = i ?? "global");
135
- return o.parse(n);
136
+ return s.parse(n);
136
137
  }
137
- var f = (e) => `legacy 'tokens' budget cap detected (${e}): cap field 'tokens' is removed; use 'context' (peak request input) or 'inputTokens'/'outputTokens' (windowed volume)`;
138
- function p(e) {
138
+ var p = (e) => `legacy 'tokens' budget cap detected (${e}): cap field 'tokens' is removed; use 'context' (peak request input) or 'inputTokens'/'outputTokens' (windowed volume)`;
139
+ function m(e) {
139
140
  let t = e ?? {};
140
- if (typeof t.maxTotalTokens == "number") throw Error(f("flat 'maxTotalTokens'"));
141
+ if (typeof t.maxTotalTokens == "number") throw Error(p("flat 'maxTotalTokens'"));
141
142
  let n = (e) => {
142
143
  if (typeof e != "object" || !e) return;
143
144
  let t = e;
144
- if (typeof t.maxTotalTokens == "number") throw Error(f("structured 'maxTotalTokens'"));
145
+ if (typeof t.maxTotalTokens == "number") throw Error(p("structured 'maxTotalTokens'"));
145
146
  let n = t.caps;
146
147
  if (Array.isArray(n)) {
147
- for (let e of n) if (typeof e == "object" && e && e.field === "tokens") throw Error(f("caps[].field === 'tokens'"));
148
+ for (let e of n) if (typeof e == "object" && e && e.field === "tokens") throw Error(p("caps[].field === 'tokens'"));
148
149
  }
149
150
  }, r = (e) => {
150
151
  if (typeof e != "object" || !e) return;
@@ -155,33 +156,33 @@ function p(e) {
155
156
  };
156
157
  n(t.defaults), r(t.agent), r(t.provider), r(t.tool);
157
158
  }
158
- function m(e) {
159
- p(e);
159
+ function h(e) {
160
+ m(e);
160
161
  let t = e;
161
162
  if (t.defaults !== void 0 || t.agent !== void 0 || t.provider !== void 0 || t.tool !== void 0) {
162
- let t = c.parse(e);
163
+ let t = l.parse(e);
163
164
  return {
164
- defaults: t.defaults ?? o.parse({}),
165
+ defaults: t.defaults ?? s.parse({}),
165
166
  agent: t.agent ?? { overrides: {} },
166
167
  provider: t.provider ?? { overrides: {} },
167
168
  tool: t.tool ?? { overrides: {} }
168
169
  };
169
170
  }
170
171
  return {
171
- defaults: d(t),
172
+ defaults: f(t),
172
173
  agent: { overrides: {} },
173
174
  provider: { overrides: {} },
174
175
  tool: { overrides: {} }
175
176
  };
176
177
  }
177
- function h(e, t, n, r) {
178
+ function g(e, t, n, r) {
178
179
  return {
179
180
  isEnforcementError: !0,
180
181
  code: "BUDGET_EXCEEDED",
181
182
  message: r ?? `${e} limit is ${t}, current value is ${Math.round(n)}`
182
183
  };
183
184
  }
184
- function g(e, t, n) {
185
+ function _(e, t, n) {
185
186
  return {
186
187
  isToolWarning: !0,
187
188
  toolName: e,
@@ -189,7 +190,7 @@ function g(e, t, n) {
189
190
  message: n
190
191
  };
191
192
  }
192
- function _(e, t) {
193
+ function v(e, t) {
193
194
  let n = 0;
194
195
  for (let t of e) {
195
196
  n += t.content?.length ?? 0;
@@ -203,7 +204,7 @@ function _(e, t) {
203
204
  }).length;
204
205
  return Math.ceil(n / 4);
205
206
  }
206
- var v = class {
207
+ var y = class {
207
208
  constructor(e, t, n = 0, r = 0, i = 0, a = 0) {
208
209
  this.db = e, this.cfg = t, this.costPerInput = n, this.costPerOutput = r, this.costPerCacheRead = i, this.costPerCacheWrite = a, this.name = "agent-mcp-budget", this.accumulators = /* @__PURE__ */ new Map();
209
210
  }
@@ -324,77 +325,104 @@ var v = class {
324
325
  scope: l.scope
325
326
  };
326
327
  }
327
- queryScopeTotals(e, t, n, r) {
328
- let i = this.accumulators.get(e), a = i ? {
329
- inputTokens: i.inputTokens,
330
- outputTokens: i.outputTokens,
331
- modelCalls: i.modelCalls,
332
- peakContextTokens: i.peakContextTokens
328
+ queryScopeTotals(e, n, r, i) {
329
+ let a = this.accumulators.get(e), o = a ? {
330
+ inputTokens: a.inputTokens,
331
+ outputTokens: a.outputTokens,
332
+ modelCalls: a.modelCalls,
333
+ peakContextTokens: a.peakContextTokens,
334
+ uncachedInputTokens: a.uncachedInputTokens,
335
+ cacheReadTokens: a.cacheReadTokens,
336
+ cacheCreationTokens: a.cacheCreationTokens
333
337
  } : {
334
338
  inputTokens: 0,
335
339
  outputTokens: 0,
336
340
  modelCalls: 0,
337
- peakContextTokens: 0
341
+ peakContextTokens: 0,
342
+ uncachedInputTokens: 0,
343
+ cacheReadTokens: 0,
344
+ cacheCreationTokens: 0
338
345
  };
339
- if (r === "task" || !this.db) return a;
346
+ if (i === "task" || !this.db) return o;
340
347
  try {
341
- let i = this.db, o;
342
- if (r === "session" && t ? o = i.prepare("SELECT\n COALESCE(SUM(tu.input_tokens), 0) AS input,\n COALESCE(SUM(tu.output_tokens), 0) AS output,\n COALESCE(SUM(tu.model_calls), 0) AS calls,\n COALESCE(MAX(tu.peak_context_tokens), 0) AS peak\n FROM task_usage tu\n JOIN tasks t ON tu.task_id = t.id\n WHERE t.session_id = ? AND tu.task_id != ?").get(t, e) : r === "agent" ? o = i.prepare("SELECT\n COALESCE(SUM(input_tokens), 0) AS input,\n COALESCE(SUM(output_tokens), 0) AS output,\n COALESCE(SUM(model_calls), 0) AS calls,\n COALESCE(MAX(peak_context_tokens), 0) AS peak\n FROM task_usage\n WHERE agent_name = ? AND task_id != ?").get(n, e) : r === "global" && (o = i.prepare("SELECT\n COALESCE(SUM(input_tokens), 0) AS input,\n COALESCE(SUM(output_tokens), 0) AS output,\n COALESCE(SUM(model_calls), 0) AS calls,\n COALESCE(MAX(peak_context_tokens), 0) AS peak\n FROM task_usage\n WHERE task_id != ?").get(e)), o) return {
343
- inputTokens: (o.input ?? 0) + a.inputTokens,
344
- outputTokens: (o.output ?? 0) + a.outputTokens,
345
- modelCalls: (o.calls ?? 0) + a.modelCalls,
346
- peakContextTokens: Math.max(o.peak ?? 0, a.peakContextTokens)
348
+ let a = this.db, s = (e) => t`
349
+ SELECT
350
+ COALESCE(SUM(tu.input_tokens), 0) AS input,
351
+ COALESCE(SUM(tu.output_tokens), 0) AS output,
352
+ COALESCE(SUM(tu.model_calls), 0) AS calls,
353
+ COALESCE(MAX(tu.peak_context_tokens), 0) AS peak,
354
+ COALESCE(SUM(COALESCE(
355
+ tu.uncached_input_tokens,
356
+ tu.input_tokens - COALESCE(tu.cache_read_input_tokens, 0)
357
+ - COALESCE(tu.cache_creation_input_tokens, 0)
358
+ )), 0) AS uncached,
359
+ COALESCE(SUM(tu.cache_read_input_tokens), 0) AS cache_read,
360
+ COALESCE(SUM(tu.cache_creation_input_tokens), 0) AS cache_write
361
+ ${e}`, c;
362
+ if (i === "session" && n ? c = a.get(s(t`FROM task_usage tu
363
+ JOIN tasks t ON tu.task_id = t.id
364
+ WHERE t.session_id = ${n} AND tu.task_id != ${e}`)) : i === "agent" ? c = a.get(s(t`FROM task_usage tu
365
+ WHERE tu.agent_name = ${r} AND tu.task_id != ${e}`)) : i === "global" && (c = a.get(s(t`FROM task_usage tu WHERE tu.task_id != ${e}`))), c) return {
366
+ inputTokens: (c.input ?? 0) + o.inputTokens,
367
+ outputTokens: (c.output ?? 0) + o.outputTokens,
368
+ modelCalls: (c.calls ?? 0) + o.modelCalls,
369
+ peakContextTokens: Math.max(c.peak ?? 0, o.peakContextTokens),
370
+ uncachedInputTokens: (c.uncached ?? 0) + o.uncachedInputTokens,
371
+ cacheReadTokens: (c.cache_read ?? 0) + o.cacheReadTokens,
372
+ cacheCreationTokens: (c.cache_write ?? 0) + o.cacheCreationTokens
347
373
  };
348
374
  } catch {}
349
- return a;
375
+ return o;
350
376
  }
351
- queryWindowTokens(e, t, n, r) {
377
+ queryWindowTokens(e, n, r, i) {
352
378
  if (!this.db) return 0;
353
379
  try {
354
- let i = this.db, a = new Date(Date.now() - n).toISOString(), o;
355
- if (e === "session") {
356
- let e = r ? " AND tu.task_id != ?" : "", n = [t, a];
357
- r && n.push(r), o = i.prepare(`SELECT COALESCE(SUM(tu.input_tokens + tu.output_tokens), 0) AS total
358
- FROM task_usage tu
359
- JOIN tasks t ON tu.task_id = t.id
360
- WHERE t.session_id = ? AND tu.created_at >= ?${e}`).get(...n);
361
- } else if (e === "agent") {
362
- let e = r ? " AND task_id != ?" : "", n = [t, a];
363
- r && n.push(r), o = i.prepare(`SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
364
- FROM task_usage
365
- WHERE agent_name = ? AND created_at >= ?${e}`).get(...n);
366
- } else if (e === "global") {
367
- let e = r ? " AND task_id != ?" : "", t = [a];
368
- r && t.push(r), o = i.prepare(`SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
369
- FROM task_usage
370
- WHERE created_at >= ?${e}`).get(...t);
371
- }
372
- return o?.total ?? 0;
380
+ let a = this.db, o = new Date(Date.now() - r).toISOString(), s;
381
+ return e === "session" ? s = a.get(t`
382
+ SELECT COALESCE(SUM(tu.input_tokens + tu.output_tokens), 0) AS total
383
+ FROM task_usage tu
384
+ JOIN tasks t ON tu.task_id = t.id
385
+ WHERE t.session_id = ${n} AND tu.created_at >= ${o}
386
+ ${i ? t` AND tu.task_id != ${i}` : t``}
387
+ `) : e === "agent" ? s = a.get(t`
388
+ SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
389
+ FROM task_usage
390
+ WHERE agent_name = ${n} AND created_at >= ${o}
391
+ ${i ? t` AND task_id != ${i}` : t``}
392
+ `) : e === "global" && (s = a.get(t`
393
+ SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total
394
+ FROM task_usage
395
+ WHERE created_at >= ${o}
396
+ ${i ? t` AND task_id != ${i}` : t``}
397
+ `)), s?.total ?? 0;
373
398
  } catch {
374
399
  return 0;
375
400
  }
376
401
  }
377
- buildSnapshot(e, t, r, i, a, o, s = 0) {
402
+ costFromTotals(e) {
403
+ return e.uncachedInputTokens * this.costPerInput + e.cacheReadTokens * this.costPerCacheRead + e.cacheCreationTokens * this.costPerCacheWrite + e.outputTokens * this.costPerOutput;
404
+ }
405
+ buildSnapshot(e, t, n, i, a, o, s = 0) {
378
406
  let c = {};
379
- c.inputTokens = t.inputTokens, c.outputTokens = t.outputTokens, c.calls = t.modelCalls, c.wallClock = Date.now() - t.startedAtMs, c.modelMs = t.totalModelMs, c.context = Math.max(t.peakContextTokens, s), c.errors = t.errors, c.consecutiveErrors = t.consecutiveErrors, c.cost = t.uncachedInputTokens * this.costPerInput + t.cacheReadTokens * this.costPerCacheRead + t.cacheCreationTokens * this.costPerCacheWrite + t.outputTokens * this.costPerOutput;
407
+ c.inputTokens = t.inputTokens, c.outputTokens = t.outputTokens, c.calls = t.modelCalls, c.wallClock = Date.now() - t.startedAtMs, c.modelMs = t.totalModelMs, c.context = Math.max(t.peakContextTokens, s), c.errors = t.errors, c.consecutiveErrors = t.consecutiveErrors, c.cost = this.costFromTotals(t);
380
408
  let l = /* @__PURE__ */ new Set(), u = /* @__PURE__ */ new Map();
381
409
  for (let t of e) {
382
410
  let e = t.scope ?? o ?? "task";
383
411
  if (e !== "task" && l.add(e), t.window) {
384
- let r = `${e}:${t.window}`;
385
- u.has(r) || u.set(r, {
412
+ let n = `${e}:${t.window}`;
413
+ u.has(n) || u.set(n, {
386
414
  scope: e,
387
- windowMs: n(t.window)
415
+ windowMs: r(t.window)
388
416
  });
389
417
  }
390
418
  }
391
419
  for (let e of l) {
392
- let t = this.queryScopeTotals(r, i, a, e);
393
- c[`${e}:inputTokens`] = t.inputTokens, c[`${e}:outputTokens`] = t.outputTokens, c[`${e}:calls`] = t.modelCalls, c[`${e}:context`] = Math.max(t.peakContextTokens, s);
420
+ let t = this.queryScopeTotals(n, i, a, e);
421
+ c[`${e}:inputTokens`] = t.inputTokens, c[`${e}:outputTokens`] = t.outputTokens, c[`${e}:calls`] = t.modelCalls, c[`${e}:context`] = Math.max(t.peakContextTokens, s), c[`${e}:cost`] = this.costFromTotals(t);
394
422
  }
395
- for (let [e, { scope: t, windowMs: n }] of u) {
423
+ for (let [e, { scope: t, windowMs: r }] of u) {
396
424
  let o = "";
397
- t === "session" ? o = i ?? "" : t === "agent" && (o = a ?? ""), c[e] = this.queryWindowTokens(t, o, n, r);
425
+ t === "session" ? o = i ?? "" : t === "agent" && (o = a ?? ""), c[e] = this.queryWindowTokens(t, o, r, n);
398
426
  }
399
427
  return c;
400
428
  }
@@ -422,7 +450,7 @@ var v = class {
422
450
  a = e.modelMs;
423
451
  break;
424
452
  case "cost":
425
- a = e.cost;
453
+ a = e[`${i}cost`] ?? e.cost;
426
454
  break;
427
455
  case "toolCalls":
428
456
  a = 0;
@@ -446,10 +474,10 @@ var v = class {
446
474
  message: a ?? n.message ?? `${n.field} limit is ${i}, current value is ${Math.round(r)}`
447
475
  });
448
476
  }
449
- resolveCapLimit(e, n) {
477
+ resolveCapLimit(e, t) {
450
478
  if (e.maximum !== void 0) return e.maximum;
451
- let r = n.agentDefinition.provider;
452
- return Math.floor(t(r.model) * (e.contextWindowFraction ?? 0));
479
+ let r = t.agentDefinition.provider;
480
+ return Math.floor(n(r.model) * (e.contextWindowFraction ?? 0));
453
481
  }
454
482
  async evaluateCap(e, t, n, r, i) {
455
483
  let a = this.resolveCapLimit(e, n), o = this.getSnapshotValue(t, e, i);
@@ -458,15 +486,15 @@ var v = class {
458
486
  await this.emitBudgetEvent("budget:warning", n, e, o, a);
459
487
  return;
460
488
  }
461
- throw await this.emitBudgetEvent("budget:block", n, e, o, a), h(e.field, a, o, e.message);
489
+ throw await this.emitBudgetEvent("budget:block", n, e, o, a), g(e.field, a, o, e.message);
462
490
  }
463
491
  }
464
492
  async enforcePreModel(e) {
465
- let { taskId: t, sessionId: n, agentName: r } = e.executionContext, a = e.executionContext.agentDefinition?.provider?.type ?? "unknown", o = this.accumulators.get(t);
493
+ let { taskId: t, sessionId: n, agentName: r } = e.executionContext, i = e.executionContext.agentDefinition?.provider?.type ?? "unknown", o = this.accumulators.get(t);
466
494
  if (!o) return;
467
- let { caps: s, mode: c, scope: l } = this.resolveCaps(r, a), u = s.filter((e) => !i.has(e.field));
495
+ let { caps: s, mode: c, scope: l } = this.resolveCaps(r, i), u = s.filter((e) => !a.has(e.field));
468
496
  if (u.length === 0) return;
469
- let d = this.buildSnapshot(u, o, t, n, r, l, _(e.messages, e.tools));
497
+ let d = this.buildSnapshot(u, o, t, n, r, l, v(e.messages, e.tools));
470
498
  for (let t of u) await this.evaluateCap(t, d, e.executionContext, c, l);
471
499
  }
472
500
  async enforcePreTool(e) {
@@ -479,7 +507,7 @@ var v = class {
479
507
  let i = this.resolveCapLimit(e, r), s = e.field === "toolCalls" ? l : this.getSnapshotValue(u, e, o);
480
508
  if (s >= i) {
481
509
  let o = e.message ?? `tool "${t}": ${e.field} limit is ${i}, current value is ${Math.round(s)}`;
482
- throw (e.mode ?? a ?? "warning") === "warning" ? (await this.emitBudgetEvent("budget:warning", r, e, s, i, o), g(t, n, o)) : (await this.emitBudgetEvent("budget:block", r, e, s, i, o), h(`tool:${t}:${e.field}`, i, s, e.message));
510
+ throw (e.mode ?? a ?? "warning") === "warning" ? (await this.emitBudgetEvent("budget:warning", r, e, s, i, o), _(t, n, o)) : (await this.emitBudgetEvent("budget:block", r, e, s, i, o), g(`tool:${t}:${e.field}`, i, s, e.message));
483
511
  }
484
512
  }
485
513
  s.toolCalls.set(t, l + 1);
@@ -534,9 +562,9 @@ var v = class {
534
562
  }
535
563
  }
536
564
  }
537
- }, y = ({ db: e, config: t }) => {
538
- let n = m(t), r = n.defaults?.costPerInputToken ?? 0;
539
- return new v(e, n, r, n.defaults?.costPerOutputToken ?? 0, n.defaults?.costPerCacheReadToken ?? r, n.defaults?.costPerCacheWriteToken ?? r);
565
+ }, b = ({ db: e, config: t }) => {
566
+ let n = h(t), r = n.defaults?.costPerInputToken ?? 0;
567
+ return new y(e, n, r, n.defaults?.costPerOutputToken ?? 0, n.defaults?.costPerCacheReadToken ?? r, n.defaults?.costPerCacheWriteToken ?? r);
540
568
  };
541
569
  //#endregion
542
- export { l as configSchema, y as createPlugin, y as default, c as pluginConfigSchema };
570
+ export { u as configSchema, b as createPlugin, b as default, l as pluginConfigSchema };
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@adhd/agent-plugin-budget",
3
- "version": "0.2.3",
3
+ "version": "0.2.4",
4
4
  "description": "Budget enforcement plugin for @adhd/agent-mcp — caps token spend, cost, wall-clock time, and model calls per task/session/agent",
5
5
  "type": "module",
6
6
  "publishConfig": {
@@ -17,10 +17,11 @@
17
17
  }
18
18
  },
19
19
  "peerDependencies": {
20
- "@adhd/agent-base-types": "^2.3.3"
20
+ "@adhd/agent-base-types": "^2.3.4"
21
21
  },
22
22
  "dependencies": {
23
- "zod": "*"
23
+ "drizzle-orm": "0.45.2",
24
+ "zod": "^4.4.3"
24
25
  },
25
26
  "keywords": [
26
27
  "agent",