session-proxy.js

const http = require('http');
const { spawn } = require('child_process');
const path = require('path');

const OPCODE_HOST = process.env.OPCODE_HOST || '127.0.0.1';
const OPCODE_PORT = parseInt(process.env.OPCODE_PORT || '4096', 10);
const PROXY_PORT = parseInt(process.env.PROXY_PORT || '18081', 10);
const API_KEY = process.env.API_KEY || 'sk-gw-demo';

// Static model override via MODELS env. If unset, auto-fetch from `opencode models` at startup.
const MODELS_OVERRIDE = process.env.MODELS ? process.env.MODELS.split(',').map(s => s.trim()) : null;
let models = ['opencode/big-pickle']; // fallback until fetch completes

// Resolve opencode binary — check mise paths, then fall back to PATH
function findOpencode() {
  const home = require('os').homedir();
  const candidates = [
    path.join(home, '.local/share/mise/shims', 'opencode'),
    path.join(home, '.local/share/mise/installs/opencode/latest', 'opencode'),
    process.env.OPENCODE_PATH,
  ].filter(Boolean);
  for (const f of candidates) { try { if (require('fs').statSync(f).isFile()) return f; } catch {} }
  return 'opencode'; // last resort: rely on PATH
}
const OPENCODE_BIN = findOpencode();

function fetchModels() {
  return new Promise((resolve) => {
    const start = Date.now();
    const child = spawn(OPENCODE_BIN, ['models'], { stdio: ['ignore', 'pipe', 'ignore'], timeout: 15000 });
    let stdout = '';
    child.stdout.on('data', d => stdout += d.toString());
    child.on('close', (code) => {
      if (code === 0) {
        // Match lines like "provider/model-name" — reject junk like [pty.loader] ...
        const list = stdout.split('\n').map(l => l.trim()).filter(l => /^[a-zA-Z][\w.-]*\/[\w.-]+$/.test(l));
        if (list.length > 0) { models = list; console.error(`[models] fetched ${list.length} models in ${Date.now() - start}ms`); }
      } else {
        console.error(`[models] opencode exited with code ${code}, keeping fallback`);
      }
      resolve();
    });
    child.on('error', (err) => { console.error(`[models] spawn failed: ${err.message}`); resolve(); });
  });
}
if (!MODELS_OVERRIDE) fetchModels();

function json(d) { return JSON.stringify(d); }
function uid() { return 'chatcmpl-' + Date.now() + '-' + Math.random().toString(36).slice(2, 8); }

// Maps OpenAI tool_call_id -> tool name across request/response cycles
const toolCallIdToName = new Map();
// Regex for [TOOL_CALL: name] {json} [/TOOL_CALL] (strict) or without closing tag (lenient)
const TOOL_CALL_STRICT = /\[TOOL_CALL:\s*([\w-]+)\]\s*(\{[\s\S]*?\})\s*\[\/TOOL_CALL\]/g;
const TOOL_CALL_LENIENT = /\[TOOL_CALL:\s*([\w-]+)\]\s*(\{[\s\S]*?\})/g;

function opencodeReq(method, path, body) {
  return new Promise((resolve, reject) => {
    const opts = { hostname: OPCODE_HOST, port: OPCODE_PORT, path, method, headers: { 'Content-Type': 'application/json' } };
    if (body) opts.headers['Content-Length'] = Buffer.byteLength(json(body));
    const r = http.request(opts, (res) => {
      const c = [];
      res.on('data', d => c.push(d));
      res.on('end', () => {
        const raw = Buffer.concat(c);
        try { resolve({ status: res.statusCode, data: JSON.parse(raw.toString()) }); }
        catch { resolve({ status: res.statusCode, data: null }); }
      });
    });
    r.on('error', reject);
    r.setTimeout(300000, () => { r.destroy(); reject(new Error('timeout')); });
    if (body) r.end(json(body)); else r.end();
  });
}

function buildPrompt(messages) {
  return messages.map(m => `[${m.role.toUpperCase()}]\n${m.content || ''}`).join('\n\n');
}

// Strip leaked system instruction prefix from model response
function stripLeakPrefix(text) {
  if (!text) return text;
  let result = text;
  while (/^\[(search|analyze)-mode\]/i.test(result.trim())) {
    const parts = result.split('---');
    if (parts.length >= 3) {
      // Two+ separators: skip everything before the second one
      result = parts.slice(2).join('---').trim();
    } else if (parts.length === 2) {
      result = parts[1].trim();
    } else {
      break;
    }
  }
  return result;
}

