mirror of
https://github.com/FuzzingLabs/fuzzforge_ai.git
synced 2026-08-16 11:20:22 +02:00
fix: inject project assets as Docker volume mounts in execute_hub_tool
This commit is contained in:
@@ -176,6 +176,7 @@ class HubClient:
|
|||||||
arguments: dict[str, Any],
|
arguments: dict[str, Any],
|
||||||
*,
|
*,
|
||||||
timeout: int | None = None,
|
timeout: int | None = None,
|
||||||
|
extra_volumes: list[str] | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
"""Execute a tool on a hub server.
|
"""Execute a tool on a hub server.
|
||||||
|
|
||||||
@@ -183,6 +184,7 @@ class HubClient:
|
|||||||
:param tool_name: Name of the tool to execute.
|
:param tool_name: Name of the tool to execute.
|
||||||
:param arguments: Tool arguments.
|
:param arguments: Tool arguments.
|
||||||
:param timeout: Execution timeout (uses default if None).
|
:param timeout: Execution timeout (uses default if None).
|
||||||
|
:param extra_volumes: Additional Docker volume mounts to inject.
|
||||||
:returns: Tool execution result.
|
:returns: Tool execution result.
|
||||||
:raises HubClientError: If execution fails.
|
:raises HubClientError: If execution fails.
|
||||||
|
|
||||||
@@ -199,7 +201,7 @@ class HubClient:
|
|||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
async with self._connect(config) as (reader, writer):
|
async with self._connect(config, extra_volumes=extra_volumes) as (reader, writer):
|
||||||
# Initialise MCP session (skip for persistent — already done)
|
# Initialise MCP session (skip for persistent — already done)
|
||||||
if not self._persistent_sessions.get(config.name):
|
if not self._persistent_sessions.get(config.name):
|
||||||
await self._initialize_session(reader, writer, config.name)
|
await self._initialize_session(reader, writer, config.name)
|
||||||
@@ -248,6 +250,7 @@ class HubClient:
|
|||||||
async def _connect(
|
async def _connect(
|
||||||
self,
|
self,
|
||||||
config: HubServerConfig,
|
config: HubServerConfig,
|
||||||
|
extra_volumes: list[str] | None = None,
|
||||||
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
|
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
|
||||||
"""Connect to an MCP server.
|
"""Connect to an MCP server.
|
||||||
|
|
||||||
@@ -256,6 +259,7 @@ class HubClient:
|
|||||||
ephemeral per-call connection logic.
|
ephemeral per-call connection logic.
|
||||||
|
|
||||||
:param config: Server configuration.
|
:param config: Server configuration.
|
||||||
|
:param extra_volumes: Additional Docker volume mounts to inject.
|
||||||
:yields: Tuple of (reader, writer) for communication.
|
:yields: Tuple of (reader, writer) for communication.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -268,7 +272,7 @@ class HubClient:
|
|||||||
|
|
||||||
# Ephemeral connection (original behaviour)
|
# Ephemeral connection (original behaviour)
|
||||||
if config.type == HubServerType.DOCKER:
|
if config.type == HubServerType.DOCKER:
|
||||||
async with self._connect_docker(config) as streams:
|
async with self._connect_docker(config, extra_volumes=extra_volumes) as streams:
|
||||||
yield streams
|
yield streams
|
||||||
elif config.type == HubServerType.COMMAND:
|
elif config.type == HubServerType.COMMAND:
|
||||||
async with self._connect_command(config) as streams:
|
async with self._connect_command(config) as streams:
|
||||||
@@ -284,10 +288,12 @@ class HubClient:
|
|||||||
async def _connect_docker(
|
async def _connect_docker(
|
||||||
self,
|
self,
|
||||||
config: HubServerConfig,
|
config: HubServerConfig,
|
||||||
|
extra_volumes: list[str] | None = None,
|
||||||
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
|
) -> AsyncGenerator[tuple[asyncio.StreamReader, asyncio.StreamWriter], None]:
|
||||||
"""Connect to a Docker-based MCP server.
|
"""Connect to a Docker-based MCP server.
|
||||||
|
|
||||||
:param config: Server configuration with image name.
|
:param config: Server configuration with image name.
|
||||||
|
:param extra_volumes: Additional volume mounts to inject (e.g. project assets).
|
||||||
:yields: Tuple of (reader, writer) for stdio communication.
|
:yields: Tuple of (reader, writer) for stdio communication.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -302,10 +308,14 @@ class HubClient:
|
|||||||
for cap in config.capabilities:
|
for cap in config.capabilities:
|
||||||
cmd.extend(["--cap-add", cap])
|
cmd.extend(["--cap-add", cap])
|
||||||
|
|
||||||
# Add volumes
|
# Add volumes from server config
|
||||||
for volume in config.volumes:
|
for volume in config.volumes:
|
||||||
cmd.extend(["-v", os.path.expanduser(volume)])
|
cmd.extend(["-v", os.path.expanduser(volume)])
|
||||||
|
|
||||||
|
# Add extra volumes (e.g. project assets injected at runtime)
|
||||||
|
for volume in (extra_volumes or []):
|
||||||
|
cmd.extend(["-v", os.path.expanduser(volume)])
|
||||||
|
|
||||||
# Add environment variables
|
# Add environment variables
|
||||||
for key, value in config.environment.items():
|
for key, value in config.environment.items():
|
||||||
cmd.extend(["-e", f"{key}={value}"])
|
cmd.extend(["-e", f"{key}={value}"])
|
||||||
|
|||||||
@@ -180,12 +180,14 @@ class HubExecutor:
|
|||||||
arguments: dict[str, Any] | None = None,
|
arguments: dict[str, Any] | None = None,
|
||||||
*,
|
*,
|
||||||
timeout: int | None = None,
|
timeout: int | None = None,
|
||||||
|
extra_volumes: list[str] | None = None,
|
||||||
) -> HubExecutionResult:
|
) -> HubExecutionResult:
|
||||||
"""Execute a hub tool.
|
"""Execute a hub tool.
|
||||||
|
|
||||||
:param identifier: Tool identifier (hub:server:tool or server:tool).
|
:param identifier: Tool identifier (hub:server:tool or server:tool).
|
||||||
:param arguments: Tool arguments.
|
:param arguments: Tool arguments.
|
||||||
:param timeout: Execution timeout.
|
:param timeout: Execution timeout.
|
||||||
|
:param extra_volumes: Additional Docker volume mounts to inject.
|
||||||
:returns: Execution result.
|
:returns: Execution result.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -232,6 +234,7 @@ class HubExecutor:
|
|||||||
tool_name_to_use or tool_name,
|
tool_name_to_use or tool_name,
|
||||||
arguments,
|
arguments,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
|
extra_volumes=extra_volumes,
|
||||||
)
|
)
|
||||||
return HubExecutionResult(
|
return HubExecutionResult(
|
||||||
success=True,
|
success=True,
|
||||||
@@ -268,6 +271,7 @@ class HubExecutor:
|
|||||||
tool.name,
|
tool.name,
|
||||||
arguments,
|
arguments,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
|
extra_volumes=extra_volumes,
|
||||||
)
|
)
|
||||||
return HubExecutionResult(
|
return HubExecutionResult(
|
||||||
success=True,
|
success=True,
|
||||||
|
|||||||
@@ -16,7 +16,7 @@ from fastmcp import FastMCP
|
|||||||
from fastmcp.exceptions import ToolError
|
from fastmcp.exceptions import ToolError
|
||||||
from fuzzforge_common.hub import HubExecutor, HubServerConfig, HubServerType
|
from fuzzforge_common.hub import HubExecutor, HubServerConfig, HubServerType
|
||||||
|
|
||||||
from fuzzforge_mcp.dependencies import get_settings
|
from fuzzforge_mcp.dependencies import get_project_path, get_settings, get_storage
|
||||||
|
|
||||||
mcp: FastMCP = FastMCP()
|
mcp: FastMCP = FastMCP()
|
||||||
|
|
||||||
@@ -180,10 +180,30 @@ async def execute_hub_tool(
|
|||||||
try:
|
try:
|
||||||
executor = _get_hub_executor()
|
executor = _get_hub_executor()
|
||||||
|
|
||||||
|
# Inject project assets as Docker volume mounts if configured.
|
||||||
|
# Mounts the assets directory at the standard paths used by hub tools:
|
||||||
|
# /app/uploads — binwalk, and other tools that use UPLOAD_DIR
|
||||||
|
# /app/samples — yara, capa, and other tools that use SAMPLES_DIR
|
||||||
|
extra_volumes: list[str] = []
|
||||||
|
try:
|
||||||
|
storage = get_storage()
|
||||||
|
project_path = get_project_path()
|
||||||
|
assets_path = storage.get_project_assets_path(project_path)
|
||||||
|
if assets_path:
|
||||||
|
assets_str = str(assets_path)
|
||||||
|
extra_volumes = [
|
||||||
|
f"{assets_str}:/app/uploads:ro",
|
||||||
|
f"{assets_str}:/app/samples:ro",
|
||||||
|
]
|
||||||
|
except Exception:
|
||||||
|
# Never block tool execution because of asset injection failure
|
||||||
|
extra_volumes = []
|
||||||
|
|
||||||
result = await executor.execute_tool(
|
result = await executor.execute_tool(
|
||||||
identifier=identifier,
|
identifier=identifier,
|
||||||
arguments=arguments or {},
|
arguments=arguments or {},
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
|
extra_volumes=extra_volumes or None,
|
||||||
)
|
)
|
||||||
|
|
||||||
return result.to_dict()
|
return result.to_dict()
|
||||||
|
|||||||
Reference in New Issue
Block a user