codeocean-mcp-server 0.7.2__tar.gz → 0.12.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/CHANGELOG.md +33 -0
  2. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/PKG-INFO +18 -4
  3. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/README.md +14 -0
  4. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/pyproject.toml +5 -3
  5. codeocean_mcp_server-0.12.0/src/codeocean_mcp_server/client.py +54 -0
  6. codeocean_mcp_server-0.12.0/src/codeocean_mcp_server/search.py +140 -0
  7. codeocean_mcp_server-0.12.0/src/codeocean_mcp_server/server.py +70 -0
  8. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/tools/capsules.py +37 -31
  9. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/tools/computations.py +33 -1
  10. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/tools/data_assets.py +15 -9
  11. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/tests/bedrock_call.py +10 -2
  12. codeocean_mcp_server-0.12.0/tests/conftest.py +23 -0
  13. codeocean_mcp_server-0.12.0/tests/test_http_transport.py +140 -0
  14. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/tests/test_multi_tool_calling.py +2 -0
  15. codeocean_mcp_server-0.12.0/tests/test_search.py +275 -0
  16. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/tests/test_tool_calling.py +84 -0
  17. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/uv.lock +71 -9
  18. codeocean_mcp_server-0.7.2/src/codeocean_mcp_server/server.py +0 -41
  19. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/.gitignore +0 -0
  20. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/.python-version +0 -0
  21. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/LICENSE +0 -0
  22. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/RELEASE.md +0 -0
  23. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/__init__.py +0 -0
  24. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/file_utils.py +0 -0
  25. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/logging_config.py +0 -0
  26. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/models.py +0 -0
  27. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/tools/__init__.py +0 -0
  28. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/src/codeocean_mcp_server/tools/custom_metadata.py +0 -0
  29. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/tests/bedrock_tools_converter.py +0 -0
  30. {codeocean_mcp_server-0.7.2 → codeocean_mcp_server-0.12.0}/tests/mcp_client.py +0 -0
@@ -1,6 +1,39 @@
1
1
  CHANGELOG
2
2
  =========
3
3
 
