Jeremiah Lowin commited on
Commit
31e3fd3
·
1 Parent(s): 3d30e82

Add context injection to all objects

Browse files
src/fastmcp/resources/resource_manager.py CHANGED
@@ -61,9 +61,16 @@ class ResourceManager:
61
  The added resource or template. If a resource or template with the same URI already exists,
62
  returns the existing resource or template.
63
  """
 
 
64
  # Check if this should be a template
65
  has_uri_params = "{" in uri and "}" in uri
66
- has_func_params = bool(inspect.signature(fn).parameters)
 
 
 
 
 
67
 
68
  if has_uri_params or has_func_params:
69
  return self.add_template_from_fn(
@@ -102,12 +109,12 @@ class ResourceManager:
102
  The added resource. If a resource with the same URI already exists,
103
  returns the existing resource.
104
  """
105
- resource = FunctionResource(
 
106
  uri=AnyUrl(uri),
107
  name=name,
108
  description=description,
109
  mime_type=mime_type or "text/plain",
110
- fn=fn,
111
  tags=tags or set(),
112
  )
113
  return self.add_resource(resource)
@@ -235,7 +242,9 @@ class ResourceManager:
235
  if params := match_uri_template(uri_str, storage_key):
236
  try:
237
  return await template.create_resource(
238
- uri_str, params, context=context
 
 
239
  )
240
  except Exception as e:
241
  raise ValueError(f"Error creating resource from template: {e}")
 
61
  The added resource or template. If a resource or template with the same URI already exists,
62
  returns the existing resource or template.
63
  """
64
+ from fastmcp.server.context import Context
65
+
66
  # Check if this should be a template
67
  has_uri_params = "{" in uri and "}" in uri
68
+ # check if the function has any parameters (other than injected context)
69
+ has_func_params = any(
70
+ p
71
+ for p in inspect.signature(fn).parameters.values()
72
+ if p.annotation is not Context
73
+ )
74
 
75
  if has_uri_params or has_func_params:
76
  return self.add_template_from_fn(
 
109
  The added resource. If a resource with the same URI already exists,
110
  returns the existing resource.
111
  """
112
+ resource = FunctionResource.from_function(
113
+ fn=fn,
114
  uri=AnyUrl(uri),
115
  name=name,
116
  description=description,
117
  mime_type=mime_type or "text/plain",
 
118
  tags=tags or set(),
119
  )
120
  return self.add_resource(resource)
 
242
  if params := match_uri_template(uri_str, storage_key):
243
  try:
244
  return await template.create_resource(
245
+ uri_str,
246
+ params=params,
247
+ context=context,
248
  )
249
  except Exception as e:
250
  raise ValueError(f"Error creating resource from template: {e}")
src/fastmcp/resources/template.py CHANGED
@@ -189,7 +189,7 @@ class ResourceTemplate(BaseModel):
189
  name=self.name,
190
  description=self.description,
191
  mime_type=self.mime_type,
192
- fn=lambda: result, # Capture result in closure
193
  tags=self.tags,
194
  context_kwarg=self.context_kwarg,
195
  )
 
189
  name=self.name,
190
  description=self.description,
191
  mime_type=self.mime_type,
192
+ fn=lambda **kwargs: result, # Capture result in closure
193
  tags=self.tags,
194
  context_kwarg=self.context_kwarg,
195
  )
src/fastmcp/resources/types.py CHANGED
@@ -15,6 +15,7 @@ import pydantic.json
15
  import pydantic_core
16
  from pydantic import Field, ValidationInfo
17
 
 
18
  from fastmcp.resources.resource import Resource
19
 
20
  if TYPE_CHECKING:
@@ -66,8 +67,23 @@ class FunctionResource(Resource):
66
  default=None, description="Name of the kwarg that should receive context"
67
  )
68
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
69
  async def read(
70
- self, context: Context[ServerSessionT, LifespanContextT] | None = None
 
71
  ) -> str | bytes:
72
  """Read the resource by calling the wrapped function."""
73
  try:
