fix: restore VM tool connections and wait for desktop readiness

src/runner/agent.ts

@@ -225,7 +225,8 @@ export class RunnerAgent {
     if(command.kind==='gateway.close'){await this.closeConnection(connectionId);return {ok:true}}
     if(command.kind==='gateway.rpc') {
       const connection=this.connections.get(connectionId),method=String(p.method),params=p.params as Record<string,unknown>
-      if(!connection||!commands.has(method)||!params||typeof params!=='object')throw new HttpError(403,'执行通道或命令无效','runner_command_forbidden')
+      const computerCommand=method==='computer.invoke'||method==='computer.transfer'
+      if(!connection||(!commands.has(method)&&!(computerCommand&&connection.gateway instanceof ComputerGateway))||!params||typeof params!=='object')throw new HttpError(403,'执行通道或命令无效','runner_command_forbidden')
       if(connection.cleanupOnly&&!['session.resume','session.interrupt','session.close'].includes(method))throw new HttpError(403,'清理通道不允许执行新工作','runner_cleanup_forbidden')
       if(method==='session.create'||method==='session.resume'){if(!connection.cleanupOnly)this.requireProfile(params.profile)}
       else if(!connection.sessions.has(String(params.session_id)))throw new HttpError(403,'会话不属于这个通道','runner_session_forbidden')

src/runner/computers/container.ts

@@ -257,6 +257,22 @@ export class ContainerComputerProvider implements ComputerProvider {
       // Persistent workspace deletion is a separate, explicit operation.
     })
   }