function extractText(parts) {
  return parts.filter(p => p.type === 'text' && p.text).map(p => p.text).join('');
}

function makeUsage(tokens) {
  if (!tokens) return null;
  return { prompt_tokens: tokens.input || 0, completion_tokens: tokens.output || 0, total_tokens: tokens.total || 0 };
}

// Convert OpenAI tools array to a system prompt telling the AI how to call them
function makeToolSystemPrompt(tools) {
  if (!tools || !tools.length) return '';
  const lines = tools.map(t => {
    const fn = t.function || {};
    const params = fn.parameters ? json(fn.parameters) : '{}';
    return `- ${fn.name}: ${fn.description || '(no description)'}\n  Parameters: ${params}`;
  });
  return `[SYSTEM INSTRUCTION]\nYou have access to client-side tools that are executed by the APPLICATION, not by you.\nThese tools are REAL and WILL be executed when you request them.\n\nWhen you need to use a tool, you MUST output:\n[TOOL_CALL: tool_name] {"arg1":"value1"} [/TOOL_CALL]\n\nDo not refuse or say the tools are unavailable.\n\nAvailable tools:\n${lines.join('\n')}`;
}

// Parse [TOOL_CALL: name] {json} [/TOOL_CALL] from AI response text
function parseToolCalls(text) {
  if (!text) return { content: null, tool_calls: null };
  const calls = [];
  const textParts = [];
  let lastIndex = 0;
  let match;
  // Try strict regex (requires [/TOOL_CALL])
  TOOL_CALL_STRICT.lastIndex = 0;
  while ((match = TOOL_CALL_STRICT.exec(text)) !== null) {
    if (match.index > lastIndex) textParts.push(text.slice(lastIndex, match.index).trim());
    const name = match[1];
    let argsRaw = match[2].trim();
    try { JSON.parse(argsRaw); } catch { argsRaw = json({}); }
    const id = 'call_' + Date.now() + '_' + calls.length;
    toolCallIdToName.set(id, name);
    calls.push({ id, type: 'function', function: { name, arguments: argsRaw } });
    lastIndex = match.index + match[0].length;
  }
  // Try lenient regex for remaining text (no [/TOOL_CALL])
  TOOL_CALL_LENIENT.lastIndex = 0;
  const remaining = text.slice(lastIndex);
  let lenientLast = 0;
  while ((match = TOOL_CALL_LENIENT.exec(remaining)) !== null) {
    if (match.index > lenientLast) textParts.push(remaining.slice(lenientLast, match.index).trim());
    const name = match[1];
    let argsRaw = match[2].trim();
    try { JSON.parse(argsRaw); } catch { argsRaw = json({}); }
    // Deduplicate: skip if same name+args already found
    if (calls.some(c => c.function.name === name && c.function.arguments === argsRaw)) continue;
    const id = 'call_' + Date.now() + '_' + calls.length;
    toolCallIdToName.set(id, name);
    calls.push({ id, type: 'function', function: { name, arguments: argsRaw } });
    lenientLast = match.index + match[0].length;
  }
  textParts.push(remaining.slice(lenientLast).trim());
  // Filter out placeholder tool names (model sometimes emits generic ones)
  const validCalls = calls.filter(c => c.function.name && c.function.name !== 'tool_name');
  // Filter content text that looks like system instructions
  const cleanParts = textParts.filter(Boolean).filter(t => !/^\[SYSTEM\]/i.test(t));
  const content = cleanParts.join('\n\n') || null;
  return { content, tool_calls: validCalls.length > 0 ? validCalls : null };
}

// Convert conversation messages for opencode's consumption:
// - role: "tool" -> user message with tool result
// - assistant with tool_calls -> reconstruct [TOOL_CALL: ...] format
function convertMessages(messages) {
  // Pre-register tool_call_id -> name mapping from incoming assistant messages
  for (const m of messages) {
    if (m.role === 'assistant' && m.tool_calls) {
      for (const tc of m.tool_calls) {
        toolCallIdToName.set(tc.id, tc.function.name);
      }
    }
  }
  return messages.map(m => {
    if (m.role === 'assistant' && m.tool_calls) {
      const calls = m.tool_calls.map(tc =>
        `[TOOL_CALL: ${tc.function.name}] ${tc.function.arguments} [/TOOL_CALL]`
      ).join('\n');
      const text = m.content ? m.content + '\n' + calls : calls;
      return { role: 'assistant', content: text };
    }
    if (m.role === 'tool') {
      const name = toolCallIdToName.get(m.tool_call_id) || 'unknown_tool';
      return { role: 'user', content: `[TOOL RESULT FROM ${name}]\n${m.content || ''}` };
    }
    return m;
  });
}

