mini-world / server.ts
victor's picture
victor HF Staff
chore: switch ENDPOINT to dedicated pod (zopt7f90d1ydgy)
a0a680b
Raw History Blame Contribute Delete
21.4 kB
/**
* server.ts β€” Bun HTTP server: owns the agent loop, broadcasts state via SSE.
*
* One AI agent runs server-side against the runpod llama.cpp endpoint with
* structured tool calling, KV-cache reuse, and LLM-summarization compaction.
* Connected clients are pure viewers β€” they GET / for the viewer HTML and
* subscribe to /stream for state updates. No client-side inference, no auth.
*/
import { serve } from 'bun';
import { readFile } from 'fs/promises';
import { SinglePlayerGame, DEFAULTS } from './engine.js';
import {
PICK_FIRST_TOOLS,
PICK_FIRST_SYSTEM_PROMPT,
buildPickFirstPrompt,
AgentMemory,
} from './agent-brain.js';
// ─── Config ──────────────────────────────────────────────────────────────────
const ENDPOINT = (process.env.ENDPOINT ?? 'https://zopt7f90d1ydgy-8080.proxy.runpod.net').replace(/\/$/, '');
const REMOTE_MODEL = process.env.MODEL ?? 'default';
const PORT = parseInt(process.env.PORT ?? '7860', 10);
const MAX_TOKENS = 100;
const TEMPERATURE = 0.6;
// ─── Singleton state ─────────────────────────────────────────────────────────
const game = new SinglePlayerGame();
const memory = new AgentMemory(PICK_FIRST_SYSTEM_PROMPT, {
tokenBudget: 4000,
triggerCompactTokens: 2000,
keepRecentTokens: 700,
});
const actionLog: Array<{ action: string; detail: string; success: boolean }> = [];
let fallbackCount = 0;
let lastUsage: any = null;
let lastTotalMs = 0;
let serverCtx = 8192;
let remoteModelId: string | null = null;
let status = 'connecting';
let paused = process.env.START_PAUSED === '1';
let lastThinking: string | null = null;
const subscribers = new Set<ReadableStreamDefaultController<Uint8Array>>();
// ─── Trace events (pi-mono `--mode json` shape) ──────────────────────────────
// One JSON object per line on stdout. Event types mirror pi-mono's
// AgentSessionEvent vocabulary; extra app-specific events (fallback_action,
// viewer_connect, world_reset) follow the same discriminated-union pattern.
type TraceEvent = Record<string, unknown>;
function trace(event: TraceEvent) {
process.stdout.write(JSON.stringify({ ts: new Date().toISOString(), ...event }) + '\n');
}
// ─── State serialization + broadcast ─────────────────────────────────────────
function buildState() {
return {
viewport: game.getViewport(),
inventory: game.inventory,
stats: game.stats,
actions: actionLog.slice(-30),
status,
thinking: lastThinking,
summary: memory.summary,
debug: {
backend: 'Remote (llama.cpp)',
endpoint: ENDPOINT.replace(/^https?:\/\//, ''),
model: remoteModelId,
serverCtx,
promptTokens: lastUsage?.prompt_tokens ?? 0,
cachedTokens: lastUsage?.prompt_tokens_details?.cached_tokens ?? 0,
completionTokens: lastUsage?.completion_tokens ?? 0,
totalMs: lastTotalMs,
memTokens: memory.estimateTokens(),
memMessages: memory.messages.length,
fallbacks: fallbackCount,
compactions: memory.compactions,
summaryLen: memory.summary?.length ?? 0,
turn: actionLog.length,
viewers: subscribers.size,
cacheRetries,
cacheCooldown,
paused,
},
};
}
const enc = new TextEncoder();
function broadcast() {
const payload = enc.encode(`data: ${JSON.stringify(buildState())}\n\n`);
for (const ctrl of subscribers) {
try {
ctrl.enqueue(payload);
} catch {
subscribers.delete(ctrl);
}
}
}
// ─── Connect: probe endpoint, derive memory budgets from n_ctx ───────────────
async function connect() {
status = 'connecting';
// Retry the health check forever β€” if the upstream pod is restarting, we
// want to wait it out, not crash the Space and fall into RUNTIME_ERROR.
while (true) {
const health = await fetch(`${ENDPOINT}/health`).then(r => r.json()).catch(() => null);
if (health && health.status === 'ok') break;
trace({ type: 'connect_waiting', endpoint: ENDPOINT });
await new Promise(r => setTimeout(r, 5000));
}
const models = await fetch(`${ENDPOINT}/v1/models`).then(r => r.json()).catch(() => null);
remoteModelId = models?.data?.[0]?.id || REMOTE_MODEL;
const props = await fetch(`${ENDPOINT}/props`).then(r => r.json()).catch(() => null);
serverCtx = props?.default_generation_settings?.n_ctx ?? 8192;
// Reserve room for: completion buffer (MAX_TOKENS), tool/template overhead (~400 tok),
// and a 1.7x margin since chars/4 undercounts real tokens (measured ratio with
// structured tool_calls + Qwen3.6 chat template was ~1.6).
const safe = Math.max(1500, Math.floor((serverCtx - MAX_TOKENS - 400) / 1.7));
memory.tokenBudget = safe;
memory.triggerCompactTokens = Math.floor(safe * 0.6);
memory.keepRecentTokens = Math.floor(safe * 0.25);
trace({
type: 'agent_start',
endpoint: ENDPOINT,
model: remoteModelId,
n_ctx: serverCtx,
memoryConfig: {
tokenBudget: memory.tokenBudget,
triggerCompactTokens: memory.triggerCompactTokens,
keepRecentTokens: memory.keepRecentTokens,
},
});
status = 'running';
}
// ─── Inference ───────────────────────────────────────────────────────────────
// llama.cpp build b9038-2ca1161bd has a regression where cache_prompt:true
// causes slots to accumulate state past n_ctx and then 500 every request to
// that slot permanently. The cooldown logic below is a vestige; in practice
// we leave caching disabled (USE_CACHE=false) and pay the re-tokenize cost
// every turn. Set ENV USE_CACHE=1 to opt back in if the build is fixed upstream.
const USE_CACHE = process.env.USE_CACHE === '1';
let cacheCooldown = USE_CACHE ? 0 : Infinity;
const COOLDOWN_TURNS = 30;
const RETRY_ATTEMPTS = 3;
const RETRY_DELAY_MS = 200;
let cacheRetries = 0;
async function callModel() {
const t0 = performance.now();
const useCache = cacheCooldown === 0;
const body: any = {
model: remoteModelId,
messages: memory.getMessages(),
tools: PICK_FIRST_TOOLS,
tool_choice: 'auto',
temperature: TEMPERATURE,
max_tokens: MAX_TOKENS,
chat_template_kwargs: { enable_thinking: false },
};
body.cache_prompt = useCache; // explicit true/false; omitting falls back to broken default
trace({ type: 'message_start' });
let res: Response;
let attempt = 0;
while (true) {
res = await fetch(`${ENDPOINT}/v1/chat/completions`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(body),
});
if (res.ok) break;
const text = await res.text();
const isContextErr = res.status === 500 && /Context size/i.test(text);
if (!isContextErr || attempt >= RETRY_ATTEMPTS) {
if (attempt > 0) {
trace({
type: 'auto_retry_end',
success: false,
attempt,
finalError: `HTTP ${res.status}: ${text}`,
});
}
throw new Error(`HTTP ${res.status}: ${text}`);
}
attempt++;
// On first context-size error, explicitly disable prompt caching for this
// and future turns until cooldown expires. Setting cache_prompt:false is
// not the same as omitting the key β€” the server default routes back to
// the poisoned cached path on this build (b9038-2ca1161bd).
if (body.cache_prompt !== false) {
if (body.cache_prompt === true) {
cacheRetries++;
cacheCooldown = COOLDOWN_TURNS;
}
body.cache_prompt = false;
}
trace({
type: 'auto_retry_start',
attempt,
maxAttempts: RETRY_ATTEMPTS,
delayMs: RETRY_DELAY_MS,
errorMessage: text,
reason: 'slot_saturation',
});
await new Promise((r) => setTimeout(r, RETRY_DELAY_MS));
}
if (attempt > 0) {
trace({ type: 'auto_retry_end', success: true, attempt });
}
if (cacheCooldown > 0) cacheCooldown--;
const data = await res.json();
const totalMs = Math.round(performance.now() - t0);
const message = data.choices?.[0]?.message ?? {};
const u = data.usage ?? {};
// Embed usage on the message in pi-mono shape (input/output/cacheRead/cacheWrite).
const usageOnMessage = {
input: u.prompt_tokens ?? 0,
output: u.completion_tokens ?? 0,
cacheRead: u.prompt_tokens_details?.cached_tokens ?? 0,
cacheWrite: 0,
};
trace({
type: 'message_end',
message: { ...message, usage: usageOnMessage },
durationMs: totalMs,
});
return { message, usage: data.usage, totalMs };
}
// ─── Summarizer ──────────────────────────────────────────────────────────────
const SUMMARY_SYS = `You compress 2D-grid-world game histories into terse, factual notes for the agent's future planning. Output 5-10 bullet points, under 200 words. Capture: directions tried and outcomes, diamonds collected (count + approximate area), obstacles encountered, exploration trajectory, and any failed actions worth remembering β€” especially repeated failures (e.g. "moved up 5 times, every attempt failed against an obstacle β€” direction up is blocked, try another"). No prose, no preamble.`;
async function callSummarizer(turnsLog: string, previousSummary: string | null) {
const userPrompt = previousSummary
? `Previous summary:\n${previousSummary}\n\nNew events to merge in:\n${turnsLog}\n\nProduce an updated unified summary.`
: `Game history so far:\n${turnsLog}\n\nProduce a summary.`;
const t0 = performance.now();
// Thinking is enabled for the summarizer (and only the summarizer) so the
// model can actually reason about what happened β€” failed-direction loops,
// unreachable diamonds, etc. β€” and produce a richer summary that helps the
// (thinking-off) turn loop break out of behavioural locks. max_tokens is
// large enough for the reasoning chain plus the bullet output.
const body: any = {
model: remoteModelId,
messages: [
{ role: 'system', content: SUMMARY_SYS },
{ role: 'user', content: userPrompt },
],
temperature: 0.2,
max_tokens: 2048,
chat_template_kwargs: { enable_thinking: true },
};
// Explicit cache_prompt; omitting it falls back to the broken default on this build.
body.cache_prompt = cacheCooldown === 0;
const res = await fetch(`${ENDPOINT}/v1/chat/completions`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(body),
});
if (!res.ok) throw new Error(`summarizer HTTP ${res.status}: ${await res.text()}`);
const data = await res.json();
const totalMs = Math.round(performance.now() - t0);
const out = (data.choices?.[0]?.message?.content || '').trim();
// Surface the summarizer call as a nested message_end so consumers see
// the same shape as a normal turn's inference.
const u = data.usage ?? {};
trace({
type: 'message_end',
source: 'summarizer',
message: {
role: 'assistant',
content: out,
usage: {
input: u.prompt_tokens ?? 0,
output: u.completion_tokens ?? 0,
cacheRead: u.prompt_tokens_details?.cached_tokens ?? 0,
cacheWrite: 0,
},
},
durationMs: totalMs,
});
return out;
}
// ─── Agent loop ──────────────────────────────────────────────────────────────
async function runAgentLoop() {
while (true) {
// Honor pause: idle in 500ms ticks, no model calls.
if (paused) {
if (status !== 'paused') {
status = 'paused';
trace({ type: 'agent_paused' });
broadcast();
}
await new Promise(r => setTimeout(r, 500));
continue;
} else if (status === 'paused') {
status = 'running';
trace({ type: 'agent_resumed' });
broadcast();
}
const turn = actionLog.length + 1;
trace({ type: 'turn_start', turn });
const viewport = game.getViewport();
// Compact before adding the next user turn so the model sees the summary
// as established prior context rather than mid-turn.
if (memory.shouldCompact()) {
status = 'compacting';
broadcast();
const tokensBefore = memory.estimateTokens();
trace({ type: 'compaction_start', reason: 'threshold', tokensBefore });
try {
const result = await memory.compactWith(callSummarizer);
if (result) {
trace({
type: 'compaction_end',
reason: 'threshold',
result: {
summary: memory.summary,
tokensBefore,
tokensAfter: memory.estimateTokens(),
droppedMessages: result.compacted,
},
});
} else {
trace({ type: 'compaction_end', reason: 'threshold', aborted: true });
}
} catch (err: any) {
trace({
type: 'compaction_end',
reason: 'threshold',
aborted: false,
willRetry: false,
errorMessage: String(err?.message ?? err),
});
}
}
// Snapshot memory length so we can roll back cleanly if this turn no-ops.
const memoryDepthBefore = memory.messages.length;
memory.addUserText(buildPickFirstPrompt(viewport, game.inventory));
let action: { action: string; detail: string; tool_call_id: string } | null = null;
status = 'thinking';
broadcast();
try {
const { message, usage, totalMs } = await callModel();
lastUsage = usage;
lastTotalMs = totalMs;
memory.addAssistant(message);
lastThinking = message.reasoning_content ?? null;
const tc = message.tool_calls?.[0];
if (tc) {
let args: any = {};
try { args = JSON.parse(tc.function.arguments); } catch {}
if (
(tc.function.name === 'move' || tc.function.name === 'pick') &&
['up', 'down', 'left', 'right'].includes(args.direction)
) {
action = { action: tc.function.name, detail: args.direction, tool_call_id: tc.id };
}
}
status = 'running';
} catch (err: any) {
status = 'fallback';
trace({
type: 'error',
source: 'inference',
message: String(err?.message ?? err),
});
}
if (!action) {
// No usable tool_call. Don't move the agent, don't pollute actionLog.
// Roll memory back to pre-turn state so the next attempt isn't biased
// by an empty assistant turn or an orphaned user message.
while (memory.messages.length > memoryDepthBefore) memory.popLast();
fallbackCount++;
trace({ type: 'fallback_action', turn, totalFallbacks: fallbackCount });
trace({ type: 'turn_end', turn });
broadcast();
await new Promise((r) => setTimeout(r, DEFAULTS.LOOP_DELAY_MS));
continue;
}
const argsObj = { direction: action.detail };
trace({
type: 'tool_execution_start',
toolCallId: action.tool_call_id,
toolName: action.action,
args: argsObj,
});
let toolResult: string;
let success = true;
try {
const result = game.executeAction(action) as { pickedUp?: string | null } | undefined;
const picked = result?.pickedUp ? ` β€” picked up ${result.pickedUp}` : '';
toolResult = `${action.action} ${action.detail} succeeded${picked}`;
} catch (err: any) {
toolResult = `${action.action} ${action.detail} failed: ${err.message}`;
success = false;
}
memory.addToolResult(action.tool_call_id, toolResult);
actionLog.push({ action: action.action, detail: action.detail, success });
trace({
type: 'tool_execution_end',
toolCallId: action.tool_call_id,
toolName: action.action,
result: toolResult,
isError: !success,
});
trace({ type: 'turn_end', turn });
broadcast();
await new Promise(r => setTimeout(r, DEFAULTS.LOOP_DELAY_MS));
}
}
// ─── HTTP server ─────────────────────────────────────────────────────────────
serve({
port: PORT,
idleTimeout: 240, // long-lived SSE connections
async fetch(req) {
const url = new URL(req.url);
if (url.pathname === '/health') {
return new Response('ok');
}
// ─── Programmatic test/debug API ────────────────────────────────────────
// GET /api/state β€” one-shot full state JSON (same shape as SSE events).
if (url.pathname === '/api/state') {
return new Response(JSON.stringify(buildState()), {
headers: { 'Content-Type': 'application/json' },
});
}
// GET /api/diag β€” endpoint health + slot status from the upstream llama.cpp.
if (url.pathname === '/api/diag') {
const props = await fetch(`${ENDPOINT}/props`).then(r => r.json()).catch(e => ({ error: String(e) }));
const slots = await fetch(`${ENDPOINT}/slots`).then(r => r.json()).catch(e => ({ error: String(e) }));
return new Response(JSON.stringify({
endpoint: ENDPOINT,
model: remoteModelId,
serverCtx,
memory: {
tokenBudget: memory.tokenBudget,
triggerCompactTokens: memory.triggerCompactTokens,
keepRecentTokens: memory.keepRecentTokens,
messages: memory.messages.length,
estimateTokens: memory.estimateTokens(),
compactions: memory.compactions,
summaryLen: memory.summary?.length ?? 0,
},
cacheRetries,
cacheCooldown,
upstream: { props, slots },
}, null, 2), { headers: { 'Content-Type': 'application/json' } });
}
// GET /api/memory β€” full conversation buffer for inspection.
if (url.pathname === '/api/memory') {
return new Response(JSON.stringify({
systemPrompt: memory.systemPrompt,
summary: memory.summary,
messages: memory.messages,
}, null, 2), { headers: { 'Content-Type': 'application/json' } });
}
// POST /api/pause β€” stop hitting the model. Loop idles, world stays intact.
if (url.pathname === '/api/pause' && req.method === 'POST') {
paused = true;
return new Response(JSON.stringify({ paused: true }), {
headers: { 'Content-Type': 'application/json' },
});
}
// POST /api/resume β€” resume the loop.
if (url.pathname === '/api/resume' && req.method === 'POST') {
paused = false;
return new Response(JSON.stringify({ paused: false }), {
headers: { 'Content-Type': 'application/json' },
});
}
// POST /api/reset β€” wipe world + memory + actionLog and start fresh.
if (url.pathname === '/api/reset' && req.method === 'POST') {
Object.assign(game, new SinglePlayerGame());
memory.messages = [];
memory.summary = null;
memory.compactions = 0;
actionLog.length = 0;
fallbackCount = 0;
lastUsage = null;
lastTotalMs = 0;
lastThinking = null;
cacheRetries = 0;
cacheCooldown = 0;
trace({ type: 'world_reset' });
broadcast();
return new Response(JSON.stringify({ ok: true }), {
headers: { 'Content-Type': 'application/json' },
});
}
if (url.pathname === '/stream') {
const stream = new ReadableStream<Uint8Array>({
start(controller) {
subscribers.add(controller);
trace({ type: 'viewer_connect', viewers: subscribers.size });
// Send immediate snapshot so clients don't wait a full turn.
controller.enqueue(enc.encode(`data: ${JSON.stringify(buildState())}\n\n`));
},
cancel(controller) {
subscribers.delete(controller);
trace({ type: 'viewer_disconnect', viewers: subscribers.size });
},
});
return new Response(stream as any, {
headers: {
'Content-Type': 'text/event-stream',
'Cache-Control': 'no-cache',
Connection: 'keep-alive',
},
});
}
if (url.pathname === '/' || url.pathname === '/index.html') {
const data = await readFile(import.meta.dir + '/index.html');
return new Response(data, { headers: { 'Content-Type': 'text/html' } });
}
return new Response('Not found', { status: 404 });
},
});
trace({ type: 'serve_listening', port: PORT, url: `http://localhost:${PORT}` });
// ─── Boot ────────────────────────────────────────────────────────────────────
await connect();
runAgentLoop().catch((err) => {
trace({
type: 'agent_end',
reason: 'crash',
errorMessage: String(err?.message ?? err),
});
status = 'crashed';
broadcast();
});