Megaraptor MCP

Instructions

Implementation Reference

• src/megaraptor_mcp/tools/flows.py:188-273 (handler)

  The `get_flow_status` handler function, decorated as an MCP tool. It validates the input client_id and flow_id, executes a VQL query, and returns the status of the flow.

  @mcp.tool()
  async def get_flow_status(
      client_id: str,
      flow_id: str,
  ) -> list[TextContent]:
      """Get the status of a specific collection flow.
  
      Args:
          client_id: The client ID (e.g., 'C.1234567890abcdef')
          flow_id: The flow ID (e.g., 'F.1234567890')
  
      Returns:
          Flow status including state, progress, and any errors.
      """
      try:
          # Input validation
          client_id = validate_client_id(client_id)
          flow_id = validate_flow_id(flow_id)
          client = get_client()
  
          vql = f"SELECT * FROM flows(client_id='{client_id}', flow_id='{flow_id}')"
          results = client.query(vql)
  
          if not results:
              return [TextContent(
                  type="text",
                  text=json.dumps({
                      "error": f"Flow {flow_id} not found for client {client_id}",
                      "hint": f"Use list_flows(client_id='{client_id}') to see available flows."
                  })
              )]
  
          flow = results[0]
  
          # Extract detailed status
          status = {
              "client_id": client_id,
              "flow_id": flow_id,
              "state": flow.get("state", ""),
              "artifacts_requested": flow.get("request", {}).get("artifacts", []),
              "artifacts_with_results": flow.get("artifacts_with_results", []),
              "create_time": flow.get("create_time", ""),
              "start_time": flow.get("start_time", ""),
              "active_time": flow.get("active_time", ""),
              "execution_duration": flow.get("execution_duration", 0),
              "total_uploaded_bytes": flow.get("total_uploaded_bytes", 0),
              "total_collected_rows": flow.get("total_collected_rows", 0),
              "outstanding_requests": flow.get("outstanding_requests", 0),
              "backtrace": flow.get("backtrace", ""),
              "status": flow.get("status", ""),
          }
  
          return [TextContent(
              type="text",
              text=json.dumps(status, indent=2, default=str)
          )]
  
      except grpc.RpcError as e:
          error_response = map_grpc_error(e, f"flow status for {flow_id}")
          # Check if it's a not-found error
          if "NOT_FOUND" in error_response.get("grpc_status", ""):
              error_response["hint"] = f"Flow {flow_id} may not exist for client {client_id}. Use list_flows(client_id='{client_id}') to see available flows."
          return [TextContent(
              type="text",
              text=json.dumps(error_response)
          )]
      except ValueError as e:
          # Validation errors
          return [TextContent(
              type="text",
              text=json.dumps({
                  "error": str(e),
                  "hint": "Provide valid client ID (C.*) and flow ID (F.*)"
              })
          )]
  
      except Exception:
          # Generic errors - don't expose internals
          return [TextContent(
              type="text",
              text=json.dumps({
                  "error": "Failed to get flow status",
                  "hint": "Check IDs and Velociraptor server connection"
              })
          )]