Author MCP Tools, Resources, and Prompts#

-- Request type (passed to your handler)
CREATE TYPE api.mcp_request AS (
    arguments jsonb,      -- Tool/prompt arguments
    uri text,             -- Resource URI
    context jsonb,        -- {"client": {...}} (a verified OAuth client) or {"connection": "cn_..."} (an approved handle)
    request_id jsonb      -- JSON-RPC id (string, integer, or null); must echo in response
);

-- Response type (JSON-RPC 2.0 envelope)
CREATE TYPE api.mcp_response AS (
    envelope jsonb        -- {jsonrpc, id, result} or {jsonrpc, id, error}
);

-- Tool with outputSchema and tags, emitting structuredContent
SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id',   'a0000001-0001-4000-8000-000000000001'::uuid,
        'type', 'tool',
        'effect', 'read',
        'name', 'weather_at',
        'description', 'Get weather for a location',
        'access', 'authenticated',
        'inputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object('location', jsonb_build_object('type', 'string')),
            'required', jsonb_build_array('location')
        ),
        'outputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object('temp_c', jsonb_build_object('type', 'number'))
        ),
        'tags', jsonb_build_array('weather', 'read-only')
    ),
    $body$
DECLARE v_structured jsonb;
BEGIN
    v_structured := jsonb_build_object('temp_c', 21.5);
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text(v_structured::text)),
        (request).request_id,
        false,
        v_structured
    );
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0001-4000-8000-000000000001',
        'type', 'tool',
        'effect', 'read',
        'name', 'database_info',
        'description', 'Get database version and connection info',
        'access', 'authenticated',
        'inputSchema', jsonb_build_object(
            'type', 'object',
            'properties', jsonb_build_object(),
            'required', jsonb_build_array()
        )
    ),
    $body$
BEGIN
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text(format(
            'Database: %s, Version: %s, User: %s',
            current_database(),
            version(),
            current_user
        ))),
        (request).request_id
    );
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0004-4000-8000-000000000001',
        'type', 'tool', 'effect', 'read',
        'name', 'list_tasks',
        'description', 'Tasks of one project, newest first',
        'access', 'authenticated',
        'inputSchema', jsonb_build_object(
            'type', 'object',
            'required', jsonb_build_array('project'),
            'properties', jsonb_build_object(
                'project', jsonb_build_object('type', 'string', 'format', 'uuid'),
                'states',  jsonb_build_object('type', 'array', 'items', jsonb_build_object('enum', jsonb_build_array('open', 'done'))),
                'detail',  jsonb_build_object('enum', jsonb_build_array('summary', 'full')),
                'limit',   jsonb_build_object('type', 'integer', 'minimum', 1, 'maximum', 100),
                'offset',  jsonb_build_object('type', 'integer', 'minimum', 0)))),
    $body$
DECLARE
    v_args     jsonb := (request).arguments;
    v_project  uuid := (v_args->>'project') ?> NULL::uuid;
    v_detail   text := COALESCE(v_args->>'detail', 'summary');
    v_limit    integer := COALESCE((v_args->>'limit') ?> NULL::integer, 20);
    v_offset   integer := COALESCE((v_args->>'offset') ?> NULL::integer, 0);
    v_problems jsonb;
    v_rows     jsonb;
BEGIN
    SELECT jsonb_agg(api.invalid_param(name, reason) ORDER BY name) INTO v_problems
    FROM (VALUES
        ('project', 'invalid',      v_project IS NULL),
        ('states',  'invalid',      EXISTS (SELECT FROM jsonb_array_elements_text(COALESCE(v_args->'states', '[]')) s
                                            WHERE s NOT IN ('open', 'done'))),
        ('detail',  'invalid',      v_detail NOT IN ('summary', 'full')),
        ('limit',   'out_of_range', v_limit NOT BETWEEN 1 AND 100),
        ('offset',  'out_of_range', v_offset < 0)
    ) p(name, reason, failed)
    WHERE failed;
    IF v_problems IS NOT NULL THEN
        RETURN api.mcp_tool_result(jsonb_build_array(api.mcp_text(jsonb_build_object(
            'code', 'pgmi.invalid_value',
            'detail', 'Fix the arguments named in invalidParams and call again',
            'invalidParams', v_problems)::text)), (request).request_id, true);
    END IF;

    SELECT jsonb_agg(CASE v_detail WHEN 'full' THEN to_jsonb(t)
                     ELSE jsonb_build_object('id', t.id, 'title', t.title, 'state', t.state) END
                     ORDER BY t.created_at DESC, t.id)
    INTO v_rows
    FROM (SELECT * FROM core.task
          WHERE project_id = v_project
            AND (v_args->'states' IS NULL OR state = ANY (ARRAY(SELECT jsonb_array_elements_text(v_args->'states'))))
          ORDER BY created_at DESC, id
          LIMIT v_limit + 1 OFFSET v_offset) t;

    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text(format('%s task(s)', LEAST(jsonb_array_length(COALESCE(v_rows, '[]')), v_limit)))),
        (request).request_id, false,
        jsonb_build_object(
            'tasks', COALESCE((SELECT jsonb_agg(e) FROM jsonb_array_elements(v_rows) WITH ORDINALITY x(e, n) WHERE n <= v_limit), '[]'),
            'hasMore', jsonb_array_length(COALESCE(v_rows, '[]')) > v_limit));
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0002-4000-8000-000000000001',
        'type', 'resource',
        'name', 'table_schema',
        'description', 'Get table schema information',
        'access', 'authenticated',
        'uriTemplate', 'postgres:///{schema}/{table}',
        'mimeType', 'application/json'
    ),
    $body$
