Skip to content

Local MCP streamable server support #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Sep 16, 2025
Merged

Local MCP streamable server support #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b87c72e
Move CodeQL MCP server to a streamable transport
Sep 15, 2025
b262b92
Fix bug in cmd resolve
Sep 15, 2025
da8a2b2
Update docs and finalize local streamable MCP server support
Sep 16, 2025
bfe5ded
Fix type check for server proc on mcp teardown
Sep 16, 2025
27b2508
Address review feedback
Sep 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
103 changes: 52 additions & 51 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,43 @@ It's primary value proposition is as a CLI tool that allows users to quickly def

Agents are defined through [personalities](personalities/), that receive a [task](taskflows/) to complete given a set of [tools](toolboxes/).

Agents can cooperate to complete sequences of tasks through so-called [Taskflows](taskflows/GRAMMAR.md).
Agents can cooperate to complete sequences of tasks through so-called [taskflows](taskflows/GRAMMAR.md).

You can find a detailed overview of the taskflow grammar [here](https://github.com/GitHubSecurityLab/seclab-taskflow-agent/blob/main/taskflows/GRAMMAR.md) and example taskflows [here](https://github.com/GitHubSecurityLab/seclab-taskflow-agent/tree/main/taskflows/examples).

## Use Cases and Examples

The Seclab Taskflow Agent framework was primarily designed to fit the iterative feedback loop driven work involved in Agentic security research workflows and vulnerability triage tasks.

Its design philosophy is centered around the belief that a prompt level focus of capturing vulnerability patterns will greatly improve and scale security research results as frontier model capabilities evolve over time.

While the maintainer himself primarily uses this framework as a code auditing tool it also serves as a more generic swiss army knife for exploring Agentic workflows. For example, the GitHub Security Lab also uses this framework for automated code scanning alert triage.

The framework includes a [CodeQL](https://codeql.github.com/) MCP server that can be used for Agentic code review, see the [CVE-2023-2283](https://github.com/GitHubSecurityLab/seclab-taskflow-agent/blob/main/taskflows/CVE-2023-2283/CVE-2023-2283.yaml) for an example of how to have an Agent review C code using a CodeQL database.

Instead of generating CodeQL queries itself, the CodeQL MCP Server is used to provide CodeQL-query based MCP tools that allow an Agent to navigate and explore code. It leverages templated CodeQL queries to provide targeted context for model driven code analysis.

## Requirements

Python >= 3.9 or Docker

# Usage
## Configuration

Provide a GitHub token for an account that is entitled to use GitHub Copilot via the `COPILOT_TOKEN` environment variable. Further configuration is use case dependent, i.e. pending which MCP servers you'd like to use in your taskflows.

You can set persisting environment variables via an `.env` file in the project root.

Provide a Copilot entitled GitHub PAT via the `COPILOT_TOKEN` environment variable.
Example:

## Source
```sh
# Tokens
COPILOT_TOKEN=<your_github_token>
# MCP configs
GITHUB_PERSONAL_ACCESS_TOKEN=<your_github_token>
CODEQL_DBS_BASE_PATH="/app/my_data/"
```

## Deploying from Source

First install the required dependencies:

Expand All @@ -48,40 +74,42 @@ Example: deploying a Taskflow:
python main.py -t example
```

## Docker
## Deploying from Docker

Alternatively you can deploy the Agent via its Docker image using `docker/run.sh`.
You can deploy the Taskflow Agent via its Docker image using `docker/run.sh`.

WARNING: the Agent Docker image is _NOT_ intended as a security boundary but strictly a deployment convenience.

The image entrypoint is `main.py` and thus it operates the same as invoking the Agent from source directly.

You can find the Docker image for the Seclab Taskflow Agent [here](https://github.com/GitHubSecurityLab/seclab-taskflow-agent/pkgs/container/seclab-taskflow-agent) and how it is built [here](release_tools/).

Note that this image is based on a public release of the Taskflow Agent, and you will have to mount any custom taskflows, personalities, or prompts into the image for them to be available to the Agent.

See [docker/run.sh](docker/run.sh) for configuration details.
Optional image mount points to supply custom data are configured via the environment:

Example: deploying a Taskflow:
- Custom data via `MY_DATA`, mounts to `/app/my_data`
- Custom personalities via `MY_PERSONALITIES`, mounts to `/app/personalities/my_personalities`
- Custom taskflows via `MY_TASKFLOWS`, mounts to `/app/taskflows/my_taskflows`
- Custom prompts via `MY_PROMPTS`, mounts to `/app/prompts/my_prompts`
- Custom toolboxes via `MY_TOOLBOXES`, mounts to `/app/toolboxes/my_toolboxes`

See [docker/run.sh](docker/run.sh) for further details.

Example: deploying a Taskflow (example.yaml):

```sh
docker/run.sh -t example
```
Example: deploying a custom taskflow:
Example: deploying a custom taskflow (custom_taskflow.yaml):

```sh
MY_TASKFLOWS=~/my_taskflows docker/run.sh -t custom_taskflow
```

Available image mount points are:

- Custom data via `MY_DATA` environment variable
- Custom personalities via `MY_PERSONALITIES` environment variable
- Custom taskflows via `MY_TASKFLOWS` environment variable
- Custom prompts via `MY_PROMPTS` environment variable
- Custom toolboxes via `MY_TOOLBOXES` environment variable

For more advanced scenarios like e.g. making custom MCP server code available, you can alter the run script to mount your custom code into the image and configure your toolboxes to use said code accordingly.

Example: custom MCP server deployment via Docker image:
Example: a custom MCP server deployment via Docker image:

```sh
export MY_MCP_SERVERS=./mcp_servers
Expand Down Expand Up @@ -109,7 +137,7 @@ docker run \

Our default run script makes the Docker socket available to the image, which contains the Docker cli, so 3rd party Docker based stdio MCP servers also function as normal.

Example: a toolbox configuration for the official GitHub MCP Server:
Example: a toolbox configuration using the official GitHub MCP Server via Docker:

```yaml
server_params:
Expand All @@ -120,23 +148,7 @@ server_params:
GITHUB_PERSONAL_ACCESS_TOKEN: "{{ env GITHUB_PERSONAL_ACCESS_TOKEN }}"
```

## Framework Configuration

Set environment variables via an `.env` file in the project root.

Example: a persistent Agent configuration with various MCP server environment variables set:

```sh
# Tokens
COPILOT_TOKEN=...
# Docker config, MY_DATA is mounted to /app/my_data
MY_DATA="/home/user/my_data"
# MCP configs
GITHUB_PERSONAL_ACCESS_TOKEN=...
CODEQL_DBS_BASE_PATH="/app/my_data/"
```

# Personalities
## Personalities

Core characteristics for a single Agent. Configured through YAML files in `personalities/`.

Expand All @@ -157,7 +169,7 @@ toolboxes:
- echo
```

# Toolboxes
## Toolboxes

MCP servers that provide tools. Configured through YAML files in `toolboxes/`.

Expand All @@ -174,18 +186,7 @@ server_params:
SOME: value
```

Example sse config:

```yaml
server_params:
kind: sse
# make sure you .env config the echo server, see echo_sse.py for example
url: http://127.0.0.1:9000/echo
headers:
SomeHeader: "{{ env USER }}"
```

# Taskflows
## Taskflows

A sequence of interdependent tasks performed by a set of Agents. Configured through a YAML based [grammar](taskflows/GRAMMAR.md) in [taskflows/](taskflows/).

Expand Down Expand Up @@ -263,6 +264,6 @@ This project is licensed under the terms of the MIT open source license. Please

[SUPPORT](./SUPPORT.md)

## Acknowledgement
## Acknowledgements

Security Lab team members @m-y-mo and @p- for contributing heavily to the testing and development of this framework, as well as the rest of the Security Lab team for helpful discussions and use cases.
Security Lab team members [Man Yue Mo](https://github.com/m-y-mo) and [Peter Stockli](https://github.com/p-) for contributing heavily to the testing and development of this framework, as well as the rest of the Security Lab team for helpful discussions and feedback.
56 changes: 41 additions & 15 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,12 @@
from typing import Any

from shell_utils import shell_tool_call
from mcp_utils import DEFAULT_MCP_CLIENT_SESSION_TIMEOUT, ReconnectingMCPServerStdio, AsyncDebugMCPServerStdio, MCPNamespaceWrap
from mcp_utils import DEFAULT_MCP_CLIENT_SESSION_TIMEOUT, ReconnectingMCPServerStdio, AsyncDebugMCPServerStdio, MCPNamespaceWrap, mcp_client_params, mcp_system_prompt, StreamableMCPThread
from render_utils import render_model_output, flush_async_output
from env_utils import TmpEnv
from yaml_parser import YamlParser
from agent import TaskAgent
from capi import list_tool_call_models
from mcp_utils import mcp_client_params
from mcp_utils import mcp_system_prompt

load_dotenv()

Expand Down Expand Up @@ -132,27 +130,47 @@ async def deploy_task_agents(agents: dict,
# XXX: auto-allow all tools if task is headless by clearing confirms
confirms = []
client_session_timeout = client_session_timeout or DEFAULT_MCP_CLIENT_SESSION_TIMEOUT
server_proc = None
match params['kind']:
# since we spawn stdio servers each time we do not expect
# new tools to appear over time so cache the tools list
case 'stdio':
if params.get('reconnecting', False):
mcp_server = ReconnectingMCPServerStdio(
name=tb,
params=params,
tool_filter=tool_filter,
client_session_timeout_seconds=client_session_timeout)
client_session_timeout_seconds=client_session_timeout,
cache_tools_list=True)
else:
mcp_server = MCPServerStdio(
name=tb,
params=params,
tool_filter=tool_filter,
client_session_timeout_seconds=client_session_timeout)
client_session_timeout_seconds=client_session_timeout,
cache_tools_list=True)
case 'sse':
mcp_server = MCPServerSse(
name=tb,
params=params,
tool_filter=tool_filter,
client_session_timeout_seconds=client_session_timeout)
case 'streamable': # XXX: needs testing
case 'streamable':
# check if we need to start this server locally as well
if 'command' in params:
def _print_out(line):
msg = f"Streamable MCP Server stdout: {line}"
logging.info(msg)
#print(msg)
def _print_err(line):
msg = f"Streamable MCP Server stderr: {line}"
logging.info(msg)
#print(msg)
server_proc = StreamableMCPThread(params['command'],
url=params['url'],
env=params['env'],
on_output=_print_out,
on_error=_print_err)
mcp_server = MCPServerStreamableHttp(
name=tb,
params=params,
Expand All @@ -161,7 +179,7 @@ async def deploy_task_agents(agents: dict,
case _:
raise ValueError(f"Unsupported MCP transport {params['kind']}")
# provide namespace and confirmation control through wrapper class
mcp_servers.append(MCPNamespaceWrap(confirms, mcp_server))
mcp_servers.append((MCPNamespaceWrap(confirms, mcp_server), server_proc))

# connect mcp servers
# https://openai.github.io/openai-agents-python/ref/mcp/server/
Expand All @@ -173,22 +191,33 @@ async def mcp_session_task(
# connects/cleanups have to happen in the same task
# but we also want to use wait_for to set a timeout
# so we use a dedicated session task to accomplish both
for server in mcp_servers:
for s in mcp_servers:
server, server_proc = s
logging.debug(f"Connecting mcp server: {server._name}")
if server_proc is not None:
server_proc.start()
await server_proc.async_wait_for_connection(poll_interval=0.1)
await server.connect()
# signal that we're connected
connected.set()
# wait until we're told to clean up
await cleanup.wait()
for server in reversed(mcp_servers):
for s in reversed(mcp_servers):
server, server_proc = s
try:
logging.debug(f"Starting cleanup for mcp server: {server._name}")
await server.cleanup()
logging.debug(f"Cleaned up mcp server: {server._name}")
if server_proc is not None:
server_proc.stop()
try:
await asyncio.to_thread(server_proc.join_and_raise)
except Exception as e:
print(f"Streamable mcp server process exception: {e}")
except asyncio.CancelledError:
logging.error(f"Timeout on cleanup for mcp server: {server._name}")
finally:
mcp_servers.remove(server)
mcp_servers.remove(s)
except RuntimeError as e:
logging.error(f"RuntimeError in mcp session task: {e}")
except asyncio.CancelledError as e:
Expand Down Expand Up @@ -233,12 +262,9 @@ async def mcp_session_task(
server_prompts=server_prompts,
important_guidelines=important_guidelines)
),
# XXX: should handoffs have handoffs?
# XXX: this would be a recursive chicken/egg problem :P
# XXX: are initial handoff functions still visible to handoff agents in the run?
handoffs=[],
exclude_from_context=exclude_from_context,
mcp_servers=mcp_servers,
mcp_servers=[s[0] for s in mcp_servers],
model=model,
model_settings=model_settings,
run_hooks=run_hooks,
Expand All @@ -257,7 +283,7 @@ async def mcp_session_task(
instructions=prompt_with_handoff_instructions(system_prompt) if handoffs else system_prompt,
handoffs=handoffs,
exclude_from_context=exclude_from_context,
mcp_servers=mcp_servers,
mcp_servers=[s[0] for s in mcp_servers],
model=model,
model_settings=model_settings,
run_hooks=run_hooks,
Expand Down
Loading