diff --git a/__tests__/unit/in-process-tool-rules-provider.test.ts b/__tests__/unit/in-process-tool-rules-provider.test.ts new file mode 100644 index 0000000..02843df --- /dev/null +++ b/__tests__/unit/in-process-tool-rules-provider.test.ts @@ -0,0 +1,113 @@ +/** + * Verifies that `toolRulesProvider` runs for the in-process session path, + * not only for HTTP requests. The provider reads rules from ctx.headers, + * which the session populates from per-call `headers` options. + * + * Before this change, callers in the same process had to supply + * `body.toolRules` (via client.execute(code, { toolRules })) — the + * provider registered on createServer was dormant for in-process calls. + */ + +import { createServer, type AgentToolProtocolServer } from '../../packages/server/src/index'; +import { AgentToolProtocolClient } from '../../packages/client/src/index'; + +describe('toolRulesProvider runs for the in-process session path', () => { + let server: AgentToolProtocolServer; + let client: AgentToolProtocolClient; + + beforeAll(async () => { + server = createServer({ + // Provider reads an opaque header and returns per-call rules. + toolRulesProvider: (ctx) => { + const raw = ctx.headers['x-test-tool-rules']; + if (!raw) return undefined; + try { + return JSON.parse(raw); + } catch { + return undefined; + } + }, + }); + + server.use({ + name: 'math', + type: 'custom', + functions: [ + { + name: 'add', + description: 'Add two numbers', + inputSchema: { + type: 'object', + properties: { a: { type: 'number' }, b: { type: 'number' } }, + required: ['a', 'b'], + }, + handler: async (input: unknown) => { + const { a, b } = input as { a: number; b: number }; + return { result: a + b }; + }, + }, + ], + }); + server.use({ + name: 'text', + type: 'custom', + functions: [ + { + name: 'uppercase', + description: 'Uppercase text', + inputSchema: { + type: 'object', + properties: { text: { type: 'string' } }, + required: ['text'], + }, + handler: async (input: unknown) => { + const { text } = input as { text: string }; + return { result: text.toUpperCase() }; + }, + }, + ], + }); + + await server.start(); + client = new AgentToolProtocolClient({ server: server as any }); + await client.init({ name: 'test-client', version: '1.0.0' }); + await client.connect(); + }); + + describe('handleExplore', () => { + test('provider applies rules from per-call headers', async () => { + // With math-only rules coming from the provider's header read, + // the text group should disappear from the explore tree. + const rules = JSON.stringify({ allowOnlyApiGroups: ['math'] }); + const r = (await client.exploreAPI('/custom', { + headers: { 'x-test-tool-rules': rules }, + } as any)) as { items: Array<{ name: string }> }; + + const groupNames = r.items.map((i) => i.name); + expect(groupNames).toContain('math'); + expect(groupNames).not.toContain('text'); + }); + + test('body.toolRules takes precedence over provider', async () => { + // Provider says math-only, but caller's explicit body.toolRules says + // text-only. Explicit wins. + const providerRules = JSON.stringify({ allowOnlyApiGroups: ['math'] }); + const r = (await client.exploreAPI('/custom', { + headers: { 'x-test-tool-rules': providerRules }, + toolRules: { allowOnlyApiGroups: ['text'] }, + } as any)) as { items: Array<{ name: string }> }; + + const groupNames = r.items.map((i) => i.name); + expect(groupNames).toContain('text'); + expect(groupNames).not.toContain('math'); + }); + + test('no provider header and no body rules: unrestricted', async () => { + const r = (await client.exploreAPI('/custom')) as { + items: Array<{ name: string }>; + }; + const groupNames = r.items.map((i) => i.name); + expect(groupNames).toEqual(expect.arrayContaining(['math', 'text'])); + }); + }); +}); diff --git a/packages/atp-compiler/__tests__/unit/api-call-analyzer.test.ts b/packages/atp-compiler/__tests__/unit/api-call-analyzer.test.ts new file mode 100644 index 0000000..2f07c2f --- /dev/null +++ b/packages/atp-compiler/__tests__/unit/api-call-analyzer.test.ts @@ -0,0 +1,387 @@ +import { describe, test, expect } from '@jest/globals'; +import { analyzeApiCalls } from '../../src/api-call-analyzer'; + +describe('analyzeApiCalls', () => { + test('extracts a single api..(...) call', () => { + const r = analyzeApiCalls(`return api.calendar.events_list({ calendarId: 'primary' });`); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('extracts nested await + multiple calls in the same group', () => { + const code = ` + const c = await api.calendar.calendars_get({ calendarId: 'primary' }); + return await api.calendar.events_list({ calendarId: c.id, maxResults: 10 }); + `; + const r = analyzeApiCalls(code); + const sorted = r.apiCalls.slice().sort((a, b) => a.operationId.localeCompare(b.operationId)); + expect(sorted).toEqual([ + { apiGroup: 'calendar', operationId: 'calendars_get' }, + { apiGroup: 'calendar', operationId: 'events_list' }, + ]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('extracts cross-group calls', () => { + const r = analyzeApiCalls(` + await api.calendar.events_list({}); + await api.gmail.messages_list({}); + `); + expect(r.apiCalls).toEqual( + expect.arrayContaining([ + { apiGroup: 'calendar', operationId: 'events_list' }, + { apiGroup: 'gmail', operationId: 'messages_list' }, + ]) + ); + }); + + test('deduplicates repeated calls to the same api..', () => { + const r = analyzeApiCalls(` + await api.calendar.events_list({ maxResults: 1 }); + await api.calendar.events_list({ maxResults: 2 }); + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('flags dynamic dispatch via computed operation member', () => { + const r = analyzeApiCalls(`return api.calendar[fn]({});`); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags dynamic dispatch via computed group member', () => { + const r = analyzeApiCalls(`return api[g].events_list({});`); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags destructured api: const { calendar } = api', () => { + const r = analyzeApiCalls(` + const { calendar } = api; + return calendar.events_list({}); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags destructure-with-rename: const { calendar: c } = api', () => { + const r = analyzeApiCalls(` + const { calendar: c } = api; + return c.events_list({}); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags aliasing: const x = api.calendar', () => { + const r = analyzeApiCalls(` + const x = api.calendar; + return x.events_list({}); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags aliasing: const x = api', () => { + const r = analyzeApiCalls(` + const x = api; + return x.calendar.events_list({}); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('returns empty calls + no dynamic flag for trivial code', () => { + const r = analyzeApiCalls(`return 42;`); + expect(r.apiCalls).toEqual([]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('fails closed with dynamicCallsDetected=true on syntax error', () => { + const r = analyzeApiCalls(`return api.calendar.events_list(`); + expect(r.dynamicCallsDetected).toBe(true); + expect(r.apiCalls).toEqual([]); + }); + + test('handles empty string input', () => { + const r = analyzeApiCalls(''); + expect(r.apiCalls).toEqual([]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Dedup — same (group, op) through different syntactic paths + // ──────────────────────────────────────────────────────────────────────── + + describe('dedup', () => { + test('same op called in both branches of a ternary', () => { + const r = analyzeApiCalls(` + return flag + ? api.calendar.events_list({ a: 1 }) + : api.calendar.events_list({ a: 2 }); + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('same op called inside a loop is reported once', () => { + const r = analyzeApiCalls(` + for (const id of ids) { + await api.calendar.events_list({ calendarId: id }); + } + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('same op across try/catch/finally reports once', () => { + const r = analyzeApiCalls(` + try { + await api.calendar.events_list({}); + } catch (e) { + await api.calendar.events_list({ retry: true }); + } finally { + api.calendar.events_list({ cleanup: true }); + } + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('same op across multiple statements is reported once', () => { + const r = analyzeApiCalls(` + const a = await api.calendar.events_list({ maxResults: 1 }); + const b = await api.calendar.events_list({ maxResults: 2 }); + const c = await api.calendar.events_list({ maxResults: 3 }); + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Complex control-flow / realistic code shapes + // ──────────────────────────────────────────────────────────────────────── + + describe('complex control flow', () => { + test('IIFE (immediately-invoked async function)', () => { + const r = analyzeApiCalls(` + return (async () => await api.calendar.events_list({}))(); + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('Promise.all with multiple awaits', () => { + const r = analyzeApiCalls(` + return await Promise.all([ + api.calendar.events_list({}), + api.gmail.messages_list({}), + api.drive.files_list({}), + ]); + `); + const sorted = r.apiCalls + .slice() + .sort((a, b) => (a.apiGroup + '.' + a.operationId).localeCompare(b.apiGroup + '.' + b.operationId)); + expect(sorted).toEqual([ + { apiGroup: 'calendar', operationId: 'events_list' }, + { apiGroup: 'drive', operationId: 'files_list' }, + { apiGroup: 'gmail', operationId: 'messages_list' }, + ]); + }); + + test('arrow inside .then() chain', () => { + const r = analyzeApiCalls(` + return api.calendar.events_list({}) + .then((list) => api.calendar.events_get({ eventId: list.items[0].id })) + .then((ev) => api.calendar.calendars_get({ calendarId: ev.calendarId })); + `); + expect(r.apiCalls.sort((a, b) => a.operationId.localeCompare(b.operationId))).toEqual([ + { apiGroup: 'calendar', operationId: 'calendars_get' }, + { apiGroup: 'calendar', operationId: 'events_get' }, + { apiGroup: 'calendar', operationId: 'events_list' }, + ]); + }); + + test('class method body', () => { + const r = analyzeApiCalls(` + class Scheduler { + async listEvents() { + return await api.calendar.events_list({ calendarId: 'primary' }); + } + } + return new Scheduler().listEvents(); + `); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('nested try/catch/finally across multiple groups', () => { + const r = analyzeApiCalls(` + try { + await api.calendar.events_list({}); + } catch (e) { + await api.gmail.messages_list({ label: 'errors' }); + } finally { + api.drive.files_list({}); + } + `); + expect(r.apiCalls.length).toBe(3); + }); + + test('switch statement with api calls per case', () => { + const r = analyzeApiCalls(` + switch (kind) { + case 'a': return api.calendar.events_list({}); + case 'b': return api.calendar.events_get({ eventId: x }); + default: return api.calendar.calendars_get({ calendarId: 'primary' }); + } + `); + expect(r.apiCalls.length).toBe(3); + }); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Multi-group realism — many sources + many ops + // ──────────────────────────────────────────────────────────────────────── + + describe('multi-group', () => { + test('realistic Google Suite workload (6 groups, 11 ops)', () => { + const r = analyzeApiCalls(` + const events = await api.calendar.events_list({ calendarId: 'primary' }); + const cal = await api.calendar.calendars_get({ calendarId: 'primary' }); + const messages = await api.gmail.messages_list({}); + const thread = await api.gmail.threads_get({ id: '1' }); + const sheet = await api.sheets.spreadsheets_get({ spreadsheetId: 'x' }); + const values = await api.sheets.spreadsheets_values_get({ spreadsheetId: 'x', range: 'A1' }); + const drive = await api.drive.files_list({}); + const file = await api.drive.files_get({ fileId: 'x' }); + const deck = await api.slides.presentations_get({ presentationId: 'x' }); + const doc = await api.docs.documents_get({ documentId: 'x' }); + const batch = await api.sheets.spreadsheets_batchUpdate({ spreadsheetId: 'x' }); + return { events, cal, messages, thread, sheet, values, drive, file, deck, doc, batch }; + `); + + const groupsTouched = new Set(r.apiCalls.map((c) => c.apiGroup)); + expect(groupsTouched).toEqual(new Set(['calendar', 'gmail', 'sheets', 'drive', 'slides', 'docs'])); + expect(r.apiCalls).toHaveLength(11); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('mixed groups with some dedup', () => { + const r = analyzeApiCalls(` + await api.calendar.events_list({}); + await api.gmail.messages_list({}); + await api.calendar.events_list({}); // dup + await api.gmail.messages_list({}); // dup + await api.calendar.events_get({ eventId: 'x' }); + `); + expect(r.apiCalls).toHaveLength(3); + expect(r.apiCalls.map((c) => c.apiGroup + '.' + c.operationId).sort()).toEqual([ + 'calendar.events_get', + 'calendar.events_list', + 'gmail.messages_list', + ]); + }); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Additional dynamic-dispatch escape patterns + // ──────────────────────────────────────────────────────────────────────── + + describe('dynamic dispatch — additional escapes', () => { + test('flags optional-chained operation member', () => { + const r = analyzeApiCalls(`return api.calendar?.events_list?.({});`); + // Detected statically OR dynamic — either outcome denies a grant + // that doesn't cover the target. This implementation DETECTS + // because the static path is resolvable; no dynamic flag needed. + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('flags .call() redirection as a regular static call', () => { + const r = analyzeApiCalls(`api.calendar.events_list.call(null, {});`); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('flags .apply() redirection', () => { + const r = analyzeApiCalls(`api.calendar.events_list.apply(null, [{}]);`); + expect(r.apiCalls).toEqual([{ apiGroup: 'calendar', operationId: 'events_list' }]); + }); + + test('flags .bind() + later call', () => { + const r = analyzeApiCalls(` + const boundList = api.calendar.events_list.bind(null); + return boundList({}); + `); + // .bind still surfaces the underlying call target. + expect(r.apiCalls).toContainEqual({ apiGroup: 'calendar', operationId: 'events_list' }); + }); + + test('flags Object.values(api) as dynamic', () => { + const r = analyzeApiCalls(` + const groups = Object.values(api); + return groups.map((g) => Object.keys(g)); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags Object.keys(api) as dynamic', () => { + const r = analyzeApiCalls(`return Object.keys(api);`); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags spread {...api} as dynamic', () => { + const r = analyzeApiCalls(`const x = { ...api }; return x;`); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags passing api as function argument', () => { + const r = analyzeApiCalls(` + const use = (a) => a.calendar.events_list({}); + return use(api); + `); + expect(r.dynamicCallsDetected).toBe(true); + }); + + test('flags returning bare api', () => { + const r = analyzeApiCalls(`return api;`); + expect(r.dynamicCallsDetected).toBe(true); + }); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Not-false-positive — identifiers named "api" that aren't the global + // ──────────────────────────────────────────────────────────────────────── + + describe('does not false-positive on unrelated "api" identifiers', () => { + test('this.api.foo.bar(...) is ignored (not the global api)', () => { + const r = analyzeApiCalls(`return this.api.foo.bar({});`); + expect(r.apiCalls).toEqual([]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('someObj.api.foo.bar(...) is ignored', () => { + const r = analyzeApiCalls(` + const wrapper = { api: null }; + return wrapper.api?.foo?.bar?.({}); + `); + expect(r.apiCalls).toEqual([]); + expect(r.dynamicCallsDetected).toBe(false); + }); + + test('deep api.x.y.z(...) — not valid ATP syntax — silently ignored', () => { + // Sandbox's runtime namespace would throw on this; static analyzer + // ignores it so we don't over-flag invalid code as dynamic. + const r = analyzeApiCalls(`return api.x.y.z({});`); + expect(r.apiCalls).toEqual([]); + expect(r.dynamicCallsDetected).toBe(false); + }); + }); + + // ──────────────────────────────────────────────────────────────────────── + // Idempotency — deterministic output across calls + // ──────────────────────────────────────────────────────────────────────── + + describe('idempotency', () => { + test('calling analyzeApiCalls twice on same code yields identical output', () => { + const code = ` + await api.calendar.events_list({}); + await api.gmail.messages_list({}); + const { calendar } = api; + `; + const a = analyzeApiCalls(code); + const b = analyzeApiCalls(code); + expect(a).toEqual(b); + }); + }); +}); diff --git a/packages/atp-compiler/src/api-call-analyzer.ts b/packages/atp-compiler/src/api-call-analyzer.ts new file mode 100644 index 0000000..d97d1f1 --- /dev/null +++ b/packages/atp-compiler/src/api-call-analyzer.ts @@ -0,0 +1,232 @@ +/** + * Pre-dispatch static analysis of ATP agent code. + * + * Extracts every `api..(...)` call chain from submitted code + * and flags patterns that defeat static analysis (dynamic dispatch, + * aliasing, destructuring). + * + * Intended caller: governance layers that need to know WHICH api groups + * and operations code will touch BEFORE dispatching it to the sandbox. + * Paired with runtime `filterApiGroups` enforcement in atp-server for + * defense-in-depth; the static pass catches unauthorized references + * up-front without paying sandbox startup cost. + * + * @example + * const { apiCalls, dynamicCallsDetected } = analyzeApiCalls(code); + * for (const call of apiCalls) { + * // check (call.apiGroup, call.operationId) against a grant + * } + * if (dynamicCallsDetected) { + * // deny unless the grant explicitly allows dynamic dispatch + * } + */ + +import { parse } from '@babel/parser'; +// @babel/traverse exports default; interop handles the CJS/ESM quirk +// (see https://github.com/babel/babel/issues/13855). +import _traverse from '@babel/traverse'; +import * as t from '@babel/types'; + +const traverse = ((_traverse as unknown as { default?: typeof _traverse }).default ?? + _traverse) as typeof _traverse; + +export interface DetectedApiCall { + apiGroup: string; + operationId: string; +} + +export interface AnalysisResult { + /** Unique `(apiGroup, operationId)` pairs statically visible in the code. */ + apiCalls: DetectedApiCall[]; + /** + * True iff the code contains patterns we cannot statically resolve to a + * concrete `(apiGroup, operationId)` — e.g. `api[varName].fn(...)`, + * destructuring (`const { calendar } = api`), or aliasing + * (`const x = api.calendar`). Governance layers should fail-closed on + * this flag unless the caller's policy opts into dynamic dispatch. + */ + dynamicCallsDetected: boolean; +} + +/** + * Analyze agent code and return its statically-visible api.* call set plus a + * dynamic-dispatch flag. Pure function, no I/O, fail-closed on parse errors + * (returns `{ apiCalls: [], dynamicCallsDetected: true }`). + */ +export function analyzeApiCalls(code: string): AnalysisResult { + const calls: DetectedApiCall[] = []; + const seen = new Set(); + let dynamicCallsDetected = false; + + let ast; + try { + ast = parse(code, { + sourceType: 'module', + allowReturnOutsideFunction: true, + plugins: ['typescript'], + }); + } catch { + // Fail-closed: syntax error → treat as dynamic so governance denies. + return { apiCalls: [], dynamicCallsDetected: true }; + } + + // Helper records a static api..(...) call, or flips + // dynamicCallsDetected when the call expression escapes static resolution. + const tryRecordCall = (calleeNode: t.Node) => { + let callee: t.Node = calleeNode; + + // Unwrap one `.call` / `.apply` / `.bind` redirection: + // api.calendar.events_list.call(null, {...}) + // has callee = MemberExpression { object: api.calendar.events_list, property: 'call' } + if ( + (t.isMemberExpression(callee) || t.isOptionalMemberExpression(callee)) && + !callee.computed && + t.isIdentifier(callee.property) && + (callee.property.name === 'call' || + callee.property.name === 'apply' || + callee.property.name === 'bind') + ) { + callee = callee.object; + } + + if (!t.isMemberExpression(callee) && !t.isOptionalMemberExpression(callee)) return; + + const groupExpr = callee.object; + if (!t.isMemberExpression(groupExpr) && !t.isOptionalMemberExpression(groupExpr)) return; + if (!t.isIdentifier(groupExpr.object, { name: 'api' })) return; + + // api[groupVar].op(...) or api.group[fnVar](...) → dynamic + if (groupExpr.computed || callee.computed) { + dynamicCallsDetected = true; + return; + } + + const groupNode = groupExpr.property; + const opNode = callee.property; + if (!t.isIdentifier(groupNode) || !t.isIdentifier(opNode)) { + dynamicCallsDetected = true; + return; + } + + const key = `${groupNode.name}.${opNode.name}`; + if (!seen.has(key)) { + seen.add(key); + calls.push({ apiGroup: groupNode.name, operationId: opNode.name }); + } + }; + + try { + traverse(ast, { + // Assignments / destructures that alias `api` or `api.` — any of + // these lets the code reach an api group via an opaque identifier later. + VariableDeclarator(path) { + const init = path.node.init; + if (!init) return; + + // const { calendar } = api + if (t.isObjectPattern(path.node.id) && t.isIdentifier(init, { name: 'api' })) { + dynamicCallsDetected = true; + return; + } + // const x = api + if (t.isIdentifier(path.node.id) && t.isIdentifier(init, { name: 'api' })) { + dynamicCallsDetected = true; + return; + } + // const x = api. OR const x = api[''] + if ( + t.isIdentifier(path.node.id) && + t.isMemberExpression(init) && + t.isIdentifier(init.object, { name: 'api' }) + ) { + dynamicCallsDetected = true; + } + }, + + // Any other mention of `api` that hands it off to an opaque consumer: + // fn(api) — alias escape via function argument + // Object.values(api) / keys(…) — enumeration + // { ...api } / [ ...api ] — spread + // return api — caller gets the alias + // api = x (reassignment) — later reads hit a different object + // + // The api..(...) pattern is recognised by the CallExpression + // visitor below; skip it here via parent-shape whitelisting. + Identifier(path) { + if (path.node.name !== 'api') return; + + // Skip the PROPERTY position of a member expression — e.g. + // `this.api`, `window.api`, `someObj.api`. That's not the + // global `api` binding we care about. + if ( + (t.isMemberExpression(path.parent) || t.isOptionalMemberExpression(path.parent)) && + path.parent.property === path.node && + !path.parent.computed + ) { + return; + } + // `api.` (non-computed member) — safe, handled by CallExpression. + if ( + (t.isMemberExpression(path.parent) || t.isOptionalMemberExpression(path.parent)) && + path.parent.object === path.node && + !path.parent.computed + ) { + return; + } + // Left-hand side of `const x = api` / `const { calendar } = api` + // — handled by VariableDeclarator above (dynamic flag set there). + if (t.isVariableDeclarator(path.parent) && path.parent.init === path.node) { + return; + } + // Skip declaration positions where `api` is a local binding name, + // not a reference: + // { api: value } — object property key + // function f(api) { ... } — param name + // class { api() {...} } — method name + // function api() {} — function name + // class api {} — class name + if ( + (t.isObjectProperty(path.parent) || t.isObjectMethod(path.parent)) && + path.parent.key === path.node && + !path.parent.computed + ) { + return; + } + if (t.isClassMethod(path.parent) && path.parent.key === path.node && !path.parent.computed) { + return; + } + if ( + (t.isFunctionDeclaration(path.parent) || + t.isFunctionExpression(path.parent) || + t.isClassDeclaration(path.parent) || + t.isClassExpression(path.parent)) && + path.parent.id === path.node + ) { + return; + } + if (path.parentPath?.isFunction() && path.listKey === 'params') { + return; + } + if (t.isImportSpecifier(path.parent) || t.isImportDefaultSpecifier(path.parent) || t.isImportNamespaceSpecifier(path.parent)) { + return; + } + + // Anything else (`Object.values(api)`, `fn(api)`, `{ ...api }`, + // `return api`, `api = ...`) escapes static resolution. + dynamicCallsDetected = true; + }, + + CallExpression(path) { + tryRecordCall(path.node.callee); + }, + OptionalCallExpression(path) { + tryRecordCall(path.node.callee); + }, + }); + } catch { + // Visitor error → fail-closed (should be unreachable under our plugin set). + return { apiCalls: [], dynamicCallsDetected: true }; + } + + return { apiCalls: calls, dynamicCallsDetected }; +} diff --git a/packages/atp-compiler/src/index.ts b/packages/atp-compiler/src/index.ts index 0820fe2..47b89b9 100644 --- a/packages/atp-compiler/src/index.ts +++ b/packages/atp-compiler/src/index.ts @@ -12,6 +12,10 @@ export * from './plugin-system/index.js'; // Checkpoint exports export * from './checkpoint/index.js'; +// Pre-dispatch static analysis of agent code (api..(...) extraction). +export { analyzeApiCalls } from './api-call-analyzer.js'; +export type { DetectedApiCall, AnalysisResult } from './api-call-analyzer.js'; + // Main exports export { ATPCompiler } from './transformer/index.js'; export { initializeRuntime, cleanupRuntime } from './runtime/index.js'; diff --git a/packages/client/src/core/in-process-session.ts b/packages/client/src/core/in-process-session.ts index 71079a0..912192f 100644 --- a/packages/client/src/core/in-process-session.ts +++ b/packages/client/src/core/in-process-session.ts @@ -212,10 +212,19 @@ export class InProcessSession extends BaseSession { async explore(path: string, options?: Record): Promise { await this.ensureInitialized(); + // Per-call `headers` pulled out of options so it can merge into ctx.headers + // (where the server's toolRulesProvider reads from). Stripped from the body + // to avoid duplicating it alongside `path` + `toolRules`. + const { headers: callerHeaders, ...body } = (options ?? {}) as { + headers?: Record; + [k: string]: unknown; + }; + const ctx = await this.createContext({ method: 'POST', path: '/api/explore', - body: { path, ...options }, + body: { path, ...body }, + headers: callerHeaders, }); return await this.server.handleExplore(ctx); @@ -224,10 +233,17 @@ export class InProcessSession extends BaseSession { async execute(code: string, config?: Record): Promise { await this.ensureInitialized(); + // `requestContext.headers` is the documented entry point for per-call + // header forwarding (consumed by openapi-loader headerProvider). Also + // merge those into ctx.headers so the server-level toolRulesProvider + // sees the same values. + const rc = (config?.requestContext ?? {}) as { headers?: Record }; + const ctx = await this.createContext({ method: 'POST', path: '/api/execute', body: { code, config }, + headers: rc.headers, }); return await this.server.handleExecute(ctx); @@ -268,6 +284,7 @@ export class InProcessSession extends BaseSession { path: string; query?: Record; body?: unknown; + headers?: Record; }): Promise { const noopLogger = { debug: () => {}, @@ -276,11 +293,21 @@ export class InProcessSession extends BaseSession { error: () => {}, }; + // Merge per-call headers (if any) on top of session-level headers. + // Session auth keys (`x-client-id`, `authorization`) still survive + // because `prepareHeaders` produces them first and the caller would + // rarely override them; when they do, the caller's choice wins here — + // match that with the RequestContext.headers convention. + const sessionHeaders = await this.prepareHeaders(options.method, options.path, options.body); + const mergedHeaders = options.headers + ? { ...sessionHeaders, ...options.headers } + : sessionHeaders; + return { method: options.method, path: options.path, query: options.query || {}, - headers: await this.prepareHeaders(options.method, options.path, options.body), + headers: mergedHeaders, body: options.body, clientId: this.clientId, clientToken: this.clientToken, diff --git a/packages/server/src/create-server.ts b/packages/server/src/create-server.ts index 6dd103b..3d16ebb 100644 --- a/packages/server/src/create-server.ts +++ b/packages/server/src/create-server.ts @@ -580,7 +580,7 @@ export class AgentToolProtocolServer { async handleExplore(ctx: RequestContext): Promise { if (!this.explorerService) ctx.throw(503, 'Explorer not initialized'); - return await handleExplore(ctx, this.explorerService); + return await handleExplore(ctx, this.explorerService, this.toolRulesProvider); } async handleExecute(ctx: RequestContext): Promise { @@ -593,7 +593,8 @@ export class AgentToolProtocolServer { this.stateManager, this.config, this.auditSink, - this.sessionManager + this.sessionManager, + this.toolRulesProvider ); } diff --git a/packages/server/src/handlers/execute.handler.ts b/packages/server/src/handlers/execute.handler.ts index bd0d097..dde7575 100644 --- a/packages/server/src/handlers/execute.handler.ts +++ b/packages/server/src/handlers/execute.handler.ts @@ -1,4 +1,4 @@ -import type { RequestContext, ResolvedServerConfig } from '../core/config.js'; +import type { RequestContext, ResolvedServerConfig, ToolRulesProvider } from '../core/config.js'; import type { SandboxExecutor } from '../executor/index.js'; import type { ExecutionStateManager } from '../execution-state/index.js'; import type { ClientSessionManager } from '../client-sessions.js'; @@ -45,7 +45,8 @@ export async function handleExecute( stateManager: ExecutionStateManager, config: ResolvedServerConfig, auditSink?: AuditSink, - sessionManager?: ClientSessionManager + sessionManager?: ClientSessionManager, + toolRulesProvider?: ToolRulesProvider ): Promise { const request = ctx.body as any; const code = request.code || ''; @@ -134,7 +135,10 @@ export async function handleExecute( }, onToolCall, eventCallback: requestConfig.eventCallback, - toolRules: requestConfig.toolRules, + // Rule source precedence: explicit requestConfig.toolRules first, then + // server-level provider (e.g. reads a header). Lets in-process callers + // and HTTP callers converge on the same provider mechanism. + toolRules: requestConfig.toolRules ?? toolRulesProvider?.(ctx), }; // Verify provenance hints if provided diff --git a/packages/server/src/handlers/explorer.handler.ts b/packages/server/src/handlers/explorer.handler.ts index 3ee0512..78881c2 100644 --- a/packages/server/src/handlers/explorer.handler.ts +++ b/packages/server/src/handlers/explorer.handler.ts @@ -1,15 +1,23 @@ import type { RequestContext } from '../core/config.js'; +import type { ToolRulesProvider } from '../core/config.js'; import type { ExplorerService } from '../explorer/index.js'; import type { ApiGroupRules } from '@mondaydotcomorg/atp-protocol'; import { runInRequestScope, getRequestScope } from '../core/request-scope.js'; export async function handleExplore( ctx: RequestContext, - explorerService: ExplorerService + explorerService: ExplorerService, + toolRulesProvider?: ToolRulesProvider ): Promise { const body = ctx.body as { path?: string; toolRules?: ApiGroupRules }; const path = body.path || '/'; - const { toolRules } = body; + + // Rule source precedence (highest to lowest): + // 1. body.toolRules — explicit per-call override + // 2. toolRulesProvider(ctx) — server-level policy (e.g. read a header) + // 3. existing request scope — already wrapped by caller + const effectiveToolRules: ApiGroupRules | undefined = + body.toolRules ?? (toolRulesProvider ? toolRulesProvider(ctx) : undefined); const executeExplore = () => { const result = explorerService.explore(path); @@ -21,8 +29,8 @@ export async function handleExplore( return result; }; - if (toolRules && !getRequestScope()?.toolRules) { - return runInRequestScope({ toolRules }, executeExplore); + if (effectiveToolRules && !getRequestScope()?.toolRules) { + return runInRequestScope({ toolRules: effectiveToolRules }, executeExplore); } return executeExplore();