@@ -80,7 +96,7 @@ class FunctionResource(Resource):
80
  result = await result
81
 
82
  if isinstance(result, Resource):
83
- return await result.read()
84
  if isinstance(result, bytes):
85
  return result
86
  if isinstance(result, str):
@@ -127,7 +143,9 @@ class FileResource(Resource):
127
  mime_type = info.data.get("mime_type", "text/plain")
128
  return not mime_type.startswith("text/")
129
 
130
- async def read(self) -> str | bytes:
 
 
131
  """Read the file content."""
132
  try:
133
  if self.is_binary:
@@ -145,7 +163,9 @@ class HttpResource(Resource):
145
  default="application/json", description="MIME type of the resource content"
146
  )
147
 
148
- async def read(self) -> str | bytes:
 
 
149
  """Read the HTTP content."""
150
  async with httpx.AsyncClient() as client:
151
  response = await client.get(self.url)
@@ -197,7 +217,9 @@ class DirectoryResource(Resource):
197
  except Exception as e:
198
  raise ValueError(f"Error listing directory {self.path}: {e}")
199
 
200
- async def read(self) -> str: # Always returns JSON string
 
 
201
  """Read the directory listing."""
202
  try:
203
  files = await anyio.to_thread.run_sync(self.list_files)
 
15
  import pydantic_core
16
  from pydantic import Field, ValidationInfo
17
 
18
+ import fastmcp
19
  from fastmcp.resources.resource import Resource
20
 
21
  if TYPE_CHECKING:
 
67
  default=None, description="Name of the kwarg that should receive context"
68
  )
69
 
70
+ @classmethod
71
+ def from_function(
72
+ cls, fn: Callable[[], Any], context_kwarg: str | None = None, **kwargs
73
+ ) -> FunctionResource:
74
+ if context_kwarg is None:
75
+ parameters = inspect.signature(fn).parameters
76
+ context_param = next(
77
+ (p for p in parameters.values() if p.annotation is fastmcp.Context),
78
+ None,
79
+ )
80
+ if context_param is not None:
81
+ context_kwarg = context_param.name
82
+ return cls(fn=fn, context_kwarg=context_kwarg, **kwargs)
83
+
84
  async def read(
85
+ self,
86
+ context: Context[ServerSessionT, LifespanContextT] | None = None,
87
  ) -> str | bytes:
88
  """Read the resource by calling the wrapped function."""
89
  try:
 
96
  result = await result
97
 
98
  if isinstance(result, Resource):
99
+ return await result.read(context=context)
100
  if isinstance(result, bytes):
101
  return result
102
  if isinstance(result, str):
 
143
  mime_type = info.data.get("mime_type", "text/plain")
144
  return not mime_type.startswith("text/")
145
 
146
+ async def read(
147
+ self, context: Context[ServerSessionT, LifespanContextT] | None = None
148
+ ) -> str | bytes:
149
  """Read the file content."""
150
  try:
151
  if self.is_binary:
 
163
  default="application/json", description="MIME type of the resource content"
164
  )
165
 
166
+ async def read(
167
+ self, context: Context[ServerSessionT, LifespanContextT] | None = None
168
+ ) -> str | bytes:
169
  """Read the HTTP content."""
170
  async with httpx.AsyncClient() as client:
171
  response = await client.get(self.url)
 
217
  except Exception as e:
218
  raise ValueError(f"Error listing directory {self.path}: {e}")
219
 
220
+ async def read(
221
+ self, context: Context[ServerSessionT, LifespanContextT] | None = None
222
+ ) -> str: # Always returns JSON string
223
  """Read the directory listing."""
224
  try:
225
  files = await anyio.to_thread.run_sync(self.list_files)
src/fastmcp/server/openapi.py CHANGED
@@ -265,7 +265,9 @@ class OpenAPIResource(Resource):
265
  self._client = client
266
  self._route = route
267
 
268
- async def read(self) -> str | bytes:
 
 
269
  """Fetch the resource data by making an HTTP request."""
270
  try:
271
  # Extract path parameters from the URI if present
 
