Skip to content

Commit 6dd5a4f

Browse files
phernandezclaude
andcommitted
refactor: replace project_session with context-based project management
- Replace global project_session with FastMCP context-based architecture - Add project_context.py with get_active_project/set_active_project functions - Add ProjectContextMiddleware for session persistence in local mode - Update all MCP tools to use context parameter and new project functions - Remove project_session.py and replace with API-driven approach - Update test fixtures to work with new context architecture This enables: - Load balancing support for cloud deployment (context can be stored in Redis) - Consistent behavior between stdio and streamable-http transports - Session isolation and thread safety - Stateless compatibility while maintaining project switching 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 838859a commit 6dd5a4f

33 files changed

Lines changed: 417 additions & 647 deletions

src/basic_memory/api/routers/project_router.py

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
)
1515

1616
# Router for resources in a specific project
17+
# The ProjectPathDep is used in the path as a prefix, so the request path is like /{project}/project/info
1718
project_router = APIRouter(prefix="/project", tags=["project"])
1819

1920
# Router for managing project resources
@@ -28,6 +29,20 @@ async def get_project_info(
2829
"""Get comprehensive information about the specified Basic Memory project."""
2930
return await project_service.get_project_info(project)
3031

32+
@project_router.get("/item", response_model=ProjectItem)
33+
async def get_project(
34+
project_service: ProjectServiceDep,
35+
project: ProjectPathDep,
36+
) -> ProjectItem:
37+
"""Get bassic info about the specified Basic Memory project."""
38+
found_project = await project_service.get_project(project)
39+
if not found_project:
40+
raise HTTPException(
41+
status_code=404, detail=f"Project: '{project}' does not exist"
42+
) # pragma: no cover
43+
44+
return ProjectItem(name=found_project.name, path=found_project.path, is_default=found_project.is_default or False)
45+
3146

3247
# Update a project
3348
@project_router.patch("/{name}", response_model=ProjectStatusResponse)
@@ -216,6 +231,26 @@ async def set_default_project(
216231
except ValueError as e: # pragma: no cover
217232
raise HTTPException(status_code=400, detail=str(e))
218233

234+
# Get the default project
235+
@project_resource_router.get("/default", response_model=ProjectItem)
236+
async def get_default_project(
237+
project_service: ProjectServiceDep,
238+
) -> ProjectItem:
239+
"""Get the default project.
240+
241+
Returns:
242+
Response with project default information
243+
"""
244+
# Get the old default project
245+
default_name = project_service.default_project
246+
default_project = await project_service.get_project(default_name)
247+
if not default_project: # pragma: no cover
248+
raise HTTPException( # pragma: no cover
249+
status_code=404, detail=f"Default Project: '{default_name}' does not exist"
250+
)
251+
252+
return ProjectItem(name=default_project.name, path=default_project.path, is_default=True)
253+
219254

220255
# Synchronize projects between config and database
221256
@project_resource_router.post("/sync", response_model=ProjectStatusResponse)

src/basic_memory/cli/app.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
import typer
44

55
from basic_memory.config import get_project_config, ConfigManager
6-
from basic_memory.mcp.project_session import session
76

87

98
def version_callback(value: bool) -> None:
@@ -28,7 +27,7 @@ def app_callback(
2827
None,
2928
"--project",
3029
"-p",
31-
help="Specify which project to use 1",
30+
help="Specify which project to use",
3231
envvar="BASIC_MEMORY_PROJECT",
3332
),
3433
version: Optional[bool] = typer.Option(
@@ -51,18 +50,11 @@ def app_callback(
5150

5251
# Initialize MCP session with the specified project or default
5352
if project: # pragma: no cover
54-
# Use the project specified via --project flag
55-
current_project_config = get_project_config(project)
56-
session.set_current_project(current_project_config.name)
57-
5853
# Update the global config to use this project
5954
from basic_memory.config import update_current_project
6055

56+
# TODO set active project via cli
6157
update_current_project(project)
62-
else:
63-
# Use the default project
64-
current_project = app_config.default_project
65-
session.set_current_project(current_project)
6658

6759

6860
# Register sub-command groups

src/basic_memory/cli/commands/mcp.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
from basic_memory.cli.app import app
77
from basic_memory.config import ConfigManager
8+
from basic_memory.mcp.middleware import ProjectContextMiddleware
89

910
# Import mcp instance
1011
from basic_memory.mcp.server import mcp as mcp_server # pragma: no cover
@@ -63,6 +64,10 @@ def run_file_sync():
6364
# Now run the MCP server (blocks)
6465
logger.info(f"Starting MCP server with {transport.upper()} transport")
6566

67+
# middleware to store the active project
68+
session_state_middleware = ProjectContextMiddleware()
69+
mcp_server.add_middleware(session_state_middleware)
70+
6671
if transport == "stdio":
6772
mcp_server.run(
6873
transport=transport,

src/basic_memory/cli/commands/project.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from rich.table import Table
1010

1111
from basic_memory.cli.app import app
12-
from basic_memory.mcp.project_session import session
12+
from basic_memory.mcp.project_context import get_active_project
1313
from basic_memory.mcp.resources.project_info import project_info
1414
import json
1515
from datetime import datetime
@@ -52,13 +52,11 @@ def list_projects() -> None:
5252
table = Table(title="Basic Memory Projects")
5353
table.add_column("Name", style="cyan")
5454
table.add_column("Path", style="green")
55-
table.add_column("Default", style="yellow")
56-
table.add_column("Active", style="magenta")
55+
table.add_column("Default", style="magenta")
5756

5857
for project in result.projects:
5958
is_default = "✓" if project.is_default else ""
60-
is_active = "✓" if session.get_current_project() == project.name else ""
61-
table.add_row(project.name, format_path(project.path), is_default, is_active)
59+
table.add_row(project.name, format_path(project.path), is_default)
6260

6361
console.print(table)
6462
except Exception as e:
@@ -161,9 +159,10 @@ def move_project(
161159
data = {"path": resolved_path}
162160

163161
project_permalink = generate_permalink(name)
164-
current_project = session.get_current_project()
162+
163+
# TODO fix route to use ProjectPathDep
165164
response = asyncio.run(
166-
call_patch(client, f"/{current_project}/project/{project_permalink}", json=data)
165+
call_patch(client, f"/{name}/project/{project_permalink}", json=data)
167166
)
168167
result = ProjectStatusResponse.model_validate(response.json())
169168

src/basic_memory/mcp/middleware.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
from typing import Any
2+
from loguru import logger
3+
from fastmcp.server.middleware import Middleware, MiddlewareContext
4+
5+
from basic_memory.config import ConfigManager
6+
from basic_memory.schemas.project_info import ProjectItem
7+
8+
9+
class ProjectContextMiddleware(Middleware):
10+
11+
def __init__(self):
12+
config = ConfigManager().config
13+
default_project_name = config.default_project
14+
default_project_path = config.projects[default_project_name]
15+
16+
self.default_project = ProjectItem(name=default_project_name, path=default_project_path, is_default=True)
17+
self.sessions: dict[str, ProjectItem] = {}
18+
19+
async def on_call_tool(self, context: MiddlewareContext, call_next) -> Any:
20+
# Initialize context with default project if not set
21+
if context.fastmcp_context:
22+
session_id = context.fastmcp_context.session_id
23+
project = self.sessions.get(session_id)
24+
25+
active_project = project or self.default_project
26+
27+
# store in sessions
28+
if not project:
29+
self.sessions[session_id] = active_project
30+
31+
# set active project on context
32+
context.fastmcp_context.set_state("active_project", active_project)
33+
logger.debug(
34+
f"project context set: session_id={session_id}, active_project={active_project}"
35+
)
36+
37+
# call tool
38+
result = await call_next(context)
39+
40+
# set global state for if the project changed during tool call
41+
if context.fastmcp_context:
42+
session_id = context.fastmcp_context.session_id
43+
active_project = context.fastmcp_context.get_state("active_project")
44+
project = self.sessions.get(session_id)
45+
if project != active_project:
46+
47+
logger.debug(f"project session saved: session_id={session_id}, active_project={active_project}")
48+
self.sessions[session_id] = active_project
49+
50+
return result
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
"""Project session management for Basic Memory MCP server.
2+
3+
Provides simple in-memory project context for MCP tools, allowing users to switch
4+
between projects during a conversation without restarting the server.
5+
6+
Session Persistence Flow:
7+
1. Pre-tool: Middleware loads self.sessions[session_id] → sets context state
8+
2. Tool execution: Uses context.get_state("active_project")
9+
3. Project switching: Tool calls set_active_project() → updates context state
10+
4. Post-tool: Middleware detects context change → saves to self.sessions[session_id]
11+
"""
12+
13+
from typing import Optional
14+
from httpx import AsyncClient
15+
from loguru import logger
16+
17+
from fastmcp import Context
18+
19+
from basic_memory.mcp.tools.utils import call_get
20+
from basic_memory.schemas.project_info import ProjectItem
21+
from basic_memory.utils import generate_permalink
22+
23+
24+
async def set_active_project(client: AsyncClient, *, context: Context | None, project: ProjectItem) -> None:
25+
"""Set the active project context.
26+
27+
Args:
28+
client (AsyncClient): The client to use for making requests to fastapi.
29+
project_name: The project to switch to
30+
context: Optional FastMCP context containing project session state
31+
"""
32+
active_project = await get_active_project(client, context= context)
33+
if active_project.name != project.name:
34+
previous_project = active_project
35+
# set project in the context
36+
if context:
37+
context.set_state("active_project", project)
38+
await context.info(f"Context {context.session_id} Switched active project: {previous_project} -> {project.name}")
39+
logger.info(f"Switched active project: {previous_project} -> {project.name}")
40+
else:
41+
logger.debug(f"No change for active project: {active_project}")
42+
43+
44+
async def get_active_project(client: AsyncClient, *, context: Context | None, project_override: Optional[str] = None) -> ProjectItem:
45+
"""
46+
Get the active project for a tool call.
47+
If no context is provided, the active project is returned from the default project
48+
49+
Args:
50+
client (AsyncClient): The client to use for making requests to fastapi.
51+
project_override: Optional explicit project name from tool parameter
52+
context: Optional FastMCP context containing project session state
53+
54+
Returns:
55+
The project to use for the tool call
56+
"""
57+
58+
# Try to get active project from context first
59+
if context:
60+
active_project = context.get_state("active_project")
61+
logger.debug(f"Context {context.session_id} found active project: {active_project}")
62+
else:
63+
response = await call_get(
64+
client,
65+
f"/projects/default",
66+
)
67+
active_project = ProjectItem.model_validate(response.json())
68+
logger.debug(f"No context provided. Using default project: {active_project}")
69+
70+
# Handle project override
71+
if project_override:
72+
logger.debug(f"overriding active project: {project_override}")
73+
permalink = generate_permalink(project_override)
74+
response = await call_get(
75+
client,
76+
f"/{permalink}/item",
77+
)
78+
active_project = ProjectItem.model_validate(response.json())
79+
80+
return active_project
81+
82+
83+
def add_project_metadata(result: str, project_name: str) -> str:
84+
"""Add project context as metadata footer for LLM awareness.
85+
86+
Args:
87+
result: The tool result string
88+
project_name: The project name that was used
89+
90+
Returns:
91+
Result with project metadata footer
92+
"""
93+
return f"{result}\n\n<!-- Project: {project_name} -->" # pragma: no cover

0 commit comments

Comments
 (0)