DECLARE
    v_uri_parts text[];
    v_schema text;
    v_table text;
    v_columns jsonb;
BEGIN
    -- Parse URI: postgres:///public/users -> ['public', 'users']
    v_uri_parts := string_to_array(
        regexp_replace((request).uri, '^postgres:///', ''), '/'
    );
    v_schema := v_uri_parts[1];
    v_table := v_uri_parts[2];

    -- Query column metadata
    SELECT jsonb_agg(jsonb_build_object(
        'column_name', column_name,
        'data_type', data_type,
        'is_nullable', is_nullable
    ))
    INTO v_columns
    FROM information_schema.columns
    WHERE table_schema = COALESCE(v_schema, 'public')
      AND table_name = v_table;

    RETURN api.mcp_resource_result(
        jsonb_build_array(jsonb_build_object(
            'uri', (request).uri,
            'mimeType', 'application/json',
            'text', COALESCE(v_columns, '[]'::jsonb)::text
        )),
        (request).request_id
    );
END;
    $body$
);

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        'id', 'e3000001-0003-4000-8000-000000000001',
        'type', 'prompt',
        'name', 'sql_assistant',
        'description', 'Generate a SQL query assistant prompt',
        'access', 'authenticated',
        'arguments', jsonb_build_array(
            jsonb_build_object('name', 'task', 'description', 'Task description', 'required', true),
            jsonb_build_object('name', 'tables', 'description', 'Relevant tables', 'required', false)
        )
    ),
    $body$
DECLARE
    v_task text;
    v_tables text;
BEGIN
    v_task := (request).arguments->>'task';
    v_tables := COALESCE((request).arguments->>'tables', 'any relevant tables');

    RETURN api.mcp_prompt_result(
        jsonb_build_array(
            jsonb_build_object(
                'role', 'user',
                'content', jsonb_build_object(
                    'type', 'text',
                    'text', format(
                        'You are a PostgreSQL expert. Help me write a SQL query for: %s. Consider using: %s.',
                        v_task, v_tables
                    )
                )
            )
        ),
        (request).request_id
    );
END;
    $body$
);

-- Gateway forwards the connection handle (Authorization: Bearer cn_... or X-Connection) as context:
SELECT api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"1","method":"tools/call",...}'::jsonb,
    '{"connection": "cn_..."}'::jsonb
);

-- Inside handlers, use the typed readers:
DECLARE
    v_principal uuid := api.current_principal_id();         -- the person (effective principal)
    v_agent     uuid := api.current_actual_principal_id();  -- the connection's agent (actual principal)
    v_kind      text := api.current_actor_kind();           -- 'agent' on MCP; 'human' for a person's own call
    v_org       uuid := api.current_organization_id();
BEGIN
    -- Use for kernel arguments, or to word an answer for an agent
END;

SELECT api.create_or_replace_mcp_handler(
    jsonb_build_object(
        ...
        'access', 'authenticated'
    ),
    $body$...
);

{"jsonrpc": "2.0", "id": "...", "error": {"code": -32602, "message": "Tool not found: my_tool (or it requires authentication)"}}

-- Discover the server (MCP 2026-07-28; legacy clients send initialize)
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"1","method":"server/discover"}'::jsonb
)).envelope;

-- Test tools/list
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"2","method":"tools/list"}'::jsonb
)).envelope;

-- Test tools/call
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"3","method":"tools/call","params":{"name":"database_info","arguments":{}}}'::jsonb
)).envelope;

-- Test as a person: approve a connection for their membership (in a
-- __test__ file, pg_temp.test_mcp_context does that and returns the context)
SELECT (api.mcp_handle_request(
    '{"jsonrpc":"2.0","id":"4","method":"tools/call","params":{"name":"execute_query","arguments":{"query":"SELECT 1"}}}'::jsonb,
    pg_temp.test_mcp_context('dev|alice', '<organization id>')
)).envelope;

DO $$
DECLARE
    v_response api.mcp_response;
    v_envelope jsonb;
BEGIN
    -- Register test handler
    PERFORM api.create_or_replace_mcp_handler(
        jsonb_build_object(
            'id', 'ffffffff-0661-4000-8000-000000000001',
            'type', 'tool',
            'effect', 'read',
            'name', 'test_tool',
            'description', 'Test',
            'inputSchema', '{"type":"object"}'::jsonb,
            'access', 'public'
        ),
        $body$
BEGIN
    RETURN api.mcp_tool_result(
        jsonb_build_array(api.mcp_text('success')),
        (request).request_id
    );
END;
        $body$
    );

    -- Invoke
    v_response := api.mcp_call_tool('test_tool', '{}'::jsonb, NULL, '"test-1"'::jsonb);
    v_envelope := (v_response).envelope;

    -- Verify
    IF v_envelope->>'jsonrpc' != '2.0' THEN
        RAISE EXCEPTION 'Missing jsonrpc 2.0';
    END IF;
    IF v_envelope->>'id' != 'test-1' THEN
        RAISE EXCEPTION 'request_id not preserved';
    END IF;
    IF v_envelope->'error' IS NOT NULL THEN
        RAISE EXCEPTION 'Unexpected error';
    END IF;

    RAISE NOTICE 'Test passed';
END $$;