Skip to content

feat: add delta-prompt optimization and enhance proxy security #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zakat0v-d0t wants to merge 1 commit into
base: main
Choose a base branch
from
+149 −36
Open

feat: add delta-prompt optimization and enhance proxy security #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 105 additions & 36 deletions server.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ const SERVER_PUBLIC_IP = (() => {

const FORGETMEAI_WATERMARK = 't.me/forgetmeai';
const PORT = Number(process.env.PORT || 9655);
const HOST = process.env.HOST || '0.0.0.0';
const HOST = process.env.HOST || '127.0.0.1';
function formatWatermark(prefix = 'ForgetMeAI') { return `${prefix}: ${FORGETMEAI_WATERMARK}`; }
function printBanner() {
console.log(`
Expand Down Expand Up @@ -145,10 +145,7 @@ function selectAccountForSession(session) {
// A DeepSeek chat_session belongs to the auth account that created it.
// If that account is rate-limited/expired, do not keep hammering it;
// reset the web session and let a healthy account take over.
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
}
session.accountId = null;
}
Expand Down Expand Up @@ -197,9 +194,42 @@ function createSession() {
messageCount: 0,
accountId: null,
history: [],
processedMessageCount: 0,
lastProcessedMessages: [],
};
}

function resetSession(session) {
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
session.processedMessageCount = 0;
session.lastProcessedMessages = [];
}

function messagesMatch(a, b) {
if (!a || !b) return false;
if (a.length !== b.length) return false;
for (let i = 0; i < a.length; i++) {
if (a[i].role !== b[i].role) return false;
if (JSON.stringify(a[i].content) !== JSON.stringify(b[i].content)) return false;
if (JSON.stringify(a[i].tool_calls) !== JSON.stringify(b[i].tool_calls)) return false;
}
return true;
}

function maskAgentId(agentId) {
if (!agentId) return 'unknown';
// Check if agentId is an IPv4 or IPv6 address
const isIp = /^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$/.test(agentId) || agentId.includes(':');
if (isIp) {
const hash = require('crypto').createHash('sha256').update(agentId).digest('hex');
return `ip_${hash.substring(0, 8)}`;
}
return agentId; // Keep custom user-supplied session names visible
}

function getOrCreateAgentSession(agentId) {
if (!sessions.has(agentId)) {
sessions.set(agentId, createSession());
Expand Down Expand Up @@ -347,20 +377,14 @@ async function askDeepSeekStream(prompt, agentId, model = 'deepseek-default') {
// Auto-reset on deep message chain
if (session.id && session.messageCount >= MAX_MESSAGE_DEPTH) {
console.log(`${agentTag} Session ${session.id} hit ${session.messageCount} messages. Auto-resetting.`);
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
// History preserved for context injection
}

// Reset expired sessions (DeepSeek web sessions last ~1-2 hours)
if (session.id && session.createdAt && (Date.now() - session.createdAt > SESSION_TTL_MS)) {
console.log(`${agentTag} Session ${session.id} expired (age: ${Math.round((Date.now() - session.createdAt) / 60000)}min). Creating new...`);
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
}

const cr = await fetch('https://chat.deepseek.com/api/v0/chat/create_pow_challenge', {
Expand Down Expand Up @@ -425,10 +449,7 @@ async function askDeepSeekStream(prompt, agentId, model = 'deepseek-default') {
console.log(`${agentTag} Session error (${resp.status}): ${errText.substring(0, 100)}`);
if (resp.status === 400 || resp.status === 404 || resp.status === 500) {
console.log(`${agentTag} Session ${session.id} expired. Creating new session...`);
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);

const sr2 = await fetch('https://chat.deepseek.com/api/v0/chat_session/create', {
method: 'POST', headers: dsHeaders, body: '{}'
Expand Down Expand Up @@ -1070,6 +1091,24 @@ const server = http.createServer(async (req, res) => {

const url = new URL(req.url, `http://${req.headers.host || 'localhost'}`);

// Verify API Key if PROXY_API_KEY is configured
if (process.env.PROXY_API_KEY) {
if (url.pathname !== '/' && url.pathname !== '/health') {
const authHeader = req.headers.authorization || '';
const expectedAuth = `Bearer ${process.env.PROXY_API_KEY}`;
if (authHeader !== expectedAuth) {
res.writeHead(401, { 'Content-Type': 'application/json' });
res.end(JSON.stringify({
error: {
message: "Unauthorized: Invalid or missing API Key (Authorization: Bearer <key>)",
type: "invalid_request_error"
}
}));
return;
}
}
}

// Health check
if (req.method === 'GET' && (url.pathname === '/' || url.pathname === '/health')) {
res.writeHead(200, { 'Content-Type': 'application/json' });
Expand All @@ -1096,7 +1135,7 @@ const server = http.createServer(async (req, res) => {
const agentList = [];
for (const [agentId, session] of sessions) {
agentList.push({
agent: agentId,
agent: maskAgentId(agentId),
session_id: session.id,
message_count: session.messageCount,
account: session.accountId,
Expand Down Expand Up @@ -1127,10 +1166,7 @@ const server = http.createServer(async (req, res) => {
}
const historyCount = session.history.length;
const historyPreview = session.history.map(e => e.user.substring(0, 40)).join(' | ');
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
res.writeHead(200, { 'Content-Type': 'application/json' });
res.end(JSON.stringify({ status: 'session_reset', agent: agentId, history_preserved: historyCount, history: historyPreview }));
return;
Expand Down Expand Up @@ -1172,10 +1208,36 @@ const server = http.createServer(async (req, res) => {
? String(requestedSession)
: ((remoteAddr === '127.0.0.1' || remoteAddr === '::1' || remoteAddr === '::ffff:127.0.0.1') ? 'dev-agent' : remoteAddr);
const agentTag = `[${agentId}]`;
const { prompt, systemPrompt } = formatMessages(messages, tools);

const session = getOrCreateAgentSession(agentId);

// Check if we can reuse the session and only send the delta
let newMessages = messages;
let isDelta = false;

if (session.id) {
const count = session.processedMessageCount || 0;
const prefix = messages.slice(0, count);
if (count > 0 && count <= messages.length && messagesMatch(prefix, session.lastProcessedMessages)) {
// Filter out any assistant messages from the new message slice
newMessages = messages.slice(count).filter(m => m.role !== 'assistant');
isDelta = true;
console.log(`${agentTag} Sending delta of ${newMessages.length} messages (reusing session ${session.id})`);
} else {
// Client restarted or history diverged, reset session
console.log(`${agentTag} History diverged or client restarted. Resetting session ${session.id}`);
resetSession(session);
}
}

if (isDelta && newMessages.length === 0) {
console.log(`${agentTag} Warning: Delta has 0 new messages after filtering. Resetting session.`);
resetSession(session);
isDelta = false;
newMessages = messages;
}

const { prompt, systemPrompt } = formatMessages(newMessages, tools);

// Build history prefix if starting fresh
let historyPrefix = '';
if (!session.id && session.history.length > 0) {
Expand All @@ -1186,9 +1248,14 @@ const server = http.createServer(async (req, res) => {
historyPrefix += '[Continue from here]\n\n';
}

const fullPrompt = systemPrompt
? `${systemPrompt}\n\n${historyPrefix}${prompt}`
: `${historyPrefix}${prompt}`;
let fullPrompt;
if (isDelta) {
fullPrompt = prompt;
} else {
fullPrompt = systemPrompt
? `${systemPrompt}\n\n${historyPrefix}${prompt}`
: `${historyPrefix}${prompt}`;
}

const startTime = Date.now();
const { resp: dsResp } = await askDeepSeekStream(fullPrompt, agentId, requestedModel);
Expand Down Expand Up @@ -1322,10 +1389,7 @@ const server = http.createServer(async (req, res) => {
return;
}
console.log(`${agentTag} Empty response (msg#${session.messageCount}, retry ${retryAttempt}/${MAX_RETRIES}). Resetting session...`);
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
// Brief delay before retry to let DeepSeek breathe
await new Promise(r => setTimeout(r, Math.min(1000 * retryAttempt, 5000)));
const { resp: retryResp } = await askDeepSeekStream(fullPrompt, agentId, requestedModel);
Expand Down Expand Up @@ -1367,10 +1431,7 @@ const server = http.createServer(async (req, res) => {
// Retry if TOOL_CALL was found but JSON was truncated/invalid
if (!toolCall && /TOOL_CALL:\s*\w/i.test(fullContent)) {
console.log(`${agentTag} TOOL_CALL detected but JSON invalid/truncated (${fullContent.length} chars). Retrying with stricter prompt...`);
session.id = null;
session.parentMessageId = null;
session.createdAt = null;
session.messageCount = 0;
resetSession(session);
await new Promise(r => setTimeout(r, 1000));
const strictPrompt = fullPrompt + '\n\n[STRICT INSTRUCTION] Your previous response had a TOOL_CALL but the arguments were too long and got cut off. Keep the arguments SHORT — no large file contents. Just use a minimal example or reference the file by name. Output ONLY: TOOL_CALL: <function>\narguments: <short JSON>';
const { resp: retryResp2 } = await askDeepSeekStream(strictPrompt, agentId, requestedModel);
Expand Down Expand Up @@ -1403,6 +1464,10 @@ const server = http.createServer(async (req, res) => {

storeHistory(agentId, prompt, fullContent, toolCall);

// Record processed messages to allow sending delta next time
session.processedMessageCount = messages.length;
session.lastProcessedMessages = messages.map(m => ({ role: m.role, content: m.content, tool_calls: m.tool_calls }));

const openaiResponse = toolCall
? buildToolCallResponse(toolCall, requestedModel, fullPrompt, reasoningContent)
: buildTextResponse(fullContent, fullPrompt, requestedModel, reasoningContent);
Expand Down Expand Up @@ -1506,4 +1571,8 @@ async function main() {
});
}

main().catch(err => { console.error('[DS-API] FATAL:', err); process.exit(1); });
if (require.main === module) {
main().catch(err => { console.error('[DS-API] FATAL:', err); process.exit(1); });
} else {
module.exports = { createSession, resetSession, messagesMatch, formatMessages, maskAgentId };
}
44 changes: 44 additions & 0 deletions tests/unit.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,3 +109,47 @@ test('chrome auth prints actionable OS instructions when Chrome is missing', ()
assert.match(out, /Linux/i);
assert.match(out, /CHROME_PATH/i);
});

test('messagesMatch correctly matches role, content, and tool_calls', () => {
const { messagesMatch } = require('../server.js');
const a = [{ role: 'user', content: 'hello', tool_calls: undefined }];
const b = [{ role: 'user', content: 'hello', tool_calls: undefined }];
assert.ok(messagesMatch(a, b));

const c = [{ role: 'user', content: 'hello', tool_calls: [{ id: 'call_1' }] }];
assert.ok(!messagesMatch(a, c));

const d = [{ role: 'assistant', content: 'hello' }];
assert.ok(!messagesMatch(a, d));
});

test('createSession and resetSession initialize and clear delta fields', () => {
const { createSession, resetSession } = require('../server.js');
const session = createSession();
assert.equal(session.processedMessageCount, 0);
assert.deepEqual(session.lastProcessedMessages, []);

session.processedMessageCount = 5;
session.lastProcessedMessages = [{ role: 'user', content: 'hi' }];
resetSession(session);
assert.equal(session.processedMessageCount, 0);
assert.deepEqual(session.lastProcessedMessages, []);
});

test('maskAgentId correctly hashes IPs and preserves user strings', () => {
const { maskAgentId } = require('../server.js');

// Check IPv4
const v4 = maskAgentId('192.168.1.1');
assert.match(v4, /^ip_[0-9a-f]{8}$/);
assert.notEqual(v4, '192.168.1.1');

// Check IPv6
const v6 = maskAgentId('2001:db8::1');
assert.match(v6, /^ip_[0-9a-f]{8}$/);
assert.notEqual(v6, '2001:db8::1');

// Check custom agent strings
assert.equal(maskAgentId('my-custom-agent'), 'my-custom-agent');
assert.equal(maskAgentId('dev-agent'), 'dev-agent');
});