Fix non-Docket background tasks for MCP servers #101
Conversation
There was a problem hiding this comment.
Pull request overview
This PR fixes a critical bug where background tasks were not executing properly in MCP (Model Context Protocol) server contexts. The root issue was that HybridBackgroundTasks was using FastAPI's BackgroundTasks.add_task() for non-Docket mode, which relies on Starlette's response lifecycle that doesn't exist in MCP's stdio/SSE communication modes.
Key changes:
- Refactored
HybridBackgroundTasksto useasyncio.create_task()instead of parent class'sadd_task()whenuse_docket=False, enabling tasks to execute in both FastAPI and MCP contexts - Simplified CLI logic for setting
use_docketflag - Updated tests to verify actual task execution rather than queue inspection
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| agent_memory_server/dependencies.py | Replaced FastAPI's response-bound background task mechanism with asyncio.create_task() and added _run_task() helper to handle both sync and async functions |
| agent_memory_server/cli.py | Simplified use_docket setting logic in MCP command (note: contains a critical bug) |
| tests/test_dependencies.py | Updated tests to verify tasks actually execute asynchronously rather than checking if they're queued in FastAPI's task list |
| agent_memory_server/init.py | Version bump to 0.12.6 |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 7 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
docs/quick-start.md
Outdated
| # Start the API server in development mode (runs on port 8000) | ||
| uv run agent-memory api --no-worker | ||
| # Start the API server in development mode (runs on port 8000, asyncio backend) | ||
| uv run agent-memory api |
There was a problem hiding this comment.
The comment says "Start the API server in development mode (runs on port 8000, asyncio backend)" but the command uv run agent-memory api does not include --task-backend=asyncio. Since the CLI default for the api command is docket (see cli.py line 281), this will use the Docket backend by default, which requires running a separate agent-memory task-worker process. This contradicts the expectation set by the comment. Either add --task-backend=asyncio to the command or update the comment to clarify that a worker is needed.
| uv run agent-memory api | |
| uv run agent-memory api --task-backend=asyncio |
| # Start the server (development mode) | ||
| uv run agent-memory api --no-worker | ||
| # Start the server (development mode, default asyncio backend) | ||
| uv run agent-memory api |
There was a problem hiding this comment.
The comment says "Start the server (development mode, default asyncio backend)" but the command uv run agent-memory api does not include --task-backend=asyncio. Since the CLI default for the api command is docket (see cli.py line 281), this will use the Docket backend by default, which requires running a separate worker. This contradicts the expectation set by the comment.
| uv run agent-memory api | |
| uv run agent-memory api --task-backend=asyncio |
| "Background task backend (asyncio, docket). " | ||
| "Default is 'asyncio' (no separate worker needed). " | ||
| "Use 'docket' for production setups with a running task worker." | ||
| "(see `agent-memory task-worker`)." |
There was a problem hiding this comment.
There's a missing space after the period in the help text. The text "(see agent-memory task-worker)." should have a space before the opening parenthesis. It should read "Use 'docket' for production setups with a running task worker (see agent-memory task-worker)."
| "(see `agent-memory task-worker`)." | |
| " (see `agent-memory task-worker`)." |
| # relies on being attached to a response object and run after the response | ||
| # is sent. In MCP mode (stdio/SSE), there's no Starlette response lifecycle, | ||
| # so tasks added via super().add_task() would never execute. | ||
| asyncio.create_task(self._run_task(func, *args, **kwargs)) |
There was a problem hiding this comment.
Calling asyncio.create_task() from a synchronous method (add_task) requires a running event loop. This will raise a RuntimeError: no running event loop if called outside of an async context. While this works fine in FastAPI (which runs in an async context) and MCP servers (also async), it could fail if add_task is called from a synchronous context, such as a synchronous endpoint or startup code. Consider using asyncio.get_event_loop().create_task() or wrapping in a try-except with a fallback to asyncio.ensure_future() for better compatibility, or document that this method must only be called from async contexts.
| asyncio.create_task(self._run_task(func, *args, **kwargs)) | |
| try: | |
| asyncio.create_task(self._run_task(func, *args, **kwargs)) | |
| except RuntimeError: | |
| # No running event loop; fallback to ensure_future (may create one) | |
| asyncio.ensure_future(self._run_task(func, *args, **kwargs)) |
| # Default to development mode using the API's default backend (Docket). For | ||
| # single-process development without a worker, add `--task-backend=asyncio` to | ||
| # the api command. | ||
| # Examples: | ||
| # Development: docker run -p 8000:8000 redislabs/agent-memory-server | ||
| # Development: docker run -p 8000:8000 redislabs/agent-memory-server agent-memory api --host 0.0.0.0 --port 8000 --task-backend=asyncio | ||
| # Production API: docker run -p 8000:8000 redislabs/agent-memory-server agent-memory api --host 0.0.0.0 --port 8000 | ||
| # Production Worker: docker run redislabs/agent-memory-server agent-memory task-worker --concurrency 10 | ||
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000", "--no-worker"] | ||
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000"] |
There was a problem hiding this comment.
The comment states "Default to development mode using the API's default backend (Docket)" but this is misleading. The CMD actually runs with Docket backend (which requires a separate worker), not asyncio. This contradicts the README.md line 37 which claims "The default image runs in development mode using the asyncio task backend (no separate worker required)".
Either the CMD should be updated to include --task-backend=asyncio to match the README's claim, or the comment and README should be updated to clarify that users need to run a separate worker when using the default Docker image.
| # Production API: docker run -p 8000:8000 redislabs/agent-memory-server:aws agent-memory api --host 0.0.0.0 --port 8000 | ||
| # Production Worker: docker run redislabs/agent-memory-server:aws agent-memory task-worker --concurrency 10 | ||
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000", "--no-worker"] | ||
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000"] |
There was a problem hiding this comment.
The comment states "Default to development mode using the API's default backend (Docket)" but this is misleading. The CMD actually runs with Docket backend (which requires a separate worker), not asyncio. This contradicts the README.md line 37 which claims "The default image runs in development mode using the asyncio task backend (no separate worker required)".
Either the CMD should be updated to include --task-backend=asyncio to match the README's claim, or the comment and README should be updated to clarify that users need to run a separate worker when using the default Docker image.
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000"] | |
| CMD ["agent-memory", "api", "--host", "0.0.0.0", "--port", "8000", "--task-backend=asyncio"] |
README.md
Outdated
| ``` | ||
|
|
||
| The default image runs in development mode (`--no-worker`), which is perfect for testing and development. | ||
| The default image runs in development mode using the **asyncio** task backend (no separate worker required), which is perfect for testing and development. |
There was a problem hiding this comment.
This line claims "The default image runs in development mode using the asyncio task backend (no separate worker required)" but the Dockerfile's CMD at line 122 and 153 does not include --task-backend=asyncio. Since the CLI default for api command is docket (see cli.py line 281), the default Docker image will actually use Docket backend and require a separate worker to process background tasks. This is misleading to users who expect a single-container development setup.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Start the server (development mode, default asyncio backend) | ||
| uv run agent-memory api |
There was a problem hiding this comment.
The README states "Start the server (development mode, default asyncio backend)" but the actual default for the API command is 'docket' (as shown in cli.py line 281). When running uv run agent-memory api without any flags, it will use the Docket backend and require a separate worker, not the asyncio backend. The comment should either be corrected to say "Start the server (development mode)" or the command should be changed to explicitly use --task-backend asyncio.
| # Start the server (development mode, default asyncio backend) | |
| uv run agent-memory api | |
| # Start the server (development mode, asyncio backend) | |
| uv run agent-memory api --task-backend=asyncio |
| """Test api command with default parameters uses Docket backend. | ||
| By default, we preserve existing behavior by enabling Docket-based | ||
| background tasks (settings.use_docket should be True). | ||
| """ | ||
| from agent_memory_server.config import settings | ||
|
|
||
| # Set initial state | ||
| settings.use_docket = True | ||
| settings.use_docket = False | ||
|
|
||
| runner = CliRunner() | ||
| result = runner.invoke(api) | ||
|
|
||
| assert result.exit_code == 0 | ||
| # Should not change use_docket when --no-worker is not specified | ||
| # Default should enable Docket-based background tasks | ||
| assert settings.use_docket is True |
There was a problem hiding this comment.
The test name says "test_api_command_defaults" and the docstring mentions "uses Docket backend" and expects "settings.use_docket is True", but the test sets the initial state to settings.use_docket = False on line 125. This is testing that the default changes the setting from False to True, which is correct, but the initial state setup comment or value seems counterintuitive. Consider either removing the initial state setup (letting it be whatever it is) or adding a comment explaining why it starts as False to verify the command properly enables Docket.
| ## Production Deployment | ||
|
|
||
| The examples above use `--no-worker` for development convenience. For production environments, you should use Docket workers for better reliability, scalability, and performance. | ||
| The examples above use asyncio task backends for simple, single-process development. For production environments, the `api` command defaults to the **Docket** backend (no flag needed), while the `mcp` command still defaults to **asyncio** for single-process MCP usage. Use `--task-backend=docket` with `mcp` when you want MCP to enqueue background work for workers. |
There was a problem hiding this comment.
Inconsistent formatting of the --task-backend flag. Line 358 uses --task-backend=docket (with equals sign) while most other documentation examples use the space format. For consistency, consider using --task-backend docket to match the pattern used in docs/cli.md and other locations.
| The examples above use asyncio task backends for simple, single-process development. For production environments, the `api` command defaults to the **Docket** backend (no flag needed), while the `mcp` command still defaults to **asyncio** for single-process MCP usage. Use `--task-backend=docket` with `mcp` when you want MCP to enqueue background work for workers. | |
| The examples above use asyncio task backends for simple, single-process development. For production environments, the `api` command defaults to the **Docket** backend (no flag needed), while the `mcp` command still defaults to **asyncio** for single-process MCP usage. Use `--task-backend docket` with `mcp` when you want MCP to enqueue background work for workers. |
| | **Web MCP clients** | Either | `uv run agent-memory mcp --mode sse [--task-backend docket]` | | ||
|
|
||
| **Recommendation**: Start with `--no-worker` for development, then graduate to worker-based deployment for production. | ||
| **Recommendation**: Start with the asyncio backend (`--task-backend asyncio`) for simple development runs, then rely on the default Docket backend for the API in production, and enable `--task-backend=docket` for MCP when you want shared workers. |
There was a problem hiding this comment.
Inconsistent formatting of the --task-backend flag. Line 499 uses both formats in the same sentence: --task-backend asyncio (with space) and --task-backend=docket (with equals). For consistency, use the same format throughout, preferably the space format to match most examples.
| **Recommendation**: Start with the asyncio backend (`--task-backend asyncio`) for simple development runs, then rely on the default Docket backend for the API in production, and enable `--task-backend=docket` for MCP when you want shared workers. | |
| **Recommendation**: Start with the asyncio backend (`--task-backend asyncio`) for simple development runs, then rely on the default Docket backend for the API in production, and enable `--task-backend docket` for MCP when you want shared workers. |
| -e OPENAI_API_KEY=your-key \ | ||
| redislabs/agent-memory-server:latest | ||
| redislabs/agent-memory-server:latest \ | ||
| agent-memory api --host 0.0.0.0 --port 8000 --task-backend=asyncio |
There was a problem hiding this comment.
Inconsistent formatting of the --task-backend flag. Line 35 uses --task-backend=asyncio (with equals sign) while most documentation uses the space format. For consistency with examples in docs/cli.md and docs/getting-started.md, consider using --task-backend asyncio.
| agent-memory api --host 0.0.0.0 --port 8000 --task-backend=asyncio | |
| agent-memory api --host 0.0.0.0 --port 8000 --task-backend asyncio |
We use a custom
HybridBackgroundTaskclass to switch between Docket tasks and async tasks with FastAPI. However, this didn't work for the MCP server because it's not using FastAPI. The best solution appears to be allowing the user to choose between two "task backends": Docket and asyncio.