Skip to content

Commit 83c88c7

Browse files
evan-onyxevan-onyx
authored
feat: mcp client1 (#5271)
* working mcp implementation v1 * attempt openapi fix * fastmcp
1 parent 2372dd4 commit 83c88c7

File tree

48 files changed

+6933
-355
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

48 files changed

+6933
-355
lines changed

‎.github/workflows/pr-integration-tests.yml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,9 @@ jobs:
7070
-i /local/openapi.json \
7171
-g python \
7272
-o /local/onyx_openapi_client \
73-
--package-name onyx_openapi_client
73+
--package-name onyx_openapi_client \
74+
--skip-validate-spec \
75+
--openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
7476
7577
- name: Set up Docker Buildx
7678
uses: docker/setup-buildx-action@v3

‎.github/workflows/pr-mit-integration-tests.yml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,9 @@ jobs:
6767
-i /local/openapi.json \
6868
-g python \
6969
-o /local/onyx_openapi_client \
70-
--package-name onyx_openapi_client
70+
--package-name onyx_openapi_client \
71+
--skip-validate-spec \
72+
--openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
7173
7274
- name: Set up Docker Buildx
7375
uses: docker/setup-buildx-action@v3

‎.github/workflows/pr-python-checks.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ jobs:
4848
-g python \
4949
-o /local/onyx_openapi_client \
5050
--package-name onyx_openapi_client \
51+
--skip-validate-spec \
52+
--openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
5153
5254
- name: Run MyPy
5355
run: |

‎backend/alembic/versions/7ed603b64d5a_add_mcp_server_and_connection_config_.py‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
"""add_mcp_server_and_connection_config_models
2+
3+
Revision ID: 7ed603b64d5a
4+
Revises: b329d00a9ea6
5+
Create Date: 2025-07-28 17:35:59.900680
6+
7+
"""
8+
9+
from alembic import op
10+
import sqlalchemy as sa
11+
from sqlalchemy.dialects import postgresql
12+
from onyx.db.enums import MCPAuthenticationType
13+
14+
# revision identifiers, used by Alembic.
15+
revision = "7ed603b64d5a"
16+
down_revision = "b329d00a9ea6"
17+
branch_labels = None
18+
depends_on = None
19+
20+
21+
def upgrade() -> None:
22+
"""Create tables and columns for MCP Server support"""
23+
24+
# 1. MCP Server main table (no FK constraints yet to avoid circular refs)
25+
op.create_table(
26+
"mcp_server",
27+
sa.Column("id", sa.Integer(), primary_key=True),
28+
sa.Column("owner", sa.String(), nullable=False),
29+
sa.Column("name", sa.String(), nullable=False),
30+
sa.Column("description", sa.String(), nullable=True),
31+
sa.Column("server_url", sa.String(), nullable=False),
32+
sa.Column(
33+
"auth_type",
34+
sa.Enum(
35+
MCPAuthenticationType,
36+
name="mcp_authentication_type",
37+
native_enum=False,
38+
),
39+
nullable=False,
40+
),
41+
sa.Column("admin_connection_config_id", sa.Integer(), nullable=True),
42+
sa.Column(
43+
"created_at",
44+
sa.DateTime(timezone=True),
45+
server_default=sa.text("now()"), # type: ignore
46+
nullable=False,
47+
),
48+
sa.Column(
49+
"updated_at",
50+
sa.DateTime(timezone=True),
51+
server_default=sa.text("now()"), # type: ignore
52+
nullable=False,
53+
),
54+
)
55+
56+
# 2. MCP Connection Config table (can reference mcp_server now that it exists)
57+
op.create_table(
58+
"mcp_connection_config",
59+
sa.Column("id", sa.Integer(), primary_key=True),
60+
sa.Column("mcp_server_id", sa.Integer(), nullable=True),
61+
sa.Column("user_email", sa.String(), nullable=False, default=""),
62+
sa.Column("config", sa.LargeBinary(), nullable=False),
63+
sa.Column(
64+
"created_at",
65+
sa.DateTime(timezone=True),
66+
server_default=sa.text("now()"), # type: ignore
67+
nullable=False,
68+
),
69+
sa.Column(
70+
"updated_at",
71+
sa.DateTime(timezone=True),
72+
server_default=sa.text("now()"), # type: ignore
73+
nullable=False,
74+
),
75+
sa.ForeignKeyConstraint(
76+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
77+
),
78+
)
79+
80+
# Helpful indexes
81+
op.create_index(
82+
"ix_mcp_connection_config_server_user",
83+
"mcp_connection_config",
84+
["mcp_server_id", "user_email"],
85+
)
86+
op.create_index(
87+
"ix_mcp_connection_config_user_email",
88+
"mcp_connection_config",
89+
["user_email"],
90+
)
91+
92+
# 3. Add the back-references from mcp_server to connection configs
93+
op.create_foreign_key(
94+
"mcp_server_admin_config_fk",
95+
"mcp_server",
96+
"mcp_connection_config",
97+
["admin_connection_config_id"],
98+
["id"],
99+
ondelete="SET NULL",
100+
)
101+
102+
# 4. Association / access-control tables
103+
op.create_table(
104+
"mcp_server__user",
105+
sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
106+
sa.Column("user_id", sa.UUID(), primary_key=True),
107+
sa.ForeignKeyConstraint(
108+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
109+
),
110+
sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
111+
)
112+
113+
op.create_table(
114+
"mcp_server__user_group",
115+
sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
116+
sa.Column("user_group_id", sa.Integer(), primary_key=True),
117+
sa.ForeignKeyConstraint(
118+
["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
119+
),
120+
sa.ForeignKeyConstraint(["user_group_id"], ["user_group.id"]),
121+
)
122+
123+
# 5. Update existing `tool` table – allow tools to belong to an MCP server
124+
op.add_column(
125+
"tool",
126+
sa.Column("mcp_server_id", sa.Integer(), nullable=True),
127+
)
128+
# Add column for MCP tool input schema
129+
op.add_column(
130+
"tool",
131+
sa.Column("mcp_input_schema", postgresql.JSONB(), nullable=True),
132+
)
133+
op.create_foreign_key(
134+
"tool_mcp_server_fk",
135+
"tool",
136+
"mcp_server",
137+
["mcp_server_id"],
138+
["id"],
139+
ondelete="CASCADE",
140+
)
141+
142+
# 6. Update persona__tool foreign keys to cascade delete
143+
# This ensures that when a tool is deleted (including via MCP server deletion),
144+
# the corresponding persona__tool rows are also deleted
145+
op.drop_constraint(
146+
"persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
147+
)
148+
op.drop_constraint(
149+
"persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
150+
)
151+
152+
op.create_foreign_key(
153+
"persona__tool_persona_id_fkey",
154+
"persona__tool",
155+
"persona",
156+
["persona_id"],
157+
["id"],
158+
ondelete="CASCADE",
159+
)
160+
op.create_foreign_key(
161+
"persona__tool_tool_id_fkey",
162+
"persona__tool",
163+
"tool",
164+
["tool_id"],
165+
["id"],
166+
ondelete="CASCADE",
167+
)
168+
169+
# 7. Update research_agent_iteration_sub_step foreign key to SET NULL on delete
170+
# This ensures that when a tool is deleted, the sub_step_tool_id is set to NULL
171+
# instead of causing a foreign key constraint violation
172+
op.drop_constraint(
173+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
174+
"research_agent_iteration_sub_step",
175+
type_="foreignkey",
176+
)
177+
op.create_foreign_key(
178+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
179+
"research_agent_iteration_sub_step",
180+
"tool",
181+
["sub_step_tool_id"],
182+
["id"],
183+
ondelete="SET NULL",
184+
)
185+
186+
187+
def downgrade() -> None:
188+
"""Drop all MCP-related tables / columns"""
189+
190+
# # # 1. Drop FK & columns from tool
191+
# op.drop_constraint("tool_mcp_server_fk", "tool", type_="foreignkey")
192+
op.execute("DELETE FROM tool WHERE mcp_server_id IS NOT NULL")
193+
194+
op.drop_constraint(
195+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
196+
"research_agent_iteration_sub_step",
197+
type_="foreignkey",
198+
)
199+
op.create_foreign_key(
200+
"research_agent_iteration_sub_step_sub_step_tool_id_fkey",
201+
"research_agent_iteration_sub_step",
202+
"tool",
203+
["sub_step_tool_id"],
204+
["id"],
205+
)
206+
207+
# Restore original persona__tool foreign keys (without CASCADE)
208+
op.drop_constraint(
209+
"persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
210+
)
211+
op.drop_constraint(
212+
"persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
213+
)
214+
215+
op.create_foreign_key(
216+
"persona__tool_persona_id_fkey",
217+
"persona__tool",
218+
"persona",
219+
["persona_id"],
220+
["id"],
221+
)
222+
op.create_foreign_key(
223+
"persona__tool_tool_id_fkey",
224+
"persona__tool",
225+
"tool",
226+
["tool_id"],
227+
["id"],
228+
)
229+
op.drop_column("tool", "mcp_input_schema")
230+
op.drop_column("tool", "mcp_server_id")
231+
232+
# 2. Drop association tables
233+
op.drop_table("mcp_server__user_group")
234+
op.drop_table("mcp_server__user")
235+
236+
# 3. Drop FK from mcp_server to connection configs
237+
op.drop_constraint("mcp_server_admin_config_fk", "mcp_server", type_="foreignkey")
238+
239+
# 4. Drop connection config indexes & table
240+
op.drop_index(
241+
"ix_mcp_connection_config_user_email", table_name="mcp_connection_config"
242+
)
243+
op.drop_index(
244+
"ix_mcp_connection_config_server_user", table_name="mcp_connection_config"
245+
)
246+
op.drop_table("mcp_connection_config")
247+
248+
# 5. Finally drop mcp_server table
249+
op.drop_table("mcp_server")

‎backend/ee/onyx/external_permissions/perm_sync_types.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,13 @@
33
from typing import Optional
44
from typing import Protocol
55

6-
from ee.onyx.db.external_perm import ExternalUserGroup
7-
from onyx.access.models import DocExternalAccess
6+
from ee.onyx.db.external_perm import ExternalUserGroup # noqa
7+
from onyx.access.models import DocExternalAccess # noqa
88
from onyx.context.search.models import InferenceChunk
9-
from onyx.db.models import ConnectorCredentialPair
9+
from onyx.db.models import ConnectorCredentialPair # noqa
1010
from onyx.db.utils import DocumentRow
1111
from onyx.db.utils import SortOrder
12-
from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
12+
from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface # noqa
1313

1414

1515
class FetchAllDocumentsFunction(Protocol):

‎backend/onyx/agents/agent_search/dr/nodes/dr_a1_orchestrator.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ def orchestrator(
253253
],
254254
)
255255

256-
# for Thoightful mode, we force a tool if requested an available
256+
# for Thoughtful mode, we force a tool if requested an available
257257
available_tools_for_decision = available_tools
258258
force_use_tool = graph_config.tooling.force_use_tool
259259
if iteration_nr == 1 and force_use_tool and force_use_tool.force_use:

‎backend/onyx/agents/agent_search/dr/nodes/dr_a3_logger.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,9 @@ def save_iteration(
120120
# lastly, insert the citations
121121
citation_dict: dict[int, int] = {}
122122
cited_doc_nrs = _extract_citation_numbers(final_answer)
123-
for cited_doc_nr in cited_doc_nrs:
124-
citation_dict[cited_doc_nr] = search_docs[cited_doc_nr - 1].id
123+
if search_docs:
124+
for cited_doc_nr in cited_doc_nrs:
125+
citation_dict[cited_doc_nr] = search_docs[cited_doc_nr - 1].id
125126

126127
# TODO: generate plan as dict in the first place
127128
plan_of_record = state.plan_of_record.plan if state.plan_of_record else ""

‎backend/onyx/agents/agent_search/dr/sub_agents/custom_tool/dr_custom_tool_2_act.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from onyx.tools.tool_implementations.custom.custom_tool import CUSTOM_TOOL_RESPONSE_ID
1919
from onyx.tools.tool_implementations.custom.custom_tool import CustomTool
2020
from onyx.tools.tool_implementations.custom.custom_tool import CustomToolCallSummary
21+
from onyx.tools.tool_implementations.mcp.mcp_tool import MCP_TOOL_RESPONSE_ID
2122
from onyx.utils.logger import setup_logger
2223

2324
logger = setup_logger()
@@ -40,7 +41,7 @@ def custom_tool_act(
4041
raise ValueError("available_tools is not set")
4142

4243
custom_tool_info = state.available_tools[state.tools_used[-1]]
43-
custom_tool_name = custom_tool_info.llm_path
44+
custom_tool_name = custom_tool_info.name
4445
custom_tool = cast(CustomTool, custom_tool_info.tool_object)
4546

4647
branch_query = state.branch_question
@@ -94,7 +95,7 @@ def custom_tool_act(
9495
# run the tool
9596
response_summary: CustomToolCallSummary | None = None
9697
for tool_response in custom_tool.run(**tool_args):
97-
if tool_response.id == CUSTOM_TOOL_RESPONSE_ID:
98+
if tool_response.id in {CUSTOM_TOOL_RESPONSE_ID, MCP_TOOL_RESPONSE_ID}:
9899
response_summary = cast(CustomToolCallSummary, tool_response.response)
99100
break
100101

‎backend/onyx/auth/users.py‎

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -572,22 +572,19 @@ async def on_after_register(
572572
logger.debug(f"Current tenant user count: {user_count}")
573573

574574
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
575-
if user_count == 1:
576-
create_milestone_and_report(
577-
user=user,
578-
distinct_id=user.email,
579-
event_type=MilestoneRecordType.USER_SIGNED_UP,
580-
properties=None,
581-
db_session=db_session,
582-
)
583-
else:
584-
create_milestone_and_report(
585-
user=user,
586-
distinct_id=user.email,
587-
event_type=MilestoneRecordType.MULTIPLE_USERS,
588-
properties=None,
589-
db_session=db_session,
590-
)
575+
event_type = (
576+
MilestoneRecordType.USER_SIGNED_UP
577+
if user_count == 1
578+
else MilestoneRecordType.MULTIPLE_USERS
579+
)
580+
create_milestone_and_report(
581+
user=user,
582+
distinct_id=user.email,
583+
event_type=event_type,
584+
properties=None,
585+
db_session=db_session,
586+
)
587+
591588
finally:
592589
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
593590

‎backend/onyx/db/enums.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,17 @@ def is_terminal(self) -> bool:
5656
return self in terminal_states
5757

5858

59+
class MCPAuthenticationType(str, PyEnum):
60+
NONE = "NONE"
61+
API_TOKEN = "API_TOKEN"
62+
OAUTH = "OAUTH"
63+
64+
65+
class MCPAuthenticationPerformer(str, PyEnum):
66+
ADMIN = "ADMIN"
67+
PER_USER = "PER_USER"
68+
69+
5970
# Consistent with Celery task statuses
6071
class TaskStatus(str, PyEnum):
6172
PENDING = "PENDING"

0 commit comments

Comments
 (0)