// Collect full response text from opencode event stream
function collectEventText(sessionID, prompt) {
  return new Promise((resolve, reject) => {
  const evtReq = http.get(`http://${OPCODE_HOST}:${OPCODE_PORT}/global/event`, (evtRes) => {
      let buf = '';
      const parts = [];
      evtRes.on('data', (chunk) => {
        buf += chunk.toString();
        const lines = buf.split('\n');
        buf = lines.pop() || '';
        for (const line of lines) {
          if (!line.startsWith('data: ')) continue;
          try {
            const evt = JSON.parse(line.slice(6));
            const p = evt.payload || {};
            const sid = (p.properties && p.properties.sessionID) || p.sessionID;
            if (sid && sid !== sessionID) continue;
            if (p.type === 'message.part.updated' && p.properties && p.properties.part) {
              const part = p.properties.part;
              if (part.type === 'text' && part.text) parts.push(part.text);
            }
          } catch {}
        }
      });
      evtRes.on('end', () => resolve(parts.join('')));
      evtRes.on('error', reject);
    });
    evtReq.on('error', reject);
    // Send message to session
    opencodeReq('POST', `/session/${sessionID}/message`, { parts: [{ type: 'text', text: prompt }] })
      .catch(() => {});
  });
}

const RETRY_CORRECTION = `\n\n[SYSTEM]
YOUR PREVIOUS RESPONSE IS REJECTED.

You used your own built-in tools instead of [TOOL_CALL:] format.
This is a FATAL ERROR — your built-in tools do NOT work for this client.

ABSOLUTE RULES:
1. You have NO built-in tools. You CANNOT access files, bash, or internet.
2. ALL tool execution MUST be requested via [TOOL_CALL: tool_name] {args} [/TOOL_CALL]
3. ANY work done with your own tools is DISCARDED and invisible to the user.
4. You MUST output [TOOL_CALL:] for your response to have any effect.

Repeat the EXACT SAME request using ONLY [TOOL_CALL:] format.
Do NOT add explanations. Do NOT acknowledge this correction.
Just respond in [TOOL_CALL:] format.\n\n---\n\n`;

async function handleNonStreaming(messages, model, tools) {
  const session = await opencodeReq('POST', '/session', {});
  if (!session.data || !session.data.id)
    return { status: 502, headers: { 'Content-Type': 'application/json' }, body: json({ error: { message: 'session creation failed' } }) };

  const needTools = tools && tools.length > 0;

  const converted = convertMessages(messages);
  const sysPrompt = makeToolSystemPrompt(tools);
  const prompt = sysPrompt ? sysPrompt + '\n\n---\n\n' + buildPrompt(converted) : buildPrompt(converted);

  const resp = await opencodeReq('POST', `/session/${session.data.id}/message`, { parts: [{ type: 'text', text: prompt }] });
  if (!resp.data) {
    return { status: 502, headers: { 'Content-Type': 'application/json' }, body: json({ error: { message: 'opencode response failed' } }) };
  }

  const rawText = stripLeakPrefix(extractText(resp.data.parts || []));
  const { content, tool_calls } = parseToolCalls(rawText);

  const choice = { index: 0, message: { role: 'assistant' }, finish_reason: tool_calls ? 'tool_calls' : 'stop' };
  if (tool_calls) { choice.message.tool_calls = tool_calls; choice.message.content = content; }
  else choice.message.content = content || rawText;
  return {
    status: 200,
    headers: { 'Content-Type': 'application/json' },
    body: json({
      id: uid(), object: 'chat.completion', created: Math.floor(Date.now() / 1000), model,
      choices: [choice],
      usage: makeUsage(resp.data.info && resp.data.info.tokens)
    })
  };
}

