import { Readable } from 'node:stream'; import { Agent, fetch as undiciFetch } from 'undici'; import { appendBalanceEvent, createSseBillingTransform } from './sse-billing.mjs'; import { developerToolsFromPolicy } from './capabilities.mjs'; import { evaluateProxyRequest } from './policies.mjs'; import { buildSandboxSessionConstraints } from './user-publish.mjs'; import { reconcileAgentSession } from './session-reconcile.mjs'; const insecureDispatcher = new Agent({ connect: { rejectUnauthorized: false }, }); function isHttpsTarget(target) { return target.startsWith('https://'); } async function apiFetch(target, apiSecret, pathname, init = {}) { const url = new URL(pathname, target); const headers = { ...(init.headers ?? {}), 'X-Secret-Key': apiSecret, }; if (init.body && !headers['Content-Type']) { headers['Content-Type'] = 'application/json'; } return undiciFetch(url, { ...init, headers, dispatcher: isHttpsTarget(target) ? insecureDispatcher : undefined, }); } async function readJsonBody(req) { if (req.method === 'GET' || req.method === 'HEAD') return null; const chunks = []; for await (const chunk of req) { chunks.push(chunk); } if (chunks.length === 0) return null; const raw = Buffer.concat(chunks).toString('utf8'); if (!raw.trim()) return null; return JSON.parse(raw); } function sendProxyResponse(res, upstream) { res.status(upstream.status); upstream.headers.forEach((value, key) => { if (key === 'transfer-encoding') return; res.setHeader(key, value); }); if (!upstream.body) { res.end(); return; } Readable.fromWeb(upstream.body).pipe(res); } function extractSessionId(req, body) { const fromParams = req.params?.sessionId ?? req.params?.id; if (fromParams) return fromParams; if (body?.session_id) return body.session_id; if (body?.sessionId) return body.sessionId; return null; } export function createTkmindProxy({ apiTarget, apiSecret, userAuth }) { const requireUser = async (req, res, next) => { try { const session = req.userSession; if (!session) { res.status(401).json({ message: '未登录' }); return; } const me = await userAuth.getMe(req.userToken); if (!me) { res.status(401).json({ message: '登录已过期' }); return; } req.currentUser = me; next(); } catch (err) { res.status(500).json({ message: err instanceof Error ? err.message : '认证失败' }); } }; const ensureChatAllowed = async (req, res, next) => { const gate = await userAuth.canUseChat(req.currentUser.id); if (!gate.ok) { res.status(402).json({ message: gate.message, code: gate.code, balanceCents: gate.balanceCents, minRechargeCents: gate.minRechargeCents, suggestedTiers: gate.suggestedTiers, }); return; } next(); }; const handlers = { 'POST /agent/start': [ requireUser, ensureChatAllowed, async (req, res) => { try { const workingDir = await userAuth.resolveWorkingDir(req.currentUser.id); const sessionPolicy = await userAuth.getAgentSessionPolicy(req.currentUser.id); const upstream = await apiFetch(apiTarget, apiSecret, '/agent/start', { method: 'POST', body: JSON.stringify({ working_dir: workingDir, enable_context_memory: sessionPolicy.enableContextMemory, ...(sessionPolicy.extensionOverrides ? { extension_overrides: sessionPolicy.extensionOverrides } : {}), ...(req.body?.recipe ? { recipe: req.body.recipe } : {}), }), }); const text = await upstream.text(); if (!upstream.ok) { res.status(upstream.status).send(text); return; } const session = JSON.parse(text); if (session?.id) { await userAuth.registerAgentSession(req.currentUser.id, session.id); if (sessionPolicy.gooseMode) { const modeRes = await apiFetch(apiTarget, apiSecret, '/agent/update_session', { method: 'POST', body: JSON.stringify({ session_id: session.id, goose_mode: sessionPolicy.gooseMode, }), }); if (!modeRes.ok) { const modeText = await modeRes.text().catch(() => ''); res.status(modeRes.status).send(modeText || '设置会话模式失败'); return; } } const publishLayout = await userAuth.getUserPublishLayout(req.currentUser.id); const api = (pathname, init) => apiFetch(apiTarget, apiSecret, pathname, init); try { await reconcileAgentSession(api, session.id, { workingDir, sessionPolicy, sandboxConstraints: publishLayout?.constraints ?? null, }); } catch (reconcileErr) { res.status(500).json({ message: reconcileErr instanceof Error ? `会话策略同步失败:${reconcileErr.message}` : '会话策略同步失败', }); return; } } res.status(upstream.status).json(session); } catch (err) { res.status(500).json({ message: err instanceof Error ? err.message : '启动会话失败' }); } }, ], 'POST /agent/resume': [ requireUser, ensureChatAllowed, async (req, res) => { try { const sessionId = req.body?.session_id; if (!sessionId) { res.status(400).json({ message: '缺少 session_id' }); return; } const owns = await userAuth.ownsSession(req.currentUser.id, sessionId); if (!owns) { res.status(403).json({ message: '无权访问该会话' }); return; } const upstream = await apiFetch(apiTarget, apiSecret, '/agent/resume', { method: 'POST', body: JSON.stringify(req.body ?? {}), }); const text = await upstream.text(); if (!upstream.ok) { res.status(upstream.status).send(text); return; } const payload = JSON.parse(text); const workingDir = await userAuth.resolveWorkingDir(req.currentUser.id); const sessionPolicy = await userAuth.getAgentSessionPolicy(req.currentUser.id); const publishLayout = await userAuth.getUserPublishLayout(req.currentUser.id); try { await reconcileAgentSession( (pathname, init) => apiFetch(apiTarget, apiSecret, pathname, init), sessionId, { workingDir, sessionPolicy, sandboxConstraints: publishLayout?.constraints ?? null, }, ); } catch (reconcileErr) { res.status(500).json({ message: reconcileErr instanceof Error ? `会话恢复后策略同步失败:${reconcileErr.message}` : '会话恢复后策略同步失败', }); return; } res.status(upstream.status).json(payload); } catch (err) { res.status(500).json({ message: err instanceof Error ? err.message : '恢复会话失败' }); } }, ], 'GET /sessions': [ requireUser, async (req, res) => { try { const upstream = await apiFetch(apiTarget, apiSecret, '/sessions', { method: 'GET', }); const text = await upstream.text(); if (!upstream.ok) { res.status(upstream.status).send(text); return; } const payload = JSON.parse(text); const owned = await userAuth.listOwnedSessionIds(req.currentUser.id); const sessions = (payload.sessions ?? []).filter((item) => owned.has(item.id)); res.json({ sessions }); } catch (err) { res.status(500).json({ message: err instanceof Error ? err.message : '读取会话失败' }); } }, ], }; const sessionScoped = (build) => [ requireUser, async (req, res, next) => { try { const body = req.body ?? (await readJsonBody(req)); req.body = body; const sessionId = extractSessionId(req, body); if (!sessionId) { res.status(400).json({ message: '缺少 session_id' }); return; } const owns = await userAuth.ownsSession(req.currentUser.id, sessionId); if (!owns) { res.status(403).json({ message: '无权访问该会话' }); return; } req.agentSessionId = sessionId; await build(req, res, next); } catch (err) { res.status(500).json({ message: err instanceof Error ? err.message : '请求失败' }); } }, ]; const proxySessionEvents = async (req, res, sessionId) => { try { const pathname = `/sessions/${sessionId}/events`; const upstream = await apiFetch(apiTarget, apiSecret, pathname, { method: 'GET', headers: { Accept: 'text/event-stream', 'Last-Event-ID': req.get('last-event-id') ?? '', }, }); if (!upstream.ok || !upstream.body) { const text = await upstream.text().catch(() => ''); res.status(upstream.status).send(text); return; } res.status(upstream.status); res.setHeader('Content-Type', 'text/event-stream'); res.setHeader('Cache-Control', 'no-cache'); res.setHeader('Connection', 'keep-alive'); let pendingBalance = null; const billingTransform = createSseBillingTransform({ onFinish: async (event) => { const result = await userAuth.billSessionUsage( req.currentUser.id, sessionId, event.token_state, null, ); if (result.ok && result.costCents > 0 && result.balanceCents != null) { pendingBalance = result.balanceCents; } }, }); const source = Readable.fromWeb(upstream.body); billingTransform.on('data', (chunk) => { res.write(chunk); if (pendingBalance != null) { res.write(appendBalanceEvent(pendingBalance)); pendingBalance = null; } }); billingTransform.on('end', () => res.end()); billingTransform.on('error', () => res.end()); source.on('error', () => res.end()); source.pipe(billingTransform); } catch (err) { res.status(502).json({ message: err instanceof Error ? err.message : 'SSE 代理失败' }); } }; const proxyFallback = async (req, res) => { try { const pathname = req.originalUrl.replace(/^\/api/, '') || '/'; const policyState = await userAuth.resolveUserPolicies(req.currentUser); const capabilityState = await userAuth.resolveUserCapabilities(req.currentUser); const gate = evaluateProxyRequest(req.method, pathname, policyState.policies, { unrestricted: policyState.unrestricted, }); if (!gate.allowed) { res.status(403).json({ message: gate.reason ?? '该 API 已被策略禁止' }); return; } if ( /^\/agent\/harness_(bootstrap|remember)$/.test(pathname) && !capabilityState.unrestricted && !capabilityState.capabilities.context_memory ) { res.status(403).json({ message: '当前账户未开通项目记忆,无法访问该 API' }); return; } const body = req.method === 'GET' || req.method === 'HEAD' ? undefined : req.body ? JSON.stringify(req.body) : undefined; const upstream = await apiFetch(apiTarget, apiSecret, pathname, { method: req.method, body, headers: { Accept: req.get('accept') ?? '*/*', 'Last-Event-ID': req.get('last-event-id') ?? '', }, }); sendProxyResponse(res, upstream); } catch (err) { res.status(502).json({ message: err instanceof Error ? err.message : '代理失败' }); } }; return { requireUser, ensureChatAllowed, handlers, sessionScoped, proxyFallback, proxySessionEvents, apiFetch: (pathname, init) => apiFetch(apiTarget, apiSecret, pathname, init), }; } export function matchHandler(handlers, method, path) { const key = `${method.toUpperCase()} ${path}`; return handlers[key] ?? null; }