Skip to content

Fix request-scoped direct tool aliases #3102

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
dezren39 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix request-scoped direct tool aliases #3102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
02459d1
Fix request-scoped direct tool aliases
dezren39 Oct 31, 2026
cab9ea9
Merge main into request-scoped tool alias PR
dezren39 Nov 3, 2026
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
71 changes: 39 additions & 32 deletions bin/muse-proxy.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1101,31 +1101,32 @@ const contentPartsForInput = (content) => {
// model's tool calls back on the way out. The CLI door never needed this —
// the CLI sends its own toolset.
const DIRECT_TOOL_ALIASES = { shell: 'run_command', bash: 'terminal' };
const DIRECT_TOOL_NAMES = Object.fromEntries(
Object.entries(DIRECT_TOOL_ALIASES).map(([client, alias]) => [alias, client])
);
const directToolAlias = (name, taken) => {
const alias = DIRECT_TOOL_ALIASES[name];
if (!alias || (taken && taken.has(alias))) return name;
return alias;
if (!alias) return name;
let candidate = alias;
let suffix = 2;
while (taken && taken.has(candidate)) candidate = `${alias}_${suffix++}`;
return candidate;
};
const clientToolName = (name) => DIRECT_TOOL_NAMES[name] || name;
const clientToolName = (name, inverse) => inverse.get(name) || name;

// The /v1/responses passthrough forwards client bodies in Responses shape;
// apply the same tool-name aliasing there, in both directions, so a native
// client (the muse-direct provider) passes the attestation gate too. An
// unparseable body forwards verbatim, exactly as before.
const aliasResponsesRequestBody = (bodyText) => {
let parsed;
try { parsed = JSON.parse(bodyText); } catch { return { text: bodyText, aliased: false }; }
if (!parsed || typeof parsed !== 'object') return { text: bodyText, aliased: false };
try { parsed = JSON.parse(bodyText); } catch { return { text: bodyText, aliased: false, inverse: new Map() }; }
if (!parsed || typeof parsed !== 'object') return { text: bodyText, aliased: false, inverse: new Map() };
const taken = new Set((Array.isArray(parsed.tools) ? parsed.tools : [])
.map(t => t && t.name).filter(Boolean));
let aliased = false;
const inverse = new Map();
const rename = (holder) => {
if (holder && typeof holder.name === 'string') {
const next = directToolAlias(holder.name, taken);
if (next !== holder.name) { holder.name = next; aliased = true; }
if (next !== holder.name) { inverse.set(next, holder.name); holder.name = next; aliased = true; }
}
};
if (Array.isArray(parsed.tools)) for (const t of parsed.tools) rename(t);
Expand All @@ -1135,32 +1136,32 @@ const aliasResponsesRequestBody = (bodyText) => {
}
}
if (parsed.tool_choice && parsed.tool_choice.type === 'function') rename(parsed.tool_choice);
return aliased ? { text: JSON.stringify(parsed), aliased } : { text: bodyText, aliased };
return aliased ? { text: JSON.stringify(parsed), aliased, inverse } : { text: bodyText, aliased, inverse };
};