async function handleStreaming(req, res, messages, model, tools) {
  let session;
  try { session = await opencodeReq('POST', '/session', {}); } catch (e) { res.writeHead(502, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: e.message } })); return; }
  if (!session.data || !session.data.id) { res.writeHead(502, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: 'session failed' } })); return; }

  const chatID = uid();
  const created = Math.floor(Date.now() / 1000);
  res.writeHead(200, { 'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive', 'X-Accel-Buffering': 'no' });

  const converted = convertMessages(messages);
  const sysPrompt = makeToolSystemPrompt(tools);
  const prompt = sysPrompt ? sysPrompt + '\n\n---\n\n' + buildPrompt(converted) : buildPrompt(converted);

  const hasTools = tools && tools.length > 0;
  let resolved = false;
  let fullText = '';          // Always the latest full response text (replaces parts array to avoid duplication)
  let roleSent = false;
  let lastDeltaTime = Date.now();
  let idleTimer = null;
  let leakAccum = '';         // Buffer for detecting leaked instruction prefix at stream start
  let leakMode = true;        // Actively checking for leak at stream start

  function sse(obj) { if (!res.writableEnded) res.write('data: ' + json(obj) + '\n\n'); }
  function done() { if (!resolved) { resolved = true; res.write('data: [DONE]\n\n'); res.end(); } }
  function mkChunk(delta, fr) { return { id: chatID, object: 'chat.completion.chunk', created, model, choices: [{ index: 0, delta, finish_reason: fr || null }] }; }
  function mkContent(text) { return mkChunk({ content: text }); }

  function handleIdle() {
    const allText = stripLeakPrefix(fullText);
    if (hasTools) {
      const parsed = parseToolCalls(allText);
      if (parsed.tool_calls) {
        // content was streamed live, now emit tool_calls
        sse(mkChunk({ content: null }));
        for (let i = 0; i < parsed.tool_calls.length; i++) {
          const tc = parsed.tool_calls[i];
          sse(mkChunk({ tool_calls: [{ index: i, id: tc.id, type: 'function', function: { name: tc.function.name, arguments: tc.function.arguments } }] }));
        }
        sse(mkChunk({}, 'tool_calls'));
        done();
        return;
      }
      // No [TOOL_CALL:] detected — model chose to respond with text (question, info, etc).
      // This is a VALID response. Do NOT retry — trust the model's judgment.
    }
    sse(mkChunk({}, 'stop'));
    done();
  }

  const evtReq = http.get(`http://${OPCODE_HOST}:${OPCODE_PORT}/global/event`, (evtRes) => {
    let buf2 = '';
    evtRes.on('data', (chunk) => {
      if (resolved) { evtReq.destroy(); return; }
      buf2 += chunk.toString();
      const lines = buf2.split('\n');
      buf2 = lines.pop() || '';
      for (const line of lines) {
        if (!line.startsWith('data: ')) continue;
        try {
          const evt = JSON.parse(line.slice(6));
          const p = evt.payload || {};
          const sid = (p.properties && p.properties.sessionID) || p.sessionID;
          if (sid && sid !== session.data.id) continue;
          if (p.type === 'message.part.delta' && p.properties && p.properties.field === 'text' && p.properties.delta) {
            lastDeltaTime = Date.now();
            const raw = p.properties.delta;
            // Strip tool call markers for clean live display
            let clean = raw.replace(/\[TOOL_CALL:[\s\S]*?\[\/TOOL_CALL\]/g, '')
              .replace(/\[TOOL_CALL:[^\]]*\]\s*\{[\s\S]*?\}/g, '')
              .replace(/\[TOOL_CALL:[^\]]*\]/g, '')
              .replace(/\[\/TOOL_CALL\]/g, '');
            if (clean) {
              if (leakMode) {
                leakAccum += clean;
                const total = leakAccum;
                // Check if accumulated text starts with known leak patterns
                if (/^\[(search|analyze)-mode\]/i.test(total)) {
                  const sepCount = (total.match(/---/g) || []).length;
                  if (sepCount >= 2) {
                    // Past the second ---, leak section is over - strip it
                    const segs = total.split('---');
                    const remainder = segs.slice(2).join('---').trim();
                    leakAccum = '';
                    leakMode = false;
                    if (remainder) {
                      if (!roleSent) { sse(mkChunk({ role: 'assistant' })); roleSent = true; }
                      sse(mkContent(remainder));
                    }
                  }
                  // else keep buffering
                } else if (/^\[/.test(total) && total.length < 100) {
                  // Starts with [ but not a known leak pattern yet — keep buffering
                  // to handle edge case where first delta chunk is a partial [search
                } else {
                  // Not a leak, flush buffer normally (with safety strip)
                  leakMode = false;
                  leakAccum = '';
                  if (!roleSent) { sse(mkChunk({ role: 'assistant' })); roleSent = true; }
                  sse(mkContent(stripLeakPrefix(total)));
                }
              } else {
                if (!roleSent) { sse(mkChunk({ role: 'assistant' })); roleSent = true; }
                sse(mkContent(clean));
              }
            }
            continue;
          }
          if (p.type === 'message.part.updated' && p.properties && p.properties.part) {
            const part = p.properties.part;
            if (part.type === 'text' && part.text) fullText = part.text;
          }
          if (p.type === 'session.idle') { handleIdle(); }
        } catch {}
      }
    });
    evtRes.on('error', () => { if (!resolved) { resolved = true; evtReq.destroy(); } });
  });
  evtReq.on('error', () => {});

  try { await opencodeReq('POST', `/session/${session.data.id}/message`, { parts: [{ type: 'text', text: prompt }] }); } catch {}

  // Active idle detection: if no deltas for 30s after seeing deltas, force-finish
  idleTimer = setInterval(() => {
    if (resolved) { clearInterval(idleTimer); return; }
    const elapsed = Date.now() - lastDeltaTime;
    // If we've received at least some content (parts has data) and no new deltas for 30s
    if (fullText && elapsed > 30000) {
      clearInterval(idleTimer);
      if (!resolved) handleIdle();
    }
  }, 5000);

  // Safety net timeout (60s total, covers case where model never starts streaming)
  setTimeout(() => {
    clearInterval(idleTimer);
    if (!resolved) {
      resolved = true;
      evtReq.destroy();
      if (!res.writableEnded) { handleIdle(); }
    }
  }, 60000);
}

