@shogo-ai/worker 2.3.12 → 2.4.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@shogo-ai/worker",
-  "version": "2.3.12",
+  "version": "2.4.0",
   "description": "Shogo Cloud Agent Worker — run Shogo agents on your own machine (laptop, devbox, CI).",
   "license": "MIT",
   "author": "Shogo Technologies, Inc.",
@@ -18,6 +18,7 @@
   "exports": {
     ".": "./src/cli.ts",
     "./tunnel": "./src/lib/tunnel.ts",
+    "./tunnel-host": "./src/lib/tunnel-host.ts",
     "./runtime-manager": "./src/lib/runtime-manager.ts",
     "./runtime-activity": "./src/lib/runtime-activity.ts",
     "./resource-limits": "./src/lib/resource-limits.ts",

package/src/lib/__tests__/helpers/fake-cloud.ts

@@ -0,0 +1,139 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026 Shogo Technologies, Inc.
+/**
+ * A fake Shogo cloud, in its own process, for the tunnel-block test.
+ *
+ * It has to be a separate process and not an in-test server: the test blocks the
+ * main thread on purpose, and an observer on that thread cannot see the pings it
+ * is supposed to be policing. Here the pings keep flowing while the API's thread
+ * is busy, which is exactly how the real cloud behaves.
+ *
+ * Speaks the real protocol: `POST /api/instances/heartbeat` (asking for a socket
+ * and advertising its URL), `GET /api/instances/ws` upgrade, `{type:'ping'}`
+ * every 40ms answered by `{type:'pong'}`, and a close (1006) for a socket that
+ * misses the deadline — the same idle check the cloud runs.
+ *
+ * Reports over stdout, one JSON object per line: `ready`, `pong`, `kill`,
+ * `request`, `close`.
+ */
+
+const PING_INTERVAL_MS = Number(process.env.FAKE_CLOUD_PING_MS ?? 40)
+const PONG_DEADLINE_MS = Number(process.env.FAKE_CLOUD_DEADLINE_MS ?? 120)
+
+const say = (line: Record<string, unknown>) => {
+  process.stdout.write(`${JSON.stringify(line)}\n`)
+}
+
+let socket: any = null
+let awaitingPongSince: number | null = null
+let kills = 0
+
+const server = Bun.serve({
+  port: 0,
+  hostname: '127.0.0.1',
+  fetch(req, srv) {
+    const url = new URL(req.url)
+    if (url.pathname === '/api/instances/ws') {
+      if (srv.upgrade(req)) return undefined as any
+      return new Response('upgrade failed', { status: 400 })
+    }
+    if (url.pathname === '/api/instances/heartbeat') {
+      return Response.json({
+        instanceId: 'inst-worker-1',
+        wsRequested: true,
+        nextPollIn: 30,
+        wsUrl: `ws://127.0.0.1:${srv.port}`,
+      })
+    }
+    return new Response('{}', { status: 404 })
+  },
+  websocket: {
+    open(ws) {
+      socket = ws
+      awaitingPongSince = null
+    },
+    message(_ws, raw) {
+      let message: any
+      try {
+        message = JSON.parse(String(raw))
+      } catch {
+        return
+      }
+      if (message.type === 'pong') {
+        awaitingPongSince = null
+        say({ event: 'pong' })
+        return
+      }
+      if (message.type === 'response') {
+        say({ event: 'request', status: message.status, body: String(message.body ?? '').slice(0, 200) })
+        return
+      }
+      if (message.type === 'stream-chunk' || message.type === 'stream-end') {
+        say({ event: 'request', status: 200, body: 'stream' })
+      }
+    },
+    close(_ws, code) {
+      socket = null
+      say({ event: 'close', code })
+    },
+  },
+})
+
+setInterval(() => {
+  if (!socket) return
+  if (awaitingPongSince === null) {
+    awaitingPongSince = Date.now()
+  } else if (Date.now() - awaitingPongSince > PONG_DEADLINE_MS) {
+    kills += 1
+    say({ event: 'kill', kills })
+    try {
+      socket.close(1006, 'no pong')
+    } catch {
+      // already gone
+    }
+    socket = null
+    awaitingPongSince = null
+    return
+  }
+  try {
+    socket.send(JSON.stringify({ type: 'ping' }))
+  } catch {
+    // socket went away
+  }
+}, PING_INTERVAL_MS)
+
+// The test drives requests over stdin (one JSON object per line), because a
+// spawned child has no IPC channel to listen on.
+void (async () => {
+  let buffer = ''
+  const decoder = new TextDecoder()
+  for await (const chunk of Bun.stdin.stream()) {
+    buffer += decoder.decode(chunk as Uint8Array, { stream: true })
+    let newline = buffer.indexOf('\n')
+    while (newline !== -1) {
+      const line = buffer.slice(0, newline).trim()
+      buffer = buffer.slice(newline + 1)
+      newline = buffer.indexOf('\n')
+      if (!line) continue
+      let message: any
+      try {
+        message = JSON.parse(line)
+      } catch {
+        continue
+      }
+      if (message?.type === 'request' && socket) {
+        socket.send(
+          JSON.stringify({
+            type: 'request',
+            requestId: message.requestId,
+            method: 'GET',
+            path: message.path,
+            headers: {},
+          }),
+        )
+      }
+    }
+  }
+})()
+
+say({ event: 'ready', port: server.port })

package/src/lib/__tests__/resource-limits-exec-budget.test.ts

@@ -0,0 +1,116 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026 Shogo Technologies, Inc.
+/**
+ * SHO-203: the agent-command budget is separate from the runtime's own leak
+ * ceiling, and the kernel cap has to sit above it so the guard (which writes a
+ * note) always trips before the kernel (which writes nothing).
+ */
+import { afterEach, describe, expect, it } from 'bun:test';
+import { mkdtempSync, rmSync, writeFileSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import {
+  applyLiveMemoryCap,
+  resolveExecMemoryMB,
+  resolveResourceLimits,
+  type RuntimeResourceLimits,
+} from '../resource-limits';
+
+const dirs: string[] = [];
+function limitsFile(contents: unknown): string {
+  const dir = mkdtempSync(join(tmpdir(), 'host-limits-'));
+  dirs.push(dir);
+  const file = join(dir, 'host-limits.json');
+  writeFileSync(file, typeof contents === 'string' ? contents : JSON.stringify(contents));
+  return file;
+}
+
+afterEach(() => {
+  while (dirs.length) rmSync(dirs.pop()!, { recursive: true, force: true });
+});
+
+const baseEnv = { RUNTIME_MEMORY_MB: '8192' };
+
+describe('resolveResourceLimits', () => {
+  it('keeps the runtime ceiling unchanged when no command budget is known', () => {
+    const limits = resolveResourceLimits({ ...baseEnv });
+    expect(limits?.memoryMB).toBe(8192);
+    expect(limits?.osMemoryMB).toBe(8192);
+  });
+
+  it('adds an explicit SHOGO_EXEC_MEMORY_MB to the OS cap only', () => {
+    const limits = resolveResourceLimits({ ...baseEnv, SHOGO_EXEC_MEMORY_MB: '131072' });
+    expect(limits?.memoryMB).toBe(8192);
+    expect(limits?.osMemoryMB).toBe(8192 + 131072);
+  });
+
+  it('reads the command budget from the watched host-limits file', () => {
+    const file = limitsFile({ execMemoryMB: 122880, totalMemMB: 131072, source: 'auto' });
+    const limits = resolveResourceLimits({ ...baseEnv, SHOGO_HOST_LIMITS_FILE: file });
+    expect(limits?.osMemoryMB).toBe(8192 + 122880);
+  });
+
+  it('prefers an explicit env budget over the file', () => {
+    const file = limitsFile({ execMemoryMB: 4096 });
+    const limits = resolveResourceLimits({
+      ...baseEnv,
+      SHOGO_HOST_LIMITS_FILE: file,
+      SHOGO_EXEC_MEMORY_MB: '2048',
+    });
+    expect(limits?.osMemoryMB).toBe(8192 + 2048);
+  });
+
+  it('ignores a missing or malformed limits file', () => {
+    expect(resolveExecMemoryMB({ SHOGO_HOST_LIMITS_FILE: join(tmpdir(), 'nope.json') })).toBe(0);
+    expect(resolveExecMemoryMB({ SHOGO_HOST_LIMITS_FILE: limitsFile('{oops') })).toBe(0);
+    expect(resolveExecMemoryMB({ SHOGO_HOST_LIMITS_FILE: limitsFile({ execMemoryMB: -5 }) })).toBe(0);
+  });
+});
+
+describe('applyLiveMemoryCap', () => {
+  const limits: RuntimeResourceLimits = { memoryMB: 8192, osMemoryMB: 8192 + 122880, cpuPercent: 0, tasksMax: 4096 };
+
+  it('raises a live systemd scope with set-property, using the combined cap', () => {
+    const calls: Array<{ command: string; args: string[] }> = [];
+    const ok = applyLiveMemoryCap(
+      { scopeName: 'shogo-rt-abc12345-m1', limits },
+      { platform: 'linux', run: (command, args) => calls.push({ command, args }) },
+    );
+    expect(ok).toBe(true);
+    expect(calls).toHaveLength(1);
+    expect(calls[0]!.command).toBe('systemctl');
+    expect(calls[0]!.args).toEqual([
+      '--user',
+      'set-property',
+      '--runtime',
+      'shogo-rt-abc12345-m1',
+      'MemoryMax=131072M',
+      'MemoryHigh=117964M',
+    ]);
+  });
+
+  it('rewrites a direct cgroup memory.max, in bytes', () => {
+    const writes: Array<{ path: string; data: string }> = [];
+    const ok = applyLiveMemoryCap(
+      { pid: 4242, limits },
+      {
+        platform: 'linux',
+        exists: (path) => path.endsWith('memory.max') || path.endsWith('memory.high'),
+        writeFile: (path, data) => writes.push({ path, data }),
+      },
+    );
+    expect(ok).toBe(true);
+    expect(writes).toEqual([
+      { path: '/sys/fs/cgroup/shogo/rt-4242/memory.max', data: String(131072 * 1024 * 1024) },
+      { path: '/sys/fs/cgroup/shogo/rt-4242/memory.high', data: String(117964 * 1024 * 1024) },
+    ]);
+  });
+
+  it('does nothing where the limit cannot be changed in place', () => {
+    expect(applyLiveMemoryCap({ scopeName: 'x', limits }, { platform: 'win32' })).toBe(false);
+    expect(applyLiveMemoryCap({ scopeName: 'x', limits }, { platform: 'darwin' })).toBe(false);
+    expect(
+      applyLiveMemoryCap({ pid: 7, limits }, { platform: 'linux', exists: () => false }),
+    ).toBe(false);
+  });
+});

package/src/lib/__tests__/runtime-manager-exec-budget-env.test.ts

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026 Shogo Technologies, Inc.
+/**
+ * SHO-203: the agent-command budget must not be borrowed from the runtime's own
+ * leak ceiling. When the desktop publishes `host-limits.json`, the runtime gets
+ * the file's path and *no* `SHOGO_EXEC_MEMORY_MB` — otherwise the env value,
+ * baked in at spawn, would win and a settings change would need a restart.
+ *
+ * `buildEnv` is private; it is driven through a typed cast, the same way
+ * `runtime-manager-tree-sitter-env.test.ts` does.
+ */
+import { afterAll, describe, expect, it } from 'bun:test';
+import { mkdtempSync, rmSync, writeFileSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { WorkerRuntimeManager } from '../runtime-manager';
+
+const dir = mkdtempSync(join(tmpdir(), 'rm-exec-budget-'));
+afterAll(() => rmSync(dir, { recursive: true, force: true }));
+
+const BIN = '/opt/shogo/runtime/shogo-agent-runtime-linux-x64';
+
+type EnvBuilder = { buildEnv(slot: unknown, runtimeBinPath: string): NodeJS.ProcessEnv };
+
+/** The fields `buildEnv` reads, plus an idempotent port for the PID file path. */
+function fakeSlot(): unknown {
+  return {
+    projectId: 'proj-exec-budget-test',
+    agentPort: 41000,
+    apiServerPort: 41100,
+    spawnConfig: { cloudUrl: 'https://cloud.test', apiKey: 'api-key' },
+  };
+}
+
+function limitsFile(): string {
+  const file = join(dir, 'host-limits.json');
+  writeFileSync(file, JSON.stringify({ execMemoryMB: 122880, totalMemMB: 131072, source: 'auto' }));
+  return file;
+}
+
+describe('WorkerRuntimeManager.buildEnv — agent-command budget', () => {
+  it('hands the desktop limits file through and leaves the env budget unset', async () => {
+    const file = limitsFile();
+    const mgr = new WorkerRuntimeManager({
+      env: { RUNTIME_MEMORY_MB: '8192', SHOGO_HOST_LIMITS_FILE: file } as NodeJS.ProcessEnv,
+    });
+    const env = (mgr as unknown as EnvBuilder).buildEnv(fakeSlot(), BIN);
+
+    expect(env.SHOGO_HOST_LIMITS_FILE).toBe(file);
+    // The bug: this used to be `String(limits.memoryMB)`.
+    expect(env.SHOGO_EXEC_MEMORY_MB).toBeUndefined();
+    expect(env.SHOGO_EXEC_PIDS_FILE).toBeTruthy();
+    await mgr.stopAll();
+  });
+
+  it('falls back to the runtime ceiling for a headless worker with no limits file', () => {
+    const mgr = new WorkerRuntimeManager({
+      env: { RUNTIME_MEMORY_MB: '8192' } as NodeJS.ProcessEnv,
+    });
+    const env = (mgr as unknown as EnvBuilder).buildEnv(fakeSlot(), BIN);
+    expect(env.SHOGO_EXEC_MEMORY_MB).toBe('8192');
+  });
+
+  it('keeps an explicit SHOGO_EXEC_MEMORY_MB in front of the file', async () => {
+    const file = limitsFile();
+    const mgr = new WorkerRuntimeManager({
+      env: {
+        RUNTIME_MEMORY_MB: '8192',
+        SHOGO_EXEC_MEMORY_MB: '4096',
+        SHOGO_HOST_LIMITS_FILE: file,
+      } as NodeJS.ProcessEnv,
+    });
+    const env = (mgr as unknown as EnvBuilder).buildEnv(fakeSlot(), BIN);
+    expect(env.SHOGO_EXEC_MEMORY_MB).toBe('4096');
+    await mgr.stopAll();
+  });
+
+  it('leaves the runtime ceiling itself alone', () => {
+    const mgr = new WorkerRuntimeManager({
+      env: { RUNTIME_MEMORY_MB: '8192' } as NodeJS.ProcessEnv,
+    });
+    const env = (mgr as unknown as EnvBuilder).buildEnv(fakeSlot(), BIN);
+    // The heap hint still comes from the runtime ceiling, not the command budget.
+    expect(env.NODE_OPTIONS).toContain('--max-old-space-size=6144');
+    expect(env.RUNTIME_MEMORY_MB).toBe('8192');
+  });
+});

package/src/lib/__tests__/tunnel-worker-block.test.ts

@@ -0,0 +1,243 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2026 Shogo Technologies, Inc.
+/**
+ * The tunnel's socket survives a busy main thread (SHO-185 §5).
+ *
+ * This is the test the whole Worker move exists for. The fake cloud lives in its
+ * own process (`helpers/fake-cloud.ts`) because the point of the test is to block
+ * *this* thread: an observer here would be blocked too and could not see the
+ * pings it is supposed to be policing.
+ *
+ * Two cases, same cloud, same block:
+ *   • on its own thread — pongs keep flowing, the socket stays open, and a
+ *     tunneled request still round-trips afterwards;
+ *   • on the API's thread — the deadline passes unanswered and the cloud closes
+ *     the socket (1006). That is the bug, and without it the first case would
+ *     only prove that a socket somewhere stayed open.
+ */
+
+import { afterAll, beforeAll, describe, expect, test } from 'bun:test'
+import { createServer, type Server } from 'node:http'
+import type { AddressInfo } from 'node:net'
+import { fileURLToPath } from 'node:url'
+import { startTunnelInWorker } from '../tunnel-host'
+import { WorkerTunnel, type RuntimeResolver } from '../tunnel'
+
+const BLOCK_MS = 400
+
+interface FakeCloud {
+  cloudUrl: string
+  pongs: number
+  kills: number
+  closeCodes: number[]
+  requests: { status: number; body: string }[]
+  sendRequest: (requestId: string, path: string) => void
+  stop: () => void
+}
+
+/**
+ * Spawn the fake cloud and read its stdout, one JSON event per line.
+ *
+ * The child is `bun --no-env-file` so the workspace's `.env` cannot point it at a
+ * real cloud, and it prints `ready` with the port it bound before anything else.
+ */
+async function startFakeCloud(): Promise<FakeCloud> {
+  // `pathname` stays percent-encoded, and this checkout's path has a space in it.
+  const entry = fileURLToPath(new URL('./helpers/fake-cloud.ts', import.meta.url))
+  const child = Bun.spawn(['bun', '--no-env-file', entry], {
+    stdout: 'pipe',
+    stdin: 'pipe',
+    stderr: 'inherit',
+    env: { ...process.env, FAKE_CLOUD_PING_MS: '40', FAKE_CLOUD_DEADLINE_MS: '120' },
+  })
+
+  const state = {
+    pongs: 0,
+    kills: 0,
+    closeCodes: [] as number[],
+    requests: [] as { status: number; body: string }[],
+  }
+
+  let port = 0
+  let buffer = ''
+  let settleReady: (port: number) => void = () => {}
+  const ready = new Promise<number>((resolve) => {
+    settleReady = resolve
+  })
+
+  const drain = (async () => {
+    const stream = child.stdout as ReadableStream<Uint8Array>
+    const reader = stream.getReader()
+    const decoder = new TextDecoder()
+    while (true) {
+      const { done, value } = await reader.read()
+      if (done) break
+      buffer += decoder.decode(value, { stream: true })
+      let newline = buffer.indexOf('\n')
+      while (newline !== -1) {
+        const line = buffer.slice(0, newline).trim()
+        buffer = buffer.slice(newline + 1)
+        newline = buffer.indexOf('\n')
+        if (!line) continue
+        let event: any
+        try {
+          event = JSON.parse(line)
+        } catch {
+          continue
+        }
+        if (event.event === 'ready') {
+          port = Number(event.port)
+          settleReady(port)
+        } else if (event.event === 'pong') state.pongs += 1
+        else if (event.event === 'kill') state.kills += 1
+        else if (event.event === 'close') state.closeCodes.push(Number(event.code))
+        else if (event.event === 'request') {
+          state.requests.push({ status: Number(event.status), body: String(event.body ?? '') })
+        }
+      }
+    }
+  })().catch(() => {})
+
+  await ready
+
+  return {
+    cloudUrl: `http://127.0.0.1:${port}`,
+    get pongs() {
+      return state.pongs
+    },
+    get kills() {
+      return state.kills
+    },
+    get closeCodes() {
+      return state.closeCodes
+    },
+    get requests() {
+      return state.requests
+    },
+    sendRequest(requestId: string, path: string) {
+      child.stdin?.write(`${JSON.stringify({ type: 'request', requestId, path })}\n`)
+    },
+    stop() {
+      try {
+        child.kill()
+      } catch {
+        // already gone
+      }
+      void drain
+    },
+  }
+}
+
+let localApi: Server
+let localPort = 0
+
+const resolver: RuntimeResolver = {
+  async resolveLocalUrl(pathWithQuery: string) {
+    return `http://127.0.0.1:${localPort}${pathWithQuery}`
+  },
+  async deriveRuntimeToken() {
+    return 'runtime-token'
+  },
+  getActiveProjects() {
+    return ['proj-1']
+  },
+  status() {
+    return { status: 'running', agentPort: 1234 }
+  },
+}
+
+beforeAll(async () => {
+  localApi = createServer((req, res) => {
+    res.setHeader('content-type', 'application/json')
+    res.end(JSON.stringify({ answeredBy: 'local-api', path: req.url }))
+  })
+  await new Promise<void>((resolve) => localApi.listen(0, '127.0.0.1', () => resolve()))
+  localPort = (localApi.address() as AddressInfo).port
+})
+
+afterAll(async () => {
+  await new Promise<void>((resolve) => localApi.close(() => resolve()))
+})
+
+/** Burn this thread — the local API's thread — past the cloud's deadline. */
+function blockMainThread(ms: number): void {
+  const until = Date.now() + ms
+  while (Date.now() < until) {
+    // A busy loop is the honest stand-in for a `lsof` or a big write.
+  }
+}
+
+const wait = (ms: number) => new Promise<void>((resolve) => setTimeout(resolve, ms))
+
+async function waitFor(predicate: () => boolean, timeoutMs = 8_000): Promise<void> {
+  const started = Date.now()
+  while (Date.now() < started + timeoutMs) {
+    if (predicate()) return
+    await wait(20)
+  }
+  throw new Error(`timed out after ${timeoutMs}ms waiting for a condition that never held`)
+}
+
+describe('a tunnel on its own thread', () => {
+  test('keeps its socket through a main thread blocked for 4× the ping interval', async () => {
+    const cloud = await startFakeCloud()
+    const host = startTunnelInWorker({
+      apiKey: 'test-key',
+      cloudUrl: cloud.cloudUrl,
+      name: 'worker-test',
+      kind: 'desktop',
+      resolver,
+      logger: { log: () => {}, warn: () => {}, error: () => {} },
+    })
+
+    try {
+      await host.ready()
+      // The socket only opens once the cloud asks for it on a heartbeat.
+      await waitFor(() => cloud.pongs > 0)
+      const before = cloud.pongs
+
+      blockMainThread(BLOCK_MS)
+
+      // The cloud never reached for the close code, and the pongs never stopped.
+      expect(cloud.kills).toBe(0)
+      expect(cloud.closeCodes).toEqual([])
+      await waitFor(() => cloud.pongs > before, 3_000)
+
+      // And the tunnel still forwards afterwards: a request pushed down the
+      // socket comes back from the local API.
+      cloud.sendRequest('req-after-block', '/api/projects')
+      await waitFor(() => cloud.requests.some((entry) => entry.status === 200), 8_000)
+      expect(cloud.requests.find((entry) => entry.status === 200)?.body).toContain('local-api')
+    } finally {
+      host.stop()
+      cloud.stop()
+    }
+  }, 40_000)
+
+  test('the same block kills an in-process tunnel — the bug §5 fixes', async () => {
+    const cloud = await startFakeCloud()
+    const tunnel = new WorkerTunnel({
+      apiKey: 'test-key',
+      cloudUrl: cloud.cloudUrl,
+      name: 'in-process',
+      kind: 'desktop',
+      resolver,
+      logger: { log: () => {}, warn: () => {}, error: () => {} },
+    })
+
+    try {
+      tunnel.start()
+      await waitFor(() => cloud.pongs > 0)
+
+      blockMainThread(BLOCK_MS)
+
+      // The cloud's deadline passed with the pong stuck behind the block, so it
+      // closed the socket: this is the 1006 the person sees as "machine offline".
+      await waitFor(() => cloud.kills > 0)
+      expect(cloud.kills).toBeGreaterThan(0)
+    } finally {
+      tunnel.stop()
+      cloud.stop()
+    }
+  }, 40_000)
+})