jevstrudel.git / tools / mcp / proxy.mjs
proxy.mjsannotatedproxy.mjssource245 lines · 10.8 KB · raw
1// The proxy's logic, apart from stdio (strudel-mcp.mjs wires it up) so tests
2// can drive it: MCP messages in, answers and notifications out, forwarded to
3// one of two jevstrudel environments.
4//
5//   dev   the Worker under `wrangler dev` (POST /mcp, worker/src/mcp.ts): the
6//         dev hub, playing into this machine's tabs. No sign-in. The default.
7//   prod  the public site's hosted MCP (POST /jev/mcp, worker/src/hosted-mcp.ts),
8//         as the user's own account, with an OAuth access token (oauth.mjs).
9//
10// The proxy answers `initialize` and `ping` itself and adds one tool of its
11// own, `use_environment`, to whichever list it serves; everything else is
12// forwarded, one HTTP request per message. Switching sends
13// notifications/tools/list_changed, since the two lists differ.
14export const LOCAL_TOOL = {
15  name: 'use_environment',
16  description:
17    "Switch which jevstrudel this MCP server talks to, or, without `environment`, say which one it is using. `dev` (the default): the dev server on this machine (`buck2 run //:dev`), playing into local tabs with no sign-in. `prod`: the public site's hosted MCP (jevstrudel.deizel.workers.dev), acting as the user's own account: their signed-in tabs on the public site, the site's and listeners' songs, and publishing listener songs under their name. The first switch to prod opens a passkey sign-in in the user's browser and waits for it; the tool list changes to match the environment.",
18  inputSchema: {
19    type: 'object',
20    properties: { environment: { type: 'string', enum: ['dev', 'prod'], description: 'Where to send tool calls from now on.' } },
21  },
22};
23
24// How long `use_environment prod` waits for the browser before answering.
25// Claude Code's tool calls wait far longer (MCP_TOOL_TIMEOUT defaults to
26// 1e8 ms in 2.1.282), and a passkey and a consent click take well under a
27// minute; past this the sign-in carries on, and the switch happens when it
28// completes, with tools/list_changed.
29export const SIGN_IN_WAIT_MS = 120_000;
30
31const text = (id, body, isError = false) => ({ jsonrpc: '2.0', id, result: { content: [{ type: 'text', text: body }], isError } });
32
33class SignInNeeded extends Error {}
34
35// dev:   { endpoint, fallbackTools(): the tool list while the Worker is down }
36// prod:  { endpoint, session: oauth.mjs's oauthSession }
37// send:  writes one message to Claude Code
38export function createProxy({ dev, prod, send, fetchFn = fetch, signInWaitMs = SIGN_IN_WAIT_MS, log = () => {} }) {
39  let current = 'dev';
40  let listed;
41  let prodTools = null;
42
43  const notifyToolsChanged = () => send({ jsonrpc: '2.0', method: 'notifications/tools/list_changed' });
44
45  async function post(endpoint, message, token) {
46    const headers = { 'Content-Type': 'application/json', Accept: 'application/json, text/event-stream' };
47    if (token) headers.Authorization = `Bearer ${token}`;
48    return fetchFn(endpoint, { method: 'POST', headers, body: JSON.stringify(message) });
49  }
50
51  async function forwardDev(message) {
52    const res = await post(dev.endpoint, message);
53    if (!res.ok) throw new Error(`the Worker answered ${res.status}`);
54    return res.json();
55  }
56
57  // One message to prod, with the access token; on a 401 (expired, revoked)
58  // one silent refresh and a retry. Throws SignInNeeded when there is no
59  // token that works, or an Error saying what prod answered.
60  async function forwardProd(message) {
61    const { session } = prod;
62    if (!session.accessToken()) throw new SignInNeeded('not signed in');
63    let res = await post(prod.endpoint, message, session.accessToken());
64    if (res.status === 401 && (await session.refresh())) res = await post(prod.endpoint, message, session.accessToken());
65    if (res.status === 401) throw new SignInNeeded('prod refused the token: it expired, was revoked, or the app was disconnected');
66    if (res.ok) return res.json();
67    const body = (await res.text()).trim().slice(0, 300);
68    if (res.status === 429) {
69      const after = res.headers.get('Retry-After');
70      throw new Error(`jevstrudel (prod) is limiting requests: ${body || 'too many'}${after ? `; retry in ${after} s` : ''}`);
71    }
72    if (res.status === 403) {
73      const challenge = res.headers.get('WWW-Authenticate') ?? '';
74      const scope = /scope="([^"]*)"/.exec(challenge)?.[1];
75      throw new Error(`jevstrudel (prod) refused: the sign-in did not grant ${scope ? `the ${scope} scope` : 'this'}. Switch to prod again after signing in fresh with every scope allowed.${body ? ` (${body})` : ''}`);
76    }
77    throw new Error(`jevstrudel (prod) answered ${res.status}${body ? `: ${body}` : ''}`);
78  }
79
80  async function devTools() {
81    try {
82      const reply = await forwardDev({ jsonrpc: '2.0', id: 'tools', method: 'tools/list' });
83      if (reply.result?.tools) return reply.result.tools;
84    } catch {}
85    return dev.fallbackTools();
86  }
87
88  // prod's list, never opening the browser: the last one seen when prod
89  // cannot be asked.
90  async function prodToolList() {
91    try {
92      const reply = await forwardProd({ jsonrpc: '2.0', id: 'tools', method: 'tools/list' });
93      if (reply.result?.tools) prodTools = reply.result.tools;
94    } catch (e) {
95      log(`tools/list from prod: ${e.message}`);
96    }
97    return prodTools ?? [];
98  }
99
100  const tools = async () => [...(current === 'prod' ? await prodToolList() : await devTools()), LOCAL_TOOL];
101
102  // The dev Worker's tools can change under a running session (wrangler
103  // reloads it); prod's change only with a deploy, and every request there
104  // counts against the user's MCP rate, so only dev is polled.
105  async function poll() {
106    if (current !== 'dev') return;
107    const now = JSON.stringify(await tools());
108    if (current !== 'dev') return;
109    if (listed !== undefined && now !== listed) notifyToolsChanged();
110    listed = now;
111  }
112
113  function switchTo(environment) {
114    if (current === environment) return false;
115    current = environment;
116    listed = undefined;
117    notifyToolsChanged();
118    return true;
119  }
120
121  const where = (environment) => (environment === 'prod' ? `prod (${prod.endpoint})` : `dev (${dev.endpoint})`);
122
123  // What prod tells a client at initialize, for the model to read once.
124  async function prodInstructions() {
125    try {
126      const reply = await forwardProd({
127        jsonrpc: '2.0',
128        id: 'init',
129        method: 'initialize',
130        params: { protocolVersion: '2025-06-18', capabilities: {}, clientInfo: { name: 'jevstrudel-proxy', version: '1.0.0' } },
131      });
132      return reply.result?.instructions ?? '';
133    } catch {
134      return '';
135    }
136  }
137
138  async function signedInToProd(id, changed) {
139    prodTools = null;
140    const instructions = await prodInstructions();
141    return text(id, `${changed ? 'Switched to' : 'Using'} ${where('prod')}, signed in. The tool list is prod's now.${instructions ? `\n\n${instructions}` : ''}`);
142  }
143
144  // Start (or rejoin) the browser sign-in; when it completes, switch to prod.
145  async function beginSignIn() {
146    const flow = await prod.session.signIn();
147    if (flow.url) {
148      log(`signing in to prod: ${flow.url}`);
149      flow.done.then(
150        () => switchTo('prod'),
151        (e) => log(`sign-in to prod failed: ${e.message}`),
152      );
153    }
154    return flow;
155  }
156
157  const signInText = (flow, what) =>
158    `${what} Sign in with a passkey and allow the app in the browser (${flow.opened}); if no tab opened, open this by hand:\n${flow.url}`;
159
160  async function useEnvironment(id, args) {
161    const environment = args?.environment;
162    if (environment === undefined) {
163      const signedIn = prod.session.accessToken() ? 'signed in' : 'not signed in';
164      return text(id, `Using ${where(current)}. prod: ${signedIn}.`);
165    }
166    if (environment === 'dev') {
167      const changed = switchTo('dev');
168      return text(id, `${changed ? 'Switched to' : 'Using'} ${where('dev')}. The tool list is dev's${changed ? ' now' : ''}.`);
169    }
170    if (environment !== 'prod') return text(id, 'environment is "dev" or "prod"', true);
171    if (prod.session.accessToken()) return signedInToProd(id, switchTo('prod'));
172    let flow;
173    try {
174      flow = await beginSignIn();
175    } catch (e) {
176      return text(id, `Could not start signing in to ${prod.endpoint}: ${e.message}`, true);
177    }
178    if (!flow.url) return signedInToProd(id, switchTo('prod'));
179    let timer;
180    const outcome = await Promise.race([
181      flow.done.then(
182        () => 'done',
183        (e) => e,
184      ),
185      new Promise((resolve) => (timer = setTimeout(() => resolve('waiting'), signInWaitMs))),
186    ]);
187    clearTimeout(timer);
188    if (outcome === 'done') return signedInToProd(id, switchTo('prod') || current === 'prod');
189    if (outcome instanceof Error) return text(id, `Signing in to prod failed: ${outcome.message}. Still using ${where(current)}.`, true);
190    return text(
191      id,
192      signInText(flow, `Still waiting for the sign-in to prod; still using ${where(current)}.`) +
193        `\nThe switch happens by itself when it completes (the tool list will change); or call use_environment again.`,
194    );
195  }
196
197  async function handle(message) {
198    const { id, method, params } = message;
199    if (method === 'initialize') {
200      return {
201        jsonrpc: '2.0',
202        id,
203        result: {
204          protocolVersion: params?.protocolVersion ?? '2025-06-18',
205          capabilities: { tools: { listChanged: true } },
206          serverInfo: { name: 'jevstrudel', version: '1.0.0' },
207        },
208      };
209    }
210    if (method === 'ping') return { jsonrpc: '2.0', id, result: {} };
211    if (method === 'tools/list') {
212      const list = await tools();
213      if (current === 'dev') listed = JSON.stringify(list);
214      return { jsonrpc: '2.0', id, result: { tools: list } };
215    }
216    if (method === 'tools/call' && params?.name === LOCAL_TOOL.name) return useEnvironment(id, params.arguments);
217
218    if (current === 'prod') {
219      try {
220        return await forwardProd(message);
221      } catch (e) {
222        if (method !== 'tools/call') return { jsonrpc: '2.0', id, error: { code: -32603, message: e.message } };
223        if (!(e instanceof SignInNeeded)) return text(id, e.message, true);
224        // the token is gone: sign in again now, and say where
225        try {
226          const flow = await beginSignIn();
227          if (!flow.url) return await forwardProd(message);
228          return text(id, signInText(flow, `jevstrudel (prod) needs a new sign-in: ${e.message}.`) + '\nThen call the tool again.', true);
229        } catch (e2) {
230          return text(id, `jevstrudel (prod) needs a new sign-in (${e.message}), and it could not start: ${e2.message}`, true);
231        }
232      }
233    }
234    try {
235      return await forwardDev(message);
236    } catch (e) {
237      if (method === 'tools/call') {
238        return text(id, `The jevstrudel Worker is not answering at ${dev.endpoint} (${e.message}). Start the dev server: \`buck2 run //:dev\`.`, true);
239      }
240      return { jsonrpc: '2.0', id, error: { code: -32603, message: e.message } };
241    }
242  }
243
244  return { handle, poll, environment: () => current };
245}