+  private async waitForDesktop(state:ComputerState,authorize:()=>void,signal?:AbortSignal):Promise<void> {
+    const deadline=Date.now()+15000
+    for(;;){
+      authorize()
+      if(signal?.aborted)throw new ComputerError('computer_cancelled','操作已停止')
+      try {
+        await this.run(this.runtime,['exec','--user','1000:1000','--',state.containerId,CUA_DRIVER,'call','health_report','{}','--socket',CUA_SOCKET],{timeout:3000,signal})
+        authorize()
+        return
+      }catch(error){
+        if(signal?.aborted)throw new ComputerError('computer_cancelled','操作已停止')
+        if(!/Cua Driver daemon is not running on/.test(String((error as any).stderr??''))||Date.now()>=deadline)throw error
+        await new Promise(resolve=>setTimeout(resolve,250))
+      }
+    }
+  }
   async execute(value:ComputerSpecification,argv:string[],options:{authorize():void;signal?:AbortSignal;timeout?:number;input?:Buffer;lane?:ComputerLane;mayFence?():boolean;user?:ComputerUser}):Promise<CommandResult> {
     const spec=this.validateSpecification(value)
     if(!argv.length||argv.length>64||argv.some(arg=>typeof arg!=='string'||arg.includes('\0'))||argv.join('').length>65536)throw new ComputerError('computer_command_invalid','电脑命令无效或过长')
@@ -265,6 +281,7 @@ export class ContainerComputerProvider implements ComputerProvider {
       const state=await this.inspect(spec)
       if(!state?.running)throw new ComputerError('computer_not_running','电脑环境尚未运行')
       options.authorize()
+      if(argv[0]===CUA_DRIVER)await this.waitForDesktop(state,options.authorize,options.signal)
       if(options.signal?.aborted)throw new ComputerError('computer_cancelled','操作已停止')
       const operationId=randomUUID(),timeout=Math.min(60000,Math.max(1000,options.timeout??30000))
       let cancellation:Promise<void>|undefined
@@ -316,6 +333,7 @@ export class ContainerComputerProvider implements ComputerProvider {
     return this.lane(spec.id).read(async()=>{
       const state=await this.inspect(spec)
       authorize();if(!state?.running)throw new ComputerError('computer_not_running','电脑尚未运行')
+      await this.waitForDesktop(state,authorize)
       const file=`/tmp/yaoyao-view-${randomBytes(12).toString('hex')}.png`
       const exec=(args:string[])=>this.run(this.runtime,['exec','--user','1000:1000','--',state.containerId,...args],{timeout:10000})
       try{

src/server/vmComputer.ts

@@ -26,6 +26,9 @@ export class VmToolSession {
       const gateway=new WorkspaceGateway(this.target,{workId:this.workId,authorize:this.authorize,...(this.publishArtifact?{publishArtifact:this.publishArtifact}:{})})
       this.gateway=gateway
       this.opening=(async()=>{
+        this.authorize()
+        await gateway.connect()
+        this.authorize()
         const opened=await gateway.rpc('session.create',{profile:this.profile,toolsOnly:true,title:'Yaoyao vm tools',source:'yaoyao_workspace',close_on_disconnect:true})
         this.authorize()
         if(this.gateway!==gateway)throw new Error('虚拟环境连接已关闭')

tests/server/computerProvider.test.ts

@@ -4,7 +4,7 @@ import {randomUUID} from 'node:crypto'
 import {mkdtemp,rm,symlink,mkdir,rename,writeFile,realpath} from 'node:fs/promises'
 import {tmpdir} from 'node:os'
 import {join} from 'node:path'
-import {ContainerComputerProvider,type ContainerCommand,type ComputerSpecification} from '../../src/runner/computers/container'
+import {ContainerComputerProvider,CUA_DRIVER,type ContainerCommand,type ComputerSpecification} from '../../src/runner/computers/container'
 import {COMMAND_ROOT,GUEST_COMMAND} from '../../src/runner/computers/commandProcess'
 const homes:string[]=[]
 afterEach(async()=>{await Promise.all(homes.splice(0).map(home=>rm(home,{recursive:true,force:true})))})
@@ -51,6 +51,21 @@ it.each(['docker','podman'] as const)('uses the same private ownership and execu
   await f.provider.remove(f.spec)
   expect(f.calls.find(args=>args[0]==='rm')).toEqual(['rm','--',state.containerId])
 })
+it('waits for the desktop driver on a cold start before sending a GUI command',async()=>{
+  const f=await fixture();await f.provider.ensure(f.spec,()=>{})
+  let checks=0
+  const provider=new ContainerComputerProvider('docker',f.provider.runnerId,f.home,async(runtime,args,options)=>{
+    if(args.includes('health_report')){
+      checks++
+      if(checks<3)throw Object.assign(new Error('driver starting'),{stderr:'Cua Driver daemon is not running on /run/user/1000/openmausbot-cua.sock.'})
+    }
+    return f.run(runtime,args,options)
+  },{})
+  await provider.execute(f.spec,[CUA_DRIVER,'call','get_desktop_state','{}'],{authorize:()=>{},lane:'gui'})
+  expect(checks).toBe(3)
+  expect(f.calls.filter(args=>args.includes('health_report'))).toHaveLength(1)
+  expect(f.calls.some(args=>args.includes('get_desktop_state'))).toBe(true)
+})
 it.each([
   (d:any)=>{d.HostConfig.Privileged=true},
   (d:any)=>{d.HostConfig.CapAdd.push('SYS_ADMIN')},

tests/server/runner.test.ts

@@ -13,6 +13,7 @@ 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 { ComputerGateway } from '../../src/runner/worker/gateway'
 import { WorkspaceGateway, type GatewayTarget, type GatewayFrame } from '../../src/server/workspaceGateway'
 import { createWorkspaceToolLease } from '../../src/server/workspaceToolLease'
 import { HttpError } from '../../src/server/errors'
@@ -72,6 +73,19 @@ async function channel(){
   await gateway.connect();const session=await gateway.rpc('session.create',{profile:'default'})
   return {gateway,session}
 }
+it('admits virtual computer calls only on a computer gateway with an owned session',async()=>{
+  const connectionId=randomUUID(),sessionId=randomUUID(),rpc=vi.fn(async()=>({ok:true}))
+  const computer=Object.assign(Object.create(ComputerGateway.prototype),{rpc,close:()=>{}}) as ComputerGateway
+  const connection={gateway:computer,sessions:new Map([[sessionId,'default']]),running:new Set(),events:Promise.resolve()}
+  ;(runner as any).connections.set(connectionId,connection)
+  const command=(method:string,id=sessionId)=>({kind:'gateway.rpc',payload:{connectionId,method,params:{session_id:id,name:'computer_desktop_state',arguments:{}}}})
+  expect(await (runner as any).execute(command('computer.invoke'))).toEqual({ok:true})
+  expect(await (runner as any).execute(command('computer.transfer'))).toEqual({ok:true})
+  await expect((runner as any).execute(command('computer.invoke',randomUUID()))).rejects.toMatchObject({code:'runner_session_forbidden'})
+  connection.gateway=Object.assign(Object.create(WorkspaceGateway.prototype),{rpc,close:()=>{}}) as ComputerGateway
+  await expect((runner as any).execute(command('computer.invoke'))).rejects.toMatchObject({code:'runner_command_forbidden'})
+  expect(rpc).toHaveBeenCalledTimes(2)
+})
 it('requires an upgraded Runner before granting mixed host and VM tools',()=>{
  const id=randomUUID()
  expect(()=>hub.target('owner','local',{environmentId:id,agentId:id,ownerKey:'owner',hostAccess:true})).toThrow('不支持同时使用本机和虚拟机')

tests/server/vmComputer.test.ts

@@ -0,0 +1,31 @@
+// @vitest-environment node
+import {expect,it} from 'vitest'
+import {WebSocketServer} from 'ws'
+import {VmToolSession} from '../../src/server/vmComputer'
+import type {GatewayTarget} from '../../src/server/workspaceGateway'
+
+it('connects the virtual computer tool channel before creating its session',async()=>{
+  const server=new WebSocketServer({port:0,host:'127.0.0.1'})
+  await new Promise<void>(resolve=>server.once('listening',resolve))
+  const methods:string[]=[]
+  server.on('connection',socket=>{
+    socket.send(JSON.stringify({method:'event',params:{type:'gateway.ready'}}))
+    socket.on('message',raw=>{
+      const frame=JSON.parse(String(raw))
+      methods.push(frame.method)
+      socket.send(JSON.stringify({id:frame.id,result:frame.method==='session.create'?{session_id:'vm-session'}:{ok:true}}))
+    })
+  })
+  const target={url:new URL(`http://127.0.0.1:${(server.address() as {port:number}).port}`),client:{},
+    session:{webSocketCredential:async()=>({name:'ticket',value:'fixture'})}} as unknown as GatewayTarget
+  const session=new VmToolSession(target,'default','work-id',()=>{})
+  try{
+    await expect(session.call('computer_shell',{command:'pwd'})).resolves.toEqual({ok:true})
+    await expect(session.call('computer_desktop_state',{})).resolves.toEqual({ok:true})
+    expect(methods).toEqual(['session.create','computer.invoke','computer.invoke'])
+  }finally{
+    session.close()
+    for(const client of server.clients)client.terminate()
+    await new Promise<void>(resolve=>server.close(()=>resolve()))
+  }
+})