pi-mono/apps/p2/plugins/mcp/contract.ts

// MCP servers' contracts: the server's manager (host-internal: every MCP connection and process, by instance and lease)
// and its view per principal (`Mcp`: a session worker sets its locations and calls tools; the owner sees and restarts
// instances).

import type { CallToolResult, ToolAnnotations } from "@earendil-works/pi-mcp";
import { arr, int, nul, obj, oneOf, opt, type Static, str, union, unknown } from "@earendil-works/pi-schema";
import { LocationSchema, literal, local, method, record, remote, serverHost, unsafe } from "@p2/app/v1";

/** An MCP tool as its server lists it (`tools/list`). */
const ToolSchema = obj({
	name: str(),
	title: opt(str()),
	description: opt(str()),
	inputSchema: unsafe<Record<string, unknown>>(unknown()),
	outputSchema: opt(unsafe<Record<string, unknown>>(unknown())),
	annotations: opt(unsafe<ToolAnnotations>(unknown())),
});
/** An MCP tool as its server lists it: name, title, description, schemas and annotations. */
export type McpTool = Static<typeof ToolSchema>;

/** Where an instance runs: the p2 host, or an environment of `envs.json`. */
const RunsOnSchema = union([literal("host"), obj({ env: str() })]);
/** Where an instance runs: `"host"` (the p2 host) or `{ env }` (an environment of `envs.json`). */
export type RunsOn = Static<typeof RunsOnSchema>;

/** One running (or failed) server process or connection. */
export const InstanceSchema = obj({
	id: str(),
	server: str(),
	source: oneOf(["owner", "project"]),
	/** The location it serves; null: shared by every session, or an HTTP server. */
	location: union([LocationSchema, nul()]),
	runsOn: RunsOnSchema,
	/** `sign-in`: the HTTP server asked for authorization. */
	status: oneOf(["starting", "connected", "sign-in", "failed"]),
	error: opt(str()),
	/** Kept while it reconnects. */
	tools: arr(ToolSchema),
});
/** One instance as the manager lists it: its server, where it runs, its status, error and tools. */
export type Instance = Static<typeof InstanceSchema>;

/** The instances one session's lease at a location uses, and why entries of its project were refused. */
const LeaseViewSchema = obj({ session: str(), location: LocationSchema, instances: arr(str()), errors: arr(str()) });

/** The largest project `mcp.json` a worker sends. */
export const PROJECT_FILE_MAX_CHARS = 256 * 1024;

/**
 * A project's `<cwd>/.pi/mcp.json` at one of the locations, as the worker read it there: its text (checked by the
 * manager: a trust boundary), or why it could not be read.
 */
const ProjectFileSchema = union([
	obj({ location: LocationSchema, text: str({ maxLength: PROJECT_FILE_MAX_CHARS }) }),
	obj({ location: LocationSchema, error: str({ maxLength: 2000 }) }),
]);
/** A project's `mcp.json` at a location: its text, or why it could not be read. */
export type ProjectFile = Static<typeof ProjectFileSchema>;

/** Where a session's conversations work, and the project files found there (absent: none). */
const Locations = { locations: arr(LocationSchema), projects: opt(arr(ProjectFileSchema)) };

/** `tools/call`'s result as the server sent it. */
const CallResultSchema = unsafe<CallToolResult>({ type: "object" });

/** Tool arguments: any JSON object (a record of any values); the MCP server checks them. */
const ArgsSchema = record(unknown());

/** What every worker request names: the session's worker it comes from (filled in by the view). */
const Caller = { session: str(), generation: int(), incarnation: str() };

/**
 * The MCP manager (server, host lifetime): the owner configuration, the instances and their leases. Every MCP process
 * and connection is here; session workers reach it only through `Mcp`.
 *
 * Authority: a lease places servers only at a location on an environment the session's conversations work on as the
 * server records them (the catalog's session location, the subagent store's own locations); the cwd there is the
 * worker's to choose, as its agent can run commands anywhere on that machine. A project's servers are code the worker
 * names (it reads the file), so this fences stale or confused workers to their machines, not a compromised worker
 * process: it runs unsandboxed on the p2 host, its subagent store entries are its own reports, and its location tool
 * moves the session to any configured environment.
 */
export const McpManager = local("p2.mcp.manager", {
	version: 1,
	params: {},
	members: {
		instances: arr(InstanceSchema),
		/** The leases of every session and location, with the instances each uses. */
		leases: arr(LeaseViewSchema),
		/** Why entries of `mcp.json` were skipped (the leases' projects' errors are in `leases`). */
		errors: arr(str()),
		/**
		 * Leases the configured servers (and the projects') at exactly `locations` for the caller's session (its leases
		 * elsewhere end); a newer generation's leases refuse an older generation.
		 */
		setLocations: method(obj({ ...Caller, ...Locations })),
		/** Ends the leases of one worker generation (its view ended). */
		releaseLeases: method(obj({ session: str(), generation: int() })),
		/** Calls `tool` of `server` at `location`, through the instance the lease there uses. */
		call: method(
			obj({ ...Caller, location: LocationSchema, server: str(), tool: str(), args: ArgsSchema }),
			CallResultSchema,
		),
		restart: method(obj({ instance: str() })),
	},
});

/**
 * MCP servers as a principal sees them: a session worker sets its locations and calls tools (its session, generation
 * and incarnation from the link), the owner sees every instance with its errors and restarts them. Nobody else gets
 * one.
 */
export const Mcp = remote("p2.mcp", {
	version: 1,
	params: {},
	hostAddress: serverHost,
	variants: {
		worker: {
			state: {
				/** The instances the session's leases use, shared ones included; `error` is the owner's only. */
				servers: arr(InstanceSchema),
				/** Why the session's leases refused entries of its projects' `mcp.json` (or a location). */
				errors: arr(str()),
			},
			methods: {
				/**
				 * Where the session's conversations work now, with the project files there: its servers are leased there,
				 * and only there.
				 */
				setLocations: method(obj(Locations)),
				call: method(
					obj({ location: LocationSchema, server: str(), tool: str(), args: ArgsSchema }),
					CallResultSchema,
				),
			},
		},
		owner: {
			/** `errors`: the owner's `mcp.json`'s, then every lease's (by session). */
			state: { instances: arr(InstanceSchema), errors: arr(str()) },
			methods: { restart: method(obj({ instance: str() })) },
		},
	},
});

/** The worker's Durable extension of MCP tools. */
export const MCP_EXTENSION = "mcp";