265
  self._client = client
266
  self._route = route
267
 
268
+ async def read(
269
+ self, context: Context[ServerSessionT, LifespanContextT] | None = None
270
+ ) -> str | bytes:
271
  """Fetch the resource data by making an HTTP request."""
272
  try:
273
  # Extract path parameters from the URI if present
src/fastmcp/server/server.py CHANGED
@@ -401,7 +401,7 @@ class FastMCP(Generic[LifespanResultT]):
401
  context = self.get_context()
402
  resource = await self._resource_manager.get_resource(uri, context=context)
403
  try:
404
- content = await resource.read()
405
  return [
406
  ReadResourceContents(content=content, mime_type=resource.mime_type)
407
  ]
@@ -427,7 +427,7 @@ class FastMCP(Generic[LifespanResultT]):
427
  if self._prompt_manager.has_prompt(name):
428
  context = self.get_context()
429
  messages = await self._prompt_manager.render_prompt(
430
- name, arguments, context=context
431
  )
432
  return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages))
433
  else:
 
401
  context = self.get_context()
402
  resource = await self._resource_manager.get_resource(uri, context=context)
403
  try:
404
+ content = await resource.read(context=context)
405
  return [
406
  ReadResourceContents(content=content, mime_type=resource.mime_type)
407
  ]
 
427
  if self._prompt_manager.has_prompt(name):
428
  context = self.get_context()
429
  messages = await self._prompt_manager.render_prompt(
430
+ name, arguments=arguments or {}, context=context
431
  )
432
  return GetPromptResult(messages=pydantic_core.to_jsonable_python(messages))
433
  else:
tests/server/test_server_interactions.py CHANGED
@@ -682,7 +682,147 @@ class TestToolParameters:
682
  assert result[0].text == "0:16:40"
683
 
684
 
685
- class TestResources:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
686
  async def test_text_resource(self):
687
  mcp = FastMCP()
688
 
@@ -756,6 +896,21 @@ class TestResources:
756
  assert result[0].blob == base64.b64encode(b"Binary file data").decode()
757
 
758
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
759
  class TestResourceTemplates:
760
  async def test_resource_with_params_not_in_uri(self):
