|
|
@@ -0,0 +1,231 @@
|
|
|
+// @vitest-environment node
|
|
|
+import { afterEach, beforeEach, expect, it, vi } from 'vitest'
|
|
|
+import Koa from 'koa'
|
|
|
+import { createServer, type Server } from 'node:http'
|
|
|
+import { mkdtempSync, rmSync, writeFileSync } from 'node:fs'
|
|
|
+import { spawn } from 'node:child_process'
|
|
|
+import { once } from 'node:events'
|
|
|
+import { tmpdir } from 'node:os'
|
|
|
+import { join } from 'node:path'
|
|
|
+import { randomUUID } from 'node:crypto'
|
|
|
+import { WebSocketServer, type WebSocket } from 'ws'
|
|
|
+import { WorkspaceStore } from '../../src/server/workspaceStore'
|
|
|
+import { LocalAuthStore } from '../../src/server/localAuth'
|
|
|
+import { RunnerHub } from '../../src/server/runnerHub'
|
|
|
+import { RunnerAgent } from '../../src/runner/agent'
|
|
|
+import { WorkspaceGateway, type GatewayTarget, type GatewayFrame } from '../../src/server/workspaceGateway'
|
|
|
+import { createWorkspaceToolLease } from '../../src/server/workspaceToolLease'
|
|
|
+import { HttpError } from '../../src/server/errors'
|
|
|
+import type { RunnerCommand, RunnerConfiguration } from '../../src/shared/runner'
|
|
|
+
|
|
|
+class Auth extends LocalAuthStore {
|
|
|
+ version=1
|
|
|
+ allowed=true
|
|
|
+ override isUserActive(){return true}
|
|
|
+ override pushAuthorizationVersion(){return this.version}
|
|
|
+ override canUseSource(){return this.allowed}
|
|
|
+}
|
|
|
+let home:string,store:WorkspaceStore,auth:Auth,hub:RunnerHub,runner:RunnerAgent,config:RunnerConfiguration
|
|
|
+let web:Server,hermes:Server,ws:WebSocketServer,controller:AbortController,running:Promise<void>
|
|
|
+let target:GatewayTarget,bindings:any[],rpc:Array<{method:string;params:any}>,frames:GatewayFrame[],sockets:WebSocket[]
|
|
|
+const gateways:WorkspaceGateway[]=[]
|
|
|
+const listen=(server:Server)=>new Promise<string>(resolve=>server.listen(0,'127.0.0.1',()=>resolve(`http://127.0.0.1:${(server.address() as {port:number}).port}`)))
|
|
|
+const close=(server:Server)=>new Promise<void>(resolve=>{server.close(()=>resolve());server.closeAllConnections()})
|
|
|
+beforeEach(async()=>{
|
|
|
+ home=mkdtempSync(join(tmpdir(),'yaoyao-runner-'));store=new WorkspaceStore(home);auth=new Auth(home)
|
|
|
+ bindings=[];rpc=[];frames=[];sockets=[]
|
|
|
+ hermes=createServer(async(req,res)=>{
|
|
|
+ let body='';for await(const part of req)body+=part
|
|
|
+ if(req.url?.endsWith('/bind'))bindings.push(JSON.parse(body))
|
|
|
+ res.setHeader('Content-Type','application/json')
|
|
|
+ res.end(JSON.stringify(req.url==='/api/auth/ws-ticket'?{ticket:'fixture'}:req.url==='/api/profiles'?{profiles:['default',{name:'writer'},{name:'secret'}]}:req.url?.endsWith('/capabilities')?{version:1,ready:true,in_process:true,native_tools:true}:{ok:true,native_tools:true,terminal:{cwd:'/fixture/workspace'}}))
|
|
|
+ })
|
|
|
+ const hermesURL=await listen(hermes)
|
|
|
+ ws=new WebSocketServer({server:hermes})
|
|
|
+ ws.on('connection',socket=>{
|
|
|
+ sockets.push(socket);socket.send(JSON.stringify({method:'event',params:{type:'gateway.ready'}}))
|
|
|
+ socket.on('message',raw=>{
|
|
|
+ const value=JSON.parse(String(raw));rpc.push(value)
|
|
|
+ const result=['session.create','session.resume'].includes(value.method)?{session_id:randomUUID(),stored_session_id:randomUUID(),info:{profile_name:value.params.profile}}:{ok:true}
|
|
|
+ socket.send(JSON.stringify({id:value.id,result}))
|
|
|
+ })
|
|
|
+ })
|
|
|
+ hub=new RunnerHub(store,auth,{url:new URL(hermesURL)} as GatewayTarget)
|
|
|
+ const registered=hub.enroll('owner',{name:'测试电脑',allowedProfiles:['default','writer']})
|
|
|
+ const app=new Koa();app.use(async(ctx,next)=>{try{await next()}catch(error){ctx.status=error instanceof HttpError?error.status:500;ctx.body={code:error instanceof HttpError?error.code:'error'}}});app.use(hub.middleware())
|
|
|
+ web=createServer(app.callback());const serverURL=await listen(web)
|
|
|
+ config={protocol:1,serverURL,runnerId:registered.runner.id,token:registered.token,hermesURL,allowedProfiles:['default','writer'],artifactRoots:[]}
|
|
|
+ runner=new RunnerAgent(config,home);controller=new AbortController();running=runner.run(controller.signal);void running.catch(()=>{})
|
|
|
+ await vi.waitFor(()=>expect(hub.summary(hub.records()[0]!).online).toBe(true))
|
|
|
+ await vi.waitFor(()=>expect((runner as any).serverEpoch).toBeTruthy())
|
|
|
+ target=hub.target('owner','local')!
|
|
|
+})
|
|
|
+afterEach(async()=>{
|
|
|
+ for(const gateway of gateways.splice(0))gateway.close()
|
|
|
+ controller.abort();await running.catch(()=>{});hub.close()
|
|
|
+ for(const socket of ws.clients)socket.terminate()
|
|
|
+ await new Promise<void>(resolve=>ws.close(()=>resolve()));await close(hermes);await close(web)
|
|
|
+ store.close();rmSync(home,{recursive:true,force:true})
|
|
|
+})
|
|
|
+async function channel(){
|
|
|
+ const gateway=new WorkspaceGateway(target);gateways.push(gateway);gateway.onEvent=frame=>frames.push(frame)
|
|
|
+ await gateway.connect();const session=await gateway.rpc('session.create',{profile:'default'})
|
|
|
+ return {gateway,session}
|
|
|
+}
|
|
|
+async function machine(path:string,body?:unknown,overrides:Record<string,string>={}){
|
|
|
+ return fetch(`${config.serverURL}/api/runner/v1/${config.runnerId}/${path}`,{
|
|
|
+ method:body===undefined?'GET':'POST',headers:{Authorization:`Bearer ${config.token}`,'x-runner-instance':runner.instance,'x-runner-protocol':'1','x-runner-epoch':(runner as any).serverEpoch,'Content-Type':'application/json',...overrides},...(body===undefined?{}:{body:JSON.stringify(body)})
|
|
|
+ })
|
|
|
+}
|
|
|
+it('connects through outbound transport, filters profiles and forwards only owned session events',async()=>{
|
|
|
+ const profiles=await target.session.request('/api/profiles')
|
|
|
+ expect(JSON.parse(profiles.body.toString()).profiles).toEqual(['default',{name:'writer'}])
|
|
|
+ const {gateway,session}=await channel()
|
|
|
+ await expect(gateway.rpc('session.create',{profile:'secret'})).rejects.toMatchObject({code:'runner_profile_forbidden'})
|
|
|
+ await expect(gateway.rpc('prompt.submit',{session_id:'someone-else',text:'no'})).rejects.toMatchObject({code:'runner_session_forbidden'})
|
|
|
+ await gateway.rpc('prompt.submit',{session_id:session.session_id,text:'hello'})
|
|
|
+ sockets[0]!.send(JSON.stringify({method:'event',params:{type:'message.delta',session_id:'someone-else',payload:{text:'private'}}}))
|
|
|
+ sockets[0]!.send(JSON.stringify({method:'event',params:{type:'message.delta',session_id:session.session_id,payload:{text:'hello'}}}))
|
|
|
+ await vi.waitFor(()=>expect(frames.some(f=>f.payload?.text==='hello')).toBe(true))
|
|
|
+ expect(frames.some(f=>f.payload?.text==='private')).toBe(false)
|
|
|
+})
|
|
|
+it('mounts loopback-only tools with stable call ids and revokes an active grant',async()=>{
|
|
|
+ const {session}=await channel(),abort=new AbortController(),call=vi.fn(async()=>({agent:{id:'created'}}))
|
|
|
+ let active=true
|
|
|
+ const lease=await createWorkspaceToolLease({target,profile:'default',workId:'work',session:()=>({runtimeId:session.session_id,storedId:session.stored_session_id}),signal:abort.signal,
|
|
|
+ assertActive:()=>{if(!active)throw new HttpError(403,'撤权','revoked')},catalog:()=>[{id:'workspace_create_agent',name:'workspace_create_agent',description:'create',inputSchema:{type:'object'}}],call,onFailure:()=>{}})
|
|
|
+ await lease.bind();const binding=bindings[0]
|
|
|
+ const toolHTTP=(path:string,body:unknown)=>fetch(binding.bridge_url+path,{method:'POST',headers:{Authorization:`Bearer ${binding.token}`,'Content-Type':'application/json'},body:JSON.stringify(body)})
|
|
|
+ const catalog=await (await toolHTTP('/tools/list',{})).json() as any
|
|
|
+ const body={callId:'same-call',toolId:catalog.tools[0].id,arguments:{name:'worker'}}
|
|
|
+ const first=await (await toolHTTP('/tools/call',body)).json()
|
|
|
+ expect(await (await toolHTTP('/tools/call',body)).json()).toEqual(first)
|
|
|
+ expect(call).toHaveBeenCalledExactlyOnceWith('workspace_create_agent',{name:'worker'},'same-call')
|
|
|
+ active=false
|
|
|
+ expect(await (await toolHTTP('/tools/call',{...body,callId:'late-call'})).json()).toMatchObject({isError:true})
|
|
|
+ expect(call).toHaveBeenCalledTimes(1)
|
|
|
+ await lease.dispose()
|
|
|
+})
|
|
|
+it('denies browser credentials, stale epochs, foreign lease ids and unauthorized HTTP operations',async()=>{
|
|
|
+ expect((await machine('admit',{id:randomUUID()},{Origin:config.serverURL})).status).toBe(403)
|
|
|
+ expect((await machine('admit',{id:randomUUID()},{Authorization:'Bearer wrong'})).status).toBe(403)
|
|
|
+ expect((await machine('admit',{id:randomUUID()},{'x-runner-epoch':'old'})).status).toBe(409)
|
|
|
+ expect((await machine('tool',{leaseId:randomUUID(),callId:'test',toolId:'workspace_create_agent',arguments:{}})).status).toBe(410)
|
|
|
+ await expect(target.session.request('/api/config',{method:'POST',body:{}})).rejects.toMatchObject({code:'runner_http_forbidden'})
|
|
|
+ await expect(target.session.request('/api/files/download',{search:new URLSearchParams({profile:'default',path:'/etc/hosts'})})).rejects.toMatchObject({code:'runner_artifact_forbidden'})
|
|
|
+})
|
|
|
+it('does not execute a queued command after account authorization changes',async()=>{
|
|
|
+ const {gateway,session}=await channel()
|
|
|
+ auth.version++
|
|
|
+ await expect(gateway.rpc('prompt.submit',{session_id:session.session_id,text:'late'})).rejects.toMatchObject({code:'runner_command_not_admitted'})
|
|
|
+ expect(rpc.filter(r=>r.method==='prompt.submit')).toHaveLength(0)
|
|
|
+})
|
|
|
+it('keeps one result for a command and refuses conflicting or unfinished replay after restart',async()=>{
|
|
|
+ const seen:RunnerCommand[]=[];const execute=(runner as any).execute.bind(runner)
|
|
|
+ vi.spyOn(runner as any,'execute').mockImplementation(async(command:any)=>{seen.push(command);return execute(command)})
|
|
|
+ await channel()
|
|
|
+ const command=seen.find(c=>c.kind==='gateway.rpc')!
|
|
|
+ const before=rpc.length
|
|
|
+ expect(await runner.handle(command)).toHaveProperty('result.session_id')
|
|
|
+ expect(rpc).toHaveLength(before)
|
|
|
+ await expect(runner.handle({...command,payload:{changed:true}})).rejects.toMatchObject({code:'idempotency_conflict'})
|
|
|
+ const incomplete={...command,id:randomUUID()}
|
|
|
+ const {createHash}=await import('node:crypto')
|
|
|
+ ;(runner as any).db.prepare("INSERT INTO commands VALUES(?,?,'started',NULL,?)").run(incomplete.id,createHash('sha256').update(JSON.stringify(incomplete)).digest('hex'),Date.now())
|
|
|
+ controller.abort();await running
|
|
|
+ runner=new RunnerAgent(config,home);controller=new AbortController();running=runner.run(controller.signal);void running.catch(()=>{})
|
|
|
+ await expect(runner.handle(incomplete)).rejects.toMatchObject({code:'runner_command_uncertain'})
|
|
|
+ expect(await runner.handle(command)).toHaveProperty('result.session_id')
|
|
|
+ expect(rpc).toHaveLength(before)
|
|
|
+})
|
|
|
+it('closes a pending gateway before a delayed handshake can revive it',async()=>{
|
|
|
+ const connect=WorkspaceGateway.prototype.connect
|
|
|
+ let release!:()=>void
|
|
|
+ const gate=new Promise<void>(resolve=>{release=resolve})
|
|
|
+ vi.spyOn(WorkspaceGateway.prototype,'connect').mockImplementationOnce(async function(this:WorkspaceGateway){await gate;return connect.call(this)})
|
|
|
+ const connectionId=randomUUID(),request=(hub as any).request.bind(hub)
|
|
|
+ // The control plane normally installs this owner record before opening.
|
|
|
+ ;(hub as any).connections.set(connectionId,{runnerId:config.runnerId,valid:()=>true,onEvent:()=>{},onDisconnect:()=>{}})
|
|
|
+ const opening=request(config.runnerId,'gateway.open',{connectionId}).catch((error:unknown)=>error)
|
|
|
+ await vi.waitFor(()=>expect((runner as any).connections.has(connectionId)).toBe(true))
|
|
|
+ await request(config.runnerId,'gateway.close',{connectionId})
|
|
|
+ release()
|
|
|
+ expect(await opening).toBeInstanceOf(Error)
|
|
|
+ expect((runner as any).connections.has(connectionId)).toBe(false)
|
|
|
+})
|
|
|
+it('fences a replaced process and does not route to the host when its runner goes offline',async()=>{
|
|
|
+ const {gateway}=await channel();const disconnected=vi.fn();gateway.onDisconnect=disconnected
|
|
|
+ const oldInstance=runner.instance
|
|
|
+ controller.abort();await running
|
|
|
+ runner=new RunnerAgent(config,home);controller=new AbortController();running=runner.run(controller.signal);void running.catch(()=>{})
|
|
|
+ await vi.waitFor(()=>expect(disconnected).toHaveBeenCalled())
|
|
|
+ expect((await machine('poll',undefined,{'x-runner-instance':oldInstance})).status).toBe(409)
|
|
|
+ hub.remove(config.runnerId)
|
|
|
+ await expect(target.session.request('/api/status')).rejects.toMatchObject({code:'runner_offline'})
|
|
|
+})
|
|
|
+
|
|
|
+it('resets the hub epoch after a transport failure and closes the obsolete gateways',async()=>{
|
|
|
+ const {gateway}=await channel(),disconnected=vi.fn();gateway.onDisconnect=disconnected
|
|
|
+ const epoch=(runner as any).serverEpoch
|
|
|
+ await (runner as any).disconnect()
|
|
|
+ // Wake the outstanding poll. Its response must be ignored after disconnect.
|
|
|
+ ;(hub as any).online.get(config.runnerId).wake?.()
|
|
|
+ await vi.waitFor(()=>expect((runner as any).serverEpoch).not.toBe(epoch))
|
|
|
+ expect(disconnected).toHaveBeenCalledTimes(1)
|
|
|
+ const next=await channel()
|
|
|
+ expect(await next.gateway.rpc('session.usage',{session_id:next.session.session_id})).toEqual({ok:true})
|
|
|
+})
|
|
|
+it('a closed account channel cannot interrupt an unrelated connection',async()=>{
|
|
|
+ const first=await channel(),second=await channel()
|
|
|
+ const entries=[...(hub as any).connections.entries()] as Array<[string,any]>
|
|
|
+ entries[0]![1].valid=()=>false
|
|
|
+ sockets[0]!.send(JSON.stringify({method:'event',params:{type:'message.delta',session_id:first.session.session_id,payload:{text:'revoked'}}}))
|
|
|
+ await vi.waitFor(()=>expect((runner as any).connections.size).toBe(1))
|
|
|
+ expect((runner as any).connected).toBe(true)
|
|
|
+ expect(await second.gateway.rpc('prompt.submit',{session_id:second.session.session_id,text:'continue'})).toEqual({ok:true})
|
|
|
+ expect(rpc.filter(r=>r.method==='session.interrupt'&&r.params.session_id===second.session.session_id)).toHaveLength(0)
|
|
|
+})
|
|
|
+it('rechecks admission after a queued prompt loses its account permission',async()=>{
|
|
|
+ const {gateway,session}=await channel()
|
|
|
+ const api=(runner as any).api.bind(runner)
|
|
|
+ let release!:()=>void,waiting=false
|
|
|
+ const gate=new Promise<void>(resolve=>{release=resolve})
|
|
|
+ vi.spyOn(runner as any,'api').mockImplementation(async(path:any,body:any)=>{if(path==='admit'){waiting=true;await gate}return api(path,body)})
|
|
|
+ const result=gateway.rpc('prompt.submit',{session_id:session.session_id,text:'must not run'}).catch(error=>error)
|
|
|
+ await vi.waitFor(()=>expect(waiting).toBe(true));auth.version++;release()
|
|
|
+ expect(await result).toMatchObject({code:'runner_command_not_admitted'})
|
|
|
+ expect(rpc.filter(r=>r.method==='prompt.submit')).toHaveLength(0)
|
|
|
+})
|
|
|
+it('does not recreate a lease when cancellation overtakes its admitted creation',async()=>{
|
|
|
+ const {session}=await channel(),abort=new AbortController()
|
|
|
+ const execute=(runner as any).execute.bind(runner)
|
|
|
+ let release!:()=>void,waiting=false
|
|
|
+ const gate=new Promise<void>(resolve=>{release=resolve})
|
|
|
+ vi.spyOn(runner as any,'execute').mockImplementation(async(command:any)=>{if(command.kind==='lease.create'){waiting=true;await gate}return execute(command)})
|
|
|
+ const result=createWorkspaceToolLease({target,profile:'default',workId:'cancel-race',session:()=>({runtimeId:session.session_id,storedId:session.stored_session_id}),signal:abort.signal,
|
|
|
+ assertActive:()=>{if(abort.signal.aborted)throw new Error('stopped')},catalog:()=>[],call:async()=>({}),onFailure:()=>{}}).catch(error=>error)
|
|
|
+ await vi.waitFor(()=>expect(waiting).toBe(true));abort.abort()
|
|
|
+ await vi.waitFor(()=>expect((runner as any).closedLeases.size).toBe(1));release()
|
|
|
+ expect(await result).toBeInstanceOf(Error)
|
|
|
+ expect((runner as any).leases.size).toBe(0)
|
|
|
+ expect(bindings).toHaveLength(0)
|
|
|
+})
|
|
|
+
|
|
|
+it.each([0o600,0o644])('runs the CLI with private configuration and rejects broad permissions (%s)',async(mode)=>{
|
|
|
+ controller.abort();await running
|
|
|
+ const file=join(home,'runner.json');writeFileSync(file,JSON.stringify(config),{mode})
|
|
|
+ const bundle=process.env.RUNNER_BUNDLE_PATH
|
|
|
+ const args=bundle?[bundle,'--config',file]:['--import','tsx','src/runner/index.ts','--config',file]
|
|
|
+ const child=spawn(process.execPath,args,{stdio:['ignore','ignore','pipe']})
|
|
|
+ let output='';child.stderr?.on('data',data=>{output+=String(data)})
|
|
|
+ const exited=once(child,'exit')
|
|
|
+ try {
|
|
|
+ if(mode===0o644){expect((await exited)[0]).not.toBe(0);expect(output).toContain('chmod 600');return}
|
|
|
+ await vi.waitFor(()=>{if(child.exitCode!==null)throw new Error(`Runner 提前退出 ${child.exitCode}: ${output}`);expect(output).toContain('执行节点已连接')},{timeout:7000})
|
|
|
+ expect((await target.session.request('/api/profiles')).status).toBe(200)
|
|
|
+ child.kill('SIGTERM')
|
|
|
+ expect((await exited)[0]).toBe(0)
|
|
|
+ expect(output).toContain('执行节点已停止')
|
|
|
+ expect(output).not.toContain(config.token)
|
|
|
+ }finally{if(child.exitCode===null&&child.signalCode===null){child.kill('SIGTERM');await exited}}
|
|
|
+},15000)
|
LSamien