46 lines
5.5 KiB
JavaScript
46 lines
5.5 KiB
JavaScript
if(!process.send)throw Error('Owned profile load client requires an IPC parent')
|
|
function requireLoopback(value){const url=new URL(value);if(url.protocol!=='http:'||!['127.0.0.1','localhost','[::1]'].includes(url.hostname)||url.username||url.password)throw Error('Owned test client only permits localhost HTTP')}
|
|
import {Client,StreamableHTTPClientTransport} from '@modelcontextprotocol/client'
|
|
import {getHeapStatistics} from 'node:v8'
|
|
const clients=new Map(),requests=new Map(),streams=new Map();let serial=0,peakRss=0,peakHeap=0
|
|
const sample=()=>{const m=process.memoryUsage();peakRss=Math.max(peakRss,m.rss);peakHeap=Math.max(peakHeap,m.heapUsed);return{memory:m,peakRss,peakHeap,maxRssBytes:process.resourceUsage().maxRSS*1024,effectiveHeapSizeLimitBytes:getHeapStatistics().heap_size_limit}}
|
|
const timer=setInterval(sample,20)
|
|
const summarize=result=>{const row=result.structuredContent?.result??{};return{isError:result.isError===true,structuredContent:{result:{dataQuality:row.dataQuality??null,collectorQualityUnavailableReason:row.collectorQualityUnavailableReason,retainedRowsBytes:row.retainedRowsBytes,analyticsBytes:row.analytics?.bytes,eventBytes:row.events?.map(x=>x.data?.bytes),exporters:row.exporters?.length,clients:row.clients?.length,month:row.month}},hiddenSourceLeaked:JSON.stringify(result).includes('HIDDEN_MEMORY_SOURCE'),encodedBytes:Buffer.byteLength(JSON.stringify(result)),error:result.isError?result.content:undefined}}
|
|
async function run(message){
|
|
const p=message.payload??{}
|
|
if(p.endpoint)requireLoopback(p.endpoint);if(p.url)requireLoopback(p.url)
|
|
if(message.operation==='connect'){
|
|
const client=new Client({name:'external-owned-memory-client',version:'1'});await client.connect(new StreamableHTTPClientTransport(new URL(p.endpoint+'/mcp'),{requestInit:{headers:{authorization:'Bearer '+p.token}}}));const id=++serial;clients.set(id,client);return{id}
|
|
}
|
|
if(message.operation==='tool')return summarize(await clients.get(p.clientId).callTool(p.input))
|
|
if(message.operation==='close'){await clients.get(p.clientId)?.close();clients.delete(p.clientId);return{closed:true}}
|
|
if(message.operation==='sse-open'){
|
|
const abort=new AbortController(),response=await fetch(p.url,{signal:abort.signal});if(response.status!==200)throw Error('SSE status '+response.status)
|
|
const reader=response.body.getReader(),decoder=new TextDecoder();let text=''
|
|
for(;;){const chunk=await reader.read();if(chunk.done)throw Error('SSE closed before sample');text+=decoder.decode(chunk.value,{stream:true});const match=/event: (sample|error)\ndata: ([^\n]+)\n\n/.exec(text);if(match){if(match[1]!=='sample')throw Error('SSE source unavailable');const body=JSON.parse(match[2]),id=++serial;streams.set(id,{abort,reader});return{id,bytes:body.bytes}}}
|
|
}
|
|
if(message.operation==='sse-close'){const stream=streams.get(p.streamId);stream?.abort.abort();await stream?.reader.cancel().catch(()=>{});streams.delete(p.streamId);return{closed:true}}
|
|
if(message.operation==='http'){
|
|
const abort=new AbortController();requests.set(message.id,abort)
|
|
try{const response=await fetch(p.url,{...p.options,signal:abort.signal}),body=await response.arrayBuffer();let summary;if(response.status===200&&(p.url.includes('/exporters')||p.url.includes('/analytics'))){const parsed=JSON.parse(new TextDecoder().decode(body));summary={exporters:parsed.exporters?.length,bytes:parsed.bytes}}return{status:response.status,bytes:body.byteLength,summary,text:body.byteLength<4096?new TextDecoder().decode(body):''}}
|
|
finally{requests.delete(message.id)}
|
|
}
|
|
if(message.operation==='batch'){
|
|
const response=await fetch(p.endpoint+'/mcp',{method:'POST',headers:{authorization:'Bearer '+p.token,'content-type':'application/json',accept:'application/json, text/event-stream','mcp-protocol-version':'2025-06-18'},body:JSON.stringify([1,2].map(id=>({jsonrpc:'2.0',id,method:'tools/call',params:{name:'mm_netflow_analytics_report_get',arguments:{serverId:1,minutes:5}}})))})
|
|
if(response.status!==200)throw Error('Batch HTTP status '+response.status)
|
|
const reader=response.body.getReader(),decoder=new TextDecoder(),messages=[];let text=''
|
|
for(;;){const chunk=await reader.read();if(chunk.done)break;text+=decoder.decode(chunk.value,{stream:true});let end;while((end=text.indexOf('\n\n'))>=0){const frame=text.slice(0,end);text=text.slice(end+2);for(const line of frame.split('\n'))if(line.startsWith('data: ')){const row=JSON.parse(line.slice(6));messages.push(row);if(row.error||row.result?.isError)process.send?.({type:'client-progress',id:message.id,refusals:1})}}}
|
|
if(messages.length!==2)throw Error('Expected two batch replies')
|
|
return{successes:messages.filter(row=>!row.error&&!row.result?.isError).length,refusals:messages.filter(row=>row.error||row.result?.isError).length}
|
|
}
|
|
throw Error('Unknown external client operation')
|
|
}
|
|
process.on('message',async message=>{
|
|
if(message.type==='abort'){requests.get(message.requestId)?.abort();return}
|
|
if(message.type==='shutdown'){for(const abort of requests.values())abort.abort();for(const stream of streams.values())stream.abort.abort();for(const client of clients.values())try{await client.close()}catch{}clearInterval(timer);process.send?.({type:'client-report',report:sample()},()=>process.exit(0));return}
|
|
if(message.type!=='client-command')return
|
|
try{const result=await run(message);process.send?.({type:'client-result',id:message.id,ok:true,result})}
|
|
catch(error){process.send?.({type:'client-result',id:message.id,ok:false,error:error instanceof Error?error.message:String(error)})}
|
|
})
|
|
process.send?.({type:'client-ready',report:sample()})
|