761
  """Test that a resource with function parameters raises an error if the URI
@@ -1026,144 +1181,19 @@ class TestResourceTemplates:
1026
  assert result[0].text == "Template resource 1: a/b"
1027
 
1028
 
1029
- class TestContextInjection:
1030
- """Test context injection in tools."""
1031
-
1032
- async def test_context_detection(self):
1033
- """Test that context parameters are properly detected."""
1034
  mcp = FastMCP()
1035
 
1036
- def tool_with_context(x: int, ctx: Context) -> str:
1037
- return f"Request {ctx.request_id}: {x}"
1038
-
1039
- mcp.add_tool(tool_with_context)
1040
- async with Client(mcp) as client:
1041
- tools = await client.list_tools()
1042
- assert len(tools) == 1
1043
- assert tools[0].name == "tool_with_context"
1044
-
1045
- async def test_context_injection(self):
1046
- """Test that context is properly injected into tool calls."""
1047
- mcp = FastMCP()
1048
-
1049
- def tool_with_context(x: int, ctx: Context) -> str:
1050
- assert ctx.request_id is not None
1051
- return f"Request {ctx.request_id}: {x}"
1052
-
1053
- mcp.add_tool(tool_with_context)
1054
- async with Client(mcp) as client:
1055
- result = await client.call_tool("tool_with_context", {"x": 42})
1056
- assert len(result) == 1
1057
- content = result[0]
1058
- assert isinstance(content, TextContent)
1059
- assert "Request" in content.text
1060
- assert "42" in content.text
1061
-
1062
- async def test_async_context(self):
1063
- """Test that context works in async functions."""
1064
- mcp = FastMCP()
1065
-
1066
- async def async_tool(x: int, ctx: Context) -> str:
1067
- assert ctx.request_id is not None
1068
- return f"Async request {ctx.request_id}: {x}"
1069
-
1070
- mcp.add_tool(async_tool)
1071
- async with Client(mcp) as client:
1072
- result = await client.call_tool("async_tool", {"x": 42})
1073
- assert len(result) == 1
1074
- content = result[0]
1075
- assert isinstance(content, TextContent)
1076
- assert "Async request" in content.text
1077
- assert "42" in content.text
1078
-
1079
- async def test_context_logging(self):
1080
- from unittest.mock import patch
1081
-
1082
- import mcp.server.session
1083
-
1084
- """Test that context logging methods work."""
1085
- mcp = FastMCP()
1086
-
1087
- async def logging_tool(msg: str, ctx: Context) -> str:
1088
- await ctx.debug("Debug message")
1089
- await ctx.info("Info message")
1090
- await ctx.warning("Warning message")
1091
- await ctx.error("Error message")
1092
- return f"Logged messages for {msg}"
1093
-
1094
- mcp.add_tool(logging_tool)
1095
-
1096
- with patch("mcp.server.session.ServerSession.send_log_message") as mock_log:
1097
- async with Client(mcp) as client:
1098
- result = await client.call_tool("logging_tool", {"msg": "test"})
1099
- assert len(result) == 1
1100
- content = result[0]
1101
- assert isinstance(content, TextContent)
1102
- assert "Logged messages for test" in content.text
1103
-
1104
- assert mock_log.call_count == 4
1105
- mock_log.assert_any_call(
1106
- level="debug", data="Debug message", logger=None
1107
- )
1108
- mock_log.assert_any_call(level="info", data="Info message", logger=None)
1109
- mock_log.assert_any_call(
1110
- level="warning", data="Warning message", logger=None
1111
- )
1112
- mock_log.assert_any_call(
1113
- level="error", data="Error message", logger=None
1114
- )
1115
-
1116
- async def test_optional_context(self):
1117
- """Test that context is optional."""
1118
- mcp = FastMCP()
1119
-
1120
- def no_context(x: int) -> int:
1121
- return x * 2
1122
-
1123
- mcp.add_tool(no_context)
1124
- async with Client(mcp) as client:
1125
- result = await client.call_tool("no_context", {"x": 21})
1126
- assert len(result) == 1
1127
- content = result[0]
1128
- assert isinstance(content, TextContent)
1129
- assert content.text == "42"
1130
-
1131
- async def test_context_resource_access(self):
1132
- """Test that context can access resources."""
1133
- mcp = FastMCP()
1134
-
1135
- @mcp.resource("test://data")
1136
- def test_resource() -> str:
1137
- return "resource data"
1138
-
1139
- @mcp.tool()
1140
- async def tool_with_resource(ctx: Context) -> str:
1141
- r_iter = await ctx.read_resource("test://data")
1142
- r_list = list(r_iter)
1143
- assert len(r_list) == 1
1144
- r = r_list[0]
1145
- return f"Read resource: {r.content} with mime type {r.mime_type}"
1146
-
1147
- async with Client(mcp) as client:
1148
- result = await client.call_tool("tool_with_resource", {})
1149
- assert len(result) == 1
1150
- content = result[0]
1151
- assert isinstance(content, TextContent)
1152
- assert "Read resource: resource data" in content.text
1153
-
1154
- async def test_tool_decorator_with_tags(self):
1155
- """Test that the tool decorator properly sets tags."""
1156
- mcp = FastMCP()
1157
-
1158
- @mcp.tool(tags={"example", "test-tag"})
1159
- def sample_tool(x: int) -> int:
1160
- return x * 2
1161
 
1162
- # Verify the tool exists
1163
  async with Client(mcp) as client:
1164
- tools = await client.list_tools()
1165
- assert len(tools) == 1
1166
- # Note: MCPTool from the client API doesn't expose tags
1167
 
1168
 
1169
  class TestPrompts:
@@ -1350,3 +1380,19 @@ class TestPrompts:
1350
  assert len(prompts_dict) == 1
1351
  prompt = prompts_dict["sample_prompt"]
1352
  assert prompt.tags == {"example", "test-tag"}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
682
  assert result[0].text == "0:16:40"
683
 
684
 
685
+ class TestToolContextInjection:
686
+ """Test context injection in tools."""
687
+
688
+ async def test_context_detection(self):
689
+ """Test that context parameters are properly detected."""
690
+ mcp = FastMCP()
691
+
692
+ def tool_with_context(x: int, ctx: Context) -> str:
693
+ return f"Request {ctx.request_id}: {x}"
694
+
695
+ mcp.add_tool(tool_with_context)
696
+ async with Client(mcp) as client:
697
+ tools = await client.list_tools()
698
+ assert len(tools) == 1
699
+ assert tools[0].name == "tool_with_context"
700
+
701
+ async def test_context_injection(self):
702
+ """Test that context is properly injected into tool calls."""
703
+ mcp = FastMCP()
704
+
705
+ @mcp.tool()
706
+ def tool_with_context(x: int, ctx: Context) -> str:
707
+ assert isinstance(ctx, Context)
708
+ assert ctx.request_id is not None
709
+ return ctx.request_id
710
+
711
+ async with Client(mcp) as client:
712
+ result = await client.call_tool("tool_with_context", {"x": 42})
713
+ assert len(result) == 1
714
+ content = result[0]
715
+ assert isinstance(content, TextContent)
716
+ assert content.text == "1"
717
+
718
+ async def test_async_context(self):
719
+ """Test that context works in async functions."""
720
+ mcp = FastMCP()
721
+
722
+ async def async_tool(x: int, ctx: Context) -> str:
723
+ assert ctx.request_id is not None
724
+ return f"Async request {ctx.request_id}: {x}"
725
+
726
+ mcp.add_tool(async_tool)
727
+ async with Client(mcp) as client:
728
+ result = await client.call_tool("async_tool", {"x": 42})
729
+ assert len(result) == 1
730
+ content = result[0]
731
+ assert isinstance(content, TextContent)
732
+ assert "Async request" in content.text
733
+ assert "42" in content.text
734
+
735
+ async def test_context_logging(self):
736
+ from unittest.mock import patch
737
+
738
+ import mcp.server.session
739
+
740
+ """Test that context logging methods work."""
741
+ mcp = FastMCP()
742
+
743
+ async def logging_tool(msg: str, ctx: Context) -> str:
744
+ await ctx.debug("Debug message")
745
+ await ctx.info("Info message")
746
+ await ctx.warning("Warning message")
747
+ await ctx.error("Error message")
748
+ return f"Logged messages for {msg}"
749
+
750
+ mcp.add_tool(logging_tool)
751
+
752
+ with patch("mcp.server.session.ServerSession.send_log_message") as mock_log:
753
+ async with Client(mcp) as client:
754
+ result = await client.call_tool("logging_tool", {"msg": "test"})
755
+ assert len(result) == 1
756
+ content = result[0]
757
+ assert isinstance(content, TextContent)
758
+ assert "Logged messages for test" in content.text
759
+
760
+ assert mock_log.call_count == 4
761
+ mock_log.assert_any_call(
762
+ level="debug", data="Debug message", logger=None
763
+ )
764
+ mock_log.assert_any_call(level="info", data="Info message", logger=None)
765
+ mock_log.assert_any_call(
766
+ level="warning", data="Warning message", logger=None
767
+ )
768
+ mock_log.assert_any_call(
769
+ level="error", data="Error message", logger=None
770
+ )
771
+
772
+ async def test_optional_context(self):
773
+ """Test that context is optional."""
774
+ mcp = FastMCP()
775
+
776
+ def no_context(x: int) -> int:
777
+ return x * 2
778
+
779
+ mcp.add_tool(no_context)
780
+ async with Client(mcp) as client:
781
+ result = await client.call_tool("no_context", {"x": 21})
782
+ assert len(result) == 1
783
+ content = result[0]
784
+ assert isinstance(content, TextContent)
785
+ assert content.text == "42"
786
+
787
+ async def test_context_resource_access(self):
788
+ """Test that context can access resources."""
789
+ mcp = FastMCP()
790
+
791
+ @mcp.resource("test://data")
792
+ def test_resource() -> str:
793
+ return "resource data"
794
+
795
+ @mcp.tool()
796
+ async def tool_with_resource(ctx: Context) -> str:
797
+ r_iter = await ctx.read_resource("test://data")
798
+ r_list = list(r_iter)
799
+ assert len(r_list) == 1
800
+ r = r_list[0]
801
+ return f"Read resource: {r.content} with mime type {r.mime_type}"
802
+
803
+ async with Client(mcp) as client:
804
+ result = await client.call_tool("tool_with_resource", {})
805
+ assert len(result) == 1
806
+ content = result[0]
807
+ assert isinstance(content, TextContent)
808
+ assert "Read resource: resource data" in content.text
809
+
810
+ async def test_tool_decorator_with_tags(self):
811
+ """Test that the tool decorator properly sets tags."""
812
+ mcp = FastMCP()
813
+
814
+ @mcp.tool(tags={"example", "test-tag"})
815
+ def sample_tool(x: int) -> int:
816
+ return x * 2
817
+
818
+ # Verify the tool exists
819
+ async with Client(mcp) as client:
820
+ tools = await client.list_tools()
821
+ assert len(tools) == 1
822
+ # Note: MCPTool from the client API doesn't expose tags
823
+
824
+
825
+ class TestResource:
826
  async def test_text_resource(self):
827
  mcp = FastMCP()
828
 
 
896
  assert result[0].blob == base64.b64encode(b"Binary file data").decode()
897
 
898
 
899
+ class TestResourceContext:
900
+ async def test_resource_with_context_annotation_gets_context(self):
901
+ mcp = FastMCP()
902
+
903
+ @mcp.resource("resource://test")
904
+ def resource_with_context(ctx: Context) -> str:
905
+ assert isinstance(ctx, Context)
906
+ return ctx.request_id
907
+
908
+ async with Client(mcp) as client:
909
+ result = await client.read_resource(AnyUrl("resource://test"))
910
+ assert isinstance(result[0], TextResourceContents)
911
+ assert result[0].text == "1"
912
+
913
+
914
  class TestResourceTemplates:
915
  async def test_resource_with_params_not_in_uri(self):
916
  """Test that a resource with function parameters raises an error if the URI
 
