Don't crash ST server on invalid streaming URL

This commit is contained in:
Cohee 2023-10-08 23:42:28 +03:00
parent 54d52a2986
commit d265179f46
2 changed files with 28 additions and 2 deletions

View File

@ -227,6 +227,10 @@ export function isAphrodite() {
return textgenerationwebui_settings.type === textgen_types.APHRODITE; return textgenerationwebui_settings.type === textgen_types.APHRODITE;
} }
export function isOoba() {
return textgenerationwebui_settings.type === textgen_types.OOBA;
}
export function getTextGenUrlSourceId() { export function getTextGenUrlSourceId() {
switch (textgenerationwebui_settings.type) { switch (textgenerationwebui_settings.type) {
case textgen_types.MANCER: case textgen_types.MANCER:
@ -327,6 +331,18 @@ async function generateTextGenWithStreaming(generate_data, signal) {
streamingUrl = api_server_textgenerationwebui; streamingUrl = api_server_textgenerationwebui;
} }
if (isMancer() || isOoba()) {
try {
const parsedUrl = new URL(streamingUrl);
if (parsedUrl.protocol !== 'ws:' && parsedUrl.protocol !== 'wss:') {
throw new Error('Invalid protocol');
}
} catch {
toastr.error('Invalid URL for streaming. Make sure it starts with ws:// or wss://');
return async function* () { throw new Error('Invalid URL for streaming.'); }
}
}
const response = await fetch('/generate_textgenerationwebui', { const response = await fetch('/generate_textgenerationwebui', {
headers: { headers: {
...getRequestHeaders(), ...getRequestHeaders(),

View File

@ -552,8 +552,18 @@ app.post("/generate_textgenerationwebui", jsonParser, async function (request, r
}); });
async function* readWebsocket() { async function* readWebsocket() {
const streamingUrl = new URL(streamingUrlString); /** @type {WebSocket} */
const websocket = new WebSocket(streamingUrl); let websocket;
/** @type {URL} */
let streamingUrl;
try {
const streamingUrl = new URL(streamingUrlString);
websocket = new WebSocket(streamingUrl);
} catch (error) {
console.log("[SillyTavern] Socket error", error);
return;
}
websocket.on('open', async function () { websocket.on('open', async function () {
console.log('WebSocket opened'); console.log('WebSocket opened');