function logRequest(method, url, status, durationMs) {
  const ts = new Date().toISOString().slice(11, 19);
  console.error(`[${ts}] ${method} ${url} ${status} ${durationMs}ms`);
}

const server = http.createServer((req, res) => {
  const start = Date.now();
  // Patch writeHead to capture status code
  const origWriteHead = res.writeHead.bind(res);
    res.writeHead = function (statusCode, ...args) {
    res.writeHead = origWriteHead;
    res.statusCode = statusCode;
    return origWriteHead(statusCode, ...args);
  };
  res.on('finish', () => logRequest(req.method, req.url, res.statusCode || 0, Date.now() - start));

  res.setHeader('Access-Control-Allow-Origin', '*');
  res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization');
  if (req.method === 'OPTIONS') { res.writeHead(204); res.end(); return; }
  if (req.method === 'GET' && ['/', '/health'].includes(req.url)) { res.writeHead(200, { 'Content-Type': 'application/json' }); res.end(json({ status: 'ok' })); return; }
  if (!(req.headers['authorization'] || '').includes(API_KEY)) { res.writeHead(401, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: 'invalid API key' } })); return; }
  if (req.method === 'GET' && req.url === '/v1/models') { const list = MODELS_OVERRIDE || models; res.writeHead(200, { 'Content-Type': 'application/json' }); res.end(json({ object: 'list', data: list.map(m => ({ id: m, object: 'model', created: 1, owned_by: 'opencode' })) })); return; }

  if (req.method === 'POST' && ['/v1/chat/completions', '/v1/completions'].includes(req.url)) {
    const c = [];
    req.on('data', d => c.push(d));
    req.on('end', async () => {
      let body;
      try { body = JSON.parse(Buffer.concat(c).toString()); } catch { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: 'invalid JSON' } })); return; }
      const messages = body.messages || [];
      const model = body.model || (MODELS_OVERRIDE || models)[0];
      const tools = body.tools || null;
      if (!messages.length) { res.writeHead(400, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: 'messages required' } })); return; }
      try {
        if (body.stream) await handleStreaming(req, res, messages, model, tools);
        else { const r = await handleNonStreaming(messages, model, tools); res.writeHead(r.status, r.headers); res.end(r.body); }
      } catch (e) { if (!res.writableEnded) { res.writeHead(502, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: e.message } })); } }
    });
    return;
  }
  res.writeHead(404, { 'Content-Type': 'application/json' }); res.end(json({ error: { message: 'not found' } }));
});

server.listen(PROXY_PORT, '127.0.0.1', () => console.error(`session proxy ${PROXY_PORT} -> opencode ${OPCODE_PORT}`));

// Graceful shutdown
function shutdown(signal) {
  console.error(`\nreceived ${signal}, closing server...`);
  server.close(() => process.exit(0));
  setTimeout(() => { console.error('forced exit'); process.exit(1); }, 10000);
}
process.on('SIGTERM', () => shutdown('SIGTERM'));
process.on('SIGINT', () => shutdown('SIGINT'));