diff --git a/AGENTS.md b/AGENTS.md index 71038bd..6605554 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,7 +25,7 @@ - `packages/utils` - `@proxy/utils`, the Result helpers shared by every package: `Do`, `requirePresent` and `parseSchema` (`src/parse.ts`), and `httpRequest`, the fetch wrapper that returns a `FetchError` union instead of throwing (`src/fetch.ts`). Also `ApiError`/`ApiErr` and `wrapDb` (`src/api-error.ts`), the error every server-side function returns, and the session cookie's name and parsing (`src/session-cookie.ts`) - `packages/integrations` - `@proxy/integrations`, the catalog the api and the web app share: each integration's actions, collections and fields (`src/catalog.ts`), the agents a login can be made for (`src/agent-providers.ts`), the email providers' servers (`src/email-providers.ts`), and access (`src/access.ts`): per action of a connection, an agent's own setting, else the connection's default, an action counting only with the one it `requires`, and tools (`src/tools.ts`): what an agent can do with a connection, generated from the catalog, one per list, open, write and command of each collection (`emails_list`, `emails_archive`), each with the action it needs and a JSON Schema of its params. A custom MCP server (`mcp` in the catalog, `custom`) has no actions of its own: `integrationFor(connection)` (`src/mcp-tools.ts`) builds its integration from the tools the connection keeps (`McpTool`, as the server lists them), one action per tool named by the tool, its risk from the tool's annotations (low when it only reads, medium when what it changes can be undone, high when it can't, and `unknown`, shown without a label, when it says nothing), and `integrationTools` passes them on as `call` tools with the server's own input schema (`ServerTool`, beside the catalog's `CollectionTool`). The web app adds the icons (`src/lib/integrations.ts`, whose `integrationOf(connection)` is what every page uses) - `packages/db` - `@proxy/db`: Prisma 6 on Postgres. `src/index.ts` is the client; each module (`auth.ts`, `user.ts`, `organization.ts`, `connection.ts`, `info.ts`, `agent.ts`, `audit.ts`) exports functions returning `Result` through `wrapDb`. A connection's `tools` column holds a custom MCP server's tools as JSON, and `name` what the person called it; both null for every other integration. `relationMode = "prisma"`: no foreign keys in the database, so Prisma runs the cascades and every relation column has its own index. Locally the database is Prisma's embedded Postgres (PGlite, `scripts/dev-db.ts`), which needs `connection_limit=1&pgbouncer=true` in `DATABASE_URL` -- `packages/api` - `@proxy/api`: Express on Bun. Routes are `handleXRoute` functions mounted in `src/server/create-api-router.ts`; JSON routes return `Result` through `withAuthResult`, which puts the signed-in user on `req.user`. Every environment variable is in `src/utils/env-schema.ts` with its docs, which the docs' Environment variables page is built from; add new ones there. Personal Agent Proxy is one origin everywhere: when `packages/web/dist` exists, as in the Docker image, the api serves the web app's build too (`src/server/serve-web.ts`), and every address the api builds starts from `APP_URL`. Browser-called writes take `requireBrowserOrigin` (CSRF); routes that read a cookie have no CORS. `src/connections/` lists, deletes and sets defaults on an organization's connections, connects a mailbox with an app password after signing in to its IMAP and SMTP servers (`email/`), and connects Granola, Notion and Linear by signing in to their MCP servers (`mcp/`, shared by all three: OAuth with PKCE against the server's authorization server, a client registered for each sign-in so nothing is set up beforehand, the handshake in an encrypted cookie, `mcpConnectRoutes` making both routes; tokens are refreshed as they run out, one refresh at a time per connection on each server, and saved only over the ones they replace, whatever else the credential holds kept beside them; each tool call is one stateless `tools/call`, `mcp-client.ts`, which also lists a server's tools and, for a server that may keep a session (`handshake`), initializes one around each call as the spec asks). Every request to an MCP server, its metadata or its authorization server goes through `utils/checked-request.ts`, which checks the address as a database's host is checked (`utils/host-check.ts`: nothing private unless `ALLOW_PRIVATE_HOSTS`, and `https` unless the address is private) before each request, since a custom server and its endpoints are wherever it said, and never follows a redirect. A custom MCP server is connected by its address (`custom/`): `POST /api/connections/mcp` with `{name, url, token?, returnTo?}` (the name is the person's, kept through the sign-in in its cookie, since a server says nothing about itself that would do) asks the server for its tools, connects a server that answers at once, with no sign-in or the token typed in as a bearer token (as written when it has its scheme in front, `Basic …`), and for one that turns a request without a token down begins its sign-in there, under the browser's origin check, answering `{kind: 'signIn', url}` with the sign-in page the browser goes to: `mcp-discovery.ts` finds the authorization server from the server's protected resource metadata (the `WWW-Authenticate` header's `resource_metadata`, else the well-known paths), its endpoints from that server's metadata (else `/authorize`, `/token` and `/register` under it) and the scopes the server's own metadata names, and `mcpSignInRoutes` (`mcp-connect.ts`, beside `mcpConnectRoutes` for the servers known in advance) keeps the endpoints in the handshake cookie and in the credential beside the tokens. The connection's account is the address without its scheme (`mcp.example.com/mcp`), connecting the same address again replaces its credential and tools, a new connection's defaults are the tools that say they only read, `POST /api/connections/:id/tools` asks the server for its tools again (replacing the tools drops the defaults and agents' settings for tools the server no longer lists, `updateConnectionTools`), and a tool is run as it is (`custom-mcp.ts`): its params go to the server as given, a tool that fails answers the agent with the server's own text as a `validation_error`, and the run is logged under the collection `tools` by the tool's name, with its params as JSON, cut to 200 characters. Granola's notes and transcripts are on by default, Notion's Read pages and Linear's Read issues alone; a Notion or Linear sign-in is for one workspace, and a Notion one lasts at most 180 days, however often it is refreshed. `src/agents/` makes, revokes and deletes agent logins (the password is sent once, only its argon2id hash is kept) and sets an agent's own settings per action of a connection where they differ from the default. `src/records/` reads and writes a collection's records through the integration's `Connector` (`connectors.ts`): Information's live in the database (`info-connector.ts`), a mailbox's are fetched over IMAP on every request (`email-connector.ts`), which merges the inbox and the folders the server marks as Drafts and Sent into one list (ids start with the folder: `inbox-…`, `drafts-…`, `sent-…`) and runs the commands: flags, moves to the folders the server marks as Archive (Gmail's All Mail) or Trash, and sending over SMTP, a draft or a new email, filing a copy in Sent where the server doesn't (Gmail does). Received emails are never edited or deleted for good; only drafts are. Granola's notes are fetched from its MCP server on every request (`granola-connector.ts`; the reading of the XML-like text Granola answers in is in `connections/granola/granola-mcp.ts`): its Notes and Transcripts collections are the same meetings by Granola's meeting id, and since Granola's tools have no pages, a page is 30 days of meetings; one that finds none skips back to the newest meeting before it, so the list ends where the meetings do. Granola's dates and attendees are shown as Granola writes them. Notion's workspace is fetched from its MCP server the same way (`notion-connector.ts`, the reading of its answers in `connections/notion/notion-mcp.ts`), as the person who signed in, as a tree in one nested collection: Notion has no list of every page, so the top is their Recents and a search (Notion's own, AI search where the workspace has it) finds the rest, one page of 50; inside a page are the pages and databases in its content, inside a database its saved views (`view-…` ids), or the rows of its only one, and inside a view its rows, paged with Notion's cursor (view queries have no plan quota; Notion's row and SQL queries do). A page reads as Notion-flavored Markdown, a row's properties beside it as "Name: value" lines (`notion-properties.ts`), and only a page's, database's or view's id is ever sent to Notion's fetch; a new page is a private draft at the top, or goes inside the page, database or view it is made in, a row's title in the database's title property; an edit sends the title, the properties written differently and the content, each only where it changed, and pages are never deleted. Notion's MCP server has no tool that deletes or trashes a page. Linear's issues are fetched from its MCP server the same way (`linear-connector.ts`, the reading of its answers in `connections/linear/linear-mcp.ts`), as the person who signed in, by their identifier (`ENG-123`): a list is the issues updated last first, 50 to a page with Linear's cursor, narrowed by the type of status (`?statusType=Started`) and searched in titles and descriptions, and an issue opens with its comments, oldest first. Teams, statuses, people, projects and labels are written by name, as Linear shows them and finds them, its refusals passed on to the agent; an edit sends only the fields written differently, an empty one clearing it. Linear's MCP server has no tool that archives or deletes an issue. PostgreSQL and MySQL (MariaDB too) are two integrations on one code path: a database is connected with its host, port, database, user, password and encryption (`connections/sql/`: `route.ts` makes both routes per dialect, `sql-client.ts` opens one connection per request with `Bun.sql`, to the address `utils/host-check.ts` resolves the host to, refusing private addresses unless `ALLOW_PRIVATE_HOSTS` is set, and `sql-statement.ts` checks an agent's query is one statement that reads), and `records/sql-connector.ts` reads and writes it through `records/sql-dialect.ts`, where the two differ: the tables are one nested collection, a table holding its rows, a row read and written as "column: value" lines and found again by its primary key (ids `public.users:42`, `pairs:3,4`; a table without a key is never written to), and a query, a second collection that is only ever run, goes through a read-only transaction with a time limit, a cursor on Postgres and `sql_select_limit` on MySQL, 200 rows at most, answered as JSON lines. MySQL's read-only transaction doesn't stop DDL, so the statement check is what keeps a query to reading; the docs tell people to make a read-only database user as well. Lists take `?search=` (only where the collection has a `searchHint`; Granola has none), `?filter=` (the collection's `filterField`, such as `?filter=Draft`), `?parent=` (only in a `nested` collection: the list inside a record, answered with its `trail`, the records above it; a record made there goes in it) and `?page=` (`list-query.ts`) and answer a page, newest first (a mailbox's 50 emails, Granola's 30 days), with the next page's token; a mailbox's token holds where each folder goes on, so paging through the merged list shows every email once, and its search uses Gmail's own syntax on Gmail (`X-GM-EXT-1`) and IMAP text search elsewhere. Agents sign in at `/agent-auth/login` with a cookie of their own (`src/agent-auth/`); `src/agent-side/` serves them: the connections an agent can use, the tools it has with each, and `runAgentTool` (`tools.ts`), the one place a tool runs, which checks it against the agent's access as it is now, validates its params against the tool's schema (an edit changes only the fields given), calls the connector, or a custom server's tool as it is, answering `{text}`, and logs it, denied ones too, to the activity log that `src/activity/` lists. Over HTTP that is `/api/agent/me`, `GET /api/agent/connections/:id/tools` and `POST /api/agent/connections/:id/tools/:tool` with `{params}`; `src/mcp/` serves the same over MCP at `/mcp` (Streamable HTTP without sessions, JSON answers, no stream): four tools whatever the login may do, `list_connections` (the entry point), `list_tools`, `read` (only tools that read, marked read-only) and `run`, and server instructions naming the agent's connections. MCP clients sign in with OAuth 2.1 (`src/oauth/`), Personal Agent Proxy being its own authorization server: metadata under `/.well-known/`, Dynamic Client Registration at `/oauth/register` (open to anyone, each registration first deleting the clients that registered over a day ago and never finished a sign-in), `/oauth/authorize` sealing the request into the web app's `/connect-agent` page, where the person chooses the agent login the client works as or makes one, and `/oauth/token` with PKCE (S256 only) and refresh tokens replaced on every use; codes and tokens are kept as their SHA-256, and revoking an agent ends its clients' tokens. These endpoints take any origin (CORS), since none reads a cookie; credentials and Information records are stored encrypted with `ENCRYPTION_KEY` (`src/utils/secret-crypto.ts`). Has its own `.prettierrc.json` with a 200-column print width +- `packages/api` - `@proxy/api`: Express on Bun. Routes are `handleXRoute` functions mounted in `src/server/create-api-router.ts`; JSON routes return `Result` through `withAuthResult`, which puts the signed-in user on `req.user`. Every environment variable is in `src/utils/env-schema.ts` with its docs, which the docs' Environment variables page is built from; add new ones there. Personal Agent Proxy is one origin everywhere: when `packages/web/dist` exists, as in the Docker image, the api serves the web app's build too (`src/server/serve-web.ts`), and every address the api builds starts from `APP_URL`. Browser-called writes take `requireBrowserOrigin` (CSRF); routes that read a cookie have no CORS. `src/connections/` lists, deletes and sets defaults on an organization's connections, connects a mailbox with an app password after signing in to its IMAP and SMTP servers (`email/`), and connects Granola, Notion and Linear by signing in to their MCP servers (`mcp/`, shared by all three: OAuth with PKCE against the server's authorization server, a client registered for each sign-in so nothing is set up beforehand, the handshake in an encrypted cookie, `mcpConnectRoutes` making both routes; tokens are refreshed as they run out, one refresh at a time per connection on each server, read again first so ones another call or server saved meanwhile are used rather than refreshed twice, and saved only over the ones they replace, whatever else the credential holds kept beside them, by `connections/token-session.ts`, which Google's connection will share; each tool call is one stateless `tools/call`, `mcp-client.ts`, which also lists a server's tools and, for a server that may keep a session (`handshake`), initializes one around each call as the spec asks). Every request to an MCP server, its metadata or its authorization server goes through `utils/checked-request.ts`, which checks the address as a database's host is checked (`utils/host-check.ts`: nothing private unless `ALLOW_PRIVATE_HOSTS`, and `https` unless the address is private) before each request, since a custom server and its endpoints are wherever it said, and never follows a redirect. A custom MCP server is connected by its address (`custom/`): `POST /api/connections/mcp` with `{name, url, token?, returnTo?}` (the name is the person's, kept through the sign-in in its cookie, since a server says nothing about itself that would do) asks the server for its tools, connects a server that answers at once, with no sign-in or the token typed in as a bearer token (as written when it has its scheme in front, `Basic …`), and for one that turns a request without a token down begins its sign-in there, under the browser's origin check, answering `{kind: 'signIn', url}` with the sign-in page the browser goes to: `mcp-discovery.ts` finds the authorization server from the server's protected resource metadata (the `WWW-Authenticate` header's `resource_metadata`, else the well-known paths), its endpoints from that server's metadata (else `/authorize`, `/token` and `/register` under it) and the scopes the server's own metadata names, and `mcpSignInRoutes` (`mcp-connect.ts`, beside `mcpConnectRoutes` for the servers known in advance) keeps the endpoints in the handshake cookie and in the credential beside the tokens. The connection's account is the address without its scheme (`mcp.example.com/mcp`), connecting the same address again replaces its credential and tools, a new connection's defaults are the tools that say they only read, `POST /api/connections/:id/tools` asks the server for its tools again (replacing the tools drops the defaults and agents' settings for tools the server no longer lists, `updateConnectionTools`), and a tool is run as it is (`custom-mcp.ts`): its params go to the server as given, a tool that fails answers the agent with the server's own text as a `validation_error`, and the run is logged under the collection `tools` by the tool's name, with its params as JSON, cut to 200 characters. Granola's notes and transcripts are on by default, Notion's Read pages and Linear's Read issues alone; a Notion or Linear sign-in is for one workspace, and a Notion one lasts at most 180 days, however often it is refreshed. `src/agents/` makes, revokes and deletes agent logins (the password is sent once, only its argon2id hash is kept) and sets an agent's own settings per action of a connection where they differ from the default. `src/records/` reads and writes a collection's records through the integration's `Connector` (`connectors.ts`): Information's live in the database (`info-connector.ts`), a mailbox's are fetched over IMAP on every request (`email-connector.ts`), which merges the inbox and the folders the server marks as Drafts and Sent into one list (ids start with the folder: `inbox-…`, `drafts-…`, `sent-…`) and runs the commands: flags, moves to the folders the server marks as Archive (Gmail's All Mail) or Trash, and sending over SMTP, a draft or a new email, filing a copy in Sent where the server doesn't (Gmail does). Received emails are never edited or deleted for good; only drafts are. Granola's notes are fetched from its MCP server on every request (`granola-connector.ts`; the reading of the XML-like text Granola answers in is in `connections/granola/granola-mcp.ts`): its Notes and Transcripts collections are the same meetings by Granola's meeting id, and since Granola's tools have no pages, a page is 30 days of meetings; one that finds none skips back to the newest meeting before it, so the list ends where the meetings do. Granola's dates and attendees are shown as Granola writes them. Notion's workspace is fetched from its MCP server the same way (`notion-connector.ts`, the reading of its answers in `connections/notion/notion-mcp.ts`), as the person who signed in, as a tree in one nested collection: Notion has no list of every page, so the top is their Recents and a search (Notion's own, AI search where the workspace has it) finds the rest, one page of 50; inside a page are the pages and databases in its content, inside a database its saved views (`view-…` ids), or the rows of its only one, and inside a view its rows, paged with Notion's cursor (view queries have no plan quota; Notion's row and SQL queries do). A page reads as Notion-flavored Markdown, a row's properties beside it as "Name: value" lines (`notion-properties.ts`), and only a page's, database's or view's id is ever sent to Notion's fetch; a new page is a private draft at the top, or goes inside the page, database or view it is made in, a row's title in the database's title property; an edit sends the title, the properties written differently and the content, each only where it changed, and pages are never deleted. Notion's MCP server has no tool that deletes or trashes a page. Linear's issues are fetched from its MCP server the same way (`linear-connector.ts`, the reading of its answers in `connections/linear/linear-mcp.ts`), as the person who signed in, by their identifier (`ENG-123`): a list is the issues updated last first, 50 to a page with Linear's cursor, narrowed by the type of status (`?statusType=Started`) and searched in titles and descriptions, and an issue opens with its comments, oldest first. Teams, statuses, people, projects and labels are written by name, as Linear shows them and finds them, its refusals passed on to the agent; an edit sends only the fields written differently, an empty one clearing it. Linear's MCP server has no tool that archives or deletes an issue. PostgreSQL and MySQL (MariaDB too) are two integrations on one code path: a database is connected with its host, port, database, user, password and encryption (`connections/sql/`: `route.ts` makes both routes per dialect, `sql-client.ts` opens one connection per request with `Bun.sql`, to the address `utils/host-check.ts` resolves the host to, refusing private addresses unless `ALLOW_PRIVATE_HOSTS` is set, and `sql-statement.ts` checks an agent's query is one statement that reads), and `records/sql-connector.ts` reads and writes it through `records/sql-dialect.ts`, where the two differ: the tables are one nested collection, a table holding its rows, a row read and written as "column: value" lines and found again by its primary key (ids `public.users:42`, `pairs:3,4`; a table without a key is never written to), and a query, a second collection that is only ever run, goes through a read-only transaction with a time limit, a cursor on Postgres and `sql_select_limit` on MySQL, 200 rows at most, answered as JSON lines. MySQL's read-only transaction doesn't stop DDL, so the statement check is what keeps a query to reading; the docs tell people to make a read-only database user as well. Lists take `?search=` (only where the collection has a `searchHint`; Granola has none), `?filter=` (the collection's `filterField`, such as `?filter=Draft`), `?parent=` (only in a `nested` collection: the list inside a record, answered with its `trail`, the records above it; a record made there goes in it) and `?page=` (`list-query.ts`) and answer a page, newest first (a mailbox's 50 emails, Granola's 30 days), with the next page's token; a mailbox's token holds where each folder goes on, so paging through the merged list shows every email once, and its search uses Gmail's own syntax on Gmail (`X-GM-EXT-1`) and IMAP text search elsewhere. Agents sign in at `/agent-auth/login` with a cookie of their own (`src/agent-auth/`); `src/agent-side/` serves them: the connections an agent can use, the tools it has with each, and `runAgentTool` (`tools.ts`), the one place a tool runs, which checks it against the agent's access as it is now, validates its params against the tool's schema (an edit changes only the fields given), calls the connector, or a custom server's tool as it is, answering `{text}`, and logs it, denied ones too, to the activity log that `src/activity/` lists. Over HTTP that is `/api/agent/me`, `GET /api/agent/connections/:id/tools` and `POST /api/agent/connections/:id/tools/:tool` with `{params}`; `src/mcp/` serves the same over MCP at `/mcp` (Streamable HTTP without sessions, JSON answers, no stream): four tools whatever the login may do, `list_connections` (the entry point), `list_tools`, `read` (only tools that read, marked read-only) and `run`, and server instructions naming the agent's connections. MCP clients sign in with OAuth 2.1 (`src/oauth/`), Personal Agent Proxy being its own authorization server: metadata under `/.well-known/`, Dynamic Client Registration at `/oauth/register` (open to anyone, each registration first deleting the clients that registered over a day ago and never finished a sign-in), `/oauth/authorize` sealing the request into the web app's `/connect-agent` page, where the person chooses the agent login the client works as or makes one, and `/oauth/token` with PKCE (S256 only) and refresh tokens replaced on every use; codes and tokens are kept as their SHA-256, and revoking an agent ends its clients' tokens. These endpoints take any origin (CORS), since none reads a cookie; credentials and Information records are stored encrypted with `ENCRYPTION_KEY` (`src/utils/secret-crypto.ts`). Has its own `.prettierrc.json` with a 200-column print width - `skills/` - Agent skills, symlinked into `.claude/skills` and `.codex/skills` by `bun run setup` ## The app diff --git a/packages/api/src/connections/custom/route.test.ts b/packages/api/src/connections/custom/route.test.ts index 6869502..edb11be 100644 --- a/packages/api/src/connections/custom/route.test.ts +++ b/packages/api/src/connections/custom/route.test.ts @@ -22,7 +22,15 @@ afterAll(() => { fetchSpy.mockRestore(); }); -mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, getConnectionWithCredential, updateConnectionTools, updateConnectionName})); +mock.module('@proxy/db/connection', () => ({ + listConnections, + createConnection, + updateConnectionCredential, + getConnectionWithCredential, + updateConnectionTools, + updateConnectionName, + getConnectionCredential: mock(async () => Ok(null)), +})); const getUserOrgId = mock(); diff --git a/packages/api/src/connections/granola/route.test.ts b/packages/api/src/connections/granola/route.test.ts index fead5dc..7024226 100644 --- a/packages/api/src/connections/granola/route.test.ts +++ b/packages/api/src/connections/granola/route.test.ts @@ -9,7 +9,7 @@ const listConnections = mock(); const createConnection = mock(); const updateConnectionCredential = mock(); -mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock()})); +mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock(), getConnectionCredential: mock(async () => Ok(null))})); const getUserOrgId = mock(); diff --git a/packages/api/src/connections/linear/route.test.ts b/packages/api/src/connections/linear/route.test.ts index 5356289..e3ccf52 100644 --- a/packages/api/src/connections/linear/route.test.ts +++ b/packages/api/src/connections/linear/route.test.ts @@ -7,7 +7,7 @@ const listConnections = mock(); const createConnection = mock(); const updateConnectionCredential = mock(); -mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock()})); +mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock(), getConnectionCredential: mock(async () => Ok(null))})); const getUserOrgId = mock(); diff --git a/packages/api/src/connections/mcp/mcp-session.test.ts b/packages/api/src/connections/mcp/mcp-session.test.ts index 7dee352..b12006c 100644 --- a/packages/api/src/connections/mcp/mcp-session.test.ts +++ b/packages/api/src/connections/mcp/mcp-session.test.ts @@ -7,7 +7,7 @@ import type {McpServer} from './mcp-oauth'; const updateConnectionCredential = mock(); -mock.module('@proxy/db/connection', () => ({updateConnectionCredential})); +mock.module('@proxy/db/connection', () => ({updateConnectionCredential, getConnectionCredential: mock(async () => Ok(null))})); mock.module('../../observability/log', () => ({ log: {debug: mock(), info: mock(), warn: mock(), error: mock()}, diff --git a/packages/api/src/connections/mcp/mcp-session.ts b/packages/api/src/connections/mcp/mcp-session.ts index fd24b36..61a9fc3 100644 --- a/packages/api/src/connections/mcp/mcp-session.ts +++ b/packages/api/src/connections/mcp/mcp-session.ts @@ -1,12 +1,9 @@ -import {Err, Result} from 'ts-results-es'; +import type {Result} from 'ts-results-es'; import type {z} from 'zod'; -import {ApiErr, type ApiError, Do, requirePresent} from '@proxy/utils'; -import {updateConnectionCredential} from '@proxy/db/connection'; +import type {ApiError} from '@proxy/utils'; -import {log} from '../../observability/log'; -import {encryptSecret} from '../../utils/secret-crypto'; -import {readStoredJson} from '../../utils/stored-json'; +import {callWithToken, type TokenRefresher} from '../token-session'; import {callMcpTool} from './mcp-client'; import {isExpiring, mcpCredentialSchema, refreshMcpTokens, type McpServer} from './mcp-oauth'; @@ -18,44 +15,13 @@ const storedCredentialSchema = mcpCredentialSchema.loose(); type StoredCredential = z.infer; -// The credential as stored, encrypted, beside what it holds. -type Session = {stored: string; credential: StoredCredential}; - -// Refreshes under way on this server, by connection and the credential they replace, so requests -// running at once share one: Notion turns down a refresh token used after it was replaced, and -// revokes the whole sign-in for it. -const renewing = new Map>>(); - -function readSession(server: McpServer, connection: StoredConnection): Result { - return Do(($) => { - const stored = $(requirePresent(connection.credential, ApiErr.internalError(new Error(`${server.name} connection without a credential`)))); - return {stored, credential: $(readStoredJson(stored, storedCredentialSchema))}; - }); -} - -// New tokens, saved for the next request unless another server saved its own first; this request -// goes on with these either way. -function refresh(server: McpServer, connectionId: string, session: Session): Promise> { - return Do(async ($) => { - const credential: StoredCredential = {...session.credential, ...$(await refreshMcpTokens(server, session.credential))}; - const stored = $(encryptSecret(JSON.stringify(credential))); - const saved = $(await updateConnectionCredential({connectionId, credential: stored, previous: session.stored})); - if (!saved) { - log.warn(`${server.name} tokens were refreshed twice at once`, {connectionId}); - } - return {stored, credential}; - }); -} - -function renew(server: McpServer, connectionId: string, session: Session): Promise> { - const key = `${connectionId}:${session.stored}`; - const running = renewing.get(key); - if (running) { - return running; - } - const renewal = refresh(server, connectionId, session).finally(() => renewing.delete(key)); - renewing.set(key, renewal); - return renewal; +function mcpRefresher(server: McpServer): TokenRefresher { + return { + name: server.name, + schema: storedCredentialSchema, + isExpiring, + refresh: async (credential) => (await refreshMcpTokens(server, credential)).map((tokens) => ({...credential, ...tokens})), + }; } /** @@ -64,20 +30,7 @@ function renew(server: McpServer, connectionId: string, session: Session): Promi * `credentials_rejected`, and has to be signed in again. */ export function withMcpToken(server: McpServer, connection: StoredConnection, use: (accessToken: string) => Promise>): Promise> { - return Do(async ($) => { - const read = $(readSession(server, connection)); - const session = isExpiring(read.credential) ? $(await renew(server, connection.id, read)) : read; - const first = await use(session.credential.accessToken); - if (first.isOk() || first.error.kind !== 'credentials_rejected') { - return $(first); - } - if (session !== read) { - return $(Err(first.error)); - } - - const renewed = $(await renew(server, connection.id, session)); - return $(await use(renewed.credential.accessToken)); - }); + return callWithToken(mcpRefresher(server), connection, (credential) => use(credential.accessToken)); } /** Calls one of a server's tools with a connection's token, refreshed as `withMcpToken` does. */ diff --git a/packages/api/src/connections/notion/route.test.ts b/packages/api/src/connections/notion/route.test.ts index 6f9315c..4d32ef7 100644 --- a/packages/api/src/connections/notion/route.test.ts +++ b/packages/api/src/connections/notion/route.test.ts @@ -7,7 +7,7 @@ const listConnections = mock(); const createConnection = mock(); const updateConnectionCredential = mock(); -mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock()})); +mock.module('@proxy/db/connection', () => ({listConnections, createConnection, updateConnectionCredential, updateConnectionTools: mock(), updateConnectionName: mock(), getConnectionCredential: mock(async () => Ok(null))})); const getUserOrgId = mock(); diff --git a/packages/api/src/connections/token-session.test.ts b/packages/api/src/connections/token-session.test.ts new file mode 100644 index 0000000..1d2bc7b --- /dev/null +++ b/packages/api/src/connections/token-session.test.ts @@ -0,0 +1,178 @@ +import {beforeEach, describe, expect, mock, test} from 'bun:test'; +import {Err, Ok, type Result} from 'ts-results-es'; +import {z} from 'zod'; + +import {ApiErr, type ApiError} from '@proxy/utils'; + +import type {TokenRefresher} from './token-session'; + +const updateConnectionCredential = mock(); +const getConnectionCredential = mock(); + +mock.module('@proxy/db/connection', () => ({getConnectionCredential, updateConnectionCredential})); + +const warn = mock(); + +mock.module('../observability/log', () => ({ + log: {debug: mock(), info: mock(), warn, error: mock()}, + serializeError: mock(() => ({})), +})); + +mock.module('../utils/secret-crypto', () => ({ + encryptSecret: (plaintext: string): Result => Ok(`enc:${plaintext}`), + decryptSecret: (stored: string): Result => Ok(stored.replace(/^enc:/, '')), +})); + +const credentialSchema = z.object({accessToken: z.string(), expiring: z.boolean()}); + +type Credential = z.infer; + +const refresh = mock<(credential: Credential) => Promise>>(); + +const REFRESHER: TokenRefresher = {name: 'Example', schema: credentialSchema, isExpiring: (credential) => credential.expiring, refresh}; + +const call = mock<(credential: Credential) => Promise>>(); + +const fresh: Credential = {accessToken: 'access-2', expiring: false}; + +function connectionWith(expiring: boolean) { + return {id: 'conn-1', credential: `enc:${JSON.stringify({accessToken: 'access-1', expiring})}`}; +} + +function tokensCalledWith(): string[] { + return call.mock.calls.map(([credential]) => credential.accessToken); +} + +beforeEach(() => { + mock.clearAllMocks(); + refresh.mockResolvedValue(Ok(fresh)); + updateConnectionCredential.mockResolvedValue(Ok(true)); + getConnectionCredential.mockResolvedValue(Ok(null)); + call.mockResolvedValue(Ok('answer')); +}); + +describe('callWithToken', () => { + test('calls with the stored credential while it lasts', async () => { + const {callWithToken} = await import('./token-session'); + const text = (await callWithToken(REFRESHER, connectionWith(false), call)).unwrap(); + + expect(text).toBe('answer'); + expect(tokensCalledWith()).toEqual(['access-1']); + expect(refresh).not.toHaveBeenCalled(); + expect(updateConnectionCredential).not.toHaveBeenCalled(); + }); + + test('refreshes a credential about to run out, and saves the new one only over the old', async () => { + const {callWithToken} = await import('./token-session'); + const connection = connectionWith(true); + await callWithToken(REFRESHER, connection, call); + + expect(tokensCalledWith()).toEqual(['access-2']); + expect(updateConnectionCredential).toHaveBeenCalledWith({connectionId: 'conn-1', credential: `enc:${JSON.stringify(fresh)}`, previous: connection.credential}); + }); + + test('goes on with the new credential when another server saved its own first', async () => { + updateConnectionCredential.mockResolvedValue(Ok(false)); + const {callWithToken} = await import('./token-session'); + expect((await callWithToken(REFRESHER, connectionWith(true), call)).unwrap()).toBe('answer'); + + expect(tokensCalledWith()).toEqual(['access-2']); + expect(warn).toHaveBeenCalledWith('Example tokens were refreshed twice at once', {connectionId: 'conn-1'}); + }); + + test('calls at once on a running-out credential share one refresh', async () => { + const {callWithToken} = await import('./token-session'); + const connection = connectionWith(true); + await Promise.all([callWithToken(REFRESHER, connection, call), callWithToken(REFRESHER, connection, call)]); + + expect(refresh).toHaveBeenCalledTimes(1); + expect(updateConnectionCredential).toHaveBeenCalledTimes(1); + expect(tokensCalledWith()).toEqual(['access-2', 'access-2']); + }); + + test('a second call in the same request uses the tokens the first one saved, without refreshing again', async () => { + let saved: string | null = null; + getConnectionCredential.mockImplementation(async () => Ok(saved)); + updateConnectionCredential.mockImplementation(async ({credential}: {credential: string}) => { + saved = credential; + return Ok(true); + }); + const {callWithToken} = await import('./token-session'); + const connection = connectionWith(true); + await callWithToken(REFRESHER, connection, call); + await callWithToken(REFRESHER, connection, call); + + expect(refresh).toHaveBeenCalledTimes(1); + expect(tokensCalledWith()).toEqual(['access-2', 'access-2']); + }); + + test('uses the tokens another server saved since, without refreshing', async () => { + getConnectionCredential.mockResolvedValue(Ok(`enc:${JSON.stringify(fresh)}`)); + const {callWithToken} = await import('./token-session'); + await callWithToken(REFRESHER, connectionWith(true), call); + + expect(refresh).not.toHaveBeenCalled(); + expect(updateConnectionCredential).not.toHaveBeenCalled(); + expect(tokensCalledWith()).toEqual(['access-2']); + }); + + test('refreshes the tokens saved since when they are running out too, saving over them', async () => { + const saved = `enc:${JSON.stringify({accessToken: 'access-3', expiring: true})}`; + getConnectionCredential.mockResolvedValue(Ok(saved)); + const {callWithToken} = await import('./token-session'); + await callWithToken(REFRESHER, connectionWith(true), call); + + expect(refresh.mock.calls.map(([credential]) => credential.accessToken)).toEqual(['access-3']); + expect(updateConnectionCredential).toHaveBeenCalledWith({connectionId: 'conn-1', credential: `enc:${JSON.stringify(fresh)}`, previous: saved}); + }); + + test('a refresh turned down goes on with the tokens another server saved meanwhile', async () => { + getConnectionCredential.mockResolvedValueOnce(Ok(null)).mockResolvedValueOnce(Ok(`enc:${JSON.stringify(fresh)}`)); + refresh.mockResolvedValue(Err(ApiErr.credentialsRejected())); + const {callWithToken} = await import('./token-session'); + expect((await callWithToken(REFRESHER, connectionWith(true), call)).unwrap()).toBe('answer'); + + expect(tokensCalledWith()).toEqual(['access-2']); + }); + + test('refreshes once when the provider turns the credential down, then gives up', async () => { + call.mockResolvedValueOnce(Err(ApiErr.credentialsRejected())); + const {callWithToken} = await import('./token-session'); + expect((await callWithToken(REFRESHER, connectionWith(false), call)).unwrap()).toBe('answer'); + expect(tokensCalledWith()).toEqual(['access-1', 'access-2']); + expect(refresh).toHaveBeenCalledTimes(1); + + mock.clearAllMocks(); + call.mockResolvedValue(Err(ApiErr.credentialsRejected())); + expect((await callWithToken(REFRESHER, connectionWith(false), call)).unwrapErr().kind).toBe('credentials_rejected'); + expect(tokensCalledWith()).toEqual(['access-1', 'access-2']); + expect(refresh).toHaveBeenCalledTimes(1); + }); + + test('does not refresh again when a just-refreshed credential is turned down', async () => { + call.mockResolvedValue(Err(ApiErr.credentialsRejected())); + const {callWithToken} = await import('./token-session'); + expect((await callWithToken(REFRESHER, connectionWith(true), call)).unwrapErr().kind).toBe('credentials_rejected'); + + expect(refresh).toHaveBeenCalledTimes(1); + expect(tokensCalledWith()).toEqual(['access-2']); + }); + + test('a refresh that fails is the answer, without a call', async () => { + refresh.mockResolvedValue(Err(ApiErr.credentialsRejected())); + const {callWithToken} = await import('./token-session'); + const result = await callWithToken(REFRESHER, connectionWith(true), call); + + expect(result.unwrapErr().kind).toBe('credentials_rejected'); + expect(call).not.toHaveBeenCalled(); + expect(updateConnectionCredential).not.toHaveBeenCalled(); + }); + + test('a connection without a credential is an internal error', async () => { + const {callWithToken} = await import('./token-session'); + const result = await callWithToken(REFRESHER, {id: 'conn-1', credential: null}, call); + + expect(result.unwrapErr().kind).toBe('internal_error'); + expect(call).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/api/src/connections/token-session.ts b/packages/api/src/connections/token-session.ts new file mode 100644 index 0000000..5986a22 --- /dev/null +++ b/packages/api/src/connections/token-session.ts @@ -0,0 +1,130 @@ +import {Err, type Result} from 'ts-results-es'; +import type {z} from 'zod'; + +import {ApiErr, type ApiError, Do, requirePresent} from '@proxy/utils'; +import {getConnectionCredential, updateConnectionCredential} from '@proxy/db/connection'; + +import {log} from '../observability/log'; +import {encryptSecret} from '../utils/secret-crypto'; +import {readStoredJson} from '../utils/stored-json'; + +/** + * A connection's OAuth tokens, kept fresh: read from the encrypted credential, refreshed as they + * run out, one refresh at a time per connection on each server, and saved only over the ones they + * replace. Before refreshing, the credential is read again: a call earlier in the same request, or + * another server, may have refreshed it already, and a refresh token is used only once. The MCP integrations use it; what a credential holds and how it is refreshed is the + * provider's. + */ + +type StoredConnection = {id: string; credential: string | null}; + +/** How one provider's credential is read and refreshed. */ +export type TokenRefresher = { + // The provider's name in errors and logs: "Notion". + name: string; + // What the decrypted credential holds. + schema: z.ZodType; + isExpiring: (credential: C) => boolean; + // New tokens; turned down once the provider revoked the sign-in. + refresh: (credential: C) => Promise>; +}; + +// The credential as stored, encrypted, beside what it holds. +type Session = {stored: string; credential: C}; + +// Refreshes under way on this server, by connection and the credential they replace, so requests +// running at once share one: Notion turns down a refresh token used after it was replaced, and +// revokes the whole sign-in for it. Each holds the new credential as stored, which every request +// sharing it reads for itself. +const renewing = new Map>>(); + +function readSession(refresher: TokenRefresher, stored: string | null): Result, ApiError> { + return Do, ApiError>(($) => { + const present = $(requirePresent(stored, ApiErr.internalError(new Error(`${refresher.name} connection without a credential`)))); + return {stored: present, credential: $(readStoredJson(present, refresher.schema))}; + }); +} + +// The credential as it is saved now: newer than the one read when another call or server +// refreshed it since. A connection deleted meanwhile keeps the one read. +function latestSession(refresher: TokenRefresher, connectionId: string, session: Session): Promise, ApiError>> { + return Do(async ($) => { + const stored = $(await getConnectionCredential(connectionId)); + if (stored === null || stored === session.stored) { + return session; + } + return $(readSession(refresher, stored)); + }); +} + +// New tokens, saved for the next request unless another server saved its own first; this request +// goes on with these either way. Tokens saved since the credential was read are used instead, and +// refreshed only when they too are running out. +function refresh(refresher: TokenRefresher, connectionId: string, session: Session): Promise> { + return Do(async ($) => { + const latest = $(await latestSession(refresher, connectionId, session)); + if (latest.stored !== session.stored && !refresher.isExpiring(latest.credential)) { + return latest.stored; + } + const refreshed = await refresher.refresh(latest.credential); + if (refreshed.isErr()) { + return $(await savedSince(refresher, connectionId, latest, refreshed.error)); + } + const stored = $(encryptSecret(JSON.stringify(refreshed.value))); + const saved = $(await updateConnectionCredential({connectionId, credential: stored, previous: latest.stored})); + if (!saved) { + log.warn(`${refresher.name} tokens were refreshed twice at once`, {connectionId}); + } + return stored; + }); +} + +// After a refresh that failed, the tokens another server saved meanwhile, else the failure. +function savedSince(refresher: TokenRefresher, connectionId: string, session: Session, error: ApiError): Promise> { + return Do(async ($) => { + const latest = $(await latestSession(refresher, connectionId, session)); + if (latest.stored === session.stored) { + return $(Err(error)); + } + return latest.stored; + }); +} + +// The refresh replacing this credential, started unless one already is. +function sharedRefresh(refresher: TokenRefresher, connectionId: string, session: Session): Promise> { + const key = `${connectionId}:${session.stored}`; + const running = renewing.get(key); + if (running) { + return running; + } + const renewal = refresh(refresher, connectionId, session).finally(() => renewing.delete(key)); + renewing.set(key, renewal); + return renewal; +} + +async function renew(refresher: TokenRefresher, connectionId: string, session: Session): Promise, ApiError>> { + const stored = await sharedRefresh(refresher, connectionId, session); + return stored.andThen((renewed) => readSession(refresher, renewed)); +} + +/** + * Makes a call with a connection's credential: refreshed first when it is running out, and once + * more when the provider turns it down anyway. A connection the provider won't refresh any more is + * `credentials_rejected`, and has to be signed in again. + */ +export function callWithToken(refresher: TokenRefresher, connection: StoredConnection, call: (credential: C) => Promise>): Promise> { + return Do(async ($) => { + const read = $(readSession(refresher, connection.credential)); + const session = refresher.isExpiring(read.credential) ? $(await renew(refresher, connection.id, read)) : read; + const first = await call(session.credential); + if (first.isOk() || first.error.kind !== 'credentials_rejected') { + return $(first); + } + if (session !== read) { + return $(Err(first.error)); + } + + const renewed = $(await renew(refresher, connection.id, session)); + return $(await call(renewed.credential)); + }); +} diff --git a/packages/db/src/connection.ts b/packages/db/src/connection.ts index f9b28f0..16e1a2f 100644 --- a/packages/db/src/connection.ts +++ b/packages/db/src/connection.ts @@ -17,7 +17,7 @@ export type ConnectionRow = { defaults: Array<{actionId: string}>; }; -// Never the credential: only `getConnectionWithCredential` reads it. +// Never the credential: only `getConnectionWithCredential` and `getConnectionCredential` read it. const connectionSelect = { id: true, integrationId: true, @@ -61,6 +61,16 @@ export async function getConnectionWithCredential( ); } +/** A connection's encrypted credential as it is now, `null` once the connection is gone. */ +export async function getConnectionCredential( + connectionId: string, +): Promise> { + const row = await wrapDb(() => + db.connection.findUnique({where: {id: connectionId}, select: {credential: true}}), + ); + return row.map((found) => found?.credential ?? null); +} + /** * `credential` is already encrypted; this package never sees it in the clear. `defaults` are the * actions agents get from the start, none when left out; `tools` a custom MCP server's, as JSON,