const unaliasResponsesOutput = (holder) => {
const unaliasResponsesOutput = (holder, inverse) => {
if (!holder || typeof holder !== 'object') return;
for (const item of holder.output || []) {
if (item && item.type === 'function_call' && typeof item.name === 'string') {
item.name = clientToolName(item.name);
item.name = clientToolName(item.name, inverse);
}
}
};

// Rewrite one SSE data payload if it carries an aliased tool name in a
// function_call item; anything else passes through untouched.
const unaliasSseData = (data) => {
if (!data.includes('run_command') && !data.includes('terminal')) return data;
const unaliasSseData = (data, inverse) => {
if (![...inverse.keys()].some(name => data.includes(name))) return data;
let evt;
try { evt = JSON.parse(data); } catch { return data; }
if (evt && evt.item && evt.item.type === 'function_call' && typeof evt.item.name === 'string') {
evt.item.name = clientToolName(evt.item.name);
evt.item.name = clientToolName(evt.item.name, inverse);
}
if (evt && evt.response) unaliasResponsesOutput(evt.response);
if (evt && evt.response) unaliasResponsesOutput(evt.response, inverse);
return JSON.stringify(evt);
};

const buildResponsesRequest = (body, baseModel, effort) => {
const buildResponsesRequest = (body, baseModel, effort, inverse) => {
const messages = Array.isArray(body.messages) ? body.messages : [];
const instructions = messages
.filter(m => m.role === 'system')
Expand All @@ -1169,6 +1170,11 @@ const buildResponsesRequest = (body, baseModel, effort) => {
.join('\n\n');
const takenNames = new Set((Array.isArray(body.tools) ? body.tools : [])
.map(t => t && t.function && t.function.name).filter(Boolean));
const alias = (name) => {
const next = directToolAlias(name, takenNames);
if (next !== name) inverse.set(next, name);
return next;
};
const input = [];
for (const m of messages) {
if (m.role === 'system') continue;
Expand All @@ -1181,7 +1187,7 @@ const buildResponsesRequest = (body, baseModel, effort) => {
input.push({
type: 'function_call',
call_id: tc.id,
name: tc.function ? directToolAlias(tc.function.name, takenNames) : undefined,
name: tc.function ? alias(tc.function.name) : undefined,
arguments: (tc.function && tc.function.arguments) || ''
});
}
Expand All @@ -1199,13 +1205,13 @@ const buildResponsesRequest = (body, baseModel, effort) => {
if (instructions) out.instructions = instructions;
if (Array.isArray(body.tools) && body.tools.length > 0) {
out.tools = body.tools.map(t => (t && t.type === 'function' && t.function)
? { type: 'function', name: directToolAlias(t.function.name, takenNames), description: t.function.description, parameters: t.function.parameters }
? { type: 'function', name: alias(t.function.name), description: t.function.description, parameters: t.function.parameters }
: t);
}
if (body.tool_choice !== undefined) {
const tc = body.tool_choice;
if (typeof tc === 'string') out.tool_choice = tc;
else if (tc && tc.type === 'function' && tc.function && tc.function.name) out.tool_choice = { type: 'function', name: directToolAlias(tc.function.name, takenNames) };
else if (tc && tc.type === 'function' && tc.function && tc.function.name) out.tool_choice = { type: 'function', name: alias(tc.function.name) };
else out.tool_choice = tc;
}
const maxOut = body.max_completion_tokens !== undefined && body.max_completion_tokens !== null
Expand All @@ -1219,7 +1225,7 @@ const buildResponsesRequest = (body, baseModel, effort) => {
return out;
};

const extractResponsesOutput = (respBody) => {
const extractResponsesOutput = (respBody, inverse) => {
let content = '';
const toolCalls = [];
for (const item of (respBody && respBody.output) || []) {
Expand All @@ -1232,7 +1238,7 @@ const extractResponsesOutput = (respBody) => {
toolCalls.push({
id: item.call_id || item.id,
type: 'function',
function: { name: clientToolName(item.name), arguments: item.arguments || '' }
function: { name: clientToolName(item.name, inverse), arguments: item.arguments || '' }
});
}
}
Expand Down Expand Up @@ -1266,8 +1272,8 @@ const chatUsageFromResponses = (usage) => {
};
};

const responsesToChatCompletion = (respBody, { chatId, created, fullName }) => {
const { content, toolCalls } = extractResponsesOutput(respBody);
const responsesToChatCompletion = (respBody, { chatId, created, fullName, inverse }) => {
const { content, toolCalls } = extractResponsesOutput(respBody, inverse);
const message = { role: 'assistant', content };
if (toolCalls.length > 0) message.tool_calls = toolCalls;
return {
Expand All @@ -1285,7 +1291,7 @@ const responsesToChatCompletion = (respBody, { chatId, created, fullName }) => {
};

// Translate an upstream Responses SSE stream into chat-completion chunks.
const streamResponsesAsChat = async (response, { res, chatId, created, fullName }) => {
const streamResponsesAsChat = async (response, { res, chatId, created, fullName, inverse }) => {
const sendChunk = (delta, finishReason = null, usage) => {
if (res.writableEnded) return;
const chunk = {
Expand Down Expand Up @@ -1326,7 +1332,7 @@ const streamResponsesAsChat = async (response, { res, chatId, created, fullName
index,
id: evt.item.call_id || evt.item.id,
type: 'function',
function: { name: clientToolName(evt.item.name || ''), arguments: '' }
function: { name: clientToolName(evt.item.name || '', inverse), arguments: '' }
}]
});
} else if (type === 'response.function_call_arguments.delta') {
Expand All @@ -1339,7 +1345,7 @@ const streamResponsesAsChat = async (response, { res, chatId, created, fullName
});
} else if (type === 'response.completed') {
const respBody = evt.response || {};
const { toolCalls } = extractResponsesOutput(respBody);
const { toolCalls } = extractResponsesOutput(respBody, inverse);
sendChunk({}, finishReasonFromResponses(respBody, toolCalls.length > 0), chatUsageFromResponses(respBody.usage));
if (!res.writableEnded) res.write('data: [DONE]\n\n');
res.end();
Expand Down Expand Up @@ -1397,7 +1403,8 @@ const tryDirectChat = async ({ body, baseModel, effort, fullName, stream, res, c
writeJsonBody(res, 502, { error: { message: keyResult.message } });
return true;
}
const upstreamBody = JSON.stringify(buildResponsesRequest(body, baseModel, effort));
const inverse = new Map();
const upstreamBody = JSON.stringify(buildResponsesRequest(body, baseModel, effort, inverse));
let outcome;
try {
outcome = await directResponsesFetch(keyResult, upstreamBody, stream ? 'text/event-stream' : 'application/json');
Expand Down Expand Up @@ -1451,7 +1458,7 @@ const tryDirectChat = async ({ body, baseModel, effort, fullName, stream, res, c
writeJsonBody(res, 500, { error: { message: `Direct Meta response incomplete: ${incompleteReason}` } });
return true;
}
writeJsonBody(res, 200, responsesToChatCompletion(respBody, { chatId, created, fullName }));
writeJsonBody(res, 200, responsesToChatCompletion(respBody, { chatId, created, fullName, inverse }));
return true;
}
res.writeHead(200, {
Expand All @@ -1460,7 +1467,7 @@ const tryDirectChat = async ({ body, baseModel, effort, fullName, stream, res, c
'Connection': 'keep-alive',
'X-Accel-Buffering': 'no'
});
await streamResponsesAsChat(response, { res, chatId, created, fullName });
await streamResponsesAsChat(response, { res, chatId, created, fullName, inverse });
return true;
};

Expand Down Expand Up @@ -1610,7 +1617,7 @@ const server = http.createServer((req, res) => {
let buffer = '';
const flushLine = (line) => {
if (line.startsWith('data: ') && line.length > 6 && !line.startsWith('data: [DONE]')) {
res.write(`data: ${unaliasSseData(line.slice(6))}\n`);
res.write(`data: ${unaliasSseData(line.slice(6), aliasedBody.inverse)}\n`);
} else {
res.write(line + '\n');
}
Expand All @@ -1633,7 +1640,7 @@ const server = http.createServer((req, res) => {
let text = await response.text().catch(() => '');
try {
const parsedJson = JSON.parse(text);
unaliasResponsesOutput(parsedJson);
unaliasResponsesOutput(parsedJson, aliasedBody.inverse);
text = JSON.stringify(parsedJson);
} catch { /* not JSON: pass through as-is */ }
res.writeHead(response.status, { 'Content-Type': contentType });
Expand Down
43 changes: 40 additions & 3 deletions bin/muse-proxy.test.mjs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -630,12 +630,12 @@ const startMetaStub = async (mode) => {
try { parsed = JSON.parse(body || '{}'); } catch { /* keep empty */ }
if (parsed.stream) {
res.writeHead(200, { 'Content-Type': 'text/event-stream' });
res.end(mode === 'shell' ? RESPONSES_SSE_FIXTURE.replaceAll('get_weather', 'run_command') : RESPONSES_SSE_FIXTURE);
res.end(['shell', 'terminal'].includes(mode) ? RESPONSES_SSE_FIXTURE.replaceAll('get_weather', mode === 'shell' ? 'run_command' : 'terminal') : RESPONSES_SSE_FIXTURE);
return;
}
res.writeHead(200, { 'Content-Type': 'application/json' });
res.end(mode === 'shell'
? JSON.stringify(RESPONSES_FIXTURE).replaceAll('get_weather', 'run_command')
res.end(['shell', 'terminal'].includes(mode)
? JSON.stringify(RESPONSES_FIXTURE).replaceAll('get_weather', mode === 'shell' ? 'run_command' : 'terminal')
: JSON.stringify(RESPONSES_FIXTURE));
return;
}
Expand Down Expand Up @@ -1023,6 +1023,43 @@ test('direct mode: contributor ids reach upstream verbatim, effort suffix parses
assert.equal(sent[2].model, 'muse-spark-1.3', 'plain base model stays plain upstream');
});

// A real tool with an alias-shaped name must not be renamed on return, even
// when another tool in the same request is translated.
for (const [upstreamName, blockedName] of [['run_command', 'shell'], ['terminal', 'bash']]) {
for (const route of ['chat', 'responses']) {
for (const stream of [false, true]) {
for (const collision of [false, true]) {
test(`direct ${route} ${stream ? 'stream' : 'json'} preserves ${upstreamName}${collision ? ' with collision' : ''}`, async (t) => {
const meta = await startMetaStub(upstreamName === 'terminal' ? 'terminal' : 'shell');
t.after(() => meta.close());
const { port, stop, home } = await startProxy({ MUSE_PROXY_META_API_BASE: meta.url });
t.after(stop);
writeStubAuth(home);
const names = [upstreamName, ...(collision ? [blockedName] : []), upstreamName === 'terminal' ? 'shell' : 'bash'];
const tools = names.map(name => route === 'chat'
? { type: 'function', function: { name, parameters: { type: 'object' } } }
: { type: 'function', name, parameters: { type: 'object' } });
const payload = route === 'chat'
? { model: 'muse-spark-1.3', stream, messages: [{ role: 'user', content: 'hi' }], tools }
: { model: 'muse-spark-1.3', stream, input: [{ role: 'user', content: 'hi' }], tools };
const res = await fetch(`http://127.0.0.1:${port}/v1/${route === 'chat' ? 'chat/completions' : 'responses'}`, {
method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(payload)
});
assert.equal(res.status, 200);
const text = await res.text();
assert.match(text, new RegExp(`"name"\\s*:\\s*"${upstreamName}"`));
const sent = JSON.parse(meta.requests.find(r => r.url === '/v1/responses').body);
assert.ok(sent.tools.some(tool => tool.name === upstreamName));
if (collision) {
assert.ok(!sent.tools.some(tool => tool.name === blockedName));
assert.ok(sent.tools.some(tool => tool.name === `${upstreamName}_2`));
}
});
}
}
}
}

test('non-object JSON body is a 400 and the proxy stays alive (#2816)', async (t) => {
const { port, stop } = await startProxy();
t.after(stop);
Expand Down