1// proxy.mjs: the environment switch and forwarding, against a fake dev
2// Worker, a fake prod and a fake OAuth session. `node --test tools/mcp`.
3import { test } from 'node:test';
4import assert from 'node:assert/strict';
5import { createProxy, LOCAL_TOOL } from './proxy.mjs';
6
7const DEV = 'http://dev.test/mcp';
8const PROD = 'http://prod.test/jev/mcp';
9const devTools = [{ name: 'play_code' }, { name: 'get_mcp_status' }];
10const prodTools = [{ name: 'play_code' }, { name: 'list_songs' }, { name: 'publish_song' }];
11
12function world({ waitMs = 1000 } = {}) {
13  const sent = [];
14  const requests = [];
15  const prod = { valid: new Set(), status: null, headers: {} };
16  const session = {
17    token: undefined,
18    refreshable: false,
19    signIns: 0,
20    finish: null,
21    accessToken() {
22      return this.token;
23    },
24    async refresh() {
25      if (!this.refreshable) {
26        this.token = undefined;
27        return false;
28      }
29      this.token = `t${Math.random()}`;
30      prod.valid.add(this.token);
31      return true;
32    },
33    // a browser sign-in that completes when the test says so
34    async signIn() {
35      if (this.flow) return this.flow;
36      this.signIns++;
37      let resolve, reject;
38      const done = new Promise((a, b) => ((resolve = a), (reject = b)));
39      this.finish = () => {
40        this.token = 'fresh';
41        prod.valid.add('fresh');
42        this.flow = null;
43        resolve();
44      };
45      this.fail = (e) => {
46        this.flow = null;
47        reject(e);
48      };
49      done.catch(() => {});
50      this.flow = { url: 'http://prod.test/jev/oauth/authorize?x=1&y=2', done, opened: 'opened with a fake' };
51      return this.flow;
52    },
53  };
54  const fetchFn = async (url, init) => {
55    const message = JSON.parse(init.body);
56    requests.push({ url, auth: init.headers.Authorization, message });
57    const answer = (result) => Response.json({ jsonrpc: '2.0', id: message.id, result });
58    if (url === DEV) {
59      if (message.method === 'tools/list') return answer({ tools: devTools });
60      return answer({ content: [{ type: 'text', text: `dev ran ${message.params.name}` }] });
61    }
62    const token = init.headers.Authorization?.replace(/^Bearer /, '');
63    if (!prod.valid.has(token)) return new Response('invalid_token', { status: 401 });
64    if (prod.status) return new Response('too many MCP requests; try again in a minute', { status: prod.status, headers: prod.headers });
65    if (message.method === 'initialize') return answer({ instructions: 'Open the site signed in.' });
66    if (message.method === 'tools/list') return answer({ tools: prodTools });
67    return answer({ content: [{ type: 'text', text: `prod ran ${message.params.name}` }] });
68  };
69  const proxy = createProxy({
70    dev: { endpoint: DEV, fallbackTools: () => [{ name: 'offline' }] },
71    prod: { endpoint: PROD, session },
72    send: (m) => sent.push(m),
73    fetchFn,
74    signInWaitMs: waitMs,
75  });
76  let n = 0;
77  const call = async (name, args = {}) => {
78    const r = await proxy.handle({ jsonrpc: '2.0', id: ++n, method: 'tools/call', params: { name, arguments: args } });
79    return { text: r.result?.content?.[0]?.text ?? '', isError: r.result?.isError ?? false, raw: r };
80  };
81  const list = async () => (await proxy.handle({ jsonrpc: '2.0', id: ++n, method: 'tools/list' })).result.tools.map((t) => t.name);
82  const changed = () => sent.filter((m) => m.method === 'notifications/tools/list_changed').length;
83  return { proxy, session, prod, requests, call, list, changed };
84}
85
86test('dev is the default, with the local tool added to its list', async () => {
87  const w = world();
88  assert.equal(w.proxy.environment(), 'dev');
89  assert.deepEqual(await w.list(), ['play_code', 'get_mcp_status', LOCAL_TOOL.name]);
90  const r = await w.call('play_code', { code: 's("bd")' });
91  assert.equal(r.text, 'dev ran play_code');
92  assert.equal(w.requests.at(-1).url, DEV);
93  assert.equal(w.requests.at(-1).auth, undefined);
94  const status = await w.call('use_environment');
95  assert.match(status.text, /Using dev \(http:\/\/dev\.test\/mcp\)\. prod: not signed in/);
96});
97
98test('switching to prod signs in, then switches, notifies, and forwards with the token', async () => {
99  const w = world();
100  const pending = w.call('use_environment', { environment: 'prod' });
101  await new Promise((r) => setTimeout(r, 20));
102  assert.equal(w.proxy.environment(), 'dev', 'no switch before the sign-in completes');
103  w.session.finish();
104  const r = await pending;
105  assert.equal(r.isError, false);
106  assert.match(r.text, /Switched to prod \(http:\/\/prod\.test\/jev\/mcp\), signed in/);
107  assert.match(r.text, /Open the site signed in\./, "prod's instructions come along");
108  assert.equal(w.proxy.environment(), 'prod');
109  assert.equal(w.changed(), 1);
110  assert.deepEqual(await w.list(), ['play_code', 'list_songs', 'publish_song', LOCAL_TOOL.name]);
111  const played = await w.call('list_songs');
112  assert.equal(played.text, 'prod ran list_songs');
113  assert.equal(w.requests.at(-1).url, PROD);
114  assert.equal(w.requests.at(-1).auth, 'Bearer fresh');
115
116  // and back, keeping the sign-in
117  const back = await w.call('use_environment', { environment: 'dev' });
118  assert.match(back.text, /Switched to dev/);
119  assert.equal(w.changed(), 2);
120  assert.equal((await w.call('play_code', { code: 'x' })).text, 'dev ran play_code');
121  const again = await w.call('use_environment', { environment: 'prod' });
122  assert.match(again.text, /Switched to prod/);
123  assert.equal(w.session.signIns, 1, 'no second sign-in');
124  assert.equal(w.changed(), 3);
125});
126
127test('a sign-in not finished in time answers with the URL, and switches when it completes', async () => {
128  const w = world({ waitMs: 30 });
129  const r = await w.call('use_environment', { environment: 'prod' });
130  assert.equal(r.isError, false);
131  assert.match(r.text, /Still waiting/);
132  assert.ok(r.text.includes('http://prod.test/jev/oauth/authorize?x=1&y=2'), 'the whole URL, to open by hand');
133  assert.equal(w.proxy.environment(), 'dev');
134  w.session.finish();
135  await new Promise((resolve) => setTimeout(resolve, 10));
136  assert.equal(w.proxy.environment(), 'prod');
137  assert.equal(w.changed(), 1);
138});
139
140test('a failed sign-in is a tool error and stays in dev', async () => {
141  const w = world();
142  const pending = w.call('use_environment', { environment: 'prod' });
143  await new Promise((r) => setTimeout(r, 10));
144  w.session.fail(new Error('access_denied'));
145  const r = await pending;
146  assert.equal(r.isError, true);
147  assert.match(r.text, /access_denied/);
148  assert.equal(w.proxy.environment(), 'dev');
149});
150
151test('a 401 refreshes silently and retries', async () => {
152  const w = world();
153  w.session.token = 'old';
154  w.prod.valid.add('old');
155  await w.call('use_environment', { environment: 'prod' });
156  w.prod.valid.delete('old'); // expired
157  w.session.refreshable = true;
158  const r = await w.call('list_songs');
159  assert.equal(r.text, 'prod ran list_songs');
160  assert.equal(w.session.signIns, 0);
161});
162
163test('a 401 with no working refresh starts a new sign-in and says where', async () => {
164  const w = world();
165  w.session.token = 'old';
166  w.prod.valid.add('old');
167  await w.call('use_environment', { environment: 'prod' });
168  w.prod.valid.delete('old'); // the app was disconnected
169  const r = await w.call('publish_song', { title: 't' });
170  assert.equal(r.isError, true);
171  assert.match(r.text, /needs a new sign-in/);
172  assert.ok(r.text.includes('http://prod.test/jev/oauth/authorize?x=1&y=2'));
173  assert.equal(w.session.signIns, 1);
174  w.session.finish();
175  await new Promise((resolve) => setTimeout(resolve, 10));
176  assert.equal((await w.call('publish_song', { title: 't' })).text, 'prod ran publish_song');
177});
178
179test('rate limits and other refusals come back as clear tool errors', async () => {
180  const w = world();
181  w.session.token = 'ok';
182  w.prod.valid.add('ok');
183  await w.call('use_environment', { environment: 'prod' });
184  w.prod.status = 429;
185  w.prod.headers = { 'Retry-After': '60' };
186  const r = await w.call('play_code', { code: 'x' });
187  assert.equal(r.isError, true);
188  assert.equal(r.text, 'jevstrudel (prod) is limiting requests: too many MCP requests; try again in a minute; retry in 60 s');
189  w.prod.status = 403;
190  w.prod.headers = { 'WWW-Authenticate': 'Bearer error="insufficient_scope", scope="publish"' };
191  const s = await w.call('publish_song');
192  assert.equal(s.isError, true);
193  assert.match(s.text, /did not grant the publish scope/);
194});
195
196test('the poll asks dev, never prod', async () => {
197  const w = world();
198  await w.list();
199  await w.proxy.poll();
200  assert.equal(w.changed(), 0);
201  w.session.token = 'ok';
202  w.prod.valid.add('ok');
203  await w.call('use_environment', { environment: 'prod' });
204  const before = w.requests.length;
205  await w.proxy.poll();
206  assert.equal(w.requests.length, before, 'no request to prod from the poll');
207});
208
209test('an unknown environment is refused', async () => {
210  const w = world();
211  const r = await w.call('use_environment', { environment: 'staging' });
212  assert.equal(r.isError, true);
213});