4
+ ## 0.12.0 (2026-08-27)
5
+
6
+ - [#38](https://github.com/codeocean/codeocean-mcp-server/pull/38) feat: serve over streamable HTTP with per-request credentials
7
+
8
+ ## 0.11.0 (2026-03-09)
9
+
10
+ - [#35](https://github.com/codeocean/codeocean-mcp-server/pull/35) feat: add `delete_computation` tool to stop and delete computations
11
+
12
+ ## 0.10.0 (2026-02-24)
13
+
14
+ - [#33](https://github.com/codeocean/codeocean-mcp-server/pull/33) feat: upgrade mcp package
15
+
16
+ ## 0.9.1 (2026-02-04)
17
+
18
+ - [#31](https://github.com/codeocean/codeocean-mcp-server/pull/31) feat: clarify data asset attach tool usage
19
+
20
+ ## 0.9.0 (2026-01-30)
21
+
22
+ - [#29](https://github.com/codeocean/codeocean-mcp-server/pull/29) feat: upgrade codeocean SDK to v0.14.0
23
+ - **Minimum Code Ocean platform version updated to `4.2.0`.**
24
+
25
+ ## 0.8.2 (2026-01-29)
26
+
27
+ - [#27](https://github.com/codeocean/codeocean-mcp-server/pull/27) fix: update documentation for search results to include pagination details
28
+
29
+ ## 0.8.1 (2026-01-28)
30
+
31
+ - fix: update documentation for CapsuleSearchResults and DataAssetSearchResults to reflect item_count
32
+
33
+ ## 0.8.0 (2026-01-27)
34
+
35
+ - [#24](https://github.com/codeocean/codeocean-mcp-server/pull/24) feat: Token efficient search responses
36
+
4
37
  ## 0.7.2 (2025-12-27)
5
38
 
6
39
  - [#20](https://github.com/codeocean/codeocean-mcp-server/pull/20) feat: custom log formatting
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: codeocean-mcp-server
3
- Version: 0.7.2
3
+ Version: 0.12.0
4
4
  Summary: Code Ocean MCP Server
5
5
  Project-URL: Homepage, https://github.com/codeocean/codeocean-mcp-server
6
6
  Project-URL: Issues, https://github.com/codeocean/codeocean-mcp-server/issues
@@ -12,8 +12,8 @@ Classifier: License :: OSI Approved :: MIT License
12
12
  Classifier: Operating System :: OS Independent
13
13
  Classifier: Programming Language :: Python :: 3
14
14
  Requires-Python: >=3.10
15
- Requires-Dist: codeocean<0.13.0,>=0.12.0
16
- Requires-Dist: mcp<1.13.0,>=1.12.0
15
+ Requires-Dist: codeocean<0.15.0,>=0.14.0
16
+ Requires-Dist: mcp<1.24.0,>=1.23.0
17
17
  Description-Content-Type: text/markdown
18
18
 
19
19
  # Code Ocean MCP Server
@@ -33,6 +33,7 @@ This MCP server provides tools to search and run capsules and pipelines, and man
33
33
  - [Roo Code](#roo-code)
34
34
  - [Cursor](#cursor)
35
35
  - [Windsurf](#windsurf)
36
+ - [Streamable HTTP Transport](#streamable-http-transport)
36
37
  - [Local Testing](#local-testing)
37
38
 
38
39
  ## Prerequisites
@@ -215,6 +216,19 @@ Windsurf (Cascade) uses mcp_config.json under ~/.codeium/windsurf/ (or via the C
215
216
 
216
217
  3. Save and restart Windsurf (or hit “Refresh” in the MCP panel).
217
218
 
219
+ ## Streamable HTTP Transport
220
+
221
+ By default the server runs over stdio and authenticates with the `CODEOCEAN_TOKEN` environment variable, as
222
+ described above. It can also serve multiple users from a single process over streamable HTTP, taking each
223
+ caller's API token from the request instead:
224
+
225
+ ```bash
226
+ CODEOCEAN_DOMAIN=https://acmecorp.codeocean.com codeocean-mcp-server --transport streamable-http --host 127.0.0.1 --port 8000
227
+ ```
228
+
229
+ Clients then pass their own token as `Authorization: Bearer <YOUR_API_KEY>` on every request; `CODEOCEAN_TOKEN`
230
+ is not used, and a request without a token is refused. The endpoint is `http://<host>:<port>/mcp`.
231
+
218
232
  ## Local Testing
219
233
 
220
234
  You can test the MCP server locally during development with [MCP Inspector](https://modelcontextprotocol.io/legacy/tools/inspector):
@@ -15,6 +15,7 @@ This MCP server provides tools to search and run capsules and pipelines, and man
15
15
  - [Roo Code](#roo-code)
16
16
  - [Cursor](#cursor)
17
17
  - [Windsurf](#windsurf)
18
+ - [Streamable HTTP Transport](#streamable-http-transport)
18
19
  - [Local Testing](#local-testing)
19
20
 
20
21
  ## Prerequisites
@@ -197,6 +198,19 @@ Windsurf (Cascade) uses mcp_config.json under ~/.codeium/windsurf/ (or via the C
197
198
 
198
199
  3. Save and restart Windsurf (or hit “Refresh” in the MCP panel).
199
200
 
201
+ ## Streamable HTTP Transport
202
+
203
+ By default the server runs over stdio and authenticates with the `CODEOCEAN_TOKEN` environment variable, as
204
+ described above. It can also serve multiple users from a single process over streamable HTTP, taking each
205
+ caller's API token from the request instead:
206
+
207
+ ```bash
208
+ CODEOCEAN_DOMAIN=https://acmecorp.codeocean.com codeocean-mcp-server --transport streamable-http --host 127.0.0.1 --port 8000
209
+ ```
210
+
211
+ Clients then pass their own token as `Authorization: Bearer <YOUR_API_KEY>` on every request; `CODEOCEAN_TOKEN`
212
+ is not used, and a request without a token is refused. The endpoint is `http://<host>:<port>/mcp`.
213
+
200
214
  ## Local Testing
201
215
 
202
216
  You can test the MCP server locally during development with [MCP Inspector](https://modelcontextprotocol.io/legacy/tools/inspector):
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "codeocean-mcp-server"
3
- version = "0.7.2"
3
+ version = "0.12.0"
4
4
  authors = [{ name = "Code Ocean", email = "dev@codeocean.com" }]
5
5
  description = "Code Ocean MCP Server"
6
6
  readme = "README.md"
@@ -11,8 +11,8 @@ classifiers = [
11
11
  "Operating System :: OS Independent",
12
12
  ]
13
13
  dependencies = [
14
- "codeocean>=0.12.0,<0.13.0",
15
- "mcp>=1.12.0,<1.13.0",
14
+ "codeocean>=0.14.0,<0.15.0",
15
+ "mcp>=1.23.0,<1.24.0",
16
16
  ]
17
17
  license = "MIT"
18
18
 
@@ -66,3 +66,5 @@ filterwarnings = [
66
66
  "ignore:datetime.datetime.utcnow() is deprecated:DeprecationWarning:botocore.auth",
67
67
  "ignore::DeprecationWarning:botocore.*",
68
68
  ]
69
+ markers = ["integration: requires external services (AWS/Bedrock)"]
70
+ asyncio_mode = "strict"
@@ -0,0 +1,54 @@
1
+ """Per-request Code Ocean client resolution for the streamable-HTTP transport."""
2
+
3
+ from functools import lru_cache
4
+
5
+ from codeocean import CodeOcean
6
+ from mcp.server.fastmcp import FastMCP
7
+ from starlette.requests import Request
8
+
9
+
10
+ @lru_cache(maxsize=32)
11
+ def _cached_client(domain: str, token: str, agent_id: str | None) -> CodeOcean:
12
+ """Return a client for the given credentials, reusing its HTTP connection pool across requests."""
13
+ return CodeOcean(domain=domain, token=token, agent_id=agent_id)
14
+
15
+
16
+ def _bearer_token(request: Request) -> str | None:
17
+ """Return the token from an 'Authorization: Bearer <token>' header, or None for any other scheme."""
18
+ scheme, _, token = request.headers.get("authorization", "").partition(" ")
19
+ if scheme.lower() != "bearer":
20
+ return None
21
+ return token.strip() or None
22
+
23
+
24
+ class RequestScopedClient:
25
+ """A `CodeOcean` stand-in that resolves to the calling request's own client on every attribute access.
26
+
27
+ Tools reach the SDK as `client.<sub_client>.<method>(...)`, evaluated when the tool runs, so wrapping
28
+ the client is enough to give each request its own credentials without touching the tools themselves.
29
+
30
+ Outside a request — tool registration, where descriptions are read from SDK docstrings — attribute
31
+ access resolves to a client built from the environment's token, which issues no network call.
32
+ """
33
+
34
+ def __init__(self, mcp: FastMCP, domain: str, token: str | None, agent_id: str | None):
35
+ """Wrap the credentials that are fixed for the process, around the token that varies per request."""
36
+ self._mcp = mcp
37
+ self._domain = domain
38
+ self._token = token
39
+ self._agent_id = agent_id
40
+
41
+ def __getattr__(self, name: str):
42
+ """Delegate to the calling request's client."""
43
+ return getattr(self._resolve(), name)
44
+
45
+ def _resolve(self) -> CodeOcean:
46
+ try:
47
+ request = self._mcp.get_context().request_context.request
48
+ except ValueError:
49
+ # Registration happens outside any request, so it needs no credential of its own.
50
+ return _cached_client(self._domain, self._token or "", self._agent_id)
51
+ token = _bearer_token(request) if request is not None else None
52
+ if not token:
53
+ raise ValueError("Missing Code Ocean API token: send it as an 'Authorization: Bearer <token>' header.")
54
+ return _cached_client(self._domain, token, self._agent_id)
@@ -0,0 +1,140 @@
1
+ import re
2
+ from typing import Any, ClassVar, Optional
3
+
4
+ from pydantic import BaseModel
5
+
6
+ # Constants
7
+ MAX_DESCRIPTION_LENGTH = 200
8
+ MAX_TAGS_COUNT = 10
9
+ TRUNCATION_SUFFIX = "...(more)"
10
+ TAGS_TRUNCATION_MARKER = "..more.."
11
+
12
+
13
+ def truncate_description(description: Optional[str], max_length: int = MAX_DESCRIPTION_LENGTH) -> Optional[str]:
14
+ """Truncate description with word-boundary-aware cutting."""
15
+ if not description:
16
+ return None
17
+
18
+ # Light whitespace normalization: collapse 3+ consecutive whitespace to single space.
19
+ normalized = re.sub(r"\s{3,}", " ", description).strip()
20
+ if not normalized:
21
+ return None
22
+
23
+ if len(normalized) <= max_length:
24
+ return normalized
25
+
26
+ truncate_at = max_length - len(TRUNCATION_SUFFIX)
27
+ if truncate_at <= 0:
28
+ return TRUNCATION_SUFFIX[:max_length]
29
+
30
+ last_space = normalized.rfind(" ", 0, truncate_at)
31
+ cut_at = last_space if last_space > max_length // 2 else truncate_at
32
+ return normalized[:cut_at].rstrip() + TRUNCATION_SUFFIX
33
+
34
+
35
+ def limit_tags(tags: Optional[list[str]], max_count: int = MAX_TAGS_COUNT) -> list[str]:
36
+ """Limit tags to maximum count with truncation marker."""
37
+ if not tags:
38
+ return []
39
+
40
+ if len(tags) <= max_count:
41
+ return list(tags)
42
+
43
+ return list(tags[: max_count - 1]) + [TAGS_TRUNCATION_MARKER]
44
+
45
+
46
+ class CompactCapsuleItem(BaseModel):
47
+ """Compact capsule/pipeline item (id kept, other fields shortened)."""
48
+
49
+ id: str
50
+ n: str
51
+ s: str
52
+ d: Optional[str] = None
53
+ t: list[str]
54
+
55
+
56
+ class CompactDataAssetItem(BaseModel):
57
+ """Compact data asset item (id kept, other fields shortened)."""
58
+
59
+ id: str
60
+ n: str
61
+ d: Optional[str] = None
62
+ t: list[str]
63
+
64
+
65
+ class SearchMeta(BaseModel):
66
+ """Base model providing search metadata fields."""
67
+
68
+ has_more: bool
69
+ next_token: Optional[str] = None
70
+ item_count: int
71
+ field_names: Optional[dict[str, str]] = None
72
+
73
+
74
+ class CapsuleSearchResults(SearchMeta):
75
+ """Compact results: {items: [{id, n, s, d, t}], has_more, next_token, item_count}.
76
+
77
+ Item fields: id=id, n=name, s=slug, d=description (truncated), t=tags (limited).
78
+ Pagination: item_count returns the number of items in the current page.
79
+ Use next_token for additional pages when has_more=true.
80
+ Set include_field_names=true to add field_names with full labels.
81
+ Use get_capsule(id) if full details needed.
82
+ """
83
+
84
+ items: list[CompactCapsuleItem]
85
+ FIELD_NAMES: ClassVar[dict[str, str]] = {"id": "id", "n": "name", "s": "slug", "d": "description", "t": "tags"}
86
+
87
+ @classmethod
88
+ def from_sdk_results(cls, sdk_results: Any, include_field_names: bool = False) -> "CapsuleSearchResults":
89
+ """Convert SDK search results to compact format."""
90
+ items = [
91
+ CompactCapsuleItem(
92
+ id=c.id,
93
+ n=c.name,
94
+ s=c.slug,
95
+ d=truncate_description(c.description),
96
+ t=limit_tags(c.tags),
97
+ )
98
+ for c in sdk_results.results
99
+ ]
100
+ return cls(
101
+ items=items,
102
+ has_more=sdk_results.has_more,
103
+ next_token=getattr(sdk_results, "next_token", None),
104
+ item_count=len(items),
105
+ field_names=cls.FIELD_NAMES if include_field_names else None,
106
+ )
107
+
108
+
109
+ class DataAssetSearchResults(SearchMeta):
110
+ """Compact results: {items: [{id, n, d, t}], has_more, next_token, item_count}.
111
+
112
+ Item fields: id=id, n=name, d=description (truncated), t=tags (limited).
113
+ Pagination: item_count returns the number of items in the current page.
114
+ Use next_token for additional pages when has_more=true.
115
+ Set include_field_names=true to add field_names with full labels.
116
+ Use get_data_asset(id) if full details needed.
117
+ """
118
+
119
+ items: list[CompactDataAssetItem]
120
+ FIELD_NAMES: ClassVar[dict[str, str]] = {"id": "id", "n": "name", "d": "description", "t": "tags"}
121
+
122
+ @classmethod
123
+ def from_sdk_results(cls, sdk_results: Any, include_field_names: bool = False) -> "DataAssetSearchResults":
124
+ """Convert SDK search results to compact format."""
125
+ items = [
126
+ CompactDataAssetItem(
127
+ id=d.id,
128
+ n=d.name,
129
+ d=truncate_description(d.description),
130
+ t=limit_tags(d.tags),
131
+ )
132
+ for d in sdk_results.results
133
+ ]
134
+ return cls(
135
+ items=items,
136
+ has_more=sdk_results.has_more,
137
+ next_token=getattr(sdk_results, "next_token", None),
138
+ item_count=len(items),
139
+ field_names=cls.FIELD_NAMES if include_field_names else None,
140
+ )
@@ -0,0 +1,70 @@
1
+ import argparse
2
+ import os
3
+
4
+ from codeocean import CodeOcean
5
+ from mcp.server.fastmcp import FastMCP
6
+
7
+ from codeocean_mcp_server.client import RequestScopedClient
8
+ from codeocean_mcp_server.logging_config import configure_logging
9
+ from codeocean_mcp_server.tools import (
10
+ capsules,
11
+ computations,
12
+ custom_metadata,
13
+ data_assets,
14
+ )
15
+
16
+
17
+ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
18
+ """Parse command line arguments."""
19
+ parser = argparse.ArgumentParser(prog="codeocean-mcp-server", description="Code Ocean MCP Server")
20
+ parser.add_argument(
21
+ "--transport",
22
+ choices=["stdio", "streamable-http"],
23
+ default="stdio",
24
+ help="Transport to serve on. Over streamable-HTTP the API token is taken per request from the "
25
+ "'Authorization: Bearer <token>' header instead of from CODEOCEAN_TOKEN.",
26
+ )
27
+ parser.add_argument("--host", default="127.0.0.1", help="Host to bind to when serving over streamable-HTTP.")
28
+ parser.add_argument("--port", type=int, default=8000, help="Port to bind to when serving over streamable-HTTP.")
29
+ return parser.parse_args(argv)
30
+
31
+
32
+ def main():
33
+ """Run the MCP server."""
34
+ configure_logging()
35
+ args = parse_args()
36
+ stdio = args.transport == "stdio"
37
+ domain = os.getenv("CODEOCEAN_DOMAIN")
38
+ token = os.getenv("CODEOCEAN_TOKEN")
39
+ if not domain:
40
+ raise ValueError("Environment variable CODEOCEAN_DOMAIN must be set.")
41
+ agent_id = os.getenv("AGENT_ID", "AI Agent")
42
+
43
+ mcp = FastMCP(
44
+ name="Code Ocean",
45
+ instructions=(
46
+ f"MCP server for Code Ocean: search & run capsules, pipelines, and assets using Code Ocean domain {domain}."
47
+ ),
48
+ host=args.host,
49
+ port=args.port,
50
+ stateless_http=True,
51
+ )
52
+
53
+ if stdio:
54
+ # Over stdio the process serves a single user, so the environment's client is used directly, as before.
55
+ if not token:
56
+ raise ValueError("Environment variable CODEOCEAN_TOKEN must be set when serving over stdio.")
57
+ client = CodeOcean(domain=domain, token=token, agent_id=agent_id)
58
+ else:
59
+ client = RequestScopedClient(mcp, domain, token, agent_id)
60
+
61
+ capsules.add_tools(mcp, client)
62
+ data_assets.add_tools(mcp, client)
63
+ computations.add_tools(mcp, client)
64
+ custom_metadata.add_tools(mcp, client)
65
+
66
+ mcp.run(transport=args.transport)
67
+
68
+
69
+ if __name__ == "__main__":
70
+ main()
@@ -3,7 +3,6 @@ from codeocean.capsule import (
3
3
  AppPanel,
4
4
  Capsule,
5
5
  CapsuleSearchParams,
6
- CapsuleSearchResults,
7
6
  Computation,
8
7
  DataAssetAttachParams,
9
8
  DataAssetAttachResults,
@@ -11,6 +10,7 @@ from codeocean.capsule import (
11
10
  from mcp.server.fastmcp import FastMCP
12
11
 
13
12
  from codeocean_mcp_server.models import dataclass_to_pydantic
13
+ from codeocean_mcp_server.search import CapsuleSearchResults
14
14
 
15
15
  AppPanelModel = dataclass_to_pydantic(AppPanel)
16
16
  CapsuleSearchParamsModel = dataclass_to_pydantic(CapsuleSearchParams)
@@ -20,37 +20,47 @@ DataAssetAttachParamsModel = dataclass_to_pydantic(DataAssetAttachParams)
20
20
  def add_tools(mcp: FastMCP, client: CodeOcean):
21
21
  """Add capsule tools to the MCP server."""
22
22
 
23
- @mcp.tool(
24
- description=(
25
- str(client.capsules.search_capsules.__doc__) + "Use only for capsule searches. "
26
- "Provide only the minimal required parameters (e.g. limit=10); "
27
- "do not include optional params "
28
- "like sort_by or sort_order unless requested."
29
- )
30
- )
31
- def search_capsules(search_params: CapsuleSearchParamsModel) -> CapsuleSearchResults:
23
+ @mcp.tool(description=(str(client.capsules.search_capsules.__doc__) + " " + str(CapsuleSearchResults.__doc__)))
24
+ def search_capsules(
25
+ search_params: CapsuleSearchParamsModel,
26
+ include_field_names: bool = False,
27
+ ) -> CapsuleSearchResults:
32
28
  """Search for capsules matching specified criteria."""
33
29
  params = CapsuleSearchParams(**search_params.model_dump(exclude_none=True))
34
- return client.capsules.search_capsules(params)
30
+ results = client.capsules.search_capsules(params)
31
+ return CapsuleSearchResults.from_sdk_results(results, include_field_names)
32
+
33
+ @mcp.tool(description=(str(client.pipelines.search_pipelines.__doc__) + " " + str(CapsuleSearchResults.__doc__)))
34
+ def search_pipelines(
35
+ search_params: CapsuleSearchParamsModel,
36
+ include_field_names: bool = False,
37
+ ) -> CapsuleSearchResults:
38
+ """Search for pipelines matching specified criteria."""
39
+ params = CapsuleSearchParams(**search_params.model_dump(exclude_none=True))
40
+ results = client.pipelines.search_pipelines(params)
41
+ return CapsuleSearchResults.from_sdk_results(results, include_field_names)
35
42
 
36
43
  @mcp.tool(
37
44
  description=(
38
- str(client.capsules.search_pipelines.__doc__) + "Use only for pipeline searches. "
39
- "Provide only the minimal required parameters (e.g. limit=10); "
40
- "do not include optional params "
41
- "like sort_by or sort_order unless requested."
45
+ str(client.capsules.get_capsule.__doc__) + "Use only to fetch metadata for a known capsule ID. "
46
+ "Do not use for searching."
42
47
  )
43
48
  )
44
- def search_pipelines(search_params: CapsuleSearchParamsModel) -> CapsuleSearchResults:
45
- """Search for pipelines matching specified criteria."""
46
- params = CapsuleSearchParams(**search_params.model_dump(exclude_none=True))
47
- return client.capsules.search_pipelines(params)
49
+ def get_capsule(capsule_id: str) -> Capsule:
50
+ """Retrieve a capsule by its ID."""
51
+ return client.capsules.get_capsule(capsule_id)
52
+
53
+ @mcp.tool(description=client.capsules.list_computations.__doc__)
54
+ def list_computations(capsule_id: str) -> list[Computation]:
55
+ """List all computations for a capsule."""
56
+ return client.capsules.list_computations(capsule_id)
48
57
 
49
58
  @mcp.tool(
50
59
  description=(
51
60
  str(client.capsules.attach_data_assets.__doc__)
52
- + "Accepts a list of parameter objects (e.g. [{'id': '...'}]), "
53
- "not just a list of IDs."
61
+ + " Use when the capsule has no running cloud workstations. For active cloud workstation sessions"
62
+ " use attach_computation_data_assets instead."
63
+ " Accepts a list of parameter objects (e.g. [{'id': '...'}]), not just a list of IDs."
54
64
  )
55
65
  )
56
66
  def attach_data_assets(
@@ -63,18 +73,14 @@ def add_tools(mcp: FastMCP, client: CodeOcean):
63
73
 
64
74
  @mcp.tool(
65
75
  description=(
66
- str(client.capsules.get_capsule.__doc__) + "Use only to fetch metadata for a known capsule ID. "
67
- "Do not use for searching."
76
+ str(client.capsules.detach_data_assets.__doc__)
77
+ + " Use when the capsule has no running cloud workstations. For active cloud workstation sessions"
78
+ " use detach_computation_data_assets instead."
68
79
  )
69
80
  )
70
- def get_capsule(capsule_id: str) -> Capsule:
71
- """Retrieve a capsule by its ID."""
72
- return client.capsules.get_capsule(capsule_id)
73
-
74
- @mcp.tool(description=client.capsules.list_computations.__doc__)
75
- def list_computations(capsule_id: str) -> list[Computation]:
76
- """List all computations for a capsule."""
77
- return client.capsules.list_computations(capsule_id)
81
+ def detach_data_assets(capsule_id: str, data_assets: list[str]) -> None:
82
+ """Remove attached data assets from a capsule."""
83
+ client.capsules.detach_data_assets(capsule_id, data_assets)
78
84
 
79
85
  @mcp.tool(description=client.capsules.get_capsule_app_panel.__doc__)
80
86
  def get_capsule_app_panel(capsule_id: str, version: int | None = None) -> AppPanelModel:
@@ -5,15 +5,17 @@ from codeocean.computation import (
5
5
  Folder,
6
6
  RunParams,
7
7
  )
8
+ from codeocean.data_asset import DataAssetAttachParams, DataAssetAttachResults
8
9
  from mcp.server.fastmcp import FastMCP
9
10
 
10
11
  from codeocean_mcp_server.file_utils import download_and_read_file
11
12
  from codeocean_mcp_server.models import dataclass_to_pydantic
12
13
 
13
14
  RunParamsModel = dataclass_to_pydantic(RunParams)
15
+ DataAssetAttachParamsModel = dataclass_to_pydantic(DataAssetAttachParams)
14
16
 
15
17
 
16
- def add_tools(mcp: FastMCP, client: CodeOcean):
18
+ def add_tools(mcp: FastMCP, client: CodeOcean): # noqa: C901
17
19
  """Add capsule tools to the MCP server."""
18
20
 
19
21
  @mcp.tool(description=client.computations.get_computation.__doc__)
@@ -59,3 +61,33 @@ def add_tools(mcp: FastMCP, client: CodeOcean):
59
61
  """Download a file using the provided URL and return its content."""
60
62
  file_urls = client.computations.get_result_file_urls(computation_id, file_path)
61
63
  return download_and_read_file(file_urls.download_url)
64
+
65
+ @mcp.tool(description=client.computations.rename_computation.__doc__)
66
+ def rename_computation(computation_id: str, name: str) -> None:
67
+ """Rename an existing computation."""
68
+ client.computations.rename_computation(computation_id, name)
69
+
70
+ @mcp.tool(description=client.computations.delete_computation.__doc__)
71
+ def delete_computation(computation_id: str) -> None:
72
+ """Delete a computation and stop it if currently running."""
73
+ client.computations.delete_computation(computation_id)
74
+
75
+ @mcp.tool(
76
+ description=(
77
+ str(client.computations.attach_data_assets.__doc__)
78
+ + " Accepts a list of parameter objects (e.g. [{'id': '...'}]). "
79
+ "Use for cloud workstation sessions."
80
+ )
81
+ )
82
+ def attach_computation_data_assets(
83
+ computation_id: str,
84
+ attach_params: list[DataAssetAttachParamsModel],
85
+ ) -> list[DataAssetAttachResults]:
86
+ """Attach data assets to a cloud workstation session."""
87
+ params = [DataAssetAttachParams(**p.model_dump(exclude_none=True)) for p in attach_params]
88
+ return client.computations.attach_data_assets(computation_id, params)
89
+
90
+ @mcp.tool(description=client.computations.detach_data_assets.__doc__)
91
+ def detach_computation_data_assets(computation_id: str, data_assets: list[str]) -> None:
92
+ """Detach data assets from a cloud workstation session."""
93
+ client.computations.detach_data_assets(computation_id, data_assets)
@@ -5,7 +5,6 @@ from codeocean.data_asset import (
5
5
  DataAsset,
6
6
  DataAssetParams,
7
7
  DataAssetSearchParams,
8
- DataAssetSearchResults,
9
8
  DataAssetUpdateParams,
10
9
  FileURLs,
11
10
  Folder,
@@ -14,6 +13,7 @@ from mcp.server.fastmcp import FastMCP
14
13
 
15
14
  from codeocean_mcp_server.file_utils import download_and_read_file
16
15
  from codeocean_mcp_server.models import dataclass_to_pydantic
16
+ from codeocean_mcp_server.search import DataAssetSearchResults
17
17
 
18
18
  DataAssetModel = dataclass_to_pydantic(DataAsset)
19
19
  DataAssetParamsModel = dataclass_to_pydantic(DataAssetParams)
@@ -25,17 +25,23 @@ def add_tools(mcp: FastMCP, client: CodeOcean):
25
25
  """Add data asset tools to the MCP server."""
26
26
 
27
27
  @mcp.tool(
28
- description=(
29
- str(client.data_assets.search_data_assets.__doc__)
30
- + "Search for data assets (external or internal). You may filter by "
31
- "fields such as `origin`, tags, or other criteria supported by the "
32
- "SDK."
33
- )
28
+ description=(str(client.data_assets.search_data_assets.__doc__) + " " + str(DataAssetSearchResults.__doc__))
34
29
  )
35
- def search_data_assets(search_params: DataAssetSearchParamsModel) -> DataAssetSearchResults:
30
+ def search_data_assets(
31
+ search_params: DataAssetSearchParamsModel,
32
+ include_field_names: bool = False,
33
+ ) -> DataAssetSearchResults:
36
34
  """Retrieve data assets matching search criteria for datasets."""
37
35
  params = DataAssetSearchParams(**search_params.model_dump(exclude_none=True))
38
- return client.data_assets.search_data_assets(params)
36
+ results = client.data_assets.search_data_assets(params)
37
+ return DataAssetSearchResults.from_sdk_results(results, include_field_names)
38
+
39
+ @mcp.tool(
40
+ description=("Get full details for a data asset by ID. Use after compact search to retrieve complete metadata.")
41
+ )
42
+ def get_data_asset(data_asset_id: str) -> DataAsset:
43
+ """Retrieve a data asset by its ID."""
44
+ return client.data_assets.get_data_asset(data_asset_id)
39
45
 
40
46
  @mcp.tool(
41
47
  description=(
@@ -7,7 +7,15 @@ from bedrock_tools_converter import convert_tool_format
7
7
 
8
8
  BEDROCK_MODEL = os.getenv("BEDROCK_MODEL") or "amazon.nova-pro-v1:0"
9
9
 
10
- client = boto3.client("bedrock-runtime")
10
+ _client = None
11
+
12
+
13
+ def _get_client():
14
+ """Lazily initialize the Bedrock client."""
15
+ global _client
16
+ if _client is None:
17
+ _client = boto3.client("bedrock-runtime")
18
+ return _client
11
19
 
12
20
 
13
21
  def call_bedrock(
@@ -26,4 +34,4 @@ def call_bedrock(
26
34
  if tools is not None:
27
35
  payload["toolConfig"] = convert_tool_format(tools, model)
28
36
 
29
- return client.converse(**payload)
37
+ return _get_client().converse(**payload)
@@ -0,0 +1,23 @@
1
+ """Pytest configuration for codeocean-mcp-server tests."""
2
+
3
+ import pytest
4
+
5
+
6
+ def pytest_addoption(parser):
7
+ """Add --integration flag to pytest."""
8
+ parser.addoption(
9
+ "--integration",
10
+ action="store_true",
11
+ default=False,
12
+ help="Run integration tests (requires AWS/Bedrock setup)",
13
+ )
14
+
15
+
16
+ def pytest_collection_modifyitems(config, items):
17
+ """Skip integration tests unless --integration flag is passed."""
18
+ if config.getoption("--integration"):
19
+ return
20
+ skip_integration = pytest.mark.skip(reason="need --integration flag to run")
21
+ for item in items:
22
+ if "integration" in item.keywords:
23
+ item.add_marker(skip_integration)