1181
  assert result[0].text == "Template resource 1: a/b"
1182
 
1183
 
1184
+ class TestResourceTemplateContext:
1185
+ async def test_resource_template_context(self):
 
 
 
1186
  mcp = FastMCP()
1187
 
1188
+ @mcp.resource("resource://{param}")
1189
+ def resource_template(param: str, ctx: Context) -> str:
1190
+ assert isinstance(ctx, Context)
1191
+ return f"Resource template: {param} {ctx.request_id}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1192
 
 
1193
  async with Client(mcp) as client:
1194
+ result = await client.read_resource(AnyUrl("resource://test"))
1195
+ assert isinstance(result[0], TextResourceContents)
1196
+ assert result[0].text == "Resource template: test 1"
1197
 
1198
 
1199
  class TestPrompts:
 
1380
  assert len(prompts_dict) == 1
1381
  prompt = prompts_dict["sample_prompt"]
1382
  assert prompt.tags == {"example", "test-tag"}
1383
+
1384
+
1385
+ class TestPromptContext:
1386
+ async def test_prompt_context(self):
1387
+ mcp = FastMCP()
1388
+
1389
+ @mcp.prompt()
1390
+ def prompt_fn(name: str, ctx: Context) -> str:
1391
+ assert isinstance(ctx, Context)
1392
+ return f"Hello, {name}! {ctx.request_id}"
1393
+
1394
+ async with Client(mcp) as client:
1395
+ result = await client.get_prompt("prompt_fn", {"name": "World"})
1396
+ assert len(result) == 1
1397
+ message = result[0]
1398
+ assert message.role == "user"