composeDesktops.test.ts 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. // @vitest-environment node
  2. import {it,expect,vi} from 'vitest'
  3. import {randomUUID} from 'node:crypto'
  4. import {mkdtempSync,rmSync} from 'node:fs'
  5. import {tmpdir} from 'node:os'
  6. import {join} from 'node:path'
  7. import {DatabaseSync} from 'node:sqlite'
  8. import Koa from 'koa'
  9. import {bodyParser} from '@koa/bodyparser'
  10. import request from 'supertest'
  11. import {WorkspaceStore} from '../../src/server/workspaceStore'
  12. import {WorkspaceNodes} from '../../src/server/workspaceGateway'
  13. import {SharedComputers} from '../../src/server/sharedComputers'
  14. import {LocalVmService} from '../../src/server/localVm'
  15. import {loadServerConfig} from '../../src/server/config'
  16. import {ComposeDesktops,parseComposeDesktops} from '../../src/server/composeDesktops'
  17. import {ComposeComputerProvider} from '../../src/runner/computers/compose'
  18. import {ComputerPool} from '../../src/runner/computers/pool'
  19. import {COMPOSE_DESKTOP_IMAGE} from '../../src/shared/composeDesktops'
  20. it('accepts only fixed deployment socket paths and rejects undeclared desktop targets',async()=>{
  21. const desktop={id:randomUUID(),name:'共享桌面',socketPath:'/run/yaoyao-desktops/desktop-1/desktop.sock'}
  22. expect(parseComposeDesktops(JSON.stringify([desktop]))).toEqual([desktop])
  23. for(const path of ['/var/run/docker.sock','http://host:1234','/run/yaoyao-desktops/../../docker.sock'])expect(()=>parseComposeDesktops(JSON.stringify([{...desktop,socketPath:path}]))).toThrow()
  24. expect(()=>parseComposeDesktops(JSON.stringify([desktop,desktop]))).toThrow()
  25. await expect(new ComposeDesktops([desktop]).call(randomUUID(),'health')).rejects.toMatchObject({code:'compose_desktop_missing'})
  26. await expect(new ComposeDesktops([desktop]).call(desktop.id,'remove')).rejects.toMatchObject({code:'compose_desktop_managed'})
  27. })
  28. it('locks Compose capacity and binds existing desktops without allowing cross-account reuse',async()=>{
  29. const home=mkdtempSync(join(tmpdir(),'yaoyao-compose-routes-')),store=new WorkspaceStore(home)
  30. const desktops=[1,2].map(n=>({id:randomUUID(),name:`共享桌面 ${n}`,socketPath:`/run/yaoyao-desktops/desktop-${n}/desktop.sock`}))
  31. const config=loadServerConfig({HERMES_YAOYAO_HOME:home,HERMES_YAOYAO_COMPOSE_DESKTOPS:JSON.stringify(desktops)})
  32. const nodes=new WorkspaceNodes(store,config,{} as any),runner={id:randomUUID(),name:'执行端',enabled:true,sourceNodeId:'local',sourceOwner:'_system'}
  33. const auth:any={require:(ctx:any)=>({id:ctx.get('x-owner')||'alice'}),requireAdmin:(ctx:any)=>{if(ctx.get('x-admin')!=='yes')throw Object.assign(new Error('管理员权限'),{status:403});return {id:'alice'}}}
  34. const hub:any={records:()=>[runner],summary:()=>({online:true,features:['compose-desktops-v1']}),composeDesktops:{status:async()=>desktops.map(d=>({...d,online:true,ready:true}))},localVm:vi.fn()}
  35. const shared=new SharedComputers(store,auth,nodes,hub),service=new LocalVmService(store,auth,nodes,hub,shared,config),app=new Koa()
  36. app.use(async(ctx,next)=>{try{await next()}catch(e){ctx.status=(e as any).status||500;ctx.body={code:(e as any).code}}});app.use(bodyParser());app.use(service.router().routes())
  37. try{
  38. expect(await service.status('alice')).toMatchObject({fixedCapacity:true,maxInstances:2,mode:'shared'})
  39. for(const [verb,path,body] of [['post','/prepare',{requestId:randomUUID()}],['put','/policy',{requestId:randomUUID(),maxInstances:4,mode:'per-bot'}],['post',`/instances/${desktops[0]!.id}/remove`,{requestId:randomUUID()}]] as const){await request(app.callback())[verb]('/api/app/admin/local-vm'+path).set('x-admin','yes').send(body).expect(409,{code:'compose_desktop_managed'})}
  40. const first=store.createAgent('alice',{name:'研究员',profile:'default',computer:'cloud',allowHostEnvironment:true}),second=store.createAgent('alice',{name:'写作者',profile:'writer'}),other=store.createAgent('bob',{name:'其他账号',profile:'default'})
  41. const bind=(id:string,desktopId:string,owner='alice')=>request(app.callback()).put(`/api/app/agents/${id}/local-vm`).set('x-owner',owner).send({enabled:true,desktopId})
  42. await bind(first.id,randomUUID()).expect(400)
  43. await bind(first.id,desktops[0]!.id).expect(200);await bind(second.id,desktops[0]!.id).expect(200)
  44. expect(store.require('alice','agent',first.id)).toMatchObject({computer:'vm',execution:'computer',allowHostEnvironment:true})
  45. await request(app.callback()).put(`/api/app/agents/${first.id}/local-vm/image`).send({imageKey:'cursor'}).expect(409,{code:'compose_desktop_managed'})
  46. for(const id of [first.id,second.id])expect(()=>nodes.requireSource('alice',store.require('alice','agent',id))).not.toThrow()
  47. await bind(other.id,desktops[0]!.id,'bob').expect(403)
  48. expect((await service.status('bob')).desktops?.[0]?.available).toBe(false)
  49. for(const op of ['create','start','stop','recreate','remove'])await request(app.callback()).post(`/api/app/agents/${first.id}/local-vm/${op}`).send({}).expect(409)
  50. const child=store.createAgent('alice',{name:'团队新成员',profile:'default',execution:'computer'},{createdByAgentId:first.id,createdFromRunId:randomUUID()})
  51. expect(child.computerEnvironmentId).toBe(desktops[0]!.id)
  52. expect(hub.localVm).not.toHaveBeenCalled()
  53. expect((await service.status('alice')).maxInstances).toBe(2)
  54. }finally{nodes.close();await service.close();store.close();rmSync(home,{recursive:true,force:true})}
  55. })
  56. it('reuses a fixed desktop, rotates authorization and never creates or deletes its container',async()=>{
  57. const id=randomUUID(),db=new DatabaseSync(':memory:'),operations:any[]=[],tokens:string[]=[]
  58. const relay=vi.fn(async(target:string,op:string,body:any)=>{operations.push({target,op,body});if(op==='list')return [{id,name:'固定桌面',ready:true}];if(target!==id)throw new Error('unknown desktop');if(op==='health')return {instance:'boot-1',ready:true};if(op==='acquire'){const token=randomUUID();tokens.push(token);return {token,instance:'boot-1'}};if(op==='execute')return {stdout:'ok',stderr:'',exitCode:0};return {ok:true}})
  59. const provider=new ComposeComputerProvider(randomUUID(),'/fixture',relay),engine=vi.spyOn(provider,'run'),pool=new ComputerPool(db,provider),spec={id,ownerKey:'owner',imageId:COMPOSE_DESKTOP_IMAGE}
  60. try{
  61. await pool.recover()
  62. const first=await pool.acquire(spec,'agent-a',()=>{})
  63. await pool.use(first,c=>provider.execute(spec,['pwd'],c))
  64. await pool.use(first,c=>provider.execute(spec,['id','-u'],{...c,user:'root'}))
  65. await pool.release(first,true)
  66. expect(()=>pool.use(first,async()=>{})).toThrow()
  67. const second=await pool.acquire(spec,'agent-b',()=>{})
  68. expect(tokens[0]).not.toBe(tokens[1])
  69. await pool.release(second)
  70. expect(operations.filter(x=>x.op==='release').map(x=>x.body.cancel)).toEqual([false,true])
  71. expect(operations.filter(x=>x.op==='execute').map(x=>x.body.user)).toEqual(['cua','root'])
  72. expect(engine).not.toHaveBeenCalled()
  73. }finally{await pool.close();db.close()}
  74. })
  75. it('serializes Compose desktop input while keeping command work concurrent and skipping cancelled queued input',async()=>{
  76. const spec={id:randomUUID(),ownerKey:'owner',imageId:COMPOSE_DESKTOP_IMAGE},calls:string[]=[],held=new Map<string,()=>void>()
  77. const relay=vi.fn(async(_target:string,op:string,body:any)=>{
  78. if(op==='health')return {instance:'boot-1',ready:true}
  79. if(op==='acquire')return {token:'lease',instance:'boot-1'}
  80. if(op==='execute'){
  81. const name=String(body.argv[0]);calls.push(name)
  82. if(name.startsWith('hold'))await new Promise<void>(resolve=>held.set(name,resolve))
  83. return {stdout:name,stderr:'',exitCode:0}
  84. }
  85. return {ok:true}
  86. })
  87. const provider=new ComposeComputerProvider(randomUUID(),'/fixture',relay),abort=new AbortController()
  88. await provider.ensure(spec,()=>{})
  89. const gui=provider.execute(spec,['hold-gui'],{authorize:()=>{},lane:'gui'})
  90. await vi.waitFor(()=>expect(calls).toContain('hold-gui'))
  91. const cancelled=provider.execute(spec,['cancelled-gui'],{authorize:()=>{},lane:'gui',signal:abort.signal}).catch(error=>error)
  92. const next=provider.execute(spec,['next-gui'],{authorize:()=>{},lane:'gui'})
  93. const shell=provider.execute(spec,['hold-shell'],{authorize:()=>{}})
  94. try{
  95. await provider.execute(spec,['parallel-shell'],{authorize:()=>{}})
  96. expect(calls).toEqual(['hold-gui','hold-shell','parallel-shell'])
  97. abort.abort();held.get('hold-gui')!()
  98. await gui;await next
  99. expect(await cancelled).toMatchObject({code:'computer_cancelled'})
  100. expect(calls).not.toContain('cancelled-gui')
  101. expect(calls).toContain('next-gui')
  102. }finally{for(const release of held.values())release();await Promise.allSettled([gui,next,cancelled,shell])}
  103. })
  104. it('does not return a successful Compose command after its holder cancels',async()=>{
  105. const spec={id:randomUUID(),ownerKey:'owner',imageId:COMPOSE_DESKTOP_IMAGE},abort=new AbortController()
  106. const relay=vi.fn(async(_target:string,op:string)=>{
  107. if(op==='health')return {instance:'boot-1',ready:true}
  108. if(op==='acquire')return {token:'lease',instance:'boot-1'}
  109. if(op==='execute'){abort.abort();return {stdout:'late result',stderr:'',exitCode:0}}
  110. if(op==='cancel')return {ok:true,stopped:true}
  111. return {ok:true}
  112. })
  113. const provider=new ComposeComputerProvider(randomUUID(),'/fixture',relay)
  114. await provider.ensure(spec,()=>{})
  115. await expect(provider.execute(spec,['work'],{authorize:()=>{},signal:abort.signal,mayFence:()=>false})).rejects.toMatchObject({code:'computer_cancelled'})
  116. expect(relay.mock.calls.some(([,op])=>op==='release')).toBe(false)
  117. })
  118. it('cancels the in-flight Compose operation by ID and waits for a confirmed stop without releasing its shared lease',async()=>{
  119. const spec={id:randomUUID(),ownerKey:'owner',imageId:COMPOSE_DESKTOP_IMAGE},abort=new AbortController()
  120. let acknowledge:()=>void=()=>{}
  121. const relay=vi.fn(async(_target:string,op:string,body:any)=>{
  122. if(op==='health')return {instance:'boot-1',ready:true}
  123. if(op==='acquire')return {token:'lease',instance:'boot-1'}
  124. if(op==='execute')return new Promise(()=>{})
  125. if(op==='cancel'){await new Promise<void>(resolve=>{acknowledge=resolve});return {ok:true,stopped:true,operationId:body.operationId}}
  126. return {ok:true}
  127. })
  128. const provider=new ComposeComputerProvider(randomUUID(),'/fixture',relay)
  129. await provider.ensure(spec,()=>{})
  130. let finished=false
  131. const execution=provider.execute(spec,['work'],{authorize:()=>{},signal:abort.signal,mayFence:()=>false}).catch(error=>{finished=true;return error})
  132. await vi.waitFor(()=>expect(relay.mock.calls.some(([,op])=>op==='execute')).toBe(true))
  133. abort.abort()
  134. await vi.waitFor(()=>expect(relay.mock.calls.some(([,op])=>op==='cancel')).toBe(true))
  135. expect(finished).toBe(false)
  136. const execute=relay.mock.calls.find(([,op])=>op==='execute')![2],cancel=relay.mock.calls.find(([,op])=>op==='cancel')![2]
  137. expect(cancel).toMatchObject({token:'lease',instance:'boot-1',operationId:execute.operationId})
  138. expect(execute.operationId).toMatch(/^[a-f0-9-]{36}$/)
  139. acknowledge()
  140. expect(await execution).toMatchObject({code:'computer_cancelled'})
  141. expect(relay.mock.calls.some(([,op])=>op==='release')).toBe(false)
  142. })
  143. it('reports an unconfirmed Compose cancellation without claiming that shared work has stopped',async()=>{
  144. const spec={id:randomUUID(),ownerKey:'owner',imageId:COMPOSE_DESKTOP_IMAGE},abort=new AbortController()
  145. const relay=vi.fn(async(_target:string,op:string)=>{
  146. if(op==='health')return {instance:'boot-1',ready:true}
  147. if(op==='acquire')return {token:'lease',instance:'boot-1'}
  148. if(op==='execute'){abort.abort();return {stdout:'late',stderr:'',exitCode:0}}
  149. if(op==='cancel')throw new Error('disconnected')
  150. return {ok:true}
  151. })
  152. const provider=new ComposeComputerProvider(randomUUID(),'/fixture',relay)
  153. await provider.ensure(spec,()=>{})
  154. await expect(provider.execute(spec,['work'],{authorize:()=>{},signal:abort.signal,mayFence:()=>false})).rejects.toMatchObject({code:'computer_cancel_uncertain'})
  155. expect(relay.mock.calls.some(([,op])=>op==='release')).toBe(false)
  156. })