bandit-runner/bandit-runner-app/scripts/patch-worker.js

#!/usr/bin/env node
/**
 * Patch the OpenNext worker to export Durable Objects
 * Directly inlines the DO code into the worker
 */

const fs = require('fs')
const path = require('path')

console.log('🔨 Patching worker to export Durable Object...')

const workerPath = path.join(__dirname, '../.open-next/worker.js')
const doPath = path.join(__dirname, '../src/lib/durable-objects/BanditAgentDO.ts')

if (!fs.existsSync(workerPath)) {
  console.error('❌ Worker file not found at:', workerPath)
  process.exit(1)
}

if (!fs.existsSync(doPath)) {
  console.error('❌ Durable Object file not found at:', doPath)
  process.exit(1)
}

// Read worker file
let workerContent = fs.readFileSync(workerPath, 'utf-8')

// Check if already patched
if (workerContent.includes('export { BanditAgentDO }')) {
  console.log('✅ Worker already patched, skipping')
  process.exit(0)
}

// Read the DO source (not used, but keep for reference)
const doSource = fs.readFileSync(doPath, 'utf-8')

// Create the DO class inline (minimal working version)
const doCode = `
// ===== Durable Object: BanditAgentDO =====

export class BanditAgentDO {
  constructor(ctx, env) {
    this.ctx = ctx;
    this.env = env;
    this.state = null;
    this.webSockets = new Set();
    this.isRunning = false;
  }

  async fetch(request) {
    try {
      const url = new URL(request.url);
      const pathname = url.pathname;

      // Handle WebSocket upgrade
      if (request.headers.get("Upgrade") === "websocket") {
        const pair = new WebSocketPair();
        const [client, server] = Object.values(pair);
        server.accept();
        this.webSockets.add(server);
        
        server.addEventListener("close", () => {
          this.webSockets.delete(server);
        });
        
        server.addEventListener("message", async (event) => {
          try {
            const data = JSON.parse(event.data);
            if (data.type === 'ping') {
              server.send(JSON.stringify({ type: 'pong', timestamp: new Date().toISOString() }));
            }
          } catch (error) {
            console.error('WebSocket message error:', error);
          }
        });
        
        return new Response(null, { status: 101, webSocket: client });
      }

      // Handle HTTP requests
      if (pathname.endsWith('/start')) {
        const body = await request.json();
        
        // Initialize state
        this.state = {
          runId: body.runId,
          modelName: body.modelName,
          status: 'running',
          currentLevel: body.startLevel || 0,
          targetLevel: body.endLevel || 33
        };
        
        // Save to storage
        await this.ctx.storage.put('state', this.state);
        
        // Broadcast to WebSocket clients
        this.broadcast({
          type: 'agent_message',
          data: {
            content: \`Run started: \${body.modelName} - Levels \${body.startLevel}-\${body.endLevel}\`,
          },
          timestamp: new Date().toISOString()
        });
        
        // Start agent execution in background
        this.runAgent().catch(err => console.error('Agent error:', err));
        
        return new Response(JSON.stringify({
          success: true,
          runId: body.runId,
          state: this.state
        }), {
          headers: { 'Content-Type': 'application/json' }
        });
      }
      
      if (pathname.endsWith('/pause')) {
        if (this.state) {
          this.state.status = 'paused';
          this.isRunning = false;
          await this.ctx.storage.put('state', this.state);
        }
        return new Response(JSON.stringify({ success: true, state: this.state }), {
          headers: { 'Content-Type': 'application/json' }
        });
      }
      
      if (pathname.endsWith('/resume')) {
        if (this.state) {
          this.state.status = 'running';
          this.isRunning = true;
          await this.ctx.storage.put('state', this.state);
          this.runAgent().catch(err => console.error('Agent error:', err));
        }
        return new Response(JSON.stringify({ success: true, state: this.state }), {
          headers: { 'Content-Type': 'application/json' }
        });
      }
      
      if (pathname.endsWith('/status')) {
        return new Response(JSON.stringify({
          state: this.state,
          isRunning: this.isRunning,
          connectedClients: this.webSockets.size
        }), {
          headers: { 'Content-Type': 'application/json' }
        });
      }

      return new Response('Not found', { status: 404 });
    } catch (error) {
      console.error('DO fetch error:', error);
      return new Response(JSON.stringify({ error: error.message }), { 
        status: 500,
        headers: { 'Content-Type': 'application/json' }
      });
    }
  }

  async runAgent() {
    if (!this.state) return;
    this.isRunning = true;

    try {
      // Call SSH proxy agent endpoint
      const response = await fetch(\`\${this.env.SSH_PROXY_URL}/agent/run\`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          runId: this.state.runId,
          modelName: this.state.modelName,
          startLevel: this.state.currentLevel,
          endLevel: this.state.targetLevel,
          apiKey: this.env.OPENROUTER_API_KEY
        })
      });

      // Stream agent events
      const reader = response.body.getReader();
      const decoder = new TextDecoder();

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        
        const chunk = decoder.decode(value);
        const lines = chunk.split('\\n').filter(l => l.trim());
        
        for (const line of lines) {
          try {
            const event = JSON.parse(line);
            this.broadcast(event);
            
            // Update state based on events
            if (event.type === 'level_complete') {
              this.state.currentLevel = event.data.level + 1;
            }
            if (event.type === 'run_complete') {
              this.state.status = 'complete';
              this.isRunning = false;
            }
            if (event.type === 'error') {
              this.state.status = 'failed';
              this.state.error = event.data.content;
              this.isRunning = false;
            }
          } catch (e) {
            // Ignore parse errors
          }
        }
      }
    } catch (error) {
      this.state.status = 'failed';
      this.state.error = error.message;
      this.isRunning = false;
      this.broadcast({
        type: 'error',
        data: { content: error.message },
        timestamp: new Date().toISOString()
      });
    }
  }

  broadcast(event) {
    const message = JSON.stringify(event);
    for (const socket of this.webSockets) {
      try {
        socket.send(message);
      } catch (error) {
        this.webSockets.delete(socket);
      }
    }
  }

  async alarm() {
    // Cleanup after 2 hours
    if (!this.isRunning && this.state) {
      const startedAt = new Date(this.state.startedAt || 0).getTime();
      if (Date.now() - startedAt > 2 * 60 * 60 * 1000) {
        await this.ctx.storage.deleteAll();
        this.state = null;
      }
    }
    await this.ctx.storage.setAlarm(Date.now() + 60 * 60 * 1000);
  }
}
// ===== End Durable Object =====
`

// Insert DO code right after the other DO exports
// Find the line with "export { BucketCachePurge }"
const bucketCacheLine = 'export { BucketCachePurge } from "./.build/durable-objects/bucket-cache-purge.js";'
const insertIndex = workerContent.indexOf(bucketCacheLine)

if (insertIndex === -1) {
  console.error('❌ Could not find insertion point in worker.js')
  process.exit(1)
}

// Insert right after that line
const insertPosition = insertIndex + bucketCacheLine.length
const patchedContent = 
  workerContent.slice(0, insertPosition) + 
  '\n' + doCode + '\n' +
  workerContent.slice(insertPosition)

// Write back
fs.writeFileSync(workerPath, patchedContent, 'utf-8')

console.log('✅ Worker patched successfully - BanditAgentDO exported')
console.log('📝 Note: Using stub DO implementation. Full LangGraph integration via SSH proxy.')